hashi_raft_node.go

package funk

import (
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/raft"
	raftboltdb "github.com/hashicorp/raft-boltdb"
	"github.com/lab5e/clusterfunk/pkg/lg"
	"github.com/lab5e/clusterfunk/pkg/toolbox"
)

// hashiRaftNode is a wrapper for the Raft library. The raw events are coalesced into
// higher level events (particularly RaftClusterSizeChanged). Coalesced events
// introduce a small (millisecond) delay on the events but everything on top of
// this library will operate in the millisecond range.
//
// In addition this type keeps track of the active nodes at all times via the
// raft events. There's no guarantee that the list of nodes in the cluster will
// be up to date or correct for the followers. The followers will only
// interact with the leader of the cluster.
type hashiRaftNode struct {
	mutex            *sync.RWMutex // Mutex for the attributes
	fsmMutex         *sync.RWMutex // Mutex for the FSM
	scheduledMutex   *sync.Mutex   // Mutex for scheduled events
	scheduled        map[RaftEventType]time.Time
	localNodeID      string                        // The local node ID
	raftEndpoint     string                        // Raft endpoint
	ra               *raft.Raft                    // Raft instance
	events           chan RaftEventType            // Coalesced events from Raft
	unfilteredEvents chan RaftEventType            // Unfiltered events from Raft
	state            map[LogMessageType]LogMessage // The internal FSM state
	nodes            toolbox.StringSet
}

func newHashcorpRaftNode() RaftNode {
	return &hashiRaftNode{
		nodes:            toolbox.NewStringSet(),
		localNodeID:      "",
		mutex:            &sync.RWMutex{},
		fsmMutex:         &sync.RWMutex{},
		scheduledMutex:   &sync.Mutex{},
		scheduled:        make(map[RaftEventType]time.Time),
		events:           make(chan RaftEventType, 2), // tiny buffer here to make multiple events feasable.
		unfilteredEvents: make(chan RaftEventType, 5), // unfiltered events that gets coalesced into one big
		state:            make(map[LogMessageType]LogMessage),
	}
}

func (r *hashiRaftNode) Nodes() *toolbox.StringSet {
	return &r.nodes
}
func (r *hashiRaftNode) Start(nodeID string, cfg RaftParameters) (bool, error) {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	if r.ra != nil {
		return false, errors.New("raft cluster is already started")
	}

	config := raft.DefaultConfig()
	config.LocalID = raft.ServerID(nodeID)
	if cfg.Verbose {
		config.LogLevel = "INFO"
		if cfg.DebugLog {
			config.LogLevel = "DEBUG"
		}
	} else {
		config.LogOutput = newMutedLogger().Writer()
	}

	addr, err := net.ResolveTCPAddr("tcp", cfg.Endpoint)
	if err != nil {
		return false, err
	}

	switch cfg.TimingMode {
	case DefaultMode:
		dc := raft.DefaultConfig()
		config.HeartbeatTimeout = dc.HeartbeatTimeout
		config.ElectionTimeout = dc.ElectionTimeout
		config.CommitTimeout = dc.CommitTimeout
		config.SnapshotInterval = dc.SnapshotInterval
		config.LeaderLeaseTimeout = dc.LeaderLeaseTimeout
	case MediumMode:
		// Half the defaults
		config.HeartbeatTimeout = 500 * time.Millisecond
		config.ElectionTimeout = 500 * time.Millisecond
		config.CommitTimeout = 25 * time.Millisecond
		config.LeaderLeaseTimeout = 250 * time.Millisecond
	default:
		//These might be too optimistic.
		config.HeartbeatTimeout = 250 * time.Millisecond
		config.ElectionTimeout = 250 * time.Millisecond
		config.CommitTimeout = 25 * time.Millisecond
		config.LeaderLeaseTimeout = 125 * time.Millisecond
	}

	// The transport logging is separate form the configuration transport. Obviously.
	logger := io.Writer(os.Stderr)
	if !cfg.DebugLog {
		// Will only log the transport log as debug
		logger = newMutedLogger().Writer()
	}
	transport, err := raft.NewTCPTransport(addr.String(), addr, 3, 500*time.Millisecond, logger)
	if err != nil {
		return false, err
	}
	r.raftEndpoint = string(transport.LocalAddr())
	var logStore raft.LogStore
	var stableStore raft.StableStore
	var snapshotStore raft.SnapshotStore

	if cfg.DiskStore != "" {
		raftdir := fmt.Sprintf("%s/%s", cfg.DiskStore, nodeID)
		lg.Debug("Using boltDB and snapshot store at %s", raftdir)
		if err := os.MkdirAll(raftdir, os.ModePerm); err != nil {
			lg.Error("Unable to create data store dir at %s", raftdir)
			return false, err
		}
		boltDB, err := raftboltdb.NewBoltStore(filepath.Join(raftdir, fmt.Sprintf("%s.db", nodeID)))
		if err != nil {
			lg.Error("Unable to create BoltDB store: %v", err)
			return false, err
		}
		logStore = boltDB
		stableStore = boltDB
		snapshotStore, err = raft.NewFileSnapshotStore(raftdir, 3, os.Stderr)
		if err != nil {
			lg.Error("Unable to create snapshot store at %s: %v", raftdir, err)
			return false, err
		}
	} else {
		logStore = raft.NewInmemStore()
		stableStore = raft.NewInmemStore()
		snapshotStore = raft.NewInmemSnapshotStore()
	}
	r.ra, err = raft.NewRaft(config, r, logStore, stableStore, snapshotStore, transport)
	if err != nil {
		return false, err
	}

	bootstrap := false
	if cfg.Bootstrap {
		lg.Info("Bootstrapping new cluster")
		configuration := raft.Configuration{
			Servers: []raft.Server{
				{
					ID:      config.LocalID,
					Address: transport.LocalAddr(),
				},
			},
		}
		f := r.ra.BootstrapCluster(configuration)
		if f.Error() != nil {
			return false, f.Error()
		}
		bootstrap = true
	}
	observerChan := make(chan raft.Observation)

	// This node will - surprise - be a member of the cluster
	r.addNode(nodeID)

	go r.observerFunc(observerChan)
	go r.coalesceEvents()
	r.ra.RegisterObserver(raft.NewObserver(observerChan, true, func(*raft.Observation) bool { return true }))
	// Trigger events matching the current state. We might have become a leader or a follower at this point.
	switch r.ra.State() {
	case raft.Leader:
		r.sendInternalEvent(RaftClusterSizeChanged)
		r.sendInternalEvent(RaftBecameLeader)
	case raft.Follower:
		r.sendInternalEvent(RaftClusterSizeChanged)
		r.sendInternalEvent(RaftBecameFollower)
	default:
		// no op since we're still a candidate
	}
	r.localNodeID = nodeID

	return bootstrap, nil
}

