From 173f01bf104cc0d826a06925cb8c5f369f8e202e Mon Sep 17 00:00:00 2001
From: Matthieu JACQUET <mjacquet@splio.com>
Date: Sat, 25 Apr 2026 21:54:00 +0200
Subject: [PATCH] fix: preserve batch fast-path when non-add/remove messages
 are present

Partition the batch so AddTagToWorkflow/RemoveTagFromWorkflow always
go through the optimized bulk path (2 storage calls), regardless of
what other message types are in the batch.

Previously, a single ContinueWorkflowTagFanout or other message type
in the batch would force ALL messages into the sequential fallback
path (O(N) storage calls).

Refs #297
---
 .../workflows/tag/WorkflowTagEngine.kt        | 24 +++++++------
 .../workflows/tag/WorkflowTagEngineTests.kt   | 35 +++++++++++++++++++
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/infinitic-workflow-tag/src/main/kotlin/io/infinitic/workflows/tag/WorkflowTagEngine.kt b/infinitic-workflow-tag/src/main/kotlin/io/infinitic/workflows/tag/WorkflowTagEngine.kt
index 286727d76..c74cc7246 100644
--- a/infinitic-workflow-tag/src/main/kotlin/io/infinitic/workflows/tag/WorkflowTagEngine.kt
+++ b/infinitic-workflow-tag/src/main/kotlin/io/infinitic/workflows/tag/WorkflowTagEngine.kt
@@ -81,19 +81,23 @@ class WorkflowTagEngine(
   ) {
     if (messages.isEmpty()) return
 
-    if (messages.all { (message, _) -> message is AddTagToWorkflow || message is RemoveTagFromWorkflow }) {
-      batchProcessAddRemoveOnly(messages)
-      return
+    val (addRemove, others) = messages.partition { (message, _) ->
+      message is AddTagToWorkflow || message is RemoveTagFromWorkflow
     }
 
-    val messagesMap: Map<Pair<WorkflowTag, WorkflowName>, List<Pair<WorkflowTagEngineMessage, MillisInstant>>> =
-        messages.groupBy { it.first.workflowTag to it.first.workflowName }
+    // Process add/remove through the optimized bulk path
+    if (addRemove.isNotEmpty()) {
+      batchProcessAddRemoveOnly(addRemove)
+    }
 
-    coroutineScope {
-      messagesMap
-          .map { (tagAndName, messages) ->
-            launch { batchProcessByTag(storage, producer, messages) }
-          }
+    // Process remaining messages (fanout, customId, getIds, continuations) per tag
+    if (others.isNotEmpty()) {
+      val messagesMap = others.groupBy { it.first.workflowTag to it.first.workflowName }
+      coroutineScope {
+        messagesMap.map { (_, groupedMessages) ->
+          launch { batchProcessByTag(storage, producer, groupedMessages) }
+        }
+      }
     }
   }
 
diff --git a/infinitic-workflow-tag/src/test/kotlin/io/infinitic/workflows/tag/WorkflowTagEngineTests.kt b/infinitic-workflow-tag/src/test/kotlin/io/infinitic/workflows/tag/WorkflowTagEngineTests.kt
index 27ba3d57d..da666c3db 100644
--- a/infinitic-workflow-tag/src/test/kotlin/io/infinitic/workflows/tag/WorkflowTagEngineTests.kt
+++ b/infinitic-workflow-tag/src/test/kotlin/io/infinitic/workflows/tag/WorkflowTagEngineTests.kt
@@ -217,6 +217,41 @@ class WorkflowTagEngineTests : StringSpec(
         (harness.clientMessages.single() as WorkflowIdsByTag).workflowIds shouldBe workflowIds
       }
 
+      "batch with mixed message types should use bulk path for add/remove" {
+        val workflowIds = setOf(WorkflowId(), WorkflowId())
+        val tag = WorkflowTag("test-tag")
+        val name = WorkflowName("test-name")
+        val harness = getHarness(tag, name, workflowIds)
+
+        val addMsg = random<AddTagToWorkflow>(
+            mapOf("workflowTag" to tag, "workflowName" to name),
+        )
+        val removeMsg = random<RemoveTagFromWorkflow>(
+            mapOf("workflowTag" to tag, "workflowName" to name),
+        )
+        val fanoutMsg = random<SendSignalByTag>(
+            mapOf("workflowTag" to tag, "workflowName" to name),
+        )
+
+        val now = MillisInstant.now()
+        harness.engine.batchProcess(
+            listOf(addMsg to now, removeMsg to now, fanoutMsg to now),
+        )
+
+        // Bulk path used for add/remove (getWorkflowIds with Set<Pair> + updateWorkflowIds)
+        coVerify(exactly = 1) {
+          harness.storage.getWorkflowIds(match<Set<Pair<WorkflowTag, WorkflowName>>> { it.isNotEmpty() })
+        }
+        coVerify(exactly = 1) { harness.storage.updateWorkflowIds(any(), any()) }
+
+        // Individual addWorkflowId/removeWorkflowId NOT called (bulk path handles them)
+        coVerify(exactly = 0) { harness.storage.addWorkflowId(any(), any(), any()) }
+        coVerify(exactly = 0) { harness.storage.removeWorkflowId(any(), any(), any()) }
+
+        // The fanout message was also processed (produced a continuation)
+        harness.tagMessages.single().shouldBeInstanceOf<ContinueWorkflowTagFanout>()
+      }
+
       "batch processing should do the same than one by one processing" {
         val n = 3
         val tags = List(n) { RandomString.make(10) }