From a95faa1676b1b3b3466e2e9b55ba906b087c7ee4 Mon Sep 17 00:00:00 2001
From: Eric Rosenberg <eric_rosenberg@apple.com>
Date: Fri, 20 Dec 2024 16:49:28 +0000
Subject: [PATCH] feedback

---
 Package.swift                                 |  2 +-
 .../HTTPDrippingDownloadHandler.swift         | 22 +++++++++++++++----
 .../HTTPReceiveDiscardHandler.swift           |  2 +-
 .../SimpleResponsivenessRequestMux.swift      |  4 ++--
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/Package.swift b/Package.swift
index 1452bc4..ca3e810 100644
--- a/Package.swift
+++ b/Package.swift
@@ -250,7 +250,7 @@ let package = Package(
         .library(name: "NIOHTTPResponsiveness", targets: ["NIOHTTPResponsiveness"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.67.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.77.0"),
         .package(url: "https://github.com/apple/swift-nio-http2.git", from: "1.27.0"),
         .package(url: "https://github.com/apple/swift-http-types.git", from: "1.3.0"),
         .package(url: "https://github.com/apple/swift-http-structured-headers.git", from: "1.1.0"),
diff --git a/Sources/NIOHTTPResponsiveness/HTTPDrippingDownloadHandler.swift b/Sources/NIOHTTPResponsiveness/HTTPDrippingDownloadHandler.swift
index 4ca0dc3..7c45ae0 100644
--- a/Sources/NIOHTTPResponsiveness/HTTPDrippingDownloadHandler.swift
+++ b/Sources/NIOHTTPResponsiveness/HTTPDrippingDownloadHandler.swift
@@ -50,6 +50,7 @@ public final class HTTPDrippingDownloadHandler: ChannelDuplexHandler {
 
     private var phase = Phase.waitingOnHead
     private var scheduled: Scheduled<Void>?
+    private var scheduledCallbackHandler: HTTPDrippingDownloadHandlerScheduledCallbackHandler?
     private var pendingRead = false
     private var pendingWrite = false
     private var activelyWritingChunk = false
@@ -156,6 +157,7 @@ public final class HTTPDrippingDownloadHandler: ChannelDuplexHandler {
     public func channelInactive(context: ChannelHandlerContext) {
         self.phase = .done
         self.scheduled?.cancel()
+        context.fireChannelInactive()
     }
 
     public func channelWritabilityChanged(context: ChannelHandlerContext) {
@@ -219,10 +221,22 @@ public final class HTTPDrippingDownloadHandler: ChannelDuplexHandler {
         // More chunks to write.. Kick off timer
         drippingState.currentChunkBytesLeft = self.size
         self.phase = .dripping(drippingState)
-        let this = NIOLoopBound(self, eventLoop: context.eventLoop)
-        let loopBoundContext = NIOLoopBound(context, eventLoop: context.eventLoop)
-        self.scheduled = context.eventLoop.scheduleTask(in: self.frequency) {
-            this.value.writeChunk(context: loopBoundContext.value)
+        if self.scheduledCallbackHandler == nil {
+            let this = NIOLoopBound(self, eventLoop: context.eventLoop)
+            let loopBoundContext = NIOLoopBound(context, eventLoop: context.eventLoop)
+            self.scheduledCallbackHandler = HTTPDrippingDownloadHandlerScheduledCallbackHandler(handler: this, context: loopBoundContext)
+        }
+        // SAFTEY: scheduling the callback only potentially throws when invoked off eventloop
+        try! context.eventLoop.scheduleCallback(in: self.frequency, handler: self.scheduledCallbackHandler!)
+    }
+
+    private struct HTTPDrippingDownloadHandlerScheduledCallbackHandler: NIOScheduledCallbackHandler & Sendable {
+        var handler: NIOLoopBound<HTTPDrippingDownloadHandler>
+        var context: NIOLoopBound<ChannelHandlerContext>
+
+        func handleScheduledCallback(eventLoop: some EventLoop) {
+            self.handler.value.writeChunk(context: self.context.value)
         }
     }
 }
+
diff --git a/Sources/NIOHTTPResponsiveness/HTTPReceiveDiscardHandler.swift b/Sources/NIOHTTPResponsiveness/HTTPReceiveDiscardHandler.swift
index 0463fe7..8fba40d 100644
--- a/Sources/NIOHTTPResponsiveness/HTTPReceiveDiscardHandler.swift
+++ b/Sources/NIOHTTPResponsiveness/HTTPReceiveDiscardHandler.swift
@@ -70,7 +70,7 @@ public final class HTTPReceiveDiscardHandler: ChannelInboundHandler {
             string:
                 "Received in excess of expectation; expected(\(expectation)) received(\(self.received))"
         )
-        self.writeSimpleResponse(context: context, status: .ok, body: body)
+        self.writeSimpleResponse(context: context, status: .badRequest, body: body)
     }
 
     private func writeSimpleResponse(
diff --git a/Sources/NIOHTTPResponsiveness/SimpleResponsivenessRequestMux.swift b/Sources/NIOHTTPResponsiveness/SimpleResponsivenessRequestMux.swift
index b93b411..b3ecb94 100644
--- a/Sources/NIOHTTPResponsiveness/SimpleResponsivenessRequestMux.swift
+++ b/Sources/NIOHTTPResponsiveness/SimpleResponsivenessRequestMux.swift
@@ -70,8 +70,8 @@ public final class SimpleResponsivenessRequestMux: ChannelInboundHandler {
                     status: .ok,
                     body: self.responsivenessConfigBuffer
                 )
-            case (.get, .some(""), .some("responsiveness"), .some("download"), .some(let count)):
-                self.addHandlerOrInternalError(context: context, handler: HTTPDrippingDownloadHandler(count: 1, size: count))
+            case (.get, .some(""), .some("responsiveness"), .some("download"), .some(let size)):
+                self.addHandlerOrInternalError(context: context, handler: HTTPDrippingDownloadHandler(count: 1, size: size))
             case (.post, .some(""), .some("responsiveness"), .some("upload"), .none):
                 // Check if we should expect a certain count
                 var expectation: Int?