func (r *hashiRaftNode) coalesceEvents() {
	for ev := range r.unfilteredEvents {
		timeout := false
		lastEvent := ev
		for !timeout {
			select {
			case ev := <-r.unfilteredEvents:
				if ev == lastEvent {
					continue
				}
				r.events <- lastEvent
				lastEvent = ev
				timeout = true
			case <-time.After(1 * time.Millisecond):
				timeout = true
			}
		}
		select {
		case r.events <- lastEvent:
		case <-time.After(5 * time.Second):
			lg.Debug("Event listener timed out after 5 seconds. Dropping event %s", ev)

		}
	}
	panic("Coalescing function has stopped")
}

func (r *hashiRaftNode) observerFunc(ch chan raft.Observation) {
	for k := range ch {
		switch v := k.Data.(type) {
		case raft.PeerObservation:
			if v.Removed {
				r.removeNode(string(v.Peer.ID))
				continue
			}
			r.addNode(string(v.Peer.ID))

		case raft.LeaderObservation:
			// This can be ignored since we're monitoring the state
			// and are getting the leader info via other channels.

		case raft.RaftState:
			switch v {
			case raft.Candidate:
				r.sendInternalEvent(RaftLeaderLost)

			case raft.Follower:
				r.sendInternalEvent(RaftBecameFollower)

			case raft.Leader:
				r.sendInternalEvent(RaftBecameLeader)
				// This might look a bit weird but the cluster size does not
				// change when there's only a single node becoming a leader
				r.scheduleInternalEvent(RaftClusterSizeChanged, 500*time.Millisecond)
			case raft.Shutdown:
				r.sendInternalEvent(RaftShutdown)
			}
		case raft.RequestVoteRequest:
			// Not using this at the moment. Just log it.
			lg.Debug("Node %s requested a vote", string(v.Candidate))

		case raft.FailedHeartbeatObservation:
			// Just log and ignoree
			lg.Debug("Node %s might be down (failed heartbeat)", v.PeerID)

		case raft.ResumedHeartbeatObservation:
			// Just log and ignoree
			lg.Debug("Node %s is back up (heartbeat OK)", v.PeerID)

		default:
			lg.Warning("Unknown Raft event: %v (%T)", k, v)
		}
	}
}

