mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: use proto streams to increase maximum module files payload (#18268)
This PR implements protobuf streaming to handle large module files by: 1. **Streaming large payloads**: When module files exceed the 4MB limit, they're streamed in chunks using a new UploadFile RPC method 2. **Database storage**: Streamed files are stored in the database and referenced by hash for deduplication 3. **Backward compatibility**: Small module files continue using the existing direct payload method
This commit is contained in:
Steven Masley
@@ -773,7 +773,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
|
||||
case database.ProvisionerStorageMethodFile:
|
||||
file, err := s.Database.GetFileByID(ctx, job.FileID)
|
||||
if err != nil {
|
||||
return nil, failJob(fmt.Sprintf("get file by hash: %s", err))
|
||||
return nil, failJob(fmt.Sprintf("get file by id: %s", err))
|
||||
}
|
||||
protoJob.TemplateSourceArchive = file.Data
|
||||
default:
|
||||
@@ -1321,6 +1321,104 @@ func (s *server) prepareForNotifyWorkspaceManualBuildFailed(ctx context.Context,
|
||||
return templateAdmins, template, templateVersion, workspaceOwner, nil
|
||||
}
|
||||
|
||||
func (s *server) UploadFile(stream proto.DRPCProvisionerDaemon_UploadFileStream) error {
|
||||
var file *sdkproto.DataBuilder
|
||||
// Always terminate the stream with an empty response.
|
||||
defer stream.SendAndClose(&proto.Empty{})
|
||||
|
||||
UploadFileStream:
|
||||
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.UploadFileRequest_DataUpload:
|
||||
if file != nil {
|
||||
return xerrors.New("unexpected file upload while waiting for file completion")
|
||||
}
|
||||
|
||||
file, err = sdkproto.NewDataBuilder(&sdkproto.DataUpload{
|
||||
UploadType: 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)
|
||||
}
|
||||
|
||||
if file.IsDone() {
|
||||
// If a file is 0 bytes, we can consider it done immediately.
|
||||
// This should never really happen in practice, but we handle it gracefully.
|
||||
break UploadFileStream
|
||||
}
|
||||
case *proto.UploadFileRequest_ChunkPiece:
|
||||
if file == nil {
|
||||
return xerrors.New("unexpected chunk piece while waiting for file upload")
|
||||
}
|
||||
|
||||
done, err := file.Add(&sdkproto.ChunkPiece{
|
||||
Data: typed.ChunkPiece.Data,
|
||||
FullDataHash: typed.ChunkPiece.FullDataHash,
|
||||
PieceIndex: typed.ChunkPiece.PieceIndex,
|
||||
})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("unable to add chunk piece: %w", err)
|
||||
}
|
||||
|
||||
if done {
|
||||
break UploadFileStream
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fileData, err := file.Complete()
|
||||
if err != nil {
|
||||
return xerrors.Errorf("complete file upload: %w", err)
|
||||
}
|
||||
|
||||
// Just rehash the data to be sure it is correct.
|
||||
hashBytes := sha256.Sum256(fileData)
|
||||
hash := hex.EncodeToString(hashBytes[:])
|
||||
|
||||
var insert database.InsertFileParams
|
||||
|
||||
switch file.Type {
|
||||
case sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES:
|
||||
insert = database.InsertFileParams{
|
||||
ID: uuid.New(),
|
||||
Hash: hash,
|
||||
CreatedAt: dbtime.Now(),
|
||||
CreatedBy: uuid.Nil,
|
||||
Mimetype: tarMimeType,
|
||||
Data: fileData,
|
||||
}
|
||||
default:
|
||||
return xerrors.Errorf("unsupported file upload type: %s", file.Type)
|
||||
}
|
||||
|
||||
//nolint:gocritic // Provisionerd actor
|
||||
_, err = s.Database.InsertFile(dbauthz.AsProvisionerd(s.lifecycleCtx), insert)
|
||||
if err != nil {
|
||||
// Duplicated files already exist in the database, so we can ignore this error.
|
||||
if !database.IsUniqueViolation(err, database.UniqueFilesHashCreatedByKey) {
|
||||
return xerrors.Errorf("insert file: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
s.Logger.Info(s.lifecycleCtx, "file uploaded to database",
|
||||
slog.F("type", file.Type.String()),
|
||||
slog.F("hash", hash),
|
||||
slog.F("size", len(fileData)),
|
||||
// new_insert indicates whether the file was newly inserted or already existed.
|
||||
slog.F("new_insert", err == nil),
|
||||
)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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())
|
||||
@@ -1606,6 +1704,20 @@ func (s *server) completeTemplateImportJob(ctx context.Context, job database.Pro
|
||||
}
|
||||
}
|
||||
|
||||
if len(jobType.TemplateImport.ModuleFilesHash) > 0 {
|
||||
hashString := hex.EncodeToString(jobType.TemplateImport.ModuleFilesHash)
|
||||
//nolint:gocritic // Acting as provisioner
|
||||
file, err := db.GetFileByHashAndCreator(dbauthz.AsProvisionerd(ctx), database.GetFileByHashAndCreatorParams{Hash: hashString, CreatedBy: uuid.Nil})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get file by hash, it should have been uploaded: %w", err)
|
||||
}
|
||||
|
||||
fileID = uuid.NullUUID{
|
||||
Valid: true,
|
||||
UUID: file.ID,
|
||||
}
|
||||
}
|
||||
|
||||
err = db.InsertTemplateVersionTerraformValuesByJobID(ctx, database.InsertTemplateVersionTerraformValuesByJobIDParams{
|
||||
JobID: jobID,
|
||||
UpdatedAt: now,
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
package provisionerdserver_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
crand "crypto/rand"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/xerrors"
|
||||
"storj.io/drpc"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/externalauth"
|
||||
"github.com/coder/coder/v2/codersdk/drpcsdk"
|
||||
proto "github.com/coder/coder/v2/provisionerd/proto"
|
||||
sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
// TestUploadFileLargeModuleFiles tests the UploadFile RPC with large module files
|
||||
func TestUploadFileLargeModuleFiles(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||
|
||||
// Create server
|
||||
server, db, _, _ := setup(t, false, &overrides{
|
||||
externalAuthConfigs: []*externalauth.Config{{}},
|
||||
})
|
||||
|
||||
testSizes := []int{
|
||||
0, // Empty file
|
||||
512, // A small file
|
||||
drpcsdk.MaxMessageSize + 1024, // Just over the limit
|
||||
drpcsdk.MaxMessageSize * 2, // 2x the limit
|
||||
sdkproto.ChunkSize*3 + 512, // Multiple chunks with partial last
|
||||
}
|
||||
|
||||
for _, size := range testSizes {
|
||||
t.Run(fmt.Sprintf("size_%d_bytes", size), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Generate test module files data
|
||||
moduleData := make([]byte, size)
|
||||
_, err := crand.Read(moduleData)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Convert to upload format
|
||||
upload, chunks := sdkproto.BytesToDataUpload(sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES, moduleData)
|
||||
|
||||
stream := newMockUploadStream(upload, chunks...)
|
||||
|
||||
// Execute upload
|
||||
err = server.UploadFile(stream)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Upload should be done
|
||||
require.True(t, stream.isDone(), "stream should be done after upload")
|
||||
|
||||
// Verify file was stored in database
|
||||
hashString := fmt.Sprintf("%x", upload.DataHash)
|
||||
file, err := db.GetFileByHashAndCreator(ctx, database.GetFileByHashAndCreatorParams{
|
||||
Hash: hashString,
|
||||
CreatedBy: uuid.Nil, // Provisionerd creates with Nil UUID
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, hashString, file.Hash)
|
||||
require.Equal(t, moduleData, file.Data)
|
||||
require.Equal(t, "application/x-tar", file.Mimetype)
|
||||
|
||||
// Try to upload it again, and it should still be successful
|
||||
stream = newMockUploadStream(upload, chunks...)
|
||||
err = server.UploadFile(stream)
|
||||
require.NoError(t, err, "re-upload should succeed without error")
|
||||
require.True(t, stream.isDone(), "stream should be done after re-upload")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestUploadFileErrorScenarios tests various error conditions in file upload
|
||||
func TestUploadFileErrorScenarios(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
//nolint:dogsled
|
||||
server, _, _, _ := setup(t, false, &overrides{
|
||||
externalAuthConfigs: []*externalauth.Config{{}},
|
||||
})
|
||||
|
||||
// Generate test data
|
||||
moduleData := make([]byte, sdkproto.ChunkSize*2)
|
||||
_, err := crand.Read(moduleData)
|
||||
require.NoError(t, err)
|
||||
|
||||
upload, chunks := sdkproto.BytesToDataUpload(sdkproto.DataUploadType_UPLOAD_TYPE_MODULE_FILES, moduleData)
|
||||
|
||||
t.Run("chunk_before_upload", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
stream := newMockUploadStream(nil, chunks[0])
|
||||
|
||||
err := server.UploadFile(stream)
|
||||
require.ErrorContains(t, err, "unexpected chunk piece while waiting for file upload")
|
||||
require.True(t, stream.isDone(), "stream should be done after error")
|
||||
})
|
||||
|
||||
t.Run("duplicate_upload", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
stream := &mockUploadStream{
|
||||
done: make(chan struct{}),
|
||||
messages: make(chan *proto.UploadFileRequest, 2),
|
||||
}
|
||||
|
||||
up := &proto.UploadFileRequest{Type: &proto.UploadFileRequest_DataUpload{DataUpload: upload}}
|
||||
|
||||
// Send it twice
|
||||
stream.messages <- up
|
||||
stream.messages <- up
|
||||
|
||||
err := server.UploadFile(stream)
|
||||
require.ErrorContains(t, err, "unexpected file upload while waiting for file completion")
|
||||
require.True(t, stream.isDone(), "stream should be done after error")
|
||||
})
|
||||
|
||||
t.Run("unsupported_upload_type", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
//nolint:govet // Ignore lock copy
|
||||
cpy := *upload
|
||||
cpy.UploadType = sdkproto.DataUploadType_UPLOAD_TYPE_UNKNOWN // Set to an unsupported type
|
||||
stream := newMockUploadStream(&cpy, chunks...)
|
||||
|
||||
err := server.UploadFile(stream)
|
||||
require.ErrorContains(t, err, "unsupported file upload type")
|
||||
require.True(t, stream.isDone(), "stream should be done after error")
|
||||
})
|
||||
}
|
||||
|
||||
type mockUploadStream struct {
|
||||
done chan struct{}
|
||||
messages chan *proto.UploadFileRequest
|
||||
}
|
||||
|
||||
func (m mockUploadStream) SendAndClose(empty *proto.Empty) error {
|
||||
close(m.done)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m mockUploadStream) Recv() (*proto.UploadFileRequest, error) {
|
||||
msg, ok := <-m.messages
|
||||
if !ok {
|
||||
return nil, xerrors.New("no more messages to receive")
|
||||
}
|
||||
return msg, nil
|
||||
}
|
||||
func (*mockUploadStream) Context() context.Context { panic(errUnimplemented) }
|
||||
func (*mockUploadStream) MsgSend(msg drpc.Message, enc drpc.Encoding) error {
|
||||
panic(errUnimplemented)
|
||||
}
|
||||
|
||||
func (*mockUploadStream) MsgRecv(msg drpc.Message, enc drpc.Encoding) error {
|
||||
panic(errUnimplemented)
|
||||
}
|
||||
func (*mockUploadStream) CloseSend() error { panic(errUnimplemented) }
|
||||
func (*mockUploadStream) Close() error { panic(errUnimplemented) }
|
||||
func (m *mockUploadStream) isDone() bool {
|
||||
select {
|
||||
case <-m.done:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func newMockUploadStream(up *sdkproto.DataUpload, chunks ...*sdkproto.ChunkPiece) *mockUploadStream {
|
||||
stream := &mockUploadStream{
|
||||
done: make(chan struct{}),
|
||||
messages: make(chan *proto.UploadFileRequest, 1+len(chunks)),
|
||||
}
|
||||
if up != nil {
|
||||
stream.messages <- &proto.UploadFileRequest{Type: &proto.UploadFileRequest_DataUpload{DataUpload: up}}
|
||||
}
|
||||
|
||||
for _, chunk := range chunks {
|
||||
stream.messages <- &proto.UploadFileRequest{Type: &proto.UploadFileRequest_ChunkPiece{ChunkPiece: chunk}}
|
||||
}
|
||||
close(stream.messages)
|
||||
return stream
|
||||
}
|
||||
Reference in New Issue
Block a user