Z-MessagePack-0.4.0.0: Z/IO/RPC/MessagePack.hs
{-|
Module : Z.Data.MessagePack
Description : Fast MessagePack serialization/deserialization
Copyright : (c) Dong Han, 2019
License : BSD
Maintainer : winterland1989@gmail.com
Stability : experimental
Portability : non-portable
This module provides <https://github.com/msgpack-rpc/msgpack-rpc/blob/master/spec.md MessagePack-RPC> implementation.
@
-- server
import Data.IORef
import Z.IO.RPC.MessagePack
import Z.IO.Network
import Z.IO
import qualified Z.Data.Text as T
newtype ServerCtx = ServerCtx { counter :: Int }
serveRPC (startTCPServer defaultTCPServerConfig) (ServerCtx 0) $ simpleRouter
[ ("foo", CallHandler $ \\ ctx (req :: Int) -> do
modifyIORef ctx (ServerCtx . (+ 1) . counter)
return (req + 1))
, ("bar", NotifyHandler $ \\ ctx (req :: T.Text) -> do
printStd (req <> "world"))
, ("qux", StreamHandler $ \\ ctx (_ :: ()) -> do
withMVar stdinBuf (pure . sourceFromBuffered))
]
-- client
import Z.IO.RPC.MessagePack
import Z.IO.Network
import Z.IO
import qualified Z.Data.Text as T
import qualified Z.Data.Vector as V
withResource (initTCPClient defaultTCPClientConfig) $ \\ uvs -> do
c <- rpcClient uvs
-- single call
call \@Int \@Int c "foo" 1
-- notify without result
notify \@T.Text c "bar" "hello"
-- streaming result
(_, src) <- callStream c "qux" ()
runBIO $ src >|> sinkToIO (\\ b -> withMVar stdoutBuf (\\ bo -> do
writeBuffer bo b
flushBuffer bo))
@
-}
module Z.IO.RPC.MessagePack where
import Control.Concurrent
import Control.Monad
import Data.Bits
import Data.IORef
import Data.Int
import Z.Data.MessagePack (MessagePack)
import qualified Z.Data.MessagePack as MP
import qualified Z.Data.MessagePack.Builder as MB
import qualified Z.Data.MessagePack.Value as MV
import qualified Z.Data.Parser as P
import Z.Data.PrimRef.PrimIORef
import qualified Z.Data.Text as T
import qualified Z.Data.Vector as V
import qualified Z.Data.Vector.FlatIntMap as FIM
import qualified Z.Data.Vector.FlatMap as FM
import Z.IO
import Z.IO.Network
data Client = Client
{ _clientSeqRef :: Counter
, _clientPipelineReqNum :: Counter
, _clientBufferedInput :: BufferedInput
, _clientBufferedOutput :: BufferedOutput
}
-- | Open a RPC client from input/output device.
rpcClient :: (Input dev, Output dev) => dev -> IO Client
rpcClient uvs = rpcClient' uvs uvs V.defaultChunkSize V.defaultChunkSize
-- | Open a RPC client with more control.
rpcClient' :: (Input i, Output o)
=> i
-> o
-> Int -- ^ recv buffer size
-> Int -- ^ send buffer size
-> IO Client
rpcClient' i o recvBufSiz sendBufSiz = do
seqRef <- newCounter 0
reqNum <- newCounter 0
bi <- newBufferedInput' recvBufSiz i
bo <- newBufferedOutput' sendBufSiz o
return (Client seqRef reqNum bi bo)
-- | Send a single RPC call and get result.
call:: (MessagePack req, MessagePack res, HasCallStack) => Client -> T.Text -> req -> IO res
call cli name req = do
msgid <- callPipeline cli name req
fetchPipeline msgid =<< execPipeline cli
-- | Send a single notification RPC call without getting result.
notify :: (MessagePack req, HasCallStack)=> Client -> T.Text -> req -> IO ()
notify c@(Client _ _ _ bo) name req = notifyPipeline c name req >> flushBuffer bo
type PipelineId = Int
type PipelineResult = FIM.FlatIntMap MV.Value
-- | Make a call inside a pipeline, which will be sent in batch when `execPipeline`.
--
-- @
-- ...
-- fooId <- callPipeline client "foo" $ ...
-- barId <- callPipeline client "bar" $ ...
-- notifyPipeline client "qux" $ ...
--
-- r <- execPipeline client
--
-- fooResult <- fetchPipeline fooId r
-- barResult <- fetchPipeline barId r
-- @
--
callPipeline :: HasCallStack => MessagePack req => Client -> T.Text -> req -> IO PipelineId
callPipeline (Client seqRef reqNum _ bo) name req = do
x <- readPrimIORef reqNum
when (x == (-1)) $ throwIO (RPCStreamUnconsumed callStack)
writePrimIORef reqNum (x+1)
msgid <- readPrimIORef seqRef
writePrimIORef seqRef (msgid+1)
let !msgid' = msgid .&. 0xFFFFFFFF -- shrink to unsiged 32bits
writeBuilder bo $ do
MB.arrayHeader 4
MB.int 0 -- type request
MB.int (fromIntegral msgid') -- msgid
MB.str name -- method name
MP.encodeMessagePack req -- param
return msgid'
-- | Make a notify inside a pipeline, which will be sent in batch when `execPipeline`.
--
-- Notify calls doesn't affect execution's result.
notifyPipeline :: HasCallStack => MessagePack req => Client -> T.Text -> req -> IO ()
notifyPipeline (Client _ reqNum _ bo) name req = do
x <- readPrimIORef reqNum
when (x == (-1)) $ throwIO (RPCStreamUnconsumed callStack)
writeBuilder bo $ do
MB.arrayHeader 3
MB.int 2 -- type notification
MB.str name -- method name
MP.encodeMessagePack req -- param
-- | Exception thrown when remote endpoint return errors.
data RPCException
= RPCStreamUnconsumed CallStack
| RPCException MV.Value CallStack
deriving Show
instance Exception RPCException
-- | Sent request in batch and get result in a map identified by 'PipelineId'.
execPipeline :: HasCallStack => Client -> IO PipelineResult
execPipeline (Client _ reqNum bi bo) = do
flushBuffer bo
x <- readPrimIORef reqNum
when (x == (-1)) $ throwIO (RPCStreamUnconsumed callStack)
writePrimIORef reqNum 0
FIM.packN x <$> replicateM x (do
(msgid, err, v) <- readParser (do
tag <- P.anyWord8
when (tag /= 0x94) (P.fail' $ "wrong response tag: " <> T.toText tag)
!typ <- MV.value
!seq_ <- MV.value
!err <- MV.value
!v <- MV.value
case typ of
MV.Int 1 -> case seq_ of
MV.Int msgid | msgid >= 0 && msgid <= 0xFFFFFFFF ->
return (msgid, err, v)
_ -> P.fail' $ "wrong msgid: " <> T.toText seq_
_ -> P.fail' $ "wrong response type: " <> T.toText typ
) bi
when (err /= MV.Nil) $ throwIO (RPCException err callStack)
return (V.IPair (fromIntegral msgid) v))
-- | Use the `PipelineId` returned when `callPipeline` to fetch call's result.
fetchPipeline :: HasCallStack => MessagePack res => PipelineId -> PipelineResult -> IO res
fetchPipeline msgid r = do
unwrap "EPARSE" . MP.convertValue =<<
unwrap' "ENOMSG" ("missing message in response: " <> T.toText msgid)
(FIM.lookup msgid r)
-- | Call a stream method, no other `call` or `notify` should be sent until
-- returned stream is consumed completely.
--
-- This is implemented by extend MessagePack-RPC protocol by adding following new message types:
--
-- @
-- -- start stream request
-- [typ 0x04, name, param]
--
-- -- stop stream request
-- [typ 0x05]
--
-- -- each stream response
-- [typ 0x06, err, value]
--
-- -- stream response end
-- [typ 0x07]
-- @
--
-- The return tuple is a pair of a stop action and a `Source`, to terminate stream early, call the
-- stop action. Please continue consuming until EOF reached,
-- otherwise the state of the `Client` will be incorrect.
callStream :: (MessagePack req, MessagePack res, HasCallStack) => Client -> T.Text -> req -> IO (IO (), Source res)
callStream (Client seqRef reqNum bi bo) name req = do
x <- readPrimIORef reqNum
when (x == (-1)) $ throwIO (RPCStreamUnconsumed callStack)
writePrimIORef reqNum (-1)
writeBuilder bo $ do
MB.arrayHeader 3
MB.int 4 -- type request
MB.str name -- method name
MP.encodeMessagePack req -- param
flushBuffer bo
return (sendEOF, sourceFromIO $ do
res <- pull (sourceParserFromBuffered (do
tag <- P.anyWord8
-- stream stop
case tag of
0x91 -> do
!typ <- MV.value
when (typ /= MV.Int 7) $
P.fail' $ "wrong response type: " <> T.toText typ
return Nothing
0x93 -> do
!typ <- MV.value
!err <- MV.value
!v <- MV.value
when (typ /= MV.Int 6) $
P.fail' $ "wrong response type: " <> T.toText typ
return (Just (err, v))
_ -> P.fail' $ "wrong response tag: " <> T.toText tag
) bi)
-- we take tcp disconnect as eof too
case join res of
Just (err, v) -> do
when (err /= MV.Nil) $ throwIO (RPCException err callStack)
unwrap "EPARSE" (MP.convertValue v)
_ -> do
writePrimIORef reqNum 0
return Nothing
)
where
sendEOF = do
writeBuilder bo $ do
MB.arrayHeader 1
MB.int 5
flushBuffer bo
--------------------------------------------------------------------------------
type ServerLoop = (UVStream -> IO ()) -> IO ()
type ServerService a = T.Text -> Maybe (ServerHandler a)
type ServerCtx a = IORef a
data ServerHandler a where
CallHandler :: (MessagePack req, MessagePack res)
=> (ServerCtx a -> req -> IO res) -> ServerHandler a
NotifyHandler :: MessagePack req
=> (ServerCtx a -> req -> IO ()) -> ServerHandler a
StreamHandler :: (MessagePack req, MessagePack res)
=> (ServerCtx a -> req -> IO (Source res)) -> ServerHandler a
-- | Simple router using `FlatMap`, lookup name in /O(log(N))/.
--
-- @
-- import Z.IO.PRC.MessagePack
-- import Z.IO.Network
-- import Z.IO
--
-- serveRPC (startTCPServer defaultTCPServerConfig) . simpleRouter $
-- [ ("foo", CallHandler $ \\ ctx req -> do
-- ... )
-- , ("bar", CallHandler $ \\ ctx req -> do
-- ... )
-- ]
--
-- @
simpleRouter :: [(T.Text, ServerHandler a)] -> ServerService a
simpleRouter handles name = FM.lookup name handleMap
where
handleMap = FM.packR handles
-- | Serve a RPC service.
serveRPC :: ServerLoop -> a -> ServerService a -> IO ()
serveRPC serve ctxData service = do
ctxRef <- newIORef ctxData
serveRPC' serve V.defaultChunkSize V.defaultChunkSize ctxRef service
data Request a
= Notify (T.Text, a)
| Call (Int64, T.Text, a)
| StreamStart (T.Text, a)
deriving Show
-- | Serve a RPC service with more control.
serveRPC' :: ServerLoop
-> Int -- ^ recv buffer size
-> Int -- ^ send buffer size
-> ServerCtx a
-> ServerService a -> IO ()
serveRPC' serve recvBufSiz sendBufSiz ctx handle = serve $ \ uvs -> do
bi <- newBufferedInput' recvBufSiz uvs
bo <- newBufferedOutput' sendBufSiz uvs
loop bi bo
where
loop bi bo = do
req <- pull (sourceParserFromBuffered (do
tag <- P.anyWord8
case tag of
-- notify or stream start
0x93 -> do
!typ <- MV.value
!name <- MV.value
!v <- MV.value
case typ of
MV.Int 2 -> case name of
MV.Str name' -> return (Notify (name', v))
_ -> P.fail' $ "wrong RPC name: " <> T.toText name
MV.Int 4 -> case name of
MV.Str name' -> return (StreamStart (name', v))
_ -> P.fail' $ "wrong RPC name: " <> T.toText name
_ -> P.fail' $ "wrong request type: " <> T.toText typ
-- call
0x94 -> do
!typ <- MV.value
!seq_ <- MV.value
!name <- MV.value
!v <- MV.value
case typ of
MV.Int 0 -> case seq_ of
MV.Int msgid | msgid >= 0 && msgid <= 0xFFFFFFFF -> case name of
MV.Str name' -> return (Call (msgid, name', v))
_ -> P.fail' $ "wrong RPC name: " <> T.toText name
_ -> P.fail' $ "wrong msgid: " <> T.toText seq_
_ -> P.fail' $ "wrong request type: " <> T.toText typ
_ -> P.fail' $ "wrong request tag: " <> T.toText tag
) bi)
case req of
Just (Notify (name, v)) -> do
case handle name of
Just (NotifyHandler f) -> do
f ctx =<< unwrap "EPARSE" (MP.convertValue v)
_ -> throwOtherError "ENOTFOUND" "notification method not found"
loop bi bo
Just (Call (msgid, name, v)) -> do
case handle name of
Just (CallHandler f) -> do
res <- try (f ctx =<< unwrap "EPARSE" (MP.convertValue v))
writeBuilder bo $ do
MB.arrayHeader 4
MB.int 1 -- type response
MB.int (fromIntegral msgid) -- msgid
case res of
Left e -> do
MB.str (T.pack $ show (e :: SomeException))
MB.nil
Right res -> do
MB.nil
MP.encodeMessagePack res
flushBuffer bo
_ -> do
writeBuilder bo $ do
MB.arrayHeader 4
MB.int 1 -- type response
MB.int (fromIntegral msgid) -- msgid
MB.str $ "request method: " <> name <> " not found"
MB.nil
flushBuffer bo
loop bi bo
Just (StreamStart (name, v)) -> do
eofRef <- newIORef False
-- fork new thread to get stream end notification
forkIO $ do
pull (sourceParserFromBuffered (do
tag <- P.anyWord8
-- stream stop
when (tag /= 0x91) $
P.fail' $ "wrong request tag: " <> T.toText tag
!typ <- MV.value
when (typ /= MV.Int 5) $
P.fail' $ "wrong request type: " <> T.toText typ
) bi)
atomicWriteIORef eofRef True
case handle name of
Just (StreamHandler f) -> do
src <- f ctx =<< unwrap "EPARSE" (MP.convertValue v)
loopSend eofRef src bo
_ -> do
writeBuilder bo $ do
MB.arrayHeader 3
MB.int 6 -- type response
MB.str $ "request method: " <> name <> " not found"
MB.nil
flushBuffer bo
loop bi bo
_ -> return ()
loopSend eofRef src bo = do
eof <- readIORef eofRef
if eof
then do
writeBuilder bo $ do
MB.arrayHeader 1
MB.int 7 -- type response
flushBuffer bo
else do
r <- pull src
case r of
Just r' -> do
writeBuilder bo $ do
MB.arrayHeader 3
MB.int 6 -- type response
MB.nil
MP.encodeMessagePack r'
flushBuffer bo
_ -> atomicWriteIORef eofRef True
loopSend eofRef src bo