franz (empty) → 0.2.1
raw patch · 10 files changed
+1015/−0 lines, 10 filesdep +basedep +bytestringdep +cerealsetup-changed
Dependencies added: base, bytestring, cereal, containers, cpu, directory, fast-builder, filepath, franz, fsnotify, network, process, sendfile, stm, stm-delay, transformers, unordered-containers, vector
Files
- ChangeLog.md +3/−0
- LICENSE +30/−0
- README.md +21/−0
- Setup.hs +2/−0
- app/client.hs +82/−0
- app/server.hs +39/−0
- franz.cabal +72/−0
- src/Database/Franz.hs +131/−0
- src/Database/Franz/Network.hs +393/−0
- src/Database/Franz/Reader.hs +242/−0
+ ChangeLog.md view
@@ -0,0 +1,3 @@+# Changelog for franz++## Unreleased changes
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright Fumiaki Kinoshita (c) 2018++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++ * Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++ * Redistributions in binary form must reproduce the above+ copyright notice, this list of conditions and the following+ disclaimer in the documentation and/or other materials provided+ with the distribution.++ * Neither the name of Fumiaki Kinoshita nor the names of other+ contributors may be used to endorse or promote products derived+ from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,21 @@+# Franz++## start a server++```+franzd .+```++## reading++Read 0th to 9th elements++```+franz test -r 0:9+```++Follow a stream++```+franz test -b _1+``
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ app/client.hs view
@@ -0,0 +1,82 @@+{-# LANGUAGE LambdaCase #-}+module Main where+import Database.Franz+import Database.Franz.Network++import Control.Monad+import Control.Concurrent.STM+import Data.Function (fix)+import Data.Functor.Identity+import Data.List (foldl')+import qualified Data.ByteString.Char8 as B+import Network.Socket (PortNumber)+import System.Environment+import System.IO+import System.Console.GetOpt+import System.Exit++parseHostPort :: String -> (String -> PortNumber -> r) -> r+parseHostPort str k = case break (==':') str of+ (host, ':' : port) -> k host (read port)+ (host, _) -> k host defaultPort++data Options = Options+ { host :: String+ , timeout :: Double+ , index :: Maybe B.ByteString+ , ranges :: [(Int, Int)]+ , beginning :: Maybe Int+ , prefixLength :: Bool+ }++readOffset :: String -> Int+readOffset ('_' : n) = -read n+readOffset n = read n++options :: [OptDescr (Options -> Options)]+options = [Option "h" ["host"] (ReqArg (\str o -> o { host = str }) "HOST:PORT") "stream input"+ , Option "r" ["range"] (ReqArg (\str o -> o { ranges = case break (==':') str of+ (begin, ':' : end) -> (readOffset begin, readOffset end) : ranges o+ _ -> (readOffset str, readOffset str) : ranges o+ }) "FROM:TO") "ranges"+ , Option "b" ["begin"] (ReqArg (\str o -> o { beginning = Just $! readOffset str }) "pos") "get all the contents from this position"+ , Option "t" ["timeout"] (ReqArg (\str o -> o { timeout = read str }) "SECONDS") "Timeout"+ , Option "i" ["index"] (ReqArg (\str o -> o { index = Just $ B.pack str }) "NAME") "Index name"+ , Option "" ["prefix-length"] (NoArg (\o -> o { prefixLength = True })) "Prefix payloads by their lengths"+ ]++defaultOptions :: Options+defaultOptions = Options+ { host = "localhost"+ , timeout = 1e12+ , ranges = []+ , beginning = Nothing+ , index = Nothing+ , prefixLength = False+ }++printBS :: Options -> (a, b, B.ByteString) -> IO ()+printBS o (_, _, bs) = do+ when (prefixLength o) $ print $ B.length bs+ B.hPutStr stdout bs+ hFlush stdout++main :: IO ()+main = getOpt Permute options <$> getArgs >>= \case+ (fs, [name], []) -> do+ let o = foldl' (flip id) defaultOptions fs+ parseHostPort (host o) withConnection mempty $ \conn -> do+ let name' = B.pack name+ let timeout' = floor $ timeout o * 1000000+ let f = maybe BySeqNum ByIndex (index o)+ let req i j = Query name' (f i) (f j) AllItems+ forM_ (reverse $ ranges o) $ \(i, j) -> fetchSimple conn timeout' (req i j)+ >>= mapM_ (printBS o)+ forM_ (beginning o) $ \start -> flip fix start $ \self i -> do+ bss <- fetchSimple conn timeout' (req i i)+ mapM_ (printBS o) bss+ unless (null bss) $ self $ let (j, _, _) = last bss in j + 1++ (_, _, es) -> do+ name <- getProgName+ die $ unlines ("franz [OPTIONS] STREAM" : es) ++ usageInfo name options
+ app/server.hs view
@@ -0,0 +1,39 @@+{-# LANGUAGE LambdaCase, RecordWildCards #-}+module Main where++import Data.List (foldl')+import Database.Franz.Network+import System.Environment+import System.Console.GetOpt+import System.Exit+import Network.Socket (PortNumber)++data Options = Options+ { port :: PortNumber+ , reaperInterval :: Double+ , streamLife :: Double+ }++defaultOptions :: Options+defaultOptions = Options+ { port = defaultPort+ , reaperInterval = 60+ , streamLife = 3600+ }++options :: [OptDescr (Options -> Options)]+options = [Option "p" ["port"] (ReqArg (\e o -> o { port = read e }) "NUM") "port number"+ , Option "l" ["life"] (ReqArg (\e o -> o { streamLife = read e }) "SECS") "lifespan of streams"]++main :: IO ()+main = getOpt Permute options <$> getArgs >>= \case+ (fs, args, []) -> do+ let Options{..} = foldl' (flip id) defaultOptions fs+ let start = startServer reaperInterval streamLife port+ case args of+ path : apath : _ -> start path (Just apath)+ path : _ -> start path Nothing+ [] -> start "." Nothing+ (_, _, es) -> do+ name <- getProgName+ die $ unlines ("franzd [-p PORT] PATH [ARCHIVE_PATH]" : es) ++ usageInfo name options
+ franz.cabal view
@@ -0,0 +1,72 @@+cabal-version: 2.0+version: 0.2.1+name: franz+description: Please see the README on GitHub at <https://github.com/fumieval/franz#readme>+homepage: https://github.com/fumieval/franz#readme+bug-reports: https://github.com/fumieval/franz/issues+author: Fumiaki Kinoshita+maintainer: fumiexcel@gmail.com+copyright: Copyright (c) 2019 Fumiaki Kinoshita+license: BSD3+license-file: LICENSE+build-type: Simple+category: Database+synopsis: Append-only database+extra-source-files:+ ChangeLog.md+ README.md++source-repository head+ type: git+ location: https://github.com/fumieval/franz++library+ exposed-modules:+ Database.Franz+ Database.Franz.Reader+ Database.Franz.Network+ hs-source-dirs:+ src+ build-depends:+ base >=4.7 && <5+ , bytestring+ , cereal+ , containers+ , cpu+ , directory+ , fast-builder ^>= 0.1.2.0+ , filepath+ , fsnotify+ , network+ , process+ , sendfile+ , stm+ , stm-delay+ , transformers+ , vector+ , unordered-containers+ default-language: Haskell2010+ ghc-options: -Wall -Wcompat++executable franz+ main-is: client.hs+ hs-source-dirs:+ app+ build-depends:+ base >=4.7 && <5+ , bytestring+ , franz+ , network+ , stm+ default-language: Haskell2010++executable franzd+ main-is: server.hs+ hs-source-dirs:+ app+ ghc-options: -threaded+ build-depends:+ base >=4.7 && <5+ , franz+ , network+ default-language: Haskell2010
+ src/Database/Franz.hs view
@@ -0,0 +1,131 @@+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+module Database.Franz (+ -- * Writer interface+ WriterHandle,+ openWriter,+ closeWriter,+ withWriter,+ write,+ flush,+ getLastSeqNo+ ) where++import Control.Concurrent+import Control.Exception+import Control.Monad+import qualified Data.ByteString.FastBuilder as BB+import qualified Data.Vector.Storable.Mutable as MV+import Data.Foldable (toList)+import Data.Int+import Data.Word (Word64)+import Data.IORef+import Data.Kind (Type)+import Foreign.Ptr (Ptr)+import Foreign.Storable (Storable(..))+import GHC.IO.Handle.FD (openFileBlocking)+import System.Directory+import System.Endian (toLE64)+import System.FilePath+import System.IO++data WriterHandle (f :: Type -> Type) = WriterHandle+ { hPayload :: Handle+ , hOffset :: Handle -- ^ Handle for offsets and indices+ , vOffset :: MVar (Int, Word64) -- ^ (next sequential number, current payload file size)+ , offsetBuf :: MV.IOVector Word64 -- ^ pending indices+ , offsetPtr :: IORef Int -- ^ the number of pending indices+ , indexCount :: Int -- ^ the number of Word64s to write for item+ }++-- | Get the sequential number of the last item item written.+getLastSeqNo :: WriterHandle f -> IO Int+getLastSeqNo = fmap (subtract 1 . fst) . readMVar . vOffset++openWriter :: Foldable f+ => f String+ -- ^ index names: a fixed-length collection of the names of indices for this stream.+ -- Use Proxy if you don't want any indices. If you want only one type of index, use `Identity ""`.+ -> FilePath+ -> IO (WriterHandle f)+openWriter idents path = do+ createDirectoryIfMissing True path+ let payloadPath = path </> "payloads"+ let offsetPath = path </> "offsets"+ let indexPath = path </> "indices"+ alreadyExists <- doesFileExist payloadPath+ vOffset <- if alreadyExists+ then do+ count <- withFile offsetPath ReadMode hFileSize+ size <- withFile payloadPath ReadMode hFileSize+ newMVar (fromIntegral count `div` 8, fromIntegral size)+ else newMVar (0,0)+ writeFile indexPath $ unlines $ toList idents+ -- Open the file in blocking mode because a write on a non-blocking+ -- FD makes use of an unsafe call to write(2), which in turn blocks other+ -- threads when GC runs.+ hPayload <- openFileBlocking payloadPath AppendMode+ hOffset <- openFileBlocking offsetPath AppendMode+ offsetBuf <- MV.new offsetBufferSize+ offsetPtr <- newIORef 0+ let indexCount = length idents + 1+ return WriterHandle{..}++-- | Flush any pending data and close a 'WriterHandle'.+closeWriter :: Foldable f => WriterHandle f -> IO ()+closeWriter h@WriterHandle{..} = do+ flush h+ hClose hPayload+ hClose hOffset++withWriter :: Foldable f => f String -> FilePath -> (WriterHandle f -> IO a) -> IO a+withWriter idents path = bracket (openWriter idents path) closeWriter++offsetBufferSize :: Int+offsetBufferSize = 256++write :: Foldable f+ => WriterHandle f+ -> f Int64 -- ^ index values+ -> BB.Builder -- ^ payload+ -> IO Int+write h@WriterHandle{..} ixs !bs = modifyMVar vOffset $ \(n, ofs) -> do+ len <- fromIntegral <$> BB.hPutBuilderLen hPayload bs+ let ofs' = ofs + len+ pos <- readIORef offsetPtr+ pos' <- if pos + indexCount >= offsetBufferSize+ then 0 <$ unsafeFlush h+ else return pos+ MV.write offsetBuf pos' $ toLE64 $ fromIntegral ofs'+ forM_ (zip [pos'+1..] (toList ixs))+ $ \(i, v) -> MV.write offsetBuf i $ toLE64 $ fromIntegral v+ writeIORef offsetPtr (pos' + indexCount)++ let !n' = n + 1+ return ((n', ofs'), n)+{-# INLINE write #-}++-- | Flush the changes.+flush :: WriterHandle f -> IO ()+flush h = withMVar (vOffset h) $ const $ unsafeFlush h++-- | Flush the change without locking 'vOffset'+unsafeFlush :: WriterHandle f -> IO ()+unsafeFlush WriterHandle{..} = do+ -- NB it's important to write the payload and indices prior to offsets+ -- because the reader watches the offset file then consume other files.+ -- Because of this, offsets are buffered in a buffer.+ len <- readIORef offsetPtr+ when (len > 0) $ do+ hFlush hPayload+ MV.unsafeWith offsetBuf $ \ptr -> hPutElems hOffset ptr len+ writeIORef offsetPtr 0+ hFlush hOffset++-- | Same as hPutBuf but with a number of elements to write instead of a number+-- of bytes.+hPutElems :: forall a. Storable a => Handle -> Ptr a -> Int -> IO ()+hPutElems hdl ptr len = hPutBuf hdl ptr (len * sizeOf (undefined :: a))+{-# INLINE hPutElems #-}
+ src/Database/Franz/Network.hs view
@@ -0,0 +1,393 @@+{-# LANGUAGE DeriveGeneric, LambdaCase, OverloadedStrings, ViewPatterns #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE DeriveFunctor #-}+module Database.Franz.Network+ ( startServer+ , defaultPort+ , Connection+ , withConnection+ , connect+ , disconnect+ , Query(..)+ , ItemRef(..)+ , RequestType(..)+ , defQuery+ , Response+ , awaitResponse+ , SomeIndexMap+ , Contents+ , fetch+ , fetchTraverse+ , fetchSimple+ , atomicallyWithin+ , FranzException(..)) where++import Control.Concurrent+import Control.Exception+import Control.Monad+import Control.Monad.Trans.Cont (ContT(..))+import Control.Concurrent.STM+import Control.Concurrent.STM.Delay+import Database.Franz.Reader+import qualified Data.IntMap.Strict as IM+import Data.IORef+import Data.Int (Int64)+import Data.Serialize+import qualified Data.ByteString.Char8 as B+import qualified Data.HashMap.Strict as HM+import qualified Data.Vector as V+import GHC.Generics (Generic)+import qualified Network.Socket.SendFile.Handle as SF+import qualified Network.Socket.ByteString as SB+import qualified Network.Socket as S+import System.Directory+import System.FilePath+import System.IO+import System.Process (callProcess)++defaultPort :: S.PortNumber+defaultPort = 1886++data RawRequest = RawRequest !ResponseId !Query+ | RawClean !ResponseId deriving Generic+instance Serialize RawRequest++type ResponseId = Int++data ResponseHeader = ResponseInstant !ResponseId+ -- ^ response ID, number of streams; there are items satisfying the query+ | ResponseWait !ResponseId -- ^ response ID; requested elements are not available right now+ | ResponseDelayed !ResponseId -- ^ response ID, number of streams; items are available+ | ResponseError !ResponseId !FranzException -- ^ something went wrong+ deriving (Show, Generic)+instance Serialize ResponseHeader++-- | Initial seqno, final seqno, base offset, index names+data PayloadHeader = PayloadHeader !Int !Int !Int ![B.ByteString]++instance Serialize PayloadHeader where+ put (PayloadHeader s t u xs) = f s *> f t *> f u *> put xs where+ f = putInt64le . fromIntegral+ get = PayloadHeader <$> f <*> f <*> f <*> get where+ f = fromIntegral <$> getInt64le++respond :: FranzReader+ -> IORef (IM.IntMap ThreadId)+ -> B.ByteString+ -> IORef B.ByteString+ -> MVar S.Socket -> IO ()+respond env refThreads (B.unpack -> path) buf vConn = do+ recvConn <- readMVar vConn+ runGetRecv buf recvConn get >>= \case+ Right (RawRequest reqId req) -> do+ (stream, query) <- handleQuery env path req+ join $ atomically $ do+ (ready, offsets) <- query+ return $ if ready+ then removeActivity stream >> send (ResponseInstant reqId) stream offsets+ else do+ m <- readIORef refThreads+ if IM.member reqId m+ then sendHeader $ ResponseError reqId $ MalformedRequest "duplicate request ID"+ else do+ sendHeader $ ResponseWait reqId+ -- Fork a thread to send a delayed response+ tid <- flip forkFinally (const $ removeActivity stream)+ $ join $ atomically $ do+ (ready', offsets') <- query+ check ready'+ return $ send (ResponseDelayed reqId) stream offsets'+ writeIORef refThreads $! IM.insert reqId tid m+ `catchSTM` \e -> return $ do+ removeActivity stream+ sendHeader $ ResponseError reqId e+ `catch` \e -> sendHeader $ ResponseError reqId e+ Right (RawClean reqId) -> do+ m <- readIORef refThreads+ mapM_ killThread $ IM.lookup reqId m+ writeIORef refThreads $! IM.delete reqId m+ Left err -> throwIO $ MalformedRequest err+ where+ sendHeader x = withMVar vConn $ \conn -> SB.sendAll conn $ encode x+ send header Stream{..} ((s0, p0), (s1, p1)) = withMVar vConn $ \conn -> do+ SB.sendAll conn $ encode (header, PayloadHeader s0 s1 p0 indexNames)+ let siz = 8 * (length indexNames + 1)+ SF.sendFile' conn indexHandle (fromIntegral $ siz * succ s0) (fromIntegral $ siz * (s1 - s0))+ SF.sendFile' conn payloadHandle (fromIntegral p0) (fromIntegral $ p1 - p0)++startServer+ :: Double -- reaping interval+ -> Double -- stream life (seconds)+ -> S.PortNumber+ -> FilePath -- live prefix+ -> Maybe FilePath -- archive prefix+ -> IO ()+startServer interval life port lprefix aprefix = withFranzReader lprefix $ \env -> do++ hSetBuffering stderr LineBuffering+ _ <- forkIO $ reaper interval life env++ vMountCount <- newTVarIO (HM.empty :: HM.HashMap B.ByteString Int)+ let hints = S.defaultHints { S.addrFlags = [S.AI_NUMERICHOST, S.AI_NUMERICSERV], S.addrSocketType = S.Stream }+ addr:_ <- S.getAddrInfo (Just hints) (Just "0.0.0.0") (Just $ show port)+ bracket (S.socket (S.addrFamily addr) S.Stream (S.addrProtocol addr)) S.close $ \sock -> do+ S.setSocketOption sock S.ReuseAddr 1+ S.setSocketOption sock S.NoDelay 1+ S.bind sock $ S.SockAddrInet (fromIntegral port) (S.tupleToHostAddress (0,0,0,0))+ S.listen sock S.maxListenQueue++ forever $ do+ (conn, connAddr) <- S.accept sock+ let respondLoop path = do+ SB.sendAll conn apiVersion+ hPutStrLn stderr $ unwords ["[server]", show connAddr, show path]+ ref <- newIORef IM.empty+ buf <- newIORef B.empty+ vConn <- newMVar conn+ forever (respond env ref path buf vConn) `finally` do+ readIORef ref >>= mapM_ killThread++ forkFinally (do+ decode <$> SB.recv conn 4096 >>= \case+ Left _ -> throwIO $ MalformedRequest "Expecting a path"+ Right path | Just apath <- aprefix -> do+ let src = apath </> B.unpack path+ let dest = lprefix </> B.unpack path+ join $ atomically $ do+ m <- readTVar vMountCount+ case HM.lookup path m of+ Nothing -> return $ do+ b <- doesFileExist src+ when b $ do+ createDirectoryIfMissing True dest+ callProcess "squashfuse" [src, dest]+ atomically $ writeTVar vMountCount $! HM.insert path 1 m+ Just n -> fmap pure $ writeTVar vMountCount $ HM.insert path (n + 1) m+ respondLoop path+ `finally` do+ join $ atomically $ do+ m <- readTVar vMountCount+ case HM.lookup path m of+ Just 1 -> return $ do+ callProcess "fusermount" ["-u", dest]+ atomically $ writeTVar vMountCount $ HM.delete path m+ Just n -> do+ writeTVar vMountCount $! HM.insert path (n - 1) m+ pure (pure ())+ Nothing -> pure (pure ())+ Right path -> respondLoop path+ )+ $ \result -> do+ case result of+ Left ex -> case fromException ex of+ Just e -> SB.sendAll conn $ encode $ ResponseError (-1) e+ Nothing -> logServer [show ex]+ Right _ -> return ()+ S.close conn+ where+ logServer = hPutStrLn stderr . unwords . (:) "[server]"++-- The protocol+--+-- Client Server+--- | ---- Archive prefix ---> | Mounts P if possible+--- | <--- apiVersion -------- |+--- | ---- RawRequest i p ---> |+--- | ---- RawRequest j q ---> |+--- | ---- RawRequest k r ---> |+--- | <--- ResponseInstant i - |+--- | <--- result for p ----- |+--- | <--- ResponseWait j ---- |+--- | <--- ResponseWait k ---- |+--- | <--- ResponseDelayed j - |+--- | <--- result for q ----- |+-- | ---- RawClean i ----> |+-- | ---- RawClean j ----> |+-- | ---- RawClean k ----> |++data Connection = Connection+ { connSocket :: MVar S.Socket+ , connReqId :: TVar Int+ , connStates :: TVar (IM.IntMap (ResponseStatus Contents))+ , connThread :: !ThreadId+ }++data ResponseStatus a = WaitingInstant+ | WaitingDelayed+ | Errored !FranzException+ | Available !a+ deriving (Show, Functor)++withConnection :: String -> S.PortNumber -> B.ByteString -> (Connection -> IO r) -> IO r+withConnection host port dir = bracket (connect host port dir) disconnect++apiVersion :: B.ByteString+apiVersion = "0"++connect :: String -> S.PortNumber -> B.ByteString -> IO Connection+connect host port dir = do+ let hints = S.defaultHints { S.addrFlags = [S.AI_NUMERICSERV], S.addrSocketType = S.Stream }+ addr:_ <- S.getAddrInfo (Just hints) (Just host) (Just $ show port)+ sock <- S.socket (S.addrFamily addr) S.Stream (S.addrProtocol addr)+ S.connect sock $ S.addrAddress addr+ SB.sendAll sock $ encode dir+ readyMsg <- SB.recv sock 4096+ unless (readyMsg == apiVersion) $ case decode readyMsg of+ Right (ResponseError _ e) -> throwIO e+ e -> throwIO $ ClientError $ "Database.Franz.Network.connect: Unexpected response: " ++ show e++ connSocket <- newMVar sock+ connReqId <- newTVarIO 0+ connStates <- newTVarIO IM.empty+ buf <- newIORef B.empty+ connThread <- flip forkFinally (either throwIO pure) $ forever+ $ (>>=either (throwIO . ClientError) atomically) $ runGetRecv buf sock $ get >>= \case+ ResponseInstant i -> do+ resp <- getResponse+ return $ do+ m <- readTVar connStates+ case IM.lookup i m of+ Nothing -> pure ()+ Just WaitingInstant -> writeTVar connStates $! IM.insert i (Available resp) m+ e -> throwSTM $ ClientError $ "Unexpected state on ResponseInstant " ++ show i ++ ": " ++ show e+ ResponseWait i -> return $ do+ m <- readTVar connStates+ case IM.lookup i m of+ Nothing -> pure ()+ Just WaitingInstant -> writeTVar connStates $! IM.insert i WaitingDelayed m+ e -> throwSTM $ ClientError $ "Unexpected state on ResponseWait " ++ show i ++ ": " ++ show e+ ResponseDelayed i -> do+ resp <- getResponse+ return $ do+ m <- readTVar connStates+ case IM.lookup i m of+ Nothing -> pure ()+ Just WaitingDelayed -> writeTVar connStates $! IM.insert i (Available resp) m+ e -> throwSTM $ ClientError $ "Unexpected state on ResponseDelayed " ++ show i ++ ": " ++ show e+ ResponseError i e -> return $ do+ m <- readTVar connStates+ case IM.lookup i m of+ Nothing -> throwSTM e+ Just _ -> writeTVar connStates $! IM.insert i (Errored e) m+ return Connection{..}++disconnect :: Connection -> IO ()+disconnect Connection{..} = do+ killThread connThread+ withMVar connSocket S.close++runGetRecv :: IORef B.ByteString -> S.Socket -> Get a -> IO (Either String a)+runGetRecv refBuf sock m = do+ lo <- readIORef refBuf+ let go (Done a lo') = do+ writeIORef refBuf lo'+ return $ Right a+ go (Partial cont) = SB.recv sock 4096 >>= go . cont+ go (Fail str lo') = do+ writeIORef refBuf lo'+ return $ Left $ show sock ++ str+ bs <- if B.null lo+ then SB.recv sock 4096+ else pure lo+ go $ runGetPartial m bs++defQuery :: B.ByteString -> Query+defQuery name = Query+ { reqStream = name+ , reqFrom = BySeqNum 0+ , reqTo = BySeqNum 0+ , reqType = AllItems+ }++type SomeIndexMap = HM.HashMap B.ByteString Int64++-- | (seqno, indices, payloads)+type Contents = [(Int, SomeIndexMap, B.ByteString)]++-- | When it is 'Right', it might block until the content arrives.+type Response = Either Contents (STM Contents)++awaitResponse :: STM (Either a (STM a)) -> STM a+awaitResponse = (>>=either pure id)++getResponse :: Get Contents+getResponse = do+ PayloadHeader s0 s1 p0 names <- get+ ixs <- V.replicateM (s1 - s0) $ (,) <$> fmap fromIntegral getInt64le <*> traverse (const getInt64le) names+ let ofss = V.cons p0 $ V.map fst ixs+ payload <- getByteString $ fromIntegral $ V.last ofss - p0+ return $ do+ i <- [0..s1-s0-1]+ let ofs0 = maybe (error "ofs0") id $ ofss V.!? i+ let ofs1 = maybe (error "ofs1") fst $ ixs V.!? i+ let indices = maybe (error "indices") snd $ ixs V.!? i+ pure (s0 + i + 1, HM.fromList $ zip names indices, B.take (ofs1 - ofs0) $ B.drop (ofs0 - p0) payload)+++-- | Fetch requested data from the server.+-- Termination of the continuation cancels the request, allowing flexible+-- control of its lifetime.+fetch :: Connection+ -> Query+ -> (STM Response -> IO r) -- ^ running the STM action blocks until the response arrives+ -> IO r+fetch Connection{..} req cont = do+ reqId <- atomically $ do+ i <- readTVar connReqId+ writeTVar connReqId $! i + 1+ modifyTVar' connStates $ IM.insert i WaitingInstant+ return i+ withMVar connSocket $ \sock -> SB.sendAll sock $ encode $ RawRequest reqId req+ let+ go = do+ m <- readTVar connStates+ case IM.lookup reqId m of+ Nothing -> return $ Left [] -- fetch ended; nothing to return+ Just WaitingInstant -> retry -- wait for an instant response+ Just (Available xs) -> do+ writeTVar connStates $! IM.delete reqId m+ return $ Left xs+ Just WaitingDelayed -> return $ Right $ do+ m' <- readTVar connStates+ case IM.lookup reqId m' of+ Nothing -> return [] -- fetch ended; nothing to return+ Just WaitingDelayed -> retry+ Just (Available xs) -> do+ writeTVar connStates $! IM.delete reqId m'+ return xs+ Just (Errored e) -> throwSTM e+ Just WaitingInstant -> throwSTM $ ClientError $ "fetch/WaitingDelayed: unexpected state WaitingInstant"+ Just (Errored e) -> throwSTM e+ cont go `finally` do+ withMVar connSocket $ \sock -> do+ atomically $ modifyTVar' connStates $ IM.delete reqId+ SB.sendAll sock $ encode $ RawClean reqId++-- | Queries in traversable @t@ form an atomic request. The response will become+-- available once all the elements are available.+--+-- Generalisation to Traversable guarantees that the response preserves the+-- shape of the request.+fetchTraverse :: Traversable t => Connection -> t Query -> (STM (Either (t Contents) (STM (t Contents))) -> IO r) -> IO r+fetchTraverse conn reqs = runContT $ do+ tresps <- traverse (ContT . fetch conn) reqs+ return $ do+ resps <- sequence tresps+ case traverse (either Just (const Nothing)) resps of+ Just instant -> return $ Left instant+ Nothing -> return $ Right $ traverse (either pure id) resps++-- | Send a single query and wait for the result. If it timeouts, it returns an empty list.+fetchSimple :: Connection+ -> Int -- ^ timeout in microseconds+ -> Query+ -> IO Contents+fetchSimple conn timeout req = fetch conn req (fmap (maybe [] id) . atomicallyWithin timeout . awaitResponse)++atomicallyWithin :: Int -- ^ timeout in microseconds+ -> STM a+ -> IO (Maybe a)+atomicallyWithin timeout m = do+ d <- newDelay timeout+ atomically $ fmap Just m `orElse` (Nothing <$ waitDelay d)
+ src/Database/Franz/Reader.hs view
@@ -0,0 +1,242 @@+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE RecordWildCards #-}+{-# LANGUAGE LambdaCase #-}+module Database.Franz.Reader where++import Control.Concurrent+import Control.Concurrent.STM+import Control.Exception+import Control.Monad+import Data.Serialize+import qualified Data.ByteString.Char8 as B+import qualified Data.HashMap.Strict as HM+import qualified Data.IntMap.Strict as IM+import qualified Data.Vector.Unboxed as U+import qualified Data.Vector as V+import Data.Void+import Data.Maybe (isJust)+import GHC.Clock (getMonotonicTime)+import GHC.Generics (Generic)+import System.Directory+import System.FilePath+import System.IO+import System.FSNotify++data RequestType = AllItems | LastItem deriving (Show, Generic)+instance Serialize RequestType++data ItemRef = BySeqNum !Int -- ^ sequential number+ | ByIndex !B.ByteString !Int -- ^ index name and value+ deriving (Show, Generic)+instance Serialize ItemRef++data Query = Query+ { reqStream :: !B.ByteString+ , reqFrom :: !ItemRef -- ^ name of the index to search+ , reqTo :: !ItemRef -- ^ name of the index to search+ , reqType :: !RequestType+ } deriving (Show, Generic)+instance Serialize Query++data Stream = Stream+ { vOffsets :: !(TVar (IM.IntMap Int))+ , indexNames :: ![B.ByteString]+ , indices :: !(HM.HashMap B.ByteString (TVar (IM.IntMap Int)))+ , vCount :: !(TVar Int)+ , vCaughtUp :: !(TVar Bool)+ , followThread :: !ThreadId+ , indexHandle :: !Handle+ , payloadHandle :: !Handle+ , vActivity :: !(TVar Activity)+ }++type Activity = Either Double Int++addActivity :: Activity -> Activity+addActivity (Left _) = Right 0+addActivity (Right n) = Right (n + 1)++removeActivity :: Stream -> IO ()+removeActivity str = do+ now <- getMonotonicTime+ atomically $ modifyTVar' (vActivity str) $ \case+ Left _ -> Left now+ Right n+ | n <= 0 -> Left now+ | otherwise -> Right (n - 1)++closeStream :: Stream -> IO ()+closeStream Stream{..} = do+ killThread followThread+ hClose payloadHandle+ hClose indexHandle++createStream :: WatchManager -> FilePath -> IO Stream+createStream man path = do+ let offsetPath = path </> "offsets"+ let payloadPath = path </> "payloads"+ exist <- doesFileExist offsetPath+ unless exist $ throwIO $ StreamNotFound offsetPath+ initialOffsetsBS <- B.readFile offsetPath+ payloadHandle <- openBinaryFile payloadPath ReadMode+ indexNames <- B.lines <$> B.readFile (path </> "indices")+ let icount = 1 + length indexNames+ let count = B.length initialOffsetsBS `div` (8 * icount)+ let getI = fromIntegral <$> getInt64le+ initialIndices <- either (throwIO . InternalError) pure+ $ runGet (V.replicateM count $ U.replicateM icount getI) initialOffsetsBS+ let initialOffsets = IM.fromList $ V.toList+ $ V.zip (V.enumFromN 0 count) $ V.map U.head initialIndices+ vOffsets <- newTVarIO $! initialOffsets+ vCaughtUp <- newTVarIO False+ vCount <- newTVarIO $! IM.size initialOffsets+ _ <- watchDir man path (\case+ Modified p _ _ | p == offsetPath -> True+ _ -> False)+ $ const $ atomically $ writeTVar vCaughtUp False++ vIndices <- forM [1..length indexNames] $ \i -> newTVarIO+ $ IM.fromList $ V.toList $ V.zip (V.map (U.! i) initialIndices) (V.enumFromN 0 count)++ indexHandle <- openFile offsetPath ReadMode++ let final :: Either SomeException Void -> IO ()+ final (Left exc) = logFollower [path, "terminated with", show exc]+ final (Right v) = absurd v++ -- TODO broadcast an exception if it exits?+ followThread <- flip forkFinally final $ do+ forM_ (IM.maxViewWithKey initialOffsets) $ \((i, _), _) -> do+ hSeek indexHandle AbsoluteSeek $ fromIntegral $ succ i * icount * 8+ forever $ do+ bs <- B.hGet indexHandle (8 * icount)+ if B.null bs+ then do+ atomically $ writeTVar vCaughtUp True+ atomically $ readTVar vCaughtUp >>= check . not+ else do+ ofs : indices <- either (throwIO . InternalError) pure $ runGet (replicateM icount getI) bs+ atomically $ do+ i <- readTVar vCount+ modifyTVar' vOffsets $ IM.insert i ofs+ forM_ (zip vIndices indices) $ \(v, x) -> modifyTVar' v $ IM.insert (fromIntegral x) i+ writeTVar vCount $! i + 1++ let indices = HM.fromList $ zip indexNames vIndices++ logFollower ["started", path]++ vActivity <- getMonotonicTime >>= newTVarIO . Left+ return Stream{..}+ where+ logFollower = hPutStrLn stderr . unwords . (:) "[follower]"++type QueryResult = ((Int, Int) -- starting SeqNo, byte offset+ , (Int, Int)) -- ending SeqNo, byte offset++range :: Int -- ^ from+ -> Int -- ^ to+ -> RequestType+ -> IM.IntMap Int -- ^ offsets+ -> (Bool, QueryResult)+range begin end rt allOffsets = case rt of+ AllItems -> (ready, (first, maybe first fst $ IM.maxViewWithKey body))+ LastItem -> case IM.maxViewWithKey body of+ Nothing -> (False, (zero, zero))+ Just (ofs', r) -> case IM.maxViewWithKey (IM.union left r) of+ Just (ofs, _) -> (ready, (ofs, ofs'))+ Nothing -> (ready, (zero, ofs'))+ where+ zero = (-1, 0)+ ready = isJust lastItem || not (null cont)+ (wing, lastItem, cont) = IM.splitLookup end allOffsets+ (left, body) = splitR begin $ maybe id (IM.insert end) lastItem wing+ first = maybe zero fst $ IM.maxViewWithKey left++splitR :: Int -> IM.IntMap a -> (IM.IntMap a, IM.IntMap a)+splitR i m = let (l, p, r) = IM.splitLookup i m in (l, maybe id (IM.insert i) p r)++data FranzException = MalformedRequest !String+ | StreamNotFound !FilePath+ | IndexNotFound !B.ByteString ![B.ByteString]+ | InternalError !String+ | ClientError !String+ deriving (Show, Generic)+instance Serialize FranzException+instance Exception FranzException++data FranzReader = FranzReader+ { watchManager :: WatchManager+ , vStreams :: TVar (HM.HashMap B.ByteString Stream)+ , prefix :: FilePath+ }++reaper :: Double -- interval+ -> Double -- lifetime+ -> FranzReader -> IO ()+reaper int life FranzReader{..} = forever $ do+ now <- getMonotonicTime+ (count, xs) <- atomically $ do+ list <- newTVar []+ m <- readTVar vStreams+ m' <- forM m $ \s -> readTVar (vActivity s) >>= \case+ Left t | now - t >= life -> Nothing <$ modifyTVar list (s:)+ _ -> pure $ Just s+ writeTVar vStreams $! HM.mapMaybe id m'+ (,) (HM.size m) <$> readTVar list+ mapM_ closeStream xs+ unless (null xs) $ hPutStrLn stderr $ unwords+ [ "[reaper] closed"+ , show (length xs)+ , "out of"+ , show count+ ]+ threadDelay $ floor $ int * 1e6++withFranzReader :: FilePath -> (FranzReader -> IO ()) -> IO ()+withFranzReader prefix k = do+ vStreams <- newTVarIO HM.empty+ withManager $ \watchManager -> k FranzReader{..}++handleQuery :: FranzReader+ -> FilePath+ -> Query+ -> IO (Stream, STM (Bool, QueryResult))+handleQuery FranzReader{..} dir (Query name begin_ end_ rt) = do+ streams <- readTVarIO vStreams+ let path = prefix </> dir </> B.unpack name+ let streamId = B.pack dir <> name+ stream@Stream{..} <- case HM.lookup streamId streams of+ Nothing -> do+ s <- createStream watchManager path+ atomically $ modifyTVar' vStreams $ HM.insert streamId s+ return s+ Just vStream -> return vStream+ atomically $ modifyTVar' vActivity addActivity+ return $ (,) stream $ do+ readTVar vCaughtUp >>= check+ allOffsets <- readTVar vOffsets+ let finalOffset = case IM.maxViewWithKey allOffsets of+ Just ((k, _), _) -> k + 1+ Nothing -> 0+ let rotate i+ | i < 0 = finalOffset + i+ | otherwise = i+ begin <- case begin_ of+ BySeqNum i -> pure $ rotate i+ ByIndex index val -> case HM.lookup index indices of+ Nothing -> throwSTM $ IndexNotFound index $ HM.keys indices+ Just v -> do+ m <- readTVar v+ let (_, wing) = splitR val m+ return $! maybe maxBound fst $ IM.minView wing+ end <- case end_ of+ BySeqNum i -> pure $ rotate i+ ByIndex index val -> case HM.lookup index indices of+ Nothing -> throwSTM $ IndexNotFound index $ HM.keys indices+ Just v -> do+ m <- readTVar v+ let (body, lastItem, _) = IM.splitLookup val m+ let body' = maybe id (IM.insert val) lastItem body+ return $! maybe minBound fst $ IM.maxView body'+ return $! range begin end rt allOffsets