with items persistent storage and speed improvements

CHANGELOG.rst

@@ -159,6 +159,30 @@ Added
 * Cherry-pick changes to runners.sh from st2-packages git repo. #6302
   Cherry-picked by @cognifloyd
 
+* Added persistent storage for "with items" (itemized) tasks via the new `TaskItemStateDB` model.
+  Each item's state is stored as a separate document, so individual item states can be read and
+  written without serializing / deserializing the entire task context for every item. This
+  significantly improves the performance of workflows that use `with: items` over large item sets.
+  Run the `st2common/bin/migrations/v3.10/st2-add-task-item-state-collection` migration to create the
+  new collection and its indexes before running itemized tasks.
+  The `st2-purge-task-executions` script and the garbage collector now also delete the associated
+  `TaskItemStateDB` records when their parent task executions are purged.
+
+  A micro benchmark (`st2common/benchmarks/micro/test_with_items_state_storage.py`) compares the old
+  inline `TaskExecutionDB.result["items"]` approach against the new per-item `TaskItemStateDB`
+  records. The old approach re-reads and rewrites the entire (growing) `items` list on every single
+  item update, so its total cost is super-linear in the item count; the new approach reads and
+  writes a single small record per item update. Measured (single threaded, ~2 KB per-item result):
+
+    * 10 items:  ~46 ms   vs ~46 ms  (about even)
+    * 100 items: ~540 ms  vs ~441 ms (~1.2x faster)
+    * 500 items: ~4756 ms vs ~2278 ms (~2.1x faster)
+
+  The gap widens further with more items, larger per-item results, and (in production) concurrent
+  item updates, since separate documents avoid the write-conflict retries the shared document
+  suffered. Note: for small item counts the two are effectively tied; the benefit is at scale.
+  Contributed by @guzzijones12
+
 * Pinned DOCKER_API_VERSION in the circleci build to make sure the docker-cli api version does not exceed what
   `cicrleci docker24 <https://circleci.com/docs/guides/execution-managed/building-docker-images/>`_ supports
 

lockfiles/st2.lock

@@ -3085,7 +3085,7 @@
           "artifacts": [
             {
               "algorithm": "sha256",
-              "hash": "491767e81c1bb11a54fb68d1a24119bdeede593a2beccca5bc09bfed36fdb35c",
+              "hash": "b9feb1769b48102061fe4fc59b2f5ad600bc2ac0b55cf12ef5fe49464ac0d230",
               "url": "git+https://github.com/StackStorm/orquesta.git"
             }
           ],

st2common/benchmarks/micro/test_with_items_state_storage.py

