From c7b6e31b8b3258f7276519f88970e9de7bb0253a Mon Sep 17 00:00:00 2001
From: Rick Ross <rick.ross@temporal.io>
Date: Fri, 16 Aug 2024 16:45:22 -0400
Subject: [PATCH] Added logic for detecting replay; fixed test

---
 .../MyCounterInterceptor.cs                   | 32 ++++++++++++++-----
 tests/CounterInterceptor/MyWorkflowTests.cs   |  2 +-
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/CounterInterceptor/MyCounterInterceptor.cs b/src/CounterInterceptor/MyCounterInterceptor.cs
index ae86ce6..22fb9fd 100644
--- a/src/CounterInterceptor/MyCounterInterceptor.cs
+++ b/src/CounterInterceptor/MyCounterInterceptor.cs
@@ -75,22 +75,34 @@ public override void Init(WorkflowOutboundInterceptor outbound) =>
 
         public override Task<object?> ExecuteWorkflowAsync(ExecuteWorkflowInput input)
         {
-            var id = Workflow.Info.WorkflowId;
-            root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowReplays));
+            // Count only if we're not replaying
+            if (!Workflow.Unsafe.IsReplaying)
+            {
+                var id = Workflow.Info.WorkflowId;
+                root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowReplays));
+            }
             return base.ExecuteWorkflowAsync(input);
         }
 
         public override Task HandleSignalAsync(HandleSignalInput input)
         {
-            var id = Workflow.Info.WorkflowId;
-            root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowSignals));
+            // Count only if we're not replaying
+            if (!Workflow.Unsafe.IsReplaying)
+            {
+                var id = Workflow.Info.WorkflowId;
+                root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowSignals));
+            }
             return base.HandleSignalAsync(input);
         }
 
         public override object? HandleQuery(HandleQueryInput input)
         {
-            var id = Workflow.Info.WorkflowId;
-            root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowQueries));
+            // Count only if we're not replaying
+            if (!Workflow.Unsafe.IsReplaying)
+            {
+                var id = Workflow.Info.WorkflowId;
+                root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowQueries));
+            }
             return base.HandleQuery(input);
         }
     }
@@ -105,8 +117,12 @@ internal WorkflowOutbound(MyCounterInterceptor root, WorkflowOutboundInterceptor
         public override Task<ChildWorkflowHandle<TWorkflow, TResult>> StartChildWorkflowAsync<TWorkflow, TResult>(
             StartChildWorkflowInput input)
         {
-            var id = Workflow.Info.WorkflowId;
-            root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowChildExecutions));
+            // Count only if we're not replaying
+            if (!Workflow.Unsafe.IsReplaying)
+            {
+                var id = Workflow.Info.WorkflowId;
+                root.Increment(id, c => Interlocked.Increment(ref root.Counts[id].WorkflowChildExecutions));
+            }
             return base.StartChildWorkflowAsync<TWorkflow, TResult>(input);
         }
     }
diff --git a/tests/CounterInterceptor/MyWorkflowTests.cs b/tests/CounterInterceptor/MyWorkflowTests.cs
index 47523bd..c99f36c 100644
--- a/tests/CounterInterceptor/MyWorkflowTests.cs
+++ b/tests/CounterInterceptor/MyWorkflowTests.cs
@@ -63,7 +63,7 @@ await worker.ExecuteAsync(async () =>
             Assert.Equal(1U, counterInterceptor.Counts[parentWorkflowId].WorkflowChildExecutions);
             Assert.Equal(0U, counterInterceptor.Counts[parentWorkflowId].WorkflowActivityExecutions);
             Assert.Equal(2U, counterInterceptor.Counts[parentWorkflowId].WorkflowSignals);
-            Assert.Equal(2U, counterInterceptor.Counts[parentWorkflowId].WorkflowQueries);
+            Assert.Equal(0U, counterInterceptor.Counts[parentWorkflowId].WorkflowQueries);
 
             // Validate the worker counters have the correct numbers for the child
             Assert.Equal(1U, counterInterceptor.Counts[childWorkflowId].WorkflowReplays);