async def poll_dynamic_task(
self,
utc_now: datetime.datetime,
task: Task,
continuation_token: Optional[ContinuationToken] = None,
) -> QueryResponse:
response = await self._table.query(
IndexName="ready_at",
Select="ALL_ATTRIBUTES",
Limit=self._batch_size,
KeyConditionExpression=Key("partition_id").eq(
self._get_partition_id(task))
& Key("ready_at").lt(f"{utc_now.isoformat()}`"),
**({
"ExclusiveStartKey": continuation_token
} if continuation_token is not None else {}),
)
return QueryResponse(
tasks=[(
task.instantiate_from_spec(json.loads(record["task_spec"])),
Lease(_TaskRecord.from_dynamo_item(record)),
) for record in response["Items"]],
continuation_token=response.get("LastEvaluatedKey", None),
)
async def poll_dynamic_task(
self,
utc_now: datetime.datetime,
task: Task,
continuation_token: Optional[ContinuationToken] = None,
) -> QueryResponse:
async with self._transaction() as cursor:
locked_by = str(uuid.uuid4())
locked_until = utc_now + task.lease_duration
await cursor.execute(
f"""
SELECT name, task_spec FROM {self._table_name}
WHERE parent_name = %s AND GREATEST(COALESCE(locked_until, 0), COALESCE(execute_after, 0)) <= %s
ORDER BY GREATEST(COALESCE(locked_until, 0), COALESCE(execute_after, 0)) ASC
LIMIT %s
FOR UPDATE SKIP LOCKED
""",
(
task.canonical_name,
_to_timestamp(utc_now),
self._batch_size,
),
)
ready_tasks = await cursor.fetchall()
await cursor.executemany(
f"""
UPDATE {self._table_name}
SET
locked_until = %s,
locked_by = %s
WHERE name = %s
""",
[(_to_timestamp(locked_until), locked_by, record["name"])
for record in ready_tasks],
)
logger.debug(f"poll_dynamic_task returned {ready_tasks}")
return QueryResponse(
tasks=[(
task.instantiate_from_spec(json.loads(
record["task_spec"])),
Lease(locked_by),
) for record in ready_tasks],
# May result in an extra round-trip if there were exactly
# batch_size tasks available, but we deem this an acceptable
# tradeoff.
continuation_token=_CONTINUATION_TOKEN
if len(ready_tasks) == self._batch_size else None,
)
async def extend_lease(self, utc_now: datetime.datetime, task: Task,
lease: Lease) -> Optional[Lease]:
assert isinstance(lease, _ManageScriptResponse)
new_locked_until = utc_now + task.lease_duration
response = await self._manage_record(task, "EXTEND", lease.version,
lease.locked_by,
new_locked_until.isoformat())
task._last_lease = response # type: ignore
if response.result == ResultType.READY:
return Lease(response)
else:
return None
async def poll_dynamic_task(
self,
utc_now: datetime.datetime,
task: Task,
continuation_token: Optional[ContinuationToken] = None,
) -> QueryResponse:
async with self._transaction() as connection:
locked_by = uuid.uuid4()
locked_until = utc_now + task.lease_duration
ready_tasks = await connection.fetch(
f"""
UPDATE {self._table_name} a
SET
locked_until = $4,
locked_by = $5
FROM (
SELECT name FROM {self._table_name}
WHERE parent_name = $1 AND GREATEST(locked_until, execute_after) <= $2
ORDER BY GREATEST(locked_until, execute_after) ASC
LIMIT $3
FOR UPDATE SKIP LOCKED
) b
WHERE a.name = b.name
RETURNING *
""",
task.canonical_name,
utc_now,
self._batch_size,
locked_until,
locked_by,
)
logger.debug(f"poll_dynamic_task returned {ready_tasks}")
return QueryResponse(
tasks=[(
task.instantiate_from_spec(json.loads(
record["task_spec"])),
Lease(locked_by),
) for record in ready_tasks],
# May result in an extra round-trip if there were exactly
# batch_size tasks available, but we deem this an acceptable
# tradeoff.
continuation_token=_CONTINUATION_TOKEN
if len(ready_tasks) == self._batch_size else None,
)
async def poll_dynamic_task(
self,
utc_now: datetime.datetime,
task: Task,
continuation_token: Optional[ContinuationToken] = None,
) -> QueryResponse:
async with self._transaction(explicit_begin=True):
locked_by = uuid.uuid4()
locked_until = utc_now + task.lease_duration
ready_tasks = await self._connection.execute_fetchall(
f"""SELECT * FROM {self._table_name}
WHERE parent_name = $1 AND MAX(COALESCE(locked_until, 0), COALESCE(execute_after, 0)) <= $2
ORDER BY MAX(COALESCE(locked_until, 0), COALESCE(execute_after, 0)) ASC
LIMIT $3
""",
(task.canonical_name, _to_timestamp(utc_now),
self._batch_size),
)
concrete_tasks = [
task.instantiate_from_spec(json.loads(record["task_spec"]))
for record in ready_tasks
]
await self._connection.executemany(
f"""
UPDATE {self._table_name}
SET
locked_until = $2,
locked_by = $3
WHERE name = $1
""",
[(
concrete_task.canonical_name,
_to_timestamp(locked_until),
str(locked_by),
) for concrete_task in concrete_tasks],
)
return QueryResponse(
tasks=[(concrete_task, Lease(locked_by))
for concrete_task in concrete_tasks],
continuation_token=_CONTINUATION_TOKEN
if len(concrete_tasks) == self._batch_size else None,
)
async def extend_lease(self, utc_now: datetime.datetime, task: Task,
lease: Lease) -> Optional[Lease]:
assert isinstance(lease, _TaskRecord)
record: Optional[_TaskRecord] = cast(_TaskRecord, copy.copy(lease))
for _ in range(MAX_OPTIMISTIC_RETRY_COUNT):
if not record:
logger.warning(f"Task {task} no longer exists, skipping.")
return None
if record.locked_by != lease.locked_by:
logger.warning(f"Lease lost on task {task}, skipping.")
return None
record.locked_until = utc_now + task.lease_duration
if await self._update_item(task, record):
return Lease(record)
# If the update fails due to version mismatch
record = await self._retreive_item(task, consistent_read=True)
raise PyncetteException(
"Unable to acquire the lock on the task due to contention")
def _create_dynamic_task(task: Task,
response_data: List[bytes]) -> Tuple[Task, Lease]:
task_data = _ManageScriptResponse.from_response(response_data)
assert task_data.task_spec is not None
return (task.instantiate_from_spec(task_data.task_spec), Lease(task_data))
async def poll_task(self,
utc_now: datetime.datetime,
task: Task,
lease: Optional[Lease] = None) -> PollResponse:
# Nominally, we need at least two round-trips to Redis since the next execute_after is calculated
# in Python code due to extra flexibility. This is why we have optimistic locking below to ensure that
# the next execution time was calculated using a correct base if another process modified it in between.
# In most cases, however, we can assume that the base time has not changed since the last invocation,
# so by caching it, we can poll a task using a single round-trip (if we are wrong, the loop below will still
# ensure corretness as the version will not match).
last_lease = getattr(task, "_last_lease", None)
if isinstance(lease, _ManageScriptResponse):
version, execute_after, locked_by = (
lease.version,
lease.execute_after,
lease.locked_by,
)
elif last_lease is not None:
logger.debug("Using cached values for execute_after")
version, execute_after, locked_by = (
last_lease.version,
last_lease.execute_after,
str(uuid.uuid4()),
)
else:
# By default we assume that the task is brand new
version, execute_after, locked_by = (
0,
None,
str(uuid.uuid4()),
)
new_locked_until = utc_now + task.lease_duration
for _ in range(5):
next_execution = task.get_next_execution(utc_now, execute_after)
response = await self._manage_record(
task,
"POLL",
task.execution_mode.name,
"REGULAR" if task.parent_task is None else "DYNAMIC",
utc_now.isoformat(),
version,
next_execution.isoformat()
if next_execution is not None else "",
new_locked_until.isoformat(),
locked_by,
)
task._last_lease = response # type: ignore
if response.result == ResultType.LEASE_MISMATCH:
logger.debug("Lease mismatch, retrying.")
execute_after = response.execute_after
version = response.version
elif response.result == ResultType.MISSING:
raise PyncetteException("Task not found")
else:
return PollResponse(
result=response.result,
scheduled_at=execute_after,
lease=Lease(response),
)
raise PyncetteException(
"Unable to acquire the lock on the task due to contention")
from typing import Any
from typing import AsyncIterator
from typing import Optional
from pyncette.model import ContinuationToken
from pyncette.model import Lease
from pyncette.model import PollResponse
from pyncette.model import QueryResponse
from pyncette.model import ResultType
from pyncette.repository import Repository
from pyncette.task import Task
logger = logging.getLogger(__name__)
_LEASE = Lease(object())
_TASK_SPEC = {
"name": "fake",
"interval": None,
"timezone": None,
"execute_at": None,
"extra_args": {},
"schedule": "* * * * * *",
}
class FakeRepository(Repository):
"""Redis-backed store for Pyncete task execution data"""
_batch_size: int
_records_per_tick: int
async def poll_task(self,
utc_now: datetime.datetime,
task: Task,
lease: Optional[Lease] = None) -> PollResponse:
last_lease = getattr(task, "_last_lease", None)
# Similar logic as in Redis repository. If we have previously processed this
# task in any manner, we try to reuse the latest state of the task we have at hand
# from cache (or lease) to avoid two roundtrips to DynamoDB in the optimistic case.
# If we are wrong, we will get a version mismatch, whereby we will load the current
# state from DB.
# However, in case the task would be PENDING or LOCKED, this will result in no requests
# being made to the DB at all. For pending this is OK, but for locked, we want to revalidate
# in order to be able to execute the task as soon as it is unlocked.
record: Optional[_TaskRecord]
potentially_stale = False
if lease is not None:
assert isinstance(lease, _TaskRecord)
record = cast(_TaskRecord, copy.copy(lease))
elif last_lease is not None:
logger.debug("Using cached values for last lease")
assert isinstance(last_lease, _TaskRecord)
record = cast(_TaskRecord, copy.copy(last_lease))
potentially_stale = True
else:
record = await self._retreive_item(task)
for _ in range(MAX_OPTIMISTIC_RETRY_COUNT):
must_revalidate = False
update = False
if record is None:
if task.parent_task is not None:
raise PyncetteException("Task not found")
record = _TaskRecord(
execute_after=task.get_next_execution(utc_now, None),
locked_until=None,
locked_by=None,
version=0,
)
update = True
assert record.execute_after is not None
scheduled_at = record.execute_after
if (record.locked_until is not None
and record.locked_until > utc_now and
(lease is None or lease.locked_by != record.locked_by)):
result = ResultType.LOCKED
if potentially_stale:
must_revalidate = True
elif (record.execute_after <= utc_now
and task.execution_mode == ExecutionMode.AT_MOST_ONCE):
result = ResultType.READY
record.execute_after = task.get_next_execution(
utc_now, record.execute_after)
record.locked_until = None
record.locked_by = None
update = True
elif (record.execute_after <= utc_now
and task.execution_mode == ExecutionMode.AT_LEAST_ONCE):
result = ResultType.READY
record.locked_until = utc_now + task.lease_duration
record.locked_by = str(uuid.uuid4())
update = True
else:
result = ResultType.PENDING
if must_revalidate or (update and
not await self._update_item(task, record)):
logger.debug("Using cached values for last lease")
record = await self._retreive_item(task, consistent_read=True)
potentially_stale = False
continue
return PollResponse(result=result,
scheduled_at=scheduled_at,
lease=Lease(record))
raise PyncetteException(
"Unable to acquire the lock on the task due to contention")