package routine import ( "context" "os" "sync" "github.com/hashicorp/go-hclog" ) type Routine func(ctx context.Context) error type routineTracker struct { cancel context.CancelFunc stoppedCh chan struct{} // closed when no longer running } func (r *routineTracker) running() bool { select { case <-r.stoppedCh: return false default: return true } } type Manager struct { lock sync.RWMutex logger hclog.Logger routines map[string]*routineTracker } func NewManager(logger hclog.Logger) *Manager { if logger == nil { logger = hclog.New(&hclog.LoggerOptions{ Output: os.Stderr, }) } return &Manager{ logger: logger, routines: make(map[string]*routineTracker), } } func (m *Manager) IsRunning(name string) bool { m.lock.Lock() defer m.lock.Unlock() if routine, ok := m.routines[name]; ok { return routine.running() } return false } func (m *Manager) Start(ctx context.Context, name string, routine Routine) error { m.lock.Lock() defer m.lock.Unlock() if instance, ok := m.routines[name]; ok && instance.running() { return nil } if ctx == nil { ctx = context.Background() } rtCtx, cancel := context.WithCancel(ctx) instance := &routineTracker{ cancel: cancel, stoppedCh: make(chan struct{}), } go m.execute(rtCtx, name, routine, instance.stoppedCh) m.routines[name] = instance m.logger.Info("started routine", "routine", name) return nil } // execute will run the given routine in the foreground and close the given channel when its done executing func (m *Manager) execute(ctx context.Context, name string, routine Routine, done chan struct{}) { defer func() { close(done) }() err := routine(ctx) if err != nil && err != context.DeadlineExceeded && err != context.Canceled { m.logger.Error("routine exited with error", "routine", name, "error", err, ) } else { m.logger.Debug("stopped routine", "routine", name) } } func (m *Manager) Stop(name string) <-chan struct{} { instance := m.stopInstance(name) if instance == nil { // Fabricate a closed channel so it won't block forever. ch := make(chan struct{}) close(ch) return ch } return instance.stoppedCh } func (m *Manager) stopInstance(name string) *routineTracker { m.lock.Lock() defer m.lock.Unlock() instance, ok := m.routines[name] if !ok { // no running instance return nil } if !instance.running() { return instance } m.logger.Debug("stopping routine", "routine", name) instance.cancel() delete(m.routines, name) return instance } func (m *Manager) StopAll() { m.lock.Lock() defer m.lock.Unlock() for name, routine := range m.routines { if !routine.running() { continue } m.logger.Debug("stopping routine", "routine", name) routine.cancel() } // just wipe out the entire map m.routines = make(map[string]*routineTracker) }