allow a fallback datacache lookup (#139)

Summary:
Pull Request resolved: https://github.com/facebook/Haxl/pull/139

This diff allows the user to insert a fallback datacache lookup, which will allow test caches to be injected into the running of a Haxl computation

Reviewed By: simonmar

Differential Revision: D26691327

fbshipit-source-id: ad901697a6bff08adddd6cc98ab9c4e64227098f

Haxl/Core/Fetch.hs

@@ -212,17 +212,30 @@ dataFetchWithInsert showFn insertFn req =
       -- Check whether the data source wants to submit requests
       -- eagerly, or batch them up.
       --
-      let blockedFetch = BlockedFetch req rvar
-      let blockedFetchI = BlockedFetchInternal fid
-      case schedulerHint userEnv :: SchedulerHint r of
-        SubmitImmediately ->
-          performFetches env [BlockedFetches [blockedFetch] [blockedFetchI]]
-        TryToBatch ->
-          -- add the request to the RequestStore and continue
-          modifyIORef' reqStoreRef $ \bs ->
-            addRequest blockedFetch blockedFetchI bs
-      --
-      return $ Blocked ivar (Return ivar)
+      let
+        blockedFetch = BlockedFetch req rvar
+        blockedFetchI = BlockedFetchInternal fid
+        submitFetch = do
+          case schedulerHint userEnv :: SchedulerHint r of
+            SubmitImmediately ->
+              performFetches env [BlockedFetches [blockedFetch] [blockedFetchI]]
+            TryToBatch ->
+              -- add the request to the RequestStore and continue
+              modifyIORef' reqStoreRef $ \bs ->
+                addRequest blockedFetch blockedFetchI bs
+          return $ Blocked ivar (Return ivar)
+
+      -- if there is a fallback configured try that,
+      -- else dispatch the fetch
+      case dataCacheFetchFallback of
+        Nothing -> submitFetch
+        Just (DataCacheLookup dcl) -> do
+          mbFallbackRes <- dcl req
+          case mbFallbackRes of
+            Nothing -> submitFetch
+            Just fallbackRes -> do
+              putIVar ivar fallbackRes env
+              done fallbackRes
 
     -- Seen before but not fetched yet.  We're blocked, but we don't have
     -- to add the request to the RequestStore.

Haxl/Core/Monad.hs

@@ -76,6 +76,8 @@ module Haxl.Core.Monad
     -- * Env
   , Env(..)
   , DataCacheItem(..)
+  , DataCacheLookup(..)
+  , HaxlDataCache
   , Caches
   , caches
   , initEnvWithData
@@ -142,6 +144,7 @@ import Data.Either (rights)
 #if __GLASGOW_HASKELL__ < 804
 import Data.Semigroup
 #endif
+import Data.Typeable
 import GHC.Exts (IsString(..))
 import Text.PrettyPrint hiding ((<>))
 import Text.Printf
@@ -153,7 +156,6 @@ import Debug.Trace (traceEventIO)
 #ifdef PROFILING
 import qualified Data.Map as Map
 import Data.Text (Text)
-import Data.Typeable
 import Foreign.Ptr (Ptr)
 import GHC.Stack
 import Haxl.Core.CallGraph
@@ -165,12 +167,18 @@ import Haxl.Core.CallGraph
 -- | The data we carry around in the Haxl monad.
 
 data DataCacheItem u w a = DataCacheItem (IVar u w a) {-# UNPACK #-} !CallId
+type HaxlDataCache u w = DataCache (DataCacheItem u w)
+newtype DataCacheLookup w =
+  DataCacheLookup
+    (forall req a . Typeable (req a)
+    => req a
+    -> IO (Maybe (ResultVal a w)))
 
 data Env u w = Env
-  { dataCache     :: {-# UNPACK #-} !(DataCache (DataCacheItem u w))
+  { dataCache     :: {-# UNPACK #-} !(HaxlDataCache u w)
       -- ^ cached data fetches
 
-  , memoCache      :: {-# UNPACK #-} !(DataCache (DataCacheItem u w))
+  , memoCache      :: {-# UNPACK #-} !(HaxlDataCache u w)
       -- ^ memoized computations
 
   , memoKey    :: {-# UNPACK #-} !CallId
@@ -234,6 +242,11 @@ data Env u w = Env
        -- ^ This is just a specialized version of @writeLogsRef@, where we put
        -- logs that user doesn't want memoized. This is a better alternative to
        -- doing arbitrary IO from a (memoized) Haxl computation.
+
+  , dataCacheFetchFallback :: !(Maybe (DataCacheLookup w))
+       -- ^ Allows you to inject a DataCache lookup just before a dataFetch is
+       -- dispatched. This is useful for injecting fetch results in testing.
+
 #ifdef PROFILING
   , callGraphRef ::  Maybe (IORef CallGraph)
        -- ^ An edge list representing the current function call graph. The type
@@ -249,12 +262,12 @@ data ProfileCurrent = ProfileCurrent
   , profCurrentLabel :: {-# UNPACK #-} !ProfileLabel
   }
 
-type Caches u w = (DataCache (DataCacheItem u w), DataCache (DataCacheItem u w))
+type Caches u w = (HaxlDataCache u w, HaxlDataCache u w)
 
 caches :: Env u w -> Caches u w
 caches env = (dataCache env, memoCache env)
 
-getMaxCallId :: DataCache (DataCacheItem u w) -> IO (Maybe Int)
+getMaxCallId :: HaxlDataCache u w -> IO (Maybe Int)
 getMaxCallId c = do
   callIds  <- rights . concatMap snd <$>
               DataCache.readCache c (\(DataCacheItem _ i) -> return i)
@@ -298,6 +311,7 @@ initEnvWithData states e (dcache, mcache) = do
     , completions = comps
     , writeLogsRef = wl
     , writeLogsRefNoMemo = wlnm
+    , dataCacheFetchFallback = Nothing
 #ifdef PROFILING
     , callGraphRef = Nothing
     , currFunction = mainFunction

tests/DataCacheTest.hs

@@ -5,6 +5,10 @@
 -- 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
@@ -18,6 +22,8 @@ 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
@@ -29,6 +35,17 @@ 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
+
 newResult :: a -> IO (IVar u w a)
 newResult a = IVar <$> newIORef (IVarFull (Ok a NilWrites))
 
@@ -85,5 +102,40 @@ dcStrictnessTest = TestLabel "DataCache strictness" $ TestCase $ do
       Left (ErrorCall "BOOM") -> True
       _other -> False
 
+dcFallbackTest :: Test
+dcFallbackTest = TestLabel "DataCache fallback" $ TestCase $ do
+  env <- addLookup <$> initEnv (stateSet TestReqState stateEmpty) ()
+  r <- runHaxl env (dataFetch req)
+  assertEqual "dcFallbackTest found" 1 r
+  rbad <- Control.Exception.try $ runHaxl env (dataFetch reqBad)
+  assertBool "dcFallbackTest not found" $
+    case rbad of
+      Left (ErrorCall "no fetch defined") -> True
+      _ -> False
+  where
+    addLookup e = e { dataCacheFetchFallback = Just (DataCacheLookup lookup) }
+    lookup
+      :: forall req a . Typeable (req a)
+      => req a
+      -> IO (Maybe (ResultVal a ()))
+    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
+      | otherwise = return Nothing
+
+    doReq :: TestReq Int -> ResultVal Int ()
+    doReq (Req r) = Ok r NilWrites
+
+    req :: TestReq Int
+    req = Req 1
+
+    reqBad :: TestReq String
+    reqBad = Req 2
+
 -- tests :: Assertion
-tests = TestList [dcSoundnessTest, dcStrictnessTest]
+tests = TestList [ dcSoundnessTest
+                 , dcStrictnessTest
+                 , dcFallbackTest
+                 ]