lbry.go/stop/stop.go

package stop

import (
	"context"
	"log"
	"sync"
)

// Chan is a receive-only channel
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
}
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{}
	ctx := context.Background()
	if len(parent) > 0 && parent[0] != nil {
		ctx = parent[0].ctx
	}
	s.ctx, s.cancel = context.WithCancel(ctx)
	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()
}

// Stop signals any listening processes to stop. After the first call, Stop() does nothing.
func (s *Group) Stop() {
	s.cancel()
}

// StopAndWait is a convenience method to close the channel and wait for goroutines to return.
func (s *Group) StopAndWait() {
	s.Stop()
	s.Wait()
}

// Child returns a new instance that will be stopped when s is stopped.
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()
}