Fix nil entry in state resource instance map from state hook (#3478)

Signed-off-by: Christian Mesh <christianmesh1@gmail.com>

internal/states/module.go

@@ -90,7 +90,11 @@ func (ms *Module) SetResourceInstance(addr addrs.ResourceInstance, inst *Resourc
 		ms.SetResourceProvider(addr.Resource, provider)
 		rs = ms.Resource(addr.Resource)
 	}
-	rs.Instances[addr.Key] = inst
+	if inst != nil {
+		rs.Instances[addr.Key] = inst
+	} else {
+		delete(rs.Instances, addr.Key)
+	}
 }
 
 // SetResourceInstanceCurrent saves the given instance object as the current

internal/tofu/update_state_hook.go

@@ -21,6 +21,9 @@ func updateStateHook(evalCtx EvalContext, addr addrs.AbsResourceInstance) error
 				// See the documentation of ResourceProvider for more details
 				s.RemoveResource(addr.ContainingResource())
 			} else {
+				// The individual instance may be nil, but that can happen when destroying
+				// some but not all instances of a resource (or when that is in-progress).
+				// SetResourceInstance handles that nil correctly and updates the state accordingly.
 				s.SetResourceInstance(addr, evalCtx.State().ResourceInstance(addr), *provider)
 			}
 		})

internal/tofu/update_state_hook_test.go

@@ -49,3 +49,60 @@ func TestUpdateStateHook(t *testing.T) {
 		t.Fatalf("wrong state passed to hook: %s", spew.Sdump(target))
 	}
 }
+
+func TestUpdateStateHookRemoved(t *testing.T) {
+	mockHook := new(MockHook)
+
+	resAddr0 := addrs.Resource{
+		Mode: addrs.ManagedResourceMode,
+		Type: "foo",
+		Name: "bar",
+	}.Instance(addrs.IntKey(0)).Absolute(addrs.RootModuleInstance)
+	resAddr1 := addrs.Resource{
+		Mode: addrs.ManagedResourceMode,
+		Type: "foo",
+		Name: "bar",
+	}.Instance(addrs.IntKey(1)).Absolute(addrs.RootModuleInstance)
+
+	providerAddr, _ := addrs.ParseAbsProviderConfigStr(`provider["registry.opentofu.org/org/foo"]`)
+	resData := &states.ResourceInstanceObjectSrc{
+		SchemaVersion: 42,
+	}
+
+	state := states.NewState()
+	state.Module(addrs.RootModuleInstance).SetResourceInstanceCurrent(resAddr0.Resource, resData, providerAddr, addrs.NoKey)
+
+	ctx := new(MockEvalContext)
+	ctx.HookHook = mockHook
+	ctx.StateState = state.SyncWrapper()
+
+	target := states.NewState()
+
+	// Write resource instance 0
+	if err := updateStateHook(ctx, resAddr0); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Flush
+	if !mockHook.PostStateUpdateCalled {
+		t.Fatal("should call PostStateUpdate")
+	}
+	mockHook.PostStateUpdateCalled = false
+	mockHook.PostStateUpdateFn(target.SyncWrapper())
+
+	// Will remove the entry if it exists
+	if err := updateStateHook(ctx, resAddr1); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Flush
+	if !mockHook.PostStateUpdateCalled {
+		t.Fatal("should call PostStateUpdate")
+	}
+	mockHook.PostStateUpdateFn(target.SyncWrapper())
+
+	// Comparison
+	if !state.ManagedResourcesEqual(target) {
+		t.Fatalf("wrong state passed to hook: %s \nExpected:\n %s", spew.Sdump(target), spew.Sdump(state))
+	}
+}