From fbee7c85bd93c03ea4799936d31c585bdbf12711 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Mon, 25 Sep 2023 12:32:57 -0500
Subject: [PATCH] Fix net server stopping when handle is garbage collected

---
 CHANGELOG.md                    |  1 +
 src/lune/builtins/net/server.rs |  8 ++++++--
 src/lune/util/futures.rs        | 18 ++++++++++++++++++
 src/lune/util/mod.rs            |  1 +
 4 files changed, 26 insertions(+), 2 deletions(-)
 create mode 100644 src/lune/util/futures.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 98a6b07..934cc67 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -61,6 +61,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed `net.serve` stopping when the returned `ServeHandle` is garbage collected
 - Fixed missing trailing newline when using the `warn` global
 - Fixed constructor for `CFrame` in the `roblox` built-in library not parsing the 12-arg overload correctly. ([#102])
 - Fixed various functions for `CFrame` in the `roblox` built-in library being incorrect, specifically row-column ordering and some flipped signs. ([#103])
diff --git a/src/lune/builtins/net/server.rs b/src/lune/builtins/net/server.rs
index 167f10b..68617d9 100644
--- a/src/lune/builtins/net/server.rs
+++ b/src/lune/builtins/net/server.rs
@@ -12,7 +12,7 @@ use tokio::sync::{mpsc, oneshot, Mutex};
 
 use crate::lune::{
     scheduler::Scheduler,
-    util::{traits::LuaEmitErrorExt, TableBuilder},
+    util::{futures::yield_forever, traits::LuaEmitErrorExt, TableBuilder},
 };
 
 use super::{
@@ -115,7 +115,11 @@ where
             .http1_keepalive(true) // Web sockets must be kept alive
             .serve(hyper_make_service)
             .with_graceful_shutdown(async move {
-                shutdown_rx.recv().await;
+                if shutdown_rx.recv().await.is_none() {
+                    // The channel was closed, meaning the serve handle
+                    // was garbage collected by lua without being used
+                    yield_forever().await;
+                }
             });
         if let Err(e) = result.await {
             eprintln!("Net serve error: {e}")
diff --git a/src/lune/util/futures.rs b/src/lune/util/futures.rs
new file mode 100644
index 0000000..40163c1
--- /dev/null
+++ b/src/lune/util/futures.rs
@@ -0,0 +1,18 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+#[derive(Debug, Clone, Copy)]
+pub struct YieldForever;
+
+impl Future for YieldForever {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
+        Poll::Pending
+    }
+}
+
+pub fn yield_forever() -> YieldForever {
+    YieldForever
+}
diff --git a/src/lune/util/mod.rs b/src/lune/util/mod.rs
index 38984af..3c9457c 100644
--- a/src/lune/util/mod.rs
+++ b/src/lune/util/mod.rs
@@ -1,6 +1,7 @@
 mod table_builder;
 
 pub mod formatting;
+pub mod futures;
 pub mod traits;
 
 pub use table_builder::TableBuilder;