Skip to content
Open
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
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
48 changes: 37 additions & 11 deletions api/oss/src/utils/caching.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,24 +40,50 @@
# HELPERS ----------------------------------------------------------------------


AGENTA_CACHE_SCOPE_WIDTH = 12 # minimum width of a scope segment


def _scope(
value: Optional[str],
legacy_truncated: Optional[bool] = False,
) -> str:
"""One scope segment (`p:` or `u:`) of a cache key.

Scope segments used to be cut down to their last 12 characters. Two projects whose
ids share that suffix then shared every cache entry in every namespace — including
`check_permissions` and `check_action_access`, which decide authorization. Ids are
server-generated UUID4s so a caller cannot steer a collision, but one project reading
another's cached permission result is not a thing to leave to chance. The id is
carried whole now.

The padding stays: an absent or short id still produces the same fixed-width segment
it always did, so key shape is unchanged everywhere the id was not being cut.

`legacy_truncated` reproduces the pre-change shape and exists only so lock holders
can span a rolling deploy — see `oss.src.utils.locking`.
"""
value = value or ""

if legacy_truncated and len(value) > AGENTA_CACHE_SCOPE_WIDTH:
value = value[-AGENTA_CACHE_SCOPE_WIDTH:]

return value + "-" * (AGENTA_CACHE_SCOPE_WIDTH - len(value))


def _pack(
namespace: Optional[str] = None,
key: Optional[Union[str, dict]] = None,
project_id: Optional[str] = None,
user_id: Optional[str] = None,
pattern: Optional[bool] = False,
legacy_truncated_scope: Optional[bool] = False,
) -> str:
if project_id:
project_id = project_id[-12:] if len(project_id) > 12 else project_id
else:
project_id = ""
project_id = project_id + "-" * (12 - len(project_id))

if user_id:
user_id = user_id[-12:] if len(user_id) > 12 else user_id
user_id = user_id + "-" * (12 - len(user_id))
else:
user_id = "*" if pattern else "-" * 12
project_id = _scope(project_id, legacy_truncated_scope)
# A pattern with no user scope matches every user, which the fixed-width padding
# cannot express — so the wildcard wins over the padded segment in that one case.
user_id = (
"*" if pattern and not user_id else _scope(user_id, legacy_truncated_scope)
)

namespace = namespace or ("" if not pattern else "*")

Expand Down
168 changes: 142 additions & 26 deletions api/oss/src/utils/locking.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,73 @@
"""


# TRANSITIONAL: SPANNING THE FULL-PROJECT-ID DEPLOY -----------------------------
#
# Scope segments in cache and lock keys used to carry only the last 12 characters of
# an id (see `caching._scope`). During a rolling deploy, pods still on the previous
# release take the truncated key, so a lock held only under the full-id key would not
# exclude them and mutual exclusion would be lost for the length of the deploy.
# Every lock operation therefore covers both keys for one release.
#
# REMOVE once no pod predating the full-id change is running: drop the second element
# of `_lock_keys` and the `legacy_key` branches below. That is the whole surface.


def _lock_keys(
namespace: str,
key: Optional[Union[str, dict]] = None,
project_id: Optional[str] = None,
user_id: Optional[str] = None,
) -> tuple[str, Optional[str]]:
"""This release's lock key, and the previous release's key when it differs."""
lock_key = _pack(
namespace=f"lock:{namespace}",
key=key,
project_id=project_id,
user_id=user_id,
)
legacy_key = _pack(
namespace=f"lock:{namespace}",
key=key,
project_id=project_id,
user_id=user_id,
legacy_truncated_scope=True,
)

# Ids no longer than the segment width were never truncated, so the two shapes
# coincide and there is no second key to cover.
return lock_key, (legacy_key if legacy_key != lock_key else None)


async def _renew_if_owner(lock_key: str, owner: Optional[str], ttl: int) -> bool:
if owner:
return bool(
await _lock_engine.eval(
_LOCK_RENEW_IF_OWNER_SCRIPT,
1,
lock_key,
owner,
str(ttl),
)
)

return bool(await _lock_engine.expire(lock_key, ttl))


async def _release_if_owner(lock_key: str, owner: Optional[str]) -> bool:
if owner:
return bool(
await _lock_engine.eval(
_LOCK_RELEASE_IF_OWNER_SCRIPT,
1,
lock_key,
owner,
)
)

return bool(await _lock_engine.delete(lock_key))


# LOCK-STORE PRIMITIVES --------------------------------------------------------
#
# Thin pass-throughs to the lock Redis client for callers that manage their own
Expand Down Expand Up @@ -116,19 +183,49 @@ async def acquire_lock(
owner=lock_owner,
)
"""
lock_owner = uuid4().hex
legacy_key = None
# Set only while this call holds the legacy key without having handed the section to
# anyone. Cleared once the caller owns it or it has already been released, so the
# `finally` below cleans up exactly the abandoned case.
legacy_claimed = False

