stop/stop.go

@@ -2,6 +2,7 @@ package stop
 
 import (
 	"context"
+	"log"
 	"sync"
 )
 
@@ -10,6 +11,8 @@ type Chan <-chan struct{}
 
 // Stopper extends sync.WaitGroup to add a convenient way to stop running goroutines
 type Group struct {
+	waitingOn map[string]int
+	l         sync.RWMutex
 	sync.WaitGroup
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -27,6 +30,15 @@ func New(parent ...*Group) *Group {
 	return s
 }
 
+// NewDebug allows you to debug the go routines the group waits on. In order to leverage this, AddNamed and DoneNamed should be used.
+func NewDebug(parent ...*Group) *Group {
+	s := New(parent...)
+	s.waitingOn = make(map[string]int)
+	s.l = sync.RWMutex{}
+
+	return s
+}
+
 // Ch returns a channel that will be closed when Stop is called.
 func (s *Group) Ch() Chan {
 	return s.ctx.Done()
@@ -47,3 +59,38 @@ func (s *Group) StopAndWait() {
 func (s *Group) Child() *Group {
 	return New(s)
 }
+
+func (s *Group) AddNamed(delta int, name string) {
+	if s.waitingOn != nil {
+		s.l.Lock()
+		defer s.l.Unlock()
+		_, ok := s.waitingOn[name]
+		if !ok {
+			s.waitingOn[name] = 1
+		} else {
+			s.waitingOn[name] = s.waitingOn[name] + 1
+		}
+	}
+	s.Add(delta)
+
+}
+
+func (s *Group) DoneNamed(name string) {
+	if s.waitingOn != nil {
+		s.l.Lock()
+		defer s.l.Unlock()
+		_, ok := s.waitingOn[name]
+		if !ok {
+			log.Printf("%s is not recorded in stop group map")
+		} else {
+			s.waitingOn[name] = s.waitingOn[name] - 1
+		}
+		log.Printf("-->> LIST WAITING ON")
+		for k, v := range s.waitingOn {
+			if v > 0 {
+				log.Printf("waiting on %d %s routines...", v, k)
+			}
+		}
+	}
+	s.Done()
+}