{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- | Implementation of lightweight profiling.  Most users should
-- import "Haxl.Core" instead.
--
module Haxl.Core.Profile
  ( withLabel
  , withFingerprintLabel
  , addProfileFetch
  , incrementMemoHitCounterFor
  , collectProfileData
  , profileCont
  ) where

import Data.IORef
import Data.Hashable
import Data.List.NonEmpty (NonEmpty(..), (<|))
#if __GLASGOW_HASKELL__ < 804
import Data.Monoid
#endif
import Data.Typeable
import qualified Data.HashMap.Strict as HashMap
import GHC.Exts
import qualified Data.Text as Text
import Haxl.Core.DataSource
import Haxl.Core.Flags
import Haxl.Core.Stats
import Haxl.Core.Monad

-- -----------------------------------------------------------------------------
-- Profiling

-- | Label a computation so profiling data is attributed to the label.
withLabel :: ProfileLabel -> GenHaxl u w a -> GenHaxl u w a
withLabel :: forall u w a. ProfileLabel -> GenHaxl u w a -> GenHaxl u w a
withLabel ProfileLabel
l (GenHaxl Env u w -> IO (Result u w a)
m) = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \Env u w
env ->
  if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ReportFlag -> ReportFlags -> Bool
testReportFlag ReportFlag
ReportProfiling (ReportFlags -> Bool) -> ReportFlags -> Bool
forall a b. (a -> b) -> a -> b
$ Flags -> ReportFlags
report (Flags -> ReportFlags) -> Flags -> ReportFlags
forall a b. (a -> b) -> a -> b
$ Env u w -> Flags
forall u w. Env u w -> Flags
flags Env u w
env
     then Env u w -> IO (Result u w a)
m Env u w
env
     else ProfileLabel
-> (Env u w -> IO (Result u w a)) -> Env u w -> IO (Result u w a)
forall u w a.
ProfileLabel
-> (Env u w -> IO (Result u w a)) -> Env u w -> IO (Result u w a)
collectProfileData ProfileLabel
l Env u w -> IO (Result u w a)
m Env u w
env

-- | Label a computation so profiling data is attributed to the label.
-- Intended only for internal use by 'memoFingerprint'.
withFingerprintLabel :: Addr# -> Addr# -> GenHaxl u w a -> GenHaxl u w a
withFingerprintLabel :: forall u w a. Addr# -> Addr# -> GenHaxl u w a -> GenHaxl u w a
withFingerprintLabel Addr#
mnPtr Addr#
nPtr (GenHaxl Env u w -> IO (Result u w a)
m) = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ \Env u w
env ->
  if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ReportFlag -> ReportFlags -> Bool
testReportFlag ReportFlag
ReportProfiling (ReportFlags -> Bool) -> ReportFlags -> Bool
forall a b. (a -> b) -> a -> b
$ Flags -> ReportFlags
report (Flags -> ReportFlags) -> Flags -> ReportFlags
forall a b. (a -> b) -> a -> b
$ Env u w -> Flags
forall u w. Env u w -> Flags
flags Env u w
env
     then Env u w -> IO (Result u w a)
m Env u w
env
     else ProfileLabel
-> (Env u w -> IO (Result u w a)) -> Env u w -> IO (Result u w a)
forall u w a.
ProfileLabel
-> (Env u w -> IO (Result u w a)) -> Env u w -> IO (Result u w a)
collectProfileData
            (Addr# -> ProfileLabel
Text.unpackCString# Addr#
mnPtr ProfileLabel -> ProfileLabel -> ProfileLabel
forall a. Semigroup a => a -> a -> a
<> ProfileLabel
"." ProfileLabel -> ProfileLabel -> ProfileLabel
forall a. Semigroup a => a -> a -> a
<> Addr# -> ProfileLabel
Text.unpackCString# Addr#
nPtr)
            Env u w -> IO (Result u w a)
m Env u w
env

-- | Collect profiling data and attribute it to given label.
collectProfileData
  :: ProfileLabel
  -> (Env u w -> IO (Result u w a))
  -> Env u w
  -> IO (Result u w a)
collectProfileData :: forall u w a.
ProfileLabel
-> (Env u w -> IO (Result u w a)) -> Env u w -> IO (Result u w a)
collectProfileData ProfileLabel
l Env u w -> IO (Result u w a)
m Env u w
env = do
  let ProfileCurrent ProfileKey
prevProfKey (ProfileLabel
prevProfLabel :| [ProfileLabel]
_) = Env u w -> ProfileCurrent
forall u w. Env u w -> ProfileCurrent
profCurrent Env u w
env
  if ProfileLabel
prevProfLabel ProfileLabel -> ProfileLabel -> Bool
forall a. Eq a => a -> a -> Bool
== ProfileLabel
l
  then
    -- do not add a new label if we are recursing
    Env u w -> IO (Result u w a)
m Env u w
env
  else do
    ProfileKey
key <- IORef Profile
-> (Profile -> (Profile, ProfileKey)) -> IO ProfileKey
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Env u w -> IORef Profile
forall u w. Env u w -> IORef Profile
profRef Env u w
env) ((Profile -> (Profile, ProfileKey)) -> IO ProfileKey)
-> (Profile -> (Profile, ProfileKey)) -> IO ProfileKey
forall a b. (a -> b) -> a -> b
$ \Profile
p ->
      case (ProfileLabel, ProfileKey)
