PSS : Add fs and inmem state storage implementations to the builtin simplev6 provider, update grpcwrap package, use PSS implementation in E2E test (#37790)

* feat: Implement `inmem` state store in provider-simple-v6
* feat: Add filesystem state store `fs` in provider-simple-v6, no locking implemented
* refactor: Move PSS chunking-related constants into the `pluggable` package, so they can be reused.
* feat: Implement PSS-related methods in grpcwrap package
* test: Add E2E test checking an init and apply (no plan) workflow is usable with both PSS implementations
* fix: Ensure state stores are configured with a suggested chunk size from Core

---------

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

internal/backend/local/backend.go

@@ -7,7 +7,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"log"
 	"os"
 	"path/filepath"
@@ -207,7 +206,7 @@ func (b *Local) Workspaces() ([]string, tfdiags.Diagnostics) {
 	// the listing always start with "default"
 	envs := []string{backend.DefaultStateName}
 
-	entries, err := ioutil.ReadDir(b.stateWorkspaceDir())
+	entries, err := os.ReadDir(b.stateWorkspaceDir())
 	// no error if there's no envs configured
 	if os.IsNotExist(err) {
 		return envs, nil

internal/backend/pluggable/chunks.go

@@ -0,0 +1,18 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package pluggable
+
+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
+)

internal/backend/pluggable/pluggable.go

@@ -105,6 +105,9 @@ func (p *Pluggable) Configure(config cty.Value) tfdiags.Diagnostics {
 	req := providers.ConfigureStateStoreRequest{
 		TypeName: p.typeName,
 		Config:   config,
+		Capabilities: providers.StateStoreClientCapabilities{
+			ChunkSize: DefaultStateStoreChunkSize,
+		},
 	}
 	resp := p.provider.ConfigureStateStore(req)
 	return resp.Diagnostics

internal/backend/remote-state/inmem/backend_test.go

@@ -65,7 +65,7 @@ func TestBackendLocked(t *testing.T) {
 	backend.TestBackendStateLocks(t, b1, b2)
 }
 
-// use the this backen to test the remote.State implementation
+// use this backend to test the remote.State implementation
 func TestRemoteState(t *testing.T) {
 	defer Reset()
 	b := backend.TestBackendConfig(t, New(), hcl.EmptyBody())

internal/builtin/providers/terraform/provider.go

@@ -16,6 +16,8 @@ import (
 // Provider is an implementation of providers.Interface
 type Provider struct{}
 
+var _ providers.Interface = &Provider{}
+
 // NewProvider returns a new terraform provider
 func NewProvider() providers.Interface {
 	return &Provider{}

internal/command/arguments/init.go

@@ -146,7 +146,7 @@ func ParseInit(args []string, experimentsEnabled bool) (*Init, tfdiags.Diagnosti
 			diags = diags.Append(tfdiags.Sourceless(
 				tfdiags.Error,
 				"Cannot use -enable-pluggable-state-storage-experiment flag without experiments enabled",
-				"Terraform cannot use the-enable-pluggable-state-storage-experiment flag (or TF_ENABLE_PLUGGABLE_STATE_STORAGE environment variable) unless experiments are enabled.",
+				"Terraform cannot use the -enable-pluggable-state-storage-experiment flag (or TF_ENABLE_PLUGGABLE_STATE_STORAGE environment variable) unless experiments are enabled.",
 			))
 		}
 		if !init.CreateDefaultWorkspace {

internal/command/e2etest/primary_test.go

@@ -4,6 +4,7 @@
 package e2etest
 
 import (
+	"os"
 	"path/filepath"
 	"reflect"
 	"sort"
@@ -12,7 +13,9 @@ import (
 
 	"github.com/davecgh/go-spew/spew"
 	"github.com/hashicorp/terraform/internal/e2e"
+	"github.com/hashicorp/terraform/internal/getproviders"
 	"github.com/hashicorp/terraform/internal/plans"
+	"github.com/hashicorp/terraform/internal/states/statefile"
 	"github.com/zclconf/go-cty/cty"
 )
 
@@ -230,3 +233,146 @@ func TestPrimaryChdirOption(t *testing.T) {
 		t.Errorf("incorrect destroy tally; want 0 destroyed:\n%s", stdout)
 	}
 }
+
+func TestPrimary_stateStore(t *testing.T) {
+
+	if !canRunGoBuild {
+		// We're running in a separate-build-then-run context, so we can't
+		// currently execute this test which depends on being able to build
+		// new executable at runtime.
+		//
+		// (See the comment on canRunGoBuild's declaration for more information.)
+		t.Skip("can't run without building a new provider executable")
+	}
+
+	t.Setenv(e2e.TestExperimentFlag, "true")
+	terraformBin := e2e.GoBuild("github.com/hashicorp/terraform", "terraform")
+
+	fixturePath := filepath.Join("testdata", "full-workflow-with-state-store-fs")
+	tf := e2e.NewBinary(t, terraformBin, fixturePath)
+
+	// In order to test integration with PSS we need a provider plugin implementing a state store.
+	// Here will build the simple6 (built with protocol v6) provider, which implements PSS.
+	simple6Provider := filepath.Join(tf.WorkDir(), "terraform-provider-simple6")
+	simple6ProviderExe := e2e.GoBuild("github.com/hashicorp/terraform/internal/provider-simple-v6/main", simple6Provider)
+
+	// Move the provider binaries into a directory that we will point terraform
+	// to using the -plugin-dir cli flag.
+	platform := getproviders.CurrentPlatform.String()
+	hashiDir := "cache/registry.terraform.io/hashicorp/"
+	if err := os.MkdirAll(tf.Path(hashiDir, "simple6/0.0.1/", platform), os.ModePerm); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Rename(simple6ProviderExe, tf.Path(hashiDir, "simple6/0.0.1/", platform, "terraform-provider-simple6")); err != nil {
+		t.Fatal(err)
+	}
+
+	//// INIT
+	stdout, stderr, err := tf.Run("init", "-enable-pluggable-state-storage-experiment=true", "-plugin-dir=cache", "-no-color")
+	if err != nil {
+		t.Fatalf("unexpected init error: %s\nstderr:\n%s", err, stderr)
+	}
+
+	if !strings.Contains(stdout, "Terraform created an empty state file for the default workspace") {
+		t.Errorf("notice about creating the default workspace is missing from init output:\n%s", stdout)
+	}
+
+	//// PLAN
+	// No separate plan step; this test lets the apply make a plan.
+
+	//// APPLY
+	stdout, stderr, err = tf.Run("apply", "-auto-approve", "-no-color")
+	if err != nil {
+		t.Fatalf("unexpected apply error: %s\nstderr:\n%s", err, stderr)
+	}
+
+	if !strings.Contains(stdout, "Resources: 1 added, 0 changed, 0 destroyed") {
+		t.Errorf("incorrect apply tally; want 1 added:\n%s", stdout)
+	}
+
+	// Check the statefile saved by the fs state store.
+	path := "terraform.tfstate.d/default/terraform.tfstate"
+	f, err := tf.OpenFile(path)
+	if err != nil {
+		t.Fatalf("unexpected error opening state file %s: %s\nstderr:\n%s", path, err, stderr)
+	}
+	defer f.Close()
+
+	stateFile, err := statefile.Read(f)
+	if err != nil {
+		t.Fatalf("unexpected error reading statefile %s: %s\nstderr:\n%s", path, err, stderr)
+	}
+
+	r := stateFile.State.RootModule().Resources
+	if len(r) != 1 {
+		t.Fatalf("expected state to include one resource, but got %d", len(r))
+	}
+	if _, ok := r["terraform_data.my-data"]; !ok {
+		t.Fatalf("expected state to include terraform_data.my-data but it's missing")
+	}
+}
+
+func TestPrimary_stateStore_inMem(t *testing.T) {
+	if !canRunGoBuild {
+		// We're running in a separate-build-then-run context, so we can't
+		// currently execute this test which depends on being able to build
+		// new executable at runtime.
+		//
+		// (See the comment on canRunGoBuild's declaration for more information.)
+		t.Skip("can't run without building a new provider executable")
+	}
+
+	t.Setenv(e2e.TestExperimentFlag, "true")
+	terraformBin := e2e.GoBuild("github.com/hashicorp/terraform", "terraform")
+
+	fixturePath := filepath.Join("testdata", "full-workflow-with-state-store-inmem")
+	tf := e2e.NewBinary(t, terraformBin, fixturePath)
+
+	// In order to test integration with PSS we need a provider plugin implementing a state store.
+	// Here will build the simple6 (built with protocol v6) provider, which implements PSS.
+	simple6Provider := filepath.Join(tf.WorkDir(), "terraform-provider-simple6")
+	simple6ProviderExe := e2e.GoBuild("github.com/hashicorp/terraform/internal/provider-simple-v6/main", simple6Provider)
+
+	// Move the provider binaries into a directory that we will point terraform
+	// to using the -plugin-dir cli flag.
+	platform := getproviders.CurrentPlatform.String()
+	hashiDir := "cache/registry.terraform.io/hashicorp/"
+	if err := os.MkdirAll(tf.Path(hashiDir, "simple6/0.0.1/", platform), os.ModePerm); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Rename(simple6ProviderExe, tf.Path(hashiDir, "simple6/0.0.1/", platform, "terraform-provider-simple6")); err != nil {
+		t.Fatal(err)
+	}
+
+	//// INIT
+	//
+	// Note - the inmem PSS implementation means that the default workspace state created during init
+	// is lost as soon as the command completes.
+	stdout, stderr, err := tf.Run("init", "-enable-pluggable-state-storage-experiment=true", "-plugin-dir=cache", "-no-color")
+	if err != nil {
+		t.Fatalf("unexpected init error: %s\nstderr:\n%s", err, stderr)
+	}
+
+	if !strings.Contains(stdout, "Terraform created an empty state file for the default workspace") {
+		t.Errorf("notice about creating the default workspace is missing from init output:\n%s", stdout)
+	}
+
+	//// PLAN
+	// No separate plan step; this test lets the apply make a plan.
+
+	//// APPLY
+	//
+	// Note - the inmem PSS implementation means that writing to the default workspace during apply
+	// is creating the default state file for the first time.
+	stdout, stderr, err = tf.Run("apply", "-auto-approve", "-no-color")
+	if err != nil {
+		t.Fatalf("unexpected apply error: %s\nstderr:\n%s", err, stderr)
+	}
+
+	if !strings.Contains(stdout, "Resources: 1 added, 0 changed, 0 destroyed") {
+		t.Errorf("incorrect apply tally; want 1 added:\n%s", stdout)
+	}
+
+	// We cannot inspect state or perform a destroy here, as the state isn't persisted between steps
+	// when we use the simple6_inmem state store.
+}

internal/command/e2etest/testdata/full-workflow-with-state-store-fs/main.tf

@@ -0,0 +1,23 @@
+terraform {
+  required_providers {
+    simple6 = {
+      source = "registry.terraform.io/hashicorp/simple6"
+    }
+  }
+
+  state_store "simple6_fs" {
+    provider "simple6" {}
+  }
+}
+
+variable "name" {
+  default = "world"
+}
+
+resource "terraform_data" "my-data" {
+  input = "hello ${var.name}"
+}
+
+output "greeting" {
+  value = resource.terraform_data.my-data.output
+}

internal/command/e2etest/testdata/full-workflow-with-state-store-inmem/main.tf

@@ -0,0 +1,23 @@
+terraform {
+  required_providers {
+    simple6 = {
+      source = "registry.terraform.io/hashicorp/simple6"
+    }
+  }
+
+  state_store "simple6_inmem" {
+    provider "simple6" {}
+  }
+}
+
+variable "name" {
+  default = "world"
+}
+
+resource "terraform_data" "my-data" {
+  input = "hello ${var.name}"
+}
+
+output "greeting" {
+  value = resource.terraform_data.my-data.output
+}

internal/command/meta_backend.go

@@ -50,20 +50,6 @@ 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
@@ -2085,7 +2071,7 @@ func (m *Meta) savedStateStore(sMgr *clistate.LocalState, factory providers.Fact
 		TypeName: s.StateStore.Type,
 		Config:   stateStoreConfigVal,
 		Capabilities: providers.StateStoreClientCapabilities{
-			ChunkSize: defaultStateStoreChunkSize,
+			ChunkSize: backendPluggable.DefaultStateStoreChunkSize,
 		},
 	})
 	diags = diags.Append(cfgStoreResp.Diagnostics)
@@ -2094,10 +2080,10 @@ func (m *Meta) savedStateStore(sMgr *clistate.LocalState, factory providers.Fact
 	}
 
 	chunkSize := cfgStoreResp.Capabilities.ChunkSize
-	if chunkSize == 0 || chunkSize > maxStateStoreChunkSize {
+	if chunkSize == 0 || chunkSize > backendPluggable.MaxStateStoreChunkSize {
 		diags = diags.Append(fmt.Errorf("Failed to negotiate acceptable chunk size. "+
 			"Expected size > 0 and <= %d bytes, provider wants %d bytes",
-			maxStateStoreChunkSize, chunkSize,
+			backendPluggable.MaxStateStoreChunkSize, chunkSize,
 		))
 		return nil, diags
 	}
@@ -2362,7 +2348,7 @@ func (m *Meta) stateStoreInitFromConfig(c *configs.StateStore, factory providers
 		TypeName: c.Type,
 		Config:   stateStoreConfigVal,
 		Capabilities: providers.StateStoreClientCapabilities{
-			ChunkSize: defaultStateStoreChunkSize,
+			ChunkSize: backendPluggable.DefaultStateStoreChunkSize,
 		},
 	})
 	diags = diags.Append(cfgStoreResp.Diagnostics)
@@ -2371,10 +2357,10 @@ func (m *Meta) stateStoreInitFromConfig(c *configs.StateStore, factory providers
 	}
 
 	chunkSize := cfgStoreResp.Capabilities.ChunkSize
-	if chunkSize == 0 || chunkSize > maxStateStoreChunkSize {
+	if chunkSize == 0 || chunkSize > backendPluggable.MaxStateStoreChunkSize {
 		diags = diags.Append(fmt.Errorf("Failed to negotiate acceptable chunk size. "+
 			"Expected size > 0 and <= %d bytes, provider wants %d bytes",
-			maxStateStoreChunkSize, chunkSize,
+			backendPluggable.MaxStateStoreChunkSize, chunkSize,
 		))
 		return nil, cty.NilVal, cty.NilVal, diags
 	}

internal/grpcwrap/provider6.go

@@ -4,8 +4,11 @@
 package grpcwrap
 
 import (
+	"bytes"
 	"context"
+	"errors"
 	"fmt"
+	"io"
 
 	"github.com/zclconf/go-cty/cty"
 	"github.com/zclconf/go-cty/cty/function"
@@ -15,6 +18,9 @@ import (
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	proto6 "github.com/hashicorp/terraform/internal/tfplugin6"
+
+	backendPluggable "github.com/hashicorp/terraform/internal/backend/pluggable"
 	"github.com/hashicorp/terraform/internal/plugin6/convert"
 	"github.com/hashicorp/terraform/internal/providers"
 	"github.com/hashicorp/terraform/internal/tfplugin6"
@@ -31,11 +37,15 @@ func Provider6(p providers.Interface) tfplugin6.ProviderServer {
 	}
 }
 
+var _ tfplugin6.ProviderServer = &provider6{}
+
 type provider6 struct {
 	provider        providers.Interface
 	schema          providers.GetProviderSchemaResponse
 	identitySchemas providers.GetResourceIdentitySchemasResponse
 
+	chunkSize int64
+
 	tfplugin6.UnimplementedProviderServer
 }
 
@@ -902,35 +912,264 @@ func (p *provider6) ListResource(req *tfplugin6.ListResource_Request, res tfplug
 }
 
 func (p *provider6) ValidateStateStoreConfig(ctx context.Context, req *tfplugin6.ValidateStateStore_Request) (*tfplugin6.ValidateStateStore_Response, error) {
-	panic("not implemented")
+	resp := &tfplugin6.ValidateStateStore_Response{}
+
+	s, ok := p.schema.StateStores[req.TypeName]
+	if !ok {
+		diag := &tfplugin6.Diagnostic{
+			Severity: tfplugin6.Diagnostic_ERROR,
+			Summary:  "Unsupported state store type",
+			Detail:   fmt.Sprintf("This provider doesn't include a state store called %q", req.TypeName),
+		}
+		resp.Diagnostics = append(resp.Diagnostics, diag)
+		return resp, nil
+	}
+	ty := s.Body.ImpliedType()
+
+	configVal, err := decodeDynamicValue6(req.Config, ty)
+	if err != nil {
+		resp.Diagnostics = convert.AppendProtoDiag(resp.Diagnostics, err)
+		return resp, nil
+	}
+
+	prepareResp := p.provider.ValidateStateStoreConfig(providers.ValidateStateStoreConfigRequest{
+		TypeName: req.TypeName,
+		Config:   configVal,
+	})
+
+	resp.Diagnostics = convert.AppendProtoDiag(resp.Diagnostics, prepareResp.Diagnostics)
+	return resp, nil
 }
 
 func (p *provider6) ConfigureStateStore(ctx context.Context, req *tfplugin6.ConfigureStateStore_Request) (*tfplugin6.ConfigureStateStore_Response, error) {
-	panic("not implemented")
+	resp := &tfplugin6.ConfigureStateStore_Response{}
+
+	s, ok := p.schema.StateStores[req.TypeName]
+	if !ok {
+		diag := &tfplugin6.Diagnostic{
+			Severity: tfplugin6.Diagnostic_ERROR,
+			Summary:  "Unsupported state store type",
+			Detail:   fmt.Sprintf("This provider doesn't include a state store called %q", req.TypeName),
+		}
+		resp.Diagnostics = append(resp.Diagnostics, diag)
+		return resp, nil
+	}
+	ty := s.Body.ImpliedType()
+
+	configVal, err := decodeDynamicValue6(req.Config, ty)
+	if err != nil {
+		resp.Diagnostics = convert.AppendProtoDiag(resp.Diagnostics, err)
+		return resp, nil
+	}
+
+	configureResp := p.provider.ConfigureStateStore(providers.ConfigureStateStoreRequest{
+		TypeName: req.TypeName,
+		Config:   configVal,
+		Capabilities: providers.StateStoreClientCapabilities{
+			ChunkSize: backendPluggable.DefaultStateStoreChunkSize,
+		},
+	})
+
+	// Validate the returned chunk size value
+	if configureResp.Capabilities.ChunkSize == 0 || configureResp.Capabilities.ChunkSize > backendPluggable.MaxStateStoreChunkSize {
+		diag := &tfplugin6.Diagnostic{
+			Severity: tfplugin6.Diagnostic_ERROR,
+			Summary:  "Failed to negotiate acceptable chunk size",
+			Detail: fmt.Sprintf("Expected size > 0 and <= %d bytes, provider wants %d bytes",
+				backendPluggable.MaxStateStoreChunkSize, configureResp.Capabilities.ChunkSize),
+		}
+		resp.Diagnostics = append(resp.Diagnostics, diag)
+		return resp, nil
+	}
+	p.chunkSize = configureResp.Capabilities.ChunkSize
+
+	resp.Diagnostics = convert.AppendProtoDiag(resp.Diagnostics, configureResp.Diagnostics)
+	resp.Capabilities = &tfplugin6.StateStoreServerCapabilities{
+		ChunkSize: configureResp.Capabilities.ChunkSize,
+	}
+	return resp, nil
 }
 
 func (p *provider6) ReadStateBytes(req *tfplugin6.ReadStateBytes_Request, srv tfplugin6.Provider_ReadStateBytesServer) error {
-	panic("not implemented")
+	stateReadResp := p.provider.ReadStateBytes(providers.ReadStateBytesRequest{
+		TypeName: req.TypeName,
+		StateId:  req.StateId,
+	})
+
+	state := stateReadResp.Bytes
+	reader := bytes.NewReader(state)
+	totalLength := reader.Size() // length in bytes
+	rangeStart := 0
+
+	for {
+		var diags []*proto6.Diagnostic
+		readBytes := make([]byte, p.chunkSize)
+		byteCount, err := reader.Read(readBytes)
+		if err != nil && !errors.Is(err, io.EOF) {
+			diags := []*proto6.Diagnostic{
+				{
+					Severity: proto6.Diagnostic_ERROR,
+					Summary:  "Error reading from state file",
+					Detail: fmt.Sprintf("State store %s experienced an error when reading from the state file for workspace %s: %s",
+						req.TypeName,
+						req.StateId,
+						err,
+					),
+				},
+			}
+			err := srv.Send(&proto6.ReadStateBytes_Response{
+				// Zero values accompany the error diagnostic
+				Bytes:       nil,
+				TotalLength: 0,
+				Range: &proto6.StateRange{
+					Start: 0,
+					End:   0,
+				},
+				Diagnostics: diags,
+			})
+			if err != nil {
+				return err
+			}
+		}
+
+		if byteCount == 0 {
+			// The previous iteration read the last byte of the data.
+			return nil
+		}
+
+		err = srv.Send(&proto6.ReadStateBytes_Response{
+			Bytes:       readBytes[0:byteCount],
+			TotalLength: int64(totalLength),
+			Range: &proto6.StateRange{
+				Start: int64(rangeStart),
+				End:   int64(rangeStart + byteCount),
+			},
+			Diagnostics: diags,
+		})
+		if err != nil {
+			return err
+		}
+
+		// Track progress to ensure Range values are correct.
+		rangeStart += byteCount
+	}
 }
 
 func (p *provider6) WriteStateBytes(srv tfplugin6.Provider_WriteStateBytesServer) error {
-	panic("not implemented")
+	var typeName string
+	var stateId string
+
+	state := bytes.Buffer{}
+	var grpcErr error
+	var totalReceivedBytes int
+	var expectedTotalLength int64
+	for {
+		chunk, err := srv.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			grpcErr = fmt.Errorf("wrapped err: %w", err)
+			break
+		}
+		if expectedTotalLength == 0 {
+			// On the first iteration
+			expectedTotalLength = chunk.TotalLength // record expected length
+			if chunk.Meta != nil {
+				// We expect the Meta to be set on the first message, only
+				typeName = chunk.Meta.TypeName
+				stateId = chunk.Meta.StateId
+			} else {
+				panic("expected Meta to be set on first chunk sent to WriteStateBytes")
+			}
+		}
+
+		n, err := state.Write(chunk.Bytes)
+		if err != nil {
+			return fmt.Errorf("error writing state: %w", err)
+		}
+		totalReceivedBytes += n
+	}
+
+	if grpcErr != nil {
+		return grpcErr
+	}
+
+	if int64(totalReceivedBytes) != expectedTotalLength {
+		return fmt.Errorf("expected to receive state in %d bytes, actually received %d bytes", expectedTotalLength, totalReceivedBytes)
+	}
+
+	if totalReceivedBytes == 0 {
+		// Even an empty state file has content; no bytes is not valid
+		return errors.New("No state data received from Terraform: No state data was received from Terraform. This is a bug and should be reported.")
+	}
+
+	resp := p.provider.WriteStateBytes(providers.WriteStateBytesRequest{
+		StateId:  stateId,
+		TypeName: typeName,
+		Bytes:    state.Bytes(),
+	})
+
+	err := srv.SendAndClose(&proto6.WriteStateBytes_Response{
+		Diagnostics: convert.AppendProtoDiag([]*proto6.Diagnostic{}, resp.Diagnostics),
+	})
+
+	return err
 }
 
 func (p *provider6) LockState(ctx context.Context, req *tfplugin6.LockState_Request) (*tfplugin6.LockState_Response, error) {
-	panic("not implemented")
+	lockResp := p.provider.LockState(providers.LockStateRequest{
+		TypeName:  req.TypeName,
+		StateId:   req.StateId,
+		Operation: req.Operation,
+	})
+
+	resp := &tfplugin6.LockState_Response{
+		LockId:      lockResp.LockId,
+		Diagnostics: convert.AppendProtoDiag([]*proto6.Diagnostic{}, lockResp.Diagnostics),
+	}
+
+	return resp, nil
 }
 
 func (p *provider6) UnlockState(ctx context.Context, req *tfplugin6.UnlockState_Request) (*tfplugin6.UnlockState_Response, error) {
-	panic("not implemented")
+	unlockResp := p.provider.UnlockState(providers.UnlockStateRequest{
+		TypeName: req.TypeName,
+		StateId:  req.StateId,
+		LockId:   req.LockId,
+	})
+
+	resp := &tfplugin6.UnlockState_Response{
+		Diagnostics: convert.AppendProtoDiag([]*proto6.Diagnostic{}, unlockResp.Diagnostics),
+	}
+
+	return resp, nil
 }
 
 func (p *provider6) GetStates(ctx context.Context, req *tfplugin6.GetStates_Request) (*tfplugin6.GetStates_Response, error) {
-	panic("not implemented")
+	getStatesResp := p.provider.GetStates(providers.GetStatesRequest{
+		TypeName: req.TypeName,
+	})
+
+	resp := &tfplugin6.GetStates_Response{
+		StateId:     getStatesResp.States,
+		Diagnostics: convert.AppendProtoDiag([]*tfplugin6.Diagnostic{}, getStatesResp.Diagnostics),
+	}
+
+	return resp, nil
 }
 
 func (p *provider6) DeleteState(ctx context.Context, req *tfplugin6.DeleteState_Request) (*tfplugin6.DeleteState_Response, error) {
-	panic("not implemented")
+	deleteStatesResp := p.provider.DeleteState(providers.DeleteStateRequest{
+		TypeName: req.TypeName,
+		StateId:  req.StateId,
+	})
+
+	resp := &tfplugin6.DeleteState_Response{
+		Diagnostics: convert.AppendProtoDiag([]*tfplugin6.Diagnostic{}, deleteStatesResp.Diagnostics),
+	}
+
+	return resp, nil
 }
 
 func (p *provider6) PlanAction(_ context.Context, req *tfplugin6.PlanAction_Request) (*tfplugin6.PlanAction_Response, error) {

internal/provider-simple-v6/provider.go

@@ -5,6 +5,7 @@
 package simple
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"log"
@@ -13,15 +14,48 @@ import (
 	"github.com/zclconf/go-cty/cty"
 	ctyjson "github.com/zclconf/go-cty/cty/json"
 
+	"github.com/hashicorp/terraform/internal/backend"
 	"github.com/hashicorp/terraform/internal/configs/configschema"
 	"github.com/hashicorp/terraform/internal/providers"
+	"github.com/hashicorp/terraform/internal/states/statefile"
 )
 
 type simple struct {
 	schema providers.GetProviderSchemaResponse
+
+	inMem *InMemStoreSingle
+	fs    *FsStore
 }
 
+// Provider returns an instance of providers.Interface
 func Provider() providers.Interface {
+	return provider()
+}
+
+// ProviderWithDefaultState returns an instance of providers.Interface,
+// where the underlying simple struct has been changed to indicate that the
+// 'default' state has already been created as an empty state file.
+func ProviderWithDefaultState() providers.Interface {
+	// Get the empty state file as bytes
+	f := statefile.New(nil, "", 0)
+
+	var buf bytes.Buffer
+	err := statefile.Write(f, &buf)
+	if err != nil {
+		panic(err)
+	}
+	emptyStateBytes := buf.Bytes()
+
+	p := provider()
+
+	p.inMem.states.m = make(map[string][]byte, 1)
+	p.inMem.states.m[backend.DefaultStateName] = emptyStateBytes
+
+	return p
+}
+
+// provider returns an instance of simple
+func provider() simple {
 	simpleResource := providers.Schema{
 		Body: &configschema.Block{
 			Attributes: map[string]*configschema.Attribute{
@@ -46,7 +80,7 @@ func Provider() providers.Interface {
 		},
 	}
 
-	return simple{
+	provider := simple{
 		schema: providers.GetProviderSchemaResponse{
 			Provider: providers.Schema{
 				Body: &configschema.Block{
@@ -75,6 +109,10 @@ func Provider() providers.Interface {
 				},
 			},
 			Actions: map[string]providers.ActionSchema{},
+			StateStores: map[string]providers.Schema{
+				inMemStoreName: stateStoreInMemGetSchema(), // simple6_inmem
+				fsStoreName:    stateStoreFsGetSchema(),    // simple6_fs
+			},
 			ServerCapabilities: providers.ServerCapabilities{
 				PlanDestroy:               true,
 				GetProviderSchemaOptional: true,
@@ -97,7 +135,13 @@ func Provider() providers.Interface {
 				},
 			},
 		},
+
+		// the "default" state doesn't exist by default here; needs explicit creation via init command
+		inMem: &InMemStoreSingle{},
+		fs:    &FsStore{},
 	}
+
+	return provider
 }
 
 func (s simple) GetProviderSchema() providers.GetProviderSchemaResponse {
@@ -310,35 +354,107 @@ func (s simple) ListResource(req providers.ListResourceRequest) (resp providers.
 }
 
 func (s simple) ValidateStateStoreConfig(req providers.ValidateStateStoreConfigRequest) providers.ValidateStateStoreConfigResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.ValidateStateStoreConfig(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.ValidateStateStoreConfig(req)
+	}
+
+	var resp providers.ValidateStateStoreConfigResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) ConfigureStateStore(req providers.ConfigureStateStoreRequest) providers.ConfigureStateStoreResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.ConfigureStateStore(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.ConfigureStateStore(req)
+	}
+
+	var resp providers.ConfigureStateStoreResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) ReadStateBytes(req providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.ReadStateBytes(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.ReadStateBytes(req)
+	}
+
+	var resp providers.ReadStateBytesResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) WriteStateBytes(req providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.WriteStateBytes(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.WriteStateBytes(req)
+	}
+
+	var resp providers.WriteStateBytesResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) LockState(req providers.LockStateRequest) providers.LockStateResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.LockState(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.LockState(req)
+	}
+
+	var resp providers.LockStateResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) UnlockState(req providers.UnlockStateRequest) providers.UnlockStateResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.UnlockState(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.UnlockState(req)
+	}
+
+	var resp providers.UnlockStateResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) GetStates(req providers.GetStatesRequest) providers.GetStatesResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.GetStates(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.GetStates(req)
+	}
+
+	var resp providers.GetStatesResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) DeleteState(req providers.DeleteStateRequest) providers.DeleteStateResponse {
-	panic("not implemented")
+	if req.TypeName == inMemStoreName {
+		return s.inMem.DeleteState(req)
+	}
+	if req.TypeName == fsStoreName {
+		return s.fs.DeleteState(req)
+	}
+
+	var resp providers.DeleteStateResponse
+	resp.Diagnostics.Append(fmt.Errorf("unsupported state store type %q", req.TypeName))
+	return resp
 }
 
 func (s simple) PlanAction(providers.PlanActionRequest) providers.PlanActionResponse {

internal/provider-simple-v6/state_store_fs.go

@@ -0,0 +1,263 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package simple
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"path"
+	"path/filepath"
+	"sort"
+
+	"github.com/hashicorp/terraform/internal/backend"
+	"github.com/hashicorp/terraform/internal/configs/configschema"
+	"github.com/hashicorp/terraform/internal/providers"
+	"github.com/hashicorp/terraform/internal/states/statemgr"
+	"github.com/hashicorp/terraform/internal/tfdiags"
+	"github.com/zclconf/go-cty/cty"
+)
+
+const fsStoreName = "simple6_fs"
+const defaultStatesDir = "terraform.tfstate.d"
+
+// FsStore allows storing state in the local filesystem.
+//
+// This state storage implementation differs from the old "local" backend in core,
+// by storing all states in the custom, or default, states directory. In the "local"
+// backend the default state was a special case and was handled differently to custom states.
+type FsStore struct {
+	// Configured values
+	statesDir string
+	chunkSize int64
+
+	states map[string]*statemgr.Filesystem
+}
+
+func stateStoreFsGetSchema() providers.Schema {
+	return providers.Schema{
+		Body: &configschema.Block{
+			Attributes: map[string]*configschema.Attribute{
+				// Named workspace_dir to match what's present in the local backend
+				"workspace_dir": {
+					Type:        cty.String,
+					Optional:    true,
+					Description: "The directory where state files will be created. When unset the value will default to terraform.tfstate.d",
+				},
+			},
+		},
+	}
+}
+
+func (f *FsStore) ValidateStateStoreConfig(req providers.ValidateStateStoreConfigRequest) providers.ValidateStateStoreConfigResponse {
+	var resp providers.ValidateStateStoreConfigResponse
+
+	attrs := req.Config.AsValueMap()
+	if v, ok := attrs["workspace_dir"]; ok {
+		if !v.IsKnown() {
+			resp.Diagnostics = resp.Diagnostics.Append(errors.New("the attribute \"workspace_dir\" cannot be an unknown value"))
+			return resp
+		}
+	}
+
+	return resp
+}
+
+func (f *FsStore) ConfigureStateStore(req providers.ConfigureStateStoreRequest) providers.ConfigureStateStoreResponse {
+	resp := providers.ConfigureStateStoreResponse{}
+
+	configVal := req.Config
+	if v := configVal.GetAttr("workspace_dir"); !v.IsNull() {
+		f.statesDir = v.AsString()
+	} else {
+		f.statesDir = defaultStatesDir
+	}
+
+	if f.states == nil {
+		f.states = make(map[string]*statemgr.Filesystem)
+	}
+
+	// We need to select return a suggested chunk size; use the value suggested by Core
+	resp.Capabilities.ChunkSize = req.Capabilities.ChunkSize
+	f.chunkSize = req.Capabilities.ChunkSize
+
+	return resp
+}
+
+func (f *FsStore) LockState(req providers.LockStateRequest) providers.LockStateResponse {
+	resp := providers.LockStateResponse{}
+	resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
+		tfdiags.Warning,
+		"Locking not implemented",
+		fmt.Sprintf("Could not lock state %q; state locking isn't implemented", req.StateId),
+	))
+	return resp
+}
+
+func (f *FsStore) UnlockState(req providers.UnlockStateRequest) providers.UnlockStateResponse {
+	resp := providers.UnlockStateResponse{}
+	resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
+		tfdiags.Warning,
+		"Unlocking not implemented",
+		fmt.Sprintf("Could not unlock state %q; state locking isn't implemented", req.StateId),
+	))
+	return resp
+}
+
+func (f *FsStore) GetStates(req providers.GetStatesRequest) providers.GetStatesResponse {
+	resp := providers.GetStatesResponse{}
+
+	entries, err := os.ReadDir(f.statesDir)
+	// no error if there's no envs configured
+	if os.IsNotExist(err) {
+		return resp
+	}
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(err)
+		return resp
+	}
+
+	var envs []string
+	for _, entry := range entries {
+		if entry.IsDir() {
+			envs = append(envs, filepath.Base(entry.Name()))
+		}
+	}
+
+	sort.Strings(envs)
+	resp.States = envs
+	return resp
+}
+
+func (f *FsStore) DeleteState(req providers.DeleteStateRequest) providers.DeleteStateResponse {
+	resp := providers.DeleteStateResponse{}
+
+	if req.StateId == "" {
+		resp.Diagnostics = resp.Diagnostics.Append(errors.New("empty state name"))
+		return resp
+	}
+
+	if req.StateId == backend.DefaultStateName {
+		resp.Diagnostics = resp.Diagnostics.Append(errors.New("cannot delete default state"))
+		return resp
+	}
+
+	delete(f.states, req.StateId)
+	err := os.RemoveAll(filepath.Join(f.statesDir, req.StateId))
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error deleting state %q: %w", req.StateId, err))
+		return resp
+	}
+
+	return resp
+}
+
+func (f *FsStore) getStatePath(stateId string) string {
+	return path.Join(f.statesDir, stateId, "terraform.tfstate")
+}
+
+func (f *FsStore) getStateDir(stateId string) string {
+	return path.Join(f.statesDir, stateId)
+}
+
+func (f *FsStore) ReadStateBytes(req providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
+	log.Printf("[DEBUG] ReadStateBytes: reading data from the %q state", req.StateId)
+	resp := providers.ReadStateBytesResponse{}
+
+	// E.g. terraform.tfstate.d/foobar/terraform.tfstate
+	path := f.getStatePath(req.StateId)
+	file, err := os.Open(path)
+
+	fileExists := true
+	if err != nil {
+		if _, ok := err.(*os.PathError); !ok {
+			// Error other than the file not existing
+			resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error opening state file %q: %w", path, err))
+			return resp
+		}
+		fileExists = false
+	}
+	defer file.Close()
+
+	buf := bytes.Buffer{}
+	var processedBytes int
+
+	if fileExists {
+		for {
+			b := make([]byte, f.chunkSize)
+			n, err := file.Read(b)
+			if err == io.EOF {
+				break
+			}
+			if err != nil {
+				resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error reading from state file %q: %w", path, err))
+				return resp
+			}
+			buf.Write(b[0:n])
+			processedBytes += n
+		}
+	}
+	log.Printf("[DEBUG] ReadStateBytes: read %d bytes of data from state file %q", processedBytes, path)
+
+	if processedBytes == 0 {
+		// Does not exist, so return no bytes
+		resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
+			tfdiags.Warning,
+			"State doesn't exist",
+			fmt.Sprintf("The %q state does not exist", req.StateId),
+		))
+	}
+
+	resp.Bytes = buf.Bytes()
+	return resp
+}
+
+func (f *FsStore) WriteStateBytes(req providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
+	log.Printf("[DEBUG] WriteStateBytes: writing data to the %q state", req.StateId)
+	resp := providers.WriteStateBytesResponse{}
+
+	// E.g. terraform.tfstate.d/foobar/terraform.tfstate
+	path := f.getStatePath(req.StateId)
+
+	// Create or open state file
+	dir := f.getStateDir(req.StateId)
+	err := os.MkdirAll(dir, 0755)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error creating state file directory %q: %w", dir, err))
+	}
+
+	file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0666)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error opening state file %q: %w", path, err))
+	}
+
+	buf := bytes.NewBuffer(req.Bytes)
+	var processedBytes int
+	if f.chunkSize == 0 {
+		panic("WriteStateBytes: chunk size zero. This is an error in Terraform and should be reported")
+	}
+	for {
+		data := buf.Next(int(f.chunkSize))
+		if len(data) == 0 {
+			break
+		}
+		n, err := file.Write(data)
+		if err != nil {
+			resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error writing to state file %q: %w", path, err))
+			return resp
+		}
+
+		processedBytes += n
+	}
+	log.Printf("[DEBUG] WriteStateBytes: wrote %d bytes of data to state file %q", processedBytes, path)
+
+	if processedBytes == 0 {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("missing state data: write action wrote %d bytes of data to file %q.", processedBytes, path))
+	}
+
+	return resp
+}

