vipman/pkg/raft/raft.go

package raft

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

	"git.giftfish.de/ston1th/raftbbolt"
	"git.giftfish.de/ston1th/vipman/pkg/config"
	"git.giftfish.de/ston1th/vipman/pkg/raft/tls"
	"github.com/go-logr/logr"
	hclog "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/raft"
	"k8s.io/klog/v2"
)

type raftLog string

func (r *raftLog) Write(p []byte) (n int, err error) {
	return os.Stdout.Write(append([]byte(*r), p...))
}

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

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

	pool := 3
	timeout := 10 * time.Second
	var transport *raft.NetworkTransport
	if cfg.TLS != nil {
		log.Info("transport: tls")
		s, err := tls.NewStream(cfg)
		if err != nil {
			return err
		}
		transport = raft.NewNetworkTransportWithLogger(s, pool, timeout, hcl)
	} else {
		log.Info("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
		}
	}

	// Create RAFT instance
	c.r, err = raft.NewRaft(rc, (*fsm)(c.kv), logStore, stableStore, snapshots, transport)
	if err != nil {
		return err
	}
	c.kv.r = c.r

	t := time.NewTicker(time.Second)
	time.Sleep(time.Second * 3)

	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 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()
			}
			close(c.done)
			return nil
		case <-t.C:
		}
	}
	return nil
}

func (c *Cluster) leader(ctx context.Context) {
	defer func() {
		handleCrash()
	}()
	c.Callbacks.Leader(ctx, c.kv)
}

func (c *Cluster) follower(ctx context.Context) {
	defer func() {
		handleCrash()
	}()
	c.Callbacks.Follower(ctx, c.kv)
}

var reallyCrash = false

func handleCrash() {
	if r := recover(); r != nil {
		logPanic(r)
		if reallyCrash {
			panic(r)
		}
	}
}

// logPanic logs the caller tree when a panic occurs (except in the special case of http.ErrAbortHandler).
func logPanic(r interface{}) {
	if r == http.ErrAbortHandler {
		// honor the http.ErrAbortHandler sentinel panic value:
		//   ErrAbortHandler is a sentinel panic value to abort a handler.
		//   While any panic from ServeHTTP aborts the response to the client,
		//   panicking with ErrAbortHandler also suppresses logging of a stack trace to the server's error log.
		return
	}

	// Same as stdlib http server code. Manually allocate stack trace buffer size
	// to prevent excessively large logs
	const size = 64 << 10
	stacktrace := make([]byte, size)
	stacktrace = stacktrace[:runtime.Stack(stacktrace, false)]
	if _, ok := r.(string); ok {
		klog.Errorf("Observed a panic: %s\n%s", r, stacktrace)
	} else {
		klog.Errorf("Observed a panic: %#v (%v)\n%s", r, r, stacktrace)
	}
}

func (c *Cluster) Stepdown() {
	c.stepdown <- struct{}{}
}

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