From 7de12dfbb7e83bf24972dde20cfa3384fc933c83 Mon Sep 17 00:00:00 2001
From: Antonio Nuno Monteiro <anmonteiro@gmail.com>
Date: Fri, 29 Nov 2024 15:55:48 -0800
Subject: [PATCH] refactor: don't store the upgrade handler in Response_state.t
 (#144)

---
 lib/reqd.ml           | 6 +++---
 lib/response_state.ml | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/lib/reqd.ml b/lib/reqd.ml
index b1929b5..8845bc5 100644
--- a/lib/reqd.ml
+++ b/lib/reqd.ml
@@ -95,14 +95,14 @@ let response { response_state; _ } =
   | Waiting -> None
   | Streaming (response, _)
   | Fixed response
-  | Upgrade (response, _) -> Some response
+  | Upgrade response -> Some response
 
 let response_exn { response_state; _ } =
   match response_state with
   | Waiting -> failwith "httpun.Reqd.response_exn: response has not started"
   | Streaming(response, _)
   | Fixed response
-  | Upgrade (response, _) -> response
+  | Upgrade response -> response
 
 let respond_with_string t response str =
   if t.error_code <> `Ok then
@@ -177,7 +177,7 @@ let unsafe_respond_with_upgrade t headers upgrade_handler =
     Writer.write_response t.writer response;
     if t.persistent then
       t.persistent <- Response.persistent_connection response;
-    t.response_state <- Upgrade (response, upgrade_handler);
+    t.response_state <- Upgrade response;
     Writer.flush t.writer (fun _reason ->
       (* TODO(anmonteiro): probably need to check `Closed here? *)
       upgrade_handler ());
diff --git a/lib/response_state.ml b/lib/response_state.ml
index 80c3ba7..58a1f47 100644
--- a/lib/response_state.ml
+++ b/lib/response_state.ml
@@ -2,7 +2,7 @@ type t =
   | Waiting
   | Fixed  of Response.t
   | Streaming of Response.t * Body.Writer.t
-  | Upgrade of Response.t * (unit -> unit)
+  | Upgrade of Response.t
 
 let output_state =
   let response_sent_state = function