@@ -0,0 +1,171 @@
+# Copyright 2025 The StackStorm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This micro benchmark compares two approaches for persisting individual item states of an
+itemized ("with items") task execution.
+
+Historically all item states were stored inline on ``TaskExecutionDB.result["items"]`` as a
+single list. Every time a single item completed, the *entire* task execution document (with
+the whole ``items`` list) had to be read, mutated and written back. As the number of items
+grows this becomes O(N) work per item update and O(N^2) work for the whole task, and the
+per-write payload keeps growing with the item count.
+
+The new approach stores each item state in its own small ``TaskItemStateDB`` record, so a
+single item update only reads and writes one small document regardless of how many items the
+task has (O(1) per item update, O(N) for the whole task).
+
+This benchmark simulates processing every item of an itemized task once and measures the total
+time for both approaches at a few different item counts so the improvement can be quantified.
+"""
+
+from st2common.util.monkey_patch import monkey_patch
+
+monkey_patch()
+
+import pytest
+
+from st2common.service_setup import db_setup
+from st2common.constants import action as ac_const
+from st2common.models.db.workflow import TaskExecutionDB
+from st2common.models.db.workflow import TaskItemStateDB
+from st2common.persistence.workflow import TaskExecution
+from st2common.persistence.workflow import TaskItemState
+
+
+# A representative per-item result payload (e.g. the stdout/stderr of an action). Action
+# results are frequently non-trivial in size; we use a ~2 KB stdout here. This matters
+# because the inline approach re-serializes and rewrites the whole ``items`` list (every
+# item's result) on every single item update, whereas the per-record approach only ever
+# writes one item's result at a time.
+ITEM_RESULT = {
+    "failed": False,
+    "succeeded": True,
+    "return_code": 0,
+    "stdout": "item processed successfully. " * 64,
+    "stderr": "",
+}
+
+ITEM_COUNTS = [10, 100, 500]
+
+
+def _create_task_execution():
+    task_ex_db = TaskExecutionDB(
+        workflow_execution="000000000000000000000000",
+        task_name="task1",
+        task_id="task1",
+        task_route=0,
+        status=ac_const.LIVEACTION_STATUS_RUNNING,
+        itemized=True,
+    )
+    return task_ex_db
+
+
+def _setup_inline_result(items_count):
+    """Old approach: pre-allocate the inline ``items`` list on the task execution."""
+    task_ex_db = _create_task_execution()
+    task_ex_db.items_count = items_count
+    task_ex_db.result = {"items": [None] * items_count}
+    task_ex_db = TaskExecution.add_or_update(task_ex_db, publish=False)
+    return str(task_ex_db.id)
+
+
+def _setup_item_state_records(items_count):
+    """New approach: one small TaskItemStateDB record per item."""
+    task_ex_db = _create_task_execution()
+    task_ex_db.items_count = items_count
+    task_ex_db.result = {"items_count": items_count}
+    task_ex_db = TaskExecution.add_or_update(task_ex_db, publish=False)
+    task_ex_id = str(task_ex_db.id)
+
+    for item_id in range(items_count):
+        item_state_db = TaskItemStateDB(
+            task_execution=task_ex_id,
+            item_id=item_id,
+            status=ac_const.LIVEACTION_STATUS_REQUESTED,
+            context={},
+        )
+        TaskItemState.insert(item_state_db, publish=False)
+
+    return task_ex_id
+
+
+def _update_inline_result(task_ex_id, items_count):
+    """Record every item state the old way: reload + rewrite the whole task document."""
+    for item_id in range(items_count):
+        task_ex_db = TaskExecution.get_by_id(task_ex_id)
+        task_ex_db.result["items"][item_id] = {
+            "status": ac_const.LIVEACTION_STATUS_SUCCEEDED,
+            "result": ITEM_RESULT,
+        }
+        TaskExecution.add_or_update(task_ex_db, publish=False)
+
+
+def _update_item_state_records(task_ex_id, items_count):
+    """Record every item state the new way: reload + rewrite a single small record."""
+    for item_id in range(items_count):
+        item_state_db = TaskItemState.get_by_task_and_item(task_ex_id, item_id)
+        item_state_db.status = ac_const.LIVEACTION_STATUS_SUCCEEDED
+        item_state_db.result = ITEM_RESULT
+        TaskItemState.add_or_update(item_state_db, publish=False)
+
+
+@pytest.mark.parametrize("items_count", ITEM_COUNTS, ids=[str(c) for c in ITEM_COUNTS])
+@pytest.mark.parametrize(
+    "approach",
+    ["inline_result", "item_state_records"],
+    ids=["inline_result", "item_state_records"],
+)
+@pytest.mark.benchmark(group="with_items_state_storage")
+def test_record_all_item_states(benchmark, approach: str, items_count: int) -> None:
+    db_setup()
+
+    if approach == "inline_result":
+        setup_fn = _setup_inline_result
+        update_fn = _update_inline_result
+    elif approach == "item_state_records":
+        setup_fn = _setup_item_state_records
+        update_fn = _update_item_state_records
+    else:
+        raise ValueError("Invalid approach: %s" % (approach,))
+
+    # The task execution / item state records are created once, outside the timed section.
+    # This is a one-time cost per task; the benchmark focuses on the repeated per-item state
+    # update, which is the hot path that runs once per item (and, in production, concurrently
+    # across items). Recording an item state the old way reads and rewrites the entire task
+    # execution document (the whole growing ``items`` list); the new way reads and rewrites a
+    # single small record.
+    task_ex_id = setup_fn(items_count)
+
+    def run_benchmark():
+        update_fn(task_ex_id, items_count)
+        return task_ex_id
+
+    # Use pedantic mode with a bounded number of rounds: each round performs items_count
+    # database reads and writes, so the default auto-calibrated benchmark would be
+    # prohibitively slow for the larger item counts.
+    benchmark.pedantic(run_benchmark, rounds=3, iterations=1, warmup_rounds=0)
+
+    # Sanity check that all item states were recorded.
+    if approach == "inline_result":
+        task_ex_db = TaskExecution.get_by_id(task_ex_id)
+        recorded = task_ex_db.result["items"]
+        assert len(recorded) == items_count
+        assert all(item is not None for item in recorded)
+    else:
+        item_state_dbs = TaskItemState.query_by_task_execution(task_ex_id)
+        assert len(item_state_dbs) == items_count
+        assert all(
+            isd.status == ac_const.LIVEACTION_STATUS_SUCCEEDED for isd in item_state_dbs
+        )

st2common/bin/migrations/v3.10/BUILD

@@ -0,0 +1,3 @@
+python_sources(
+    sources=["st2*"],
+)

st2common/bin/migrations/v3.10/st2-add-task-item-state-collection

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# Copyright 2025 The StackStorm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Migration which creates the collection and indexes for the new TaskItemStateDB model.
+
+This model backs the persistent storage for "with items" (itemized) tasks. It stores
+each item's state as a separate document so that individual item states can be read and
+written without serializing / deserializing the entire task context for every item.
+
+MongoDB creates collections lazily, so this migration simply ensures the indexes for
+the model exist. That both creates the collection and builds its indexes up front so the
+first itemized task run does not pay that cost. Re-running this migration is safe: it is
+idempotent and will not recreate indexes that already exist.
+"""
+
+import sys
+import traceback
+
+from st2common import config
+from st2common.models.db import db_ensure_indexes
+from st2common.models.db.workflow import TaskItemStateDB
+from st2common.service_setup import db_setup
+from st2common.service_setup import db_teardown
+
+
+def create_task_item_state_collection():
+    print("Ensuring collection and indexes for TaskItemStateDB...")
+    db_ensure_indexes([TaskItemStateDB])
+    print("Done.")
+
+
+def main():
+    config.parse_args()
+    db_setup()
+
+    try:
+        create_task_item_state_collection()
+        exit_code = 0
+    except Exception as e:
+        print("ABORTED: Migration aborted on first failure: %s" % (str(e)))
+        traceback.print_exc()
+        exit_code = 1
+
+    db_teardown()
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

