{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}

{-|
Module      : Recalc.Server
Description : Generic JSON-RPC backend with a Servant-like API.

This module implements a reactor pattern for `Recalc`, ensuring asynchronous
processing of spreadsheet updates and computations. A single threaded loop
listens for JSON-RPC messages and maintains jobs.

Jobs are picked up by multiple "worker" threads such that long-running
computations do not block messaging.

=== Example

@
type Protocol

type Api = ToApi Protocol

initialState = ...

main :: IO ()
main = runHandler @Api initialState $ \state ->
  hoist @Api (`runReaderT` state) (namedHandlers server)
 where
  server :: Protocol (AsServerT (Handler EngineState))
  server = Protocol { ... }
@
-}
module Recalc.Server
  ( Handler
  , liftEngine
  , sendIO
  , debug
  , dumpEngineState
  , scheduleJob
  , runHandler
  , module Recalc.Server.Generic
  , module Recalc.Server.Types
  , aesonOptions
  ) where

import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception qualified as Exception
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State.Strict qualified as State
import Data.Aeson qualified as Json
import Data.ByteString qualified as BS
import Data.ByteString.Lazy qualified as LB
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import GHC.IO qualified as Unsafe (unsafePerformIO)
import System.Exit (exitFailure, exitSuccess)
import System.IO
import System.IO.Error (isEOFError)
import Text.Read (readMaybe)

import Recalc.Server.Generic
import Recalc.Server.Json (aesonOptions)
import Recalc.Server.Types

data State engine = State
  { forall engine. State engine -> TChan (Job engine)
jobs :: TChan (Job engine)
  , forall engine. State engine -> TVar engine
engine :: TVar engine
  }

newtype Job engine = Job (ReaderT (State engine) IO ())

runJob :: Job engine -> State engine -> IO ()
runJob :: forall engine. Job engine -> State engine -> IO ()
runJob (Job ReaderT (State engine) IO ()
t) = ReaderT (State engine) IO () -> State engine -> IO ()
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (State engine) IO ()
t

{- "Handler" monad & utilities -}

-- | handler threads are in IO and have access to the state
type Handler engine = ReaderT (State engine) IO

scheduleJob :: ReaderT (State engine) IO () -> Handler engine ()
scheduleJob :: forall engine.
ReaderT (State engine) IO () -> ReaderT (State engine) IO ()
scheduleJob ReaderT (State engine) IO ()
x = IO () -> ReaderT (State engine) IO ()
forall a. IO a -> ReaderT (State engine) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT (State engine) IO ())
-> (TChan (Job engine) -> IO ())
-> TChan (Job engine)
-> ReaderT (State engine) IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ())
-> (TChan (Job engine) -> STM ()) -> TChan (Job engine) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TChan (Job engine) -> Job engine -> STM ()
forall a. TChan a -> a -> STM ()
`writeTChan` ReaderT (State engine) IO () -> Job engine
forall engine. ReaderT (State engine) IO () -> Job engine
Job ReaderT (State engine) IO ()
x) (TChan (Job engine) -> ReaderT (State engine) IO ())
-> ReaderT (State engine) IO (TChan (Job engine))
-> ReaderT (State engine) IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (State engine -> TChan (Job engine))
-> ReaderT (State engine) IO (TChan (Job engine))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks State engine -> TChan (Job engine)
forall engine. State engine -> TChan (Job engine)
jobs

liftEngine :: State.State engine a -> Handler engine a
liftEngine :: forall engine a. State engine a -> Handler engine a
liftEngine State engine a
f = IO a -> ReaderT (State engine) IO a
forall a. IO a -> ReaderT (State engine) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> ReaderT (State engine) IO a)
-> (TVar engine -> IO a)
-> TVar engine
-> ReaderT (State engine) IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM a -> IO a
forall a. STM a -> IO a
atomically (STM a -> IO a) -> (TVar engine -> STM a) -> TVar engine -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TVar engine -> (engine -> (a, engine)) -> STM a
forall s a. TVar s -> (s -> (a, s)) -> STM a
`stateTVar` State engine a -> engine -> (a, engine)
forall s a. State s a -> s -> (a, s)
State.runState State engine a
f) (TVar engine -> ReaderT (State engine) IO a)
-> ReaderT (State engine) IO (TVar engine)
-> ReaderT (State engine) IO a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (State engine -> TVar engine)
-> ReaderT (State engine) IO (TVar engine)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks State engine -> TVar engine
forall engine. State engine -> TVar engine
engine

