Fixed node controller assignemnts update after pods restart (#368)

If a storage node restarts and its IP stop responding, the coordinator
health check will detect it in a ~2seconds period, though the component
that sends the cluster assignements updates will not detect the issue.

One of the reason is that it's using a grpc stream, so there is no
timeout on the send operation (because it would apply to the lifetime of
the stream itself, which instead needs to be long-lived).

The storage servers need to receive the new cluster assignments after
restart, otherwise they cannot pass it back to their clients.

To fix the issue, we close the assignments stream whenever the
health-check failure is triggered. This will ensure that coordinator
will retry after the restart.

coordinator/impl/node_controller.go

@@ -27,7 +27,6 @@ import (
 	"oxia/coordinator/model"
 	"oxia/proto"
 	"sync"
-	"sync/atomic"
 	"time"
 )
 
@@ -60,11 +59,13 @@ type nodeController struct {
 	nodeAvailabilityListener NodeAvailabilityListener
 	rpc                      RpcProvider
 	log                      zerolog.Logger
-	closed                   atomic.Bool
 	ctx                      context.Context
 	cancel                   context.CancelFunc
 	initialRetryBackoff      time.Duration
 
+	sendAssignmentsCtx    context.Context
+	sendAssignmentsCancel context.CancelFunc
+
 	nodeIsRunningGauge metrics.Gauge
 	failedHealthChecks metrics.Counter
 }
@@ -142,6 +143,10 @@ func (n *nodeController) healthCheckWithRetries() {
 			n.failedHealthChecks.Inc()
 			n.nodeAvailabilityListener.NodeBecameUnavailable(n.addr)
 		}
+
+		// To avoid the send assignments stream to miss the notification about the current
+		// node went down, we interrupt the current stream when the ping on the node fails
+		n.sendAssignmentsCancel()
 	})
 }
 
@@ -231,41 +236,56 @@ func (n *nodeController) sendAssignmentsUpdatesWithRetries() {
 }
 
 func (n *nodeController) sendAssignmentsUpdates(backoff backoff.BackOff) error {
-	stream, err := n.rpc.PushShardAssignments(n.ctx, n.addr)
+	n.Lock()
+	n.sendAssignmentsCtx, n.sendAssignmentsCancel = context.WithCancel(n.ctx)
+	n.Unlock()
+	defer n.sendAssignmentsCancel()
+
+	stream, err := n.rpc.PushShardAssignments(n.sendAssignmentsCtx, n.addr)
 	if err != nil {
 		return err
 	}
 
 	var assignments *proto.ShardAssignments
-	for !n.closed.Load() {
 
-		assignments, err = n.shardAssignmentsProvider.WaitForNextUpdate(stream.Context(), assignments)
-		if err != nil {
-			return err
-		}
+	for {
+		select {
+		case <-n.ctx.Done():
+			return nil
+
+		default:
+			n.log.Debug().
+				Interface("current-assignments", assignments).
+				Msg("Waiting for next assignments update")
+			assignments, err = n.shardAssignmentsProvider.WaitForNextUpdate(stream.Context(), assignments)
+			if err != nil {
+				return err
+			}
 
-		if assignments == nil {
-			continue
-		}
+			if assignments == nil {
+				n.log.Debug().
+					Msg("Assignments are nil")
+				continue
+			}
 
-		n.log.Debug().
-			Interface("assignments", assignments).
-			Msg("Sending assignments")
+			n.log.Debug().
+				Interface("assignments", assignments).
+				Msg("Sending assignments")
 
-		if err := stream.Send(assignments); err != nil {
-			n.log.Debug().Err(err).
-				Msg("Failed to send assignments")
-			return err
-		}
+			if err := stream.Send(assignments); err != nil {
+				n.log.Debug().Err(err).
+					Msg("Failed to send assignments")
+				return err
+			}
 
-		backoff.Reset()
+			n.log.Debug().
+				Msg("Send assignments completed successfully")
+			backoff.Reset()
+		}
 	}
-
-	return nil
 }
 
 func (n *nodeController) Close() error {
-	n.closed.Store(true)
 	n.nodeIsRunningGauge.Unregister()
 	n.cancel()
 	return nil