mantisdk 0.1.0__py3-none-any.whl

mantisdk/store/collection/memory.py

@@ -0,0 +1,970 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+import weakref
+from collections import deque
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import (
+    Any,
+    Deque,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+import aiologic
+from pydantic import BaseModel
+
+from mantisdk.types import (
+    Attempt,
+    FilterField,
+    FilterOptions,
+    PaginatedResult,
+    ResourcesUpdate,
+    Rollout,
+    SortOptions,
+    Span,
+    Worker,
+)
+from mantisdk.utils.metrics import MetricsBackend
+
+from .base import (
+    AtomicLabels,
+    AtomicMode,
+    Collection,
+    DuplicatedPrimaryKeyError,
+    FilterMap,
+    KeyValue,
+    LightningCollections,
+    Queue,
+    ensure_numeric,
+    normalize_filter_options,
+    resolve_sort_options,
+    tracked,
+)
+
+T = TypeVar("T")  # Recommended to be a BaseModel, not a dict
+K = TypeVar("K")
+V = TypeVar("V")
+
+logger = logging.getLogger(__name__)
+
+# Nested structure type:
+# dict[pk1] -> dict[pk2] -> ... -> item
+ListBasedCollectionItemType = Union[
+    Dict[Any, "ListBasedCollectionItemType[T]"],  # intermediate node
+    Dict[Any, T],  # leaf node dictionary
+]
+
+MutationMode = Literal["insert", "update", "upsert", "delete"]
+
+
+def _item_matches_filters(
+    item: object,
+    filters: Optional[FilterMap],
+    filter_logic: Literal["and", "or"],
+    must_filters: Optional[FilterMap] = None,
+) -> bool:
+    """Check whether an item matches the provided filter definition.
+
+    Filter format:
+
+    ```json
+    {
+        "_aggregate": "or",
+        "field_name": {
+            "exact": <value>,
+            "within": <iterable_of_allowed_values>,
+            "contains": <substring_or_element>,
+        },
+        ...
+    }
+    ```
+
+    Operators within the same field are stored in a unified pool and combined using
+    a universal logical operator.
+    """
+    if must_filters and not _item_matches_filters(item, must_filters, "and"):
+        return False
+
+    if not filters:
+        return True
+
+    all_conditions_match: List[bool] = []
+
+    for field_name, ops in filters.items():
+        item_value = getattr(item, field_name, None)
+
+        for op_name, expected in ops.items():
+            # Ignore no-op filters
+            if expected is None:
+                continue
+
+            if op_name == "exact":
+                all_conditions_match.append(item_value == expected)
+
+            elif op_name == "within":
+                try:
+                    all_conditions_match.append(item_value in expected)  # type: ignore[arg-type]
+                except TypeError:
+                    all_conditions_match.append(False)
+
+            elif op_name == "contains":
+                if item_value is None:
+                    all_conditions_match.append(False)
+                elif isinstance(item_value, str) and isinstance(expected, str):
+                    all_conditions_match.append(expected in item_value)
+                else:
+                    # Fallback: treat as generic iterable containment.
+                    try:
+                        all_conditions_match.append(expected in item_value)  # type: ignore[arg-type]
+                    except TypeError:
+                        all_conditions_match.append(False)
+            else:
+                raise ValueError(f"Unsupported filter operator '{op_name}' for field '{field_name}'")
+
+    return all(all_conditions_match) if filter_logic == "and" else any(all_conditions_match)
+
+
+def _get_sort_value(item: object, sort_by: str) -> Any:
+    """Get a sort key for the given item/field.
+
+    - If the field name ends with '_time', values are treated as comparable timestamps.
+    - For other fields we try to infer a safe default from the Pydantic model annotation.
+    """
+    value = getattr(item, sort_by, None)
+
+    if sort_by.endswith("_time"):
+        # For *_time fields, push missing values to the end.
+        return float("inf") if value is None else value
+
+    if value is None:
+        # Introspect model field type to choose a reasonable default for None.
+        model_fields = getattr(item.__class__, "model_fields", {})
+        if sort_by not in model_fields:
+            raise ValueError(
+                f"Failed to sort items by '{sort_by}': field does not exist " f"on {item.__class__.__name__}"
+            )
+
+        field_type_str = str(model_fields[sort_by].annotation)
+        if "str" in field_type_str or "Literal" in field_type_str:
+            return ""
+        if "int" in field_type_str:
+            return 0
+        if "float" in field_type_str:
+            return 0.0
+        raise ValueError(f"Failed to sort items by '{sort_by}': unsupported field type {field_type_str!r}")
+
+    return value
+
+
+class ListBasedCollection(Collection[T]):
+    """In-memory implementation of Collection using a nested dict for O(1) primary-key lookup.
+
+    The internal structure is:
+
+        {
+            pk1_value: {
+                pk2_value: {
+                    ...
+                        pkN_value: item
+                }
+            }
+        }
+
+    where the nesting depth equals the number of primary keys.
+
+    Sorting behavior:
+
+    1. If no sort_by is provided, the items are returned in the order of insertion.
+    2. If sort_by is provided, the items are sorted by the value of the sort_by field.
+    3. If the sort_by field is a timestamp, the null values are treated as infinity.
+    4. If the sort_by field is not a timestamp, the null values are treated as empty string
+       if the field is str-like, 0 if the field is int-like, 0.0 if the field is float-like.
+    """
+
+    def __init__(
+        self,
+        items: List[T],
+        item_type: Type[T],
+        primary_keys: Sequence[str],
+        id: Optional[str] = None,
+        tracker: Optional[MetricsBackend] = None,
+    ):
+        super().__init__(tracker=tracker)
+        if not primary_keys:
+            raise ValueError("primary_keys must be non-empty")
+
+        self._id = id if id is not None else str(uuid.uuid4())
+        self._items: Dict[Any, Any] = {}
+        self._size: int = 0
+        if issubclass(item_type, dict):
+            raise TypeError(f"Expect item to be not a dict, got {item_type.__name__}")
+        self._item_type: Type[T] = item_type
+        self._primary_keys: Tuple[str, ...] = tuple(primary_keys)
+
+        # Pre-populate the collection with the given items.
+        for item in items or []:
+            self._mutate_single(item, mode="insert")
+
+    @property
+    def collection_name(self) -> str:
+        return self._id
+
+    def primary_keys(self) -> Sequence[str]:
+        """Return the primary key field names for this collection."""
+        return self._primary_keys
+
+    def item_type(self) -> Type[T]:
+        """Return the Pydantic model type of items stored in this collection."""
+        return self._item_type
+
+    async def size(self) -> int:
+        """Return the number of items stored in the collection."""
+        return self._size
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({self._size})>"
+
+    # -------------------------------------------------------------------------
+    # Internal helpers
+    # -------------------------------------------------------------------------
+
+    def _ensure_item_type(self, item: T) -> None:
+        """Validate that the item matches the declared item_type."""
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, " f"got {type(item).__name__}")
+
+    def _extract_primary_key_values(self, item: T) -> Tuple[Any, ...]:
+        """Extract the primary key values from an item.
+
+        Raises:
+            ValueError: If any primary key is missing on the item.
+        """
+        values: List[Any] = []
+        for key in self._primary_keys:
+            if not hasattr(item, key):
+                raise ValueError(f"Item {item} does not have primary key field '{key}'")
+            values.append(getattr(item, key))
+        return tuple(values)
+
+    def _render_key_values(self, key_values: Sequence[Any]) -> str:
+        return ", ".join(f"{name}={value!r}" for name, value in zip(self._primary_keys, key_values))
+
+    def _locate_node(
+        self,
+        key_values: Sequence[Any],
+        create_missing: bool,
+    ) -> Tuple[MutableMapping[Any, Any], Any]:
+        """Locate the parent mapping and final key for an item path.
+
+        Args:
+            key_values: The sequence of primary key values.
+            create_missing: Whether to create intermediate dictionaries as needed.
+
+        Returns:
+            (parent_mapping, final_key)
+
+        Raises:
+            KeyError: If the path does not exist and create_missing is False.
+            ValueError: If the internal structure is corrupted (non-dict where dict is expected).
+        """
+        if not key_values:
+            raise ValueError("key_values must be non-empty")
+
+        current: MutableMapping[Any, Any] = self._items
+        for idx, value in enumerate(key_values):
+            is_last = idx == len(key_values) - 1
+            if is_last:
+                # At the final level, current[value] is the item (or will be).
+                return current, value  # type: ignore
+
+            # Intermediate level: current[value] must be a dict.
+            if value not in current:
+                if not create_missing:
+                    raise KeyError(f"Path does not exist for given primary keys: {self._render_key_values(key_values)}")
+                current[value] = {}
+            next_node = current[value]  # type: ignore
+            if not isinstance(next_node, dict):
+                raise ValueError(f"Internal structure corrupted: expected dict, got {type(next_node)!r}")  # type: ignore
+            current = next_node  # type: ignore
+
+        # We should always return inside the loop.
+        raise RuntimeError("Unreachable")
+
+    def _mutate_single(self, item: T, mode: MutationMode, update_fields: Sequence[str] | None = None) -> Optional[T]:
+        """Core mutation logic shared by insert, update, upsert, and delete."""
+        self._ensure_item_type(item)
+        key_values = self._extract_primary_key_values(item)
+
+        if mode in ("insert", "upsert"):
+            parent, final_key = self._locate_node(key_values, create_missing=True)
+            exists = final_key in parent
+
+            if mode == "insert":
+                if exists:
+                    raise DuplicatedPrimaryKeyError(
+                        f"Item already exists with primary key(s): {self._render_key_values(key_values)}"
+                    )
+                parent[final_key] = item
+                self._size += 1
+            else:  # upsert
+                if not exists:
+                    self._size += 1
+                    parent[final_key] = item
+
+                elif update_fields is None:
+                    # update_or_insert: update all fields
+                    parent[final_key] = item
+
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+
+                    # Try to fetch the existing item
+                    existing = parent[final_key]
+                    if not isinstance(existing, self._item_type):
+                        raise ValueError(
+                            f"Internal structure corrupted: expected {self._item_type.__name__}, got {type(existing)!r}"
+                        )
+
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+
+            return parent[final_key]
+
+        elif mode in ("update", "delete"):
+            # For update/delete we must not create missing paths.
+            try:
+                parent, final_key = self._locate_node(key_values, create_missing=False)
+            except KeyError:
+                raise ValueError(
+                    f"Item does not exist with primary key(s): {self._render_key_values(key_values)}"
+                ) from None
+
+            if final_key not in parent:
+                raise ValueError(f"Item does not exist with primary key(s): {self._render_key_values(key_values)}")
+
+            if mode == "update":
+                if update_fields is None:
+                    # replace the entire item
+                    parent[final_key] = item
+                else:
+                    if not issubclass(self._item_type, BaseModel):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {self._item_type.__name__}"
+                        )
+                    if not isinstance(item, self._item_type):
+                        raise TypeError(
+                            f"When using update_fields, the item type must be a Pydantic BaseModel, got {type(item).__name__}"
+                        )
+                    parent[final_key] = parent[final_key].model_copy(
+                        update={field: getattr(item, field) for field in update_fields}
+                    )
+                return parent[final_key]
+            else:  # delete
+                del parent[final_key]
+                self._size -= 1
+        else:
+            raise ValueError(f"Unknown mutation mode: {mode}")
+
+    def _iter_items(
+        self,
+        root: Optional[Mapping[Any, Any]] = None,
+        filters: Optional[FilterMap] = None,
+        must_filters: Optional[FilterMap] = None,
+        filter_logic: Literal["and", "or"] = "and",
+    ) -> Iterable[T]:
+        """Iterate over all items in the nested dictionary structure, optionally applying filters."""
+        if root is None:
+            root = self._items
+        if not root:
+            return
+        stack: List[Mapping[Any, Any]] = [root]
+        while stack:
+            node = stack.pop()
+            for value in node.values():
+                # Leaf nodes contain items; intermediate nodes are dicts.
+                if isinstance(value, self._item_type):
+                    if _item_matches_filters(value, filters, filter_logic, must_filters):
+                        yield value
+                elif isinstance(value, dict):
+                    stack.append(value)  # type: ignore
+                else:
+                    raise ValueError(
+                        f"Internal structure corrupted: expected dict or {self._item_type.__name__}, "
+                        f"got {type(value)!r}"
+                    )
+
+    def _iter_matching_items(
+        self,
+        filters: Optional[FilterMap],
+        must_filters: Optional[FilterMap],
+        filter_logic: Literal["and", "or"],
+    ) -> Iterable[T]:
+        """Efficiently iterate over items matching filters, using primary-key prefix when possible."""
+        # Fast path: when optional filters can't form a prefix, fall back to scanning.
+        if filter_logic != "and" and must_filters is None:
+            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)
+
+        # Try to derive a primary-key prefix from exact filters.
+        pk_values_prefix: List[Any] = []
+        prefix_sources: List[FilterMap] = []
+        if must_filters:
+            prefix_sources.append(must_filters)
+        if filter_logic == "and" and filters:
+            prefix_sources.append(filters)
+
+        for pk in self._primary_keys:
+            # combined_ops are: [{"exact": value}, {"within": [...]}, ...]
+            combined_ops: List[FilterField] = []
+            for source in prefix_sources:
+                field_ops = source.get(pk)  # type: ignore[union-attr]
+                if field_ops:
+                    combined_ops.append(field_ops)
+            if not combined_ops:
+                break
+            # Only allow a pure {"exact": value} constraint.
+            exact_value: Any | None = None
+            allow_prefix = True
+            for ops in combined_ops:
+                if set(ops.keys()) != {"exact"}:
+                    allow_prefix = False
+                    break
+                candidate = ops.get("exact")
+                if candidate is None:
+                    allow_prefix = False
+                    break
+                if exact_value is not None and candidate != exact_value:
+                    # Contradictory exact filters mean no items can match.
+                    logger.warning(f"Contradictory exact filters for field '{pk}': {exact_value} != {candidate}")
+                    return ()
+                exact_value = candidate
+
+            if not allow_prefix:
+                break
+
+            value = exact_value
+            if value is None:
+                break
+            pk_values_prefix.append(value)
+
+        if not pk_values_prefix:
+            return self._iter_items(filters=filters, must_filters=must_filters, filter_logic=filter_logic)
+
+        try:
+            if len(pk_values_prefix) == len(self._primary_keys):
+                # All primary keys specified -> at most a single item.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                single_item = parent.get(final_key)
+                if isinstance(single_item, self._item_type) and _item_matches_filters(
+                    single_item,
+                    filters,
+                    filter_logic,
+                    must_filters,
+                ):
+                    return (single_item,)
+                return ()
+            else:
+                # Prefix of primary keys specified -> iterate only the subtree below that prefix.
+                parent, final_key = self._locate_node(pk_values_prefix, create_missing=False)
+                subtree = parent.get(final_key)
+                if isinstance(subtree, dict):
+                    return self._iter_items(
+                        subtree,  # type: ignore
+                        filters=filters,
+                        must_filters=must_filters,
+                        filter_logic=filter_logic,
+                    )
+                return ()
+        except KeyError:
+            # No items exist for this primary-key prefix.
+            return ()
+
+    @tracked("query")
+    async def query(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+        limit: int = -1,
+        offset: int = 0,
+    ) -> PaginatedResult[T]:
+        """Query the collection with filters, sort order, and pagination.
+
+        Args:
+            filter: Mapping of field name to operator dict along with the optional `_aggregate` logic.
+            sort: Options describing which field to sort by and in which order.
+            limit: Max number of items to return. Use -1 for "no limit".
+            offset: Number of items to skip from the start of the *matching* items.
+        """
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)
+
+        # No sorting: stream through items and apply pagination on the fly.
+        if not sort_by:
+            matched_items: List[T] = []
+            total_matched = 0
+
+            for item in items_iter:
+                # Count every match for 'total'
+                total_matched += 1
+
+                # Apply offset/limit window
+                if total_matched <= offset:
+                    continue
+                if limit != -1 and len(matched_items) >= limit:
+                    # Still need to finish iteration to get accurate total_matched.
+                    continue
+
+                matched_items.append(item)
+
+            return PaginatedResult(
+                items=matched_items,
+                limit=limit,
+                offset=offset,
+                total=total_matched,
+            )
+
+        # With sorting: we must materialize all matching items to sort them.
+        all_matches: List[T] = list(items_iter)
+
+        total_matched = len(all_matches)
+        reverse = sort_order == "desc"
+        all_matches.sort(key=lambda x: _get_sort_value(x, sort_by), reverse=reverse)
+
+        if limit == -1:
+            paginated_items = all_matches[offset:]
+        else:
+            paginated_items = all_matches[offset : offset + limit]
+
+        return PaginatedResult(
+            items=paginated_items,
+            limit=limit,
+            offset=offset,
+            total=total_matched,
+        )
+
+    @tracked("get")
+    async def get(
+        self,
+        filter: Optional[FilterOptions] = None,
+        sort: Optional[SortOptions] = None,
+    ) -> Optional[T]:
+        """Return the first (or best-sorted) item that matches the given filters, or None."""
+        filters, must_filters, filter_logic = normalize_filter_options(filter)
+        sort_by, sort_order = resolve_sort_options(sort)
+        items_iter: Iterable[T] = self._iter_matching_items(filters, must_filters, filter_logic)
+
+        if not sort_by:
+            # Just return the first matching item, if any.
+            for item in items_iter:
+                return item
+            return None
+
+        # Single-pass min/max according to sort_order.
+        best_item: Optional[T] = None
+        best_key: Any = None
+
+        for item in items_iter:
+            key = _get_sort_value(item, sort_by)
+            if best_item is None:
+                best_item = item
+                best_key = key
+                continue
+
+            if sort_order == "asc":
+                if key < best_key:
+                    best_item, best_key = item, key
+            else:
+                if key > best_key:
+                    best_item, best_key = item, key
+
+        return best_item
+
+    @tracked("insert")
+    async def insert(self, items: Sequence[T]) -> None:
+        """Insert the given items.
+
+        Raises:
+            DuplicatedPrimaryKeyError: If any item with the same primary keys already exists.
+        """
+        seen_keys: set[Tuple[Any, ...]] = set()
+        prepared: List[T] = []
+        for item in items:
+            self._ensure_item_type(item)
+            key_values = self._extract_primary_key_values(item)
+            if key_values in seen_keys:
+                raise DuplicatedPrimaryKeyError(
+                    f"Insert payload contains duplicated primary key(s): {self._render_key_values(key_values)}"
+                )
+            seen_keys.add(key_values)
+            prepared.append(item)
+
+        for item in prepared:
+            self._mutate_single(item, mode="insert")
+
+    @tracked("update")
+    async def update(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+        """Update the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        updated_items: List[T] = []
+        for item in items:
+            updated = self._mutate_single(item, mode="update", update_fields=update_fields)
+            if updated is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            updated_items.append(updated)
+        return updated_items
+
+    @tracked("upsert")
+    async def upsert(self, items: Sequence[T], update_fields: Sequence[str] | None = None) -> Sequence[T]:
+        """Upsert the given items (insert if missing, otherwise update)."""
+        upserted_items: List[T] = []
+        for item in items:
+            upserted = self._mutate_single(item, mode="upsert", update_fields=update_fields)
+            if upserted is None:
+                raise RuntimeError(f"_mutate_single returned None for item {item}. This should never happen.")
+            upserted_items.append(upserted)
+        return upserted_items
+
+    @tracked("delete")
+    async def delete(self, items: Sequence[T]) -> None:
+        """Delete the given items.
+
+        Raises:
+            ValueError: If any item with the given primary keys does not exist.
+        """
+        # We use a two-phase approach to avoid partial deletion if one fails:
+        # first compute key_values to validate, then perform deletions.
+        for item in items:
+            # _mutate_single will validate existence and update size.
+            self._mutate_single(item, mode="delete")
+
+
+class DequeBasedQueue(Queue[T]):
+    """Queue implementation backed by collections.deque.
+
+    Provides O(1) amortized enqueue (append) and dequeue (popleft).
+    """
+
+    def __init__(
+        self,
+        item_type: Type[T],
+        items: Optional[Sequence[T]] = None,
+        id: Optional[str] = None,
+        tracker: Optional[MetricsBackend] = None,
+    ):
+        super().__init__(tracker=tracker)
+        self._items: Deque[T] = deque()
+        self._item_type: Type[T] = item_type
+        self._id = id if id is not None else str(uuid.uuid4())
+        if items:
+            self._items.extend(items)
+
+    def item_type(self) -> Type[T]:
+        return self._item_type
+
+    @property
+    def collection_name(self) -> str:
+        return self._id
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}[{self.item_type().__name__}] ({len(self._items)})>"
+
+    @tracked("has")
+    async def has(self, item: T) -> bool:
+        if not isinstance(item, self._item_type):
+            raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+        return item in self._items
+
+    @tracked("enqueue")
+    async def enqueue(self, items: Sequence[T]) -> Sequence[T]:
+        for item in items:
+            if not isinstance(item, self._item_type):
+                raise TypeError(f"Expected item of type {self._item_type.__name__}, got {type(item).__name__}")
+            self._items.append(item)
+        return items
+
+    @tracked("dequeue")
+    async def dequeue(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        out: List[T] = []
+        for _ in range(min(limit, len(self._items))):
+            out.append(self._items.popleft())
+        return out
+
+    @tracked("peek")
+    async def peek(self, limit: int = 1) -> Sequence[T]:
+        if limit <= 0:
+            return []
+        result: List[T] = []
+        count = min(limit, len(self._items))
+        for idx, item in enumerate(self._items):
+            if idx >= count:
+                break
+            result.append(item)
+        return result
+
+    @tracked("size")
+    async def size(self) -> int:
+        return len(self._items)
+
+
+class DictBasedKeyValue(KeyValue[K, V]):
+    """KeyValue implementation backed by a plain dictionary."""
+
+    def __init__(
+        self, data: Optional[Mapping[K, V]] = None, id: Optional[str] = None, tracker: Optional[MetricsBackend] = None
+    ):
+        super().__init__(tracker=tracker)
+        self._values: Dict[K, V] = dict(data) if data else {}
+        self._id = id if id is not None else str(uuid.uuid4())
+
+    @property
+    def collection_name(self) -> str:
+        return self._id
+
+    @tracked("has")
+    async def has(self, key: K) -> bool:
+        return key in self._values
+
+    @tracked("get")
+    async def get(self, key: K, default: V | None = None) -> V | None:
+        return self._values.get(key, default)
+
+    @tracked("set")
+    async def set(self, key: K, value: V) -> None:
+        self._values[key] = value
+
+    @tracked("inc")
+    async def inc(self, key: K, amount: V) -> V:
+        assert ensure_numeric(amount, description="amount")
+        if key in self._values:
+            current_value = self._values[key]
+            assert ensure_numeric(current_value, description=f"value for key {key!r}")
+            new_value = cast(V, current_value + amount)
+            self._values[key] = new_value
+        else:
+            new_value = amount
+            self._values[key] = new_value
+        return new_value
+
+    @tracked("chmax")
+    async def chmax(self, key: K, value: V) -> V:
+        assert ensure_numeric(value, description="value")
+        if key in self._values:
+            current_value = self._values[key]
+            assert ensure_numeric(current_value, description=f"value for key {key!r}")
+            if value > current_value:
+                self._values[key] = value
+                return value
+            return current_value
+        else:
+            self._values[key] = value
+            return value
+
+    @tracked("pop")
+    async def pop(self, key: K, default: V | None = None) -> V | None:
+        return self._values.pop(key, default)
+
+    @tracked("size")
+    async def size(self) -> int:
+        return len(self._values)
+
+
+class InMemoryLightningCollections(LightningCollections):
+    """In-memory implementation of LightningCollections using Python data structures.
+
+    Serves as the storage base for [`InMemoryLightningStore`][mantisdk.InMemoryLightningStore].
+    """
+
+    def __init__(self, lock_type: Literal["thread", "asyncio"], tracker: MetricsBackend | None = None):
+        super().__init__(tracker=tracker)
+        self._lock: Mapping[AtomicLabels, _LoopAwareAsyncLock | _ThreadSafeAsyncLock] = {
+            "rollouts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "attempts": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "spans": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "resources": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "workers": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "rollout_queue": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "span_sequence_ids": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+            "generic": _LoopAwareAsyncLock() if lock_type == "asyncio" else _ThreadSafeAsyncLock(),
+        }
+        self._rollouts = ListBasedCollection(
+            items=[], item_type=Rollout, primary_keys=["rollout_id"], id="rollouts", tracker=tracker
+        )
+        self._attempts = ListBasedCollection(
+            items=[], item_type=Attempt, primary_keys=["rollout_id", "attempt_id"], id="attempts", tracker=tracker
+        )
+        self._spans = ListBasedCollection(
+            items=[], item_type=Span, primary_keys=["rollout_id", "attempt_id", "span_id"], id="spans", tracker=tracker
+        )
+        self._resources = ListBasedCollection(
+            items=[], item_type=ResourcesUpdate, primary_keys=["resources_id"], id="resources", tracker=tracker
+        )
+        self._workers = ListBasedCollection(
+            items=[], item_type=Worker, primary_keys=["worker_id"], id="workers", tracker=tracker
+        )
+        self._rollout_queue = DequeBasedQueue(items=[], item_type=str, id="rollout_queue", tracker=tracker)
+        self._span_sequence_ids = DictBasedKeyValue[str, int](
+            data={}, id="span_sequence_ids", tracker=tracker
+        )  # rollout_id -> sequence_id
+
+    @property
+    def collection_name(self) -> str:
+        return "router"
+
+    @property
+    def rollouts(self) -> ListBasedCollection[Rollout]:
+        return self._rollouts
+
+    @property
+    def attempts(self) -> ListBasedCollection[Attempt]:
+        return self._attempts
+
+    @property
+    def spans(self) -> ListBasedCollection[Span]:
+        return self._spans
+
+    @property
+    def resources(self) -> ListBasedCollection[ResourcesUpdate]:
+        return self._resources
+
+    @property
+    def workers(self) -> ListBasedCollection[Worker]:
+        return self._workers
+
+    @property
+    def rollout_queue(self) -> DequeBasedQueue[str]:
+        return self._rollout_queue
+
+    @property
+    def span_sequence_ids(self) -> DictBasedKeyValue[str, int]:
+        return self._span_sequence_ids
+
+    @asynccontextmanager
+    async def atomic(
+        self,
+        *,
+        mode: AtomicMode = "rw",
+        snapshot: bool = False,
+        labels: Optional[Sequence[AtomicLabels]] = None,
+        **kwargs: Any,
+    ):
+        """In-memory collections apply a lock outside. It doesn't need to manipulate the collections inside.
+
+        Skip the locking if mode is "r" and snapshot is False.
+
+        This collection implementation does NOT support rollback / commit.
+        """
+        if mode == "r" and not snapshot:
+            yield self
+            return
+        if not labels:
+            # If no labels are provided, use all locks.
+            labels = list(self._lock.keys())
+
+        # IMPORTANT: Sort the labels to ensure consistent locking order.
+        # This is necessary to avoid deadlocks when multiple threads/coroutines
+        # are trying to acquire the same locks in different orders.
+        labels = sorted(labels)
+
+        async with self.tracking_context(operation="atomic", collection=self.collection_name):
+            managers = [(label, self._lock[label]) for label in labels]
+            async with AsyncExitStack() as stack:
+                for label, manager in managers:
+                    async with self.tracking_context(operation="lock", collection=label):
+                        await stack.enter_async_context(manager)
+                yield self
+
+    @tracked("evict_spans_for_rollout")
+    async def evict_spans_for_rollout(self, rollout_id: str) -> None:
+        """Evict all spans for a given rollout ID.
+
+        Uses private API for efficiency.
+        """
+        self._spans._items.pop(rollout_id, [])  # pyright: ignore[reportPrivateUsage]
+
+
+class _LoopAwareAsyncLock:
+    """Async lock that transparently rebinds to the current event loop.
+
+    The lock intentionally remains *thread-unsafe*: callers must only use it from
+    one thread at a time. If multiple threads interact with the store, each
+    thread gets its own event loop specific lock.
+    """
+
+    def __init__(self) -> None:
+        self._locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock] = weakref.WeakKeyDictionary()
+
+    # When serializing and deserializing, we don't need to serialize the locks.
+    # Because another process will have its own set of event loops and its own lock.
+    def __getstate__(self) -> dict[str, Any]:
+        return {}
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self._locks = weakref.WeakKeyDictionary()
+
+    def _get_lock_for_current_loop(self) -> asyncio.Lock:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._locks[loop] = lock
+        return lock
+
+    async def __aenter__(self) -> asyncio.Lock:
+        lock = self._get_lock_for_current_loop()
+        await lock.acquire()
+        return lock
+
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any) -> None:
+        loop = asyncio.get_running_loop()
+        lock = self._locks.get(loop)
+        if lock is None or not lock.locked():
+            raise RuntimeError("Lock released without being acquired")
+        lock.release()
+
+
+class _ThreadSafeAsyncLock:
+    """A thread lock powered by aiologic that can be used in both async and sync contexts.
+
+    aiologic claims itself to be a thread-safe asyncio lock.
+    """
+
+    def __init__(self):
+        self._lock = aiologic.Lock()
+
+    async def __aenter__(self):
+        await self._lock.async_acquire()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any):
+        # .release() is non-blocking, so we can call it directly
+        self._lock.async_release()