chore: implement first part of file streaming

Emyrk · Emyrk · commit fbfa08ba4323 · 2025-06-06T09:19:23.000-05:00
diff --git a/provisioner/terraform/executor.go b/provisioner/terraform/executor.go
@@ -19,13 +19,11 @@ import (
 	tfjson "github.com/hashicorp/terraform-json"
 	"go.opentelemetry.io/otel/attribute"
 	"golang.org/x/xerrors"
-	protobuf "google.golang.org/protobuf/proto"
 
 	"cdr.dev/slog"
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/tracing"
-	"github.com/coder/coder/v2/codersdk/drpcsdk"
 	"github.com/coder/coder/v2/provisionersdk/proto"
 )
 
@@ -357,10 +355,10 @@ func (e *executor) plan(ctx, killCtx context.Context, env, vars []string, logr l
 		ModuleFiles:           moduleFiles,
 	}
 
-	if protobuf.Size(msg) > drpcsdk.MaxMessageSize {
-		e.logger.Warn(ctx, "cannot persist terraform modules, message payload too big", slog.F("archive_size", len(msg.ModuleFiles)))
-		msg.ModuleFiles = nil
-	}
+	//if protobuf.Size(msg) > drpcsdk.MaxMessageSize {
+	//	e.logger.Warn(ctx, "cannot persist terraform modules, message payload too big", slog.F("archive_size", len(msg.ModuleFiles)))
+	//	msg.ModuleFiles = nil
+	//}
 
 	return msg, nil
 }
diff --git a/provisioner/terraform/modules.go b/provisioner/terraform/modules.go
@@ -13,6 +13,7 @@ import (
 
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/v2/coderd/util/xio"
 	"github.com/coder/coder/v2/provisionersdk/proto"
 )
 
@@ -85,7 +86,11 @@ func GetModulesArchive(root fs.FS) ([]byte, error) {
 
 	empty := true
 	var b bytes.Buffer
-	w := tar.NewWriter(&b)
+
+	// Limit to 20MB for now.
+	// TODO: Determine what a reasonable limit is for modules
+	lw := xio.NewLimitWriter(&b, 20<<20)
+	w := tar.NewWriter(lw)
 
 	for _, it := range m.Modules {
 		// Check to make sure that the module is a remote module fetched by
diff --git a/provisionerd/provisionerd_internal_test.go b/provisionerd/provisionerd_internal_test.go
@@ -0,0 +1 @@
+package provisionerd
diff --git a/provisionersdk/proto/dataupload.go b/provisionersdk/proto/dataupload.go
diff --git a/provisionersdk/session.go b/provisionersdk/session.go
@@ -17,6 +17,9 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
+	"github.com/coder/coder/v2/codersdk/drpcsdk"
+
+	protobuf "google.golang.org/protobuf/proto"
 
 	"github.com/coder/coder/v2/provisionersdk/proto"
 )
@@ -161,6 +164,29 @@ func (s *Session) handleRequests() error {
 				return err
 			}
 			resp.Type = &proto.Response_Plan{Plan: complete}
+
+			if protobuf.Size(resp) > drpcsdk.MaxMessageSize {
+				// Send the modules over as a stream
+				s.Logger.Info(s.Context(), "plan response too large, sending modules as stream")
+				dataUp, chunks := proto.BytesToDataUpload(complete.ModuleFiles)
+
+				complete.ModuleFiles = nil // sent over the stream
+				resp.Type = &proto.Response_Plan{Plan: complete}
+
+				err := s.stream.Send(&proto.Response{Type: &proto.Response_DataUpload{DataUpload: dataUp}})
+				if err != nil {
+					complete.Error = fmt.Sprintf("send data upload: %s", err.Error())
+				} else {
+					for i, chunk := range chunks {
+						err := s.stream.Send(&proto.Response{Type: &proto.Response_ChunkPiece{ChunkPiece: chunk}})
+						if err != nil {
+							complete.Error = fmt.Sprintf("send data piece upload %d/%d: %s", i, dataUp.Chunks, err.Error())
+							break
+						}
+					}
+				}
+			}
+
 			if complete.Error == "" {
 				planned = true
 			}
