-- 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 ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE BangPatterns #-}
-- | Bucketing requests by 'DataSource'.
--
-- When a request is issued by the client via 'dataFetch', it is placed
-- in the 'RequestStore'. When we are ready to fetch the current batch
-- of requests, the 'contents' operation extracts the fetches, bucketed
-- by 'DataSource'.
--
-- This module is provided for access to Haxl internals only; most
-- users should not need to import it.
--
module Haxl.Core.RequestStore
  ( BlockedFetches(..)
  , BlockedFetchInternal(..)
  , RequestStore
  , isEmpty
  , noRequests
  , addRequest
  , contents
  , getSize
  , ReqCountMap(..)
  , emptyReqCounts
  , filterRCMap
  , getMapFromRCMap
  , getSummaryMapFromRCMap
  , addToCountMap
  , subFromCountMap
  ) where

import Haxl.Core.DataSource
import Haxl.Core.Stats
import Data.Map (Map)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Map.Strict as Map
import Data.Proxy
import Data.Text (Text)
import Data.Kind (Type)
import Data.Typeable
import Unsafe.Coerce

-- | A container for multiple 'BlockedFetch' objects.
newtype RequestStore u = RequestStore (Map TypeRep (BlockedFetches u))
  -- Since we don't know which data sources we will be using, the store
  -- is dynamically-typed.  It maps the TypeRep of the request to the
  -- 'BlockedFetches' for that 'DataSource'.

newtype BlockedFetchInternal = BlockedFetchInternal CallId

-- | A batch of 'BlockedFetch' objects for a single 'DataSource'
data BlockedFetches u =
  forall r. (DataSource u r) =>
        BlockedFetches [BlockedFetch r] [BlockedFetchInternal]

isEmpty :: RequestStore u -> Bool
isEmpty (RequestStore m) = Map.null m

-- | A new empty 'RequestStore'.
noRequests :: RequestStore u
noRequests = RequestStore Map.empty

-- | Adds a 'BlockedFetch' to a 'RequestStore'.
addRequest
  :: forall u r. (DataSource u r)
  => BlockedFetch r -> BlockedFetchInternal -> RequestStore u -> RequestStore u
addRequest bf bfi (RequestStore m) =
  RequestStore $ Map.insertWith combine ty (BlockedFetches [bf] [bfi]) m
 where
  combine :: BlockedFetches u -> BlockedFetches u -> BlockedFetches u
  combine _ (BlockedFetches bfs bfis)
    | typeOf1 (getR bfs) == ty = BlockedFetches (unsafeCoerce bf:bfs) (bfi:bfis)
    | otherwise                = error "RequestStore.insert"
         -- the dynamic type check here should be unnecessary, but if
         -- there are bugs in `Typeable` or `Map` then we'll get an
         -- error instead of a crash.  The overhead is negligible.

  -- a type conversion only, so we can get the type of the reqeusts from
  -- the list of BlockedFetch.
  getR :: [BlockedFetch r1] -> r1 a
  getR _ = undefined

  -- The TypeRep of requests for this data source
  ty :: TypeRep
  !ty = typeOf1 (undefined :: r a)

-- | Retrieves the whole contents of the 'RequestStore'.
contents :: RequestStore u -> [BlockedFetches u]
contents (RequestStore m) = Map.elems m

getSize :: RequestStore u -> Int
getSize (RequestStore m) = Map.size m

-- A counter to keep track of outgone requests. Entries are added to this
-- map as we send requests to datasources, and removed as these fetches
-- are completed.
-- This is a 2 level map: the 1st level stores requests for a particular
-- datasource, the 2nd level stores count of requests per type.
newtype ReqCountMap = ReqCountMap (Map Text (Map TypeRep Int))
  deriving (Show)

emptyReqCounts :: ReqCountMap
emptyReqCounts = ReqCountMap Map.empty

addToCountMap
  :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
  => Proxy r
  -> Int -- type and number of requests
  -> ReqCountMap
  -> ReqCountMap
addToCountMap = updateCountMap (+)

subFromCountMap
  :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
  => Proxy r
  -> Int -- type and number of requests
  -> ReqCountMap
  -> ReqCountMap
subFromCountMap = updateCountMap (-)

updateCountMap
  :: forall (r :: Type -> Type). (DataSourceName r, Typeable r)
  => (Int -> Int -> Int)
  -> Proxy r
  -> Int -- type and number of requests
  -> ReqCountMap
  -> ReqCountMap
updateCountMap op p n (ReqCountMap m) = ReqCountMap $ Map.insertWith
  (flip (Map.unionWith op)) -- flip is important as "op" is not commutative
  (dataSourceName p) (Map.singleton ty n)
  m
  where
    -- The TypeRep of requests for this data source
    -- The way this is implemented, all elements in the 2nd level map will be
    -- mapped to the same key, as all requests to a datasource have the same
    -- "type". It will be more beneficial to be able to instead map requests
    -- to their names (ie, data constructor) - but there's no cheap way of doing
    -- that.
    ty :: TypeRep
    !ty = typeOf1 (undefined :: r a)

-- Filter all keys with 0 fetches. Since ReqCountMap is a 2-level map, we need
-- nested filter operations.
filterRCMap :: ReqCountMap -> ReqCountMap
filterRCMap (ReqCountMap m) = ReqCountMap $
  Map.filter ((> 0) . Map.size) (Map.filter (> 0) <$> m)

-- Filters the ReqCountMap by default
getMapFromRCMap :: ReqCountMap -> Map Text (Map TypeRep Int)
getMapFromRCMap r
  | ReqCountMap m <- filterRCMap r = m

getSummaryMapFromRCMap :: ReqCountMap -> HashMap.HashMap Text Int
getSummaryMapFromRCMap (ReqCountMap m) = HashMap.fromList
  [ (k, s)
  | (k, v) <- Map.toList m
  , not $ Map.null v
  , let s = sum $ Map.elems v
  , s > 0
  ]


-- 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 RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}

-- | Defines 'runHaxl'.  Most users should import "Haxl.Core" instead.
--
module Haxl.Core.Run
  ( runHaxl
  , runHaxlWithWrites
  ) where

import Control.Concurrent.STM
import Control.Exception as Exception
import Control.Monad
import Data.IORef
import Data.Maybe
import Text.Printf
import Unsafe.Coerce

import Haxl.Core.DataCache
import Haxl.Core.Exception
import Haxl.Core.Flags
import Haxl.Core.Monad
import Haxl.Core.Fetch
import Haxl.Core.Profile
import Haxl.Core.RequestStore as RequestStore
import Haxl.Core.Stats
import Haxl.Core.Util

import qualified Data.HashTable.IO as H

-- -----------------------------------------------------------------------------
-- runHaxl

-- | Runs a 'Haxl' computation in the given 'Env'.
--
-- Note: to make multiple concurrent calls to 'runHaxl', each one must
-- have a separate 'Env'. A single 'Env' must /not/ be shared between
-- multiple concurrent calls to 'runHaxl', otherwise deadlocks or worse
-- will likely ensue.
--
-- However, multiple 'Env's may share a single 'StateStore', and thereby
-- use the same set of datasources.
runHaxl:: forall u w a. Monoid w => Env u w -> GenHaxl u w a -> IO a
runHaxl env haxl = fst <$> runHaxlWithWrites env haxl

runHaxlWithWrites :: forall u w a. Monoid w => Env u w -> GenHaxl u w a -> IO (a, w)
runHaxlWithWrites env@Env{..} haxl = do
  result@IVar{ivarRef = resultRef} <- newIVar -- where to put the final result
  ifTraceLog <- do
    if trace flags < 3
    then return $ \_ -> return ()
    else do
      start <- getTimestamp
      return $ \s -> do
        now <- getTimestamp
        let t = fromIntegral (now - start) / 1000.0 :: Double
        printf "%.1fms: %s" t (s :: String)
  let
    -- Run a job, and put its result in the given IVar
    schedule :: Env u w -> JobList u w -> GenHaxl u w b -> IVar u w b -> IO ()
    schedule env@Env{..} rq (GenHaxl run) ivar@IVar{ivarRef = !ref} = do
      ifTraceLog $ printf "schedule: %d\n" (1 + lengthJobList rq)
      let {-# INLINE result #-}
          result r = do
            e <- readIORef ref
            case e of
              IVarFull _ ->
                -- 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)
                reschedule env rq
              IVarEmpty haxls -> do
                writeIORef ref (IVarFull r)
                -- Have we got the final result now?
                if ref == unsafeCoerce resultRef
                        -- comparing IORefs of different types is safe, it's
                        -- pointer-equality on the MutVar#.
                   then
                     -- We have a result, but don't discard unfinished
                     -- computations in the run queue. See
                     -- Note [runHaxl and unfinished requests].
                     -- Nothing can depend on the final IVar, so haxls must
                     -- be empty.
                     case rq of
                       JobNil -> return ()
                       _ -> modifyIORef' runQueueRef (appendJobList rq)
                   else reschedule env (appendJobList haxls rq)
      r <-
        if testReportFlag ReportProfiling $ report flags  -- withLabel unfolded
          then Exception.try $ profileCont run env
          else Exception.try $ run env
      case r of
        Left e -> do
          rethrowAsyncExceptions e
          result (ThrowIO e)
        Right (Done a) -> do
          wt <- readIORef writeLogsRef
          result $ Ok a (Just wt)
        Right (Throw ex) -> do
          wt <- readIORef writeLogsRef
          result $ ThrowHaxl ex (Just wt)
        Right (Blocked i fn) -> do
          addJob env (toHaxl fn) ivar i
          reschedule env rq

    -- Here we have a choice:
    --   - If the requestStore is non-empty, we could submit those
    --     requests right away without waiting for more.  This might
    --     be good for latency, especially if the data source doesn't
    --     support batching, or if batching is pessimal.
    --   - To optimise the batch sizes, we want to execute as much as
    --     we can and only submit requests when we have no more
    --     computation to do.
    --   - compromise: wait at least Nms for an outstanding result
    --     before giving up and submitting new requests.
    --
    -- For now we use the batching strategy in the scheduler, but
    -- individual data sources can request that their requests are
    -- sent eagerly by using schedulerHint.
    --
    reschedule :: Env u w -> JobList u w -> IO ()
    reschedule env@Env{..} haxls = do
      case haxls of
        JobNil -> do
          rq <- readIORef runQueueRef
          case rq of
            JobNil -> emptyRunQueue env
            JobCons env' a b c -> do
              writeIORef runQueueRef JobNil
              schedule env' c a b
        JobCons env' a b c ->
          schedule env' c a b

    emptyRunQueue :: Env u w -> IO ()
    emptyRunQueue env = do
      ifTraceLog $ printf "emptyRunQueue\n"
      haxls <- checkCompletions env
      case haxls of
        JobNil -> checkRequestStore env
        _ -> reschedule env haxls

    checkRequestStore :: Env u w -> IO ()
    checkRequestStore env@Env{..} = do
      ifTraceLog $ printf "checkRequestStore\n"
      reqStore <- readIORef reqStoreRef
      if RequestStore.isEmpty reqStore
        then waitCompletions env
        else do
          ifTraceLog $ printf "performFetches %d\n" (RequestStore.getSize reqStore)
          writeIORef reqStoreRef noRequests
          performRequestStore env reqStore
          -- empty the cache if we're not caching.  Is this the best
          -- place to do it?  We do get to de-duplicate requests that
          -- happen simultaneously.
          when (caching flags == 0) $ do
            let DataCache dc = dataCache
            H.foldM (\_ (k, _) -> H.delete dc k) () dc
          emptyRunQueue env

    checkCompletions :: Env u w -> IO (JobList u w)
    checkCompletions Env{..} = do
      ifTraceLog $ printf "checkCompletions\n"
      comps <- atomicallyOnBlocking (LogicBug ReadingCompletionsFailedRun) $ do
        c <- readTVar completions
        writeTVar completions []
        return c
      case comps of
        [] -> return JobNil
        _ -> do
          ifTraceLog $ printf "%d complete\n" (length comps)
          let
              getComplete (CompleteReq a IVar{ivarRef = !cr} allocs) = do
                when (allocs < 0) $ do
                  cur <- getAllocationCounter
                  setAllocationCounter (cur + allocs)
                r <- readIORef cr
                case r of
                  IVarFull _ -> do
                    ifTraceLog $ printf "existing result\n"
                    return JobNil
                    -- this happens if a data source reports a result,
                    -- and then throws an exception.  We call putResult
                    -- a second time for the exception, which comes
                    -- ahead of the original request (because it is
                    -- pushed on the front of the completions list) and
                    -- therefore overrides it.
                  IVarEmpty cv -> do
                    writeIORef cr (IVarFull a)
                    return cv
          jobs <- mapM getComplete comps
          return (foldr appendJobList JobNil jobs)

    waitCompletions :: Env u w -> IO ()
    waitCompletions env@Env{..} = do
      ifTraceLog $ printf "waitCompletions\n"
      let
        wrapped = atomicallyOnBlocking (LogicBug ReadingCompletionsFailedRun)
        doWait = wrapped $ do
          c <- readTVar completions
          when (null c) retry
        doWaitProfiled = do
          queueEmpty <- null <$> wrapped (readTVar completions)
          when queueEmpty $ do
            -- Double check the queue as we want to make sure that
            -- submittedReqsRef is copied before waiting on the queue but as a
            -- fast path do not want to copy it if the queue is empty.
            -- There is still a race oppoortunity as submittedReqsRef is
            -- decremented in whatever thread the completion happens, and so it
            -- is possible for waitingOn to be empty while queueEmpty2 is True.
            waitingOn <- readIORef submittedReqsRef
            queueEmpty2 <- null <$> wrapped (readTVar completions)
            when queueEmpty2 $ do
              start <- getTimestamp
              doWait
              end <- getTimestamp
              let fw = FetchWait
                        { fetchWaitReqs = getSummaryMapFromRCMap waitingOn
                        , fetchWaitStart = start
                        , fetchWaitDuration = (end-start)
                        }
              modifyIORef' statsRef $ \(Stats s) -> Stats (fw:s)
      if testReportFlag ReportFetchStats $ report flags
        then doWaitProfiled
        else doWait
      emptyRunQueue env

  --
  schedule env JobNil haxl result
  r <- readIORef resultRef
  writeIORef writeLogsRef mempty
  wtNoMemo <- atomicModifyIORef' writeLogsRefNoMemo
    (\old_wrts -> (mempty, old_wrts))
  case r of
    IVarEmpty _ -> throwIO (CriticalError "runHaxl: missing result")
    IVarFull (Ok a wt) -> do
      return (a, fromMaybe mempty wt <> wtNoMemo)
    IVarFull (ThrowHaxl e _wt)  -> throwIO e
      -- The written logs are discarded when there's a Haxl exception. We
      -- can change this behavior if we need to get access to partial logs.
    IVarFull (ThrowIO e)  -> throwIO e

