feat: count a parent's child workflows, optionally by status

src/main/java/com/airbnb/skipper/WorkflowsService.kt

@@ -45,6 +45,40 @@ open class WorkflowsService
             return workflowStore.countWorkflowsByStatus(io.vavr.collection.List.ofAll(statuses))
         }
 
+        /**
+         * Count the child workflows of a parent workflow, i.e. the workflow instances started with
+         * the given parent workflow ID, in any status.
+         *
+         * This reads the store directly, so the result is not deterministic across replays. Call it
+         * from client code, or from an action, never directly from a workflow method.
+         *
+         * @param parentWorkflowId the ID of the parent workflow.
+         * @return the number of child workflows, or 0 if the parent has none (or does not exist).
+         */
+        fun countChildWorkflows(parentWorkflowId: String): Long = workflowStore.countChildWorkflows(parentWorkflowId, emptySet())
+
+        /**
+         * Count the child workflows of a parent workflow that are in one of the specified statuses.
+         * See [countChildWorkflows] for the replay caveat.
+         *
+         * @param parentWorkflowId the ID of the parent workflow.
+         * @param statusViews the statuses to count. Must not be empty; use the overload without
+         *     statuses to count children in any status.
+         * @return the number of child workflows in the specified statuses.
+         */
+        fun countChildWorkflows(
+            parentWorkflowId: String,
+            statusViews: Set<@JvmSuppressWildcards WorkflowInstanceStatusView>
+        ): Long {
+            // An empty set is rejected rather than read as "any status": a caller whose computed
+            // set ends up empty would otherwise silently get every child counted.
+            require(statusViews.isNotEmpty()) {
+                "statusViews must not be empty; use countChildWorkflows(parentWorkflowId) to count all children."
+            }
+            val statuses = statusViews.map { v -> WorkflowInstance.Status.fromView(v) }.toSet()
+            return workflowStore.countChildWorkflows(parentWorkflowId, statuses)
+        }
+
         /**
          * Find all the workflow instances that have exhausted their retries and are waiting for
          * manual intervention.

src/main/java/com/airbnb/skipper/internal/storage/WorkflowStore.java

@@ -9,6 +9,7 @@
 import io.vavr.Tuple3;
 import io.vavr.collection.List;
 import io.vavr.control.Option;
+import java.util.Set;
 
 public interface WorkflowStore {
   /**
@@ -49,6 +50,27 @@ public interface WorkflowStore {
    */
   long countWorkflowsByStatus(List<WorkflowInstance.Status> statuses);
 
