cooler job tracking

1 files changed, 88 insertions, 0 deletions

diff --git a/jwg.go b/jwg.go
new file mode 100644
index 0000000..6ef06d0
--- /dev/null
+++ b/jwg.go
@@ -0,0 +1,88 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+)
+
+type JobWaitGroup struct {
+	lock sync.RWMutex
+	jobs map[string]bool
+	wg   sync.WaitGroup
+}
+
+func (jwg *JobWaitGroup) Do(name string, fn func()) {
+	jwg.lock.Lock()
+	defer jwg.lock.Unlock()
+
+	if jwg.jobs == nil {
+		jwg.jobs = make(map[string]bool)
+	}
+	if _, dup := jwg.jobs[name]; dup {
+		panic(fmt.Sprintf("job %q already exists", name))
+	}
+	jwg.jobs[name] = true
+	//fmt.Fprintf(os.Stderr, "add %p %v\n", &jwg.wg, jwg.wg)
+	jwg.wg.Add(1)
+	go func() {
+		defer func() {
+			jwg.lock.Lock()
+			defer jwg.lock.Unlock()
+			jwg.jobs[name] = false
+			jwg.wg.Done()
+			//fmt.Fprintf(os.Stderr, "del %p %v\n", &jwg.wg, jwg.wg)
+		}()
+		fn()
+	}()
+}
+
+func (jwg *JobWaitGroup) Wait() {
+	jwg.wg.Wait()
+}
+
+func (jwg *JobWaitGroup) Status() (int, []string) {
+	jwg.lock.RLock()
+
+	n := len(jwg.jobs)
+	var jobs []string
+	for job, active := range jwg.jobs {
+		if active {
+			jobs = append(jobs, job)
+		}
+	}
+	defer jwg.lock.RUnlock()
+	sort.Strings(jobs)
+
+	return n, jobs
+}
+
+func (jwg *JobWaitGroup) Watch(d time.Duration) {
+	ticker := time.NewTicker(d)
+	done := make(chan struct{})
+	go func() {
+		jwg.Wait()
+		ticker.Stop()
+		done <- struct{}{}
+	}()
+	for {
+		select {
+		case <-ticker.C:
+			n, jobs := jwg.Status()
+			if len(jobs) == 0 {
+				panic("no active jobs, but wg still waiting")
+			}
+			line := fmt.Sprintf("%d/%d : %v", len(jobs), n, strings.Join(jobs, ", "))
+			if len(line) > 70 {
+				line = line[:67] + "..."
+			}
+			fmt.Fprintf(os.Stderr, "\r%-70s", line)
+		case <-done:
+			fmt.Fprintf(os.Stderr, "\r%-70s\n", "done")
+			return
+		}
+	}
+}