From e9014142d5ce84638ee6926732fb2a399a964be8 Mon Sep 17 00:00:00 2001
From: Kazu Yamamoto <kazu@iij.ad.jp>
Date: Thu, 31 Oct 2024 11:51:23 +0900
Subject: [PATCH] labeling threads

---
 warp-quic/Network/Wai/Handler/WarpQUIC.hs |  5 ++++-
 warp/Network/Wai/Handler/Warp/HTTP2.hs    | 10 +++++++---
 warp/Network/Wai/Handler/Warp/Run.hs      |  8 +++++++-
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/warp-quic/Network/Wai/Handler/WarpQUIC.hs b/warp-quic/Network/Wai/Handler/WarpQUIC.hs
index 2436d7613..7aa4c6f06 100644
--- a/warp-quic/Network/Wai/Handler/WarpQUIC.hs
+++ b/warp-quic/Network/Wai/Handler/WarpQUIC.hs
@@ -59,4 +59,7 @@ quicApp settings app ii conn = do
             let runX
                     | "h3" `BS.isPrefixOf` appProto = H3.run
                     | otherwise = HQ.run
-            runX conn conf $ http2server settings ii transport addr app
+                label
+                    | "h3" `BS.isPrefixOf` appProto = "Warp HTTP/3"
+                    | otherwise = "Warp HQ"
+            runX conn conf $ http2server label settings ii transport addr app
diff --git a/warp/Network/Wai/Handler/Warp/HTTP2.hs b/warp/Network/Wai/Handler/Warp/HTTP2.hs
index c2d327dcc..0e6705635 100644
--- a/warp/Network/Wai/Handler/Warp/HTTP2.hs
+++ b/warp/Network/Wai/Handler/Warp/HTTP2.hs
@@ -12,6 +12,7 @@ module Network.Wai.Handler.Warp.HTTP2 (
 import qualified Data.ByteString as BS
 import Data.IORef (readIORef)
 import qualified Data.IORef as I
+import GHC.Conc.Sync (labelThread, myThreadId)
 import qualified Network.HTTP2.Frame as H2
 import qualified Network.HTTP2.Server as H2
 import Network.Socket (SockAddr)
@@ -69,7 +70,7 @@ http2 settings ii conn transport app peersa th bs = do
     checkTLS
     setConnHTTP2 conn True
     H2.run H2.defaultServerConfig conf $
-        http2server settings ii transport peersa app
+        http2server "Warp HTTP/2" settings ii transport peersa app
   where
     checkTLS = case transport of
         TCP -> return () -- direct
@@ -80,13 +81,16 @@ http2 settings ii conn transport app peersa th bs = do
 --
 -- Since 3.3.11
 http2server
-    :: S.Settings
+    :: String
+    -> S.Settings
     -> InternalInfo
     -> Transport
     -> SockAddr
     -> Application
     -> H2.Server
-http2server settings ii transport addr app h2req0 aux0 response = do
+http2server label settings ii transport addr app h2req0 aux0 response = do
+    tid <- myThreadId
+    labelThread tid (label ++ " http2server " ++ show addr)
     req <- toWAIRequest h2req0 aux0
     ref <- I.newIORef Nothing
     eResponseReceived <- UnliftIO.tryAny $ app req $ \rsp -> do
diff --git a/warp/Network/Wai/Handler/Warp/Run.hs b/warp/Network/Wai/Handler/Warp/Run.hs
index cef4b1c0f..5a941bf1d 100644
--- a/warp/Network/Wai/Handler/Warp/Run.hs
+++ b/warp/Network/Wai/Handler/Warp/Run.hs
@@ -15,6 +15,7 @@ import qualified Data.ByteString as S
 import Data.IORef (newIORef, readIORef)
 import Data.Streaming.Network (bindPortTCP)
 import Foreign.C.Error (Errno (..), eCONNABORTED, eMFILE)
+import GHC.Conc.Sync (labelThread, myThreadId)
 import GHC.IO.Exception (IOErrorType (..), IOException (..))
 import Network.Socket (
     SockAddr,
@@ -327,7 +328,9 @@ fork
     -> Counter
     -> InternalInfo
     -> IO ()
-fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
+fork set mkConn addr app counter ii = settingsFork set $ \unmask -> do
+    tid <- myThreadId
+    labelThread tid "Warp just forked"
     -- Call the user-supplied on exception code if any
     -- exceptions are thrown.
     --
@@ -386,6 +389,7 @@ serveConnection
     -> IO ()
 serveConnection conn ii th origAddr transport settings app = do
     -- fixme: Upgrading to HTTP/2 should be supported.
+    tid <- myThreadId
     (h2, bs) <-
         if isHTTP2 transport
             then return (True, "")
@@ -396,8 +400,10 @@ serveConnection conn ii th origAddr transport settings app = do
                     else return (False, bs0)
     if settingsHTTP2Enabled settings && h2
         then do
+            labelThread tid ("Warp HTTP/2 " ++ show origAddr)
             http2 settings ii conn transport app origAddr th bs
         else do
+            labelThread tid ("Warp HTTP/1.1 " ++ show origAddr)
             http1 settings ii conn transport app origAddr th bs
   where
     recv4 bs0 = do