vipman/pkg/cluster/raft.go

package cluster

import (
	"context"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"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.ProtocolVersion = 3
	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()
	snapshots, err := raft.NewFileSnapshotStore(cfg.RaftDir, retainSnapshotCount, os.Stderr)
	if err != nil {
		return fmt.Errorf("file snapshot store: %s", err)
	}

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

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

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

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

	// Add Local Peer
	configuration.Servers = append(configuration.Servers, raft.Server{
		Suffrage: raft.Voter,
		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")
	//}
	for _, p := range cfg.RemotePeers {
		peerAddress := fmt.Sprintf("%s:%d", p.Address, p.Port)
		if localAddress != peerAddress {
			configuration.Servers = append(configuration.Servers, raft.Server{
				Suffrage: raft.Voter,
				ID:       raft.ServerID(p.ID),
				Address:  raft.ServerAddress(peerAddress)})
		}
	}

	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
	c.r, err = raft.NewRaft(rc, (*fsm)(c.kv), logStore, stableStore, snapshots, transport)
	if err != nil {
		return err
	}
	c.kv.r = c.r

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

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	leading := false
	following := false
	for {
		if !leading && localAddress == 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
					//cf := c.r.GetConfiguration()
					//cf.Error()
					//log.Info(fmt.Sprintf("[%s] %#v", cfg.LocalPeer.ID, cf.Configuration()))
					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
					//cf := c.r.GetConfiguration()
					//cf.Error()
					//log.Info(fmt.Sprintf("[%s] %#v", cfg.LocalPeer.ID, cf.Configuration()))
					go c.follower(ctx)
				}
			}
		case <-ticker.C:
		case <-c.stop:
			cancel()
			if leading {
				c.r.LeadershipTransfer().Error()
			}
			close(c.done)
			return nil
		}
	}
	return nil
}

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

func (c *Cluster) follower(ctx context.Context) {
	defer func() {
		handleCrash()
		//c.Callbacks.OnStoppedLeading()
	}()
	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) 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):
	}
}