Add Shutdown Notice from Compute Nodes (#4769)

When compute nodes shut down, they should notify the orchestrator with
their final sequence numbers. While this info is already shared in
heartbeats, a dedicated shutdown message confirms intentional shutdown
vs connection failure and guarantees the orchestrator receives latest
sequence numbers. This lets the orchestrator clean up node state
immediately rather than waiting for missing heartbeats.

Changes:
- Added `sendShutdownNotification` method to ControlPlane for notifying
orchestrator before shutdown
- Modified Stop() to send notification if node is connected and context
isn't cancelled
- Added test cases covering successful notification, skipped
notifications (when disconnected/cancelled), and error handling


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced constants for shutdown notice request and response message
types.
	- Added structures for handling shutdown notice requests and responses.
- Implemented a method for processing shutdown notifications in the node
manager.
- Enhanced control plane to send shutdown notifications to the
orchestrator.
	- Updated the compute manager to handle shutdown requests from nodes.

- **Bug Fixes**
- Improved error handling and state management during shutdown
operations.

- **Tests**
- Added comprehensive tests for shutdown functionality in node manager
and control plane.
- Enhanced mock responder to simulate shutdown notifications for testing
purposes.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: wdbaruni

pkg/models/messages/constants.go

@@ -13,8 +13,10 @@ const (
 	HandshakeRequestMessageType      = "transport.HandshakeRequest"
 	HeartbeatRequestMessageType      = "transport.HeartbeatRequest"
 	NodeInfoUpdateRequestMessageType = "transport.UpdateNodeInfoRequest"
+	ShutdownNoticeRequestMessageType = "transport.ShutdownNoticeRequest"
 
 	HandshakeResponseType      = "transport.HandshakeResponse"
 	HeartbeatResponseType      = "transport.HeartbeatResponse"
 	NodeInfoUpdateResponseType = "transport.UpdateNodeInfoResponse"
+	ShutdownNoticeResponseType = "transport.ShutdownNoticeResponse"
 )

pkg/models/messages/node.go pkg/models/messages/control_plane.go

@@ -40,3 +40,15 @@ type UpdateNodeInfoResponse struct {
 	Accepted bool   `json:"accepted"`
 	Reason   string `json:"reason,omitempty"`
 }
+
+// ShutdownNoticeRequest tells the orchestrator that this node is shutting down
+type ShutdownNoticeRequest struct {
+	NodeID                 string `json:"NodeID"`
+	Reason                 string `json:"Reason,omitempty"`
+	LastOrchestratorSeqNum uint64 `json:"LastOrchestratorSeqNum"` // Last seq received from orchestrator
+}
+
+// ShutdownNoticeResponse sends any final instructions back to the shutting-down node
+type ShutdownNoticeResponse struct {
+	LastComputeSeqNum uint64 `json:"LastComputeSeqNum"` // Last seq received from compute node
+}

pkg/orchestrator/nodes/manager.go

@@ -582,6 +582,62 @@ func (n *nodesManager) Heartbeat(
 	return messages.HeartbeatResponse{}, NewErrConcurrentModification()
 }
 
+// ShutdownNotice processes a shutdown notification from a node and updates its state.
+// It updates:
+//   - Final sequence numbers
+//   - Connection state to disconnected
+//   - Preserves the sequence numbers in persistent storage
+//
+// Returns ShutdownNoticeResponse with the last sequence number processed from that node.
+func (n *nodesManager) ShutdownNotice(
+	ctx context.Context, request ExtendedShutdownNoticeRequest) (messages.ShutdownNoticeResponse, error) {
+	// Get existing live state
+	existingEntry, ok := n.liveState.Load(request.NodeID)
+	if !ok {
+		return messages.ShutdownNoticeResponse{}, NewErrHandshakeRequired(request.NodeID)
+	}
+
+	existing := existingEntry.(*trackedLiveState)
+	if existing.connectionState.Status != models.NodeStates.CONNECTED {
+		return messages.ShutdownNoticeResponse{}, NewErrHandshakeRequired(request.NodeID)
+	}
+
+	// Update connection state with final sequence numbers
+	updated := existing.connectionState
+	updated.Status = models.NodeStates.DISCONNECTED
+	updated.DisconnectedSince = n.clock.Now().UTC()
+	n.updateSequenceNumbers(&updated, request.LastOrchestratorSeqNum, request.LastComputeSeqNum)
+	updated.LastError = "graceful shutdown"
+
+	// Attempt atomic update
+	if !n.liveState.CompareAndSwap(request.NodeID, existingEntry, &trackedLiveState{
+		connectionState:   updated,
+		availableCapacity: models.Resources{}, // Clear capacity since node is shutting down
+		queueUsedCapacity: models.Resources{},
+	}) {
+		return messages.ShutdownNoticeResponse{}, NewErrConcurrentModification()
+	}
+
+	log.Info().
+		Str("node", request.NodeID).
+		Str("reason", request.Reason).
+		Uint64("lastOrchestratorSeq", updated.LastOrchestratorSeqNum).
+		Uint64("lastComputeSeq", updated.LastComputeSeqNum).
+		Msg("Node shutdown notice received")
+
+	// Notify about state change
+	n.notifyConnectionStateChange(NodeConnectionEvent{
+		NodeID:    request.NodeID,
+		Previous:  models.NodeStates.CONNECTED,
+		Current:   models.NodeStates.DISCONNECTED,
+		Timestamp: updated.DisconnectedSince,
+	})
+
+	return messages.ShutdownNoticeResponse{
+		LastComputeSeqNum: updated.LastComputeSeqNum,
+	}, nil
+}
+
 // ApproveNode approves a node for cluster participation.
 // The node must be in PENDING state. The operation updates
 // both persistent and live state.

pkg/orchestrator/nodes/manager_test.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -1072,6 +1073,184 @@ func (s *NodeManagerTestSuite) TestPersistenceWithContextCancellation() {
 	assert.Equal(s.T(), lastComputeSeqNum, state.ConnectionState.LastComputeSeqNum)
 }
 
+func (s *NodeManagerTestSuite) TestShutdownNotice() {
+	lastOrchestratorSeqNum := uint64(42)
+	lastComputeSeqNum := uint64(24)
+
+	// Test cases covering different shutdown scenarios
+	tests := []struct {
+		name           string
+		setupNode      bool // whether to setup node with handshake first
+		disconnect     bool // whether to disconnect node before shutdown
+		reason         string
+		expectedError  string
+		validateState  func(*testing.T, models.NodeState)
+		validateEvents func(*testing.T, []nodes.NodeConnectionEvent)
+	}{
+		{
+			name:      "successful shutdown",
+			setupNode: true,
+			reason:    "maintenance",
+			validateState: func(t *testing.T, state models.NodeState) {
+				assert.Equal(t, models.NodeStates.DISCONNECTED, state.ConnectionState.Status)
+				assert.Equal(t, "graceful shutdown", state.ConnectionState.LastError)
+				assert.False(t, state.ConnectionState.DisconnectedSince.IsZero())
+			},
+			validateEvents: func(t *testing.T, events []nodes.NodeConnectionEvent) {
+				require.Len(t, events, 2) // connect + disconnect
+				assert.Equal(t, models.NodeStates.CONNECTED, events[1].Previous)
+				assert.Equal(t, models.NodeStates.DISCONNECTED, events[1].Current)
+			},
+		},
+		{
+			name:          "shutdown without handshake",
+			setupNode:     false,
+			reason:        "testing",
+			expectedError: "handshake required",
+		},
+		{
+			name:          "shutdown already disconnected node",
+			setupNode:     true,
+			disconnect:    true,
+			reason:        "testing",
+			expectedError: "handshake required",
+			validateState: func(t *testing.T, state models.NodeState) {
+				assert.Equal(t, models.NodeStates.DISCONNECTED, state.ConnectionState.Status)
+			},
+		},
+		{
+			name:      "shutdown preserves sequence numbers",
+			setupNode: true,
+			reason:    "testing",
+			validateState: func(t *testing.T, state models.NodeState) {
+				assert.Equal(t, lastOrchestratorSeqNum, state.ConnectionState.LastOrchestratorSeqNum)
+				assert.Equal(t, lastComputeSeqNum, state.ConnectionState.LastComputeSeqNum)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		s.Run(tt.name, func() {
+			// Track connection events
+			var events []nodes.NodeConnectionEvent
+			eventsMu := sync.Mutex{}
+			s.manager.OnConnectionStateChange(func(event nodes.NodeConnectionEvent) {
+				eventsMu.Lock()
+				events = append(events, event)
+				eventsMu.Unlock()
+			})
+
+			nodeInfo := s.createNodeInfo("shutdown-test")
+
+			// Setup node if required
+			if tt.setupNode {
+				_, err := s.manager.Handshake(s.ctx, messages.HandshakeRequest{NodeInfo: nodeInfo})
+				s.Require().NoError(err)
+
+				// Update sequence numbers
+				_, err = s.manager.Heartbeat(s.ctx, nodes.ExtendedHeartbeatRequest{
+					HeartbeatRequest: messages.HeartbeatRequest{
+						NodeID:                 nodeInfo.ID(),
+						LastOrchestratorSeqNum: 42,
+					},
+					LastComputeSeqNum: 24,
+				})
+				s.Require().NoError(err)
+
+				if tt.disconnect {
+					s.clock.Add(s.disconnected + time.Second)
+					s.Eventually(func() bool {
+						state, err := s.manager.Get(s.ctx, nodeInfo.ID())
+						s.Require().NoError(err)
+						return state.ConnectionState.Status == models.NodeStates.DISCONNECTED
+					}, 500*time.Millisecond, 20*time.Millisecond)
+				}
+			}
+
+			// Send shutdown notice
+			req := nodes.ExtendedShutdownNoticeRequest{
+				ShutdownNoticeRequest: messages.ShutdownNoticeRequest{
+					NodeID:                 nodeInfo.ID(),
+					Reason:                 tt.reason,
+					LastOrchestratorSeqNum: lastOrchestratorSeqNum,
+				},
+				LastComputeSeqNum: lastComputeSeqNum,
+			}
+
+			_, err := s.manager.ShutdownNotice(s.ctx, req)
+			if tt.expectedError != "" {
+				s.Assert().Error(err)
+				s.Assert().Contains(err.Error(), tt.expectedError)
+				return
+			}
+			s.Assert().NoError(err)
+
+			// Validate final state
+			state, err := s.manager.Get(s.ctx, nodeInfo.ID())
+			s.Require().NoError(err)
+
+			if tt.validateState != nil {
+				tt.validateState(s.T(), state)
+			}
+
+			if tt.validateEvents != nil {
+				eventsMu.Lock()
+				tt.validateEvents(s.T(), events)
+				eventsMu.Unlock()
+			}
+		})
+	}
+}
+
+func (s *NodeManagerTestSuite) TestConcurrentShutdown() {
+	nodeInfo := s.createNodeInfo("concurrent-shutdown")
+	lastOrchestratorSeqNum := uint64(42)
+	lastComputeSeqNum := uint64(24)
+
+	// Connect node
+	_, err := s.manager.Handshake(s.ctx, messages.HandshakeRequest{NodeInfo: nodeInfo})
+	s.Require().NoError(err)
+
+	// Track successful shutdowns
+	var wg sync.WaitGroup
+	successCount := int32(0)
+	const numConcurrent = 10
+
+	for i := 0; i < numConcurrent; i++ {
+		wg.Add(1)
+		go func(attempt int) {
+			defer wg.Done()
+
+			req := nodes.ExtendedShutdownNoticeRequest{
+				ShutdownNoticeRequest: messages.ShutdownNoticeRequest{
+					NodeID:                 nodeInfo.ID(),
+					Reason:                 fmt.Sprintf("concurrent shutdown %d", attempt),
+					LastOrchestratorSeqNum: lastOrchestratorSeqNum,
+				},
+				LastComputeSeqNum: lastComputeSeqNum,
+			}
+
+			_, err := s.manager.ShutdownNotice(s.ctx, req)
+			if err == nil {
+				atomic.AddInt32(&successCount, 1)
+			}
+		}(i)
+	}
+
+	wg.Wait()
+
+	// Exactly one shutdown should succeed
+	s.Assert().Equal(int32(1), successCount)
+
+	// Verify final state
+	state, err := s.manager.Get(s.ctx, nodeInfo.ID())
+	s.Require().NoError(err)
+	s.Assert().Equal(models.NodeStates.DISCONNECTED, state.ConnectionState.Status)
+	s.Assert().Equal("graceful shutdown", state.ConnectionState.LastError)
+	s.Assert().Equal(lastOrchestratorSeqNum, state.ConnectionState.LastOrchestratorSeqNum)
+	s.Assert().Equal(lastComputeSeqNum, state.ConnectionState.LastComputeSeqNum)
+}
+
 type mockNodeInfoProvider struct {
 	info models.NodeInfo
 }

pkg/orchestrator/nodes/types.go

@@ -64,6 +64,10 @@ type Manager interface {
 	// The node must be registered and not rejected.
 	UpdateNodeInfo(ctx context.Context, request messages.UpdateNodeInfoRequest) (messages.UpdateNodeInfoResponse, error)
 
+	// ShutdownNotice handles a node's graceful shutdown notification.
+	// It updates sequence numbers and marks the node as cleanly disconnected.
+	ShutdownNotice(ctx context.Context, request ExtendedShutdownNoticeRequest) (messages.ShutdownNoticeResponse, error)
+
 	// Heartbeat processes a node's heartbeat message and updates its state.
 	// It returns the last known sequence numbers for synchronization.
 	Heartbeat(ctx context.Context, request ExtendedHeartbeatRequest) (messages.HeartbeatResponse, error)
@@ -136,3 +140,11 @@ type ExtendedHeartbeatRequest struct {
 	// LastComputeSeqNum is the last processed compute message sequence
 	LastComputeSeqNum uint64
 }
+
+// ExtendedShutdownNoticeRequest represents a shutdown message with additional metadata.
+type ExtendedShutdownNoticeRequest struct {
+	messages.ShutdownNoticeRequest
+
+	// LastComputeSeqNum is the last processed compute message sequence
+	LastComputeSeqNum uint64
+}

pkg/transport/nclprotocol/compute/controlplane.go

@@ -186,6 +186,25 @@ func (cp *ControlPlane) updateNodeInfo(ctx context.Context) error {
 	return nil
 }
 
+// sendShutdownNotification informs the orchestrator that this node is gracefully shutting down.
+// It includes the last message sequence numbers to help prevent message duplication on reconnect.
+// The notification is best-effort - we don't wait for or retry the response since we're shutting down.
+func (cp *ControlPlane) sendShutdownNotification(ctx context.Context) error {
+	ctx, cancel := context.WithTimeout(ctx, cp.cfg.RequestTimeout)
+	defer cancel()
+
+	msg := envelope.NewMessage(messages.ShutdownNoticeRequest{
+		NodeID:                 cp.cfg.NodeID,
+		LastOrchestratorSeqNum: cp.incomingSeqTracker.GetLastSeqNum(),
+	}).WithMetadataValue(envelope.KeyMessageType, messages.ShutdownNoticeRequestMessageType)
+
+	_, err := cp.requester.Request(ctx, ncl.NewPublishRequest(msg))
+	if err != nil {
+		return fmt.Errorf("shutdown notification failed: %w", err)
+	}
+	return nil
+}
+
 // checkpointProgress saves the latest processed message sequence number if it has
 // changed since the last checkpoint. This allows for resuming message processing
 // from the last known point after node restarts.
@@ -211,10 +230,18 @@ func (cp *ControlPlane) Stop(ctx context.Context) error {
 		return nil
 	}
 
+	// Prevent new operations
 	cp.running = false
 	close(cp.stopCh)
 	cp.mu.Unlock()
 
+	// Send shutdown notification before closing
+	if cp.healthTracker.GetHealth().CurrentState == nclprotocol.Connected && ctx.Err() == nil {
+		if err := cp.sendShutdownNotification(ctx); err != nil {
+			log.Error().Err(err).Msg("Failed to send shutdown notification")
+		}
+	}
+
 	// Wait for graceful shutdown
 	done := make(chan struct{})
 	go func() {