haproxy-lb/pkg/etcd/etcd.go

package etcd

import (
	"context"
	"crypto/tls"
	"fmt"
	"time"

	"git.giftfish.de/ston1th/vipman/pkg/config"
	"git.giftfish.de/ston1th/vipman/pkg/util"
	"github.com/coreos/etcd/clientv3"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	//"go.etcd.io/etcd/clientv3"
	"go.etcd.io/etcd/clientv3/concurrency"
)

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
	}

	ctxSsession, cancelSession := context.WithCancel(context.Background())
	c.cancelSession = cancelSession
	s, err := concurrency.NewSession(c.cli,
		concurrency.WithContext(ctxSsession),
		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)
	ctxCampaign, cancelCampaign := context.WithCancel(context.Background())
	go func() {
		errc := make(chan error, 1)
		for {
			go func() {
				errc <- e.Campaign(ctxCampaign, "")
			}()
			select {
			case err := <-errc:
				if err != nil {
					time.Sleep(time.Second)
				}
			case <-ctxCampaign.Done():
				return
			}
		}
	}()

	t := time.NewTicker(time.Second)
	ctx, cancel := context.WithCancel(context.Background())
	ctxObserve, cancelObserve := context.WithCancel(context.Background())
	leading := false
	following := false
	leaderChan := c.observe(ctxObserve, e.Observe(ctxObserve))
	for {
		if !leading && c.leaderID(e.Leader(ctxObserve)) == 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())
			ctxResign, _ := context.WithTimeout(context.Background(), time.Second*2)
			e.Resign(ctxResign)
		case <-c.stop:
			t.Stop()
			cancel()
			cancelObserve()
			cancelCampaign()
			if leading {
				ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
				e.Resign(ctx)
			}
			if c.callbacks.Cleanup != nil {
				ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
				c.callbacks.Cleanup(ctx, c)
			}
			s.Close()
			close(c.done)
			return nil
		}
	}
	return nil
}

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) Stop() {
	if !c.session {
		c.cancelSession()
		return
	}
	close(c.stop)
	<-c.done
	c.cli.Close()
}

func (c *Cluster) leader(ctx context.Context) {
	defer func() {
		util.HandleCrash()
	}()
	c.callbacks.Leader(ctx, c)
}

func (c *Cluster) follower(ctx context.Context) {
	defer func() {
		util.HandleCrash()
	}()
	c.callbacks.Follower(ctx, c)
}