+  /**
+   * Count the child workflows of a parent workflow, i.e. the workflows created with the given
+   * {@code parentWorkflowId}.
+   *
+   * @param parentWorkflowId The ID of the parent workflow.
+   * @param statuses Only count children in one of these statuses. An empty set counts children in
+   *     any status.
+   * @return The number of matching child workflows, or 0 if the parent has none (or does not
+   *     exist).
+   */
+  default long countChildWorkflows(String parentWorkflowId, Set<WorkflowInstance.Status> statuses) {
+    if (parentWorkflowId == null) {
+      throw new java.lang.NullPointerException("parentWorkflowId is marked non-null but is null");
+    }
+    if (statuses == null) {
+      throw new java.lang.NullPointerException("statuses is marked non-null but is null");
+    }
+    throw new UnsupportedOperationException(
+        "countChildWorkflows is not supported by this WorkflowStore implementation");
+  }
+
   /**
    * Find workflows with exhausted retries.
    *

src/main/java/com/airbnb/skipper/internal/storage/mysql/MySqlWorkflowStore.kt

@@ -383,6 +383,51 @@ class MySqlWorkflowStore
                 .toLong()
         }
 
+        override fun countChildWorkflows(
+            parentWorkflowId: String,
+            statuses: Set<WorkflowInstance.Status>,
+        ): Long {
+            // Served by the (owner, parent_workflow_id, status) index.
+            val statusClause =
+                if (statuses.isEmpty()) {
+                    ""
+                } else {
+                    " AND status IN (" + statuses.joinToString(",") { "?" } + ")"
+                }
+            val sql =
+                "SELECT COUNT(*) FROM $workflowInstancesTable WHERE owner = ? AND parent_workflow_id = ?" +
+                    statusClause
+            return transactionManager.execute { conn ->
+                try {
+                    conn.prepareStatement(sql).use { ps ->
+                        metrics
+                            .timer(METRICS_COMPONENT, "countChildWorkflows.latency")
+                            .time()
+                            .use { _ ->
+                                ps.setString(1, owner)
+                                ps.setString(2, parentWorkflowId)
+                                statuses.forEachIndexed { i, status ->
+                                    ps.setString(i + 3, status.name)
+                                }
+                                val rs = ps.executeQuery()
+                                // COUNT(*) without GROUP BY always yields exactly one row.
+                                rs.next()
+                                rs.getLong(1)
+                            }
+                    }
+                } catch (e: SQLException) {
+                    metrics
+                        .counter(
+                            ImmutableMap.of("source", "countChildWorkflows"),
+                            METRICS_COMPONENT,
+                            "errors",
+                        )
+                        .inc()
+                    throw InternalError(e)
+                }
+            }
+        }
+
         override fun findWorkflowsWithExhaustedRetries(
             limit: Int,
             sortField: WorkflowSortField?,

src/main/java/com/airbnb/skipper/internal/storage/sqlite/SqliteWorkflowStore.kt

@@ -387,6 +387,48 @@ class SqliteWorkflowStore
                 .toLong()
         }
 
+        override fun countChildWorkflows(
+            parentWorkflowId: String,
+            statuses: Set<WorkflowInstance.Status>
+        ): Long {
+            // Served by the (owner, parent_workflow_id, status) index.
+            val statusClause =
+                if (statuses.isEmpty()) {
+                    ""
+                } else {
+                    " AND status IN (" + statuses.joinToString(",") { "?" } + ")"
+                }
+            val sql =
+                "SELECT COUNT(*) FROM $workflowInstancesTable WHERE owner = ? AND parent_workflow_id = ?" +
+                    statusClause
+            return transactionManager.execute { conn ->
+                try {
+                    conn.prepareStatement(sql).use { ps ->
+                        metrics.timer(METRICS_COMPONENT, "countChildWorkflows.latency").time().use { _ ->
+                            ps.setString(1, owner)
+                            ps.setString(2, parentWorkflowId)
+                            statuses.forEachIndexed { i, status ->
+                                ps.setString(i + 3, status.name)
+                            }
+                            val rs = ps.executeQuery()
+                            // COUNT(*) without GROUP BY always yields exactly one row.
+                            rs.next()
+                            rs.getLong(1)
+                        }
+                    }
+                } catch (e: SQLException) {
+                    metrics
+                        .counter(
+                            ImmutableMap.of("source", "countChildWorkflows"),
+                            METRICS_COMPONENT,
+                            "errors"
+                        )
+                        .inc()
+                    throw InternalError(e)
+                }
+            }
+        }
+
         override fun findWorkflowsWithExhaustedRetries(
             limit: Int,
             sortField: WorkflowSortField?,

src/main/resources/db/migration/V7__add_parent_workflow_id_index.sql

@@ -0,0 +1,5 @@
+ALTER TABLE `${tablePrefix}workflow_instances`
+  ADD INDEX `${tablePrefix}workflow_instance_index_by_parent` (`owner`, `parent_workflow_id`, `status`);
+-- DDL ENDS HERE. STOP COPYING.
+-- Update the schema migrations table to record this migration
+INSERT INTO `${tablePrefix}schema_migrations` (`migration_id`, `created_at`) VALUES (7, NOW());

src/main/resources/db/sqlite/V3__add_parent_workflow_id_index.sql

@@ -0,0 +1,7 @@
+-- SQLite-dialect counterpart of db/migration/V7__add_parent_workflow_id_index.sql.
+--
+-- Lets child workflows be counted per parent (optionally by status) without scanning every
+-- instance of the owner.
+
+CREATE INDEX IF NOT EXISTS ${tablePrefix}workflow_instance_index_by_parent
+  ON ${tablePrefix}workflow_instances (owner, parent_workflow_id, status);

src/test/java/com/airbnb/skipper/WorkflowsServiceTest.kt

@@ -1,5 +1,6 @@
 package com.airbnb.skipper
 
+import com.airbnb.skipper.api.WorkflowInstanceStatusView
 import com.airbnb.skipper.internal.CheckpointTag
 import com.airbnb.skipper.internal.SkipperEngine
 import com.airbnb.skipper.internal.TestUtils.EXTRA_REQUEST_DATA
@@ -46,6 +47,39 @@ class WorkflowsServiceTest {
         workflowsService = WorkflowsService(workflowStore, skipperEngine, scheduler)
     }
 
+    @Test
+    fun countChildWorkflows_withoutStatuses_shouldCountChildrenInAnyStatus() {
+        whenever(workflowStore.countChildWorkflows("parent-1", emptySet())).thenReturn(5L)
+
+        assertEquals(5L, workflowsService.countChildWorkflows("parent-1"))
+    }
+
+    @Test
+    fun countChildWorkflows_withStatuses_shouldPassThemToTheStore() {
+        whenever(
+            workflowStore.countChildWorkflows(
+                "parent-1",
+                setOf(WorkflowInstance.Status.RUNNING, WorkflowInstance.Status.WAITING)
+            )
+        ).thenReturn(2L)
+
+        val count =
+            workflowsService.countChildWorkflows(
+                "parent-1",
+                setOf(WorkflowInstanceStatusView.RUNNING, WorkflowInstanceStatusView.WAITING)
+            )
+
+        assertEquals(2L, count)
+    }
+
+    @Test
+    fun countChildWorkflows_withEmptyStatuses_shouldRejectInsteadOfCountingEverything() {
+        assertThrows(IllegalArgumentException::class.java) {
+            workflowsService.countChildWorkflows("parent-1", emptySet())
+        }
+        verify(workflowStore, never()).countChildWorkflows(any(), any())
+    }
+
     @Test
     fun resetWorkflowsFromError_shouldSuccessfullyResetSingleWorkflow() {
         // Arrange

src/test/java/com/airbnb/skipper/integtest/BaseWorkflowIntegTest.kt

@@ -993,6 +993,23 @@ abstract class BaseWorkflowIntegTest {
             assertEquals("parent-id", workflow.getWorkflowInstanceView().parentWorkflowId)
         }
 
+        @Test
+        fun testCountChildWorkflowsThroughWorkflowsService() {
+            val parentId = "$workflowId-parent"
+            for (i in 1..2) {
+                workflowFactory.builder(SampleWorkflow::class.java, "$workflowId-child-$i")
+                    .parentWorkflowId(parentId).build()
+                    .completableFutureWorkflowMethod().join()
+            }
+            helper.waitForExecutionFlowToFinish()
+            val workflowsService = deps.runtime.workflowsService.get()
+
+            assertEquals(2, workflowsService.countChildWorkflows(parentId))
+            assertEquals(2, workflowsService.countChildWorkflows(parentId, setOf(WorkflowInstanceStatusView.COMPLETED)))
+            assertEquals(0, workflowsService.countChildWorkflows(parentId, setOf(WorkflowInstanceStatusView.RUNNING)))
+            assertEquals(0, workflowsService.countChildWorkflows(workflowId))
+        }
+
         @Test
         fun testNonBlockingInvokeCompletableFutureWorkflowMethod() {
             val workflow = workflowFactory<SampleWorkflow>(workflowId)

src/test/java/com/airbnb/skipper/internal/storage/BaseWorkflowStoreTest.kt

@@ -10,6 +10,7 @@ import com.airbnb.skipper.Workflow
 import com.airbnb.skipper.WorkflowCallbackHandler
 import com.airbnb.skipper.WorkflowInstance
 import com.airbnb.skipper.internal.CheckpointTag
+import com.airbnb.skipper.internal.InternalError
 import com.airbnb.skipper.internal.TestUtils.EXTRA_REQUEST_DATA
 import com.airbnb.skipper.internal.api.ActionCheckpoint
 import com.airbnb.skipper.internal.api.PersistedSignal
@@ -61,6 +62,12 @@ abstract class BaseWorkflowStoreTest {
 
     protected abstract fun createWorkflowStore(owner: String): WorkflowStore
 
+    /** A store over tables named with [tablePrefix], which need not exist. */
+    protected abstract fun createWorkflowStore(
+        owner: String,
+        tablePrefix: String
+    ): WorkflowStore
+
     @Test
     @Throws(Exception::class)
     fun testCreateWorkflow() {
@@ -665,6 +672,76 @@ abstract class BaseWorkflowStoreTest {
         assertEquals(3, runningAndCompleted)
     }
 
+    @Test
+    fun countChildWorkflows_shouldCountOnlyTheParentsChildrenInTheGivenStatuses() {
+        createChildWorkflow("child-1", "parent-a")
+        createChildWorkflow("child-2", "parent-a")
+        val completedChild = createChildWorkflow("child-3", "parent-a")
+        workflowStore().updateWorkflow(
+            WorkflowUpdateRequest.builder()
+                .workflowInstance(completedChild)
+                .newStatus(WorkflowInstance.Status.COMPLETED)
+                .result(Either.right("done"))
+                .resultIsAsync(false)
+                .build()
+        )
+        // Neither a sibling parent's child, a top-level workflow, nor a child of the same parent ID
+        // under another owner is counted.
+        createChildWorkflow("child-4", "parent-b")
+        createInstanceForTimer("top-level")
+        createChildWorkflow("child-5", "parent-a", createWorkflowStore("test-owner2"))
+
+        assertEquals(3, workflowStore().countChildWorkflows("parent-a", emptySet()))
+        assertEquals(
+            2,
+            workflowStore().countChildWorkflows("parent-a", setOf(WorkflowInstance.Status.RUNNING))
+        )
+        assertEquals(
+            1,
+            workflowStore().countChildWorkflows("parent-a", setOf(WorkflowInstance.Status.COMPLETED))
+        )
+        assertEquals(
+            3,
+            workflowStore()
+                .countChildWorkflows(
+                    "parent-a",
+                    setOf(WorkflowInstance.Status.RUNNING, WorkflowInstance.Status.COMPLETED)
+                )
+        )
+        assertEquals(
+            0,
+            workflowStore().countChildWorkflows("parent-a", setOf(WorkflowInstance.Status.ERROR))
+        )
+        assertEquals(1, workflowStore().countChildWorkflows("parent-b", emptySet()))
+        assertEquals(0, workflowStore().countChildWorkflows("no-such-parent", emptySet()))
+    }
+
+    @Test
+    fun countChildWorkflows_shouldWrapDatabaseErrorsInInternalError() {
+        val storeWithoutTables = createWorkflowStore("test-owner", "missing_tables_")
+
+        assertThrows(InternalError::class.java) {
+            storeWithoutTables.countChildWorkflows("parent-a", setOf(WorkflowInstance.Status.RUNNING))
+        }
+    }
+
+    private fun createChildWorkflow(
+        workflowId: String,
+        parentWorkflowId: String,
+        store: WorkflowStore = workflowStore()
+    ): WorkflowInstance =
+        store.createWorkflow(
+            WorkflowCreationRequest.builder()
+                .workflowClass(Workflow::class.java)
+                .workflowId(workflowId)
+                .workflowMethod("execute")
+                .input("test")
+                .requestContext(REQUEST_CONTEXT)
+                .extraRequestData(EXTRA_REQUEST_DATA)
+                .parentWorkflowId(parentWorkflowId)
+                .build()
+        )
+
     @Test
     @Throws(InterruptedException::class)
     fun testFindWorkflowsWithExhaustedRetries() {

src/test/java/com/airbnb/skipper/internal/storage/WorkflowStoreTest.kt

@@ -0,0 +1,29 @@
+package com.airbnb.skipper.internal.storage
+
+import com.airbnb.skipper.WorkflowInstance
+import org.junit.jupiter.api.Assertions.assertThrows
+import org.junit.jupiter.api.Test
+import org.mockito.Mockito
+import org.mockito.kotlin.mock
+
+/** Covers the [WorkflowStore] default methods seen by implementations that don't override them. */
+class WorkflowStoreTest {
+    private val store = mock<WorkflowStore>(defaultAnswer = Mockito.CALLS_REAL_METHODS)
+
+    @Test
+    fun countChildWorkflows_byDefault_shouldBeUnsupported() {
+        assertThrows(UnsupportedOperationException::class.java) {
+            store.countChildWorkflows("parent-1", setOf(WorkflowInstance.Status.RUNNING))
+        }
+    }
+
+    @Test
+    fun countChildWorkflows_byDefault_shouldRejectNullArguments() {
+        assertThrows(NullPointerException::class.java) {
+            store.countChildWorkflows(null, emptySet())
+        }
+        assertThrows(NullPointerException::class.java) {
+            store.countChildWorkflows("parent-1", null)
+        }
+    }
+}