From 4c5efe16684508f0f1818069e5491498f99a244d Mon Sep 17 00:00:00 2001
From: ston1th <ston1th@giftfish.de>
Date: Sat, 11 Dec 2021 18:59:25 +0100
Subject: [PATCH] fix

---
 pkg/etcd/etcd.go | 228 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 228 insertions(+)
 create mode 100644 pkg/etcd/etcd.go

diff --git a/pkg/etcd/etcd.go b/pkg/etcd/etcd.go
new file mode 100644
index 0000000..f1cb14a
--- /dev/null
+++ b/pkg/etcd/etcd.go
@@ -0,0 +1,228 @@
+package etcd
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+	"time"
+
+	"git.giftfish.de/ston1th/haproxy-lb/pkg/cluster"
+	"git.giftfish.de/ston1th/haproxy-lb/pkg/config"
+	"git.giftfish.de/ston1th/haproxy-lb/pkg/util"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.etcd.io/etcd/client/v3/concurrency"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+func (c *Cluster) Start(etcdcfg *config.Config) error {
+	err := c.checkConfig()
+	if err != nil {
+		return err
+	}
+	cfg := etcdcfg.Cluster.Etcd
+	var level zapcore.Level
+	err = level.Set(cfg.LogLevel)
+	if err != nil {
+		return err
+	}
+	al := zap.NewAtomicLevel()
+	al.SetLevel(level)
+	var tlsConfig *tls.Config
+	if cfg.TLS != nil {
+		ca, err := config.LoadCertPool(cfg.TLS.CA)
+		if err != nil {
+			return err
+		}
+		tlsConfig = &tls.Config{
+			MinVersion:         tls.VersionTLS12,
+			InsecureSkipVerify: cfg.TLS.Insecure,
+			RootCAs:            ca,
+		}
+	}
+
+	c.cli, err = clientv3.New(clientv3.Config{
+		Endpoints:   cfg.Endpoints,
+		TLS:         tlsConfig,
+		DialTimeout: time.Second * 5,
+		Username:    cfg.Username,
+		Password:    cfg.Password,
+		LogConfig: &zap.Config{
+			Level:             al,
+			Encoding:          "console",
+			DisableCaller:     true,
+			DisableStacktrace: true,
+			EncoderConfig:     zap.NewDevelopmentEncoderConfig(),
+			OutputPaths:       []string{"stderr"},
+			ErrorOutputPaths:  []string{"stderr"},
+		},
+	})
+	if err != nil {
+		return err
+	}
+
+	log := c.Logger()
+	log.Info("waiting for etcd")
+	sessionCtx, sessionCancel := context.WithCancel(context.Background())
+	c.sessionCancel = sessionCancel
+	s, err := concurrency.NewSession(c.cli,
+		concurrency.WithContext(sessionCtx),
+		concurrency.WithTTL(10),
+	)
+	if err != nil {
+		return err
+	}
+	c.session = true
+
+	prefix := cfg.Prefix + "/" + cfg.ClusterName
+	c.kvPrefix = prefix + "/kv"
+	ePrefix := prefix + "/e"
+	c.id = fmt.Sprintf("%x", s.Lease())
+	c.keyID = ePrefix + "/" + c.id
+
+	e := concurrency.NewElection(s, ePrefix)
+	campaignCtx, campaignCancel := context.WithCancel(context.Background())
+	go func() {
+		errc := make(chan error, 1)
+		for {
+			go func() {
+				errc <- e.Campaign(campaignCtx, "")
+			}()
+			select {
+			case err := <-errc:
+				if err != nil {
+					log.Error(err, "campaign error")
+					c.Fatal(cluster.ErrRestart)
+					return
+				}
+			case <-campaignCtx.Done():
+				return
+			}
+		}
+	}()
+
+	t := time.NewTicker(time.Second * 2)
+	ctx, cancel := context.WithCancel(context.Background())
+	observeCtx, observeCancel := context.WithCancel(context.Background())
+	leading := false
+	following := false
+	leaderChan := c.observe(observeCtx, e.Observe(observeCtx))
+	for {
+		if !leading && c.leaderID(e.Leader(observeCtx)) == c.keyID {
+			leading = true
+			go c.leader(ctx)
+		} else if !following && !leading {
+			following = true
+			go c.follower(ctx)
+		}
+
+		select {
+		case <-t.C:
+		case leader := <-leaderChan:
+			if leader {
+				if following {
+					cancel()
+					following = false
+					ctx, cancel = context.WithCancel(context.Background())
+				}
+				if !leading {
+					leading = true
+					go c.leader(ctx)
+				}
+			} else {
+				if leading {
+					cancel()
+					leading = false
+					ctx, cancel = context.WithCancel(context.Background())
+				}
+				if !following {
+					following = true
+					go c.follower(ctx)
+				}
+			}
+		case <-c.stepdown:
+			cancel()
+			ctx, cancel = context.WithCancel(context.Background())
+			resignCtx, _ := context.WithTimeout(context.Background(), time.Second*2)
+			e.Resign(resignCtx)
+		case <-c.stop:
+			t.Stop()
+			cancel()
+			observeCancel()
+			campaignCancel()
+			if leading {
+				ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
+				e.Resign(ctx)
+			}
+			if c.callbacks.Cleanup != nil {
+				c.callbacks.Cleanup(context.Background(), c)
+			}
+			s.Close()
+			close(c.done)
+			return c.fatal
+		}
+	}
+	return c.fatal
+}
+
+func (c *Cluster) leaderID(r *clientv3.GetResponse, err error) string {
+	if err != nil {
+		return ""
+	}
+	if len(r.Kvs) == 0 {
+		return ""
+	}
+	return string(r.Kvs[0].Key)
+}
+
+func (c *Cluster) observe(ctx context.Context, resp <-chan clientv3.GetResponse) <-chan bool {
+	leader := make(chan bool)
+	go func() {
+		for {
+			select {
+			case r := <-resp:
+				if c.leaderID(&r, nil) == c.keyID {
+					leader <- true
+				} else {
+					leader <- false
+				}
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+	return leader
+}
+
+func (c *Cluster) Stepdown() {
+	select {
+	case c.stepdown <- struct{}{}:
+	case <-time.After(time.Second * 2):
+	}
+}
+
+func (c *Cluster) Fatal(err error) {
+	c.fatal = err
+	c.Stop()
+}
+
+func (c *Cluster) Stop() {
+	if !c.session {
+		c.sessionCancel()
+		return
+	}
+	c.stop <- struct{}{}
+	close(c.stop)
+	<-c.done
+	c.cli.Close()
+}
+
+func (c *Cluster) leader(ctx context.Context) {
+	defer util.HandleCrash()
+	c.callbacks.Leader(ctx, c)
+}
+
+func (c *Cluster) follower(ctx context.Context) {
+	defer util.HandleCrash()
+	c.callbacks.Follower(ctx, c)
+}