PSS: Add chunk size negotiation to `savedStateStore` method, update mock provider for tests. (#37726)

* Add SetStateStoreChunkSize to the mock provider for tests

* Implement configurable state chunk size

* Add chunk size negotiation to `savedStateStore`, update happy path test to assert it's set

* Update `savedStateStore` to return diagnostic if a nil factory is passed in, add unhappy path tests

* Fix test error message

* Apply suggestions from code review

Co-authored-by: Radek Simko <radeksimko@users.noreply.github.com>

* Fix rename

---------

Co-authored-by: Radek Simko <radek.simko@gmail.com>
Co-authored-by: Radek Simko <radeksimko@users.noreply.github.com>

internal/command/meta_backend.go

@@ -44,6 +44,20 @@ import (
 	tfversion "github.com/hashicorp/terraform/version"
 )
 
+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
@@ -1529,7 +1543,7 @@ func (m *Meta) updateSavedBackendHash(cHash int, sMgr *clistate.LocalState) tfdi
 }
 
 // Initializing a saved state store from the backend state file (aka 'cache file', aka 'legacy state file')
-func (m *Meta) savedStateStore(sMgr *clistate.LocalState, providerFactory providers.Factory) (backend.Backend, tfdiags.Diagnostics) {
+func (m *Meta) savedStateStore(sMgr *clistate.LocalState, factory providers.Factory) (backend.Backend, tfdiags.Diagnostics) {
 	// We're preparing a state_store version of backend.Backend.
 	//
 	// The provider and state store will be configured using the backend state file.
@@ -1537,9 +1551,15 @@ func (m *Meta) savedStateStore(sMgr *clistate.LocalState, providerFactory provid
 	var diags tfdiags.Diagnostics
 	var b backend.Backend
 
-	s := sMgr.State()
-
-	provider, err := providerFactory()
+	if factory == nil {
+		diags = diags.Append(&hcl.Diagnostic{
+			Severity: hcl.DiagError,
+			Summary:  "Missing provider details when configuring state store",
+			Detail:   "Terraform attempted to configure a state store and no provider factory was available to launch it. This is a bug in Terraform and should be reported.",
+		})
+		return nil, diags
+	}
+	provider, err := factory()
 	if err != nil {
 		diags = diags.Append(fmt.Errorf("error when obtaining provider instance during state store initialization: %w", err))
 		return nil, diags
@@ -1548,6 +1568,7 @@ func (m *Meta) savedStateStore(sMgr *clistate.LocalState, providerFactory provid
 	// running provider instance inside the returned backend.Backend instance.
 	// Stopping the provider process is the responsibility of the calling code.
 
+	s := sMgr.State()
 	resp := provider.GetProviderSchema()
 
 	if len(resp.StateStores) == 0 {
@@ -1653,12 +1674,32 @@ func (m *Meta) savedStateStore(sMgr *clistate.LocalState, providerFactory provid
 	cfgStoreResp := provider.ConfigureStateStore(providers.ConfigureStateStoreRequest{
 		TypeName: s.StateStore.Type,
 		Config:   stateStoreConfigVal,
+		Capabilities: providers.StateStoreClientCapabilities{
+			ChunkSize: defaultStateStoreChunkSize,
+		},
 	})
 	diags = diags.Append(cfgStoreResp.Diagnostics)
 	if diags.HasErrors() {
 		return nil, diags
 	}
 
+	chunkSize := cfgStoreResp.Capabilities.ChunkSize
+	if chunkSize == 0 || chunkSize > maxStateStoreChunkSize {
+		diags = diags.Append(fmt.Errorf("Failed to negotiate acceptable chunk size. "+
+			"Expected size > 0 and <= %d bytes, provider wants %d bytes",
+			maxStateStoreChunkSize, chunkSize,
+		))
+		return nil, diags
+	}
+
+	p, ok := provider.(providers.StateStoreChunkSizeSetter)
+	if !ok {
+		msg := fmt.Sprintf("Unable to set chunk size for provider %s; this is a bug in Terraform - please report it", s.StateStore.Type)
+		panic(msg)
+	}
+	// casting to int here is okay because the number should never exceed int32
+	p.SetStateStoreChunkSize(s.StateStore.Type, int(chunkSize))
+
 	// Now we have a fully configured state store, ready to be used.
 	// To make it usable we need to return it in a backend.Backend interface.
 	b, err = backendPluggable.NewPluggable(provider, s.StateStore.Type)

internal/command/meta_backend_test.go

@@ -2440,6 +2440,7 @@ func TestSavedBackend(t *testing.T) {
 func TestSavedStateStore(t *testing.T) {
 	t.Run("the returned state store is configured with the backend state and not the current config", func(t *testing.T) {
 		// Create a temporary working directory
+		chunkSize := 42
 		td := t.TempDir()
 		testCopyDir(t, testFixturePath("state-store-changed"), td) // Fixtures with config that differs from backend state file
 		t.Chdir(td)
@@ -2473,7 +2474,21 @@ func TestSavedStateStore(t *testing.T) {
 			if config["value"].AsString() != "old-value" {
 				t.Fatalf("expected the state store to be configured with values from the backend state file (the string \"old-value\"), not the config. Got: %#v", config)
 			}
-			return providers.ConfigureStateStoreResponse{}
+			return providers.ConfigureStateStoreResponse{
+				Capabilities: providers.StateStoreServerCapabilities{
+					ChunkSize: int64(chunkSize),
+				},
+			}
+		}
+		mock.SetStateStoreChunkSizeFn = func(storeType string, size int) {
+			if storeType != "test_store" || size != chunkSize {
+				t.Fatalf("expected SetStateStoreChunkSize to be passed store type %q and chunk size %v, but got %q and %v",
+					"test_store",
+					chunkSize,
+					storeType,
+					size,
+				)
+			}
 		}
 
 		// Code under test
@@ -2489,10 +2504,104 @@ func TestSavedStateStore(t *testing.T) {
 				b,
 			)
 		}
+
+		if !mock.SetStateStoreChunkSizeCalled {
+			t.Fatal("expected configuring the pluggable state store to include a call to SetStateStoreChunkSize on the provider")
+		}
+	})
+
+	t.Run("error - no provider factory", func(t *testing.T) {
+		// sMgr pointing to a file that doesn't exist is sufficient setup for this test
+		sMgr := &clistate.LocalState{Path: "foobar.tfstate"}
+
+		m := testMetaBackend(t, nil)
+		_, diags := m.savedStateStore(sMgr, nil)
+		if !diags.HasErrors() {
+			t.Fatal("expected errors but got none")
+		}
+
+		expectedErr := "Missing provider details when configuring state store"
+		if !strings.Contains(diags.Err().Error(), expectedErr) {
+			t.Fatalf("expected the returned error to include %q, got: %s",
+				expectedErr,
+				diags.Err(),
+			)
+		}
+	})
+
+	t.Run("error - when there's no state stores in provider", func(t *testing.T) {
+		// Create a temporary working directory
+		td := t.TempDir()
+		testCopyDir(t, testFixturePath("state-store-changed"), td) // Fixtures with config that differs from backend state file
+		t.Chdir(td)
+
+		// Make a state manager for accessing the backend state file,
+		// and read the backend state from file
+		m := testMetaBackend(t, nil)
+		statePath := filepath.Join(m.DataDir(), DefaultStateFilename)
+		sMgr := &clistate.LocalState{Path: statePath}
+		err := sMgr.RefreshState()
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		mock := testStateStoreMock(t)
+		delete(mock.GetProviderSchemaResponse.StateStores, "test_store") // Remove the only state store impl.
+
+		_, diags := m.savedStateStore(sMgr, providers.FactoryFixed(mock))
+		if !diags.HasErrors() {
+			t.Fatal("expected errors but got none")
+		}
+		expectedErr := "Provider does not support pluggable state storage"
+		if !strings.Contains(diags.Err().Error(), expectedErr) {
+			t.Fatalf("expected the returned error to include %q, got: %s",
+				expectedErr,
+				diags.Err(),
+			)
+		}
 	})
 
-	// NOTE: the mock's functions include assertions about the values passed to
-	// the ConfigureProvider and ConfigureStateStore methods
+	t.Run("error - when there's no matching state store in provider Terraform suggests different identifier", func(t *testing.T) {
+		// Create a temporary working directory
+		td := t.TempDir()
+		testCopyDir(t, testFixturePath("state-store-changed"), td) // Fixtures with config that differs from backend state file
+		t.Chdir(td)
+
+		// Make a state manager for accessing the backend state file,
+		// and read the backend state from file
+		m := testMetaBackend(t, nil)
+		statePath := filepath.Join(m.DataDir(), DefaultStateFilename)
+		sMgr := &clistate.LocalState{Path: statePath}
+		err := sMgr.RefreshState()
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+
+		mock := testStateStoreMock(t)
+		testStore := mock.GetProviderSchemaResponse.StateStores["test_store"]
+		delete(mock.GetProviderSchemaResponse.StateStores, "test_store")
+		// Make the provider contain a "test_bore" impl., while the config specifies a "test_store" impl.
+		mock.GetProviderSchemaResponse.StateStores["test_bore"] = testStore
+
+		_, diags := m.savedStateStore(sMgr, providers.FactoryFixed(mock))
+		if !diags.HasErrors() {
+			t.Fatal("expected errors but got none")
+		}
+		expectedErr := "State store not implemented by the provider"
+		if !strings.Contains(diags.Err().Error(), expectedErr) {
+			t.Fatalf("expected the returned error to include %q, got: %s",
+				expectedErr,
+				diags.Err(),
+			)
+		}
+		expectedMsg := `Did you mean "test_bore"?`
+		if !strings.Contains(diags.Err().Error(), expectedMsg) {
+			t.Fatalf("expected the returned error to include %q, got: %s",
+				expectedMsg,
+				diags.Err(),
+			)
+		}
+	})
 }
 
 func TestMetaBackend_GetStateStoreProviderFactory(t *testing.T) {

internal/providers/testing/provider_mock.go

@@ -16,6 +16,7 @@ import (
 )
 
 var _ providers.Interface = (*MockProvider)(nil)
+var _ providers.StateStoreChunkSizeSetter = (*MockProvider)(nil)
 
 // MockProvider implements providers.Interface but mocks out all the
 // calls for testing purposes.
@@ -156,6 +157,10 @@ type MockProvider struct {
 	WriteStateBytesFn       func(providers.WriteStateBytesRequest) providers.WriteStateBytesResponse
 	WriteStateBytesResponse providers.WriteStateBytesResponse
 
+	// See providers.StateStoreChunkSizeSetter interface
+	SetStateStoreChunkSizeCalled bool
+	SetStateStoreChunkSizeFn     func(string, int)
+
 	LockStateCalled   bool
 	LockStateResponse providers.LockStateResponse
 	LockStateRequest  providers.LockStateRequest
@@ -1181,3 +1186,12 @@ func (p *MockProvider) ValidateActionConfig(r providers.ValidateActionConfigRequ
 
 	return resp
 }
+
+func (p *MockProvider) SetStateStoreChunkSize(storeType string, chunkSize int) {
+	p.SetStateStoreChunkSizeCalled = true
+	if p.SetStateStoreChunkSizeFn != nil {
+		p.SetStateStoreChunkSizeFn(storeType, chunkSize)
+	}
+
+	// If there's no function to use above we do nothing
+}