Reduce calls to state DeepCopy() by half during apply (#3011)

Signed-off-by: Christian Mesh <christianmesh1@gmail.com>
Co-authored-by: Andrei Ciobanu <andrei.ciobanu@opentofu.org>
Co-authored-by: Martin Atkins <mart@degeneration.co.uk>

internal/backend/local/backend_apply.go

@@ -108,6 +108,12 @@ func (b *Local) opApply(
 		op.ReportResult(runningOp, diags)
 		return
 	}
+	// Setup the state hook with the current input state that will be modified throughout the run
+	if lr.InputState == nil {
+		stateHook.workingCopy = states.NewState().SyncWrapper()
+	} else {
+		stateHook.workingCopy = lr.InputState.DeepCopy().SyncWrapper()
+	}
 	// stateHook uses schemas for when it periodically persists state to the
 	// persistent storage backend.
 	stateHook.Schemas = schemas

internal/backend/local/hook_state.go

@@ -22,7 +22,8 @@ type StateHook struct {
 	tofu.NilHook
 	sync.Mutex
 
-	StateMgr statemgr.Writer
+	StateMgr    statemgr.Writer
+	workingCopy *states.SyncState
 
 	// If PersistInterval is nonzero then for any new state update after
 	// the duration has elapsed we'll try to persist a state snapshot
@@ -64,7 +65,7 @@ type IntermediateStatePersistInfo struct {
 
 var _ tofu.Hook = (*StateHook)(nil)
 
-func (h *StateHook) PostStateUpdate(new *states.State) (tofu.HookAction, error) {
+func (h *StateHook) PostStateUpdate(mutate func(*states.SyncState)) (tofu.HookAction, error) {
 	h.Lock()
 	defer h.Unlock()
 
@@ -77,9 +78,14 @@ func (h *StateHook) PostStateUpdate(new *states.State) (tofu.HookAction, error)
 	}
 
 	if h.StateMgr != nil {
-		if err := h.StateMgr.WriteState(new); err != nil {
+		mutate(h.workingCopy)
+
+		if err := h.StateMgr.WriteState(h.workingCopy.Lock()); err != nil {
+			h.workingCopy.Unlock()
 			return tofu.HookActionHalt, err
 		}
+		h.workingCopy.Unlock()
+
 		if mgrPersist, ok := h.StateMgr.(statemgr.Persister); ok && h.PersistInterval != 0 && h.Schemas != nil {
 			if h.shouldPersist() {
 				err := mgrPersist.PersistState(context.TODO(), h.Schemas)

internal/backend/local/hook_state_test.go

@@ -12,28 +12,42 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/opentofu/opentofu/internal/addrs"
 	"github.com/opentofu/opentofu/internal/states"
 	"github.com/opentofu/opentofu/internal/states/statemgr"
 	"github.com/opentofu/opentofu/internal/tofu"
+	"github.com/zclconf/go-cty/cty"
 )
 
 func TestStateHook_impl(t *testing.T) {
 	var _ tofu.Hook = new(StateHook)
 }
 
+func stateHookExpected() *states.State {
+	expected := states.NewState()
+	expected.RootModule().SetOutputValue("sensitive_output", cty.StringVal("it's a secret"), true, "")
+	return expected
+}
+func stateHookMutator(state *states.SyncState) {
+	state.SetOutputValue(addrs.AbsOutputValue{OutputValue: addrs.OutputValue{Name: "sensitive_output"}}, cty.StringVal("it's a secret"), true, "")
+}
+
 func TestStateHook(t *testing.T) {
 	is := statemgr.NewTransientInMemory(nil)
-	var hook tofu.Hook = &StateHook{StateMgr: is}
+	var hook tofu.Hook = &StateHook{
+		StateMgr:    is,
+		workingCopy: states.NewState().SyncWrapper(),
+	}
 
-	s := statemgr.TestFullInitialState()
-	action, err := hook.PostStateUpdate(s)
+	action, err := hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
 	if action != tofu.HookActionContinue {
 		t.Fatalf("bad: %v", action)
 	}
-	if !is.State().Equal(s) {
+
+	if !is.State().Equal(stateHookExpected()) {
 		t.Fatalf("bad state: %#v", is.State())
 	}
 }
@@ -47,10 +61,11 @@ func TestStateHookStopping(t *testing.T) {
 		intermediatePersist: IntermediateStatePersistInfo{
 			LastPersist: time.Now(),
 		},
+		workingCopy: states.NewState().SyncWrapper(),
 	}
 
-	s := statemgr.TestFullInitialState()
-	action, err := hook.PostStateUpdate(s)
+	s := stateHookExpected()
+	action, err := hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -67,7 +82,7 @@ func TestStateHookStopping(t *testing.T) {
 	// We'll now force lastPersist to be long enough ago that persisting
 	// should be due on the next call.
 	hook.intermediatePersist.LastPersist = time.Now().Add(-5 * time.Hour)
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -77,7 +92,7 @@ func TestStateHookStopping(t *testing.T) {
 	if is.Persisted == nil || !is.Persisted.Equal(s) {
 		t.Fatalf("mismatching state persisted")
 	}
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -115,7 +130,7 @@ func TestStateHookStopping(t *testing.T) {
 	}
 
 	is.Persisted = nil
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -123,7 +138,7 @@ func TestStateHookStopping(t *testing.T) {
 		t.Fatalf("mismatching state persisted")
 	}
 	is.Persisted = nil
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -158,10 +173,11 @@ func TestStateHookCustomPersistRule(t *testing.T) {
 		intermediatePersist: IntermediateStatePersistInfo{
 			LastPersist: time.Now(),
 		},
+		workingCopy: states.NewState().SyncWrapper(),
 	}
 
-	s := statemgr.TestFullInitialState()
-	action, err := hook.PostStateUpdate(s)
+	s := stateHookExpected()
+	action, err := hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -178,7 +194,7 @@ func TestStateHookCustomPersistRule(t *testing.T) {
 	// We'll now force lastPersist to be long enough ago that persisting
 	// should be due on the next call.
 	hook.intermediatePersist.LastPersist = time.Now().Add(-5 * time.Hour)
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -188,7 +204,7 @@ func TestStateHookCustomPersistRule(t *testing.T) {
 	if is.Persisted != nil {
 		t.Fatalf("has a persisted state, but shouldn't")
 	}
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -231,7 +247,7 @@ func TestStateHookCustomPersistRule(t *testing.T) {
 	}
 
 	is.Persisted = nil
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}
@@ -239,7 +255,7 @@ func TestStateHookCustomPersistRule(t *testing.T) {
 		t.Fatalf("mismatching state persisted")
 	}
 	is.Persisted = nil
-	_, err = hook.PostStateUpdate(s)
+	_, err = hook.PostStateUpdate(stateHookMutator)
 	if err != nil {
 		t.Fatalf("unexpected error from PostStateUpdate: %s", err)
 	}

internal/states/module.go

@@ -79,6 +79,20 @@ func (ms *Module) RemoveResource(addr addrs.Resource) {
 	delete(ms.Resources, addr.String())
 }
 
+// SetResourceInstance saves the given full resource instance data within the
+// specified resource.  This both ensures that the resource exists and overwrites
+// any existing value for the resource instance.
+// Any value for the instance data should not be modified after being passed to
+// this function (deepcopy'd ahead of time).
+func (ms *Module) SetResourceInstance(addr addrs.ResourceInstance, inst *ResourceInstance, provider addrs.AbsProviderConfig) {
+	rs := ms.Resource(addr.Resource)
+	if rs == nil {
+		ms.SetResourceProvider(addr.Resource, provider)
+		rs = ms.Resource(addr.Resource)
+	}
+	rs.Instances[addr.Key] = inst
+}
+
 // SetResourceInstanceCurrent saves the given instance object as the current
 // generation of the resource instance with the given address, simultaneously
 // updating the recorded provider configuration address and dependencies.

internal/states/state.go

@@ -285,6 +285,25 @@ func (s *State) AllResourceInstanceObjectAddrs() []struct {
 	return ret
 }
 
+// ResourceProvider returns the provider required by the resource if it exists, or
+// null if no such resource exists
+func (s *State) ResourceProvider(addr addrs.AbsResource) *addrs.AbsProviderConfig {
+	if s == nil {
+		panic("State.ResourceProvider on nil *State")
+	}
+	ms := s.Module(addr.Module)
+	if ms == nil {
+		return nil
+	}
+	rs := ms.Resource(addr.Resource)
+	if rs == nil {
+		return nil
+	}
+	// Intentionally copy the struct
+	pc := rs.ProviderConfig
+	return &pc
+}
+
 // ResourceInstance returns the state for the resource instance with the given
 // address, or nil if no such resource is tracked in the state.
 func (s *State) ResourceInstance(addr addrs.AbsResourceInstance) *ResourceInstance {

internal/states/sync.go

@@ -171,6 +171,15 @@ func (s *SyncState) Resource(addr addrs.AbsResource) *Resource {
 	return ret
 }
 
+// ResourceProvider returns the provider required by the resource if it exists, or
+// null if no such resource exists
+func (s *SyncState) ResourceProvider(addr addrs.AbsResource) *addrs.AbsProviderConfig {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	return s.state.ResourceProvider(addr)
+}
+
 // ResourceInstance returns a snapshot of the state the resource instance with
 // the given address, or nil if no such instance is tracked.
 //
@@ -200,7 +209,7 @@ func (s *SyncState) ResourceInstanceObject(addr addrs.AbsResourceInstance, gen G
 	return inst.GetGeneration(gen).DeepCopy()
 }
 
-// SetResourceMeta updates the resource-level metadata for the resource at
+// SetResourceProvider updates the resource-level metadata for the resource at
 // the given address, creating the containing module state and resource state
 // as a side-effect if not already present.
 func (s *SyncState) SetResourceProvider(addr addrs.AbsResource, provider addrs.AbsProviderConfig) {
@@ -211,6 +220,20 @@ func (s *SyncState) SetResourceProvider(addr addrs.AbsResource, provider addrs.A
 	ms.SetResourceProvider(addr.Resource, provider)
 }
 
+// SetResourceInstance saves the given full resource instance data within the
+// specified resource.  This both ensures that the resource exists and overwrites
+// any existing value for the resource instance.
+// It may create the module and resource as a side effect. The value of inst shall
+// not be modified during the function call.
+func (s *SyncState) SetResourceInstance(addr addrs.AbsResourceInstance, inst *ResourceInstance, provider addrs.AbsProviderConfig) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	ms := s.state.EnsureModule(addr.Module)
+	ms.SetResourceInstance(addr.Resource, inst.DeepCopy(), provider)
+	s.maybePruneModule(addr.Module)
+}
+
 // RemoveResource removes the entire state for the given resource, taking with
 // it any instances associated with the resource. This should generally be
 // called only for resource objects whose instances have all been destroyed,

internal/tofu/context_apply_test.go

@@ -199,7 +199,7 @@ func TestContext2Apply_stop(t *testing.T) {
 		{"PostDiff", "indefinite.foo"},
 		{"PreApply", "indefinite.foo"},
 		{"PostApply", "indefinite.foo"},
-		{"PostStateUpdate", ""}, // State gets updated one more time to include the apply result.
+		{"PostStateUpdate", ""},
 	}
 	// The "Stopping" event gets sent to the hook asynchronously from the others
 	// because it is triggered in the ctx.Stop call above, rather than from

internal/tofu/hook.go

@@ -128,10 +128,15 @@ type Hook interface {
 	// function is called.
 	Stopping()
 
-	// PostStateUpdate is called each time the state is updated. It receives
-	// a deep copy of the state, which it may therefore access freely without
-	// any need for locks to protect from concurrent writes from the caller.
-	PostStateUpdate(new *states.State) (HookAction, error)
+	// PostStateUpdate is called each time a portion of the state is updated. It receives
+	// a function that mutates the syncronized state object.
+	//
+	// This should return an error only if the hook tried to persist the updated state
+	// to the primary state storage location and failed to do so, because this method
+	// is called as a side-effect of updating resource instances in the state and any
+	// error returned will be treated as a failure to update the state, returned to the
+	// end-user.
+	PostStateUpdate(func(*states.SyncState)) (HookAction, error)
 }
 
 // NilHook is a Hook implementation that does nothing. It exists only to
@@ -248,6 +253,6 @@ func (*NilHook) Stopping() {
 	// Does nothing at all by default
 }
 
-func (*NilHook) PostStateUpdate(new *states.State) (HookAction, error) {
+func (*NilHook) PostStateUpdate(func(*states.SyncState)) (HookAction, error) {
 	return HookActionContinue, nil
 }

internal/tofu/hook_mock.go

@@ -182,7 +182,7 @@ type MockHook struct {
 	StoppingCalled bool
 
 	PostStateUpdateCalled bool
-	PostStateUpdateState  *states.State
+	PostStateUpdateFn     func(*states.SyncState)
 	PostStateUpdateReturn HookAction
 	PostStateUpdateError  error
 }
@@ -464,11 +464,11 @@ func (h *MockHook) Stopping() {
 	h.StoppingCalled = true
 }
 
-func (h *MockHook) PostStateUpdate(new *states.State) (HookAction, error) {
+func (h *MockHook) PostStateUpdate(fn func(*states.SyncState)) (HookAction, error) {
 	h.Lock()
 	defer h.Unlock()
 
 	h.PostStateUpdateCalled = true
-	h.PostStateUpdateState = new
+	h.PostStateUpdateFn = fn
 	return h.PostStateUpdateReturn, h.PostStateUpdateError
 }

internal/tofu/hook_stop.go

@@ -130,7 +130,7 @@ func (h *stopHook) PostClose(_ addrs.AbsResourceInstance, _ error) (HookAction,
 
 func (h *stopHook) Stopping() {}
 
-func (h *stopHook) PostStateUpdate(new *states.State) (HookAction, error) {
+func (h *stopHook) PostStateUpdate(func(*states.SyncState)) (HookAction, error) {
 	return h.hook()
 }
 

internal/tofu/hook_test.go

@@ -226,7 +226,7 @@ func (h *testHook) Stopping() {
 	h.Calls = append(h.Calls, &testHookCall{"Stopping", ""})
 }
 
-func (h *testHook) PostStateUpdate(new *states.State) (HookAction, error) {
+func (h *testHook) PostStateUpdate(fn func(*states.SyncState)) (HookAction, error) {
 	h.mu.Lock()
 	defer h.mu.Unlock()
 	h.Calls = append(h.Calls, &testHookCall{"PostStateUpdate", ""})

internal/tofu/node_resource_apply_instance.go

@@ -256,7 +256,7 @@ func (n *NodeApplyableResourceInstance) dataResourceExecute(ctx context.Context,
 
 	diags = diags.Append(n.writeChange(ctx, evalCtx, nil, ""))
 
-	diags = diags.Append(updateStateHook(evalCtx))
+	diags = diags.Append(updateStateHook(evalCtx, n.Addr))
 
 	// Post-conditions might block further progress. We intentionally do this
 	// _after_ writing the state/diff because we want to check against
@@ -439,7 +439,7 @@ func (n *NodeApplyableResourceInstance) managedResourceExecute(ctx context.Conte
 	}
 
 	diags = diags.Append(n.postApplyHook(evalCtx, state, diags.Err()))
-	diags = diags.Append(updateStateHook(evalCtx))
+	diags = diags.Append(updateStateHook(evalCtx, n.Addr))
 
 	// Post-conditions might block further progress. We intentionally do this
 	// _after_ writing the state because we want to check against

internal/tofu/node_resource_deposed.go

@@ -331,7 +331,7 @@ func (n *NodeDestroyDeposedResourceInstanceObject) Execute(ctx context.Context,
 
 	diags = diags.Append(n.postApplyHook(evalCtx, state, diags.Err()))
 
-	return diags.Append(updateStateHook(evalCtx))
+	return diags.Append(updateStateHook(evalCtx, n.Addr))
 }
 
 // GraphNodeDeposer is an optional interface implemented by graph nodes that
@@ -459,5 +459,5 @@ func (n *NodeForgetDeposedResourceInstanceObject) Execute(ctx context.Context, e
 	contextState := evalCtx.State()
 	contextState.ForgetResourceInstanceDeposed(n.Addr, n.DeposedKey)
 
-	return diags.Append(updateStateHook(evalCtx))
+	return diags.Append(updateStateHook(evalCtx, n.Addr))
 }

internal/tofu/node_resource_destroy.go

@@ -269,14 +269,14 @@ func (n *NodeDestroyResourceInstance) managedResourceExecute(ctx context.Context
 
 	// create the err value for postApplyHook
 	diags = diags.Append(n.postApplyHook(evalCtx, state, diags.Err()))
-	diags = diags.Append(updateStateHook(evalCtx))
+	diags = diags.Append(updateStateHook(evalCtx, n.Addr))
 	return diags
 }
 
 func (n *NodeDestroyResourceInstance) dataResourceExecute(_ context.Context, evalCtx EvalContext) (diags tfdiags.Diagnostics) {
 	log.Printf("[TRACE] NodeDestroyResourceInstance: removing state object for %s", n.Addr)
 	evalCtx.State().SetResourceInstanceCurrent(n.Addr, nil, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
-	return diags.Append(updateStateHook(evalCtx))
+	return diags.Append(updateStateHook(evalCtx, n.Addr))
 }
 
 // ephemeralResourceExecute for NodeDestroyResourceInstance is only here to return an error.

internal/tofu/node_resource_forget.go

@@ -75,7 +75,7 @@ func (n *NodeForgetResourceInstance) Execute(ctx context.Context, evalCtx EvalCo
 	contextState := evalCtx.State()
 	contextState.ForgetResourceInstanceAll(n.Addr)
 
-	diags = diags.Append(updateStateHook(evalCtx))
+	diags = diags.Append(updateStateHook(evalCtx, n.Addr))
 
 	return diags
 }

internal/tofu/update_state_hook.go

@@ -5,20 +5,24 @@
 
 package tofu
 
-// updateStateHook calls the PostStateUpdate hook with the current state.
-func updateStateHook(ctx EvalContext) error {
-	// In principle we could grab the lock here just long enough to take a
-	// deep copy and then pass that to our hooks below, but we'll instead
-	// hold the hook for the duration to avoid the potential confusing
-	// situation of us racing to call PostStateUpdate concurrently with
-	// different state snapshots.
-	stateSync := ctx.State()
-	state := stateSync.Lock().DeepCopy()
-	defer stateSync.Unlock()
+import (
+	"github.com/opentofu/opentofu/internal/addrs"
+	"github.com/opentofu/opentofu/internal/states"
+)
 
+// updateState calls the PostStateUpdate hook with the state modification function
+func updateStateHook(evalCtx EvalContext, addr addrs.AbsResourceInstance) error {
 	// Call the hook
-	err := ctx.Hook(func(h Hook) (HookAction, error) {
-		return h.PostStateUpdate(state)
+	return evalCtx.Hook(func(h Hook) (HookAction, error) {
+		return h.PostStateUpdate(func(s *states.SyncState) {
+			provider := evalCtx.State().ResourceProvider(addr.ContainingResource())
+			if provider == nil {
+				// If there is no provider currently defined for the resource, it has been removed
+				// See the documentation of ResourceProvider for more details
+				s.RemoveResource(addr.ContainingResource())
+			} else {
+				s.SetResourceInstance(addr, evalCtx.State().ResourceInstance(addr), *provider)
+			}
+		})
 	})
-	return err
 }

internal/tofu/update_state_hook_test.go

@@ -9,7 +9,6 @@ import (
 	"testing"
 
 	"github.com/davecgh/go-spew/spew"
-	"github.com/zclconf/go-cty/cty"
 
 	"github.com/opentofu/opentofu/internal/addrs"
 	"github.com/opentofu/opentofu/internal/states"
@@ -18,21 +17,35 @@ import (
 func TestUpdateStateHook(t *testing.T) {
 	mockHook := new(MockHook)
 
+	resAddr := addrs.Resource{
+		Mode: addrs.ManagedResourceMode,
+		Type: "foo",
+		Name: "bar",
+	}.Instance(addrs.NoKey).Absolute(addrs.RootModuleInstance)
+	providerAddr, _ := addrs.ParseAbsProviderConfigStr(`provider["registry.opentofu.org/org/foo"]`)
+	resData := &states.ResourceInstanceObjectSrc{
+		SchemaVersion: 42,
+	}
+
 	state := states.NewState()
-	state.Module(addrs.RootModuleInstance).SetLocalValue("foo", cty.StringVal("hello"))
+	state.Module(addrs.RootModuleInstance).SetResourceInstanceCurrent(resAddr.Resource, resData, providerAddr, addrs.NoKey)
 
 	ctx := new(MockEvalContext)
 	ctx.HookHook = mockHook
 	ctx.StateState = state.SyncWrapper()
 
-	if err := updateStateHook(ctx); err != nil {
+	if err := updateStateHook(ctx, resAddr); err != nil {
 		t.Fatalf("err: %s", err)
 	}
 
 	if !mockHook.PostStateUpdateCalled {
 		t.Fatal("should call PostStateUpdate")
 	}
-	if mockHook.PostStateUpdateState.LocalValue(addrs.LocalValue{Name: "foo"}.Absolute(addrs.RootModuleInstance)) != cty.StringVal("hello") {
-		t.Fatalf("wrong state passed to hook: %s", spew.Sdump(mockHook.PostStateUpdateState))
+
+	target := states.NewState()
+	mockHook.PostStateUpdateFn(target.SyncWrapper())
+
+	if !state.ManagedResourcesEqual(target) {
+		t.Fatalf("wrong state passed to hook: %s", spew.Sdump(target))
 	}
 }