319 lines
10 KiB
Python
319 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import json
|
|
from typing import Any, Protocol
|
|
|
|
import redis.asyncio as redis
|
|
from redis.exceptions import ResponseError
|
|
|
|
from app.core.config import settings
|
|
from app.core.logging import get_logger
|
|
|
|
|
|
logger = get_logger(__name__, service="earth_news")
|
|
|
|
TARGET_LOCATION_STREAM = "earth_news:target_location:jobs"
|
|
TARGET_LOCATION_PRIORITY_STREAM = "earth_news:target_location:priority"
|
|
TARGET_LOCATION_GROUP = "earth_news_target_location"
|
|
TARGET_LOCATION_DEAD_LETTER_STREAM = "earth_news:target_location:dead"
|
|
TARGET_LOCATION_RESULT_TTL_SECONDS = 60 * 60 * 12
|
|
TARGET_LOCATION_JOB_DEDUP_TTL_SECONDS = 60 * 60 * 6
|
|
TARGET_LOCATION_PRIORITY_JOB_DEDUP_TTL_SECONDS = 60 * 5
|
|
TARGET_LOCATION_PENDING_RECLAIM_IDLE_MS = 2 * 60 * 1000
|
|
TARGET_LOCATION_PRIORITY_READ_BLOCK_MS = 1
|
|
TARGET_LOCATION_MAX_ATTEMPTS = 3
|
|
|
|
_redis_client: redis.Redis | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NewsTargetLocationMessage:
|
|
message_id: str
|
|
item_id: str
|
|
payload: dict[str, Any]
|
|
stream_name: str = TARGET_LOCATION_STREAM
|
|
attempts: int = 0
|
|
|
|
|
|
class NewsTargetLocationQueue(Protocol):
|
|
async def enqueue(self, *, item_id: str, payload: dict[str, Any], force: bool = False) -> bool:
|
|
...
|
|
|
|
async def consume_batch(
|
|
self,
|
|
*,
|
|
consumer_name: str,
|
|
count: int,
|
|
block_ms: int,
|
|
) -> list[NewsTargetLocationMessage]:
|
|
...
|
|
|
|
async def ack(self, message: NewsTargetLocationMessage) -> None:
|
|
...
|
|
|
|
async def retry_or_dead_letter(
|
|
self,
|
|
message: NewsTargetLocationMessage,
|
|
*,
|
|
error: str,
|
|
) -> None:
|
|
...
|
|
|
|
|
|
def _get_redis_client() -> redis.Redis:
|
|
global _redis_client
|
|
if _redis_client is None:
|
|
_redis_client = redis.from_url(settings.REDIS_URL, decode_responses=True)
|
|
return _redis_client
|
|
|
|
|
|
def _result_key(item_id: str) -> str:
|
|
return f"earth_news:target_location:result:{item_id}"
|
|
|
|
|
|
def _queued_key(item_id: str) -> str:
|
|
return f"earth_news:target_location:queued:{item_id}"
|
|
|
|
|
|
def _priority_queued_key(item_id: str) -> str:
|
|
return f"earth_news:target_location:priority_queued:{item_id}"
|
|
|
|
|
|
class RedisStreamsNewsTargetLocationQueue:
|
|
def __init__(self, client: redis.Redis | None = None) -> None:
|
|
self.client = client or _get_redis_client()
|
|
self._group_ready = False
|
|
|
|
async def _ensure_group(self) -> None:
|
|
if self._group_ready:
|
|
return
|
|
for stream_name in (TARGET_LOCATION_PRIORITY_STREAM, TARGET_LOCATION_STREAM):
|
|
try:
|
|
await self.client.xgroup_create(
|
|
stream_name,
|
|
TARGET_LOCATION_GROUP,
|
|
id="0",
|
|
mkstream=True,
|
|
)
|
|
except ResponseError as exc:
|
|
if "BUSYGROUP" not in str(exc):
|
|
raise
|
|
self._group_ready = True
|
|
|
|
async def enqueue(self, *, item_id: str, payload: dict[str, Any], force: bool = False) -> bool:
|
|
await self._ensure_group()
|
|
if force:
|
|
await self.client.delete(_result_key(item_id))
|
|
queued_key = _priority_queued_key(item_id)
|
|
elif await self.client.exists(_result_key(item_id)):
|
|
return False
|
|
else:
|
|
queued_key = _queued_key(item_id)
|
|
dedup_ttl = (
|
|
TARGET_LOCATION_PRIORITY_JOB_DEDUP_TTL_SECONDS
|
|
if force
|
|
else TARGET_LOCATION_JOB_DEDUP_TTL_SECONDS
|
|
)
|
|
queued = await self.client.set(
|
|
queued_key,
|
|
"1",
|
|
nx=True,
|
|
ex=dedup_ttl,
|
|
)
|
|
if not queued:
|
|
return bool(await self.client.exists(queued_key))
|
|
stream_name = TARGET_LOCATION_PRIORITY_STREAM if force else TARGET_LOCATION_STREAM
|
|
await self.client.xadd(
|
|
stream_name,
|
|
{
|
|
"item_id": item_id,
|
|
"attempts": "0",
|
|
"payload": json.dumps(payload, ensure_ascii=False),
|
|
},
|
|
)
|
|
return True
|
|
|
|
async def consume_batch(
|
|
self,
|
|
*,
|
|
consumer_name: str,
|
|
count: int,
|
|
block_ms: int,
|
|
) -> list[NewsTargetLocationMessage]:
|
|
await self._ensure_group()
|
|
streams = []
|
|
priority_claimed = await self._claim_stale_messages(
|
|
stream_name=TARGET_LOCATION_PRIORITY_STREAM,
|
|
consumer_name=consumer_name,
|
|
count=count,
|
|
)
|
|
if priority_claimed:
|
|
return priority_claimed
|
|
|
|
priority_messages = await self.client.xreadgroup(
|
|
TARGET_LOCATION_GROUP,
|
|
consumer_name,
|
|
{TARGET_LOCATION_PRIORITY_STREAM: ">"},
|
|
count=count,
|
|
block=TARGET_LOCATION_PRIORITY_READ_BLOCK_MS,
|
|
)
|
|
if priority_messages:
|
|
streams = priority_messages
|
|
else:
|
|
regular_claimed = await self._claim_stale_messages(
|
|
stream_name=TARGET_LOCATION_STREAM,
|
|
consumer_name=consumer_name,
|
|
count=count,
|
|
)
|
|
if regular_claimed:
|
|
return regular_claimed
|
|
streams = await self.client.xreadgroup(
|
|
TARGET_LOCATION_GROUP,
|
|
consumer_name,
|
|
{TARGET_LOCATION_STREAM: ">"},
|
|
count=count,
|
|
block=block_ms,
|
|
)
|
|
messages: list[NewsTargetLocationMessage] = []
|
|
for stream_name, stream_messages in streams:
|
|
for message_id, fields in stream_messages:
|
|
message = await self._message_from_fields(stream_name, message_id, fields)
|
|
if message is not None:
|
|
messages.append(message)
|
|
return messages
|
|
|
|
async def _claim_stale_messages(
|
|
self,
|
|
*,
|
|
stream_name: str,
|
|
consumer_name: str,
|
|
count: int,
|
|
) -> list[NewsTargetLocationMessage]:
|
|
try:
|
|
_next_id, claimed, _deleted = await self.client.xautoclaim(
|
|
stream_name,
|
|
TARGET_LOCATION_GROUP,
|
|
consumer_name,
|
|
TARGET_LOCATION_PENDING_RECLAIM_IDLE_MS,
|
|
start_id="0-0",
|
|
count=count,
|
|
)
|
|
except ResponseError:
|
|
return []
|
|
messages: list[NewsTargetLocationMessage] = []
|
|
for message_id, fields in claimed:
|
|
message = await self._message_from_fields(stream_name, message_id, fields)
|
|
if message is not None:
|
|
messages.append(message)
|
|
return messages
|
|
|
|
async def _message_from_fields(
|
|
self,
|
|
stream_name: str,
|
|
message_id: str,
|
|
fields: dict[str, str],
|
|
) -> NewsTargetLocationMessage | None:
|
|
raw_payload = fields.get("payload")
|
|
item_id = fields.get("item_id")
|
|
if not raw_payload or not item_id:
|
|
await self._discard_message(stream_name, message_id)
|
|
return None
|
|
try:
|
|
payload = json.loads(raw_payload)
|
|
except json.JSONDecodeError:
|
|
await self._discard_message(stream_name, message_id)
|
|
return None
|
|
attempts = int(fields.get("attempts") or 0)
|
|
return NewsTargetLocationMessage(
|
|
message_id=message_id,
|
|
item_id=item_id,
|
|
payload=payload,
|
|
stream_name=stream_name,
|
|
attempts=attempts,
|
|
)
|
|
|
|
async def ack(self, message: NewsTargetLocationMessage) -> None:
|
|
await self.client.xack(message.stream_name, TARGET_LOCATION_GROUP, message.message_id)
|
|
await self.client.xdel(message.stream_name, message.message_id)
|
|
|
|
async def _discard_message(self, stream_name: str, message_id: str) -> None:
|
|
await self.client.xack(stream_name, TARGET_LOCATION_GROUP, message_id)
|
|
await self.client.xdel(stream_name, message_id)
|
|
|
|
async def retry_or_dead_letter(
|
|
self,
|
|
message: NewsTargetLocationMessage,
|
|
*,
|
|
error: str,
|
|
) -> None:
|
|
await self.ack(message)
|
|
if message.attempts + 1 >= TARGET_LOCATION_MAX_ATTEMPTS:
|
|
await self.client.xadd(
|
|
TARGET_LOCATION_DEAD_LETTER_STREAM,
|
|
{
|
|
"item_id": message.item_id,
|
|
"attempts": str(message.attempts + 1),
|
|
"error": error,
|
|
"payload": json.dumps(message.payload, ensure_ascii=False),
|
|
},
|
|
)
|
|
return
|
|
await self.client.xadd(
|
|
message.stream_name,
|
|
{
|
|
"item_id": message.item_id,
|
|
"attempts": str(message.attempts + 1),
|
|
"payload": json.dumps(message.payload, ensure_ascii=False),
|
|
},
|
|
)
|
|
|
|
|
|
def get_news_target_location_queue() -> NewsTargetLocationQueue:
|
|
return RedisStreamsNewsTargetLocationQueue()
|
|
|
|
|
|
async def enqueue_target_location_job(payload: dict[str, Any], *, force: bool = False) -> bool:
|
|
item_id = str(payload.get("id") or "")
|
|
if not item_id:
|
|
return False
|
|
try:
|
|
queue = get_news_target_location_queue()
|
|
return await queue.enqueue(item_id=item_id, payload=payload, force=force)
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Failed to enqueue Earth news target location job",
|
|
event="earth_news.target_location.enqueue_failed",
|
|
context={"item_id": item_id, "error": str(exc)},
|
|
)
|
|
return False
|
|
|
|
|
|
async def get_cached_target_location_patch(item_id: str) -> dict[str, Any] | None:
|
|
try:
|
|
raw_value = await _get_redis_client().get(_result_key(item_id))
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Failed to read Earth news target location cache",
|
|
event="earth_news.target_location.cache_read_failed",
|
|
context={"item_id": item_id, "error": str(exc)},
|
|
)
|
|
return None
|
|
if not raw_value:
|
|
return None
|
|
try:
|
|
value = json.loads(raw_value)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return value if isinstance(value, dict) else None
|
|
|
|
|
|
async def save_target_location_patch(item_id: str, patch: dict[str, Any]) -> None:
|
|
client = _get_redis_client()
|
|
await client.setex(
|
|
_result_key(item_id),
|
|
TARGET_LOCATION_RESULT_TTL_SECONDS,
|
|
json.dumps(patch, ensure_ascii=False),
|
|
)
|
|
await client.delete(_queued_key(item_id), _priority_queued_key(item_id))
|