disable thread local locking in callbacks (#3319)

This commit is contained in:
rkuo-danswer
2024-12-03 14:32:56 -08:00
committed by GitHub
parent 66f47d294c
commit 13f6e8a6b4
3 changed files with 12 additions and 2 deletions

View File

@ -39,7 +39,6 @@ from danswer.redis.redis_usergroup import RedisUserGroup
from danswer.utils.logger import setup_logger from danswer.utils.logger import setup_logger
from shared_configs.configs import MULTI_TENANT from shared_configs.configs import MULTI_TENANT
logger = setup_logger() logger = setup_logger()
celery_app = Celery(__name__) celery_app = Celery(__name__)
@ -117,9 +116,13 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
# it is planned to use this lock to enforce singleton behavior on the primary # it is planned to use this lock to enforce singleton behavior on the primary
# worker, since the primary worker does redis cleanup on startup, but this isn't # worker, since the primary worker does redis cleanup on startup, but this isn't
# implemented yet. # implemented yet.
# set thread_local=False since we don't control what thread the periodic task might
# reacquire the lock with
lock: RedisLock = r.lock( lock: RedisLock = r.lock(
DanswerRedisLocks.PRIMARY_WORKER, DanswerRedisLocks.PRIMARY_WORKER,
timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT, timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
thread_local=False,
) )
logger.info("Primary worker lock: Acquire starting.") logger.info("Primary worker lock: Acquire starting.")

View File

@ -789,9 +789,12 @@ def connector_indexing_task(
) )
break break
# set thread_local=False since we don't control what thread the indexing/pruning
# might run our callback with
lock: RedisLock = r.lock( lock: RedisLock = r.lock(
redis_connector_index.generator_lock_key, redis_connector_index.generator_lock_key,
timeout=CELERY_INDEXING_LOCK_TIMEOUT, timeout=CELERY_INDEXING_LOCK_TIMEOUT,
thread_local=False,
) )
acquired = lock.acquire(blocking=False) acquired = lock.acquire(blocking=False)

View File

@ -8,6 +8,7 @@ from celery import shared_task
from celery import Task from celery import Task
from celery.exceptions import SoftTimeLimitExceeded from celery.exceptions import SoftTimeLimitExceeded
from redis import Redis from redis import Redis
from redis.lock import Lock as RedisLock
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
from danswer.background.celery.apps.app_base import task_logger from danswer.background.celery.apps.app_base import task_logger
@ -239,9 +240,12 @@ def connector_pruning_generator_task(
r = get_redis_client(tenant_id=tenant_id) r = get_redis_client(tenant_id=tenant_id)
lock = r.lock( # set thread_local=False since we don't control what thread the indexing/pruning
# might run our callback with
lock: RedisLock = r.lock(
DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{redis_connector.id}", DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{redis_connector.id}",
timeout=CELERY_PRUNING_LOCK_TIMEOUT, timeout=CELERY_PRUNING_LOCK_TIMEOUT,
thread_local=False,
) )
acquired = lock.acquire(blocking=False) acquired = lock.acquire(blocking=False)