debug :: Show a => String -> a -> Handler engine ()
debug :: forall a engine. Show a => String -> a -> Handler engine ()
debug String
message a
dat = IO () -> ReaderT (State engine) IO ()
forall a. IO a -> ReaderT (State engine) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT (State engine) IO ())
-> IO () -> ReaderT (State engine) IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> String -> IO ()
hPutStrLn Handle
stderr (String
message String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> a -> String
forall a. Show a => a -> String
show a
dat)

dumpEngineState :: Show engine => Handler engine ()
dumpEngineState :: forall engine. Show engine => Handler engine ()
dumpEngineState = IO () -> ReaderT (State engine) IO ()
forall a. IO a -> ReaderT (State engine) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT (State engine) IO ())
-> (TVar engine -> IO ())
-> TVar engine
-> ReaderT (State engine) IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Handle -> engine -> IO ()
forall a. Show a => Handle -> a -> IO ()
hPrint Handle
stderr (engine -> IO ())
-> (TVar engine -> IO engine) -> TVar engine -> IO ()
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< TVar engine -> IO engine
forall a. TVar a -> IO a
readTVarIO) (TVar engine -> ReaderT (State engine) IO ())
-> ReaderT (State engine) IO (TVar engine)
-> ReaderT (State engine) IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (State engine -> TVar engine)
-> ReaderT (State engine) IO (TVar engine)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks State engine -> TVar engine
forall engine. State engine -> TVar engine
engine

-- | reads lines from stdin and maintains a channel of reactor inputs with new requests,
-- handler threads deal with requests by reading from the request channel.
runHandler
  :: forall api engine. HasHandler api => engine -> (State engine -> HandlerT api IO) -> IO ()
runHandler :: forall {k} (api :: k) engine.
HasHandler api =>
engine -> (State engine -> HandlerT api IO) -> IO ()
runHandler engine
engine0 State engine -> HandlerT api IO
handlers = (IO () -> [Handler ()] -> IO ()
forall a. IO a -> [Handler a] -> IO a
`Exception.catches` [Handler ()]
forall {a}. [Handler a]
failureHandlers) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
  (Handle -> IO ()) -> [Handle] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Handle -> BufferMode -> IO ()
`hSetBuffering` BufferMode
NoBuffering) [Handle
stdin, Handle
stdout]
  Handle -> BufferMode -> IO ()
hSetBuffering Handle
stderr BufferMode
LineBuffering

  Handle -> String -> IO ()
hPutStrLn Handle
stderr String
"entered Haskell.main"

  State engine
state <- TChan (Job engine) -> TVar engine -> State engine
forall engine. TChan (Job engine) -> TVar engine -> State engine
State (TChan (Job engine) -> TVar engine -> State engine)
-> IO (TChan (Job engine)) -> IO (TVar engine -> State engine)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM (TChan (Job engine)) -> IO (TChan (Job engine))
forall a. STM a -> IO a
atomically STM (TChan (Job engine))
forall a. STM (TChan a)
newTChan IO (TVar engine -> State engine)
-> IO (TVar engine) -> IO (State engine)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> engine -> IO (TVar engine)
forall a. a -> IO (TVar a)
newTVarIO engine
engine0

  Int -> IO ThreadId -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
