Send create connection after updating credentials (nginx#1011)

Author: aphralG

internal/command/command_plugin.go

@@ -31,7 +31,7 @@ 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)
+		UpdateClient(ctx context.Context, client mpi.CommandServiceClient) error
 		Subscribe(ctx context.Context)
 		IsConnected() bool
 		CreateConnection(ctx context.Context, resource *mpi.Resource) (*mpi.CreateConnectionResponse, error)
@@ -68,7 +68,7 @@ func (cp *CommandPlugin) Init(ctx context.Context, messagePipe bus.MessagePipeIn
 }
 
 func (cp *CommandPlugin) Close(ctx context.Context) error {
-	slog.InfoContext(ctx, "Canceling subscribe context")
+	slog.InfoContext(ctx, "Closing command plugin")
 
 	cp.subscribeMutex.Lock()
 	if cp.subscribeCancel != nil {
@@ -178,12 +178,16 @@ 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.GrpcConnectionInterface); ok {
-		err := cp.conn.Close(ctx)
-		if err != nil {
-			slog.ErrorContext(ctx, "Command plugin: unable to close connection", "error", err)
+		connectionErr := cp.conn.Close(ctx)
+		if connectionErr != nil {
+			slog.ErrorContext(ctx, "Command plugin: unable to close connection", "error", connectionErr)
 		}
 		cp.conn = newConnection
-		cp.commandService.UpdateClient(cp.conn.CommandServiceClient())
+		err := cp.commandService.UpdateClient(ctx, cp.conn.CommandServiceClient())
+		if err != nil {
+			slog.ErrorContext(ctx, "Failed to reset connection", "error", err)
+			return
+		}
 		slog.DebugContext(ctx, "Command service client reset successfully")
 	}
 }

internal/command/command_service.go

@@ -254,10 +254,19 @@ func (cs *CommandService) CreateConnection(
 	return response, nil
 }
 
-func (cs *CommandService) UpdateClient(client mpi.CommandServiceClient) {
+func (cs *CommandService) UpdateClient(ctx context.Context, client mpi.CommandServiceClient) error {
 	cs.subscribeClientMutex.Lock()
-	defer cs.subscribeClientMutex.Unlock()
 	cs.commandServiceClient = client
+	cs.subscribeClientMutex.Unlock()
+
+	cs.isConnected.Store(false)
+	resp, err := cs.CreateConnection(ctx, cs.resource)
+	if err != nil {
+		return err
+	}
+	slog.InfoContext(ctx, "Successfully sent create connection request", "response", resp)
+
+	return nil
 }
 
 // Retry callback for sending a data plane response to the Management Plane.

internal/command/command_service_test.go

@@ -200,13 +200,15 @@ func TestCommandService_CreateConnection(t *testing.T) {
 
 func TestCommandService_UpdateClient(t *testing.T) {
 	commandServiceClient := &v1fakes.FakeCommandServiceClient{}
+	ctx := context.Background()
 
 	commandService := NewCommandService(
 		commandServiceClient,
 		types.AgentConfig(),
 		make(chan *mpi.ManagementPlaneRequest),
 	)
-	commandService.UpdateClient(commandServiceClient)
+	err := commandService.UpdateClient(ctx, commandServiceClient)
+	require.NoError(t, err)
 	assert.NotNil(t, commandService.commandServiceClient)
 }
 

internal/command/commandfakes/fake_command_service.go

@@ -51,6 +51,7 @@ func (fp *FilePlugin) Init(ctx context.Context, messagePipe bus.MessagePipeInter
 }
 
 func (fp *FilePlugin) Close(ctx context.Context) error {
+	slog.InfoContext(ctx, "Closing file plugin")
 	return fp.conn.Close(ctx)
 }
 

internal/file/file_plugin.go

@@ -115,12 +115,11 @@ func (gc *GrpcConnection) FileServiceClient() mpi.FileServiceClient {
 }
 
 func (gc *GrpcConnection) Close(ctx context.Context) error {
-	slog.InfoContext(ctx, "Closing grpc connection")
-
 	gc.mutex.Lock()
 	defer gc.mutex.Unlock()
 
 	if gc.conn != nil {
+		slog.InfoContext(ctx, "Closing grpc connection")
 		err := gc.conn.Close()
 		gc.conn = nil
 		if err != nil {
@@ -354,7 +353,8 @@ func ProtoValidatorStreamClientInterceptor() (grpc.StreamClientInterceptor, erro
 func ValidateGrpcError(err error) error {
 	if err != nil {
 		if statusError, ok := status.FromError(err); ok {
-			if statusError.Code() == codes.InvalidArgument || statusError.Code() == codes.Unimplemented {
+			if statusError.Code() == codes.InvalidArgument || statusError.Code() == codes.Unimplemented ||
+				statusError.Code() == codes.Canceled {
 				return backoff.Permanent(err)
 			}
 		}

internal/grpc/grpc.go

@@ -60,7 +60,7 @@ func (r *Resource) Init(ctx context.Context, messagePipe bus.MessagePipeInterfac
 }
 
 func (*Resource) Close(ctx context.Context) error {
-	slog.DebugContext(ctx, "Closing resource plugin")
+	slog.InfoContext(ctx, "Closing resource plugin")
 	return nil
 }
 

internal/resource/resource_plugin.go

@@ -112,7 +112,7 @@ func (w *Watcher) Init(ctx context.Context, messagePipe bus.MessagePipeInterface
 // nolint: unparam
 // error is always nil
 func (w *Watcher) Close(ctx context.Context) error {
-	slog.DebugContext(ctx, "Closing watcher plugin")
+	slog.InfoContext(ctx, "Closing watcher plugin")
 
 	w.cancel()