Distributed Coordination

package main

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

// --- Types ---

type Lease struct {
	Resource  string
	Holder    string
	Token     uint64 // monotonically increasing fencing token
	ExpiresAt time.Time
}

type LeaseErrorCode string

const (
	CodeHeldByOther LeaseErrorCode = "HELD_BY_OTHER"
	CodeNotHolder   LeaseErrorCode = "NOT_HOLDER"
	CodeExpired     LeaseErrorCode = "EXPIRED"
	CodeStaleToken  LeaseErrorCode = "STALE_TOKEN"
)

type LeaseError struct {
	Code LeaseErrorCode
	Msg  string
}

func (e *LeaseError) Error() string { return string(e.Code) + ": " + e.Msg }

// --- Lease store ---

type LeaseManager struct {
	mu        sync.Mutex
	leases    map[string]*Lease
	nextToken uint64
}

func NewLeaseManager() *LeaseManager {
	return &LeaseManager{leases: make(map[string]*Lease), nextToken: 1}
}

// Acquire grants a lease and always issues a fresh fencing token.
func (lm *LeaseManager) Acquire(resource, holder string, ttl time.Duration) (Lease, error) {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	now := time.Now()
	if existing, ok := lm.leases[resource]; ok {
		if existing.ExpiresAt.After(now) && existing.Holder != holder {
			return Lease{}, &LeaseError{
				Code: CodeHeldByOther,
				Msg:  fmt.Sprintf("%s held by %s", resource, existing.Holder),
			}
		}
	}

	lease := &Lease{
		Resource:  resource,
		Holder:    holder,
		Token:     lm.nextToken,
		ExpiresAt: now.Add(ttl),
	}
	lm.nextToken++
	lm.leases[resource] = lease
	return *lease, nil
}

// Renew extends the deadline and keeps the same token.
func (lm *LeaseManager) Renew(resource, holder string, token uint64, ttl time.Duration) (Lease, error) {
	lm.mu.Lock()
	defer lm.mu.Unlock()

	existing, ok := lm.leases[resource]
	if !ok || existing.Holder != holder || existing.Token != token {
		return Lease{}, &LeaseError{Code: CodeNotHolder, Msg: holder + " does not hold " + resource}
	}
	if !existing.ExpiresAt.After(time.Now()) {
		return Lease{}, &LeaseError{Code: CodeExpired, Msg: resource + " already expired"}
	}

	existing.ExpiresAt = time.Now().Add(ttl)
	return *existing, nil
}

func (lm *LeaseManager) Release(resource, holder string, token uint64) {
	lm.mu.Lock()
	defer lm.mu.Unlock()
	if existing, ok := lm.leases[resource]; ok {
		if existing.Holder == holder && existing.Token == token {
			delete(lm.leases, resource)
		}
	}
}

func (lm *LeaseManager) Inspect(resource string) (Lease, bool) {
	lm.mu.Lock()
	defer lm.mu.Unlock()
	existing, ok := lm.leases[resource]
	if !ok || !existing.ExpiresAt.After(time.Now()) {
		return Lease{}, false
	}
	return *existing, true
}

// expireNow is a test hook used to simulate a stalled holder.
func (lm *LeaseManager) expireNow(resource string) {
	lm.mu.Lock()
	defer lm.mu.Unlock()
	if existing, ok := lm.leases[resource]; ok {
		existing.ExpiresAt = time.Now().Add(-time.Second)
	}
}

// --- The resource that enforces safety ---

type FencedStore struct {
	mu      sync.Mutex
	data    map[string]string
	highest map[string]uint64
}

func NewFencedStore() *FencedStore {
	return &FencedStore{data: make(map[string]string), highest: make(map[string]uint64)}
}

func (fs *FencedStore) Write(key, value string, token uint64) error {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	if seen, ok := fs.highest[key]; ok && token < seen {
		return &LeaseError{
			Code: CodeStaleToken,
			Msg:  fmt.Sprintf("token %d for %s, already saw %d", token, key, seen),
		}
	}
	fs.highest[key] = token
	fs.data[key] = value
	return nil
}

func (fs *FencedStore) Read(key string) (string, bool) {
	fs.mu.Lock()
	defer fs.mu.Unlock()
	v, ok := fs.data[key]
	return v, ok
}

// --- Leader worker with background renewal ---

type LeaderWorker struct {
	manager    *LeaseManager
	resource   string
	nodeID     string
	ttl        time.Duration
	renewEvery time.Duration

	mu    sync.Mutex
	lease *Lease
	stop  chan struct{}
}

