Move CallGraphRef to Haxl Env

Summary:
Adds the ability for one to record the function callstack within the `Haxl` monad.
The main function is `withCallGraph`, and supporting types are in `Haxl.Core.CallGraph`.

Reviewed By: watashi

Differential Revision: D14857851

fbshipit-source-id: daab97e8144ff63104ad9e79a3a792a4471e5588

Haxl/Core.hs

@@ -7,6 +7,7 @@
 -- | Everything needed to define data sources and to invoke the
 -- engine.
 --
+{-# LANGUAGE CPP #-}
 module Haxl.Core (
     -- * The monad and operations
     GenHaxl (..), runHaxl, runHaxlWithWrites
@@ -95,8 +96,12 @@ module Haxl.Core (
 
     -- * Exceptions
   , module Haxl.Core.Exception
+
+    -- * Recording the function callgraph
+  , module Haxl.Core.CallGraph
   ) where
 
+import Haxl.Core.CallGraph
 import Haxl.Core.DataSource
 import Haxl.Core.Flags
 import Haxl.Core.Memo

Haxl/Core/CallGraph.hs

@@ -0,0 +1,38 @@
+-- Copyright 2004-present Facebook. All Rights Reserved.
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE CPP #-}
+
+module Haxl.Core.CallGraph where
+
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
+#if __GLASGOW_HASKELL__ < 804
+import Data.Monoid
+#endif
+import Data.Text (Text)
+import qualified Data.Text as Text
+
+type ModuleName = Text
+
+-- | An unqualified function
+type Function = Text
+
+-- | A qualified function
+data QualFunction = QualFunction ModuleName Function deriving (Eq, Ord)
+
+instance Show QualFunction where
+  show (QualFunction mn nm) = Text.unpack $ mn <> Text.pack "." <> nm
+
+-- | Represents an edge between a parent function which calls a child function
+-- in the call graph
+type FunctionCall = (QualFunction, QualFunction)
+
+-- | An edge list which represents the dependencies between function calls
+type CallGraph = ([FunctionCall], Map QualFunction Text)
+
+-- | Used as the root of all function calls
+mainFunction :: QualFunction
+mainFunction = QualFunction "MAIN" "main"
+
+emptyCallGraph :: CallGraph
+emptyCallGraph = ([], Map.empty)

Haxl/Core/Monad.hs

@@ -100,6 +100,11 @@ module Haxl.Core.Monad
   , dumpCacheAsHaskell
   , dumpCacheAsHaskellFn
 
+    -- * CallGraph
+#ifdef PROFILING
+  , withCallGraph
+#endif
+
     -- * Unsafe operations
   ,  unsafeLiftIO, unsafeToHaxlException
   ) where
@@ -135,7 +140,11 @@ import Debug.Trace (traceEventIO)
 #endif
 
 #ifdef PROFILING
+import qualified Data.Map as Map
+import Data.Text (Text)
+import Data.Typeable
 import GHC.Stack
+import Haxl.Core.CallGraph
 #endif
 
 
@@ -207,6 +216,14 @@ data Env u w = Env
        -- ^ A log of all writes done as part of this haxl computation. Any
        -- haxl computation that needs to be memoized runs in its own
        -- environment so
+#ifdef PROFILING
+  , callGraphRef ::  Maybe (IORef CallGraph)
+       -- ^ An edge list representing the current function call graph. The type
+       -- is wrapped in a Maybe to avoid changing the existing callsites.
+
+  , currFunction :: QualFunction
+       -- ^ The most recent function call.
+#endif
   }
 
 type Caches u w = (IORef (DataCache (IVar u w)), IORef (DataCache (IVar u w)))
@@ -241,6 +258,10 @@ initEnvWithData states e (cref, mref) = do
     , pendingWaits = []
     , speculative = 0
     , writeLogsRef = wl
+#ifdef PROFILING
+    , callGraphRef = Nothing
+    , currFunction = mainFunction
+#endif
     }
 
 -- | Initializes an environment with 'StateStore' and an input map.
@@ -729,6 +750,36 @@ withEnv newEnv (GenHaxl m) = GenHaxl $ \_env -> do
     Blocked ivar k ->
       return (Blocked ivar (Cont (withEnv newEnv (toHaxl k))))
 
+#ifdef PROFILING
+-- -----------------------------------------------------------------------------
+-- CallGraph recording
+
+-- | Returns a version of the Haxl computation which records function calls in
+-- an edge list which is the function call graph. Each function that is to be
+-- recorded must be wrapped with a call to @withCallGraph@.
+withCallGraph
+  :: Typeable a
+  => (a -> Maybe Text)
+  -> QualFunction
+  -> GenHaxl u () a
+  -> GenHaxl u () a
+withCallGraph toText f a = do
+  coreEnv <- env id
+  -- TODO: Handle exceptions
+  value <- withEnv coreEnv{currFunction = f} a
+  case callGraphRef coreEnv of
+    Just graph -> unsafeLiftIO $ modifyIORef' graph
+      (updateCallGraph (f, currFunction coreEnv) (toText value))
+    _ -> throw $ CriticalError
+      "withCallGraph called without an IORef CallGraph"
+  return value
+  where
+    updateCallGraph :: FunctionCall -> Maybe Text -> CallGraph -> CallGraph
+    updateCallGraph fnCall@(childQFunc, _) (Just value) (edgeList, valueMap) =
+      (fnCall : edgeList, Map.insert childQFunc value valueMap)
+    updateCallGraph fnCall Nothing (edgeList, valueMap) =
+      (fnCall : edgeList, valueMap)
+#endif
 
 -- -----------------------------------------------------------------------------
 -- Exceptions

changelog.md

@@ -1,5 +1,10 @@
-# Changes in version 2.1.1.0
+# Changes in version 2.1.2.0
 
+  * Add a callgraph reference to 'Env' to record the function callgraph during a
+    computation. The callgraph is stored as an edge list in the Env through the
+    use of `withCallGraph` and enables users to debug a Haxl computation.
+
+# Changes in version 2.1.1.0
   * Adds feature to track outgone datasource fetches. This is only turned on
     for report level greater than 1. The fetches are stored as a running Map
     in the env ('submittedReqsRef').

haxl.cabal

@@ -1,5 +1,5 @@
 name:                haxl
-version:             2.0.1.1
+version:             2.1.2.0
 synopsis:            A Haskell library for efficient, concurrent,
                      and concise data access.
 homepage:            https://github.com/facebook/Haxl
@@ -68,6 +68,7 @@ library
 
   exposed-modules:
     Haxl.Core,
+    Haxl.Core.CallGraph,
     Haxl.Core.DataCache,
     Haxl.Core.DataSource,
     Haxl.Core.Exception,