Created
May 24, 2025 05:12
-
-
Save rjmholt/a22bfb302a7a8452275f438741659614 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| module CacheVar where | |
| data CacheBehavior | |
| = SingleEvaluation | |
| | AllowInitializationRetry | |
| deriving stock (Show, Eq) | |
| data CacheState a | |
| = Proceed | |
| | Done a | |
| | Wait | |
| deriving stock (Show, Eq) | |
| data CacheVar a = CacheVar | |
| { cacheState :: TVar (CacheState a) | |
| , awaitSignal :: TMVar () | |
| , runGetValue :: IO a | |
| , cacheBehavior :: CacheBehavior | |
| } | |
| newCacheVar :: CacheBehavior -> IO a -> IO (CacheVar a) | |
| newCacheVar cacheBehavior runGetValue = do | |
| cacheState <- newTVarIO Proceed | |
| awaitSignal <- newTMVarIO () | |
| pure $ | |
| CacheVar | |
| { cacheState | |
| , awaitSignal | |
| , runGetValue | |
| , cacheBehavior | |
| } | |
| readCacheVar :: RequireCallStack => CacheVar a -> IO a | |
| readCacheVar cacheVar = | |
| case cacheVar.cacheBehavior of | |
| AllowInitializationRetry -> | |
| readCacheVarWithRetry cacheVar | |
| SingleEvaluation -> | |
| readCacheVarEvaluateOnce cacheVar | |
| readCacheVarWithRetry :: forall a. RequireCallStack => CacheVar a -> IO a | |
| readCacheVarWithRetry cacheVar = | |
| readCacheVarWithEvaluation cacheVar $ | |
| -- Try to run the action and populate the cache | |
| -- If we fail, roll back the cache | |
| ( do | |
| v <- cacheVar.runGetValue | |
| atomically $ populateCache cacheVar v | |
| pure v | |
| ) | |
| `catch` \(e :: SomeException) -> atomically (resetCache cacheVar) >> throwWithCallStack e | |
| readCacheVarEvaluateOnce :: forall a. RequireCallStack => CacheVar a -> IO a | |
| readCacheVarEvaluateOnce cacheVar = | |
| readCacheVarWithEvaluation cacheVar $ do | |
| eV <- unsafeTryAsync cacheVar.runGetValue | |
| let v = unwrap eV | |
| atomically $ populateCache cacheVar v | |
| pure v | |
| where | |
| unwrap :: Either SomeException a -> a | |
| unwrap = \case | |
| Left e -> | |
| unsafeThrowImpurely e | |
| Right v -> | |
| v | |
| readCacheVarWithEvaluation :: RequireCallStack => CacheVar a -> IO a -> IO a | |
| readCacheVarWithEvaluation cacheVar evaluateAndCache = do | |
| acquiredState <- atomically $ tryGetCachedValue cacheVar | |
| case acquiredState of | |
| Proceed -> | |
| evaluateAndCache | |
| Wait -> | |
| -- We've waited, now we need to try again | |
| readCacheVar cacheVar | |
| Done v -> | |
| -- The cache has a value, so just return it | |
| pure v | |
| tryGetCachedValue :: CacheVar a -> STM (CacheState a) | |
| tryGetCachedValue cacheVar = do | |
| s <- readTVar cacheVar.cacheState | |
| case s of | |
| Done _ -> | |
| pure () | |
| Wait -> | |
| awaitValueReady cacheVar | |
| Proceed -> do | |
| -- Tell other threads to wait | |
| writeTVar cacheVar.cacheState Wait | |
| -- Mark the wait signal | |
| takeEvaluationLock cacheVar | |
| pure s | |
| populateCache :: CacheVar a -> a -> STM () | |
| populateCache cacheVar v = do | |
| -- Put the value in the cache | |
| writeTVar cacheVar.cacheState $ Done v | |
| -- Wake the other threads to find it | |
| signalValueReady cacheVar | |
| resetCache :: CacheVar s -> STM () | |
| resetCache cacheVar = do | |
| -- Reinitialize the cache | |
| writeTVar cacheVar.cacheState Proceed | |
| -- Unblock the other threads so one can attempt to compute | |
| signalValueReady cacheVar | |
| takeEvaluationLock :: CacheVar a -> STM () | |
| takeEvaluationLock cacheVar = takeTMVar cacheVar.awaitSignal | |
| awaitValueReady :: CacheVar a -> STM () | |
| awaitValueReady cacheVar = readTMVar cacheVar.awaitSignal | |
| signalValueReady :: CacheVar a -> STM () | |
| signalValueReady cacheVar = void $ tryPutTMVar cacheVar.awaitSignal () |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment