fix: Clean shutdown for Sink Async servers using asyncio.Event

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

Author: BulkBeing

packages/pynumaflow/pynumaflow/sinker/async_server.py

@@ -1,8 +1,12 @@
+import asyncio
+import contextlib
 import os
+import sys
 
 import aiorun
 import grpc
 
+from pynumaflow.info.server import write as info_server_write
 from pynumaflow.info.types import ContainerType, ServerInfo, MINIMUM_NUMAFLOW_VERSION
 from pynumaflow.sinker.servicer.async_servicer import AsyncSinkServicer
 from pynumaflow.proto.sinker import sink_pb2_grpc
@@ -24,7 +28,7 @@
     MAX_NUM_THREADS,
 )
 
-from pynumaflow.shared.server import NumaflowServer, start_async_server
+from pynumaflow.shared.server import NumaflowServer
 from pynumaflow.sinker._dtypes import SinkAsyncCallable
 
 
@@ -118,13 +122,17 @@ def __init__(
         ]
 
         self.servicer = AsyncSinkServicer(sinker_instance)
+        self._error: BaseException | None = None
 
     def start(self):
         """
         Starter function for the Async server class, need a separate caller
         so that all the async coroutines can be started from a single context
         """
         aiorun.run(self.aexec(), use_uvloop=True, shutdown_callback=self.shutdown_callback)
+        if self._error:
+            _LOGGER.critical("Server exiting due to UDF error: %s", self._error)
+            sys.exit(1)
 
     async def aexec(self):
         """
@@ -133,17 +141,41 @@ async def aexec(self):
         # As the server is async, we need to create a new server instance in the
         # same thread as the event loop so that all the async calls are made in the
         # same context
-        # Create a new server instance, add the servicer to it and start the server
         server = grpc.aio.server(options=self._server_options)
         server.add_insecure_port(self.sock_path)
+
+        # The asyncio.Event must be created here (inside aexec) rather than in __init__,
+        # because it must be bound to the running event loop that aiorun creates.
+        # At __init__ time no event loop exists yet.
+        shutdown_event = asyncio.Event()
+        self.servicer.set_shutdown_event(shutdown_event)
+
         sink_pb2_grpc.add_SinkServicer_to_server(self.servicer, server)
+
         serv_info = ServerInfo.get_default_server_info()
         serv_info.minimum_numaflow_version = MINIMUM_NUMAFLOW_VERSION[ContainerType.Sinker]
-        await start_async_server(
-            server_async=server,
-            sock_path=self.sock_path,
-            max_threads=self.max_threads,
-            cleanup_coroutines=list(),
-            server_info_file=self.server_info_file,
-            server_info=serv_info,
+
+        await server.start()
+        info_server_write(server_info=serv_info, info_file=self.server_info_file)
+
+        _LOGGER.info(
+            "Async GRPC Server listening on: %s with max threads: %s",
+            self.sock_path,
+            self.max_threads,
         )
+
+        async def _watch_for_shutdown():
+            """Wait for the shutdown event and stop the server with a grace period."""
+            await shutdown_event.wait()
+            _LOGGER.info("Shutdown signal received, stopping server gracefully...")
+            await server.stop(5)
+
+        shutdown_task = asyncio.create_task(_watch_for_shutdown())
+        await server.wait_for_termination()
+
+        # Propagate error so start() can exit with a non-zero code
+        self._error = self.servicer._error
+
+        shutdown_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await shutdown_task

packages/pynumaflow/pynumaflow/sinker/servicer/async_servicer.py

@@ -4,7 +4,7 @@
 from google.protobuf import empty_pb2 as _empty_pb2
 from pynumaflow.shared.asynciter import NonBlockingIterator
 
-from pynumaflow.shared.server import handle_async_error
+from pynumaflow.shared.server import update_context_err
 from pynumaflow.sinker._dtypes import Datum, SinkAsyncCallable
 from pynumaflow.proto.sinker import sink_pb2_grpc, sink_pb2
 from pynumaflow.sinker.servicer.utils import (
@@ -30,6 +30,12 @@ def __init__(
         self.background_tasks = set()
         self.__sink_handler: SinkAsyncCallable = handler
         self.cleanup_coroutines = []
+        self._shutdown_event: asyncio.Event | None = None
+        self._error: BaseException | None = None
+
+    def set_shutdown_event(self, event: asyncio.Event):
+        """Wire up the shutdown event created by the server's aexec() coroutine."""
+        self._shutdown_event = event
 
     async def SinkFn(
         self,
@@ -82,10 +88,14 @@ async def SinkFn(
                 datum = datum_from_sink_req(d)
                 await req_queue.put(datum)
         except BaseException as err:
-            # if there is an exception, we will mark all the responses as a failure
-            err_msg = f"UDSinkError: {repr(err)}"
+            err_msg = f"UDSinkError, {ERR_UDF_EXCEPTION_STRING}: {repr(err)}"
             _LOGGER.critical(err_msg, exc_info=True)
-            await handle_async_error(context, err, ERR_UDF_EXCEPTION_STRING)
+            update_context_err(context, err, err_msg)
+            # Store the error and signal the server to shut down gracefully
+            # instead of killing the process via SIGKILL.
+            self._error = err
+            if self._shutdown_event:
+                self._shutdown_event.set()
             return
 
     async def __invoke_sink(

packages/pynumaflow/tests/sink/test_async_sink.py

@@ -3,7 +3,6 @@
 import threading
 import unittest
 from collections.abc import AsyncIterable
-from unittest.mock import patch
 
 import grpc
 from google.protobuf import empty_pb2 as _empty_pb2
@@ -32,7 +31,7 @@
     mock_fallback_message,
     mockenv,
 )
-from tests.testing_utils import get_time_args, mock_terminate_on_stop
+from tests.testing_utils import get_time_args
 
 LOGGER = setup_logging(__name__)
 
@@ -128,7 +127,6 @@ async def start_server():
 
 
 # We are mocking the terminate function from the psutil to not exit the program during testing
-@patch("psutil.Process.kill", mock_terminate_on_stop)
 class TestAsyncSink(unittest.TestCase):
     @classmethod
     def setUpClass(cls) -> None: