124 lines
2.2 KiB
Go
124 lines
2.2 KiB
Go
package leader
|
|
|
|
import (
|
|
"context"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"go.etcd.io/etcd/client/v3/concurrency"
|
|
|
|
"novit.tech/go/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
|
|
}
|