vipman/pkg/cluster/raft.go

package cluster

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"runtime"
	"time"

	"git.giftfish.de/ston1th/vipman/pkg/config"
	"github.com/go-logr/logr"
	"github.com/hashicorp/raft"
	raftboltdb "github.com/hashicorp/raft-boltdb"
	"k8s.io/klog"
)

type raftLog string

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

var _ = ioutil.Discard

// StartRaftCluster - Begins a running instance of the Raft cluster
func (c *Cluster) StartRaftCluster(log logr.Logger, cfg *config.Config) error {

	// Create local configuration address
	localAddress := fmt.Sprintf("%s:%d", cfg.LocalPeer.Address, cfg.LocalPeer.Port)

	// Begin the Raft configuration
	rc := raft.DefaultConfig()
	rc.LocalID = raft.ServerID(cfg.LocalPeer.ID)
	rl := raftLog(cfg.LocalPeer.ID + "  ")
	rc.LogOutput = &rl
	//rc.LogOutput = ioutil.Discard

	// Initialize communication
	address, err := net.ResolveTCPAddr("tcp", localAddress)
	if err != nil {
		return err
	}

	// Create transport
	transport, err := raft.NewTCPTransport(localAddress, address, 3, 10*time.Second, os.Stdout)
	if err != nil {
		return err
	}

	// Create Raft structures
	snapshots := raft.NewInmemSnapshotStore()
	//logStore := raft.NewInmemStore()
	bootstrap := true
	if _, err := os.Stat(cfg.DBFile); err == nil {
		bootstrap = false
	}
	stableStore, err := raftboltdb.NewBoltStore(cfg.DBFile)
	if err != nil {
		return err
	}
	logStore, err := raftboltdb.NewBoltStore(cfg.DBFile + ".log")
	if err != nil {
		return err
	}

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

	// Add Local Peer
	configuration.Servers = append(configuration.Servers, raft.Server{
		ID:      raft.ServerID(cfg.LocalPeer.ID),
		Address: raft.ServerAddress(localAddress),
	})

	// If we want to start a node as leader then we will not add any remote peers, this will leave this as a cluster of one
	// The remotePeers will add themselves to the cluster as they're added
	if !cfg.Leader {
		for _, p := range cfg.RemotePeers {
			peerAddress := fmt.Sprintf("%s:%d", p.Address, p.Port)
			if localAddress != peerAddress {
				configuration.Servers = append(configuration.Servers, raft.Server{
					ID:      raft.ServerID(p.ID),
					Address: raft.ServerAddress(peerAddress)})
			}
		}
		log.Info("This node will attempt to start as Follower")
	} else {
		log.Info("This node will attempt to start as Leader")
	}

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

	// Create RAFT instance
	raftServer, err := raft.NewRaft(rc, c.stateMachine, logStore, stableStore, snapshots, transport)
	if err != nil {
		return err
	}

	ticker := time.NewTicker(time.Second)
	time.Sleep(time.Second * 5)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	lead := raftServer.Leader()
	leading := false
	for {
		l := raftServer.Leader()
		if lead != l && string(l) != "" {
			go c.Callbacks.OnNewLeader(string(l))
			lead = l
		}
		if !leading && localAddress == string(raftServer.Leader()) {
			leading = true
			go c.Run(ctx)
		}

		select {
		case leader := <-raftServer.LeaderCh():
			if leader {
				if !leading {
					leading = true
					go c.Run(ctx)
				}
			} else {
				if leading {
					cancel()
					leading = false
					ctx, cancel = context.WithCancel(context.Background())
				}
			}
		case <-ticker.C:
		case <-c.stop:
			if leading {
				cancel()
				raftServer.LeadershipTransfer()
			}
			close(c.done)
			return nil
		}
	}
	return nil
}

func (c *Cluster) Run(ctx context.Context) {
	defer func() {
		handleCrash()
		c.Callbacks.OnStoppedLeading()
	}()
	c.Callbacks.OnStartedLeading(ctx)
}

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) Stop() {
	close(c.stop)
	<-c.done
}