fix: Clean shutdown for Sink threaded server using threading.Event

Signed-off-by: Sreekanth <prsreekanth920@gmail.com>

Author: BulkBeing

packages/pynumaflow/pynumaflow/shared/server.py

@@ -4,6 +4,7 @@
 import multiprocessing
 import os
 import socket
+import threading
 import traceback
 
 from google.protobuf import any_pb2
@@ -18,6 +19,7 @@
 from pynumaflow._constants import (
     _LOGGER,
     MULTIPROC_MAP_SOCK_ADDR,
+    NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS,
     UDFType,
 )
 from pynumaflow.exceptions import SocketError
@@ -57,6 +59,7 @@ def sync_server_start(
     server_options=None,
     server_info: ServerInfo | None = None,
     udf_type: str = UDFType.Map,
+    shutdown_event: threading.Event | None = None,
 ):
     """
     Utility function to start a sync grpc server instance.
@@ -75,6 +78,7 @@ def sync_server_start(
         udf_type=udf_type,
         server_info_file=server_info_file,
         server_info=server_info,
+        shutdown_event=shutdown_event,
     )
 
 
@@ -86,10 +90,15 @@ def _run_server(
     udf_type: str,
     server_info_file: str | None = None,
     server_info: ServerInfo | None = None,
+    shutdown_event: threading.Event | None = None,
 ) -> None:
     """
     Starts the Synchronous server instance on the given UNIX socket
     with given max threads. Wait for the server to terminate.
+
+    If *shutdown_event* is provided, a background daemon thread will wait
+    on it and then call ``server.stop(NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS)``
+    for a cooperative graceful shutdown (no process kill).
     """
     server = grpc.server(
         ThreadPoolExecutor(
@@ -115,10 +124,21 @@ def _run_server(
     server.add_insecure_port(bind_address)
     # start the gRPC server
     server.start()
+
     # Add the server information to the server info file if provided
     if server_info and server_info_file:
         info_server_write(server_info=server_info, info_file=server_info_file)
 
+    if shutdown_event is not None:
+
+        def _watch_for_shutdown():
+            shutdown_event.wait()
+            _LOGGER.info("Shutdown signal received, stopping server gracefully...")
+            server.stop(NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS)
+
+        watcher = threading.Thread(target=_watch_for_shutdown, daemon=True)
+        watcher.start()
+
     _LOGGER.info("GRPC Server listening on: %s %d", bind_address, os.getpid())
     server.wait_for_termination()
 

packages/pynumaflow/pynumaflow/shared/thread_with_return.py

@@ -5,6 +5,7 @@ class ThreadWithReturnValue(Thread):
     """
     A custom Thread class that allows the target function to return a value.
     This class extends the built-in threading.Thread class.
+    Exceptions raised by the target are captured and re-raised on join().
     """
 
     def __init__(self, group=None, target=None, name=None, args=(), kwargs={}, verbose=None):
@@ -23,32 +24,43 @@ def __init__(self, group=None, target=None, name=None, args=(), kwargs={}, verbo
         Thread.__init__(self, group, target, name, args, kwargs)
         # Variable to store the return value of the target function
         self._return = None
+        self._exception: BaseException | None = None
 
     def run(self):
         """
         Run the thread.
 
         This method is overridden from the Thread class.
         It calls the target function and saves the return value.
+        If the target raises, the exception is captured for re-raising on join().
         """
         if self._target is not None:
-            # Execute target and store the result
-            self._return = self._target(*self._args, **self._kwargs)
+            try:
+                # Execute target and store the result
+                self._return = self._target(*self._args, **self._kwargs)
+            except BaseException as exc:
+                self._exception = exc
 
     def join(self, *args):
         """
         Wait for the thread to complete and return the result.
 
         This method is overridden from the Thread class.
-        It calls the parent class's join() method and then returns the stored return value.
+        It calls the parent class's join() method, re-raises any captured
+        exception, and then returns the stored return value.
 
         Parameters:
         *args: Variable length argument list to pass to the join() method.
 
         Returns:
         The return value from the target function.
+
+        Raises:
+        BaseException: If the target function raised during run().
         """
         # Call the parent class's join() method to wait for the thread to finish
         Thread.join(self, *args)
+        if self._exception is not None:
+            raise self._exception
         # Return the result of the target function
         return self._return

packages/pynumaflow/pynumaflow/sinker/server.py

@@ -1,4 +1,5 @@
 import os
+import sys
 
 from pynumaflow.info.types import ServerInfo, ContainerType, MINIMUM_NUMAFLOW_VERSION
 from pynumaflow.sinker.servicer.sync_servicer import SyncSinkServicer
@@ -120,6 +121,7 @@ def start(self):
         )
         serv_info = ServerInfo.get_default_server_info()
         serv_info.minimum_numaflow_version = MINIMUM_NUMAFLOW_VERSION[ContainerType.Sinker]
+
         # Start the server
         sync_server_start(
             servicer=self.servicer,
@@ -129,4 +131,9 @@ def start(self):
             server_options=self._server_options,
             udf_type=UDFType.Sink,
             server_info=serv_info,
+            shutdown_event=self.servicer._shutdown_event,
         )
+
+        if self.servicer._error:
+            _LOGGER.critical("Server exiting due to UDF error: %s", self.servicer._error)
+            sys.exit(1)

packages/pynumaflow/pynumaflow/sinker/servicer/sync_servicer.py

@@ -1,9 +1,10 @@
+import threading
 from collections.abc import Iterator
 
 
-from pynumaflow._constants import _LOGGER, STREAM_EOF
+from pynumaflow._constants import _LOGGER, STREAM_EOF, ERR_UDF_EXCEPTION_STRING
 from pynumaflow.proto.sinker import sink_pb2_grpc, sink_pb2
-from pynumaflow.shared.server import exit_on_error
+from pynumaflow.shared.server import update_context_err
 from pynumaflow.shared.synciter import SyncIterator
 from pynumaflow.shared.thread_with_return import ThreadWithReturnValue
 from pynumaflow.sinker._dtypes import SinkSyncCallable
@@ -24,6 +25,8 @@ class SyncSinkServicer(sink_pb2_grpc.SinkServicer):
 
     def __init__(self, handler: SinkSyncCallable):
         self.handler: SinkSyncCallable = handler
+        self._shutdown_event: threading.Event = threading.Event()
+        self._error: BaseException | None = None
 
     def SinkFn(
         self, request_iterator: Iterator[sink_pb2.SinkRequest], context: NumaflowServicerContext
@@ -79,10 +82,11 @@ def SinkFn(
                 cur_task.join()
 
         except BaseException as err:
-            # Handle exceptions
-            err_msg = f"UDSinkError: {repr(err)}"
+            err_msg = f"UDSinkError, {ERR_UDF_EXCEPTION_STRING}: {repr(err)}"
             _LOGGER.critical(err_msg, exc_info=True)
-            exit_on_error(context, err_msg)
+            update_context_err(context, err, err_msg)
+            self._error = err
+            self._shutdown_event.set()
             return
 
     def _invoke_sink(self, request_queue: SyncIterator, context: NumaflowServicerContext):
@@ -93,7 +97,6 @@ def _invoke_sink(self, request_queue: SyncIterator, context: NumaflowServicerCon
         except BaseException as err:
             err_msg = f"UDSinkError: {repr(err)}"
             _LOGGER.critical(err_msg, exc_info=True)
-            exit_on_error(context, err_msg)
             raise err
 
     def IsReady(self, request, context: NumaflowServicerContext) -> sink_pb2.ReadyResponse: