Improving actor write performance by skipping runtime "exists" checks on state writes

src/Dapr.Actors/Runtime/ActorStateManager.cs

@@ -222,14 +222,13 @@ public async Task SetStateAsync<T>(string stateName, T value, CancellationToken
                 stateMetadata.ChangeKind = StateChangeKind.Update;
             }
         }
-        else if (await this.actor.Host.StateProvider.ContainsStateAsync(this.actorTypeName, this.actor.Id.ToString(), stateName, cancellationToken))
-        {
-            stateChangeTracker.Add(stateName, StateMetadata.Create(value, StateChangeKind.Update));
-        }
         else
         {
-            stateChangeTracker[stateName] = StateMetadata.Create(value, StateChangeKind.Add);
+            // Add and Update are both upserts, while Update ensures a later remove stages a delete.
+            stateChangeTracker[stateName] = StateMetadata.Create(value, StateChangeKind.Update);
         }
+
+        await Task.CompletedTask;
     }
 
     public async Task SetStateAsync<T>(string stateName, T value, TimeSpan ttl, CancellationToken cancellationToken)
@@ -257,14 +256,13 @@ public async Task SetStateAsync<T>(string stateName, T value, TimeSpan ttl, Canc
                 stateMetadata.ChangeKind = StateChangeKind.Update;
             }
         }
-        else if (await this.actor.Host.StateProvider.ContainsStateAsync(this.actorTypeName, this.actor.Id.ToString(), stateName, cancellationToken))
-        {
-            stateChangeTracker.Add(stateName, StateMetadata.Create(value, StateChangeKind.Update, ttl: ttl));
-        }
         else
         {
-            stateChangeTracker[stateName] = StateMetadata.Create(value, StateChangeKind.Add, ttl: ttl);
+            // Add and Update are both upserts, while Update ensures a later remove stages a delete.
+            stateChangeTracker[stateName] = StateMetadata.Create(value, StateChangeKind.Update, ttl: ttl);
         }
+
+        await Task.CompletedTask;
     }
 
     public async Task RemoveStateAsync(string stateName, CancellationToken cancellationToken)

test/Dapr.Actors.Test/ActorStateManagerTest.cs

@@ -180,22 +180,20 @@ public async Task ReentrantSaveRefreshesDefaultTrackerWithoutRoundTrip()
         Assert.Equal("value1", await mngr.GetStateAsync<string>("key1", token));
 
         // A reentrancy-scoped call writes the same key through its own (fresh, empty)
-        // tracker. SetStateAsync's own ContainsStateAsync existence check against the
-        // runtime, to decide Add vs Update, is call #2 - unrelated to defaultTracker and
-        // unavoidable, since the reentrant tracker starts empty and has no local record of
-        // the key yet.
+        // tracker. SetStateAsync stages an upsert without checking the runtime, so this
+        // does not add another GetStateAsync call.
         await mngr.SetStateContext("ctx1");
         await mngr.SetStateAsync("key1", "value2", token);
         await mngr.SaveStateAsync(token);
         await mngr.SetStateContext(null);
 
         // The default tracker must reflect the new value without a further call to the
         // runtime - it already has the value the save above just confirmed was persisted.
-        // Total call count stays at 2; a naive reload-on-evict fix would make this 3.
+        // Total call count stays at 1; a naive reload-on-evict fix would make this 2.
         Assert.Equal("value2", await mngr.GetStateAsync<string>("key1", token));
         interactor.Verify(
             d => d.GetStateAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
-            Times.Exactly(2));
+            Times.Once);
     }
 
     [Fact]
@@ -844,6 +842,89 @@ public async Task SetStateAsync_WithTTL_UpdatesTTLOnCachedNoneEntry()
         Assert.Contains("ttlInSeconds", capturedData);
     }
 
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public async Task SetStateAsync_UnknownKeyDoesNotReadFromStateStore(bool useTtl)
+    {
+        var interactor = new Mock<TestDaprInteractor>();
+        var host = ActorHost.CreateForTest<TestActor>();
+        host.StateProvider = new DaprStateProvider(interactor.Object, new JsonSerializerOptions());
+        var mngr = new ActorStateManager(new TestActor(host));
+        var token = CancellationToken.None;
+        string capturedData = null;
+
+        interactor
+            .Setup(d => d.SaveStateTransactionallyAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+            .Callback<string, string, string, CancellationToken>((_, _, data, _) => capturedData = data)
+            .Returns(Task.CompletedTask);
+
+        if (useTtl)
+        {
+            await mngr.SetStateAsync("k1", "value", TimeSpan.FromMinutes(5), token);
+        }
+        else
+        {
+            await mngr.SetStateAsync("k1", "value", token);
+        }
+
+        Assert.Equal("value", await mngr.GetStateAsync<string>("k1", token));
+        Assert.True(await mngr.ContainsStateAsync("k1", token));
+        Assert.False(await mngr.TryAddStateAsync("k1", "replacement", token));
+
+        await mngr.SaveStateAsync(token);
+
+        Assert.Contains("\"operation\":\"upsert\"", capturedData);
+        Assert.Contains("\"value\":\"value\"", capturedData);
+        Assert.Equal("value", await mngr.GetStateAsync<string>("k1", token));
+        Assert.True(await mngr.ContainsStateAsync("k1", token));
+
+        await mngr.SaveStateAsync(token);
+
+        interactor.Verify(
+            d => d.SaveStateTransactionallyAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
+            Times.Once);
+        interactor.Verify(
+            d => d.GetStateAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
+            Times.Never);
+    }
+
+    [Theory]
+    [InlineData(true)]
+    [InlineData(false)]
+    public async Task SetThenRemoveUnknownKeyStagesDeleteWithoutReadingStateStore(bool useTtl)
+    {
+        var interactor = new Mock<TestDaprInteractor>();
+        var host = ActorHost.CreateForTest<TestActor>();
+        host.StateProvider = new DaprStateProvider(interactor.Object, new JsonSerializerOptions());
+        var mngr = new ActorStateManager(new TestActor(host));
+        var token = CancellationToken.None;
+        string capturedData = null;
+
+        interactor
+            .Setup(d => d.SaveStateTransactionallyAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+            .Callback<string, string, string, CancellationToken>((_, _, data, _) => capturedData = data)
+            .Returns(Task.CompletedTask);
+
+        if (useTtl)
+        {
+            await mngr.SetStateAsync("k1", "value", TimeSpan.FromMinutes(5), token);
+        }
+        else
+        {
+            await mngr.SetStateAsync("k1", "value", token);
+        }
+
+        Assert.True(await mngr.TryRemoveStateAsync("k1", token));
+        await mngr.SaveStateAsync(token);
+
+        Assert.Contains("\"operation\":\"delete\"", capturedData);
+        Assert.DoesNotContain("\"operation\":\"upsert\"", capturedData);
+        interactor.Verify(
+            d => d.GetStateAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()),
+            Times.Never);
+    }
+
     // ----- SetStateContext (reentrancy) -----
 
     [Fact]

test/Dapr.Actors.Test/DaprStateProviderTest.cs

@@ -198,6 +198,38 @@ public async Task SaveStateAsync_Update_EmitsUpsertOperation()
             capturedContent);
     }
 
+    [Fact]
+    public async Task SaveStateAsync_AddAndUpdateEmitEquivalentUpserts()
+    {
+        var interactor = new Mock<TestDaprInteractor>();
+        var provider = new DaprStateProvider(interactor.Object, new JsonSerializerOptions());
+        var token = CancellationToken.None;
+        var capturedContents = new List<string>();
+
+        interactor
+            .Setup(d => d.SaveStateTransactionallyAsync(
+                It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+            .Callback<string, string, string, CancellationToken>((_, _, data, _) => capturedContents.Add(data))
+            .Returns(Task.CompletedTask);
+
+        await provider.SaveStateAsync(
+            "actorType",
+            "actorId",
+            new[] { new ActorStateChange("key1", typeof(string), "value", StateChangeKind.Add, null) },
+            token);
+        await provider.SaveStateAsync(
+            "actorType",
+            "actorId",
+            new[] { new ActorStateChange("key1", typeof(string), "value", StateChangeKind.Update, null) },
+            token);
+
+        Assert.Equal(2, capturedContents.Count);
+        Assert.Equal(capturedContents[0], capturedContents[1]);
+        Assert.Equal(
+            "[{\"operation\":\"upsert\",\"request\":{\"key\":\"key1\",\"value\":\"value\"}}]",
+            capturedContents[0]);
+    }
+
     [Fact]
     public async Task TryLoadStateAsync_ReturnsFalseWhenTTLExpireTimeIsExactlyNow()
     {