Allow any Monoid to be used as write storage in Haxl monad

Summary:
Currently, the `Haxl` monad stores writes using a hardcoded sequential data structure, a `WriteTree`. This makes it efficient to append new writes and merge two collections of writes, operations needed for `Haxl` to perform memoization efficiently.

However, in some cases storing and reading logs this way is inefficient in terms of space and time. A very simple example is a counter that can be incremented during `Haxl` computations. A `WriteTree Int` would store one entry per increment, while this case could be trivially represented using a single `Int`.

On this diff I'm changing the core monad implementation to accept any monoid as write structure. Rather than dealing with individual write objects inserted into a collection, writes are provided as monoids which get merged on each write.

**This is a breaking change**, as Haxl's `w` now represents the whole write log rather than a single write. However, `WriteTree` was kept and made a monoid, so `WriteTree w` can be used by clients in place of `w` to keep the same behavior.

Reviewed By: josefs

Differential Revision: D39305208

fbshipit-source-id: 7b488682280cda170e8b156765f69c3377bcadfc

Haxl/Core/Memo.hs

@@ -71,7 +71,8 @@ cachedComputation
  :: 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 @@ preCacheComputation
  :: 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 @@ newMemoWith memoCmp = do
 -- > 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 @@ runMemo (MemoVar memoRef) = GenHaxl $ \env -> do
  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 @@ execMemoNowProfiled envOuter cont ivar cid =
  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 @@ execMemoNow env cont ivar = do
  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 u w a b -> (a -> GenHaxl u w b) -> GenHaxl u w ()
 prepareMemo1 (MemoVar1 r) f
  = unsafeLiftIO $ writeIORef r (MemoTbl1 f HashMap.empty)
 
-runMemo1 :: (Eq a, Hashable a) => 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 u w a b c -> (a -> b -> GenHaxl u w c) -> GenHaxl u w (
 prepareMemo2 (MemoVar2 r) f
  = unsafeLiftIO $ writeIORef r (MemoTbl2 f HashMap.empty)
 
-runMemo2 :: (Eq a, Hashable a, Eq b, Hashable b)
+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 @@ runMemo2 (MemoVar2 r) k1 k2 = unsafeLiftIO (readIORef r) >>= \case
 -- 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 @@ memo key = cachedComputation (MemoKey key)
 -- | 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 @@ deriving instance Eq (MemoTextKey a)
 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 @@ instance Hashable (MemoFingerprintKey a) where
 --
 {-# 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 @@ memoFingerprint key@(MemoFingerprintKey _ _ mnPtr nPtr) =
 -- 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 @@ memoize a = runMemo <$> newMemoWith a
 --
 -- 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 @@ memoize1 f = runMemo1 <$> newMemoWith1 f
 -- 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

Haxl/Core/Monad.hs

@@ -140,6 +140,7 @@ import qualified Control.Exception as Exception
 import Data.Either (rights)
 import Data.IORef
 import Data.Int
+import Data.Maybe
 import Data.List.NonEmpty (NonEmpty(..))
 import qualified Data.List.NonEmpty as NonEmpty
 #if __GLASGOW_HASKELL__ < 804
@@ -235,12 +236,12 @@ data Env u w = Env
  -- 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 +281,7 @@ getMaxCallId c = do
 
 -- | 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 +294,8 @@ initEnvWithData states e (dcache, mcache) = do
  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 +322,14 @@ initEnvWithData states e (dcache, mcache) = do
  }
 
 -- | 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 +377,12 @@ data WriteTree w
  | MergeWrites (WriteTree w) (WriteTree w)
  deriving (Show)
 
+instance Semigroup (WriteTree w) where
+ (<>) = appendWTs
+
+instance Monoid (WriteTree w) where
+ mempty = NilWrites
+
 appendWTs :: WriteTree w -> WriteTree w -> WriteTree w
 appendWTs NilWrites w = w
 appendWTs w NilWrites = w
@@ -431,20 +438,20 @@ mapWriteTree f (MergeWrites wt1 wt2) =
 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 +564,15 @@ getIVarApply i@IVar{ivarRef = !ref} a = GenHaxl $ \env -> do
  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 +611,8 @@ addJobPanic = error "addJob: not empty"
 -- 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 +622,14 @@ done env (ThrowHaxl e _) = raise env e
 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
 
 
 -- -----------------------------------------------------------------------------
@@ -905,7 +912,7 @@ nextCallId env = atomicModifyIORef' (callIdRef env) $ \x -> (x+1,x+1)
 -- 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 :: (w -> w) -> GenHaxl u (WriteTree w) a -> GenHaxl u (WriteTree w) a
 mapWrites f action = GenHaxl $ \curEnv -> do
  wlogs <- newIORef NilWrites
  wlogsNoMemo <- newIORef NilWrites

Haxl/Core/Run.hs

@@ -22,6 +22,7 @@ import Control.Concurrent.STM
 import Control.Exception as Exception
 import Control.Monad
 import Data.IORef
+import Data.Maybe
 import Text.Printf
 import Unsafe.Coerce
 
@@ -49,10 +50,10 @@ import qualified Data.HashTable.IO as H
 --
 -- 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 @@ runHaxlWithWrites env@Env{..} haxl = do
  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 @@ runHaxlWithWrites env@Env{..} haxl = do
  --
  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.

haxl.cabal

@@ -1,5 +1,5 @@
 name: haxl
-version: 2.4.0.0
+version: 2.5.0.0
 synopsis: A Haskell library for efficient, concurrent,
  and concise data access.
 homepage: https:/facebook/Haxl