diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Revision history for Z-Redis +## 0.3.0.1 -- 2021-02-25 + +* Add `callStream` and `StreamHandler`. + ## 0.3.0.0 -- 2021-02-25 * Add `Z.IO.RPC.MessagePack` module. diff --git a/Z-MessagePack.cabal b/Z-MessagePack.cabal --- a/Z-MessagePack.cabal +++ b/Z-MessagePack.cabal @@ -1,6 +1,6 @@ cabal-version: >=1.10 name: Z-MessagePack -version: 0.3.0.0 +version: 0.3.0.1 synopsis: MessagePack description: MessagePack binary serialization format. license: BSD3 @@ -10,8 +10,8 @@ copyright: (c) Hideyuki Tanaka, 2009-2015, (c) Dong Han, 2020 category: Data build-type: Simple -homepage: https://github.com/haskell-Z/z-redis -bug-reports: https://github.com/haskell-Z/z-redis/issues +homepage: https://github.com/ZHaskell/z-messagepack +bug-reports: https://github.com/ZHaskell/z-messagepack/issues extra-source-files: CHANGELOG.md diff --git a/Z/IO/RPC/MessagePack.hs b/Z/IO/RPC/MessagePack.hs --- a/Z/IO/RPC/MessagePack.hs +++ b/Z/IO/RPC/MessagePack.hs @@ -21,6 +21,8 @@ return (req + 1)) , ("bar", NotifyHandler $ \\ (req :: T.Text) -> do printStd (req <> "world")) + , ("qux", StreamHandler $ \\ (_ :: ()) -> do + withMVar stdinBuf (pure . sourceFromBuffered)) ] -- client @@ -28,19 +30,31 @@ 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 - call \@T.Text \@T.Text c "bar" "hello" + -- 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.Int +import Data.IORef import Z.Data.PrimRef.PrimIORef import qualified Z.Data.MessagePack.Builder as MB import qualified Z.Data.MessagePack.Value as MV @@ -80,13 +94,13 @@ return (Client seqRef reqNum bi bo) -- | Send a single RPC call and get result. -call:: (MessagePack req, MessagePack res) => Client -> T.Text -> req -> IO res +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 => Client -> T.Text -> req -> IO () +notify :: (MessagePack req, HasCallStack)=> Client -> T.Text -> req -> IO () notify c@(Client _ _ _ bo) name req = notifyPipeline c name req >> flushBuffer bo type PipelineId = Int @@ -108,9 +122,11 @@ -- 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) - modifyPrimIORef reqNum (+1) let !msgid' = msgid .&. 0xFFFFFFFF -- shrink to unsiged 32bits writeBuilder bo $ do MB.arrayHeader 4 @@ -124,7 +140,9 @@ -- -- Notify calls doesn't affect execution's result. notifyPipeline :: HasCallStack => MessagePack req => Client -> T.Text -> req -> IO () -notifyPipeline (Client _ _ _ bo) name req = do +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 @@ -132,16 +150,20 @@ MP.encodeMessagePack req -- param -- | Exception thrown when remote endpoint return errors. -data RPCException = RPCException MV.Value CallStack deriving Show +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 - n <- readPrimIORef reqNum + x <- readPrimIORef reqNum + when (x == (-1)) $ throwIO (RPCStreamUnconsumed callStack) writePrimIORef reqNum 0 - FIM.packN n <$> replicateM n (do + 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) @@ -166,6 +188,75 @@ 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 () @@ -173,6 +264,7 @@ data ServerHandler where CallHandler :: (MessagePack req, MessagePack res) => (req -> IO res) -> ServerHandler NotifyHandler :: MessagePack req => (req -> IO ()) -> ServerHandler + StreamHandler :: (MessagePack req, MessagePack res) => (req -> IO (Source res)) -> ServerHandler -- | Simple router using `FlatMap`, lookup name in /O(log(N))/. -- @@ -198,6 +290,12 @@ serveRPC :: ServerLoop -> ServerService -> IO () serveRPC serve = serveRPC' serve V.defaultChunkSize V.defaultChunkSize +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 @@ -212,16 +310,19 @@ req <- pull (sourceParserFromBuffered (do tag <- P.anyWord8 case tag of - -- notify + -- 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 (Left (name', v)) + MV.Str name' -> return (Notify (name', v)) _ -> P.fail' $ "wrong RPC name: " <> T.toText name - _ -> P.fail' $ "wrong notification type: " <> T.toText typ + 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 @@ -231,20 +332,21 @@ case typ of MV.Int 0 -> case seq of MV.Int msgid | msgid >= 0 && msgid <= 0xFFFFFFFF -> case name of - MV.Str name' -> return (Right (msgid, name', v)) + 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) + print req case req of - Just (Left (name, v)) -> do + Just (Notify (name, v)) -> do case handle name of Just (NotifyHandler f) -> do f =<< unwrap "EPARSE" (MP.convertValue v) _ -> throwOtherError "ENOTFOUND" "notification method not found" loop bi bo - Just (Right (msgid, name, v)) -> do + Just (Call (msgid, name, v)) -> do case handle name of Just (CallHandler f) -> do res <- try (f =<< unwrap "EPARSE" (MP.convertValue v)) @@ -269,4 +371,53 @@ 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 =<< 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