internal/provider-simple-v6/state_store_fs_test.go

@@ -0,0 +1,75 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package simple
+
+import (
+	"testing"
+
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/terraform/internal/addrs"
+	"github.com/hashicorp/terraform/internal/backend"
+	"github.com/hashicorp/terraform/internal/backend/pluggable"
+	"github.com/hashicorp/terraform/internal/states"
+	"github.com/zclconf/go-cty/cty"
+)
+
+// TODO: Testing of locking with 2 clients once locking is fully implemented.
+
+func TestFsStoreRemoteState(t *testing.T) {
+	td := t.TempDir()
+	t.Chdir(td)
+
+	provider := Provider()
+
+	plug, err := pluggable.NewPluggable(provider, fsStoreName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b := backend.TestBackendConfig(t, plug, hcl.EmptyBody())
+
+	// The "default" state doesn't exist by default
+	// (Note that this depends on the factory method used to get the provider above)
+	stateIds, wDiags := b.Workspaces()
+	if wDiags.HasErrors() {
+		t.Fatal(wDiags.Err())
+	}
+	if len(stateIds) != 0 {
+		t.Fatalf("unexpected response from Workspaces method: %#v", stateIds)
+	}
+
+	// create a new state using this backend
+	newStateId := "foobar"
+	emptyState := states.NewState()
+
+	sMgr, sDiags := b.StateMgr(newStateId)
+	if sDiags.HasErrors() {
+		t.Fatal(sDiags.Err())
+	}
+	if err := sMgr.WriteState(emptyState); err != nil {
+		t.Fatal(err)
+	}
+	if err := sMgr.PersistState(nil); err != nil {
+		t.Fatal(err)
+	}
+
+	// force overwriting the remote state
+	newState := states.NewState()
+	newState.SetOutputValue(
+		addrs.OutputValue{Name: "foo"}.Absolute(addrs.RootModuleInstance),
+		cty.StringVal("bar"),
+		false)
+
+	if err := sMgr.WriteState(newState); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := sMgr.PersistState(nil); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := sMgr.RefreshState(); err != nil {
+		t.Fatal(err)
+	}
+}

internal/provider-simple-v6/state_store_inmem.go

@@ -0,0 +1,226 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package simple
+
+import (
+	"errors"
+	"fmt"
+	"math/big"
+	"sort"
+	"sync"
+
+	"github.com/hashicorp/go-uuid"
+	"github.com/hashicorp/terraform/internal/backend"
+	"github.com/hashicorp/terraform/internal/configs/configschema"
+	"github.com/hashicorp/terraform/internal/providers"
+	"github.com/hashicorp/terraform/internal/tfdiags"
+	"github.com/zclconf/go-cty/cty"
+)
+
+const inMemStoreName = "simple6_inmem"
+
+// InMemStoreSingle allows 'storing' state in memory for the purpose of testing.
+//
+// "Single" reflects the fact that this implementation does not use any global scope vars
+// in its implementation, unlike the current inmem backend. HOWEVER, you can test whether locking
+// blocks multiple clients trying to access the same state at once by creating multiple instances
+// of backend.Backend that wrap the same provider.Interface instance.
+type InMemStoreSingle struct {
+	states stateMap
+	locks  lockMap
+}
+
+func stateStoreInMemGetSchema() providers.Schema {
+	return providers.Schema{
+		Body: &configschema.Block{
+			Attributes: map[string]*configschema.Attribute{
+				"lock_id": {
+					Type:        cty.String,
+					Optional:    true,
+					Description: "initializes the state in a locked configuration",
+				},
+			},
+		},
+	}
+}
+
+func (m *InMemStoreSingle) ValidateStateStoreConfig(req providers.ValidateStateStoreConfigRequest) providers.ValidateStateStoreConfigResponse {
+	var resp providers.ValidateStateStoreConfigResponse
+
+	attrs := req.Config.AsValueMap()
+
+	// This is completely arbitrary validation included here to avoid this method being empty. It is not here for a purpose,
+	// but could be used if an E2E test wants to trigger a validation error.
+	if v, ok := attrs["lock_id"]; ok && !v.IsNull() {
+		cutoff := cty.NumberVal(big.NewFloat(3))
+		if v.Length().LessThan(cutoff) == cty.True {
+			resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("when set, the attribute \"lock_id\" must have a length equal or greater than %s", cutoff.AsString()))
+			return resp
+		}
+	}
+
+	return resp
+}
+
+func (m *InMemStoreSingle) ConfigureStateStore(req providers.ConfigureStateStoreRequest) providers.ConfigureStateStoreResponse {
+	resp := providers.ConfigureStateStoreResponse{}
+
+	m.states.Lock()
+	defer m.states.Unlock()
+
+	// set the default client lock info per the test config
+	configVal := req.Config
+	if v := configVal.GetAttr("lock_id"); !v.IsNull() {
+		m.locks.lock(backend.DefaultStateName, v.AsString())
+	}
+
+	// We need to return a suggested chunk size; use the value suggested by Core
+	resp.Capabilities.ChunkSize = req.Capabilities.ChunkSize
+	return resp
+}
+
+func (m *InMemStoreSingle) ReadStateBytes(req providers.ReadStateBytesRequest) providers.ReadStateBytesResponse {
+	resp := providers.ReadStateBytesResponse{}
+
+	v, ok := m.states.m[req.StateId]
+	if !ok {
+		// Does not exist, so return no bytes
+
+		resp.Diagnostics = resp.Diagnostics.Append(tfdiags.Sourceless(
+			tfdiags.Warning,
+			"State doesn't exist",
+			fmt.Sprintf("The %q state does not exist", req.StateId),
+		))
+		return resp
+	}
+
+	resp.Bytes = v
+	return resp
+}
+
+func (m *InMemStoreSingle) WriteStateBytes(req providers.WriteStateBytesRequest) providers.WriteStateBytesResponse {
+	resp := providers.WriteStateBytesResponse{}
+
+	if m.states.m == nil {
+		m.states.m = make(map[string][]byte, 1)
+	}
+
+	m.states.m[req.StateId] = req.Bytes
+
+	return resp
+}
+
+func (m *InMemStoreSingle) LockState(req providers.LockStateRequest) providers.LockStateResponse {
+	resp := providers.LockStateResponse{}
+
+	lockIdBytes, err := uuid.GenerateUUID()
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error creating random lock uuid: %w", err))
+		return resp
+	}
+
+	lockId := string(lockIdBytes)
+	returnedLockId, err := m.locks.lock(req.StateId, lockId)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(err)
+	}
+
+	resp.LockId = string(returnedLockId)
+	return resp
+}
+
+func (m *InMemStoreSingle) UnlockState(req providers.UnlockStateRequest) providers.UnlockStateResponse {
+	resp := providers.UnlockStateResponse{}
+
+	err := m.locks.unlock(req.StateId, req.LockId)
+	if err != nil {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("error when unlocking state %q: %w", req.StateId, err))
+		return resp
+	}
+
+	return resp
+}
+
+func (m *InMemStoreSingle) GetStates(req providers.GetStatesRequest) providers.GetStatesResponse {
+	m.states.Lock()
+	defer m.states.Unlock()
+
+	resp := providers.GetStatesResponse{}
+
+	var stateIds []string
+
+	for s := range m.states.m {
+		stateIds = append(stateIds, s)
+	}
+
+	sort.Strings(stateIds)
+	resp.States = stateIds
+	return resp
+}
+
+func (m *InMemStoreSingle) DeleteState(req providers.DeleteStateRequest) providers.DeleteStateResponse {
+	m.states.Lock()
+	defer m.states.Unlock()
+
+	resp := providers.DeleteStateResponse{}
+
+	if req.StateId == backend.DefaultStateName || req.StateId == "" {
+		resp.Diagnostics = resp.Diagnostics.Append(fmt.Errorf("can't delete default state"))
+		return resp
+	}
+
+	delete(m.states.m, req.StateId)
+	return resp
+}
+
+type stateMap struct {
+	sync.Mutex
+	m map[string][]byte // key=state id, value=state
+}
+
+type lockMap struct {
+	sync.Mutex
+	m map[string]string // key=state id, value=lock_id
+}
+
+func (l *lockMap) lock(name string, lockId string) (string, error) {
+	l.Lock()
+	defer l.Unlock()
+
+	lock, ok := l.m[name]
+	if ok {
+		// Error; lock already exists for that state id
+		return "", fmt.Errorf("state %q is already locked with lock id %q", name, lock)
+	}
+
+	if l.m == nil {
+		l.m = make(map[string]string, 1)
+	}
+
+	l.m[name] = lockId
+
+	return lockId, nil
+}
+
+func (l *lockMap) unlock(name, id string) error {
+	l.Lock()
+	defer l.Unlock()
+
+	lockId, ok := l.m[name]
+
+	if !ok {
+		return errors.New("state not locked")
+	}
+
+	if id != lockId {
+		return fmt.Errorf("invalid lock id: state %q was locked with lock id %q, but tried to unlock with lock id %q",
+			name,
+			lockId,
+			id,
+		)
+	}
+
+	delete(l.m, name)
+	return nil
+}

