rethinkdb-client-driver 0.0.21 → 0.0.22
raw patch · 4 files changed
+195/−79 lines, 4 filesdep +stm
Dependencies added: stm
Files
- rethinkdb-client-driver.cabal +2/−1
- src/Database/RethinkDB.hs +166/−66
- src/Database/RethinkDB/Messages.hs +13/−12
- test/Test.hs +14/−0
rethinkdb-client-driver.cabal view
@@ -1,5 +1,5 @@ name: rethinkdb-client-driver-version: 0.0.21+version: 0.0.22 license: MIT license-file: LICENSE author: Tomas Carnecky@@ -43,6 +43,7 @@ , network , old-locale , scientific+ , stm , template-haskell , text , time >= 1.4
src/Database/RethinkDB.hs view
@@ -3,11 +3,16 @@ module Database.RethinkDB ( Handle- , defaultPort, newHandle, handleDatabase, close- , run, nextChunk, collect, continue, stop, wait, serverInfo+ , defaultPort, newHandle, handleDatabase, close, serverInfo - , Error(..), Response(..), ChangeNotification(..)+ -- * High-level query API+ , run, nextChunk, collect + -- * Low-level query API+ , start, continue, stop, wait, nextResult++ , Token, Error(..), Response(..), ChangeNotification(..)+ -- * The Datum type , Datum(..), Array, Object, ToDatum(..), FromDatum(..) , (.=), (.:), (.:?), object@@ -27,17 +32,30 @@ ) where +import Control.Monad+import Control.Concurrent+import Control.Concurrent.STM+ import Data.Monoid ((<>)) import Data.Text (Text) import qualified Data.Text as T +import Data.Vector (Vector) import qualified Data.Vector as V++import Data.Map.Strict (Map)+import qualified Data.Map.Strict as M+ import qualified Data.Aeson.Types as A -import Network.Socket (Socket)+import Data.Sequence (Seq, ViewR(..))+import qualified Data.Sequence as S+ import Data.IORef +import Network.Socket (Socket)+ import Database.RethinkDB.Types import Database.RethinkDB.Types.Datum import Database.RethinkDB.Messages@@ -47,9 +65,35 @@ -- Handle data Handle = Handle- { hSocket :: !Socket+ { hSocket :: !(MVar Socket)+ -- ^ Any thread can write to the socket. In theory. But I don't think+ -- Haskell allows atomic writes to a socket, so it is protected inside+ -- an 'MVar'.+ --+ -- When too many threads write to the socket, this may cause resource+ -- contention. Users are encouraged to use a resource pool to alleviate+ -- that.+ , hTokenRef :: !(IORef Token)+ -- ^ This is used to allocate new tokens. We use 'atomicModifyIORef' to+ -- efficiently allocate new tokens.+ --+ -- RethinkDB seems to expect the token to never be zero. So we need to+ -- start with one and then count up. + , hError :: !(TVar (Maybe Error))+ -- ^ If there was a fatal error while reading from the socket, it will+ -- be stored here. If this is set then no further replies will be+ -- processed. The user needs to close and re-open the handle to recover.++ , hResponses :: !(TVar (Map Token (Seq (Either Error Response))))+ -- ^ Responses to queries. A thread reads the responses from the socket+ -- and pushes them into the queues.++ , hReader :: !ThreadId+ -- ^ Thread which reads from the socket and copies responses into the+ -- queues in 'hResponses'.+ , hDatabase :: !(Exp Database) -- ^ The database which should be used when the 'Table' expression -- doesn't specify one.@@ -71,13 +115,35 @@ sendMessage sock (handshakeMessage mbAuth) _reply <- recvMessage sock handshakeReplyParser - -- RethinkDB seems to expect the token to never be null. So we start with- -- one and then count up.- ref <- newIORef 1+ err <- newTVarIO Nothing+ responses <- newTVarIO M.empty - return $ Handle sock ref db+ readerThreadId <- forkIO $ forever $ do+ res <- recvMessage sock responseMessageParser+ case res of+ Left e -> atomically $ do+ mbError <- readTVar err+ case mbError of+ Nothing -> writeTVar err (Just e)+ Just _ -> pure () + Right (Left (token, msg)) -> atomically $ modifyTVar' responses $+ M.insertWith mappend token (S.singleton $ Left $ ProtocolError $ T.pack msg) + Right (Right r) -> atomically $ modifyTVar' responses $+ M.insertWith mappend (responseToken r) (S.singleton $ Right r)++ return ()++ Handle+ <$> newMVar sock+ <*> newIORef 1+ <*> pure err+ <*> pure responses+ <*> pure readerThreadId+ <*> pure db++ -- | The 'Database' which some expressions will use when not explicitly given -- one (eg. 'Table'). handleDatabase :: Handle -> Exp Database@@ -86,41 +152,52 @@ -- | Close the given handle. You MUST NOT use the handle after this. close :: Handle -> IO ()-close handle = closeSocket (hSocket handle)+close handle = do+ withMVar (hSocket handle) closeSocket+ killThread (hReader handle) +serverInfo :: Handle -> IO (Either Error ServerInfo)+serverInfo handle = do+ token <- atomicModifyIORef' (hTokenRef handle) (\x -> (x + 1, x))+ withMVar (hSocket handle) $ \socket ->+ sendMessage socket (queryMessage token $ singleElementArray 5)+ nextResult handle token++++--------------------------------------------------------------------------------+-- * High-level query API+--+-- These are synchronous functions, they make it really easy to run a query and+-- immediately get its results.+--+-- If the result is a sequence, you can either manually iterate through the+-- chunks ('nextChunk') or fetch the whole sequence at once ('collect').++ -- | Start a new query and wait for its (first) result. If the result is an--- single value ('Datum'), then three will be no further results. If it is+-- single value ('Datum'), then there will be no further results. If it is -- a sequence, then you must consume results until the sequence ends.-run :: (FromResponse (Result a))- => Handle -> Exp a -> IO (Res a)+run :: (FromResponse (Result a)) => Handle -> Exp a -> IO (Res a) run handle expr = do- _token <- start handle expr- reply <- getResponse handle+ token <- start handle expr+ nextResult handle token - case reply of- Left e -> return $ Left e- Right response -> case responseType response of- ClientErrorType -> mkError response ClientError- CompileErrorType -> mkError response CompileError- RuntimeErrorType -> mkError response RuntimeError- _ -> return $ parseMessage parseResponse response Right -parseMessage :: (a -> A.Parser b) -> a -> (b -> Either Error c) -> Either Error c-parseMessage parser value f = case A.parseEither parser value of- Left e -> Left $ ProtocolError $ T.pack e- Right v -> f v--mkError :: Response -> (T.Text -> Error) -> IO (Either Error a)-mkError r e = return $ case V.toList (responseResult r) of- [a] -> parseMessage A.parseJSON a (Left . e)- _ -> Left $ ProtocolError $ "mkError: Could not parse error" <> T.pack (show (responseResult r))-+-- | Get the next chunk of a sequence. It is an error to request the next chunk+-- if the sequence is already 'Done',+nextChunk :: (FromResponse (Sequence a))+ => Handle -> Sequence a -> IO (Either Error (Sequence a))+nextChunk _ (Done _) = return $ Left $ ProtocolError ""+nextChunk handle (Partial token _) = do+ continue handle token+ nextResult handle token -- | Collect all the values in a sequence and make them available as -- a 'Vector a'.-collect :: (FromDatum a) => Handle -> Sequence a -> IO (Either Error (V.Vector a))+collect :: (FromDatum a) => Handle -> Sequence a -> IO (Either Error (Vector a)) collect _ (Done x) = return $ Right x collect handle s@(Partial _ x) = do chunk <- nextChunk handle s@@ -134,12 +211,21 @@ +--------------------------------------------------------------------------------+-- * Low-level query API+--+-- These functions map almost verbatim to the wire protocol messages. They are+-- asynchronous, you can send multiple queries and get the corresponding+-- responses sometime later.++ -- | Start a new query. Returns the 'Token' which can be used to track its -- progress. start :: Handle -> Exp a -> IO Token start handle term = do- token <- atomicModifyIORef (hTokenRef handle) (\x -> (x + 1, x))- sendMessage (hSocket handle) (queryMessage token msg)+ token <- atomicModifyIORef' (hTokenRef handle) (\x -> (x + 1, x))+ withMVar (hSocket handle) $ \socket ->+ sendMessage socket (queryMessage token msg) return token where@@ -159,50 +245,64 @@ -- | Let the server know that it can send the next response corresponding to -- the given token. continue :: Handle -> Token -> IO ()-continue handle token = sendMessage- (hSocket handle)- (queryMessage token $ singleElementArray 2)+continue handle token = withMVar (hSocket handle) $ \socket ->+ sendMessage socket (queryMessage token $ singleElementArray 2) -- | Stop (abort?) a query. stop :: Handle -> Token -> IO ()-stop handle token = sendMessage- (hSocket handle)- (queryMessage token $ singleElementArray 3)+stop handle token = withMVar (hSocket handle) $ \socket ->+ sendMessage socket (queryMessage token $ singleElementArray 3) -- | Wait until a previous query (which was started with the 'noreply' option) -- finishes. wait :: Handle -> Token -> IO ()-wait handle token = sendMessage- (hSocket handle)- (queryMessage token $ singleElementArray 4)+wait handle token = withMVar (hSocket handle) $ \socket ->+ sendMessage socket (queryMessage token $ singleElementArray 4) -serverInfo :: Handle -> IO (Either Error ServerInfo)-serverInfo handle = do- token <- atomicModifyIORef (hTokenRef handle) (\x -> (x + 1, x))- sendMessage (hSocket handle) (queryMessage token $ singleElementArray 5)- reply <- getResponse handle- case reply of- Left e -> return $ Left e- Right res -> return $ parseMessage parseResponse res Right+-- | This function blocks until there is a response ready for the query with+-- the given token. It may block indefinitely if the token refers to a query+-- which has already finished or does not exist yet!+responseForToken :: Handle -> Token -> IO (Either Error Response)+responseForToken h token = atomically $ do+ m <- readTVar (hResponses h)+ case M.lookup token m of+ Nothing -> retry+ Just s -> case S.viewr s of+ EmptyR -> retry+ rest :> a -> do+ modifyTVar' (hResponses h) $ if S.null rest+ then M.delete token+ else M.insert token rest + pure a --- | Get the next chunk of a sequence. It is an error to request the next chunk--- if the sequence is already 'Done',-nextChunk :: (FromResponse (Sequence a))- => Handle -> Sequence a -> IO (Either Error (Sequence a))-nextChunk _ (Done _) = return $ Left $ ProtocolError ""-nextChunk handle (Partial token _) = do- continue handle token- reply <- getResponse handle- case reply of- Left e -> return $ Left e- Right response -> return $ parseMessage parseResponse response Right +nextResult :: (FromResponse a) => Handle -> Token -> IO (Either Error a)+nextResult h token = do+ mbError <- atomically $ readTVar (hError h)+ case mbError of+ Just err -> return $ Left err+ Nothing -> do+ errorOrResponse <- responseForToken h token+ case errorOrResponse of+ Left err -> return $ Left err+ Right response -> case responseType response of+ ClientErrorType -> mkError response ClientError+ CompileErrorType -> mkError response CompileError+ RuntimeErrorType -> mkError response RuntimeError+ _ -> return $ parseMessage parseResponse response Right -getResponse :: Handle -> IO (Either Error Response)-getResponse handle = do- recvMessage (hSocket handle) responseMessageParser++parseMessage :: (a -> A.Parser b) -> a -> (b -> Either Error c) -> Either Error c+parseMessage parser value f = case A.parseEither parser value of+ Left e -> Left $ ProtocolError $ T.pack e+ Right v -> f v++mkError :: Response -> (Text -> Error) -> IO (Either Error a)+mkError r e = return $ case V.toList (responseResult r) of+ [a] -> parseMessage A.parseJSON a (Left . e)+ _ -> Left $ ProtocolError $ "mkError: Could not parse error" <> T.pack (show (responseResult r))
src/Database/RethinkDB/Messages.hs view
@@ -62,14 +62,16 @@ handshakeMessage :: Maybe Text -> ByteString handshakeMessage mbAuth = runPut $ do- -- Protocol version: V0_3- putWord32le 0x5f75e83e+ -- Protocol version: V0_4+ putWord32le 0x400c2d20 -- Authentication- flip (maybe (putWord32le 0)) mbAuth $ \auth -> do- let key = T.encodeUtf8 auth- putWord32le $ fromIntegral $ SBS.length key- putByteString $ key+ case mbAuth of+ Nothing -> putWord32le 0+ Just auth -> do+ let key = T.encodeUtf8 auth+ putWord32le $ fromIntegral $ SBS.length key+ putByteString key -- Protocol type: JSON putWord32le 0x7e6970c7@@ -89,15 +91,14 @@ buf = A.encode msg -responseMessageParser :: Get Response+responseMessageParser :: Get (Either (Token, String) Response) responseMessageParser = do token <- getWord64host len <- getWord32le buf <- getByteString (fromIntegral len) case A.eitherDecodeStrict buf of- Left e -> fail $ "responseMessageParser: response is not a JSON value (" ++ e ++ ")"- Right value -> do- case A.parseEither (responseParser token) value of- Left e -> fail $ "responseMessageParser: could not parse response (" ++ e ++ ")"- Right x -> return x+ Left e -> pure $ Left (token, "responseMessageParser: response is not a JSON value (" ++ e ++ ")")+ Right value -> case A.parseEither (responseParser token) value of+ Left e -> pure $ Left (token, "responseMessageParser: could not parse response (" ++ e ++ ")")+ Right x -> pure $ Right x
test/Test.hs view
@@ -56,6 +56,11 @@ main :: IO () main = do h <- newHandle "localhost" defaultPort Nothing (Database "test")++ si <- serverInfo h+ putStrLn "Server Info:"+ print si+ hspec $ spec h @@ -147,3 +152,12 @@ describe "lifts" $ do it "[Exp a]" $ property $ \(xs :: [Datum]) -> monadic $ do expectSuccess h (lift (map lift xs :: [Exp Datum])) (V.fromList xs)++ describe "concurrent queries" $ do+ it "should read all responses from the socket" $ property $ \(a :: Double) -> monadic $ do+ token1 <- start h $ call2 (+) (lift a) (lift a)+ token2 <- start h $ call2 (*) (lift a) (lift a)+ Right res2 <- nextResult h token2+ Right res1 <- nextResult h token1++ return $ res1 == a+a && res2 == a*a