From d728aab20f7a362581f068e46b96f44b1d6a6b3f Mon Sep 17 00:00:00 2001
From: thomasjm <tom@codedown.io>
Date: Mon, 20 May 2024 15:42:31 -0700
Subject: [PATCH] Switch to more efficient IO loop

---
 app/Main.hs    | 10 ++++-----
 app/Streams.hs | 58 ++++++++++++++++++++++++++++----------------------
 2 files changed, 37 insertions(+), 31 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 1f61749..db93ea1 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -187,9 +187,9 @@ main = do
 
   flip runLoggingT logFn $ filterLogger logFilterFn $ flip runReaderT transformerState $
     flip withException (\(e :: SomeException) -> logErrorN [i|HNLS overall exception: #{e}|]) $
-      withAsync (ioLoop (BS.hGetSome hlsOut defaultChunkSize) (readWrappedOut optDebugHlsReads clientReqMap serverReqMap sendToStdoutWithLogging)) $ \_hlsOutAsync ->
+      withAsync (ioLoop "from-hls" (BS.hGetSome hlsOut defaultChunkSize) (handleWrappedOut optDebugHlsReads clientReqMap serverReqMap sendToStdoutWithLogging)) $ \_hlsOutAsync ->
         withAsync (readWrappedErr hlsErr) $ \_hlsErrAsync ->
-          withAsync (ioLoop (BS.hGetSome stdin defaultChunkSize) (handleStdin optDebugHlsWrites optDebugClientReads hlsIn clientReqMap serverReqMap)) $ \_stdinAsync -> do
+          withAsync (ioLoop "from-client" (BS.hGetSome stdin defaultChunkSize) (handleStdin optDebugHlsWrites optDebugClientReads hlsIn clientReqMap serverReqMap)) $ \_stdinAsync -> do
             waitForProcess p >>= \case
               ExitFailure n -> logErrorN [i|haskell-language-server subprocess exited with code #{n}|]
               ExitSuccess -> logInfoN [i|haskell-language-server subprocess exited successfully|]
@@ -226,10 +226,10 @@ handleStdin debugHlsWrites debugClientReads wrappedIn clientReqMap serverReqMap
       logDebugN [i|Sending extra notification: #{A.encode msg}|]
       writeToHlsHandle debugHlsWrites wrappedIn $ A.encode msg
 
-readWrappedOut :: (
+handleWrappedOut :: (
   MonadUnliftIO m, MonadLoggerIO m, MonadReader TransformerState m, MonadFail m
-  ) => Bool -> MVar ClientRequestMap -> MVar ServerRequestMap -> (forall a. ToJSON a => a -> m ()) -> BL8.ByteString -> m b
-readWrappedOut debugHlsReads clientReqMap serverReqMap sendToStdout bytes = forever $ do
+  ) => Bool -> MVar ClientRequestMap -> MVar ServerRequestMap -> (forall a. ToJSON a => a -> m ()) -> BL8.ByteString -> m ()
+handleWrappedOut debugHlsReads clientReqMap serverReqMap sendToStdout bytes = do
   when debugHlsReads $ logDebugN [i|Read from HLS: #{bytes}|]
 
   case A.eitherDecode bytes of
diff --git a/app/Streams.hs b/app/Streams.hs
index 0d0cd2b..71a7136 100644
--- a/app/Streams.hs
+++ b/app/Streams.hs
@@ -1,25 +1,51 @@
+{-# LANGUAGE TypeFamilies #-}
 
-module Streams where
+module Streams (ioLoop) where
 
 import Control.Applicative ((<|>))
 import Control.Monad.IO.Class
+import Control.Monad.IO.Unlift
+import Control.Monad.Logger
 import qualified Data.Attoparsec.ByteString as Attoparsec
 import Data.Attoparsec.ByteString.Char8
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Lazy as BL
+import Data.String.Interpolate
+import qualified Data.Text as T
+import UnliftIO.Exception (finally)
 
 
-ioLoop :: MonadIO m => IO BS.ByteString -> (BL.ByteString -> m ()) -> m ()
-ioLoop clientIn cb = go (parse parser mempty)
+ioLoop :: (MonadLoggerIO m, MonadUnliftIO m) => T.Text -> IO BS.ByteString -> (BL.ByteString -> m ()) -> m ()
+ioLoop loopName clientIn cb = finally (go (parse parser "")) (logDebugN [i|#{loopName}: exited|])
   where
     go r = do
-      res <- parseOne clientIn r
+      res <- parseOne loopName clientIn r
       case res of
-        Nothing -> pure ()
-        Just (msg, remainder) -> do
+        Left err -> do
+          logErrorN $ [i|#{loopName}: failed to parse: #{err}|]
+          pure ()
+        Right (msg, remainder) -> do
           cb $ BL.fromStrict msg
           go (parse parser remainder)
 
+parseOne ::
+  MonadLoggerIO m
+  => T.Text
+  -> IO BS.ByteString
+  -> Result BS.ByteString
+  -> m (Either T.Text (BS.ByteString, BS.ByteString))
+parseOne _loopName clientIn = go
+  where
+    go (Fail _ ctxs err) = do
+      pure $ Left [i|Header parse fail. Ctxs: #{ctxs}. Err: #{err}.|]
+    go (Partial c) = do
+      bs <- liftIO clientIn
+      if BS.null bs
+        then pure $ Left [i|Got null bytestring.|]
+        else go (c bs)
+    go (Done remainder msg) = do
+      pure $ Right (msg,remainder)
+
 parser :: Parser BS.ByteString
 parser = do
   try contentType <|> return ()
@@ -46,23 +72,3 @@ _ONE_CRLF :: BS.ByteString
 _ONE_CRLF = "\r\n"
 _TWO_CRLF :: BS.ByteString
 _TWO_CRLF = "\r\n\r\n"
-
-parseOne ::
-  MonadIO m
-  => IO BS.ByteString
-  -> Result BS.ByteString
-  -> m (Maybe (BS.ByteString, BS.ByteString))
-parseOne clientIn = go
-  where
-    go (Fail _ _ctxs _err) = do
-      -- logger <& HeaderParseFail ctxs err `WithSeverity` Error
-      pure Nothing
-    go (Partial c) = do
-      bs <- liftIO clientIn
-      if BS.null bs
-        then do
-          -- logger <& EOF `WithSeverity` Error
-          pure Nothing
-        else go (c bs)
-    go (Done remainder msg) = do
-      pure $ Just (msg,remainder)