Extract run with process lock logic into func. Use for content reindexing

This commit is contained in:
Debanjum Singh Solanky
2024-04-16 23:52:16 +05:30
parent e9f608174b
commit 8e0bae894d
2 changed files with 39 additions and 23 deletions

View File

@@ -307,30 +307,20 @@ def configure_middleware(app):
app.add_middleware(SessionMiddleware, secret_key=os.environ.get("KHOJ_DJANGO_SECRET_KEY", "!secret"))
@schedule.repeat(schedule.every(22).to(25).hours)
def update_content_index():
try:
if ProcessLockAdapters.is_process_locked(ProcessLock.Operation.UPDATE_EMBEDDINGS):
logger.info("🔒 Skipping update content index due to lock")
return
ProcessLockAdapters.set_process_lock(
ProcessLock.Operation.UPDATE_EMBEDDINGS, max_duration_in_seconds=60 * 60 * 2
)
for user in get_all_users():
all_files = collect_files(user=user)
success = configure_content(all_files, user=user)
all_files = collect_files(user=None)
success = configure_content(all_files, user=None)
if not success:
raise RuntimeError("Failed to update content index")
logger.info("📪 Content index updated via Scheduler")
with timer("📬 Updating content index via Scheduler"):
for user in get_all_users():
all_files = collect_files(user=user)
success = configure_content(all_files, user=user)
all_files = collect_files(user=None)
success = configure_content(all_files, user=None)
if not success:
raise RuntimeError("Failed to update content index")
logger.info("📪 Content index updated via Scheduler")
ProcessLockAdapters.remove_process_lock(ProcessLock.Operation.UPDATE_EMBEDDINGS)
except Exception as e:
logger.error(f"🚨 Error updating content index via Scheduler: {e}", exc_info=True)
@schedule.repeat(schedule.every(22).to(25).hours)
def update_content_index_regularly():
ProcessLockAdapters.run_with_lock(update_content_index, ProcessLock.Operation.UPDATE_EMBEDDINGS)
def configure_search_types():

View File

@@ -5,7 +5,7 @@ import secrets
import sys
from datetime import date, datetime, timedelta, timezone
from enum import Enum
from typing import List, Optional, Type
from typing import Callable, List, Optional, Type
from asgiref.sync import sync_to_async
from django.contrib.sessions.backends.db import SessionStore
@@ -46,7 +46,7 @@ from khoj.search_filter.file_filter import FileFilter
from khoj.search_filter.word_filter import WordFilter
from khoj.utils import state
from khoj.utils.config import OfflineChatProcessorModel
from khoj.utils.helpers import generate_random_name, is_none_or_empty
from khoj.utils.helpers import generate_random_name, is_none_or_empty, timer
logger = logging.getLogger(__name__)
@@ -421,6 +421,7 @@ class ProcessLockAdapters:
tz=timezone.utc
):
process_lock.delete()
logger.info(f"🔓 Deleted stale {process_name} process lock on timeout")
return False
return True
@@ -428,6 +429,31 @@ class ProcessLockAdapters:
def remove_process_lock(process_name: str):
return ProcessLock.objects.filter(name=process_name).delete()
@staticmethod
def run_with_lock(func: Callable, operation: ProcessLock.Operation, max_duration_in_seconds: int = 600):
# Exit early if process lock is already taken
if ProcessLockAdapters.is_process_locked(operation):
logger.info(f"🔒 Skip executing {func} as {operation} lock is already taken")
return
success = False
try:
# Set process lock
ProcessLockAdapters.set_process_lock(operation, max_duration_in_seconds)
logger.info(f"🔐 Locked {operation} to execute {func}")
# Execute Function
with timer(f"🔒 Run {func} with {operation} process lock", logger):
func()
success = True
except Exception as e:
logger.error(f"🚨 Error executing {func} with {operation} process lock: {e}", exc_info=True)
success = False
finally:
# Remove Process Lock
ProcessLockAdapters.remove_process_lock(operation)
logger.info(f"🔓 Unlocked {operation} process after executing {func} {'Succeeded' if success else 'Failed'}")
class ClientApplicationAdapters:
@staticmethod