Net serve now returns a handle to stop serving requests safely

.lune/examples/http_server.luau

@@ -25,19 +25,9 @@ local function notFound(request: NetRequest): NetResponse
 	}
 end
 
--- Exit our example after a small delay, if you copy this
--- example just remove this part to keep the server running
-
-task.delay(2, function()
-	print("Shutting down...")
-	task.wait(1)
-	process.exit(0)
-end)
-
 -- Run the server on port 8080
 
-print(`Listening on port {PORT} 🚀`)
-net.serve(PORT, function(request)
+local handle = net.serve(PORT, function(request)
 	if string.sub(request.path, 1, 5) == "/ping" then
 		return pong(request)
 	elseif string.sub(request.path, 1, 7) == "/teapot" then
@@ -46,3 +36,14 @@ net.serve(PORT, function(request)
 		return notFound(request)
 	end
 end)
+
+print(`Listening on port {PORT} 🚀`)
+
+-- Exit our example after a small delay, if you copy this
+-- example just remove this part to keep the server running
+
+task.delay(2, function()
+	print("Shutting down...")
+	task.wait(1)
+	handle.stop()
+end)

CHANGELOG.md

@@ -9,8 +9,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- `net.serve` now returns a `NetServeHandle` which can be used to stop serving requests safely.
+
+  Example usage:
+
+  ```lua
+  local handle = net.serve(8080, function()
+      return "Hello, world!"
+  end)
+
+  print("Shutting down after 1 second...")
+  handle.stop()
+  print("Shut down succesfully")
+  ```
+
 - Setting `cwd` in the options for `process.spawn` to a path starting with a tilde (`~`) will now use a path relative to the platform-specific home / user directory.
-- Added a global type `ProcessSpawnOptions` for the return type of `process.spawn`
+- Added a global type `ProcessSpawnOptions` for the third and optional argument of `process.spawn`
 
 ### Changed
 

luneDocs.json

@@ -231,7 +231,7 @@
   },
   "@roblox/global/net.serve": {
     "code_sample": "",
-    "documentation": "Creates an HTTP server that listens on the given `port`.\n\nThe call to this function will block indefinitely and if\nput inside `task.spawn` will not be cancellable using `task.cancel`, to\nstop the script from running it must be terminated manually or using `process.exit`.",
+    "documentation": "Creates an HTTP server that listens on the given `port`.\n\nThis will ***not*** block and will keep listening for requests on the given `port`\nuntil the `stop` function on the returned `NetServeHandle` has been called.",
     "learn_more_link": "",
     "params": [
       {

luneTypes.d.luau

@@ -153,6 +153,10 @@ export type NetResponse = {
 	body: string?,
 }
 
+export type NetServeHandle = {
+	stop: () -> (),
+}
+
 --[=[
 	@class net
 
@@ -175,14 +179,13 @@ declare net: {
 
 		Creates an HTTP server that listens on the given `port`.
 
-		The call to this function will block indefinitely and if
-		put inside `task.spawn` will not be cancellable using `task.cancel`, to
-		stop the script from running it must be terminated manually or using `process.exit`.
+		This will ***not*** block and will keep listening for requests on the given `port`
+		until the `stop` function on the returned `NetServeHandle` has been called.
 
 		@param port The port to use for the server
 		@param handler The handler function to use for the server
 	]=]
-	serve: (port: number, handler: (request: NetRequest) -> (string | NetResponse)) -> (),
+	serve: (port: number, handler: (request: NetRequest) -> (string | NetResponse)) -> NetServeHandle,
 	--[=[
 		@within net
 

packages/lib/src/globals/net.rs

@@ -11,7 +11,10 @@ use mlua::prelude::*;
 use hyper::{body::to_bytes, http::HeaderValue, server::conn::AddrStream, service::Service};
 use hyper::{Body, HeaderMap, Request, Response, Server};
 use reqwest::{ClientBuilder, Method};
-use tokio::{sync::mpsc::Sender, task};
+use tokio::{
+    sync::mpsc::{self, Sender},
+    task,
+};
 
 use crate::utils::{
     message::LuneMessage,
@@ -147,29 +150,62 @@ async fn net_request<'a>(lua: &'static Lua, config: LuaValue<'a>) -> LuaResult<L
         .build_readonly()
 }
 
-async fn net_serve(lua: &'static Lua, (port, callback): (u16, LuaFunction<'_>)) -> LuaResult<()> {
+async fn net_serve<'a>(
+    lua: &'static Lua,
+    (port, callback): (u16, LuaFunction<'a>), // TODO: Parse options as either callback or table with request callback + websocket callback
+) -> LuaResult<LuaTable<'a>> {
+    // Note that we need to use a mpsc here and not
+    // a oneshot channel since we move the sender
+    // into our table with the stop function
+    let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
+    let websocket_callback = Arc::new(None); // TODO: Store websocket callback, if given
+    let server_callback = Arc::new(lua.create_registry_value(callback)?);
+    let server = Server::bind(&([127, 0, 0, 1], port).into())
+        .executor(LocalExec)
+        .serve(MakeNetService(lua, server_callback, websocket_callback))
+        .with_graceful_shutdown(async move {
+            shutdown_rx.recv().await.unwrap();
+            shutdown_rx.close();
+        });
+    // Make sure we register the thread properly by sending messages
+    // when the server starts up and when it shuts down or errors
     let server_sender = lua
         .app_data_ref::<Weak<Sender<LuneMessage>>>()
         .unwrap()
         .upgrade()
         .unwrap();
-    let server_callback = lua.create_registry_value(callback)?;
-    let server = Server::bind(&([127, 0, 0, 1], port).into())
-        .executor(LocalExec)
-        .serve(MakeNetService(lua, server_callback.into()));
-    if let Err(err) = server.await.map_err(LuaError::external) {
-        server_sender
-            .send(LuneMessage::LuaError(err))
-            .await
-            .map_err(LuaError::external)?;
-    }
-    Ok(())
+    let _ = server_sender.send(LuneMessage::Spawned).await;
+    task::spawn_local(async move {
+        let res = server.await.map_err(LuaError::external);
+        let _ = match res {
+            Err(e) => server_sender.try_send(LuneMessage::LuaError(e)),
+            Ok(_) => server_sender.try_send(LuneMessage::Finished),
+        };
+    });
+    // Create a new read-only table that contains methods
+    // for manipulating server behavior and shutting it down
+    let handle_stop = move |_, _: ()| {
+        if shutdown_tx.try_send(()).is_err() {
+            Err(LuaError::RuntimeError(
+                "Server has already been stopped".to_string(),
+            ))
+        } else {
+            Ok(())
+        }
+    };
+    TableBuilder::new(lua)?
+        .with_function("stop", handle_stop)?
+        .build_readonly()
 }
 
 // Hyper service implementation for net, lots of boilerplate here
 // but make_svc and make_svc_function do not work for what we need
 
-pub struct NetService(&'static Lua, Arc<LuaRegistryKey>);
+pub struct NetService(
+    &'static Lua,
+    Arc<LuaRegistryKey>,
+    Arc<Option<LuaRegistryKey>>,
+);
 
 impl Service<Request<Body>> for NetService {
     type Response = Response<Body>;
@@ -182,13 +218,14 @@ impl Service<Request<Body>> for NetService {
 
     fn call(&mut self, req: Request<Body>) -> Self::Future {
         let lua = self.0;
-        let key = self.1.clone();
+        let key1 = self.1.clone();
+        let _key2 = self.2.clone(); // TODO: This is the web socket callback
         let (parts, body) = req.into_parts();
         Box::pin(async move {
             // Convert request body into bytes, extract handler
             // function & lune message sender to use later
             let bytes = to_bytes(body).await.map_err(LuaError::external)?;
-            let handler: LuaFunction = lua.registry_value(&key)?;
+            let handler: LuaFunction = lua.registry_value(&key1)?;
             let sender = lua
                 .app_data_ref::<Weak<Sender<LuneMessage>>>()
                 .unwrap()
@@ -283,7 +320,11 @@ impl Service<Request<Body>> for NetService {
     }
 }
 
-struct MakeNetService(&'static Lua, Arc<LuaRegistryKey>);
+struct MakeNetService(
+    &'static Lua,
+    Arc<LuaRegistryKey>,
+    Arc<Option<LuaRegistryKey>>,
+);
 
 impl Service<&AddrStream> for MakeNetService {
     type Response = NetService;
@@ -296,8 +337,9 @@ impl Service<&AddrStream> for MakeNetService {
 
     fn call(&mut self, _: &AddrStream) -> Self::Future {
         let lua = self.0;
-        let key = self.1.clone();
-        Box::pin(async move { Ok(NetService(lua, key)) })
+        let key1 = self.1.clone();
+        let key2 = self.2.clone();
+        Box::pin(async move { Ok(NetService(lua, key1, key2)) })
     }
 }
 

tests/net/serve.luau

@@ -1,22 +1,53 @@
 local RESPONSE = "Hello, lune!"
 
-task.spawn(function()
-	net.serve(8080, function(request)
-		-- info("Request:", request)
-		-- info("Responding with", RESPONSE)
-		assert(request.path == "/some/path")
-		assert(request.query.key == "param2")
-		return RESPONSE
-	end)
+local handle = net.serve(8080, function(request)
+	-- info("Request:", request)
+	-- info("Responding with", RESPONSE)
+	assert(request.path == "/some/path")
+	assert(request.query.key == "param2")
+	assert(request.query.key2 == "param3")
+	return RESPONSE
 end)
 
-local response = net.request("http://127.0.0.1:8080/some/path?key=param1&key=param2").body
+local response =
+	net.request("http://127.0.0.1:8080/some/path?key=param1&key=param2&key2=param3").body
 assert(response == RESPONSE, "Invalid response from server")
 
-task.delay(1, function()
-	process.exit(0)
-end)
+handle.stop()
 
-task.delay(2, function()
-	error("Process did not exit")
-end)
+-- Stopping is not guaranteed to happen instantly since it is async, but
+-- it should happen on the next yield, so we wait the minimum amount here
+task.wait()
+
+-- Sending a net request may error if there was
+-- a connection issue, we should handle that here
+local success, response2 = pcall(net.request, "http://127.0.0.1:8080/")
+if not success then
+	local message = tostring(response2)
+	assert(
+		string.find(message, "Connection reset")
+			or string.find(message, "Connection closed")
+			or string.find(message, "Connection refused"),
+		"Server did not stop responding to requests"
+	)
+else
+	assert(not response2.ok, "Server did not stop responding to requests")
+end
+
+--[[
+	Trying to stop the server again should error and
+	mention that the server has already been stopped
+
+	Note that we cast pcall to any because of a
+	Luau limitation where it throws a type error for
+	`err` because handle.stop doesn't return any value
+]]
+local success2, err = (pcall :: any)(handle.stop)
+assert(not success2, "Calling stop twice on the net serve handle should error")
+local message = tostring(err)
+assert(
+	string.find(message, "stop")
+		or string.find(message, "shutdown")
+		or string.find(message, "shut down"),
+	"The error message for calling stop twice on the net serve handle should be descriptive"
+)