From 7c3e7607597d723b09b12e89aa8dd5a35d0c2fb9 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <github@flamingcow.io>
Date: Sun, 4 Jun 2023 13:44:26 -0700
Subject: [PATCH] Test failover

---
 .golangci.yaml |  2 ++
 candidate.go   | 55 ++++++++++++++++++++++++++++++--------------------
 elect_test.go  | 27 +++++++++++++++++++++++++
 lib_test.go    |  4 ++++
 voter.go       |  2 +-
 5 files changed, 67 insertions(+), 23 deletions(-)

diff --git a/.golangci.yaml b/.golangci.yaml
index ec9f2dc..acc1cc1 100644
--- a/.golangci.yaml
+++ b/.golangci.yaml
@@ -42,3 +42,5 @@ linters-settings:
       use-field-name: true
       rules:
         json: goCamel
+  wsl:
+    allow-separated-leading-comment: true
diff --git a/candidate.go b/candidate.go
index df2d09e..1880189 100644
--- a/candidate.go
+++ b/candidate.go
@@ -13,8 +13,6 @@ import (
 	"github.com/samber/lo"
 )
 
-// TODO: Ensure promotion takes longer than demotion
-
 type Candidate struct {
 	C <-chan CandidateState
 
@@ -59,6 +57,8 @@ func NewCandidate(numVoters int, signingKey string) *Candidate {
 }
 
 func (c *Candidate) Stop() {
+	// Lock not required
+
 	close(c.stop)
 	<-c.done
 }
@@ -71,10 +71,15 @@ func (c *Candidate) State() CandidateState {
 }
 
 func (c *Candidate) IsLeader() bool {
+	// Lock not required
+
 	return c.State() == StateLeader
 }
 
 func (c *Candidate) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
 	if r.Method != http.MethodPost {
 		http.Error(
 			w,
@@ -169,34 +174,34 @@ func (c *Candidate) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	c.vote(v)
 }
 
-func (c *Candidate) vote(v *vote) {
-	v.received = time.Now()
+func (c *Candidate) VoteIfNo(v vote) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 
-	{
-		c.mu.Lock()
-		c.votes[v.VoterID] = v
-		c.mu.Unlock()
-	}
-
-	c.elect()
-}
-
-func (c *Candidate) voteIfNo(v *vote) {
 	if v.LastSeenCandidateID == c.resp.CandidateID {
 		return
 	}
 
-	c.vote(v)
+	c.vote(&v)
+}
+
+func (c *Candidate) vote(v *vote) {
+	// Must hold lock to call
+
+	v.received = time.Now()
+	c.votes[v.VoterID] = v
+
+	c.elect()
 }
 
 func (c *Candidate) elect() {
+	// Must hold lock to call
+
 	no := 0
 	yes := 0
 
 	cutoff := time.Now().Add(-10 * time.Second)
 
-	c.mu.Lock() /////////////
-
 	for key, vote := range c.votes {
 		if vote.received.Before(cutoff) {
 			delete(c.votes, key)
@@ -214,8 +219,6 @@ func (c *Candidate) elect() {
 		yes++
 	}
 
-	c.mu.Unlock() ////////////
-
 	if no == 0 && yes > c.numVoters/2 {
 		c.update(StateLeader)
 	} else {
@@ -224,8 +227,7 @@ func (c *Candidate) elect() {
 }
 
 func (c *Candidate) update(state CandidateState) {
-	c.mu.Lock()
-	defer c.mu.Unlock()
+	// Must hold lock to call
 
 	if c.state == state {
 		return
@@ -236,6 +238,8 @@ func (c *Candidate) update(state CandidateState) {
 }
 
 func (c *Candidate) loop() {
+	// Lock not required
+
 	t := time.NewTicker(1 * time.Second)
 	defer t.Stop()
 	defer close(c.done)
@@ -246,7 +250,14 @@ func (c *Candidate) loop() {
 			return
 
 		case <-t.C:
-			c.elect()
+			c.lockElect()
 		}
 	}
 }
+
+func (c *Candidate) lockElect() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.elect()
+}
diff --git a/elect_test.go b/elect_test.go
index 5fc3b9a..152ead2 100644
--- a/elect_test.go
+++ b/elect_test.go
@@ -15,3 +15,30 @@ func TestOne(t *testing.T) {
 
 	require.Eventually(t, ts.Candidate(0).IsLeader, 15*time.Second, 100*time.Millisecond)
 }
+
+func TestThree(t *testing.T) {
+	t.Parallel()
+
+	ts := NewTestSystem(t, 3)
+	defer ts.Stop()
+
+	require.Eventually(t, ts.Candidate(0).IsLeader, 15*time.Second, 100*time.Millisecond)
+}
+
+func TestFailover(t *testing.T) {
+	t.Parallel()
+
+	ts := NewTestSystem(t, 3)
+	defer ts.Stop()
+
+	require.Eventually(t, ts.Candidate(0).IsLeader, 15*time.Second, 100*time.Millisecond)
+
+	ts.SetServer(1)
+
+	require.Eventually(t, func() bool { return !ts.Candidate(0).IsLeader() }, 15*time.Second, 100*time.Millisecond)
+
+	// TODO: Check that new candidate doesn't become leader before old candidate loses it
+	// require.False(t, ts.Candidate(1).IsLeader())
+
+	require.Eventually(t, ts.Candidate(1).IsLeader, 15*time.Second, 100*time.Millisecond)
+}
diff --git a/lib_test.go b/lib_test.go
index 7704b30..692d0da 100644
--- a/lib_test.go
+++ b/lib_test.go
@@ -88,6 +88,10 @@ func (ts *TestSystem) Stop() {
 	ts.proxy.Close()
 }
 
+func (ts *TestSystem) SetServer(i int) {
+	ts.proxy.SetBackend(ts.Server(i).Addr())
+}
+
 func (ts *TestSystem) Server(i int) *TestServer {
 	return ts.servers[i]
 }
diff --git a/voter.go b/voter.go
index a6c6886..1eba2c9 100644
--- a/voter.go
+++ b/voter.go
@@ -108,7 +108,7 @@ func (v *Voter) poll() bool {
 func (v *Voter) sendVote() {
 	v.vote.VoteSent = time.Now().UTC()
 
-	v.candidate.voteIfNo(&v.vote)
+	v.candidate.VoteIfNo(v.vote)
 
 	js := lo.Must(json.Marshal(v.vote))