improve stopper debugging

stop/stop.go

@@ -11,17 +11,18 @@ 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
+
+	mu        *sync.Mutex
+	waitingOn map[string]int
 }
 type Stopper = Group
 
 // New allocates and returns a new instance. Use New(parent) to create an instance that is stopped when parent is stopped.
 func New(parent ...*Group) *Group {
-	s := &Group{}
+	s := &Group{mu: &sync.Mutex{}}
 	ctx := context.Background()
 	if len(parent) > 0 && parent[0] != nil {
 		ctx = parent[0].ctx
@@ -30,15 +31,6 @@ 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()
@@ -60,37 +52,44 @@ 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
-		}
-	}
+func (s *Group) DebugAdd(delta int, name string) {
 	s.Add(delta)
 
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if s.waitingOn == nil {
+		s.waitingOn = make(map[string]int)
+	}
+
+	if current, ok := s.waitingOn[name]; ok {
+		s.waitingOn[name] = current + 1
+	} else {
+		s.waitingOn[name] = 1
+	}
 }
 
-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")
+func (s *Group) DebugDone(name string) {
+	defer s.Done()
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if current, ok := s.waitingOn[name]; ok {
+		if current <= 1 {
+			delete(s.waitingOn, name)
 		} else {
-			s.waitingOn[name] = s.waitingOn[name] - 1
+			s.waitingOn[name] = current - 1
 		}
-		log.Printf("-->> LIST WAITING ON")
-		for k, v := range s.waitingOn {
-			if v > 0 {
-				log.Printf("waiting on %d %s routines...", v, k)
-			}
+	} else {
+		log.Printf("%s is not recorded in stop group map", name)
+	}
+
+	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()
 }