diff --git a/Haxl/Core/DataSource.hs b/Haxl/Core/DataSource.hs
--- a/Haxl/Core/DataSource.hs
+++ b/Haxl/Core/DataSource.hs
@@ -107,6 +107,9 @@
   schedulerHint :: u -> SchedulerHint req
   schedulerHint _ = TryToBatch
 
+  schedulerHintState :: Maybe (State req) -> u -> SchedulerHint req
+  schedulerHintState _ u = schedulerHint u
+
   classifyFailure :: u -> req a -> SomeException -> FailureClassification
   classifyFailure _ _ _ = StandardFailure
 
diff --git a/Haxl/Core/Fetch.hs b/Haxl/Core/Fetch.hs
--- a/Haxl/Core/Fetch.hs
+++ b/Haxl/Core/Fetch.hs
@@ -264,7 +264,9 @@
         blockedFetch = BlockedFetch req rvar
         blockedFetchI = BlockedFetchInternal fid
         submitFetch = do
-          case schedulerHint userEnv :: SchedulerHint r of
+          let hint :: SchedulerHint r
+              hint = schedulerHintState (stateGet states) userEnv
+          case hint of
             SubmitImmediately ->
               performFetches env [BlockedFetches [blockedFetch] [blockedFetchI]]
             TryToBatch ->
@@ -532,13 +534,15 @@
 
 #if __GLASGOW_HASKELL__ >= 804
 instance Semigroup FailureCount where
-  (<>) = mappend
+  FailureCount s1 i1 <> FailureCount s2 i2 = FailureCount (s1+s2) (i1+i2)
 #endif
 
 instance Monoid FailureCount where
   mempty = FailureCount 0 0
+#if __GLASGOW_HASKELL__ < 804
   mappend (FailureCount s1 i1) (FailureCount s2 i2)
     = FailureCount (s1+s2) (i1+i2)
+#endif
 
 wrapFetchInStats
   :: DataSource u req
diff --git a/Haxl/Core/Memo.hs b/Haxl/Core/Memo.hs
--- a/Haxl/Core/Memo.hs
+++ b/Haxl/Core/Memo.hs
@@ -71,7 +71,8 @@
    :: forall req u w a.
       ( Eq (req a)
       , Hashable (req a)
-      , Typeable (req a))
+      , Typeable (req a)
+      , Monoid w)
    => req a -> GenHaxl u w a -> GenHaxl u w a
 
 cachedComputation req haxl = GenHaxl $ \env@Env{..} -> do
@@ -101,7 +102,8 @@
   :: forall req u w a.
      ( Eq (req a)
      , Hashable (req a)
-     , Typeable (req a))
+     , Typeable (req a)
+     , Monoid w)
   => req a -> GenHaxl u w a -> GenHaxl u w a
 preCacheComputation req haxl = GenHaxl $ \env@Env{..} -> do
   mbRes <- DataCache.lookup req memoCache
@@ -196,7 +198,7 @@
 -- >   b <- g
 -- >   return (a + b)
 --
-runMemo :: MemoVar u w a -> GenHaxl u w a
+runMemo :: Monoid w => MemoVar u w a -> GenHaxl u w a
 runMemo (MemoVar memoRef) = GenHaxl $ \env -> do
   stored <- readIORef memoRef
   case stored of
@@ -215,7 +217,8 @@
       unHaxl (getIVarWithWrites ivar) env
 
 execMemoNowProfiled
-  :: Env u w
+  :: Monoid w
+  => Env u w
   -> GenHaxl u w a
   -> IVar u w a
   -> CallId
@@ -246,9 +249,9 @@
       setAllocationCounter a1
       return ret
 
-execMemoNow :: Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
+execMemoNow :: Monoid w => Env u w -> GenHaxl u w a -> IVar u w a -> IO (Result u w a)
 execMemoNow env cont ivar = do
-  wlogs <- newIORef NilWrites
+  wlogs <- newIORef mempty
   let
     !menv = env { writeLogsRef = wlogs }
     -- use an env with empty writes, so we can memoize the extra
@@ -262,13 +265,13 @@
       throwIO e
     Right (Done a) -> trace_ "execMemoNow: Done" $ do
       wt <- readIORef wlogs
-      putIVar ivar (Ok a wt) env
-      mbModifyWLRef wt (writeLogsRef env)
+      putIVar ivar (Ok a (Just wt)) env
+      modifyIORef' (writeLogsRef env) (<> wt)
       return (Done a)
     Right (Throw ex) -> trace_ ("execMemoNow: Throw" ++ show ex) $ do
       wt <- readIORef wlogs
