Distributed Consensus

import dgram from 'node:dgram';

// --- Types ---
type NodeState = 'follower' | 'candidate' | 'leader';

interface VoteRequest {
	type: 'vote_request';
	term: number;
	candidateId: string;
	lastLogIndex: number;
	lastLogTerm: number;
}

interface VoteResponse {
	type: 'vote_response';
	term: number;
	voteGranted: boolean;
	voterId: string;
}

interface Heartbeat {
	type: 'heartbeat';
	term: number;
	leaderId: string;
}

interface HeartbeatResponse {
	type: 'heartbeat_response';
	term: number;
	nodeId: string;
	success: boolean;
}

type RaftMessage = VoteRequest | VoteResponse | Heartbeat | HeartbeatResponse;

interface Peer {
	id: string;
	host: string;
	port: number;
}

// --- Raft Node ---
class RaftNode {
	private state: NodeState = 'follower';
	private currentTerm = 0;
	private votedFor: string | null = null;
	private leaderId: string | null = null;
	private votesReceived = new Set<string>();

	private electionTimeout: ReturnType<typeof setTimeout> | null = null;
	private heartbeatInterval: ReturnType<typeof setInterval> | null = null;

	private socket: dgram.Socket;

	// Timing (randomized to prevent split votes)
	private readonly ELECTION_TIMEOUT_MIN = 1500; // ms
	private readonly ELECTION_TIMEOUT_MAX = 3000;
	private readonly HEARTBEAT_INTERVAL = 500;

	constructor(
		private readonly id: string,
		private readonly port: number,
		private readonly peers: Peer[],
		private readonly logIndex: number = 0,
		private readonly logTerm: number = 0
	) {
		this.socket = dgram.createSocket('udp4');
		this.socket.on('message', (data) => this.handleMessage(data));
		this.socket.bind(port);
	}

	start(): void {
		console.log(`[${this.id}] Starting as follower (term ${this.currentTerm})`);
		this.resetElectionTimeout();
	}

	private randomElectionTimeout(): number {
		return (
			this.ELECTION_TIMEOUT_MIN +
			Math.random() * (this.ELECTION_TIMEOUT_MAX - this.ELECTION_TIMEOUT_MIN)
		);
	}

	private resetElectionTimeout(): void {
		if (this.electionTimeout) clearTimeout(this.electionTimeout);
		this.electionTimeout = setTimeout(() => this.startElection(), this.randomElectionTimeout());
	}

	// --- State transitions ---
	private startElection(): void {
		this.state = 'candidate';
		this.currentTerm++;
		this.votedFor = this.id;
		this.votesReceived.clear();
		this.votesReceived.add(this.id); // vote for self

		console.log(`[${this.id}] Starting election for term ${this.currentTerm}`);

		// Request votes from all peers
		const request: VoteRequest = {
			type: 'vote_request',
			term: this.currentTerm,
			candidateId: this.id,
			lastLogIndex: this.logIndex,
			lastLogTerm: this.logTerm
		};

		for (const peer of this.peers) {
			this.send(peer, request);
		}

		this.resetElectionTimeout();
	}

	private becomeLeader(): void {
		if (this.state !== 'candidate') return;

		this.state = 'leader';
		this.leaderId = this.id;
		console.log(`[${this.id}] Became LEADER for term ${this.currentTerm}`);

		// Stop election timer
		if (this.electionTimeout) clearTimeout(this.electionTimeout);

		// Start sending heartbeats
		this.sendHeartbeats();
		this.heartbeatInterval = setInterval(() => this.sendHeartbeats(), this.HEARTBEAT_INTERVAL);
	}

	private stepDown(newTerm: number): void {
		console.log(`[${this.id}] Stepping down. Old term: ${this.currentTerm}, new term: ${newTerm}`);
		this.state = 'follower';
		this.currentTerm = newTerm;
		this.votedFor = null;
		this.leaderId = null;

		if (this.heartbeatInterval) {
			clearInterval(this.heartbeatInterval);
			this.heartbeatInterval = null;
		}

		this.resetElectionTimeout();
	}