st2common/st2common/garbage_collection/workflows.py

@@ -25,6 +25,7 @@
 from st2common.constants import action as action_constants
 from st2common.persistence.workflow import WorkflowExecution
 from st2common.persistence.workflow import TaskExecution
+from st2common.persistence.workflow import TaskItemState
 
 
 __all__ = ["purge_workflow_executions", "purge_task_executions"]
@@ -132,6 +133,42 @@ def purge_task_executions(logger, timestamp, purge_incomplete=False):
         filters["status"] = {"$in": DONE_STATES}
 
     exec_filters = copy.copy(filters)
+
+    # Collect the ids of the task executions that are about to be deleted so we can also
+    # purge the per-item state records (TaskItemStateDB) that belong to them. These
+    # records are keyed by the string form of the task execution id.
+    to_delete_task_execution_ids = [
+        str(task_ex_db.id)
+        for task_ex_db in TaskExecution.query(
+            only_fields=["id"], no_dereference=True, **exec_filters
+        )
+    ]
+
+    # 1. Delete the per-item state records associated with these task executions.
+    if to_delete_task_execution_ids:
+        try:
+            deleted_item_count = TaskItemState.delete_by_query(
+                task_execution__in=to_delete_task_execution_ids
+            )
+        except InvalidQueryError as e:
+            msg = (
+                "Bad query (%s) used to delete task item state instances: %s"
+                "Please contact support."
+                % (
+                    {"task_execution__in": to_delete_task_execution_ids},
+                    six.text_type(e),
+                )
+            )
+            raise InvalidQueryError(msg)
+        except:
+            logger.exception(
+                "Deletion of task item state models failed for task executions: %s.",
+                to_delete_task_execution_ids,
+            )
+        else:
+            logger.info("Deleted %s task item state objects" % deleted_item_count)
+
+    # 2. Delete the task execution objects.
     try:
         deleted_count = TaskExecution.delete_by_query(**exec_filters)
     except InvalidQueryError as e:

st2common/st2common/models/db/workflow.py

@@ -25,7 +25,7 @@
 from st2common.util import date as date_utils
 
 
-__all__ = ["WorkflowExecutionDB", "TaskExecutionDB"]
+__all__ = ["WorkflowExecutionDB", "TaskExecutionDB", "TaskItemStateDB"]
 
 
 LOG = logging.getLogger(__name__)
@@ -85,4 +85,31 @@ class TaskExecutionDB(stormbase.StormFoundationDB, stormbase.ChangeRevisionField
     }
 
 
-MODELS = [WorkflowExecutionDB, TaskExecutionDB]
+class TaskItemStateDB(stormbase.StormFoundationDB, stormbase.ChangeRevisionFieldMixin):
+    """
+    Model for storing individual item states for tasks with items (itemized tasks).
+    This allows efficient storage and retrieval of individual item states without
+    serializing/deserializing the entire task context for each item.
+    """
+
+    RESOURCE_TYPE = types.ResourceType.EXECUTION
+
+    task_execution = me.StringField(required=True)
+    item_id = me.IntField(required=True)
+    status = me.StringField(required=True)
+    result = JSONDictEscapedFieldCompatibilityField()
+    context = JSONDictEscapedFieldCompatibilityField()
+    start_timestamp = db_field_types.ComplexDateTimeField(
+        default=date_utils.get_datetime_utc_now
+    )
+    end_timestamp = db_field_types.ComplexDateTimeField()
+
+    meta = {
+        "indexes": [
+            {"fields": ["task_execution"]},
+            {"fields": ["task_execution", "item_id"], "unique": True},
+        ]
+    }
+
+
+MODELS = [WorkflowExecutionDB, TaskExecutionDB, TaskItemStateDB]