diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history of strict-checked-vars + +## 0.1.0.0 — 2024-10-22 + +- First release, extracted from [`ouroboros-consensus`](https://github.com/IntersectMBO/ouroboros-consensus). diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/NOTICE b/NOTICE new file mode 100644 --- /dev/null +++ b/NOTICE @@ -0,0 +1,13 @@ +Copyright 2019-2023 Input Output Global Inc (IOG), INTERSECT 2023-2024. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,7 @@ +# `resource-registry` + +In some cases, the lifetime of a resource is not suitable for CPS and hence +can't use `bracket`-like functions, for example, the resource might be tracked +inside some other data structure. In this case, the container data structure can +be allocated in a resource registry together with the resources so that an +exception will deallocate the resources in a proper order. diff --git a/resource-registry.cabal b/resource-registry.cabal new file mode 100644 --- /dev/null +++ b/resource-registry.cabal @@ -0,0 +1,95 @@ +cabal-version: 3.0 +name: resource-registry +version: 0.1.0.0 +synopsis: Track allocated resources +description: + When the scope of a @bracket@ doesn't enclose all uses of the resource, a + 'ResourceRegistry' can be used instead to capture the lifetime of those + resources. + +license: Apache-2.0 +license-files: + LICENSE + NOTICE + +author: IOG Engineering Team +maintainer: operations@iohk.io +copyright: + 2019-2023 Input Output Global Inc (IOG) + 2023-2024 INTERSECT + +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 +extra-doc-files: + CHANGELOG.md + README.md + +source-repository head + type: git + location: https://github.com/IntersectMBO/io-classes-extra + subdir: resource-registry + +source-repository this + type: git + location: https://github.com/IntersectMBO/io-classes-extra + subdir: resource-registry + tag: resource-registry-0.1.0.0 + +common warnings + ghc-options: + -Wall + -Wcompat + -Wincomplete-uni-patterns + -Wincomplete-record-updates + -Wpartial-fields + -Widentities + -Wredundant-constraints + -Wmissing-export-lists + -Wunused-packages + -Wno-unticked-promoted-constructors + +library + import: warnings + exposed-modules: Control.ResourceRegistry + build-depends: + base >=4.14 && <4.21, + bimap ^>=0.5, + containers >=0.6 && <0.8, + io-classes ^>=1.5, + mtl >=2.2 && <2.4, + nothunks ^>=0.2, + strict-stm ^>=1.5, + + hs-source-dirs: src + default-language: Haskell2010 + default-extensions: ImportQualifiedPost + +test-suite resource-registry-test + import: warnings + default-language: Haskell2010 + default-extensions: ImportQualifiedPost + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + other-modules: + Test.Util.QSM + Test.Util.SOP + Test.Util.ToExpr + + build-depends: + QuickCheck, + base, + containers, + generics-sop, + io-classes, + 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 new file mode 100644 --- /dev/null +++ b/src/Control/ResourceRegistry.hs @@ -0,0 +1,1497 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveAnyClass #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DerivingVia #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE QuantifiedConstraints #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# 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 + wNoThunks ctxt = noThunksInKeysAndValues ctxt . Bimap.toList diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,629 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} + +-- | Tests for the resource registry +-- +-- The tests for the registry are model based. The model records which resources +-- we expect to be alive and which we expect to have been deallocated. The only +-- resources we are modelling here are threads; the commands we then execute are +-- +-- * Fork a thread from some other thread +-- * Terminate a thread +-- * Have a thread crash +-- * Collect all live threads +-- +-- We then verify that the resource registry behaves like the model, cleaning +-- up resources as threads terminate or crash. +module Main (main) where + +import Control.Concurrent.Class.MonadMVar.Strict +import Control.Concurrent.Class.MonadSTM.Strict +import Control.Monad +import Control.Monad.Class.MonadAsync +import Control.Monad.Class.MonadFork +import Control.Monad.Class.MonadThrow +import Control.Monad.Class.MonadTimer.SI +import Control.Monad.Except +#if __GLASGOW_HASKELL__ > 904 +import Control.Monad.IO.Class +#endif +import Control.ResourceRegistry +import Data.Foldable +import Data.Function +import Data.Functor.Classes +import Data.Kind +import Data.List (delete, sort) +import Data.Map.Strict (Map) +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 Test.QuickCheck hiding (forAll) +import Test.QuickCheck.Monadic hiding (run) +import Test.StateMachine +import Test.StateMachine.Types qualified as QSM +import Test.StateMachine.Types.Rank2 qualified as Rank2 +import Test.Tasty hiding (after) +import Test.Tasty.QuickCheck hiding (forAll) +import Test.Util.QSM +import Test.Util.SOP +import Test.Util.ToExpr () + +main :: IO () +main = defaultMain + $ testGroup "ResourceRegistry" [ + testProperty "sequential" prop_sequential + ] + +{------------------------------------------------------------------------------- + Mock implementaton +-------------------------------------------------------------------------------} + +-- | Mock thread IDs record thread pedigree +-- +-- > [t] top-level thread t +-- > [t, t'] child t' of top-level thread t +-- > [t, t', t''] child t'' of thread t', itself a child of t +-- +-- NOTE: All thread IDs will be unique. If both threads 1 and 2 spawn a child, +-- they would be @[1,3]@ and @[2,4]@. +type MockThread = [Int] + +-- | Threads and their subthreads +-- +-- 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 } + deriving (Show, Generic) + +-- | State of a mock thread +data MockState = MS { + alive :: Bool + , kids :: MockThreads + } + deriving (Show, Generic) + +-- | All known threads, and whether or not they are alive +-- +-- TODO: Perhaps it would be better to have an invariant that in 'MockThreads' +-- threads must be recorded as dead if any of their parents are, rather than +-- 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] + + 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 [] _ _ = + 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) + +-- 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 + } + +mockKill :: MockThread -> MockThreads -> Except Err MockThreads +mockKill t = alterThreadF t $ \case + Nothing -> throwError $ ErrInvalidThread (show t) + Just st -> return st { alive = False } + +data Mock = Mock { + nextId :: Int + , threads :: MockThreads + , links :: Map MockThread (Link MockThread) + } + deriving (Show, Generic) + +emptyMock :: Mock +emptyMock = Mock { + nextId = 1 + , threads = MTs Map.empty + , links = Map.empty + } + +{------------------------------------------------------------------------------- + Commands +-------------------------------------------------------------------------------} + +-- | Should we link a new thread to its parent? +-- +-- If yes, we need some information about that parent. +data Link a = LinkFromParent a | DontLink + deriving (Show, Functor, Generic) + +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 + deriving (Show, Functor, Foldable, Traversable, Generic) + +data Success t = + Unit () + | Spawned t + | Threads [t] + deriving (Show, Eq, Functor, Foldable, Traversable) + +data Err = + ErrTimeout + | ErrInvalidThread String + deriving (Show, Eq) + +newtype Resp t = Resp (Either Err (Success t)) + deriving (Show, Eq, Functor, Foldable, Traversable) + +{------------------------------------------------------------------------------- + "Up-to" comparison for responses + + This is used in 'postcondition'. +-------------------------------------------------------------------------------} + +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) + +{------------------------------------------------------------------------------- + Run against the mock implementation +-------------------------------------------------------------------------------} + +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) + + 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' + +{------------------------------------------------------------------------------- + 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) + } + +-- | Instructions to a thread +-- +-- The type argument indicates the result of the instruction +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 () + +-- | Instruction along with an MVar for the result +data QueuedInstr m = forall a. QueuedInstr (ThreadInstr m a) (StrictMVar m a) + +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 + +instance (MonadThread m) => Show (TestThread m) where + show TestThread{..} = "" + +instance (MonadThread m) => Eq (TestThread m) where + (==) = (==) `on` (threadId . testThread) + +-- | Create a new thread in the given registry +-- +-- 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 alive parentReg = \shouldLink -> do + comms <- atomically $ newTQueue + spawned <- newEmptyMVar + + 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 + } + + -- 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 $ + case cmd of + Fork -> + Spawned <$> newThread alive reg DontLink + ForkFrom thread shouldLink -> do + Spawned <$> runInThread thread (ThreadFork shouldLink) + Terminate thread -> do + runInThread thread ThreadTerminate + Unit <$> waitForTermination thread + Crash thread -> do + runInThread thread ThreadCrash + Unit <$> waitForTermination thread + LiveThreads -> + atomically $ Threads <$> readTVar alive + 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 () + +{------------------------------------------------------------------------------- + QSM wrappers +-------------------------------------------------------------------------------} + +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 Resp r) + +{------------------------------------------------------------------------------- + Relate model to IO +-------------------------------------------------------------------------------} + +-- TODO: Use RefEnv? +type Refs m r = [(Reference (TestThread m) r, MockThread)] + +(!) :: (Eq k, Show k) => [(k, a)] -> k -> a +env ! r = case lookup r env of + Just a -> a + Nothing -> error $ "Unknown reference: " ++ show r + +data Model m r = Model Mock (Refs m r) + deriving (Show, Generic) + +initModel :: Model m r +initModel = Model emptyMock [] + +{------------------------------------------------------------------------------- + Events +-------------------------------------------------------------------------------} + +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 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 + } + +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') + + 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 + ] + + 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 + +shrinker :: Model m Symbolic -> At m Cmd Symbolic -> [At m Cmd Symbolic] +shrinker _ _ = [] + +{------------------------------------------------------------------------------- + QSM required instances +-------------------------------------------------------------------------------} + +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 ())) + +instance Rank2.Foldable (At m Cmd) +instance Rank2.Functor (At m Cmd) +instance Rank2.Traversable (At m Cmd) + +instance Rank2.Foldable (At m Resp) + +instance ToExpr MockState +instance ToExpr MockThreads +instance ToExpr Mock +instance ToExpr (Link MockThread) +instance ToExpr (Model IO Concrete) + +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 alive reg (At 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 m c = after . lockstep m c + +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) + +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 + +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 + +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 + } + +prop_sequential :: Property +prop_sequential = forAllCommands (sm unused unused) Nothing prop_sequential' + +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 + +unused :: a +unused = error "not used during command generation" + +{------------------------------------------------------------------------------- + For running things from ghci +-------------------------------------------------------------------------------} + +_forkCount :: QSM.Commands (At IO Cmd) (At IO Resp) +_forkCount = example (sm unused unused) $ do + 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 + +_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 + +_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 diff --git a/test/Test/Util/QSM.hs b/test/Test/Util/QSM.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/QSM.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module Test.Util.QSM ( + Example + -- opaque + , example + , run + , run' + ) where + +import Control.Monad +import Control.Monad.Fail qualified as Fail +import Data.Typeable +import Test.StateMachine.Logic qualified as Logic +import Test.StateMachine.Sequential +import Test.StateMachine.Types +import Test.StateMachine.Types.Rank2 qualified as Rank2 + +data Example cmd a = + Done a + | Run (cmd Symbolic) ([Var] -> Example cmd a) + | Fail String + +instance Functor (Example cmd) where + fmap = liftM + +instance Applicative (Example cmd) where + pure = Done + (<*>) = ap + +instance Monad (Example cmd) where + 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 + fail = Fail + +-- | Run a command, and capture its references +run :: Typeable a => cmd Symbolic -> Example cmd [Reference a Symbolic] +run cmd = Run cmd (Done . map (Reference . Symbolic)) + +-- | Run a command, ignoring its references +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 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 + + let m' :: model Symbolic + m' = transition sm m cmd resp + + vars :: [Var] + vars = getUsedVars resp + + cmd' :: Command cmd resp + cmd' = Command cmd resp vars + + (cmd' :) <$> go m' (k vars) diff --git a/test/Test/Util/SOP.hs b/test/Test/Util/SOP.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/SOP.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +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.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 + + 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 diff --git a/test/Test/Util/ToExpr.hs b/test/Test/Util/ToExpr.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/ToExpr.hs @@ -0,0 +1,36 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +{-# OPTIONS_GHC -Wno-orphans #-} + +-- | This module implements QSM's @CanDiff@ typeclass using @tree-diff@'s +-- @ToExpr@. +module Test.Util.ToExpr () where + +import Data.TreeDiff as T +import Test.StateMachine qualified as QSM +import Test.StateMachine.Diffing (CanDiff (..)) +import Test.StateMachine.Types.References qualified as QSM + +instance ToExpr x => CanDiff x where + type ADiff x = Edit EditExpr + type AnExpr x = Expr + + toDiff = toExpr + exprDiff _ = T.exprDiff + diffToDocCompact _ = ansiWlBgEditExprCompact + diffToDoc _ = ansiWlBgEditExpr + exprToDoc _ = ansiWlBgExpr + +{------------------------------------------------------------------------------- + QSM's References instances +-------------------------------------------------------------------------------} + +instance ToExpr (r k) => ToExpr (QSM.Reference k r) + +instance ToExpr a => ToExpr (QSM.Concrete a) where + toExpr (QSM.Concrete x) = toExpr x + +instance ToExpr (QSM.Opaque a) where + toExpr _ = App "Opaque" []