Delay starting sources until the consumer get assigned partitions (#495)

Author: ovv

quixstreams/app.py

@@ -862,6 +862,7 @@ def _run_dataframe(self, dataframe):
 
     def _run_sources(self):
         self._running = True
+        self._source_manager.start_sources()
         while self._running:
             self._source_manager.raise_for_error()
 
@@ -947,6 +948,12 @@ def _on_assign(self, _, topic_partitions: List[TopicPartition]):
             return
         logger.debug(f"Rebalancing: assigning partitions")
 
+        # Only start the sources once the consumer is assigned. Otherwise a source
+        # can produce data before the consumer starts. If that happens on a new
+        # consumer with `auto_offset_reset` set to `latest` the consumer will not
+        # get the source data.
+        self._source_manager.start_sources()
+
         # First commit everything processed so far because assignment can take a while
         # and fail
         self._processing_context.commit_checkpoint(force=True)

quixstreams/sources/manager.py

@@ -28,6 +28,7 @@ def __init__(self, source):
         self.source: BaseSource = source
 
         self._exceptions: List[Exception] = []
+        self._started = False
         self._stopping = False
 
         # copy parent process log level to the child process
@@ -36,6 +37,10 @@ def __init__(self, source):
         # reader and writer pipe used to communicate from the child to the parent process
         self._rpipe, self._wpipe = multiprocessing.Pipe(duplex=False)
 
+    @property
+    def started(self):
+        return self._started
+
     # --- CHILD PROCESS METHODS --- #
 
     def _setup_signal_handlers(self):
@@ -54,6 +59,7 @@ def run(self) -> None:
             * Execution of the source `run` method
             * Reporting the source exceptions to the parent process
         """
+        self._started = True
         self._setup_signal_handlers()
         configure_logging(self._loglevel, str(self.source), pid=True)
         logger.info("Source started")
@@ -103,6 +109,7 @@ def _report_exception(self, err: BaseException) -> None:
 
     def start(self) -> "SourceProcess":
         logger.info("Starting source %s", self.source)
+        self._started = True
         return super().start()
 
     def raise_for_error(self) -> None:
@@ -186,7 +193,8 @@ def topics(self) -> List[Topic]:
 
     def start_sources(self) -> None:
         for process in self.processes:
-            process.start()
+            if not process.started:
+                process.start()
 
     def stop_sources(self) -> None:
         for process in self.processes:
@@ -221,7 +229,6 @@ def is_alive(self) -> bool:
         return False
 
     def __enter__(self):
-        self.start_sources()
         return self
 
     def __exit__(self, *args, **kwargs):