{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module Main where
import Data.Functor ((<&>))
import Data.List qualified as List
import Data.Maybe (mapMaybe)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Effectful
import Effectful.Exception (AsyncException (..), ErrorCall (..), SomeException)
import Effectful.Exception qualified as Effectful
import Effectful.FileSystem (FileSystem, runFileSystem)
import Effectful.FileSystem.IO.ByteString (readFile)
import Effectful.Haxl as Haxl
import Effectful.Hspec
import Effectful.Prim (Prim, runPrim)
import Effectful.Prim.IORef (modifyIORef', newIORef, readIORef)
import ExampleDataSource (Id (..))
import ExampleDataSource qualified
import LoadCache (loadCache)
import System.FilePath (dropFileName, (</>))
import Prelude hiding (readFile)
initialState :: Eff es StateStore
initialState = flip stateSet stateEmpty <$> ExampleDataSource.initGlobalState
withFetchStats :: Flags
withFetchStats =
defaultFlags
{ report = setReportFlag ReportFetchStats defaultReportFlags
}
main :: IO ()
main =
runEff
. runPrim
. runFileSystem
. runHspec
. runHaxlWith @() @() initialState () withFetchStats
$ do
describe "env" do
envTest
withEnvTest
describe "batching" do
describe "single round" do
applicativeBatchingTest
numBatchingTest
largeBatchTest
describe "multi-round" do
sequentialRoundsTest
dependentFanOutTest
deepChainTest
describe "caching" do
preCacheTest
preCacheExceptionTest
cachedComputationTest
cacheResultTest
memoTest
describe "exceptions" do
syncExceptionTest
asyncExceptionTest
describe "cache dump" do
dumpCacheTest
-- * Helpers
-- Run a 'GenHaxl' action and return the result alongside the per-round
-- 'FetchStats'.
withStats
:: forall es a
. (Haxl () () :> es)
=> GenHaxl () () es a
-> Eff es (a, [FetchStats])
withStats g = do
(a, Stats stats) <- inject $ haxlWithStats @() @() g
pure (a, stats)
rounds :: [FetchStats] -> Int
rounds = length
fetchBatchSizes :: [FetchStats] -> [Int]
fetchBatchSizes =
reverse . mapMaybe \case
FetchStats{..} -> Just fetchBatchSize
_ -> Nothing
totalFetches :: [FetchStats] -> Int
totalFetches = sum . fetchBatchSizes
-- Assert that a computation completes in exactly @n@ fetch rounds.
shouldUseRounds
:: (HasCallStack, Haxl () () :> es, Hspec :> es)
=> GenHaxl () () es a -> Int -> Eff es ()
g `shouldUseRounds` n = do
(_, stats) <- withStats g
rounds stats `shouldBe` n
-- Assert that a computation issues exactly @n@ fetches in total across all
-- rounds.
shouldUseFetches
:: (HasCallStack, Haxl () () :> es, Hspec :> es)
=> GenHaxl () () es a -> Int -> Eff es ()
g `shouldUseFetches` n = do
(_, stats) <- withStats g
totalFetches stats `shouldBe` n
-- Assert that a computation completes in exactly @r@ rounds issuing exactly
-- @f@ fetches in total.
shouldUseRoundsAndFetches
:: (HasCallStack, Haxl () () :> es, Hspec :> es)
=> GenHaxl () () es a -> (Int, Int) -> Eff es ()
g `shouldUseRoundsAndFetches` (r, f) = do
(_, stats) <- withStats g
rounds stats `shouldBe` r
totalFetches stats `shouldBe` f
deriving stock instance Eq Flags
instance Eq ReportFlags where
r1 == r2 = show r1 == show r2
deriving stock instance Show Flags
-- * Env
-- 'env' should extract data from the current 'Env'.
envTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
envTest = it "extracts data from the Env" do
flags <- inject . haxl @() @() $ env Haxl.flags
flags `shouldBe` withFetchStats
-- 'withEnv' should override the 'Env' for the duration of a computation,
-- then restore it afterwards.
withEnvTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
withEnvTest = it "overrides the Env for the duration of a computation" do
newEnv <- initEnv <&> \env -> env{Haxl.flags = defaultFlags}
(inner, outer) <- inject . haxl @() @() $ do
let newEnv' = newEnv{Haxl.flags = defaultFlags}
inner <- withEnv newEnv' $ env Haxl.flags
outer <- env Haxl.flags
pure (inner, outer)
inner `shouldBe` defaultFlags
outer `shouldBe` withFetchStats
-- * Single-round batching
-- Two independent applicative fetches should be batched into one round.
applicativeBatchingTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
applicativeBatchingTest = it "batches two independent applicative fetches" do
let g :: GenHaxl () () es Int
g =
ExampleDataSource.countAardvarks "abcabc"
+ (length <$> ExampleDataSource.listWombats 3)
(result, stats) <- withStats g
result `shouldBe` 2 + 3
rounds stats `shouldBe` 1
totalFetches stats `shouldBe` 2
-- The 'Num' instance uses 'liftA2', so arithmetic over fetched values should
-- also batch into one round.
numBatchingTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
numBatchingTest = it "batches fetches combined with Num arithmetic" do
let g :: GenHaxl () () es Int
g =
ExampleDataSource.countAardvarks "abcabc"
+ (length <$> ExampleDataSource.listWombats 3)
g `shouldUseRoundsAndFetches` (1, 2)
-- Many independent fetches should all land in a single round regardless of
-- how many there are.
largeBatchTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
largeBatchTest = it "batches many independent fetches into one round" do
let queries = take 8 $ flip replicate 'a' <$> [1 ..]
let g :: GenHaxl () () es Int
g = sum <$> traverse ExampleDataSource.countAardvarks queries
g `shouldUseRoundsAndFetches` (1, 8)
-- * Multi-round batching
-- A monadic bind forces a second round: the second fetch cannot be issued
-- until the first result is known.
sequentialRoundsTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
sequentialRoundsTest = it "uses two rounds for a monadic sequence" do
let g :: GenHaxl () () es Int
g = do
n <- ExampleDataSource.countAardvarks "abcabc"
length <$> ExampleDataSource.listWombats (Id n)
g `shouldUseRoundsAndFetches` (2, 2)
-- After one blocking fetch, multiple independent fetches in the next round
-- should still be batched together.
dependentFanOutTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
dependentFanOutTest = it "batches a fan-out of fetches in round 2" do
let g :: GenHaxl () () es Int
g = do
n <- ExampleDataSource.countAardvarks "abcabc" -- round 1
-- Both listWombats calls are independent of each other, so they
-- should be batched into a single round 2.
(length <$> ExampleDataSource.listWombats (Id n))
+ (length <$> ExampleDataSource.listWombats (Id $ n + 1))
(result, stats) <- withStats g
result `shouldBe` 2 + 3
rounds stats `shouldBe` 2
fetchBatchSizes stats `shouldBe` [1, 2]
-- A large fan-out after an initial fetch: all N dependent fetches should
-- land in round 2, not spread across N rounds.
largeFanOutTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
largeFanOutTest = it "batches a large fan-out into round 2" do
let g :: GenHaxl () () es Int
g = do
n <- ExampleDataSource.countAardvarks "abcabc" -- round 1, n == 2
-- Fan out to n+6 = 8 independent fetches in round 2
let queries = take (n + 6) $ flip replicate 'a' <$> [1 ..]
sum <$> traverse ExampleDataSource.countAardvarks queries
(_, stats) <- withStats g
rounds stats `shouldBe` 2
fetchBatchSizes stats `shouldBe` [1, 8]
-- A chain of dependent fetches should take exactly one round per step, with
-- no spurious extra rounds.
deepChainTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
deepChainTest = it "uses one round per step in a dependent chain" do
let g :: GenHaxl () () es Int
g = do
n1 <- ExampleDataSource.countAardvarks "ab" -- round 1
n2 <- ExampleDataSource.countAardvarks (replicate n1 'a') -- round 2
length <$> ExampleDataSource.listWombats (Id n2) -- round 3
g `shouldUseRounds` 3
-- Mixing applicative and monadic style: applicative branches within each
-- monadic step should still batch.
mixedStyleTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
mixedStyleTest = it "batches applicative branches within each monadic step" do
let g :: GenHaxl () () es Int
g = do
-- Round 1: two fetches batched
(n1, n2) <-
(,)
<$> ExampleDataSource.countAardvarks "ab"
<*> ExampleDataSource.countAardvarks "abc"
-- Round 2: two fetches batched
(w1, w2) <-
(,)
<$> ExampleDataSource.listWombats (Id n1)
<*> ExampleDataSource.listWombats (Id n2)
pure (length w1 + length w2)
(_, stats) <- withStats g
rounds stats `shouldBe` 2
fetchBatchSizes stats `shouldBe` [2, 2]
-- * Caching
-- Results pre-seeded with 'cacheRequest' are returned without issuing any
-- fetches.
preCacheTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
preCacheTest = it "returns pre-cached results without fetching" do
let g :: GenHaxl () () es Int
g = do
cacheRequest (ExampleDataSource.CountAardvarks "xxx") (Right 3)
cacheRequest (ExampleDataSource.ListWombats 1000000) (Right [1, 2, 3])
ExampleDataSource.countAardvarks "xxx" + (length <$> ExampleDataSource.listWombats 1000000)
(result, stats) <- withStats g
result `shouldBe` 6
rounds stats `shouldBe` 0
preCacheExceptionTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
preCacheExceptionTest = it "re-throws pre-cached exceptions" do
r <- Effectful.try . inject . haxl @() @() $ do
cacheRequest (ExampleDataSource.CountAardvarks "yyy") $
except (NotFound "yyy")
ExampleDataSource.countAardvarks "yyy"
r `shouldSatisfy` \case
Left (NotFound "yyy") -> True
_ -> False
-- 'cachedComputation' deduplicates a sub-computation across uses, reducing
-- the total number of fetches.
cachedComputationTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
cachedComputationTest = it "deduplicates cachedComputation across uses" do
let g :: GenHaxl () () es Int
g =
cachedComputation (ExampleDataSource.CountAardvarks "ababa") $
(length <$> ExampleDataSource.listWombats 10)
+ (length <$> ExampleDataSource.listWombats 20)
-- x used twice: second use hits the cache, so only 2 + 1 = 3 fetches total
-- rather than 2 + 2 + 1 = 5.
(r, stats) <- withStats $ g + g + ExampleDataSource.countAardvarks "baba"
r `shouldBe` 62
totalFetches stats `shouldBe` 3
-- 'cacheResult' runs its 'Eff' action exactly once even when used
-- applicatively, which would otherwise execute it twice.
cacheResultTest
:: forall es
. (Haxl () () :> es, Hspec :> es, Prim :> es)
=> Expectation es
cacheResultTest = it "executes cacheResult action exactly once" do
ref <- newIORef (0 :: Int)
let request = cacheResult (ExampleDataSource.CountAardvarks "ababa") $ do
modifyIORef' ref (+ 1)
readIORef ref
r <- inject . haxl @() @() $ request + request
-- The action increments and reads ref; it should run once, giving 1.
-- Both uses of `request` return 1, so the sum is 2.
r `shouldBe` 2
-- 'memo' deduplicates across uses in the same run, same as
-- 'cachedComputation'.
memoTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
memoTest = it "deduplicates memo'd computations" do
let g :: GenHaxl () () es Int
g =
memo (ExampleDataSource.CountAardvarks "ababa") $
(length <$> ExampleDataSource.listWombats 10)
+ (length <$> ExampleDataSource.listWombats 20)
(r, stats) <- withStats $ g + g + ExampleDataSource.countAardvarks "baba"
r `shouldBe` 62
totalFetches stats `shouldBe` 3
-- * Exceptions
syncExceptionTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
syncExceptionTest = do
it "catches synchronous data-source exceptions with Haxl.try" do
r <-
inject . haxl @() @() . Haxl.try $
ExampleDataSource.countAardvarks "BANG"
r `shouldSatisfy` \case
Left (ErrorCall "BANG") -> True
_ -> False
it "catches a second independent sync exception" do
r <-
inject . haxl @() @() . Haxl.try $
ExampleDataSource.countAardvarks "BANG2"
r `shouldSatisfy` \case
Left (ErrorCall "BANG2") -> True
_ -> False
it "still batches both fetches when one branch throws" do
(r, Stats stats) <-
inject . haxlWithStats @() @() $
Haxl.try $
(<>)
<$> ExampleDataSource.listWombats 1000000
<*> ExampleDataSource.listWombats 1000001
r `shouldSatisfy` \case
Left FetchError{} -> True
Right _ -> False
sum [fetchBatchSize | FetchStats{..} <- stats] `shouldBe` 2
sum [fetchFailures | FetchStats{..} <- stats] `shouldBe` 2
-- Asynchronous exceptions (e.g. 'ThreadKilled') escape 'Haxl.try' and must
-- be caught outside 'runHaxl'.
asyncExceptionTest :: forall es. (Haxl () () :> es, Hspec :> es) => Expectation es
asyncExceptionTest = it "propagates async exceptions past Haxl.try" do
r :: Either AsyncException (Either SomeException Int) <-
Effectful.try $
inject . haxl @() @() . Haxl.try $
(length <$> ExampleDataSource.listWombats 100)
+ ExampleDataSource.countAardvarks "BANG3"
r `shouldSatisfy` \case
Left ThreadKilled -> True
_ -> False
-- * Cache dump
-- Load a known cache then dump it; the output should match the fixture file.
-- Both operations run in the same 'haxl' call so they share the 'Env'.
dumpCacheTest
:: forall es
. (FileSystem :> es, Haxl () () :> es, Hspec :> es)
=> Expectation es
dumpCacheTest = it "round-trips the cache through dumpCacheAsHaskell" do
str <- inject . haxl @() @() $ do
loadCache
dumpCacheAsHaskell
fixture <-
Text.unpack . Text.decodeUtf8
<$> readFile (dropFileName __FILE__ </> "LoadCache.txt")
-- Line order is non-deterministic across GHC versions, so sort before
-- comparing.
List.sort (lines str) `shouldBe` List.sort (lines fixture)