Move org size inc into update_dedupe_index_info

Author: tw4l

backend/btrixcloud/colls.py

@@ -731,7 +731,7 @@ async def create_dedupe_index(self, coll: Collection, org: Organization):
         # enable index by setting indexState to a non-null value
         # and setting stats to zeroed out default
         await self.update_dedupe_index_info(
-            coll.id, state="initing" if coll.crawlCount else "idle"
+            coll.id, org.id, state="initing" if coll.crawlCount else "idle"
         )
 
         await self.update_dedupe_index_stats(coll.id, DedupeIndexStats())
@@ -793,7 +793,7 @@ async def run_index_import_job(
         if job_type in ("import", "purge"):
             # if job created, update state here so its reflected in the UI more quickly
             await self.update_dedupe_index_info(
-                coll_id, state="purging" if job_type == "purge" else "importing"
+                coll_id, oid, state="purging" if job_type == "purge" else "importing"
             )
 
     async def delete_dedupe_index(
@@ -856,6 +856,7 @@ async def update_dedupe_index_stats(
     async def update_dedupe_index_info(
         self,
         coll_id: UUID,
+        oid: UUID,
         state: TYPE_DEDUPE_INDEX_STATES,
         index_file: Optional[DedupeIndexFile] = None,
         dt: Optional[datetime] = None,
@@ -871,7 +872,19 @@ async def update_dedupe_index_info(
             {"$set": query},
             return_document=pymongo.ReturnDocument.BEFORE,
         )
-        return Collection.from_dict(res) if res else None
+
+        if index_file:
+            size_diff = index_file.size
+
+            prev_coll = Collection.from_dict(res) if res else None
+            if prev_coll and prev_coll.indexFile:
+                size_diff = index_file.size - prev_coll.indexFile.size
+
+            await self.orgs.inc_org_bytes_stored_field(
+                oid, "bytesStoredDedupeIndexes", size_diff
+            )
+
+        return res is not None
 
     async def get_dedupe_index_saved(self, coll_id: UUID) -> Optional[datetime]:
         """return datetime for when index was last saved, if any"""

backend/btrixcloud/operator/collindexes.py

@@ -147,6 +147,7 @@ async def sync_index(self, data: MCSyncData):
         status = CollIndexStatus(**data.parent.get("status", {}))
 
         coll_id = spec.id
+        oid = spec.oid
         redis_name = f"redis-coll-{coll_id}"
         new_children = []
 
@@ -159,7 +160,7 @@ async def sync_index(self, data: MCSyncData):
         if data.finalizing:
             is_done = False
             if status.state == "saved" and status.index.savedAt:
-                await self.set_state("idle", status, coll_id)
+                await self.set_state("idle", status, coll_id, oid)
                 is_done = True
             elif status.state == "idle" and status.index.notFound:
                 is_done = True
@@ -168,33 +169,33 @@ async def sync_index(self, data: MCSyncData):
                 is_done = True
             else:
                 try:
-                    await self.coll_ops.get_collection_raw(spec.id, spec.oid)
+                    await self.coll_ops.get_collection_raw(coll_id, oid)
                 # pylint: disable=bare-except
                 except:
                     # collection not found, delete index
                     is_done = True
 
             if is_done:
-                print(f"CollIndex removed: {spec.id}")
+                print(f"CollIndex removed: {coll_id}")
                 return {"status": status.dict(), "children": [], "finalized": True}
 
         try:
             # determine if index was previously saved before initing redis
             if not redis_pod:
                 if not status.indexLastSavedAt:
-                    res = await self.coll_ops.get_dedupe_index_saved(spec.id)
+                    res = await self.coll_ops.get_dedupe_index_saved(coll_id)
                     if res:
                         status.indexLastSavedAt = date_to_str(res)
 
             if self.is_expired(status) or data.finalizing:
                 # do actual deletion here
                 if not data.finalizing:
-                    self.run_task(self.do_delete(spec.id))
+                    self.run_task(self.do_delete(coll_id))
 
                 # Saving process
                 # 1. run bgsave while redis is active
                 if status.index.running:
-                    await self.do_save_redis(spec.id, status)
+                    await self.do_save_redis(coll_id, oid, status)
 
                 elif status.index.finished and not status.index.savedAt:
                     await self.k8s.send_signal_to_pod(redis_name, "SIGUSR1", "save")
@@ -204,7 +205,7 @@ async def sync_index(self, data: MCSyncData):
                     await self.mark_index_saved(redis_name, spec, status)
 
             else:
-                await self.update_state(data, spec.id, status)
+                await self.update_state(data, coll_id, oid, status)
 
         # pylint: disable=broad-exception-caught
         except Exception as e:
@@ -276,7 +277,9 @@ def sync_redis_pod_status(self, pod, status: CollIndexStatus):
             except:
                 pass
 
-    async def update_state(self, data, coll_id: UUID, status: CollIndexStatus):
+    async def update_state(
+        self, data, coll_id: UUID, oid: UUID, status: CollIndexStatus
+    ):
         """update state"""
         desired_state = status.state
         if not status.index.loaded:
@@ -304,7 +307,7 @@ async def update_state(self, data, coll_id: UUID, status: CollIndexStatus):
                 desired_state = "initing"
 
         if desired_state != status.state:
-            await self.set_state(desired_state, status, coll_id)
+            await self.set_state(desired_state, status, coll_id, oid)
 
     def is_expired(self, status: CollIndexStatus):
         """return true if collindex is considered expired and should be deleted"""
@@ -330,22 +333,26 @@ def is_last_active_exceeds(
         return False
 
     async def set_state(
-        self, state: TYPE_DEDUPE_INDEX_STATES, status: CollIndexStatus, coll_id: UUID
+        self,
+        state: TYPE_DEDUPE_INDEX_STATES,
+        status: CollIndexStatus,
+        coll_id: UUID,
+        oid: UUID,
     ):
         """set state after updating db"""
         print(f"Setting coll index state {status.state} -> {state} {coll_id}")
         status.state = state
         status.lastStateChangeAt = date_to_str(dt_now())
 
-        # self.run_task(self.coll_ops.update_dedupe_index_info(coll_id, state))
-        await self.coll_ops.update_dedupe_index_info(coll_id, state)
+        # self.run_task(self.coll_ops.update_dedupe_index_info(coll_id, oid, state))
+        await self.coll_ops.update_dedupe_index_info(coll_id, oid, state)
 
     async def do_delete(self, coll_id: UUID):
         """delete the CollIndex object"""
         print(f"Deleting collindex {coll_id}")
         await self.k8s.delete_custom_object(f"collindex-{coll_id}", "collindexes")
 
-    async def do_save_redis(self, coll_id: UUID, status: CollIndexStatus):
+    async def do_save_redis(self, coll_id: UUID, oid: UUID, status: CollIndexStatus):
         """shutdown save redis"""
         try:
             redis = await self.k8s.get_redis_connected(f"coll-{coll_id}")
@@ -355,14 +362,14 @@ async def do_save_redis(self, coll_id: UUID, status: CollIndexStatus):
             if status.state not in ("saving", "saved"):
                 await redis.bgsave(False)
 
-                await self.set_state("saving", status, coll_id)
+                await self.set_state("saving", status, coll_id, oid)
 
             if await self.is_bgsave_done(redis):
                 await redis.shutdown()
 
         # pylint: disable=broad-exception-caught
         except Exception:
-            await self.set_state("ready", status, coll_id)
+            await self.set_state("ready", status, coll_id, oid)
             traceback.print_exc()
 
     async def is_bgsave_done(self, redis: Redis) -> bool:
@@ -563,16 +570,6 @@ async def update_saved_dedupe_index_state_in_db(
             storage=org.storage,
         )
 
-        prev_coll = await self.coll_ops.update_dedupe_index_info(
-            coll_id, "idle", index_file, finished_at
-        )
-
-        # Update org storage totals with size of index file, or difference
-        # from previous saved index file if one existed
-        size_diff = size
-        if prev_coll and prev_coll.indexFile:
-            size_diff = size - prev_coll.indexFile.size
-
-        await self.coll_ops.orgs.inc_org_bytes_stored_field(
-            oid, "bytesStoredDedupeIndexes", size_diff
+        await self.coll_ops.update_dedupe_index_info(
+            coll_id, oid, "idle", index_file, finished_at
         )