func (r *hashiRaftNode) Stop(removeWhenStopping bool) error {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	if r.ra == nil {
		return errors.New("raft cluster is already stopped")
	}

	if err := r.ra.Shutdown().Error(); err != nil {
		lg.Error("Got error on shutdown: %v", err)
	}
	r.ra = nil
	r.localNodeID = ""
	r.raftEndpoint = ""
	return nil
}

func (r *hashiRaftNode) LocalNodeID() string {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	return r.localNodeID
}

func (r *hashiRaftNode) AddClusterNode(nodeID string, endpoint string) error {
	r.mutex.RLock()
	defer r.mutex.RUnlock()

	if r.ra == nil {
		return errors.New("raft cluster is not started")
	}

	if err := r.ra.VerifyLeader().Error(); err != nil {
		// Not the leader so can't add node
		return errors.New("must be leader to add a new member")
	}

	configFuture := r.ra.GetConfiguration()
	if err := configFuture.Error(); err != nil {
		return err
	}

	for _, srv := range configFuture.Configuration().Servers {
		if srv.ID == raft.ServerID(nodeID) && srv.Address == raft.ServerAddress(endpoint) {
			// it's already joined
			return nil
		}
	}

	f := r.ra.AddVoter(raft.ServerID(nodeID), raft.ServerAddress(endpoint), 0, 0)
	if f.Error() != nil {
		return f.Error()
	}
	r.addNode(nodeID)
	return nil
}

func (r *hashiRaftNode) RemoveClusterNode(nodeID string, endpoint string) error {
	r.mutex.RLock()
	defer r.mutex.RUnlock()

	if r.ra == nil {
		return errors.New("raft cluster isn't started")
	}
	if r.ra.VerifyLeader().Error() != nil {
		return errors.New("must be leader to remove ndoe")
	}

	configFuture := r.ra.GetConfiguration()
	if err := configFuture.Error(); err != nil {
		return err
	}

	r.removeNode(nodeID)
	for _, srv := range configFuture.Configuration().Servers {
		if srv.ID == raft.ServerID(nodeID) && srv.Address == raft.ServerAddress(endpoint) {
			return r.ra.RemoveServer(raft.ServerID(nodeID), 0, 0).Error()
		}
	}

	// The server does not exist in the cluster - *technically* an error but
	// it's no longer in the cluster so we're good.
	return nil
}

func (r *hashiRaftNode) Endpoint() string {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	return r.raftEndpoint
}

func (r *hashiRaftNode) Leader() bool {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	if r.ra == nil {
		return false
	}
	return (r.ra.VerifyLeader().Error() == nil)
}

func (r *hashiRaftNode) AppendLogEntry(data []byte) (uint64, error) {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	if r.ra == nil {
		return 0, errors.New("raft node not started")
	}
	f := r.ra.Apply(data, time.Second*2)
	if err := f.Error(); err != nil {
		return 0, err
	}
	return f.Index(), nil
}

func (r *hashiRaftNode) LastLogIndex() uint64 {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	if r.ra == nil {
		return 0
	}
	return r.ra.AppliedIndex()
}

func (r *hashiRaftNode) Events() <-chan RaftEventType {
	return r.events
}

func (r *hashiRaftNode) GetLogMessages(startingIndex uint64) []LogMessage {
	r.fsmMutex.Lock()
	defer r.fsmMutex.Unlock()
	ret := make([]LogMessage, 0)
	for _, v := range r.state {
		if v.Index > startingIndex {
			ret = append(ret, v)
		}
	}
	return ret
}

func (r *hashiRaftNode) StepDown() error {
	if !r.Leader() {
		return errors.New("not the leader")
	}
	return r.ra.LeadershipTransfer().Error()
}

func (r *hashiRaftNode) addNode(id string) {
	if r.Nodes().Add(id) {
		r.sendInternalEvent(RaftClusterSizeChanged)
	}
}

func (r *hashiRaftNode) removeNode(id string) {
	if r.Nodes().Remove(id) {
		r.sendInternalEvent(RaftClusterSizeChanged)
	}
}

func (r *hashiRaftNode) RefreshNodes() {
	cfg := r.ra.GetConfiguration()
	if cfg.Error() != nil {
		lg.Error("Unable to update nodes: %v", cfg.Error())
		return
	}
	list := []string{}
	for _, v := range cfg.Configuration().Servers {
		list = append(list, string(v.ID))
	}
	r.Nodes().Sync(list...)
}

