Implement Raft leader heartbeats (eea78512-6817-44f0-849c-52f5003b0111)

2 files changed, 141 insertions, 0 deletions

diff --git a/src/main/java/protocols/implementations/VSRaftProtocol.java b/src/main/java/protocols/implementations/VSRaftProtocol.java
index 8b919f6..56e758d 100644
--- a/src/main/java/protocols/implementations/VSRaftProtocol.java
+++ b/src/main/java/protocols/implementations/VSRaftProtocol.java
@@ -76,6 +76,7 @@ public class VSRaftProtocol extends VSAbstractProtocol {
      * @see protocols.VSAbstractProtocol#onServerStart()
      */
     public void onServerStart() {
+        becomeLeader();
     }
 
     /* (non-Javadoc)
@@ -94,6 +95,9 @@ public class VSRaftProtocol extends VSAbstractProtocol {
      * @see protocols.VSAbstractProtocol#onServerSchedule()
      */
     public void onServerSchedule() {
+        if (isLeader) {
+            sendHeartbeat();
+        }
     }
 
     /* (non-Javadoc)
@@ -136,4 +140,29 @@ public class VSRaftProtocol extends VSAbstractProtocol {
             ackPids.clear();
         }
     }
+
+    /**
+     * Transitions this process into the leader role and starts heartbeats.
+     */
+    private void becomeLeader() {
+        isLeader = true;
+        isCandidate = false;
+        leaderId = process.getProcessID();
+        lastHeartbeatTime = process.getTime();
+        sendHeartbeat();
+    }
+
+    /**
+     * Sends a heartbeat and schedules the next leader heartbeat interval.
+     */
+    private void sendHeartbeat() {
+        VSMessage heartbeat = new VSMessage();
+        heartbeat.setString("type", "heartbeat");
+        heartbeat.setInteger("term", currentTerm);
+        heartbeat.setInteger("leaderId", leaderId);
+        sendMessage(heartbeat);
+
+        lastHeartbeatTime = process.getTime();
+        scheduleAt(process.getTime() + getLong("heartbeatInterval"));
+    }
 }
diff --git a/src/test/java/protocols/implementations/VSRaftProtocolTest.java b/src/test/java/protocols/implementations/VSRaftProtocolTest.java
new file mode 100644
index 0000000..6410a4e
--- /dev/null
+++ b/src/test/java/protocols/implementations/VSRaftProtocolTest.java
@@ -0,0 +1,112 @@
+package protocols.implementations;
+
+import core.VSInternalProcess;
+import core.VSMessage;
+import core.VSTask;
+import core.VSTaskManager;
+import core.time.VSVectorTime;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import prefs.VSPrefs;
+import simulator.VSSimulatorVisualization;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+/**
+ * Unit tests for VSRaftProtocol heartbeat behavior.
+ */
+class VSRaftProtocolTest {
+
+    @Mock
+    private VSInternalProcess mockProcess;
+
+    @Mock
+    private VSSimulatorVisualization mockCanvas;
+
+    @Mock
+    private VSTaskManager mockTaskManager;
+
+    @Mock
+    private VSPrefs mockPrefs;
+
+    @Mock
+    private VSVectorTime mockVectorTime;
+
+    private VSRaftProtocol protocol;
+
+    @BeforeEach
+    void setUp() {
+        MockitoAnnotations.openMocks(this);
+        protocol = new VSRaftProtocol();
+        protocol.process = mockProcess;
+        protocol.prefs = mockPrefs;
+        protocol.isServer(true);
+
+        when(mockProcess.getSimulatorCanvas()).thenReturn(mockCanvas);
+        when(mockCanvas.getTaskManager()).thenReturn(mockTaskManager);
+        when(mockProcess.getPrefs()).thenReturn(mockPrefs);
+        when(mockProcess.getVectorTime()).thenReturn(mockVectorTime);
+        when(mockVectorTime.getCopy()).thenReturn(mockVectorTime);
+        when(mockPrefs.getString(anyString())).thenReturn("TestString");
+        when(mockProcess.getTime()).thenReturn(100L);
+        when(mockProcess.getProcessID()).thenReturn(7);
+    }
+
+    @Test
+    void testOnStartBecomesLeaderAndSendsHeartbeat() {
+        ArgumentCaptor<VSMessage> messageCaptor =
+            ArgumentCaptor.forClass(VSMessage.class);
+        ArgumentCaptor<VSTask> taskCaptor = ArgumentCaptor.forClass(VSTask.class);
+
+        protocol.onStart();
+
+        verify(mockProcess).sendMessage(messageCaptor.capture());
+        verify(mockTaskManager).addTask(taskCaptor.capture());
+
+        VSMessage heartbeat = messageCaptor.getValue();
+        assertEquals("heartbeat", heartbeat.getString("type"));
+        assertEquals(0, heartbeat.getInteger("term"));
+        assertEquals(7, heartbeat.getInteger("leaderId"));
+        assertEquals(1600L, taskCaptor.getValue().getTaskTime());
+    }
+
+    @Test
+    void testServerScheduleSendsHeartbeatWhenLeader() {
+        ArgumentCaptor<VSMessage> messageCaptor =
+            ArgumentCaptor.forClass(VSMessage.class);
+        ArgumentCaptor<VSTask> taskCaptor = ArgumentCaptor.forClass(VSTask.class);
+
+        protocol.onStart();
+        protocol.onServerScheduleStart();
+
+        verify(mockProcess, times(2)).sendMessage(messageCaptor.capture());
+        verify(mockTaskManager, times(2)).addTask(taskCaptor.capture());
+
+        assertEquals(2, messageCaptor.getAllValues().size());
+        assertEquals(2, taskCaptor.getAllValues().size());
+
+        VSMessage scheduledHeartbeat = messageCaptor.getAllValues().get(1);
+        assertEquals("heartbeat", scheduledHeartbeat.getString("type"));
+        assertEquals(0, scheduledHeartbeat.getInteger("term"));
+        assertEquals(7, scheduledHeartbeat.getInteger("leaderId"));
+        assertEquals(1600L, taskCaptor.getAllValues().get(1).getTaskTime());
+    }
+
+    @Test
+    void testServerScheduleDoesNothingWhenNotLeader() {
+        protocol.currentContextIsServer(true);
+
+        protocol.onServerSchedule();
+
+        verify(mockProcess, never()).sendMessage(org.mockito.ArgumentMatchers.any());
+        verify(mockTaskManager, never()).addTask(org.mockito.ArgumentMatchers.any());
+    }
+}