change upload behavior

Emyrk · Emyrk · commit f547103056c7 · 2025-06-06T09:19:23.000-05:00
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -1321,6 +1321,39 @@ func (s *server) prepareForNotifyWorkspaceManualBuildFailed(ctx context.Context,
 	return templateAdmins, template, templateVersion, workspaceOwner, nil
 }
 
+func (s *server) CompleteJobWithFiles(stream proto.DRPCProvisionerDaemon_CompleteJobWithFilesStream) error {
+	var file *sdkproto.DataBuilder
+
+	// stream expects files first
+	for {
+		msg, err := stream.Recv()
+		if err != nil {
+			return xerrors.Errorf("receive complete job with files: %w", err)
+		}
+
+		switch typed := msg.Type.(type) {
+		case *proto.CompleteWithFilesRequest_Complete:
+		case *proto.CompleteWithFilesRequest_ChunkPiece:
+
+		case *proto.CompleteWithFilesRequest_DataUpload:
+			if file != nil {
+				return xerrors.New("unexpected file upload while waiting for file completion")
+			}
+
+			file, err = sdkproto.NewDataBuilder(&sdkproto.DataUpload{
+				UploadType: sdkproto.DataUploadType(typed.DataUpload.UploadType),
+				DataHash:   typed.DataUpload.DataHash,
+				FileSize:   typed.DataUpload.FileSize,
+				Chunks:     typed.DataUpload.Chunks,
+			})
+			if err != nil {
+				return xerrors.Errorf("unable to create file upload: %w", err)
+			}
+		}
+	}
+
+}
+
 // CompleteJob is triggered by a provision daemon to mark a provisioner job as completed.
 func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob) (*proto.Empty, error) {
 	ctx, span := s.startTrace(ctx, tracing.FuncName())
diff --git a/provisionerd/proto/provisionerd.proto b/provisionerd/proto/provisionerd.proto
@@ -181,11 +181,10 @@ message CommitQuotaResponse {
 
 message CancelAcquire {}
 
-message CompleteWithFilesRequest {
+message UploadFileRequest {
     oneof type {
-        CompletedJob complete = 1;
-        DataUpload data_upload = 2;
-        ChunkPiece chunk_piece = 3;
+        DataUpload data_upload = 1;
+        ChunkPiece chunk_piece = 2;
     }
 }
 
@@ -216,5 +215,8 @@ service ProvisionerDaemon {
 
     // CompleteJob indicates a job has been completed.
     rpc CompleteJob(CompletedJob) returns (Empty);
-    rpc CompleteJobWithFiles(stream CompleteWithFilesRequest) returns (Empty);
+
+    // UploadFile streams files to be inserted into the database.
+    // The file upload_type should be used to determine how to handle the file.
+    rpc UploadFile(stream UploadFileRequest) returns (Empty);
 }
diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go
@@ -18,8 +18,10 @@ import (
 	semconv "go.opentelemetry.io/otel/semconv/v1.14.0"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/xerrors"
+	protobuf "google.golang.org/protobuf/proto"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
 	"github.com/coder/retry"
 
 	"github.com/coder/coder/v2/coderd/tracing"
@@ -515,7 +517,59 @@ func (p *Server) FailJob(ctx context.Context, in *proto.FailedJob) error {
 	return err
 }
 
+func (p *Server) CompleteJobWithModuleFiles(ctx context.Context, in *proto.CompletedJob, moduleFiles []byte) error {
+	return nil
+	// Send the files separately if the message size is too large.
+	//_, err := clientDoWithRetries(ctx, p.client, func(ctx context.Context, client proto.DRPCProvisionerDaemonClient) (*proto.Empty, error) {
+	//	stream, err := client.CompleteJobWithFiles(ctx)
+	//	if err != nil {
+	//		return nil, xerrors.Errorf("failed to start CompleteJobWithFiles stream: %w", err)
+	//	}
+	//
+	//	dataUp, chunks := sdkproto.BytesToDataUpload(moduleFiles)
+	//
+	//	err = stream.Send(&proto.CompleteWithFilesRequest{Type: &proto.CompleteWithFilesRequest_DataUpload{DataUpload: &proto.DataUpload{
+	//		UploadType: proto.DataUploadType(dataUp.UploadType),
+	//		DataHash:   dataUp.DataHash,
+	//		FileSize:   dataUp.FileSize,
+	//		Chunks:     dataUp.Chunks,
+	//	}}})
+	//	if err != nil {
+	//		return nil, xerrors.Errorf("send data upload: %w", err)
+	//	}
+	//
+	//	for i, chunk := range chunks {
+	//		err = stream.Send(&proto.CompleteWithFilesRequest{Type: &proto.CompleteWithFilesRequest_ChunkPiece{ChunkPiece: &proto.ChunkPiece{
+	//			Data:         chunk.Data,
+	//			FullDataHash: chunk.FullDataHash,
+	//			PieceIndex:   chunk.PieceIndex,
+	//		}}})
+	//		if err != nil {
+	//			return nil, xerrors.Errorf("send chunk piece %d: %w", i, err)
+	//		}
+	//	}
+	//
+	//	err = stream.Send(&proto.CompleteWithFilesRequest{Type: &proto.CompleteWithFilesRequest_Complete{Complete: in}})
+	//	if err != nil {
+	//		return nil, xerrors.Errorf("send complete job: %w", err)
+	//	}
+	//
+	//	return &proto.Empty{}, nil
+	//})
+	//return err
+}
+
 func (p *Server) CompleteJob(ctx context.Context, in *proto.CompletedJob) error {
+	if ti, ok := in.Type.(*proto.CompletedJob_TemplateImport_); ok {
+		messageSize := protobuf.Size(in)
+		if messageSize > drpcsdk.MaxMessageSize &&
+			messageSize-len(ti.TemplateImport.ModuleFiles) < drpcsdk.MaxMessageSize {
+			moduleFiles := ti.TemplateImport.ModuleFiles
+			ti.TemplateImport.ModuleFiles = nil // Clear the files in the final message
+			return p.CompleteJobWithModuleFiles(ctx, in, moduleFiles)
+		}
+	}
+
 	_, err := clientDoWithRetries(ctx, p.client, func(ctx context.Context, client proto.DRPCProvisionerDaemonClient) (*proto.Empty, error) {
 		return client.CompleteJob(ctx, in)
 	})
