|
7 | 7 | from durabletask.azuremanaged.internal import sandbox_service_pb2 as pb |
8 | 8 | from durabletask.azuremanaged.preview.sandboxes.helpers import ( |
9 | 9 | SandboxActivity, |
10 | | - activities_overlap, |
11 | 10 | format_activity, |
12 | 11 | normalize_required, |
13 | 12 | resolve_activities, |
@@ -91,22 +90,109 @@ def _build_sandbox_worker_profile( |
91 | 90 | return worker_profile |
92 | 91 |
|
93 | 92 |
|
| 93 | +class _ActivityOwnerSlot: |
| 94 | + """Earliest worker profiles recorded for a single activity overlap bucket. |
| 95 | +
|
| 96 | + Only two owners ever need to be retained to answer "which worker profile |
| 97 | + first claimed an activity in this bucket, ignoring `worker_profile_id`": |
| 98 | + the very first owner, plus the first owner that differs from it. |
| 99 | + """ |
| 100 | + |
| 101 | + __slots__ = ("_first", "_first_other") |
| 102 | + |
| 103 | + def __init__(self) -> None: |
| 104 | + self._first: Optional[tuple[int, str]] = None |
| 105 | + self._first_other: Optional[tuple[int, str]] = None |
| 106 | + |
| 107 | + def add(self, order: int, worker_profile_id: str) -> None: |
| 108 | + if self._first is None: |
| 109 | + self._first = (order, worker_profile_id) |
| 110 | + elif self._first_other is None and worker_profile_id != self._first[1]: |
| 111 | + self._first_other = (order, worker_profile_id) |
| 112 | + |
| 113 | + def first_owner_other_than(self, worker_profile_id: str) -> Optional[tuple[int, str]]: |
| 114 | + if self._first is not None and self._first[1] != worker_profile_id: |
| 115 | + return self._first |
| 116 | + return self._first_other |
| 117 | + |
| 118 | + |
| 119 | +class _ActivityOwnerIndex: |
| 120 | + """Indexes activity ownership so overlap checks cost O(1) per activity. |
| 121 | +
|
| 122 | + Reproduces the semantics of |
| 123 | + :func:`durabletask.azuremanaged.preview.sandboxes.helpers.activities_overlap` |
| 124 | + exactly: activity names are compared case insensitively, an unversioned |
| 125 | + activity overlaps every version of the same name, and two activities with |
| 126 | + the same name overlap when their explicit versions are equal. Registration |
| 127 | + order is tracked so the reported conflict matches the first overlapping |
| 128 | + activity, exactly as a linear scan would report it. |
| 129 | + """ |
| 130 | + |
| 131 | + def __init__(self) -> None: |
| 132 | + self._registration_count = 0 |
| 133 | + self._by_name: dict[str, _ActivityOwnerSlot] = {} |
| 134 | + self._by_name_and_version: dict[tuple[str, Optional[str]], _ActivityOwnerSlot] = {} |
| 135 | + |
| 136 | + def find_conflicting_profile( |
| 137 | + self, |
| 138 | + activity: SandboxActivity, |
| 139 | + worker_profile_id: str) -> Optional[str]: |
| 140 | + """Return the profile that first claimed an overlapping activity, if any.""" |
| 141 | + name_key = activity.name.casefold() |
| 142 | + if activity.version is None: |
| 143 | + # An unversioned activity overlaps every version of the same name. |
| 144 | + owner = _first_owner_other_than(self._by_name.get(name_key), worker_profile_id) |
| 145 | + else: |
| 146 | + # A versioned activity overlaps the same version plus any |
| 147 | + # unversioned registration of the same name. |
| 148 | + owner = _earlier_owner( |
| 149 | + _first_owner_other_than( |
| 150 | + self._by_name_and_version.get((name_key, None)), worker_profile_id), |
| 151 | + _first_owner_other_than( |
| 152 | + self._by_name_and_version.get((name_key, activity.version)), worker_profile_id)) |
| 153 | + return None if owner is None else owner[1] |
| 154 | + |
| 155 | + def add(self, activity: SandboxActivity, worker_profile_id: str) -> None: |
| 156 | + """Record `worker_profile_id` as an owner of `activity`.""" |
| 157 | + name_key = activity.name.casefold() |
| 158 | + order = self._registration_count |
| 159 | + self._registration_count += 1 |
| 160 | + self._by_name.setdefault(name_key, _ActivityOwnerSlot()).add(order, worker_profile_id) |
| 161 | + self._by_name_and_version.setdefault( |
| 162 | + (name_key, activity.version), _ActivityOwnerSlot()).add(order, worker_profile_id) |
| 163 | + |
| 164 | + |
| 165 | +def _first_owner_other_than( |
| 166 | + slot: Optional[_ActivityOwnerSlot], |
| 167 | + worker_profile_id: str) -> Optional[tuple[int, str]]: |
| 168 | + return None if slot is None else slot.first_owner_other_than(worker_profile_id) |
| 169 | + |
| 170 | + |
| 171 | +def _earlier_owner( |
| 172 | + left: Optional[tuple[int, str]], |
| 173 | + right: Optional[tuple[int, str]]) -> Optional[tuple[int, str]]: |
| 174 | + if left is None: |
| 175 | + return right |
| 176 | + if right is None: |
| 177 | + return left |
| 178 | + return left if left[0] <= right[0] else right |
| 179 | + |
| 180 | + |
94 | 181 | def build_sandbox_worker_profiles() -> list[pb.SandboxWorkerProfile]: |
95 | 182 | """Build sandbox worker_profiles from worker profile configuration.""" |
96 | 183 | worker_profiles: list[pb.SandboxWorkerProfile] = [] |
97 | | - activity_owners: list[tuple[SandboxActivity, str]] = [] |
| 184 | + activity_owners = _ActivityOwnerIndex() |
98 | 185 | for profile in registered_sandbox_worker_profiles(): |
99 | 186 | activities = resolve_activities(profile.activities) |
100 | 187 |
|
101 | 188 | for activity in activities: |
102 | | - existing_profile = next((owner_profile for owner_activity, owner_profile in activity_owners |
103 | | - if activities_overlap(owner_activity, activity) |
104 | | - and owner_profile != profile.worker_profile_id), None) |
| 189 | + existing_profile = activity_owners.find_conflicting_profile( |
| 190 | + activity, profile.worker_profile_id) |
105 | 191 | if existing_profile: |
106 | 192 | raise ValueError( |
107 | 193 | f"Sandbox activity '{format_activity(activity)}' is assigned to both worker profile " |
108 | 194 | f"'{existing_profile}' and '{profile.worker_profile_id}'.") |
109 | | - activity_owners.append((activity, profile.worker_profile_id)) |
| 195 | + activity_owners.add(activity, profile.worker_profile_id) |
110 | 196 |
|
111 | 197 | worker_profiles.append(_build_sandbox_worker_profile( |
112 | 198 | activities=activities, |
|
0 commit comments