From 549c14fb26d5c9c947e84fb02f54aa2276fb6e16 Mon Sep 17 00:00:00 2001
From: Simon Marlow <smarlow@fb.com>
Date: Tue, 5 Jul 2016 02:30:36 -0700
Subject: [PATCH] Track stack traces of dataFetch calls when profiling

Summary: This diff collects the stack traces of `dataFetch` calls, when `reportLevel` >= 5 and profiling is on.  Zero overhead for non-profiled code.

Reviewed By: niteria

Differential Revision: D2535947

fbshipit-source-id: fd43c20edd5455bd5e41113059fc69206b998e44
---
 Haxl/Core/Monad.hs | 18 ++++++++++++++++--
 Haxl/Core/Types.hs | 36 +++++++++++++++++++++++++++---------
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/Haxl/Core/Monad.hs b/Haxl/Core/Monad.hs
index a45a95e..73c2835 100644
--- a/Haxl/Core/Monad.hs
+++ b/Haxl/Core/Monad.hs
@@ -521,6 +521,19 @@ cachedWithInsert showFn insertFn env req = do
             Right _ -> "Cached request: " ++ showFn req
           return (Cached r)
 
+-- | Record the call stack for a data fetch in the Stats.  Only useful
+-- when profiling.
+logFetch :: Env u -> (r a -> String) -> r a -> IO ()
+#ifdef PROFILING
+logFetch env showFn req = do
+  ifReport (flags env) 5 $ do
+    stack <- currentCallStack
+    modifyIORef' (statsRef env) $ \(Stats s) ->
+      Stats (FetchCall (showFn req) stack : s)
+#else
+logFetch _ _ _ = return ()
+#endif
+
 -- | Performs actual fetching of data for a 'Request' from a 'DataSource'.
 dataFetch :: (DataSource u r, Request r a) => r a -> GenHaxl u a
 dataFetch = dataFetchWithInsert show DataCache.insert
@@ -547,13 +560,14 @@ dataFetchWithInsert showFn insertFn req = GenHaxl $ \env ref -> do
     -- Not seen before: add the request to the RequestStore, so it
     -- will be fetched in the next round.
     Uncached rvar -> do
+      logFetch env showFn req
       modifyIORef' ref $ \bs -> addRequest (BlockedFetch req rvar) bs
       return $ Blocked (Cont (continueFetch showFn req rvar))
 
     -- Seen before but not fetched yet.  We're blocked, but we don't have
     -- to add the request to the RequestStore.
-    CachedNotFetched rvar -> return
-      $ Blocked (Cont (continueFetch showFn req rvar))
+    CachedNotFetched rvar ->
+      return (Blocked (Cont (continueFetch showFn req rvar)))
 
     -- Cached: either a result, or an exception
     Cached (Left ex) -> return (Throw ex)
diff --git a/Haxl/Core/Types.hs b/Haxl/Core/Types.hs
index 4446b17..1a1c108 100644
--- a/Haxl/Core/Types.hs
+++ b/Haxl/Core/Types.hs
@@ -121,7 +121,7 @@ data Flags = Flags
     -- ^ Tracing level (0 = quiet, 3 = very verbose).
   , report :: {-# UNPACK #-} !Int
     -- ^ Report level (0 = quiet, 1 = # of requests, 2 = time, 3 = # of errors,
-    -- 4 = profiling)
+    -- 4 = profiling, 5 = log stack traces of dataFetch calls)
   , caching :: {-# UNPACK #-} !Int
     -- ^ Non-zero if caching is enabled.  If caching is disabled, then
     -- we still do batching and de-duplication within a round, but do
@@ -160,17 +160,29 @@ newtype Stats = Stats [RoundStats]
 -- | Pretty-print Stats.
 ppStats :: Stats -> String
 ppStats (Stats rss) =
-  intercalate "\n" [ "Round: " ++ show i ++ " - " ++ ppRoundStats rs
-                   | (i, rs) <- zip [(1::Round)..] (reverse rss) ]
+  intercalate "\n"
+     [ "Round: " ++ show i ++ " - " ++ ppRoundStats rs
+     | (i, rs) <- zip [(1::Int)..] (filter isRoundStats (reverse rss)) ]
+ where
+  isRoundStats RoundStats{} = True
+  isRoundStats _ = False
 
 -- | Maps data source name to the number of requests made in that round.
 -- The map only contains entries for sources that made requests in that
 -- round.
-data RoundStats = RoundStats
-  { roundTime :: Microseconds
-  , roundAllocation :: Int
-  , roundDataSources :: HashMap Text DataSourceRoundStats
-  }
+data RoundStats
+    -- | Timing stats for a round of data fetching
+  = RoundStats
+    { roundTime :: Microseconds
+    , roundAllocation :: Int
+    , roundDataSources :: HashMap Text DataSourceRoundStats
+    }
+    -- | The stack trace of a call to 'dataFetch'.  These are collected
+    -- only when profiling and reportLevel is 5 or greater.
+  | FetchCall
+    { fetchReq :: String
+    , fetchStack :: [String]
+    }
 
 -- | Pretty-print RoundStats.
 ppRoundStats :: RoundStats -> String
@@ -178,6 +190,7 @@ ppRoundStats (RoundStats t a dss) =
     show t ++ "us " ++ show a ++ " bytes\n"
       ++ unlines [ "  " ++ unpack nm ++ ": " ++ ppDataSourceRoundStats dsrs
                  | (nm, dsrs) <- sortBy (compare `on` fst) (toList dss) ]
+ppRoundStats (FetchCall r ss) = show r ++ '\n':show ss
 
 instance ToJSON RoundStats where
   toJSON RoundStats{..} = object
@@ -185,6 +198,10 @@ instance ToJSON RoundStats where
     , "allocation" .= roundAllocation
     , "dataSources" .= roundDataSources
     ]
+  toJSON (FetchCall req strs) = object
+    [ "request" .= req
+    , "stack" .= strs
+    ]
 
 -- | Detailed stats of each data source in each round.
 data DataSourceRoundStats = DataSourceRoundStats
@@ -213,12 +230,13 @@ instance ToJSON DataSourceRoundStats where
 fetchesInRound :: RoundStats -> Int
 fetchesInRound (RoundStats _ _ hm) =
   sum $ map dataSourceFetches $ HashMap.elems hm
+fetchesInRound _ = 0
 
 emptyStats :: Stats
 emptyStats = Stats []
 
 numRounds :: Stats -> Int
-numRounds (Stats rs) = length rs
+numRounds (Stats rs) = length [ s | s@RoundStats{} <- rs ]
 
 numFetches :: Stats -> Int
 numFetches (Stats rs) = sum (map fetchesInRound rs)