package leader import ( "context" "flag" "fmt" "log" "os" "os/signal" "sync" "syscall" "time" "go.etcd.io/etcd/client/v3/concurrency" "novit.tech/direktil/etcdb" ) var leaderName string func init() { hostname, err := os.Hostname() if err != nil { panic("failed to get hostname: " + err.Error()) } leaderName = hostname flag.StringVar(&leaderName, "leader-name", leaderName, "leader name in leader election") } type Spec struct { Etcd *etcdb.ClientSpec Election string OnLeader func(ctx context.Context) } // RunProcess is design for processes that are a single leader. Binds signals to exit gracefully. func (spec *Spec) RunProcess() { ctx, cancel := context.WithCancel(context.Background()) go func() { c := make(chan os.Signal, 1) signal.Notify(c, syscall.SIGINT, syscall.SIGTERM) s := <-c log.Print("got signal ", s) cancel() <-c log.Print("2nd signal, forced exit") os.Exit(0) }() spec.Run(ctx) } func (spec *Spec) Run(ctx context.Context) { for { err := spec.runSession(ctx) if err != nil { log.Print("leader election: failed: ", err) } if ctx.Err() != nil { break } time.Sleep(time.Second) } } func (spec *Spec) runSession(ctx context.Context) (err error) { cli := spec.Etcd.Client() runCtx, runCancel := context.WithCancel(context.Background()) defer runCancel() s, err := concurrency.NewSession( cli, concurrency.WithTTL(5), concurrency.WithContext(runCtx), ) if err != nil { return fmt.Errorf("failed to start leader election: %v", err) } defer s.Close() election := concurrency.NewElection(s, spec.Election) log.Print("leader election: waiting to become leader") if err := election.Campaign(ctx, leaderName); err != nil { return fmt.Errorf("leader election: campaign failed: %v", err) } log.Print("leader election: got leader status") leaderCtx, cancel := context.WithCancel(ctx) wg := new(sync.WaitGroup) wg.Add(1) go func() { defer wg.Done() spec.OnLeader(leaderCtx) }() select { case <-s.Done(): log.Print("leader election: session done") case <-ctx.Done(): log.Print("leader election: context cancelled") } cancel() log.Print("leader election: waiting for routine") wg.Wait() log.Print("leader election: routine exited") return }