Merge pull request #141 from bkeroack/store-cluster-state

Store methods to expose node Raft state

store/store.go

@@ -91,6 +91,18 @@ const (
 	Strong
 )
 
+// ClusterState defines the possible Raft states the current node can be in
+type ClusterState int
+
+// Represents the Raft cluster states
+const (
+	Leader ClusterState = iota
+	Follower
+	Candidate
+	Shutdown
+	Unknown
+)
+
 // clusterMeta represents cluster meta which must be kept in consensus.
 type clusterMeta struct {
 	APIPeers map[string]string // Map from Raft address to API address
@@ -254,6 +266,28 @@ func (s *Store) Close(wait bool) error {
 	return nil
 }
 
+// IsLeader is used to determine if the current node is cluster leader
+func (s *Store) IsLeader() bool {
+	return s.raft.State() == raft.Leader
+}
+
+// State returns the current node's Raft state
+func (s *Store) State() ClusterState {
+	state := s.raft.State()
+	switch state {
+	case raft.Leader:
+		return Leader
+	case raft.Candidate:
+		return Candidate
+	case raft.Follower:
+		return Follower
+	case raft.Shutdown:
+		return Shutdown
+	default:
+		return Unknown
+	}
+}
+
 // JoinRequired returns whether the node needs to join a cluster after being opened.
 func (s *Store) JoinRequired() bool {
 	return s.joinRequired

store/store_test.go

@@ -459,6 +459,37 @@ func Test_APIPeers(t *testing.T) {
 	}
 }
 
+func Test_IsLeader(t *testing.T) {
+	s := mustNewStore(true)
+	defer os.RemoveAll(s.Path())
+
+	if err := s.Open(true); err != nil {
+		t.Fatalf("failed to open single-node store: %s", err.Error())
+	}
+	defer s.Close(true)
+	s.WaitForLeader(10 * time.Second)
+
+	if !s.IsLeader() {
+		t.Fatalf("single node is not leader!")
+	}
+}
+
+func Test_State(t *testing.T) {
+	s := mustNewStore(true)
+	defer os.RemoveAll(s.Path())
+
+	if err := s.Open(true); err != nil {
+		t.Fatalf("failed to open single-node store: %s", err.Error())
+	}
+	defer s.Close(true)
+	s.WaitForLeader(10 * time.Second)
+
+	state := s.State()
+	if state != Leader {
+		t.Fatalf("single node returned incorrect state (not Leader): %v", s)
+	}
+}
+
 func mustNewStore(inmem bool) *Store {
 	path := mustTempDir()
 	defer os.RemoveAll(path)