-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}

-- | Implementation of data-fetching operations.  Most users should
-- import "Haxl.Core" instead.
--
module Haxl.Core.Fetch
  ( dataFetch
  , dataFetchWithShow
  , dataFetchWithInsert
  , uncachedRequest
  , cacheResult
  , dupableCacheRequest
  , cacheResultWithShow
  , cacheRequest
  , performFetches
  , performRequestStore
  , ShowReq
  ) where

import Control.Concurrent.STM
import Control.Exception as Exception
import Control.Monad
import Data.Either
import Data.Hashable
import Data.IORef
import Data.Int
import Data.List
#if __GLASGOW_HASKELL__ < 804
import Data.Monoid
#endif
import Data.Proxy
import Data.Typeable
import Data.Text (Text)
import Data.Kind (Type)
import qualified Data.Text as Text
import Text.Printf
#ifdef PROFILING
import GHC.Stack
#endif

import Haxl.Core.DataSource
import Haxl.Core.DataCache as DataCache
import Haxl.Core.Exception
import Haxl.Core.Flags
import Haxl.Core.Monad
import Haxl.Core.Profile
import Haxl.Core.RequestStore
import Haxl.Core.ShowP
import Haxl.Core.Stats
import Haxl.Core.StateStore
import Haxl.Core.Util

-- -----------------------------------------------------------------------------
-- Data fetching and caching

