Avoid duplicate in-flight RPC slot fetches

src/EVM/Fetch.hs

@@ -19,9 +19,10 @@ module EVM.Fetch
   , saveCache
   , RPCContract (..)
   , makeContractFromRPC
-  -- Below 4 are needed for Echidna
+  -- Below 5 are needed for Echidna and deterministic fetch tests
   , fetchSlotWithSession
   , fetchSlotWithCache
+  , fetchSlotWithCacheUsing
   , fetchWithSession
   , getCacheState
   , FetchStatus(..)
@@ -47,7 +48,7 @@ import System.Directory (createDirectoryIfMissing, doesFileExist)
 import Data.Aeson.Encode.Pretty (encodePretty)
 import qualified Data.ByteString.Lazy as BSL
 import Data.Bifunctor (first)
-import Control.Exception (SomeException, SomeAsyncException(..), try, catches, Handler(..), throwIO)
+import Control.Exception (SomeException, SomeAsyncException(..), try, catches, Handler(..), throwIO, mask)
 
 import Data.Aeson hiding (Error)
 import Data.Aeson.Optics
@@ -70,7 +71,7 @@ import Control.Monad (when)
 import EVM.Effects
 import qualified EVM.Expr as Expr
 import Control.Concurrent (threadDelay)
-import Control.Concurrent.MVar (MVar, newMVar, readMVar, modifyMVar_)
+import Control.Concurrent.MVar (MVar, newEmptyMVar, newMVar, putMVar, readMVar, modifyMVar, modifyMVar_)
 import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef')
 import Data.Time.Clock (UTCTime, getCurrentTime, diffUTCTime, addUTCTime)
 import System.Random (randomRIO)
@@ -98,8 +99,17 @@ data Session = Session
   -- Shared rate limit cooldown across workers. When any worker hits a
   -- rate limit, it sets a deadline; other workers wait before retrying.
   , rpcThrottle     :: IORef (Maybe UTCTime)
+  , inFlightSlots   :: MVar (Map.Map SlotFetchKey (MVar SlotFetchResult))
   }
 
+type SlotFetchKey = (BlockNumber, Addr, W256)
+
+type SlotFetchResult = Either SomeException (FetchResult W256)
+
+data InFlightFetch a
+  = InFlightOwner (MVar (Either SomeException a))
+  | InFlightWaiter (MVar (Either SomeException a))
+
 data FetchCache = FetchCache
   { contractCache :: Map.Map Addr RPCContract
   , slotCache     :: Map.Map (Addr, W256) W256
@@ -146,7 +156,7 @@ data RpcQuery a where
   QueryChainId ::                 RpcQuery W256
 
 data BlockNumber = Latest | BlockNumber W256
-  deriving (Show, Eq)
+  deriving (Show, Eq, Ord)
 
 deriving instance Show (RpcQuery a)
 
@@ -420,33 +430,110 @@ makeContractFromRPC (RPCContract (ByteStringS code) nonce balance) =
 
 -- Needed for Echidna only
 fetchSlotWithCache :: Config -> Session -> BlockNumber -> Text -> Addr -> W256 -> IO (FetchResult W256)
-fetchSlotWithCache conf sess nPre url addr slot = do
+fetchSlotWithCache conf sess nPre url addr slot =
+  fetchSlotWithCacheUsing
+    (\n a s -> fetchQuery n (fetchWithRetry conf.debug url sess) (QuerySlot a s))
+    conf
+    sess
+    nPre
+    url
+    addr
+    slot
+
+fetchSlotWithCacheUsing
+  :: (BlockNumber -> Addr -> W256 -> IO (Either Text W256))
+  -> Config
+  -> Session
+  -> BlockNumber
+  -> Text
+  -> Addr
+  -> W256
+  -> IO (FetchResult W256)
+fetchSlotWithCacheUsing fetch conf sess nPre url addr slot = do
   n <- getLatestBlockNum conf sess nPre url
-  -- Check successful cache
+  cached <- lookupSlotCache conf sess addr slot
+  case cached of
+    Just res -> pure res
+    Nothing ->
+      withInFlight sess.inFlightSlots cachedSuccess (n, addr, slot) $
+        fetchFreshSlot fetch conf sess n addr slot
+
+lookupSlotCache :: Config -> Session -> Addr -> W256 -> IO (Maybe (FetchResult W256))
+lookupSlotCache conf sess addr slot = do
   cache <- readMVar sess.sharedCache
   case Map.lookup (addr, slot) cache.slotCache of
     Just s -> do
       when (conf.debug) $ putStrLn $ "-> Using cached slot value for slot " <> show slot <> " at " <> show addr
-      pure (FetchSuccess s Cached)
+      pure (Just (FetchSuccess s Cached))
     Nothing -> do
-      -- Check failure cache
       failures <- readMVar sess.failedSlots
       if Set.member (addr, slot) failures
       then do
         when (conf.debug) $ putStrLn $ "-> Skipping previously failed slot " <> show slot <> " at " <> show addr
-        pure (FetchFailure Cached)
-      else do
-        -- Attempt fetch
-        when (conf.debug) $ putStrLn $ "-> Fetching slot " <> show slot <> " at " <> show addr
-        ret <- fetchQuery n (fetchWithRetry conf.debug url sess) (QuerySlot addr slot)
-        case ret of
-          Right val -> do
-            -- Success: cache it
-            modifyMVar_ sess.sharedCache $ \c ->
-              pure $ c { slotCache = Map.insert (addr, slot) val c.slotCache }
-            pure (FetchSuccess val Fresh)
-          Left err -> do
-            pure (FetchError err)
+        pure (Just (FetchFailure Cached))
+      else
+        pure Nothing
+
+fetchFreshSlot
+  :: (BlockNumber -> Addr -> W256 -> IO (Either Text W256))
+  -> Config
+  -> Session
+  -> BlockNumber
+  -> Addr
+  -> W256
+  -> IO (FetchResult W256)
+fetchFreshSlot fetch conf sess n addr slot = do
+  cached <- lookupSlotCache conf sess addr slot
+  case cached of
+    Just res -> pure res
+    Nothing -> do
+      when (conf.debug) $ putStrLn $ "-> Fetching slot " <> show slot <> " at " <> show addr
+      ret <- fetch n addr slot
+      case ret of
+        Right val -> do
+          modifyMVar_ sess.sharedCache $ \c ->
+            pure $ c { slotCache = Map.insert (addr, slot) val c.slotCache }
+          pure (FetchSuccess val Fresh)
+        Left err ->
+          pure (FetchError err)
+
+cachedSuccess :: FetchResult W256 -> FetchResult W256
+cachedSuccess (FetchSuccess val _) = FetchSuccess val Cached
+cachedSuccess res = res
+
+claimInFlight
+  :: Ord key
+  => MVar (Map.Map key (MVar (Either SomeException a)))
+  -> key
+  -> IO (InFlightFetch a)
+claimInFlight inFlight fetchKey =
+  modifyMVar inFlight $ \entries ->
+    case Map.lookup fetchKey entries of
+      Just result ->
+        pure (entries, InFlightWaiter result)
+      Nothing -> do
+        result <- newEmptyMVar
+        pure (Map.insert fetchKey result entries, InFlightOwner result)
+
+withInFlight
+  :: Ord key
+  => MVar (Map.Map key (MVar (Either SomeException a)))
+  -> (a -> a)
+  -> key
+  -> IO a
+  -> IO a
+withInFlight inFlight waiterResult fetchKey action = mask $ \restore -> do
+  claimed <- claimInFlight inFlight fetchKey
+  case claimed of
+    InFlightWaiter result -> do
+      outcome <- restore $ readMVar result
+      either throwIO (pure . waiterResult) outcome
+    InFlightOwner result -> do
+      outcome <- try (restore action)
+      putMVar result outcome
+      modifyMVar_ inFlight $ \entries ->
+        pure $ Map.delete fetchKey entries
+      either throwIO pure outcome
 
 -- | Get the complete cache state including both successes and failures
 -- Returns in the format expected by Echidna's UI:
@@ -555,7 +642,8 @@ mkSession cacheDir mblock = do
   failedContracts <- liftIO $ newMVar Set.empty
   failedSlots <- liftIO $ newMVar Set.empty
   rpcThrottle <- liftIO $ newIORef Nothing
-  pure $ Session sess latestBlockNum cache cacheDir failedContracts failedSlots rpcThrottle
+  inFlightSlots <- liftIO $ newMVar Map.empty
+  pure $ Session sess latestBlockNum cache cacheDir failedContracts failedSlots rpcThrottle inFlightSlots
 
 mkSessionWithoutCache :: App m => m Session
 mkSessionWithoutCache = mkSession Nothing Nothing
@@ -626,22 +714,12 @@ oracle solvers preSess rpcInfo q = do
       | otherwise -> do
         let sess = fromMaybe (internalError $ "oracle: no session provided for fetch addr: " ++ show addr) preSess
         conf <- readConfig
-        cache <- liftIO $ readMVar sess.sharedCache
-        case Map.lookup (addr, slot) cache.slotCache of
-          Just s -> do
-            when (conf.debug) $ liftIO $ putStrLn $ "-> Using cached slot value for slot " <> show slot <> " at " <> show addr
-            pure $ continue s
-          Nothing -> do
-            when (conf.debug) $ liftIO $ putStrLn $ "Fetching slot " <> (show slot) <> " at " <> (show addr)
-            let (block, url) = fromJust rpcInfo.blockNumURL
-            n <- liftIO $ getLatestBlockNum conf sess block url
-            ret <- liftIO $ fetchQuery n (fetchWithRetry conf.debug url sess) (QuerySlot addr slot)
-            case ret of
-              Right val -> do
-                liftIO $ modifyMVar_ sess.sharedCache $ \c ->
-                  pure $ c { slotCache = Map.insert (addr, slot) val c.slotCache }
-                pure $ continue val
-              Left err -> internalError $ "oracle error: " ++ show err
+        let (block, url) = fromJust rpcInfo.blockNumURL
+        res <- liftIO $ fetchSlotWithCache conf sess block url addr slot
+        case res of
+          FetchSuccess val _ -> pure $ continue val
+          FetchFailure _ -> internalError $ "oracle error: " ++ show q
+          FetchError err -> internalError $ "oracle error: " ++ show err
 
     PleaseReadEnv variable continue -> do
       value <- liftIO $ lookupEnv variable

test/test.hs

@@ -8,6 +8,9 @@ module Main where
 import Prelude hiding (LT, GT)
 
 import GHC.TypeLits
+import Control.Concurrent (forkIO, threadDelay)
+import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, readMVar, tryPutMVar)
+import Control.Exception (SomeException, try)
 import Control.Monad
 import Control.Monad.ST (stToIO)
 import Control.Monad.State.Strict