{- Note [runHaxl and unfinished requests]

runHaxl returns immediately when the supplied computation has returned
a result.  This doesn't necessarily mean that the whole computation
graph has completed, however.  In particular, when using pAnd and pOr,
we might have created some data fetches that have not completed, but
weren't required, because the other branch of the pAnd/pOr subsumed
the result.

When runHaxl returns, it might be that:
- reqStoreRef contains some unsubmitted requests
- runQueueRef contains some jobs
- there are in-flight BackgroundFetch requests, that will return their
  results to the completions queue in due course.
- there are various unfilled IVars in the cache and/or memo tables

This should be all safe, we can even restart runHaxl with the same Env
after it has stopped and the in-progress computations will
continue. But don't discard the contents of
reqStoreRef/runQueueRef/completions, because then we'll deadlock if we
discover one of the unfilled IVars in the cache or memo table.
-}

{- TODO: later
data SchedPolicy
  = SubmitImmediately
  | WaitAtLeast Int{-ms-}
  | WaitForAllPendingRequests
-}

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

instance Exception ReadingCompletionsFailedRun


-- 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.

-- |
-- Most users should import "Haxl.Core" instead of importing this
-- module directly.
--
module Haxl.Core.ShowP
  ( ShowP(..)
  ) where

-- | A class of type constructors for which we can show all
-- parameterizations.
class ShowP f where
  showp :: f a -> String


-- 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 ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}

-- |
-- Most users should import "Haxl.Core" instead of importing this
-- module directly.
--
module Haxl.Core.StateStore
  ( StateKey(..)
  , StateStore
  , stateGet
  , stateSet
  , stateEmpty
  ) where

import Data.Map (Map)
import Data.Kind (Type)
import qualified Data.Map.Strict as Map
#if __GLASGOW_HASKELL__ < 804
import Data.Monoid
#endif
import Data.Typeable
import Unsafe.Coerce

-- | 'StateKey' maps one type to another type. A type that is an
-- instance of 'StateKey' can store and retrieve information from a
-- 'StateStore'.
--
class Typeable f => StateKey (f :: Type -> Type) where
  data State f

  -- | We default this to typeOf1, but if f is itself a complex type that is
  -- already applied to some paramaters, we want to be able to use the same
  -- state by using typeOf2, etc
  getStateType :: Proxy f -> TypeRep
  getStateType = typeRep

-- | The 'StateStore' maps a 'StateKey' to the 'State' for that type.
newtype StateStore = StateStore (Map TypeRep StateStoreData)

#if __GLASGOW_HASKELL__ >= 804
instance Semigroup StateStore where
  -- Left-biased union
  StateStore m1 <> StateStore m2 = StateStore $ m1 <> m2
#endif

instance Monoid StateStore where
  mempty = stateEmpty
#if __GLASGOW_HASKELL__ < 804
  mappend (StateStore m1) (StateStore m2) = StateStore $ m1 <> m2
#endif

-- | Encapsulates the type of 'StateStore' data so we can have a
-- heterogeneous collection.
data StateStoreData = forall f. StateKey f => StateStoreData (State f)

-- | A `StateStore` with no entries.
stateEmpty :: StateStore
stateEmpty = StateStore Map.empty

-- | Inserts a `State` in the `StateStore` container.
stateSet :: forall f . StateKey f => State f -> StateStore -> StateStore
stateSet st (StateStore m) =
  StateStore (Map.insert (getStateType (Proxy :: Proxy f)) (StateStoreData st) m)

-- | Retrieves a `State` from the `StateStore` container.
stateGet :: forall r . StateKey r => StateStore -> Maybe (State r)
stateGet (StateStore m) =
  case Map.lookup ty m of
     Nothing -> Nothing
     Just (StateStoreData (st :: State f))
       | getStateType (Proxy :: Proxy f) == ty -> Just (unsafeCoerce st)
       | otherwise             -> Nothing
          -- the dynamic type check here should be unnecessary, but if
          -- there are bugs in `Typeable` or `Map` then we'll get an
          -- error instead of a crash.  The overhead is a few percent.
 where
  ty = getStateType (Proxy :: Proxy r)


-- 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 OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE CPP #-}

-- |
-- Types and operations for statistics and profiling.  Most users
-- should import "Haxl.Core" instead of importing this module
-- directly.
--
module Haxl.Core.Stats
  (
  -- * Data-source stats
    Stats(..)
  , CallId
  , FetchStats(..)
  , Microseconds
  , Timestamp
  , DataSourceStats(..)
  , getTimestamp
  , emptyStats
  , numFetches
  , ppStats
  , ppFetchStats
  , aggregateFetchBatches

  -- * Profiling
  , Profile(..)
  , ProfileMemo(..)
  , ProfileFetch(..)
  , emptyProfile
  , ProfileKey
  , ProfileLabel
  , ProfileData(..)
  , emptyProfileData
  , AllocCount
  , LabelHitCount

  -- * Allocation
  , getAllocationCounter
  , setAllocationCounter
  ) where

import Data.Aeson
import Data.Function (on)
import Data.Maybe (mapMaybe)
import Data.HashMap.Strict (HashMap)
import Data.Int
import Data.List (intercalate, sortOn, groupBy)
#if __GLASGOW_HASKELL__ < 804
import Data.Semigroup (Semigroup)
#endif
import Data.Ord (Down(..))
import Data.Text (Text)
import Data.Time.Clock.POSIX
import Data.Typeable
import Text.Printf
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Text as Text

import GHC.Conc (getAllocationCounter, setAllocationCounter)

-- ---------------------------------------------------------------------------
-- Measuring time

type Microseconds = Int64
type Timestamp = Microseconds -- since an epoch

getTimestamp :: IO Timestamp
getTimestamp = do
  t <- getPOSIXTime -- for now, TODO better
  return (round (t * 1000000))

-- ---------------------------------------------------------------------------
-- Stats

data DataSourceStats =
  forall a. (Typeable a, Show a, Eq a, ToJSON a) => DataSourceStats a

instance Show DataSourceStats where
  show (DataSourceStats x) = printf "DataSourceStats %s" (show x)

instance Eq DataSourceStats where
  (==) (DataSourceStats a) (DataSourceStats b) =
    cast a == Just b

-- | Stats that we collect along the way.
newtype Stats = Stats [FetchStats]
  deriving (Show, ToJSON, Semigroup, Monoid)

-- | Pretty-print Stats.
ppStats :: Stats -> String
ppStats (Stats rss) =
  intercalate "\n"
    [ "["
    ++ [
      if fetchWasRunning rs
          (minStartTime + (t - 1) * usPerDash)
          (minStartTime + t * usPerDash)
        then fetchSymbol rs
        else '-'
      | t <- [1..numDashes]
      ]
    ++ "] " ++ show i ++ " - " ++ ppFetchStats rs
    | (i, rs) <- zip [(1::Int)..] validFetchStats ]
  where
    isFetchStats FetchStats{} = True
    isFetchStats FetchWait{} = True
    isFetchStats FetchDataSourceStats{} = True
    isFetchStats _ = False
    validFetchStats = filter isFetchStats (reverse rss)
    numDashes = 50
    getStart FetchStats{..} = Just fetchStart
    getStart FetchWait{..} = Just fetchWaitStart
    getStart _ = Nothing
    getEnd FetchStats{..} = Just $ fetchStart + fetchDuration
    getEnd FetchWait{..} = Just $ fetchWaitStart + fetchWaitDuration
    getEnd _ = Nothing
    minStartTime = minimum $ mapMaybe getStart validFetchStats
    endTime = maximum $ mapMaybe getEnd validFetchStats
    usPerDash = (endTime - minStartTime) `div` numDashes
    fetchSymbol FetchStats{} = '*'
    fetchSymbol FetchWait{} = '.'
    fetchSymbol _ = '?'
    fetchWasRunning :: FetchStats -> Timestamp -> Timestamp -> Bool
    fetchWasRunning fs@FetchStats{} t1 t2 =
      (fetchStart fs + fetchDuration fs) >= t1 && fetchStart fs < t2
    fetchWasRunning fw@FetchWait{} t1 t2 =
      (fetchWaitStart fw + fetchWaitDuration fw) >= t1 && fetchWaitStart fw < t2
    fetchWasRunning _ _ _ = False

type CallId = Int

