Validate permission locks (#3799)

* WIP for external group sync lock fixes

* prototyping permissions validation

* validate permission sync tasks in celery

* mypy

* cleanup and wire off external group sync checks for now

* add active key to reset

* improve logging

* reset on payload format change

* return False on exception

* missed a return

* add count of tasks scanned

* add comment

* better logging

* add return

* more return

* catch payload exceptions

* code review fixes

* push to restart test

---------

Co-authored-by: Richard Kuo (Danswer) <rkuo@onyx.app>

backend/onyx/redis/redis_connector_doc_perm_sync.py

@@ -17,6 +17,8 @@ from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 
 
 class RedisConnectorPermissionSyncPayload(BaseModel):
+    id: str
+    submitted: datetime
     started: datetime | None
     celery_task_id: str | None
 
@@ -41,6 +43,12 @@ class RedisConnectorPermissionSync:
     TASKSET_PREFIX = f"{PREFIX}_taskset"  # connectorpermissions_taskset
     SUBTASK_PREFIX = f"{PREFIX}+sub"  # connectorpermissions+sub
 
+    # used to signal the overall workflow is still active
+    # it's impossible to get the exact state of the system at a single point in time
+    # so we need a signal with a TTL to bridge gaps in our checks
+    ACTIVE_PREFIX = PREFIX + "_active"
+    ACTIVE_TTL = 3600
+
     def __init__(self, tenant_id: str | None, id: int, redis: redis.Redis) -> None:
         self.tenant_id: str | None = tenant_id
         self.id = id
@@ -54,6 +62,7 @@ class RedisConnectorPermissionSync:
         self.taskset_key = f"{self.TASKSET_PREFIX}_{id}"
 
         self.subtask_prefix: str = f"{self.SUBTASK_PREFIX}_{id}"
+        self.active_key = f"{self.ACTIVE_PREFIX}_{id}"
 
     def taskset_clear(self) -> None:
         self.redis.delete(self.taskset_key)
@@ -107,6 +116,20 @@ class RedisConnectorPermissionSync:
 
         self.redis.set(self.fence_key, payload.model_dump_json())
 
+    def set_active(self) -> None:
+        """This sets a signal to keep the permissioning flow from getting cleaned up within
+        the expiration time.
+
+        The slack in timing is needed to avoid race conditions where simply checking
+        the celery queue and task status could result in race conditions."""
+        self.redis.set(self.active_key, 0, ex=self.ACTIVE_TTL)
+
+    def active(self) -> bool:
+        if self.redis.exists(self.active_key):
+            return True
+
+        return False
+
     @property
     def generator_complete(self) -> int | None:
         """the fence payload is an int representing the starting number of
@@ -173,6 +196,7 @@ class RedisConnectorPermissionSync:
         return len(async_results)
 
     def reset(self) -> None:
+        self.redis.delete(self.active_key)
         self.redis.delete(self.generator_progress_key)
         self.redis.delete(self.generator_complete_key)
         self.redis.delete(self.taskset_key)
@@ -187,6 +211,9 @@ class RedisConnectorPermissionSync:
     @staticmethod
     def reset_all(r: redis.Redis) -> None:
         """Deletes all redis values for all connectors"""
+        for key in r.scan_iter(RedisConnectorPermissionSync.ACTIVE_PREFIX + "*"):
+            r.delete(key)
+
         for key in r.scan_iter(RedisConnectorPermissionSync.TASKSET_PREFIX + "*"):
             r.delete(key)
 

backend/onyx/redis/redis_connector_ext_group_sync.py

@@ -11,6 +11,7 @@ from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 
 
 class RedisConnectorExternalGroupSyncPayload(BaseModel):
+    submitted: datetime
     started: datetime | None
     celery_task_id: str | None
 
@@ -135,6 +136,12 @@ class RedisConnectorExternalGroupSync:
     ) -> int | None:
         pass
 
+    def reset(self) -> None:
+        self.redis.delete(self.generator_progress_key)
+        self.redis.delete(self.generator_complete_key)
+        self.redis.delete(self.taskset_key)
+        self.redis.delete(self.fence_key)
+
     @staticmethod
     def remove_from_taskset(id: int, task_id: str, r: redis.Redis) -> None:
         taskset_key = f"{RedisConnectorExternalGroupSync.TASKSET_PREFIX}_{id}"

backend/onyx/redis/redis_connector_index.py

@@ -33,8 +33,8 @@ class RedisConnectorIndex:
     TERMINATE_TTL = 600
 
     # used to signal the overall workflow is still active
-    # there are gaps in time between states where we need some slack
-    # to correctly transition
+    # it's impossible to get the exact state of the system at a single point in time
+    # so we need a signal with a TTL to bridge gaps in our checks
     ACTIVE_PREFIX = PREFIX + "_active"
     ACTIVE_TTL = 3600
 

backend/onyx/redis/redis_pool.py

@@ -122,7 +122,7 @@ class TenantRedis(redis.Redis):
             "ttl",
         ]  # Regular methods that need simple prefixing
 
-        if item == "scan_iter":
+        if item == "scan_iter" or item == "sscan_iter":
             return self._prefix_scan_iter(original_attr)
         elif item in methods_to_wrap and callable(original_attr):
             return self._prefix_method(original_attr)