Add file chunk size config parameter. Add unit tests.

Author: dhurley

internal/config/config.go

@@ -287,13 +287,19 @@ func registerClientFlags(fs *flag.FlagSet) {
 	fs.Int(
 		ClientGRPCMaxMessageReceiveSizeKey,
 		DefMaxMessageRecieveSize,
-		"Updates the client grpc setting MaxRecvMsgSize with the specific value in MB.",
+		"Updates the client grpc setting MaxRecvMsgSize with the specific value in bytes.",
 	)
 
 	fs.Int(
 		ClientGRPCMaxMessageSendSizeKey,
 		DefMaxMessageSendSize,
-		"Updates the client grpc setting MaxSendMsgSize with the specific value in MB.",
+		"Updates the client grpc setting MaxSendMsgSize with the specific value in bytes.",
+	)
+
+	fs.Uint32(
+		ClientGRPCFileChunkSizeKey,
+		DefFileChunkSize,
+		"File chunk size in bytes.",
 	)
 }
 
@@ -611,6 +617,7 @@ func resolveClient() *Client {
 			MaxMessageSize:        viperInstance.GetInt(ClientGRPCMaxMessageSizeKey),
 			MaxMessageReceiveSize: viperInstance.GetInt(ClientGRPCMaxMessageReceiveSizeKey),
 			MaxMessageSendSize:    viperInstance.GetInt(ClientGRPCMaxMessageSendSizeKey),
+			FileChunkSize:         viperInstance.GetUint32(ClientGRPCFileChunkSizeKey),
 		},
 		Backoff: &BackOff{
 			InitialInterval:     viperInstance.GetDuration(ClientBackoffInitialIntervalKey),

internal/config/config_test.go

@@ -100,6 +100,7 @@ func checkDefaultsClientValues(t *testing.T, viperInstance *viper.Viper) {
 	assert.Equal(t, DefMaxMessageSize, viperInstance.GetInt(ClientGRPCMaxMessageSizeKey))
 	assert.Equal(t, DefMaxMessageRecieveSize, viperInstance.GetInt(ClientGRPCMaxMessageReceiveSizeKey))
 	assert.Equal(t, DefMaxMessageSendSize, viperInstance.GetInt(ClientGRPCMaxMessageSendSizeKey))
+	assert.Equal(t, DefFileChunkSize, viperInstance.GetUint32(ClientGRPCFileChunkSizeKey))
 	assert.Equal(t, make(map[string]string), viperInstance.GetStringMapString(LabelsRootKey))
 }
 
@@ -746,6 +747,7 @@ func createConfig() *Config {
 				MaxMessageSize:        1048575,
 				MaxMessageReceiveSize: 1048575,
 				MaxMessageSendSize:    1048575,
+				FileChunkSize:         48575,
 			},
 			Backoff: &BackOff{
 				InitialInterval:     200 * time.Millisecond,

internal/config/defaults.go

@@ -27,9 +27,10 @@ const (
 	DefCommandTLServerNameKey  = ""
 
 	// Client GRPC Settings
-	DefMaxMessageSize        = 0       // 0 = unset
-	DefMaxMessageRecieveSize = 4194304 // default 4 MB
-	DefMaxMessageSendSize    = 4194304 // default 4 MB
+	DefMaxMessageSize               = 0       // 0 = unset
+	DefMaxMessageRecieveSize        = 4194304 // default 4 MB
+	DefMaxMessageSendSize           = 4194304 // default 4 MB
+	DefFileChunkSize         uint32 = 2097152 // 2MB
 
 	// Client HTTP Settings
 	DefHTTPTimeout = 10 * time.Second

internal/config/flags.go

@@ -36,6 +36,7 @@ var (
 	ClientGRPCMaxMessageSendSizeKey    = pre(ClientRootKey) + "grpc_max_message_send_size"
 	ClientGRPCMaxMessageReceiveSizeKey = pre(ClientRootKey) + "grpc_max_message_receive_size"
 	ClientGRPCMaxMessageSizeKey        = pre(ClientRootKey) + "grpc_max_message_size"
+	ClientGRPCFileChunkSizeKey         = pre(ClientRootKey) + "grpc_file_chunk_size"
 
 	ClientBackoffInitialIntervalKey     = pre(ClientRootKey) + "backoff_initial_interval"
 	ClientBackoffMaxIntervalKey         = pre(ClientRootKey) + "backoff_max_interval"

internal/config/testdata/nginx-agent.conf

@@ -35,6 +35,7 @@ client:
        max_message_size: 1048575
        max_message_receive_size: 1048575
        max_message_send_size: 1048575
+       file_chunk_size: 48575
     backoff:
         initial_interval: 200ms
         max_interval: 10s

internal/config/types.go

@@ -80,12 +80,11 @@ type (
 	}
 
 	GRPC struct {
-		KeepAlive *KeepAlive `yaml:"-" mapstructure:"target"`
-		// if MaxMessageSize is size set then we use that value,
-		// otherwise MaxMessageRecieveSize and MaxMessageSendSize for individual settings
-		MaxMessageSize        int `yaml:"-" mapstructure:"max_message_size"`
-		MaxMessageReceiveSize int `yaml:"-" mapstructure:"max_message_receive_size"`
-		MaxMessageSendSize    int `yaml:"-" mapstructure:"max_message_send_size"`
+		KeepAlive             *KeepAlive `yaml:"-"               mapstructure:"target"`
+		MaxMessageSize        int        `yaml:"-"               mapstructure:"max_message_size"`
+		MaxMessageReceiveSize int        `yaml:"-"               mapstructure:"max_message_receive_size"`
+		MaxMessageSendSize    int        `yaml:"-"               mapstructure:"max_message_send_size"`
+		FileChunkSize         uint32     `yaml:"file_chunk_size" mapstructure:"file_chunk_size"`
 	}
 
 	KeepAlive struct {

internal/file/file_manager_service.go

@@ -241,9 +241,7 @@ func (fms *FileManagerService) UpdateFile(
 		return fms.sendUpdateFileRequest(ctx, fileToUpdate)
 	}
 
-	// TODO: Add config parameter for chunk size
-	var chunkSize uint32 = 2097152 // 2MB
-	return fms.sendUpdateFileStream(ctx, fileToUpdate, chunkSize)
+	return fms.sendUpdateFileStream(ctx, fileToUpdate, fms.agentConfig.Client.Grpc.FileChunkSize)
 }
 
 func (fms *FileManagerService) sendUpdateFileRequest(
@@ -314,6 +312,10 @@ func (fms *FileManagerService) sendUpdateFileStream(
 	fileToUpdate *mpi.File,
 	chunkSize uint32,
 ) error {
+	if chunkSize == 0 {
+		return fmt.Errorf("file chunk size must be greater than zero")
+	}
+
 	updateFileStreamClient, err := fms.fileServiceClient.UpdateFileStream(ctx)
 	if err != nil {
 		return err
@@ -441,10 +443,12 @@ func (fms *FileManagerService) readChunk(
 			return mpi.FileDataChunk_Content{}, fmt.Errorf("failed to read chunk: %w", err)
 		}
 
+		slog.DebugContext(ctx, "No more data to read from file")
+
 		return mpi.FileDataChunk_Content{}, nil
 	}
 
-	slog.DebugContext(ctx, "Read chunk", "chunkID", chunkID, "chunk_size", len(buf))
+	slog.DebugContext(ctx, "Read file chunk", "chunk_id", chunkID, "chunk_size", len(buf))
 
 	chunk := mpi.FileDataChunk_Content{
 		Content: &mpi.FileDataChunkContent{
@@ -628,38 +632,51 @@ func (fms *FileManagerService) executeFileActions(ctx context.Context) error {
 }
 
 func (fms *FileManagerService) fileUpdate(ctx context.Context, file *mpi.File) error {
+	slog.DebugContext(ctx, "Updating file", "file", file.GetFileMeta().GetName())
 	if file.GetFileMeta().GetSize() <= int64(fms.agentConfig.Client.Grpc.MaxMessageReceiveSize) {
-		backOffCtx, backoffCancel := context.WithTimeout(ctx, fms.agentConfig.Client.Backoff.MaxElapsedTime)
-		defer backoffCancel()
-
-		getFile := func() (*mpi.GetFileResponse, error) {
-			return fms.fileServiceClient.GetFile(ctx, &mpi.GetFileRequest{
-				MessageMeta: &mpi.MessageMeta{
-					MessageId:     id.GenerateMessageID(),
-					CorrelationId: logger.GetCorrelationID(ctx),
-					Timestamp:     timestamppb.Now(),
-				},
-				FileMeta: file.GetFileMeta(),
-			})
-		}
+		return fms.getFile(ctx, file)
+	}
 
-		getFileResp, getFileErr := backoff.RetryWithData(
-			getFile,
-			backoffHelpers.Context(backOffCtx, fms.agentConfig.Client.Backoff),
-		)
+	return fms.getChunkedFile(ctx, file)
+}
 
-		if getFileErr != nil {
-			return fmt.Errorf("error getting file data for %s: %w", file.GetFileMeta(), getFileErr)
-		}
+func (fms *FileManagerService) getFile(ctx context.Context, file *mpi.File) error {
+	slog.DebugContext(ctx, "Getting file", "file", file.GetFileMeta().GetName())
 
-		if writeErr := fms.fileOperator.Write(ctx, getFileResp.GetContents().GetContents(),
-			file.GetFileMeta()); writeErr != nil {
-			return writeErr
-		}
+	backOffCtx, backoffCancel := context.WithTimeout(ctx, fms.agentConfig.Client.Backoff.MaxElapsedTime)
+	defer backoffCancel()
+
+	getFile := func() (*mpi.GetFileResponse, error) {
+		return fms.fileServiceClient.GetFile(ctx, &mpi.GetFileRequest{
+			MessageMeta: &mpi.MessageMeta{
+				MessageId:     id.GenerateMessageID(),
+				CorrelationId: logger.GetCorrelationID(ctx),
+				Timestamp:     timestamppb.Now(),
+			},
+			FileMeta: file.GetFileMeta(),
+		})
+	}
+
+	getFileResp, getFileErr := backoff.RetryWithData(
+		getFile,
+		backoffHelpers.Context(backOffCtx, fms.agentConfig.Client.Backoff),
+	)
+
+	if getFileErr != nil {
+		return fmt.Errorf("error getting file data for %s: %w", file.GetFileMeta(), getFileErr)
+	}
 
-		return fms.validateFileHash(file.GetFileMeta().GetName())
+	if writeErr := fms.fileOperator.Write(ctx, getFileResp.GetContents().GetContents(),
+		file.GetFileMeta()); writeErr != nil {
+		return writeErr
 	}
 
+	return fms.validateFileHash(file.GetFileMeta().GetName())
+}
+
+func (fms *FileManagerService) getChunkedFile(ctx context.Context, file *mpi.File) error {
+	slog.DebugContext(ctx, "Getting chunked file", "file", file.GetFileMeta().GetName())
+
 	stream, err := fms.fileServiceClient.GetFileStream(ctx, &mpi.GetFileRequest{
 		MessageMeta: &mpi.MessageMeta{
 			MessageId:     id.GenerateMessageID(),
@@ -673,28 +690,52 @@ func (fms *FileManagerService) fileUpdate(ctx context.Context, file *mpi.File) e
 	}
 
 	// Get header chunk first
-	chunk := &mpi.FileDataChunk{}
-	recvError := stream.RecvMsg(chunk)
-	if recvError != nil {
-		return recvError
+	headerChunk, recvHeaderChunkError := stream.Recv()
+	if recvHeaderChunkError != nil {
+		return recvHeaderChunkError
+	}
+
+	slog.DebugContext(ctx, "File header chunk received", "header_chunk", headerChunk)
+
+	header := headerChunk.GetHeader()
+
+	writeChunkedFileError := fms.writeChunkedFile(ctx, file, header, stream)
+	if writeChunkedFileError != nil {
+		return writeChunkedFileError
 	}
 
-	header := chunk.GetHeader()
+	return nil
+}
 
-	filePermissions := files.FileMode(header.GetFileMeta().GetPermissions())
-	createFileDirectoriesError := fms.fileOperator.CreateFileDirectories(ctx, header.GetFileMeta(), filePermissions)
+func (fms *FileManagerService) writeChunkedFile(
+	ctx context.Context,
+	file *mpi.File,
+	header *mpi.FileDataChunkHeader,
+	stream grpc2.ServerStreamingClient[mpi.FileDataChunk],
+) error {
+	filePermissions := files.FileMode(file.GetFileMeta().GetPermissions())
+	createFileDirectoriesError := fms.fileOperator.CreateFileDirectories(ctx, file.GetFileMeta(), filePermissions)
 	if createFileDirectoriesError != nil {
 		return createFileDirectoriesError
 	}
 
-	fileToWrite, openError := os.OpenFile(header.GetFileMeta().GetName(), os.O_WRONLY, filePermissions)
-	if openError != nil {
-		return openError
+	fileToWrite, createError := os.Create(file.GetFileMeta().GetName())
+	defer func() {
+		closeError := fileToWrite.Close()
+		if closeError != nil {
+			slog.WarnContext(
+				ctx, "Failed to close file",
+				"file", file.GetFileMeta().GetName(),
+				"error", closeError,
+			)
+		}
+	}()
+	if createError != nil {
+		return createError
 	}
 
 	for i := uint32(0); i < header.GetChunks(); i++ {
-		chunk := &mpi.FileDataChunk{}
-		recvError := stream.RecvMsg(chunk)
+		chunk, recvError := stream.Recv()
 		if recvError != nil {
 			return recvError
 		}