diff --git a/fluent-logger.cabal b/fluent-logger.cabal --- a/fluent-logger.cabal +++ b/fluent-logger.cabal @@ -1,5 +1,5 @@ name: fluent-logger -version: 0.2.1.0 +version: 0.2.2.0 synopsis: A structured logger for Fluentd (Haskell) description: A structured logger for Fluentd (Haskell) license: Apache-2.0 diff --git a/src/Network/Fluent/Logger.hs b/src/Network/Fluent/Logger.hs --- a/src/Network/Fluent/Logger.hs +++ b/src/Network/Fluent/Logger.hs @@ -34,14 +34,15 @@ import Data.Monoid ( mconcat ) import qualified Network.Socket as NS import Network.Socket.Options ( setRecvTimeout, setSendTimeout ) -import Network.Socket.ByteString.Lazy ( sendAll ) +import Network.Socket.ByteString.Lazy ( sendAll, recv ) import Control.Monad ( void, forever, when ) import Control.Applicative ( (<$>) ) import Control.Concurrent ( ThreadId, forkIO, killThread, threadDelay ) -import Control.Concurrent.STM ( atomically +import Control.Concurrent.STM ( atomically, orElse , TChan, newTChanIO, readTChan, peekTChan, writeTChan - , TVar, newTVarIO, readTVar, modifyTVar ) -import Control.Exception ( SomeException, handle, bracket, throwIO ) + , TVar, newTVarIO, readTVar, modifyTVar + , TMVar, tryPutTMVar, newEmptyTMVarIO, takeTMVar ) +import Control.Exception ( SomeException, AsyncException, Handler(Handler), handle, bracket, throwIO, catches, onException ) import Data.MessagePack import Data.Serialize hiding (label) import Data.Int ( Int64 ) @@ -117,8 +118,28 @@ NS.connect sock $ NS.addrAddress addr return sock +type CloseFlag = TMVar () + +setFlagWhenClose :: NS.Socket -> CloseFlag -> IO () +setFlagWhenClose sock flag = do + received <- recv sock 256 `onException` setFlag + if LBS.length received == 0 + then setFlag + else setFlagWhenClose sock flag + where + setFlag = atomically $ void $ tryPutTMVar flag () + runSender :: FluentLoggerSender -> IO () -runSender logger = forever $ bracket (connectFluent logger) close (sendFluent logger) +runSender logger = forever $ filterException $ bracket (connectFluent logger) close handleSocket where + passAsyncException :: AsyncException -> IO a + passAsyncException e = throwIO e + dropOtherExceptions :: SomeException -> IO () + dropOtherExceptions _ = return () + filterException :: IO () -> IO () + filterException action = catches action [Handler passAsyncException, Handler dropOtherExceptions] + handleSocket sock = do + flag <- newEmptyTMVarIO + bracket (forkIO $ setFlagWhenClose sock flag) killThread (const $ sendFluent logger sock flag) connectFluent :: FluentLoggerSender -> IO NS.Socket connectFluent logger = exponentialBackoff $ getSocket host port timeout where @@ -137,17 +158,24 @@ threadDelay delay handle (retry $ min 60000000 $ interval * 3 `div` 2) action -sendFluent :: FluentLoggerSender -> NS.Socket -> IO () -sendFluent logger sock = toSender where + +data SenderEvent = SenderNewData LBS.ByteString + | SenderSocketClose + +sendFluent :: FluentLoggerSender -> NS.Socket -> CloseFlag -> IO () +sendFluent logger sock flag = toSender where chan = fluentLoggerSenderChan logger buffered = fluentLoggerSenderBuffered logger toSender = do - entry <- atomically $ peekTChan chan - sendAll sock entry - atomically $ do - void $ readTChan chan - modifyTVar buffered (subtract $ LBS.length entry) - sendFluent logger sock + event <- atomically $ (takeTMVar flag >> return SenderSocketClose) `orElse` (fmap SenderNewData $ peekTChan chan) + case event of + SenderSocketClose -> return () + SenderNewData entry -> do + sendAll sock entry + atomically $ do + void $ readTChan chan + modifyTVar buffered (subtract $ LBS.length entry) + sendFluent logger sock flag -- | Create a fluent logger -- diff --git a/test/MockServer.hs b/test/MockServer.hs --- a/test/MockServer.hs +++ b/test/MockServer.hs @@ -20,13 +20,14 @@ import Data.Serialize (Serialize, get) import Data.MessagePack import Data.Monoid +import qualified Data.Set as S import Network.Fluent.Logger.Packable (pack) import Network.Fluent.Logger.Unpackable (Unpackable, unpack) data MockServer a = MockServer { mockServerChan :: Chan a , mockServerThread :: ThreadId - , mockServerConnectionCount :: MVar Int + , mockServerConnectionThreads :: MVar (S.Set ThreadId) } mockServerHost :: ByteString @@ -38,17 +39,19 @@ mockServerSettings :: ServerSettings mockServerSettings = serverSettings mockServerPort "*" -app :: (MonadIO m, MonadThrow m, Serialize a) => Chan a -> MVar Int -> AppData -> m () +app :: (MonadIO m, MonadThrow m, Serialize a) => Chan a -> MVar (S.Set ThreadId) -> AppData -> m () app chan mvconn ad = do - liftIO $ modifyMVar_ mvconn (return . (+ 1)) - appSource ad $= conduitGet get $$ sinkChan chan mvconn + myid <- liftIO myThreadId + liftIO $ modifyMVar_ mvconn (return . S.insert myid) + appSource ad $= conduitGet get $$ (addCleanup (unregister myid) $ sinkChan chan) where + unregister myid _ = liftIO $ modifyMVar_ mvconn (return . S.delete myid) -sinkChan :: (MonadIO m, Serialize a) => Chan a -> MVar Int -> Sink a m () -sinkChan chan mvconn = do +sinkChan :: (MonadIO m, Serialize a) => Chan a -> Sink a m () +sinkChan chan = do mx <- await case mx of - Nothing -> liftIO $ modifyMVar_ mvconn (return . (subtract 1)) - Just x -> liftIO (writeChan chan x) >> sinkChan chan mvconn + Nothing -> return () + Just x -> liftIO (writeChan chan x) >> sinkChan chan withMockServer :: Serialize a => (MockServer a -> IO ()) -> IO () withMockServer = bracket runMockServer stopMockServer @@ -59,18 +62,19 @@ runMockServer :: Serialize a => IO (MockServer a) runMockServer = do chan <- newChan - mvconn <- newMVar 0 + mvconn <- newMVar S.empty tid <- forkIO $ runTCPServer mockServerSettings $ app chan mvconn threadDelay 10000 return MockServer { mockServerChan = chan , mockServerThread = tid - , mockServerConnectionCount = mvconn + , mockServerConnectionThreads = mvconn } stopMockServer :: Serialize a => MockServer a -> IO () stopMockServer server = do killThread $ mockServerThread server + (mapM_ killThread . S.toList) =<< (swapMVar (mockServerConnectionThreads server) S.empty) threadDelay 10000 getConnectionCount :: MockServer a -> IO Int -getConnectionCount = readMVar . mockServerConnectionCount +getConnectionCount s = fmap S.size $ readMVar $ mockServerConnectionThreads s diff --git a/test/Network/Fluent/LoggerSpec.hs b/test/Network/Fluent/LoggerSpec.hs --- a/test/Network/Fluent/LoggerSpec.hs +++ b/test/Network/Fluent/LoggerSpec.hs @@ -18,7 +18,7 @@ it "posts a message" postPostsMessage it "keeps message order" postKeepsMessageOrder it "buffers message if server down" postBuffersMessageIfServerDown - -- it "buffers message if lost connection" postBuffersMessageIfLostConnection + it "buffers message if lost connection" postBuffersMessageIfLostConnection -- it "losts message if buffer is over" postLostsMessageIfBufferIsOver describe "postWithTime" $ do it "posts a message with given time" postWithTimePostsMessageWithGivenTime @@ -41,6 +41,9 @@ getCurrentEpochTime :: IO Int getCurrentEpochTime = round <$> getPOSIXTime +shouldHaveConns :: MockServer a -> Int -> IO () +shouldHaveConns server exp = getConnectionCount server >>= (`shouldBe` exp) + postPostsMessage :: IO () postPostsMessage = withFluentLogger postSettings $ \logger -> do @@ -84,13 +87,16 @@ post logger label ( 1 :: Int ) (_, _, content) <- recvMockServer server :: IO (ByteString, Int, Int) content `shouldBe` 1 + server `shouldHaveConns` 1 post logger label ( 2 :: Int ) withMockServer $ \server -> do post logger label ( 3 :: Int ) (_, _, content) <- recvMockServer server :: IO (ByteString, Int, Int) content `shouldBe` 2 + server `shouldHaveConns` 1 (_, _, content) <- recvMockServer server :: IO (ByteString, Int, Int) content `shouldBe` 3 + server `shouldHaveConns` 1 postLostsMessageIfBufferIsOver :: IO () postLostsMessageIfBufferIsOver = @@ -125,5 +131,4 @@ content `shouldBe` "foobar" threadDelay 5000 server `shouldHaveConns` 0 - where - shouldHaveConns server exp = getConnectionCount server >>= (`shouldBe` exp) +