content: {-
    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 l (GenHaxl m) = GenHaxl $ \env ->
    if not $ testReportFlag ReportProfiling $ report $ flags env
       then m env
       else collectProfileData l m 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 mnPtr nPtr (GenHaxl m) = GenHaxl $ \env ->
    if not $ testReportFlag ReportProfiling $ report $ flags env
       then m env
       else collectProfileData
              (Text.unpackCString# mnPtr <> "." <> Text.unpackCString# nPtr)
              m 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 l m env = do
    let ProfileCurrent prevProfKey (prevProfLabel :| _) = profCurrent env
    if prevProfLabel == l
    then
      -- do not add a new label if we are recursing
      m env
    else do
      key <- atomicModifyIORef' (profRef env) $ \p ->
        case HashMap.lookup (l, prevProfKey) (profileTree p) of
          Just k -> (p, k)
          Nothing -> (p
            { profileTree = HashMap.insert
              (l, prevProfKey)
              (profileNextKey p)
              (profileTree p)
            , profileNextKey = profileNextKey p + 1 }, profileNextKey p)
      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 isCont env = do
    t0 <- getTimestamp
    a0 <- getAllocationCounter
    let
      ProfileCurrent caller stack = profCurrent env
      nextCurrent = ProfileCurrent
        { profCurrentKey = key
        , profLabelStack = l <| stack
        }
      runCont (GenHaxl h) = GenHaxl $ runProfileData l key h True
  
    r <- m 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 <- case r of
      Done !a -> return (Done a)
      Throw !e -> return (Throw e)
      Blocked ivar k -> return (Blocked ivar (Cont $ runCont (toHaxl k)))
  
    a1 <- getAllocationCounter
    t1 <- 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
    modifyProfileData env key caller (a0 - a1) (t1-t0) (if isCont then 0 else 1)
  
    -- So we do not count the allocation overhead of modifyProfileData
    setAllocationCounter a1
    return result
  {-# INLINE runProfileData #-}
  
  modifyProfileData
    :: Env u w
    -> ProfileKey
    -> ProfileKey
    -> AllocCount
    -> Microseconds
    -> LabelHitCount
    -> IO ()
  modifyProfileData env key caller allocs t labelIncrement = do
    modifyIORef' (profRef env) $ \ p ->
      p { profile =
            HashMap.insertWith updEntry key newEntry .
            HashMap.insertWith updCaller caller newCaller $
            profile p }
    where newEntry =
            emptyProfileData
              { profileAllocs = allocs
              , profileLabelHits = labelIncrement
              , profileTime = t
              }
          updEntry _ old =
            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 =
            emptyProfileData { profileAllocs = -allocs
                             , profileTime = -t
                             }
          updCaller _ old =
            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 m env = do
    t0 <- getTimestamp
    a0 <- getAllocationCounter
    r <- m env
    a1 <- getAllocationCounter
    t1 <- getTimestamp
    let
      allocs = a0 - a1
      t = t1 - t0
      newEntry = emptyProfileData
        { profileAllocs = allocs
        , profileTime = t
        }
      updEntry _ old = old
        { profileAllocs = profileAllocs old + allocs
        , profileTime = profileTime old + t
        }
      profKey = profCurrentKey (profCurrent env)
    modifyIORef' (profRef env) $ \ p ->
      p { profile =
           HashMap.insertWith updEntry profKey newEntry $
           profile p }
    -- So we do not count the allocation overhead of modifyProfileData
    setAllocationCounter a1
    return r
  {-# INLINE profileCont #-}
  
  incrementMemoHitCounterFor :: Env u w -> CallId -> Bool -> IO ()
  incrementMemoHitCounterFor env callId wasCached = do
    modifyIORef' (profRef env) $ \p ->  p {
      profile = HashMap.insertWith
                  upd
                  (profCurrentKey $ profCurrent env)
                  (emptyProfileData { profileMemos = [val] })
                  (profile p)
      }
    where
      val = ProfileMemo callId wasCached
      upd _ old = 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 env _req cid wasCached = do
    c <- getAllocationCounter
    let (ProfileCurrent profKey _) = profCurrent env
    modifyIORef' (profRef env) $ \ p ->
      let
        val = ProfileFetch cid (memoKey env) wasCached
        upd _ old = old { profileFetches = val : profileFetches old }
  
      in p { profile =
             HashMap.insertWith
               upd
               profKey
               (emptyProfileData { profileFetches = [val] })
               (profile p)
           }
    -- So we do not count the allocation overhead of addProfileFetch
    setAllocationCounter c
  