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 NoImplicitPrelude #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE CPP #-}
  
  module ProfileTests where
  
  import Haxl.Prelude
  
  import Haxl.Core
  import Haxl.Core.Monad
  import Haxl.Core.Stats
  
  import Test.HUnit
  
  import Control.DeepSeq (force)
  import Control.Exception (evaluate)
  import Data.Aeson
  import Data.IORef
  import qualified Data.HashMap.Strict as HashMap
  #if MIN_VERSION_aeson(2,0,0)
  import qualified Data.Aeson.KeyMap as KeyMap
  #else
  import qualified Data.HashMap.Strict as KeyMap
  #endif
  import Data.Int
  
  import TestTypes
  import TestUtils
  import WorkDataSource
  import SleepDataSource
  
  mkProfilingEnv :: IO HaxlEnv
  mkProfilingEnv = do
    env <- makeTestEnv False
    return env { flags = (flags env) { report = profilingReportFlags } }
  
  -- expects only one label to be shown
  labelToDataMap :: Profile -> HashMap.HashMap ProfileLabel ProfileData
  labelToDataMap Profile{..} = HashMap.fromList hashKeys
    where
      labelKeys = HashMap.fromList [
        (k, l) | ((l, _), k) <- HashMap.toList profileTree]
      hashKeys = [ (l, v)
        | (k, v) <- HashMap.toList profile
        , Just l <- [HashMap.lookup k labelKeys]]
  
  collectsdata :: Assertion
  collectsdata = do
    e <- mkProfilingEnv
    _x <- runHaxl e $
            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
                -- first sleep though in order to force a Blocked result
                sleep slp `andThen` case fromJSON <$> KeyMap.lookup "A" u of
                  Just (Success n) | sum [n .. 1000::Integer] > 0 -> return 5
                  _otherwise -> return (4::Int)
    profCopy <- readIORef (profRef e)
    let
      profData = profile profCopy
      labelKeys = HashMap.fromList [
        (l, k) | ((l, _), k) <- HashMap.toList (profileTree profCopy)]
      getData k = do
        k2 <- HashMap.lookup k labelKeys
        HashMap.lookup k2 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
        _otherwise -> False
    let fooParents = case HashMap.lookup "foo" labelKeys of
          Nothing -> []
          Just kfoo ->
            [ kparent
            | ((_, kparent), k) <- HashMap.toList (profileTree profCopy)
            , k == kfoo]
    assertEqual "foo's parent" 1 (length fooParents)
    assertEqual "foo's parent is bar" (Just (head fooParents)) $
      HashMap.lookup ("bar", 0) (profileTree profCopy)
  
  
  collectsLazyData :: Assertion
  collectsLazyData = do
    e <- mkProfilingEnv
    _x <- runHaxl e $ withLabel "bar" $ do
            u <- env userEnv
            withLabel "foo" $ do
               let start = if KeyMap.member "A" u
                           then 10
                           else 1
               return $ sum [start..10000::Integer]
    profCopy <- readIORef (profRef e)
    -- check the allocations are attributed to foo
    assertBool "foo has allocations" $
      case profileAllocs <$> HashMap.lookup "foo" (labelToDataMap profCopy) of
        Just x -> x > 10000
        Nothing -> False
  
  exceptions :: Assertion
  exceptions = do
    env <- mkProfilingEnv
    _x <- runHaxl env $
            withLabel "outer" $
              tryToHaxlException $ withLabel "inner" $
                unsafeLiftIO $ evaluate $ force (error "pure exception" :: Int)
    profData <- labelToDataMap <$> readIORef (profRef env)
    assertBool "inner label not added" $
      not $ HashMap.member "inner" profData
  
    env2 <- mkProfilingEnv
    _x <- runHaxl env2 $
            withLabel "outer" $
              tryToHaxlException $ withLabel "inner" $
                throw $ NotFound "haxl exception"
    profData <- labelToDataMap <$> readIORef (profRef env2)
    assertBool "inner label added" $
      HashMap.member "inner" profData
  
  
  -- Test that we correctly attribute work done in child threads when
  -- using BackgroundFetch to the caller of runHaxl. This is important
  -- for correct accounting when relying on allocation limits.
  threadAlloc :: Integer -> Assertion
  threadAlloc batches = do
    env' <- initEnv (stateSet mkWorkState stateEmpty) () :: IO (Env () ())
    let env = env'  { flags = (flags env') {
      report = setReportFlag ReportFetchStats defaultReportFlags } }
    a0 <- getAllocationCounter
    let
      wsize = 100000
      w = forM [wsize..(wsize+batches-1)] work
    _x <- runHaxl env $ sum <$> w
    a1 <- getAllocationCounter
    let
      lower = fromIntegral $ 1000000 * batches
      upper = fromIntegral $ 25000000 * batches
    assertBool "threadAlloc lower bound" $ (a0 - a1) > lower
    assertBool "threadAlloc upper bound" $ (a0 - a1) < upper
      -- the result was 16MB on 64-bit, or around 25KB if we miss the allocs
      -- in the child thread. For batched it should be similarly scaled.
      -- When we do not reset the counter for each batch was
      -- scaled again by number of batches.
  
    stats <- readIORef (statsRef env)
    assertEqual
      "threadAlloc: batches"
      [fromIntegral batches]
      (aggregateFetchBatches length stats)
    -- if we actually do more than 1 batch then the above test is not useful
  
  data MemoType = Global | Local
  
  -- Test that we correctly attribute memo work
  memos:: MemoType -> Assertion
  memos memoType = do
    env <- mkProfilingEnv
    let
      memoAllocs = 10000000 :: Int64
      doWork = unsafeLiftIO $ do
        a0 <- getAllocationCounter
        setAllocationCounter $ a0 - memoAllocs
        return (5 :: Int)
      mkWork
        | Global <- memoType = return (memo (1 :: Int) doWork)
        | Local <- memoType = memoize doWork
    _ <- runHaxl env $ do
      work <- mkWork
      andThen
        (withLabel "do" work)
        (withLabel "cached" work)
    profData <- labelToDataMap <$> readIORef (profRef env)
    case HashMap.lookup "do" profData of
      Nothing -> assertFailure "do not in data"
      Just ProfileData{..} -> do
        assertEqual "has correct memo id" profileMemos [ProfileMemo 1 False]
        assertBool "allocs are included in 'do'" (profileAllocs >= memoAllocs)
    case HashMap.lookup "cached" profData of
      Nothing -> assertFailure "cached not in data"
      Just ProfileData{..} -> do
        assertEqual "has correct memo id" profileMemos [ProfileMemo 1 True]
        assertBool "allocs are *not* included in 'cached'" (profileAllocs < 50000)
    (Stats memoStats) <- readIORef (statsRef env)
    assertEqual "exactly 1 memo/fetch" 1 (length memoStats)
    let memoStat = head memoStats
    putStrLn $ "memoStat=" ++ show memoStat
    assertEqual "correct call id" 1 (memoStatId memoStat)
    assertBool "allocs are big enough" $ memoSpace memoStat >= memoAllocs
    assertBool "allocs are not too big" $ memoSpace memoStat < memoAllocs + 100000
  
  
  tests = TestList
    [ TestLabel "collectsdata" $ TestCase collectsdata
    , TestLabel "collectsdata - lazy" $ TestCase collectsLazyData
    , TestLabel "exceptions" $ TestCase exceptions
    , TestLabel "threads" $ TestCase (threadAlloc 1)
    , TestLabel "threads with batch" $ TestCase (threadAlloc 50)
    , TestLabel "memos - Global" $ TestCase (memos Global)
    , TestLabel "memos - Local" $ TestCase (memos Local)
    ]
  