Modularise the server input and output

lsp-test/lsp-test.cabal

@@ -94,12 +94,13 @@ test-suite tests
   main-is:          Test.hs
   default-language: Haskell2010
   default-extensions: ImportQualifiedPost
-  ghc-options:      -W
+  ghc-options:      -W -threaded -rtsopts -with-rtsopts=-N
   other-modules:    DummyServer
   build-depends:
     , aeson
     , base          >=4.10 && <5
     , containers
+    , co-log-core
     , data-default
     , directory
     , filepath
@@ -120,6 +121,7 @@ test-suite func-test
   default-language: Haskell2010
   default-extensions: ImportQualifiedPost
   main-is:          FuncTest.hs
+  ghc-options:      -threaded -rtsopts -with-rtsopts=-N
   build-depends:
     , base
     , co-log-core
@@ -138,6 +140,7 @@ test-suite example
   default-language:   Haskell2010
   default-extensions: ImportQualifiedPost
   main-is:            Test.hs
+  ghc-options:        -threaded -rtsopts -with-rtsopts=-N
   build-depends:
     , base
     , lsp-test
@@ -150,7 +153,7 @@ benchmark simple-bench
   default-language: Haskell2010
   default-extensions: ImportQualifiedPost
   main-is:          SimpleBench.hs
-  ghc-options:      -Wall -O2 -eventlog -rtsopts
+  ghc-options:      -Wall -O2 -eventlog -threaded -rtsopts -with-rtsopts=-N
   build-depends:
     , base
     , extra

lsp-test/test/DummyServer.hs

@@ -5,6 +5,7 @@
 
 module DummyServer where
 
+import Colog.Core qualified as L
 import Control.Monad
 import Control.Monad.Reader
 import Data.Aeson hiding (Null, defaultOptions)
@@ -26,6 +27,7 @@ import UnliftIO.Concurrent
 
 withDummyServer :: ((Handle, Handle) -> IO ()) -> IO ()
 withDummyServer f = do
+  let logger = L.cmap show L.logStringStderr
   (hinRead, hinWrite) <- createPipe
   (houtRead, houtWrite) <- createPipe
 
@@ -47,7 +49,7 @@ withDummyServer f = do
         }
 
   bracket
-    (forkIO $ void $ runServerWithHandles mempty mempty hinRead houtWrite definition)
+    (forkIO $ void $ runServerWithHandles logger (L.hoistLogAction liftIO logger) hinRead houtWrite definition)
     killThread
     (const $ f (hinWrite, houtRead))
 

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

@@ -2,6 +2,7 @@
 
 module Language.LSP.Server (
   module Language.LSP.Server.Control,
+  module Language.LSP.Server.IO,
   VFSData (..),
   ServerDefinition (..),
 
@@ -63,3 +64,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