Modularise the server input and output

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.

Author: michaelpj

lsp/lsp.cabal

@@ -48,6 +48,7 @@ library
     Language.LSP.Server.Control
     Language.LSP.Server.Core
     Language.LSP.Server.Processing
+    Language.LSP.Server.IO
 
   ghc-options:        -Wall
   build-depends:

lsp/src/Language/LSP/Server.hs

@@ -1,9 +1,9 @@
 {-# LANGUAGE TypeOperators #-}
-
-module Language.LSP.Server (
-  module Language.LSP.Server.Control,
-  VFSData (..),
-  ServerDefinition (..),
+module Language.LSP.Server
+  ( module Language.LSP.Server.Control
+  , module Language.LSP.Server.IO
+  , VFSData(..)
+  , ServerDefinition(..)
 
   -- * Handlers
   Handlers (..),
@@ -63,3 +63,4 @@ module Language.LSP.Server (
 
 import Language.LSP.Server.Control
 import Language.LSP.Server.Core
+import Language.LSP.Server.IO

lsp/src/Language/LSP/Server/Control.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -14,58 +13,39 @@ module Language.LSP.Server.Control (
   LspServerLog (..),
 ) where
 
-import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&))
+import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), cmap, (<&))
 import Colog.Core qualified as L
 import Control.Applicative ((<|>))
-import Control.Concurrent
+import Control.Concurrent.Async qualified as Async
 import Control.Concurrent.STM.TChan
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.STM
 import Data.Aeson qualified as J
-import Data.Attoparsec.ByteString qualified as Attoparsec
-import Data.Attoparsec.ByteString.Char8
 import Data.ByteString qualified as BS
 import Data.ByteString.Builder.Extra (defaultChunkSize)
-import Data.ByteString.Lazy qualified as BSL
-import Data.List
 import Data.Text qualified as T
-import Data.Text.Lazy qualified as TL
-import Data.Text.Lazy.Encoding qualified as TL
 import Data.Text.Prettyprint.Doc
 import Language.LSP.Logging (defaultClientLogger)
 import Language.LSP.Protocol.Message
 import Language.LSP.Server.Core
+import Language.LSP.Server.IO qualified as IO
 import Language.LSP.Server.Processing qualified as Processing
 import Language.LSP.VFS
 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"
 
 -- ---------------------------------------------------------------------
 
@@ -115,7 +95,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
@@ -131,129 +111,32 @@ runServerWith ::
   -- | Client input.
   IO BS.ByteString ->
   -- | Function to provide output to.
-  (BSL.ByteString -> IO ()) ->
+  (BS.ByteString -> IO ()) ->
   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
+      processingLoop = initVFS $ \vfs ->
+        Processing.processingLoop
+          (cmap (fmap LspProcessingLog) ioLogger)
+          (cmap (fmap LspProcessingLog) logger)
+          vfs
+          serverDefinition
+          sendMsg
+          recvMsg
 
--- ---------------------------------------------------------------------
-
-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
+  -- Bind all the threads together so that any of them terminating will terminate everything
+  serverOut `Async.race_` serverIn `Async.race_` processingLoop
 
-_ONE_CRLF :: BS.ByteString
-_ONE_CRLF = "\r\n"
-_TWO_CRLF :: BS.ByteString
-_TWO_CRLF = "\r\n\r\n"
+  ioLogger <& Stopping `WithSeverity` Info
+  return 0

lsp/src/Language/LSP/Server/IO.hs

@@ -0,0 +1,122 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RankNTypes #-}
+
+module Language.LSP.Server.IO (serverOut, serverIn, LspIoLog) where
+
+import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&))
+import Control.Applicative ((<|>))
+import Control.Monad
+import Data.Aeson qualified as J
+import Data.Attoparsec.ByteString qualified as Attoparsec
+import Data.Attoparsec.ByteString.Char8
+import Data.ByteString qualified as BS
+import Data.ByteString.Lazy qualified as BSL
+import Data.List
+import Data.Text qualified as T
+import Data.Text.Encoding qualified as T
+import Data.Text.Prettyprint.Doc
+
+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) ->
+  -- | Channel to send out messages on.
+  (J.Value -> IO ()) ->
+  -- | Action to pull in new messages (e.g. from a handle).
+  IO BS.ByteString ->
+  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) ->
+  -- | Channel to receive messages on.
+  IO J.Value ->
+  -- | Action to send messages out on (e.g. via a handle).
+  (BS.ByteString -> IO ()) ->
+  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"