try:
lock_key = _pack(
namespace=f"lock:{namespace}",
lock_key, legacy_key = _lock_keys(
namespace=namespace,
key=key,
project_id=project_id,
user_id=user_id,
)
lock_owner = uuid4().hex

# The legacy key is claimed first: a pod on the previous release sets only that
# one, so taking it is what makes the two generations exclude each other. Claiming
# it second would let both generations hold their own key and enter together.
if legacy_key is not None:
# Marked before the await rather than after it. Cancellation can arrive once
# Redis has applied the SET but before its reply reaches us, and a flag set
# afterwards would never record that this call owns the key — leaving it held
# for its whole TTL with nothing to release it. Claiming an obligation that
# may not exist is the safe direction: `_release_if_owner` is ownership
# checked, so it does nothing when the token is not ours.
legacy_claimed = True

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

🔎 Supported by static analysis

🤖 get_repo_knowledge executed:

get_repo_knowledge Agenta-AI/agenta /tmp/coderabbit-repo-knowledge/agenta-ai-agenta-4b53879a/conventions /tmp/coderabbit-repo-knowledge/agenta-ai-agenta-4b53879a/learnings

Length of output: 14349


🏁 Script executed:

#!/bin/bash
set -eu
printf '%s\n' '--- changed file context ---'
sed -n '150,255p' api/oss/src/utils/locking.py
printf '%s\n' '--- relevant tests ---'
sed -n '330,440p' api/oss/tests/pytest/unit/utils/test_cache_key_tenancy.py
printf '%s\n' '--- lock symbol references ---'
rg -n "_release_if_owner|acquire_lock|legacy_claimed|primary_claimed|lock_key" api/oss/src/utils/locking.py api/oss/tests/pytest/unit/utils

Repository: Agenta-AI/agenta

Length of output: 17434


🏁 Script executed:

#!/bin/bash
set -eu
sed -n '245,282p' api/oss/src/utils/locking.py
sed -n '40,112p' api/oss/src/utils/locking.py
git diff -- api/oss/src/utils/locking.py | sed -n '1,180p'

Repository: Agenta-AI/agenta

Length of output: 4012


Track cleanup for the primary claim.

If cancellation occurs after Redis writes lock_key but before _lock_engine.set() returns, finally releases only legacy_key. The primary lock remains until its TTL expires, so new callers can receive a blocked result while no caller owns the critical section.

Set a primary cleanup obligation before the primary SET. Clear it after a confirmed refusal or successful ownership transfer. Release lock_key with _release_if_owner in finally, using the existing best-effort cleanup pattern. Add a regression test that suspends after writing lock_key, then cancels the task.

if not await _lock_engine.set(legacy_key, lock_owner, nx=True, ex=ttl):
Comment thread
coderabbitai[bot] marked this conversation as resolved.
# A confirmed refusal — the key belongs to someone else, so there is
# nothing of ours to drop.
legacy_claimed = False
if LOCK_DEBUG:
log.debug(
"[lock] BLOCKED",
key=legacy_key,
)
return None

# Atomic SET NX: Returns True if lock acquired, False if already held
acquired = await _lock_engine.set(lock_key, lock_owner, nx=True, ex=ttl)

if acquired:
# The caller owns both keys from here; `release_lock` clears them together.
legacy_claimed = False
if LOCK_DEBUG:
log.debug(
"[lock] ACQUIRED",
Expand All @@ -137,6 +234,12 @@ async def acquire_lock(
)
return lock_owner
else:
# This caller is not entering the critical section, so it must not leave the
# legacy key held until its TTL — that would block everyone for `ttl`.
if legacy_key is not None:
await _release_if_owner(legacy_key, lock_owner)
legacy_claimed = False

if LOCK_DEBUG:
log.debug(
"[lock] BLOCKED",
Expand All @@ -153,6 +256,21 @@ async def acquire_lock(
raise
return None

finally:
# Reached when claiming the primary key raised or the task was cancelled between
# the two sets. The section went to nobody, so leaving the legacy key held would
# block both generations for its full TTL. Cancellation matters as much as an
# exception here, which is why this is a `finally` rather than an except branch.
if legacy_claimed and legacy_key is not None:
try:
await _release_if_owner(legacy_key, lock_owner)
except Exception as cleanup_error: # pragma: no cover - best effort
log.error(
f"[lock] LEGACY CLEANUP ERROR: namespace={namespace} "
f"key={key} error={cleanup_error}",
exc_info=True,
)


async def renew_lock(
namespace: str,
Expand Down Expand Up @@ -180,23 +298,24 @@ async def renew_lock(
True if lock was renewed, False if lock has already expired or on error
"""
try:
lock_key = _pack(
namespace=f"lock:{namespace}",
lock_key, legacy_key = _lock_keys(
namespace=namespace,
key=key,
project_id=project_id,
user_id=user_id,
)

if owner:
renewed = await _lock_engine.eval(
_LOCK_RENEW_IF_OWNER_SCRIPT,
1,
lock_key,
owner,
str(ttl),
)
else:
renewed = await _lock_engine.expire(lock_key, ttl)
renewed = await _renew_if_owner(lock_key, owner, ttl)

# Held for as long as the lock itself, or a pod on the previous release would
# take it the moment it lapsed while this holder was still inside the section.
# Its renewal has to count: if the legacy key is gone while the primary survives,
# the section is no longer mutually exclusive across releases, and reporting
# success would leave the caller believing otherwise. Renew it either way rather
# than short-circuiting, so the primary is still extended when the legacy key is
# what failed.
if legacy_key is not None:
renewed = await _renew_if_owner(legacy_key, owner, ttl) and renewed

if renewed:
if LOCK_DEBUG:
Expand Down Expand Up @@ -250,22 +369,19 @@ async def release_lock(
await release_lock(namespace="account-creation", key=email)
"""
try:
lock_key = _pack(
namespace=f"lock:{namespace}",
lock_key, legacy_key = _lock_keys(
namespace=namespace,
key=key,
project_id=project_id,
user_id=user_id,
)

if owner:
deleted = await _lock_engine.eval(
_LOCK_RELEASE_IF_OWNER_SCRIPT,
1,
lock_key,
owner,
)
else:
deleted = await _lock_engine.delete(lock_key)
deleted = await _release_if_owner(lock_key, owner)

# Released even when the primary was already gone: the two were taken together,
# so leaving this one behind would block the section for the rest of its TTL.
if legacy_key is not None:
await _release_if_owner(legacy_key, owner)

if deleted:
if LOCK_DEBUG:
Expand Down
Loading
Loading