Performance [P3]: Index sandbox activity overlap validation

durabletask-azuremanaged/durabletask/azuremanaged/preview/sandboxes/profile_builder.py

@@ -7,7 +7,6 @@
 from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb
 from durabletask.azuremanaged.preview.sandboxes.helpers import (
     SandboxActivity,
-    activities_overlap,
     format_activity,
     normalize_required,
     resolve_activities,
@@ -91,22 +90,109 @@ def _build_sandbox_worker_profile(
     return worker_profile
 
 
+class _ActivityOwnerSlot:
+    """Earliest worker profiles recorded for a single activity overlap bucket.
+
+    Only two owners ever need to be retained to answer "which worker profile
+    first claimed an activity in this bucket, ignoring `worker_profile_id`":
+    the very first owner, plus the first owner that differs from it.
+    """
+
+    __slots__ = ("_first", "_first_other")
+
+    def __init__(self) -> None:
+        self._first: Optional[tuple[int, str]] = None
+        self._first_other: Optional[tuple[int, str]] = None
+
+    def add(self, order: int, worker_profile_id: str) -> None:
+        if self._first is None:
+            self._first = (order, worker_profile_id)
+        elif self._first_other is None and worker_profile_id != self._first[1]:
+            self._first_other = (order, worker_profile_id)
+
+    def first_owner_other_than(self, worker_profile_id: str) -> Optional[tuple[int, str]]:
+        if self._first is not None and self._first[1] != worker_profile_id:
+            return self._first
+        return self._first_other
+
+
+class _ActivityOwnerIndex:
+    """Indexes activity ownership so overlap checks cost O(1) per activity.
+
+    Reproduces the semantics of
+    :func:`durabletask.azuremanaged.preview.sandboxes.helpers.activities_overlap`
+    exactly: activity names are compared case insensitively, an unversioned
+    activity overlaps every version of the same name, and two activities with
+    the same name overlap when their explicit versions are equal. Registration
+    order is tracked so the reported conflict matches the first overlapping
+    activity, exactly as a linear scan would report it.
+    """
+
+    def __init__(self) -> None:
+        self._registration_count = 0
+        self._by_name: dict[str, _ActivityOwnerSlot] = {}
+        self._by_name_and_version: dict[tuple[str, Optional[str]], _ActivityOwnerSlot] = {}
+
+    def find_conflicting_profile(
+            self,
+            activity: SandboxActivity,
+            worker_profile_id: str) -> Optional[str]:
+        """Return the profile that first claimed an overlapping activity, if any."""
+        name_key = activity.name.casefold()
+        if activity.version is None:
+            # An unversioned activity overlaps every version of the same name.
+            owner = _first_owner_other_than(self._by_name.get(name_key), worker_profile_id)
+        else:
+            # A versioned activity overlaps the same version plus any
+            # unversioned registration of the same name.
+            owner = _earlier_owner(
+                _first_owner_other_than(
+                    self._by_name_and_version.get((name_key, None)), worker_profile_id),
+                _first_owner_other_than(
+                    self._by_name_and_version.get((name_key, activity.version)), worker_profile_id))
+        return None if owner is None else owner[1]
+
+    def add(self, activity: SandboxActivity, worker_profile_id: str) -> None:
+        """Record `worker_profile_id` as an owner of `activity`."""
+        name_key = activity.name.casefold()
+        order = self._registration_count
+        self._registration_count += 1
+        self._by_name.setdefault(name_key, _ActivityOwnerSlot()).add(order, worker_profile_id)
+        self._by_name_and_version.setdefault(
+            (name_key, activity.version), _ActivityOwnerSlot()).add(order, worker_profile_id)
+
+
+def _first_owner_other_than(
+        slot: Optional[_ActivityOwnerSlot],
+        worker_profile_id: str) -> Optional[tuple[int, str]]:
+    return None if slot is None else slot.first_owner_other_than(worker_profile_id)
+
+
+def _earlier_owner(
+        left: Optional[tuple[int, str]],
+        right: Optional[tuple[int, str]]) -> Optional[tuple[int, str]]:
+    if left is None:
+        return right
+    if right is None:
+        return left
+    return left if left[0] <= right[0] else right
+
+
 def build_sandbox_worker_profiles() -> list[pb.SandboxWorkerProfile]:
     """Build sandbox worker_profiles from worker profile configuration."""
     worker_profiles: list[pb.SandboxWorkerProfile] = []
-    activity_owners: list[tuple[SandboxActivity, str]] = []
+    activity_owners = _ActivityOwnerIndex()
     for profile in registered_sandbox_worker_profiles():
         activities = resolve_activities(profile.activities)
 
         for activity in activities:
-            existing_profile = next((owner_profile for owner_activity, owner_profile in activity_owners
-                                     if activities_overlap(owner_activity, activity)
-                                     and owner_profile != profile.worker_profile_id), None)
+            existing_profile = activity_owners.find_conflicting_profile(
+                activity, profile.worker_profile_id)
             if existing_profile:
                 raise ValueError(
                     f"Sandbox activity '{format_activity(activity)}' is assigned to both worker profile "
                     f"'{existing_profile}' and '{profile.worker_profile_id}'.")
-            activity_owners.append((activity, profile.worker_profile_id))
+            activity_owners.add(activity, profile.worker_profile_id)
 
         worker_profiles.append(_build_sandbox_worker_profile(
             activities=activities,

tests/durabletask-azuremanaged/test_sandboxes_extension.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT License.
 
 import inspect
+import random
 import threading
 
 import grpc
@@ -20,6 +21,7 @@
     SandboxWorkerProfileImageOptions,
 )
 from durabletask.azuremanaged.preview.sandboxes.profile_builder import (
+    _ActivityOwnerIndex,
     _build_sandbox_worker_profile,
     build_sandbox_worker_profiles,
 )
@@ -28,6 +30,7 @@
     build_sandbox_worker_start,
 )
 from durabletask.azuremanaged.preview.sandboxes.helpers import resolve_activities
+from durabletask.azuremanaged.preview.sandboxes.helpers import activities_overlap
 from durabletask.azuremanaged.preview.sandboxes.helpers import SandboxActivity
 from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb
 from durabletask.azuremanaged.internal import sandbox_service_pb2_grpc as stubs
@@ -206,6 +209,9 @@ def configure(self, options: SandboxWorkerProfileOptions) -> None:
         try:
             build_sandbox_worker_profiles()
         except ValueError as ex:
+            assert str(ex) == (
+                "Sandbox activity 'pytestoverlapremotehello' is assigned to both worker profile "
+                "'pytest-overlap-profile-a' and 'pytest-overlap-profile-b'.")
             assert "pytestoverlapremotehello" in str(ex)
             assert "pytest-overlap-profile-a" in str(ex)
             assert "pytest-overlap-profile-b" in str(ex)
@@ -252,6 +258,128 @@ def configure(self, options: SandboxWorkerProfileOptions) -> None:
         sandbox_worker_profiles._worker_profiles.pop("pytest-version-profile-b", None)
 
 
+def test_build_sandbox_worker_profiles_rejects_versioned_activity_overlapping_unversioned_owner() -> None:
+    @sandbox_worker_profile("pytest-unversioned-owner-profile-a")
+    class PytestUnversionedOwnerProfileA(SandboxWorkerProfile):
+        def configure(self, options: SandboxWorkerProfileOptions) -> None:
+            options.image.image_ref = "example.azurecr.io/python-worker-a:v1"
+            options.image.managed_identity_client_id = "image-pull-client-id"
+            options.scheduler_managed_identity_client_id = "scheduler-client-id"
+            options.add_activity("PytestUnversionedOwner", version=None)
+
+    @sandbox_worker_profile("pytest-unversioned-owner-profile-b")
+    class PytestUnversionedOwnerProfileB(SandboxWorkerProfile):
+        def configure(self, options: SandboxWorkerProfileOptions) -> None:
+            options.image.image_ref = "example.azurecr.io/python-worker-b:v1"
+            options.image.managed_identity_client_id = "image-pull-client-id"
+            options.scheduler_managed_identity_client_id = "scheduler-client-id"
+            options.add_activity("pytestunversionedowner", version="v9")
+
+    try:
+        try:
+            build_sandbox_worker_profiles()
+        except ValueError as ex:
+            assert str(ex) == (
+                "Sandbox activity 'pytestunversionedowner@v9' is assigned to both worker profile "
+                "'pytest-unversioned-owner-profile-a' and 'pytest-unversioned-owner-profile-b'.")
+        else:
+            raise AssertionError(
+                "Expected an unversioned sandbox activity to overlap every version of the same name.")
+    finally:
+        sandbox_worker_profiles._worker_profiles.pop("pytest-unversioned-owner-profile-a", None)
+        sandbox_worker_profiles._worker_profiles.pop("pytest-unversioned-owner-profile-b", None)
+
+
+def test_build_sandbox_worker_profiles_rejects_unversioned_activity_overlapping_versioned_owner() -> None:
+    @sandbox_worker_profile("pytest-versioned-owner-profile-a")
+    class PytestVersionedOwnerProfileA(SandboxWorkerProfile):
+        def configure(self, options: SandboxWorkerProfileOptions) -> None:
+            options.image.image_ref = "example.azurecr.io/python-worker-a:v1"
+            options.image.managed_identity_client_id = "image-pull-client-id"
+            options.scheduler_managed_identity_client_id = "scheduler-client-id"
+            options.add_activity("PytestVersionedOwner", version="v9")
+
+    @sandbox_worker_profile("pytest-versioned-owner-profile-b")
+    class PytestVersionedOwnerProfileB(SandboxWorkerProfile):
+        def configure(self, options: SandboxWorkerProfileOptions) -> None:
+            options.image.image_ref = "example.azurecr.io/python-worker-b:v1"
+            options.image.managed_identity_client_id = "image-pull-client-id"
+            options.scheduler_managed_identity_client_id = "scheduler-client-id"
+            options.add_activity("pytestversionedowner", version=None)
+
+    try:
+        try:
+            build_sandbox_worker_profiles()
+        except ValueError as ex:
+            assert str(ex) == (
+                "Sandbox activity 'pytestversionedowner' is assigned to both worker profile "
+                "'pytest-versioned-owner-profile-a' and 'pytest-versioned-owner-profile-b'.")
+        else:
+            raise AssertionError(
+                "Expected an unversioned sandbox activity to overlap an existing versioned owner.")
+    finally:
+        sandbox_worker_profiles._worker_profiles.pop("pytest-versioned-owner-profile-a", None)
+        sandbox_worker_profiles._worker_profiles.pop("pytest-versioned-owner-profile-b", None)
+
+
+def test_build_sandbox_worker_profiles_rejects_identical_activity_versions() -> None:
+    @sandbox_worker_profile("pytest-same-version-profile-a")
+    class PytestSameVersionProfileA(SandboxWorkerProfile):
+        def configure(self, options: SandboxWorkerProfileOptions) -> None:
+            options.image.image_ref = "example.azurecr.io/python-worker-a:v1"
+            options.image.managed_identity_client_id = "image-pull-client-id"
+            options.scheduler_managed_identity_client_id = "scheduler-client-id"
+            options.add_activity("PytestSameVersionActivity", version="v3")
+
+    @sandbox_worker_profile("pytest-same-version-profile-b")
+    class PytestSameVersionProfileB(SandboxWorkerProfile):
+        def configure(self, options: SandboxWorkerProfileOptions) -> None:
+            options.image.image_ref = "example.azurecr.io/python-worker-b:v1"
+            options.image.managed_identity_client_id = "image-pull-client-id"
+            options.scheduler_managed_identity_client_id = "scheduler-client-id"
+            options.add_activity("pytestsameversionactivity", version="v3")
+
+    try:
+        try:
+            build_sandbox_worker_profiles()
+        except ValueError as ex:
+            assert str(ex) == (
+                "Sandbox activity 'pytestsameversionactivity@v3' is assigned to both worker profile "
+                "'pytest-same-version-profile-a' and 'pytest-same-version-profile-b'.")
+        else:
+            raise AssertionError("Expected identical sandbox activity versions to overlap.")
+    finally:
+        sandbox_worker_profiles._worker_profiles.pop("pytest-same-version-profile-a", None)
+        sandbox_worker_profiles._worker_profiles.pop("pytest-same-version-profile-b", None)
+
+
+def test_activity_owner_index_matches_linear_overlap_scan() -> None:
+    registrations = [
+        (SandboxActivity(name, version), worker_profile_id)
+        for name in ("Alpha", "alpha", "BETA", "Gamma")
+        for version in (None, "v1", "v2", "V1")
+        for worker_profile_id in ("profile-a", "profile-b", "profile-c")
+    ]
+
+    for seed in range(25):
+        shuffled = list(registrations)
+        random.Random(seed).shuffle(shuffled)
+
+        index = _ActivityOwnerIndex()
+        owners: list[tuple[SandboxActivity, str]] = []
+        for activity, worker_profile_id in shuffled:
+            expected = next(
+                (owner_profile for owner_activity, owner_profile in owners
+                 if activities_overlap(owner_activity, activity)
+                 and owner_profile != worker_profile_id),
+                None)
+
+            assert index.find_conflicting_profile(activity, worker_profile_id) == expected
+
+            owners.append((activity, worker_profile_id))
+            index.add(activity, worker_profile_id)
+
+
 def test_profile_options_add_activity_accepts_callable() -> None:
     def pytest_callable_remote_hello(_ctx, value):
         return value