Implement configurable state chunk size

Author: radeksimko

docs/plugin-protocol/tfplugin6.proto

@@ -915,12 +915,22 @@ message ConfigureStateStore {
     message Request {
         string type_name = 1;
         DynamicValue config = 2;
+        StateStoreClientCapabilities capabilities = 3;
     }
     message Response {
         repeated Diagnostic diagnostics = 1;
+        StateStoreServerCapabilities capabilities = 2;
     }
 }
 
+message StateStoreClientCapabilities {
+    int64 chunk_size = 1; // suggested chunk size by Core
+}
+
+message StateStoreServerCapabilities {
+    int64 chunk_size = 1; // chosen chunk size by plugin
+}
+
 message ReadStateBytes {
     message Request {
         string type_name = 1;

internal/command/meta_backend.go

@@ -41,6 +41,20 @@ import (
 	"github.com/hashicorp/terraform/internal/tfdiags"
 )
 
+const (
+	// defaultStateStoreChunkSize is the default chunk size proposed
+	// to the provider.
+	// This can be tweaked but should provide reasonable performance
+	// trade-offs for average network conditions and state file sizes.
+	defaultStateStoreChunkSize int64 = 8 << 20 // 8 MB
+
+	// maxStateStoreChunkSize is the highest chunk size provider may choose
+	// which we still consider reasonable/safe.
+	// This reflects terraform-plugin-go's max. RPC message size of 256MB
+	// and leaves plenty of space for other variable data like diagnostics.
+	maxStateStoreChunkSize int64 = 128 << 20 // 128 MB
+)
+
 // BackendOpts are the options used to initialize a backendrun.OperationsBackend.
 type BackendOpts struct {
 	// BackendConfig is a representation of the backend configuration block given in

internal/plugin6/grpc_provider.go

@@ -48,6 +48,11 @@ func (p *GRPCProviderPlugin) GRPCServer(broker *plugin.GRPCBroker, s *grpc.Serve
 	return nil
 }
 
+// grpcMaxMessageSize is the maximum gRPC send and receive message sizes
+// This matches the maximum set by a server implemented via terraform-plugin-go.
+// See https://github.com/hashicorp/terraform-plugin-go/blob/a361c9bf/tfprotov6/tf6server/server.go#L88
+const grpcMaxMessageSize = 256 << 20
+
 // GRPCProvider handles the client, or core side of the plugin rpc connection.
 // The GRPCProvider methods are mostly a translation layer between the
 // terraform providers types and the grpc proto types, directly converting
@@ -75,8 +80,12 @@ type GRPCProvider struct {
 
 	// schema stores the schema for this provider. This is used to properly
 	// serialize the requests for schemas.
-	mu     sync.Mutex
 	schema providers.GetProviderSchemaResponse
+	// stateChunkSize stores the negotiated chunk size for any implemented
+	// state store (keyed by type name) that have gone through successful configuration.
+	stateChunkSize map[string]int
+
+	mu sync.Mutex
 }
 
 func (p *GRPCProvider) GetProviderSchema() providers.GetProviderSchemaResponse {
@@ -1503,25 +1512,38 @@ func (p *GRPCProvider) ConfigureStateStore(r providers.ConfigureStateStoreReques
 		return resp
 	}
 
+	clientCapabilities := stateStoreClientCapabilitiesToProto(r.Capabilities)
+	logger.Trace("GRPCProvider.v6: ConfigureStateStore: proposing client capabilities", clientCapabilities)
+
 	protoReq := &proto6.ConfigureStateStore_Request{
 		TypeName: r.TypeName,
 		Config: &proto6.DynamicValue{
 			Msgpack: mp,
 		},
+		Capabilities: clientCapabilities,
 	}
 
 	protoResp, err := p.client.ConfigureStateStore(p.ctx, protoReq)
 	if err != nil {
 		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
 		return resp
 	}
+	resp.Capabilities = stateStoreServerCapabilitiesFromProto(protoResp.Capabilities)
+	logger.Trace("GRPCProvider.v6: ConfigureStateStore: received server capabilities", resp.Capabilities)
+
 	resp.Diagnostics = resp.Diagnostics.Append(convert.ProtoToDiagnostics(protoResp.Diagnostics))
 	return resp
 }
 
 func (p *GRPCProvider) ReadStateBytes(r providers.ReadStateBytesRequest) (resp providers.ReadStateBytesResponse) {
 	logger.Trace("GRPCProvider.v6: ReadStateBytes")
 
+	// ReadStateBytes can be more sensitive to message sizes
+	// so we ensure it aligns with (the lower) terraform-plugin-go.
+	opts := grpc.MaxRecvMsgSizeCallOption{
+		MaxRecvMsgSize: grpcMaxMessageSize,
+	}
+
 	schema := p.GetProviderSchema()
 	if schema.Diagnostics.HasErrors() {
 		resp.Diagnostics = schema.Diagnostics
@@ -1544,14 +1566,15 @@ func (p *GRPCProvider) ReadStateBytes(r providers.ReadStateBytesRequest) (resp p
 	ctx, cancel := context.WithCancel(p.ctx)
 	defer cancel()
 
-	client, err := p.client.ReadStateBytes(ctx, protoReq)
+	client, err := p.client.ReadStateBytes(ctx, protoReq, opts)
 	if err != nil {
 		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
 		return resp
 	}
 
 	buf := &bytes.Buffer{}
 	var expectedTotalLength int
+	// TODO: Send warning if client misbehaves and uses (lower) chunk size that we didn't agree on
 	for {
 		chunk, err := client.Recv()
 		if err == io.EOF {
@@ -1613,6 +1636,18 @@ func (p *GRPCProvider) ReadStateBytes(r providers.ReadStateBytesRequest) (resp p
 func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) (resp providers.WriteStateBytesResponse) {
 	logger.Trace("GRPCProvider.v6: WriteStateBytes")
 
+	// WriteStateBytes can be more sensitive to message sizes
+	// so we ensure it aligns with (the lower) terraform-plugin-go.
+	opts := grpc.MaxSendMsgSizeCallOption{
+		MaxSendMsgSize: grpcMaxMessageSize,
+	}
+
+	chunkSize, ok := p.stateChunkSize[r.TypeName]
+	if !ok {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("Unable to determine chunk size for provider %s; this is a bug in Terraform - please report it", r.TypeName))
+		return resp
+	}
+
 	schema := p.GetProviderSchema()
 	if schema.Diagnostics.HasErrors() {
 		resp.Diagnostics = schema.Diagnostics
@@ -1630,10 +1665,7 @@ func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) (resp
 	ctx, cancel := context.WithCancel(p.ctx)
 	defer cancel()
 
-	// TODO: Configurable chunk size
-	chunkSize := 4 * 1_000_000 // 4MB
-
-	client, err := p.client.WriteStateBytes(ctx)
+	client, err := p.client.WriteStateBytes(ctx, opts)
 	if err != nil {
 		resp.Diagnostics = resp.Diagnostics.Append(grpcErr(err))
 		return resp
@@ -1683,6 +1715,15 @@ func (p *GRPCProvider) WriteStateBytes(r providers.WriteStateBytesRequest) (resp
 	return resp
 }
 
+func (p *GRPCProvider) SetStateStoreChunkSize(typeName string, size int) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.stateChunkSize == nil {
+		p.stateChunkSize = make(map[string]int, 1)
+	}
+	p.stateChunkSize[typeName] = size
+}
+
 func (p *GRPCProvider) GetStates(r providers.GetStatesRequest) (resp providers.GetStatesResponse) {
 	logger.Trace("GRPCProvider.v6: GetStates")
 
@@ -1936,3 +1977,15 @@ func clientCapabilitiesToProto(c providers.ClientCapabilities) *proto6.ClientCap
 		WriteOnlyAttributesAllowed: c.WriteOnlyAttributesAllowed,
 	}
 }
+
+func stateStoreClientCapabilitiesToProto(c providers.StateStoreClientCapabilities) *proto6.StateStoreClientCapabilities {
+	return &proto6.StateStoreClientCapabilities{
+		ChunkSize: c.ChunkSize,
+	}
+}
+
+func stateStoreServerCapabilitiesFromProto(c *proto6.StateStoreServerCapabilities) providers.StateStoreServerCapabilities {
+	return providers.StateStoreServerCapabilities{
+		ChunkSize: c.ChunkSize,
+	}
+}

internal/providers/provider.go

@@ -146,6 +146,10 @@ type Interface interface {
 	Close() error
 }
 
+type StateStoreChunkSizeSetter interface {
+	SetStateStoreChunkSize(typeName string, size int)
+}
+
 // GetProviderSchemaResponse is the return type for GetProviderSchema, and
 // should only be used when handling a value for that method. The handling of
 // of schemas in any other context should always use ProviderSchema, so that
@@ -839,11 +843,23 @@ type ConfigureStateStoreRequest struct {
 
 	// Config is the configuration value to configure the store with.
 	Config cty.Value
+
+	Capabilities StateStoreClientCapabilities
+}
+
+type StateStoreClientCapabilities struct {
+	ChunkSize int64
 }
 
 type ConfigureStateStoreResponse struct {
 	// Diagnostics contains any warnings or errors from the method call.
 	Diagnostics tfdiags.Diagnostics
+
+	Capabilities StateStoreServerCapabilities
+}
+
+type StateStoreServerCapabilities struct {
+	ChunkSize int64
 }
 
 type ReadStateBytesRequest struct {