diff --git a/stop/stop.go b/stop/stop.go index 0da5f51..0fb0473 100644 --- a/stop/stop.go +++ b/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() +}