allow rocksdb db to be shared amongst workers (#96)

Author: patkivikram

faust/app/base.py

@@ -1561,6 +1561,7 @@ async def _stop_consumer(self) -> None:
                     self.flow_control.clear()
                     await self._stop_fetcher()
                     await self._consumer_wait_empty(consumer, self.log)
+                    await self.tables.stop()
 
     async def _consumer_wait_empty(self, consumer: ConsumerT, logger: Any) -> None:
         if self.conf.stream_wait_empty:
@@ -1604,7 +1605,7 @@ def on_rebalance_end(self) -> None:
         try:
             if not sensor_state:
                 self.log.warning(
-                    "Missing sensor state for rebalance #%s", self.rebalancing_count
+                    "Missing sensor state for rebalance end #%s", self.rebalancing_count
                 )
             else:
                 self.sensors.on_rebalance_end(self, sensor_state)
@@ -1671,7 +1672,9 @@ async def _stop_fetcher(self) -> None:
     def _on_rebalance_when_stopped(self) -> None:
         self.consumer.close()
 
-    async def _on_partitions_assigned(self, assigned: Set[TP]) -> None:
+    async def _on_partitions_assigned(
+        self, assigned: Set[TP], generation_id: int = 0
+    ) -> None:
         """Handle new topic partition assignment.
 
         This is called during a rebalance after :meth:`on_partitions_revoked`.
@@ -1687,10 +1690,9 @@ async def _on_partitions_assigned(self, assigned: Set[TP]) -> None:
         # (Kafka does not send error, it just logs)
         session_timeout = self.conf.broker_session_timeout * 0.95
         self.unassigned = not assigned
-
+        logger.info("Executing _on_partitions_assigned")
         revoked, newly_assigned = self._update_assignment(assigned)
         await asyncio.sleep(0)
-
         with flight_recorder(self.log, timeout=session_timeout) as on_timeout:
             consumer = self.consumer
             try:
@@ -1715,7 +1717,9 @@ async def _on_partitions_assigned(self, assigned: Set[TP]) -> None:
                     )
                 on_timeout.info("tables.on_rebalance()")
                 await asyncio.sleep(0)
-                await T(self.tables.on_rebalance)(assigned, revoked, newly_assigned)
+                await T(self.tables.on_rebalance)(
+                    assigned, revoked, newly_assigned, generation_id
+                )
                 on_timeout.info("+send signal: on_partitions_assigned")
                 await T(self.on_partitions_assigned.send)(assigned)
                 on_timeout.info("-send signal: on_partitions_assigned")

faust/stores/rocksdb.py

@@ -1,5 +1,6 @@
 """RocksDB storage."""
 import asyncio
+import gc
 import math
 import shutil
 import typing
@@ -332,6 +333,7 @@ async def on_rebalance(
         assigned: Set[TP],
         revoked: Set[TP],
         newly_assigned: Set[TP],
+        generation_id: int = 0,
     ) -> None:
         """Rebalance occurred.
 
@@ -345,11 +347,14 @@ async def on_rebalance(
         self.rebalance_ack = False
         async with self.db_lock:
             self.revoke_partitions(table, revoked)
-            await self.assign_partitions(table, newly_assigned)
+            await self.assign_partitions(table, newly_assigned, generation_id)
 
     async def stop(self) -> None:
-        for db in self._dbs.values():
-            db.close()
+        self.logger.info("Closing rocksdb on stop")
+        # for db in self._dbs.values():
+        #     db.close()
+        self._dbs.clear()
+        gc.collect()
 
     def revoke_partitions(self, table: CollectionT, tps: Set[TP]) -> None:
         """De-assign partitions used on this worker instance.
@@ -363,9 +368,13 @@ def revoke_partitions(self, table: CollectionT, tps: Set[TP]) -> None:
             if tp.topic in table.changelog_topic.topics:
                 db = self._dbs.pop(tp.partition, None)
                 if db is not None:
-                    db.close()
+                    self.logger.info(f"closing db {tp.topic} partition {tp.partition}")
+                    # db.close()
+        gc.collect()
 
-    async def assign_partitions(self, table: CollectionT, tps: Set[TP]) -> None:
+    async def assign_partitions(
+        self, table: CollectionT, tps: Set[TP], generation_id: int = 0
+    ) -> None:
         """Assign partitions to this worker instance.
 
         Arguments:
@@ -377,26 +386,42 @@ async def assign_partitions(self, table: CollectionT, tps: Set[TP]) -> None:
         my_topics = table.changelog_topic.topics
         for tp in tps:
             if tp.topic in my_topics and tp not in standby_tps and self.rebalance_ack:
-                await self._try_open_db_for_partition(tp.partition)
+                await self._try_open_db_for_partition(
+                    tp.partition, generation_id=generation_id
+                )
                 await asyncio.sleep(0)
 
     async def _try_open_db_for_partition(
-        self, partition: int, max_retries: int = 30, retry_delay: float = 1.0
+        self,
+        partition: int,
+        max_retries: int = 30,
+        retry_delay: float = 1.0,
+        generation_id: int = 0,
     ) -> DB:
         for i in range(max_retries):
             try:
                 # side effect: opens db and adds to self._dbs.
+                self.logger.info(
+                    f"opening partition {partition} for gen id "
+                    f"{generation_id} app id {self.app.consumer_generation_id}"
+                )
                 return self._db_for_partition(partition)
             except rocksdb.errors.RocksIOError as exc:
                 if i == max_retries - 1 or "lock" not in repr(exc):
                     # release all the locks and crash
+                    self.log.warning(
+                        "DB for partition %r retries timed out ", partition
+                    )
                     await self.stop()
                     raise
                 self.log.info(
                     "DB for partition %r is locked! Retry in 1s...", partition
                 )
-                if not self.rebalance_ack:
-                    self.log.info("Rebalanced again giving up partition", partition)
+                if generation_id != self.app.consumer_generation_id:
+                    self.log.info(
+                        f"Rebalanced again giving up partition {partition} gen id"
+                        f" {generation_id} app {self.app.consumer_generation_id}"
+                    )
                     return
                 await self.sleep(retry_delay)
         else:  # pragma: no cover

faust/tables/base.py

@@ -561,10 +561,16 @@ def _windowed_delta(self, key: Any, d: Seconds, event: EventT = None) -> Any:
         )
 
     async def on_rebalance(
-        self, assigned: Set[TP], revoked: Set[TP], newly_assigned: Set[TP]
+        self,
+        assigned: Set[TP],
+        revoked: Set[TP],
+        newly_assigned: Set[TP],
+        generation_id: int = 0,
     ) -> None:
         """Call when cluster is rebalancing."""
-        await self.data.on_rebalance(self, assigned, revoked, newly_assigned)
+        await self.data.on_rebalance(
+            self, assigned, revoked, newly_assigned, generation_id
+        )
 
     async def on_recovery_completed(
         self, active_tps: Set[TP], standby_tps: Set[TP]

faust/tables/manager.py

@@ -181,18 +181,26 @@ def on_partitions_revoked(self, revoked: Set[TP]) -> None:
         T(self.recovery.on_partitions_revoked)(revoked)
 
     async def on_rebalance(
-        self, assigned: Set[TP], revoked: Set[TP], newly_assigned: Set[TP]
+        self,
+        assigned: Set[TP],
+        revoked: Set[TP],
+        newly_assigned: Set[TP],
+        generation_id: int = 0,
     ) -> None:
         """Call when the cluster is rebalancing."""
         self._recovery_started.set()  # cannot add more tables.
         T = traced_from_parent_span()
         for table in self.values():
-            await T(table.on_rebalance)(assigned, revoked, newly_assigned)
+            await T(table.on_rebalance)(
+                assigned, revoked, newly_assigned, generation_id
+            )
 
         await asyncio.sleep(0)
         await T(self._update_channels)()
         await asyncio.sleep(0)
-        await T(self.recovery.on_rebalance)(assigned, revoked, newly_assigned)
+        await T(self.recovery.on_rebalance)(
+            assigned, revoked, newly_assigned, generation_id
+        )
 
     async def wait_until_recovery_completed(self) -> bool:
         if (