-- | Maps data source name to the number of requests made in that round.
-- The map only contains entries for sources that made requests in that
-- round.
data FetchStats
    -- | Timing stats for a (batched) data fetch
  = FetchStats
    { fetchDataSource :: Text
    , fetchBatchSize :: {-# UNPACK #-} !Int
    , fetchStart :: {-# UNPACK #-} !Timestamp
    , fetchDuration :: {-# UNPACK #-} !Microseconds
    , fetchSpace :: {-# UNPACK #-} !Int64
    , fetchFailures :: {-# UNPACK #-} !Int
    , fetchIgnoredFailures :: {-# UNPACK #-} !Int
    , fetchBatchId :: {-# UNPACK #-} !Int
    , fetchIds :: [CallId]
    }

    -- | The stack trace of a call to 'dataFetch'.  These are collected
    -- only when profiling and reportLevel is 5 or greater.
  | FetchCall
    { fetchReq :: String
    , fetchStack :: [String]
    , fetchStatId :: {-# UNPACK #-} !CallId
    }
  | MemoCall
    { memoStatId :: {-# UNPACK #-} !CallId
    , memoSpace :: {-# UNPACK #-} !Int64
    }
  | FetchWait
    { fetchWaitReqs :: HashMap Text Int
       -- ^ What DataSources had requests that were being waited for
    , fetchWaitStart :: {-# UNPACK #-} !Timestamp
    , fetchWaitDuration :: {-# UNPACK #-} !Microseconds
    }
  | FetchDataSourceStats
    { fetchDsStatsCallId :: CallId
    , fetchDsStatsDataSource :: Text
    , fetchDsStatsStats :: DataSourceStats
    , fetchBatchId :: {-# UNPACK #-} !Int
    }
  deriving (Eq, Show)

-- | Pretty-print RoundStats.
ppFetchStats :: FetchStats -> String
ppFetchStats FetchStats{..} =
  printf "%s: %d fetches (%.2fms, %d bytes, %d failures)"
    (Text.unpack fetchDataSource) fetchBatchSize
    (fromIntegral fetchDuration / 1000 :: Double)  fetchSpace fetchFailures
ppFetchStats (FetchCall r ss _) = show r ++ '\n':show ss
ppFetchStats MemoCall{} = ""
ppFetchStats FetchWait{..}
  | HashMap.size fetchWaitReqs == 0 = msg "unexpected: Blocked on nothing"
  | HashMap.size fetchWaitReqs <= 2 =
    msg $ printf "Blocked on %s"
      (intercalate "," [printf "%s (%d reqs)" ds c
                       | (ds,c) <- HashMap.toList fetchWaitReqs])
  | otherwise = msg $ printf "Blocked on %d sources (%d reqs)"
                        (HashMap.size fetchWaitReqs)
                        (sum $ HashMap.elems fetchWaitReqs)
  where
    msg :: String -> String
    msg x = printf "%s (%.2fms)"
                x
                (fromIntegral fetchWaitDuration / 1000 :: Double)
ppFetchStats FetchDataSourceStats{..} =
  printf "%s (stats): %s" (Text.unpack fetchDsStatsDataSource)
    (show fetchDsStatsStats)

-- | Aggregate stats merging FetchStats from the same dispatched batch into one.
aggregateFetchBatches :: ([FetchStats] -> a) -> Stats -> [a]
aggregateFetchBatches agg (Stats fetches) =
      map agg $
      groupBy ((==) `on` fetchBatchId) $
      sortOn (Down . fetchBatchId)
      [f | f@FetchStats{} <- fetches]

instance ToJSON FetchStats where
  toJSON FetchStats{..} = object
    [ "type" .= ("FetchStats" :: Text)
    , "datasource" .= fetchDataSource
    , "fetches" .= fetchBatchSize
    , "start" .= fetchStart
    , "duration" .= fetchDuration
    , "allocation" .= fetchSpace
    , "failures" .= fetchFailures
    , "ignoredFailures" .= fetchIgnoredFailures
    , "batchid" .= fetchBatchId
    , "fetchids" .= fetchIds
    ]
  toJSON (FetchCall req strs fid) = object
    [ "type" .= ("FetchCall" :: Text)
    , "request" .= req
    , "stack" .= strs
    , "fetchid" .= fid
    ]
  toJSON (MemoCall cid allocs) = object
    [ "type" .= ("MemoCall" :: Text)
    , "callid" .= cid
    , "allocation" .= allocs
    ]
  toJSON FetchWait{..} = object
    [ "type" .= ("FetchWait" :: Text)
    , "duration" .= fetchWaitDuration
    ]
  toJSON FetchDataSourceStats{..} = object
    [ "type" .= ("FetchDataSourceStats" :: Text)
    , "datasource" .= fetchDsStatsDataSource
    , "stats" .= sjson fetchDsStatsStats
    , "batchid" .= fetchBatchId
    ]
    where
      sjson (DataSourceStats s) = toJSON s

emptyStats :: Stats
emptyStats = Stats []

numFetches :: Stats -> Int
numFetches (Stats rs) = sum [ fetchBatchSize | FetchStats{..} <- rs ]

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

type ProfileLabel = Text
type AllocCount = Int64
type LabelHitCount = Int64
type ProfileKey = Int64

data ProfileFetch = ProfileFetch
  { profileFetchFetchId :: {-# UNPACK #-} !CallId
  , profileFetchMemoId ::  {-# UNPACK #-} !CallId
  , profileFetchWasCached :: !Bool
  }
  deriving (Show, Eq)

data ProfileMemo = ProfileMemo
  { profileMemoId :: {-# UNPACK #-} !CallId
  , profileMemoWasCached :: !Bool
  }
  deriving (Show, Eq)

data Profile = Profile
  { profile      :: HashMap ProfileKey ProfileData
     -- ^ Data per key (essentially per call stack)
  , profileTree :: HashMap (ProfileLabel, ProfileKey) ProfileKey
     -- ^ (label, parent) -> current. The exception is the root which will have
     -- ("MAIN", 0) -> 0
  , profileNextKey :: ProfileKey
     -- ^ Provides a unique key per callstack
  }

emptyProfile :: Profile
emptyProfile = Profile HashMap.empty (HashMap.singleton ("MAIN", 0) 0) 1

data ProfileData = ProfileData
  { profileAllocs :: {-# UNPACK #-} !AllocCount
     -- ^ allocations made by this label
  , profileFetches :: [ProfileFetch]
     -- ^ fetches made in this label
  , profileLabelHits :: {-# UNPACK #-} !LabelHitCount
     -- ^ number of hits at this label
  , profileMemos :: [ProfileMemo]
     -- ^ memo and a boolean representing if it was cached at the time
  , profileTime :: {-# UNPACK #-} !Microseconds
     -- ^ amount of time spent in computation at this label
  }
  deriving Show

emptyProfileData :: ProfileData
emptyProfileData = ProfileData 0 [] 0 [] 0


-- 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.

-- | Internal utilities only.
--
{-# OPTIONS_GHC -fno-warn-unused-imports #-}
module Haxl.Core.Util
  ( atomicallyOnBlocking
  , compose
  , textShow
  , trace_
  ) where

import Data.Text (Text)
import Debug.Trace (trace)
import qualified Data.Text as Text

import Control.Concurrent.STM
import Control.Exception

atomicallyOnBlocking :: Exception e => e -> STM a -> IO a
atomicallyOnBlocking e stm =
  catch (atomically stm)
        (\BlockedIndefinitelyOnSTM -> throw e)

-- | Composes a list of endofunctions.
compose :: [a -> a] -> a -> a
compose = foldr (.) id

textShow :: (Show a) => a -> Text
textShow = Text.pack . show

-- | This function can be used to trace a bunch of lines to stdout when
-- debugging haxl core.
trace_ :: String -> a -> a
trace_ _ = id
--trace_ = Debug.Trace.trace


-- 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 DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- A generic Haxl datasource for performing arbitrary IO concurrently.
-- Every IO operation will be performed in a separate thread.
-- You can use this with any kind of IO, but each different operation
-- requires an instance of the 'ConcurrentIO' class.
--
-- For example, to make a concurrent sleep operation:
--
-- > sleep :: Int -> GenHaxl u w Int
-- > sleep n = dataFetch (Sleep n)
-- >
-- > data Sleep
-- > instance ConcurrentIO Sleep where
-- >   data ConcurrentIOReq Sleep a where
-- >     Sleep :: Int -> ConcurrentIOReq Sleep Int
-- >
-- >   performIO (Sleep n) = threadDelay (n*1000) >> return n
-- >
-- > deriving instance Eq (ConcurrentIOReq Sleep a)
-- > deriving instance Show (ConcurrentIOReq Sleep a)
-- >
-- > instance ShowP (ConcurrentIOReq Sleep) where showp = show
-- >
-- > instance Hashable (ConcurrentIOReq Sleep a) where
-- >   hashWithSalt s (Sleep n) = hashWithSalt s n
--
-- Note that you can have any number of constructors in your
-- ConcurrentIOReq GADT, so most of the boilerplate only needs to be
-- written once.

module Haxl.DataSource.ConcurrentIO
  ( mkConcurrentIOState
  , ConcurrentIO(..)
  ) where

import Control.Concurrent
import Control.Monad
import qualified Data.Text as Text
import Data.Typeable

import Haxl.Core

class ConcurrentIO tag where
  data ConcurrentIOReq tag a
  performIO :: ConcurrentIOReq tag a -> IO a

deriving instance Typeable ConcurrentIOReq -- not needed by GHC 7.10 and later

instance (Typeable tag) => StateKey (ConcurrentIOReq tag) where
  data State (ConcurrentIOReq tag) = ConcurrentIOState
  getStateType _ = typeRep (Proxy :: Proxy ConcurrentIOReq)

mkConcurrentIOState :: IO (State (ConcurrentIOReq ()))
mkConcurrentIOState = return ConcurrentIOState

instance Typeable tag => DataSourceName (ConcurrentIOReq tag) where
  dataSourceName _ =
    Text.pack (show (typeRepTyCon (typeRep (Proxy :: Proxy tag))))

instance
  (Typeable tag, ShowP (ConcurrentIOReq tag), ConcurrentIO tag)
  => DataSource u (ConcurrentIOReq tag)
 where
  fetch _state _flags _u = BackgroundFetch $ \bfs -> do
    forM_ bfs $ \(BlockedFetch req rv) ->
      forkFinally (performIO req) (putResultFromChildThread rv)


-- 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 RebindableSyntax, OverloadedStrings, ApplicativeDo #-}
module AdoTests (tests) where

import TestUtils
import MockTAO

import Control.Applicative
import Test.HUnit

import Prelude()
import Haxl.Prelude

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

--
-- Test ApplicativeDo batching
--
ado1 = expectResult 12 ado1_

ado1_ = do
  a <- friendsOf =<< id1
  b <- friendsOf =<< id2
  return (length (a ++ b))

ado2 = expectResult 12 ado2_

ado2_ = do
  x <- id1
  a <- friendsOf x
  y <- id2
  b <- friendsOf y
  return (length (a ++ b))

ado3 = expectResult 11 ado3_

ado3_ = do
  x <- id1
  a <- friendsOf x
  a' <- friendsOf =<< if null a then id3 else id4
  y <- id2
  b <- friendsOf y
  b' <- friendsOf  =<< if null b then id4 else id3
  return (length (a' ++ b'))

tests future = TestList
  [ TestLabel "ado1" $ TestCase (ado1 future)
  , TestLabel "ado2" $ TestCase (ado2 future)
  , TestLabel "ado3" $ TestCase (ado3 future)
  ]


-- 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 OverloadedStrings #-}
module AllTests (allTests) where

import TestExampleDataSource
import BatchTests
import CoreTests
import DataCacheTest
import ExceptionStackTests
import AdoTests
import OutgoneFetchesTests
import ProfileTests
import MemoizationTests
import MonadAsyncTest
import TestBadDataSource
import FullyAsyncTest
import WriteTests
import ParallelTests
import StatsTests
import DataSourceDispatchTests

import Test.HUnit

allTests :: Test
allTests = TestList
  [ TestLabel "ExampleDataSource" TestExampleDataSource.tests
  , TestLabel "BatchTests-future" $ BatchTests.tests True
  , TestLabel "BatchTests-sync" $ BatchTests.tests False
  , TestLabel "CoreTests" CoreTests.tests
  , TestLabel "DataCacheTests" DataCacheTest.tests
  , TestLabel "ExceptionStackTests" ExceptionStackTests.tests
  , TestLabel "AdoTests" $ AdoTests.tests False
  , TestLabel "OutgoneFetchesTest" OutgoneFetchesTests.tests
  , TestLabel "ProfileTests" ProfileTests.tests
  , TestLabel "MemoizationTests" MemoizationTests.tests
  , TestLabel "MonadAsyncTests" MonadAsyncTest.tests
  , TestLabel "BadDataSourceTests" TestBadDataSource.tests
  , TestLabel "FullyAsyncTest" FullyAsyncTest.tests
  , TestLabel "WriteTest" WriteTests.tests
  , TestLabel "ParallelTest" ParallelTests.tests
  , TestLabel "StatsTests" StatsTests.tests
  , TestLabel "DataSourceDispatchTests" DataSourceDispatchTests.tests
  ]


-- 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 DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards #-}

-- | A data source that can be made to fail in various ways, for testing

module BadDataSource (
    -- * initialise the state
    State(..), initGlobalState, FetchImpl(..),

    -- * requests for this data source
    FailAfter(..)
  ) where

import Haxl.Prelude
import Prelude ()

import Haxl.Core

import Control.Exception
import Data.Typeable
import Data.Hashable
import Control.Concurrent
import Control.Monad (void)

import GHC.Conc ( PrimMVar )
import Foreign.StablePtr
import Foreign.C.Types ( CInt(..) )

foreign import ccall safe
  hs_try_putmvar :: CInt -> StablePtr PrimMVar -> IO ()

data FetchImpl =
  Async
  | Background
  | BackgroundMVar
  | BackgroundSeq
  | BackgroundPar

data FailAfter a where
  FailAfter :: Int -> FailAfter Int
  deriving Typeable

deriving instance Eq (FailAfter a)
deriving instance Show (FailAfter a)
instance ShowP FailAfter where showp = show

instance Hashable (FailAfter a) where
   hashWithSalt s (FailAfter a) = hashWithSalt s (0::Int,a)

instance StateKey FailAfter where
  data State FailAfter = FailAfterState
         { failAcquireDelay :: Int
         , failAcquire :: IO ()
         , failReleaseDelay :: Int
         , failRelease :: IO ()
         , failDispatchDelay :: Int
         , failDispatch :: IO ()
         , failWaitDelay :: Int
         , failWait :: IO ()
         , failImpl :: FetchImpl
        }

instance DataSourceName FailAfter where
  dataSourceName _ = "BadDataSource"

instance DataSource u FailAfter where
  fetch state@FailAfterState{..}
    | BackgroundSeq <- failImpl = backgroundFetchSeq runOne state
    | BackgroundPar <- failImpl = backgroundFetchPar runOne state
    | Background <- failImpl = backgroundFetchAcquireRelease
        acquire release dispatchbg wait
        submit state
    | BackgroundMVar <- failImpl = backgroundFetchAcquireReleaseMVar
        acquire release dispatchbgMVar wait
        submit state
    | Async <- failImpl = asyncFetchAcquireRelease
       acquire release dispatch wait
       submit state
   where
     acquire = do threadDelay failAcquireDelay; failAcquire
     release _ = do threadDelay failReleaseDelay; failRelease
     dispatch _ = do threadDelay failDispatchDelay; failDispatch
     dispatchBase put = (do
                          failDispatch
                          _ <- mask_ $ forkIO $ finally
                            (threadDelay failDispatchDelay)
                            put
                          return ()) `onException` put
     dispatchbg _ c m = dispatchBase (hs_try_putmvar (fromIntegral c) m)
     dispatchbgMVar _ _ m = dispatchBase (void $ tryPutMVar m ())
     wait _ = do threadDelay failWaitDelay; failWait
     submit :: () -> FailAfter a -> IO (IO (Either SomeException a))
     submit _ (FailAfter t) = do
       threadDelay t
       return (return (Left (toException (FetchError "failed request"))))
     runOne :: FailAfter a -> IO (Either SomeException a)
     runOne r = do
       bracket acquire release $ \s -> do
         dispatch s
         getRes <- submit s r
         wait s
         getRes

initGlobalState :: FetchImpl -> IO (State FailAfter)
initGlobalState impl = do
  return FailAfterState
    { failAcquireDelay = 0
    , failAcquire = return ()
    , failReleaseDelay = 0
    , failRelease = return ()
    , failDispatchDelay = 0
    , failDispatch = return ()
    , failWaitDelay = 0
    , failWait = return ()
    , failImpl = impl
    }


-- 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 RebindableSyntax, OverloadedStrings #-}
module BatchTests (tests) where

import TestTypes
import TestUtils
import MockTAO

import Control.Applicative
import Test.HUnit

import Haxl.Core

import Prelude()
import Haxl.Prelude
import Data.IORef

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

--
-- Test batching over multiple arguments in liftA2
--
batching1 = expectResult 12 batching1_

batching1_ = do
  a <- id1
  b <- id2
  length <$> liftA2 (++) (friendsOf a) (friendsOf b)

--
-- Test batching in mapM (which is really traverse)
--
batching2 = expectResult 12 batching2_

batching2_ = do
  a <- id1
  b <- id2
  fs <- mapM friendsOf [a,b]
  return (sum (map length fs))

--
-- Test batching when we have a monadic bind in each branch
--
batching3 = expectResult 12 batching3_

batching3_ = do
  let a = id1 >>= friendsOf
      b = id2 >>= friendsOf
  length <$> a .++ b

--
-- Test batching over both arguments of (+)
--
batching4 = expectResult 12 batching4_

batching4_ = do
  let a = length <$> (id1 >>= friendsOf)
      b = length <$> (id2 >>= friendsOf)
  a + b

--
-- Test batching over both arguments of (+)
--
batching5 = expectResult 2 batching5_

batching5_ :: Haxl Int
batching5_ = if a .> b then 1 else 2
 where
  a = length <$> (id1 >>= friendsOf)
  b = length <$> (id2 >>= friendsOf)

--
-- Test batching when we perform all batching tests together with sequence
--
batching6 = expectResult [12,12,12,12,2] batching6_

batching6_ = sequence [batching1_,batching2_,batching3_,batching4_,batching5_]

--
-- Ensure if/then/else and bool operators break batching
--
batching7 = expectResult 12 batching7_

batching7_ :: Haxl Int
batching7_ = if a .> 0 then a+b else 0
 where
  a = length <$> (id1 >>= friendsOf)
  b = length <$> (id2 >>= friendsOf)

-- We expect 3 rounds here due to boolean operators
batching8 = expectResult 12 batching8_

batching8_ :: Haxl Int
batching8_ = if (c .== 0) .|| (a .> 0 .&& b .> 0) then a+b else 0
 where
  a = length <$> (id1 >>= friendsOf)
  b = length <$> (id2 >>= friendsOf)
  c = length <$> (id3 >>= friendsOf)

-- (>>) should batch, so we expect one round
batching9 = expectResult 6 batching9_

batching9_ :: Haxl Int
batching9_ = (id1 >>= friendsOf) >> (length <$> (id2 >>= friendsOf))

--
-- Test data caching, numFetches
--

-- simple (one cache hit)
caching1 = expectFetches 3 caching1_
caching1_ = nf id1 + nf id2 + nf id3 + nf id3
 where
  nf id = length <$> (id >>= friendsOf)

-- simple, in rounds (no cache hits)
caching2 = expectFetches 3 caching2_
caching2_ = if nf id1 .> 0 then nf id2 + nf id3 else 0
 where
  nf id = length <$> (id >>= friendsOf)

-- rounds (one cache hit)
caching3 = expectFetches 3 caching3_
caching3_ = if nf id1 .> 0 then nf id1 + nf id2 + nf id3 else 0
 where
  nf id = length <$> (id >>= friendsOf)

--
-- Basic sanity check on data-cache re-use
--
cacheReuse future = do
  env <- makeTestEnv future
  expectResultWithEnv 12 batching7_ env

  -- make a new env
  tao <- MockTAO.initGlobalState future
  let st = stateSet tao stateEmpty
  env2 <- initEnvWithData st testinput (caches env)
  cid <- readIORef (callIdRef env2)
  assertBool "callId is unique" (cid > 0)

  -- ensure no more data fetching rounds needed
  expectResultWithEnv 12 batching7_ env2

noCaching future = do
  env <- makeTestEnv future
  let env' = env{ flags = (flags env){caching = 0} }
  result <- runHaxl env' caching3_
  assertEqual "result" result 18
  stats <- readIORef (statsRef env)
  assertEqual "fetches" 4 (numFetches stats)

exceptionTest1 = expectResult []
  $ withDefault [] $ friendsOf 101

exceptionTest2 = expectResult [7..12] $ liftA2 (++)
  (withDefault [] (friendsOf 101))
  (withDefault [] (friendsOf 2))

deterministicExceptions future = do
  env <- makeTestEnv future
  let haxl =
        sequence [ do _ <- friendsOf =<< id1; throw (NotFound "xxx")
                 , throw (NotFound "yyy")
                 ]
  -- the first time, friendsOf should block, but we should still get the
  -- "xxx" exception.
  r <- runHaxl env $ try haxl
  assertBool "exceptionTest3" $
    case r of
     Left (NotFound "xxx") -> True
     _ -> False
  -- the second time, friendsOf will be cached, and we should get the "xxx"
  -- exception as before.
  r <- runHaxl env $ try haxl
  assertBool "exceptionTest3" $
    case r of
     Left (NotFound "xxx") -> True
     _ -> False

pOrTests future = do
  env <- makeTestEnv future

  -- Test semantics
  r <- runHaxl env $ do
        a <- return False `pOr` return False
        b <- return False `pOr` return True
        c <- return True `pOr` return False
        d <- return True `pOr` return True
        return (not a && b && c && d)
  assertBool "pOr0" r

  -- pOr is left-biased with respect to exceptions:
  r <- runHaxl env $ try $ return True `pOr` throw (NotFound "foo")
  assertBool "pOr1" $
    case (r :: Either NotFound Bool) of
      Right True -> True
      _ -> False
  r <- runHaxl env $ try $ throw (NotFound "foo") `pOr` return True
  assertBool "pOr2" $
    case (r :: Either NotFound Bool) of
      Left (NotFound "foo") -> True
      _ -> False

  -- pOr is non-deterministic (see also Note [tricky pOr/pAnd])
  let nondet = (do _ <- friendsOf 1; throw (NotFound "foo")) `pOr` return True
  r <- runHaxl env $ try nondet
  assertBool "pOr3" $
    case (r :: Either NotFound Bool) of
      Right True -> True
      _ -> False
  -- next we populate the cache
  _ <- runHaxl env $ friendsOf 1
  -- and now exactly the same pOr again will throw this time:
  r <- runHaxl env $ try nondet
  assertBool "pOr4" $
    case (r :: Either NotFound Bool) of
      Left (NotFound "foo") -> True
      _ -> False

  -- One more test: Blocked/False => Blocked
  r <- runHaxl env $ try $
    (do _ <- friendsOf 2; throw (NotFound "foo")) `pOr` return False
  assertBool "pOr5" $
    case (r :: Either NotFound Bool) of
      Left (NotFound _) -> True
      _ -> False

pAndTests future = do
  env <- makeTestEnv future

  -- Test semantics
  r <- runHaxl env $ do
        a <- return False `pAnd` return False
        b <- return False `pAnd` return True
        c <- return True `pAnd` return False
        d <- return True `pAnd` return True
        return (not a && not b && not c && d)
  assertBool "pAnd0" r

  -- pAnd is left-biased with respect to exceptions:
  r <- runHaxl env $ try $ return False `pAnd` throw (NotFound "foo")
  assertBool "pAnd1" $
    case (r :: Either NotFound Bool) of
      Right False -> True
      _ -> False
  r <- runHaxl env $ try $ throw (NotFound "foo") `pAnd` return False
  assertBool "pAnd2" $
    case (r :: Either NotFound Bool) of
      Left (NotFound "foo") -> True
      _ -> False

  -- pAnd is non-deterministic (see also Note [tricky pOr/pAnd])
  let nondet =
        (do _ <- friendsOf 1; throw (NotFound "foo")) `pAnd` return False
  r <- runHaxl env $ try nondet
  assertBool "pAnd3" $
    case (r :: Either NotFound Bool) of
      Right False -> True
      _ -> False
  -- next we populate the cache
  _ <- runHaxl env $ friendsOf 1
  -- and now exactly the same pAnd again will throw this time:
  r <- runHaxl env $ try nondet
  assertBool "pAnd4" $
    case (r :: Either NotFound Bool) of
      Left (NotFound "foo") -> True
      _ -> False

  -- One more test: Blocked/True => Blocked
  r <- runHaxl env $ try $
    (do _ <- friendsOf 2; throw (NotFound "foo")) `pAnd` return True
  assertBool "pAnd5" $
    case (r :: Either NotFound Bool) of
      Left (NotFound _) -> True
      _ -> False

tests :: Bool -> Test
tests future = TestList
  [ TestLabel "batching1" $ TestCase (batching1 future)
  , TestLabel "batching2" $ TestCase (batching2 future)
  , TestLabel "batching3" $ TestCase (batching3 future)
  , TestLabel "batching4" $ TestCase (batching4 future)
  , TestLabel "batching5" $ TestCase (batching5 future)
  , TestLabel "batching6" $ TestCase (batching6 future)
  , TestLabel "batching7" $ TestCase (batching7 future)
  , TestLabel "batching8" $ TestCase (batching8 future)
  , TestLabel "batching9" $ TestCase (batching9 future)
  , TestLabel "caching1" $ TestCase (caching1 future)
  , TestLabel "caching2" $ TestCase (caching2 future)
  , TestLabel "caching3" $ TestCase (caching3 future)
  , TestLabel "CacheReuse" $ TestCase (cacheReuse future)
  , TestLabel "NoCaching" $ TestCase (noCaching future)
  , TestLabel "exceptionTest1" $ TestCase (exceptionTest1 future)
  , TestLabel "exceptionTest2" $ TestCase (exceptionTest2 future)
  , TestLabel "deterministicExceptions" $
      TestCase (deterministicExceptions future)
  , TestLabel "pOrTest" $ TestCase (pOrTests future)
  , TestLabel "pAndTest" $ TestCase (pAndTests future)
  ]


-- 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 RankNTypes, GADTs, BangPatterns, DeriveDataTypeable,
    StandaloneDeriving #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind -fno-warn-type-defaults #-}

module Bench where

import Haxl.Core.DataCache as DataCache

import Prelude hiding (mapM)

import Data.Hashable
import Data.IORef
import Data.Time.Clock
import Data.Traversable
import Data.Typeable
import System.Environment
import Text.Printf

data TestReq a where
  ReqInt    :: {-# UNPACK #-} !Int -> TestReq Int
  ReqDouble :: {-# UNPACK #-} !Int -> TestReq Double
  ReqBool   :: {-# UNPACK #-} !Int -> TestReq Bool
  deriving Typeable

deriving instance Eq (TestReq a)
deriving instance Show (TestReq a)

instance Hashable (TestReq a) where
  hashWithSalt salt (ReqInt i) = hashWithSalt salt (0::Int, i)
  hashWithSalt salt (ReqDouble i) = hashWithSalt salt (1::Int, i)
  hashWithSalt salt (ReqBool i) = hashWithSalt salt (2::Int, i)

main = do
  [n] <- fmap (fmap read) getArgs
  t0 <- getCurrentTime
  cache <- emptyDataCache
  let
     f 0 = return ()
     f !n = do
       m <- newIORef 0
       DataCache.insert (ReqInt n) m cache
       f (n-1)
  --
  f n
  m <- DataCache.lookup (ReqInt (n `div` 2)) cache
  print =<< mapM readIORef m
  t1 <- getCurrentTime
  printf "insert: %.2fs\n" (realToFrac (t1 `diffUTCTime` t0) :: Double)

  t0 <- getCurrentTime
  let
     f 0  !m = return m
     f !n !m = do
      mbRes <- DataCache.lookup (ReqInt n) cache
      case mbRes of
        Nothing -> f (n-1) m
        Just _  -> f (n-1) (m+1)
  f n 0 >>= print
  t1 <- getCurrentTime
  printf "lookup: %.2fs\n" (realToFrac (t1 `diffUTCTime` t0) :: Double)


-- 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 OverloadedStrings #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RebindableSyntax #-}

module CoreTests where

import Haxl.Prelude
import Prelude ()

import Haxl.Core

import Test.HUnit

import Data.Aeson
import qualified Data.ByteString.Lazy.Char8 as BS
import Data.List

import Control.Exception (Exception(..))

import ExampleDataSource

testEnv = do
  -- To use a data source, we need to initialize its state:
  exstate <- ExampleDataSource.initGlobalState

  -- And create a StateStore object containing the states we need:
  let st = stateSet exstate stateEmpty

  -- Create the Env:
  initEnv st () :: IO (Env () ())

useless :: String -> GenHaxl u w Bool
useless _ = throw (NotFound "ha ha")

exceptions :: Assertion
exceptions =
  do
    en <- emptyEnv () :: IO (Env () ())
    a <- runHaxl en $ try (useless "input")
    assertBool "NotFound -> HaxlException" $
      isLeft (a :: Either HaxlException Bool)

    b <- runHaxl en $ try (useless "input")
    assertBool "NotFound -> Logic Error" $
      isLeft (b :: Either LogicError Bool)

    c <- runHaxl en $ try (useless "input")
    assertBool "NotFound -> NotFound" $
      isLeft (c :: Either NotFound Bool)

    -- Make sure TransientError -doesn't- catch our NotFound
    d <- runHaxl en $
             (useless "input"
              `catch` \TransientError{} -> return False)
             `catch` \LogicError{} -> return True
    assertBool "Transient != NotFound" d

    -- test catch
    e <- runHaxl en $
         throw (NotFound "haha") `catch` \NotFound{} -> return True
    assertBool "catch1" e

    f <- runHaxl en $
         throw (NotFound "haha") `catch` \LogicError{} -> return True
    assertBool "catch2" f

    -- test catchIf
    let transientOrNotFound e
          | Just TransientError{} <- fromException e  = True
          | Just NotFound{} <- fromException e  = True
          | otherwise = False

    e <- runHaxl en $
         catchIf transientOrNotFound (throw (NotFound "haha")) $ \_ ->
           return True
    assertBool "catchIf1" e

    e <- runHaxl en $
         catchIf transientOrNotFound (throw (FetchError "haha")) $ \_ ->
           return True
    assertBool "catchIf2" e

    e <- runHaxl en $
           (catchIf transientOrNotFound (throw (CriticalError "haha")) $ \_ ->
              return True)
            `catch` \InternalError{} -> return False
    assertBool "catchIf2" (not e)

    -- test tryToHaxlException
    e <- runHaxl en $ tryToHaxlException $ head []
    assertBool "tryToHaxlException1" $
      case e of
        Left ex | Just NonHaxlException{} <- fromException (toException ex)
          -> True
        _ -> False

    env <- testEnv
    e <- runHaxl env $ tryToHaxlException $ do
      xs <- listWombats 3
      return $! length xs `quot` 0
    print e
    assertBool "tryToHaxlException1" $
      case e of
        Left ex | Just NonHaxlException{} <- fromException (toException ex)
          -> True
        _ -> False

    env <- testEnv
    e <- runHaxl env $ mapM tryToHaxlException
      [ do xs <- listWombats 3; return $! length xs `quot` 0
      , head []
      ]
    print e
    assertBool "tryToHaxlException2" $
      case e of
        [Left ex1, Left ex2]
           | "divide" `isInfixOf` show ex1
           , "head" `isInfixOf` show ex2 -> True
        _ -> False
  where
  isLeft Left{} = True
  isLeft _ = False

-- This is mostly a compile test, to make sure all the plumbing
-- makes the compiler happy.
base :: (Exception a) => a -> IO HaxlException
base e = do
  en <- emptyEnv () :: IO (Env () ())
  runHaxl en $ throw e `catch` \x -> return x

printing :: Assertion
printing = do
  a <- base $ NotFound "notfound!"
  print a

  b <- base $ CriticalError "ohthehumanity!"
  print b

  c <- base $ FetchError "timeout!"
  print c

  BS.putStrLn $ encode a
  BS.putStrLn $ encode b
  BS.putStrLn $ encode c

withEnvTest :: Test
withEnvTest = TestLabel "withEnvTest" $ TestCase $ do
  exstate <- ExampleDataSource.initGlobalState
  e <- initEnv (stateSet exstate stateEmpty) False :: IO (Env Bool ())
  b <- runHaxl e $ withEnv e { userEnv = True } $ env userEnv
  assertBool "withEnv1" b
  e <- initEnv (stateSet exstate stateEmpty) False :: IO (Env Bool ())
  b <- runHaxl e $ withEnv e { userEnv = True } $ do
    _ <- countAardvarks "aaa"
    env userEnv
  assertBool "withEnv2" b
  e <- initEnv (stateSet exstate stateEmpty) False :: IO (Env Bool ())
  b <- runHaxl e $ withEnv e { userEnv = True } $ do
    memo ("xxx" :: Text) $ do
      _ <- countAardvarks "aaa"
      env userEnv
  assertBool "withEnv3" b
  e <- initEnv (stateSet exstate stateEmpty) False :: IO (Env Bool ())
  b <- runHaxl e $
    withEnv e { userEnv = True } $ do
      memo ("yyy" :: Text) $ do
        _ <- countAardvarks "aaa"
        _ <- countAardvarks "bbb"
        env userEnv
  assertBool "withEnv4" b

tests = TestList
  [ TestLabel "exceptions" $ TestCase exceptions,
    TestLabel "print_stuff" $ TestCase printing,
    TestLabel "withEnv" $ withEnvTest
  ]


-- 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 StandaloneDeriving, GADTs, DeriveDataTypeable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
module DataCacheTest (tests, newResult, takeResult) where

import Haxl.Core.DataCache as DataCache
import Haxl.Core.Monad
import Haxl.Core

import Control.Exception
import Data.Either
import Data.Hashable
import Data.Traversable
import Data.Typeable
import Prelude hiding (mapM)
import Test.HUnit
import Data.IORef
import Data.Text
import Unsafe.Coerce

data TestReq a where
   Req :: Int -> TestReq a -- polymorphic result
  deriving Typeable

deriving instance Eq (TestReq a)
deriving instance Show (TestReq a)

instance Hashable (TestReq a) where
  hashWithSalt salt (Req i) = hashWithSalt salt i

instance DataSource u TestReq where
  fetch = error "no fetch defined"

instance DataSourceName TestReq where
  dataSourceName _ = pack "TestReq"

instance StateKey TestReq where
  data State TestReq = TestReqState

instance ShowP TestReq where showp = show

data CacheableReq x where CacheableInt :: Int -> CacheableReq Int
  deriving Typeable
deriving instance Eq (CacheableReq x)
deriving instance Show (CacheableReq x)
instance Hashable (CacheableReq x) where
  hashWithSalt s (CacheableInt val) = hashWithSalt s (0::Int, val)

newResult :: Monoid w => a -> IO (IVar u w a)
newResult a = newFullIVar (Ok a mempty)

takeResult :: IVar u w a -> IO (ResultVal a w)
takeResult IVar{ivarRef = ref} = do
  e <- readIORef ref
  case e of
    IVarFull a -> return a
    _ -> error "takeResult"

dcSoundnessTest :: Test
dcSoundnessTest = TestLabel "DataCache soundness" $ TestCase $ do
  m1 <- newResult 1
  m2 <- newResult "hello"
  cache <- emptyDataCache
  DataCache.insert (Req 2 :: TestReq String) m2 cache
  DataCache.insert (Req 1 :: TestReq Int) m1 cache

  -- "Req 1" has a result of type Int, so if we try to look it up
  -- with a result of type String, we should get Nothing, not a crash.
  r <- mapM takeResult =<< DataCache.lookup (Req 1) cache
  assertBool "dcSoundness1" $
    case r :: Maybe (ResultVal String ()) of
     Nothing -> True
     _something_else -> False

  r <- mapM takeResult =<< DataCache.lookup (Req 1) cache
  assertBool "dcSoundness2" $
    case r :: Maybe (ResultVal Int ()) of
     Just (Ok 1 Nothing) -> True
     _something_else -> False

  r <- mapM takeResult =<< DataCache.lookup (Req 2) cache
  assertBool "dcSoundness3" $
    case r :: Maybe (ResultVal String ()) of
      Just (Ok "hello" Nothing) -> True
      _something_else -> False

  r <- mapM takeResult =<< DataCache.lookup (Req 2) cache
  assertBool "dcSoundness4" $
    case r :: Maybe (ResultVal Int ()) of
      Nothing -> True
      _something_else -> False

dcStrictnessTest :: Test
dcStrictnessTest = TestLabel "DataCache strictness" $ TestCase $ do
  env <- initEnv stateEmpty () :: IO (Env () ())
  r <- Control.Exception.try $ runHaxl env $
    cachedComputation (Req (error "BOOM")) $ return "OK"
  assertBool "dcStrictnessTest" $
    case r of
      Left (ErrorCall "BOOM") -> True
      _other -> False

dcFallbackTest :: Test
dcFallbackTest = TestLabel "DataCache fallback" $ TestList
  [ TestLabel "Base" $ TestCase $ do
      env <- mkEnv
      (r,cached) <- runHaxl env (do
                           a <- dataFetch req
                           b <- cacheResult (CacheableInt 1234) (return 99999)
                           return (a,b))
      (Stats stats) <- readIORef (statsRef env)
      assertEqual "fallback still has stats" 1
        (Prelude.length [x | x@FetchStats{} <- stats])
      assertEqual "dcFallbackTest found" 1 r
      assertEqual "dcFallbackTest cached" 1234 cached
  , TestLabel "Exception" $ TestCase $ do
      env <- mkEnv
      rbad <- Control.Exception.try $ runHaxl env (dataFetch reqBad)
      assertBool "dcFallbackTest not found" $
        case rbad of
          Left (ErrorCall "no fetch defined") -> True
          _ -> False
  , TestLabel "Completions" $ TestCase $ do
      -- check applicative still runs as it would have without a fallback
      env <- mkEnv
      let
        fetchA = dataFetch reqEx
        fetchB = tellWrite 7 >> dataFetch req
      (rbad, writes) <- runHaxlWithWrites env $
        Haxl.Core.try $ (,) <$> fetchA <*> fetchB
      fetches <- countFetches env
      assertEqual "dispatched 2 fetches" 2 fetches
      assertBool "exception propogates" $
        case rbad of
          Left (NotFound _) -> True
          _ -> False
      assertEqual "write side effects happen" [7] (flattenWT writes)
  ]
  where

    mkEnv = addLookup <$> initEnv (stateSet TestReqState stateEmpty) ()

    countFetches env = do
      (Stats stats) <- readIORef (statsRef env)
      let
        c = sum [ fetchBatchSize x
                | x@FetchStats{} <- stats
                ]
      return c

    addLookup :: Env () (WriteTree Int) -> Env () (WriteTree Int)
    addLookup e = e { dataCacheFetchFallback = Just (DataCacheLookup lookup)
                    , flags = (flags e) { report = profilingReportFlags }
                    }
    lookup
      :: forall req a . Typeable (req a)
      => req a
      -> IO (Maybe (ResultVal a (WriteTree Int)))
    lookup r
      | typeOf r == typeRep (Proxy :: Proxy (TestReq Int)) =
        -- have to coerce on the way out as results are not Typeable
        -- so you better be sure you do it right!
        return $ unsafeCoerce . doReq <$> cast r
      | typeOf r == typeRep (Proxy :: Proxy (CacheableReq Int)) =
          return $ unsafeCoerce . doCache <$> cast r
      | otherwise = return Nothing

    doReq :: TestReq Int -> ResultVal Int (WriteTree Int)
    doReq (Req 999) = ThrowHaxl (toException $ NotFound empty) Nothing
    doReq (Req r) = Ok r Nothing

    doCache :: CacheableReq Int -> ResultVal Int (WriteTree Int)
    doCache (CacheableInt i) = Ok i Nothing

    req :: TestReq Int
    req = Req 1

    reqEx :: TestReq Int
    reqEx = Req 999

    reqBad :: TestReq String
    reqBad = Req 2

-- tests :: Assertion
tests = TestList [ dcSoundnessTest
                 , dcStrictnessTest
                 , dcFallbackTest
                 ]


-- 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 FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ApplicativeDo #-}

module DataSourceDispatchTests (tests) where
import Test.HUnit hiding (State)
import Control.Monad
import Haxl.Core
import Data.Hashable

data DataSourceDispatch ty where
    GetBatchSize :: Int -> DataSourceDispatch Int

deriving instance Eq (DataSourceDispatch ty)
deriving instance Show (DataSourceDispatch ty)

instance DataSourceName DataSourceDispatch where
    dataSourceName _ = "DataSourceDispatch"

instance StateKey DataSourceDispatch where
    data State DataSourceDispatch = DataSourceDispatchState

instance ShowP DataSourceDispatch where showp = show

instance Hashable (DataSourceDispatch a) where
  hashWithSalt s (GetBatchSize n) = hashWithSalt s n

initDataSource :: IO (State DataSourceDispatch)
initDataSource = return DataSourceDispatchState

instance DataSource UserEnv DataSourceDispatch where
    fetch _state _flags _u = SyncFetch $ \bfs -> forM_ bfs (fill $ length bfs)
      where
      fill :: Int -> BlockedFetch DataSourceDispatch -> IO ()
      fill l (BlockedFetch (GetBatchSize _ ) rv) = putResult rv (Right l)

    schedulerHint Batching = TryToBatch
    schedulerHint NoBatching = SubmitImmediately

data UserEnv = Batching | NoBatching deriving (Eq)

makeTestEnv :: UserEnv -> IO (Env UserEnv ())
makeTestEnv testUsrEnv = do
  st <- initDataSource
  e <- initEnv (stateSet st stateEmpty) testUsrEnv
  return e { flags = (flags e) {
    report = setReportFlag ReportFetchStats defaultReportFlags } }

schedulerTest:: Test
schedulerTest = TestCase $ do
    let
        fet = do
          x <- dataFetch (GetBatchSize 0)
          y <- dataFetch (GetBatchSize 1)
          return [x,y]

    e <- makeTestEnv Batching
    r1 :: [Int] <- runHaxl e fet
    assertEqual "Failed to create batches for data fetch" [2,2] r1

    eNoBatching <- makeTestEnv NoBatching
    r2 :: [Int] <- runHaxl eNoBatching fet
    assertEqual "Unexpexted batches in SubmitImmediately" [1,1] r2

    return ()

tests :: Test
tests = TestList
  [ TestLabel "schedulerTest" schedulerTest
  ]


-- 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 DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}

module ExampleDataSource (
    -- * initialise the state
    initGlobalState,

    -- * requests for this data source
    Id(..), ExampleReq(..),
    countAardvarks,
    listWombats,
  ) where

import Haxl.Prelude
import Prelude ()

import Haxl.Core

import Data.Typeable
import Data.Hashable
import Control.Concurrent
import qualified Control.Exception as E
import System.IO

-- Here is an example minimal data source.  Our data source will have
-- two requests:
--
--   countAardvarks :: String -> Haxl Int
--   listWombats    :: Id     -> Haxl [Id]
--
-- First, the data source defines a request type, with one constructor
-- for each request:

newtype Id = Id Int
  deriving (Eq, Ord, Enum, Num, Integral, Real, Hashable, Typeable)

instance Show Id where
  show (Id i) = show i

data ExampleReq a where
  CountAardvarks :: String -> ExampleReq Int
  ListWombats    :: Id     -> ExampleReq [Id]
  deriving Typeable -- requests must be Typeable

-- The request type (ExampleReq) is parameterized by the result type of
-- each request.  Each request might have a different result, so we use a
-- GADT - a data type in which each constructor may have different type
-- parameters. Here CountAardvarks is a request that takes a String
-- argument and its result is Int, whereas ListWombats takes an Id
-- argument and returns a [Id].

-- The request type needs instances for 'Eq1' and 'Hashable1'.  These
-- are like 'Eq' and 'Hashable', but for types with one parameter
-- where the parameter is irrelevant for hashing and equality.
-- These two instances are used to support caching of requests.

-- We need Eq, but we have to derive it with a standalone declaration
-- like this, because plain deriving doesn't work with GADTs.
deriving instance Eq (ExampleReq a)

deriving instance Show (ExampleReq a)

instance ShowP ExampleReq where showp = show

instance Hashable (ExampleReq a) where
   hashWithSalt s (CountAardvarks a) = hashWithSalt s (0::Int,a)
   hashWithSalt s (ListWombats a)    = hashWithSalt s (1::Int,a)

instance StateKey ExampleReq where
  data State ExampleReq = ExampleState {
        -- in here you can put any state that the
        -- data source needs to maintain throughout the
        -- run.
        }

-- Next we need to define an instance of DataSourceName:

instance DataSourceName ExampleReq where
  dataSourceName _ = "ExampleDataSource"

-- Next we need to define an instance of DataSource:

instance DataSource u ExampleReq where
  -- I'll define exampleFetch below
  fetch = exampleFetch

  -- we don't want to treat NotFound as an exception for stats purposes
  classifyFailure _ _ e
    | Just NotFound{} <- E.fromException e = IgnoredForStatsFailure
    | otherwise = StandardFailure

-- Every data source should define a function 'initGlobalState' that
-- initialises the state for that data source.  The arguments to this
-- function might vary depending on the data source - we might need to
-- pass in resources from the environment, or parameters to set up the
-- data source.
initGlobalState :: IO (State ExampleReq)
initGlobalState = do
  -- initialize the state here.
  return ExampleState { }

-- The most important bit: fetching the data.  The fetching function
-- takes a list of BlockedFetch, which is defined as
--
-- data BlockedFetch r
--   = forall a . BlockedFetch (r a) (ResultVar a)
--
-- That is, each BlockedFetch is a pair of
--
--   - the request to fetch (with result type a)
--   - a ResultVar to store either the result or an error
--
-- The job of fetch is to fetch the data and fill in all the ResultVars.
--
exampleFetch :: State ExampleReq             -- current state
             -> Flags                        -- tracing verbosity, etc.
             -> u                            -- user environment
             -> PerformFetch ExampleReq      -- tells the framework how to fetch

exampleFetch _state _flags _user = SyncFetch $ mapM_ fetch1

  -- There are two ways a data source can fetch data: synchronously or
  -- asynchronously.  See the type 'PerformFetch' in "Haxl.Core.Types" for
  -- details.

fetch1 :: BlockedFetch ExampleReq -> IO ()
fetch1 (BlockedFetch (CountAardvarks "BANG") _) =
  error "BANG"  -- data sources should not throw exceptions, but in
                -- the event that one does, the framework will
                -- propagate the exception to the call site of
                -- dataFetch.
fetch1 (BlockedFetch (CountAardvarks "BANG2") m) = do
  putSuccess m 1
  error "BANG2" -- the exception is propagated even if we have already
                -- put the result with putSuccess
fetch1 (BlockedFetch (CountAardvarks "BANG3") _) = do
  hPutStr stderr "BANG3"
  killThread =<< myThreadId -- an asynchronous exception
fetch1 (BlockedFetch (CountAardvarks "BANG4") r) = do
  putFailure r $ NotFound "BANG4"
fetch1 (BlockedFetch (CountAardvarks str) m) =
  putSuccess m (length (filter (== 'a') str))
fetch1 (BlockedFetch (ListWombats a) r) =
  if a > 999999
    then putFailure r $ FetchError "too large"
    else putSuccess r $ take (fromIntegral a) [1..]

-- Normally a data source will provide some convenient wrappers for
-- its requests:

countAardvarks :: String -> GenHaxl u w Int
countAardvarks str = dataFetch (CountAardvarks str)

listWombats :: Id -> GenHaxl u w [Id]
listWombats i = dataFetch (ListWombats i)


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

module ExceptionStackTests (tests) where

import Prelude ()
import Haxl.Core
import Haxl.Prelude

import Test.HUnit

import qualified ExampleDataSource

testEnv :: ReportFlags -> IO (Env () ())
testEnv report = do
  exstate <- ExampleDataSource.initGlobalState
  let st = stateSet exstate stateEmpty
  env <- initEnv st ()
  return env{ flags = (flags env){ report = report } }

reportFlags :: ReportFlags
reportFlags = setReportFlag ReportExceptionLabelStack defaultReportFlags

runHaxlTest
  :: ReportFlags
  -> String
  -> (Int -> Int -> GenHaxl () () Int)
  -> IO (Maybe [Text])
runHaxlTest report str func = do
  env <- testEnv report
  result <- runHaxl env $
    withLabel "try" $ tryToHaxlException $ withLabel "test" $ do
      x <- withLabel "dummy" $ pure 1
      y <- withLabel "fetch" $ ExampleDataSource.countAardvarks str
      withLabel "func" $ func x y
  case result of
    Left (HaxlException stk _) -> return stk
    Right{} -> assertFailure "expected: HaxlException"

fetchException :: Test
fetchException = TestCase $ do
  result <- runHaxlTest reportFlags "BANG4" $ \i j -> return $ i + j
  assertEqual "stack" (Just ["fetch", "test", "try", "MAIN"]) result

userException :: Test
userException = TestCase $ do
  result <- runHaxlTest reportFlags "aaa" $ \_ _ -> withLabel "throw" $
    throw $ InvalidParameter "throw"
  assertEqual "stack" (Just ["throw", "func", "test", "try", "MAIN"]) result

#ifndef PROFILING
disabledExceptionStack :: Test
disabledExceptionStack = TestCase $ do
  result <- runHaxlTest defaultReportFlags "BANG4" $ \i j -> return $ i + j
  assertEqual "stack" Nothing result
#endif

tests :: Test
tests = TestList
  [ TestLabel "FetchException" fetchException
  , TestLabel "UserException" userException
#ifndef PROFILING
  , TestLabel "DisabledExceptionStack"  disabledExceptionStack
#endif
  ]


module FBMain (main) where

import Facebook.Init
import TestRunner
import AllTests

main :: IO ()
main = withFacebookUnitTest $ testRunner $ allTests


-- 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.

module FullyAsyncTest where

import Haxl.Prelude as Haxl
import Prelude()

import SleepDataSource
import Haxl.DataSource.ConcurrentIO

import Haxl.Core
import Test.HUnit
import Data.IORef
import Haxl.Core.Monad (unsafeLiftIO)

tests :: Test
tests = sleepTest

testEnv :: IO (Env () ())
testEnv = do
  st <- mkConcurrentIOState
  env <- initEnv (stateSet st stateEmpty) ()
  return env { flags = (flags env) {
    report = setReportFlag ReportFetchStats defaultReportFlags } }

sleepTest :: Test
sleepTest = TestCase $ do
  env <- testEnv

  ref <- newIORef ([] :: [Int])
  let tick n = unsafeLiftIO (modifyIORef ref (n:))

  -- simulate running a selection of data fetches that complete at
  -- different times, overlapping them as much as possible.
  runHaxl env $
    sequence_
       [ sequence_ [sleep 100, sleep 400] `andThen` tick 5     -- A
       , sleep 100 `andThen` tick 2 `andThen` sleep 200 `andThen` tick 4    -- B
       , sleep 50 `andThen` tick 1 `andThen` sleep 150 `andThen` tick 3     -- C
       ]

  ys <- readIORef ref
  assertEqual "FullyAsyncTest: ordering" [1,2,3,4,5] (reverse ys)

  stats <- readIORef (statsRef env)
  print stats
  assertEqual "FullyAsyncTest: stats" 5 (numFetches stats)

{-
           A         B         C
50          |        |       tick 1
100         |     tick 2       |
150         |        |         |
200         |        |       tick 3
250         |        |
300         |     tick 4
350         |
400         |
450         |
500      tick 5
-}


-- 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, OverloadedStrings #-}
module LoadCache where

import Haxl.Core
import ExampleDataSource

#include "LoadCache.txt"


-- 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.

module MemoizationTests (tests) where

import Data.IORef

import Test.HUnit

import Haxl.Core
import Haxl.Core.Monad (unsafeLiftIO)

import ExampleDataSource

memoSoundness :: Test
memoSoundness = TestCase $ do
  iEnv <- do
    exState <- ExampleDataSource.initGlobalState
    initEnv (stateSet exState stateEmpty) () :: IO (Env () ())

  unMemoizedWombats <- runHaxl iEnv $ listWombats 100

  (initialGet, subsequentGet) <- runHaxl iEnv $ do
    wombatsMemo <- newMemoWith (listWombats 100)
    let memoizedWombats = runMemo wombatsMemo

    initialGet <- memoizedWombats
    subsequentGet <- memoizedWombats

    return (initialGet, subsequentGet)

  assertBool "Memo Soundness 1" $ initialGet == unMemoizedWombats
  assertBool "Memo Soundness 2" $ subsequentGet == unMemoizedWombats

  let impure runCounterRef = unsafeLiftIO $ do
        modifyIORef runCounterRef succ
        readIORef runCounterRef

      initialRunCounter = 0 :: Int

  runCounterRef <- newIORef initialRunCounter

  (initialImpureGet, subsequentImpureGet) <- runHaxl iEnv $ do
    impureMemo <- newMemoWith (impure runCounterRef)
    let memoizedImpure = runMemo impureMemo

    initialImpureGet <- memoizedImpure
    subsequentImpureGet <- memoizedImpure

    return (initialImpureGet, subsequentImpureGet)

  assertBool "Memo Soundness 3" $ initialImpureGet == succ initialRunCounter
  assertBool "Memo Soundness 4" $ subsequentImpureGet == initialImpureGet

  let fMemoVal = 42 :: Int

  dependentResult <- runHaxl iEnv $ do
    fMemoRef <- newMemo
    gMemoRef <- newMemo

    let f = runMemo fMemoRef
        g = runMemo gMemoRef

    prepareMemo fMemoRef $ return fMemoVal
    prepareMemo gMemoRef $ succ <$> f

    a <- f
    b <- g
    return (a + b)

  assertBool "Memo Soundness 5" $ dependentResult == fMemoVal + succ fMemoVal

tests = TestList [TestLabel "Memo Soundness" memoSoundness]


-- 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 DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards #-}

module MockTAO (
    Id(..),
    initGlobalState,
    assocRangeId2s,
    friendsAssoc,
    friendsOf,
  ) where

import Data.Hashable
import Data.Map (Map)
import Data.Typeable
import Prelude ()
import qualified Data.Map as Map
import qualified Data.Text as Text
import Control.Concurrent
import Control.Exception
import Control.Monad (void)

import Haxl.Prelude
import Haxl.Core

import TestTypes

-- -----------------------------------------------------------------------------
-- Minimal mock TAO

data TAOReq a where
  AssocRangeId2s :: Id -> Id -> TAOReq [Id]
  deriving Typeable

deriving instance Show (TAOReq a)
deriving instance Eq (TAOReq a)

instance ShowP TAOReq where showp = show

instance Hashable (TAOReq a) where
  hashWithSalt s (AssocRangeId2s a b) = hashWithSalt s (a,b)

instance StateKey TAOReq where
  data State TAOReq = TAOState { future :: Bool }

instance DataSourceName TAOReq where
  dataSourceName _ = "MockTAO"

instance DataSource UserEnv TAOReq where
  fetch TAOState{..} _flags _user
    | future = BackgroundFetch $ \f -> do
        mask_ $ void . forkIO $  mapM_ (doFetch True) f
    | otherwise = SyncFetch $ mapM_ (doFetch False)

initGlobalState :: Bool -> IO (State TAOReq)
initGlobalState future = return TAOState { future=future }

doFetch ::  Bool -> BlockedFetch TAOReq -> IO ()
doFetch bg (BlockedFetch req@(AssocRangeId2s a b) r) = put result
  where put = if bg then putResultFromChildThread r else putResult r
        result = case Map.lookup (a, b) assocs of
          Nothing -> except . NotFound . Text.pack $ show req
          Just result -> Right result

assocs :: Map (Id,Id) [Id]
assocs = Map.fromList [
  ((friendsAssoc, 1), [5..10]),
  ((friendsAssoc, 2), [7..12]),
  ((friendsAssoc, 3), [10..15]),
  ((friendsAssoc, 4), [15..19])
 ]

friendsAssoc :: Id
friendsAssoc = 167367433327742

assocRangeId2s :: Id -> Id -> Haxl [Id]
assocRangeId2s a b = dataFetch (AssocRangeId2s a b)

friendsOf :: Id -> Haxl [Id]
friendsOf = assocRangeId2s friendsAssoc


-- 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 OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module MonadAsyncTest (tests) where
import Haxl.Core
import Test.HUnit hiding (State)
import Control.Concurrent
import Control.Exception as Exception
import Control.Monad
import Haxl.Core.Monad (unsafeLiftIO, WriteTree)
import System.IO.Unsafe
import Data.Hashable
import Data.IORef
import Data.Text (Text)

newtype SimpleWrite = SimpleWrite Text deriving (Eq, Show)

{-# NOINLINE shouldThrowRef #-}
shouldThrowRef :: IORef Bool
shouldThrowRef = unsafePerformIO (newIORef False)

-- | This datasource contains "bad" code which can throw at the wrong
-- moment.
data ThrowableSleep a where
  Sleep :: Int -> ThrowableSleep Int

deriving instance Eq (ThrowableSleep a)
deriving instance Show (ThrowableSleep a)

instance ShowP ThrowableSleep where showp = show

instance Hashable (ThrowableSleep a) where
  hashWithSalt s (Sleep n) = hashWithSalt s n

instance StateKey ThrowableSleep where
  data State ThrowableSleep = ThrowableSleepState

initDataSource :: IO (State ThrowableSleep)
initDataSource = return ThrowableSleepState

instance DataSourceName ThrowableSleep where
  dataSourceName _ = "ThrowableSleep"

instance DataSource u ThrowableSleep where
  fetch _state _flags _u = BackgroundFetch $ \bfs -> forM_ bfs fill
    where
    fill :: BlockedFetch ThrowableSleep -> IO ()
    fill (BlockedFetch (Sleep n) rv) = do
      _ <- forkFinally
        (do
          threadDelay (n*1000)
          return n
        )
        (\res -> do
          shouldThrow <- atomicModifyIORef' shouldThrowRef (\s -> (False, s))
          -- Simulate case when datasource throws before putting Result into
          -- completions queue.
          when shouldThrow $ do
            throwIO $ ErrorCall "datasource threw an exception"
          -- In case the datasource throws before this point, there'll be
          -- nothing to put the result to the queue of 'completions', and
          -- therefore Haxl would block indefinitely.
          --
          -- Note that Haxl tries to catch datasource exceptions and put the
          -- "exception result" into `completions` using `wrapFetchInCatch`
          -- function. However that doesn't work in this case because the
          -- datasource throws in a separate thread.
          putResultFromChildThread rv res
        )
      return ()

tests :: Test
tests = TestList
  [ TestLabel "exceptionTest" exceptionTest
  ]

mkTestEnv :: IO (Env () (WriteTree SimpleWrite))
mkTestEnv = do
  st <- initDataSource
  initEnv (stateSet st stateEmpty) ()

exceptionTest :: Test
exceptionTest = TestCase $ do
  e <- mkTestEnv

  let
      fet (n :: Int) (st :: Bool )= do
        x <- dataFetch (Sleep (fromIntegral n))
        unsafeLiftIO $ writeIORef shouldThrowRef st
        y <- dataFetch (Sleep (fromIntegral x*2))
        return (x+y)

  r1 :: (Either Exception.SomeException Int)
    <- Exception.try $ runHaxl e $ fet 10 True

  -- Datasources are responsible for putting the fetched result into the
  -- completions queue. If for some reason they fail to do so, Haxl throws a
  -- LogicBug since the scheduler is still expecting some request(s) to
  -- be completed.
  case r1 of
    Left ex | Just (LogicBug _) <- Exception.fromException ex -> return ()
    _ -> assertFailure "r1 computation did not fail with Logic Bug!"

  -- Sanitize the env to get rid of all empty IVars
  -- While this test examines the case when there's an exception in the Haxl
  -- datasource itself, a similar behavior will occur in case an async
  -- exception is thrown to the Haxl scheduler thread.
  e' <- sanitizeEnv e

  r2 :: (Either Exception.SomeException Int)
    <- Exception.try $ runHaxl e' $ fet 10 False
  case r2 of
    Right _ -> return ()
    Left ex | Just (LogicBug _) <- Exception.fromException ex -> do
                assertFailure $ "bad exception in r2: " ++ show ex
    Left _ -> return ()


-- 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.

-- | Benchmarking tool for core performance characteristics of the Haxl monad.
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ApplicativeDo, RecordWildCards #-}
module MonadBench (main) where

import Control.Monad
import Data.List as List
import Data.Maybe
import Data.Time.Clock
import Options.Applicative
import System.Exit
import System.IO
import Text.Printf

import Haxl.Prelude as Haxl
import Prelude()

import Haxl.Core
import Haxl.Core.Monad (WriteTree)
import Haxl.Core.Util

import ExampleDataSource

newtype SimpleWrite = SimpleWrite Text deriving (Eq, Show)

testEnv :: ReportFlags -> IO (Env () (WriteTree SimpleWrite))
testEnv report = do
  exstate <- ExampleDataSource.initGlobalState
  let st = stateSet exstate stateEmpty
  env <- initEnv st ()
  return env { flags = (flags env) { report = report } }

type Test = (String, Int, Int -> GenHaxl () (WriteTree SimpleWrite) ())

testName :: Test -> String
testName (t,_,_) = t

allTests :: [Test]
allTests =
    -- parallel, identical queries
  [ ("par1", large, \n -> Haxl.sequence_ (replicate n (listWombats 3)))
    -- parallel, distinct queries
  , ("par2", medium, \n ->
        Haxl.sequence_ (map listWombats [1..fromIntegral n]))
    -- sequential, identical queries
  , ("seqr", huge, \n ->
        foldr andThen (return ()) (replicate n (listWombats 3)))
    -- sequential, left-associated, distinct queries
  , ("seql", medium, \n -> do
       _ <- foldl andThen (return []) (map listWombats [1.. fromIntegral n])
       return ())
    -- No memoization
  , ("memo0", small, \n -> Haxl.sequence_ [unionWombats | _ <- [1..n]])
    -- One put, N gets.
  , ("memo1", medium_large, \n ->
        Haxl.sequence_ [memo (42 :: Int) unionWombats | _ <- [1..n]])
    -- N puts, N gets.
  , ("memo2", small, \n ->
        Haxl.sequence_ [memo (i :: Int) unionWombats | i <- [1..n]])
  , ("memo3", medium_large, \n -> do
        ref <- newMemoWith unionWombats
        let c = runMemo ref
        Haxl.sequence_ [c | _ <- [1..n]])
  , ("memo4", small, \n -> do
        let f = unionWombatsTo
        Haxl.sequence_ [f x | x <- take n $ cycle [100, 200 .. 1000]])
  , ("memo5", medium_large, \n -> do
        f <- memoize1 unionWombatsTo
        Haxl.sequence_ [f x | x <- take n $ cycle [100, 200 .. 1000]])
  , ("memo6", small, \n -> do
        let f = unionWombatsFromTo
        Haxl.sequence_ [ f x y
                       | x <- take n $ cycle [100, 200 .. 1000]
                       , let y = x + 1000
                       ])
  , ("memo7", medium_large, \n -> do
        f <- memoize2 unionWombatsFromTo
        Haxl.sequence_ [ f x y
                       | x <- take n $ cycle [100, 200 .. 1000]
                       , let y = x + 1000
                       ])
  , ("cc1", medium_large, \n ->
        Haxl.sequence_ [ cachedComputation (ListWombats 1000) unionWombats
                       | _ <- [1..n]
                       ])
  , ("tree", 20, \n -> void $ tree n (\_ act -> act))
  , ("tree_labels", 20, \n -> void $
      tree n (\n act -> withLabel (textShow n) act))
    -- parallel writes
  , ("write1", large, \n ->
        Haxl.sequence_ (replicate n (tellWrite (SimpleWrite "haha"))))
    -- sequential writes
  , ("write2", huge, \n -> foldr
        andThen
        (return ())
        (replicate n (tellWrite (SimpleWrite "haha"))))
  ]
  where
    huge = large * 10
    large =  medium * 10
    medium_large =  medium * 4
    medium = 200000
    small = 1000

data Options = Options
  { test :: String
  , nOverride :: Maybe Int
  , reportFlag :: ReportFlags
  }

runTest :: Options -> Test -> IO ()
runTest Options{..} (t, nDef, act) = do
  let n = fromMaybe nDef nOverride
  env <- testEnv reportFlag
  t0 <- getCurrentTime
  runHaxl env $ act n
  t1 <- getCurrentTime
  printf "%12s: %10d reqs: %.2fs\n"
    t n (realToFrac (t1 `diffUTCTime` t0) :: Double)

optionsParser :: Parser Options
optionsParser = do
  test <- argument str (metavar "TEST")
  reportFlag <- reportFlagParser
  nOverride <- optional $ argument auto (metavar "NUM")
  return Options{..}
  where
    reportFlagParser = foldl' (flip ($)) defaultReportFlags <$> sequenceA
      [ flag id (setReportFlag i) $ long $ show i
      | i <- enumFrom minBound
      ]

main :: IO ()
main = do
  opts@Options{..} <- execParser $ info optionsParser mempty
  let tests = if test == "all"
        then allTests
        else filter ((==) test . testName) allTests
  when
    (null tests)
    (do
        hPutStrLn stderr $ "syntax: monadbench [all|" ++
          intercalate "|" (map testName allTests) ++ "]"
        exitWith (ExitFailure 1))
  Control.Monad.mapM_ (runTest opts) tests

tree
  :: Int
  -> (Int -> GenHaxl () (WriteTree SimpleWrite) [Id]
  -> GenHaxl () (WriteTree SimpleWrite) [Id])
  -> GenHaxl () (WriteTree SimpleWrite) [Id]
tree 0 wrap = wrap 0 $ listWombats 0
tree n wrap = wrap n $ concat <$> Haxl.sequence
  [ tree (n-1) wrap
  , listWombats (fromIntegral n), tree (n-1) wrap
  ]

unionWombats :: GenHaxl () (WriteTree SimpleWrite) [Id]
unionWombats = foldl List.union [] <$> Haxl.mapM listWombats [1..1000]

unionWombatsTo :: Id -> GenHaxl () (WriteTree SimpleWrite) [Id]
unionWombatsTo x = foldl List.union [] <$> Haxl.mapM listWombats [1..x]

unionWombatsFromTo :: Id -> Id -> GenHaxl () (WriteTree SimpleWrite) [Id]
unionWombatsFromTo x y = foldl List.union [] <$> Haxl.mapM listWombats [x..y]


-- 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 OverloadedStrings #-}
{-# LANGUAGE ApplicativeDo #-}
module OutgoneFetchesTests (tests) where

import Haxl.Prelude as Haxl
import Prelude()

import Haxl.Core
import Haxl.DataSource.ConcurrentIO

import Data.IORef
import qualified Data.Map as Map
import Data.Typeable
import Test.HUnit
import System.Timeout

import ExampleDataSource
import SleepDataSource

testEnv :: IO (Env () ())
testEnv = do
  exstate <- ExampleDataSource.initGlobalState
  sleepState <- mkConcurrentIOState
  let st = stateSet exstate $ stateSet sleepState stateEmpty
  e <- initEnv st ()
  return e { flags = (flags e) {
    report = setReportFlag ReportOutgoneFetches defaultReportFlags } }
    -- report=1 to enable fetches tracking

-- A cheap haxl computation we interleave b/w the @sleep@ fetches.
wombats :: GenHaxl () () Int
wombats = length <$> listWombats 3

outgoneFetchesTest :: String -> Int -> GenHaxl () () a -> Test
outgoneFetchesTest label unfinished haxl = TestLabel label $ TestCase $ do
  env <- testEnv
  _ <- timeout (100*1000) $ runHaxl env haxl -- 100ms
  actual <- getMapFromRCMap <$> readIORef (submittedReqsRef env)
  assertEqual "fetchesMap" expected actual
  where
  expected = if unfinished == 0 then Map.empty else
    Map.singleton (dataSourceName (Proxy :: Proxy (ConcurrentIOReq Sleep))) $
      Map.singleton (typeOf1 (undefined :: ConcurrentIOReq Sleep a)) unfinished

tests :: Test
tests = TestList
  [ outgoneFetchesTest "finished" 0 $ do
      -- test that a completed datasource fetch doesn't show up in Env
      _ <- sleep 1  -- finished
      _ <- sleep 1  -- cached/finished
      _ <- sleep 1  -- cached/finished
      wombats
  , outgoneFetchesTest "unfinished" 2 $ do
      -- test that unfinished datasource fetches shows up in Env
      _ <- sleep 200 -- unfinished
      _ <- wombats
      _ <- sleep 300 -- unfinished
      _  <- wombats
      return ()
  , outgoneFetchesTest "mixed" 2 $ do
      -- test for finished/unfinished fetches from the same datasource
      _ <- sleep 1   -- finished
      _ <- sleep 200  -- unfinished
      _ <- sleep 300  -- unfinished
      return ()
  , outgoneFetchesTest "cached" 1 $ do
      -- test for cached requests not showing up twice in ReqCountMap
      _ <- sleep 200  -- unfinished
      _ <- sleep 200  -- cached/unfinished
      return ()
  , outgoneFetchesTest "unsent" 1 $
      -- test for unsent requests not showing up in ReqCountMap
      sleep 200 `andThen` sleep 300 -- second req should never be sent
  ]


module ParallelTests where

import Haxl.Prelude
import Haxl.Core

import Haxl.DataSource.ConcurrentIO
import SleepDataSource

import Data.Time.Clock

import Test.HUnit

testEnv :: IO (Env () ())
testEnv = do
  sleepState <- mkConcurrentIOState
  let st = stateSet sleepState stateEmpty
  initEnv st ()

sync_test :: IO ()
sync_test = do
  env <- testEnv
  -- This computation tests that the two arguments of the pOr can fire
  -- without causing an error. The reason we test for this is that the
  -- synchronization involved in this case is a little fragile.
  False <- runHaxl env $ do
    (fmap (const False) (sleep 50)
      `pOr` fmap (const False) (sleep 100))
      `pOr` fmap (const False) (sleep 200)
  return ()

semantics_when_computation_is_blocked_test :: IO ()
semantics_when_computation_is_blocked_test = do
  env <- testEnv
    -- Test semantics of blocking
  let sleepReturn bool t = do
        _ <- sleep t
        return bool
  r <- runHaxl env $ do
        -- All sleep times are different so that they're not cached
        a <- sleepReturn False 10 `pOr` sleepReturn False 11
        b <- sleepReturn False 12 `pOr` sleepReturn True 13
        c <- sleepReturn True 14 `pOr` sleepReturn False 15
        d <- sleepReturn True 16 `pOr` sleepReturn True 17
        return (not a && b && c && d)
  assertBool "pOr blocked semantics" r

timing_test = do
  env <- testEnv
  t0 <- getCurrentTime
  True <- runHaxl env $
    fmap (const True) (sleep 200) `pOr` fmap (const True) (sleep 100)
  t1 <- getCurrentTime
  True <- runHaxl env $
    fmap (const True) (sleep 100) `pOr` fmap (const True) (sleep 200)
  t2 <- getCurrentTime
  False <- runHaxl env $
    fmap (const False) (sleep 200) `pOr` fmap (const False) (sleep 100)
  t3 <- getCurrentTime
  False <- runHaxl env $
    fmap (const False) (sleep 100) `pOr` fmap (const False) (sleep 200)
  t4 <- getCurrentTime
  -- diffUTCTime returns the difference in seconds,
  -- while sleep expects milliseconds
  assert (t4 `diffUTCTime` t3 < 0.2)
  assert (t3 `diffUTCTime` t2 < 0.2)
  assert (t2 `diffUTCTime` t1 < 0.2)
  assert (t1 `diffUTCTime` t0 < 0.2)

tests = TestList [TestLabel "sync_test" (TestCase sync_test)
                 ,TestLabel "timing_test" (TestCase timing_test)
                 ,TestLabel "semantics_when_computation_is_blocked_test" (TestCase semantics_when_computation_is_blocked_test)
                 ]