Modularise the server input and output

michaelpj · michaelpj · commit ad69ecf3e398 · 2023-08-27T11:05:11.000+01:00
The goal here is to make the `Control` module as boring and dispensible
as possible, so that users can put the pieces together as they like.
Thisi s a step in that direction, tackling the server in/out threads.
diff --git a/lsp/lsp.cabal b/lsp/lsp.cabal
@@ -47,6 +47,7 @@ library
     Language.LSP.Server.Control
     Language.LSP.Server.Core
     Language.LSP.Server.Processing
+    Language.LSP.Server.IO
 
   ghc-options:        -Wall
   build-depends:
diff --git a/lsp/src/Language/LSP/Server.hs b/lsp/src/Language/LSP/Server.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE TypeOperators #-}
 module Language.LSP.Server
   ( module Language.LSP.Server.Control
+  , module Language.LSP.Server.IO
   , VFSData(..)
   , ServerDefinition(..)
 
@@ -66,3 +67,4 @@ module Language.LSP.Server
 
 import Language.LSP.Server.Control
 import Language.LSP.Server.Core
+import Language.LSP.Server.IO
diff --git a/lsp/src/Language/LSP/Server/Control.hs b/lsp/src/Language/LSP/Server/Control.hs
@@ -1,7 +1,6 @@
 {-# LANGUAGE OverloadedStrings   #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE LambdaCase #-}
 
 -- So we can keep using the old prettyprinter modules (which have a better
 -- compatibility range) for now.
@@ -17,57 +16,38 @@ module Language.LSP.Server.Control
   ) where
 
 import qualified Colog.Core as L
-import           Colog.Core (LogAction (..), WithSeverity (..), Severity (..), (<&))
-import           Control.Concurrent
+import           Colog.Core (LogAction (..), WithSeverity (..), Severity (..), (<&), cmap)
+import qualified Control.Concurrent.Async as Async
 import           Control.Concurrent.STM.TChan
 import           Control.Applicative((<|>))
 import           Control.Monad
 import           Control.Monad.STM
 import           Control.Monad.IO.Class
 import qualified Data.Aeson as J
-import qualified Data.Attoparsec.ByteString as Attoparsec
-import Data.Attoparsec.ByteString.Char8
 import qualified Data.ByteString as BS
 import Data.ByteString.Builder.Extra (defaultChunkSize)
-import qualified Data.ByteString.Lazy as BSL
-import qualified Data.Text.Lazy as TL
-import qualified Data.Text.Lazy.Encoding as TL
 import qualified Data.Text as T
 import           Data.Text.Prettyprint.Doc
-import           Data.List
 import           Language.LSP.Server.Core
 import qualified Language.LSP.Server.Processing as Processing
 import           Language.LSP.Protocol.Message
 import           Language.LSP.VFS
+import qualified Language.LSP.Server.IO as IO
 import Language.LSP.Logging (defaultClientLogger)
 import           System.IO
 
 data LspServerLog =
   LspProcessingLog Processing.LspProcessingLog
-  | DecodeInitializeError String
-  | HeaderParseFail [String] String
-  | EOF
+  | LspIoLog IO.LspIoLog
   | Starting
-  | ParsedMsg T.Text
-  | SendMsg TL.Text
+  | Stopping
   deriving (Show)
 
 instance Pretty LspServerLog where
   pretty (LspProcessingLog l) = pretty l
-  pretty (DecodeInitializeError err) =
-    vsep [
-      "Got error while decoding initialize:"
-      , pretty err
-      ]
-  pretty (HeaderParseFail ctxs err) =
-    vsep [
-      "Failed to parse message header:"
-      , pretty (intercalate " > " ctxs) <> ": " <+> pretty err
-      ]
-  pretty EOF = "Got EOF"
+  pretty (LspIoLog l) = pretty l
   pretty Starting = "Starting server"
-  pretty (ParsedMsg msg) = "---> " <> pretty msg
-  pretty (SendMsg msg) = "<--2-- " <> pretty msg
+  pretty Stopping = "Stopping server"
 
 -- ---------------------------------------------------------------------
 
@@ -116,7 +96,7 @@ runServerWithHandles ioLogger logger hin hout serverDefinition = do
     clientIn = BS.hGetSome hin defaultChunkSize
 
     clientOut out = do
-      BSL.hPut hout out
+      BS.hPut hout out
       hFlush hout
 
   runServerWith ioLogger logger clientIn clientOut serverDefinition
@@ -130,134 +110,34 @@ runServerWith ::
     -- ^ The logger to use once the server has started and can successfully send messages.
     -> IO BS.ByteString
     -- ^ Client input.
-    -> (BSL.ByteString -> IO ())
+    -> (BS.ByteString -> IO ())
     -- ^ Function to provide output to.
     -> ServerDefinition config
     -> IO Int         -- exit code
 runServerWith ioLogger logger clientIn clientOut serverDefinition = do
 
   ioLogger <& Starting `WithSeverity` Info
 
-  cout <- atomically newTChan :: IO (TChan J.Value)
-  _rhpid <- forkIO $ sendServer ioLogger cout clientOut
+  cout <- atomically newTChan
+  cin <- atomically newTChan
 
-  let sendMsg msg = atomically $ writeTChan cout $ J.toJSON msg
+  let serverOut = IO.serverOut (cmap (fmap LspIoLog) ioLogger) (atomically $ readTChan cout) clientOut
+      serverIn = IO.serverIn (cmap (fmap LspIoLog) ioLogger) (atomically . writeTChan cin) clientIn
 
-  initVFS $ \vfs -> do
-    ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg
+      sendMsg msg = atomically $ writeTChan cout $ J.toJSON msg
+      recvMsg = atomically $ readTChan cin
 
-  return 1
-
--- ---------------------------------------------------------------------
-
-ioLoop ::
-  forall config
-  .  LogAction IO (WithSeverity LspServerLog)
-  -> LogAction (LspM config) (WithSeverity LspServerLog)
-  -> IO BS.ByteString
-  -> ServerDefinition config
-  -> VFS
-  -> (FromServerMessage -> IO ())
-  -> IO ()
-ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do
-  minitialize <- parseOne ioLogger clientIn (parse parser "")
-  case minitialize of
-    Nothing -> pure ()
-    Just (msg,remainder) -> do
-      case J.eitherDecode $ BSL.fromStrict msg of
-        Left err -> ioLogger <& DecodeInitializeError err `WithSeverity` Error
-        Right initialize -> do
-          mInitResp <- Processing.initializeRequestHandler pioLogger serverDefinition vfs sendMsg initialize
-          case mInitResp of
-            Nothing -> pure ()
-            Just env -> runLspT env $ loop (parse parser remainder)
-  where
-
-    pioLogger =  L.cmap (fmap LspProcessingLog) ioLogger
-    pLogger =  L.cmap (fmap LspProcessingLog) logger
-
-    loop :: Result BS.ByteString -> LspM config ()
-    loop = go
-      where
-        go r = do
-          res <- parseOne logger clientIn r
-          case res of
-            Nothing -> pure ()
-            Just (msg,remainder) -> do
-              Processing.processMessage pLogger $ BSL.fromStrict msg
-              go (parse parser remainder)
-
-    parser = do
-      try contentType <|> (return ())
-      len <- contentLength
-      try contentType <|> (return ())
-      _ <- string _ONE_CRLF
-      Attoparsec.take len
-
-    contentLength = do
-      _ <- string "Content-Length: "
-      len <- decimal
-      _ <- string _ONE_CRLF
-      return len
-
-    contentType = do 
-      _ <- string "Content-Type: "
-      skipWhile (/='\r')
-      _ <- string _ONE_CRLF
-      return ()
-
-parseOne ::
-  MonadIO m
-  => LogAction m (WithSeverity LspServerLog)
-  -> IO BS.ByteString
-  -> Result BS.ByteString
-  -> m (Maybe (BS.ByteString,BS.ByteString))
-parseOne logger 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
-      -- TODO: figure out how to re-enable
-      -- This can lead to infinite recursion in logging, see https://github.com/haskell/lsp/issues/447
-      -- logger <& ParsedMsg (T.decodeUtf8 msg) `WithSeverity` Debug
-      pure $ Just (msg,remainder)
-
--- ---------------------------------------------------------------------
-
--- | Simple server to make sure all output is serialised
-sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.ByteString -> IO ()) -> IO ()
-sendServer _logger msgChan clientOut = do
-  forever $ do
-    msg <- atomically $ readTChan msgChan
-
-    -- We need to make sure we only send over the content of the message,
-    -- and no other tags/wrapper stuff
-    let str = J.encode msg
-
-    let out = BSL.concat
-                [ TL.encodeUtf8 $ TL.pack $ "Content-Length: " ++ show (BSL.length str)
-                , BSL.fromStrict _TWO_CRLF
-                , str ]
-
-    clientOut out
-    -- TODO: figure out how to re-enable
-    -- This can lead to infinite recursion in logging, see https://github.com/haskell/lsp/issues/447
-    -- logger <& SendMsg (TL.decodeUtf8 str) `WithSeverity` Debug
-
--- |
---
---
-_ONE_CRLF :: BS.ByteString
-_ONE_CRLF = "\r\n"
-_TWO_CRLF :: BS.ByteString
-_TWO_CRLF = "\r\n\r\n"
+      processingLoop = initVFS $ \vfs ->
+        Processing.processingLoop
+        (cmap (fmap LspProcessingLog) ioLogger)
+        (cmap (fmap LspProcessingLog) logger)
+        vfs
+        serverDefinition
+        sendMsg
+        recvMsg
 
