initial commit

pkg/raft/raft.go

@@ -0,0 +1,202 @@
+package raft
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"path/filepath"
+	"time"
+
+	"git.giftfish.de/ston1th/raftbbolt"
+	"git.giftfish.de/ston1th/vipman/pkg/config"
+	"git.giftfish.de/ston1th/vipman/pkg/raft/tls"
+	"git.giftfish.de/ston1th/vipman/pkg/util"
+	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
+	}
+	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 {
+		c.V(1).Info("setup", "transport", "tls")
+		s, err := tls.NewStream(cfg)
+		if err != nil {
+			return err
+		}
+		transport = raft.NewNetworkTransportWithLogger(s, pool, timeout, hcl)
+	} else {
+		c.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 nil
+		}
+	}
+	return nil
+}
+
+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) 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):
+	}
+}