func NewLeaderWorker(m *LeaseManager, resource, nodeID string) *LeaderWorker {
	return &LeaderWorker{
		manager:    m,
		resource:   resource,
		nodeID:     nodeID,
		ttl:        9 * time.Second,
		renewEvery: 3 * time.Second,
		stop:       make(chan struct{}),
	}
}

func (w *LeaderWorker) Start(onLead func(uint64), onLose func()) {
	tick := func() {
		w.mu.Lock()
		defer w.mu.Unlock()

		if w.lease != nil {
			renewed, err := w.manager.Renew(w.resource, w.nodeID, w.lease.Token, w.ttl)
			if err != nil {
				w.lease = nil
				onLose()
				return
			}
			w.lease = &renewed
			return
		}

		acquired, err := w.manager.Acquire(w.resource, w.nodeID, w.ttl)
		if err != nil {
			var le *LeaseError
			if errors.As(err, &le) {
				fmt.Printf("[lease] %s not leader (%s)\n", w.nodeID, le.Code)
			}
			return
		}
		w.lease = &acquired
		onLead(acquired.Token)
	}

	go func() {
		tick()
		ticker := time.NewTicker(w.renewEvery)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				tick()
			case <-w.stop:
				return
			}
		}
	}()
}

func (w *LeaderWorker) Token() (uint64, bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.lease == nil {
		return 0, false
	}
	return w.lease.Token, true
}

func (w *LeaderWorker) Stop() {
	close(w.stop)
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.lease != nil {
		w.manager.Release(w.resource, w.nodeID, w.lease.Token)
		w.lease = nil
	}
}

// --- Demonstration ---

func main() {
	manager := NewLeaseManager()
	store := NewFencedStore()

	first, _ := manager.Acquire("warehouse-samarkand", "al-kindi", 5*time.Second)
	fmt.Printf("al-kindi holds token %d\n", first.Token)

	// al-kindi stalls long enough that the lease expires
	manager.expireNow("warehouse-samarkand")

	second, _ := manager.Acquire("warehouse-samarkand", "al-razi", 5*time.Second)
	fmt.Printf("al-razi holds token %d\n", second.Token)

	if err := store.Write("manifest", "written-by-al-razi", second.Token); err != nil {
		fmt.Println("unexpected:", err)
	}

	// al-kindi wakes up believing it is still the leader
	if err := store.Write("manifest", "written-by-al-kindi", first.Token); err != nil {
		fmt.Println("rejected:", err)
	}

	if v, ok := store.Read("manifest"); ok {
		fmt.Println("final value:", v)
	}
}

package main

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

type NodeState string

const (
	Follower  NodeState = "follower"
	Candidate NodeState = "candidate"
	Leader    NodeState = "leader"
)

type LogEntry struct {
	Term    int
	Command string
}

type RequestVoteArgs struct {
	Term         int
	CandidateID  string
	LastLogIndex int
	LastLogTerm  int
}

type RequestVoteReply struct {
	Term        int
	VoteGranted bool
}

type AppendEntriesArgs struct {
	Term     int
	LeaderID string
	Entries  []LogEntry
}

type AppendEntriesReply struct {
	Term    int
	Success bool
}

// --- In-memory transport ---

type Cluster struct {
	mu          sync.RWMutex
	nodes       map[string]*RaftNode
	partitioned map[string]bool
}

func NewCluster() *Cluster {
	return &Cluster{nodes: make(map[string]*RaftNode), partitioned: make(map[string]bool)}
}

func (c *Cluster) Register(n *RaftNode) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.nodes[n.id] = n
}

func (c *Cluster) IDs() []string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]string, 0, len(c.nodes))
	for id := range c.nodes {
		out = append(out, id)
	}
	return out
}

func (c *Cluster) Partition(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.partitioned[id] = true
}

func (c *Cluster) reachable(from, to string) (*RaftNode, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if c.partitioned[from] || c.partitioned[to] {
		return nil, false
	}
	n, ok := c.nodes[to]
	return n, ok
}

func (c *Cluster) RequestVote(from, to string, args RequestVoteArgs) (RequestVoteReply, bool) {
	target, ok := c.reachable(from, to)
	if !ok {
		return RequestVoteReply{}, false
	}
	time.Sleep(time.Duration(2+rand.Intn(8)) * time.Millisecond)
	return target.HandleRequestVote(args), true
}

