From 6e90a7f007e579b800ca717cee60b214068a3754 Mon Sep 17 00:00:00 2001
From: Reactor Scram <ReactorScram@users.noreply.github.com>
Date: Fri, 27 Sep 2024 09:10:48 -0500
Subject: [PATCH] refactor(headless-client): re-arrange main (#6835)

Extracted from #6782

This moves more of `main` inside the async block, which makes it easier
to set up telemetry in the future.

We also log errors for the DNS notifier, which was overlooked before.
---
 rust/headless-client/src/main.rs | 48 ++++++++++++++++----------------
 1 file changed, 24 insertions(+), 24 deletions(-)

diff --git a/rust/headless-client/src/main.rs b/rust/headless-client/src/main.rs
index 2d892a9cd..c87987492 100644
--- a/rust/headless-client/src/main.rs
+++ b/rust/headless-client/src/main.rs
@@ -184,26 +184,25 @@ fn main() -> Result<()> {
         private_key,
         callbacks,
     };
-    let _guard = rt.enter(); // Constructing `PhoenixChannel` requires a runtime context.
 
-    // The Headless Client will bail out here if there's no Internet, because `PhoenixChannel` will try to
-    // resolve the portal host and fail. This is intentional behavior. The Headless Client should always be running under a manager like `systemd` or Windows' Service Controller,
-    // so when it fails it will be restarted with backoff. `systemd` can additionally make us wait
-    // for an Internet connection if it launches us at startup.
-    // When running interactively, it is useful for the user to see that we can't reach the portal.
-    let portal = PhoenixChannel::connect(
-        Secret::new(url),
-        get_user_agent(None, env!("CARGO_PKG_VERSION")),
-        "client",
-        (),
-        ExponentialBackoffBuilder::default()
-            .with_max_elapsed_time(max_partition_time)
-            .build(),
-        Arc::new(tcp_socket_factory),
-    )?;
-    let session = Session::connect(args, portal, rt.handle().clone());
+    rt.block_on(async {
+        // The Headless Client will bail out here if there's no Internet, because `PhoenixChannel` will try to
+        // resolve the portal host and fail. This is intentional behavior. The Headless Client should always be running under a manager like `systemd` or Windows' Service Controller,
+        // so when it fails it will be restarted with backoff. `systemd` can additionally make us wait
+        // for an Internet connection if it launches us at startup.
+        // When running interactively, it is useful for the user to see that we can't reach the portal.
+        let portal = PhoenixChannel::connect(
+            Secret::new(url),
+            get_user_agent(None, env!("CARGO_PKG_VERSION")),
+            "client",
+            (),
+            ExponentialBackoffBuilder::default()
+                .with_max_elapsed_time(max_partition_time)
+                .build(),
+            Arc::new(tcp_socket_factory),
+        )?;
+        let session = Session::connect(args, portal, rt.handle().clone());
 
-    let result = rt.block_on(async {
         let mut terminate = signals::Terminate::new()?;
         let mut hangup = signals::Hangup::new()?;
         let mut terminate = pin!(terminate.recv().fuse());
@@ -290,16 +289,17 @@ fn main() -> Result<()> {
             }
         };
 
+        if let Err(error) = dns_notifier.close() {
+            tracing::error!(?error, "DNS notifier")
+        }
         if let Err(error) = network_notifier.close() {
-            tracing::error!(?error, "network listener");
+            tracing::error!(?error, "network notifier");
         }
 
+        session.disconnect();
+
         result
-    });
-
-    session.disconnect();
-
-    result
+    })
 }
 
 /// Read the token from disk if it was not in the environment