8 (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> (IO () -> IO ()) -> IO () -> IO ThreadId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever
    (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ (Job engine -> State engine -> IO ()
forall engine. Job engine -> State engine -> IO ()
`runJob` State engine
state) (Job engine -> IO ()) -> IO (Job engine) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< STM (Job engine) -> IO (Job engine)
forall a. STM a -> IO a
atomically (TChan (Job engine) -> STM (Job engine)
forall a. TChan a -> STM a
readTChan (State engine -> TChan (Job engine)
forall engine. State engine -> TChan (Job engine)
jobs State engine
state))

  IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    -- Content-Length: XX\n\r\n\r{....}
    (String
ln0, String
ln1) <- (,) (String -> String -> (String, String))
-> IO String -> IO (String -> (String, String))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO String
getLine' IO (String -> (String, String)) -> IO String -> IO (String, String)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO String
getLine
    case (String -> Maybe Int
forall a. Read a => String -> Maybe a
readMaybe (String -> Maybe Int) -> String -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> String -> String
forall a. Int -> [a] -> [a]
drop (forall (t :: * -> *) a. Foldable t => t a -> Int
length @[] String
"Content-Length: ") String
ln0, String
ln1) of
      (Just Int
k, String
"\r") -> do
        ByteString
reqRaw <- Handle -> Int -> IO ByteString
LB.hGet Handle
stdin Int
k
        case do JsonRpcRequest Value
x <- ByteString -> Either String (JsonRpcRequest Value)
forall a. FromJSON a => ByteString -> Either String a
Json.eitherDecode ByteString
reqRaw; forall (api :: k).
HasHandler api =>
JsonRpcRequest Value -> Handler api -> Either String (IO Value)
forall {k} (api :: k).
HasHandler api =>
JsonRpcRequest Value -> Handler api -> Either String (IO Value)
handle @api JsonRpcRequest Value
x (State engine -> HandlerT api IO
handlers State engine
state) of
          Left String
err -> Handle -> String -> IO ()
hPutStrLn Handle
stderr String
err
          Right IO Value
io -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> (Job engine -> STM ()) -> Job engine -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TChan (Job engine) -> Job engine -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan (State engine -> TChan (Job engine)
forall engine. State engine -> TChan (Job engine)
jobs State engine
state) (Job engine -> IO ()) -> Job engine -> IO ()
forall a b. (a -> b) -> a -> b
$ ReaderT (State engine) IO () -> Job engine
forall engine. ReaderT (State engine) IO () -> Job engine
Job (IO () -> ReaderT (State engine) IO ()
forall a. IO a -> ReaderT (State engine) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Value
io IO Value -> (Value -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> IO ()
forall a. ToJSON a => a -> IO ()
sendIO))
      (Maybe Int, String)
_ -> do
        Handle -> String -> IO ()
hPutStrLn Handle
stderr (String
"unexpected inputs: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
ln0 String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
ln1)
 where
  failureHandlers :: [Handler a]
failureHandlers = [(IOException -> IO a) -> Handler a
forall a e. Exception e => (e -> IO a) -> Handler a
Exception.Handler IOException -> IO a
forall {b}. IOException -> IO b
ioExcept, (SomeException -> IO a) -> Handler a
forall a e. Exception e => (e -> IO a) -> Handler a
Exception.Handler SomeException -> IO a
forall {b}. SomeException -> IO b
someExcept]

  ioExcept :: IOException -> IO b
ioExcept (IOException
e :: Exception.IOException) = Handle -> IOException -> IO ()
forall a. Show a => Handle -> a -> IO ()
hPrint Handle
stderr IOException
e IO () -> IO b -> IO b
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO b
forall a. IO a
exitFailure
  someExcept :: SomeException -> IO b
someExcept (SomeException
e :: Exception.SomeException) = Handle -> SomeException -> IO ()
forall a. Show a => Handle -> a -> IO ()
hPrint Handle
stderr SomeException
e IO () -> IO b -> IO b
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO b
forall a. IO a
exitFailure

  getLine' :: IO String
getLine' =
    IO String
getLine IO String -> (IOException -> IO String) -> IO String
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`Exception.catch` \case
      IOException
err
        | IOException -> Bool
isEOFError IOException
err -> IO String
forall a. IO a
exitSuccess
        | Bool
otherwise -> IO String
forall a. IO a
exitFailure

{-# NOINLINE stdoutLock #-}
stdoutLock :: MVar ()
stdoutLock :: MVar ()
stdoutLock = IO (MVar ()) -> MVar ()
forall a. IO a -> a
Unsafe.unsafePerformIO (() -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ())

putLbs :: LB.ByteString -> IO ()
putLbs :: ByteString -> IO ()
putLbs ByteString
lbs = MVar () -> (() -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar ()
stdoutLock ((() -> IO ()) -> IO ()) -> (() -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \()
_ -> ByteString -> IO ()
BS.putStr (ByteString -> ByteString
BS.toStrict ByteString
lbs)

-- | send JSON-RPC message on stdout (locked to prevent interleaving)
sendIO :: Json.ToJSON a => a -> IO ()
sendIO :: forall a. ToJSON a => a -> IO ()
sendIO = ByteString -> IO ()
putLbs (ByteString -> IO ()) -> (a -> ByteString) -> a -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
toRpc (ByteString -> ByteString) -> (a -> ByteString) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. ToJSON a => a -> ByteString
Json.encode
 where
  toRpc :: ByteString -> ByteString
toRpc ByteString
bs =
    [ByteString] -> ByteString
LB.concat
      [ ByteString
"Content-Length: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Int64 -> ByteString
showLB (ByteString -> Int64
LB.length ByteString
bs)
      , ByteString
"\r\n\r\n"
      , ByteString
bs
      ]

  showLB :: Int64 -> ByteString
showLB = ByteString -> ByteString
LB.fromStrict (ByteString -> ByteString)
-> (Int64 -> ByteString) -> Int64 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
Text.encodeUtf8 (Text -> ByteString) -> (Int64 -> Text) -> Int64 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Text.pack (String -> Text) -> (Int64 -> String) -> Int64 -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> String
forall a. Show a => a -> String
show