func (c *Cluster) AppendEntries(from, to string, args AppendEntriesArgs) (AppendEntriesReply, bool) {
	target, ok := c.reachable(from, to)
	if !ok {
		return AppendEntriesReply{}, false
	}
	time.Sleep(time.Duration(2+rand.Intn(8)) * time.Millisecond)
	return target.HandleAppendEntries(args), true
}

// --- Raft node ---

type RaftNode struct {
	id      string
	cluster *Cluster

	mu sync.Mutex
	// persistent state
	currentTerm int
	votedFor    string
	log         []LogEntry
	// volatile state
	state           NodeState
	leaderID        string
	lastHeartbeat   time.Time
	electionTimeout time.Duration

	stopCh chan struct{}
}

func NewRaftNode(id string, c *Cluster) *RaftNode {
	n := &RaftNode{
		id:            id,
		cluster:       c,
		state:         Follower,
		lastHeartbeat: time.Now(),
		stopCh:        make(chan struct{}),
	}
	n.resetElectionTimeout()
	c.Register(n)
	return n
}

// Randomised timeout prevents perpetual split votes.
func (n *RaftNode) resetElectionTimeout() {
	n.electionTimeout = time.Duration(300+rand.Intn(300)) * time.Millisecond
}

func (n *RaftNode) Start() {
	go func() {
		ticker := time.NewTicker(50 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				n.tick()
			case <-n.stopCh:
				return
			}
		}
	}()
}

func (n *RaftNode) Stop() { close(n.stopCh) }

func (n *RaftNode) tick() {
	n.mu.Lock()
	if n.state == Leader || time.Since(n.lastHeartbeat) < n.electionTimeout {
		n.mu.Unlock()
		return
	}
	n.mu.Unlock()
	n.startElection()
}

// caller must hold n.mu
func (n *RaftNode) observeTermLocked(term int) {
	if term > n.currentTerm {
		n.currentTerm = term
		n.votedFor = ""
		n.state = Follower
		n.leaderID = ""
		n.lastHeartbeat = time.Now()
		n.resetElectionTimeout()
	}
}

func (n *RaftNode) lastLogIndexLocked() int { return len(n.log) - 1 }

func (n *RaftNode) lastLogTermLocked() int {
	if len(n.log) == 0 {
		return 0
	}
	return n.log[len(n.log)-1].Term
}

func (n *RaftNode) startElection() {
	n.mu.Lock()
	n.state = Candidate
	n.currentTerm++
	n.votedFor = n.id
	n.lastHeartbeat = time.Now()
	n.resetElectionTimeout()
	termAtStart := n.currentTerm
	args := RequestVoteArgs{
		Term:         termAtStart,
		CandidateID:  n.id,
		LastLogIndex: n.lastLogIndexLocked(),
		LastLogTerm:  n.lastLogTermLocked(),
	}
	n.mu.Unlock()

	fmt.Printf("[raft] %s starts election for term %d\n", n.id, termAtStart)

	peers := []string{}
	for _, id := range n.cluster.IDs() {
		if id != n.id {
			peers = append(peers, id)
		}
	}
	majority := (len(peers)+1)/2 + 1

	var wg sync.WaitGroup
	var voteMu sync.Mutex
	votes := 1

	for _, peer := range peers {
		wg.Add(1)
		go func(p string) {
			defer wg.Done()
			reply, ok := n.cluster.RequestVote(n.id, p, args)
			if !ok {
				return
			}

			n.mu.Lock()
			if reply.Term > n.currentTerm {
				n.observeTermLocked(reply.Term)
				n.mu.Unlock()
				return
			}
			stale := n.state != Candidate || n.currentTerm != termAtStart
			n.mu.Unlock()
			if stale || !reply.VoteGranted {
				return
			}

			voteMu.Lock()
			votes++
			won := votes == majority
			voteMu.Unlock()
			if won {
				n.becomeLeader(termAtStart)
			}
		}(peer)
	}
	wg.Wait()
}

func (n *RaftNode) becomeLeader(term int) {
	n.mu.Lock()
	if n.state != Candidate || n.currentTerm != term {
		n.mu.Unlock()
		return
	}
	n.state = Leader
	n.leaderID = n.id
	n.mu.Unlock()

	fmt.Printf("[raft] %s is leader for term %d\n", n.id, term)

	go func() {
		ticker := time.NewTicker(120 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if !n.sendHeartbeats() {
					return
				}
			case <-n.stopCh:
				return
			}
		}
	}()
}

