handle command and file service client updates after grpc reset

Author: sean-breen

internal/command/command_plugin.go

@@ -31,6 +31,8 @@ type (
 		UpdateDataPlaneStatus(ctx context.Context, resource *mpi.Resource) error
 		UpdateDataPlaneHealth(ctx context.Context, instanceHealths []*mpi.InstanceHealth) error
 		SendDataPlaneResponse(ctx context.Context, response *mpi.DataPlaneResponse) error
+		UpdateClient(client mpi.CommandServiceClient)
+		Resource() *mpi.Resource
 		Subscribe(ctx context.Context)
 		IsConnected() bool
 		CreateConnection(ctx context.Context, resource *mpi.Resource) (*mpi.CreateConnectionResponse, error)
@@ -177,8 +179,19 @@ func (cp *CommandPlugin) processDataPlaneResponse(ctx context.Context, msg *bus.
 func (cp *CommandPlugin) processConnectionReset(ctx context.Context, msg *bus.Message) {
 	slog.DebugContext(ctx, "Command plugin received connection reset")
 	if newConnection, ok := msg.Data.(*grpc.GrpcConnection); ok {
-		_ = cp.conn.Close(ctx)
+		if !cp.commandService.IsConnected() {
+			slog.DebugContext(ctx, "Command plugin: service is not connected")
+			return
+		}
+		err := cp.conn.Close(ctx)
+		if err != nil {
+			slog.ErrorContext(ctx, "Unable to close connection", "error", err)
+		}
 		cp.conn = newConnection
+		cp.commandService.UpdateClient(cp.conn.CommandServiceClient())
+		//go cp.monitorSubscribeChannel(ctx)
+		//cp.createConnection(ctx, cp.commandService.Resource())
+		slog.DebugContext(ctx, "Command plugin: client reset successfully")
 	}
 }
 

internal/command/command_service.go

@@ -41,6 +41,7 @@ type (
 		subscribeClient              mpi.CommandService_SubscribeClient
 		agentConfig                  *config.Config
 		isConnected                  *atomic.Bool
+		clientUpdated                *atomic.Bool
 		subscribeChannel             chan *mpi.ManagementPlaneRequest
 		configApplyRequestQueue      map[string][]*mpi.ManagementPlaneRequest // key is the instance ID
 		resource                     *mpi.Resource
@@ -58,10 +59,14 @@ func NewCommandService(
 	isConnected := &atomic.Bool{}
 	isConnected.Store(false)
 
+	clientUpdated := &atomic.Bool{}
+	clientUpdated.Store(false)
+
 	commandService := &CommandService{
 		commandServiceClient:    commandServiceClient,
 		agentConfig:             agentConfig,
 		isConnected:             isConnected,
+		clientUpdated:           clientUpdated,
 		subscribeChannel:        subscribeChannel,
 		configApplyRequestQueue: make(map[string][]*mpi.ManagementPlaneRequest),
 		resource:                &mpi.Resource{},
@@ -210,6 +215,10 @@ func (cs *CommandService) CreateConnection(
 		slog.InfoContext(ctx, "No Data Plane Instance found")
 	}
 
+	if cs.isConnected.Load() {
+		return nil, errors.New("command service already connected")
+	}
+
 	request := &mpi.CreateConnectionRequest{
 		MessageMeta: &mpi.MessageMeta{
 			MessageId:     id.GenerateMessageID(),
@@ -228,7 +237,6 @@ func (cs *CommandService) CreateConnection(
 	}
 
 	slog.DebugContext(ctx, "Sending create connection request", "request", request)
-
 	response, err := backoff.RetryWithData(
 		cs.connectCallback(ctx, request),
 		backoffHelpers.Context(ctx, commonSettings),
@@ -249,6 +257,17 @@ func (cs *CommandService) CreateConnection(
 	return response, nil
 }
 
+func (cs *CommandService) UpdateClient(client mpi.CommandServiceClient) {
+	cs.commandServiceClient = client
+	cs.clientUpdated.Store(true)
+}
+
+func (cs *CommandService) Resource() *mpi.Resource {
+	cs.resourceMutex.Lock()
+	defer cs.resourceMutex.Unlock()
+	return cs.resource
+}
+
 // Retry callback for sending a data plane response to the Management Plane.
 func (cs *CommandService) sendDataPlaneResponseCallback(
 	ctx context.Context,

internal/file/file_plugin.go

@@ -97,8 +97,12 @@ func (fp *FilePlugin) Subscriptions() []string {
 func (fp *FilePlugin) handleConnectionReset(ctx context.Context, msg *bus.Message) {
 	slog.DebugContext(ctx, "File plugin received connection reset message")
 	if newConnection, ok := msg.Data.(*grpc.GrpcConnection); ok {
-		_ = fp.conn.Close(ctx)
+		err := fp.conn.Close(ctx)
+		if err != nil {
+			slog.ErrorContext(ctx, "Error closing file plugin", "error", err)
+		}
 		fp.conn = newConnection
+		fp.fileManagerService = NewFileManagerService(fp.conn.FileServiceClient(), fp.config)
 	}
 }