diff --git a/Data/Conduit/Branching.hs b/Data/Conduit/Branching.hs --- a/Data/Conduit/Branching.hs +++ b/Data/Conduit/Branching.hs @@ -1,38 +1,39 @@ --- | Branching conduits ... +{-| Branching conduits + +WARNING: executables using this function must be compiled with -threaded +-} module Data.Conduit.Branching (mkBranchingConduit, branchConduits) where import Data.Conduit import qualified Data.Conduit.List as CL import Control.Monad import Control.Monad.IO.Class (liftIO) -import Control.Concurrent (MVar, putMVar, takeMVar, newEmptyMVar) import qualified Data.IntMap as IntMap import Data.Maybe (mapMaybe) import Control.Concurrent.ParallelIO +import Data.Conduit.TMChan +import GHC.Conc (atomically) +mkSink :: (MonadResource m) + => (a -> [Int]) + -> [TBMChan a] + -> Sink a m () +mkSink brfunc chans = + let cmap = IntMap.fromList $ zip [0..] chans -- creates an intmap with all output channels + querymap x = IntMap.lookup x cmap -- query the intmap to get a Maybe (TBMChan a) + cleanup = mapM_ (liftIO . atomically . closeTBMChan) chans -- this is the cleanup function that closes all chans + inject input = + let outchans = mapMaybe querymap (brfunc input) -- compiles the list of output channels + in mapM_ (\c -> liftIO $ atomically $ writeTBMChan c input) outchans -- and write into them + in addCleanup (const cleanup) (CL.mapM_ inject) + mkBranchingConduit :: (MonadResource m) => Int -- ^ Number of branches -> (a -> [Int]) -- ^ Branching function, where 0 is the first branch -> IO (Sink a m (), [Source m a]) -- ^ Returns a sink and N sources mkBranchingConduit nbbranches brfunction = do - mvars <- replicateM nbbranches newEmptyMVar - return (mvarSink brfunction mvars, map mvarSource mvars) - -mvarSink :: (MonadResource m) => (a -> [Int]) -> [MVar (Maybe a)] -> Sink a m () -mvarSink brfunc mvs = - let mvarmap = IntMap.fromList (zip [0..] mvs) - doBranch input = - let channels = brfunc input - mvars = mapMaybe (\x -> IntMap.lookup x mvarmap) channels - in mapM_ (\mv -> liftIO $ putMVar mv (Just input)) mvars - in bracketP (return ()) (const $ mapM_ (\mv -> putMVar mv Nothing) mvs) (const $ CL.mapM_ doBranch) - -mvarSource :: (MonadResource m) => MVar (Maybe a) -> Source m a -mvarSource mv = do - v <- liftIO $ takeMVar mv - case v of - Just x -> yield x >> mvarSource mv - Nothing -> return () + chans <- replicateM nbbranches (newTBMChanIO 16) + return (mkSink brfunction chans, map sourceTBMChan chans) branchConduits :: Source (ResourceT IO) a -- ^ The source to branch from -> (a -> [Int]) -- ^ The branching function (0 is the first sink) diff --git a/Data/Conduit/ElasticSearch.hs b/Data/Conduit/ElasticSearch.hs --- a/Data/Conduit/ElasticSearch.hs +++ b/Data/Conduit/ElasticSearch.hs @@ -44,30 +44,30 @@ , method = "POST" , checkStatus = (\_ _ -> Nothing) } - prepareBS :: LogstashMessage -> Either (LogstashMessage, Value) (LogstashMessage, Value) + prepareBS :: LogstashMessage -> Either (LogstashMessage, Value) (LogstashMessage, Value, Value) prepareBS input = case logstashTime input of Nothing -> Left (input, object [ "error" .= String "Time was not supplied" ]) Just (UTCTime day _) -> let (y,m,d) = toGregorian day index = BSL.toStrict (E.encodeUtf8 (format "logstash-{}.{}.{}" (y, left 2 '0' m, left 2 '0' d))) - in Right (input, object [ "index" .= object [ "_index" .= index, "_type" .= logstashType input, "_source" .= toJSON input ] ]) - sendBulk :: (MonadResource m) => [Either (LogstashMessage, Value) (LogstashMessage, Value)] -> m [Either (LogstashMessage, Value) Value] + in Right (input, object [ "index" .= object [ "_index" .= index, "_type" .= logstashType input ] ], toJSON input) + sendBulk :: (MonadResource m) => [Either (LogstashMessage, Value) (LogstashMessage, Value, Value)] -> m [Either (LogstashMessage, Value) Value] sendBulk input = let (errors, tosend) = partitionEithers input lerrors = map Left errors - body = BSL.intercalate "\n{}\n" (map (encode . snd) tosend) `BSL.append` "\n{}\n" + body = BSL.unlines $ concatMap (\(_,x,y) -> [encode x, encode y]) tosend req = defR2 { requestBody = RequestBodyLBS body } in do - liftIO $ BSL.putStrLn body res <- fmap (responseBody) $ liftIO $ safeQuery req let genericError er = return (Left (emptyLSMessage "error", object [ "error" .= (T.pack er), "data" .= res ]) : lerrors) getObject st (Object hh) = HM.lookup st hh getObject _ _ = Nothing + fst3 (a,_,_) = a items = decode res >>= getObject "items" extractErrors x = if (getObject "create" (snd x) >>= getObject "ok") == Just (Bool True) then Right (snd x) else Left x case items of - Just (Array v) -> return $ map extractErrors (zip (map fst tosend) (V.toList v)) ++ lerrors + Just (Array v) -> return $ map extractErrors (zip (map fst3 tosend) (V.toList v)) ++ lerrors _ -> genericError "Can't find items" diff --git a/Data/Conduit/Misc.hs b/Data/Conduit/Misc.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Misc.hs @@ -0,0 +1,44 @@ +module Data.Conduit.Misc where + +import Data.Conduit +import Control.Monad + +-- | Converts a stream of lists into a stream of single elements. +concat :: (Monad m) => Conduit [a] m a +concat = awaitForever (mapM_ yield) + +-- | Converts a stream of [a] into a stream of (Flush a). This is done by +-- sending a Flush when the input is the empty list, or that we reached +-- a certain threshold +concatFlush :: (Monad m) => Integer -> Conduit [a] m (Flush a) +concatFlush mx = concatFlush' 0 + where + concatFlush' x = awaitE >>= either return (\input -> if null input + then yield Flush >> concatFlush' 0 + else foldM sendf x input >>= concatFlush' + ) + sendf curx ev = do + yield (Chunk ev) + if curx >= mx + then yield Flush >> return 0 + else return (curx+1) + +-- | Regroup a stream of (Flush a) into a stream of lists, using "Flush" as +-- the separator +groupFlush :: (Monad m) => Conduit (Flush a) m [a] +groupFlush = grouper [] + where + grouper lst = awaitE >>= either return (handle lst) + handle lst Flush = do + unless (null lst) (yield (reverse lst)) + grouper [] + handle lst (Chunk x) = grouper (x:lst) + +-- | Analogous to maybe, but for chunks +mchunk :: b -> (a -> b) -> Flush a -> b +mchunk n _ Flush = n +mchunk _ f (Chunk x) = f x + +-- | Like mapMaybe, but in a Flush. Will not touch the Flush values. +mapFlushMaybe :: (Monad m) => (a -> Maybe b) -> Conduit (Flush a) m (Flush b) +mapFlushMaybe f = awaitForever $ mchunk (yield Flush) (maybe (return ()) (yield . Chunk) . f) diff --git a/Data/Conduit/Redis.hs b/Data/Conduit/Redis.hs --- a/Data/Conduit/Redis.hs +++ b/Data/Conduit/Redis.hs @@ -15,9 +15,12 @@ import Control.Exception import Control.Concurrent hiding (yield) +import Debug.Trace + mp :: Either Reply (Maybe (BS.ByteString, BS.ByteString)) -> Maybe BS.ByteString mp (Right (Just (_, x))) = Just x -mp _ = Nothing +mp (Right Nothing) = Nothing +mp x = trace (show x) Nothing safePush :: BS.ByteString -> MVar Connection -> ConnectInfo -> BS.ByteString -> IO () safePush list mconn cinfo input = catch mypush (\SomeException{} -> resetMVar) @@ -36,22 +39,26 @@ Right _ -> return () err -> BS.putStrLn ("retrying ... " `BS.append` BS.pack (show err)) >> threadDelay 500000 >> safePush list mconn cinfo input -popN :: BS.ByteString -> Int -> Redis [BS.ByteString] -popN l n = do - f <- fmap mp (brpop [l] 0) - nx <- fmap rights $ replicateM (n-1) (rpop l) - return $ catMaybes (f:nx) +popN :: BS.ByteString -> Int -> Integer -> Redis [BS.ByteString] +popN l n to = do + f <- fmap mp (brpop [l] to) + if f == Nothing + then return [] -- short circuiting + else do + nx <- fmap rights $ replicateM (n-1) (rpop l) + return $ catMaybes (f:nx) redisSource :: (MonadResource m) => HostName -- ^ Hostname of the Redis server -> Int -- ^ Port of the Redis server (usually 6379) -> BS.ByteString -- ^ Name of the list -> Int -- ^ Number of elements to pop at once + -> Integer -- ^ Timeout of the brpop function in seconds, useful for getting Flush events into your conduit. Set to 0 for no timeout. -> Source m [BS.ByteString] -redisSource h p list nb = +redisSource h p list nb to = let cinfo = defaultConnectInfo { connectHost = h, connectPort = PortNumber $ fromIntegral p } myPipe :: (MonadResource m) => Connection -> Source m [BS.ByteString] - myPipe conn = forever (liftIO (runRedis conn (popN list nb)) >>= yield) + myPipe conn = forever (liftIO (runRedis conn (popN list nb to)) >>= yield) in bracketP (connect cinfo) (\conn -> runRedis conn (void quit)) myPipe -- | Warning, this outputs strings when things go wrong! diff --git a/hslogstash.cabal b/hslogstash.cabal --- a/hslogstash.cabal +++ b/hslogstash.cabal @@ -2,7 +2,7 @@ -- documentation, see http://haskell.org/cabal/users-guide/ name: hslogstash -version: 0.2.2 +version: 0.2.3 synopsis: A library to work with, or as, a logstash server description: This library contains a few modules that let you work with Logstash messages, read them from a Redis list, store them into Elasticsearch, and more. license: BSD3 @@ -19,8 +19,8 @@ location: git://github.com/bartavelle/hslogstash.git library - exposed-modules: Logstash.Message, Logstash.IO, Data.Conduit.Redis, Data.Conduit.ElasticSearch, Data.Conduit.Logstash, Data.Conduit.Network.Retry, Data.Conduit.Branching + exposed-modules: Logstash.Message, Logstash.IO, Data.Conduit.Redis, Data.Conduit.ElasticSearch, Data.Conduit.Logstash, Data.Conduit.Network.Retry, Data.Conduit.Branching, Data.Conduit.Misc extensions: OverloadedStrings, BangPatterns ghc-options: -Wall -- other-modules: - build-depends: base <5, aeson, network, bytestring, text, vector, unordered-containers, time, text-format, attoparsec, hedis, conduit, transformers, http-conduit, iconv, network-conduit, containers, parallel-io + build-depends: base <5, aeson, network, bytestring, text, vector, unordered-containers, time, text-format, attoparsec, hedis, conduit, transformers, http-conduit, iconv, network-conduit, containers, parallel-io, stm-conduit