1 files changed, 28 insertions, 4 deletions

diff --git a/ext/http/slab.rs b/ext/http/slab.rs
index 8dd562cc2..4718aded1 100644
--- a/ext/http/slab.rs
+++ b/ext/http/slab.rs
@@ -20,6 +20,10 @@ pub type Request = hyper1::Request<Incoming>;
 pub type Response = hyper1::Response<ResponseBytes>;
 pub type SlabId = u32;
 
+#[repr(transparent)]
+#[derive(Clone, Default)]
+pub struct RefCount(pub Rc<()>);
+
 enum RequestBodyState {
   Incoming(Incoming),
   Resource(HttpRequestBodyAutocloser),
@@ -50,24 +54,27 @@ pub struct HttpSlabRecord {
   request_info: HttpConnectionProperties,
   request_parts: Parts,
   request_body: Option<RequestBodyState>,
-  // The response may get taken before we tear this down
+  /// The response may get taken before we tear this down
   response: Option<Response>,
   promise: CompletionHandle,
   trailers: Rc<RefCell<Option<HeaderMap>>>,
   been_dropped: bool,
+  /// Use a `Rc` to keep track of outstanding requests. We don't use this, but
+  /// when it drops, it decrements the refcount of the server itself.
+  refcount: Option<RefCount>,
   #[cfg(feature = "__zombie_http_tracking")]
   alive: bool,
 }
 
 thread_local! {
-  static SLAB: RefCell<Slab<HttpSlabRecord>> = const { RefCell::new(Slab::new()) };
+  pub(crate) static SLAB: RefCell<Slab<HttpSlabRecord>> = const { RefCell::new(Slab::new()) };
 }
 
 macro_rules! http_trace {
   ($index:expr, $args:tt) => {
     #[cfg(feature = "__http_tracing")]
     {
-      let total = SLAB.with(|x| x.try_borrow().map(|x| x.len()));
+      let total = $crate::slab::SLAB.with(|x| x.try_borrow().map(|x| x.len()));
       if let Ok(total) = total {
         println!("HTTP id={} total={}: {}", $index, total, format!($args));
       } else {
@@ -77,6 +84,8 @@ macro_rules! http_trace {
   };
 }
 
+pub(crate) use http_trace;
+
 /// Hold a lock on the slab table and a reference to one entry in the table.
 pub struct SlabEntry(
   NonNull<HttpSlabRecord>,
@@ -121,6 +130,7 @@ fn slab_insert_raw(
   request_parts: Parts,
   request_body: Option<Incoming>,
   request_info: HttpConnectionProperties,
+  refcount: RefCount,
 ) -> SlabId {
   let index = SLAB.with(|slab| {
     let mut slab = slab.borrow_mut();
@@ -135,6 +145,7 @@ fn slab_insert_raw(
       trailers,
       been_dropped: false,
       promise: CompletionHandle::default(),
+      refcount: Some(refcount),
       #[cfg(feature = "__zombie_http_tracking")]
       alive: true,
     })
@@ -146,9 +157,10 @@ fn slab_insert_raw(
 pub fn slab_insert(
   request: Request,
   request_info: HttpConnectionProperties,
+  refcount: RefCount,
 ) -> SlabId {
   let (request_parts, request_body) = request.into_parts();
-  slab_insert_raw(request_parts, Some(request_body), request_info)
+  slab_insert_raw(request_parts, Some(request_body), request_info, refcount)
 }
 
 pub fn slab_drop(index: SlabId) {
@@ -159,10 +171,21 @@ pub fn slab_drop(index: SlabId) {
     !record.been_dropped,
     "HTTP state error: Entry has already been dropped"
   );
+
+  // The logic here is somewhat complicated. A slab record cannot be expunged until it has been dropped by Rust AND
+  // the promise has been completed (indicating that JavaScript is done processing). However, if Rust has finished
+  // dealing with this entry, we DO want to clean up some of the associated items -- namely the request body, which
+  // might include actual resources, and the refcount, which is keeping the server alive.
   record.been_dropped = true;
   if record.promise.is_completed() {
     drop(entry);
     slab_expunge(index);
+  } else {
+    // Take the request body, as the future has been dropped and this will allow some resources to close
+    record.request_body.take();
+    // Take the refcount keeping the server alive. The future is no longer alive, which means this request
+    // is toast.
+    record.refcount.take();
   }
 }
 
@@ -318,6 +341,7 @@ mod tests {
         local_port: None,
         stream_type: NetworkStreamType::Tcp,
       },
+      RefCount::default(),
     );
     let entry = slab_get(id);
     entry.complete();