Fix unix socket cleanup.

Author: alistairjevans

ext/hyper_ruby/src/lib.rs

@@ -51,6 +51,7 @@ struct Server {
     config: RefCell<ServerConfig>,
     work_rx: RefCell<Option<crossbeam_channel::Receiver<RequestWithCompletion>>>,
     work_tx: RefCell<Option<Arc<crossbeam_channel::Sender<RequestWithCompletion>>>>,
+    runtime: RefCell<Option<Arc<tokio::runtime::Runtime>>>,
 }
 
 impl Server {
@@ -62,6 +63,7 @@ impl Server {
             config: RefCell::new(ServerConfig::new()),
             work_rx: RefCell::new(Some(work_rx)),
             work_tx: RefCell::new(Some(Arc::new(work_tx))),
+            runtime: RefCell::new(None),
         }
     }
 
@@ -143,10 +145,13 @@ impl Server {
             .ok_or_else(|| MagnusError::new(magnus::exception::runtime_error(), "Work channel not initialized"))?
             .clone();
 
-        let rt = tokio::runtime::Builder::new_multi_thread()
+        let rt = Arc::new(tokio::runtime::Builder::new_multi_thread()
             .enable_all()
             .build()
-            .map_err(|e| MagnusError::new(magnus::exception::runtime_error(), e.to_string()))?;
+            .map_err(|e| MagnusError::new(magnus::exception::runtime_error(), e.to_string()))?);
+
+        // Store the runtime
+        *self.runtime.borrow_mut() = Some(rt.clone());
 
         rt.block_on(async {
             let work_tx = work_tx.clone();
@@ -167,7 +172,7 @@ impl Server {
                     let incoming = UnixListenerStream::new(listener);
                     warp::serve(any_route)
                         .run_incoming(incoming)
-                        .await
+                        .await;
                 } else {
                     let addr: SocketAddr = config.bind_address.parse()
                         .expect("invalid address format");
@@ -183,31 +188,31 @@ impl Server {
             Ok::<(), MagnusError>(())
         })?;
 
-        // Keep the runtime alive
-        std::thread::spawn(move || {
-            rt.block_on(async {
-                loop {
-                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
-                }
-            });
-        });
-
         Ok(())
     }
 
     pub fn stop(&self) -> Result<(), MagnusError> {
-        let rt = tokio::runtime::Runtime::new()
-            .map_err(|e| MagnusError::new(magnus::exception::runtime_error(), e.to_string()))?;
-
-        rt.block_on(async {
-            let mut handle = self.server_handle.lock().await;
-            if let Some(task) = handle.take() {
-                task.abort();
-            }
-        });
+        // Use the stored runtime instead of creating a new one
+        if let Some(rt) = self.runtime.borrow().as_ref() {
+            rt.block_on(async {
+                let mut handle = self.server_handle.lock().await;
+                if let Some(task) = handle.take() {
+                    task.abort();
+                }
+            });
+        }
 
-        // Drop the channel to signal workers to shut down
+        // Drop the channel and runtime
         self.work_tx.borrow_mut().take();
+        self.runtime.borrow_mut().take();
+
+        let bind_address = self.config.borrow().bind_address.clone();
+        if bind_address.starts_with("unix:") {
+            let path = bind_address.trim_start_matches("unix:");
+            std::fs::remove_file(path).unwrap_or_else(|e| {
+                println!("Failed to remove socket file: {:?}", e);
+            });
+        }
 
         Ok(())
     }

test/test_hyper_ruby.rb

@@ -35,6 +35,22 @@ def test_large_post
     end
   end
 
+  def test_unix_socket_cleans_up_socket
+    with_unix_socket_server(-> (request) { handler_simple(request) }) do |client|
+      response = client.get("/")
+      assert_equal 200, response.status
+      assert_equal "text/plain", response.headers["content-type"]
+      assert_equal 'GET', response.body
+    end
+
+    with_unix_socket_server(-> (request) { handler_simple(request) }) do |client|
+      response = client.get("/")
+      assert_equal 200, response.status
+      assert_equal "text/plain", response.headers["content-type"]
+      assert_equal 'GET', response.body
+    end
+  end
+
   # def test_blocking
   #   with_server(-> (request) { handler_simple(request) }) do |client|
   #     gets
@@ -65,6 +81,31 @@ def with_server(request_handler, &block)
     worker.join if worker
   end
 
+  def with_unix_socket_server(request_handler, &block)
+    server = HyperRuby::Server.new
+    server.configure({ bind_address: "unix:/tmp/hyper_ruby_test.sock" })
+    server.start
+    
+    # Create ruby worker threads that process requests;
+    # 1 is usually enough, and generally handles better than multiple threads 
+    # if there's no IO (because of the GIL)
+    worker = Thread.new do
+      server.run_worker do |request|
+        # Process the request in Ruby
+        # request is a hash with :method, :path, :headers, and :body keys
+        request_handler.call(request)
+      end
+    end
+
+    client = HTTPX.with(transport: "unix", addresses: ["/tmp/hyper_ruby_test.sock"], origin: "http://host")
+	
+    block.call(client)
+
+  ensure
+    server.stop if server
+    worker.join if worker
+  end
+
   def handler_simple(request)
     HyperRuby::Response.new(200, { 'Content-Type' => 'text/plain' }, request.http_method)
   end