Skip to content

Feature delete expired threads in mongo aio checkpointer #81

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Closed
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
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
import asyncio
import builtins
import logging
import sys
import threading
import time
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import Any, Optional
Expand Down Expand Up @@ -75,6 +78,10 @@ def __init__(
db_name: str = "checkpointing_db",
checkpoint_collection_name: str = "checkpoints_aio",
writes_collection_name: str = "checkpoint_writes_aio",
auto_delete_expired_threads: bool = False,
thread_status_collection_name: str = "checkpoint_thread_status_aio",
thread_expire_time_second: int = 2592000, #30 days
thread_expire_check_time_second: int = 86400, # 1 day
**kwargs: Any,
) -> None:
super().__init__()
Expand All @@ -83,6 +90,14 @@ def __init__(
self.checkpoint_collection = self.db[checkpoint_collection_name]
self.writes_collection = self.db[writes_collection_name]
self.loop = asyncio.get_running_loop()
self.auto_delete_expired_threads = auto_delete_expired_threads
if self.auto_delete_expired_threads:
self.thread_expire_time_second = thread_expire_time_second
self.thread_expire_check_time_second = thread_expire_check_time_second
self.thread_status_collection = self.db[thread_status_collection_name]
self.last_thread_expire_check_time_second = 0
self.thread_expire_check_lock = threading.RLock()
self.try_delete_expired_threads_from_checkpoints()

@classmethod
@asynccontextmanager
Expand Down Expand Up @@ -290,6 +305,18 @@ async def aput(
await self.checkpoint_collection.update_one(
upsert_query, {"$set": doc}, upsert=True
)
if self.auto_delete_expired_threads:
thread_status_collection_doc = {
"thread_id": thread_id,
"last_update": int(time.time())
}
thread_status_collection_upsert_query = {
"thread_id": thread_id
}
await self.thread_status_collection.update_one(
thread_status_collection_upsert_query, {"$set": thread_status_collection_doc}, upsert=True
)
await self.try_delete_expired_threads_from_checkpoints()
return {
"configurable": {
"thread_id": thread_id,
Expand Down Expand Up @@ -450,3 +477,34 @@ def put_writes(
return asyncio.run_coroutine_threadsafe(
self.aput_writes(config, writes, task_id), self.loop
).result()

async def try_delete_expired_threads_from_checkpoints(self):
self.thread_expire_check_lock.acquire()
if self.last_thread_expire_check_time_second + self.thread_expire_check_time_second < int(time.time()):
try:
thread_status_collection_result = self.thread_status_collection.find(
{
"last_update" : {
"$lt": int(time.time()) - self.thread_expire_time_second
}
}
)
expired_thread_ids = []
async for doc in thread_status_collection_result:
if doc["last_update"] + self.thread_expire_time_second < int(time.time()):
thread_id = doc["thread_id"]
expired_thread_ids.append(thread_id)
await self.checkpoint_collection.delete_many({
"thread_id": {
"$in": expired_thread_ids
}
})
await self.thread_status_collection.delete_many({
"thread_id": {
"$in": expired_thread_ids
}
})
except Exception as e:
logging.error("delete_expired_threads_from_checkpoints error:", e)
self.last_thread_expire_check_time_second = int(time.time())
self.thread_expire_check_lock.release()