Always try to send heartbeats to all nodes (not just followers) if leadership expiration is disabled

Otherwise, the following situation may happen:
- Node 1 is the leader and its term is 24
- Node 2 is a follower of 1, its term is 24 too
- Node 3 also thinks he's the leader, his term is 23, and he doesn't know that a new leader is elected

README.en.md

@@ -0,0 +1,106 @@
+# TinyRaft
+
+Raft leader election isolated from the rest of the algorithm.
+
+## Description
+
+TinyRaft can be used as a building block for the standard Raft with a K/V
+database and log-based replication or for other variations of "weak consensus".
+
+Leader is not required to replicate logs - he just has to support a recovery
+procedure making all replicas consistent. For example, he can use a CRDT or
+perform erasure coding on entries. He may use term number as logical time.
+The only requirement is to guarantee preservation of entries confirmed by
+all hosts participating in consensus.
+
+Supports leader expiration like in NuRaft:
+https://github.com/eBay/NuRaft/blob/master/docs/leadership_expiration.md
+
+## Why?
+
+Raft is a popular consensus algorithm. "Consensus" means making sure that
+multiple members of the cluster have the same view on some data and its changes
+over time. Some nodes may fall behind and have an older version of the data,
+but they shouldn't be able to see diverging changes.
+
+However, Raft achieves consensus in a pretty strict way called log replication:
+1. Raft applies all changes *strictly sequentially* to the leader's log,
+   sends that log to followers and applies it to persistent state when
+   all followers confirm it.
+2. If a previous leader becomes a follower and happens to have changes newer
+   than the current leader, these changes are thrown away. It's safe to do
+   because Raft doesn't confirm such changes to clients as "stable".
+3. If a follower has too old state which is impossible to "catch up" using
+   the log, full synchronisation is used and the whole state of the follower
+   is copied from the leader.
+
+That means Raft basically consists of 2 parts:
+- leader election
+- log replication (and state storage)
+
+But in some applications you may want to use leader election without log
+replication. Maybe because you want to replace it with your own state
+reconcilication algorithm, maybe because you don't want any persistence at
+all.
+
+There are other algorithms for leader election like [Bully](https://en.wikipedia.org/wiki/Bully_algorithm)
+but it seems it's partition-prone and thus unpopular. And if you try to modify
+it to make it more stable you quickly end up with something very similar to
+Raft's leader election. :-)
+
+But all Raft implementations I've seen contain both leader election and log
+replication as a whole and you can't untie them. This includes:
+- eBay's [NuRaft](https://github.com/eBay/NuRaft)
+- Baidu's [braft](https://github.com/brpc/braft)
+- Yandex's [Hydra](https://github.com/ytsaurus/ytsaurus/tree/main/yt/yt/server/lib/hydra)
+- [Hashicorp Raft](https://github.com/hashicorp/raft)
+- [etcd Raft](https://github.com/etcd-io/raft)
+- [Skiff](https://github.com/pgte/skiff-algorithm)
+- [Liferaft](https://github.com/unshiftio/liferaft)
+- Other implementations from https://raft.github.io/#implementations
+
+Some of these libraries don't even allow to supply a user-defined network layer
+or external event loop. Some require a full-featured K/V database underneath
+to work. Maybe you could trick these libraries and make them think that they
+always replicate an empty state, but it seems like a hack - libraries
+definitely weren't meant to be used that way.
+
+That's where TinyRaft comes to the help.
+
+## Ideas For Replacing Log Replication
+
+### CRDT
+
+Associate Raft's term number and an additional integer version number with
+each key. Increase version number on every write. Save a "tombstone" with an
+increased version number instead of deleting the key. Garbage collect deleted
+keys when all nodes confirm deletion. Make the newest (term, version) pair
+of each key win during synchronisation between nodes.
+
+It'll be slightly less strict than log replication but it's still OK.
+The only problem is possible short non-repeatable read for unconfirmed writes:
+- Client 1 requests a write of version 1 of key X during term 1
+- Node A is the leader, it applies the change, now it has X=(1,1)
+- Node A sends write to the network, but network fails and node A comes offline
+- Client 1 doesn't get the confirmation
+- A new leader is elected, now it's node B
+- Client 2 requests a read of key X from node B and gets nothing as the
+  key is not yet replicated
+- Node A comes back online, joins the leader (B) and synchronizes key X
+- Now node B also has X=(1,1)
+- Client 2 requests a read of key X from node B and now gets version (1,1)
+
+But it's fine for key/value databases with eventual consistency or for an
+object storage which doesn't require atomic multi-key transactions.
+
+### Erasure coding
+
+Split each value into chunks, calculate additional parity chunks using
+erasure codes (Reed-Solomon or similar), and instead of replicating the value
+to each node as is, make each node only store the corresponding chunk of it.
+
+## Usage example
+
+Full-featured usage example would require network code.
+
+Examples without real network layer can be seen in `tinyraft_test.go`.

README.md

@@ -0,0 +1,118 @@
+# TinyRaft
+
+Выбор лидера, как в алгоритме Raft, изолированный от остальной части алгоритма.
+
+## Описание
+
+TinyRaft можно использовать как готовый блок для реализации стандартного
+алгоритма Raft с БД "ключ-значение" и репликацией лога или для других вариаций
+"слабого консенсуса".
+
+Идея заключается в том, что, строго говоря, лидеру не обязательно реплицировать
+логи - он просто должен поддерживать некоторую процедуру приведения узлов
+в консистентное состояние. Например, он может использовать CRDT или использовать
+коды коррекции ошибок над записями. При этом лидер может использовать номер
+периода (term) как логическое время. Единственное требование к процедуре
+синхронизации - не терять изменения, подтверждённые всеми узлами кластера.
+
+В TinyRaft есть поддержка истечения лидерства, как в NuRaft:
+https://github.com/eBay/NuRaft/blob/master/docs/leadership_expiration.md
+
+## Зачем?!
+
+Raft - популярный алгоритм консенсуса. "Консенсус" означает гарантию единого
+представления некоторых данных и их изменений во времени всеми узлами кластера.
+Некоторые узлы могут отставать и иметь более старую версию данных, но они не
+должны видеть расходящуюся историю.
+
+Однако, Raft достигает консенсуса довольно строгими мерами - репликацией лога:
+1. Все изменения применяются *строго последовательно* к логу одного узла,
+   называемого лидером, он рассылает этот лог последователям и окончательно
+   применяет его только тогда, когда все последователи подтверждают получение.
+2. Если лидер меняется и оказывается, что бывший лидер успел применить
+   изменения, которые ещё не успел применить новый лидер, ранее бывший
+   последователем, эти изменения отбрасываются, чтобы лог нового последователя
+   совпадал с логом нового лидера. Это безопасно, так как Raft не подтверждает
+   такие изменения клиентам как "стабильные".
+3. Если у последователя оказывается слишком старое состояние, которое невозможно
+   довести до актуального применением лога, используется полная синхронизация
+   через снимок и полное состояние последователя копируется с лидера.
+
+А так как смена лидера - явление штатное и не редкое, такой способ репликации
+накладывает значительные ограничения на размер БД, синхронизируемой через Raft.
+По сути, если база будет большая, а любая реплика отстанет чуть больше, чем на
+незначительное время, то следующая синхронизация займёт вечность. :-)
+
+Поэтому, по сути, никто не использует Raft для синхронизации хоть сколько-нибудь
+больших БД. Например, в etcd максимальный размер БД по умолчанию - всего 2 ГБ
+и его не рекомендуется поднимать выше 8 ГБ. Если хранить данных надо побольше,
+все изобретают какое-то подобие PG (Ceph), например, в TiKV это называется
+Region.
+
+Таким образом, репликацию лога иногда использовать не хочется. Что остаётся в
+Raft, если её исключить? Остаётся алгоритм выбора лидера. При этом существуют
+и более легковесные алгоритмы выбора лидера типа [Bully](https://en.wikipedia.org/wiki/Bully_algorithm),
+но оные уязвимы к сетевым разделениям и непопулярны. Если же попытаться
+исправить Bully, вы быстро заметите, что изобрели что-то очень похожее на
+алгоритм выбора лидера Raft. :-)
+
+Все реализации Raft, которые я видел, содержат и выбор лидера, и репликацию
+логов, в нераздельном виде. Некоторые при этом сразу требуют наличия K/V базы
+данных, а некоторые не позволяют заменить сетевой слой на собственный или не
+интегрируются с внешним событийным циклом. Примеры:
+- eBay [NuRaft](https://github.com/eBay/NuRaft)
+- Baidu [braft](https://github.com/brpc/braft)
+- Yandex [Hydra](https://github.com/ytsaurus/ytsaurus/tree/main/yt/yt/server/lib/hydra)
+- [Hashicorp Raft](https://github.com/hashicorp/raft)
+- [etcd Raft](https://github.com/etcd-io/raft)
+- [Skiff](https://github.com/pgte/skiff-algorithm)
+- [Liferaft](https://github.com/unshiftio/liferaft)
+- Другие реализации со страницы https://raft.github.io/#implementations
+
+Возможно, можно было бы попробовать взять какую-то из библиотек и подсунуть ей
+пустую реализацию репликации логов и машины состояний, но скорее всего, это было
+бы неудобно и не факт, что вообще бы заработало - использовать библиотеки таким
+образом явно никто не предполагал.
+
+Поэтому и сделан TinyRaft.
+
+## Идеи для замены репликации лога
+
+### CRDT
+
+Связываем с каждым ключом пару: номер периода Raft (term) и дополнительное
+целое число - номер версии ключа. Синхронизацию лидера с последователями
+делаем, просто сравнивая пары (период, версия) всех ключей и считая, что
+побеждает ключ с максимальным значением пары. Вместо удаления - тоже
+увеличиваем версию и помечаем ключ, как удалённый, а фактически удаляем
+только когда все узлы подтвердят логическое удаление ключа.
+
+Подход чуть менее строгий, чем репликация логов, но всё равно нормальный.
+Из проблем в нём возможно только неповторяемое чтение для неподтверждённых
+записей:
+- Клиент 1 пишет версию 1 ключа X в периоде 1
+- Узел A - лидер, он получает версию X=(1,1)
+- Узел A рассылает изменение другим, но сеть отваливается и они его не получают
+- Клиент 1 тоже не получает подтверждения
+- Выбирается новый лидер, теперь это B
+- Клиент 2 читает ключ X с узла B и получает пустоту - у B ключа ещё нет
+- Узел А возвращается в сеть, присоединяется к B и синхронизирует данные
+- Клиент 2 снова читает ключ X с узла B и на этот раз получает версию (1,1)
+
+Однако, для баз данных ключ/значение или для объектных хранилищ, где не нужны
+атомарные транзакции поверх нескольких ключей - это не является проблемой.
+
+Да и для транзакционных, при наличии CAS-семантики, вполне возможно, тоже.
+
+### Коды коррекции ошибок
+
+Разделяем каждое значение на части, вычисляем дополнительные части с кодами
+избыточности (Рида-Соломона, erasure code) и вместо копирования значения на
+каждый узел целиком сохраняем только соответствующую его часть. При
+синхронизации восстанавливаем недостающие части из целого.
+
+## Пример использования
+
+Полноценный пример использования требует написания сетевого кода.
+
+Примеры без настоящего сетевого слоя можно посмотреть в файле `tinyraft_test.go`.

tinyraft.go

@@ -11,11 +11,18 @@
 //
 // Supports leader expiration like in NuRaft:
 // https://github.com/eBay/NuRaft/blob/master/docs/leadership_expiration.md
+//
+// Also supports leader priorities, similar to NuRaft but even simpler:
+// If a node receives a VoteRequest message with larger term but with smaller
+// priority than its own, it immediately starts a new voting round.
+// It guarantees that a node with non-maximum priority can't become leader
+// without being re-elected.
+// If all priorities are equal (or just zero), the election algorithm
+// becomes identical to the basic algorithm without priorities.
 
 package tinyraft
 
 import (
-    "context"
     "fmt"
     "math/rand"
     "sync"
@@ -41,48 +48,81 @@ const (
     StateFollower
 )
 
-type NodeID uint64
+// Node ID. Any value type can be used as node ID in TinyRaft
+type NodeID interface{}
 
-const NoLeader = NodeID(0)
+// Empty node ID
+var NoLeader = NodeID(nil)
 
+// Raft message
 type RaftMsg struct {
-    Type   MsgType
-    Term   uint64
-    Leader NodeID
+    Type     MsgType
+    Term     uint64
+    Leader   NodeID
+    Priority uint64
 }
 
+// TinyRaft configuration
 type RaftConfig struct {
+    // Local node ID
     NodeID            NodeID
+    // IDs of all nodes participating in this Raft group
     Nodes             []NodeID
+    // Term number to start with
     InitialTerm       uint64
+    // Election round starts when a follower doesn't see heartbeats (pings)
+    // from the leader in ElectionTimeout
     ElectionTimeout   time.Duration
+    // A node waits for an additional random time between 0 and RandomTimeout
+    // before starting election to prevent election collisions
     RandomTimeout     time.Duration
+    // Interval to send heartbeats to followers for the leader. Should be at
+    // most 1/2 of ElectionTimeout
     HeartbeatTimeout  time.Duration
+    // If set to a non-zero value, a follower responds with "pong" messages to
+    // the leader, and the leader cancels its leadership if it doesn't receive
+    // "pong" responses in LeadershipTimeout
     LeadershipTimeout time.Duration
+    // Disable automatic heartbeats if you do them yourself and call
+    // MarkAlive() manually
+    DisableHeartbeat  bool
+    // Leader priority of this node if non-zero
+    // When zero, leader priorities are disabled
+    LeaderPriority    uint64
+    // Function to send message to a given node from this one. TinyRaft
+    // contains a mutex, but this function is never called with a held mutex
     SendFunction      func(to NodeID, msg RaftMsg)
 }
 
+type RaftOutMsg struct {
+    To  NodeID
+    Msg RaftMsg
+}
+
 type Raft struct {
-    ctx      context.Context
     config   RaftConfig
     send     func(to NodeID, msg RaftMsg)
     onChange func(state NodeState, term uint64, leader NodeID, followers []NodeID)
 
-    mu sync.Mutex
+    mu       sync.Mutex
+    shutdown bool
 
     term      uint64
     state     NodeState
     leader    NodeID
     followers []NodeID
+    changed   bool
+    outbox    []RaftOutMsg
 
-    electionTimer   *time.Timer
-    electionTimerCh chan time.Duration
+    electionTimer  *time.Timer
+    heartbeatTimer *time.Timer
 
     voted int
     votes map[NodeID][]NodeID
 }
 
-func NewRaft(ctx context.Context, config RaftConfig) *Raft {
+// Create new instance of TinyRaft
+func NewRaft(config RaftConfig) *Raft {
     if config.ElectionTimeout == 0 || config.NodeID == 0 || config.SendFunction == nil {
         return nil
     }
@@ -90,61 +130,66 @@ func NewRaft(ctx context.Context, config RaftConfig) *Raft {
         config.Nodes = append(config.Nodes, config.NodeID)
     }
     r := &Raft{
-        ctx:     ctx,
         config:  config,
         send:    config.SendFunction,
         term:    config.InitialTerm,
     }
-    r.electionTimerCh = make(chan time.Duration, 1)
-    go func() {
-        ok := true
-        dur := time.Duration(0)
-        for ok {
-            if dur == 0 {
-                select {
-                case <-r.ctx.Done():
-                    return
-                case dur, ok = <-r.electionTimerCh:
-                }
-            } else {
-                select {
-                case <-r.ctx.Done():
-                    return
-                case dur, ok = <-r.electionTimerCh:
-                case <-time.After(dur):
-                    r.Start()
-                }
-            }
-        }
-    }()
-    go func() {
-        for {
-            select {
-            case <-r.ctx.Done():
-                return
-            case <-time.After(r.config.HeartbeatTimeout):
-                r.heartbeat()
-            }
-        }
-    }()
     return r
 }
 
+// Shutdown all background tasks and event handlers of a TinyRaft instance
+func (r *Raft) Shutdown() {
+    defer r.enter()()
+    r.shutdown = true
+    r.noNextTerm()
+    if r.heartbeatTimer != nil {
+        r.heartbeatTimer.Stop()
+        r.heartbeatTimer = nil
+    }
+    r.onChange = nil
+}
+
 func (r *Raft) noNextTerm() {
-    r.electionTimerCh <- 0
+    if r.electionTimer != nil {
+        r.electionTimer.Stop()
+        r.electionTimer = nil
+    }
 }
 
 func (r *Raft) nextTerm(after time.Duration) {
-    if r.config.RandomTimeout == 0 && after == 0 {
-        r.electionTimerCh <- time.Millisecond
-        return
-    }
-    r.electionTimerCh <- after + time.Duration(float64(r.config.RandomTimeout) * rand.Float64())
+    r.noNextTerm()
+    after += time.Duration(float64(r.config.RandomTimeout) * rand.Float64())
+    r.electionTimer = time.AfterFunc(after, r.Start)
 }
 
-func (r *Raft) Start() {
+func (r *Raft) enter() func() {
     r.mu.Lock()
-    defer r.mu.Unlock()
+    return func() {
+        if r.changed && r.onChange != nil {
+            // Call onChange without holding r.mu
+            defer r.onChange(r.state, r.term, r.leader, r.followers)
+        }
+        if len(r.outbox) > 0 {
+            // Call send without holding r.mu
+            defer func(outbox []RaftOutMsg) {
+                for _, o := range outbox {
+                    r.send(o.To, o.Msg)
+                }
+            }(r.outbox)
+            r.outbox = nil
+        }
+        r.changed = false
+        r.mu.Unlock()
+    }
+}
+
+// Start initial term (when instance is not initialized yet) or next term
+// (when it's already initialized)
+func (r *Raft) Start() {
+    defer r.enter()()
+    if r.shutdown {
+        return
+    }
     r.startUnlocked()
 }
 
@@ -157,99 +202,174 @@ func (r *Raft) startUnlocked() {
     r.leader = r.config.NodeID
     for _, node := range r.config.Nodes {
         if node != r.config.NodeID {
-            r.send(node, RaftMsg{Type: MsgVoteRequest, Term: r.term, Leader: r.leader})
+            r.outbox = append(r.outbox, RaftOutMsg{To: node, Msg: RaftMsg{Type: MsgVoteRequest, Term: r.term, Leader: r.leader, Priority: r.config.LeaderPriority}})
         }
     }
     // Next term will start right after this one times out
     r.nextTerm(r.config.ElectionTimeout)
-    r.changed()
+    r.changed = true
 }
 
 func (r *Raft) heartbeat() {
-    r.mu.Lock()
-    defer r.mu.Unlock()
-    if r.state == StateLeader {
-        for _, node := range r.votes[r.config.NodeID] {
-            if node != r.config.NodeID {
-                r.send(node, RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader})
-            }
+    defer r.enter()()
+    r.heartbeatUnlocked()
+    r.scheduleHeartbeat()
+}
+
+func (r *Raft) scheduleHeartbeat() {
+    if r.config.DisableHeartbeat {
+        return
+    }
+    if r.state != StateLeader {
+        return
+    }
+    if r.heartbeatTimer != nil {
+        r.heartbeatTimer.Stop()
+        r.heartbeatTimer = nil
+    }
+    r.heartbeatTimer = time.AfterFunc(r.config.HeartbeatTimeout, r.heartbeat)
+}
+
+func (r *Raft) heartbeatUnlocked() {
+    if r.state != StateLeader {
+        return
+    }
+    to := r.config.Nodes
+    if r.config.LeadershipTimeout > 0 {
+        to = r.followers
+    }
+    for _, node := range to {
+        if node != r.config.NodeID {
+            r.outbox = append(r.outbox, RaftOutMsg{To: node, Msg: RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader, Priority: r.config.LeaderPriority}})
         }
     }
 }
 
-func (r *Raft) changed() {
-    if r.onChange != nil {
-        r.onChange(r.state, r.term, r.leader, r.followers)
+// This method should be called for every incoming message for this node
+func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
+    defer r.enter()()
+    if r.shutdown {
+        return
+    }
+    switch msg.Type {
+    case MsgVoteRequest:
+        r.onReceiveVoteRequest(from, msg)
+    case MsgVote:
+        r.onReceiveVote(from, msg)
+    case MsgPing:
+        r.onReceivePing(from, msg)
+    case MsgPong:
+        r.onReceivePong(from, msg)
     }
 }
 
-func (r *Raft) OnReceive(from NodeID, msg RaftMsg) {
-    r.mu.Lock()
-    defer r.mu.Unlock()
-    if msg.Type == MsgVoteRequest {
-        if msg.Term > r.term && msg.Leader != NoLeader {
+func (r *Raft) onReceiveVoteRequest(from NodeID, msg RaftMsg) {
+    if msg.Term >= r.term && msg.Leader != NoLeader {
+        if msg.Priority < r.config.LeaderPriority {
+            r.term = msg.Term
+            r.startUnlocked()
+        } else if msg.Term > r.term {
             r.leader = msg.Leader
             r.term = msg.Term
             r.state = StateCandidate
             r.followers = nil
             r.nextTerm(r.config.HeartbeatTimeout*2 + r.config.ElectionTimeout)
-            r.changed()
+            r.changed = true
         }
-        r.send(from, RaftMsg{Type: MsgVote, Term: r.term, Leader: r.leader})
-    } else if msg.Type == MsgVote && msg.Leader != NoLeader && msg.Term == r.term {
+    }
+    var prio uint64
+    if r.leader == r.config.NodeID {
+        prio = r.config.LeaderPriority
+    }
+    r.outbox = append(r.outbox, RaftOutMsg{To: from, Msg: RaftMsg{Type: MsgVote, Term: r.term, Leader: r.leader, Priority: prio}})
+}
+
+func (r *Raft) onReceiveVote(from NodeID, msg RaftMsg) {
+    if msg.Leader == NoLeader || msg.Term < r.term {
+        return
+    }
+    if msg.Term > r.term {
+        r.term = msg.Term
+        r.state = StateFollower
+        r.followers = nil
+        r.leader = msg.Leader
+        r.votes = make(map[NodeID][]NodeID)
+        // Repeat MsgVote to the leader to join it
+        r.outbox = append(r.outbox, RaftOutMsg{To: r.leader, Msg: RaftMsg{Type: MsgVote, Term: r.term, Leader: r.leader, Priority: msg.Priority}})
+        r.changed = true
+        return
+    }
+    // add <from> as voter for <msg.Leader>
+    found := false
+    for _, voter := range r.votes[msg.Leader] {
+        if voter == from {
+            found = true
+            break
+        }
+    }
+    if !found {
         r.voted++
         r.votes[msg.Leader] = append(r.votes[msg.Leader], from)
-        n := len(r.votes[msg.Leader])
-        if n == 1 + len(r.config.Nodes)/2 {
-            if msg.Leader == r.config.NodeID {
-                r.state = StateLeader
-                r.leader = r.config.NodeID
-                r.followers = r.votes[r.config.NodeID]
-                for _, node := range r.followers {
-                    if node != r.config.NodeID {
-                        // Send a heartbeat to confirm leadership
-                        r.send(node, RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader})
-                    }
-                }
-                if r.config.LeadershipTimeout > 0 {
-                    r.nextTerm(r.config.LeadershipTimeout)
-                } else {
-                    r.noNextTerm()
-                }
-                r.changed()
-            } else {
-                r.nextTerm(0)
-            }
-        } else if n > 1 + len(r.config.Nodes)/2 && r.state == StateLeader && msg.Leader == r.config.NodeID {
+    }
+    n := len(r.votes[msg.Leader])
+    if n == 1 + len(r.config.Nodes)/2 {
+        if msg.Leader == r.config.NodeID {
+            r.state = StateLeader
+            r.leader = r.config.NodeID
             r.followers = r.votes[r.config.NodeID]
             // Send a heartbeat to confirm leadership
-            r.send(from, RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader})
-            r.changed()
-        } else if r.isVotingFailed() {
+            r.heartbeatUnlocked()
+            r.scheduleHeartbeat()
+            if r.config.LeadershipTimeout > 0 {
+                r.nextTerm(r.config.LeadershipTimeout)
+            } else {
+                r.noNextTerm()
+            }
+            r.changed = true
+        } else {
             r.nextTerm(0)
         }
-    } else if msg.Type == MsgPing {
-        if r.state == StateCandidate && r.term == msg.Term && from == r.leader {
-            r.state = StateFollower
-            r.changed()
-        }
-        if r.state == StateFollower && from == r.leader {
-            r.markAliveUnlocked()
-        }
-        if r.config.LeadershipTimeout > 0 {
-            r.send(from, RaftMsg{Type: MsgPong, Term: r.term, Leader: r.leader})
-        }
-    } else if msg.Type == MsgPong && r.state == StateLeader {
-        r.noNextTerm()
-        if msg.Leader != r.config.NodeID {
-            r.startUnlocked()
-        }
+    } else if n > 1 + len(r.config.Nodes)/2 && r.state == StateLeader && msg.Leader == r.config.NodeID {
+        r.followers = r.votes[r.config.NodeID]
+        // Send a heartbeat to confirm leadership
+        r.outbox = append(r.outbox, RaftOutMsg{To: from, Msg: RaftMsg{Type: MsgPing, Term: r.term, Leader: r.leader, Priority: r.config.LeaderPriority}})
+        r.changed = true
+    } else if r.isVotingFailed() {
+        r.nextTerm(0)
     }
 }
 
+func (r *Raft) onReceivePing(from NodeID, msg RaftMsg) {
+    if r.state == StateCandidate && r.term == msg.Term && from == r.leader {
+        r.state = StateFollower
+        r.changed = true
+    }
+    if r.state == StateFollower && from == r.leader {
+        r.markAliveUnlocked()
+    }
+    if r.config.LeadershipTimeout > 0 {
+        r.outbox = append(r.outbox, RaftOutMsg{To: from, Msg: RaftMsg{Type: MsgPong, Term: r.term, Leader: r.leader}})
+    }
+}
+
+func (r *Raft) onReceivePong(from NodeID, msg RaftMsg) {
+    if r.state != StateLeader {
+        return
+    }
+    if msg.Leader != r.config.NodeID {
+        r.startUnlocked()
+    } else if r.config.LeadershipTimeout > 0 {
+        r.nextTerm(r.config.LeadershipTimeout)
+    } else {
+        r.noNextTerm()
+    }
+}
+
+// Manually mark node alive, i.e. postpone the election timer.
+// You can call this method when a follower receives a replication message
+// from the leader so it doesn't rely on pings.
 func (r *Raft) MarkAlive() {
-    r.mu.Lock()
-    defer r.mu.Unlock()
+    defer r.enter()()
     r.markAliveUnlocked()
 }
 
@@ -266,7 +386,7 @@ func (r *Raft) isVotingFailed() bool {
     return true
 }
 
-func (r *Raft) hasNode(nodes []NodeID, n NodeID) bool {
+func hasNode(nodes []NodeID, n NodeID) bool {
     for _, x := range nodes {
         if x == n {
             return true
@@ -275,28 +395,41 @@ func (r *Raft) hasNode(nodes []NodeID, n NodeID) bool {
     return false
 }
 
+// Update the list of nodes participating in this Raft group.
+// Use to add or remove nodes.
 func (r *Raft) SetNodes(nodes []NodeID) {
-    r.mu.Lock()
-    defer r.mu.Unlock()
-    if len(nodes) == 0 {
+    defer r.enter()()
+    if !hasNode(nodes, r.config.NodeID) {
         nodes = append(nodes, r.config.NodeID)
     }
     r.config.Nodes = nodes
     if r.state == StateLeader {
         newFollowers := []NodeID(nil)
         for _, f := range r.followers {
-            if r.hasNode(r.config.Nodes, f) {
+            if hasNode(r.config.Nodes, f) {
                 newFollowers = append(newFollowers, f)
             }
         }
-        r.followers = newFollowers
-        r.changed()
-    } else if r.state == StateFollower && !r.hasNode(r.config.Nodes, r.leader) || r.state == StateCandidate {
+        if len(newFollowers) < 1+len(r.config.Nodes)/2 {
+            r.noNextTerm()
+            r.startUnlocked()
+        } else {
+            r.votes[r.config.NodeID] = newFollowers
+            r.followers = newFollowers
+            r.changed = true
+        }
+    } else if r.state == StateFollower && !hasNode(r.config.Nodes, r.leader) || r.state == StateCandidate {
         r.noNextTerm()
         r.startUnlocked()
     }
 }
 
+// Conversion function for MsgType and fmt.Printf("%v")
+func (t MsgType) String() string {
+    return t.GoString()
+}
+
+// Conversion function for MsgType and fmt.Printf("%#v")
 func (t MsgType) GoString() string {
     switch t {
     case MsgVoteRequest:
@@ -311,6 +444,12 @@ func (t MsgType) GoString() string {
     return fmt.Sprintf("unknown MsgType(%d)", t)
 }
 
+// Conversion function for NodeState and fmt.Printf("%v")
+func (t NodeState) String() string {
+    return t.GoString()
+}
+
+// Conversion function for NodeState and fmt.Printf("%#v")
 func (t NodeState) GoString() string {
     switch t {
     case StateCandidate:
@@ -323,36 +462,42 @@ func (t NodeState) GoString() string {
     return fmt.Sprintf("unknown NodeState(%d)", t)
 }
 
+// Get ID of the local node
 func (r *Raft) NodeID() NodeID {
     r.mu.Lock()
     defer r.mu.Unlock()
     return r.config.NodeID
 }
 
+// Get state of the local node
 func (r *Raft) State() NodeState {
     r.mu.Lock()
     defer r.mu.Unlock()
     return r.state
 }
 
+// Get current term. Valid when node is in StateLeader or StateFollower
 func (r *Raft) Term() uint64 {
     r.mu.Lock()
     defer r.mu.Unlock()
     return r.term
 }
 
+// Get current leader. Leader() equals NodeID() if the node is leader
 func (r *Raft) Leader() NodeID {
     r.mu.Lock()
     defer r.mu.Unlock()
     return r.leader
 }
 
+// Get the list of current followers. Valid when node is in StateLeader
 func (r *Raft) Followers() []NodeID {
     r.mu.Lock()
     defer r.mu.Unlock()
     return r.followers
 }
 
+// Set or clear node state change listener
 func (r *Raft) OnChange(cb func(state NodeState, term uint64, leader NodeID, followers []NodeID)) {
     r.mu.Lock()
     defer r.mu.Unlock()

tinyraft_test.go

@@ -1,47 +1,58 @@
-package tinyraft_test
+package tinyraft
 
 import (
-    "context"
     "fmt"
     "testing"
     "time"
-
-    . "git.yourcmc.ru/vitalif/tinyraft"
 )
 
-func addTestNode(ctx context.Context, nodes []*Raft, i int, partitions map[int]bool) {
+func addTestNode(nodes []*Raft, i int, partitions map[int]bool, mod func(cfg *RaftConfig)) {
     from := NodeID(i + 1)
-    n := NewRaft(ctx, RaftConfig{
+    cfg := RaftConfig{
         NodeID: from,
         Nodes: []NodeID{ 1, 2, 3, 4, 5 },
         HeartbeatTimeout: 100 * time.Millisecond,
         ElectionTimeout: 500 * time.Millisecond,
         RandomTimeout: 500 * time.Millisecond,
-        LeadershipTimeout: 0,
+        LeadershipTimeout: 1500 * time.Millisecond,
         SendFunction: func(to NodeID, msg RaftMsg) {
-            if !partitions[int((to-1)*100 + from-1)] && nodes[to-1] != nil {
+            intTo := to.(int)
+            intFrom := from.(int)
+            if !partitions[(intTo-1)*100 + intFrom-1] && nodes[intTo-1] != nil {
                 go func() {
                     fmt.Printf("received from %v to %v: %#v\n", from, to, msg)
-                    nodes[to-1].OnReceive(from, msg)
+                    if nodes[intTo-1] != nil {
+                        nodes[intTo-1].OnReceive(from, msg)
+                    }
                 }()
             }
         },
-    })
+    }
+    if mod != nil {
+        mod(&cfg)
+    }
+    n := NewRaft(cfg)
     n.OnChange(func(state NodeState, term uint64, leader NodeID, followers []NodeID) {
         fmt.Printf("node %v: %#v, term %v, leader: %v, followers: %v\n", from, state, term, leader, followers)
     })
     nodes[i] = n
 }
 
-func checkQuorum(t *testing.T, nodes []*Raft, count int, quorum int) {
+func checkQuorum(t *testing.T, nodes []*Raft, count int, quorum []NodeID) {
     leaders := 0
     for i := 0; i < count; i++ {
         if nodes[i] == nil {
             // Removed node
         } else if nodes[i].State() == StateLeader {
             leaders++
-            if len(nodes[i].Followers()) != quorum {
-                t.Fatalf("Leader (node %v) has less followers than %v: %v", i+1, quorum, nodes[i].Followers())
+            followers := nodes[i].Followers()
+            for _, qnode := range quorum {
+                if !hasNode(followers, qnode) {
+                    t.Fatalf("Leader (node %v) followers != %v: %v", i+1, quorum, followers)
+                }
+            }
+            if len(followers) != len(quorum) {
+                t.Fatalf("Leader (node %v) followers != %v: %v", i+1, quorum, followers)
             }
         } else if nodes[i].State() != StateFollower {
             t.Fatalf("Node %v state is neither leader nor follower: %#v", i+1, nodes[i].State())
@@ -68,13 +79,10 @@ func checkNoQuorum(t *testing.T, nodes []*Raft, count int) {
 }
 
 func TestStartThenRemoveNode(t *testing.T) {
-    var cancels []func()
     partitions := map[int]bool{}
     nodes := make([]*Raft, 5)
     for i := 0; i < 5; i++ {
-        ctx, cancel := context.WithCancel(context.Background())
-        cancels = append(cancels, cancel)
-        addTestNode(ctx, nodes, i, partitions)
+        addTestNode(nodes, i, partitions, nil)
     }
     ch1 := make(chan NodeID, 100)
     prevLeader1 := NoLeader
@@ -90,7 +98,8 @@ func TestStartThenRemoveNode(t *testing.T) {
     }
     // Check that 5 nodes are in quorum after 2000ms
     time.Sleep(2000 * time.Millisecond)
-    checkQuorum(t, nodes, 5, 5)
+    quorum := []NodeID{1, 2, 3, 4, 5}
+    checkQuorum(t, nodes, 5, quorum)
     leaderChanges := 0
     ok := true
     for ok {
@@ -105,61 +114,73 @@ func TestStartThenRemoveNode(t *testing.T) {
         t.Fatalf("leaderChanges = %v (expected < 3)", leaderChanges)
     }
     // Stop the leader
-    leader := nodes[0].Leader()
-    cancels[leader-1]()
+    leader := nodes[0].Leader().(int)
+    nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
-    cancels[leader-1] = nil
     // Check quorum after 2000ms
     time.Sleep(2000 * time.Millisecond)
-    checkQuorum(t, nodes, 5, 4)
+    for p := 0; p < len(quorum); p++ {
+        if quorum[p] == NodeID(leader) {
+            copy(quorum[p:], quorum[p+1:])
+            quorum = quorum[0:len(quorum)-1]
+            p--
+        }
+    }
+    checkQuorum(t, nodes, 5, quorum)
     // Stop the leader again
     for i := 0; i < 5; i++ {
         if nodes[i] != nil {
-            leader = nodes[i].Leader()
+            leader = nodes[i].Leader().(int)
         }
     }
-    cancels[leader-1]()
+    nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
-    cancels[leader-1] = nil
     // Check quorum after 2000ms
     time.Sleep(2000 * time.Millisecond)
-    checkQuorum(t, nodes, 5, 3)
+    for p := 0; p < len(quorum); p++ {
+        if quorum[p] == NodeID(leader) {
+            copy(quorum[p:], quorum[p+1:])
+            quorum = quorum[0:len(quorum)-1]
+            p--
+        }
+    }
+    checkQuorum(t, nodes, 5, quorum)
     // Stop the leader again
     for i := 0; i < 5; i++ {
         if nodes[i] != nil {
-            leader = nodes[i].Leader()
+            leader = nodes[i].Leader().(int)
         }
     }
-    cancels[leader-1]()
+    nodes[leader-1].Shutdown()
     nodes[leader-1] = nil
-    cancels[leader-1] = nil
     // Check for the absence of quorum
     time.Sleep(2000 * time.Millisecond)
     checkNoQuorum(t, nodes, 5)
     // Cleanup
     close(ch1)
-    for _, cancel := range cancels {
-        if cancel != nil {
-            cancel()
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
         }
     }
 }
 
 func TestAddNode(t *testing.T) {
-    ctx, cancel := context.WithCancel(context.Background())
     partitions := map[int]bool{}
     nodes := make([]*Raft, 6)
     for i := 0; i < 5; i++ {
-        addTestNode(ctx, nodes, i, partitions)
+        addTestNode(nodes, i, partitions, func(cfg *RaftConfig) {
+            cfg.InitialTerm = 1000
+        })
     }
     for i := 0; i < 5; i++ {
         nodes[i].Start()
     }
     // Check that 5 nodes are in quorum after 2000ms
     time.Sleep(2000 * time.Millisecond)
-    checkQuorum(t, nodes, 5, 5)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4, 5})
     // Add a node
-    addTestNode(ctx, nodes, 5, partitions)
+    addTestNode(nodes, 5, partitions, nil)
     nodes[5].Start()
     nodes[0].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     nodes[1].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
@@ -169,6 +190,242 @@ func TestAddNode(t *testing.T) {
     nodes[5].SetNodes([]NodeID{1, 2, 3, 4, 5, 6})
     // Check that 6 nodes are in quorum after 2000ms
     time.Sleep(2000 * time.Millisecond)
-    checkQuorum(t, nodes, 6, 6)
-    cancel()
+    checkQuorum(t, nodes, 6, []NodeID{1, 2, 3, 4, 5, 6})
+    for _, node := range nodes {
+        node.Shutdown()
+    }
+}
+
+func TestLeadershipExpiration(t *testing.T) {
+    partitions := map[int]bool{}
+    nodes := make([]*Raft, 5)
+    for i := 0; i < 5; i++ {
+        addTestNode(nodes, i, partitions, nil)
+    }
+    for i := 0; i < 5; i++ {
+        nodes[i].Start()
+    }
+    // Check that 5 nodes are in quorum after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4, 5})
+    // Break network on the leader
+    leader := nodes[0].Leader().(int)
+    fmt.Printf("Stopping the leader's (%v) network\n", leader)
+    for i := 0; i < 5; i++ {
+        nodes[i].mu.Lock()
+    }
+    for i := 0; i < 5; i++ {
+        partitions[(leader-1)*100 + i] = true
+        partitions[i*100 + leader-1] = true
+    }
+    for i := 0; i < 5; i++ {
+        nodes[i].mu.Unlock()
+    }
+    // Check that the leader loses leadership after 2 * LeadershipTimeout
+    time.Sleep(3000 * time.Millisecond)
+    if nodes[leader-1].State() != StateCandidate {
+        t.Fatalf("leadership expiration doesn't work")
+    }
+    // Cleanup
+    for _, node := range nodes {
+        node.Shutdown()
+    }
+}
+
+func TestRestart(t *testing.T) {
+    partitions := map[int]bool{}
+    nodes := make([]*Raft, 5)
+    for i := 0; i < 5; i++ {
+        addTestNode(nodes, i, partitions, func(cfg *RaftConfig) {
+            cfg.InitialTerm = 1000
+        })
+    }
+    leaderChanges := 0
+    prevLeader1 := NoLeader
+    nodes[1].OnChange(func(state NodeState, term uint64, leader NodeID, followers []NodeID) {
+        fmt.Printf("node %v: %#v, term %v, leader: %v, followers: %v\n", 2, state, term, leader, followers)
+        if state == StateCandidate {
+            leader = NoLeader
+        }
+        if leader != prevLeader1 {
+            prevLeader1 = leader
+            leaderChanges++
+        }
+    })
+    for i := 0; i < 5; i++ {
+        nodes[i].Start()
+    }
+    // Check that 5 nodes are in quorum after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4, 5})
+    if leaderChanges >= 3 {
+        t.Fatalf("leaderChanges = %v (expected < 3)", leaderChanges)
+    }
+    // Stop a follower
+    leader := nodes[0].Leader().(int)
+    restarted := 1 + leader % 5
+    if restarted == 2 {
+        restarted = 1 + (leader+1) % 5
+    }
+    fmt.Printf("stopping a follower (node %v)\n", restarted)
+    nodes[restarted-1].Shutdown()
+    nodes[restarted-1] = nil
+    // Wait 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    // Restart a follower
+    fmt.Printf("restarting a follower (node %v)\n", restarted)
+    leaderChanges = 0
+    addTestNode(nodes, restarted-1, partitions, nil)
+    nodes[restarted-1].Start()
+    // Check quorum and the fact that the leader didn't change after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4, 5})
+    if leaderChanges > 0 {
+        t.Fatalf("leader changed after restart of a follower")
+    }
+    // Cleanup
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
+        }
+    }
+}
+
+func TestChangeNodes(t *testing.T) {
+    partitions := map[int]bool{}
+    fmt.Printf("starting nodes 1-5\n")
+    nodes := make([]*Raft, 6)
+    for i := 0; i < 5; i++ {
+        addTestNode(nodes, i, partitions, func(cfg *RaftConfig) {
+            cfg.InitialTerm = 1000
+        })
+    }
+    for i := 0; i < 5; i++ {
+        nodes[i].Start()
+    }
+    // Check that 5 nodes are in quorum after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4, 5})
+    // Stop node 4
+    fmt.Printf("stopping node 4\n")
+    nodes[3].Shutdown()
+    nodes[3] = nil
+    time.Sleep(1000 * time.Millisecond)
+    // Change nodes from 1 2 3 4 5 to 1 2 3 5 6
+    fmt.Printf("starting node 6\n")
+    addTestNode(nodes, 5, partitions, nil)
+    nodes[5].Start()
+    nodes[0].SetNodes([]NodeID{1, 2, 3, 5, 6})
+    nodes[1].SetNodes([]NodeID{1, 2, 3, 5, 6})
+    nodes[2].SetNodes([]NodeID{1, 2, 3, 5, 6})
+    nodes[4].SetNodes([]NodeID{1, 2, 3, 5, 6})
+    nodes[5].SetNodes([]NodeID{1, 2, 3, 5, 6})
+    // Check that 5 nodes are in quorum after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 6, []NodeID{1, 2, 3, 5, 6})
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
+        }
+    }
+}
+
+func TestLeaderPriority(t *testing.T) {
+    partitions := map[int]bool{}
+    fmt.Printf("starting nodes 1-5\n")
+    nodes := make([]*Raft, 5)
+    for i := 0; i < 5; i++ {
+        addTestNode(nodes, i, partitions, func(cfg *RaftConfig) {
+            cfg.LeaderPriority = uint64(i+1)
+        })
+    }
+    for i := 0; i < 5; i++ {
+        nodes[i].Start()
+    }
+    // Check that 5 nodes are in quorum after 2000ms
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4, 5})
+    leader := nodes[0].Leader().(int)
+    if leader != 5 {
+        t.Fatalf("leader is not 5")
+    }
+    // Stop node 5
+    fmt.Printf("stopping node 5\n")
+    nodes[4].Shutdown()
+    nodes[4] = nil
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3, 4})
+    leader = nodes[0].Leader().(int)
+    if leader != 4 {
+        t.Fatalf("leader is not 4")
+    }
+    // Stop node 4
+    fmt.Printf("stopping node 4\n")
+    nodes[3].Shutdown()
+    nodes[3] = nil
+    time.Sleep(2000 * time.Millisecond)
+    checkQuorum(t, nodes, 5, []NodeID{1, 2, 3})
+    leader = nodes[0].Leader().(int)
+    if leader != 3 {
+        t.Fatalf("leader is not 3")
+    }
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
+        }
+    }
+}
+
+func TestPartition1_3(t *testing.T) {
+    partitions := map[int]bool{}
+    // Partition between 1 and 3
+    partitions[0*100 + 2] = true
+    partitions[2*100 + 0] = true
+    fmt.Printf("starting nodes 1-3\n")
+    nodes := make([]*Raft, 3)
+    for i := 0; i < 3; i++ {
+        addTestNode(nodes, i, partitions, func(cfg *RaftConfig) {
+            cfg.Nodes = []NodeID{ 1, 2, 3 }
+        })
+    }
+    ch1 := make(chan NodeID, 100)
+    prevLeader1 := NoLeader
+    nodes[1].OnChange(func(state NodeState, term uint64, leader NodeID, followers []NodeID) {
+        fmt.Printf("node %v: %#v, term %v, leader: %v, followers: %v\n", 2, state, term, leader, followers)
+        if state != StateCandidate && leader != NoLeader && leader != prevLeader1 {
+            prevLeader1 = leader
+            ch1 <- leader
+        }
+    })
+    for i := 0; i < 3; i++ {
+        nodes[i].Start()
+    }
+    // Check that 3 nodes are in quorum after 2000ms (and 2 should be the leader).
+    // 2 is always elected as leader because it's the first node that can determine
+    // that the first voting round has failed: it receives a vote for 1 from 1 and
+    // a vote for 3 from 3. So it knows that no node can get 2 votes. 1 and 3, at
+    // the same time, don't know it for sure because they don't receive failures
+    // from each other due to the network partition.
+    time.Sleep(2000 * time.Millisecond)
+    quorum := []NodeID{1, 2, 3}
+    checkQuorum(t, nodes, 3, quorum)
+    leaderChanges := 0
+    ok := true
+    for ok {
+        select {
+        case <-ch1:
+            leaderChanges++
+        default:
+            ok = false
+        }
+    }
+    if leaderChanges >= 3 {
+        t.Fatalf("leaderChanges = %v (expected < 3)", leaderChanges)
+    }
+    // Cleanup
+    for _, node := range nodes {
+        if node != nil {
+            node.Shutdown()
+        }
+    }
 }