From 576ba9c522ace725406f66f6b755cc76e86fb2c0 Mon Sep 17 00:00:00 2001
From: Ibraheem Ahmed <ibraheem@ibraheem.ca>
Date: Wed, 3 Jul 2024 20:05:44 -0400
Subject: [PATCH] Avoid hangs before exiting CLI (#4793)

## Summary

The resolver sometimes starts HTTP requests that end up not being
necessary. When dropping the Tokio runtime before exiting we currently
wait for those to complete. This can cause noticeable hangs in the CLI,
particularly when the runtime is blocked on slow DNS resolution.

Resolves https://github.com/astral-sh/uv/issues/4599.

## Test Plan

This change resolves any reproducible hangs for me locally.
---
 crates/uv/src/main.rs | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/crates/uv/src/main.rs b/crates/uv/src/main.rs
index bea85ce68..0edf14338 100644
--- a/crates/uv/src/main.rs
+++ b/crates/uv/src/main.rs
@@ -951,12 +951,19 @@ fn main() -> ExitCode {
         // https://learn.microsoft.com/en-us/cpp/build/reference/stack-stack-allocations?view=msvc-170
         let stack_size = stack_size.parse().expect("Invalid stack size");
         let tokio_main = move || {
-            tokio::runtime::Builder::new_multi_thread()
+            let runtime = tokio::runtime::Builder::new_multi_thread()
                 .enable_all()
                 .thread_stack_size(stack_size)
                 .build()
-                .expect("Failed building the Runtime")
-                .block_on(run())
+                .expect("Failed building the Runtime");
+            let result = runtime.block_on(run());
+            // Avoid waiting for pending tasks to complete.
+            //
+            // The resolver may have kicked off HTTP requests during resolution that
+            // turned out to be unnecessary. Waiting for those to complete can cause
+            // the CLI to hang before exiting.
+            runtime.shutdown_background();
+            result
         };
         std::thread::Builder::new()
             .stack_size(stack_size)
@@ -965,11 +972,13 @@ fn main() -> ExitCode {
             .join()
             .expect("Tokio executor failed, was there a panic?")
     } else {
-        tokio::runtime::Builder::new_multi_thread()
+        let runtime = tokio::runtime::Builder::new_multi_thread()
             .enable_all()
             .build()
-            .expect("Failed building the Runtime")
-            .block_on(run())
+            .expect("Failed building the Runtime");
+        let result = runtime.block_on(run());
+        runtime.shutdown_background();
+        result
     };
 
     match result {