HITL - Change async network process loop.

habitat-hitl/habitat_hitl/_internal/networking/networking_process.py

@@ -365,7 +365,7 @@ async def networking_main_async(
 
     network_mgr = NetworkManager(interprocess_record)
 
-    # Start servers
+    # Start servers.
     websocket_server = await start_websocket_server(
         network_mgr, networking_config
     )
@@ -375,36 +375,47 @@ async def networking_main_async(
         else None
     )
 
-    check_keyframe_queue_task = asyncio.ensure_future(
-        network_mgr.check_keyframe_queue()
+    # Define tasks (concurrent looping coroutines).
+    tasks = []
+    tasks.append(asyncio.create_task(network_mgr.check_keyframe_queue()))
+    tasks.append(
+        asyncio.create_task(network_mgr.check_close_broken_connection())
     )
 
-    check_close_broken_connection_task = asyncio.ensure_future(
-        network_mgr.check_close_broken_connection()
-    )
-
-    # Handle SIGTERM. We should get this signal when we do networking_process.terminate(). See terminate_networking_process.
+    # Handle termination signals.
+    # We should get SIGTERM when we do networking_process.terminate(). See terminate_networking_process.
     stop: asyncio.Future = asyncio.Future()
     loop = asyncio.get_event_loop()
-    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
-
-    # This await essentially means "wait forever" (or until we get SIGTERM). Meanwhile, the other tasks we've started above (websocket server, http server, check_keyframe_queue_task) will also run forever in the asyncio event loop.
+    stop_signals = [
+        signal.SIGTERM,
+        signal.SIGQUIT,
+        signal.SIGINT,
+        signal.SIGHUP,
+    ]
+    for stop_signal in stop_signals:
+        loop.add_signal_handler(stop_signal, stop.set_result, None)
+
+    # Run tasks.
+    tasks = asyncio.gather(*tasks)
+
+    # Wait for cancellation (from termination signals).
     await stop
+    print("Networking process terminating.")
 
-    # Do cleanup code after we've received SIGTERM: close both servers and cancel check_keyframe_queue_task.
+    # Close servers.
     websocket_server.close()
     await websocket_server.wait_closed()
 
     if http_runner:
         await http_runner.cleanup()
 
-    check_keyframe_queue_task.cancel()
-    check_close_broken_connection_task.cancel()
+    # Cancel running tasks.
+    tasks.cancel()
 
 
 def networking_main(interprocess_record: InterprocessRecord) -> None:
     # Set up the event loop and run the main coroutine
     loop = asyncio.get_event_loop()
     loop.run_until_complete(networking_main_async(interprocess_record))
     loop.close()
-    print("networking_main finished")
+    print("Networking process terminated.")