Make exporter survive connection interruption

Author: NickCao

packages/jumpstarter-cli-admin/jumpstarter_cli_admin/get_test.py

@@ -764,6 +764,7 @@ async def test_get_exporters(_load_kube_config_mock, list_exporters_mock: AsyncM
 
 EXPORTERS_DEVICES_LIST_NAME = EXPORTERS_LIST_NAME
 
+
 @pytest.mark.anyio
 @patch.object(ExportersV1Alpha1Api, "list_exporters")
 @patch.object(ExportersV1Alpha1Api, "_load_kube_config")
@@ -1141,6 +1142,7 @@ async def test_get_lease(_load_kube_config_mock, get_lease_mock: AsyncMock):
 lease.jumpstarter.dev/82a8ac0d-d7ff-4009-8948-18a3c5c607b2
 """
 
+
 @pytest.mark.anyio
 @patch.object(LeasesV1Alpha1Api, "list_leases")
 @patch.object(LeasesV1Alpha1Api, "_load_kube_config")

packages/jumpstarter/jumpstarter/exporter/exporter.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 
 import grpc
-from anyio import connect_unix, create_task_group
+from anyio import connect_unix, create_memory_object_stream, create_task_group, sleep
 from google.protobuf import empty_pb2
 from jumpstarter_protocol import (
     jumpstarter_pb2,
@@ -37,9 +37,12 @@ async def __aexit__(self, exc_type, exc_value, traceback):
         )
 
     async def __handle(self, path, endpoint, token, tls_config):
-        async with await connect_unix(path) as stream:
-            async with connect_router_stream(endpoint, token, stream, tls_config):
-                pass
+        try:
+            async with await connect_unix(path) as stream:
+                async with connect_router_stream(endpoint, token, stream, tls_config):
+                    pass
+        except Exception as e:
+            logger.info("failed to handle connection: {}".format(e))
 
     @asynccontextmanager
     async def session(self):
@@ -64,21 +67,47 @@ async def session(self):
                 yield path
 
     async def handle(self, lease_name, tg):
-        controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel_factory())
         logger.info("Listening for incoming connection requests on lease %s", lease_name)
+
+        listen_tx, listen_rx = create_memory_object_stream()
+
+        async def listen():
+            while True:
+                try:
+                    controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel_factory())
+                    async for request in controller.Listen(jumpstarter_pb2.ListenRequest(lease_name=lease_name)):
+                        await listen_tx.send(request)
+                except Exception as e:
+                    logger.info("Listen stream interrupted, restarting: {}".format(e))
+                    await sleep(3)
+
+        tg.start_soon(listen)
+
         async with self.session() as path:
-            async for request in controller.Listen(jumpstarter_pb2.ListenRequest(lease_name=lease_name)):
+            async for request in listen_rx:
                 logger.info("Handling new connection request on lease %s", lease_name)
                 tg.start_soon(self.__handle, path, request.router_endpoint, request.router_token, self.tls)
 
     async def serve(self):
-        controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel_factory())
         # initial registration
         async with self.session():
             pass
         started = False
+        status_tx, status_rx = create_memory_object_stream()
+
+        async def status():
+            while True:
+                try:
+                    controller = jumpstarter_pb2_grpc.ControllerServiceStub(self.channel_factory())
+                    async for status in controller.Status(jumpstarter_pb2.StatusRequest()):
+                        await status_tx.send(status)
+                except Exception as e:
+                    logger.info("Status stream interrupted, restarting: {}".format(e))
+                    await sleep(3)
+
         async with create_task_group() as tg:
-            async for status in controller.Status(jumpstarter_pb2.StatusRequest()):
+            tg.start_soon(status)
+            async for status in status_rx:
                 if self.lease_name != "" and self.lease_name != status.lease_name:
                     self.lease_name = status.lease_name
                     logger.info("Lease status changed, killing existing connections")