diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,12 +1,140 @@ # Franz -## start a server +![Haskell CI](https://github.com/fumieval/franz/workflows/Haskell%20CI/badge.svg) +Franz is an append-only container format, forked from liszt. + +Each stream is stored as a pair of concatenated payloads with an array of their +byte offsets. + +## Design requirements + +* The writer must be integrated so that no server failure blocks the application. +* There's a way to archive streams into one file. +* There's a way to fetch data in a period of time efficiently. + * In particular, the server should be able to search by timestamps, rather than performing binary search by the client. +* The server must not take too long to restart. + +## Usecase + +* Instances of franzd are running on a remote server and a local gateway. +* The application produces franz files locally using the writer API. +* On the local gateway, a proxy connects to the remote server and downsamples the file. +* Clients can connect to the gateway. When needed, they may also connect directly to the remote server. + +## Format details + +The on-disk representation of a franz stream comprises the following files: + +* `payloads`: concatenated payloads +* `offsets`: A sequence of N-tuples of 64-bit little endian integers representing + * 0th: byte offsets of payloads + * nth, n ∈ [1..N]: the value of nth index, where N is the number of index names +* `indices`: Line-separated list of index names. An index represents a 64 bit little-endian integer attached to a payload. + +A stream is stored as a directory containing the files above. + +The Franz reader also supports a squashfs image, provided that the content is a valid franz stream. + +## franzd + +franzd is a read-only server which follows franz files and gives access on wire. +Where to look for streams can be specified as a command-line argument, separately for live streams and squashfs images. + +Each stream is stored as a pair of concatenated payloads with an array of their +byte offsets. + +## Why not Kafka + +- None of us want to debug/contribute to kafka. +- Trying to read from a stream creates the stream (this is a problem due to the way we name our streams and rely on `latest`) +- Can't delete a stream as long as there is a reader existing +- Lack of understanding of it (but there is a lot of good documentation out there. [recommended][1]) +- Kafka takes a long time to start up after an abnormal shutdown on the server side +- Supports clustering but sometimes makes the reliability of the whole system worse + +[1]: https://kafka.apache.org/documentation/#design + +## Design requirements + +* The writer must be integrated so that no external process make logging fail. +* There's a way to archive streams into one file. +* There's a way to read streams without a server. +* There's a way to fetch data in a period of time efficiently. + * In particular, the server should be able to search by timestamps or gen nums, rather than performing binary search by the client. +* The server must not take too long to restart. + +## Usecase + +* Instances of franzd are running on a remote server and a local gateway. +* The application produces franz files locally using the writer API. +* On the local gateway, a proxy connects to the remote server and downsamples the file. +* Clients can connect to the gateway. When needed, they may also connect directly to the remote server. + +## Format details + +The on-disk representation of a franz stream comprises the following files: + +* `payloads`: concatenated payloads +* `offsets`: A sequence of N-tuples of 64-bit little endian integers representing + * 0th: byte offsets of payloads + * nth, n ∈ [1..N]: the value of nth index, where N is the number of index names +* `indices`: Line-separated list of index names. An index represents a 64 bit little-endian integer attached to a payload. + +A stream is stored as a directory containing the files above. + +The Franz reader also supports a squashfs image, provided that the content is a valid franz stream. + +## franzd + +franzd is a read-only server which follows franz files and gives access on wire. +Where to look for streams can be specified as a command-line argument, separately for live streams and squashfs images. + ``` -franzd . +franzd --live /path/to/live --archive /path/to/archive ``` -## reading +You can obtain a `Connection` to a remote franz file with `withConnection`. +It tries to mount a squashfs image at `path`. This is shared between connections, and unmounts when the last client closes the connection. + +```haskell +withConnection :: (MonadIO m, MonadMask m) + => String -- host + -> Int -- port + -> ByteString -- path + -> (Connection -> m r) -> m r +``` + +`fetch` returns a list of triples of offsets, tags, and payloads. + +```haskell +data RequestType = AllItems | LastItem deriving (Show, Generic) + +data ItemRef = BySeqNum !Int -- ^ sequential number + | ByIndex !B.ByteString Int -- ^ index name and value + +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) + +type SomeIndexMap = HM.HashMap B.ByteString Int64 + +type Contents = [(Int, SomeIndexMap, B.ByteString)] + +-- | When it is 'Right', it blocks until the content is available on the server. +type Response = Either Contents (STM Contents) + +fetch :: Connection + -> Query + -> (STM Response -> IO r) + -- ^ running the STM action blocks until the response arrives + -> IO r +``` + +## franz CLI: reading Read 0th to 9th elements diff --git a/app/client.hs b/app/client.hs --- a/app/client.hs +++ b/app/client.hs @@ -8,6 +8,8 @@ import Data.Function (fix) import Data.Functor.Identity import Data.List (foldl') +import qualified Data.Vector as V +import qualified Data.ByteString.Builder as BB import qualified Data.ByteString.Char8 as B import Network.Socket (PortNumber) import System.Environment @@ -22,11 +24,11 @@ data Options = Options { host :: String + , prefix :: String , timeout :: Double , index :: Maybe B.ByteString , ranges :: [(Int, Int)] , beginning :: Maybe Int - , prefixLength :: Bool } readOffset :: String -> Int @@ -41,8 +43,8 @@ }) "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 "p" ["prefix"] (ReqArg (\str o -> o { prefix = str }) "PREFIX") "Archive prefix" , 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 @@ -52,12 +54,12 @@ , ranges = [] , beginning = Nothing , index = Nothing - , prefixLength = False + , prefix = "" } printBS :: Options -> (a, b, B.ByteString) -> IO () printBS o (_, _, bs) = do - when (prefixLength o) $ print $ B.length bs + BB.hPutBuilder stdout $ BB.word64LE $ fromIntegral $ B.length bs B.hPutStr stdout bs hFlush stdout @@ -65,7 +67,7 @@ main = getOpt Permute options <$> getArgs >>= \case (fs, [name], []) -> do let o = foldl' (flip id) defaultOptions fs - parseHostPort (host o) withConnection mempty $ \conn -> do + parseHostPort (host o) withConnection (B.pack $ prefix o) $ \conn -> do let name' = B.pack name let timeout' = floor $ timeout o * 1000000 let f = maybe BySeqNum ByIndex (index o) @@ -75,7 +77,7 @@ 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 + unless (null bss) $ self $ let (j, _, _) = V.last bss in j + 1 (_, _, es) -> do name <- getProgName diff --git a/app/server.hs b/app/server.hs --- a/app/server.hs +++ b/app/server.hs @@ -1,39 +1,17 @@ {-# 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 - } +import Database.Franz.Server +import Options.Applicative -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"] +options :: Parser Settings +options = Settings + <$> option auto (long "reap-interval" <> value 60 <> metavar "SECONDS" <> help "Stream reaping interval") + <*> option auto (long "stream-lifetime" <> value 3600 <> metavar "SECONDS" <> help "Number of seconds to leave stream open with no readers") + <*> option auto (long "port" <> value defaultPort <> help "Port number") + <*> strOption (long "live" <> value "." <> metavar "DIR" <> help "Live prefix") + <*> optional (strOption (long "archive" <> metavar "DIR" <> help "Search for squashfs archives in this directory. If none found, search for live streams instead.")) + <*> strOption (long "mount" <> value "/tmp/franz" <> metavar "DIR" <> help "Mount prefix") 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 +main = execParser (info options mempty) >>= startServer diff --git a/franz.cabal b/franz.cabal --- a/franz.cabal +++ b/franz.cabal @@ -1,5 +1,5 @@ cabal-version: 2.0 -version: 0.2.1 +version: 0.3 name: franz description: Please see the README on GitHub at homepage: https://github.com/fumieval/franz#readme @@ -23,8 +23,11 @@ library exposed-modules: Database.Franz + Database.Franz.Internal Database.Franz.Reader + Database.Franz.Server Database.Franz.Network + Database.Franz.Protocol hs-source-dirs: src build-depends: @@ -32,17 +35,22 @@ , bytestring , cereal , containers + , concurrent-resource-map ^>=0.2 , cpu + , deepseq , directory , fast-builder ^>= 0.1.2.0 , filepath , fsnotify + , mtl , network , process + , retry , sendfile , stm , stm-delay , transformers + , unboxed-ref , vector , unordered-containers default-language: Haskell2010 @@ -58,15 +66,17 @@ , franz , network , stm + , vector default-language: Haskell2010 executable franzd main-is: server.hs hs-source-dirs: app - ghc-options: -threaded + ghc-options: -threaded -O2 -Wall build-depends: base >=4.7 && <5 , franz , network + , optparse-applicative default-language: Haskell2010 diff --git a/src/Database/Franz.hs b/src/Database/Franz.hs --- a/src/Database/Franz.hs +++ b/src/Database/Franz.hs @@ -14,6 +14,7 @@ ) where import Control.Concurrent +import Control.DeepSeq (NFData(..)) import Control.Exception import Control.Monad import qualified Data.ByteString.FastBuilder as BB @@ -32,13 +33,15 @@ 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 + { 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 } +instance NFData (WriterHandle f) where + rnf WriterHandle{} = () -- | Get the sequential number of the last item item written. getLastSeqNo :: WriterHandle f -> IO Int @@ -74,7 +77,7 @@ return WriterHandle{..} -- | Flush any pending data and close a 'WriterHandle'. -closeWriter :: Foldable f => WriterHandle f -> IO () +closeWriter :: WriterHandle f -> IO () closeWriter h@WriterHandle{..} = do flush h hClose hPayload diff --git a/src/Database/Franz/Internal.hs b/src/Database/Franz/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Database/Franz/Internal.hs @@ -0,0 +1,28 @@ +module Database.Franz.Internal (getInt64le, runGetRecv) where + +import qualified Data.ByteString as B +import Data.IORef +import Data.Serialize hiding (getInt64le) +import Network.Socket as S +import Network.Socket.ByteString as SB +import System.Endian (fromLE64) + +-- | Better implementation of 'Data.Serialize.getInt64le' +getInt64le :: Num a => Get a +getInt64le = fromIntegral . fromLE64 <$> getWord64host +{-# INLINE getInt64le #-} + +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 diff --git a/src/Database/Franz/Network.hs b/src/Database/Franz/Network.hs --- a/src/Database/Franz/Network.hs +++ b/src/Database/Franz/Network.hs @@ -1,10 +1,10 @@ -{-# LANGUAGE DeriveGeneric, LambdaCase, OverloadedStrings, ViewPatterns #-} +{-# LANGUAGE LambdaCase, OverloadedStrings #-} {-# LANGUAGE RecordWildCards #-} -{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE TypeFamilies #-} module Database.Franz.Network - ( startServer - , defaultPort + ( defaultPort , Connection , withConnection , connect @@ -18,175 +18,30 @@ , SomeIndexMap , Contents , fetch - , fetchTraverse , fetchSimple , atomicallyWithin , FranzException(..)) where +import Control.Arrow ((&&&)) import Control.Concurrent +import Control.Concurrent.STM +import Control.Concurrent.STM.Delay (newDelay, waitDelay) 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 Data.ConcurrentResourceMap import qualified Data.HashMap.Strict as HM +import Data.IORef +import Data.IORef.Unboxed +import Data.Int (Int64) +import qualified Data.IntMap.Strict as IM +import Data.Serialize hiding (getInt64le) 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 Data.Vector.Generic.Mutable as VGM +import Database.Franz.Internal +import Database.Franz.Protocol 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]" +import qualified Network.Socket.ByteString as SB -- The protocol -- @@ -206,10 +61,21 @@ -- | ---- RawClean j ----> | -- | ---- RawClean k ----> | +newtype ConnStateMap v = ConnStateMap (IM.IntMap v) + +instance ResourceMap ConnStateMap where + type Key ConnStateMap = Int + empty = ConnStateMap IM.empty + delete k (ConnStateMap m) = ConnStateMap (IM.delete k m) + insert k v (ConnStateMap m) = ConnStateMap (IM.insert k v m) + lookup k (ConnStateMap m) = IM.lookup k m + data Connection = Connection { connSocket :: MVar S.Socket - , connReqId :: TVar Int - , connStates :: TVar (IM.IntMap (ResponseStatus Contents)) + , connReqId :: !Counter + , connStates :: !(ConcurrentResourceMap + ConnStateMap + (TVar (ResponseStatus Contents))) , connThread :: !ThreadId } @@ -217,19 +83,19 @@ | WaitingDelayed | Errored !FranzException | Available !a + -- | The user cancelled the request. + | RequestFinished 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.setSocketOption sock S.NoDelay 1 S.connect sock $ S.addrAddress addr SB.sendAll sock $ encode dir readyMsg <- SB.recv sock 4096 @@ -238,38 +104,38 @@ e -> throwIO $ ClientError $ "Database.Franz.Network.connect: Unexpected response: " ++ show e connSocket <- newMVar sock - connReqId <- newTVarIO 0 - connStates <- newTVarIO IM.empty + connReqId <- newCounter 0 + connStates <- newResourceMap 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 + let -- Get a reference to shared state for the request if it exists. + withRequest i f = withInitialisedResource connStates i (\_ -> pure ()) $ \case + Nothing -> + -- If it throws an exception on no value, great, it will + -- float out here. If it returns a value, it'll just be + -- ignore as we can't do anything with it anyway. + void $ atomically (f Nothing) + Just reqVar -> atomically $ readTVar reqVar >>= \case + -- If request is finished, do nothing to the content. + RequestFinished -> void $ f Nothing + s -> f (Just s) >>= mapM_ (writeTVar reqVar) + + runGetThrow :: Get a -> IO a + runGetThrow g = runGetRecv buf sock g + >>= either (throwIO . ClientError) pure + + connThread <- flip forkFinally (either throwIO pure) $ forever $ runGetThrow get >>= \case + Response i -> do + resp <- runGetThrow getResponse + withRequest i . traverse $ \case + WaitingInstant -> pure (Available resp) + WaitingDelayed -> pure (Available resp) + e -> throwSTM $ ClientError $ "Unexpected state on ResponseInstant " ++ show i ++ ": " ++ show e + ResponseWait i -> withRequest i . traverse $ \case + WaitingInstant -> pure WaitingDelayed + e -> throwSTM $ ClientError $ "Unexpected state on ResponseWait " ++ show i ++ ": " ++ show e + ResponseError i e -> withRequest i $ \case + Nothing -> throwSTM e + Just{} -> pure $ Just (Errored e) return Connection{..} disconnect :: Connection -> IO () @@ -277,21 +143,6 @@ 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 @@ -300,10 +151,10 @@ , reqType = AllItems } -type SomeIndexMap = HM.HashMap B.ByteString Int64 +type SomeIndexMap = HM.HashMap IndexName Int64 -- | (seqno, indices, payloads) -type Contents = [(Int, SomeIndexMap, B.ByteString)] +type Contents = V.Vector (Int, SomeIndexMap, B.ByteString) -- | When it is 'Right', it might block until the content arrives. type Response = Either Contents (STM Contents) @@ -313,77 +164,93 @@ 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) - + PayloadHeader s0 s1 p0 names <- get + let df = s1 - s0 + if df <= 0 + then pure mempty + else do + ixs <- V.replicateM df $ (,) <$> getInt64le <*> traverse (const getInt64le) names + payload <- getByteString $ fst (V.unsafeLast ixs) - p0 + pure $ V.create $ do + vres <- VGM.unsafeNew df + let go i ofs0 + | i >= df = pure () + | otherwise = do + let (ofs1, indices) = V.unsafeIndex ixs i + !m = HM.fromList $ zip names indices + !bs = B.take (ofs1 - ofs0) $ B.drop (ofs0 - p0) payload + !num = s0 + i + 1 + VGM.unsafeWrite vres i (num, m, bs) + go (i + 1) ofs1 + go 0 p0 + return vres -- | Fetch requested data from the server. --- Termination of the continuation cancels the request, allowing flexible --- control of its lifetime. -fetch :: Connection +-- +-- Termination of 'fetch' 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 + -> (STM Response -> IO r) + -- ^ Wait for the response in a blocking manner. You should only run + -- the continuation inside a 'fetch' block: leaking the STM action + -- and running it outside will result in a 'ClientError' exception. -> 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 + reqId <- atomicAddCounter connReqId 1 --- | 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 + let -- When we exit the scope of the request, ensure that we cancel any + -- outstanding request and set the appropriate state, lest the user + -- leaks the resource and tries to re-run the provided action. + cleanupRequest reqVar = do + let inFlight WaitingInstant = True + inFlight WaitingDelayed = True + inFlight _ = False + -- Check set the internal state to RequestFinished while + -- noting if there's possibly a request still in flight. + requestInFlight <- atomically $ + stateTVar reqVar $ inFlight &&& const RequestFinished + when requestInFlight $ withMVar connSocket $ \sock -> + SB.sendAll sock $ encode $ RawClean reqId + -- We use a shared resource map here to ensure that we only hold + -- onto the share connection state TVar for the duration of making a + -- fetch request. If anything goes wrong in the middle, we're + -- certain it'll get removed. + withSharedResource connStates reqId + (newTVarIO WaitingInstant) + cleanupRequest $ \reqVar -> do + -- Send the user request. + withMVar connSocket $ \sock -> SB.sendAll sock $ encode + $ RawRequest reqId req + + let + requestFinished = ClientError "request already finished" + + getDelayed = readTVar reqVar >>= \case + RequestFinished -> throwSTM requestFinished + WaitingDelayed -> retry + Available xs -> return xs + Errored e -> throwSTM e + WaitingInstant -> throwSTM $ ClientError $ + "fetch/WaitingDelayed: unexpected state WaitingInstant" + + -- Run the user's continuation. 'withSharedResource' takes care of + -- any clean-up necessary. + cont $ readTVar reqVar >>= \case + RequestFinished -> throwSTM requestFinished + Errored e -> throwSTM e + WaitingInstant -> retry -- wait for an instant response + Available xs -> pure $ Left xs + WaitingDelayed -> pure $ Right getDelayed + -- | 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) +fetchSimple conn timeout req = fetch conn req (fmap (maybe mempty id) . atomicallyWithin timeout . awaitResponse) atomicallyWithin :: Int -- ^ timeout in microseconds -> STM a diff --git a/src/Database/Franz/Protocol.hs b/src/Database/Franz/Protocol.hs new file mode 100644 --- /dev/null +++ b/src/Database/Franz/Protocol.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE DeriveGeneric #-} +module Database.Franz.Protocol + ( apiVersion + , defaultPort + , IndexName + , FranzException(..) + , RequestType(..) + , ItemRef(..) + , Query(..) + , RawRequest(..) + , ResponseId + , ResponseHeader(..) + , PayloadHeader(..)) where + +import Control.Exception (Exception) +import qualified Data.ByteString.Char8 as B +import Data.Serialize hiding (getInt64le) +import Database.Franz.Internal (getInt64le) +import Network.Socket (PortNumber) +import GHC.Generics (Generic) + +apiVersion :: B.ByteString +apiVersion = B.pack "0" + +defaultPort :: PortNumber +defaultPort = 1886 + +type IndexName = B.ByteString + +data FranzException = MalformedRequest !String + | StreamNotFound !FilePath + | IndexNotFound !IndexName ![IndexName] + | InternalError !String + | ClientError !String + deriving (Show, Generic) +instance Serialize FranzException +instance Exception FranzException + +data RequestType = AllItems | LastItem deriving (Show, Generic) +instance Serialize RequestType + +data ItemRef = BySeqNum !Int -- ^ sequential number + | ByIndex !IndexName !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 RawRequest + = RawRequest !ResponseId !Query + | RawClean !ResponseId deriving Generic +instance Serialize RawRequest + +type ResponseId = Int + +data ResponseHeader = Response !ResponseId + -- ^ response ID, number of streams; there are items satisfying the query + | ResponseWait !ResponseId -- ^ response ID; requested elements are not available right now + | 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 <$> getInt64le <*> getInt64le <*> getInt64le <*> get diff --git a/src/Database/Franz/Reader.hs b/src/Database/Franz/Reader.hs --- a/src/Database/Franz/Reader.hs +++ b/src/Database/Franz/Reader.hs @@ -1,13 +1,17 @@ {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE BangPatterns #-} module Database.Franz.Reader where import Control.Concurrent import Control.Concurrent.STM import Control.Exception import Control.Monad +import Control.Monad.State.Strict +import Control.Monad.Trans.Maybe import Data.Serialize +import Database.Franz.Protocol import qualified Data.ByteString.Char8 as B import qualified Data.HashMap.Strict as HM import qualified Data.IntMap.Strict as IM @@ -16,28 +20,11 @@ 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] @@ -52,9 +39,10 @@ type Activity = Either Double Int -addActivity :: Activity -> Activity -addActivity (Left _) = Right 0 -addActivity (Right n) = Right (n + 1) +addActivity :: Stream -> STM () +addActivity str = modifyTVar' (vActivity str) $ \case + Left _ -> Right 0 + Right n -> Right (n + 1) removeActivity :: Stream -> IO () removeActivity str = do @@ -101,6 +89,7 @@ indexHandle <- openFile offsetPath ReadMode let final :: Either SomeException Void -> IO () + final (Left exc) | Just ThreadKilled <- fromException exc = pure () final (Left exc) = logFollower [path, "terminated with", show exc] final (Right v) = absurd v @@ -124,8 +113,6 @@ let indices = HM.fromList $ zip indexNames vIndices - logFollower ["started", path] - vActivity <- getMonotonicTime >>= newTVarIO . Left return Stream{..} where @@ -140,7 +127,7 @@ -> IM.IntMap Int -- ^ offsets -> (Bool, QueryResult) range begin end rt allOffsets = case rt of - AllItems -> (ready, (first, maybe first fst $ IM.maxViewWithKey body)) + AllItems -> (ready, (firstItem, maybe firstItem 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 @@ -151,69 +138,111 @@ 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 + firstItem = 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 + , vStreams :: TVar (HM.HashMap FilePath (HM.HashMap B.ByteString Stream)) } +data ReaperState = ReaperState + { -- | How many streams we pruned. + prunedStreams :: !Int + -- | How many streams we saw in total. + , totalStreams :: !Int + } + 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 + -- Check if stream's activity indicates that we should prune it. + let shouldPrune (Left t) = now - t >= life + shouldPrune _ = False + + -- Try prunning stream at given filepath. Checks if stream + -- should really be pruned first. + tryPrune :: FilePath -> B.ByteString -> STM (Maybe Stream) + tryPrune mPath sPath = runMaybeT $ do + currentAllStreams <- lift $ readTVar vStreams + currentStreams <- MaybeT . pure $ HM.lookup mPath currentAllStreams + currentStream <- MaybeT . pure $ HM.lookup sPath currentStreams + currentAct <- lift $ readTVar (vActivity currentStream) + guard $ shouldPrune currentAct + let newStreams = HM.delete sPath currentStreams + lift . writeTVar vStreams $ if HM.null newStreams + -- Stream we're deleting was the + -- last one around. + then HM.delete mPath currentAllStreams + -- Still have some other streams left for this mount path. + -- Keep those only. + else HM.insert mPath newStreams currentAllStreams + pure currentStream + + -- Take a snapshots of all streams. + allStreams <- readTVarIO vStreams + -- Traverse the snapshot, looking for streams that currently seem + -- out of date. If we find an out-of-date stream, take outer lock + -- too, check again, delete it if necessary. Close stream promptly + -- after deletion. + -- + -- We could first traverse the whole snapshot, gather potential + -- streams for deletion, take lock and delete all these streams from + -- the map. However, on an assumption that we normally reaps streams + -- and much lower rate than we use them and in favour of locking at + -- as small time intervals as possible, we simply traverse the + -- snapshot and if we find an out of date stream in the snapshot, we + -- take the lock on the whole stream data structure and delete the + -- stream. We also check that we aren't leaving an empty map entry + -- behind and if we are, we delete it straight away. This stops us + -- from having to traverse the whole map later to clean things up as + -- well as ensuring that we never leave empty values in the map for + -- others to see. + -- + -- While doing all this, keep track of how many streams we saw and + -- how many we have closed which is used later to report statistics + -- to the user: this is in contrast to doing linear-time traversals + -- just to get counts of things we've already traversed. + stats <- flip execStateT (ReaperState 0 0) $ + forM_ (HM.toList allStreams) $ \(mPath, streams) -> do + forM_ (HM.toList streams) $ \(sPath, stream) -> do + modify' $ \s -> s { totalStreams = totalStreams s + 1 } + snapAct <- lift $ readTVarIO (vActivity stream) + when (shouldPrune snapAct) $ do + -- The stream indicates that it's old and should be pruned. Take + -- a lock on the outer map, check again inside a transaction and + -- amend the maps as necessary. + deletedStream'm <- lift . atomically $ tryPrune mPath sPath + forM_ deletedStream'm $ \prunedStream -> do + lift $ closeStream prunedStream + modify' $ \s -> s { prunedStreams = prunedStreams s + 1 } + + when (prunedStreams stats > 0) $ hPutStrLn stderr $ unwords [ "[reaper] closed" - , show (length xs) + , show (prunedStreams stats) , "out of" - , show count + , show (totalStreams stats) ] + threadDelay $ floor $ int * 1e6 -withFranzReader :: FilePath -> (FranzReader -> IO ()) -> IO () -withFranzReader prefix k = do +withFranzReader :: (FranzReader -> IO ()) -> IO () +withFranzReader k = do vStreams <- newTVarIO HM.empty withManager $ \watchManager -> k FranzReader{..} -handleQuery :: FranzReader - -> FilePath +handleQuery :: FilePath -- ^ prefix + -> FranzReader + -> FilePath -- ^ directory -> 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 + -> (Stream -> STM (Bool, QueryResult) -> IO r) -> IO r +handleQuery prefix FranzReader{..} dir (Query name begin_ end_ rt) cont + = bracket acquire removeActivity + $ \stream@Stream{..} -> cont stream $ do readTVar vCaughtUp >>= check allOffsets <- readTVar vOffsets let finalOffset = case IM.maxViewWithKey allOffsets of @@ -240,3 +269,18 @@ let body' = maybe id (IM.insert val) lastItem body return $! maybe minBound fst $ IM.maxView body' return $! range begin end rt allOffsets + where + acquire = join $ atomically $ do + allStreams <- readTVar vStreams + let !path = prefix dir B.unpack name + let !streams = maybe mempty id $ HM.lookup dir allStreams + case HM.lookup name streams of + Nothing -> pure $ bracketOnError + (createStream watchManager path) + closeStream $ \s -> atomically $ do + addActivity s + modifyTVar' vStreams $ HM.insert dir $ HM.insert name s streams + pure s + Just s -> do + addActivity s + pure (pure s) diff --git a/src/Database/Franz/Server.hs b/src/Database/Franz/Server.hs new file mode 100644 --- /dev/null +++ b/src/Database/Franz/Server.hs @@ -0,0 +1,218 @@ +{-# LANGUAGE DeriveGeneric, LambdaCase, OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE TypeFamilies #-} +module Database.Franz.Server + ( Settings(..) + , startServer + , defaultPort + ) where + +import Control.Concurrent +import Control.Exception +import Control.Monad +import Control.Retry +import Control.Concurrent.STM +import Database.Franz.Internal +import Database.Franz.Protocol +import Database.Franz.Reader +import Data.ConcurrentResourceMap +import Data.Serialize +import qualified Data.IntMap.Strict as IM +import Data.IORef +import qualified Data.ByteString.Char8 as B +import qualified Data.HashMap.Strict as HM +import Data.Tuple (swap) +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 (ProcessHandle, spawnProcess, cleanupProcess, + waitForProcess, getProcessExitCode, getPid) + +respond :: FilePath + -> FranzReader + -> IORef (IM.IntMap ThreadId) -- ^ thread pool + -> FilePath -- ^ path + -> IORef B.ByteString -- ^ buffer + -> MVar S.Socket + -> IO () +respond prefix env refThreads path buf vConn = do + recvConn <- readMVar vConn + runGetRecv buf recvConn get >>= \case + Right (RawRequest reqId req) -> do + let pop result = do + case result of + Left ex | Just e <- fromException ex -> sendHeader $ ResponseError reqId e + _ -> pure () + `finally` popThread reqId + handleQuery prefix env path req $ \stream query -> do + atomically (fmap Right query `catchSTM` (pure . Left)) >>= \case + Left e -> sendHeader $ ResponseError reqId e + Right (ready, offsets) + | ready -> send (Response reqId) stream offsets + | otherwise -> do + tid <- flip forkFinally pop $ bracket_ + (atomically $ addActivity stream) + (removeActivity stream) $ do + sendHeader $ ResponseWait reqId + offsets' <- atomically $ do + (ready', offsets') <- query + check ready' + pure offsets' + send (Response reqId) stream offsets' + -- Store the thread ID of the thread yielding a future + -- response such that we can kill it mid-way if user + -- sends a cancel request or we're killed with an + -- exception. + atomicModifyIORef' refThreads $ \m -> (IM.insert reqId tid m, ()) + -- Response is not ready but the user indicated that they + -- are not interested in waiting either. While we have no + -- work left to do, we do want to send a message back + -- saying the response would be a delayed one so that the + -- user can give up waiting. + | otherwise -> sendHeader $ ResponseWait reqId + `catch` \e -> sendHeader $ ResponseError reqId e + Right (RawClean reqId) -> do + tid <- popThread reqId + mapM_ killThread tid + Left err -> throwIO $ MalformedRequest err + where + popThread reqId = atomicModifyIORef' refThreads + $ swap . IM.updateLookupWithKey (\_ _ -> Nothing) reqId + + 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) + -- byte offset + number of indices + let siz = 8 * (length indexNames + 1) + -- Send byte offsets and indices + SF.sendFile' conn indexHandle (fromIntegral $ siz * succ s0) (fromIntegral $ siz * (s1 - s0)) + -- Send payloads + SF.sendFile' conn payloadHandle (fromIntegral p0) (fromIntegral $ p1 - p0) + +data Settings = Settings + { reapInterval :: Double + , streamLifetime :: Double + , port :: S.PortNumber + , livePrefix :: FilePath + , archivePrefix :: Maybe FilePath + , mountPrefix :: FilePath + } + +newtype MountMap v = MountMap (HM.HashMap FilePath v) + +instance ResourceMap MountMap where + type Key MountMap = FilePath + empty = MountMap mempty + delete k (MountMap m) = MountMap (HM.delete k m) + insert k v (MountMap m) = MountMap (HM.insert k v m) + lookup k (MountMap m) = HM.lookup k m + +newMountMap :: IO (ConcurrentResourceMap MountMap ProcessHandle) +newMountMap = newResourceMap + +startServer + :: Settings + -> IO () +startServer Settings{..} = withFranzReader $ \franzReader -> do + + forM_ archivePrefix $ \path -> do + e <- doesDirectoryExist path + unless e $ error $ "archive prefix " ++ path ++ " doesn't exist" + + hSetBuffering stderr LineBuffering + _ <- forkIO $ reaper reapInterval streamLifetime franzReader + + vMounts <- newMountMap + + 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.addrAddress addr + S.listen sock S.maxListenQueue + logServer ["Listening on", show port] + + forever $ do + (conn, connAddr) <- S.accept sock + buf <- newIORef B.empty + let respondLoop prefix path = do + SB.sendAll conn apiVersion + logServer [show connAddr, show path] + ref <- newIORef IM.empty + vConn <- newMVar conn + forever (respond prefix franzReader ref path buf vConn) `finally` do + readIORef ref >>= mapM_ killThread + + forkFinally (runGetRecv buf conn get >>= \case + Left _ -> throwIO $ MalformedRequest "Expecting a path" + Right pathBS -> do + let path = B.unpack pathBS + case archivePrefix of + -- just start a session without thinking about archives + Nothing -> respondLoop livePrefix path + -- Mount a squashfs image and increment the counter + Just prefix | src <- prefix path -> do + -- ^ check if an archive exists + exist <- doesFileExist src + if exist + then withSharedResource vMounts path + (mountFuse src (mountPrefix path)) + (\fuse -> do + -- close the last client's streams + streams <- atomically $ do + streams <- readTVar $ vStreams franzReader + writeTVar (vStreams franzReader) $ HM.delete path streams + pure streams + forM_ (HM.lookup path streams) $ mapM_ closeStream + killFuse fuse (mountPrefix path) `finally` do + pid <- getPid fuse + forM_ pid $ \p -> logServer ["Undead squashfuse detected:", show p]) + (const $ respondLoop mountPrefix path) + else do + logServer ["Archive", src, "doesn't exist; falling back to live streams"] + respondLoop livePrefix 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 + logServer [show connAddr, "disconnected"] + +logServer :: [String] -> IO () +logServer = hPutStrLn stderr . unwords . (:) "[server]" + +mountFuse :: FilePath -> FilePath -> IO ProcessHandle +mountFuse src dest = do + createDirectoryIfMissing True dest + logServer ["squashfuse", "-f", src, dest] + bracketOnError (spawnProcess "squashfuse" ["-f", src, dest]) (flip killFuse dest) $ \fuse -> do + -- It keeps process handles so that mounted directories are cleaned up + -- but there's no easy way to tell when squashfuse finished mounting. + -- Wait until the destination becomes non-empty. + notMounted <- retrying (limitRetries 5 <> exponentialBackoff 100000) (const pure) + $ \status -> getProcessExitCode fuse >>= \case + Nothing -> do + logServer ["Waiting for squashfuse to mount", src, ":", show status] + null <$> listDirectory dest + Just e -> do + removeDirectory dest + throwIO $ InternalError $ "squashfuse exited with " <> show e + when notMounted $ throwIO $ InternalError $ "Failed to mount " <> src + return fuse + +killFuse :: ProcessHandle -> FilePath -> IO () +killFuse fuse path = do + cleanupProcess (Nothing, Nothing, Nothing, fuse) + e <- waitForProcess fuse + logServer ["squashfuse:", show e] + removeDirectory path