open-consul/agent/consul/leader_routine_manager.go

121 lines
2.3 KiB
Go
Raw Normal View History

package consul
import (
"context"
"log"
"os"
"sync"
)
type LeaderRoutine func(ctx context.Context) error
type leaderRoutine struct {
running bool
cancel context.CancelFunc
}
type LeaderRoutineManager struct {
lock sync.RWMutex
logger *log.Logger
routines map[string]*leaderRoutine
}
func NewLeaderRoutineManager(logger *log.Logger) *LeaderRoutineManager {
if logger == nil {
logger = log.New(os.Stderr, "", log.LstdFlags)
}
return &LeaderRoutineManager{
logger: logger,
routines: make(map[string]*leaderRoutine),
}
}
func (m *LeaderRoutineManager) 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 *LeaderRoutineManager) Start(name string, routine LeaderRoutine) error {
return m.StartWithContext(nil, name, routine)
}
func (m *LeaderRoutineManager) StartWithContext(parentCtx context.Context, name string, routine LeaderRoutine) error {
m.lock.Lock()
defer m.lock.Unlock()
if instance, ok := m.routines[name]; ok && instance.running {
return nil
}
if parentCtx == nil {
parentCtx = context.Background()
}
ctx, cancel := context.WithCancel(parentCtx)
instance := &leaderRoutine{
running: true,
cancel: cancel,
}
go func() {
err := routine(ctx)
if err != nil && err != context.DeadlineExceeded && err != context.Canceled {
m.logger.Printf("[ERROR] leader: %s routine exited with error: %v", name, err)
} else {
m.logger.Printf("[DEBUG] leader: stopped %s routine", name)
}
m.lock.Lock()
instance.running = false
m.lock.Unlock()
}()
m.routines[name] = instance
m.logger.Printf("[INFO] leader: started %s routine", name)
return nil
}
func (m *LeaderRoutineManager) Stop(name string) error {
m.lock.Lock()
defer m.lock.Unlock()
instance, ok := m.routines[name]
if !ok {
// no running instance
return nil
}
if !instance.running {
return nil
}
m.logger.Printf("[DEBUG] leader: stopping %s routine", name)
instance.cancel()
delete(m.routines, name)
return nil
}
func (m *LeaderRoutineManager) StopAll() {
m.lock.Lock()
defer m.lock.Unlock()
for name, routine := range m.routines {
if !routine.running {
continue
}
m.logger.Printf("[DEBUG] leader: stopping %s routine", name)
routine.cancel()
}
// just whipe out the entire map
m.routines = make(map[string]*leaderRoutine)
}