diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Revision history of strict-checked-vars +## 0.1.1.0 — 2025-05-15 + +* Use `io-classes-1.8`. + ## 0.1.0.0 — 2024-10-22 - First release, extracted from [`ouroboros-consensus`](https://github.com/IntersectMBO/ouroboros-consensus). diff --git a/resource-registry.cabal b/resource-registry.cabal --- a/resource-registry.cabal +++ b/resource-registry.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: resource-registry -version: 0.1.0.0 +version: 0.1.1.0 synopsis: Track allocated resources description: When the scope of a @bracket@ doesn't enclose all uses of the resource, a @@ -21,7 +21,7 @@ category: Control build-type: Simple bug-reports: https://github.com/IntersectMBO/io-classes-extra/issues -tested-with: ghc ==8.10 || ==9.2 || ==9.4 || ==9.6 || ==9.8 || ==9.10 +tested-with: ghc ==9.6 || ==9.8 || ==9.10 || ==9.12 extra-doc-files: CHANGELOG.md README.md @@ -35,7 +35,7 @@ type: git location: https://github.com/IntersectMBO/io-classes-extra subdir: resource-registry - tag: resource-registry-0.1.0.0 + tag: resource-registry-0.1.1.0 common warnings ghc-options: @@ -54,13 +54,12 @@ import: warnings exposed-modules: Control.ResourceRegistry build-depends: - base >=4.14 && <4.21, + base >=4.14 && <4.22, bimap ^>=0.5, - containers >=0.6 && <0.8, - io-classes ^>=1.5, - mtl >=2.2 && <2.4, + containers >=0.6.7 && <0.8, + io-classes:{io-classes, strict-stm} ^>=1.8, + mtl ^>=2.3, nothunks ^>=0.2, - strict-stm ^>=1.5, hs-source-dirs: src default-language: Haskell2010 @@ -83,13 +82,10 @@ base, containers, generics-sop, - io-classes, + io-classes:{io-classes, si-timers, strict-mvar, strict-stm}, mtl, quickcheck-state-machine:no-vendored-treediff, resource-registry, - si-timers, - strict-mvar, - strict-stm, tasty, tasty-quickcheck, tree-diff, diff --git a/src/Control/ResourceRegistry.hs b/src/Control/ResourceRegistry.hs --- a/src/Control/ResourceRegistry.hs +++ b/src/Control/ResourceRegistry.hs @@ -12,1486 +12,1519 @@ {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE UndecidableInstances #-} - -{-# OPTIONS_GHC -Wno-orphans #-} - --- | Note on terminology: when thread A forks thread B, we will say that thread A --- is the \"parent\" and thread B is the \"child\". No further relationship --- between the two threads is implied by this terminology. In particular, note --- that the child may outlive the parent. We will use \"fork\" and \"spawn\" --- interchangeably. --- --- = Motivation --- --- Whenever we allocate resources, we must keep track of them so that we can --- deallocate them when they are no longer required. The most important tool we --- have to achieve this is 'bracket': --- --- > bracket allocateResource releaseResource $ \r -> --- > .. use r .. --- --- Often 'bracket' comes in the guise of a with-style combinator --- --- > withResource $ \r -> --- > .. use r .. --- --- Where this pattern is applicable, it should be used and there is no need to --- use the 'ResourceRegistry'. However, 'bracket' introduces strict lexical --- scoping: the resource is available inside the scope of the bracket, and --- will be deallocated once we leave that scope. That pattern is sometimes --- hard to use. --- --- For example, suppose we have this interface to an SQL server --- --- > query :: Query -> IO QueryHandle --- > close :: QueryHandle -> IO () --- > next :: QueryHandle -> IO Row --- --- and suppose furthermore that we are writing a simple webserver that allows a --- client to send multiple SQL queries, get rows from any open query, and close --- queries when no longer required: --- --- > server :: IO () --- > server = go Map.empty --- > where --- > go :: Map QueryId QueryHandle -> IO () --- > go handles = getRequest >>= \case --- > New q -> do --- > h <- query q -- allocate --- > qId <- generateQueryId --- > sendResponse qId --- > go $ Map.insert qId h handles --- > Close qId -> do --- > close (handles ! qId) -- release --- > go $ Map.delete qId handles --- > Next qId -> do --- > sendResponse =<< next (handles ! qId) --- > go handles --- --- The server opens and closes query handles in response to client requests. --- Restructuring this code to use 'bracket' would be awkward, but as it stands --- this code does not ensure that resources get deallocated; for example, if --- the server thread is killed ('killThread'), resources will be leaked. --- --- Another, perhaps simpler, example is spawning threads. Threads too should --- be considered to be resources that we should keep track of and deallocate --- when they are no longer required, primarily because when we deallocate --- (terminate) those threads they too will have a chance to deallocate /their/ --- resources. As for other resources, we have a with-style combinator for this --- --- > withAsync $ \thread -> .. --- --- Lexical scoping of threads is often inconvenient, however, more so than for --- regular resources. The temptation is therefore to simply fork a thread and --- forget about it, but if we are serious about resource deallocation this is --- not an acceptable solution. --- --- = The resource registry --- --- The resource registry is essentially a piece of state tracking which --- resources have been allocated. The registry itself is allocated with a --- with-style combinator 'withRegistry', and when we leave that scope any --- resources not yet deallocated will be released at that point. Typically --- the registry is only used as a fall-back, ensuring that resources will --- deallocated even in the presence of exceptions. For example, here's how --- we might rewrite the above server example using a registry: --- --- > server' :: IO () --- > server' = --- > withRegistry $ \registry -> go registry Map.empty --- > where --- > go :: ResourceRegistry IO --- > -> Map QueryId (ResourceKey, QueryHandle) --- > -> IO () --- > go registry handles = getRequest >>= \case --- > New q -> do --- > (key, h) <- allocate registry (query q) close -- allocate --- > qId <- generateQueryId --- > sendResponse qId --- > go registry $ Map.insert qId (key, h) handles --- > Close qId -> do --- > release registry (fst (handles ! qId)) -- release --- > go registry $ Map.delete qId handles --- > Next qId -> do --- > sendResponse =<< next (snd (handles ! qId)) --- > go registry handles --- --- We allocate the query with the help of the registry, providing the registry --- with the means to deallocate the query should that be required. We can /and --- should/ still manually release resources also: in this particular example, --- the (lexical) scope of the registry is the entire server thread, so delaying --- releasing queries until we exit that scope will probably mean we hold on to --- resources for too long. The registry is only there as a fall-back. --- --- = Spawning threads --- --- We already observed in the introduction that insisting on lexical scoping --- for threads is often inconvenient, and that simply using --- 'Control.Monad.Class.MonadFork.forkIO' is no solution as it means we might --- leak resources. There is however another problem with --- 'Control.Monad.Class.MonadFork.forkIO'. Consider this snippet: --- --- > withRegistry $ \registry -> --- > r <- allocate registry allocateResource releaseResource --- > forkIO $ .. use r .. --- --- It is easy to see that this code is problematic: we allocate a resource @r@, --- then spawn a thread that uses @r@, and finally leave the scope of --- 'withRegistry', thereby deallocating @r@ -- leaving the thread to run with --- a now deallocated resource. --- --- It is /only/ safe for threads to use a given registry, and/or its registered --- resources, if the lifetime of those threads is tied to the lifetime of the --- registry. There would be no problem with the example above if the thread --- would be terminated when we exit the scope of 'withRegistry'. --- --- The 'forkThread' combinator provided by the registry therefore does two --- things: it allocates the thread as a resource in the registry, so that it can --- kill the thread when releasing all resources in the registry. It also records --- the thread ID in a set of known threads. Whenever the registry is accessed --- from a thread /not/ in this set, the registry throws a runtime exception, --- since such a thread might outlive the registry and hence its contents. The --- intention is that this guards against dangerous patterns like the one above. --- --- = Linking --- --- When thread A spawns thread B using 'withAsync', the lifetime of B is tied --- to the lifetime of A: --- --- > withAsync .. $ \threadB -> .. --- --- After all, when A exits the scope of the 'withAsync', thread B will be --- killed. The reverse is however not true: thread B can terminate before --- thread A. It is often useful for thread A to be able to declare a dependency --- on thread B: if B somehow fails, that is, terminates with an exception, we --- want that exception to be rethrown in thread A as well. A can achieve this --- by /linking/ to B: --- --- > withAsync .. $ \threadB -> do --- > link threadB --- > .. --- --- Linking a parent to a child is however of limited value if the lifetime of --- the child is not limited by the lifetime of the parent. For example, if A --- does --- --- > threadB <- async $ .. --- > link threadB --- --- and A terminates before B does, any exception thrown by B might be send to a --- thread that no longer exists. This is particularly problematic when we start --- chaining threads: if A spawns-and-links-to B which spawns-and-links-to C, and --- C throws an exception, perhaps the intention is that this gets rethrown to B, --- and then rethrown to A, terminating all three threads; however, if B has --- terminated before the exception is thrown, C will throw the exception to a --- non-existent thread and A is never notified. --- --- For this reason, the registry's 'linkToRegistry' combinator does not link the --- specified thread to the thread calling 'linkToRegistry', but rather to the --- thread that created the registry. After all, the lifetime of threads spawned --- with 'forkThread' can certainly exceed the lifetime of their parent threads, --- but the lifetime of /all/ threads spawned using the registry will be limited --- by the scope of that registry, and hence the lifetime of the thread that --- created it. So, when we call 'linkToRegistry', the exception will be thrown --- the thread that created the registry, which (if not caught) will cause that --- that to exit the scope of 'withRegistry', thereby terminating all threads in --- that registry. --- --- = Combining the registry and with-style allocation --- --- It is perfectly possible (indeed, advisable) to use 'bracket' and --- bracket-like allocation functions alongside the registry, but note that the --- usual caveats with 'bracket' and forking threads still applies. In --- particular, spawning threads inside the 'bracket' that make use of the --- bracketed resource is problematic; this is of course true whether or not a --- registry is used. --- --- In principle this also includes 'withAsync'; however, since 'withAsync' --- results in a thread that is not known to the registry, such a thread will not --- be able to use the registry (the registry would throw an unknown thread --- exception, as described above). For this purpose we provide 'withThread'; --- 'withThread' (as opposed to 'forkThread') should be used when a parent thread --- wants to handle exceptions in the child thread; see 'withThread' for --- detailed discussion. --- --- It is /also/ fine to includes nested calls to 'withRegistry'. Since the --- lifetime of such a registry (and all resources within) is tied to the thread --- calling 'withRegistry', which itself is tied to the "parent registry" in --- which it was created, this creates a hierarchy of registries. It is of course --- essential for compositionality that we should be able to create local --- registries, but even if we do have easy access to a parent regisry, creating --- a local one where possibly is useful as it limits the scope of the resources --- created within, and hence their maximum lifetimes. - -module Control.ResourceRegistry ( - -- * The resource registry proper - Context - , ResourceId - , ResourceRegistry - -- * Exceptions - , RegistryClosedException (..) - , ResourceRegistryThreadException - -- * Creating and releasing the registry itself - , bracketWithPrivateRegistry - , registryThread - , withRegistry - -- * Allocating and releasing regular resources - , ResourceKey - , allocate - , allocateEither - , release - , releaseAll - , unsafeRelease - , unsafeReleaseAll - -- * Threads - , Thread - , cancelThread - , forkLinkedThread - , forkThread - , linkToRegistry - , threadId - , waitAnyThread - , waitThread - , withThread - -- * Temporary registry - , TempRegistryException (..) - , WithTempRegistry - , allocateTemp - , modifyWithTempRegistry - , runInnerWithTempRegistry - , runWithTempRegistry - -- * Unsafe combinators primarily for testing - , closeRegistry - , countResources - , unsafeNewRegistry - ) where - -import Control.Applicative ((<|>)) -import Control.Concurrent.Class.MonadSTM.Strict -import Control.Exception (asyncExceptionFromException) -import Control.Monad -import Control.Monad.Class.MonadAsync -import Control.Monad.Class.MonadFork -import Control.Monad.Class.MonadThrow -import Control.Monad.Reader -import Control.Monad.State.Strict -import Data.Bifunctor -import Data.Bimap (Bimap) -import Data.Bimap qualified as Bimap -import Data.Either (partitionEithers) -import Data.Map.Strict (Map) -import Data.Map.Strict qualified as Map -import Data.Maybe (catMaybes, listToMaybe) -import Data.Set (Set) -import Data.Set qualified as Set -import Data.Void -import Data.Word (Word64) -import GHC.Generics (Generic) -import GHC.Stack (CallStack, HasCallStack) -import GHC.Stack qualified as GHC -import NoThunks.Class hiding (Context) - --- | Tracks resources during their lifetime. -data ResourceRegistry m = ResourceRegistry { - -- | Context in which the registry was created - registryContext :: !(Context m) - - -- | Registry state - , registryState :: !(StrictTVar m (RegistryState m)) - } - deriving (Generic) - -deriving instance (forall a. NoThunks a => NoThunks (StrictTVar m a)) - => NoThunks (ResourceRegistry m) - -{------------------------------------------------------------------------------- - Internal: registry state --------------------------------------------------------------------------------} - --- | The age of a resource --- --- Age here is represented by an meaningless number. The one and only property --- that matters is that the age of resource A that was successfully allocated --- before resource B was (in the same registry) will be greater than the age of --- resource B. --- --- For the current implementation, that property will be true unless the --- registry lives long enough to have contained 2^64 separately allocated --- resources. --- --- This data is not exposed by the 'ResourceRegistry' interface. -newtype Age = Age Word64 - deriving stock (Show) - deriving newtype (Eq, Ord) - deriving NoThunks via InspectHeapNamed "Age" Age - --- | The age of the first resource successfully allocated in a fresh registry -ageOfFirstResource :: Age -ageOfFirstResource = Age maxBound - --- | Map the age of the latest resource to be successfully allocated to the age --- of the next resource to be successfully allocated in the same registry -nextYoungerAge :: Age -> Age -nextYoungerAge (Age n) = Age (n - 1) - --- | Internal registry state -data RegistryState m = RegistryState { - -- | Forked threads - registryThreads :: !(KnownThreads m) - - -- | Currently allocated resources - -- - -- INVARIANT: We record exactly the ages of currently allocated resources, - -- @'Bimap.keys' . 'registryAges' = 'Map.keys' . 'registryResources'@. - , registryResources :: !(Map ResourceId (Resource m)) - - -- | Next available resource key - , registryNextKey :: !ResourceId - - -- | The age of each currently allocated resource - -- - -- We use a 'Bimap' so we can maintain the keys in sorted order by age, - -- which is necessary when closing the registry. - , registryAges :: !(Bimap ResourceId Age) - - -- | The age of the next resource - , registryNextAge :: !Age - - -- | Does the registry still accept new allocations? - -- - -- See 'RegistryClosedException' for discussion. - , registryStatus :: !RegistryStatus - } - deriving (Generic, NoThunks) - --- | The currently allocated keys in youngest-to-oldest order -getYoungestToOldest :: RegistryState m -> [ResourceId] -getYoungestToOldest = map snd . Bimap.toAscListR . registryAges - --- | Threads known to the registry --- --- This is the set of threads spawned using 'forkThread'. The lifetimes of all --- of these threads are limited by the lifetime of the registry. --- --- Does not include the thread ID of the thread that created the registry. After --- all, this thread may well outlive the registry (though the registry cannot --- outlive it). --- --- Invariant (informal): the set of registered threads is a subset of the --- registered resources ('registryResources'). (This invariant is temporarily --- broken when we start a new thread in 'forkThread' but will be re-established --- before that thread starts execution proper.) -newtype KnownThreads m = KnownThreads (Set (ThreadId m)) - deriving NoThunks via InspectHeapNamed "KnownThreads" (KnownThreads m) - --- | Status of the registry (open or closed) -data RegistryStatus = - RegistryOpen - - -- | We record the 'CallStack' to the call to 'close - | RegistryClosed !PrettyCallStack - deriving (Generic, NoThunks) - --- | Resource key --- --- Resource keys are tied to a particular registry. -data ResourceKey m = ResourceKey !(ResourceRegistry m) !ResourceId - deriving Generic - -deriving instance NoThunks (ResourceRegistry m) - => NoThunks (ResourceKey m) - --- | Return the 'ResourceId' of a 'ResourceKey'. -resourceKeyId :: ResourceKey m -> ResourceId -resourceKeyId (ResourceKey _rr rid) = rid - --- | Resource ID --- --- This uniquifying data is not exposed by the 'ResourceRegistry' interface. -newtype ResourceId = ResourceId Int - deriving stock (Show, Eq, Ord) - deriving newtype (Enum, NoThunks) - --- | Information about a resource -data Resource m = Resource { - -- | Context in which the resource was created - resourceContext :: !(Context m) - - -- | Deallocate the resource - , resourceRelease :: !(Release m) - } - deriving (Generic, NoThunks) - --- | Release the resource, return 'True' when the resource was actually --- released, return 'False' when the resource was already released. --- --- If unsure, returning 'True' is always fine. -newtype Release m = Release (m Bool) - deriving NoThunks via OnlyCheckWhnfNamed "Release" (Release m) - -releaseResource :: Resource m -> m Bool -releaseResource Resource{resourceRelease = Release f} = f - -instance Show (Release m) where - show _ = "<>" - -{------------------------------------------------------------------------------- - Internal: pure functions on the registry state --------------------------------------------------------------------------------} - -modifyKnownThreads :: - (Set (ThreadId m) -> Set (ThreadId m)) - -> KnownThreads m - -> KnownThreads m -modifyKnownThreads f (KnownThreads ts) = KnownThreads (f ts) - --- | Auxiliary for functions that should be disallowed when registry is closed -unlessClosed :: - State (RegistryState m) a - -> State (RegistryState m) (Either PrettyCallStack a) -unlessClosed f = do - status <- gets registryStatus - case status of - RegistryClosed closed -> return $ Left closed - RegistryOpen -> Right <$> f - --- | Allocate key for new resource -allocKey :: State (RegistryState m) (Either PrettyCallStack ResourceId) -allocKey = unlessClosed $ do - nextKey <- gets registryNextKey - modify $ \st -> st {registryNextKey = succ nextKey} - return nextKey - --- | Insert new resource -insertResource :: - ResourceId - -> Resource m - -> State (RegistryState m) (Either PrettyCallStack ()) -insertResource key r = unlessClosed $ do - modify $ \st -> st { - registryResources = Map.insert key r (registryResources st) - , registryAges = Bimap.insert - key - (registryNextAge st) - (registryAges st) - , registryNextAge = nextYoungerAge (registryNextAge st) - } - --- | Remove resource from the registry (if it exists) -removeResource :: ResourceId -> State (RegistryState m) (Maybe (Resource m)) -removeResource key = state $ \st -> - let (mbResource, resources') = Map.updateLookupWithKey - (\_ _ -> Nothing) - key - (registryResources st) - - st' = st { - registryResources = resources' - , registryAges = Bimap.delete key (registryAges st) - } - in (mbResource, st') - --- | Insert thread into the set of known threads -insertThread :: MonadThread m => ThreadId m -> State (RegistryState m) () -insertThread tid = - modify $ \st -> st { - registryThreads = modifyKnownThreads (Set.insert tid) $ - registryThreads st - } - --- | Remove thread from set of known threads -removeThread :: MonadThread m => ThreadId m -> State (RegistryState m) () -removeThread tid = - modify $ \st -> st { - registryThreads = modifyKnownThreads (Set.delete tid) $ - registryThreads st - } - --- | Close the registry --- --- Returns the keys currently allocated if the registry is not already closed. --- --- POSTCONDITION: They are returned in youngest-to-oldest order. -close :: - PrettyCallStack - -> State (RegistryState m) (Either PrettyCallStack [ResourceId]) -close closeCallStack = unlessClosed $ do - modify $ \st -> st {registryStatus = RegistryClosed closeCallStack} - gets getYoungestToOldest - --- | Convenience function for updating the registry state -updateState :: - forall m a. - MonadSTM m - => ResourceRegistry m - -> State (RegistryState m) a - -> m a -updateState rr f = - atomically $ stateTVar (registryState rr) (runState f) - --- | Attempt to allocate a resource in a registry which is closed --- --- When calling 'closeRegistry' (typically, leaving the scope of --- 'withRegistry'), all resources in the registry must be released. If a --- concurrent thread is still allocating resources, we end up with a race --- between the thread trying to allocate new resources and the registry trying --- to free them all. To avoid this, before releasing anything, the registry will --- record itself as closed. Any attempt by a concurrent thread to allocate a new --- resource will then result in a 'RegistryClosedException'. --- --- It is probably not particularly useful for threads to try and catch this --- exception (apart from in a generic handler that does local resource cleanup). --- The thread will anyway soon receive a 'Control.Exception.ThreadKilled' --- exception. -data RegistryClosedException = - forall m. MonadThread m => RegistryClosedException { - -- | The context in which the registry was created - registryClosedRegistryContext :: !(Context m) - - -- | Callstack to the call to 'closeRegistry' - -- - -- Note that 'closeRegistry' can only be called from the same thread - -- that created the registry. - , registryClosedCloseCallStack :: !PrettyCallStack - - -- | Context of the call resulting in the exception - , registryClosedAllocContext :: !(Context m) - } - -deriving instance Show RegistryClosedException -instance Exception RegistryClosedException - -{------------------------------------------------------------------------------- - Creating and releasing the registry itself --------------------------------------------------------------------------------} - --- | Create a new registry --- --- You are strongly encouraged to use 'withRegistry' instead. --- Exported primarily for the benefit of tests. -unsafeNewRegistry :: - (MonadSTM m, MonadThread m, HasCallStack) - => m (ResourceRegistry m) -unsafeNewRegistry = do - context <- captureContext - stateVar <- newTVarIO initState - return ResourceRegistry { - 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 - } - --- | Close the registry --- --- This can only be called from the same thread that created the registry. --- This is a no-op if the registry is already closed. --- --- This entire function runs with exceptions masked, so that we are not --- interrupted while we release all resources. --- --- Resources will be allocated from young to old, so that resources allocated --- later can safely refer to resources created earlier. --- --- The release functions are run in the scope of an exception handler, so that --- if releasing one resource throws an exception, we still attempt to release --- the other resources. Should we catch an exception whilst we close the --- registry, we will rethrow it after having attempted to release all resources. --- If there is more than one, we will pick a random one to rethrow, though we --- will prioritize asynchronous exceptions over other exceptions. This may be --- important for exception handlers that catch all-except-asynchronous --- exceptions. -closeRegistry :: - (MonadMask m, MonadThread m, MonadSTM m, HasCallStack) - => ResourceRegistry m - -> m () -closeRegistry rr = mask_ $ do - context <- captureContext - unless (contextThreadId context == contextThreadId (registryContext rr)) $ - throwIO $ ResourceRegistryClosedFromWrongThread { - resourceRegistryCreatedIn = registryContext rr - , resourceRegistryUsedIn = context - } - - -- Close the registry so that we cannot allocate any further resources - alreadyClosed <- updateState rr $ close (contextCallStack context) - case alreadyClosed of - Left _ -> - return () - Right keys -> do - -- At this point we have not /removed/ any elements from the map, - -- allowing concurrent threads to do their own cleanup of resources - -- (this may for instance be important if a thread deallocates its - -- resources in a particular order -- note that cancelling a thread - -- is a synchronous operation, so we will wait for it to finish - -- releasing its resources.) - -- /If/ a concurrent thread does some cleanup, then some of the calls - -- to 'release' that we do here might be no-ops. - void $ releaseResources rr keys release - --- | Helper for 'closeRegistry', 'releaseAll', and 'unsafeReleaseAll': release --- the resources allocated with the given 'ResourceId's. --- --- Returns the contexts of the resources that were actually released. -releaseResources :: - MonadCatch m - => ResourceRegistry m - -> [ResourceId] - -- ^ PRECONDITION: The currently allocated keys, - -- youngest-to-oldest - -> (ResourceKey m -> m (Maybe (Context m))) - -- ^ How to release the resource, e.g., 'release' or - -- 'unsafeRelease'. - -> m [Context m] -releaseResources rr sortedKeys releaser = do - (exs, mbContexts) <- fmap partitionEithers $ - forM sortedKeys $ try . releaser . ResourceKey rr - - case prioritize exs of - Nothing -> return (catMaybes mbContexts) - Just e -> throwIO e - where - prioritize :: [SomeException] -> Maybe SomeException - prioritize = - (\(asyncEx, otherEx) -> listToMaybe asyncEx <|> listToMaybe otherEx) - . first catMaybes - . unzip - . map (\e -> (asyncExceptionFromException e, e)) - --- | Create a new registry --- --- See documentation of 'ResourceRegistry' for a detailed discussion. -withRegistry :: - (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) - => (ResourceRegistry m -> m a) - -> m a -withRegistry = bracket unsafeNewRegistry closeRegistry - --- | Create a new private registry for use by a bracketed resource --- --- Use this combinator as a more specific and easier-to-maintain alternative to --- the following. --- --- > 'withRegistry' $ \rr -> --- > 'bracket' (newFoo rr) closeFoo $ \foo -> --- > (... rr does not occur in this scope ...) --- --- NB The scoped body can use `withRegistry` if it also needs its own, separate --- registry. --- --- Use this combinator to emphasize that the registry is private to (ie only --- used by and/or via) the bracketed resource and that it thus has nearly the --- same lifetime. This combinator ensures the following specific invariants --- regarding lifetimes and order of releases. --- --- o The registry itself is older than the bracketed resource. --- --- o The only registered resources older than the bracketed resource were --- allocated in the registry by the function that allocated the bracketed --- resource. --- --- o Because of the older resources, the bracketed resource is itself also --- registered in the registry; that's the only way we can be sure to release --- all resources in the right order. --- --- NB Because the registry is private to the resource, the @a@ type could save --- the handle to @registry@ and safely close the registry if the scoped body --- calls @closeA@ before the bracket ends. Though we have not used the type --- system to guarantee that the interface of the @a@ type cannot leak the --- registry to the body, this combinator does its part to keep the registry --- private to the bracketed resource. --- --- See documentation of 'ResourceRegistry' for a more general discussion. -bracketWithPrivateRegistry :: - (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) - => (ResourceRegistry m -> m a) - -> (a -> m ()) -- ^ Release the resource - -> (a -> m r) - -> m r -bracketWithPrivateRegistry newA closeA body = - withRegistry $ \registry -> do - (_key, a) <- allocate registry (\_key -> newA registry) closeA - body a - -{------------------------------------------------------------------------------- - Temporary registry --------------------------------------------------------------------------------} - --- | Run an action with a temporary resource registry. --- --- When allocating resources that are meant to end up in some final state, --- e.g., stored in a 'Control.Monad.Class.MonadSTM.TVar', after which they are --- guaranteed to be released correctly, it is possible that an exception is --- thrown after allocating such a resource, but before it was stored in the --- final state. In that case, the resource would be leaked. --- 'runWithTempRegistry' solves that problem. --- --- When no exception is thrown before the end of 'runWithTempRegistry', the --- user must have transferred all the resources it allocated to their final --- state. This means that these resources don't have to be released by the --- temporary registry anymore, the final state is now in charge of releasing --- them. --- --- In case an exception is thrown before the end of 'runWithTempRegistry', --- /all/ resources allocated in the temporary registry will be released. --- --- Resources must be allocated using 'allocateTemp'. --- --- To make sure that the user doesn't forget to transfer a resource to the --- final state @st@, the user must pass a function to 'allocateTemp' that --- checks whether a given @st@ contains the resource, i.e., whether the --- resource was successfully transferred to its final destination. --- --- When no exception is thrown before the end of 'runWithTempRegistry', we --- check whether all allocated resources have been transferred to the final --- state @st@. If there's a resource that hasn't been transferred to the final --- state /and/ that hasn't be released or closed before (see the release --- function passed to 'allocateTemp'), a 'TempRegistryRemainingResource' --- exception will be thrown. --- --- For that reason, 'WithTempRegistry' is parameterised over the final state --- type @st@ and the given 'WithTempRegistry' action must return the final --- state. --- --- NOTE: we explicitly don't let 'runWithTempRegistry' return the final state, --- because the state /must/ have been stored somewhere safely, transferring --- the resources, before the temporary registry is closed. -runWithTempRegistry :: - (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) - => WithTempRegistry st m (a, st) - -> m a -runWithTempRegistry m = withRegistry $ \rr -> do - varTransferredTo <- newTVarIO mempty - let tempRegistry = TempRegistry { - tempResourceRegistry = rr - , tempTransferredTo = varTransferredTo - } - (a, st) <- runReaderT (unWithTempRegistry m) tempRegistry - -- We won't reach this point if an exception is thrown, so we won't check - -- for remaining resources in that case. - -- - -- No need to mask here, whether we throw the async exception or - -- 'TempRegistryRemainingResource' doesn't matter. - transferredTo <- readTVarIO varTransferredTo - untrackTransferredTo rr transferredTo st - - context <- captureContext - remainingResources <- releaseAllHelper rr context release - - whenJust (listToMaybe remainingResources) $ \remainingResource -> - throwIO $ TempRegistryRemainingResource { - tempRegistryContext = registryContext rr - , tempRegistryResource = remainingResource - } - return a - - where - whenJust (Just x) f = f x - whenJust Nothing _ = pure () - --- | Embed a self-contained 'WithTempRegistry' computation into a larger one. --- --- The internal 'WithTempRegistry' is effectively passed to --- 'runWithTempRegistry'. It therefore must have no dangling resources, for --- example. This is the meaning of /self-contained/ above. --- --- The key difference beyond 'runWithTempRegistry' is that the resulting --- composite resource is also guaranteed to be registered in the outer --- 'WithTempRegistry' computation's registry once the inner registry is closed. --- Combined with the following assumption, this establishes the invariant that --- all resources are (transitively) in a temporary registry. --- --- As the resource might require some implementation details to be closed, the --- function to close it will also be provided by the inner computation. --- --- ASSUMPTION: closing @res@ closes every resource contained in @innerSt@ --- --- NOTE: In the current implementation, there will be a brief moment where the --- inner registry still contains the inner computation's resources and also the --- outer registry simultaneously contains the new composite resource. If an --- async exception is received at that time, then the inner resources will be --- closed and then the composite resource will be closed. This means there's a --- risk of /double freeing/, which can be harmless if anticipated. -runInnerWithTempRegistry :: - forall innerSt st m res a. - (MonadSTM m, MonadMask m, MonadThread m) - => WithTempRegistry innerSt m (a, innerSt, res) - -- ^ The embedded computation; see ASSUMPTION above - -> (res -> m Bool) - -- ^ How to free; same as for 'allocateTemp' - -> (st -> res -> Bool) - -- ^ How to check; same as for 'allocateTemp' - -> WithTempRegistry st m a -runInnerWithTempRegistry inner free isTransferred = do - outerTR <- WithTempRegistry ask - - lift $ runWithTempRegistry $ do - (a, innerSt, res) <- inner - - -- Allocate in the outer layer. - _ <- withFixedTempRegistry outerTR - $ allocateTemp (return res) free isTransferred - - -- TODO This point here is where an async exception could cause both the - -- inner resources to be closed and the outer resource to be closed later. - -- - -- If we want to do better than that, we'll need a variant of - -- 'runWithTempRegistry' that lets us perform some action with async - -- exceptions masked "at the same time" it closes its registry. - - -- Note that everything in `inner` allocated via `allocateTemp` must - -- either be closed or else present in `innerSt` by this point -- - -- `runWithTempRegistry` would have thrown if not. - pure (a, innerSt) - where - withFixedTempRegistry :: - TempRegistry st m - -> WithTempRegistry st m res - -> WithTempRegistry innerSt m res - withFixedTempRegistry env (WithTempRegistry (ReaderT f)) = - WithTempRegistry $ ReaderT $ \_ -> f env - --- | When 'runWithTempRegistry' exits successfully while there are still --- resources remaining in the temporary registry that haven't been transferred --- to the final state. -data TempRegistryException = - forall m. MonadThread m => TempRegistryRemainingResource { - -- | The context in which the temporary registry was created. - tempRegistryContext :: !(Context m) - - -- | The context in which the resource was allocated that was not - -- transferred to the final state. - , tempRegistryResource :: !(Context m) - } - -deriving instance Show TempRegistryException -instance Exception TempRegistryException - --- | Given a final state, return the 'ResourceId's of the resources that have --- been /transferred to/ that state. -newtype TransferredTo st = TransferredTo { - runTransferredTo :: st -> Set ResourceId - } - deriving newtype (Semigroup, Monoid) - deriving NoThunks via OnlyCheckWhnfNamed "TransferredTo" (TransferredTo st) - --- | The environment used to run a 'WithTempRegistry' action. -data TempRegistry st m = TempRegistry { - tempResourceRegistry :: !(ResourceRegistry m) - , tempTransferredTo :: !(StrictTVar m (TransferredTo st)) - -- ^ Used as a @Writer@. - } - --- | An action with a temporary registry in scope, see 'runWithTempRegistry' --- for more details. --- --- The most important function to run in this monad is 'allocateTemp'. -newtype WithTempRegistry st m a = WithTempRegistry { - unWithTempRegistry :: ReaderT (TempRegistry st m) m a - } - deriving newtype ( Functor - , Applicative - , Monad - , MonadThrow - , MonadCatch - , MonadMask - ) - -instance MonadTrans (WithTempRegistry st) where - lift = WithTempRegistry . lift - -instance MonadState s m => MonadState s (WithTempRegistry st m) where - state = WithTempRegistry . state - --- | Untrack all the resources from the registry that have been transferred to --- the given state. --- --- Untracking a resource means removing it from the registry without releasing --- it. --- --- NOTE: does not check that it's called by the same thread that allocated the --- resources, as it's an internal function only used in 'runWithTempRegistry'. -untrackTransferredTo :: - MonadSTM m - => ResourceRegistry m - -> TransferredTo st - -> st - -> m () -untrackTransferredTo rr transferredTo st = - updateState rr $ mapM_ removeResource rids - where - rids = runTransferredTo transferredTo st - --- | Allocate a resource in a temporary registry until it has been transferred --- to the final state @st@. See 'runWithTempRegistry' for more details. -allocateTemp :: - (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) - => m a - -- ^ Allocate the resource - -> (a -> m Bool) - -- ^ Release the resource, return 'True' when the resource was actually - -- released, return 'False' when the resource was already released. - -- - -- Note that it is safe to always return 'True' when unsure. - -> (st -> a -> Bool) - -- ^ Check whether the resource is in the given state - -> WithTempRegistry st m a -allocateTemp alloc free isTransferred = WithTempRegistry $ do - TempRegistry rr varTransferredTo <- ask - (key, a) <- lift (mustBeRight <$> - allocateEither rr (fmap Right . const alloc) free) - lift $ atomically $ modifyTVar varTransferredTo $ mappend $ - TransferredTo $ \st -> - if isTransferred st a - then Set.singleton (resourceKeyId key) - else Set.empty - return a - --- | Higher level API on top of 'runWithTempRegistry': modify the given @st@, --- allocating resources in the process that will be transferred to the --- returned @st@. -modifyWithTempRegistry :: - forall m st a. - (MonadSTM m, MonadMask m, MonadThread m) - => m st -- ^ Get the state - -> (st -> ExitCase st -> m ()) -- ^ Store the new state - -> StateT st (WithTempRegistry st m) a -- ^ Modify the state - -> m a -modifyWithTempRegistry getSt putSt modSt = runWithTempRegistry $ - fst <$> generalBracket (lift getSt) transfer mutate - where - transfer :: st -> ExitCase (a, st) -> WithTempRegistry st m () - transfer initSt ec = lift $ putSt initSt (snd <$> ec) - - mutate :: st -> WithTempRegistry st m (a, st) - mutate = runStateT modSt - -{------------------------------------------------------------------------------- - Simple queries on the registry --------------------------------------------------------------------------------} - --- | The thread that created the registry -registryThread :: ResourceRegistry m -> ThreadId m -registryThread = contextThreadId . registryContext - --- | Number of currently allocated resources --- --- Primarily for the benefit of testing. -countResources :: MonadSTM m => ResourceRegistry m -> m Int -countResources rr = atomically $ aux <$> readTVar (registryState rr) - where - aux :: RegistryState m -> Int - aux = Map.size . registryResources - -{------------------------------------------------------------------------------- - Allocating resources --------------------------------------------------------------------------------} - --- | Allocate new resource --- --- The allocation function will be run with asynchronous exceptions masked. This --- means that the resource allocation must either be fast or else interruptible; --- see "Dealing with Asynchronous Exceptions during Resource Acquisition" --- for details. -allocate :: - forall m a. - (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) - => ResourceRegistry m - -> (ResourceId -> m a) - -> (a -> m ()) -- ^ Release the resource - -> m (ResourceKey m, a) -allocate rr alloc free = mustBeRight <$> - allocateEither rr (fmap Right . alloc) (\a -> free a >> return True) - --- | Generalization of 'allocate' for allocation functions that may fail -allocateEither :: - forall m e a. - (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) - => ResourceRegistry m - -> (ResourceId -> m (Either e a)) - -> (a -> m Bool) - -- ^ Release the resource, return 'True' when the resource - -- hasn't been released or closed before. - -> m (Either e (ResourceKey m, a)) -allocateEither rr alloc free = do - context <- captureContext - ensureKnownThread rr context - -- We check if the registry has been closed when we allocate the key, so - -- that we can avoid needlessly allocating the resource. - mKey <- updateState rr allocKey - case mKey of - Left closed -> - throwRegistryClosed rr context closed - Right key -> mask_ $ do - ma <- alloc key - case ma of - Left e -> return $ Left e - Right a -> do - -- TODO: Might want to have an exception handler around this call to - -- 'updateState' just in case /that/ throws an exception. - inserted <- updateState rr $ - insertResource key (mkResource context a) - case inserted of - Left closed -> do - -- Despite the earlier check, it's possible that the registry - -- got closed after we allocated a new key but before we got a - -- chance to register the resource. In this case, we must - -- deallocate the resource again before throwing the exception. - void $ free a - throwRegistryClosed rr context closed - Right () -> - return $ Right (ResourceKey rr key, a) - where - mkResource :: Context m -> a -> Resource m - mkResource context a = Resource { - resourceContext = context - , resourceRelease = Release $ free a - } - -throwRegistryClosed :: - (MonadThrow m, MonadThread m) - => ResourceRegistry m - -> Context m - -> PrettyCallStack - -> m x -throwRegistryClosed rr context closed = throwIO RegistryClosedException { - registryClosedRegistryContext = registryContext rr - , registryClosedCloseCallStack = closed - , registryClosedAllocContext = context - } - --- | Release resource --- --- This deallocates the resource and removes it from the registry. It will be --- the responsibility of the caller to make sure that the resource is no longer --- used in any thread. --- --- The deallocation function is run with exceptions masked, so that we are --- guaranteed not to remove the resource from the registry without releasing it. --- --- Releasing an already released resource is a no-op. --- --- When the resource has not been released before, its context is returned. -release :: - (MonadMask m, MonadSTM m, MonadThread m, HasCallStack) - => ResourceKey m - -> m (Maybe (Context m)) -release key@(ResourceKey rr _) = do - context <- captureContext - ensureKnownThread rr context - unsafeRelease key - --- | Unsafe version of 'release' --- --- The only difference between 'release' and 'unsafeRelease' is that the latter --- does not insist that it is called from a thread that is known to the --- registry. This is dangerous, because it implies that there is a thread with --- access to a resource which may be deallocated before that thread is --- terminated. Of course, we can't detect all such situations (when the thread --- merely uses a resource but does not allocate or release we can't tell), but --- normally when we /do/ detect this we throw an exception. --- --- This function should only be used if the above situation can be ruled out --- or handled by other means. -unsafeRelease :: - (MonadMask m, MonadSTM m) - => ResourceKey m - -> m (Maybe (Context m)) -unsafeRelease (ResourceKey rr rid) = do - mask_ $ do - mResource <- updateState rr $ removeResource rid - case mResource of - Nothing -> return Nothing - Just resource -> do - actuallyReleased <- releaseResource resource - return $ - if actuallyReleased - then Just (resourceContext resource) - else Nothing - --- | Release all resources in the 'ResourceRegistry' without closing. --- --- See 'closeRegistry' for more details. -releaseAll :: - (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 - --- | 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 --- registry. See 'unsafeRelease' for why this is dangerous. -unsafeReleaseAll :: - (MonadMask m, MonadSTM m, MonadThread m, HasCallStack) - => ResourceRegistry m - -> m () -unsafeReleaseAll rr = do - context <- captureContext - void $ releaseAllHelper rr context unsafeRelease - --- | Internal helper used by 'releaseAll' and 'unsafeReleaseAll'. -releaseAllHelper :: - (MonadMask m, MonadSTM m, MonadThread m) - => ResourceRegistry m - -> Context m - -> (ResourceKey m -> m (Maybe (Context m))) - -- ^ How to release a resource - -> m [Context m] -releaseAllHelper rr context releaser = mask_ $ do - mKeys <- updateState rr $ unlessClosed $ gets getYoungestToOldest - case mKeys of - Left closed -> throwRegistryClosed rr context closed - Right keys -> releaseResources rr keys releaser - -{------------------------------------------------------------------------------- - Threads --------------------------------------------------------------------------------} - --- | Thread --- --- The internals of this type are not exported. -data Thread m a = MonadThread m => Thread { - -- | The underlying @async@ thread id - threadId :: !(ThreadId m) - , threadResourceId :: !ResourceId - , threadAsync :: !(Async m a) - , threadRegistry :: !(ResourceRegistry m) - } - deriving NoThunks via OnlyCheckWhnfNamed "Thread" (Thread m a) - --- | 'Eq' instance for 'Thread' compares 'threadId' only. -instance MonadThread m => Eq (Thread m a) where - Thread{threadId = a} == Thread{threadId = b} = a == b - --- | Cancel a thread --- --- This is a synchronous operation: the thread will have terminated when this --- function returns. --- --- Uses 'uninterruptibleCancel' because that's what 'withAsync' does. -cancelThread :: MonadAsync m => Thread m a -> m () -cancelThread = uninterruptibleCancel . threadAsync - --- | Wait for thread to terminate and return its result. --- --- If the thread throws an exception, this will rethrow that exception. --- --- NOTE: If A waits on B, and B is linked to the registry, and B throws an --- exception, then A might /either/ receive the exception thrown by B /or/ --- the 'Control.Exception.ThreadKilled' exception thrown by the registry. -waitThread :: MonadAsync m => Thread m a -> m a -waitThread = wait . threadAsync - --- | Lift 'waitAny' to 'Thread' -waitAnyThread :: forall m a. MonadAsync m => [Thread m a] -> m a -waitAnyThread ts = snd <$> waitAny (map threadAsync ts) - --- | Fork a new thread -forkThread :: - forall m a. - (MonadMask m, MonadAsync m, HasCallStack) - => ResourceRegistry m - -> String -- ^ Label for the thread - -> m a - -> m (Thread m a) -forkThread rr label body = snd <$> - allocate rr (\key -> mkThread key <$> async (body' key)) cancelThread - where - mkThread :: ResourceId -> Async m a -> Thread m a - mkThread rid child = Thread { - threadId = asyncThreadId child - , threadResourceId = rid - , threadAsync = child - , threadRegistry = rr - } - - body' :: ResourceId -> m a - body' rid = do - me <- myThreadId - labelThread me label - (registerThread me >> body) `finally` unregisterThread me rid - - -- Register the thread - -- - -- We must add the thread to the list of known threads before the thread - -- will start to use the registry. - registerThread :: ThreadId m -> m () - registerThread tid = updateState rr $ insertThread tid - - -- Unregister the thread - -- - -- Threads are the only kinds of resources that "deallocate themselves". - -- We remove the thread from the resources as well as the set of known - -- threads, so that these datastructures do not grow without bound. - -- - -- This runs with asynchronous exceptions masked (due to 'finally'), - -- though for the current implementation of 'unregisterThread' this - -- makes no difference. - unregisterThread :: ThreadId m -> ResourceId -> m () - unregisterThread tid rid = - updateState rr $ do - removeThread tid - void $ removeResource rid - --- | Bracketed version of 'forkThread' --- --- The analogue of 'withAsync' for the registry. --- --- Scoping thread lifetime using 'withThread' is important when a parent --- thread wants to link to a child thread /and handle any exceptions arising --- from the link/: --- --- > let handleLinkException :: ExceptionInLinkedThread -> m () --- > handleLinkException = .. --- > in handle handleLinkException $ --- > withThread registry codeInChild $ \child -> --- > .. --- --- instead of --- --- > handle handleLinkException $ do -- PROBABLY NOT CORRECT! --- > child <- forkThread registry codeInChild --- > .. --- --- where the parent may exit the scope of the exception handler before the child --- terminates. If the lifetime of the child cannot be limited to the lifetime of --- the parent, the child should probably be linked to the registry instead and --- the thread that spawned the registry should handle any exceptions. --- --- Note that in /principle/ there is no problem in using 'withAsync' alongside a --- registry. After all, in a pattern like --- --- > withRegistry $ \registry -> --- > .. --- > withAsync (.. registry ..) $ \async -> --- > .. --- --- the async will be cancelled when leaving the scope of 'withAsync' and so --- that reference to the registry, or indeed any of the resources inside the --- registry, is safe. However, the registry implements a sanity check that the --- registry is only used from known threads. This is useful: when a thread that --- is not known to the registry (in other words, whose lifetime is not tied to --- the lifetime of the registry) spawns a resource in that registry, that --- resource may well be deallocated before the thread terminates, leading to --- undefined and hard to debug behaviour (indeed, whether or not this results in --- problems may well depend on precise timing); an exception that is thrown when --- /allocating/ the resource is (more) deterministic and easier to debug. --- Unfortunately, it means that the above pattern is not applicable, as the --- thread spawned by 'withAsync' is not known to the registry, and so if it were --- to try to use the registry, the registry would throw an error (even though --- this pattern is actually safe). This situation is not ideal, but for now we --- merely provide an alternative to 'withAsync' that /does/ register the thread --- with the registry. --- --- NOTE: Threads that are spawned out of the user's control but that must still --- make use of the registry can use the unsafe API. This should be used with --- caution, however. -withThread :: - (MonadMask m, MonadAsync m) - => ResourceRegistry m - -> String -- ^ Label for the thread - -> m a - -> (Thread m a -> m b) - -> m b -withThread rr label body = bracket (forkThread rr label body) cancelThread - --- | Link specified 'Thread' to the (thread that created) the registry -linkToRegistry :: (MonadAsync m, MonadFork m, MonadMask m) => Thread m a -> m () -linkToRegistry t = linkTo (registryThread $ threadRegistry t) (threadAsync t) - --- | Fork a thread and link to it to the registry. --- --- This function is just a convenience. -forkLinkedThread :: - (MonadAsync m, MonadFork m, MonadMask m, HasCallStack) - => ResourceRegistry m - -> String -- ^ Label for the thread - -> m a - -> m (Thread m a) -forkLinkedThread rr label body = do - t <- forkThread rr label body - -- There is no race condition here between the new thread throwing an - -- exception and the 'linkToRegistry': if the thread /already/ threw the - -- exception when we link it, the exception will be raised immediately - -- (see 'linkTo' for details). - linkToRegistry t - return t - -{------------------------------------------------------------------------------- - Check that registry is used from known thread --------------------------------------------------------------------------------} - -ensureKnownThread :: - forall m. - (MonadThrow m, MonadThread m, MonadSTM m) - => ResourceRegistry m - -> Context m - -> m () -ensureKnownThread rr context = do - isKnown <- checkIsKnown - unless isKnown $ - throwIO $ ResourceRegistryUsedFromUntrackedThread { - resourceRegistryCreatedIn = registryContext rr - , resourceRegistryUsedIn = context - } - where - checkIsKnown :: m Bool - checkIsKnown - | contextThreadId context == contextThreadId (registryContext rr) = - return True - | otherwise = atomically $ do - KnownThreads ts <- registryThreads <$> readTVar (registryState rr) - return $ contextThreadId context `Set.member` ts - --- | Registry used from untracked threads --- --- If this exception is raised, it indicates a bug in the caller. -data ResourceRegistryThreadException = - -- | If the registry is used from an untracked thread, we cannot do proper - -- reference counting. The following threads are /tracked/: the thread - -- that spawned the registry and all threads spawned by the registry. - forall m. MonadThread m => ResourceRegistryUsedFromUntrackedThread { - -- | Information about the context in which the registry was created - resourceRegistryCreatedIn :: !(Context m) - - -- | The context in which it was used - , resourceRegistryUsedIn :: !(Context m) - } - - -- | Registry closed from different threat than that created it - | forall m. MonadThread m => ResourceRegistryClosedFromWrongThread { - -- | Information about the context in which the registry was created - resourceRegistryCreatedIn :: !(Context m) - - -- | The context in which it was used - , resourceRegistryUsedIn :: !(Context m) - } - -deriving instance Show ResourceRegistryThreadException -instance Exception ResourceRegistryThreadException - -{------------------------------------------------------------------------------- - Auxiliary: context --------------------------------------------------------------------------------} - --- | The internal context of a resource registry, recording a 'PrettyCallStack' --- of its creation and the creator's 'ThreadId' -data Context m = MonadThread m => Context { - -- | CallStack in which it was created - contextCallStack :: !PrettyCallStack - - -- | Thread that created the registry or resource - , contextThreadId :: !(ThreadId m) - } - --- Existential type; we can't use generics -instance NoThunks (Context m) where - showTypeOf _ = "Context" - wNoThunks ctxt (Context cs tid) = allNoThunks - [ noThunks ctxt cs - , noThunks ctxt (InspectHeapNamed @"ThreadId" tid) - ] - -deriving instance Show (Context m) - -captureContext :: MonadThread m => HasCallStack => m (Context m) -captureContext = Context prettyCallStack <$> myThreadId - -{------------------------------------------------------------------------------- - Misc utilities --------------------------------------------------------------------------------} - --- | Generalization of 'link' that links an async to an arbitrary thread. --- --- Non standard (not in 'async' library) --- -linkTo :: - (MonadAsync m, MonadFork m, MonadMask m) - => ThreadId m - -> Async m a - -> m () -linkTo tid = linkToOnly tid (not . isCancel) - --- | Generalization of 'linkOnly' that links an async to an arbitrary thread. --- --- Non standard (not in 'async' library). --- -linkToOnly :: - forall m a. - (MonadAsync m, MonadFork m, MonadMask m) - => ThreadId m - -> (SomeException -> Bool) - -> Async m a - -> m () -linkToOnly tid shouldThrow a = do - void $ forkRepeat ("linkToOnly " <> show linkedThreadId) $ do - r <- waitCatch a - case r of - Left e | shouldThrow e -> throwTo tid (exceptionInLinkedThread e) - _otherwise -> return () - where - linkedThreadId :: ThreadId m - linkedThreadId = asyncThreadId a - - exceptionInLinkedThread :: SomeException -> ExceptionInLinkedThread - exceptionInLinkedThread = - ExceptionInLinkedThread (show linkedThreadId) - -isCancel :: SomeException -> Bool -isCancel e - | Just AsyncCancelled <- fromException e = True - | otherwise = False - -forkRepeat :: (MonadFork m, MonadMask m) => String -> m a -> m (ThreadId m) -forkRepeat label action = - mask $ \restore -> - let go = do r <- tryAll (restore action) - case r of - Left _ -> go - _ -> return () - in forkIO (labelThisThread label >> go) - -tryAll :: MonadCatch m => m a -> m (Either SomeException a) -tryAll = try - -mustBeRight :: Either Void a -> a -mustBeRight (Left v) = absurd v -mustBeRight (Right a) = a - -{------------------------------------------------------------------------------- - Auxiliary: CallStack with different Show instance --------------------------------------------------------------------------------} - --- | CallStack with 'Show' instance using 'prettyCallStack' -newtype PrettyCallStack = PrettyCallStack CallStack - deriving newtype (NoThunks) - -instance Show PrettyCallStack where - show (PrettyCallStack cs) = GHC.prettyCallStack cs - --- | Capture a 'PrettyCallStack' -prettyCallStack :: HasCallStack => PrettyCallStack -prettyCallStack = PrettyCallStack GHC.callStack - -{------------------------------------------------------------------------------- - Orphan instance --------------------------------------------------------------------------------} - -instance (NoThunks k, NoThunks v) - => NoThunks (Bimap k v) where +{-# OPTIONS_GHC -Wno-orphans #-} + +-- | Note on terminology: when thread A forks thread B, we will say that thread A +-- is the \"parent\" and thread B is the \"child\". No further relationship +-- between the two threads is implied by this terminology. In particular, note +-- that the child may outlive the parent. We will use \"fork\" and \"spawn\" +-- interchangeably. +-- +-- = Motivation +-- +-- Whenever we allocate resources, we must keep track of them so that we can +-- deallocate them when they are no longer required. The most important tool we +-- have to achieve this is 'bracket': +-- +-- > bracket allocateResource releaseResource $ \r -> +-- > .. use r .. +-- +-- Often 'bracket' comes in the guise of a with-style combinator +-- +-- > withResource $ \r -> +-- > .. use r .. +-- +-- Where this pattern is applicable, it should be used and there is no need to +-- use the 'ResourceRegistry'. However, 'bracket' introduces strict lexical +-- scoping: the resource is available inside the scope of the bracket, and +-- will be deallocated once we leave that scope. That pattern is sometimes +-- hard to use. +-- +-- For example, suppose we have this interface to an SQL server +-- +-- > query :: Query -> IO QueryHandle +-- > close :: QueryHandle -> IO () +-- > next :: QueryHandle -> IO Row +-- +-- and suppose furthermore that we are writing a simple webserver that allows a +-- client to send multiple SQL queries, get rows from any open query, and close +-- queries when no longer required: +-- +-- > server :: IO () +-- > server = go Map.empty +-- > where +-- > go :: Map QueryId QueryHandle -> IO () +-- > go handles = getRequest >>= \case +-- > New q -> do +-- > h <- query q -- allocate +-- > qId <- generateQueryId +-- > sendResponse qId +-- > go $ Map.insert qId h handles +-- > Close qId -> do +-- > close (handles ! qId) -- release +-- > go $ Map.delete qId handles +-- > Next qId -> do +-- > sendResponse =<< next (handles ! qId) +-- > go handles +-- +-- The server opens and closes query handles in response to client requests. +-- Restructuring this code to use 'bracket' would be awkward, but as it stands +-- this code does not ensure that resources get deallocated; for example, if +-- the server thread is killed ('killThread'), resources will be leaked. +-- +-- Another, perhaps simpler, example is spawning threads. Threads too should +-- be considered to be resources that we should keep track of and deallocate +-- when they are no longer required, primarily because when we deallocate +-- (terminate) those threads they too will have a chance to deallocate /their/ +-- resources. As for other resources, we have a with-style combinator for this +-- +-- > withAsync $ \thread -> .. +-- +-- Lexical scoping of threads is often inconvenient, however, more so than for +-- regular resources. The temptation is therefore to simply fork a thread and +-- forget about it, but if we are serious about resource deallocation this is +-- not an acceptable solution. +-- +-- = The resource registry +-- +-- The resource registry is essentially a piece of state tracking which +-- resources have been allocated. The registry itself is allocated with a +-- with-style combinator 'withRegistry', and when we leave that scope any +-- resources not yet deallocated will be released at that point. Typically +-- the registry is only used as a fall-back, ensuring that resources will +-- deallocated even in the presence of exceptions. For example, here's how +-- we might rewrite the above server example using a registry: +-- +-- > server' :: IO () +-- > server' = +-- > withRegistry $ \registry -> go registry Map.empty +-- > where +-- > go :: ResourceRegistry IO +-- > -> Map QueryId (ResourceKey, QueryHandle) +-- > -> IO () +-- > go registry handles = getRequest >>= \case +-- > New q -> do +-- > (key, h) <- allocate registry (query q) close -- allocate +-- > qId <- generateQueryId +-- > sendResponse qId +-- > go registry $ Map.insert qId (key, h) handles +-- > Close qId -> do +-- > release registry (fst (handles ! qId)) -- release +-- > go registry $ Map.delete qId handles +-- > Next qId -> do +-- > sendResponse =<< next (snd (handles ! qId)) +-- > go registry handles +-- +-- We allocate the query with the help of the registry, providing the registry +-- with the means to deallocate the query should that be required. We can /and +-- should/ still manually release resources also: in this particular example, +-- the (lexical) scope of the registry is the entire server thread, so delaying +-- releasing queries until we exit that scope will probably mean we hold on to +-- resources for too long. The registry is only there as a fall-back. +-- +-- = Spawning threads +-- +-- We already observed in the introduction that insisting on lexical scoping +-- for threads is often inconvenient, and that simply using +-- 'Control.Monad.Class.MonadFork.forkIO' is no solution as it means we might +-- leak resources. There is however another problem with +-- 'Control.Monad.Class.MonadFork.forkIO'. Consider this snippet: +-- +-- > withRegistry $ \registry -> +-- > r <- allocate registry allocateResource releaseResource +-- > forkIO $ .. use r .. +-- +-- It is easy to see that this code is problematic: we allocate a resource @r@, +-- then spawn a thread that uses @r@, and finally leave the scope of +-- 'withRegistry', thereby deallocating @r@ -- leaving the thread to run with +-- a now deallocated resource. +-- +-- It is /only/ safe for threads to use a given registry, and/or its registered +-- resources, if the lifetime of those threads is tied to the lifetime of the +-- registry. There would be no problem with the example above if the thread +-- would be terminated when we exit the scope of 'withRegistry'. +-- +-- The 'forkThread' combinator provided by the registry therefore does two +-- things: it allocates the thread as a resource in the registry, so that it can +-- kill the thread when releasing all resources in the registry. It also records +-- the thread ID in a set of known threads. Whenever the registry is accessed +-- from a thread /not/ in this set, the registry throws a runtime exception, +-- since such a thread might outlive the registry and hence its contents. The +-- intention is that this guards against dangerous patterns like the one above. +-- +-- = Linking +-- +-- When thread A spawns thread B using 'withAsync', the lifetime of B is tied +-- to the lifetime of A: +-- +-- > withAsync .. $ \threadB -> .. +-- +-- After all, when A exits the scope of the 'withAsync', thread B will be +-- killed. The reverse is however not true: thread B can terminate before +-- thread A. It is often useful for thread A to be able to declare a dependency +-- on thread B: if B somehow fails, that is, terminates with an exception, we +-- want that exception to be rethrown in thread A as well. A can achieve this +-- by /linking/ to B: +-- +-- > withAsync .. $ \threadB -> do +-- > link threadB +-- > .. +-- +-- Linking a parent to a child is however of limited value if the lifetime of +-- the child is not limited by the lifetime of the parent. For example, if A +-- does +-- +-- > threadB <- async $ .. +-- > link threadB +-- +-- and A terminates before B does, any exception thrown by B might be send to a +-- thread that no longer exists. This is particularly problematic when we start +-- chaining threads: if A spawns-and-links-to B which spawns-and-links-to C, and +-- C throws an exception, perhaps the intention is that this gets rethrown to B, +-- and then rethrown to A, terminating all three threads; however, if B has +-- terminated before the exception is thrown, C will throw the exception to a +-- non-existent thread and A is never notified. +-- +-- For this reason, the registry's 'linkToRegistry' combinator does not link the +-- specified thread to the thread calling 'linkToRegistry', but rather to the +-- thread that created the registry. After all, the lifetime of threads spawned +-- with 'forkThread' can certainly exceed the lifetime of their parent threads, +-- but the lifetime of /all/ threads spawned using the registry will be limited +-- by the scope of that registry, and hence the lifetime of the thread that +-- created it. So, when we call 'linkToRegistry', the exception will be thrown +-- the thread that created the registry, which (if not caught) will cause that +-- that to exit the scope of 'withRegistry', thereby terminating all threads in +-- that registry. +-- +-- = Combining the registry and with-style allocation +-- +-- It is perfectly possible (indeed, advisable) to use 'bracket' and +-- bracket-like allocation functions alongside the registry, but note that the +-- usual caveats with 'bracket' and forking threads still applies. In +-- particular, spawning threads inside the 'bracket' that make use of the +-- bracketed resource is problematic; this is of course true whether or not a +-- registry is used. +-- +-- In principle this also includes 'withAsync'; however, since 'withAsync' +-- results in a thread that is not known to the registry, such a thread will not +-- be able to use the registry (the registry would throw an unknown thread +-- exception, as described above). For this purpose we provide 'withThread'; +-- 'withThread' (as opposed to 'forkThread') should be used when a parent thread +-- wants to handle exceptions in the child thread; see 'withThread' for +-- detailed discussion. +-- +-- It is /also/ fine to includes nested calls to 'withRegistry'. Since the +-- lifetime of such a registry (and all resources within) is tied to the thread +-- calling 'withRegistry', which itself is tied to the "parent registry" in +-- which it was created, this creates a hierarchy of registries. It is of course +-- essential for compositionality that we should be able to create local +-- registries, but even if we do have easy access to a parent regisry, creating +-- a local one where possibly is useful as it limits the scope of the resources +-- created within, and hence their maximum lifetimes. +module Control.ResourceRegistry + ( -- * The resource registry proper + Context + , ResourceId + , ResourceRegistry + + -- * Exceptions + , RegistryClosedException (..) + , ResourceRegistryThreadException + + -- * Creating and releasing the registry itself + , bracketWithPrivateRegistry + , registryThread + , withRegistry + + -- * Allocating and releasing regular resources + , ResourceKey + , allocate + , allocateEither + , release + , releaseAll + , unsafeRelease + , unsafeReleaseAll + + -- * Threads + , Thread + , cancelThread + , forkLinkedThread + , forkThread + , linkToRegistry + , threadId + , waitAnyThread + , waitThread + , withThread + + -- * Temporary registry + , TempRegistryException (..) + , WithTempRegistry + , allocateTemp + , modifyWithTempRegistry + , runInnerWithTempRegistry + , runWithTempRegistry + + -- * Unsafe combinators primarily for testing + , closeRegistry + , countResources + , unsafeNewRegistry + ) where + +import Control.Applicative ((<|>)) +import Control.Concurrent.Class.MonadSTM.Strict +import Control.Exception (asyncExceptionFromException) +import Control.Monad +import Control.Monad.Class.MonadAsync +import Control.Monad.Class.MonadFork +import Control.Monad.Class.MonadThrow +import Control.Monad.Reader +import Control.Monad.State.Strict +import Data.Bifunctor +import Data.Bimap (Bimap) +import Data.Bimap qualified as Bimap +import Data.Either (partitionEithers) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Maybe (catMaybes, listToMaybe) +import Data.Set (Set) +import Data.Set qualified as Set +import Data.Void +import Data.Word (Word64) +import GHC.Generics (Generic) +import GHC.Stack (CallStack, HasCallStack) +import GHC.Stack qualified as GHC +import NoThunks.Class hiding (Context) + +-- | Tracks resources during their lifetime. +data ResourceRegistry m = ResourceRegistry + { registryContext :: !(Context m) + -- ^ Context in which the registry was created + , registryState :: !(StrictTVar m (RegistryState m)) + -- ^ Registry state + } + deriving Generic + +deriving instance + (forall a. NoThunks a => NoThunks (StrictTVar m a)) => + NoThunks (ResourceRegistry m) + +{------------------------------------------------------------------------------- + Internal: registry state +-------------------------------------------------------------------------------} + +-- | The age of a resource +-- +-- Age here is represented by an meaningless number. The one and only property +-- that matters is that the age of resource A that was successfully allocated +-- before resource B was (in the same registry) will be greater than the age of +-- resource B. +-- +-- For the current implementation, that property will be true unless the +-- registry lives long enough to have contained 2^64 separately allocated +-- resources. +-- +-- This data is not exposed by the 'ResourceRegistry' interface. +newtype Age = Age Word64 + deriving stock Show + deriving newtype (Eq, Ord) + deriving NoThunks via InspectHeapNamed "Age" Age + +-- | The age of the first resource successfully allocated in a fresh registry +ageOfFirstResource :: Age +ageOfFirstResource = Age maxBound + +-- | Map the age of the latest resource to be successfully allocated to the age +-- of the next resource to be successfully allocated in the same registry +nextYoungerAge :: Age -> Age +nextYoungerAge (Age n) = Age (n - 1) + +-- | Internal registry state +data RegistryState m = RegistryState + { registryThreads :: !(KnownThreads m) + -- ^ Forked threads + , registryResources :: !(Map ResourceId (Resource m)) + -- ^ Currently allocated resources + -- + -- INVARIANT: We record exactly the ages of currently allocated resources, + -- @'Bimap.keys' . 'registryAges' = 'Map.keys' . 'registryResources'@. + , registryNextKey :: !ResourceId + -- ^ Next available resource key + , registryAges :: !(Bimap ResourceId Age) + -- ^ The age of each currently allocated resource + -- + -- We use a 'Bimap' so we can maintain the keys in sorted order by age, + -- which is necessary when closing the registry. + , registryNextAge :: !Age + -- ^ The age of the next resource + , registryStatus :: !RegistryStatus + -- ^ Does the registry still accept new allocations? + -- + -- See 'RegistryClosedException' for discussion. + } + deriving (Generic, NoThunks) + +-- | The currently allocated keys in youngest-to-oldest order +getYoungestToOldest :: RegistryState m -> [ResourceId] +getYoungestToOldest = map snd . Bimap.toAscListR . registryAges + +-- | Threads known to the registry +-- +-- This is the set of threads spawned using 'forkThread'. The lifetimes of all +-- of these threads are limited by the lifetime of the registry. +-- +-- Does not include the thread ID of the thread that created the registry. After +-- all, this thread may well outlive the registry (though the registry cannot +-- outlive it). +-- +-- Invariant (informal): the set of registered threads is a subset of the +-- registered resources ('registryResources'). (This invariant is temporarily +-- broken when we start a new thread in 'forkThread' but will be re-established +-- before that thread starts execution proper.) +newtype KnownThreads m = KnownThreads (Set (ThreadId m)) + deriving NoThunks via InspectHeapNamed "KnownThreads" (KnownThreads m) + +-- | Status of the registry (open or closed) +data RegistryStatus + = RegistryOpen + | -- | We record the 'CallStack' to the call to 'close + RegistryClosed !PrettyCallStack + deriving (Generic, NoThunks) + +-- | Resource key +-- +-- Resource keys are tied to a particular registry. +data ResourceKey m = ResourceKey !(ResourceRegistry m) !ResourceId + deriving Generic + +deriving instance + NoThunks (ResourceRegistry m) => + NoThunks (ResourceKey m) + +-- | Return the 'ResourceId' of a 'ResourceKey'. +resourceKeyId :: ResourceKey m -> ResourceId +resourceKeyId (ResourceKey _rr rid) = rid + +-- | Resource ID +-- +-- This uniquifying data is not exposed by the 'ResourceRegistry' interface. +newtype ResourceId = ResourceId Int + deriving stock (Show, Eq, Ord) + deriving newtype (Enum, NoThunks) + +-- | Information about a resource +data Resource m = Resource + { resourceContext :: !(Context m) + -- ^ Context in which the resource was created + , resourceRelease :: !(Release m) + -- ^ Deallocate the resource + } + deriving (Generic, NoThunks) + +-- | Release the resource, return 'True' when the resource was actually +-- released, return 'False' when the resource was already released. +-- +-- If unsure, returning 'True' is always fine. +newtype Release m = Release (m Bool) + deriving NoThunks via OnlyCheckWhnfNamed "Release" (Release m) + +releaseResource :: Resource m -> m Bool +releaseResource Resource{resourceRelease = Release f} = f + +instance Show (Release m) where + show _ = "<>" + +{------------------------------------------------------------------------------- + Internal: pure functions on the registry state +-------------------------------------------------------------------------------} + +modifyKnownThreads :: + (Set (ThreadId m) -> Set (ThreadId m)) -> + KnownThreads m -> + KnownThreads m +modifyKnownThreads f (KnownThreads ts) = KnownThreads (f ts) + +-- | Auxiliary for functions that should be disallowed when registry is closed +unlessClosed :: + State (RegistryState m) a -> + State (RegistryState m) (Either PrettyCallStack a) +unlessClosed f = do + status <- gets registryStatus + case status of + RegistryClosed closed -> return $ Left closed + RegistryOpen -> Right <$> f + +-- | Allocate key for new resource +allocKey :: State (RegistryState m) (Either PrettyCallStack ResourceId) +allocKey = unlessClosed $ do + nextKey <- gets registryNextKey + modify $ \st -> st{registryNextKey = succ nextKey} + return nextKey + +-- | Insert new resource +insertResource :: + ResourceId -> + Resource m -> + State (RegistryState m) (Either PrettyCallStack ()) +insertResource key r = unlessClosed $ do + modify $ \st -> + st + { registryResources = Map.insert key r (registryResources st) + , registryAges = + Bimap.insert + key + (registryNextAge st) + (registryAges st) + , registryNextAge = nextYoungerAge (registryNextAge st) + } + +-- | Remove resource from the registry (if it exists) +removeResource :: ResourceId -> State (RegistryState m) (Maybe (Resource m)) +removeResource key = state $ \st -> + let (mbResource, resources') = + Map.updateLookupWithKey + (\_ _ -> Nothing) + key + (registryResources st) + + st' = + st + { registryResources = resources' + , registryAges = Bimap.delete key (registryAges st) + } + in (mbResource, st') + +-- | Insert thread into the set of known threads +insertThread :: MonadThread m => ThreadId m -> State (RegistryState m) () +insertThread tid = + modify $ \st -> + st + { registryThreads = + modifyKnownThreads (Set.insert tid) $ + registryThreads st + } + +-- | Remove thread from set of known threads +removeThread :: MonadThread m => ThreadId m -> State (RegistryState m) () +removeThread tid = + modify $ \st -> + st + { registryThreads = + modifyKnownThreads (Set.delete tid) $ + registryThreads st + } + +-- | Close the registry +-- +-- Returns the keys currently allocated if the registry is not already closed. +-- +-- POSTCONDITION: They are returned in youngest-to-oldest order. +close :: + PrettyCallStack -> + State (RegistryState m) (Either PrettyCallStack [ResourceId]) +close closeCallStack = unlessClosed $ do + modify $ \st -> st{registryStatus = RegistryClosed closeCallStack} + gets getYoungestToOldest + +-- | Convenience function for updating the registry state +updateState :: + forall m a. + MonadSTM m => + ResourceRegistry m -> + State (RegistryState m) a -> + m a +updateState rr f = + atomically $ stateTVar (registryState rr) (runState f) + +-- | Attempt to allocate a resource in a registry which is closed +-- +-- When calling 'closeRegistry' (typically, leaving the scope of +-- 'withRegistry'), all resources in the registry must be released. If a +-- concurrent thread is still allocating resources, we end up with a race +-- between the thread trying to allocate new resources and the registry trying +-- to free them all. To avoid this, before releasing anything, the registry will +-- record itself as closed. Any attempt by a concurrent thread to allocate a new +-- resource will then result in a 'RegistryClosedException'. +-- +-- It is probably not particularly useful for threads to try and catch this +-- exception (apart from in a generic handler that does local resource cleanup). +-- The thread will anyway soon receive a 'Control.Exception.ThreadKilled' +-- exception. +data RegistryClosedException + = forall m. MonadThread m => RegistryClosedException + { registryClosedRegistryContext :: !(Context m) + -- ^ The context in which the registry was created + , registryClosedCloseCallStack :: !PrettyCallStack + -- ^ Callstack to the call to 'closeRegistry' + -- + -- Note that 'closeRegistry' can only be called from the same thread + -- that created the registry. + , registryClosedAllocContext :: !(Context m) + -- ^ Context of the call resulting in the exception + } + +deriving instance Show RegistryClosedException +instance Exception RegistryClosedException + +{------------------------------------------------------------------------------- + Creating and releasing the registry itself +-------------------------------------------------------------------------------} + +-- | Create a new registry +-- +-- You are strongly encouraged to use 'withRegistry' instead. +-- Exported primarily for the benefit of tests. +unsafeNewRegistry :: + (MonadSTM m, MonadThread m, HasCallStack) => + m (ResourceRegistry m) +unsafeNewRegistry = do + context <- captureContext + stateVar <- newTVarIO initState + return + ResourceRegistry + { 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 + } + +-- | Close the registry +-- +-- This can only be called from the same thread that created the registry. +-- This is a no-op if the registry is already closed. +-- +-- This entire function runs with exceptions masked, so that we are not +-- interrupted while we release all resources. +-- +-- Resources will be allocated from young to old, so that resources allocated +-- later can safely refer to resources created earlier. +-- +-- The release functions are run in the scope of an exception handler, so that +-- if releasing one resource throws an exception, we still attempt to release +-- the other resources. Should we catch an exception whilst we close the +-- registry, we will rethrow it after having attempted to release all resources. +-- If there is more than one, we will pick a random one to rethrow, though we +-- will prioritize asynchronous exceptions over other exceptions. This may be +-- important for exception handlers that catch all-except-asynchronous +-- exceptions. +closeRegistry :: + (MonadMask m, MonadThread m, MonadSTM m, HasCallStack) => + ResourceRegistry m -> + m () +closeRegistry rr = mask_ $ do + context <- captureContext + unless (contextThreadId context == contextThreadId (registryContext rr)) $ + throwIO $ + ResourceRegistryClosedFromWrongThread + { resourceRegistryCreatedIn = registryContext rr + , resourceRegistryUsedIn = context + } + + -- Close the registry so that we cannot allocate any further resources + alreadyClosed <- updateState rr $ close (contextCallStack context) + case alreadyClosed of + Left _ -> + return () + Right keys -> do + -- At this point we have not /removed/ any elements from the map, + -- allowing concurrent threads to do their own cleanup of resources + -- (this may for instance be important if a thread deallocates its + -- resources in a particular order -- note that cancelling a thread + -- is a synchronous operation, so we will wait for it to finish + -- releasing its resources.) + -- /If/ a concurrent thread does some cleanup, then some of the calls + -- to 'release' that we do here might be no-ops. + void $ releaseResources rr keys release + +-- | Helper for 'closeRegistry', 'releaseAll', and 'unsafeReleaseAll': release +-- the resources allocated with the given 'ResourceId's. +-- +-- Returns the contexts of the resources that were actually released. +releaseResources :: + MonadCatch m => + ResourceRegistry m -> + -- | PRECONDITION: The currently allocated keys, + -- youngest-to-oldest + [ResourceId] -> + -- | How to release the resource, e.g., 'release' or + -- 'unsafeRelease'. + (ResourceKey m -> m (Maybe (Context m))) -> + m [Context m] +releaseResources rr sortedKeys releaser = do + (exs, mbContexts) <- + fmap partitionEithers $ + forM sortedKeys $ + try . releaser . ResourceKey rr + + case prioritize exs of + Nothing -> return (catMaybes mbContexts) + Just e -> throwIO e + where + prioritize :: [SomeException] -> Maybe SomeException + prioritize = + (\(asyncEx, otherEx) -> listToMaybe asyncEx <|> listToMaybe otherEx) + . first catMaybes + . unzip + . map (\e -> (asyncExceptionFromException e, e)) + +-- | Create a new registry +-- +-- See documentation of 'ResourceRegistry' for a detailed discussion. +withRegistry :: + (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) => + (ResourceRegistry m -> m a) -> + m a +withRegistry = bracket unsafeNewRegistry closeRegistry + +-- | Create a new private registry for use by a bracketed resource +-- +-- Use this combinator as a more specific and easier-to-maintain alternative to +-- the following. +-- +-- > 'withRegistry' $ \rr -> +-- > 'bracket' (newFoo rr) closeFoo $ \foo -> +-- > (... rr does not occur in this scope ...) +-- +-- NB The scoped body can use `withRegistry` if it also needs its own, separate +-- registry. +-- +-- Use this combinator to emphasize that the registry is private to (ie only +-- used by and/or via) the bracketed resource and that it thus has nearly the +-- same lifetime. This combinator ensures the following specific invariants +-- regarding lifetimes and order of releases. +-- +-- o The registry itself is older than the bracketed resource. +-- +-- o The only registered resources older than the bracketed resource were +-- allocated in the registry by the function that allocated the bracketed +-- resource. +-- +-- o Because of the older resources, the bracketed resource is itself also +-- registered in the registry; that's the only way we can be sure to release +-- all resources in the right order. +-- +-- NB Because the registry is private to the resource, the @a@ type could save +-- the handle to @registry@ and safely close the registry if the scoped body +-- calls @closeA@ before the bracket ends. Though we have not used the type +-- system to guarantee that the interface of the @a@ type cannot leak the +-- registry to the body, this combinator does its part to keep the registry +-- private to the bracketed resource. +-- +-- See documentation of 'ResourceRegistry' for a more general discussion. +bracketWithPrivateRegistry :: + (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) => + (ResourceRegistry m -> m a) -> + -- | Release the resource + (a -> m ()) -> + (a -> m r) -> + m r +bracketWithPrivateRegistry newA closeA body = + withRegistry $ \registry -> do + (_key, a) <- allocate registry (\_key -> newA registry) closeA + body a + +{------------------------------------------------------------------------------- + Temporary registry +-------------------------------------------------------------------------------} + +-- | Run an action with a temporary resource registry. +-- +-- When allocating resources that are meant to end up in some final state, +-- e.g., stored in a 'Control.Monad.Class.MonadSTM.TVar', after which they are +-- guaranteed to be released correctly, it is possible that an exception is +-- thrown after allocating such a resource, but before it was stored in the +-- final state. In that case, the resource would be leaked. +-- 'runWithTempRegistry' solves that problem. +-- +-- When no exception is thrown before the end of 'runWithTempRegistry', the +-- user must have transferred all the resources it allocated to their final +-- state. This means that these resources don't have to be released by the +-- temporary registry anymore, the final state is now in charge of releasing +-- them. +-- +-- In case an exception is thrown before the end of 'runWithTempRegistry', +-- /all/ resources allocated in the temporary registry will be released. +-- +-- Resources must be allocated using 'allocateTemp'. +-- +-- To make sure that the user doesn't forget to transfer a resource to the +-- final state @st@, the user must pass a function to 'allocateTemp' that +-- checks whether a given @st@ contains the resource, i.e., whether the +-- resource was successfully transferred to its final destination. +-- +-- When no exception is thrown before the end of 'runWithTempRegistry', we +-- check whether all allocated resources have been transferred to the final +-- state @st@. If there's a resource that hasn't been transferred to the final +-- state /and/ that hasn't be released or closed before (see the release +-- function passed to 'allocateTemp'), a 'TempRegistryRemainingResource' +-- exception will be thrown. +-- +-- For that reason, 'WithTempRegistry' is parameterised over the final state +-- type @st@ and the given 'WithTempRegistry' action must return the final +-- state. +-- +-- NOTE: we explicitly don't let 'runWithTempRegistry' return the final state, +-- because the state /must/ have been stored somewhere safely, transferring +-- the resources, before the temporary registry is closed. +runWithTempRegistry :: + (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) => + WithTempRegistry st m (a, st) -> + m a +runWithTempRegistry m = withRegistry $ \rr -> do + varTransferredTo <- newTVarIO mempty + let tempRegistry = + TempRegistry + { tempResourceRegistry = rr + , tempTransferredTo = varTransferredTo + } + (a, st) <- runReaderT (unWithTempRegistry m) tempRegistry + -- We won't reach this point if an exception is thrown, so we won't check + -- for remaining resources in that case. + -- + -- No need to mask here, whether we throw the async exception or + -- 'TempRegistryRemainingResource' doesn't matter. + transferredTo <- readTVarIO varTransferredTo + untrackTransferredTo rr transferredTo st + + context <- captureContext + remainingResources <- releaseAllHelper rr context release + + whenJust (listToMaybe remainingResources) $ \remainingResource -> + throwIO $ + TempRegistryRemainingResource + { tempRegistryContext = registryContext rr + , tempRegistryResource = remainingResource + } + return a + where + whenJust (Just x) f = f x + whenJust Nothing _ = pure () + +-- | Embed a self-contained 'WithTempRegistry' computation into a larger one. +-- +-- The internal 'WithTempRegistry' is effectively passed to +-- 'runWithTempRegistry'. It therefore must have no dangling resources, for +-- example. This is the meaning of /self-contained/ above. +-- +-- The key difference beyond 'runWithTempRegistry' is that the resulting +-- composite resource is also guaranteed to be registered in the outer +-- 'WithTempRegistry' computation's registry once the inner registry is closed. +-- Combined with the following assumption, this establishes the invariant that +-- all resources are (transitively) in a temporary registry. +-- +-- As the resource might require some implementation details to be closed, the +-- function to close it will also be provided by the inner computation. +-- +-- ASSUMPTION: closing @res@ closes every resource contained in @innerSt@ +-- +-- NOTE: In the current implementation, there will be a brief moment where the +-- inner registry still contains the inner computation's resources and also the +-- outer registry simultaneously contains the new composite resource. If an +-- async exception is received at that time, then the inner resources will be +-- closed and then the composite resource will be closed. This means there's a +-- risk of /double freeing/, which can be harmless if anticipated. +runInnerWithTempRegistry :: + forall innerSt st m res a. + (MonadSTM m, MonadMask m, MonadThread m) => + -- | The embedded computation; see ASSUMPTION above + WithTempRegistry innerSt m (a, innerSt, res) -> + -- | How to free; same as for 'allocateTemp' + (res -> m Bool) -> + -- | How to check; same as for 'allocateTemp' + (st -> res -> Bool) -> + WithTempRegistry st m a +runInnerWithTempRegistry inner free isTransferred = do + outerTR <- WithTempRegistry ask + + lift $ runWithTempRegistry $ do + (a, innerSt, res) <- inner + + -- Allocate in the outer layer. + _ <- + withFixedTempRegistry outerTR $ + allocateTemp (return res) free isTransferred + + -- TODO This point here is where an async exception could cause both the + -- inner resources to be closed and the outer resource to be closed later. + -- + -- If we want to do better than that, we'll need a variant of + -- 'runWithTempRegistry' that lets us perform some action with async + -- exceptions masked "at the same time" it closes its registry. + + -- Note that everything in `inner` allocated via `allocateTemp` must + -- either be closed or else present in `innerSt` by this point -- + -- `runWithTempRegistry` would have thrown if not. + pure (a, innerSt) + where + withFixedTempRegistry :: + TempRegistry st m -> + WithTempRegistry st m res -> + WithTempRegistry innerSt m res + withFixedTempRegistry env (WithTempRegistry (ReaderT f)) = + WithTempRegistry $ ReaderT $ \_ -> f env + +-- | When 'runWithTempRegistry' exits successfully while there are still +-- resources remaining in the temporary registry that haven't been transferred +-- to the final state. +data TempRegistryException + = forall m. MonadThread m => TempRegistryRemainingResource + { tempRegistryContext :: !(Context m) + -- ^ The context in which the temporary registry was created. + , tempRegistryResource :: !(Context m) + -- ^ The context in which the resource was allocated that was not + -- transferred to the final state. + } + +deriving instance Show TempRegistryException +instance Exception TempRegistryException + +-- | Given a final state, return the 'ResourceId's of the resources that have +-- been /transferred to/ that state. +newtype TransferredTo st = TransferredTo + { runTransferredTo :: st -> Set ResourceId + } + deriving newtype (Semigroup, Monoid) + deriving NoThunks via OnlyCheckWhnfNamed "TransferredTo" (TransferredTo st) + +-- | The environment used to run a 'WithTempRegistry' action. +data TempRegistry st m = TempRegistry + { tempResourceRegistry :: !(ResourceRegistry m) + , tempTransferredTo :: !(StrictTVar m (TransferredTo st)) + -- ^ Used as a @Writer@. + } + +-- | An action with a temporary registry in scope, see 'runWithTempRegistry' +-- for more details. +-- +-- The most important function to run in this monad is 'allocateTemp'. +newtype WithTempRegistry st m a = WithTempRegistry + { unWithTempRegistry :: ReaderT (TempRegistry st m) m a + } + deriving newtype + ( Functor + , Applicative + , Monad + , MonadThrow + , MonadCatch + , MonadMask + ) + +instance MonadTrans (WithTempRegistry st) where + lift = WithTempRegistry . lift + +instance MonadState s m => MonadState s (WithTempRegistry st m) where + state = WithTempRegistry . state + +-- | Untrack all the resources from the registry that have been transferred to +-- the given state. +-- +-- Untracking a resource means removing it from the registry without releasing +-- it. +-- +-- NOTE: does not check that it's called by the same thread that allocated the +-- resources, as it's an internal function only used in 'runWithTempRegistry'. +untrackTransferredTo :: + MonadSTM m => + ResourceRegistry m -> + TransferredTo st -> + st -> + m () +untrackTransferredTo rr transferredTo st = + updateState rr $ mapM_ removeResource rids + where + rids = runTransferredTo transferredTo st + +-- | Allocate a resource in a temporary registry until it has been transferred +-- to the final state @st@. See 'runWithTempRegistry' for more details. +allocateTemp :: + (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) => + -- | Allocate the resource + m a -> + -- | Release the resource, return 'True' when the resource was actually + -- released, return 'False' when the resource was already released. + -- + -- Note that it is safe to always return 'True' when unsure. + (a -> m Bool) -> + -- | Check whether the resource is in the given state + (st -> a -> Bool) -> + WithTempRegistry st m a +allocateTemp alloc free isTransferred = WithTempRegistry $ do + TempRegistry rr varTransferredTo <- ask + (key, a) <- + lift + ( mustBeRight + <$> allocateEither rr (fmap Right . const alloc) free + ) + lift $ + atomically $ + modifyTVar varTransferredTo $ + mappend $ + TransferredTo $ \st -> + if isTransferred st a + then Set.singleton (resourceKeyId key) + else Set.empty + return a + +-- | Higher level API on top of 'runWithTempRegistry': modify the given @st@, +-- allocating resources in the process that will be transferred to the +-- returned @st@. +modifyWithTempRegistry :: + forall m st a. + (MonadSTM m, MonadMask m, MonadThread m) => + -- | Get the state + m st -> + -- | Store the new state + (st -> ExitCase st -> m ()) -> + -- | Modify the state + StateT st (WithTempRegistry st m) a -> + m a +modifyWithTempRegistry getSt putSt modSt = + runWithTempRegistry $ + fst <$> generalBracket (lift getSt) transfer mutate + where + transfer :: st -> ExitCase (a, st) -> WithTempRegistry st m () + transfer initSt ec = lift $ putSt initSt (snd <$> ec) + + mutate :: st -> WithTempRegistry st m (a, st) + mutate = runStateT modSt + +{------------------------------------------------------------------------------- + Simple queries on the registry +-------------------------------------------------------------------------------} + +-- | The thread that created the registry +registryThread :: ResourceRegistry m -> ThreadId m +registryThread = contextThreadId . registryContext + +-- | Number of currently allocated resources +-- +-- Primarily for the benefit of testing. +countResources :: MonadSTM m => ResourceRegistry m -> m Int +countResources rr = atomically $ aux <$> readTVar (registryState rr) + where + aux :: RegistryState m -> Int + aux = Map.size . registryResources + +{------------------------------------------------------------------------------- + Allocating resources +-------------------------------------------------------------------------------} + +-- | Allocate new resource +-- +-- The allocation function will be run with asynchronous exceptions masked. This +-- means that the resource allocation must either be fast or else interruptible; +-- see "Dealing with Asynchronous Exceptions during Resource Acquisition" +-- for details. +allocate :: + forall m a. + (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) => + ResourceRegistry m -> + (ResourceId -> m a) -> + -- | Release the resource + (a -> m ()) -> + m (ResourceKey m, a) +allocate rr alloc free = + mustBeRight + <$> allocateEither rr (fmap Right . alloc) (\a -> free a >> return True) + +-- | Generalization of 'allocate' for allocation functions that may fail +allocateEither :: + forall m e a. + (MonadSTM m, MonadMask m, MonadThread m, HasCallStack) => + ResourceRegistry m -> + (ResourceId -> m (Either e a)) -> + -- | Release the resource, return 'True' when the resource + -- hasn't been released or closed before. + (a -> m Bool) -> + m (Either e (ResourceKey m, a)) +allocateEither rr alloc free = do + context <- captureContext + ensureKnownThread rr context + -- We check if the registry has been closed when we allocate the key, so + -- that we can avoid needlessly allocating the resource. + mKey <- updateState rr allocKey + case mKey of + Left closed -> + throwRegistryClosed rr context closed + Right key -> mask_ $ do + ma <- alloc key + case ma of + Left e -> return $ Left e + Right a -> do + -- TODO: Might want to have an exception handler around this call to + -- 'updateState' just in case /that/ throws an exception. + inserted <- + updateState rr $ + insertResource key (mkResource context a) + case inserted of + Left closed -> do + -- Despite the earlier check, it's possible that the registry + -- got closed after we allocated a new key but before we got a + -- chance to register the resource. In this case, we must + -- deallocate the resource again before throwing the exception. + void $ free a + throwRegistryClosed rr context closed + Right () -> + return $ Right (ResourceKey rr key, a) + where + mkResource :: Context m -> a -> Resource m + mkResource context a = + Resource + { resourceContext = context + , resourceRelease = Release $ free a + } + +throwRegistryClosed :: + (MonadThrow m, MonadThread m) => + ResourceRegistry m -> + Context m -> + PrettyCallStack -> + m x +throwRegistryClosed rr context closed = + throwIO + RegistryClosedException + { registryClosedRegistryContext = registryContext rr + , registryClosedCloseCallStack = closed + , registryClosedAllocContext = context + } + +-- | Release resource +-- +-- This deallocates the resource and removes it from the registry. It will be +-- the responsibility of the caller to make sure that the resource is no longer +-- used in any thread. +-- +-- The deallocation function is run with exceptions masked, so that we are +-- guaranteed not to remove the resource from the registry without releasing it. +-- +-- Releasing an already released resource is a no-op. +-- +-- When the resource has not been released before, its context is returned. +release :: + (MonadMask m, MonadSTM m, MonadThread m, HasCallStack) => + ResourceKey m -> + m (Maybe (Context m)) +release key@(ResourceKey rr _) = do + context <- captureContext + ensureKnownThread rr context + unsafeRelease key + +-- | Unsafe version of 'release' +-- +-- The only difference between 'release' and 'unsafeRelease' is that the latter +-- does not insist that it is called from a thread that is known to the +-- registry. This is dangerous, because it implies that there is a thread with +-- access to a resource which may be deallocated before that thread is +-- terminated. Of course, we can't detect all such situations (when the thread +-- merely uses a resource but does not allocate or release we can't tell), but +-- normally when we /do/ detect this we throw an exception. +-- +-- This function should only be used if the above situation can be ruled out +-- or handled by other means. +unsafeRelease :: + (MonadMask m, MonadSTM m) => + ResourceKey m -> + m (Maybe (Context m)) +unsafeRelease (ResourceKey rr rid) = do + mask_ $ do + mResource <- updateState rr $ removeResource rid + case mResource of + Nothing -> return Nothing + Just resource -> do + actuallyReleased <- releaseResource resource + return $ + if actuallyReleased + then Just (resourceContext resource) + else Nothing + +-- | Release all resources in the 'ResourceRegistry' without closing. +-- +-- See 'closeRegistry' for more details. +releaseAll :: + (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 + +-- | 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 +-- registry. See 'unsafeRelease' for why this is dangerous. +unsafeReleaseAll :: + (MonadMask m, MonadSTM m, MonadThread m, HasCallStack) => + ResourceRegistry m -> + m () +unsafeReleaseAll rr = do + context <- captureContext + void $ releaseAllHelper rr context unsafeRelease + +-- | Internal helper used by 'releaseAll' and 'unsafeReleaseAll'. +releaseAllHelper :: + (MonadMask m, MonadSTM m, MonadThread m) => + ResourceRegistry m -> + Context m -> + -- | How to release a resource + (ResourceKey m -> m (Maybe (Context m))) -> + m [Context m] +releaseAllHelper rr context releaser = mask_ $ do + mKeys <- updateState rr $ unlessClosed $ gets getYoungestToOldest + case mKeys of + Left closed -> throwRegistryClosed rr context closed + Right keys -> releaseResources rr keys releaser + +{------------------------------------------------------------------------------- + Threads +-------------------------------------------------------------------------------} + +-- | Thread +-- +-- The internals of this type are not exported. +data Thread m a = MonadThread m => Thread + { threadId :: !(ThreadId m) + -- ^ The underlying @async@ thread id + , threadResourceId :: !ResourceId + , threadAsync :: !(Async m a) + , threadRegistry :: !(ResourceRegistry m) + } + deriving NoThunks via OnlyCheckWhnfNamed "Thread" (Thread m a) + +-- | 'Eq' instance for 'Thread' compares 'threadId' only. +instance MonadThread m => Eq (Thread m a) where + Thread{threadId = a} == Thread{threadId = b} = a == b + +-- | Cancel a thread +-- +-- This is a synchronous operation: the thread will have terminated when this +-- function returns. +-- +-- Uses 'uninterruptibleCancel' because that's what 'withAsync' does. +cancelThread :: MonadAsync m => Thread m a -> m () +cancelThread = uninterruptibleCancel . threadAsync + +-- | Wait for thread to terminate and return its result. +-- +-- If the thread throws an exception, this will rethrow that exception. +-- +-- NOTE: If A waits on B, and B is linked to the registry, and B throws an +-- exception, then A might /either/ receive the exception thrown by B /or/ +-- the 'Control.Exception.ThreadKilled' exception thrown by the registry. +waitThread :: MonadAsync m => Thread m a -> m a +waitThread = wait . threadAsync + +-- | Lift 'waitAny' to 'Thread' +waitAnyThread :: forall m a. MonadAsync m => [Thread m a] -> m a +waitAnyThread ts = snd <$> waitAny (map threadAsync ts) + +-- | Fork a new thread +forkThread :: + forall m a. + (MonadMask m, MonadAsync m, HasCallStack) => + ResourceRegistry m -> + -- | Label for the thread + String -> + m a -> + m (Thread m a) +forkThread rr label body = + snd + <$> allocate rr (\key -> mkThread key <$> async (body' key)) cancelThread + where + mkThread :: ResourceId -> Async m a -> Thread m a + mkThread rid child = + Thread + { threadId = asyncThreadId child + , threadResourceId = rid + , threadAsync = child + , threadRegistry = rr + } + + body' :: ResourceId -> m a + body' rid = do + me <- myThreadId + labelThread me label + (registerThread me >> body) `finally` unregisterThread me rid + + -- Register the thread + -- + -- We must add the thread to the list of known threads before the thread + -- will start to use the registry. + registerThread :: ThreadId m -> m () + registerThread tid = updateState rr $ insertThread tid + + -- Unregister the thread + -- + -- Threads are the only kinds of resources that "deallocate themselves". + -- We remove the thread from the resources as well as the set of known + -- threads, so that these datastructures do not grow without bound. + -- + -- This runs with asynchronous exceptions masked (due to 'finally'), + -- though for the current implementation of 'unregisterThread' this + -- makes no difference. + unregisterThread :: ThreadId m -> ResourceId -> m () + unregisterThread tid rid = + updateState rr $ do + removeThread tid + void $ removeResource rid + +-- | Bracketed version of 'forkThread' +-- +-- The analogue of 'withAsync' for the registry. +-- +-- Scoping thread lifetime using 'withThread' is important when a parent +-- thread wants to link to a child thread /and handle any exceptions arising +-- from the link/: +-- +-- > let handleLinkException :: ExceptionInLinkedThread -> m () +-- > handleLinkException = .. +-- > in handle handleLinkException $ +-- > withThread registry codeInChild $ \child -> +-- > .. +-- +-- instead of +-- +-- > handle handleLinkException $ do -- PROBABLY NOT CORRECT! +-- > child <- forkThread registry codeInChild +-- > .. +-- +-- where the parent may exit the scope of the exception handler before the child +-- terminates. If the lifetime of the child cannot be limited to the lifetime of +-- the parent, the child should probably be linked to the registry instead and +-- the thread that spawned the registry should handle any exceptions. +-- +-- Note that in /principle/ there is no problem in using 'withAsync' alongside a +-- registry. After all, in a pattern like +-- +-- > withRegistry $ \registry -> +-- > .. +-- > withAsync (.. registry ..) $ \async -> +-- > .. +-- +-- the async will be cancelled when leaving the scope of 'withAsync' and so +-- that reference to the registry, or indeed any of the resources inside the +-- registry, is safe. However, the registry implements a sanity check that the +-- registry is only used from known threads. This is useful: when a thread that +-- is not known to the registry (in other words, whose lifetime is not tied to +-- the lifetime of the registry) spawns a resource in that registry, that +-- resource may well be deallocated before the thread terminates, leading to +-- undefined and hard to debug behaviour (indeed, whether or not this results in +-- problems may well depend on precise timing); an exception that is thrown when +-- /allocating/ the resource is (more) deterministic and easier to debug. +-- Unfortunately, it means that the above pattern is not applicable, as the +-- thread spawned by 'withAsync' is not known to the registry, and so if it were +-- to try to use the registry, the registry would throw an error (even though +-- this pattern is actually safe). This situation is not ideal, but for now we +-- merely provide an alternative to 'withAsync' that /does/ register the thread +-- with the registry. +-- +-- NOTE: Threads that are spawned out of the user's control but that must still +-- make use of the registry can use the unsafe API. This should be used with +-- caution, however. +withThread :: + (MonadMask m, MonadAsync m) => + ResourceRegistry m -> + -- | Label for the thread + String -> + m a -> + (Thread m a -> m b) -> + m b +withThread rr label body = bracket (forkThread rr label body) cancelThread + +-- | Link specified 'Thread' to the (thread that created) the registry +linkToRegistry :: (MonadAsync m, MonadFork m, MonadMask m) => Thread m a -> m () +linkToRegistry t = linkTo (registryThread $ threadRegistry t) (threadAsync t) + +-- | Fork a thread and link to it to the registry. +-- +-- This function is just a convenience. +forkLinkedThread :: + (MonadAsync m, MonadFork m, MonadMask m, HasCallStack) => + ResourceRegistry m -> + -- | Label for the thread + String -> + m a -> + m (Thread m a) +forkLinkedThread rr label body = do + t <- forkThread rr label body + -- There is no race condition here between the new thread throwing an + -- exception and the 'linkToRegistry': if the thread /already/ threw the + -- exception when we link it, the exception will be raised immediately + -- (see 'linkTo' for details). + linkToRegistry t + return t + +{------------------------------------------------------------------------------- + Check that registry is used from known thread +-------------------------------------------------------------------------------} + +ensureKnownThread :: + forall m. + (MonadThrow m, MonadThread m, MonadSTM m) => + ResourceRegistry m -> + Context m -> + m () +ensureKnownThread rr context = do + isKnown <- checkIsKnown + unless isKnown $ + throwIO $ + ResourceRegistryUsedFromUntrackedThread + { resourceRegistryCreatedIn = registryContext rr + , resourceRegistryUsedIn = context + } + where + checkIsKnown :: m Bool + checkIsKnown + | contextThreadId context == contextThreadId (registryContext rr) = + return True + | otherwise = atomically $ do + KnownThreads ts <- registryThreads <$> readTVar (registryState rr) + return $ contextThreadId context `Set.member` ts + +-- | Registry used from untracked threads +-- +-- If this exception is raised, it indicates a bug in the caller. +data ResourceRegistryThreadException + = -- | If the registry is used from an untracked thread, we cannot do proper + -- reference counting. The following threads are /tracked/: the thread + -- that spawned the registry and all threads spawned by the registry. + forall m. MonadThread m => ResourceRegistryUsedFromUntrackedThread + { resourceRegistryCreatedIn :: !(Context m) + -- ^ Information about the context in which the registry was created + , resourceRegistryUsedIn :: !(Context m) + -- ^ The context in which it was used + } + | -- | Registry closed from different threat than that created it + forall m. MonadThread m => ResourceRegistryClosedFromWrongThread + { resourceRegistryCreatedIn :: !(Context m) + -- ^ Information about the context in which the registry was created + , resourceRegistryUsedIn :: !(Context m) + -- ^ The context in which it was used + } + +deriving instance Show ResourceRegistryThreadException +instance Exception ResourceRegistryThreadException + +{------------------------------------------------------------------------------- + Auxiliary: context +-------------------------------------------------------------------------------} + +-- | The internal context of a resource registry, recording a 'PrettyCallStack' +-- of its creation and the creator's 'ThreadId' +data Context m = MonadThread m => Context + { contextCallStack :: !PrettyCallStack + -- ^ CallStack in which it was created + , contextThreadId :: !(ThreadId m) + -- ^ Thread that created the registry or resource + } + +-- Existential type; we can't use generics +instance NoThunks (Context m) where + showTypeOf _ = "Context" + wNoThunks ctxt (Context cs tid) = + allNoThunks + [ noThunks ctxt cs + , noThunks ctxt (InspectHeapNamed @"ThreadId" tid) + ] + +deriving instance Show (Context m) + +captureContext :: MonadThread m => HasCallStack => m (Context m) +captureContext = Context prettyCallStack <$> myThreadId + +{------------------------------------------------------------------------------- + Misc utilities +-------------------------------------------------------------------------------} + +-- | Generalization of 'link' that links an async to an arbitrary thread. +-- +-- Non standard (not in 'async' library) +linkTo :: + (MonadAsync m, MonadFork m, MonadMask m) => + ThreadId m -> + Async m a -> + m () +linkTo tid = linkToOnly tid (not . isCancel) + +-- | Generalization of 'linkOnly' that links an async to an arbitrary thread. +-- +-- Non standard (not in 'async' library). +linkToOnly :: + forall m a. + (MonadAsync m, MonadFork m, MonadMask m) => + ThreadId m -> + (SomeException -> Bool) -> + Async m a -> + m () +linkToOnly tid shouldThrow a = do + void $ forkRepeat ("linkToOnly " <> show linkedThreadId) $ do + r <- waitCatch a + case r of + Left e | shouldThrow e -> throwTo tid (exceptionInLinkedThread e) + _otherwise -> return () + where + linkedThreadId :: ThreadId m + linkedThreadId = asyncThreadId a + + exceptionInLinkedThread :: SomeException -> ExceptionInLinkedThread + exceptionInLinkedThread = + ExceptionInLinkedThread (show linkedThreadId) + +isCancel :: SomeException -> Bool +isCancel e + | Just AsyncCancelled <- fromException e = True + | otherwise = False + +forkRepeat :: (MonadFork m, MonadMask m) => String -> m a -> m (ThreadId m) +forkRepeat label action = + mask $ \restore -> + let go = do + r <- tryAll (restore action) + case r of + Left _ -> go + _ -> return () + in forkIO (labelThisThread label >> go) + +tryAll :: MonadCatch m => m a -> m (Either SomeException a) +tryAll = try + +mustBeRight :: Either Void a -> a +mustBeRight (Left v) = absurd v +mustBeRight (Right a) = a + +{------------------------------------------------------------------------------- + Auxiliary: CallStack with different Show instance +-------------------------------------------------------------------------------} + +-- | CallStack with 'Show' instance using 'prettyCallStack' +newtype PrettyCallStack = PrettyCallStack CallStack + deriving newtype NoThunks + +instance Show PrettyCallStack where + show (PrettyCallStack cs) = GHC.prettyCallStack cs + +-- | Capture a 'PrettyCallStack' +prettyCallStack :: HasCallStack => PrettyCallStack +prettyCallStack = PrettyCallStack GHC.callStack + +{------------------------------------------------------------------------------- + Orphan instance +-------------------------------------------------------------------------------} + +instance + (NoThunks k, NoThunks v) => + NoThunks (Bimap k v) + where wNoThunks ctxt = noThunksInKeysAndValues ctxt . Bimap.toList diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -47,9 +47,8 @@ import Data.Map.Strict qualified as Map import Data.TreeDiff import Data.Typeable -import Generics.SOP qualified as SOP import GHC.Generics (Generic, Generic1) -import Prelude +import Generics.SOP qualified as SOP import Test.QuickCheck hiding (forAll) import Test.QuickCheck.Monadic hiding (run) import Test.StateMachine @@ -60,12 +59,15 @@ import Test.Util.QSM import Test.Util.SOP import Test.Util.ToExpr () +import Prelude main :: IO () -main = defaultMain - $ testGroup "ResourceRegistry" [ - testProperty "sequential" prop_sequential - ] +main = + defaultMain $ + testGroup + "ResourceRegistry" + [ testProperty "sequential" prop_sequential + ] {------------------------------------------------------------------------------- Mock implementaton @@ -85,14 +87,14 @@ -- -- Once created, threads are never removed from this map. Instead, when they are -- killed their 'alive' status is set to 'False'. -newtype MockThreads = MTs { mockThreadsMap :: Map Int MockState } +newtype MockThreads = MTs {mockThreadsMap :: Map Int MockState} deriving (Show, Generic) -- | State of a mock thread -data MockState = MS { - alive :: Bool - , kids :: MockThreads - } +data MockState = MS + { alive :: Bool + , kids :: MockThreads + } deriving (Show, Generic) -- | All known threads, and whether or not they are alive @@ -102,69 +104,75 @@ -- computing that here. mockThreads :: MockThreads -> [(MockThread, Bool)] mockThreads = go [] True - where - go :: [Int] -> Bool -> MockThreads -> [(MockThread, Bool)] - go prefix parentAlive = - concatMap aux . Map.toList . mockThreadsMap - where - aux :: (Int, MockState) -> [(MockThread, Bool)] - aux (tid, MS{..}) = - (t, parentAlive && alive) : go t alive' kids - where - t :: [Int] - t = prefix ++ [tid] + where + go :: [Int] -> Bool -> MockThreads -> [(MockThread, Bool)] + go prefix parentAlive = + concatMap aux . Map.toList . mockThreadsMap + where + aux :: (Int, MockState) -> [(MockThread, Bool)] + aux (tid, MS{..}) = + (t, parentAlive && alive) : go t alive' kids + where + t :: [Int] + t = prefix ++ [tid] - alive' :: Bool - alive' = parentAlive && alive + alive' :: Bool + alive' = parentAlive && alive mockLiveThreads :: MockThreads -> [MockThread] mockLiveThreads = map fst . filter snd . mockThreads -alterThreadF :: forall m. MonadError Err m - => MockThread - -> (Maybe MockState -> m MockState) - -> MockThreads -> m MockThreads +alterThreadF :: + forall m. + MonadError Err m => + MockThread -> + (Maybe MockState -> m MockState) -> + MockThreads -> + m MockThreads alterThreadF [] _ _ = - error "alterThreadF: invalid thread" + error "alterThreadF: invalid thread" alterThreadF [t] f (MTs m) = - MTs <$> Map.alterF (fmap Just . f) t m -alterThreadF thread@(t:ts) f (MTs m) = - MTs <$> Map.alterF (fmap Just . f') t m - where - f' :: Maybe MockState -> m MockState - f' Nothing = throwError $ ErrInvalidThread (show thread) - f' (Just ms) = (\kids' -> ms { kids = kids' }) <$> - alterThreadF ts f (kids ms) + MTs <$> Map.alterF (fmap Just . f) t m +alterThreadF thread@(t : ts) f (MTs m) = + MTs <$> Map.alterF (fmap Just . f') t m + where + f' :: Maybe MockState -> m MockState + f' Nothing = throwError $ ErrInvalidThread (show thread) + f' (Just ms) = + (\kids' -> ms{kids = kids'}) + <$> alterThreadF ts f (kids ms) -- Create thread with the given ID mockFork :: MockThread -> MockThreads -> Except Err MockThreads mockFork t = alterThreadF t $ \case - Just _ -> error "fork: thread already exists (bug in runMock)" - Nothing -> return newState - where - newState :: MockState - newState = MS { - alive = True - , kids = MTs Map.empty - } + Just _ -> error "fork: thread already exists (bug in runMock)" + Nothing -> return newState + where + newState :: MockState + newState = + MS + { alive = True + , kids = MTs Map.empty + } mockKill :: MockThread -> MockThreads -> Except Err MockThreads mockKill t = alterThreadF t $ \case - Nothing -> throwError $ ErrInvalidThread (show t) - Just st -> return st { alive = False } + Nothing -> throwError $ ErrInvalidThread (show t) + Just st -> return st{alive = False} -data Mock = Mock { - nextId :: Int - , threads :: MockThreads - , links :: Map MockThread (Link MockThread) - } +data Mock = Mock + { nextId :: Int + , threads :: MockThreads + , links :: Map MockThread (Link MockThread) + } deriving (Show, Generic) emptyMock :: Mock -emptyMock = Mock { - nextId = 1 +emptyMock = + Mock + { nextId = 1 , threads = MTs Map.empty - , links = Map.empty + , links = Map.empty } {------------------------------------------------------------------------------- @@ -177,34 +185,30 @@ data Link a = LinkFromParent a | DontLink deriving (Show, Functor, Generic) -data Cmd t = - -- | Fork a new top-level thread +data Cmd t + = -- | Fork a new top-level thread -- -- We don't allow linking here, because we don't want an exception in one -- of these threads to kill the thread running the tests. Fork - - -- | Fork a child thread - | ForkFrom t (Link ()) - - -- | Cause a thread to terminate normally - | Terminate t - - -- | Cause a thread to terminate abnormally - | Crash t - - -- | Get all live threads - | LiveThreads + | -- | Fork a child thread + ForkFrom t (Link ()) + | -- | Cause a thread to terminate normally + Terminate t + | -- | Cause a thread to terminate abnormally + Crash t + | -- | Get all live threads + LiveThreads deriving (Show, Functor, Foldable, Traversable, Generic) -data Success t = - Unit () +data Success t + = Unit () | Spawned t | Threads [t] deriving (Show, Eq, Functor, Foldable, Traversable) -data Err = - ErrTimeout +data Err + = ErrTimeout | ErrInvalidThread String deriving (Show, Eq) @@ -219,11 +223,11 @@ normalize :: Resp MockThread -> Resp MockThread normalize (Resp r) = Resp $ aux <$> r - where - aux :: Success MockThread -> Success MockThread - aux (Unit ()) = Unit () - aux (Spawned t) = Spawned t - aux (Threads ts) = Threads (sort ts) + where + aux :: Success MockThread -> Success MockThread + aux (Unit ()) = Unit () + aux (Spawned t) = Spawned t + aux (Threads ts) = Threads (sort ts) {------------------------------------------------------------------------------- Run against the mock implementation @@ -231,51 +235,52 @@ runMock :: Cmd MockThread -> Mock -> (Resp MockThread, Mock) runMock cmd m@Mock{..} = - case runExcept (go cmd) of - Left err -> (Resp (Left err), m) - Right (success, m') -> (Resp (Right success), m') - where - go :: Cmd MockThread -> Except Err (Success MockThread, Mock) - go Fork = createThread DontLink [nextId] - go (ForkFrom t linked) = createThread (const t <$> linked) (t ++ [nextId]) - go (Terminate t) = (\x -> (Unit (), m { threads = x })) <$> mockKill t threads - go (Crash t) = (\x -> (Unit (), m { threads = x })) <$> killAll t threads - go LiveThreads = return (Threads $ mockLiveThreads threads, m) + case runExcept (go cmd) of + Left err -> (Resp (Left err), m) + Right (success, m') -> (Resp (Right success), m') + where + go :: Cmd MockThread -> Except Err (Success MockThread, Mock) + go Fork = createThread DontLink [nextId] + go (ForkFrom t linked) = createThread (const t <$> linked) (t ++ [nextId]) + go (Terminate t) = (\x -> (Unit (), m{threads = x})) <$> mockKill t threads + go (Crash t) = (\x -> (Unit (), m{threads = x})) <$> killAll t threads + go LiveThreads = return (Threads $ mockLiveThreads threads, m) - createThread :: Link MockThread -- Thread to link to (if any) - -> MockThread -> Except Err (Success MockThread, Mock) - createThread shouldLink t = do - threads' <- mockFork t threads - return ( - Spawned t - , m { nextId = succ nextId - , threads = threads' - , links = Map.insert t shouldLink links - } - ) + createThread :: + Link MockThread -> -- Thread to link to (if any) + MockThread -> + Except Err (Success MockThread, Mock) + createThread shouldLink t = do + threads' <- mockFork t threads + return + ( Spawned t + , m + { nextId = succ nextId + , threads = threads' + , links = Map.insert t shouldLink links + } + ) - killAll :: MockThread -> MockThreads -> Except Err MockThreads - killAll t = - mockKill t >=> killParent (Map.findWithDefault DontLink t links) - where - killParent :: Link MockThread -> MockThreads -> Except Err MockThreads - killParent DontLink = return - killParent (LinkFromParent t') = killAll t' + killAll :: MockThread -> MockThreads -> Except Err MockThreads + killAll t = + mockKill t >=> killParent (Map.findWithDefault DontLink t links) + where + killParent :: Link MockThread -> MockThreads -> Except Err MockThreads + killParent DontLink = return + killParent (LinkFromParent t') = killAll t' {------------------------------------------------------------------------------- Run in IO (possibly simulated) -------------------------------------------------------------------------------} -data TestThread m = TestThread { - -- | The underlying 'Thread' - testThread :: Thread m () - - -- | Parent thread this thread is linked to (if any) - , threadLinked :: Link (TestThread m) - - -- | Send the thread instructions (see 'ThreadInstr') - , threadComms :: StrictTQueue m (QueuedInstr m) - } +data TestThread m = TestThread + { testThread :: Thread m () + -- ^ The underlying 'Thread' + , threadLinked :: Link (TestThread m) + -- ^ Parent thread this thread is linked to (if any) + , threadComms :: StrictTQueue m (QueuedInstr m) + -- ^ Send the thread instructions (see 'ThreadInstr') + } -- | Instructions to a thread -- @@ -283,10 +288,8 @@ data ThreadInstr m :: Type -> Type where -- | Have the thread spawn a child thread ThreadFork :: Link () -> ThreadInstr m (TestThread m) - -- | Have the thread terminate normally ThreadTerminate :: ThreadInstr m () - -- | Raise an exception in the thread ThreadCrash :: ThreadInstr m () @@ -295,14 +298,14 @@ runInThread :: (MonadMVar m, MonadSTM m) => TestThread m -> ThreadInstr m a -> m a runInThread TestThread{..} instr = do - result <- newEmptyMVar - atomically $ writeTQueue threadComms (QueuedInstr instr result) - takeMVar result + result <- newEmptyMVar + atomically $ writeTQueue threadComms (QueuedInstr instr result) + takeMVar result -instance (MonadThread m) => Show (TestThread m) where +instance MonadThread m => Show (TestThread m) where show TestThread{..} = "" -instance (MonadThread m) => Eq (TestThread m) where +instance MonadThread m => Eq (TestThread m) where (==) = (==) `on` (threadId . testThread) -- | Create a new thread in the given registry @@ -310,61 +313,69 @@ -- In order to be able to see which threads are alive, we have threads -- register and unregister themselves. We do not reuse the registry for this, -- to avoid circular reasoning in the tests. -newThread :: forall m. (MonadMVar m, MonadMask m, MonadAsync m, MonadFork m) - => StrictTVar m [TestThread m] - -> ResourceRegistry m - -> Link (TestThread m) - -> m (TestThread m) +newThread :: + forall m. + (MonadMVar m, MonadMask m, MonadAsync m, MonadFork m) => + StrictTVar m [TestThread m] -> + ResourceRegistry m -> + Link (TestThread m) -> + m (TestThread m) newThread alive parentReg = \shouldLink -> do - comms <- atomically $ newTQueue - spawned <- newEmptyMVar + comms <- atomically $ newTQueue + spawned <- newEmptyMVar - thread <- forkThread parentReg "newThread" $ - withRegistry $ \childReg -> - threadBody childReg spawned comms - case shouldLink of - LinkFromParent _ -> linkToRegistry thread - DontLink -> return () + thread <- forkThread parentReg "newThread" $ + withRegistry $ \childReg -> + threadBody childReg spawned comms + case shouldLink of + LinkFromParent _ -> linkToRegistry thread + DontLink -> return () - let testThread :: TestThread m - testThread = TestThread { - testThread = thread - , threadLinked = shouldLink - , threadComms = comms - } + let testThread :: TestThread m + testThread = + TestThread + { testThread = thread + , threadLinked = shouldLink + , threadComms = comms + } - -- Make sure to register thread before starting it - atomically $ modifyTVar alive (testThread:) - putMVar spawned testThread - return testThread - where - threadBody :: ResourceRegistry m - -> StrictMVar m (TestThread m) - -> StrictTQueue m (QueuedInstr m) - -> m () - threadBody childReg spawned comms = do - us <- readMVar spawned - loop us `finally` (atomically $ modifyTVar alive (delete us)) - where - loop :: TestThread m -> m () - loop us = do - QueuedInstr instr result <- atomically $ readTQueue comms - case instr of - ThreadFork linked -> do - child <- newThread alive childReg (const us <$> linked) - putMVar result child - loop us - ThreadTerminate -> do - putMVar result () - ThreadCrash -> do - putMVar result () - error "crashing" + -- Make sure to register thread before starting it + atomically $ modifyTVar alive (testThread :) + putMVar spawned testThread + return testThread + where + threadBody :: + ResourceRegistry m -> + StrictMVar m (TestThread m) -> + StrictTQueue m (QueuedInstr m) -> + m () + threadBody childReg spawned comms = do + us <- readMVar spawned + loop us `finally` (atomically $ modifyTVar alive (delete us)) + where + loop :: TestThread m -> m () + loop us = do + QueuedInstr instr result <- atomically $ readTQueue comms + case instr of + ThreadFork linked -> do + child <- newThread alive childReg (const us <$> linked) + putMVar result child + loop us + ThreadTerminate -> do + putMVar result () + ThreadCrash -> do + putMVar result () + error "crashing" -runIO :: forall m. (MonadMVar m, MonadTimer m, MonadMask m, MonadAsync m, MonadFork m) - => StrictTVar m [TestThread m] - -> ResourceRegistry m - -> Cmd (TestThread m) -> m (Resp (TestThread m)) -runIO alive reg cmd = catchEx $ timeout 1 $ +runIO :: + forall m. + (MonadMVar m, MonadTimer m, MonadMask m, MonadAsync m, MonadFork m) => + StrictTVar m [TestThread m] -> + ResourceRegistry m -> + Cmd (TestThread m) -> + m (Resp (TestThread m)) +runIO alive reg cmd = catchEx $ + timeout 1 $ case cmd of Fork -> Spawned <$> newThread alive reg DontLink @@ -378,19 +389,19 @@ Unit <$> waitForTermination thread LiveThreads -> atomically $ Threads <$> readTVar alive - where - catchEx :: m (Maybe (Success a)) -> m (Resp a) - catchEx = fmap (Resp . maybe (Left ErrTimeout) Right) + where + catchEx :: m (Maybe (Success a)) -> m (Resp a) + catchEx = fmap (Resp . maybe (Left ErrTimeout) Right) - -- For the thread and all of its linked parents to have terminated - waitForTermination :: TestThread m -> m () - waitForTermination t = do - result <- try $ waitThread (testThread t) - case (result, threadLinked t) of - (Left (_ :: SomeException), LinkFromParent t') -> - waitForTermination t' - _otherwise -> - return () + -- For the thread and all of its linked parents to have terminated + waitForTermination :: TestThread m -> m () + waitForTermination t = do + result <- try $ waitThread (testThread t) + case (result, threadLinked t) of + (Left (_ :: SomeException), LinkFromParent t') -> + waitForTermination t' + _otherwise -> + return () {------------------------------------------------------------------------------- QSM wrappers @@ -398,7 +409,7 @@ newtype At m f r = At (f (Reference (TestThread m) r)) -deriving instance (MonadThread m, Show1 r) => Show (At m Cmd r) +deriving instance (MonadThread m, Show1 r) => Show (At m Cmd r) deriving instance (MonadThread m, Show1 r) => Show (At m Resp r) {------------------------------------------------------------------------------- @@ -410,8 +421,8 @@ (!) :: (Eq k, Show k) => [(k, a)] -> k -> a env ! r = case lookup r env of - Just a -> a - Nothing -> error $ "Unknown reference: " ++ show r + Just a -> a + Nothing -> error $ "Unknown reference: " ++ show r data Model m r = Model Mock (Refs m r) deriving (Show, Generic) @@ -423,72 +434,81 @@ Events -------------------------------------------------------------------------------} -toMock :: forall m f r. (Functor f, Eq1 r, Show1 r, MonadThread m) - => Model m r -> At m f r -> f MockThread +toMock :: + forall m f r. + (Functor f, Eq1 r, Show1 r, MonadThread m) => + Model m r -> At m f r -> f MockThread toMock (Model _ hs) (At fr) = (hs !) <$> fr -step :: (Eq1 r, Show1 r, MonadThread m) - => Model m r -> At m Cmd r -> (Resp MockThread, Mock) +step :: + (Eq1 r, Show1 r, MonadThread m) => + Model m r -> At m Cmd r -> (Resp MockThread, Mock) step m@(Model mock _) c = runMock (toMock m c) mock -data Event m r = Event { - before :: Model m r - , cmd :: At m Cmd r - , after :: Model m r - , mockResp :: Resp MockThread - } +data Event m r = Event + { before :: Model m r + , cmd :: At m Cmd r + , after :: Model m r + , mockResp :: Resp MockThread + } -lockstep :: (Eq1 r, Show1 r, MonadThread m) - => Model m r - -> At m Cmd r - -> At m Resp r - -> Event m r -lockstep m@(Model _ hs) c (At resp) = Event { - before = m - , cmd = c - , after = Model mock' (hs <> hs') +lockstep :: + (Eq1 r, Show1 r, MonadThread m) => + Model m r -> + At m Cmd r -> + At m Resp r -> + Event m r +lockstep m@(Model _ hs) c (At resp) = + Event + { before = m + , cmd = c + , after = Model mock' (hs <> hs') , mockResp = resp' } - where - (resp', mock') = step m c - hs' = zip (newHandles resp) (newHandles resp') + where + (resp', mock') = step m c + hs' = zip (newHandles resp) (newHandles resp') - newHandles :: Resp r -> [r] - newHandles (Resp (Left _)) = [] - newHandles (Resp (Right (Unit ()))) = [] - newHandles (Resp (Right (Spawned t))) = [t] - newHandles (Resp (Right (Threads _))) = [] + newHandles :: Resp r -> [r] + newHandles (Resp (Left _)) = [] + newHandles (Resp (Right (Unit ()))) = [] + newHandles (Resp (Right (Spawned t))) = [t] + newHandles (Resp (Right (Threads _))) = [] {------------------------------------------------------------------------------- Generator -------------------------------------------------------------------------------} generator :: forall m. Model m Symbolic -> Maybe (Gen (At m Cmd Symbolic)) -generator (Model _ hs) = Just $ oneof $ concat [ - withoutHandle - , if null hs then [] else withHandle (elements (map fst hs)) - ] - where - withoutHandle :: [Gen (At m Cmd Symbolic)] - withoutHandle = [ - fmap At $ return Fork - , fmap At $ return LiveThreads +generator (Model _ hs) = + Just $ + oneof $ + concat + [ withoutHandle + , if null hs then [] else withHandle (elements (map fst hs)) ] + where + withoutHandle :: [Gen (At m Cmd Symbolic)] + withoutHandle = + [ fmap At $ return Fork + , fmap At $ return LiveThreads + ] - withHandle :: Gen (Reference (TestThread m) Symbolic) - -> [Gen (At m Cmd Symbolic)] - withHandle pickThread = [ - fmap At $ Terminate <$> pickThread - , fmap At $ Crash <$> pickThread - , fmap At $ ForkFrom <$> pickThread <*> genLink - ] + withHandle :: + Gen (Reference (TestThread m) Symbolic) -> + [Gen (At m Cmd Symbolic)] + withHandle pickThread = + [ fmap At $ Terminate <$> pickThread + , fmap At $ Crash <$> pickThread + , fmap At $ ForkFrom <$> pickThread <*> genLink + ] - genLink :: Gen (Link ()) - genLink = aux <$> arbitrary - where - aux :: Bool -> Link () - aux True = LinkFromParent () - aux False = DontLink + genLink :: Gen (Link ()) + genLink = aux <$> arbitrary + where + aux :: Bool -> Link () + aux True = LinkFromParent () + aux False = DontLink shrinker :: Model m Symbolic -> At m Cmd Symbolic -> [At m Cmd Symbolic] shrinker _ _ = [] @@ -497,18 +517,18 @@ QSM required instances -------------------------------------------------------------------------------} -instance SOP.Generic (Cmd t) +instance SOP.Generic (Cmd t) instance SOP.HasDatatypeInfo (Cmd t) deriving instance Generic1 (At m Cmd) deriving instance Generic1 (At m Resp) instance CommandNames (At m Cmd) where - cmdName (At cmd) = constrName cmd - cmdNames _ = constrNames (Proxy @(Cmd ())) + cmdName (At cmd) = constrName cmd + cmdNames _ = constrNames (Proxy @(Cmd ())) -instance Rank2.Foldable (At m Cmd) -instance Rank2.Functor (At m Cmd) +instance Rank2.Foldable (At m Cmd) +instance Rank2.Functor (At m Cmd) instance Rank2.Traversable (At m Cmd) instance Rank2.Foldable (At m Resp) @@ -519,69 +539,78 @@ instance ToExpr (Link MockThread) instance ToExpr (Model IO Concrete) -instance (MonadThread m) => ToExpr (TestThread m) where +instance MonadThread m => ToExpr (TestThread m) where toExpr = defaultExprViaShow {------------------------------------------------------------------------------- QSM toplevel -------------------------------------------------------------------------------} -semantics :: (MonadMVar m, MonadMask m, MonadAsync m, MonadFork m, MonadTimer m, Typeable m) - => StrictTVar m [TestThread m] - -> ResourceRegistry m - -> At m Cmd Concrete -> m (At m Resp Concrete) +semantics :: + (MonadMVar m, MonadMask m, MonadAsync m, MonadFork m, MonadTimer m, Typeable m) => + StrictTVar m [TestThread m] -> + ResourceRegistry m -> + At m Cmd Concrete -> + m (At m Resp Concrete) semantics alive reg (At c) = - (At . fmap reference) <$> - runIO alive reg (concrete <$> c) + (At . fmap reference) + <$> runIO alive reg (concrete <$> c) -transition :: (Eq1 r, Show1 r, MonadThread m) - => Model m r -> At m Cmd r -> At m Resp r -> Model m r +transition :: + (Eq1 r, Show1 r, MonadThread m) => + Model m r -> At m Cmd r -> At m Resp r -> Model m r transition m c = after . lockstep m c -precondition :: forall m. (MonadThread m) - => Model m Symbolic -> At m Cmd Symbolic -> Logic +precondition :: + forall m. + MonadThread m => + Model m Symbolic -> At m Cmd Symbolic -> Logic precondition (Model mock hs) (At c) = - forAll (toList c) checkRef - where - checkRef :: Reference (TestThread m) Symbolic -> Logic - checkRef r = - case lookup r hs of - Nothing -> Bot - Just r' -> r' `member` mockLiveThreads (threads mock) + forAll (toList c) checkRef + where + checkRef :: Reference (TestThread m) Symbolic -> Logic + checkRef r = + case lookup r hs of + Nothing -> Bot + Just r' -> r' `member` mockLiveThreads (threads mock) -postcondition :: (MonadThread m) - => Model m Concrete - -> At m Cmd Concrete - -> At m Resp Concrete - -> Logic +postcondition :: + MonadThread m => + Model m Concrete -> + At m Cmd Concrete -> + At m Resp Concrete -> + Logic postcondition m c r = - normalize (toMock (after e) r) .== normalize (mockResp e) - where - e = lockstep m c r + normalize (toMock (after e) r) .== normalize (mockResp e) + where + e = lockstep m c r -symbolicResp :: (MonadThread m, Typeable m) - => Model m Symbolic - -> At m Cmd Symbolic - -> GenSym (At m Resp Symbolic) +symbolicResp :: + (MonadThread m, Typeable m) => + Model m Symbolic -> + At m Cmd Symbolic -> + GenSym (At m Resp Symbolic) symbolicResp m c = At <$> traverse (const genSym) resp - where - (resp, _mock') = step m c + where + (resp, _mock') = step m c -sm :: (MonadMVar m, MonadMask m, MonadAsync m, MonadFork m, MonadTimer m, Typeable m) - => StrictTVar m [TestThread m] - -> ResourceRegistry m - -> StateMachine (Model m) (At m Cmd) m (At m Resp) -sm alive reg = StateMachine { - initModel = initModel - , transition = transition - , precondition = precondition +sm :: + (MonadMVar m, MonadMask m, MonadAsync m, MonadFork m, MonadTimer m, Typeable m) => + StrictTVar m [TestThread m] -> + ResourceRegistry m -> + StateMachine (Model m) (At m Cmd) m (At m Resp) +sm alive reg = + StateMachine + { initModel = initModel + , transition = transition + , precondition = precondition , postcondition = postcondition - , invariant = Nothing - , generator = generator - , shrinker = shrinker - , semantics = semantics alive reg - , mock = symbolicResp - , cleanup = noCleanup + , invariant = Nothing + , generator = generator + , shrinker = shrinker + , semantics = semantics alive reg + , mock = symbolicResp + , cleanup = noCleanup } prop_sequential :: Property @@ -589,13 +618,13 @@ prop_sequential' :: QSM.Commands (At IO Cmd) (At IO Resp) -> Property prop_sequential' cmds = monadicIO $ do - alive <- liftIO $ newTVarIO [] - reg <- liftIO $ unsafeNewRegistry - let sm' = sm alive reg - (hist, _model, res) <- runCommands sm' cmds - prettyCommands sm' hist - $ checkCommandNames cmds - $ res === Ok + alive <- liftIO $ newTVarIO [] + reg <- liftIO $ unsafeNewRegistry + let sm' = sm alive reg + (hist, _model, res) <- runCommands sm' cmds + prettyCommands sm' hist $ + checkCommandNames cmds $ + res === Ok unused :: a unused = error "not used during command generation" @@ -606,24 +635,24 @@ _forkCount :: QSM.Commands (At IO Cmd) (At IO Resp) _forkCount = example (sm unused unused) $ do - run' $ At $ Fork - run' $ At $ LiveThreads + run' $ At $ Fork + run' $ At $ LiveThreads _forkKillCount :: QSM.Commands (At IO Cmd) (At IO Resp) _forkKillCount = example (sm unused unused) $ do - [t] <- run $ At $ Fork - run' $ At $ Terminate t - run' $ At $ LiveThreads + [t] <- run $ At $ Fork + run' $ At $ Terminate t + run' $ At $ LiveThreads _forkFromKillCount :: QSM.Commands (At IO Cmd) (At IO Resp) _forkFromKillCount = example (sm unused unused) $ do - [t] <- run $ At $ Fork - run' $ At $ ForkFrom t DontLink - run' $ At $ Terminate t - run' $ At $ LiveThreads + [t] <- run $ At $ Fork + run' $ At $ ForkFrom t DontLink + run' $ At $ Terminate t + run' $ At $ LiveThreads _invalidForkFrom :: QSM.Commands (At IO Cmd) (At IO Resp) _invalidForkFrom = example (sm unused unused) $ do - [t] <- run $ At $ Fork - run' $ At $ Terminate t - run' $ At $ ForkFrom t DontLink + [t] <- run $ At $ Fork + run' $ At $ Terminate t + run' $ At $ ForkFrom t DontLink diff --git a/test/Test/Util/QSM.hs b/test/Test/Util/QSM.hs --- a/test/Test/Util/QSM.hs +++ b/test/Test/Util/QSM.hs @@ -1,8 +1,8 @@ {-# LANGUAGE ScopedTypeVariables #-} -module Test.Util.QSM ( - Example - -- opaque +module Test.Util.QSM + ( Example + -- opaque , example , run , run' @@ -16,8 +16,8 @@ import Test.StateMachine.Types import Test.StateMachine.Types.Rank2 qualified as Rank2 -data Example cmd a = - Done a +data Example cmd a + = Done a | Run (cmd Symbolic) ([Var] -> Example cmd a) | Fail String @@ -25,13 +25,13 @@ fmap = liftM instance Applicative (Example cmd) where - pure = Done + pure = Done (<*>) = ap instance Monad (Example cmd) where - return = pure - Done a >>= f = f a - Run c k >>= f = Run c (k >=> f) + return = pure + Done a >>= f = f a + Run c k >>= f = Run c (k >=> f) Fail err >>= _ = Fail err instance Fail.MonadFail (Example cmd) where @@ -45,30 +45,32 @@ run' :: cmd Symbolic -> Example cmd () run' cmd = Run cmd (\_vars -> Done ()) -example :: forall model cmd m resp. (Rank2.Foldable resp, Show (cmd Symbolic)) - => StateMachine model cmd m resp - -> Example cmd () - -> Commands cmd resp +example :: + forall model cmd m resp. + (Rank2.Foldable resp, Show (cmd Symbolic)) => + StateMachine model cmd m resp -> + Example cmd () -> + Commands cmd resp example sm = - Commands . fst . flip runGenSym newCounter . go (initModel sm) - where - go :: model Symbolic -> Example cmd () -> GenSym [Command cmd resp] - go _ (Done ()) = return [] - go _ (Fail err) = error $ "example: " ++ err - go m (Run cmd k) = do - case Logic.logic (precondition sm m cmd) of - Logic.VFalse counterexample -> - error $ "Invalid command " ++ show cmd ++ ": " ++ show counterexample - Logic.VTrue -> do - resp <- mock sm m cmd + Commands . fst . flip runGenSym newCounter . go (initModel sm) + where + go :: model Symbolic -> Example cmd () -> GenSym [Command cmd resp] + go _ (Done ()) = return [] + go _ (Fail err) = error $ "example: " ++ err + go m (Run cmd k) = do + case Logic.logic (precondition sm m cmd) of + Logic.VFalse counterexample -> + error $ "Invalid command " ++ show cmd ++ ": " ++ show counterexample + Logic.VTrue -> do + resp <- mock sm m cmd - let m' :: model Symbolic - m' = transition sm m cmd resp + let m' :: model Symbolic + m' = transition sm m cmd resp - vars :: [Var] - vars = getUsedVars resp + vars :: [Var] + vars = getUsedVars resp - cmd' :: Command cmd resp - cmd' = Command cmd resp vars + cmd' :: Command cmd resp + cmd' = Command cmd resp vars - (cmd' :) <$> go m' (k vars) + (cmd' :) <$> go m' (k vars) diff --git a/test/Test/Util/SOP.hs b/test/Test/Util/SOP.hs --- a/test/Test/Util/SOP.hs +++ b/test/Test/Util/SOP.hs @@ -1,31 +1,32 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -module Test.Util.SOP ( - constrName +module Test.Util.SOP + ( constrName , constrNames ) where import Data.Proxy import Generics.SOP qualified as SOP -constrInfo :: SOP.HasDatatypeInfo a - => proxy a - -> SOP.NP SOP.ConstructorInfo (SOP.Code a) +constrInfo :: + SOP.HasDatatypeInfo a => + proxy a -> + SOP.NP SOP.ConstructorInfo (SOP.Code a) constrInfo = SOP.constructorInfo . SOP.datatypeInfo constrName :: forall a. SOP.HasDatatypeInfo a => a -> String constrName a = - SOP.hcollapse $ SOP.hliftA2 go (constrInfo p) (SOP.unSOP (SOP.from a)) - where - go :: SOP.ConstructorInfo b -> SOP.NP SOP.I b -> SOP.K String b - go nfo _ = SOP.K $ SOP.constructorName nfo + SOP.hcollapse $ SOP.hliftA2 go (constrInfo p) (SOP.unSOP (SOP.from a)) + where + go :: SOP.ConstructorInfo b -> SOP.NP SOP.I b -> SOP.K String b + go nfo _ = SOP.K $ SOP.constructorName nfo - p = Proxy @a + p = Proxy @a constrNames :: SOP.HasDatatypeInfo a => proxy a -> [String] constrNames p = - SOP.hcollapse $ SOP.hmap go (constrInfo p) - where - go :: SOP.ConstructorInfo a -> SOP.K String a - go nfo = SOP.K $ SOP.constructorName nfo + SOP.hcollapse $ SOP.hmap go (constrInfo p) + where + go :: SOP.ConstructorInfo a -> SOP.K String a + go nfo = SOP.K $ SOP.constructorName nfo diff --git a/test/Test/Util/ToExpr.hs b/test/Test/Util/ToExpr.hs --- a/test/Test/Util/ToExpr.hs +++ b/test/Test/Util/ToExpr.hs @@ -1,7 +1,6 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} - {-# OPTIONS_GHC -Wno-orphans #-} -- | This module implements QSM's @CanDiff@ typeclass using @tree-diff@'s @@ -14,14 +13,14 @@ import Test.StateMachine.Types.References qualified as QSM instance ToExpr x => CanDiff x where - type ADiff x = Edit EditExpr + type ADiff x = Edit EditExpr type AnExpr x = Expr - toDiff = toExpr - exprDiff _ = T.exprDiff + toDiff = toExpr + exprDiff _ = T.exprDiff diffToDocCompact _ = ansiWlBgEditExprCompact - diffToDoc _ = ansiWlBgEditExpr - exprToDoc _ = ansiWlBgExpr + diffToDoc _ = ansiWlBgEditExpr + exprToDoc _ = ansiWlBgExpr {------------------------------------------------------------------------------- QSM's References instances