	// --- Message handling ---
	private handleMessage(data: Buffer): void {
		const msg: RaftMessage = JSON.parse(data.toString());

		switch (msg.type) {
			case 'vote_request':
				this.handleVoteRequest(msg);
				break;
			case 'vote_response':
				this.handleVoteResponse(msg);
				break;
			case 'heartbeat':
				this.handleHeartbeat(msg);
				break;
			case 'heartbeat_response':
				this.handleHeartbeatResponse(msg);
				break;
		}
	}

	private handleVoteRequest(req: VoteRequest): void {
		// If request has higher term, step down
		if (req.term > this.currentTerm) {
			this.stepDown(req.term);
		}

		let granted = false;

		if (
			req.term >= this.currentTerm &&
			(this.votedFor === null || this.votedFor === req.candidateId) &&
			(req.lastLogTerm > this.logTerm ||
				(req.lastLogTerm === this.logTerm && req.lastLogIndex >= this.logIndex))
		) {
			granted = true;
			this.votedFor = req.candidateId;
			this.resetElectionTimeout(); // reset since we granted a vote
		}

		const response: VoteResponse = {
			type: 'vote_response',
			term: this.currentTerm,
			voteGranted: granted,
			voterId: this.id
		};

		const peer = this.peers.find((p) => p.id === req.candidateId);
		if (peer) this.send(peer, response);
	}

	private handleVoteResponse(res: VoteResponse): void {
		if (res.term > this.currentTerm) {
			this.stepDown(res.term);
			return;
		}

		if (this.state !== 'candidate' || res.term !== this.currentTerm) return;

		if (res.voteGranted) {
			this.votesReceived.add(res.voterId);
			const majority = Math.floor((this.peers.length + 1) / 2) + 1;

			console.log(
				`[${this.id}] Got vote from ${res.voterId} (${this.votesReceived.size}/${majority} needed)`
			);

			if (this.votesReceived.size >= majority) {
				this.becomeLeader();
			}
		}
	}

	private handleHeartbeat(hb: Heartbeat): void {
		if (hb.term >= this.currentTerm) {
			if (this.state !== 'follower') {
				this.stepDown(hb.term);
			}
			this.currentTerm = hb.term;
			this.leaderId = hb.leaderId;
			this.resetElectionTimeout();
		}

		const peer = this.peers.find((p) => p.id === hb.leaderId);
		if (peer) {
			const response: HeartbeatResponse = {
				type: 'heartbeat_response',
				term: this.currentTerm,
				nodeId: this.id,
				success: hb.term >= this.currentTerm
			};
			this.send(peer, response);
		}
	}

	private handleHeartbeatResponse(_res: HeartbeatResponse): void {
		// In a full implementation, track which followers are up-to-date
	}

	// --- Communication ---
	private sendHeartbeats(): void {
		const hb: Heartbeat = {
			type: 'heartbeat',
			term: this.currentTerm,
			leaderId: this.id
		};
		for (const peer of this.peers) {
			this.send(peer, hb);
		}
	}

	private send(peer: Peer, msg: RaftMessage): void {
		const data = Buffer.from(JSON.stringify(msg));
		this.socket.send(data, peer.port, peer.host);
	}

	getStatus(): object {
		return {
			id: this.id,
			state: this.state,
			term: this.currentTerm,
			leader: this.leaderId,
			votedFor: this.votedFor
		};
	}
}

// --- Run a node ---
const nodeId = process.argv[2] || 'node-1';
const port = parseInt(process.argv[3] || '4001');

const allNodes: Peer[] = [
	{ id: 'node-1', host: '127.0.0.1', port: 4001 },
	{ id: 'node-2', host: '127.0.0.1', port: 4002 },
	{ id: 'node-3', host: '127.0.0.1', port: 4003 }
];

const peers = allNodes.filter((n) => n.id !== nodeId);
const node = new RaftNode(nodeId, port, peers);
node.start();

// Print status periodically
setInterval(() => {
	console.log(`[${nodeId}] Status:`, node.getStatus());
}, 5000);

process.on('SIGINT', () => process.exit(0));