func (r *hashiRaftNode) sendInternalEvent(ev RaftEventType) {
	select {
	case r.unfilteredEvents <- ev:
		// Remove aync scheduled events of this type.
		r.scheduledMutex.Lock()
		delete(r.scheduled, ev)
		r.scheduledMutex.Unlock()
	case <-time.After(10 * time.Second):
		// This might be caused by poor sync between the Raft internal state and
		// the library state. Just log the error.
		lg.Warning("Unable to send internal event %s. Channel full?", ev.String())
	}
}

func (r *hashiRaftNode) scheduleInternalEvent(ev RaftEventType, timeout time.Duration) {
	r.scheduledMutex.Lock()
	r.scheduled[ev] = time.Now().Add(timeout)
	r.scheduledMutex.Unlock()
	go func() {
		time.Sleep(timeout)
		r.scheduledMutex.Lock()
		_, ok := r.scheduled[ev]
		r.scheduledMutex.Unlock()
		if ok {
			// send the event. Log for now since it happens only in certain
			// circumstances. In most circumstances a change in leadership
			// happens because a node goes down or fails and then a
			// cluster size notification is sent but on rare occasions when
			// a node silently fails it won't trigger a size change event.
			lg.Warning("Did not get the expected event %s in %s. Generating it.", ev, timeout/time.Millisecond)
			r.sendInternalEvent(ev)
		}
	}()
}

func (r *hashiRaftNode) EnableNode(id string) {
	if !r.Leader() {
		return
	}
	r.addNode(id)
}

func (r *hashiRaftNode) DisableNode(id string) {
	if !r.Leader() {
		return
	}
	r.removeNode(id)
}

func (r *hashiRaftNode) LeaderNodeID() string {
	if r.Leader() {
		return r.LocalNodeID()
	}

	list, err := r.MemberList()
	if err != nil {
		// TODO: Handle gracefully (sort of - this will end up as an error later)
		return ""
	}
	for _, v := range list {
		if v.Leader {
			return v.ID
		}
	}
	return ""
}

// The raft.FSM implementation. Right now the implementation looks a lot more
// like a storage layer but technically it's a FSM

func (r *hashiRaftNode) Apply(l *raft.Log) interface{} {
	msg := LogMessage{}
	if err := msg.UnmarshalBinary(l.Data); err != nil {
		panic(fmt.Sprintf(" ***** Error decoding log message: %v", err))
	}
	r.fsmMutex.Lock()
	defer r.fsmMutex.Unlock()
	msg.Index = l.Index
	r.state[msg.MessageType] = msg
	r.sendInternalEvent(RaftReceivedLog)
	return l.Data
}

// Snapshot is used to support log compaction. This call should
// return an FSMSnapshot which can be used to save a point-in-time
// snapshot of the FSM. Apply and Snapshot are not called in multiple
// threads, but Apply will be called concurrently with Persist. This means
// the FSM should be implemented in a fashion that allows for concurrent
// updates while a snapshot is happening.
func (r *hashiRaftNode) Snapshot() (raft.FSMSnapshot, error) {
	return &raftSnapshot{}, nil
}

// Restore is used to restore an FSM from a snapshot. It is not called
// concurrently with any other command. The FSM must discard all previous
// state.
func (r *hashiRaftNode) Restore(io.ReadCloser) error {
	lg.Debug("FSMSnapshot Restore")
	return nil
}

// memberList returns a list of nodes in the raft cluster.
func (r *hashiRaftNode) MemberList() ([]Node, error) {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	if r.ra == nil {
		return nil, errors.New("raft cluster is not started")
	}
	config := r.ra.GetConfiguration()
	if err := config.Error(); err != nil {
		return nil, err
	}
	leader := r.ra.Leader()

	members := config.Configuration().Servers
	ret := make([]Node, len(members))
	for i, v := range members {
		ret[i] = Node{
			ID:     string(v.ID),
			State:  v.Suffrage.String(),
			Leader: (v.Address == leader),
		}
	}
	return ret, nil
}

type raftSnapshot struct {
}

// Persist should dump all necessary state to the WriteCloser 'sink',
// and call sink.Close() when finished or call sink.Cancel() on error.
func (r *raftSnapshot) Persist(sink raft.SnapshotSink) error {
	lg.Debug("FSMSnapshot Persist")
	sink.Close()
	return nil
}

// Release is invoked when we are finished with the snapshot.
func (r *raftSnapshot) Release() {
	// nothing happens here.
	lg.Debug("FSMSnapshot Release")
}