-> HashMap (ProfileLabel, ProfileKey) ProfileKey
-> Maybe ProfileKey
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup (ProfileLabel
l, ProfileKey
prevProfKey) (Profile -> HashMap (ProfileLabel, ProfileKey) ProfileKey
profileTree Profile
p) of
        Just ProfileKey
k -> (Profile
p, ProfileKey
k)
        Maybe ProfileKey
Nothing -> (Profile
p
          { profileTree = HashMap.insert
            (l, prevProfKey)
            (profileNextKey p)
            (profileTree p)
          , profileNextKey = profileNextKey p + 1 }, Profile -> ProfileKey
profileNextKey Profile
p)
    ProfileLabel
-> ProfileKey
-> (Env u w -> IO (Result u w a))
-> Bool
-> Env u w
-> IO (Result u w a)
forall u w a.
ProfileLabel
-> ProfileKey
-> (Env u w -> IO (Result u w a))
-> Bool
-> Env u w
-> IO (Result u w a)
runProfileData ProfileLabel
l ProfileKey
key Env u w -> IO (Result u w a)
m Bool
False Env u w
env
{-# INLINE collectProfileData #-}

runProfileData
  :: ProfileLabel
  -> ProfileKey
  -> (Env u w -> IO (Result u w a))
  -> Bool
  -> Env u w
  -> IO (Result u w a)
runProfileData :: forall u w a.
ProfileLabel
-> ProfileKey
-> (Env u w -> IO (Result u w a))
-> Bool
-> Env u w
-> IO (Result u w a)
runProfileData ProfileLabel
l ProfileKey
key Env u w -> IO (Result u w a)
m Bool
isCont Env u w
env = do
  ProfileKey
t0 <- IO ProfileKey
getTimestamp
  ProfileKey
a0 <- IO ProfileKey
getAllocationCounter
  let
    ProfileCurrent ProfileKey
caller NonEmpty ProfileLabel
stack = Env u w -> ProfileCurrent
forall u w. Env u w -> ProfileCurrent
profCurrent Env u w
env
    nextCurrent :: ProfileCurrent
nextCurrent = ProfileCurrent
      { profCurrentKey :: ProfileKey
profCurrentKey = ProfileKey
key
      , profLabelStack :: NonEmpty ProfileLabel
profLabelStack = ProfileLabel
l ProfileLabel -> NonEmpty ProfileLabel -> NonEmpty ProfileLabel
forall a. a -> NonEmpty a -> NonEmpty a
<| NonEmpty ProfileLabel
stack
      }
    runCont :: GenHaxl u w a -> GenHaxl u w a
runCont (GenHaxl Env u w -> IO (Result u w a)
h) = (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall u w a. (Env u w -> IO (Result u w a)) -> GenHaxl u w a
GenHaxl ((Env u w -> IO (Result u w a)) -> GenHaxl u w a)
-> (Env u w -> IO (Result u w a)) -> GenHaxl u w a
forall a b. (a -> b) -> a -> b
$ ProfileLabel
-> ProfileKey
-> (Env u w -> IO (Result u w a))
-> Bool
-> Env u w
-> IO (Result u w a)
forall u w a.
ProfileLabel
-> ProfileKey
-> (Env u w -> IO (Result u w a))
-> Bool
-> Env u w
-> IO (Result u w a)
runProfileData ProfileLabel
l ProfileKey
key Env u w -> IO (Result u w a)
h Bool
True

  Result u w a
r <- Env u w -> IO (Result u w a)
m Env u w
env{profCurrent=nextCurrent} -- what if it throws?

  -- Make the result strict in Done/Throw so that if the user code
  -- returns (force a), the force is evaluated *inside* the profile.
  Result u w a
result <- case Result u w a
r of
    Done !a
a -> Result u w a -> IO (Result u w a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Result u w a
forall u w a. a -> Result u w a
Done a
a)
    Throw !SomeException
e -> Result u w a -> IO (Result u w a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (SomeException -> Result u w a
forall u w a. SomeException -> Result u w a
Throw SomeException
e)
    Blocked IVar u w b
ivar Cont u w a
k -> Result u w a -> IO (Result u w a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (IVar u w b -> Cont u w a -> Result u w a
forall u w a b. IVar u w b -> Cont u w a -> Result u w a
Blocked IVar u w b
ivar (GenHaxl u w a -> Cont u w a
forall u w a. GenHaxl u w a -> Cont u w a
Cont (GenHaxl u w a -> Cont u w a) -> GenHaxl u w a -> Cont u w a
forall a b. (a -> b) -> a -> b
$ GenHaxl u w a -> GenHaxl u w a
runCont (Cont u w a -> GenHaxl u w a
forall u w a. Cont u w a -> GenHaxl u w a
toHaxl Cont u w a
k)))

  ProfileKey
a1 <- IO ProfileKey
getAllocationCounter
  ProfileKey
t1 <- IO ProfileKey
getTimestamp

  -- caller might not be the actual caller of this function
  -- for example MAIN may be continuing a function from the middle of the stack.
  -- 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
  Env u w
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> IO ()
forall u w.
Env u w
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> IO ()
modifyProfileData Env u w
env ProfileKey
key ProfileKey
caller (ProfileKey
a0 ProfileKey -> ProfileKey -> ProfileKey
forall a. Num a => a -> a -> a
- ProfileKey
a1) (ProfileKey
t1ProfileKey -> ProfileKey -> ProfileKey
forall a. Num a => a -> a -> a
-ProfileKey
t0) (if Bool
isCont then ProfileKey
0 else ProfileKey
1)

  -- So we do not count the allocation overhead of modifyProfileData
  ProfileKey -> IO ()
setAllocationCounter ProfileKey
a1
  Result u w a -> IO (Result u w a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Result u w a
result
{-# INLINE runProfileData #-}

modifyProfileData
  :: Env u w
  -> ProfileKey
  -> ProfileKey
  -> AllocCount
  -> Microseconds
  -> LabelHitCount
  -> IO ()
modifyProfileData :: forall u w.
Env u w
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> ProfileKey
-> IO ()
modifyProfileData Env u w
env ProfileKey
key ProfileKey
caller ProfileKey
allocs ProfileKey
t ProfileKey
labelIncrement = do
  IORef Profile -> (Profile -> Profile) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Env u w -> IORef Profile
forall u w. Env u w -> IORef Profile
profRef Env u w
env) ((Profile -> Profile) -> IO ()) -> (Profile -> Profile) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Profile
p ->
    Profile
p { profile =
          HashMap.insertWith updEntry key newEntry .
          HashMap.insertWith updCaller caller newCaller $
          profile p }
  where newEntry :: ProfileData
newEntry =
          ProfileData
emptyProfileData
            { profileAllocs = allocs
            , profileLabelHits = labelIncrement
            , profileTime = t
            }
        updEntry :: ProfileData -> ProfileData -> ProfileData
updEntry ProfileData
_ ProfileData
old =
          ProfileData
old
            { profileAllocs = profileAllocs old + allocs
            , profileLabelHits = profileLabelHits old + labelIncrement
            , profileTime = profileTime old + t
            }
        -- subtract allocs/time from caller, so they are not double counted
        -- we don't know the caller's caller, but it will get set on
        -- the way back out, so an empty hashset is fine for now
        newCaller :: ProfileData
newCaller =
          ProfileData
emptyProfileData { profileAllocs = -allocs
                           , profileTime = -t
                           }
        updCaller :: ProfileData -> ProfileData -> ProfileData
updCaller ProfileData
_ ProfileData
old =
          ProfileData
old { profileAllocs = profileAllocs old - allocs
              , profileTime = profileTime old - t
              }


-- Like collectProfileData, but intended to be run from the scheduler.
--
-- * doesn't add a dependency (the original withLabel did this)
--
-- * doesn't subtract allocs from the caller (we're evaluating this
--   cont from the top level, so we don't need this)
--
-- * doesn't wrap a Blocked continuation in withLabel (the scheduler
--   will call profileCont the next time this cont runs)
--
profileCont
  :: (Env u w -> IO (Result u w a))
  -> Env u w
  -> IO (Result u w a)
profileCont :: forall u w a.
(Env u w -> IO (Result u w a)) -> Env u w -> IO (Result u w a)
profileCont Env u w -> IO (Result u w a)
m Env u w
env = do
  ProfileKey
t0 <- IO ProfileKey
getTimestamp
  ProfileKey
a0 <- IO ProfileKey
getAllocationCounter
  Result u w a
r <- Env u w -> IO (Result u w a)
m Env u w
env
  ProfileKey
a1 <- IO ProfileKey
getAllocationCounter
  ProfileKey
t1 <- IO ProfileKey
getTimestamp
  let
    allocs :: ProfileKey
allocs = ProfileKey
a0 ProfileKey -> ProfileKey -> ProfileKey
forall a. Num a => a -> a -> a
- ProfileKey
a1
    t :: ProfileKey
t = ProfileKey
t1 ProfileKey -> ProfileKey -> ProfileKey
forall a. Num a => a -> a -> a
- ProfileKey
t0
    newEntry :: ProfileData
newEntry = ProfileData
emptyProfileData
      { profileAllocs = allocs
      , profileTime = t
      }
    updEntry :: ProfileData -> ProfileData -> ProfileData
updEntry ProfileData
_ ProfileData
old = ProfileData
old
      { profileAllocs = profileAllocs old + allocs
      , profileTime = profileTime old + t
      }
    profKey :: ProfileKey
profKey = ProfileCurrent -> ProfileKey
profCurrentKey (Env u w -> ProfileCurrent
forall u w. Env u w -> ProfileCurrent
profCurrent Env u w
env)
  IORef Profile -> (Profile -> Profile) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Env u w -> IORef Profile
forall u w. Env u w -> IORef Profile
profRef Env u w
env) ((Profile -> Profile) -> IO ()) -> (Profile -> Profile) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Profile
p ->
    Profile
p { profile =
         HashMap.insertWith updEntry profKey newEntry $
         profile p }
  -- So we do not count the allocation overhead of modifyProfileData
  ProfileKey -> IO ()
setAllocationCounter ProfileKey
a1
  Result u w a -> IO (Result u w a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Result u w a
r
{-# INLINE profileCont #-}

incrementMemoHitCounterFor :: Env u w -> CallId -> Bool -> IO ()
incrementMemoHitCounterFor :: forall u w. Env u w -> CallId -> Bool -> IO ()
incrementMemoHitCounterFor Env u w
env CallId
callId Bool
wasCached = do
  IORef Profile -> (Profile -> Profile) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Env u w -> IORef Profile
forall u w. Env u w -> IORef Profile
profRef Env u w
env) ((Profile -> Profile) -> IO ()) -> (Profile -> Profile) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Profile
p ->  Profile
p {
    profile = HashMap.insertWith
                upd
                (profCurrentKey $ profCurrent env)
                (emptyProfileData { profileMemos = [val] })
                (profile p)
    }
  where
    val :: ProfileMemo
val = CallId -> Bool -> ProfileMemo
ProfileMemo CallId
callId Bool
wasCached
    upd :: ProfileData -> ProfileData -> ProfileData
upd ProfileData
_ ProfileData
old = ProfileData
old { profileMemos = val : profileMemos old }

{-# NOINLINE addProfileFetch #-}
addProfileFetch
  :: forall r u w a . (DataSourceName r, Eq (r a), Hashable (r a), Typeable (r a))
  => Env u w -> r a -> CallId -> Bool -> IO ()
addProfileFetch :: forall (r :: * -> *) u w a.
(DataSourceName r, Eq (r a), Hashable (r a), Typeable (r a)) =>
Env u w -> r a -> CallId -> Bool -> IO ()
addProfileFetch Env u w
env r a
_req CallId
cid Bool
wasCached = do
  ProfileKey
c <- IO ProfileKey
getAllocationCounter
  let (ProfileCurrent ProfileKey
profKey NonEmpty ProfileLabel
_) = Env u w -> ProfileCurrent
forall u w. Env u w -> ProfileCurrent
profCurrent Env u w
env
  IORef Profile -> (Profile -> Profile) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Env u w -> IORef Profile
forall u w. Env u w -> IORef Profile
profRef Env u w
env) ((Profile -> Profile) -> IO ()) -> (Profile -> Profile) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Profile
p ->
    let
      val :: ProfileFetch
val = CallId -> CallId -> Bool -> ProfileFetch
ProfileFetch CallId
cid (Env u w -> CallId
forall u w. Env u w -> CallId
memoKey Env u w
env) Bool
wasCached
      upd :: ProfileData -> ProfileData -> ProfileData
upd ProfileData
_ ProfileData
old = ProfileData
old { profileFetches = val : profileFetches old }

    in Profile
p { profile =
           HashMap.insertWith
             upd
             profKey
             (emptyProfileData { profileFetches = [val] })
             (profile p)
         }
  -- So we do not count the allocation overhead of addProfileFetch
  ProfileKey -> IO ()
setAllocationCounter ProfileKey
c