diff --git a/lightstep-haskell.cabal b/lightstep-haskell.cabal --- a/lightstep-haskell.cabal +++ b/lightstep-haskell.cabal @@ -1,7 +1,7 @@ cabal-version: 2.4 name: lightstep-haskell -version: 0.1.7 +version: 0.2.0 synopsis: LightStep OpenTracing client library description: LightStep OpenTracing client library. Uses GRPC transport via proto-lens. category: Tools @@ -25,6 +25,7 @@ LambdaCase MultiParamTypeClasses MultiWayIf + NamedFieldPuns NumericUnderscores RecordWildCards ScopedTypeVariables diff --git a/src/LightStep/HighLevel/IO.hs b/src/LightStep/HighLevel/IO.hs --- a/src/LightStep/HighLevel/IO.hs +++ b/src/LightStep/HighLevel/IO.hs @@ -1,28 +1,30 @@ {-# LANGUAGE OverloadedStrings #-} -module LightStep.HighLevel.IO - ( module LightStep.HighLevel.IO - , module LightStep.LowLevel - ) where +module LightStep.HighLevel.IO + ( module LightStep.HighLevel.IO, + module LightStep.LowLevel, + ) +where + import Control.Concurrent import Control.Concurrent.Async import Control.Concurrent.STM -import Control.Lens import Control.Exception.Safe +import Control.Lens import Control.Monad.Except -import GHC.Conc +import qualified Data.HashMap.Strict as HM import Data.Maybe import Data.ProtoLens.Message (defMessage) +import qualified Data.Text as T +import GHC.Conc import LightStep.Internal.Debug import LightStep.LowLevel import Proto.Collector import Proto.Collector_Fields import System.IO.Unsafe -import qualified Data.HashMap.Strict as HM -import qualified Data.Text as T import System.Timeout -{-# noinline globalSharedMutableSpanStacks #-} +{-# NOINLINE globalSharedMutableSpanStacks #-} globalSharedMutableSpanStacks :: MVar (HM.HashMap ThreadId [Span]) globalSharedMutableSpanStacks = unsafePerformIO (newMVar mempty) @@ -30,7 +32,7 @@ withSpan opName action = bracket (pushSpan opName) - popSpan + popSpan (const (onException action (setTag "error" "true"))) pushSpan :: MonadIO m => T.Text -> m () @@ -42,82 +44,83 @@ [] -> do pure $! HM.insert tId [sp] stacks (psp : _) -> - let !sp' = sp - & references .~ [defMessage & relationship .~ Reference'CHILD_OF & spanContext .~ (psp ^. spanContext)] - & spanContext.traceId .~ (psp ^. spanContext.traceId) - in pure $! HM.update (Just . (sp' :)) tId stacks + let !sp' = + sp + & references .~ [defMessage & relationship .~ Reference'CHILD_OF & spanContext .~ (psp ^. spanContext)] + & spanContext . traceId .~ (psp ^. spanContext . traceId) + in pure $! HM.update (Just . (sp' :)) tId stacks popSpan :: MonadIO m => () -> m () popSpan () = liftIO $ do tId <- myThreadId - sp <- modifyMVar globalSharedMutableSpanStacks - (\stacks -> - let (sp : sps) = stacks HM.! tId - !stacks' = HM.insert tId sps stacks - in pure (stacks', sp)) + sp <- + modifyMVar + globalSharedMutableSpanStacks + ( \stacks -> + let (sp : sps) = stacks HM.! tId + !stacks' = HM.insert tId sps stacks + in pure (stacks', sp) + ) sp' <- finishSpan sp submitSpan sp' modifyCurrentSpan :: MonadIO m => (Span -> Span) -> m () modifyCurrentSpan f = liftIO $ do tId <- myThreadId - modifyMVar_ globalSharedMutableSpanStacks - (\stacks -> - let (sp : sps) = stacks HM.! tId - !stacks' = HM.insert tId (f sp : sps) stacks - in pure stacks') + modifyMVar_ + globalSharedMutableSpanStacks + ( \stacks -> + let (sp : sps) = stacks HM.! tId + !stacks' = HM.insert tId (f sp : sps) stacks + in pure stacks' + ) setTag :: MonadIO m => T.Text -> T.Text -> m () setTag k v = modifyCurrentSpan (tags %~ (<> [defMessage & key .~ k & stringValue .~ v])) - {-# NOINLINE globalSharedMutableSingletonState #-} globalSharedMutableSingletonState :: TBQueue Span globalSharedMutableSingletonState = unsafePerformIO $ newTBQueueIO 100 withSingletonLightStep :: LightStepConfig -> IO () -> IO () withSingletonLightStep cfg action = do - runExceptT (mkClient cfg) >>= \case - Left err -> do - d_ $ "Failed to start LightStep client: " <> show err - action - Right client -> do - d_ $ "Connected to LightStep " <> lsHostName cfg <> ":" <> show (lsPort cfg) - doneVar <- newEmptyMVar - let work = do - d_ "Getting more spans" - sps <- liftIO $ atomically $ do - x <- readTBQueue globalSharedMutableSingletonState - xs <- flushTBQueue globalSharedMutableSingletonState - pure (x : xs) - d_ $ "Got " <> show (length sps) <> " spans" - reportSpansRes <- tryAny (reportSpans client sps) - case reportSpansRes of - Right () -> - d_ $ "Reported " <> show (length sps) <> " spans" - Left err -> - d_ $ "Error while reporting spans: " <> show err - shutdown = do - d_ "Getting the last spans before shutdown" - sps <- liftIO $ atomically $ flushTBQueue globalSharedMutableSingletonState - when (not $ null sps) $ do - d_ $ "Got " <> show (length sps) <> " spans" - reportSpans client sps - d_ $ "Reported " <> show (length sps) <> " spans" - d_ "No more spans" - closeClient client - d_ "Client closed" - liftIO $ putMVar doneVar () - race_ action $ do - tid <- myThreadId - labelThread tid "LightStep reporter" - runExceptT . fix $ \loop -> do - work - loop - race_ - (runExceptT shutdown) - (waitUntilDone (lsGracefulShutdownTimeoutSeconds cfg) doneVar) + client <- mkClient cfg + d_ $ "Connected to LightStep " <> lsHostName cfg <> ":" <> show (lsPort cfg) + doneVar <- newEmptyMVar + let work = do + d_ "Getting more spans" + sps <- atomically $ do + x <- readTBQueue globalSharedMutableSingletonState + xs <- flushTBQueue globalSharedMutableSingletonState + pure (x : xs) + d_ $ "Got " <> show (length sps) <> " spans" + reportSpansRes <- tryAny (reportSpans client sps) + case reportSpansRes of + Right () -> + d_ $ "Reported " <> show (length sps) <> " spans" + Left err -> + d_ $ "Error while reporting spans: " <> show err + shutdown = do + d_ "Getting the last spans before shutdown" + sps <- atomically $ flushTBQueue globalSharedMutableSingletonState + when (not $ null sps) $ do + d_ $ "Got " <> show (length sps) <> " spans" + reportSpans client sps + d_ $ "Reported " <> show (length sps) <> " spans" + d_ "No more spans" + closeClient client + d_ "Client closed" + putMVar doneVar () + race_ action $ do + tid <- myThreadId + labelThread tid "LightStep reporter" + fix $ \loop -> do + work + loop + race_ + shutdown + (waitUntilDone (lsGracefulShutdownTimeoutSeconds cfg) doneVar) submitSpan :: Span -> IO () submitSpan sp = do diff --git a/src/LightStep/LowLevel.hs b/src/LightStep/LowLevel.hs --- a/src/LightStep/LowLevel.hs +++ b/src/LightStep/LowLevel.hs @@ -7,10 +7,9 @@ where import Chronos -import System.Timeout +import Control.Concurrent import Control.Exception.Safe import Control.Lens hiding (op) -import Control.Monad.IO.Class import Data.ProtoLens.Message (defMessage) import qualified Data.Text as T import LightStep.Internal.Debug @@ -20,8 +19,15 @@ import Proto.Collector import Proto.Collector_Fields import Proto.Google.Protobuf.Timestamp_Fields +import System.Timeout -data LightStepClient = LightStepClient GrpcClient T.Text Reporter +data LightStepClient + = LightStepClient + { lscGrpcVar :: MVar GrpcClient, + lscToken :: T.Text, + lscReporter :: Reporter, + lscConfig :: LightStepConfig + } data LightStepConfig = LightStepConfig @@ -32,35 +38,40 @@ lsGracefulShutdownTimeoutSeconds :: Int } -reportSpans :: LightStepClient -> [Span] -> ExceptT ClientError IO () -reportSpans (LightStepClient grpc token rep) sps = do - let req = liftIO . timeout 3_000_000 . runExceptT $ - rawUnary - (RPC :: RPC CollectorService "report") - grpc - ( defMessage - & auth .~ (defMessage & accessToken .~ token) - & spans .~ sps - & reporter .~ rep - ) - ret <- req - `withException` (\err -> d_ $ "reportSpans failed: " <> show (err :: SomeException)) - d_ $ show ret - -- FIXME: handle errors +reportSpans :: LightStepClient -> [Span] -> IO () +reportSpans client@LightStepClient {..} sps = do + let tryOnce = do + grpc <- readMVar lscGrpcVar + let req = + timeout 3_000_000 . runExceptT $ + rawUnary + (RPC :: RPC CollectorService "report") + grpc + ( defMessage + & auth .~ (defMessage & accessToken .~ lscToken) + & spans .~ sps + & reporter .~ lscReporter + ) + req `withException` (\err -> d_ $ "reportSpans failed: " <> show (err :: SomeException)) + ret <- tryOnce + ret2 <- case ret of + Nothing -> do + d_ "GRPC client is stuck, trying to reconnect" + reconnectClient client + -- one retry after reconnect + tryOnce + _ -> pure ret + d_ $ show ret2 pure () -mkClient :: LightStepConfig -> ClientIO LightStepClient -mkClient LightStepConfig {..} = do - grpc <- - setupGrpcClient - ( (grpcClientConfigSimple lsHostName lsPort True) - { _grpcClientConfigCompression = compression, - _grpcClientConfigTimeout = Timeout 5 -- seconds - } - ) - pure $ LightStepClient grpc lsToken rep +mkClient :: LightStepConfig -> IO LightStepClient +mkClient cfg@LightStepConfig {..} = do + grpcVar <- newEmptyMVar + let client = LightStepClient grpcVar lsToken rep cfg + grpc <- makeGrpcClient client + putMVar grpcVar grpc + pure client where - compression = if False then gzip else uncompressed rep = defMessage & reporterId .~ 2 @@ -70,8 +81,37 @@ defMessage & key .~ "lightstep.tracer_platform" & stringValue .~ "haskell" ] -closeClient :: LightStepClient -> ClientIO () -closeClient (LightStepClient grpc _ _) = close grpc +makeGrpcClient :: LightStepClient -> IO GrpcClient +makeGrpcClient client = do + let LightStepConfig {..} = lscConfig client + newGrpcOrError <- + runExceptT $ + setupGrpcClient + ( (grpcClientConfigSimple lsHostName lsPort True) + { _grpcClientConfigCompression = compression, + _grpcClientConfigTimeout = Timeout 5, -- seconds + _grpcClientConfigGoAwayHandler = \_ -> d_ "GoAway handler fired" + } + ) + case newGrpcOrError of + Right newGrpc -> pure newGrpc + Left err -> throwIO err + where + compression = if False then gzip else uncompressed + +reconnectClient :: LightStepClient -> IO () +reconnectClient client@LightStepClient {lscGrpcVar} = do + d_ "reconnectClient begin" + newClient <- makeGrpcClient client + oldClient <- swapMVar lscGrpcVar newClient + runExceptT $ close oldClient + d_ "reconnectClient end" + +closeClient :: LightStepClient -> IO () +closeClient LightStepClient {lscGrpcVar} = do + grpc <- readMVar lscGrpcVar + runExceptT $ close grpc + pure () startSpan :: T.Text -> IO Span startSpan op = do