@@ -21,6 +24,7 @@ import Data.ByteString.Lazy qualified as BSLazy
 import Data.Binary.Put (runPut)
 import Data.Binary.Get (runGetOrFail)
 import Data.Either
+import Data.IORef (atomicModifyIORef', newIORef, readIORef)
 import Data.List qualified as List
 import Data.Map.Strict qualified as Map
 import Data.Maybe
@@ -40,6 +44,7 @@ import Test.Tasty.HUnit
 import Test.Tasty.Runners hiding (Failure, Success)
 import Test.Tasty.ExpectedFailure
 import Text.ParserCombinators.ReadP (readP_to_S)
+import System.Timeout (timeout)
 import Witch (unsafeInto, into)
 
 import Optics.Core hiding (pre, re, elements)
@@ -123,6 +128,59 @@ withCVC5Solver = withSolvers CVC5 3 Nothing defMemLimit
 withBitwuzlaSolver :: App m => (SolverGroup -> m a) -> m a
 withBitwuzlaSolver = withSolvers Bitwuzla 3 Nothing defMemLimit
 
+waitForTest :: String -> IO a -> IO a
+waitForTest name action = do
+  timeout 5_000_000 action >>= \case
+    Just result -> pure result
+    Nothing -> assertFailure $ name <> " timed out"
+
+runConcurrentSlotFetches
+  :: Fetch.Session
+  -> (Fetch.BlockNumber -> Addr -> W256 -> IO (Either Text W256))
+  -> Int
+  -> Addr
+  -> W256
+  -> Maybe (MVar (), MVar ())
+  -> IO [Either SomeException (Fetch.FetchResult W256)]
+runConcurrentSlotFetches sess fetcher count addr slot gate = do
+  start <- newEmptyMVar
+  doneVars <- replicateM count newEmptyMVar
+  forM_ doneVars $ \done -> do
+    _ <- forkIO $ do
+      readMVar start
+      result <- try $ Fetch.fetchSlotWithCacheUsing fetcher defaultConfig sess (Fetch.BlockNumber 1) (T.pack "unused") addr slot
+      putMVar done result
+    pure ()
+  putMVar start ()
+  case gate of
+    Nothing -> pure ()
+    Just (started, release) -> do
+      waitForTest "owner fetch start" (readMVar started)
+      threadDelay 100_000
+      putMVar release ()
+  waitForTest "slot fetches" (mapM readMVar doneVars)
+
+runConcurrentSlotFetchesForSlots
+  :: Fetch.Session
+  -> (Fetch.BlockNumber -> Addr -> W256 -> IO (Either Text W256))
+  -> Addr
+  -> [W256]
+  -> MVar ()
+  -> IO [Either SomeException (Fetch.FetchResult W256)]
+runConcurrentSlotFetchesForSlots sess fetcher addr slots release = do
+  start <- newEmptyMVar
+  doneVars <- forM slots $ \slot -> do
+    done <- newEmptyMVar
+    _ <- forkIO $ do
+      readMVar start
+      result <- try $ Fetch.fetchSlotWithCacheUsing fetcher defaultConfig sess (Fetch.BlockNumber 1) (T.pack "unused") addr slot
+      putMVar done result
+    pure done
+  putMVar start ()
+  threadDelay 100_000
+  putMVar release ()
+  waitForTest "slot fetches" (mapM readMVar doneVars)
+
 
 main :: IO ()
 main = defaultMain tests
@@ -191,6 +249,67 @@ tests = testGroup "hevm"
 
         -- there won't be query now as accessStorage uses fetch cache
         assertBoolM (show vm4.result) (isNothing vm4.result)
+    , testCase "fetchSlotWithCache single-flights concurrent same slot" $ runEnv testEnv $ do
+        sess <- Fetch.mkSessionWithoutCache
+        counter <- liftIO $ newIORef (0 :: Int)
+        started <- liftIO newEmptyMVar
+        release <- liftIO newEmptyMVar
+        let fetcher _ _ _ = do
+              _ <- atomicModifyIORef' counter $ \n -> (n + 1, n + 1)
+              _ <- tryPutMVar started ()
+              readMVar release
+              pure (Right 0x1234)
+        results <- liftIO $ runConcurrentSlotFetches sess fetcher 8 0x1000 0x1 (Just (started, release))
+        assertEqualM "underlying fetch count" 1 =<< liftIO (readIORef counter)
+        let successes = rights results
+        assertEqualM "all calls should succeed" 8 (length successes)
+        assertEqualM "values" (replicate 8 0x1234) [val | Fetch.FetchSuccess val _ <- successes]
+        assertEqualM "fresh result count" 1 (length [() | Fetch.FetchSuccess _ Fetch.Fresh <- successes])
+        assertEqualM "cached result count" 7 (length [() | Fetch.FetchSuccess _ Fetch.Cached <- successes])
+    , testCase "fetchSlotWithCache does not merge different slots" $ runEnv testEnv $ do
+        sess <- Fetch.mkSessionWithoutCache
+        counter <- liftIO $ newIORef (0 :: Int)
+        release <- liftIO newEmptyMVar
+        let slots = [0x1, 0x2, 0x3, 0x4]
+            fetcher _ _ slot = do
+              _ <- atomicModifyIORef' counter $ \n -> (n + 1, n + 1)
+              readMVar release
+              pure (Right slot)
+        results <- liftIO $ runConcurrentSlotFetchesForSlots sess fetcher 0x1000 slots release
+        assertEqualM "underlying fetch count" (length slots) =<< liftIO (readIORef counter)
+        assertEqualM "values" slots [val | Fetch.FetchSuccess val _ <- rights results]
+    , testCase "fetchSlotWithCache owner error wakes waiters and clears in-flight entry" $ runEnv testEnv $ do
+        sess <- Fetch.mkSessionWithoutCache
+        counter <- liftIO $ newIORef (0 :: Int)
+        started <- liftIO newEmptyMVar
+        release <- liftIO newEmptyMVar
+        let fetcher _ _ _ = do
+              _ <- atomicModifyIORef' counter $ \n -> (n + 1, n + 1)
+              _ <- tryPutMVar started ()
+              readMVar release
+              pure (Left (T.pack "boom"))
+        results <- liftIO $ runConcurrentSlotFetches sess fetcher 8 0x1000 0x2 (Just (started, release))
+        assertEqualM "all calls should complete with the same fetch error"
+          (replicate 8 (Fetch.FetchError (T.pack "boom")))
+          (rights results)
+        retry <- liftIO $ Fetch.fetchSlotWithCacheUsing fetcher defaultConfig sess (Fetch.BlockNumber 1) (T.pack "unused") 0x1000 0x2
+        assertEqualM "retry should fetch again" (Fetch.FetchError (T.pack "boom")) retry
+        assertEqualM "underlying fetch count after retry" 2 =<< liftIO (readIORef counter)
+    , testCase "fetchSlotWithCache owner exception wakes waiters and clears in-flight entry" $ runEnv testEnv $ do
+        sess <- Fetch.mkSessionWithoutCache
+        counter <- liftIO $ newIORef (0 :: Int)
+        started <- liftIO newEmptyMVar
+        release <- liftIO newEmptyMVar
+        let fetcher _ _ _ = do
+              _ <- atomicModifyIORef' counter $ \n -> (n + 1, n + 1)
+              _ <- tryPutMVar started ()
+              readMVar release
+              ioError (userError "boom")
+        results <- liftIO $ runConcurrentSlotFetches sess fetcher 8 0x1000 0x3 (Just (started, release))
+        assertEqualM "all calls should complete with exceptions" 8 (length (lefts results))
+        retry <- liftIO $ try $ Fetch.fetchSlotWithCacheUsing fetcher defaultConfig sess (Fetch.BlockNumber 1) (T.pack "unused") 0x1000 0x3
+        assertBoolM "retry should throw again" (isLeft (retry :: Either SomeException (Fetch.FetchResult W256)))
+        assertEqualM "underlying fetch count after retry" 2 =<< liftIO (readIORef counter)
     ]
   , testGroup "ABI"
     [ testProperty "Put/get inverse" $ \x ->
@@ -3699,4 +3818,3 @@ expectedConcVals nm val = case val of
   _ -> internalError $ "unsupported Abi type " <> show nm <> " val: " <> show val <> " val type: " <> showAlter val
   where
     mkWord = word . encodeAbiValue
-