internal/provider-simple-v6/state_store_inmem_test.go

@@ -0,0 +1,90 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package simple
+
+import (
+	"testing"
+
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/terraform/internal/addrs"
+	"github.com/hashicorp/terraform/internal/backend"
+	"github.com/hashicorp/terraform/internal/backend/pluggable"
+	"github.com/hashicorp/terraform/internal/states"
+	"github.com/zclconf/go-cty/cty"
+)
+
+func TestInMemStoreLocked(t *testing.T) {
+	// backend.TestBackendStateLocks assumes the "default" state exists
+	// by default, so we need to make it exist using the method below.
+	provider := ProviderWithDefaultState()
+
+	plug1, err := pluggable.NewPluggable(provider, inMemStoreName)
+	if err != nil {
+		t.Fatal(err)
+	}
+	plug2, err := pluggable.NewPluggable(provider, inMemStoreName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b1 := backend.TestBackendConfig(t, plug1, nil)
+	b2 := backend.TestBackendConfig(t, plug2, nil)
+
+	backend.TestBackendStateLocks(t, b1, b2)
+}
+
+func TestInMemStoreRemoteState(t *testing.T) {
+	provider := Provider()
+
+	plug, err := pluggable.NewPluggable(provider, inMemStoreName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	b := backend.TestBackendConfig(t, plug, hcl.EmptyBody())
+
+	// The "default" state doesn't exist by default
+	// (Note that this depends on the factory method used to get the provider above)
+	stateIds, wDiags := b.Workspaces()
+	if wDiags.HasErrors() {
+		t.Fatal(wDiags.Err())
+	}
+	if len(stateIds) != 0 {
+		t.Fatalf("unexpected response from Workspaces method: %#v", stateIds)
+	}
+
+	// create a new state using this backend
+	newStateId := "foobar"
+	emptyState := states.NewState()
+
+	sMgr, sDiags := b.StateMgr(newStateId)
+	if sDiags.HasErrors() {
+		t.Fatal(sDiags.Err())
+	}
+	if err := sMgr.WriteState(emptyState); err != nil {
+		t.Fatal(err)
+	}
+	if err := sMgr.PersistState(nil); err != nil {
+		t.Fatal(err)
+	}
+
+	// force overwriting the remote state
+	newState := states.NewState()
+	newState.SetOutputValue(
+		addrs.OutputValue{Name: "foo"}.Absolute(addrs.RootModuleInstance),
+		cty.StringVal("bar"),
+		false)
+
+	if err := sMgr.WriteState(newState); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := sMgr.PersistState(nil); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := sMgr.RefreshState(); err != nil {
+		t.Fatal(err)
+	}
+}