haproxy-lb/pkg/raft/raft.go

package raft

import (
	"context"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"time"

	"git.giftfish.de/ston1th/haproxy-lb/pkg/config"
	"git.giftfish.de/ston1th/haproxy-lb/pkg/raft/tls"
	"git.giftfish.de/ston1th/haproxy-lb/pkg/util"
	"git.giftfish.de/ston1th/raftbbolt"
	hclog "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/raft"
)

func (c *Cluster) Start(raftcfg *config.Config) error {
	err := c.checkConfig()
	if err != nil {
		return err
	}
	log := c.Logger()
	cfg := raftcfg.Cluster.Raft
	hcl := hclog.New(&hclog.LoggerOptions{
		Name:       "raft",
		Level:      hclog.LevelFromString(cfg.LogLevel),
		Output:     hclog.DefaultOutput,
		TimeFormat: "I0102 15:04:05.000000",
	})

	rc := raft.DefaultConfig()
	rc.ProtocolVersion = 3
	rc.LocalID = raft.ServerID(cfg.ID)
	rc.Logger = hcl
	c.id = cfg.ID

	pool := 3
	timeout := 10 * time.Second
	var transport *raft.NetworkTransport
	if cfg.TLS != nil {
		log.V(1).Info("setup", "transport", "tls")
		s, err := tls.NewStream(cfg)
		if err != nil {
			return err
		}
		transport = raft.NewNetworkTransportWithLogger(s, pool, timeout, hcl)
	} else {
		log.V(1).Info("setup", "transport", "tcp")
		address, err := net.ResolveTCPAddr("tcp", cfg.Address)
		if err != nil {
			return err
		}
		transport, err = raft.NewTCPTransportWithLogger(cfg.Address, address, pool, timeout, hcl)
		if err != nil {
			return err
		}
	}

	// Create Raft structures
	//snapshots := raft.NewInmemSnapshotStore()
	snapshots, err := raft.NewFileSnapshotStoreWithLogger(cfg.Dir, retainSnapshotCount, hcl)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

	store := filepath.Join(cfg.Dir, "store.db")

	//logStore := raft.NewInmemStore()
	bootstrap := true
	if _, err := os.Stat(store); err == nil {
		bootstrap = false
	}

	stableStore, err := raftbbolt.NewStore(store)
	if err != nil {
		return err
	}
	logStore := stableStore

	// Cluster configuration
	configuration := raft.Configuration{}

	// Add Local Peer
	configuration.Servers = append(configuration.Servers, raft.Server{
		ID:      raft.ServerID(cfg.ID),
		Address: raft.ServerAddress(cfg.Address),
	})
	for _, p := range cfg.Peers {
		if cfg.Address != p.Address {
			configuration.Servers = append(configuration.Servers, raft.Server{
				ID:      raft.ServerID(p.ID),
				Address: raft.ServerAddress(p.Address)})
		}
	}

	if bootstrap {
		if err := raft.BootstrapCluster(rc, logStore, stableStore, snapshots, transport, configuration); err != nil {
			return err
		}
	}

	c.r, err = raft.NewRaft(rc, (*fsm)(c.kvm), logStore, stableStore, snapshots, transport)
	if err != nil {
		return err
	}
	c.kvm.r = c.r

	t := time.NewTicker(time.Second)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	leading := false
	following := false
	for {
		if !leading && cfg.Address == string(c.r.Leader()) {
			leading = true
			go c.leader(ctx)
		} else if !following && !leading {
			following = true
			go c.follower(ctx)
		}

		select {
		case <-t.C:
		case leader := <-c.r.LeaderCh():
			if leader {
				if following {
					cancel()
					following = false
					ctx, cancel = context.WithCancel(context.Background())
				}
				if !leading {
					leading = true
					go c.leader(ctx)
				}
			} else if c.r.State() == raft.Follower {
				if leading {
					cancel()
					leading = false
					ctx, cancel = context.WithCancel(context.Background())
				}
				if !following {
					following = true
					go c.follower(ctx)
				}
			}
		case <-c.stepdown:
			if leading {
				cancel()
				c.r.LeadershipTransfer().Error()
			}
		case <-c.stop:
			t.Stop()
			cancel()
			if leading {
				c.r.LeadershipTransfer().Error()
			}
			if c.callbacks.Cleanup != nil {
				ctx, _ := context.WithTimeout(context.Background(), time.Second*2)
				c.callbacks.Cleanup(ctx, c)
			}
			close(c.done)
			return c.fatal
		}
	}
	return c.fatal
}

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)
}

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() {
	close(c.stop)
	<-c.done
	s := make(chan struct{})
	go func() {
		c.r.Shutdown().Error()
		close(s)
	}()
	select {
	case <-s:
	case <-time.After(time.Second * 5):
	}
}