+  -- Bind all the threads together so that any of them terminating will terminate everything
+  serverOut `Async.race_` serverIn `Async.race_` processingLoop
 
+  ioLogger <& Stopping `WithSeverity` Info
+  return 0
diff --git a/lsp/src/Language/LSP/Server/IO.hs b/lsp/src/Language/LSP/Server/IO.hs
@@ -0,0 +1,116 @@
+{-# LANGUAGE OverloadedStrings   #-}
+{-# LANGUAGE RankNTypes #-}
+
+module Language.LSP.Server.IO (serverOut, serverIn, LspIoLog) where
+
+import           Colog.Core (LogAction (..), WithSeverity (..), Severity (..), (<&))
+import           Control.Monad
+import qualified Data.Aeson as J
+import qualified Data.Attoparsec.ByteString as Attoparsec
+import Data.Attoparsec.ByteString.Char8
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Lazy as BSL
+import qualified Data.Text as T
+import qualified Data.Text.Encoding as T
+import           Data.Text.Prettyprint.Doc
+import           Data.List
+import Control.Applicative ((<|>))
+
+data LspIoLog =
+  HeaderParseFail [String] String
+  | BodyParseFail String
+  | RecvMsg BS.ByteString
+  | SendMsg BS.ByteString
+  | EOF
+  deriving (Show)
+
+instance Pretty LspIoLog where
+  pretty (HeaderParseFail ctxs err) =
+    vsep [
+      "Failed to parse message header:"
+      , pretty (intercalate " > " ctxs) <> ": " <+> pretty err
+      ]
+  pretty (BodyParseFail err) =
+    vsep [
+      "Failed to parse message body:"
+      , pretty err
+      ]
+  pretty (RecvMsg msg) = "---> " <> pretty (T.decodeUtf8 msg)
+  pretty (SendMsg msg) = "<--- " <> pretty (T.decodeUtf8 msg)
+  pretty EOF = "Got EOF"
+
+-- | Process which receives messages and sends them. Output queue of messages ensures they are serialised.
+serverIn ::
+  LogAction IO (WithSeverity LspIoLog)
+  -> (J.Value -> IO ()) -- ^ Channel to send out messages on.
+  -> IO BS.ByteString -- ^ Action to pull in new messages (e.g. from a handle).
+  -> IO ()
+serverIn logger msgOut clientIn = do
+  bs <- clientIn
+  loop (parse parser bs)
+  where
+    loop :: Result BS.ByteString -> IO ()
+    loop (Fail _ ctxs err) = do
+      logger <& HeaderParseFail ctxs err `WithSeverity` Error
+      pure ()
+    loop (Partial c) = do
+      bs <- clientIn
+      if BS.null bs
+      then do
+        logger <& EOF `WithSeverity` Error
+        pure ()
+      else loop (c bs)
+    loop (Done remainder parsed) = do
+      logger <& RecvMsg parsed `WithSeverity` Debug
+      case J.eitherDecode (BSL.fromStrict parsed) of
+        -- Note: this is recoverable, because we can just discard the
+        -- message and keep going, whereas a header parse failure is
+        -- not recoverable
+        Left err -> logger <& BodyParseFail err `WithSeverity` Error
+        Right msg -> msgOut msg
+      loop (parse parser remainder)
+
+    parser = do
+      try contentType <|> (return ())
+      len <- contentLength
+      try contentType <|> (return ())
+      _ <- string _ONE_CRLF
+      Attoparsec.take len
+
+    contentLength = do
+      _ <- string "Content-Length: "
+      len <- decimal
+      _ <- string _ONE_CRLF
+      return len
+
+    contentType = do
+      _ <- string "Content-Type: "
+      skipWhile (/='\r')
+      _ <- string _ONE_CRLF
+      return ()
+
+-- | Process which receives messages and sends them. Input queue of messages ensures they are serialised.
+serverOut
+  :: LogAction IO (WithSeverity LspIoLog)
+  -> IO J.Value -- ^ Channel to receive messages on.
+  -> (BS.ByteString -> IO ()) -- ^ Action to send messages out on (e.g. via a handle).
+  -> IO ()
+serverOut logger msgIn clientOut = forever $ do
+  msg <- msgIn
+
+  -- We need to make sure we only send over the content of the message,
+  -- and no other tags/wrapper stuff
+  let str = J.encode msg
+
+  let out = BS.concat
+              [ T.encodeUtf8 $ T.pack $ "Content-Length: " ++ show (BSL.length str)
+              , _TWO_CRLF
+              , BSL.toStrict str ]
+
+  clientOut out
+  logger <& SendMsg out `WithSeverity` Debug
+
+_ONE_CRLF :: BS.ByteString
+_ONE_CRLF = "\r\n"
+_TWO_CRLF :: BS.ByteString
+_TWO_CRLF = "\r\n\r\n"
diff --git a/lsp/src/Language/LSP/Server/Processing.hs b/lsp/src/Language/LSP/Server/Processing.hs
