43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import time
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class RefIndex:
|
||
|
|
def __init__(self, ttl_seconds: int = 3600):
|
||
|
|
self._index: dict[str, dict] = {}
|
||
|
|
self._ttl = ttl_seconds
|
||
|
|
|
||
|
|
def add(self, msg_id: str, msg_seq: int | None, reply_to_id: str | None = None) -> None:
|
||
|
|
self._index[msg_id] = {
|
||
|
|
"msg_seq": msg_seq,
|
||
|
|
"reply_to_id": reply_to_id,
|
||
|
|
"timestamp": time.time(),
|
||
|
|
}
|
||
|
|
self._cleanup()
|
||
|
|
|
||
|
|
def get(self, msg_id: str) -> dict | None:
|
||
|
|
entry = self._index.get(msg_id)
|
||
|
|
if entry and time.time() - entry["timestamp"] < self._ttl:
|
||
|
|
return entry
|
||
|
|
if entry:
|
||
|
|
del self._index[msg_id]
|
||
|
|
return None
|
||
|
|
|
||
|
|
def find_by_reply(self, msg_id: str) -> list[str]:
|
||
|
|
result = []
|
||
|
|
for k, v in self._index.items():
|
||
|
|
if v.get("reply_to_id") == msg_id:
|
||
|
|
if time.time() - v["timestamp"] < self._ttl:
|
||
|
|
result.append(k)
|
||
|
|
return result
|
||
|
|
|
||
|
|
def _cleanup(self) -> None:
|
||
|
|
now = time.time()
|
||
|
|
expired = [k for k, v in self._index.items() if now - v["timestamp"] > self._ttl]
|
||
|
|
for k in expired:
|
||
|
|
del self._index[k]
|