resource-registry 0.1.1.0 → 0.2.0.0
raw patch · 3 files changed
+159/−31 lines, 3 filesdep ~basedep ~containersdep ~io-classesnew-uploader
Dependency ranges changed: base, containers, io-classes, nothunks
Files
- CHANGELOG.md +17/−1
- resource-registry.cabal +2/−2
- src/Control/ResourceRegistry.hs +140/−28
CHANGELOG.md view
@@ -1,4 +1,20 @@-# Revision history of strict-checked-vars+# Revision history of `resource-registry`++## 0.2.0.0 — 2025-10-23++* Define `transferRegistry` for moving all resources from one registry to a+ different one.++* Cancel registered threads in the registry before closing the registry. This+ prevents race conditions where the registry gets closed but the running async+ thread still tries to allocate something in the registry before being+ cancelled, which would result in an exception.++* Expose `allocateThread` to register another thread in a registry, such that it+ is cancelled before closing the registry, even if it belongs to a separate+ registry.++* Add the label of the thread to `Context`. ## 0.1.1.0 — 2025-05-15
resource-registry.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: resource-registry-version: 0.1.1.0+version: 0.2.0.0 synopsis: Track allocated resources description: When the scope of a @bracket@ doesn't enclose all uses of the resource, a@@ -35,7 +35,7 @@ type: git location: https://github.com/IntersectMBO/io-classes-extra subdir: resource-registry- tag: resource-registry-0.1.1.0+ tag: resource-registry-0.2.0.0 common warnings ghc-options:
src/Control/ResourceRegistry.hs view
@@ -239,6 +239,7 @@ -- * Allocating and releasing regular resources , ResourceKey , allocate+ , allocateThread , allocateEither , release , releaseAll@@ -263,11 +264,13 @@ , modifyWithTempRegistry , runInnerWithTempRegistry , runWithTempRegistry+ , transferRegistry -- * Unsafe combinators primarily for testing , closeRegistry , countResources , unsafeNewRegistry+ , resourceKeyId ) where import Control.Applicative ((<|>))@@ -342,6 +345,23 @@ data RegistryState m = RegistryState { registryThreads :: !(KnownThreads m) -- ^ Forked threads+ , registryReleaseThreads :: ![ReleaseThread m]+ -- ^ The list of releasing actions for threads that were forked from this+ -- registry.+ --+ -- We will cancel these when closing a registry but before actually closing+ -- the registry. This will guard the case in which a forked thread allocates+ -- resources in the registry. We would face a race condition among:+ --+ -- - The registry killing the thread+ --+ -- - The thread allocating something in the now closed registry.+ --+ -- The latter case will throw an exception.+ --+ -- Note this is separate from 'registryThreads' because one can add threads to+ -- cancel via 'allocateThread' and still those threads should not be directly+ -- added to the set of known threads. , registryResources :: !(Map ResourceId (Resource m)) -- ^ Currently allocated resources --@@ -433,6 +453,10 @@ instance Show (Release m) where show _ = "<<release>>" +-- | Release a thread when closing a registry.+newtype ReleaseThread m = ReleaseThread {releaseThread :: m ()}+ deriving NoThunks via OnlyCheckWhnfNamed "ReleaseThread" (ReleaseThread m)+ {------------------------------------------------------------------------------- Internal: pure functions on the registry state -------------------------------------------------------------------------------}@@ -455,11 +479,18 @@ -- | Allocate key for new resource allocKey :: State (RegistryState m) (Either PrettyCallStack ResourceId)-allocKey = unlessClosed $ do+allocKey = unlessClosed $ unsafeAllocKey++unsafeAllocKey :: State (RegistryState m) ResourceId+unsafeAllocKey = do nextKey <- gets registryNextKey modify $ \st -> st{registryNextKey = succ nextKey} return nextKey +-- | Allocate multiple keys for resources+allocNKeys :: Int -> State (RegistryState m) (Either PrettyCallStack [ResourceId])+allocNKeys n = unlessClosed $ replicateM n $ unsafeAllocKey+ -- | Insert new resource insertResource :: ResourceId ->@@ -584,18 +615,19 @@ { registryContext = context , registryState = stateVar }- where- initState :: RegistryState m- initState =- RegistryState- { registryThreads = KnownThreads Set.empty- , registryResources = Map.empty- , registryNextKey = ResourceId 1- , registryAges = Bimap.empty- , registryNextAge = ageOfFirstResource- , registryStatus = RegistryOpen- } +initState :: RegistryState m+initState =+ RegistryState+ { registryThreads = KnownThreads Set.empty+ , registryReleaseThreads = []+ , registryResources = Map.empty+ , registryNextKey = ResourceId 1+ , registryAges = Bimap.empty+ , registryNextAge = ageOfFirstResource+ , registryStatus = RegistryOpen+ }+ -- | Close the registry -- -- This can only be called from the same thread that created the registry.@@ -619,7 +651,17 @@ (MonadMask m, MonadThread m, MonadSTM m, HasCallStack) => ResourceRegistry m -> m ()-closeRegistry rr = mask_ $ do+closeRegistry rr = mask_ $ releaseAllBy close rr++-- | Release all the resources and perform another action while doing so. This+-- is to be used both by 'closeRegistry' which will 'close' the registry, as+-- well as 'releaseAll' which will not actually close the registry.+releaseAllBy ::+ (MonadMask m, MonadThread m, MonadSTM m, HasCallStack) =>+ (PrettyCallStack -> State (RegistryState m) (Either PrettyCallStack [ResourceId])) ->+ ResourceRegistry m ->+ m ()+releaseAllBy action rr = do context <- captureContext unless (contextThreadId context == contextThreadId (registryContext rr)) $ throwIO $@@ -627,9 +669,21 @@ { resourceRegistryCreatedIn = registryContext rr , resourceRegistryUsedIn = context }+ unsafeReleaseAllBy action context rr +-- | Unsafe version of 'releaseAllBy'.+unsafeReleaseAllBy ::+ (MonadMask m, MonadThread m, MonadSTM m, HasCallStack) =>+ (PrettyCallStack -> State (RegistryState m) (Either PrettyCallStack [ResourceId])) ->+ Context m ->+ ResourceRegistry m ->+ m ()+unsafeReleaseAllBy action context rr = do+ ts <- updateState rr $ gets registryReleaseThreads+ mapM_ releaseThread ts+ -- Close the registry so that we cannot allocate any further resources- alreadyClosed <- updateState rr $ close (contextCallStack context)+ alreadyClosed <- updateState rr $ action (contextCallStack context) case alreadyClosed of Left _ -> return ()@@ -1152,15 +1206,7 @@ (MonadMask m, MonadSTM m, MonadThread m, HasCallStack) => ResourceRegistry m -> m ()-releaseAll rr = do- context <- captureContext- unless (contextThreadId context == contextThreadId (registryContext rr)) $- throwIO $- ResourceRegistryClosedFromWrongThread- { resourceRegistryCreatedIn = registryContext rr- , resourceRegistryUsedIn = context- }- void $ releaseAllHelper rr context release+releaseAll rr = releaseAllBy (\_ -> unlessClosed $ gets getYoungestToOldest) rr -- | This is to 'releaseAll' what 'unsafeRelease' is to 'release': we do not -- insist that this funciton is called from a thread that is known to the@@ -1171,9 +1217,9 @@ m () unsafeReleaseAll rr = do context <- captureContext- void $ releaseAllHelper rr context unsafeRelease+ unsafeReleaseAllBy (\_ -> unlessClosed $ gets getYoungestToOldest) context rr --- | Internal helper used by 'releaseAll' and 'unsafeReleaseAll'.+-- | Internal helper used by 'runWithTempRegistry'. releaseAllHelper :: (MonadMask m, MonadSTM m, MonadThread m) => ResourceRegistry m ->@@ -1230,6 +1276,19 @@ waitAnyThread :: forall m a. MonadAsync m => [Thread m a] -> m a waitAnyThread ts = snd <$> waitAny (map threadAsync ts) +-- | Allocate a thread in a registry. This will ensure that such a thread is+-- cancelled before the registry is closed. Useful for threads that belong to a+-- different registry but will try to allocate resources in this registry.+allocateThread ::+ (MonadMask m, MonadAsync m, HasCallStack) =>+ ResourceRegistry m -> (ResourceId -> m (Thread m a)) -> m (ResourceKey m, Thread m a)+allocateThread rr alloc = do+ (k, t) <- allocate rr alloc cancelThread+ updateState rr $+ modify+ (\s -> s{registryReleaseThreads = ReleaseThread (void (release k)) : registryReleaseThreads s})+ pure (k, t)+ -- | Fork a new thread forkThread :: forall m a.@@ -1241,7 +1300,7 @@ m (Thread m a) forkThread rr label body = snd- <$> allocate rr (\key -> mkThread key <$> async (body' key)) cancelThread+ <$> allocateThread rr (\key -> mkThread key <$> async (body' key)) where mkThread :: ResourceId -> Async m a -> Thread m a mkThread rid child =@@ -1428,21 +1487,27 @@ -- ^ CallStack in which it was created , contextThreadId :: !(ThreadId m) -- ^ Thread that created the registry or resource+ , contextThreadLabel :: !(Maybe String)+ -- ^ The label of the thread that created the registry, if it is set } -- Existential type; we can't use generics instance NoThunks (Context m) where showTypeOf _ = "Context"- wNoThunks ctxt (Context cs tid) =+ wNoThunks ctxt (Context cs tid lbl) = allNoThunks [ noThunks ctxt cs , noThunks ctxt (InspectHeapNamed @"ThreadId" tid)+ , noThunks ctxt lbl ] deriving instance Show (Context m) captureContext :: MonadThread m => HasCallStack => m (Context m)-captureContext = Context prettyCallStack <$> myThreadId+captureContext = do+ tid <- myThreadId+ lbl <- threadLabel tid+ pure $ Context prettyCallStack tid lbl {------------------------------------------------------------------------------- Misc utilities@@ -1528,3 +1593,50 @@ NoThunks (Bimap k v) where wNoThunks ctxt = noThunksInKeysAndValues ctxt . Bimap.toList++-- | Move all the resources from the origin registry to the destination+-- registry and return the list of allocated 'ResourceKey's.+--+-- Transferring individual resources between registries is risky because a+-- later resource in the origin registry might depend on a resource that we are+-- trying to transfer. However, transferring whole registries is fine.+transferRegistry ::+ (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) =>+ ResourceRegistry m ->+ ResourceRegistry m ->+ m [ResourceKey m]+transferRegistry fromReg toReg = do+ context <- captureContext++ -- The calling thread is known to the origin registry+ ensureKnownThread fromReg context++ -- Alloc all the needed keys+ mKeys <- updateState toReg . allocNKeys =<< countResources fromReg++ case mKeys of+ -- If the destination registry is closed, throw+ Left closed -> throwRegistryClosed toReg context closed+ Right keys -> mask_ $ do+ -- Get the resources out of the origin registry and empty it+ regState <- atomically $ swapTVar (registryState fromReg) initState++ forM_+ ( zip keys $+ Map.elems (registryResources regState)+ )+ ( \(k, res) -> do+ -- Insert the resources into the new registry+ inserted <- updateState toReg (insertResource k res)++ case inserted of+ -- If the destination registry is closed, throw+ Left closed -> do+ let Release rel = resourceRelease res+ void rel+ throwRegistryClosed toReg context closed+ Right () ->+ pure ()+ )++ pure $ map (ResourceKey toReg) keys