diff --git a/Database/RethinkDB/Functions.hs b/Database/RethinkDB/Functions.hs --- a/Database/RethinkDB/Functions.hs +++ b/Database/RethinkDB/Functions.hs @@ -586,13 +586,13 @@ -- -- > run' h $ table "users" # indexStatus [] indexStatus :: Expr table => [ReQL] -> table -> ReQL -indexStatus ixes tbl = op INDEX_STATUS (tbl, ixes) +indexStatus ixes tbl = op INDEX_STATUS (tbl, op ARGS [ixes]) -- | Wait for an index to be built -- -- > run' h $ table "users" # indexWait [] indexWait :: Expr table => [ReQL] -> table -> ReQL -indexWait ixes tbl = op INDEX_STATUS (tbl, ixes) +indexWait ixes tbl = op INDEX_STATUS (tbl, op ARGS [ixes]) indexRename :: Expr table => ReQL -> ReQL -> table -> ReQL indexRename from to tbl = op INDEX_RENAME (tbl, from, to) diff --git a/Database/RethinkDB/MapReduce.hs b/Database/RethinkDB/MapReduce.hs --- a/Database/RethinkDB/MapReduce.hs +++ b/Database/RethinkDB/MapReduce.hs @@ -200,7 +200,7 @@ 0 -> None $ wrap t -- Don't rewrite an operation that can be chained - 1 | (arg1 : _) <- args', notConst arg1 -> do + 1 | (arg1 : _) <- args', notConst arg1 -> fromMaybe rewrite $ mrChain type' arg1 (tail args) optargs -- Default to rewriting the term @@ -259,6 +259,8 @@ Just . ConcatMapFun $ \x -> branch (op' HAS_FIELDS (x : map wrap sel) [noRecurse]) [op' PLUCK (x : map wrap sel) [noRecurse]] () +mapMRF BRACKET [k] [] = + Just . MapFun $ \s -> op' BRACKET (s, k) [noRecurse] mapMRF _ _ _ = Nothing -- | Convert some of the built-in operations into a map/reduce @@ -307,7 +309,7 @@ rewrite1 ttype args optargs = MRF maps red mbase finals where (finally2, [mr]) = extract Nothing ttype args optargs MRF maps red mbase fin1 = mr - finals = finally2 . fin1 + finals = finally2 . return . fin1 -- | Rewrite a command that combines the result of multiple map/reduce -- operations into a single map/reduce operation @@ -318,7 +320,7 @@ maps = MapFun $ \x -> expr $ map (($ x) . getMapFun) mrs reduces a b = expr $ map (uncurry $ mkReduce a b) . index $ map getReduceFun mrs finallys = let fs = map getFinallyFun mrs in - \x -> finally . expr . map (uncurry $ mkFinally x) $ index fs + \x -> finally . map (uncurry $ mkFinally x) $ index fs mkReduce a b i f = f (a!i) (b!i) mkFinally x i f = f (x!i) getMapFun (MRF (MapFun f) _ _ _) = f @@ -337,7 +339,7 @@ -- the result of the given command extract :: Maybe Int -> TermType -> [Chain] -> [(Key, Chain)] - -> (ReQL -> ReQL, [MRF]) + -> ([ReQL] -> ReQL, [MRF]) extract st tt args optargs = fst $ flip runState st $ runWriterT $ do args' <- sequence $ map extractOne args optargvs' <- sequence $ map extractOne (map snd optargs) @@ -345,11 +347,13 @@ return $ \v -> op' tt (map ($ v) args') (zipWith (:=) optargks $ map ($ v) optargvs') where extractOne chain = either (return . const) go $ chainToMRF chain + + go :: MRF -> WriterT [MRF] (State (Maybe Int)) ([ReQL] -> ReQL) go mrf = do tell [mrf] st' <- get case st' of - Nothing -> return id + Nothing -> return head Just n -> do put $ Just $ n + 1 - return $ \v -> v ! expr n + return $ \v -> v !! n diff --git a/Database/RethinkDB/Network.hs b/Database/RethinkDB/Network.hs --- a/Database/RethinkDB/Network.hs +++ b/Database/RethinkDB/Network.hs @@ -26,16 +26,22 @@ import Control.Monad (when, forever, forM_) import Data.Typeable (Typeable) -import Network (HostName, connectTo, PortID(PortNumber)) -import System.IO (Handle, hClose, hIsEOF, hSetBuffering, BufferMode(..)) -import Data.ByteString.Lazy (hPut, hGet, ByteString) +import Network (HostName) +import Network.Socket ( + socket, Family(AF_INET), SocketType(Stream), sClose, SockAddr(SockAddrInet), setSocketOption, SocketOption(NoDelay), + Socket) +import qualified Network.Socket as Socket +import Network.BSD (getProtocolNumber, getHostByName, hostAddress) +import Network.Socket.ByteString.Lazy (sendAll) +import Network.Socket.ByteString (recv) +import Data.ByteString.Lazy (ByteString) import qualified Data.ByteString.Lazy as B import qualified Data.ByteString.UTF8 as BS (fromString) +import qualified Data.ByteString as BS import Control.Concurrent ( writeChan, MVar, Chan, modifyMVar, takeMVar, forkIO, readChan, myThreadId, newMVar, ThreadId, newChan, killThread, newEmptyMVar, putMVar, mkWeakMVar) -import Data.Monoid((<>)) import Control.Exception (catch, Exception, throwIO, SomeException(..)) import Data.IORef (IORef, newIORef, atomicModifyIORef', readIORef, writeIORef) import Data.Map (Map) @@ -46,8 +52,9 @@ import System.Mem.Weak (finalize) import Data.Binary.Get (runGet, getWord32le, getWord64le) import Data.Binary.Put (runPut, putWord32le, putWord64le, putLazyByteString) -import Data.Word (Word64, Word32) +import Data.Word (Word64, Word32, Word16) import qualified Data.HashMap.Strict as HM +import Control.Exception (bracketOnError) import Database.RethinkDB.Wire import Database.RethinkDB.Wire.Response @@ -74,7 +81,7 @@ -- | A connection to the database server data RethinkDBHandle = RethinkDBHandle { - rdbHandle :: Handle, + rdbSocket :: Socket, rdbWriteLock :: MVar (Maybe SomeException), rdbToken :: IORef Token, -- ^ The next token to use rdbDatabase :: Database, -- ^ The default database @@ -91,7 +98,7 @@ fmap f Cursor{ .. } = Cursor { cursorMap = fmap f . cursorMap, .. } instance Show RethinkDBHandle where - show RethinkDBHandle{ rdbHandle } = "RethinkDB Connection " ++ show rdbHandle + show RethinkDBHandle{ rdbSocket } = "RethinkDB Connection " ++ show rdbSocket newToken :: RethinkDBHandle -> IO Token newToken RethinkDBHandle{rdbToken} = @@ -102,6 +109,16 @@ deriving (Show, Typeable) instance Exception RethinkDBConnectionError +connectTo :: HostName -> Word16 -> IO Socket +connectTo host port = do + proto <- getProtocolNumber "tcp" + bracketOnError (socket AF_INET Stream proto) sClose $ \sock -> do + -- TODO: ipv6 + he <- getHostByName host + Socket.connect sock (SockAddrInet (fromIntegral port) (hostAddress he)) + setSocketOption sock NoDelay 1 + return sock + -- | Create a new connection to the database server -- -- /Example:/ connect using the default port with no passphrase @@ -111,41 +128,48 @@ connect :: HostName -> Integer -> Maybe String -> IO RethinkDBHandle connect host port mauth = do let auth = B.fromChunks . return . BS.fromString $ fromMaybe "" mauth - h <- connectTo host (PortNumber (fromInteger port)) - hSetBuffering h NoBuffering - hPut h $ runPut $ do + s <- connectTo host (fromInteger port) + sendAll s $ runPut $ do putWord32le magicNumber putWord32le (fromIntegral $ B.length auth) putLazyByteString auth putWord32le $ fromIntegral $ toWire Protocol.JSON - res <- hGetNullTerminatedString h + res <- sGetNullTerminatedString s when (res /= "SUCCESS") $ throwIO (RethinkDBConnectionError $ show res) r <- newIORef 1 let db' = Database "test" wlock <- newMVar Nothing waits <- newIORef M.empty - let rdb = RethinkDBHandle h wlock r db' waits + let rdb = RethinkDBHandle s wlock r db' waits tid <- forkIO $ readResponses rdb return $ rdb tid -hGetNullTerminatedString :: Handle -> IO ByteString -hGetNullTerminatedString h = go "" where +recvAll :: Socket -> Int -> IO ByteString +recvAll s n_ = go [] n_ where + go acc 0 = return $ B.fromChunks $ reverse acc + go acc n = do + d <- recv s n + if BS.null d + then throwIO $ RethinkDBConnectionError "Connection closed unexpectedly" + else go (d : acc) (n - BS.length d) + +sGetNullTerminatedString :: Socket -> IO ByteString +sGetNullTerminatedString s = go [] where go acc = do - end <- hIsEOF h - if end then return acc else do - c <- B.hGet h 1 - if c == B.pack [0] then return acc else - go (acc <> c) + c <- recv s 1 + if BS.null c || c == BS.pack [0] + then return (B.fromChunks (reverse acc)) + else go (c : acc) magicNumber :: Word32 magicNumber = fromIntegral $ toWire V0_3 -withHandle :: RethinkDBHandle -> (Handle -> IO a) -> IO a -withHandle RethinkDBHandle{ rdbHandle, rdbWriteLock } f = +withSocket :: RethinkDBHandle -> (Socket -> IO a) -> IO a +withSocket RethinkDBHandle{ rdbSocket, rdbWriteLock } f = modifyMVar rdbWriteLock $ \mex -> case mex of Nothing -> do - a <- f rdbHandle + a <- f rdbSocket return (Nothing, a) Just ex -> throwIO ex @@ -284,8 +308,8 @@ sendQLQuery :: RethinkDBHandle -> Token -> WireQuery -> IO () sendQLQuery h tok query = do let queryS = encode $ queryJSON query - withHandle h $ \s -> do - hPut s $ runPut $ do + withSocket h $ \s -> do + sendAll s $ runPut $ do putWord64le tok putWord32le (fromIntegral $ B.length queryS) putLazyByteString queryS @@ -300,22 +324,22 @@ tid <- myThreadId let h = h' tid let handler e@SomeException{} = do - hClose $ rdbHandle h + sClose $ rdbSocket h modifyMVar (rdbWriteLock h) $ \_ -> return (Just e, ()) writeIORef (rdbWait h) M.empty flip catch handler $ forever $ readSingleResponse h readSingleResponse :: RethinkDBHandle -> IO () readSingleResponse h = do - tokenString <- hGet (rdbHandle h) 8 + tokenString <- recvAll (rdbSocket h) 8 when (B.length tokenString /= 8) $ throwIO $ RethinkDBConnectionError "RethinkDB connection closed unexpectedly" let token = runGet getWord64le tokenString - header <- hGet (rdbHandle h) 4 + header <- recvAll (rdbSocket h) 4 when (B.length header /= 4) $ throwIO $ RethinkDBConnectionError "RethinkDB connection closed unexpectedly" let replyLength = runGet getWord32le header - rawResponse <- hGet (rdbHandle h) (fromIntegral replyLength) + rawResponse <- recvAll (rdbSocket h) (fromIntegral replyLength) let parsedResponse = eitherDecode rawResponse case parsedResponse of Left errMsg -> do @@ -348,10 +372,10 @@ -- | Close an open connection close :: RethinkDBHandle -> IO () -close h@RethinkDBHandle{ rdbHandle, rdbThread } = do +close h@RethinkDBHandle{ rdbSocket, rdbThread } = do noReplyWait h killThread rdbThread - hClose rdbHandle + sClose rdbSocket closeToken :: RethinkDBHandle -> Token -> IO () closeToken h tok = do diff --git a/bench/Bench.hs b/bench/Bench.hs --- a/bench/Bench.hs +++ b/bench/Bench.hs @@ -9,21 +9,20 @@ import qualified Database.RethinkDB as R import Criterion.Main import Control.Monad +import Control.Concurrent.Async main :: IO () main = do h <- prepare let test name = bench name . nfIO . void . run' h - let testn n name q = bench ("[" ++ show n ++ "x] " ++ name) . nfIO . (mapM_ next =<<) . sequence . replicate n $ runCursor h q + let testn n name q = bench (name ++ "-" ++ show n) $ nfIO $ mapM_ wait =<< replicateM n (async $ run' h q) defaultMain [ - test "nil" $ nil, - testn 10 "nil" $ nil, - testn 100 "nil" $ nil, - test "point get" $ table "bench" # get (num 0) + test "nil" $ expr Null, + testn 10 "nil" $ expr [Null], + testn 100 "nil" $ expr [Null], + testn 1000 "nil" $ expr [Null], + test "point-get" $ table "bench" # get (num 0) ] - -runCursor :: RethinkDBHandle -> ReQL -> IO (Cursor Datum) -runCursor = run prepare :: IO RethinkDBHandle prepare = do diff --git a/rethinkdb.cabal b/rethinkdb.cabal --- a/rethinkdb.cabal +++ b/rethinkdb.cabal @@ -1,5 +1,5 @@ name: rethinkdb -version: 1.15.1.0 +version: 1.15.2.0 cabal-version: >=1.8 build-type: Simple license: Apache @@ -39,7 +39,7 @@ base64-bytestring ==1.0.* if flag(dev) - exposed-modules: + other-modules: Debug exposed: True buildable: True @@ -82,9 +82,10 @@ criterion, rethinkdb, text, - aeson + aeson, + async hs-source-dirs: bench type: exitcode-stdio-1.0 main-is: Bench.hs ghc-options: -O2 -threaded -rtsopts -with-rtsopts=-N - ghc-prof-options: -fprof-auto "-with-rtsopts=-p -s -h -i0.1 -N" + ghc-prof-options: "-with-rtsopts=-p -s -h -i0.1 -N"