-- | Possible responses when checking the cache.
data CacheResult u w a
  -- | The request hadn't been seen until now.
  = Uncached
       (ResultVar a)
       {-# UNPACK #-} !(IVar u w a)
       {-# UNPACK #-} !CallId

  -- | The request has been seen before, but its result has not yet been
  -- fetched.
  | CachedNotFetched
      {-# UNPACK #-} !(IVar u w a)
       {-# UNPACK #-} !CallId

  -- | The request has been seen before, and its result has already been
  -- fetched.
  | Cached (ResultVal a w)
           {-# UNPACK #-} !CallId

-- | Show functions for request and its result.
type ShowReq r a = (r a -> String, a -> String)

-- Note [showFn]
--
-- Occasionally, for tracing purposes or generating exceptions, we need to
-- call 'show' on the request in a place where we *cannot* have a Show
-- dictionary. (Because the function is a worker which is called by one of
-- the *WithShow variants that take explicit show functions via a ShowReq
-- argument.) None of the functions that does this is exported, so this is
-- hidden from the Haxl user.

cachedWithInsert
  :: forall r a u w.
     (DataSource u r, Typeable (r a))
  => (r a -> String)    -- See Note [showFn]
  -> (r a -> DataCacheItem u w a -> DataCache (DataCacheItem u w) -> IO ())
  -> Env u w -> r a -> IO (CacheResult u w a)
cachedWithInsert showFn insertFn env@Env{..} req = do
  let
    doFetch = do
      ivar <- newIVar
      k <- nextCallId env
      let !rvar = stdResultVar env ivar (Proxy :: Proxy r)
      insertFn req (DataCacheItem ivar k) dataCache
      return (Uncached rvar ivar k)
  mbRes <- DataCache.lookup req dataCache
  case mbRes of
    Nothing -> doFetch
    Just (DataCacheItem i@IVar{ivarRef = cr} k) -> do
      e <- readIORef cr
      case e of
        IVarEmpty _ -> do
          ivar <- withCurrentCCS i
          return (CachedNotFetched ivar k)
        IVarFull r -> do
          ifTrace flags 3 $ putStrLn $ case r of
            ThrowIO{} -> "Cached error: " ++ showFn req
            ThrowHaxl{} -> "Cached error: " ++ showFn req
            Ok{} -> "Cached request: " ++ showFn req
          return (Cached r k)

-- | Make a ResultVar with the standard function for sending a CompletionReq
-- to the scheduler. This is the function will be executed when the fetch
-- completes.
stdResultVar
  :: forall r a u w. (DataSourceName r, Typeable r)
  => Env u w
  -> IVar u w a
  -> Proxy r
  -> ResultVar a
stdResultVar Env{..} ivar p =
  mkResultVar $ \r isChildThread _ -> do
    allocs <- if isChildThread
      then
        -- In a child thread, return the current allocation counter too,
        -- for correct tracking of allocation.
        getAllocationCounter
      else
        return 0
    atomicallyOnBlocking
      (LogicBug (ReadingCompletionsFailedFetch (dataSourceName p))) $ do
      cs <- readTVar completions
      writeTVar completions (CompleteReq (eitherToResult r) ivar allocs : cs)
    -- Decrement the counter as request has finished. Do this after updating the
    -- completions TVar so that if the scheduler is tracking what was being
    -- waited on it gets a consistent view.
    ifReport flags ReportOutgoneFetches $
      atomicModifyIORef' submittedReqsRef (\m -> (subFromCountMap p 1 m, ()))
{-# INLINE stdResultVar #-}

-- | Record the call stack for a data fetch in the Stats.  Only useful
-- when profiling.
logFetch :: Env u w -> (r a -> String) -> r a -> CallId -> IO ()
#ifdef PROFILING
logFetch env showFn req fid = do
  ifReport (flags env) ReportFetchStack $ do
    stack <- currentCallStack
    modifyIORef' (statsRef env) $ \(Stats s) ->
      Stats (FetchCall (showFn req) stack fid : s)
#else
logFetch _ _ _ _ = return ()
#endif

calcFailure
  :: forall u req a . DataSource u req
  => u
  -> req a
  -> Either SomeException a
  -> FailureCount
calcFailure _u _r Right{} = mempty
calcFailure u r (Left e) = case classifyFailure u r e of
  StandardFailure -> mempty { failureCountStandard = 1 }
  IgnoredForStatsFailure -> mempty { failureCountIgnored = 1 }

addFallbackFetchStats
  :: forall u w req a . DataSource u req
  => Env u w
  -> CallId
  -> req a
  -> ResultVal a w
  -> IO ()
addFallbackFetchStats Env{..} fid req res = do
  bid <- atomicModifyIORef' statsBatchIdRef $ \x -> (x+1,x+1)
  start <- getTimestamp
  let
    dsName = dataSourceName (Proxy :: Proxy req)
    FailureCount{..} = case res of
      Ok{} -> mempty
      (ThrowHaxl e _) -> calcFailure userEnv req (Left e)
      (ThrowIO e) -> calcFailure userEnv req (Left e)
    this = FetchStats { fetchDataSource = dsName
                      , fetchBatchSize = 1
                      , fetchStart = start
                      , fetchDuration = 0
                      , fetchSpace = 0
                      , fetchFailures = failureCountStandard
                      , fetchIgnoredFailures = failureCountIgnored
                      , fetchBatchId = bid
                      , fetchIds = [fid] }
  atomicModifyIORef' statsRef $ \(Stats fs) -> (Stats (this : fs), ())

addFallbackResult
  :: Env u w
  -> ResultVal a w
  -> IVar u w a
  -> IO ()
addFallbackResult Env{..} res ivar = do
  atomicallyOnBlocking
    (LogicBug (ReadingCompletionsFailedFetch "addFallbackResult")) $ do
    cs <- readTVar completions
    writeTVar completions (CompleteReq res ivar 0 : cs)

-- | Performs actual fetching of data for a 'Request' from a 'DataSource'.
dataFetch :: (DataSource u r, Request r a) => r a -> GenHaxl u w a
dataFetch = dataFetchWithInsert show DataCache.insert

-- | Performs actual fetching of data for a 'Request' from a 'DataSource', using
-- the given show functions for requests and their results.
dataFetchWithShow
  :: (DataSource u r, Eq (r a), Hashable (r a), Typeable (r a))
  => ShowReq r a
  -> r a -> GenHaxl u w a
dataFetchWithShow (showReq, showRes) = dataFetchWithInsert showReq
  (DataCache.insertWithShow showReq showRes)

-- | Performs actual fetching of data for a 'Request' from a 'DataSource', using
-- the given function to insert requests in the cache.
dataFetchWithInsert
  :: forall u w r a
   . (DataSource u r, Eq (r a), Hashable (r a), Typeable (r a))
  => (r a -> String)    -- See Note [showFn]
  -> (r a -> DataCacheItem u w a -> DataCache (DataCacheItem u w) -> IO ())
  -> r a
  -> GenHaxl u w a
dataFetchWithInsert showFn insertFn req =
  GenHaxl $ \env@Env{..} -> do
  -- First, check the cache
  res <- cachedWithInsert showFn insertFn env req
  case res of
    -- This request has not been seen before
    Uncached rvar ivar fid -> do
      logFetch env showFn req fid
      ifProfiling flags $ addProfileFetch env req fid False
      --
      -- Check whether the data source wants to submit requests
      -- eagerly, or batch them up.
      --
      let
        blockedFetch = BlockedFetch req rvar
        blockedFetchI = BlockedFetchInternal fid
        submitFetch = do
          case schedulerHint userEnv :: SchedulerHint r of
            SubmitImmediately ->
              performFetches env [BlockedFetches [blockedFetch] [blockedFetchI]]
            TryToBatch ->
              -- add the request to the RequestStore and continue
              modifyIORef' reqStoreRef $ \bs ->
                addRequest blockedFetch blockedFetchI bs
          return $ Blocked ivar (Return ivar)

      -- if there is a fallback configured try that,
      -- else dispatch the fetch
      case dataCacheFetchFallback of
        Nothing -> submitFetch
        Just (DataCacheLookup dcl) -> do
          mbFallbackRes <- dcl req
          case mbFallbackRes of
            Nothing -> submitFetch
            Just fallbackRes -> do
              addFallbackResult env fallbackRes ivar
              ifReport flags ReportFetchStats $ addFallbackFetchStats
                env
                fid
                req
                fallbackRes
              return $ Blocked ivar (Return ivar)

    -- Seen before but not fetched yet.  We're blocked, but we don't have
    -- to add the request to the RequestStore.
    CachedNotFetched ivar fid -> do
      ifProfiling flags $ addProfileFetch env req fid True
      return $ Blocked ivar (Return ivar)

    -- Cached: either a result, or an exception
    Cached r fid -> do
      ifProfiling flags $ addProfileFetch env req fid True
      done env r

-- | A data request that is not cached.  This is not what you want for
-- normal read requests, because then multiple identical requests may
-- return different results, and this invalidates some of the
-- properties that we expect Haxl computations to respect: that data
-- fetches can be arbitrarily reordered, and identical requests can be
-- commoned up, for example.
--
-- 'uncachedRequest' is useful for performing writes, provided those
-- are done in a safe way - that is, not mixed with reads that might
-- conflict in the same Haxl computation.
--
-- if we are recording or running a test, we fallback to using dataFetch
-- This allows us to store the request in the cache when recording, which
-- allows a transparent run afterwards. Without this, the test would try to
-- call the datasource during testing and that would be an exception.
uncachedRequest
 :: forall a u w (r :: Type -> Type). (DataSource u r, Request r a)
 => r a -> GenHaxl u w a
uncachedRequest req = do
  flg <- env flags
  if recording flg /= 0
    then dataFetch req
    else GenHaxl $ \e@Env{..} -> do
      ivar <- newIVar
      k <- nextCallId e
      let !rvar = stdResultVar e ivar (Proxy :: Proxy r)
      modifyIORef' reqStoreRef $ \bs ->
        addRequest (BlockedFetch req rvar) (BlockedFetchInternal k) bs
      return $ Blocked ivar (Return ivar)

-- | Transparently provides caching. Useful for datasources that can
-- return immediately, but also caches values.  Exceptions thrown by
-- the IO operation (except for asynchronous exceptions) are
-- propagated into the Haxl monad and can be caught by 'catch' and
-- 'try'.
cacheResult :: Request r a => r a -> IO a -> GenHaxl u w a
cacheResult = cacheResultWithInsert show DataCache.insert

-- | Transparently provides caching in the same way as 'cacheResult', but uses
-- the given functions to show requests and their results.
cacheResultWithShow
  :: (Eq (r a), Hashable (r a), Typeable (r a))
  => ShowReq r a -> r a -> IO a -> GenHaxl u w a
cacheResultWithShow (showReq, showRes) = cacheResultWithInsert showReq
  (DataCache.insertWithShow showReq showRes)

-- Transparently provides caching, using the given function to insert requests
-- into the cache.
cacheResultWithInsert
  :: Typeable (r a)
  => (r a -> String)    -- See Note [showFn]
  -> (r a -> DataCacheItem u w a -> DataCache (DataCacheItem u w) -> IO ())
  -> r a
  -> IO a
  -> GenHaxl u w a
cacheResultWithInsert showFn insertFn req val = GenHaxl $ \env@Env{..} -> do
  mbRes <- DataCache.lookup req dataCache
  case mbRes of
    Nothing -> do
      let
        getResult = do
          eitherResult <- Exception.try val
          case eitherResult of
            Left e -> rethrowAsyncExceptions e
            _ -> return ()
          return $ eitherToResultThrowIO eitherResult
      -- if there is a fallback configured try that
      result <- case dataCacheFetchFallback of
        Nothing -> getResult
        Just (DataCacheLookup dcl) -> do
          mbFallbackRes <- dcl req
          maybe getResult return mbFallbackRes
      ivar <- newFullIVar result
      k <- nextCallId env
      insertFn req (DataCacheItem ivar k) dataCache
      done env result
    Just (DataCacheItem IVar{ivarRef = cr} _) -> do
      e <- readIORef cr
      case e of
        IVarEmpty _ -> raise env corruptCache
        IVarFull r -> done env r
  where
    corruptCache = DataSourceError $ Text.concat
      [ Text.pack (showFn req)
      , " has a corrupted cache value: these requests are meant to"
      , " return immediately without an intermediate value. Either"
      , " the cache was updated incorrectly, or you're calling"
      , " cacheResult on a query that involves a blocking fetch."
      ]

-- | Inserts a request/result pair into the cache. Throws an exception
-- if the request has already been issued, either via 'dataFetch' or
-- 'cacheRequest'.
--
-- This can be used to pre-populate the cache when running tests, to
-- avoid going to the actual data source and ensure that results are
-- deterministic.
--
cacheRequest
  :: Request req a => req a -> Either SomeException a -> GenHaxl u w ()
cacheRequest request result = GenHaxl $ \e@Env{..} -> do
  mbRes <- DataCache.lookup request dataCache
  case mbRes of
    Nothing -> do
      cr <- newFullIVar (eitherToResult result)
      k <- nextCallId e
      DataCache.insert request (DataCacheItem cr k) dataCache
      return (Done ())

    -- It is an error if the request is already in the cache.
    -- We can't test whether the cached result is the same without adding an
    -- Eq constraint, and we don't necessarily have Eq for all results.
    _other -> raise e $
      DataSourceError "cacheRequest: request is already in the cache"

-- | Similar to @cacheRequest@ but doesn't throw an exception if the key
-- already exists in the cache.
-- If this function is called twice to cache the same Haxl request, the first
-- value will be discarded and overwritten with the second value.
-- Useful e.g. for unit tests
dupableCacheRequest
  :: Request req a => req a -> Either SomeException a -> GenHaxl u w ()
dupableCacheRequest request result = GenHaxl $ \e@Env{..} -> do
  cr <- newFullIVar (eitherToResult result)
  k <- nextCallId e
  DataCache.insert request (DataCacheItem cr k) dataCache
  return (Done ())

performRequestStore
   :: forall u w. Env u w -> RequestStore u -> IO ()
performRequestStore env reqStore =
  performFetches env (contents reqStore)

-- | Issues a batch of fetches in a 'RequestStore'. After
-- 'performFetches', all the requests in the 'RequestStore' are
-- complete, and all of the 'ResultVar's are full.
performFetches
  :: forall u w. Env u w -> [BlockedFetches u] -> IO ()
performFetches env@Env{flags=f, statsRef=sref, statsBatchIdRef=sbref} jobs = do
  t0 <- getTimestamp

  ifTrace f 3 $
    forM_ jobs $ \(BlockedFetches reqs _) ->
      forM_ reqs $ \(BlockedFetch r _) -> putStrLn (showp r)

  let
    applyFetch i bfs@(BlockedFetches (reqs :: [BlockedFetch r]) _) =
      case stateGet (states env) of
        Nothing ->
          return (FetchToDo reqs (SyncFetch (mapM_ (setError e))))
         where
           e :: ShowP req => req a -> DataSourceError
           e req = DataSourceError $ "data source not initialized: " <> dsName
                  <> ": "
                  <> Text.pack (showp req)
        Just state ->
          return $ FetchToDo reqs
            $ (if testReportFlag ReportFetchStats $ report f
                then wrapFetchInStats
                        (userEnv env)
                        sref
                        sbref
                        dsName
                        (length reqs)
                        bfs
                else id)
            $ wrapFetchInTrace i (length reqs) dsName
            $ wrapFetchInCatch reqs
            $ fetch state f (userEnv env)
      where
        dsName = dataSourceName (Proxy :: Proxy r)

  fetches <- zipWithM applyFetch [0..] jobs

  scheduleFetches fetches (submittedReqsRef env) (flags env)

  t1 <- getTimestamp
  let roundtime = fromIntegral (t1 - t0) / 1000000 :: Double

  ifTrace f 1 $
    printf "Batch data fetch done (%.4fs)\n" (realToFrac roundtime :: Double)

data FetchToDo where
  FetchToDo
    :: forall (req :: Type -> Type). (DataSourceName req, Typeable req)
    => [BlockedFetch req] -> PerformFetch req -> FetchToDo

-- Catch exceptions arising from the data source and stuff them into
-- the appropriate requests.  We don't want any exceptions propagating
-- directly from the data sources, because we want the exception to be
-- thrown by dataFetch instead.
--
wrapFetchInCatch :: [BlockedFetch req] -> PerformFetch req -> PerformFetch req
wrapFetchInCatch reqs fetch =
  case fetch of
    SyncFetch f ->
      SyncFetch $ \reqs -> f reqs `Exception.catch` handler
    AsyncFetch f ->
      AsyncFetch $ \reqs io -> f reqs io `Exception.catch` handler
      -- this might be wrong: if the outer 'fio' throws an exception,
      -- then we don't know whether we have executed the inner 'io' or
      -- not.  If not, then we'll likely get some errors about "did
      -- not set result var" later, because we haven't executed some
      -- data fetches.  But we can't execute 'io' in the handler,
      -- because we might have already done it.  It isn't possible to
      -- do it completely right here, so we have to rely on data
      -- sources themselves to catch (synchronous) exceptions.  Async
      -- exceptions aren't a problem because we're going to rethrow
      -- them all the way to runHaxl anyway.
    BackgroundFetch f ->
      BackgroundFetch $ \reqs -> f reqs `Exception.catch` handler
  where
    handler :: SomeException -> IO ()
    handler e = do
      rethrowAsyncExceptions e
      mapM_ (forceError e) reqs

    -- Set the exception even if the request already had a result.
    -- Otherwise we could be discarding an exception.
    forceError e (BlockedFetch _ rvar) =
      putResult rvar (except e)

data FailureCount = FailureCount
  { failureCountStandard :: {-# UNPACK #-} !Int
  , failureCountIgnored :: {-# UNPACK #-} !Int
  }

#if __GLASGOW_HASKELL__ >= 804
instance Semigroup FailureCount where
  FailureCount s1 i1 <> FailureCount s2 i2 = FailureCount (s1+s2) (i1+i2)
#endif

instance Monoid FailureCount where
  mempty = FailureCount 0 0
#if __GLASGOW_HASKELL__ < 804
  mappend (FailureCount s1 i1) (FailureCount s2 i2)
    = FailureCount (s1+s2) (i1+i2)
#endif

wrapFetchInStats
  :: DataSource u req
  => u
  -> IORef Stats
  -> IORef Int
  -> Text
  -> Int
  -> BlockedFetches u
  -> PerformFetch req
  -> PerformFetch req
wrapFetchInStats
  u
  !statsRef
  !batchIdRef
  dataSource
  batchSize
  (BlockedFetches _reqs reqsI)
  perform = do
  case perform of
    SyncFetch f ->
      SyncFetch $ \reqs -> do
        bid <- newBatchId
        fail_ref <- newIORef mempty
        (t0,t,alloc,_) <- statsForIO (f (map (addFailureCount u fail_ref)
          (reqsWithFetchDsStats bid reqs)))
        failures <- readIORef fail_ref
        updateFetchStats bid allFids t0 t alloc batchSize failures
    AsyncFetch f -> do
      AsyncFetch $ \reqs inner -> do
        bid <- newBatchId
        inner_r <- newIORef (0, 0)
        fail_ref <- newIORef mempty
        let inner' = do
              (_,t,alloc,_) <- statsForIO inner
              writeIORef inner_r (t,alloc)
            reqs' = map (addFailureCount u fail_ref) reqs
            reqs'' = reqsWithFetchDsStats bid reqs'
        (t0, totalTime, totalAlloc, _) <- statsForIO (f reqs'' inner')
        (innerTime, innerAlloc) <- readIORef inner_r
        failures <- readIORef fail_ref
        updateFetchStats bid allFids t0 (totalTime - innerTime)
          (totalAlloc - innerAlloc) batchSize failures
    BackgroundFetch io -> do
      BackgroundFetch $ \reqs -> do
        bid <- newBatchId
        startTime <- getTimestamp
        io (reqsWithFetchDsStats bid
          (zipWith (addTimer u bid startTime) reqs reqsI))
  where
    allFids = map (\(BlockedFetchInternal k) -> k) reqsI
    newBatchId = atomicModifyIORef' batchIdRef $ \x -> (x+1,x+1)
    statsForIO io = do
      prevAlloc <- getAllocationCounter
      (t0,t,a) <- time io
      postAlloc <- getAllocationCounter
      return (t0,t, fromIntegral $ prevAlloc - postAlloc, a)
    reqsWithFetchDsStats = \bid reqs
      -> zipWith (addFetchDatasourceStats bid) reqs reqsI
    addTimer
      u
      bid
      t0
      (BlockedFetch req (ResultVar fn))
      (BlockedFetchInternal fid) =
        BlockedFetch req $ ResultVar $ \result isChildThread stats -> do
          t1 <- getTimestamp
          -- We cannot measure allocation easily for BackgroundFetch. Here we
          -- just attribute all allocation to the last
          -- `putResultFromChildThread` and use 0 for the others.
          -- While the individual allocations may not be correct,
          -- the total sum and amortized allocation are still meaningful.
          -- see Note [tracking allocation in child threads]
          allocs <- if isChildThread then getAllocationCounter else return 0
          updateFetchStats bid [fid] t0 (t1 - t0)
            (negate allocs)
            1 -- batch size: we don't know if this is a batch or not
            (calcFailure u req result) -- failures
          fn result isChildThread stats

    addFetchDatasourceStats
      :: Int
      -> BlockedFetch r
      -> BlockedFetchInternal
      -> BlockedFetch r
    addFetchDatasourceStats bid
      (BlockedFetch req (ResultVar fn))
      (BlockedFetchInternal fid) = BlockedFetch req $ ResultVar
        $ \result isChildThread stats -> do
          let mkStats dss = FetchDataSourceStats
                { fetchDsStatsCallId = fid
                , fetchDsStatsDataSource = dataSource
                , fetchDsStatsStats = dss
                , fetchBatchId = bid
                }
          case stats of
            Just dss -> atomicModifyIORef' statsRef
              $ \(Stats fs) -> (Stats (mkStats dss : fs), ())
            Nothing -> return ()
          fn result isChildThread stats

    updateFetchStats
      :: Int
      -> [CallId]
      -> Timestamp
      -> Microseconds
      -> Int64
      -> Int
      -> FailureCount
      -> IO ()
    updateFetchStats bid fids start time space batch FailureCount{..} = do
      let this = FetchStats { fetchDataSource = dataSource
                            , fetchBatchSize = batch
                            , fetchStart = start
                            , fetchDuration = time
                            , fetchSpace = space
                            , fetchFailures = failureCountStandard
                            , fetchIgnoredFailures = failureCountIgnored
                            , fetchBatchId = bid
                            , fetchIds = fids }
      atomicModifyIORef' statsRef $ \(Stats fs) -> (Stats (this : fs), ())

    addFailureCount :: DataSource u r
      => u -> IORef FailureCount -> BlockedFetch r -> BlockedFetch r
    addFailureCount u ref (BlockedFetch req (ResultVar fn)) =
      BlockedFetch req $ ResultVar $ \result isChildThread stats -> do
        let addFailures r = (r <> calcFailure u req result, ())
        when (isLeft result) $ atomicModifyIORef' ref addFailures
        fn result isChildThread stats

wrapFetchInTrace
  :: Int
  -> Int
  -> Text
  -> PerformFetch req
  -> PerformFetch req
#ifdef EVENTLOG
wrapFetchInTrace i n dsName f =
  case f of
    SyncFetch io -> SyncFetch (wrapF "Sync" io)
    AsyncFetch fio -> AsyncFetch (wrapF "Async" . fio . unwrapF "Async")
  where
    d = Text.unpack dsName
    wrapF :: String -> IO a -> IO a
    wrapF ty = bracket_ (traceEventIO $ printf "START %d %s (%d %s)" i d n ty)
                        (traceEventIO $ printf "STOP %d %s (%d %s)" i d n ty)
    unwrapF :: String -> IO a -> IO a
    unwrapF ty = bracket_ (traceEventIO $ printf "STOP %d %s (%d %s)" i d n ty)
                          (traceEventIO $ printf "START %d %s (%d %s)" i d n ty)
#else
wrapFetchInTrace _ _ _ f = f
#endif

time :: IO a -> IO (Timestamp,Microseconds,a)
time io = do
  t0 <- getTimestamp
  a <- io
  t1 <- getTimestamp
  return (t0, t1 - t0, a)

-- | Start all the async fetches first, then perform the sync fetches before
-- getting the results of the async fetches.
scheduleFetches :: [FetchToDo] -> IORef ReqCountMap -> Flags -> IO ()
scheduleFetches fetches ref flags = do
  -- update ReqCountmap for these fetches
  ifReport flags ReportOutgoneFetches $ sequence_
    [ atomicModifyIORef' ref $
        \m -> (addToCountMap (Proxy :: Proxy r) (length reqs) m, ())
    | FetchToDo (reqs :: [BlockedFetch r]) _f <- fetches
    ]

  ifTrace flags 1 $ printf "Batch data fetch round: %s\n" $
     intercalate (", "::String) $
        map (\(c, n, ds) -> printf "%s %s %d" n ds c) stats

  fully_async_fetches
  async_fetches sync_fetches
 where

  fetchName :: forall req . PerformFetch req -> String
  fetchName (BackgroundFetch _) = "background"
  fetchName (AsyncFetch _) = "async"
  fetchName (SyncFetch _) = "sync"

  srcName :: forall req . (DataSourceName req) => [BlockedFetch req] -> String
  srcName _ = Text.unpack $ dataSourceName (Proxy :: Proxy req)

  stats = [(length reqs, fetchName f, srcName reqs)
          | FetchToDo reqs f <- fetches]

  fully_async_fetches :: IO ()
  fully_async_fetches = sequence_
    [f reqs | FetchToDo reqs (BackgroundFetch f) <- fetches]

  async_fetches :: IO () -> IO ()
  async_fetches = compose
    [f reqs | FetchToDo reqs (AsyncFetch f) <- fetches]

  sync_fetches :: IO ()
  sync_fetches = sequence_
    [f reqs | FetchToDo reqs (SyncFetch f) <- fetches]

-- | An exception thrown when reading from datasources fails
data ReadingCompletionsFailedFetch = ReadingCompletionsFailedFetch Text
  deriving Show

instance Exception ReadingCompletionsFailedFetch


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE BangPatterns #-}

-- |
-- The 'Flags' type and related functions.  This module is provided
-- for access to Haxl internals only; most users should import
-- "Haxl.Core" instead.
--
module Haxl.Core.Flags
  (
    -- * Report flags
    ReportFlag(..)
  , ReportFlags
  , defaultReportFlags
  , profilingReportFlags
  , setReportFlag
  , clearReportFlag
  , testReportFlag
    -- * Flags
  , Flags(..)
  , defaultFlags
  , ifTrace
  , ifReport
  , ifProfiling
  ) where

import Control.Monad
import Data.Bits
import Data.List (foldl')
import Text.Printf (printf)

-- ---------------------------------------------------------------------------
-- ReportFlags
data ReportFlag
  = ReportOutgoneFetches  -- ^ outgone fetches, for debugging eg: timeouts
  | ReportFetchStats  -- ^ data fetch stats & errors
  | ReportProfiling   -- ^ enabling label stack and profiling
  | ReportExceptionLabelStack  -- ^ include label stack in HaxlException
  | ReportFetchStack  -- ^ log cost-center stack traces of dataFetch calls
  deriving (Bounded, Enum, Eq, Show)

profilingDependents :: [ReportFlag]
profilingDependents =
  [ ReportExceptionLabelStack
  , ReportFetchStack
  ]

newtype ReportFlags = ReportFlags Int

instance Show ReportFlags where
  show (ReportFlags fs) = printf "%0*b" (fromEnum maxReportFlag + 1) fs
    where
      maxReportFlag = maxBound :: ReportFlag

defaultReportFlags :: ReportFlags
defaultReportFlags = ReportFlags 0

profilingReportFlags :: ReportFlags
profilingReportFlags = foldl' (flip setReportFlag) defaultReportFlags
  [ ReportOutgoneFetches
  , ReportFetchStats
  , ReportProfiling
  ]

setReportFlag :: ReportFlag -> ReportFlags -> ReportFlags
setReportFlag f (ReportFlags fs) =
  ReportFlags $ setDependencies $ setBit fs $ fromEnum f
  where
    setDependencies
      | f `elem` profilingDependents = flip setBit $ fromEnum ReportProfiling
      | otherwise = id

clearReportFlag :: ReportFlag -> ReportFlags -> ReportFlags
clearReportFlag f (ReportFlags fs) =
  ReportFlags $ clearDependents $ clearBit fs $ fromEnum f
  where
    clearDependents z = case f of
      ReportProfiling -> foldl' clearBit z $ map fromEnum profilingDependents
      _ -> z

{-# INLINE testReportFlag #-}
testReportFlag :: ReportFlag -> ReportFlags -> Bool
testReportFlag !f (ReportFlags !fs) = testBit fs $ fromEnum f

-- ---------------------------------------------------------------------------
-- Flags

-- | Flags that control the operation of the engine.
data Flags = Flags
  { trace :: {-# UNPACK #-} !Int
    -- ^ Tracing level (0 = quiet, 3 = very verbose).
  , report :: {-# UNPACK #-} !ReportFlags
    -- ^ Report flags
  , caching :: {-# UNPACK #-} !Int
    -- ^ Non-zero if caching is enabled.  If caching is disabled, then
    -- we still do batching and de-duplication, but do not cache
    -- results.
  , recording :: {-# UNPACK #-} !Int
    -- ^ Non-zero if recording is enabled. This allows tests to record cache
    -- calls for datasources by making uncachedRequest behave like dataFetch
  }

defaultFlags :: Flags
defaultFlags = Flags
  { trace = 0
  , report = defaultReportFlags
  , caching = 1
  , recording = 0
  }

-- | Runs an action if the tracing level is above the given threshold.
ifTrace :: Monad m => Flags -> Int -> m a -> m ()
ifTrace flags i = when (trace flags >= i) . void

-- | Runs an action if the ReportFlag is set.
ifReport :: Monad m => Flags -> ReportFlag -> m a -> m ()
ifReport flags i = when (testReportFlag i $ report flags) . void

ifProfiling :: Monad m => Flags -> m a -> m ()
ifProfiling flags = ifReport flags ReportProfiling


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

-- |
-- Memoization support. This module is provided for access to Haxl
-- internals only; most users should import "Haxl.Core" instead.
--
module Haxl.Core.Memo
  (
    -- * Basic memoization
    cachedComputation
  , preCacheComputation

    -- * High-level memoization
  , memo
  , memoFingerprint
  , MemoFingerprintKey(..)
  , memoize, memoize1, memoize2
  , memoUnique

    -- * Local memoization
  , MemoVar
  , newMemo
  , newMemoWith
  , prepareMemo
  , runMemo
  ) where

import Control.Exception as Exception hiding (throw)
import Data.IORef
import qualified Data.HashMap.Strict as HashMap
import Data.Text (Text)
import Data.Typeable
import Data.Hashable
import Data.Int
import Data.Word

import GHC.Prim (Addr#)

import Haxl.Core.Exception
import Haxl.Core.DataCache as DataCache
import Haxl.Core.Flags
import Haxl.Core.Monad
import Haxl.Core.Stats
import Haxl.Core.Profile
import Haxl.Core.Util (trace_)

-- -----------------------------------------------------------------------------
-- Memoization

-- | 'cachedComputation' memoizes a Haxl computation.  The key is a
-- request.
--
-- /Note:/ These cached computations will /not/ be included in the output
-- of 'dumpCacheAsHaskell'.
--
cachedComputation
   :: forall req u w a.
      ( Eq (req a)
      , Hashable (req a)
      , Typeable (req a)
      , Monoid w)
   => req a -> GenHaxl u w a -> GenHaxl u w a

cachedComputation req haxl = GenHaxl $ \env@Env{..} -> do
  mbRes <- DataCache.lookup req memoCache
  case mbRes of
    Just (DataCacheItem ivar k) -> do
      ifProfiling flags $ do
        incrementMemoHitCounterFor env k True
      unHaxl (getIVarWithWrites ivar) env
    Nothing -> do
      ivar <- newIVar
      k <- nextCallId env
      -- no need to incremenetMemoHitCounter as execMemo will do it
      DataCache.insertNotShowable req (DataCacheItem ivar k) memoCache
      execMemoNowProfiled env haxl ivar k

-- | Like 'cachedComputation', but fails if the cache is already
-- populated.
--
-- Memoization can be (ab)used to "mock" a cached computation, by
-- pre-populating the cache with an alternative implementation. In
-- that case we don't want the operation to populate the cache to
-- silently succeed if the cache is already populated.
--
preCacheComputation
  :: forall req u w a.
     ( Eq (req a)
     , Hashable (req a)
     , Typeable (req a)
     , Monoid w)
  => req a -> GenHaxl u w a -> GenHaxl u w a
preCacheComputation req haxl = GenHaxl $ \env@Env{..} -> do
  mbRes <- DataCache.lookup req memoCache
  case mbRes of
    Just _ -> return $ Throw $ toException $ InvalidParameter
      "preCacheComputation: key is already cached"
    Nothing -> do
      ivar <- newIVar
      k <- nextCallId env
      DataCache.insertNotShowable req (DataCacheItem ivar k) memoCache
      execMemoNowProfiled env haxl ivar k

-- -----------------------------------------------------------------------------
-- Memoization

newtype MemoVar u w a = MemoVar (IORef (MemoStatus u w a))

data MemoStatus u w a
  = MemoEmpty
  | MemoReady (GenHaxl u w a) CallId
  | MemoRun {-# UNPACK #-} !(IVar u w a) {-# UNPACK #-} !CallId

-- | Create a new @MemoVar@ for storing a memoized computation. The created
-- @MemoVar@ is initially empty, not tied to any specific computation. Running
-- this memo (with @runMemo@) without preparing it first (with @prepareMemo@)
-- will result in an exception.
newMemo :: GenHaxl u w (MemoVar u w a)
newMemo = unsafeLiftIO $ MemoVar <$> newIORef MemoEmpty

-- | Store a computation within a supplied @MemoVar@. Any memo stored within the
-- @MemoVar@ already (regardless of completion) will be discarded, in favor of
-- the supplied computation. A @MemoVar@ must be prepared before it is run.
prepareMemo :: MemoVar u w a -> GenHaxl u w a -> GenHaxl u w ()
prepareMemo (MemoVar memoRef) memoCmp
  = GenHaxl $ \env -> do
      k <- nextCallId env
      writeIORef memoRef (MemoReady memoCmp k)
      return (Done ())

-- | Convenience function, combines @newMemo@ and @prepareMemo@.
newMemoWith :: GenHaxl u w a -> GenHaxl u w (MemoVar u w a)
newMemoWith memoCmp = do
  memoVar <- newMemo
  prepareMemo memoVar memoCmp
  return memoVar

-- | Continue the memoized computation within a given @MemoVar@.
-- Notes:
--
--   1. If the memo contains a complete result, return that result.
--   2. If the memo contains an in-progress computation, continue it as far as
--      possible for this round.
--   3. If the memo is empty (it was not prepared), throw an error.
--
-- For example, to memoize the computation @one@ given by:
--
-- > one :: Haxl Int
-- > one = return 1
--
-- use:
--
-- > do
-- >   oneMemo <- newMemoWith one
-- >   let memoizedOne = runMemo aMemo one
-- >   oneResult <- memoizedOne
--
-- To memoize mutually dependent computations such as in:
--
-- > h :: Haxl Int
-- > h = do
-- >   a <- f
-- >   b <- g
-- >   return (a + b)
-- >  where
-- >   f = return 42
-- >   g = succ <$> f
--
-- without needing to reorder them, use:
--
-- > h :: Haxl Int
-- > h = do
-- >   fMemoRef <- newMemo
-- >   gMemoRef <- newMemo
-- >
-- >   let f = runMemo fMemoRef
-- >       g = runMemo gMemoRef
-- >
-- >   prepareMemo fMemoRef $ return 42
-- >   prepareMemo gMemoRef $ succ <$> f
-- >
-- >   a <- f
-- >   b <- g
-- >   return (a + b)
--
runMemo :: Monoid w => MemoVar u w a -> GenHaxl u w a
runMemo (MemoVar memoRef) = GenHaxl $ \env -> do
  stored <- readIORef memoRef
  case stored of
    -- Memo was not prepared first; throw an exception.
    MemoEmpty -> trace_ "MemoEmpty " $
      raise env $ CriticalError "Attempting to run empty memo."
    -- Memo has been prepared but not run yet
    MemoReady cont k -> trace_ "MemoReady" $ do
      ivar <- newIVar
      writeIORef memoRef (MemoRun ivar k)
      execMemoNowProfiled env cont ivar k
    -- The memo has already been run, get (or wait for) for the result
    MemoRun ivar k -> trace_ "MemoRun" $ do
      ifProfiling (flags env) $ do
        incrementMemoHitCounterFor env k True
      unHaxl (getIVarWithWrites ivar) env

execMemoNowProfiled
  :: Monoid w
  => Env u w
  -> GenHaxl u w a
  -> IVar u w a
  -> CallId
  -> IO (Result u w a)
execMemoNowProfiled envOuter cont ivar cid =
  if not $ testReportFlag ReportProfiling $ report $ flags envOuter
  then execMemoNow envOuter cont ivar
  else do
    incrementMemoHitCounterFor envOuter cid False
    unHaxl
      (collectMemoData 0 $ GenHaxl $ \e -> execMemoNow e cont ivar)
      envOuter
  where
    addStats :: Env u w -> Int64 -> IO ()
    addStats env acc = modifyIORef' (statsRef env) $ \(Stats s) ->
      Stats (MemoCall cid acc : s)
    collectMemoData :: Int64 -> GenHaxl u w a -> GenHaxl u w a
    collectMemoData acc f = GenHaxl $ \env -> do
      a0 <- getAllocationCounter
      r <- unHaxl f env{memoKey=cid}
      a1 <- getAllocationCounter
      let newTotal = acc + (a0 - a1)
      ret <- case r of
        Done a -> do addStats env newTotal; return (Done a)
        Throw e -> do addStats env newTotal; return (Throw e)
        Blocked ivar k ->
          return (Blocked ivar (Cont (collectMemoData newTotal (toHaxl k))))
      setAllocationCounter a1
      return ret

execMemoNow :: Monoid w => Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
execMemoNow env cont ivar = do
  wlogs <- newIORef mempty
  let
    !menv = env { writeLogsRef = wlogs }
    -- use an env with empty writes, so we can memoize the extra
    -- writes done as part of 'cont'
  r <- Exception.try $ unHaxl cont menv

  case r of
    Left e -> trace_ ("execMemoNow: Left " ++ show e) $ do
      rethrowAsyncExceptions e
      putIVar ivar (ThrowIO e) env
      throwIO e
    Right (Done a) -> trace_ "execMemoNow: Done" $ do
      wt <- readIORef wlogs
      putIVar ivar (Ok a (Just wt)) env
      modifyIORef' (writeLogsRef env) (<> wt)
      return (Done a)
    Right (Throw ex) -> trace_ ("execMemoNow: Throw" ++ show ex) $ do
      wt <- readIORef wlogs
      putIVar ivar (ThrowHaxl ex (Just wt)) env
      modifyIORef' (writeLogsRef env) (<> wt)
      return (Throw ex)
    Right (Blocked ivar' cont) -> trace_ "execMemoNow: Blocked" $ do
      -- We "block" this memoized computation in the new environment 'menv', so
      -- that when it finishes, we can store all the write logs from the env
      -- in the IVar.
      addJob menv (toHaxl cont) ivar ivar'
      -- Now we call @getIVarWithWrites@ to populate the writes in the original
      -- environment 'env'.
      return (Blocked ivar (Cont (getIVarWithWrites ivar)))

-- -----------------------------------------------------------------------------
-- 1-ary and 2-ary memo functions

newtype MemoVar1 u w a b = MemoVar1 (IORef (MemoStatus1 u w a b))
newtype MemoVar2 u w a b c = MemoVar2 (IORef (MemoStatus2 u w a b c))

data MemoStatus1 u w a b
  = MemoEmpty1
  | MemoTbl1 (a -> GenHaxl u w b) (HashMap.HashMap a (MemoVar u w b))

data MemoStatus2 u w a b c
  = MemoEmpty2
  | MemoTbl2
      (a -> b -> GenHaxl u w c)
      (HashMap.HashMap a (HashMap.HashMap b (MemoVar u w c)))

newMemo1 :: GenHaxl u w (MemoVar1 u w a b)
newMemo1 = unsafeLiftIO $ MemoVar1 <$> newIORef MemoEmpty1

newMemoWith1 :: (a -> GenHaxl u w b) -> GenHaxl u w (MemoVar1 u w a b)
newMemoWith1 f = newMemo1 >>= \r -> prepareMemo1 r f >> return r

prepareMemo1 :: MemoVar1 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
prepareMemo1 (MemoVar1 r) f
  = unsafeLiftIO $ writeIORef r (MemoTbl1 f HashMap.empty)

runMemo1 :: (Eq a, Hashable a, Monoid w) => MemoVar1 u w a b -> a -> GenHaxl u w b
runMemo1 (MemoVar1 r) k = unsafeLiftIO (readIORef r) >>= \case
  MemoEmpty1 -> throw $ CriticalError "Attempting to run empty memo."
  MemoTbl1 f h -> case HashMap.lookup k h of
    Nothing -> do
      x <- newMemoWith (f k)
      unsafeLiftIO $ writeIORef r (MemoTbl1 f (HashMap.insert k x h))
      runMemo x
    Just v -> runMemo v

newMemo2 :: GenHaxl u w (MemoVar2 u w a b c)
newMemo2 = unsafeLiftIO $ MemoVar2 <$> newIORef MemoEmpty2

newMemoWith2 :: (a -> b -> GenHaxl u w c) -> GenHaxl u w (MemoVar2 u w a b c)
newMemoWith2 f = newMemo2 >>= \r -> prepareMemo2 r f >> return r

prepareMemo2 :: MemoVar2 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w ()
prepareMemo2 (MemoVar2 r) f
  = unsafeLiftIO $ writeIORef r (MemoTbl2 f HashMap.empty)

runMemo2 :: (Eq a, Hashable a, Eq b, Hashable b, Monoid w)
         => MemoVar2 u w a b c
         -> a -> b -> GenHaxl u w c
runMemo2 (MemoVar2 r) k1 k2 = unsafeLiftIO (readIORef r) >>= \case
  MemoEmpty2 -> throw $ CriticalError "Attempting to run empty memo."
  MemoTbl2 f h1 -> case HashMap.lookup k1 h1 of
    Nothing -> do
      v <- newMemoWith (f k1 k2)
      unsafeLiftIO $ writeIORef r
        (MemoTbl2 f (HashMap.insert k1 (HashMap.singleton k2 v) h1))
      runMemo v
    Just h2 -> case HashMap.lookup k2 h2 of
      Nothing -> do
        v <- newMemoWith (f k1 k2)
        unsafeLiftIO $ writeIORef r
          (MemoTbl2 f (HashMap.insert k1 (HashMap.insert k2 v h2) h1))
        runMemo v
      Just v -> runMemo v

-- -----------------------------------------------------------------------------
-- A key type that can be used for memoizing computations by a Text key

-- | Memoize a computation using an arbitrary key.  The result will be
-- calculated once; the second and subsequent time it will be returned
-- immediately.  It is the caller's responsibility to ensure that for
-- every two calls @memo key haxl@, if they have the same @key@ then
-- they compute the same result.
memo
  :: (Typeable a, Typeable k, Hashable k, Eq k, Monoid w)
  => k -> GenHaxl u w a -> GenHaxl u w a
memo key = cachedComputation (MemoKey key)

{-# RULES
"memo/Text" memo = memoText :: (Typeable a, Monoid w) =>
            Text -> GenHaxl u w a -> GenHaxl u w a
 #-}

{-# NOINLINE memo #-}

-- | Memoize a computation using its location and a Fingerprint. This ensures
-- uniqueness across computations.
memoUnique
  :: (Typeable a, Typeable k, Hashable k, Eq k, Monoid w)
  => MemoFingerprintKey a -> Text -> k -> GenHaxl u w a -> GenHaxl u w a
memoUnique fp label key = withLabel label . memo (fp, key)

{-# NOINLINE memoUnique #-}

data MemoKey k a where
  MemoKey :: (Typeable k, Hashable k, Eq k) => k -> MemoKey k a
  deriving Typeable

deriving instance Eq (MemoKey k a)

instance Hashable (MemoKey k a) where
  hashWithSalt s (MemoKey t) = hashWithSalt s t

-- An optimised memo key for Text keys.  This is used automatically
-- when the key is Text, due to the RULES pragma above.

data MemoTextKey a where
  MemoText :: Text -> MemoTextKey a
  deriving Typeable

deriving instance Eq (MemoTextKey a)

instance Hashable (MemoTextKey a) where
  hashWithSalt s (MemoText t) = hashWithSalt s t

memoText :: (Typeable a, Monoid w) => Text -> GenHaxl u w a -> GenHaxl u w a
memoText key = withLabel key . cachedComputation (MemoText key)

-- | A memo key derived from a 128-bit MD5 hash.  Do not use this directly,
-- it is for use by automatically-generated memoization.
data MemoFingerprintKey a where
  MemoFingerprintKey
    :: {-# UNPACK #-} !Word64
    -> {-# UNPACK #-} !Word64
    -> Addr# -> Addr#
    -> MemoFingerprintKey a
  deriving Typeable

deriving instance Eq (MemoFingerprintKey a)

instance Hashable (MemoFingerprintKey a) where
  hashWithSalt s (MemoFingerprintKey x _ _ _) =
    hashWithSalt s (fromIntegral x :: Int)

-- This is optimised for cheap call sites: when we have a call
--
--   memoFingerprint (MemoFingerprintKey 1234 5678 "module"# "name"#) e
--
-- then the MemoFingerprintKey constructor will be statically
-- allocated (with two 64-bit fields and pointers to cstrings for the names),
-- and shared by all calls to memo. So the memo call will not allocate,
-- unlike memoText.
--
{-# NOINLINE memoFingerprint #-}
memoFingerprint
  :: (Typeable a, Monoid w) => MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
memoFingerprint key@(MemoFingerprintKey _ _ mnPtr nPtr) =
  withFingerprintLabel mnPtr nPtr . cachedComputation key

-- * Generic memoization machinery.

-- | Transform a Haxl computation into a memoized version of itself.
--
-- Given a Haxl computation, @memoize@ creates a version which stores its result
-- in a @MemoVar@ (which @memoize@ creates), and returns the stored result on
-- subsequent invocations. This permits the creation of local memos, whose
-- lifetimes are scoped to the current function, rather than the entire request.
memoize :: Monoid w => GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
memoize a = runMemo <$> newMemoWith a

-- | Transform a 1-argument function returning a Haxl computation into a
-- memoized version of itself.
--
-- Given a function @f@ of type @a -> GenHaxl u w b@, @memoize1@ creates a version
-- which memoizes the results of @f@ in a table keyed by its argument, and
-- returns stored results on subsequent invocations with the same argument.
--
-- e.g.:
--
-- > allFriends :: [Int] -> GenHaxl u w [Int]
-- > allFriends ids = do
-- >   memoizedFriendsOf <- memoize1 friendsOf
-- >   concat <$> mapM memoizeFriendsOf ids
--
-- The above implementation will not invoke the underlying @friendsOf@
-- repeatedly for duplicate values in @ids@.
memoize1 :: (Eq a, Hashable a, Monoid w)
         => (a -> GenHaxl u w b)
         -> GenHaxl u w (a -> GenHaxl u w b)
memoize1 f = runMemo1 <$> newMemoWith1 f

-- | Transform a 2-argument function returning a Haxl computation, into a
-- memoized version of itself.
--
-- The 2-ary version of @memoize1@, see its documentation for details.
memoize2 :: (Eq a, Hashable a, Eq b, Hashable b, Monoid w)
         => (a -> b -> GenHaxl u w c)
         -> GenHaxl u w (a -> b -> GenHaxl u w c)
memoize2 f = runMemo2 <$> newMemoWith2 f


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{- TODO

- do EVENTLOG stuff, track the data fetch numbers for performFetch

- timing: we should be using clock_gettime(CLOCK_MONOTONIC) instead of
  getCurrentTime, which will be affected by NTP and leap seconds.

- write different scheduling policies

-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- The implementation of the 'Haxl' monad.  Most users should
-- import "Haxl.Core" instead of importing this module directly.
--
module Haxl.Core.Monad
  (
    -- * The monad
    GenHaxl(..)
  , Result(..)

    -- * Writes (for debugging only)
  , WriteTree(..)
  , tellWrite
  , tellWriteNoMemo
  , write
  , writeNoMemo
  , flattenWT
  , appendWTs
  , mbModifyWLRef
  , mapWrites
  , mapWriteTree

    -- * Cont
  , Cont(..)
  , toHaxl

    -- * IVar
  , IVar(..)
  , IVarContents(..)
  , newIVar
  , newFullIVar
  , withCurrentCCS
  , getIVar
  , getIVarWithWrites
  , putIVar

    -- * ResultVal
  , ResultVal(..)
  , done
  , eitherToResult
  , eitherToResultThrowIO

    -- * CompleteReq
  , CompleteReq(..)

    -- * Env
  , Env(..)
  , DataCacheItem(..)
  , DataCacheLookup(..)
  , HaxlDataCache
  , Caches
  , caches
  , initEnvWithData
  , initEnv
  , emptyEnv
  , env, withEnv
  , nextCallId
  , sanitizeEnv

    -- * Profiling
  , ProfileCurrent(..)

    -- * JobList
  , JobList(..)
  , appendJobList
  , lengthJobList
  , addJob

    -- * Exceptions
  , throw
  , raise
  , catch
  , catchIf
  , try
  , tryToHaxlException

    -- * Dumping the cache
  , dumpCacheAsHaskell
  , dumpCacheAsHaskellFn

    -- * CallGraph
#ifdef PROFILING
  , withCallGraph
#endif

    -- * Unsafe operations
  ,  unsafeLiftIO, unsafeToHaxlException
  ) where

import Haxl.Core.Flags
import Haxl.Core.Stats
import Haxl.Core.StateStore
import Haxl.Core.Exception
import Haxl.Core.RequestStore as RequestStore
import Haxl.Core.DataCache as DataCache
import Haxl.Core.Util (trace_)

import Control.Applicative (liftA2)
import Control.Arrow (left)
import Control.Concurrent.STM
import qualified Control.Monad.Catch as Catch
import Control.Exception (Exception(..), SomeException, throwIO)
#if __GLASGOW_HASKELL__ >= 808
import Control.Monad hiding (MonadFail)
import qualified Control.Monad as CTL
#else
import Control.Monad
#endif
import qualified Control.Exception as Exception
import Data.Either (rights)
import Data.IORef
import Data.Int
import Data.Maybe
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty
#if __GLASGOW_HASKELL__ < 804
import Data.Semigroup
#endif
import qualified Data.Text as Text
import Data.Typeable
import GHC.Exts (IsString(..))
import Text.PrettyPrint hiding ((<>))
import Text.Printf
#ifdef EVENTLOG
import Control.Exception (bracket_)
import Debug.Trace (traceEventIO)
#endif

#ifdef PROFILING
import qualified Data.Map as Map
import Data.Text (Text)
import Foreign.Ptr (Ptr)
import GHC.Stack
import Haxl.Core.CallGraph
#endif

-- -----------------------------------------------------------------------------
-- The environment

-- | The data we carry around in the Haxl monad.

data DataCacheItem u w a = DataCacheItem (IVar u w a) {-# UNPACK #-} !CallId
type HaxlDataCache u w = DataCache (DataCacheItem u w)
newtype DataCacheLookup w =
  DataCacheLookup
    (forall req a . Typeable (req a)
    => req a
    -> IO (Maybe (ResultVal a w)))

data Env u w = Env
  { dataCache     :: {-# UNPACK #-} !(HaxlDataCache u w)
      -- ^ cached data fetches

  , memoCache      :: {-# UNPACK #-} !(HaxlDataCache u w)
      -- ^ memoized computations

  , memoKey    :: {-# UNPACK #-} !CallId
      -- ^ current running memo key

  , flags        :: !Flags
      -- conservatively not unpacking, because this is passed
      -- to 'fetch' and would need to be rebuilt.

  , userEnv      :: u
      -- ^ user-supplied data, retrievable with 'env'

  , statsRef     :: {-# UNPACK #-} !(IORef Stats)
      -- ^ statistics, collected according to the 'report' level in 'flags'.

  , statsBatchIdRef :: {-# UNPACK #-} !(IORef Int)
     -- ^ keeps track of a Unique ID for each batch dispatched with stats
     -- enabled, for aggregating after.

  , callIdRef :: {-# UNPACK #-} !(IORef CallId)
     -- ^ keeps track of a Unique ID for each fetch/memo.

  , profCurrent    :: ProfileCurrent
     -- ^ current profiling label, see 'withLabel'

  , profRef      :: {-# UNPACK #-} !(IORef Profile)
      -- ^ profiling data, collected according to the 'report' level in 'flags'.

  , states       :: StateStore
      -- ^ Data sources and other components can store their state in
      -- here. Items in this store must be instances of 'StateKey'.

  , reqStoreRef :: {-# UNPACK #-} !(IORef (RequestStore u))
       -- ^ The set of requests that we have not submitted to data sources yet.
       -- Owned by the scheduler.

  , runQueueRef :: {-# UNPACK #-} !(IORef (JobList u w))
       -- ^ runnable computations. Things get added to here when we wake up
       -- a computation that was waiting for something.  When the list is
       -- empty, either we're finished, or we're waiting for some data fetch
       -- to return.

  , submittedReqsRef :: {-# UNPACK #-} !(IORef ReqCountMap)
       -- ^ all outgone fetches which haven't yet returned. Entries are
       -- removed from this map as the fetches finish. This field is
       -- useful for tracking outgone fetches to detect downstream
       -- failures.

  , completions :: {-# UNPACK #-} !(TVar [CompleteReq u w])
       -- ^ Requests that have completed.  Modified by data sources
       -- (via putResult) and the scheduler.  Waiting for this list to
       -- become non-empty is how the scheduler blocks waiting for
       -- data fetches to return.

  , writeLogsRef :: {-# UNPACK #-} !(IORef w)
       -- ^ 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 that we can get a hold of those writes and put them
       -- in the IVar associated with the compuatation.
  , writeLogsRefNoMemo :: {-# UNPACK #-} !(IORef w)
       -- ^ This is just a specialized version of @writeLogsRef@, where we put
       -- logs that user doesn't want memoized. This is a better alternative to
       -- doing arbitrary IO from a (memoized) Haxl computation.

  , dataCacheFetchFallback :: !(Maybe (DataCacheLookup w))
       -- ^ Allows you to inject a DataCache lookup just before a dataFetch is
       -- dispatched. This is useful for injecting fetch results in testing.

#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
  }

data ProfileCurrent = ProfileCurrent
  { profCurrentKey ::  {-# UNPACK #-} !ProfileKey
  , profLabelStack :: {-# UNPACK #-} !(NonEmpty ProfileLabel)
  }

type Caches u w = (HaxlDataCache u w, HaxlDataCache u w)

caches :: Env u w -> Caches u w
caches env = (dataCache env, memoCache env)

getMaxCallId :: HaxlDataCache u w -> IO (Maybe Int)
getMaxCallId c = do
  callIds  <- rights . concatMap snd <$>
              DataCache.readCache c (\(DataCacheItem _ i) -> return i)
  case callIds of
    [] -> return Nothing
    vals -> return $ Just (maximum vals)

-- | Initialize an environment with a 'StateStore', an input map, a
-- preexisting 'DataCache', and a seed for the random number generator.
initEnvWithData :: Monoid w => StateStore -> u -> Caches u w -> IO (Env u w)
initEnvWithData states e (dcache, mcache) = do
  newCid <- max <$>
    (maybe 0 ((+) 1) <$> getMaxCallId dcache) <*>
    (maybe 0 ((+) 1) <$> getMaxCallId mcache)
  ciref<- newIORef newCid
  sref <- newIORef emptyStats
  sbref <- newIORef 0
  pref <- newIORef emptyProfile
  rs <- newIORef noRequests          -- RequestStore
  rq <- newIORef JobNil              -- RunQueue
  sr <- newIORef emptyReqCounts      -- SubmittedReqs
  comps <- newTVarIO []              -- completion queue
  wl <- newIORef mempty
  wlnm <- newIORef mempty
  return Env
    { dataCache = dcache
    , memoCache = mcache
    , memoKey = (-1)
    , flags = defaultFlags
    , userEnv = e
    , states = states
    , statsRef = sref
    , statsBatchIdRef = sbref
    , profCurrent = ProfileCurrent 0 $ "MAIN" :| []
    , callIdRef = ciref
    , profRef = pref
    , reqStoreRef = rs
    , runQueueRef = rq
    , submittedReqsRef = sr
    , completions = comps
    , writeLogsRef = wl
    , writeLogsRefNoMemo = wlnm
    , dataCacheFetchFallback = Nothing
#ifdef PROFILING
    , callGraphRef = Nothing
    , currFunction = mainFunction
#endif
    }

-- | Initializes an environment with 'StateStore' and an input map.
initEnv :: Monoid w => StateStore -> u -> IO (Env u w)
initEnv states e = do
  dcache <- emptyDataCache
  mcache <- emptyDataCache
  initEnvWithData states e (dcache, mcache)

-- | A new, empty environment.
emptyEnv :: Monoid w => u -> IO (Env u w)
emptyEnv = initEnv stateEmpty

-- | If you're using the env from a failed Haxl computation in a second Haxl
-- computation, it is recommended to sanitize the Env to remove all empty
-- IVars - especially if it's possible the first Haxl computation could've
-- been interrupted via an async exception. This is because if the Haxl
-- computation was interrupted by an exception, it's possible that there are
-- entries in the cache which are still blocked, while the results from
-- outgone fetches have been discarded.
sanitizeEnv :: Env u w -> IO (Env u w)
sanitizeEnv env@Env{..} = do
  sanitizedDC <- DataCache.filter isIVarFull dataCache
  sanitizedMC <- DataCache.filter isIVarFull memoCache
  rs <- newIORef noRequests          -- RequestStore
  rq <- newIORef JobNil              -- RunQueue
  comps <- newTVarIO []              -- completion queue
  sr <- newIORef emptyReqCounts      -- SubmittedReqs
  return env
    { dataCache = sanitizedDC
    , memoCache = sanitizedMC
    , reqStoreRef = rs
    , runQueueRef = rq
    , completions = comps
    , submittedReqsRef = sr
    }
  where
  isIVarFull (DataCacheItem IVar{..} _) = do
    ivarContents <- readIORef ivarRef
    case ivarContents of
      IVarFull _ -> return True
      _ -> return False

-- -----------------------------------------------------------------------------
-- WriteTree

-- | A tree of writes done during a Haxl computation. We could use a simple
-- list, but this allows us to avoid multiple mappends when concatenating
-- writes from two haxl computations.
--
-- Users should try to treat this data type as opaque, and prefer
-- to use @flattenWT@ to get a simple list of writes from a @WriteTree@.
data WriteTree w
  = NilWrites
  | SomeWrite w
  | MergeWrites (WriteTree w) (WriteTree w)
  deriving (Show)

instance Semigroup (WriteTree w) where
  (<>) = appendWTs

instance Monoid (WriteTree w) where
  mempty = NilWrites

appendWTs :: WriteTree w -> WriteTree w -> WriteTree w
appendWTs NilWrites w = w
appendWTs w NilWrites = w
appendWTs w1 w2 = MergeWrites w1 w2

-- This function must be called at the end of the Haxl computation to get
-- a list of writes.
-- Haxl provides no guarantees on the order of the returned logs.
flattenWT :: WriteTree w -> [w]
flattenWT = go []
  where
    go !ws NilWrites = ws
    go !ws (SomeWrite w) = w : ws
    go !ws (MergeWrites w1 w2) = go (go ws w2) w1

-- This is a convenience wrapper over modifyIORef, which only modifies
-- writeLogsRef IORef, for non NilWrites.
mbModifyWLRef :: WriteTree w -> IORef (WriteTree w) -> IO ()
mbModifyWLRef NilWrites _ = return ()
mbModifyWLRef !wt ref = modifyIORef' ref (`appendWTs` wt)

mapWriteTree :: (w -> w) -> WriteTree w -> WriteTree w
mapWriteTree _ NilWrites = NilWrites
mapWriteTree f (SomeWrite w) = SomeWrite (f w)
mapWriteTree f (MergeWrites wt1 wt2) =
  MergeWrites (mapWriteTree f wt1) (mapWriteTree f wt2)

-- -----------------------------------------------------------------------------
-- | The Haxl monad, which does several things:
--
--  * It is a reader monad for 'Env', which contains the current state
--    of the scheduler, including unfetched requests and the run queue
--    of computations.
--
--  * It is a writer monad for 'WriteTree'. These can be used to do
--    arbitrary "logs" from any Haxl computation. These are better than
--    doing arbitrary IO from a Haxl computation as these writes also get
--    memoized if the Haxl computation associated with them is memoized.
--    Now if this memoized computation is run again, you'll get the writes
--    twice.

--  * It is a concurrency, or resumption, monad. A computation may run
--    partially and return 'Blocked', in which case the framework should
--    perform the outstanding requests in the 'RequestStore', and then
--    resume the computation.
--
--  * The Applicative combinator '<*>' explores /both/ branches in the
--    event that the left branch is 'Blocked', so that we can collect
--    multiple requests and submit them as a batch.
--
--  * It contains IO, so that we can perform real data fetching.
--
newtype GenHaxl u w a = GenHaxl
  { unHaxl :: Env u w -> IO (Result u w a) }

tellWrite :: w -> GenHaxl u (WriteTree w) ()
tellWrite = write . SomeWrite

write :: Monoid w => w -> GenHaxl u w ()
write wt = GenHaxl $ \Env{..} -> do
  modifyIORef' writeLogsRef (<> wt)
  return $ Done ()

tellWriteNoMemo :: w -> GenHaxl u (WriteTree w) ()
tellWriteNoMemo = writeNoMemo . SomeWrite

writeNoMemo :: Monoid w => w -> GenHaxl u w ()
writeNoMemo wt = GenHaxl $ \Env{..} -> do
  modifyIORef' writeLogsRefNoMemo (<> wt)
  return $ Done ()

instance IsString a => IsString (GenHaxl u w a) where
  fromString s = return (fromString s)

-- -----------------------------------------------------------------------------
-- JobList

-- | A list of computations together with the IVar into which they
-- should put their result.
--
-- This could be an ordinary list, but the optimised representation
-- saves space and time.
--
data JobList u w
 = JobNil
 | forall a . JobCons
     (Env u w)          -- See Note [make withEnv work] below.
     (GenHaxl u w a)
     {-# UNPACK #-} !(IVar u w a)
     (JobList u w)

-- Note [make withEnv work]
--
-- The withEnv operation supplies a new Env for the scope of a GenHaxl
-- computation.  The problem is that the computation might be split
-- into pieces and put onto various JobLists, so we have to be sure to
-- use the correct Env when we execute the pieces. Furthermore, if one
-- of these pieces blocks and gets run again later, we must ensure to
-- restart it with the correct Env.  So we stash the Env along with
-- the continuation in the JobList.

appendJobList :: JobList u w -> JobList u w -> JobList u w
appendJobList JobNil c = c
appendJobList c JobNil = c
appendJobList (JobCons a b c d) e = JobCons a b c $! appendJobList d e

lengthJobList :: JobList u w -> Int
lengthJobList JobNil = 0
lengthJobList (JobCons _ _ _ j) = 1 + lengthJobList j

-- -----------------------------------------------------------------------------
-- IVar

-- | A synchronisation point.  It either contains a value, or a list
-- of computations waiting for the value.
#ifdef PROFILING
data IVar u w a = IVar
  { ivarRef :: {-# UNPACK #-} !(IORef (IVarContents u w a))
  , ivarCCS :: {-# UNPACK #-} !(Ptr CostCentreStack)
#else
newtype IVar u w a = IVar
  { ivarRef :: IORef (IVarContents u w a)
#endif
  }

data IVarContents u w a
  = IVarFull (ResultVal a w)
  | IVarEmpty (JobList u w)
    -- morally this is a list of @a -> GenHaxl u w ()@, but instead of
    -- using a function, each computation begins with `getIVar` to grab
    -- the value it is waiting for.  This is less type safe but a little
    -- faster (benchmarked with tests/MonadBench.hs).

newIVar :: IO (IVar u w a)
newIVar = do
  ivarRef <- newIORef (IVarEmpty JobNil)
#ifdef PROFILING
  ivarCCS <- getCurrentCCS ivarRef
#endif
  return IVar{..}

newFullIVar :: ResultVal a w -> IO (IVar u w a)
newFullIVar r = do
  ivarRef <- newIORef (IVarFull r)
#ifdef PROFILING
  ivarCCS <- getCurrentCCS ivarRef
#endif
  return IVar{..}

withCurrentCCS :: IVar u w a -> IO (IVar u w a)
#ifdef PROFILING
withCurrentCCS ivar = do
  ccs <- getCurrentCCS ivar
  return ivar{ivarCCS = ccs}
#else
withCurrentCCS = return
#endif

getIVar :: IVar u w a -> GenHaxl u w a
getIVar i@IVar{ivarRef = !ref} = GenHaxl $ \env -> do
  e <- readIORef ref
  case e of
    IVarFull (Ok a _wt) -> return (Done a)
    IVarFull (ThrowHaxl e _wt) -> raiseFromIVar env i e
    IVarFull (ThrowIO e) -> throwIO e
    IVarEmpty _ -> return (Blocked i (Return i))

-- Just a specialised version of getIVar, for efficiency in <*>
getIVarApply :: IVar u w (a -> b) -> a -> GenHaxl u w b
getIVarApply i@IVar{ivarRef = !ref} a = GenHaxl $ \env -> do
  e <- readIORef ref
  case e of
    IVarFull (Ok f _wt) -> return (Done (f a))
    IVarFull (ThrowHaxl e _wt) -> raiseFromIVar env i e
    IVarFull (ThrowIO e) -> throwIO e
    IVarEmpty _ ->
      return (Blocked i (Cont (getIVarApply i a)))

-- Another specialised version of getIVar, for efficiency in cachedComputation
getIVarWithWrites :: Monoid w => IVar u w a -> GenHaxl u w a
getIVarWithWrites i@IVar{ivarRef = !ref} = GenHaxl $ \env@Env{..} -> do
  e <- readIORef ref
  case e of
    IVarFull (Ok a wt) -> do
      modifyIORef' writeLogsRef (<> fromMaybe mempty wt)
      return (Done a)
    IVarFull (ThrowHaxl e wt) -> do
      modifyIORef' writeLogsRef (<> fromMaybe mempty wt)
      raiseFromIVar env i e
    IVarFull (ThrowIO e) -> throwIO e
    IVarEmpty _ ->
      return (Blocked i (Cont (getIVarWithWrites i)))

putIVar :: IVar u w a -> ResultVal a w -> Env u w -> IO ()
putIVar IVar{ivarRef = !ref} a Env{..} = do
  e <- readIORef ref
  case e of
    IVarEmpty jobs -> do
      writeIORef ref (IVarFull a)
      modifyIORef' runQueueRef (appendJobList jobs)
      -- An IVar is typically only meant to be written to once
      -- so it would make sense to throw an error here. But there
      -- are legitimate use-cases for writing several times.
      -- (See Haxl.Core.Parallel)
    IVarFull{} -> return ()

{-# INLINE addJob #-}
addJob :: Env u w -> GenHaxl u w b -> IVar u w b -> IVar u w a -> IO ()
addJob env !haxl !resultIVar IVar{ivarRef = !ref} =
  modifyIORef' ref $ \contents ->
    case contents of
      IVarEmpty list -> IVarEmpty (JobCons env haxl resultIVar list)
      _ -> addJobPanic

addJobPanic :: forall a . a
addJobPanic = error "addJob: not empty"

-- -----------------------------------------------------------------------------
-- ResultVal

-- | The contents of a full IVar.  We have to distinguish exceptions
-- thrown in the IO monad from exceptions thrown in the Haxl monad, so
-- that when the result is fetched using getIVar, we can throw the
-- exception in the right way.
data ResultVal a w
  = Ok a (Maybe w)
  | ThrowHaxl SomeException (Maybe w)
  | ThrowIO SomeException
    -- we get no write logs when an IO exception occurs

done :: Env u w -> ResultVal a w -> IO (Result u w a)
done _ (Ok a _) = return (Done a)
done env (ThrowHaxl e _) = raise env e
done _ (ThrowIO e) = throwIO e

eitherToResultThrowIO :: Either SomeException a -> ResultVal a w
eitherToResultThrowIO (Right a) = Ok a Nothing
eitherToResultThrowIO (Left e)
  | Just HaxlException{} <- fromException e = ThrowHaxl e Nothing
  | otherwise = ThrowIO e

eitherToResult :: Either SomeException a -> ResultVal a w
eitherToResult (Right a) = Ok a Nothing
eitherToResult (Left e) = ThrowHaxl e Nothing

-- -----------------------------------------------------------------------------
-- CompleteReq

-- | A completed request from a data source, containing the result,
-- and the 'IVar' representing the blocked computations.  The job of a
-- data source is just to add these to a queue ('completions') using
-- 'putResult'; the scheduler collects them from the queue and unblocks
-- the relevant computations.
data CompleteReq u w
  = forall a . CompleteReq
      (ResultVal a w)
      !(IVar u w a)  -- IVar because the result is cached
      {-# UNPACK #-} !Int64 -- see Note [tracking allocation in child threads]

{- Note [tracking allocation in child threads]

For a BackgroundFetch, we might be doing some of the work in a
separate thread, but we want to make sure that the parent thread gets
charged for the allocation, so that allocation limits still work.

The design is a bit tricky here.  We want to track the allocation
accurately but without adding much overhead.

The best way to propagate the allocation back from the child thread is
through putResult.  If we had some other method, we would also need a
way to synchronise it with the main runHaxl loop; the advantage of
putResult is that this is already a synchronisation method, because
runHaxl is waiting for the result of the dataFetch.

We need a special version of putResult for child threads
(putResultFromChildThread), because we don't want to propagate any
allocation from the runHaxl thread back to itself and count it twice.

We also want to capture the allocation as late as possible, so that we
count everything.  For that reason, we pass a Bool down from putResult
into the function in the ResultVar, and it reads the allocation
counter as the last thing before adding the result to the completions
TVar.

The other problem to consider is how to capture the allocation when
the child thread is doing multiple putResults.  Our solution here is
to ensure that the *last* one is a putResultFromChildThread, so it
captures all the allocation from everything leading up to it.

Why not reset the counter each time, so we could do multiple
putResultFromChildThreads?  Because the child thread might be using an
allocation limit itself, and changing the counter would mess it up.
-}

-- -----------------------------------------------------------------------------
-- Result

-- | The result of a computation is either 'Done' with a value, 'Throw'
-- with an exception, or 'Blocked' on the result of a data fetch with
-- a continuation.
data Result u w a
  = Done a
  | Throw SomeException
  | forall b . Blocked
      {-# UNPACK #-} !(IVar u w b)
      (Cont u w a)
         -- ^ The 'IVar' is what we are blocked on; 'Cont' is the
         -- continuation.  This might be wrapped further if we're
         -- nested inside multiple '>>=', before finally being added
         -- to the 'IVar'.  Morally @b -> GenHaxl u w a@, but see
         -- 'IVar',

instance (Show a) => Show (Result u w a) where
  show (Done a) = printf "Done(%s)" $ show a
  show (Throw e) = printf "Throw(%s)" $ show e
  show Blocked{} = "Blocked"

instance Functor (Result u w) where
  fmap f (Done a) = Done (f a)
  fmap _ (Throw exc) = Throw exc
  fmap f (Blocked ivar cont) = Blocked ivar (f :<$> cont)

{- Note [Exception]

How do we want to represent Haxl exceptions (those that are thrown by
"throw" in the Haxl monad)?

1) Explicitly via a Throw constructor in the Result type
2) Using throwIO in the IO monad

If we did (2), we would have to use an exception handler in <*>,
because an exception in the right-hand argument of <*> should not
necessarily be thrown by the whole computation - an exception on the
left should get priority, and the left might currently be Blocked.

We must be careful about turning IO monad exceptions into Haxl
exceptions.  An IO monad exception will normally propagate right
out of runHaxl and terminate the whole computation, whereas a Haxl
exception can get dropped on the floor, if it is on the right of
<*> and the left side also throws, for example.  So turning an IO
monad exception into a Haxl exception is a dangerous thing to do.
In particular, we never want to do it for an asynchronous exception
(AllocationLimitExceeded, ThreadKilled, etc.), because these are
supposed to unconditionally terminate the computation.

There are three places where we take an arbitrary IO monad exception and
turn it into a Haxl exception:

 * wrapFetchInCatch.  Here we want to propagate a failure of the
   data source to the callers of the data source, but if the
   failure came from elsewhere (an asynchronous exception), then we
   should just propagate it

 * cacheResult (cache the results of IO operations): again,
   failures of the IO operation should be visible to the caller as
   a Haxl exception, but we exclude asynchronous exceptions from
   this.

 * unsafeToHaxlException: assume the caller knows what they're
   doing, and just wrap all exceptions.
-}

-- -----------------------------------------------------------------------------
-- Cont

-- | A data representation of a Haxl continuation.  This is to avoid
-- repeatedly traversing a left-biased tree in a continuation, leading
-- O(n^2) complexity for some pathalogical cases - see the "seql" benchmark
-- in tests/MonadBench.hs.
-- See "A Smart View on Datatypes", Jaskelioff/Rivas, ICFP'15
data Cont u w a
  = Cont (GenHaxl u w a)
  | forall b. Cont u w b :>>= (b -> GenHaxl u w a)
  | forall b. (b -> a) :<$> (Cont u w b)
  | Return (IVar u w a)

toHaxl :: Cont u w a -> GenHaxl u w a
toHaxl (Cont haxl) = haxl
toHaxl (m :>>= k) = toHaxlBind m k
toHaxl (f :<$> x) = toHaxlFmap f x
toHaxl (Return i) = getIVar i

toHaxlBind :: Cont u w b -> (b -> GenHaxl u w a) -> GenHaxl u w a
toHaxlBind (m :>>= k) k2 = toHaxlBind m (k >=> k2)
toHaxlBind (Cont haxl) k = haxl >>= k
toHaxlBind (f :<$> x) k = toHaxlBind x (k . f)
toHaxlBind (Return i) k = getIVar i >>= k

toHaxlFmap :: (a -> b) -> Cont u w a -> GenHaxl u w b
toHaxlFmap f (m :>>= k) = toHaxlBind m (k >=> return . f)
toHaxlFmap f (Cont haxl) = f <$> haxl
toHaxlFmap f (g :<$> x) = toHaxlFmap (f . g) x
toHaxlFmap f (Return i) = f <$> getIVar i

-- -----------------------------------------------------------------------------
-- Monad/Applicative instances

instance Monad (GenHaxl u w) where
  return = pure
  GenHaxl m >>= k = GenHaxl $ \env -> do
    e <- m env
    case e of
      Done a -> unHaxl (k a) env
      Throw e -> return (Throw e)
      Blocked ivar cont -> trace_ ">>= Blocked" $
        return (Blocked ivar (cont :>>= k))

  -- We really want the Applicative version of >>
  (>>) = (*>)

#if __GLASGOW_HASKELL__ >= 808
instance CTL.MonadFail (GenHaxl u w) where
#endif
  fail msg = GenHaxl $ \_env ->
    return $ Throw $ toException $ MonadFail $ Text.pack msg

instance Functor (GenHaxl u w) where
  fmap f (GenHaxl m) = GenHaxl $ \env -> do
    r <- m env
    case r of
      Done a -> return (Done (f a))
      Throw e -> return (Throw e)
      Blocked ivar cont -> trace_ "fmap Blocked" $
        return (Blocked ivar (f :<$> cont))

instance Applicative (GenHaxl u w) where
  pure a = GenHaxl $ \_env -> pure (Done a)
  GenHaxl ff <*> GenHaxl aa = GenHaxl $ \env -> do
    rf <- ff env
    case rf of
      Done f -> do
        ra <- aa env
        case ra of
          Done a -> trace_ "Done/Done" $ return (Done (f a))
          Throw e -> trace_ "Done/Throw" $ return (Throw e)
          Blocked ivar fcont -> trace_ "Done/Blocked" $
            return (Blocked ivar (f :<$> fcont))
      Throw e -> trace_ "Throw" $ return (Throw e)
      Blocked ivar1 fcont -> do
        ra <- aa env
        case ra of
          Done a -> trace_ "Blocked/Done" $
            return (Blocked ivar1 (($ a) :<$> fcont))
          Throw e -> trace_ "Blocked/Throw" $
            return (Blocked ivar1 (fcont :>>= (\_ -> throw e)))
          Blocked ivar2 acont -> trace_ "Blocked/Blocked" $
            blockedBlocked env ivar1 fcont ivar2 acont
             -- Note [Blocked/Blocked]

instance Semigroup a => Semigroup (GenHaxl u w a) where
  (<>) = liftA2 (<>)

instance Monoid a => Monoid (GenHaxl u w a) where
  mempty = pure mempty
#if __GLASGOW_HASKELL__ < 804
  mappend = liftA2 mappend
#endif

blockedBlocked
  :: Env u w
  -> IVar u w c
  -> Cont u w (a -> b)
  -> IVar u w d
  -> Cont u w a
  -> IO (Result u w b)
blockedBlocked _ _ (Return i) ivar2 acont =
  return (Blocked ivar2 (acont :>>= getIVarApply i))
blockedBlocked _ _ (g :<$> Return i) ivar2 acont =
  return (Blocked ivar2 (acont :>>= \ a -> (\f -> g f a) <$> getIVar i))
blockedBlocked env ivar1 fcont ivar2 acont = do
  i <- newIVar
  addJob env (toHaxl fcont) i ivar1
  let cont = acont :>>= \a -> getIVarApply i a
  return (Blocked ivar2 cont)

-- Note [Blocked/Blocked]
--
-- This is the tricky case: we're blocked on both sides of the <*>.
-- We need to divide the computation into two pieces that may continue
-- independently when the resources they are blocked on become
-- available.  Moreover, the computation as a whole depends on the two
-- pieces.  It works like this:
--
--   ff <*> aa
--
-- becomes
--
--   (ff >>= putIVar i) <*> (a <- aa; f <- getIVar i; return (f a)
--
-- where the IVar i is a new synchronisation point.  If the right side
-- gets to the `getIVar` first, it will block until the left side has
-- called 'putIVar'.
--
-- We can also do it the other way around:
--
--   (do ff <- f; getIVar i; return (ff a)) <*> (a >>= putIVar i)
--
-- The first was slightly faster according to tests/MonadBench.hs.

-- -----------------------------------------------------------------------------
-- Env utils

-- | Extracts data from the 'Env'.
env :: (Env u w -> a) -> GenHaxl u w a
env f = GenHaxl $ \env -> return (Done (f env))

-- | Returns a version of the Haxl computation which always uses the
-- provided 'Env', ignoring the one specified by 'runHaxl'.
withEnv :: Env u w -> GenHaxl u w a -> GenHaxl u w a
withEnv newEnv (GenHaxl m) = GenHaxl $ \_env -> do
  r <- m newEnv
  case r of
    Done a -> return (Done a)
    Throw e -> return (Throw e)
    Blocked ivar k ->
      return (Blocked ivar (Cont (withEnv newEnv (toHaxl k))))

nextCallId :: Env u w -> IO CallId
nextCallId env = atomicModifyIORef' (callIdRef env) $ \x -> (x+1,x+1)

-- | Runs the Haxl computation, transforming the writes within using the
-- provided function.
--
-- Memoization behavior is unchanged, meaning if a memoized computation is run
-- once inside @mapWrites@ and then once without, the writes from the second run
-- will NOT be transformed.
mapWrites :: Monoid w => (w -> w) -> GenHaxl u w a -> GenHaxl u w a
mapWrites f action = GenHaxl $ \curEnv -> do
  wlogs <- newIORef mempty
  wlogsNoMemo <- newIORef mempty
  let
    !newEnv = curEnv { writeLogsRef = wlogs, writeLogsRefNoMemo = wlogsNoMemo }
  unHaxl (mapWritesImpl curEnv newEnv action) newEnv
  where
    mapWritesImpl oldEnv curEnv (GenHaxl m) = GenHaxl $ \_ -> do
      let
        pushTransformedWrites = do
          wt <- readIORef $ writeLogsRef curEnv
          modifyIORef' (writeLogsRef oldEnv) (<> f wt)
          wtNoMemo <- readIORef $ writeLogsRefNoMemo curEnv
          modifyIORef' (writeLogsRefNoMemo oldEnv) (<> f wtNoMemo)

      r <- m curEnv

      case r of
        Done a -> pushTransformedWrites >> return (Done a)
        Throw e -> pushTransformedWrites >> return (Throw e)
        Blocked ivar k -> return
          (Blocked ivar (Cont (mapWritesImpl oldEnv curEnv (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 w a
  -> GenHaxl u w 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

-- | Throw an exception in the Haxl monad
throw :: Exception e => e -> GenHaxl u w a
throw e = GenHaxl $ \env -> raise env e

raise :: Exception e => Env u w -> e -> IO (Result u w a)
raise env e = raiseImpl env (toException e)
#ifdef PROFILING
  currentCallStack
#endif

raiseFromIVar :: Exception e => Env u w -> IVar u w a -> e -> IO (Result u w b)
#ifdef PROFILING
raiseFromIVar env IVar{..} e =
  raiseImpl env (toException e) (ccsToStrings ivarCCS)
#else
raiseFromIVar env _ivar e = raiseImpl env (toException e)
#endif

{-# INLINE raiseImpl #-}
#ifdef PROFILING
raiseImpl :: Env u w -> SomeException -> IO [String] -> IO (Result u w b)
raiseImpl Env{..} e getCostCentreStack
#else
raiseImpl :: Env u w -> SomeException -> IO (Result u w b)
raiseImpl Env{..} e
#endif
  | testReportFlag ReportExceptionLabelStack $ report flags
  , Just (HaxlException Nothing h) <- fromException e = do
    let stk = NonEmpty.toList $ profLabelStack profCurrent
    return $ Throw $ toException $ HaxlException (Just stk) h
#ifdef PROFILING
  | Just (HaxlException Nothing h) <- fromException e = do
    stk <- reverse . map Text.pack <$> getCostCentreStack
    return $ Throw $ toException $ HaxlException (Just stk) h
#endif
  | otherwise = return $ Throw e

-- | Catch an exception in the Haxl monad
catch :: Exception e => GenHaxl u w a -> (e -> GenHaxl u w a) -> GenHaxl u w a
catch (GenHaxl m) h = GenHaxl $ \env -> do
   r <- m env
   case r of
     Done a    -> return (Done a)
     Throw e | Just e' <- fromException e -> unHaxl (h e') env
             | otherwise -> return (Throw e)
     Blocked ivar k -> return (Blocked ivar (Cont (catch (toHaxl k) h)))

-- | Catch exceptions that satisfy a predicate
catchIf
  :: Exception e => (e -> Bool) -> GenHaxl u w a -> (e -> GenHaxl u w a)
  -> GenHaxl u w a
catchIf cond haxl handler =
  catch haxl $ \e -> if cond e then handler e else throw e

-- | Returns @'Left' e@ if the computation throws an exception @e@, or
-- @'Right' a@ if it returns a result @a@.
try :: Exception e => GenHaxl u w a -> GenHaxl u w (Either e a)
try haxl = (Right <$> haxl) `catch` (return . Left)

-- | @since 0.3.1.0
instance Catch.MonadThrow (GenHaxl u w) where throwM = Haxl.Core.Monad.throw
-- | @since 0.3.1.0
instance Catch.MonadCatch (GenHaxl u w) where catch = Haxl.Core.Monad.catch

-- -----------------------------------------------------------------------------
-- Unsafe operations

-- | Under ordinary circumstances this is unnecessary; users of the Haxl
-- monad should generally /not/ perform arbitrary IO.
unsafeLiftIO :: IO a -> GenHaxl u w a
unsafeLiftIO m = GenHaxl $ \_env -> Done <$> m

-- | Convert exceptions in the underlying IO monad to exceptions in
-- the Haxl monad.  This is morally unsafe, because you could then
-- catch those exceptions in Haxl and observe the underlying execution
-- order.  Not to be exposed to user code.
--
-- Note: this function does not catch async exceptions. This is a flaw in Haxl
-- where it can sometimes leave the environment in a bad state when async
-- exceptions are thrown (for example the cache may think a fetch is happening
-- but the exception has stopped it). TODO would be to make Haxl async exception
-- safe and then remove the rethrowAsyncExceptions below, but for now this is
-- safer to avoid bugs. Additionally this would not protect you from async
-- exceptions thrown while executing code in the scheduler, and so relying on
-- this function to catch all async exceptions would be ambitious at best.
unsafeToHaxlException :: GenHaxl u w a -> GenHaxl u w a
unsafeToHaxlException (GenHaxl m) = GenHaxl $ \env -> do
  r <- m env `Exception.catch` \e -> do
    rethrowAsyncExceptions e
    return (Throw e)
  case r of
    Blocked cvar c ->
      return (Blocked cvar (Cont (unsafeToHaxlException (toHaxl c))))
    other -> return other

-- | Like 'try', but lifts all exceptions into the 'HaxlException'
-- hierarchy.  Uses 'unsafeToHaxlException' internally.  Typically
-- this is used at the top level of a Haxl computation, to ensure that
-- all exceptions are caught.
tryToHaxlException :: GenHaxl u w a -> GenHaxl u w (Either HaxlException a)
tryToHaxlException h = left asHaxlException <$> try (unsafeToHaxlException h)

-- -----------------------------------------------------------------------------

-- | Dump the contents of the cache as Haskell code that, when
-- compiled and run, will recreate the same cache contents.  For
-- example, the generated code looks something like this:
--
-- > loadCache :: GenHaxl u w ()
-- > loadCache = do
-- >   cacheRequest (ListWombats 3) (Right ([1,2,3]))
-- >   cacheRequest (CountAardvarks "abcabc") (Right (2))
--
dumpCacheAsHaskell :: GenHaxl u w String
dumpCacheAsHaskell =
    dumpCacheAsHaskellFn "loadCache" "GenHaxl u w ()" "cacheRequest"

-- | Dump the contents of the cache as Haskell code that, when
-- compiled and run, will recreate the same cache contents.
-- Does not take into account the writes done as part of the computation.
--
-- Takes the name and type for the resulting function as arguments.
-- Also takes the cacheFn to use, we can use either @cacheRequest@ or
-- @dupableCacheRequest@.
dumpCacheAsHaskellFn :: String -> String -> String -> GenHaxl u w String
dumpCacheAsHaskellFn fnName fnType cacheFn = do
  cache <- env dataCache  -- NB. dataCache, not memoCache.  We ignore memoized
                       -- results when dumping the cache.
  let
    readIVar (DataCacheItem IVar{ivarRef = !ref} _) = do
      r <- readIORef ref
      case r of
        IVarFull (Ok a _) -> return (Just (Right a))
        IVarFull (ThrowHaxl e _) -> return (Just (Left e))
        IVarFull (ThrowIO e) -> return (Just (Left e))
        IVarEmpty _ -> return Nothing

    mk_cr (req, res) =
      text cacheFn <+> parens (text req) <+> parens (result res)
    result (Left e) = text "except" <+> parens (text (show e))
    result (Right s) = text "Right" <+> parens (text s)

  entries <- unsafeLiftIO $ do
    showCache cache readIVar

  let
    body = if null entries
      then text "return ()"
      else vcat (map mk_cr (concatMap snd entries))

  return $ show $
    text (fnName ++ " :: " ++ fnType) $$
    text (fnName ++ " = do") $$
      nest 2 body $$
    text "" -- final newline


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE CPP #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | Psuedo-parallel operations.  Most users should import "Haxl.Core"
-- instead.
--
module Haxl.Core.Parallel
  ( -- * Parallel operations
    biselect
  , pAnd
  , pOr
  , unsafeChooseFirst
  ) where

import Haxl.Core.Monad hiding (catch, throw)
import Haxl.Core.Exception

import Control.Exception (throw)

-- -----------------------------------------------------------------------------
-- Parallel operations

-- Bind more tightly than .&&, .||
infixr 5 `pAnd`
infixr 4 `pOr`

biselect :: GenHaxl u w (Either a b)
         -> GenHaxl u w (Either a c)
         -> GenHaxl u w (Either a (b,c))
biselect haxla haxlb = biselect_opt id id Left Right haxla haxlb

{-# INLINE biselect_opt #-}
biselect_opt :: (l -> Either a b)
             -> (r -> Either a c)
             -> (a -> t)
             -> ((b,c) -> t)
             -> GenHaxl u w l
             -> GenHaxl u w r
             -> GenHaxl u w t
biselect_opt discrimA discrimB left right haxla haxlb =
  let go (GenHaxl haxla) (GenHaxl haxlb) = GenHaxl $ \env -> do
        ra <- haxla env
        case ra of
          Done ea ->
            case discrimA ea of
              Left a -> return (Done (left a))
              Right b -> do
                  rb <- haxlb env
                  case rb of
                    Done eb ->
                      case discrimB eb of
                        Left a -> return (Done (left a))
                        Right c -> return (Done (right (b,c)))
                    Throw e -> return (Throw e)
                    Blocked ib haxlb' ->
                      return (Blocked ib
                              (haxlb' :>>= \b' -> go_right b b'))
          Throw e -> return (Throw e)
          Blocked ia haxla' -> do
            rb <- haxlb env
            case rb of
              Done eb ->
                case discrimB eb of
                  Left a -> return (Done (left a))
                  Right c ->
                     return (Blocked ia
                             (haxla' :>>= \a' -> go_left a' c))
              Throw e -> return (Throw e)
              Blocked ib haxlb' -> do
                i <- newIVar
                addJob env (return ()) i ia
                addJob env (return ()) i ib
                return (Blocked i (Cont (go (toHaxl haxla') (toHaxl haxlb'))))
                -- The code above makes sure that the computation
                -- wakes up whenever either 'ia' or 'ib' is filled.
                -- The ivar 'i' is used as a synchronisation point
                -- for the whole computation, and we make sure that
                -- whenever 'ia' or 'ib' are filled in then 'i' will
                -- also be filled.

      go_right b eb =
        case discrimB eb of
          Left a -> return (left a)
          Right c -> return (right (b,c))
      go_left ea c =
        case discrimA ea of
          Left a -> return (left a)
          Right b -> return (right (b,c))
  in go haxla haxlb

-- | Parallel version of '(.||)'.  Both arguments are evaluated in
-- parallel, and if either returns 'True' then the other is
-- not evaluated any further.
--
-- WARNING: exceptions may be unpredictable when using 'pOr'.  If one
-- argument returns 'True' before the other completes, then 'pOr'
-- returns 'True' immediately, ignoring a possible exception that
-- the other argument may have produced if it had been allowed to
-- complete.
pOr :: GenHaxl u w Bool -> GenHaxl u w Bool -> GenHaxl u w Bool
pOr x y = biselect_opt discrim discrim left right x y
  where
    discrim True = Left ()
    discrim False = Right ()
    left _ = True
    right _ = False

-- | Parallel version of '(.&&)'.  Both arguments are evaluated in
-- parallel, and if either returns 'False' then the other is
-- not evaluated any further.
--
-- WARNING: exceptions may be unpredictable when using 'pAnd'.  If one
-- argument returns 'False' before the other completes, then 'pAnd'
-- returns 'False' immediately, ignoring a possible exception that
-- the other argument may have produced if it had been allowed to
-- complete.
pAnd :: GenHaxl u w Bool -> GenHaxl u w Bool -> GenHaxl u w Bool
pAnd x y = biselect_opt discrim discrim left right x y
  where
    discrim False = Left ()
    discrim True = Right ()
    left _ = False
    right _ = True

-- | This function takes two haxl computations as input, and returns the
-- output of whichever computation finished first. This is clearly
-- non-deterministic in its output and exception behavior, be careful when
-- using it.
unsafeChooseFirst
  :: GenHaxl u w a
  -> GenHaxl u w b
  -> GenHaxl u w (Either a b)
unsafeChooseFirst x y = biselect_opt discrimx discrimy id right x y
  where
    discrimx :: a -> Either (Either a b) ()
    discrimx a = Left (Left a)

    discrimy :: b -> Either (Either a b) ()
    discrimy b = Left (Right b)

    right _ = throw $ CriticalError
      "unsafeChooseFirst: We should never have a 'Right ()'"


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# 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