-      putIVar ivar (ThrowHaxl ex wt) env
-      mbModifyWLRef wt (writeLogsRef env)
+      putIVar ivar (ThrowHaxl ex (Just wt)) env
+      modifyIORef' (writeLogsRef env) (<> wt)
       return (Throw ex)
     Right (Blocked ivar' cont) -> trace_ "execMemoNow: Blocked" $ do
       -- We "block" this memoized computation in the new environment 'menv', so
@@ -305,7 +308,7 @@
 prepareMemo1 (MemoVar1 r) f
   = unsafeLiftIO $ writeIORef r (MemoTbl1 f HashMap.empty)
 
-runMemo1 :: (Eq a, Hashable a) => MemoVar1 u w a b -> a -> GenHaxl u w b
+runMemo1 :: (Eq a, Hashable a, Monoid w) => MemoVar1 u w a b -> a -> GenHaxl u w b
 runMemo1 (MemoVar1 r) k = unsafeLiftIO (readIORef r) >>= \case
   MemoEmpty1 -> throw $ CriticalError "Attempting to run empty memo."
   MemoTbl1 f h -> case HashMap.lookup k h of
@@ -325,7 +328,7 @@
 prepareMemo2 (MemoVar2 r) f
   = unsafeLiftIO $ writeIORef r (MemoTbl2 f HashMap.empty)
 
-runMemo2 :: (Eq a, Hashable a, Eq b, Hashable b)
+runMemo2 :: (Eq a, Hashable a, Eq b, Hashable b, Monoid w)
          => MemoVar2 u w a b c
          -> a -> b -> GenHaxl u w c
 runMemo2 (MemoVar2 r) k1 k2 = unsafeLiftIO (readIORef r) >>= \case
@@ -353,12 +356,12 @@
 -- every two calls @memo key haxl@, if they have the same @key@ then
 -- they compute the same result.
 memo
-  :: (Typeable a, Typeable k, Hashable k, Eq k)
+  :: (Typeable a, Typeable k, Hashable k, Eq k, Monoid w)
   => k -> GenHaxl u w a -> GenHaxl u w a
 memo key = cachedComputation (MemoKey key)
 
 {-# RULES
-"memo/Text" memo = memoText :: (Typeable a) =>
+"memo/Text" memo = memoText :: (Typeable a, Monoid w) =>
             Text -> GenHaxl u w a -> GenHaxl u w a
  #-}
 
@@ -367,7 +370,7 @@
 -- | Memoize a computation using its location and a Fingerprint. This ensures
 -- uniqueness across computations.
 memoUnique
-  :: (Typeable a, Typeable k, Hashable k, Eq k)
+  :: (Typeable a, Typeable k, Hashable k, Eq k, Monoid w)
   => MemoFingerprintKey a -> Text -> k -> GenHaxl u w a -> GenHaxl u w a
 memoUnique fp label key = withLabel label . memo (fp, key)
 
@@ -394,7 +397,7 @@
 instance Hashable (MemoTextKey a) where
   hashWithSalt s (MemoText t) = hashWithSalt s t
 
-memoText :: (Typeable a) => Text -> GenHaxl u w a -> GenHaxl u w a
+memoText :: (Typeable a, Monoid w) => Text -> GenHaxl u w a -> GenHaxl u w a
 memoText key = withLabel key . cachedComputation (MemoText key)
 
 -- | A memo key derived from a 128-bit MD5 hash.  Do not use this directly,
@@ -424,7 +427,7 @@
 --
 {-# NOINLINE memoFingerprint #-}
 memoFingerprint
-  :: Typeable a => MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
+  :: (Typeable a, Monoid w) => MemoFingerprintKey a -> GenHaxl u w a -> GenHaxl u w a
 memoFingerprint key@(MemoFingerprintKey _ _ mnPtr nPtr) =
   withFingerprintLabel mnPtr nPtr . cachedComputation key
 
@@ -436,7 +439,7 @@
 -- in a @MemoVar@ (which @memoize@ creates), and returns the stored result on
 -- subsequent invocations. This permits the creation of local memos, whose
 -- lifetimes are scoped to the current function, rather than the entire request.
-memoize :: GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
+memoize :: Monoid w => GenHaxl u w a -> GenHaxl u w (GenHaxl u w a)
 memoize a = runMemo <$> newMemoWith a
 
 -- | Transform a 1-argument function returning a Haxl computation into a
@@ -455,7 +458,7 @@
 --
 -- The above implementation will not invoke the underlying @friendsOf@
 -- repeatedly for duplicate values in @ids@.
-memoize1 :: (Eq a, Hashable a)
+memoize1 :: (Eq a, Hashable a, Monoid w)
          => (a -> GenHaxl u w b)
          -> GenHaxl u w (a -> GenHaxl u w b)
 memoize1 f = runMemo1 <$> newMemoWith1 f
@@ -464,7 +467,7 @@
 -- memoized version of itself.
 --
 -- The 2-ary version of @memoize1@, see its documentation for details.
-memoize2 :: (Eq a, Hashable a, Eq b, Hashable b)
+memoize2 :: (Eq a, Hashable a, Eq b, Hashable b, Monoid w)
          => (a -> b -> GenHaxl u w c)
          -> GenHaxl u w (a -> b -> GenHaxl u w c)
 memoize2 f = runMemo2 <$> newMemoWith2 f
diff --git a/Haxl/Core/Monad.hs b/Haxl/Core/Monad.hs
--- a/Haxl/Core/Monad.hs
+++ b/Haxl/Core/Monad.hs
@@ -49,6 +49,7 @@
   , appendWTs
   , mbModifyWLRef
   , mapWrites
+  , mapWriteTree
 
     -- * Cont
   , Cont(..)
@@ -140,6 +141,7 @@
 import Data.Either (rights)
 import Data.IORef
 import Data.Int
+import Data.Maybe
 import Data.List.NonEmpty (NonEmpty(..))
 import qualified Data.List.NonEmpty as NonEmpty
 #if __GLASGOW_HASKELL__ < 804
@@ -235,12 +237,12 @@
        -- become non-empty is how the scheduler blocks waiting for
        -- data fetches to return.
 
-  , writeLogsRef :: {-# UNPACK #-} !(IORef (WriteTree w))
+  , writeLogsRef :: {-# UNPACK #-} !(IORef w)
        -- ^ A log of all writes done as part of this haxl computation. Any
        -- haxl computation that needs to be memoized runs in its own
        -- environment so that we can get a hold of those writes and put them
        -- in the IVar associated with the compuatation.
-  , writeLogsRefNoMemo :: {-# UNPACK #-} !(IORef (WriteTree w))
+  , writeLogsRefNoMemo :: {-# UNPACK #-} !(IORef w)
        -- ^ This is just a specialized version of @writeLogsRef@, where we put
        -- logs that user doesn't want memoized. This is a better alternative to
        -- doing arbitrary IO from a (memoized) Haxl computation.
@@ -280,7 +282,7 @@
 
 -- | Initialize an environment with a 'StateStore', an input map, a
 -- preexisting 'DataCache', and a seed for the random number generator.
-initEnvWithData :: StateStore -> u -> Caches u w -> IO (Env u w)
+initEnvWithData :: Monoid w => StateStore -> u -> Caches u w -> IO (Env u w)
 initEnvWithData states e (dcache, mcache) = do
   newCid <- max <$>
     (maybe 0 ((+) 1) <$> getMaxCallId dcache) <*>
@@ -293,8 +295,8 @@
   rq <- newIORef JobNil              -- RunQueue
   sr <- newIORef emptyReqCounts      -- SubmittedReqs
   comps <- newTVarIO []              -- completion queue
-  wl <- newIORef NilWrites
-  wlnm <- newIORef NilWrites
+  wl <- newIORef mempty
+  wlnm <- newIORef mempty
   return Env
     { dataCache = dcache
     , memoCache = mcache
@@ -321,14 +323,14 @@
     }
 
 -- | Initializes an environment with 'StateStore' and an input map.
-initEnv :: StateStore -> u -> IO (Env u w)
+initEnv :: Monoid w => StateStore -> u -> IO (Env u w)
 initEnv states e = do
   dcache <- emptyDataCache
   mcache <- emptyDataCache
   initEnvWithData states e (dcache, mcache)
 
 -- | A new, empty environment.
-emptyEnv :: u -> IO (Env u w)
+emptyEnv :: Monoid w => u -> IO (Env u w)
 emptyEnv = initEnv stateEmpty
 
 -- | If you're using the env from a failed Haxl computation in a second Haxl
@@ -376,6 +378,12 @@
   | MergeWrites (WriteTree w) (WriteTree w)
   deriving (Show)
 
+instance Semigroup (WriteTree w) where
+  (<>) = appendWTs
+
+instance Monoid (WriteTree w) where
+  mempty = NilWrites
+
 appendWTs :: WriteTree w -> WriteTree w -> WriteTree w
 appendWTs NilWrites w = w
 appendWTs w NilWrites = w
@@ -431,20 +439,20 @@
 newtype GenHaxl u w a = GenHaxl
   { unHaxl :: Env u w -> IO (Result u w a) }
 
-tellWrite :: w -> GenHaxl u w ()
+tellWrite :: w -> GenHaxl u (WriteTree w) ()
 tellWrite = write . SomeWrite
 
-write :: WriteTree w -> GenHaxl u w ()
+write :: Monoid w => w -> GenHaxl u w ()
 write wt = GenHaxl $ \Env{..} -> do
-  mbModifyWLRef wt writeLogsRef
+  modifyIORef' writeLogsRef (<> wt)
   return $ Done ()
 
-tellWriteNoMemo :: w -> GenHaxl u w ()
+tellWriteNoMemo :: w -> GenHaxl u (WriteTree w) ()
 tellWriteNoMemo = writeNoMemo . SomeWrite
 
-writeNoMemo :: WriteTree w -> GenHaxl u w ()
+writeNoMemo :: Monoid w => w -> GenHaxl u w ()
 writeNoMemo wt = GenHaxl $ \Env{..} -> do
-  mbModifyWLRef wt writeLogsRefNoMemo
+  modifyIORef' writeLogsRefNoMemo (<> wt)
   return $ Done ()
 
 
@@ -557,15 +565,15 @@
       return (Blocked i (Cont (getIVarApply i a)))
 
 -- Another specialised version of getIVar, for efficiency in cachedComputation
-getIVarWithWrites :: IVar u w a -> GenHaxl u w a
+getIVarWithWrites :: Monoid w => IVar u w a -> GenHaxl u w a
 getIVarWithWrites i@IVar{ivarRef = !ref} = GenHaxl $ \env@Env{..} -> do
   e <- readIORef ref
   case e of
     IVarFull (Ok a wt) -> do
-      mbModifyWLRef wt writeLogsRef
+      modifyIORef' writeLogsRef (<> fromMaybe mempty wt)
       return (Done a)
     IVarFull (ThrowHaxl e wt) -> do
-      mbModifyWLRef wt writeLogsRef
+      modifyIORef' writeLogsRef (<> fromMaybe mempty wt)
       raiseFromIVar env i e
     IVarFull (ThrowIO e) -> throwIO e
     IVarEmpty _ ->
@@ -604,8 +612,8 @@
 -- that when the result is fetched using getIVar, we can throw the
 -- exception in the right way.
 data ResultVal a w
-  = Ok a (WriteTree w)
-  | ThrowHaxl SomeException (WriteTree w)
+  = Ok a (Maybe w)
+  | ThrowHaxl SomeException (Maybe w)
   | ThrowIO SomeException
     -- we get no write logs when an IO exception occurs
 
@@ -615,14 +623,14 @@
 done _ (ThrowIO e) = throwIO e
 
 eitherToResultThrowIO :: Either SomeException a -> ResultVal a w
-eitherToResultThrowIO (Right a) = Ok a NilWrites
+eitherToResultThrowIO (Right a) = Ok a Nothing
 eitherToResultThrowIO (Left e)
-  | Just HaxlException{} <- fromException e = ThrowHaxl e NilWrites
+  | Just HaxlException{} <- fromException e = ThrowHaxl e Nothing
   | otherwise = ThrowIO e
 
 eitherToResult :: Either SomeException a -> ResultVal a w
-eitherToResult (Right a) = Ok a NilWrites
-eitherToResult (Left e) = ThrowHaxl e NilWrites
+eitherToResult (Right a) = Ok a Nothing
+eitherToResult (Left e) = ThrowHaxl e Nothing
 
 
 -- -----------------------------------------------------------------------------
@@ -698,6 +706,11 @@
   show (Throw e) = printf "Throw(%s)" $ show e
   show Blocked{} = "Blocked"
 
+instance Functor (Result u w) where
+  fmap f (Done a) = Done (f a)
+  fmap _ (Throw exc) = Throw exc
+  fmap f (Blocked ivar cont) = Blocked ivar (f :<$> cont)
+
 {- Note [Exception]
 
 How do we want to represent Haxl exceptions (those that are thrown by
@@ -775,7 +788,7 @@
 -- Monad/Applicative instances
 
 instance Monad (GenHaxl u w) where
-  return a = GenHaxl $ \_env -> return (Done a)
+  return = pure
   GenHaxl m >>= k = GenHaxl $ \env -> do
     e <- m env
     case e of
@@ -803,7 +816,7 @@
         return (Blocked ivar (f :<$> cont))
 
 instance Applicative (GenHaxl u w) where
-  pure = return
+  pure a = GenHaxl $ \_env -> pure (Done a)
   GenHaxl ff <*> GenHaxl aa = GenHaxl $ \env -> do
     rf <- ff env
     case rf of
@@ -831,7 +844,9 @@
 
 instance Monoid a => Monoid (GenHaxl u w a) where
   mempty = pure mempty
+#if __GLASGOW_HASKELL__ < 804
   mappend = liftA2 mappend
+#endif
 
 blockedBlocked
   :: Env u w
@@ -903,10 +918,10 @@
 -- Memoization behavior is unchanged, meaning if a memoized computation is run
 -- once inside @mapWrites@ and then once without, the writes from the second run
 -- will NOT be transformed.
-mapWrites :: (w -> w) -> GenHaxl u w a -> GenHaxl u w a
+mapWrites :: Monoid w => (w -> w) -> GenHaxl u w a -> GenHaxl u w a
 mapWrites f action = GenHaxl $ \curEnv -> do
-  wlogs <- newIORef NilWrites
-  wlogsNoMemo <- newIORef NilWrites
+  wlogs <- newIORef mempty
+  wlogsNoMemo <- newIORef mempty
   let
     !newEnv = curEnv { writeLogsRef = wlogs, writeLogsRefNoMemo = wlogsNoMemo }
   unHaxl (mapWritesImpl curEnv newEnv action) newEnv
@@ -915,9 +930,9 @@
       let
         pushTransformedWrites = do
           wt <- readIORef $ writeLogsRef curEnv
-          mbModifyWLRef (mapWriteTree f wt) (writeLogsRef oldEnv)
+          modifyIORef' (writeLogsRef oldEnv) (<> f wt)
           wtNoMemo <- readIORef $ writeLogsRefNoMemo curEnv
-          mbModifyWLRef (mapWriteTree f wtNoMemo) (writeLogsRefNoMemo oldEnv)
+          modifyIORef' (writeLogsRefNoMemo oldEnv) (<> f wtNoMemo)
 
       r <- m curEnv
 
diff --git a/Haxl/Core/Run.hs b/Haxl/Core/Run.hs
--- a/Haxl/Core/Run.hs
+++ b/Haxl/Core/Run.hs
@@ -22,6 +22,7 @@
 import Control.Exception as Exception
 import Control.Monad
 import Data.IORef
+import Data.Maybe
 import Text.Printf
 import Unsafe.Coerce
 
@@ -49,10 +50,10 @@
 --
 -- However, multiple 'Env's may share a single 'StateStore', and thereby
 -- use the same set of datasources.
-runHaxl:: forall u w a. Env u w -> GenHaxl u w a -> IO a
+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. Env u w -> GenHaxl u w a -> IO (a, [w])
+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
@@ -105,10 +106,10 @@
           result (ThrowIO e)
         Right (Done a) -> do
           wt <- readIORef writeLogsRef
-          result (Ok a wt)
+          result $ Ok a (Just wt)
         Right (Throw ex) -> do
           wt <- readIORef writeLogsRef
-          result (ThrowHaxl ex wt)
+          result $ ThrowHaxl ex (Just wt)
         Right (Blocked i fn) -> do
           addJob env (toHaxl fn) ivar i
           reschedule env rq
@@ -237,13 +238,13 @@
   --
   schedule env JobNil haxl result
   r <- readIORef resultRef
-  writeIORef writeLogsRef NilWrites
+  writeIORef writeLogsRef mempty
   wtNoMemo <- atomicModifyIORef' writeLogsRefNoMemo
-    (\old_wrts -> (NilWrites , old_wrts))
+    (\old_wrts -> (mempty, old_wrts))
   case r of
     IVarEmpty _ -> throwIO (CriticalError "runHaxl: missing result")
     IVarFull (Ok a wt) -> do
-      return (a, flattenWT (wt `appendWTs` wtNoMemo))
+      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.
diff --git a/Haxl/Core/StateStore.hs b/Haxl/Core/StateStore.hs
--- a/Haxl/Core/StateStore.hs
+++ b/Haxl/Core/StateStore.hs
@@ -49,13 +49,15 @@
 
 #if __GLASGOW_HASKELL__ >= 804
 instance Semigroup StateStore where
-  (<>) = mappend
+  -- Left-biased union
+  StateStore m1 <> StateStore m2 = StateStore $ m1 <> m2
 #endif
 
 instance Monoid StateStore where
   mempty = stateEmpty
-  -- Left-biased union
+#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.
diff --git a/changelog.md b/changelog.md
--- a/changelog.md
+++ b/changelog.md
@@ -1,3 +1,6 @@
+# Changes in version 2.5.1.0
+  * Add schedulerHintState method to DataSource
+
 # Changes in version 2.4.0.0
   * Added fetchBatchId to FetchStats
   * Profiling now tracks full stacks and links each label to memos/fetches
diff --git a/haxl.cabal b/haxl.cabal
--- a/haxl.cabal
+++ b/haxl.cabal
@@ -1,5 +1,5 @@
 name:                haxl
-version:             2.4.0.0
+version:             2.5.1.0
 synopsis:            A Haskell library for efficient, concurrent,
                      and concise data access.
 homepage:            https://github.com/facebook/Haxl
@@ -14,11 +14,12 @@
 stability:           alpha
 cabal-version:       >= 1.10
 tested-with:
-  GHC==8.2.2,
   GHC==8.4.4,
   GHC==8.6.5,
-  GHC==8.8.3,
-  GHC==8.10.1
+  GHC==8.8.4,
+  GHC==8.10.7
+  GHC==9.0.1
+  GHC==9.2.2
 
 description:
   Haxl is a library and EDSL for efficient scheduling of concurrent data
@@ -46,18 +47,18 @@
     aeson >= 0.6 && < 2.1,
     base >= 4.10 && < 5,
     binary >= 0.7 && < 0.10,
-    bytestring >= 0.9 && < 0.11,
+    bytestring >= 0.9 && < 0.12,
     containers >= 0.5 && < 0.7,
     deepseq,
     exceptions >=0.8 && <0.11,
     filepath >= 1.3 && < 1.5,
     ghc-prim,
-    hashable >= 1.2 && < 1.4,
+    hashable >= 1.2 && < 1.5,
     hashtables >= 1.2.3.1,
     pretty == 1.1.*,
     -- text 1.2.1.0 required for instance Binary Text
     text >= 1.2.1.0 && < 1.3,
-    time >= 1.4 && < 1.10,
+    time >= 1.4 && < 1.12,
     stm >= 2.4 && < 2.6,
     transformers,
     unordered-containers == 0.2.*,
diff --git a/tests/CoreTests.hs b/tests/CoreTests.hs
--- a/tests/CoreTests.hs
+++ b/tests/CoreTests.hs
@@ -33,7 +33,7 @@
   let st = stateSet exstate stateEmpty
 
   -- Create the Env:
-  initEnv st ()
+  initEnv st () :: IO (Env () ())
 
 useless :: String -> GenHaxl u w Bool
 useless _ = throw (NotFound "ha ha")
@@ -41,7 +41,7 @@
 exceptions :: Assertion
 exceptions =
   do
-    en <- emptyEnv ()
+    en <- emptyEnv () :: IO (Env () ())
     a <- runHaxl en $ try (useless "input")
     assertBool "NotFound -> HaxlException" $
       isLeft (a :: Either HaxlException Bool)
@@ -132,7 +132,7 @@
 -- makes the compiler happy.
 base :: (Exception a) => a -> IO HaxlException
 base e = do
-  en <- emptyEnv ()
+  en <- emptyEnv () :: IO (Env () ())
   runHaxl en $ throw e `catch` \x -> return x
 
 printing :: Assertion
@@ -154,21 +154,21 @@
 withEnvTest :: Test
 withEnvTest = TestLabel "withEnvTest" $ TestCase $ do
   exstate <- ExampleDataSource.initGlobalState
-  e <- initEnv (stateSet exstate stateEmpty) False
+  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
+  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
+  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
+  e <- initEnv (stateSet exstate stateEmpty) False :: IO (Env Bool ())
   b <- runHaxl e $
     withEnv e { userEnv = True } $ do
       memo ("yyy" :: Text) $ do
diff --git a/tests/DataCacheTest.hs b/tests/DataCacheTest.hs
--- a/tests/DataCacheTest.hs
+++ b/tests/DataCacheTest.hs
@@ -55,8 +55,8 @@
   hashWithSalt s (CacheableInt val) = hashWithSalt s (0::Int, val)
 
 
-newResult :: a -> IO (IVar u w a)
-newResult a = newFullIVar (Ok a NilWrites)
+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
@@ -85,13 +85,13 @@
   r <- mapM takeResult =<< DataCache.lookup (Req 1) cache
   assertBool "dcSoundness2" $
     case r :: Maybe (ResultVal Int ()) of
-     Just (Ok 1 NilWrites) -> True
+     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" NilWrites) -> True
+      Just (Ok "hello" Nothing) -> True
       _something_else -> False
 
   r <- mapM takeResult =<< DataCache.lookup (Req 2) cache
@@ -103,7 +103,7 @@
 
 dcStrictnessTest :: Test
 dcStrictnessTest = TestLabel "DataCache strictness" $ TestCase $ do
-  env <- initEnv stateEmpty ()
+  env <- initEnv stateEmpty () :: IO (Env () ())
   r <- Control.Exception.try $ runHaxl env $
     cachedComputation (Req (error "BOOM")) $ return "OK"
   assertBool "dcStrictnessTest" $
@@ -145,7 +145,7 @@
         case rbad of
           Left (NotFound _) -> True
           _ -> False
-      assertEqual "write side effects happen" [7] writes
+      assertEqual "write side effects happen" [7] (flattenWT writes)
   ]
   where
 
@@ -159,14 +159,14 @@
                 ]
       return c
 
-    addLookup :: Env () Int -> Env () Int
+    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 Int))
+      -> 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
@@ -176,12 +176,12 @@
           return $ unsafeCoerce . doCache <$> cast r
       | otherwise = return Nothing
 
-    doReq :: TestReq Int -> ResultVal Int Int
-    doReq (Req 999) = ThrowHaxl (toException $ NotFound empty) NilWrites
-    doReq (Req r) = Ok r NilWrites
+    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 Int
-    doCache (CacheableInt i) = Ok i NilWrites
+    doCache :: CacheableReq Int -> ResultVal Int (WriteTree Int)
+    doCache (CacheableInt i) = Ok i Nothing
 
     req :: TestReq Int
     req = Req 1
diff --git a/tests/FullyAsyncTest.hs b/tests/FullyAsyncTest.hs
--- a/tests/FullyAsyncTest.hs
+++ b/tests/FullyAsyncTest.hs
@@ -20,6 +20,7 @@
 tests :: Test
 tests = sleepTest
 
+testEnv :: IO (Env () ())
 testEnv = do
   st <- mkConcurrentIOState
   env <- initEnv (stateSet st stateEmpty) ()
diff --git a/tests/MemoizationTests.hs b/tests/MemoizationTests.hs
--- a/tests/MemoizationTests.hs
+++ b/tests/MemoizationTests.hs
@@ -19,7 +19,7 @@
 memoSoundness = TestCase $ do
   iEnv <- do
     exState <- ExampleDataSource.initGlobalState
-    initEnv (stateSet exState stateEmpty) ()
+    initEnv (stateSet exState stateEmpty) () :: IO (Env () ())
 
   unMemoizedWombats <- runHaxl iEnv $ listWombats 100
 
diff --git a/tests/MonadAsyncTest.hs b/tests/MonadAsyncTest.hs
--- a/tests/MonadAsyncTest.hs
+++ b/tests/MonadAsyncTest.hs
@@ -19,7 +19,7 @@
 import Control.Concurrent
 import Control.Exception as Exception
 import Control.Monad
-import Haxl.Core.Monad (unsafeLiftIO)
+import Haxl.Core.Monad (unsafeLiftIO, WriteTree)
 import System.IO.Unsafe
 import Data.Hashable
 import Data.IORef
@@ -86,7 +86,7 @@
   [ TestLabel "exceptionTest" exceptionTest
   ]
 
-mkTestEnv :: IO (Env () SimpleWrite)
+mkTestEnv :: IO (Env () (WriteTree SimpleWrite))
 mkTestEnv = do
   st <- initDataSource
   initEnv (stateSet st stateEmpty) ()
diff --git a/tests/MonadBench.hs b/tests/MonadBench.hs
--- a/tests/MonadBench.hs
+++ b/tests/MonadBench.hs
@@ -22,20 +22,21 @@
 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 () SimpleWrite)
+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 () SimpleWrite ())
+type Test = (String, Int, Int -> GenHaxl () (WriteTree SimpleWrite) ())
 
 testName :: Test -> String
 testName (t,_,_) = t
@@ -151,19 +152,20 @@
 
 tree
   :: Int
-  -> (Int -> GenHaxl () SimpleWrite [Id] -> GenHaxl () SimpleWrite [Id])
-  -> GenHaxl () SimpleWrite [Id]
+  -> (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 () SimpleWrite [Id]
+unionWombats :: GenHaxl () (WriteTree SimpleWrite) [Id]
 unionWombats = foldl List.union [] <$> Haxl.mapM listWombats [1..1000]
 
-unionWombatsTo :: Id -> GenHaxl () SimpleWrite [Id]
+unionWombatsTo :: Id -> GenHaxl () (WriteTree SimpleWrite) [Id]
 unionWombatsTo x = foldl List.union [] <$> Haxl.mapM listWombats [1..x]
 
-unionWombatsFromTo :: Id -> Id -> GenHaxl () SimpleWrite [Id]
+unionWombatsFromTo :: Id -> Id -> GenHaxl () (WriteTree SimpleWrite) [Id]
 unionWombatsFromTo x y = foldl List.union [] <$> Haxl.mapM listWombats [x..y]
diff --git a/tests/ParallelTests.hs b/tests/ParallelTests.hs
--- a/tests/ParallelTests.hs
+++ b/tests/ParallelTests.hs
@@ -10,6 +10,7 @@
 
 import Test.HUnit
 
+testEnv :: IO (Env () ())
 testEnv = do
   sleepState <- mkConcurrentIOState
   let st = stateSet sleepState stateEmpty
diff --git a/tests/ProfileTests.hs b/tests/ProfileTests.hs
--- a/tests/ProfileTests.hs
+++ b/tests/ProfileTests.hs
@@ -31,10 +31,12 @@
 #endif
 import Data.Int
 
+import TestTypes
 import TestUtils
 import WorkDataSource
 import SleepDataSource
 
+mkProfilingEnv :: IO HaxlEnv
 mkProfilingEnv = do
   env <- makeTestEnv False
   return env { flags = (flags env) { report = profilingReportFlags } }
@@ -137,7 +139,7 @@
 -- for correct accounting when relying on allocation limits.
 threadAlloc :: Integer -> Assertion
 threadAlloc batches = do
-  env' <- initEnv (stateSet mkWorkState stateEmpty) ()
+  env' <- initEnv (stateSet mkWorkState stateEmpty) () :: IO (Env () ())
   let env = env'  { flags = (flags env') {
     report = setReportFlag ReportFetchStats defaultReportFlags } }
   a0 <- getAllocationCounter
diff --git a/tests/StatsTests.hs b/tests/StatsTests.hs
--- a/tests/StatsTests.hs
+++ b/tests/StatsTests.hs
@@ -66,7 +66,7 @@
   assertEqual
     "Grouping works as expected" expectedResultInterspersed aggInterspersedBatch
 
-
+testEnv :: IO (Env () ())
 testEnv = do
   -- To use a data source, we need to initialize its state:
   exstate <- ExampleDataSource.initGlobalState
diff --git a/tests/TestExampleDataSource.hs b/tests/TestExampleDataSource.hs
--- a/tests/TestExampleDataSource.hs
+++ b/tests/TestExampleDataSource.hs
@@ -24,6 +24,7 @@
 import ExampleDataSource
 import LoadCache
 
+testEnv :: IO (Env () ())
 testEnv = do
   -- To use a data source, we need to initialize its state:
   exstate <- ExampleDataSource.initGlobalState
diff --git a/tests/WriteTests.hs b/tests/WriteTests.hs
--- a/tests/WriteTests.hs
+++ b/tests/WriteTests.hs
@@ -15,6 +15,7 @@
 
 import Test.HUnit
 
+import Control.Arrow
 import Control.Concurrent
 import Data.Either
 import Data.Foldable
@@ -22,7 +23,7 @@
 import Data.IORef
 import qualified Data.Text as Text
 
-import Haxl.Core.Monad (mapWrites, flattenWT)
+import Haxl.Core.Monad (mapWrites, mapWriteTree, flattenWT, WriteTree)
 import Haxl.Core
 import Haxl.Prelude as Haxl
 
@@ -58,12 +59,12 @@
 assertEqualIgnoreOrder msg lhs rhs =
   assertEqual msg (sort lhs) (sort rhs)
 
-doInnerWrite :: GenHaxl u SimpleWrite Int
+doInnerWrite :: GenHaxl u (WriteTree SimpleWrite) Int
 doInnerWrite = do
   tellWrite $ SimpleWrite "inner"
   return 0
 
-doOuterWrite :: GenHaxl u SimpleWrite Int
+doOuterWrite :: GenHaxl u (WriteTree SimpleWrite) Int
 doOuterWrite = do
   tellWrite $ SimpleWrite "outer1"
 
@@ -76,19 +77,22 @@
 
   return 1
 
-doNonMemoWrites :: GenHaxl u SimpleWrite Int
+doNonMemoWrites :: GenHaxl u (WriteTree SimpleWrite) Int
 doNonMemoWrites = do
   tellWrite $ SimpleWrite "inner"
   tellWriteNoMemo $ SimpleWrite "inner not memo"
   return 0
 
+runHaxlWithWriteList :: Env u (WriteTree w) -> GenHaxl u (WriteTree w) a -> IO (a, [w])
+runHaxlWithWriteList env haxl = second flattenWT <$> runHaxlWithWrites env haxl
+
 writeSoundness :: Test
 writeSoundness = TestCase $ do
   let numReps = 4
 
   -- do writes without memoization
   env1 <- emptyEnv ()
-  (allRes, allWrites) <- runHaxlWithWrites env1 $
+  (allRes, allWrites) <- runHaxlWithWriteList env1 $
     Haxl.sequence (replicate numReps doInnerWrite)
 
   assertBool "Write Soundness 1" $
@@ -98,7 +102,7 @@
   -- do writes with memoization
   env2 <- emptyEnv ()
 
-  (memoRes, memoWrites) <- runHaxlWithWrites env2 $ do
+  (memoRes, memoWrites) <- runHaxlWithWriteList env2 $ do
     doWriteMemo <- newMemoWith doInnerWrite
     let memoizedWrite = runMemo doWriteMemo
 
@@ -111,7 +115,7 @@
   -- do writes with interleaved memo
   env3 <- emptyEnv ()
 
-  (ilRes, ilWrites) <- runHaxlWithWrites env3 $ do
+  (ilRes, ilWrites) <- runHaxlWithWriteList env3 $ do
     doWriteMemo <- newMemoWith doInnerWrite
     let memoizedWrite = runMemo doWriteMemo
 
@@ -124,7 +128,7 @@
   -- do writes with nested memo
   env4 <- emptyEnv ()
 
-  (nestRes, nestWrites) <- runHaxlWithWrites env4 $ do
+  (nestRes, nestWrites) <- runHaxlWithWriteList env4 $ do
     doWriteMemo' <- newMemoWith doOuterWrite
     let memoizedWrite' = runMemo doWriteMemo'
 
@@ -142,7 +146,7 @@
 
   -- do both kinds of writes without memoization
   env5 <- emptyEnv ()
-  (allRes, allWrites) <- runHaxlWithWrites env5 $
+  (allRes, allWrites) <- runHaxlWithWriteList env5 $
     Haxl.sequence (replicate numReps doNonMemoWrites)
 
   assertBool "Write Soundness 9" $
@@ -153,7 +157,7 @@
   -- do both kinds of writes with memoization
   env6 <- emptyEnv ()
 
-  (memoRes, memoWrites) <- runHaxlWithWrites env6 $ do
+  (memoRes, memoWrites) <- runHaxlWithWriteList env6 $ do
     doWriteMemo <- newMemoWith doNonMemoWrites
     let memoizedWrite = runMemo doWriteMemo
 
@@ -168,7 +172,7 @@
 writeLogsCorrectnessTest :: Test
 writeLogsCorrectnessTest = TestLabel "writeLogs_correctness" $ TestCase $ do
   e <- emptyEnv ()
-  (_ , wrts) <- runHaxlWithWrites e doNonMemoWrites
+  (_ , wrts) <- runHaxlWithWriteList e doNonMemoWrites
   assertEqualIgnoreOrder "Expected writes" [SimpleWrite "inner",
     SimpleWrite "inner not memo"] wrts
   wrtsNoMemo <- readIORef $ writeLogsRefNoMemo e
@@ -178,16 +182,17 @@
 
 mapWritesTest :: Test
 mapWritesTest = TestLabel "mapWrites" $ TestCase $ do
-  let func (SimpleWrite s) = SimpleWrite $ Text.toUpper s
+  let funcSingle (SimpleWrite s) = SimpleWrite $ Text.toUpper s
+      func = mapWriteTree funcSingle
   env0 <- emptyEnv ()
-  (res0, wrts0) <- runHaxlWithWrites env0 $ mapWrites func doNonMemoWrites
+  (res0, wrts0) <- runHaxlWithWriteList env0 $ mapWrites func doNonMemoWrites
   assertEqual "Expected computation result" 0 res0
   assertEqualIgnoreOrder "Writes correctly transformed" [SimpleWrite "INNER",
     SimpleWrite "INNER NOT MEMO"] wrts0
 
   -- Writes should behave the same inside and outside mapWrites
   env1 <- emptyEnv ()
-  (res1, wrts1) <- runHaxlWithWrites env1 $ do
+  (res1, wrts1) <- runHaxlWithWriteList env1 $ do
     outer <- doOuterWrite
     outerMapped <- mapWrites func doOuterWrite
     return $ outer == outerMapped
@@ -203,7 +208,7 @@
 
   -- Memoization behaviour should be unaffected
   env2 <- emptyEnv ()
-  (_res2, wrts2) <- runHaxlWithWrites env2 $ do
+  (_res2, wrts2) <- runHaxlWithWriteList env2 $ do
     writeMemo <- newMemoWith doNonMemoWrites
     let doWriteMemo = runMemo writeMemo
     _ <- mapWrites func doWriteMemo
@@ -220,7 +225,7 @@
 
   -- Same as previous, but the non-mapped computation is run first
   env3 <- emptyEnv ()
-  (_res3, wrts3) <- runHaxlWithWrites env3 $ do
+  (_res3, wrts3) <- runHaxlWithWriteList env3 $ do
     writeMemo <- newMemoWith doNonMemoWrites
     let doWriteMemo = runMemo writeMemo
     _ <- doWriteMemo
@@ -237,14 +242,14 @@
 
   -- inner computation performs no writes
   env4 <- emptyEnv ()
-  (res4, wrts4) <- runHaxlWithWrites env4 $
+  (res4, wrts4) <- runHaxlWithWriteList env4 $
     mapWrites func (return (0 :: Int))
   assertEqual "No Writes: Expected computation result" 0 res4
   assertEqualIgnoreOrder "No writes" [] wrts4
 
   -- inner computation throws an exception
   env5 <- emptyEnv ()
-  (res5, wrts5) <- runHaxlWithWrites env5 $ mapWrites func $ try $ do
+  (res5, wrts5) <- runHaxlWithWriteList env5 $ mapWrites func $ try $ do
     _ <- doNonMemoWrites
     _ <- throw (NotFound "exception")
     return 0
@@ -259,7 +264,7 @@
 
   -- inner computation calls a datasource
   env6 <- initEnv (stateSet DSState stateEmpty) ()
-  (res6, wrts6) <- runHaxlWithWrites env6 $ mapWrites func $ do
+  (res6, wrts6) <- runHaxlWithWriteList env6 $ mapWrites func $ do
     _ <- doNonMemoWrites
     dataFetch GetNumber
 
@@ -273,7 +278,7 @@
 
   -- inner computation calls a datasource, flipped calls
   env7 <- initEnv (stateSet DSState stateEmpty) ()
-  (res7, wrts7) <- runHaxlWithWrites env7 $ mapWrites func $ do
+  (res7, wrts7) <- runHaxlWithWriteList env7 $ mapWrites func $ do
     df <- dataFetch GetNumber
     _ <- doNonMemoWrites
     return df