Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions sandboxes/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ def __init__(self, pool_config: PoolConfig | None = None):
self._busy_sandboxes: set[str] = set()

# Label index for fast lookup
self._label_index: dict[str, set[str]] = {}
self._label_index: dict[tuple[str, str], set[str]] = {}

# Locks for thread-safe operations
self._lock = asyncio.Lock()
Expand Down Expand Up @@ -288,7 +288,7 @@ async def _create_sandbox(self, provider: Any, config: SandboxConfig) -> Sandbox

# Update label index
for key, value in entry.labels.items():
label_key = f"{key}:{value}"
label_key = (key, value)
if label_key not in self._label_index:
self._label_index[label_key] = set()
self._label_index[label_key].add(sandbox.id)
Expand All @@ -304,7 +304,7 @@ async def _create_sandbox(self, provider: Any, config: SandboxConfig) -> Sandbox
async def _find_reusable_sandbox(self, labels: dict[str, str]) -> Sandbox | None:
"""Find an idle sandbox with matching labels."""
# Build label keys
label_keys = [f"{k}:{v}" for k, v in labels.items()]
label_keys = list(labels.items())

# Find sandboxes with all matching labels
matching_ids = None
Expand Down Expand Up @@ -333,7 +333,7 @@ async def find_by_labels(self, labels: dict[str, str]) -> list[Sandbox]:
return []

async with self._lock:
label_keys = [f"{k}:{v}" for k, v in labels.items()]
label_keys = list(labels.items())
matching_ids: set[str] | None = None

for label_key in label_keys:
Expand Down Expand Up @@ -435,7 +435,7 @@ def _remove_from_pool_locked(self, sandbox_id: str) -> SandboxPoolEntry | None:
entry = self._pool.pop(sandbox_id)

for key, value in entry.labels.items():
label_key = f"{key}:{value}"
label_key = (key, value)
if label_key in self._label_index:
self._label_index[label_key].discard(sandbox_id)
if not self._label_index[label_key]:
Expand Down
25 changes: 25 additions & 0 deletions tests/test_pool_label_keys.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
import pytest

from sandboxes.base import SandboxConfig
from sandboxes.pool import PoolConfig, SandboxPool
from tests.test_pool import MockProvider


@pytest.mark.asyncio
async def test_colons_do_not_collapse_distinct_label_pairs():
provider = MockProvider()
pool = SandboxPool(PoolConfig(auto_cleanup=False))
try:
first = await pool.acquire(provider, SandboxConfig(labels={"scope:project": "a"}))
await pool.release(first.id)
assert await pool.find_by_labels({"scope": "project:a"}) == []
second = await pool.acquire(provider, SandboxConfig(labels={"scope": "project:a"}))
assert second.id != first.id
assert [sandbox.id for sandbox in await pool.find_by_labels({"scope:project": "a"})] == [
first.id
]
assert [sandbox.id for sandbox in await pool.find_by_labels({"scope": "project:a"})] == [
second.id
]
finally:
await pool.stop()