Distributed tracing (#206)

Author: serras

cabal.project

@@ -17,7 +17,8 @@ packages: compendium-client/
           grpc/client/
           grpc/server/
           graphql/
-          instrumentation/prometheus
+          instrumentation/prometheus/
+          instrumentation/tracing/
 
 source-repository-package
     type: git

core/rpc/mu-rpc.cabal

@@ -30,11 +30,13 @@ library
   build-depends:
       base              >=4.12  && <5
     , conduit           >=1.3.2 && <1.4
+    , http-types        >=0.12  && <0.13
     , mtl               >=2.2   && <2.3
     , mu-schema         ==0.3.*
     , sop-core          >=0.5   && <0.6
     , template-haskell  >=2.14  && <2.16
     , text              >=1.2   && <1.3
+    , wai               >=3.2   && <4
 
   hs-source-dirs:   src
   default-language: Haskell2010

core/rpc/src/Mu/Rpc.hs

@@ -28,10 +28,11 @@ module Mu.Rpc (
 ) where
 
 import           Data.Kind
-import           Data.Text           (Text)
-import qualified Data.Text           as T
+import           Data.Text                 (Text)
+import qualified Data.Text                 as T
 import           GHC.TypeLits
-import qualified Language.Haskell.TH as TH
+import qualified Language.Haskell.TH       as TH
+import           Network.HTTP.Types.Header
 import           Type.Reflection
 
 import           Mu.Schema
@@ -136,6 +137,7 @@ data RpcInfo i
   | RpcInfo { packageInfo :: Package Text Text Text TyInfo
             , serviceInfo :: Service Text Text Text TyInfo
             , methodInfo  :: Method  Text Text Text TyInfo
+            , headers     :: RequestHeaders
             , extraInfo   :: i
             }
 
@@ -148,13 +150,13 @@ data TyInfo
 instance Show (RpcInfo i) where
   show NoRpcInfo
     = "<no info>"
-  show (RpcInfo (Package Nothing _) (Service s _) (Method m _ _) _)
+  show (RpcInfo (Package Nothing _) (Service s _) (Method m _ _) _ _)
     = T.unpack (s <> ":" <> m)
-  show (RpcInfo (Package (Just p) _) (Service s _) (Method m _ _) _)
+  show (RpcInfo (Package (Just p) _) (Service s _) (Method m _ _) _ _)
     = T.unpack (p <> ":" <> s <> ":" <> m)
 
 class ReflectRpcInfo (p :: Package') (s :: Service') (m :: Method') where
-  reflectRpcInfo :: Proxy p -> Proxy s -> Proxy m -> i -> RpcInfo i
+  reflectRpcInfo :: Proxy p -> Proxy s -> Proxy m -> RequestHeaders -> i -> RpcInfo i
 class ReflectService (s :: Service') where
   reflectService :: Proxy s -> Service Text Text Text TyInfo
 class ReflectMethod (m :: Method') where
@@ -199,10 +201,10 @@ instance (ReflectArg m, ReflectArgs ms)
 
 instance (KnownMaySymbol pname, ReflectServices ss, ReflectService s, ReflectMethod m)
          => ReflectRpcInfo ('Package pname ss) s m where
-  reflectRpcInfo _ ps pm
+  reflectRpcInfo _ ps pm req extra
     = RpcInfo (Package (maySymbolVal (Proxy @pname))
                        (reflectServices (Proxy @ss)))
-              (reflectService ps) (reflectMethod pm)
+              (reflectService ps) (reflectMethod pm) req extra
 
 instance (KnownSymbol sname, ReflectMethods ms)
          => ReflectService ('Service sname ms) where

default.nix

@@ -1,5 +1,5 @@
 let
-  haskellNix = import (builtins.fetchTarball https://github.com/input-output-hk/haskell.nix/archive/9d491b5.tar.gz) {};
+  haskellNix = import (builtins.fetchTarball https://github.com/input-output-hk/haskell.nix/archive/d3edb6e.tar.gz) {};
   nixpkgsSrc = haskellNix.sources.nixpkgs-2003;
   nixpkgsArgs = haskellNix.nixpkgsArgs;
 in
@@ -31,4 +31,5 @@ in {
   mu-protobuf = hnPkgs.mu-protobuf.components.all;
   mu-rpc = hnPkgs.mu-rpc.components.library;
   mu-schema = hnPkgs.mu-schema.components.library;
+  mu-tracing = hnPkgs.mu-tracing.components.library;
 }

examples/health-check/mu-example-health-check.cabal

@@ -31,10 +31,13 @@ executable health-server
     , mu-protobuf     >=0.4.0
     , mu-rpc          >=0.4.0
     , mu-schema       >=0.3.0
+    , mu-tracing      >=0.4.0
+    , prometheus-client >= 1  && <2
     , stm             >=2.5   && <3
     , stm-conduit     >=4     && <5
     , stm-containers  >=1.1   && <2
     , text            >=1.2   && <2
+    , tracing-control >=0.0.6
     , wai             >=3.2   && <4
     , warp            >=3.3   && <4
 

examples/health-check/src/Server.hs

@@ -1,6 +1,8 @@
 {-# language DataKinds             #-}
+{-# language FlexibleContexts      #-}
 {-# language OverloadedStrings     #-}
 {-# language PartialTypeSignatures #-}
+{-# language PolyKinds             #-}
 {-# language TypeApplications      #-}
 {-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}
 
@@ -9,34 +11,48 @@ module Main where
 import           Control.Concurrent.Async
 import           Control.Concurrent.STM
 import           Control.Monad.IO.Class
+import           Control.Monad.Trace
 import           Data.Conduit
 import qualified Data.Conduit.Combinators      as C
 import           Data.Conduit.TMChan
 import           Data.Maybe                    (fromMaybe)
 import           Data.Proxy
 import qualified Data.Text                     as T
 import           DeferredFolds.UnfoldlM
+import           Monitor.Tracing.Zipkin        (Endpoint (..))
 import           Network.Wai.Handler.Warp
+import           Prometheus
 import qualified StmContainers.Map             as M
 
 import           Mu.GraphQL.Server
 import           Mu.GRpc.Server
 import           Mu.Instrumentation.Prometheus
+import           Mu.Instrumentation.Tracing
 import           Mu.Server
 
 import           Definition
 
 main :: IO ()
 main = do
+  -- Initialize prometheus
+  met <- initPrometheus "health"
+  -- Initialize zipkin
+  zpk <- newZipkin defaultZipkinSettings
+           { settingsPublishPeriod = Just 1
+           , settingsEndpoint = Just $ Endpoint (Just "me") Nothing Nothing Nothing }
+  let rootInfo = MuTracing alwaysSampled "health-check"
+  -- Initialize app
   m <- M.newIO
   upd <- newTBMChanIO 100
-  met <- initPrometheus "health"
+  -- Put together the server
+  let s = zipkin rootInfo $ prometheus met $ server m upd
+  -- Run the app
   putStrLn "running health check application"
-  let s = prometheus met (server m upd)
   runConcurrently $ (\_ _ _ -> ())
-    <$> Concurrently (runner 50051 (gRpcApp msgProtoBuf s))
-    <*> Concurrently (runner 50052 (gRpcApp msgAvro s))
-    <*> Concurrently (runner 50053 (graphQLAppQuery s (Proxy @"HealthCheckServiceFS2")))
+    <$> Concurrently (runner 50051 (gRpcAppTrans msgProtoBuf (runZipkin zpk) s))
+    <*> Concurrently (runner 50052 (gRpcAppTrans msgAvro     (runZipkin zpk) s))
+    <*> Concurrently (runner 50053 (graphQLAppTransQuery (runZipkin zpk) s
+                                      (Proxy @"HealthCheckServiceFS2")))
   where runner p app = run p (prometheusWai ["metrics"] app)
 
 -- Server implementation
@@ -45,7 +61,9 @@ main = do
 type StatusMap = M.Map T.Text T.Text
 type StatusUpdates = TBMChan HealthStatusMsg
 
-server :: StatusMap -> StatusUpdates -> ServerIO info HealthCheckService _
+server :: (MonadServer m, MonadTrace m)
+       => StatusMap -> StatusUpdates
+       -> ServerT '[] info HealthCheckService m _
 server m upd
   = wrapServer (\info h -> liftIO (print info) >> h) $
     singleService ( method @"setStatus"   $ setStatus_ m upd
@@ -55,30 +73,36 @@ server m upd
                   , method @"cleanAll"    $ cleanAll_ m
                   , method @"watch"       $ watch_ upd)
 
-setStatus_ :: StatusMap -> StatusUpdates -> HealthStatusMsg -> ServerErrorIO ()
+setStatus_ :: (MonadServer m, MonadTrace m)
+           => StatusMap -> StatusUpdates -> HealthStatusMsg
+           -> m ()
 setStatus_ m upd
            s@(HealthStatusMsg (Just (HealthCheckMsg nm)) (Just (ServerStatusMsg ss)))
-  = alwaysOk $ do
+  = childSpan "setStatus" $ alwaysOk $ do
       putStr "setStatus: " >> print (nm, ss)
       atomically $ do
         M.insert ss nm m
         writeTBMChan upd s
 setStatus_ _ _ _ = serverError (ServerError Invalid "name or status missing")
 
-checkH_ :: StatusMap -> HealthCheckMsg -> ServerErrorIO ServerStatusMsg
+checkH_ :: (MonadServer m, MonadTrace m)
+        => StatusMap -> HealthCheckMsg
+        -> m ServerStatusMsg
 checkH_ _ (HealthCheckMsg "") = serverError (ServerError Invalid "no server name given")
 checkH_ m (HealthCheckMsg nm) = alwaysOk $ do
   putStr "check: " >> print nm
   ss <- atomically $ M.lookup nm m
   pure $ ServerStatusMsg (fromMaybe "" ss)
 
-clearStatus_ :: StatusMap -> HealthCheckMsg -> ServerErrorIO ()
+clearStatus_ :: (MonadServer m, MonadTrace m)
+             => StatusMap -> HealthCheckMsg -> m ()
 clearStatus_ _ (HealthCheckMsg "") = serverError (ServerError Invalid "no server name given")
 clearStatus_ m (HealthCheckMsg nm) = alwaysOk $ do
   putStr "clearStatus: " >> print nm
   atomically $ M.delete nm m
 
-checkAll_ :: StatusMap -> ServerErrorIO AllStatusMsg
+checkAll_ :: (MonadServer m, MonadTrace m)
+          => StatusMap -> m AllStatusMsg
 checkAll_ m = alwaysOk $ do
     putStrLn "checkAll"
     AllStatusMsg <$> atomically (consumeValues kvToStatus (M.unfoldlM m))
@@ -87,15 +111,17 @@ checkAll_ m = alwaysOk $ do
     consumeValues f = foldlM' (\xs (x,y) -> pure (f x y:xs)) []
     kvToStatus k v = HealthStatusMsg (Just (HealthCheckMsg k)) (Just (ServerStatusMsg v))
 
-cleanAll_ :: StatusMap -> ServerErrorIO ()
+cleanAll_ :: (MonadServer m, MonadTrace m)
+          => StatusMap -> m ()
 cleanAll_ m = alwaysOk $ do
   putStrLn "cleanAll"
   atomically $ M.reset m
 
-watch_ :: StatusUpdates
+watch_ :: (MonadServer m, MonadTrace m)
+       => StatusUpdates
        -> HealthCheckMsg
-       -> ConduitT ServerStatusMsg Void ServerErrorIO ()
-       -> ServerErrorIO ()
+       -> ConduitT ServerStatusMsg Void m ()
+       -> m ()
 watch_ upd hcm@(HealthCheckMsg nm) sink = do
   alwaysOk (putStr "watch: " >> print nm)
   runConduit $ sourceTBMChan upd
@@ -109,3 +135,5 @@ watch_ upd hcm@(HealthCheckMsg nm) sink = do
                       Just (Just y) -> yield y >> catMaybesC
                       Just Nothing  -> catMaybesC
                       Nothing       -> pure ()
+
+instance MonadMonitor m => MonadMonitor (TraceT m)