fix acocunting of label hits (#130)

Summary:
Pull Request resolved: #130

Label hits were incorrectly counted when a Blocked computation was hit

Reviewed By: anubhav94N

Differential Revision: D23573141

fbshipit-source-id: 79f4e17fe5fa0002a9e334e401ca1620ce07ff5b

Haxl/Core/Profile.hs

@@ -78,16 +78,17 @@ collectProfileData l m env = do
             (profileNextKey p)
             (profileTree p)
           , profileNextKey = profileNextKey p + 1 }, profileNextKey p)
-    runProfileData l key m env
+    runProfileData l key m False env
 {-# INLINE collectProfileData #-}
 
 runProfileData
   :: ProfileLabel
   -> ProfileKey
   -> (Env u w -> IO (Result u w a))
+  -> Bool
   -> Env u w
   -> IO (Result u w a)
-runProfileData l key m env = do
+runProfileData l key m isCont env = do
   a0 <- getAllocationCounter
   let
     nextCurrent = ProfileCurrent
@@ -103,7 +104,7 @@ runProfileData l key m env = do
   -- But this is what we want as we need to account for allocations.
   -- So do not be tempted to pass through prevProfKey (from collectProfileData)
   -- which is the original caller
-  modifyProfileData env key caller (a0 - a1)
+  modifyProfileData env key caller (a0 - a1) (if isCont then 0 else 1)
 
   -- So we do not count the allocation overhead of modifyProfileData
   setAllocationCounter a1
@@ -112,16 +113,17 @@ runProfileData l key m env = do
     Throw e -> return (Throw e)
     Blocked ivar k -> return (Blocked ivar (Cont $ runCont (toHaxl k)))
   where
-    runCont (GenHaxl h) = GenHaxl $ \env -> runProfileData l key h env
+    runCont (GenHaxl h) = GenHaxl $ runProfileData l key h True
 {-# INLINE runProfileData #-}
 
 modifyProfileData
   :: Env u w
   -> ProfileKey
   -> ProfileKey
   -> AllocCount
+  -> LabelHitCount
   -> IO ()
-modifyProfileData env key caller allocs = do
+modifyProfileData env key caller allocs labelIncrement = do
   modifyIORef' (profRef env) $ \ p ->
     p { profile =
           HashMap.insertWith updEntry key newEntry .
@@ -130,12 +132,12 @@ modifyProfileData env key caller allocs = do
   where newEntry =
           emptyProfileData
             { profileAllocs = allocs
-            , profileLabelHits = 1
+            , profileLabelHits = labelIncrement
             }
         updEntry _ old =
           old
             { profileAllocs = profileAllocs old + allocs
-            , profileLabelHits = profileLabelHits old + 1
+            , profileLabelHits = profileLabelHits old + labelIncrement
             }
         -- subtract allocs from caller, so they are not double counted
         -- we don't know the caller's caller, but it will get set on

tests/ProfileTests.hs

@@ -27,6 +27,7 @@ import Data.Int
 
 import TestUtils
 import WorkDataSource
+import SleepDataSource
 
 mkProfilingEnv = do
   env <- makeTestEnv False
@@ -49,8 +50,10 @@ collectsdata = do
           withLabel "bar" $
             withLabel "foo" $ do
               u <- env userEnv
+              slp <- sum <$> mapM (\x -> withLabel "baz" $ return x) [1..5]
               -- do some non-trivial work that can't be lifted out
-              case fromJSON <$> HashMap.lookup "A" u of
+              -- first sleep though in order to force a Blocked result
+              sleep slp `andThen` case fromJSON <$> HashMap.lookup "A" u of
                 Just (Success n) | sum [n .. 1000::Integer] > 0 -> return 5
                 _otherwise -> return (4::Int)
   profCopy <- readIORef (profRef e)
@@ -61,11 +64,15 @@ collectsdata = do
     getData k = do
       k2 <- HashMap.lookup k labelKeys
       HashMap.lookup k2 profData
-  assertEqual "has data" 3 $ HashMap.size profData
+  assertEqual "has data" 4 $ HashMap.size profData
   assertBool "foo allocates" $
     case profileAllocs <$> getData "foo" of
       Just x -> x > 10000
       Nothing -> False
+  assertEqual "foo is only called once" (Just 1) $
+    profileLabelHits <$> getData "foo"
+  assertEqual "baz is called 5 times" (Just 5) $
+    profileLabelHits <$> getData "baz"
   assertBool "bar does not allocate (much)" $
     case profileAllocs <$> getData "bar" of
       Just n -> n < 5000  -- getAllocationCounter can be off by +/- 4K

tests/TestUtils.hs

@@ -16,6 +16,7 @@ module TestUtils
 
 import TestTypes
 import MockTAO
+import Haxl.DataSource.ConcurrentIO
 
 import Data.IORef
 import Data.Aeson
@@ -49,7 +50,8 @@ id4 = lookupInput "D"
 makeTestEnv :: Bool -> IO HaxlEnv
 makeTestEnv future = do
   tao <- MockTAO.initGlobalState future
-  let st = stateSet tao stateEmpty
+  stio <- mkConcurrentIOState
+  let st = stateSet stio $ stateSet tao stateEmpty
   env <- initEnv st testinput
   return env { flags = (flags env) { report = 2 } }