// sendHeartbeats returns false once this node is no longer leader.
func (n *RaftNode) sendHeartbeats() bool {
	n.mu.Lock()
	if n.state != Leader {
		n.mu.Unlock()
		return false
	}
	args := AppendEntriesArgs{Term: n.currentTerm, LeaderID: n.id}
	n.mu.Unlock()

	peers := []string{}
	for _, id := range n.cluster.IDs() {
		if id != n.id {
			peers = append(peers, id)
		}
	}

	var wg sync.WaitGroup
	var mu sync.Mutex
	reachable := 1

	for _, peer := range peers {
		wg.Add(1)
		go func(p string) {
			defer wg.Done()
			reply, ok := n.cluster.AppendEntries(n.id, p, args)
			if !ok {
				return
			}
			mu.Lock()
			reachable++
			mu.Unlock()
			n.mu.Lock()
			n.observeTermLocked(reply.Term)
			n.mu.Unlock()
		}(peer)
	}
	wg.Wait()

	majority := (len(peers)+1)/2 + 1
	if reachable < majority {
		fmt.Printf("[raft] %s lost quorum, stepping down\n", n.id)
		n.mu.Lock()
		n.state = Follower
		n.leaderID = ""
		n.lastHeartbeat = time.Now()
		n.resetElectionTimeout()
		n.mu.Unlock()
		return false
	}
	return true
}

// --- RPC handlers ---

func (n *RaftNode) HandleRequestVote(args RequestVoteArgs) RequestVoteReply {
	n.mu.Lock()
	defer n.mu.Unlock()
	n.observeTermLocked(args.Term)

	if args.Term < n.currentTerm {
		return RequestVoteReply{Term: n.currentTerm, VoteGranted: false}
	}
	if n.votedFor != "" && n.votedFor != args.CandidateID {
		return RequestVoteReply{Term: n.currentTerm, VoteGranted: false}
	}

	// Election restriction keeps committed entries safe.
	upToDate := args.LastLogTerm > n.lastLogTermLocked() ||
		(args.LastLogTerm == n.lastLogTermLocked() && args.LastLogIndex >= n.lastLogIndexLocked())
	if !upToDate {
		return RequestVoteReply{Term: n.currentTerm, VoteGranted: false}
	}

	n.votedFor = args.CandidateID
	n.lastHeartbeat = time.Now()
	return RequestVoteReply{Term: n.currentTerm, VoteGranted: true}
}

func (n *RaftNode) HandleAppendEntries(args AppendEntriesArgs) AppendEntriesReply {
	n.mu.Lock()
	defer n.mu.Unlock()
	n.observeTermLocked(args.Term)

	if args.Term < n.currentTerm {
		return AppendEntriesReply{Term: n.currentTerm, Success: false}
	}

	n.state = Follower
	n.leaderID = args.LeaderID
	n.lastHeartbeat = time.Now()
	n.resetElectionTimeout()
	n.log = append(n.log, args.Entries...)
	return AppendEntriesReply{Term: n.currentTerm, Success: true}
}

func (n *RaftNode) Snapshot() string {
	n.mu.Lock()
	defer n.mu.Unlock()
	leader := n.leaderID
	if leader == "" {
		leader = "none"
	}
	return fmt.Sprintf("%s state=%s term=%d leader=%s", n.id, n.state, n.currentTerm, leader)
}

func (n *RaftNode) IsLeader() bool {
	n.mu.Lock()
	defer n.mu.Unlock()
	return n.state == Leader
}

func main() {
	cluster := NewCluster()
	names := []string{"baghdad", "cordoba", "damascus", "samarkand", "bukhara"}
	nodes := make([]*RaftNode, 0, len(names))
	for _, name := range names {
		nodes = append(nodes, NewRaftNode(name, cluster))
	}
	for _, n := range nodes {
		n.Start()
	}

	time.Sleep(1500 * time.Millisecond)
	for _, n := range nodes {
		fmt.Println(n.Snapshot())
	}

	for _, n := range nodes {
		if n.IsLeader() {
			fmt.Printf("--- partitioning %s ---\n", n.id)
			cluster.Partition(n.id)
			break
		}
	}

	time.Sleep(2 * time.Second)
	for _, n := range nodes {
		fmt.Println(n.Snapshot())
	}
	for _, n := range nodes {
		n.Stop()
	}
}