diff --git a/rethinkdb-client-driver.cabal b/rethinkdb-client-driver.cabal --- a/rethinkdb-client-driver.cabal +++ b/rethinkdb-client-driver.cabal @@ -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 diff --git a/src/Database/RethinkDB.hs b/src/Database/RethinkDB.hs --- a/src/Database/RethinkDB.hs +++ b/src/Database/RethinkDB.hs @@ -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)) diff --git a/src/Database/RethinkDB/Messages.hs b/src/Database/RethinkDB/Messages.hs --- a/src/Database/RethinkDB/Messages.hs +++ b/src/Database/RethinkDB/Messages.hs @@ -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 diff --git a/test/Test.hs b/test/Test.hs --- a/test/Test.hs +++ b/test/Test.hs @@ -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