diff --git a/Data/Conduit/Branching.hs b/Data/Conduit/Branching.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Branching.hs @@ -0,0 +1,46 @@ +-- | Branching conduits ... +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 + +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 () + +branchConduits :: Source (ResourceT IO) a -- ^ The source to branch from + -> (a -> [Int]) -- ^ The branching function (0 is the first sink) + -> [Sink a (ResourceT IO) ()] -- ^ The destination sinks + -> IO () -- ^ Results of the sinks +branchConduits src brfunc sinks = do + (newsink, sources) <- mkBranchingConduit (length sinks) brfunc + let srcconduit = src $$ newsink + dstconduits = map (uncurry ($$)) (zip sources sinks) + actions = map runResourceT (srcconduit : dstconduits) + parallel_ actions diff --git a/Data/Conduit/ElasticSearch.hs b/Data/Conduit/ElasticSearch.hs --- a/Data/Conduit/ElasticSearch.hs +++ b/Data/Conduit/ElasticSearch.hs @@ -1,7 +1,7 @@ {-| This module exports "Conduit" interfaces to ElasticSearch. It is totally experimental. -} -module Data.Conduit.ElasticSearch (esSink) where +module Data.Conduit.ElasticSearch (esConduit) where import Prelude hiding (catch) import Control.Exception @@ -10,48 +10,64 @@ import Network.HTTP.Conduit import Data.Aeson import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BSL +import qualified Data.ByteString.Lazy.Char8 as BSL import Data.Time import qualified Data.Text.Lazy.Encoding as E import Data.Text.Format (format,left) import Logstash.Message import Control.Monad.IO.Class import Control.Concurrent (threadDelay) +import Data.Either (partitionEithers) + import qualified Data.HashMap.Strict as HM +import qualified Data.Vector as V +import qualified Data.Text as T safeQuery :: Request (ResourceT IO) -> IO (Response BSL.ByteString) safeQuery req = catch (withManager $ httpLbs req) (\e -> print (e :: SomeException) >> threadDelay 500000 >> safeQuery req) --- | Takes JSONifiable values, and returns the result of the ES request +-- | Takes a "LogstashMessage", and returns the result of the ES request -- along with the value in case of errors, or ES's values in case of -- success -esSink :: (MonadResource m) => Maybe (Request m) -- ^ Defaults parameters for the http request to ElasticSearch. Use "Nothing" for defaults. +esConduit :: (MonadResource m) => Maybe (Request m) -- ^ Defaults parameters for the http request to ElasticSearch. Use "Nothing" for defaults. -> BS.ByteString -- ^ Hostname of the ElasticSearch server -> Int -- ^ Port of the HTTP interface (usually 9200) - -> Conduit LogstashMessage m (Either (LogstashMessage, Value) Value) -esSink r h p = CL.mapM doIndexA + -> Conduit [LogstashMessage] m [Either (LogstashMessage, Value) Value] +esConduit r h p = CL.map (map prepareBS) =$= CL.mapM sendBulk where defR1 = case r of Just x -> x Nothing -> def defR2 = defR1 { host = h , port = p + , path = "/_bulk" , method = "POST" , checkStatus = (\_ _ -> Nothing) } - doIndexA :: (MonadResource m) => LogstashMessage -> m (Either (LogstashMessage, Value) Value) - doIndexA input = + prepareBS :: LogstashMessage -> Either (LogstashMessage, Value) (LogstashMessage, Value) + prepareBS input = case logstashTime input of - Nothing -> return $! Left (input, object [ "error" .= String "Time was not supplied" ]) - Just (UTCTime day _) -> do + Nothing -> Left (input, object [ "error" .= String "Time was not supplied" ]) + Just (UTCTime day _) -> let (y,m,d) = toGregorian day - req = defR2 { path = BSL.toStrict (E.encodeUtf8 (format "/logstash-{}.{}.{}/{}/" (y, left 2 '0' m, left 2 '0' d, logstashType input))) - , requestBody = RequestBodyLBS (encode input) - } - res <- liftIO $ safeQuery req - case decode (responseBody res) of - Just (Object hh) -> case HM.lookup "ok" hh of - Just (Bool True) -> return $! Right (Object hh) - _ -> return $! Left (input, Object hh) - Just j -> return $! Left (input, j) - Nothing -> return $! Left (input, object [ "error" .= String "Could not decode", "content" .= responseBody res ]) + 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] + sendBulk input = + let (errors, tosend) = partitionEithers input + lerrors = map Left errors + body = BSL.intercalate "\n{}\n" (map (encode . snd) tosend) `BSL.append` "\n{}\n" + 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 + 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 + _ -> genericError "Can't find items" diff --git a/Data/Conduit/Logstash.hs b/Data/Conduit/Logstash.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Logstash.hs @@ -0,0 +1,26 @@ +module Data.Conduit.Logstash (logstashListener) where + +import Data.Conduit +import qualified Data.Conduit.Binary as CB +import qualified Data.Conduit.List as CL +import Data.Conduit.Network +import Data.Aeson +import qualified Data.ByteString.Lazy as BSL +import qualified Data.ByteString as BS +import Codec.Text.IConv +import Data.Text.Encoding +import Logstash.Message + +tryDecode :: (FromJSON a) => BS.ByteString -> Either BS.ByteString a +tryDecode i = + let latin1 = convert "LATIN1" "UTF-8" li + li = BSL.fromStrict i + o = case decodeUtf8' i of + Left _ -> decode latin1 + Right _ -> decode li + in case o of + Just x -> Right x + Nothing -> Left i + +logstashListener :: Int -> Sink (Either BS.ByteString LogstashMessage) (ResourceT IO) () -> IO () +logstashListener port sink = runResourceT $ runTCPServer (serverSettings port HostAny) (\app -> appSource app $= CB.lines $= CL.map tryDecode $$ sink) diff --git a/Data/Conduit/Network/Retry.hs b/Data/Conduit/Network/Retry.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Network/Retry.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +module Data.Conduit.Network.Retry where + +import Prelude hiding (catch) +import Data.Conduit +import Data.Conduit.Network +import Network.Socket (Socket, close) +import Network.Socket.ByteString (sendAll) +import Control.Monad.IO.Class (MonadIO (liftIO)) +import Control.Exception +import Data.ByteString (ByteString) +import Control.Concurrent (threadDelay) +import Control.Concurrent.MVar +import Control.Monad ((>=>)) +import Control.Monad.Trans.Class (lift) + +{-| Tentative /safe/ "Sink" for a "Socket". It should try reopening the "Socket" +every time the call to 'sendAll' fails. This means that some bytes might be sent +multiple times, if the socket fails in the middle of the sendAll call. This is +targeted at protocols where only a full message makes sense. + +This is used to send a full JSON object to Logstash. +-} +sinkSocketRetry :: MonadResource m => IO Socket -> Int -> IO () -> GInfSink ByteString m +sinkSocketRetry mkSocket delay exeptionCallback = + let + safeMkSocket :: IO Socket + safeMkSocket = catch mkSocket (\SomeException{} -> exeptionCallback >> threadDelay delay >> safeMkSocket) + safeSend :: MVar Socket -> ByteString -> IO () + safeSend s o = do + sock <- takeMVar s + catch (sendAll sock o >> putMVar s sock) $ \SomeException{} -> do + close sock + safeMkSocket >>= putMVar s + threadDelay delay + safeSend s o + push :: MonadResource m => MVar Socket -> GInfSink ByteString m + push s = awaitE >>= either return (\bs -> lift (liftIO $ safeSend s bs) >> push s) + in bracketP (safeMkSocket >>= newMVar) (takeMVar >=> close) push + +-- | A specialization of the previous Sink that opens a TCP connection. +tcpSinkRetry :: MonadResource m => ByteString -> Int -> Int -> IO () -> GInfSink ByteString m +tcpSinkRetry host port = sinkSocketRetry (fmap fst (getSocket host port)) + diff --git a/Data/Conduit/Redis.hs b/Data/Conduit/Redis.hs --- a/Data/Conduit/Redis.hs +++ b/Data/Conduit/Redis.hs @@ -1,30 +1,66 @@ {-| Quick conduit for reading from Redis lists. Not tested much, and probably quite slow. -} -module Data.Conduit.Redis where +module Data.Conduit.Redis (redisSource, redisSink) where import Data.Conduit -import Data.Conduit.Util +import qualified Data.Conduit.List as CL import qualified Data.ByteString.Char8 as BS import Network +import Control.Concurrent.MVar import Database.Redis hiding (String, decode) -import Control.Monad (void) -import Control.Monad.IO.Class (liftIO) +import Control.Monad (void,replicateM, forever) +import Control.Monad.IO.Class (liftIO, MonadIO) +import Data.Either (rights) +import Data.Maybe (catMaybes) +import Control.Exception +import Control.Concurrent hiding (yield) -redisSource :: (MonadResource m) => HostName -- ^ Hostname of the Redis server - -> Int -- ^ Port of the Redis server (usually 6379) - -> BS.ByteString -- ^ Name of the list - -> Source m BS.ByteString -redisSource h p list = +mp :: Either Reply (Maybe (BS.ByteString, BS.ByteString)) -> Maybe BS.ByteString +mp (Right (Just (_, x))) = Just x +mp _ = Nothing + +safePush :: BS.ByteString -> MVar Connection -> ConnectInfo -> BS.ByteString -> IO () +safePush list mconn cinfo input = catch mypush (\SomeException{} -> resetMVar) + where + resetMVar = do + void $ takeMVar mconn + connect cinfo >>= putMVar mconn + BS.putStrLn "reconnecting to redis server ..." + threadDelay 500000 + safePush list mconn cinfo input + mypush = do + conn <- readMVar mconn + x <- runRedis conn (lpush list [input]) + case x of + Left (SingleLine "OK") -> return () + 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) + +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 + -> Source m [BS.ByteString] +redisSource h p list nb = let cinfo = defaultConnectInfo { connectHost = h, connectPort = PortNumber $ fromIntegral p } - pull = do - o <- blpop [list] 0 - case o of - Right (Just (_,k)) -> return k - _ -> pull - in sourceStateIO (connect cinfo) - (\conn -> runRedis conn (void quit)) - (\conn -> do - o <- liftIO $ runRedis conn pull - return (StateOpen conn o) - ) + myPipe :: (MonadResource m) => Connection -> Source m [BS.ByteString] + myPipe conn = forever (liftIO (runRedis conn (popN list nb)) >>= yield) + in bracketP (connect cinfo) (\conn -> runRedis conn (void quit)) myPipe + +-- | Warning, this outputs strings when things go wrong! +redisSink :: (MonadResource m) + => HostName -- ^ Hostname of the Redis server + -> Int -- ^ Port of the Redis server (usually 6379) + -> BS.ByteString -- ^ Name of the list + -> Sink BS.ByteString m () +redisSink h p list = + let cinfo = defaultConnectInfo { connectHost = h, connectPort = PortNumber $ fromIntegral p } + in bracketP (connect cinfo >>= newMVar) (const $ return ()) (\mconn -> CL.mapM_ (liftIO . safePush list mconn cinfo)) diff --git a/Logstash/Message.hs b/Logstash/Message.hs --- a/Logstash/Message.hs +++ b/Logstash/Message.hs @@ -26,12 +26,12 @@ instance FromJSON LogstashMessage where parseJSON (Object v) = LogstashMessage - <$> v .: "@type" - <*> v .: "@source" - <*> v .: "@tags" - <*> v .: "@fields" - <*> v .: "@message" - <*> v .: "@timestamp" + <$> v .: "@type" + <*> v .: "@source" + <*> v .: "@tags" + <*> v .: "@fields" + <*> v .: "@message" + <*> v .:? "@timestamp" parseJSON _ = mzero {-| As the name implies, this creates a dummy Logstash message, only @@ -115,3 +115,11 @@ (Just (String t), Just (String s), Just tags) -> Just $ LogstashMessage t s tags mflds mmsg mts _ -> Nothing value2logstash _ = Nothing + +-- | Adds the current timestamp if it is not provided. +addLogstashTime :: LogstashMessage -> IO LogstashMessage +addLogstashTime msg = case logstashTime msg of + Just _ -> return msg + Nothing -> do + curtime <- getCurrentTime + return msg { logstashTime = Just curtime } 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.1 +version: 0.2.2 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 + exposed-modules: Logstash.Message, Logstash.IO, Data.Conduit.Redis, Data.Conduit.ElasticSearch, Data.Conduit.Logstash, Data.Conduit.Network.Retry, Data.Conduit.Branching 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 + 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