lsm-tree (empty) → 1.0.0.0
raw patch · 174 files changed
+56989/−0 lines, 174 filesdep +QuickCheckdep +ansi-terminaldep +async
Dependencies added: QuickCheck, ansi-terminal, async, barbies, base, binary, bitvec, blockio, bloomfilter-blocked, bytestring, cborg, clock, constraints, containers, contra-tracer, crc32c, criterion, cryptohash-sha256, data-elevator, deepseq, directory, filepath, fs-api, fs-sim, heaps, indexed-traversable, io-classes, io-sim, lsm-tree, mtl, nonempty-containers, nothunks, optparse-applicative, pretty-show, primes, primitive, quickcheck-classes, quickcheck-dynamic, quickcheck-instances, quickcheck-lockstep, random, safe-wild-cards, semialign, serialise, split, splitmix, tasty, tasty-bench, tasty-golden, tasty-hunit, tasty-quickcheck, temporary, text, these, time, transformers, utf8-string, vector, vector-algorithms, wide-word
Files
- CHANGELOG.md +5/−0
- LICENSE +201/−0
- NOTICE +13/−0
- app/Database/LSMTree/Demo.hs +221/−0
- app/Main.hs +15/−0
- bench-unions/Bench/Unions.hs +568/−0
- bench-unions/Main.hs +6/−0
- bench/macro/lsm-tree-bench-bloomfilter.hs +286/−0
- bench/macro/lsm-tree-bench-lookups.hs +572/−0
- bench/macro/utxo-bench.hs +1125/−0
- bench/macro/utxo-rocksdb-bench.hs +423/−0
- bench/micro/Bench/Database/LSMTree.hs +469/−0
- bench/micro/Bench/Database/LSMTree/Internal/BloomFilter.hs +93/−0
- bench/micro/Bench/Database/LSMTree/Internal/Index.hs +121/−0
- bench/micro/Bench/Database/LSMTree/Internal/Index/Compact.hs +154/−0
- bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs +281/−0
- bench/micro/Bench/Database/LSMTree/Internal/Merge.hs +453/−0
- bench/micro/Bench/Database/LSMTree/Internal/RawPage.hs +68/−0
- bench/micro/Bench/Database/LSMTree/Internal/Serialise.hs +20/−0
- bench/micro/Bench/Database/LSMTree/Internal/WriteBuffer.hs +224/−0
- bench/micro/Main.hs +32/−0
- lsm-tree.cabal +1180/−0
- src-control/Control/ActionRegistry.hs +644/−0
- src-control/Control/Concurrent/Class/MonadSTM/RWVar.hs +179/−0
- src-control/Control/RefCount.hs +621/−0
- src-core/Database/LSMTree/Internal/Arena.hs +252/−0
- src-core/Database/LSMTree/Internal/Assertions.hs +54/−0
- src-core/Database/LSMTree/Internal/BitMath.hs +189/−0
- src-core/Database/LSMTree/Internal/BlobFile.hs +129/−0
- src-core/Database/LSMTree/Internal/BlobRef.hs +221/−0
- src-core/Database/LSMTree/Internal/BloomFilter.hs +326/−0
- src-core/Database/LSMTree/Internal/ByteString.hs +138/−0
- src-core/Database/LSMTree/Internal/CRC32C.hs +501/−0
- src-core/Database/LSMTree/Internal/ChecksumHandle.hs +250/−0
- src-core/Database/LSMTree/Internal/Chunk.hs +133/−0
- src-core/Database/LSMTree/Internal/Config.hs +443/−0
- src-core/Database/LSMTree/Internal/Config/Override.hs +205/−0
- src-core/Database/LSMTree/Internal/Cursor.hs +131/−0
- src-core/Database/LSMTree/Internal/Entry.hs +138/−0
- src-core/Database/LSMTree/Internal/IncomingRun.hs +369/−0
- src-core/Database/LSMTree/Internal/Index.hs +211/−0
- src-core/Database/LSMTree/Internal/Index/Compact.hs +743/−0
- src-core/Database/LSMTree/Internal/Index/CompactAcc.hs +295/−0
- src-core/Database/LSMTree/Internal/Index/Ordinary.hs +244/−0
- src-core/Database/LSMTree/Internal/Index/OrdinaryAcc.hs +135/−0
- src-core/Database/LSMTree/Internal/Lookup.hs +344/−0
- src-core/Database/LSMTree/Internal/Map/Range.hs +91/−0
- src-core/Database/LSMTree/Internal/Merge.hs +500/−0
- src-core/Database/LSMTree/Internal/MergeSchedule.hs +998/−0
- src-core/Database/LSMTree/Internal/MergingRun.hs +1031/−0
- src-core/Database/LSMTree/Internal/MergingTree.hs +549/−0
- src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs +141/−0
- src-core/Database/LSMTree/Internal/Page.hs +67/−0
- src-core/Database/LSMTree/Internal/PageAcc.hs +404/−0
- src-core/Database/LSMTree/Internal/PageAcc1.hs +137/−0
- src-core/Database/LSMTree/Internal/Paths.hs +392/−0
- src-core/Database/LSMTree/Internal/Primitive.hs +180/−0
- src-core/Database/LSMTree/Internal/Range.hs +28/−0
- src-core/Database/LSMTree/Internal/RawBytes.hs +312/−0
- src-core/Database/LSMTree/Internal/RawOverflowPage.hs +173/−0
- src-core/Database/LSMTree/Internal/RawPage.hs +474/−0
- src-core/Database/LSMTree/Internal/Readers.hs +441/−0
- src-core/Database/LSMTree/Internal/Run.hs +351/−0
- src-core/Database/LSMTree/Internal/RunAcc.hs +355/−0
- src-core/Database/LSMTree/Internal/RunBuilder.hs +250/−0
- src-core/Database/LSMTree/Internal/RunNumber.hs +21/−0
- src-core/Database/LSMTree/Internal/RunReader.hs +345/−0
- src-core/Database/LSMTree/Internal/Serialise.hs +184/−0
- src-core/Database/LSMTree/Internal/Serialise/Class.hs +536/−0
- src-core/Database/LSMTree/Internal/Snapshot.hs +809/−0
- src-core/Database/LSMTree/Internal/Snapshot/Codec.hs +809/−0
- src-core/Database/LSMTree/Internal/StrictArray.hs +60/−0
- src-core/Database/LSMTree/Internal/Types.hs +231/−0
- src-core/Database/LSMTree/Internal/UniqCounter.hs +51/−0
- src-core/Database/LSMTree/Internal/Unsafe.hs +2229/−0
- src-core/Database/LSMTree/Internal/Unsliced.hs +82/−0
- src-core/Database/LSMTree/Internal/Vector.hs +138/−0
- src-core/Database/LSMTree/Internal/Vector/Growing.hs +131/−0
- src-core/Database/LSMTree/Internal/WriteBuffer.hs +132/−0
- src-core/Database/LSMTree/Internal/WriteBufferBlobs.hs +270/−0
- src-core/Database/LSMTree/Internal/WriteBufferReader.hs +190/−0
- src-core/Database/LSMTree/Internal/WriteBufferWriter.hs +242/−0
- src-extras/Database/LSMTree/Extras.hs +59/−0
- src-extras/Database/LSMTree/Extras/Generators.hs +573/−0
- src-extras/Database/LSMTree/Extras/Index.hs +101/−0
- src-extras/Database/LSMTree/Extras/MergingRunData.hs +215/−0
- src-extras/Database/LSMTree/Extras/MergingTreeData.hs +431/−0
- src-extras/Database/LSMTree/Extras/NoThunks.hs +903/−0
- src-extras/Database/LSMTree/Extras/Orphans.hs +176/−0
- src-extras/Database/LSMTree/Extras/Random.hs +120/−0
- src-extras/Database/LSMTree/Extras/ReferenceImpl.hs +310/−0
- src-extras/Database/LSMTree/Extras/RunData.hs +350/−0
- src-extras/Database/LSMTree/Extras/UTxO.hs +150/−0
- src-kmerge/KMerge/Heap.hs +180/−0
- src-kmerge/KMerge/LoserTree.hs +206/−0
- src-mcg/MCG.hs +105/−0
- src-prototypes/FormatPage.hs +884/−0
- src-prototypes/ScheduledMerges.hs +1981/−0
- src-rocksdb/RocksDB.hs +313/−0
- src-rocksdb/RocksDB/FFI.hs +245/−0
- src/Database/LSMTree.hs +2819/−0
- src/Database/LSMTree/Simple.hs +1634/−0
- test-control/Main.hs +13/−0
- test-control/Test/Control/ActionRegistry.hs +45/−0
- test-control/Test/Control/Concurrent/Class/MonadSTM/RWVar.hs +84/−0
- test-control/Test/Control/RefCount.hs +252/−0
- test-prototypes/Main.hs +16/−0
- test-prototypes/Test/FormatPage.hs +229/−0
- test-prototypes/Test/ScheduledMerges.hs +576/−0
- test-prototypes/Test/ScheduledMerges/RunSizes.hs +118/−0
- test-prototypes/Test/ScheduledMergesQLS.hs +463/−0
- test/Database/LSMTree/Class.hs +299/−0
- test/Database/LSMTree/Class/Common.hs +94/−0
- test/Database/LSMTree/Model.hs +21/−0
- test/Database/LSMTree/Model/IO.hs +100/−0
- test/Database/LSMTree/Model/Session.hs +875/−0
- test/Database/LSMTree/Model/Table.hs +323/−0
- test/Main.hs +113/−0
- test/Test/Database/LSMTree.hs +314/−0
- test/Test/Database/LSMTree/Class.hs +737/−0
- test/Test/Database/LSMTree/Generators.hs +200/−0
- test/Test/Database/LSMTree/Internal.hs +122/−0
- test/Test/Database/LSMTree/Internal/Arena.hs +36/−0
- test/Test/Database/LSMTree/Internal/BlobFile/FS.hs +67/−0
- test/Test/Database/LSMTree/Internal/BloomFilter.hs +170/−0
- test/Test/Database/LSMTree/Internal/CRC32C.hs +172/−0
- test/Test/Database/LSMTree/Internal/Chunk.hs +168/−0
- test/Test/Database/LSMTree/Internal/Entry.hs +221/−0
- test/Test/Database/LSMTree/Internal/Index/Compact.hs +492/−0
- test/Test/Database/LSMTree/Internal/Index/Ordinary.hs +569/−0
- test/Test/Database/LSMTree/Internal/Lookup.hs +606/−0
- test/Test/Database/LSMTree/Internal/Merge.hs +238/−0
- test/Test/Database/LSMTree/Internal/MergingRun.hs +53/−0
- test/Test/Database/LSMTree/Internal/MergingTree.hs +278/−0
- test/Test/Database/LSMTree/Internal/PageAcc.hs +206/−0
- test/Test/Database/LSMTree/Internal/PageAcc1.hs +49/−0
- test/Test/Database/LSMTree/Internal/RawBytes.hs +120/−0
- test/Test/Database/LSMTree/Internal/RawOverflowPage.hs +57/−0
- test/Test/Database/LSMTree/Internal/RawPage.hs +537/−0
- test/Test/Database/LSMTree/Internal/Readers.hs +487/−0
- test/Test/Database/LSMTree/Internal/Run.hs +296/−0
- test/Test/Database/LSMTree/Internal/RunAcc.hs +158/−0
- test/Test/Database/LSMTree/Internal/RunBloomFilterAlloc.hs +304/−0
- test/Test/Database/LSMTree/Internal/RunBuilder.hs +95/−0
- test/Test/Database/LSMTree/Internal/RunReader.hs +201/−0
- test/Test/Database/LSMTree/Internal/Serialise.hs +37/−0
- test/Test/Database/LSMTree/Internal/Serialise/Class.hs +130/−0
- test/Test/Database/LSMTree/Internal/Snapshot/Codec.hs +473/−0
- test/Test/Database/LSMTree/Internal/Snapshot/Codec/Golden.hs +528/−0
- test/Test/Database/LSMTree/Internal/Snapshot/FS.hs +230/−0
- test/Test/Database/LSMTree/Internal/Unsliced.hs +51/−0
- test/Test/Database/LSMTree/Internal/Vector.hs +120/−0
- test/Test/Database/LSMTree/Internal/Vector/Growing.hs +222/−0
- test/Test/Database/LSMTree/Internal/WriteBufferBlobs/FS.hs +87/−0
- test/Test/Database/LSMTree/Internal/WriteBufferReader/FS.hs +76/−0
- test/Test/Database/LSMTree/Model/Table.hs +91/−0
- test/Test/Database/LSMTree/Resolve.hs +40/−0
- test/Test/Database/LSMTree/StateMachine.hs +2937/−0
- test/Test/Database/LSMTree/StateMachine/DL.hs +295/−0
- test/Test/Database/LSMTree/StateMachine/Op.hs +134/−0
- test/Test/Database/LSMTree/Tracer/Golden.hs +159/−0
- test/Test/Database/LSMTree/UnitTests.hs +292/−0
- test/Test/FS.hs +210/−0
- test/Test/Util/Arbitrary.hs +46/−0
- test/Test/Util/FS.hs +634/−0
- test/Test/Util/FS/Error.hs +226/−0
- test/Test/Util/Orphans.hs +23/−0
- test/Test/Util/PrettyProxy.hs +19/−0
- test/Test/Util/QC.hs +48/−0
- test/Test/Util/QLS.hs +83/−0
- test/Test/Util/RawPage.hs +103/−0
- test/Test/Util/TypeFamilyWrappers.hs +62/−0
- test/kmerge-test.hs +468/−0
- test/map-range-test.hs +77/−0
@@ -0,0 +1,5 @@+# Revision history for `lsm-tree`++## 1.0.0.0 -- 2025-08-06++* First released version.
@@ -0,0 +1,201 @@+ 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.
@@ -0,0 +1,13 @@+Copyright (c) 2023-2025 Cardano Development Foundation++ 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.
@@ -0,0 +1,221 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE OverloadedLists #-}+{-# LANGUAGE OverloadedStrings #-}+{-# OPTIONS_GHC -Wno-orphans #-}+{-# OPTIONS_GHC -Wno-unused-do-bind #-}++{- HLINT ignore "Redundant pure" -}++module Database.LSMTree.Demo (demo) where++import Control.Exception (SomeException, try)+import Control.Monad (when)+import Control.Monad.Class.MonadST (MonadST (..))+import qualified Control.Monad.IOSim as IOSim+import Control.Monad.Primitive (RealWorld)+import Control.Monad.ST.Unsafe (unsafeIOToST)+import Control.Tracer (nullTracer)+import Data.Functor (void)+import Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,+ writePrimVar)+import Data.Typeable (Typeable)+import qualified Data.Vector as V+import Data.Word (Word64)+import Database.LSMTree as LSMT+import qualified System.Directory as IO (createDirectoryIfMissing,+ doesDirectoryExist, removeDirectoryRecursive)+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.IO as FS+import qualified System.FS.BlockIO.Sim as FSSim+import qualified System.FS.Sim.MockFS as FSSim+import System.IO.Unsafe (unsafePerformIO)++-- | Interactive demo showing functional requiremens for the @lsm-tree@ library+-- are met.+--+-- The functional requirements are discussed in this document: "Storing the+-- Cardano ledger state on disk: final report for high-performance backend"+--+-- Sections of the demo code are headed by the number of the corresponding+-- functional requirement.+demo :: Bool -> IO ()+demo interactive = do+ freshDirectory "_demo"+ withOpenSessionIO tracer "_demo" $ \session -> do+ withTableWith config session $ \(table :: Table IO K V B) -> do+ pause interactive -- [0]++ -- 2. basic key-value store operations++ inserts table $ V.fromList [ (K i, V i, Just (B i)) | i <- [1 .. 10_000] ]+ as <- lookups table $ V.fromList [ K 1, K 2, K 3, K 4 ]+ print (fmap getValue as)+ pause interactive -- [1]++ deletes table $ V.fromList [ K i | i <- [1 .. 10_000], even i ]+ bs <- lookups table $ V.fromList [ K 1, K 2, K 3, K 4 ]+ print (fmap getValue bs)+ pause interactive -- [2]++ -- 2. Intermezzo: blob retrieval++ cs <- try @SomeException $ retrieveBlobs session $ V.mapMaybe getBlob as+ print cs+ pause interactive -- [3]++ ds <- try @SomeException $ retrieveBlobs session $ V.mapMaybe getBlob bs+ print ds+ pause interactive -- [4]++ -- 3. range lookups and cursors++ es <- rangeLookup table $ FromToIncluding (K 1) (K 4)+ print (fmap getEntryValue es)+ pause interactive -- [5]++ withCursorAtOffset table (K 1) $ \cursor -> do+ fs <- LSMT.take 2 cursor+ print (fmap getEntryValue fs)+ pause interactive -- [6]++ -- 4. upserts (or monoidal updates)++ -- better than lookup followed by insert+ upserts table $ V.fromList [ (K i, V 1) | i <- [1 .. 10_000] ]+ gs <- lookups table $ V.fromList [ K 1, K 2, K 3, K 4 ]+ print (fmap getValue gs)+ pause interactive -- [7]++ -- 5. multiple independently writable references++ withDuplicate table $ \dupliTable -> do+ inserts dupliTable $ V.fromList [ (K i, V 1, Nothing) | i <- [1 .. 10_000] ]+ hs <- lookups dupliTable $ V.fromList [ K 1, K 2, K 3, K 4 ]+ print (fmap getValue hs)+ pause interactive -- [8]++ is <- lookups table $ V.fromList [ K 1, K 2, K 3, K 4]+ print (fmap getValue is)+ pause interactive -- [9]++ -- 6. snapshots++ saveSnapshot "odds_evens" label table+ saveSnapshot "all_ones" label dupliTable+ js <- listSnapshots session+ print js+ pause interactive -- [10]++ -- 6. snapshots continued++ withTableFromSnapshot session "odds_evens" label $ \(table :: Table IO K V B) -> do+ withTableFromSnapshot session "all_ones" label $ \(dupliTable :: Table IO K V B) -> do+ pause interactive -- [11]++ -- 7. table unions++ withUnion table dupliTable $ \uniTable -> do+ ks <- lookups uniTable $ V.fromList [ K 1, K 2, K 3, K 4]+ print (fmap getValue ks)+ pause interactive -- [12]++ withIncrementalUnion table dupliTable $ \uniTable -> do+ ls <- lookups uniTable $ V.fromList [ K 1, K 2, K 3, K 4]+ print (fmap getValue ls)+ pause interactive -- [13]++ m@(UnionDebt m') <- remainingUnionDebt uniTable+ supplyUnionCredits uniTable (UnionCredits (m' `div` 2))+ print m+ pause interactive -- [14]++ ns <- lookups uniTable $ V.fromList [ K 1, K 2, K 3, K 4]+ print (fmap getValue ns)+ pause interactive -- [15]++ -- 8. simulation++ let+ simpleAction ::+ (LSMT.IOLike m, Typeable h)+ => FS.HasFS m h -> FS.HasBlockIO m h -> m ()+ simpleAction hasFS hasBlockIO = do+ let sessionDir = FS.mkFsPath ["_demo"]+ FS.createDirectoryIfMissing hasFS False sessionDir+ withOpenSession tracer hasFS hasBlockIO 17 sessionDir $ \session -> do+ withTableWith config session $ \(table :: Table m K V B) -> do+ inserts table $ V.fromList [ (K i, V i, Just (B i)) | i <- [1 .. 10_000] ]+ os <- lookups table $ V.fromList [ K 1, K 2, K 3, K 4 ]+ print' (fmap getValue os)++ do+ FS.withIOHasBlockIO (FS.MountPoint "") FS.defaultIOCtxParams $ \hasFS hasBlockIO -> do+ simpleAction hasFS hasBlockIO+ pause interactive -- [16]++ do+ pure $! IOSim.runSimOrThrow $ do+ (hasFS, hasBlockIO) <- FSSim.simHasBlockIO' FSSim.empty+ simpleAction hasFS hasBlockIO+ pause interactive -- [17]++{-------------------------------------------------------------------------------+ Types+-------------------------------------------------------------------------------}++newtype K = K Word64+ deriving stock (Show, Eq)+ deriving newtype SerialiseKey++newtype V = V Word64+ deriving stock (Show, Eq)+ deriving newtype (Num, SerialiseValue)+instance ResolveValue V where+ resolve = (+)++newtype B = B Word64+ deriving stock (Show, Eq)+ deriving newtype (Num, SerialiseValue)++config :: TableConfig+config = defaultTableConfig {+ confWriteBufferAlloc = AllocNumEntries 172+ }++tracer :: Monad m => Tracer m LSMTreeTrace+tracer = nullTracer++label :: SnapshotLabel+label = "KVB"++{-------------------------------------------------------------------------------+ Utils+-------------------------------------------------------------------------------}++{-# NOINLINE pauseRef #-}+pauseRef :: PrimVar RealWorld Int+pauseRef = unsafePerformIO $ newPrimVar 0++incrPauseRef :: IO Int+incrPauseRef = do+ x <- readPrimVar pauseRef+ writePrimVar pauseRef $! x + 1+ pure x++pause :: Bool -> IO ()+pause interactive = do+ x <- incrPauseRef+ putStr ("[" <> show x <> "] " <> "press ENTER to continue...")+ if interactive+ then void $ getLine+ else putStrLn ""++freshDirectory :: FilePath -> IO ()+freshDirectory path = do+ b <- IO.doesDirectoryExist path+ when b $ IO.removeDirectoryRecursive path+ IO.createDirectoryIfMissing False path++print' :: (Show a, MonadST m) => a -> m ()+print' x = stToIO $ unsafeIOToST $ print x
@@ -0,0 +1,15 @@+module Main (main) where++import Database.LSMTree.Demo (demo)+import System.Environment (getArgs)+import System.IO (BufferMode (..), hSetBuffering, stdout)++main :: IO ()+main = do+ args <- getArgs+ let isInteractive = args == ["Interactive"]+ if isInteractive+ then putStrLn "Running in Interactive mode"+ else putStrLn "Running in NonInteractive mode"+ hSetBuffering stdout NoBuffering+ demo isInteractive
@@ -0,0 +1,568 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DuplicateRecordFields #-}+{-# LANGUAGE NondecreasingIndentation #-}+{-# LANGUAGE OverloadedStrings #-}++{-# OPTIONS_GHC -Wno-orphans #-}++{- | Benchmarks for table unions of an LSM tree.++Here is a summary of the table union benchmark.+Note that /all/ function calls are made through+the "Database.LSMTree.Simple" module's API.++__Phase 1: Setup__++The benchmark will setup an initial set of @--table-count@ tables+to be unioned together during "__Phase 2__."+The size of each generated table is the same+and is equal to @--initial-size@.+Each created table has @--initial-size@ insertions operations performed on it+before being written out to disk as a snapshot.+The @--initial-size@ inserted keys in each table are randomly selected+from the following range.+Each key is unique, meaning that keys are randomly sampled+from the range without replacement.++\[+\left[\quad 0,\quad 2 * initialSize \quad\right)+\]++Additionally, the directory in which to isolate the benchmark environment+is specified via the @--bench-dir@ command line option.++__Phase 2: Measurement__++When generating measurements for the table unions,+the benchmark will reload the snapshots of the tables generated+in __Phase 1__ from disk.+Subsequently, the tables will be "incrementally unioned" together.++Once the tables have been loaded and the union initiated,+serveral iterations of lookups will be performed.+One iteration involves performing a @--batch-count@ number of batches+of @--batch-size@ lookups each.+We measure the time spent on running an iteration,+and we compute how many lookups per second were performed during the iteration.++First, 50 iterations are performed /without/ supplying any credits to the unioned table.+This establishes a base-line performance picture.+Iterations \( \left[ 0, 50 \right) \) measure lookups per seconds+on the unioned table with \(100\%\) of the debt remaining.++Subsequently, 100 more iterations are performed.+Before each of these iterations,+a fixed number of credits are supplied to the incremental union table.+The series of measurements allows reasoning about table performance over time+as the tables debt decreases (at a uniform rate).+The number of credits supplied before each iteration is+\(1%\) of the total starting debt.+After 100 steps, \(100\%\) of the debt will be paid off.+Iterations \( \left[ 50, 100 \right) \) measure lookups per second+on the unioned table as the remaining debt decreases.++Finally, 50 concluding iterations are performed.+Since no debt is remaining, no credits are supplied.+Rather, these measurements create a "post-payoff" performance picture.+Iterations \( \left[ 150, 200 \right) \) measure lookups per seconds+on to the unioned table with \(0\%\) of the debt remaining.++__Results__++An informative gnuplot script and data file of the benchmark measurements is+generated and placed in the @bench-unions@ directory.+Run the following command in a shell to generate a PNG of the graph.++@+ cd bench-unions && gnuplot unions-bench.gnuplot && cd ..+@++TODO: explain the baseline table++TODO: explain the seed++TODO: explain collisions analysis+-}+module Bench.Unions (main) where++import Control.Applicative ((<**>))+import Control.Concurrent.Async (forConcurrently_)+import Control.Monad (forM_, void, (>=>))+import Control.Monad.State.Strict (MonadState (..), runState)+import qualified Data.ByteString.Short as BS+import Data.Foldable (traverse_)+import Data.IORef+import qualified Data.List as List+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import Data.Monoid+import qualified Data.Primitive as P+import qualified Data.Set as Set+import qualified Data.Vector as V+import Data.Word (Word64)+import qualified Options.Applicative as O+import Prelude hiding (lookup)+import qualified System.Clock as Clock+import System.Directory (createDirectoryIfMissing)+import System.IO+import System.Mem (performMajorGC)+import qualified System.Random as Random+import System.Random (StdGen)+import Text.Printf (printf)+import qualified Text.Read as Read++import Database.LSMTree.Extras (groupsOfN)+import qualified Database.LSMTree.Extras.Random as Random++import qualified Database.LSMTree.Simple as LSM++-------------------------------------------------------------------------------+-- Constant Values+-------------------------------------------------------------------------------++baselineTableID :: Int+baselineTableID = 0++baselineTableName :: LSM.SnapshotName+baselineTableName = makeTableName baselineTableID++defaultBenchDir :: FilePath+defaultBenchDir = "_bench_unions"++defaultInitialSize :: Int+defaultInitialSize = 1_000_000++defaultTableCount :: Int+defaultTableCount = 10++-- | The default seed is the first 20 digits of Pi.+defaultSeed :: Int+defaultSeed = 1415926535897932384++-------------------------------------------------------------------------------+-- Keys and values+-------------------------------------------------------------------------------++type K = BS.ShortByteString+type V = BS.ShortByteString++label :: LSM.SnapshotLabel+label = LSM.SnapshotLabel "K V B"++-- | We generate 34 byte keys by using a PRNG to extend a word64 to 32 bytes+-- and then appending two constant bytes. This corresponds relatively closely+-- to UTxO keys, which are 32 byte cryptographic hashes, followed by two bytes+-- which are typically the 16bit value 0 or 1 (a transaction output index).+--+makeKey :: Word64 -> K+makeKey seed =+ case P.runPrimArray $ do+ v <- P.newPrimArray 5+ let g0 = Random.mkStdGen (fromIntegral seed)+ let (!w0, !g1) = Random.uniform g0+ P.writePrimArray v 0 w0+ let (!w1, !g2) = Random.uniform g1+ P.writePrimArray v 1 w1+ let (!w2, !g3) = Random.uniform g2+ P.writePrimArray v 2 w2+ let (!w3, _g4) = Random.uniform g3+ P.writePrimArray v 3 w3+ P.writePrimArray v 4 0x3d3d3d3d3d3d3d3d -- ========+ case v of+ P.MutablePrimArray mba -> do+ _ <- P.resizeMutableByteArray (P.MutableByteArray mba) 34+ pure v++ of (P.PrimArray ba :: P.PrimArray Word64) ->+ byteArrayToSBS (P.ByteArray ba)++-- | \( O(1) \) conversion.+byteArrayToSBS :: P.ByteArray -> BS.ShortByteString+#if MIN_VERSION_bytestring(0,12,0)+byteArrayToSBS ba = BS.ShortByteString ba+#else+byteArrayToSBS (P.ByteArray ba) = BS.SBS ba+#endif++-- We use constant value. This shouldn't affect anything.+theValue :: V+theValue = BS.replicate 60 120 -- 'x'+{-# NOINLINE theValue #-}++-------------------------------------------------------------------------------+-- Options and commands+-------------------------------------------------------------------------------++data GlobalOpts = GlobalOpts+ { rootDir :: !FilePath -- ^ Session directory.+ , tableCount :: !Int -- ^ Number of tables in the benchmark+ , initialSize :: !Int -- ^ Initial size of each table in the benchmark+ , seed :: !Int -- ^ The seed for the RNG for deterministic behavior,+ -- use system entropy when not explicitly provided.+ }+ deriving stock Show++data RunOpts = RunOpts+ { batchCount :: !Int+ , batchSize :: !Int+ }+ deriving stock Show++data Cmd+ -- | Setup benchmark: generate initial LSM trees etc.+ = CmdSetup++ -- | Run collision analysis.+ | CmdCollisions++ -- | Run the actual benchmark+ | CmdRun RunOpts+ deriving stock Show++-------------------------------------------------------------------------------+-- command line interface+-------------------------------------------------------------------------------++globalOptsP :: O.Parser GlobalOpts+globalOptsP = pure GlobalOpts+ <*> O.option O.str (O.long "bench-dir" <> O.value defaultBenchDir <> O.showDefault <> O.help "Benchmark directory to put files in")+ <*> O.option (positiveParser "number of tables") (O.long "table-count" <> O.value defaultTableCount <> O.showDefault <> O.help "Number of tables to benchmark")+ <*> O.option (positiveParser "size of tables") (O.long "initial-size" <> O.value defaultInitialSize <> O.showDefault <> O.help "Initial size of each table")+ <*> O.option O.auto (O.long "seed" <> O.value defaultSeed <> O.showDefault <> O.help "Random seed, defaults to the first 20 digits of Pi")+ where+ positiveParser str =+ let validator v+ | v >= 1 = Right v+ | otherwise = Left $ unwords+ [ "Non-positive number for", str <> ":", show v ]+ in O.eitherReader $ Read.readEither >=> validator++cmdP :: O.Parser Cmd+cmdP = O.subparser $ mconcat+ [ O.command "setup" $ O.info+ (CmdSetup <$ O.helper)+ (O.progDesc "Setup benchmark by generating required tables")+ , O.command "collisions" $ O.info+ (CmdCollisions <$ O.helper)+ (O.progDesc "Collision analysis, compute shared keys between tables")+ , O.command "run" $ O.info+ (CmdRun <$> runOptsP <**> O.helper)+ (O.progDesc "Proper run, measuring performance and generating a benchmark report")+ ]++runOptsP :: O.Parser RunOpts+runOptsP = pure RunOpts+ <*> O.option O.auto (O.long "batch-count" <> O.value 200 <> O.showDefault <> O.help "Batch count")+ <*> O.option O.auto (O.long "batch-size" <> O.value 256 <> O.showDefault <> O.help "Batch size")++-------------------------------------------------------------------------------+-- measurements+-------------------------------------------------------------------------------+++-- | Returns number of seconds elapsed+timed :: IO a -> IO (a, Double)+timed action = do+ performMajorGC+ t1 <- Clock.getTime Clock.Monotonic+ x <- action+ t2 <- Clock.getTime Clock.Monotonic+ performMajorGC+ let !t = fromIntegral (Clock.toNanoSecs (Clock.diffTimeSpec t2 t1)) * 1e-9+ pure (x, t)++-- | Returns number of seconds elapsed+timed_ :: IO () -> IO Double+timed_ action = do+ ((), t) <- timed action+ pure t++-------------------------------------------------------------------------------+-- Setup+-------------------------------------------------------------------------------++doSetup :: GlobalOpts -> IO ()+doSetup gopts = do+ -- Define some constants+ let populationBatchSize = 256+ entryCount = initialSize gopts+ -- The key size is twice the specified size because we will delete half+ -- of the keys in the domain of each table uniformly at random.+ keyMax = 2 * entryCount - 1+ keyMin = 0+ tableIDs = tableRange gopts++ -- Setup RNG+ let tableRNGs = deriveSetupRNGs gopts $ length tableIDs++ -- Ensure that our mount point exists on the real file system+ createDirectoryIfMissing True $ rootDir gopts++ -- Populate the specified number of tables+ LSM.withOpenSession (rootDir gopts) $ \session -> do+ -- Create a "baseline" table+ --+ -- We create a single table that *already has* all the same key value pairs+ -- which exist in the union of all tables *which are going* to be unioned.+ -- This way we can compare performance of the union of tables to a+ -- "baseline" table since both share all the same key value pairs.+ table_0 <- LSM.newTable @K @V session++ forConcurrently_ (NE.zip tableIDs tableRNGs) $ \(tID, tRNG) -> do+ -- Create a new table+ table_n <- LSM.newTable @K @V session+ -- Populate the table in batches+ forM_ (groupsOfN (populationBatchSize * 2) [ keyMin .. keyMax ]) $ \batch -> do+ let prunedBatch =+ Random.sampleUniformWithoutReplacement+ tRNG (NE.length batch `div` 2) $ NE.toList batch+ let keyInserts = V.fromList [+ (makeKey (fromIntegral k), theValue)+ | k <- prunedBatch+ ]+ -- Insert the batch of the randomly selected keys+ -- into both the baseline table (0) and the current table+ LSM.inserts table_0 keyInserts+ LSM.inserts table_n keyInserts+ LSM.saveSnapshot (makeTableName tID) label table_n++ -- Finally, save the baseline table+ LSM.saveSnapshot baselineTableName label table_0++makeTableName :: Show a => a -> LSM.SnapshotName+makeTableName n = LSM.toSnapshotName $ "bench_" <> show n++tableRange :: GlobalOpts -> NonEmpty Int+tableRange gopts =+ let n1 = succ baselineTableID+ in NE.fromList [ n1 .. tableCount gopts + baselineTableID ]++-------------------------------------------------------------------------------+-- Collision analysis+-------------------------------------------------------------------------------++-- | Count duplicate keys in all tables that will be unioned together+doCollisionAnalysis :: GlobalOpts -> IO ()+doCollisionAnalysis gopts = do+ LSM.withOpenSession (rootDir gopts) $ \session -> do+ seenRef <- newIORef Set.empty+ dupRef <- newIORef Set.empty++ forM_ (tableRange gopts) $ \tID -> do+ let name = makeTableName tID+ LSM.withTableFromSnapshot session name label $ \(table :: LSM.Table K V) -> do+ LSM.withCursor table $ \cursor -> do+ streamCursor cursor $ \(k, _) -> do+ seen <- readIORef seenRef+ if Set.member k seen then+ modifyIORef dupRef $ Set.insert k+ else+ modifyIORef seenRef $ Set.insert k++ seen <- readIORef seenRef+ dups <- readIORef dupRef+ printf "Keys seen at least once: %d\n" $ Set.size seen+ printf "Keys seen at least twice: %d\n" $ Set.size dups++streamCursor :: LSM.Cursor K V -> ((K, V) -> IO ()) -> IO ()+streamCursor cursor f = go+ where+ go = LSM.next cursor >>= \case+ Nothing -> pure ()+ Just kv -> f kv >> go++-------------------------------------------------------------------------------+-- run+-------------------------------------------------------------------------------++doRun :: GlobalOpts -> RunOpts -> IO ()+doRun gopts opts = do+ -- Perform 3 measurement phases+ -- * Phase 1: Measure performance before supplying any credits.+ -- * Phase 2: Measure performance as credits are incrementally supplied and debt is repaid.+ -- * Phase 3: Measure performance when debt is 0.++ let rng = deriveRunRNG gopts+ dataPath = "bench-unions/unions-bench.dat"++ withFile dataPath WriteMode $ \h -> do+ hPutStrLn h "# iteration \t baseline (ops/sec) \t union (ops/sec) \t union debt"++ LSM.withOpenSession (rootDir gopts) $ \session -> do+ -- Load the baseline table+ LSM.withTableFromSnapshot session baselineTableName label+ $ \baselineTable -> do+ -- Load the union tables+ withTablesFromSnapshots session label (makeTableName <$> tableRange gopts)+ $ \inputTables -> do+ -- Start the incremental union+ LSM.withIncrementalUnions inputTables $ \unionedTable -> do+ let measurePerformance :: Int -> Maybe LSM.UnionCredits -> IO ()+ measurePerformance iteration mayCredits = do+ LSM.supplyUnionCredits unionedTable `traverse_` mayCredits+ LSM.UnionDebt currDebt <- LSM.remainingUnionDebt unionedTable++ baselineOpsSec <- timeOpsPerSecond gopts opts baselineTable rng+ unionOpsSec <- timeOpsPerSecond gopts opts unionedTable rng++ printf "iteration: %d, baseline: %7.01f ops/sec, union: %7.01f ops/sec, debt: %d\n"+ iteration baselineOpsSec unionOpsSec currDebt++ hPutStrLn h $ unwords [ show iteration, show baselineOpsSec+ , show unionOpsSec, show currDebt ]++ LSM.UnionDebt totalDebt <- LSM.remainingUnionDebt unionedTable++ -- Phase 1 measurements: Debt = 100%+ forM_ [0..50-1] $ \step -> do+ measurePerformance step Nothing++ -- Phase 2 measurements: Debt ∈ [0%, 99%]+ forM_ [50..150-1] $ \step -> do+ let creditsPerIteration = LSM.UnionCredits ((totalDebt + 99) `div` 100)+ measurePerformance step (Just creditsPerIteration)++ -- Phase 3 measurements: Debt = 0%+ forM_ [150..200-1] $ \step -> do+ measurePerformance step Nothing++-- | Exception-safe opening of multiple snapshots+withTablesFromSnapshots ::+ LSM.Session+ -> LSM.SnapshotLabel+ -> (NonEmpty LSM.SnapshotName)+ -> (NonEmpty (LSM.Table K V) -> IO a)+ -> IO a+withTablesFromSnapshots session snapLabel (x0 :| xs0) k = do+ LSM.withTableFromSnapshot session x0 snapLabel $ \t0 -> go' (NE.singleton t0) xs0+ where+ go' acc [] = k (NE.reverse acc)+ go' acc (x:xs) =+ LSM.withTableFromSnapshot session x snapLabel $ \t ->+ go' (t NE.<| acc) xs++-- | Returns operations per second+timeOpsPerSecond :: GlobalOpts -> RunOpts -> LSM.Table K V -> StdGen -> IO Double+timeOpsPerSecond gopts opts table g = do+ t <- timed_ $+ sequentialIterations+ (initialSize gopts) (batchSize opts) (batchCount opts) table g++ let ops = batchCount opts * batchSize opts+ opsPerSec = fromIntegral ops / t++ pure opsPerSec++-------------------------------------------------------------------------------+-- PRNG initialisation+-------------------------------------------------------------------------------++deriveSetupRNGs :: GlobalOpts -> Int -> NonEmpty Random.StdGen+deriveSetupRNGs gOpts amount =+ let -- g1 is reserved for the run command+ (_g1, !g2) = deriveInitialForkedRNGs gOpts+ in NE.fromList $ take amount $ List.unfoldr (Just . Random.splitGen) g2++deriveRunRNG :: GlobalOpts -> Random.StdGen+deriveRunRNG gOpts =+ let -- g2 and its splits are reserved for the setup command+ (!g1, _g2) = deriveInitialForkedRNGs gOpts+ in g1++deriveInitialForkedRNGs :: GlobalOpts -> (Random.StdGen, Random.StdGen)+deriveInitialForkedRNGs = Random.splitGen . Random.mkStdGen . seed++-------------------------------------------------------------------------------+-- Batch generation+-------------------------------------------------------------------------------++generateBatch ::+ Int -- ^ initial size of the table+ -> Int -- ^ batch size+ -> StdGen+ -> (StdGen, V.Vector K)+generateBatch initialSize batchSize g =+ (g', V.map makeKey lookups)+ where+ (!g', lookups) = generateBatch' initialSize batchSize g++{-# INLINE generateBatch' #-}+generateBatch' ::+ Int -- ^ initial size of the table+ -> Int -- ^ batch size+ -> StdGen+ -> (StdGen, V.Vector Word64)+generateBatch' initialSize batchSize g = (g', lookups)+ where+ randomKey :: StdGen -> (Word64, StdGen)+ randomKey = Random.uniformR (0, 2 * fromIntegral initialSize - 1)++ lookups :: V.Vector Word64+ (lookups, !g') =+ runState (V.replicateM batchSize (state randomKey)) g++-------------------------------------------------------------------------------+-- sequential+-------------------------------------------------------------------------------++{-# INLINE sequentialIteration #-}+sequentialIteration ::+ Int -- ^ initial size of the table+ -> Int -- ^ batch size+ -> LSM.Table K V+ -> StdGen+ -> IO StdGen+sequentialIteration !initialSize !batchSize !tbl !g = do+ let (!g', ls) = generateBatch initialSize batchSize g++ -- lookups+ !_ <- LSM.lookups tbl ls++ -- continue to the next batch+ pure g'++sequentialIterations ::+ Int -- ^ initial size of the table+ -> Int -- ^ batch size+ -> Int -- ^ batch count+ -> LSM.Table K V+ -> StdGen+ -> IO ()+sequentialIterations !initialSize !batchSize !batchCount !tbl !g0 = do+ void $ forFoldM_ g0 [ 0 .. batchCount - 1 ] $ \_b g ->+ sequentialIteration initialSize batchSize tbl g++-------------------------------------------------------------------------------+-- main+-------------------------------------------------------------------------------++main :: IO ()+main = do+ hSetBuffering stdout NoBuffering+#ifdef NO_IGNORE_ASSERTS+ putStrLn "WARNING: Benchmarking in debug mode."+ putStrLn " To benchmark in release mode, pass:"+ putStrLn " --project-file=cabal.project.release"+#endif+ (gopts, cmd) <- O.customExecParser prefs cliP+ print gopts+ print cmd+ case cmd of+ CmdCollisions -> doCollisionAnalysis gopts+ CmdRun opts -> doRun gopts opts+ CmdSetup -> doSetup gopts+ where+ cliP = O.info ((,) <$> globalOptsP <*> cmdP <**> O.helper) O.fullDesc+ prefs = O.prefs $ O.showHelpOnEmpty <> O.helpShowGlobals <> O.subparserInline++-------------------------------------------------------------------------------+-- general utils+-------------------------------------------------------------------------------++forFoldM_ :: Monad m => s -> [a] -> (a -> s -> m s) -> m s+forFoldM_ !s0 xs0 f = go s0 xs0+ where+ go !s [] = pure s+ go !s (x:xs) = do+ !s' <- f x s+ go s' xs
@@ -0,0 +1,6 @@+module Main (main) where++import qualified Bench.Unions++main :: IO ()+main = Bench.Unions.main
@@ -0,0 +1,286 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE NumericUnderscores #-}++module Main ( main ) where++import Control.Exception (evaluate)+import Control.Monad+import Control.Monad.ST+import Control.Monad.ST.Unsafe+import Data.Bits ((.&.))+import Data.BloomFilter.Blocked (Bloom, BloomSize)+import qualified Data.BloomFilter.Blocked as Bloom+import Data.Time+import Data.Vector (Vector)+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import Data.WideWord.Word256 (Word256)+import GHC.Stats+import Numeric+import System.IO+import System.Mem (performMajorGC)+import System.Random+import Text.Printf (printf)++import Database.LSMTree.Extras.Orphans ()+import Database.LSMTree.Internal.Assertions (fromIntegralChecked)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Serialise (SerialisedKey,+ serialiseKey)++main :: IO ()+main = do+ hSetBuffering stdout NoBuffering+ benchmarks++-- Benchmark parameters you can tweak. The defaults are rather small, so the+-- runtime is short and do not show the effects of data sizes no longer fitting+-- into the CPU cache.++-- | The number of entries in the filter in the smallest LSM runs is+-- @2^benchmarkSizeBase@+benchmarkSizeBase :: SizeBase+benchmarkSizeBase = 16++-- | The number of lookups to do. This has to be smaller than the total size of+-- all the filters (otherwise we will not get true positive probes, which is+-- part of the point of this benchmark).+benchmarkNumLookups :: Integer+benchmarkNumLookups = 25_000_000++-- | The number of lookups to do in a single batch.+benchmarkBatchSize :: Int+benchmarkBatchSize = 256++benchmarkNumBitsPerEntry :: RequestedBitsPerEntry+benchmarkNumBitsPerEntry = 10++benchmarks :: IO ()+benchmarks = do+#ifdef NO_IGNORE_ASSERTS+ putStrLn "WARNING: Benchmarking in debug mode."+ putStrLn " To benchmark in release mode, pass:"+ putStrLn " --project-file=cabal.project.release"+#endif++ enabled <- getRTSStatsEnabled+ unless enabled $ fail "Need RTS +T statistics enabled"+ let filterSizes = lsmStyleBloomFilters benchmarkSizeBase+ benchmarkNumBitsPerEntry+ putStrLn "Bloom filter stats:"+ putStrLn "(numEntries, sizeFactor, BloomSize { sizeBits, sizeHashes })"+ mapM_ print filterSizes+ putStrLn $ "total number of entries:\t " ++ show (totalNumEntries filterSizes)+ putStrLn $ "total filter size in bytes:\t " ++ show (totalNumBytes filterSizes)+ putStrLn $ "total number of key lookups:\t " ++ show benchmarkNumLookups++ unless (totalNumEntriesSanityCheck benchmarkSizeBase filterSizes) $+ fail "totalNumEntriesSanityCheck failed"+ unless (totalNumEntries filterSizes >= benchmarkNumLookups) $+ fail "number of key lookups is more than number of entries"++ putStrLn "Generating bloom filters..."+ let rng0 = mkStdGen 42+ vbs <- elemManyEnv filterSizes rng0+ putStrLn " finished."+ putStrLn ""++ hashcost <-+ benchmark "makeHashes"+ "(This baseline is the cost of computing and hashing the keys)"+ (benchInBatches benchmarkBatchSize rng0+ (benchMakeHashes vbs))+ (fromIntegralChecked benchmarkNumLookups)+ (0, 0)+ 289++ _ <-+ benchmark "elemHashes"+ "(this is the simple one-by-one lookup, less the cost of computing and hashing the keys)"+ (benchInBatches benchmarkBatchSize rng0+ (benchElemHashes vbs))+ (fromIntegralChecked benchmarkNumLookups)+ hashcost+ 0++ _ <-+ benchmark "bloomQueries"+ "(this is the batch lookup, less the cost of computing and hashing the keys)"+ (benchInBatches benchmarkBatchSize rng0+ (\ks -> Bloom.bloomQueries benchSalt vbs ks `seq` ()))+ (fromIntegralChecked benchmarkNumLookups)+ hashcost+ 0++ pure ()++type Alloc = Int++benchmark :: String+ -> String+ -> (Int -> ())+ -> Int+ -> (NominalDiffTime, Alloc)+ -> Int+ -> IO (NominalDiffTime, Alloc)+benchmark name description action n (subtractTime, subtractAlloc) expectedAlloc = do+ putStrLn $ "Benchmarking " ++ name ++ " ... "+ putStrLn description+ performMajorGC+ allocBefore <- allocated_bytes <$> getRTSStats+ timeBefore <- getCurrentTime+ evaluate (action n)+ timeAfter <- getCurrentTime+ performMajorGC+ allocAfter <- allocated_bytes <$> getRTSStats+ putStrLn "Finished."+ let allocTotal :: Alloc+ timeTotal = timeAfter `diffUTCTime` timeBefore+ allocTotal = fromIntegral allocAfter - fromIntegral allocBefore+ timeNet = timeTotal - subtractTime+ allocNet = allocTotal - subtractAlloc++ timePerKey, allocPerKey :: Double+ timePerKey = realToFrac timeNet / fromIntegral n+ allocPerKey = fromIntegral allocNet / fromIntegral n+ let printStat :: String -> Double -> String -> IO ()+ printStat label v unit =+ putStrLn $ label ++ showGFloat (Just 2) v (' ':unit)+ printStat "Time total: " (realToFrac timeTotal) "seconds"+ printStat "Alloc total: " (fromIntegral allocTotal) "bytes"+ printStat "Time net: " (realToFrac timeNet) "seconds"+ printStat "Alloc net: " (fromIntegral allocNet) "bytes"+ printStat "Time net per key: " timePerKey "seconds"+ printStat "Alloc net per key: " allocPerKey "bytes"++ unless (truncate allocPerKey == expectedAlloc) $ do+ printf "WARNING: expecting %d, got %d bytes allocated per key\n"+ expectedAlloc+ (truncate allocPerKey :: Int)++ putStrLn ""+ pure (timeNet, allocNet)++-- | (numEntries, sizeFactor, (BloomSize numBits numHashFuncs))+type BloomFilterSizeInfo = (Integer, Integer, BloomSize)+type SizeBase = Int+type RequestedBitsPerEntry = Double++-- | Calculate the sizes of a realistic LSM style set of Bloom filters, one+-- for each LSM run. This uses base 4, with 4 disk levels, using tiering+-- for internal levels and leveling for the final (biggest) level.+--+-- Due to the incremental merging, each level actually has (in the worst case)+-- 2x the number of runs, hence 8 per level for tiering levels.+--+lsmStyleBloomFilters :: SizeBase -> RequestedBitsPerEntry -> [BloomFilterSizeInfo]+lsmStyleBloomFilters l1 requestedBitsPerEntry =+ [ (numEntries, sizeFactor, bsize)+ | (numEntries, sizeFactor)+ <- replicate 8 (2^(l1+0), 1) -- 8 runs at level 1 (tiering)+ ++ replicate 8 (2^(l1+2), 4) -- 8 runs at level 2 (tiering)+ ++ replicate 8 (2^(l1+4),16) -- 8 runs at level 3 (tiering)+ ++ [(2^(l1+8),256)] -- 1 run at level 4 (leveling)+ , let bsize = Bloom.sizeForBits requestedBitsPerEntry (fromIntegral numEntries)+ ]++totalNumEntries, totalNumBytes :: [BloomFilterSizeInfo] -> Integer+totalNumEntries filterSizes =+ sum [ numEntries | (numEntries, _, _) <- filterSizes ]++totalNumBytes filterSizes =+ sum [ toInteger (Bloom.sizeBits bsize)+ | (_,_,bsize) <- filterSizes ]+ `div` 8++totalNumEntriesSanityCheck :: SizeBase -> [BloomFilterSizeInfo] -> Bool+totalNumEntriesSanityCheck l1 filterSizes =+ totalNumEntries filterSizes+ ==+ sum [ 2^l1 * sizeFactor | (_, sizeFactor, _) <- filterSizes ]++benchSalt :: Bloom.Salt+benchSalt = 4++-- | Input environment for benchmarking 'Bloom.elemMany'.+--+-- The idea here is to have a collection of bloom filters corresponding to+-- the sizes used in a largeish LSM. In particular, the sizes are in increasing+-- powers of 4, and the largest ones should be bigger than the CPU L3 cache.+-- Furthermore, the keys in the filters are non-overlapping, and lookups will+-- be true positives in only one filter. Thus most lookups will be true+-- negatives.+--+-- The goal is to benchmark the benefits of optimisations for the LSM situation:+--+-- * where the same key is being looked up in many filters,+-- * with a hit in only one filter expected, and+-- * where the total size of the filters is too large to fully fit in cache+-- (though the smaller ones may fit in the caches).+--+elemManyEnv :: [BloomFilterSizeInfo]+ -> StdGen+ -> IO (Vector (Bloom SerialisedKey))+elemManyEnv filterSizes rng0 =+ stToIO $ do+ -- create the filters+ mbs <- sequence+ [ Bloom.new bsize benchSalt+ | (_, _, bsize) <- filterSizes+ ]+ -- add elements+ foldM_+ (\rng (i, mb) -> do+ -- progress+ when (i .&. 0xFFFF == 0) (unsafeIOToST $ putStr ".")+ -- insert n elements into filter b+ let k :: Word256+ (!k, !rng') = uniform rng+ Bloom.insert mb (serialiseKey k)+ pure rng'+ )+ rng0+ (zip [0 .. totalNumEntries filterSizes - 1]+ (cycle [ mb'+ | (mb, (_, sizeFactor, _)) <- zip mbs filterSizes+ , mb' <- replicate (fromIntegralChecked sizeFactor) mb ]))+ V.fromList <$> mapM Bloom.unsafeFreeze mbs++type BatchBench = V.Vector SerialisedKey -> ()++{-# NOINLINE benchInBatches #-}+benchInBatches :: Int -> StdGen -> BatchBench -> Int -> ()+benchInBatches !b !rng0 !action =+ go rng0+ where+ go !rng !n+ | n <= 0 = ()+ | otherwise =+ let (!rng'', !rng') = splitGen rng+ ks :: VP.Vector Word256+ !ks = VP.unfoldrExactN b uniform rng'+ ks' :: V.Vector SerialisedKey+ !ks' = V.map serialiseKey (V.convert ks)+ in action ks' `seq` go rng'' (n-b)++-- | This gives us a combined cost of calculating the series of keys and their+-- hashes (when used with 'benchInBatches').+benchMakeHashes :: Vector (Bloom SerialisedKey) -> BatchBench+benchMakeHashes !_bs !ks =+ let khs :: VP.Vector (Bloom.Hashes SerialisedKey)+ !khs = V.convert (V.map (Bloom.hashesWithSalt benchSalt) ks)+ in khs `seq` ()++-- | This gives us a combined cost of calculating the series of keys, their+-- hashes, and then using 'Bloom.elemHashes' with each filter (when used+-- with 'benchInBatches').+benchElemHashes :: Vector (Bloom SerialisedKey) -> BatchBench+benchElemHashes !bs !ks =+ let khs :: VP.Vector (Bloom.Hashes SerialisedKey)+ !khs = V.convert (V.map (Bloom.hashesWithSalt benchSalt) ks)+ in V.foldl'+ (\_ b -> VP.foldl'+ (\_ kh -> Bloom.elemHashes b kh `seq` ())+ () khs)+ () bs
@@ -0,0 +1,572 @@+{-# LANGUAGE CPP #-}++module Main ( main ) where++import Control.DeepSeq+import Control.Exception (bracket)+import Control.Monad+import Control.Monad.Class.MonadST+import Control.Monad.Primitive+import Control.Monad.ST.Strict (ST, runST)+import Control.RefCount+import Data.Bits ((.&.))+import Data.BloomFilter.Blocked (Bloom)+import qualified Data.BloomFilter.Blocked as Bloom+import Data.Time+import qualified Data.Vector as V+import Data.Vector.Algorithms.Merge as Merge+import qualified Data.Vector.Generic.Mutable as VGM+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed.Mutable as VUM+import Database.LSMTree.Extras.Orphans ()+import Database.LSMTree.Extras.UTxO+import Database.LSMTree.Internal.Arena (ArenaManager, newArenaManager,+ withArena)+import Database.LSMTree.Internal.Entry (Entry (Insert),+ NumEntries (..))+import Database.LSMTree.Internal.Index (Index)+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Compact))+import Database.LSMTree.Internal.Lookup+import Database.LSMTree.Internal.Paths (RunFsPaths (RunFsPaths))+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunAcc (RunBloomFilterAlloc (..))+import Database.LSMTree.Internal.RunBuilder (RunParams (..))+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue,+ SerialisedKey, serialiseKey, serialiseValue)+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import Debug.Trace (traceMarkerIO)+import GHC.Stats+import Numeric+import System.Environment (getArgs)+import System.Exit (exitFailure)+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.IO as FS+import qualified System.FS.IO as FS+import System.IO+import System.IO.Unsafe (unsafePerformIO)+import System.Mem (performMajorGC)+import System.Random++main :: IO ()+main = do+ hSetBuffering stdout NoBuffering+ args <- getArgs+ case args of+ [arg] | let cache = read arg ->+ benchmarks (if cache then Run.CacheRunData else Run.NoCacheRunData)+ _ -> do+ putStrLn "Wrong usage, pass in [True] or [False] for the caching flag."+ exitFailure++-- | The number of entries in the smallest LSM runs is @2^benchmarkSizeBase@.+--+-- This is currently set to however many UTXO entries fit into 2MB of disk+-- pages.+--+-- >>> benchmarkSizeBase+-- 14+benchmarkSizeBase :: SizeBase+benchmarkSizeBase = floor @Double $+ logBase 2 ( fromIntegral @Int ((2 * 1024 * 1024 `div` 4096)+ * (floor @Double numEntriesFitInPage)) )++-- | The number of lookups to do. This has to be smaller than the total size of+-- all the runs (otherwise we will not get true positive probes, which is part+-- of the point of this benchmark).+--+-- The benchmark might do slightly more lookups, because it generates batches of+-- keys of size 'benchmarkGenBatchSize'. This shouldn't affect the results+-- significantly, unless 'benchmarkNumLookups' approaches+-- 'benchmarkGenBatchSize'.+benchmarkNumLookups :: Int+benchmarkNumLookups = 1_000_000 -- 10 * the stretch target++-- | The size of batches as they are generated by the benchmark.+benchmarkGenBatchSize :: Int+benchmarkGenBatchSize = 256++benchmarkNumBitsPerEntry :: Double+benchmarkNumBitsPerEntry = 10++benchmarkResolveSerialisedValue :: ResolveSerialisedValue+benchmarkResolveSerialisedValue = const++-- >>> pageBits+-- 32768+pageBits :: Int+pageBits = 4096 * 8 -- page size in bits++-- >>> unusedPageBits+-- 32688+unusedPageBits :: Int+unusedPageBits = pageBits -- page size in bits+ - 8 * 8 -- directory+ - 16 -- last value offset++-- >>> entryBits+-- 752+entryBits :: Int+entryBits = 34 * 8 -- key size+ + 60 * 8 -- value size++-- >>> entryBitsWithOverhead+-- 787+entryBitsWithOverhead :: Int+entryBitsWithOverhead = entryBits -- key and value size+ + 1 -- blobref indicator+ + 2 -- operation type+ + 16 -- key offset+ + 16 -- value offset++-- >>> numEntriesFitInPage+-- 41.53494282083863+numEntriesFitInPage :: Fractional a => a+numEntriesFitInPage = fromIntegral unusedPageBits / fromIntegral entryBitsWithOverhead++benchSalt :: Bloom.Salt+benchSalt = 4++benchmarks :: Run.RunDataCaching -> IO ()+benchmarks !caching = withFS $ \hfs hbio -> do+#ifdef NO_IGNORE_ASSERTS+ putStrLn "WARNING: Benchmarking in debug mode."+ putStrLn " To benchmark in release mode, pass:"+ putStrLn " --project-file=cabal.project.release"+#endif+ arenaManager <- newArenaManager+ enabled <- getRTSStatsEnabled+ unless enabled $ fail "Need RTS +T statistics enabled"+ let runSizes = lsmStyleRuns benchmarkSizeBase+ putStrLn "Precomputed run stats:"+ putStrLn "(numEntries, sizeFactor)"+ mapM_ print runSizes+ putStrLn $ "total number of entries:\t " ++ show (totalNumEntries runSizes)+ putStrLn $ "total number of key lookups:\t " ++ show benchmarkNumLookups++ unless (totalNumEntriesSanityCheck benchmarkSizeBase runSizes) $+ fail "totalNumEntriesSanityCheck failed"+ unless (totalNumEntries runSizes >= benchmarkNumLookups) $+ fail "number of key lookups is more than number of entries"++ traceMarkerIO "Generating runs"+ putStr "<Generating runs>"+ -- This initial key RNG is used for both generating the runs and generating+ -- lookup batches. This ensures that we only only perform true positive+ -- lookups. Also, lookupsEnv shuffles the generated keys into the different+ -- runs, such that the generated lookups access the runs in random places+ -- instead of sequentially.+ let keyRng0 = mkStdGen 17++ (!runs, !blooms, !indexes, !handles) <- lookupsEnv runSizes keyRng0 hfs hbio caching+ putStrLn "<finished>"++ traceMarkerIO "Computing statistics for generated runs"+ let numEntries = V.map Run.size runs+ numPages = V.map Run.sizeInPages runs+ nhashes = V.map (Bloom.sizeHashes . Bloom.size) blooms+ bitsPerEntry = V.zipWith+ (\b (NumEntries n) ->+ fromIntegral (Bloom.sizeBits (Bloom.size b))+ / fromIntegral n :: Double)+ blooms+ numEntries+ stats = V.zip4 numEntries numPages nhashes bitsPerEntry+ putStrLn "Actual stats for generated runs:"+ putStrLn "(numEntries, numPages, numHashes, bits per entry)"+ mapM_ print stats++ _ <- putStr "Pausing. Drop caches now! When ready, press enter." >> getLine++ traceMarkerIO "Running benchmark"+ putStrLn ""+ bgenKeyBatches@(x1, y1) <-+ benchmark "benchGenKeyBatches"+ "Calculate batches of keys. This serves as a baseline for later benchmark runs to compare against."+ (pure . benchGenKeyBatches blooms keyRng0) benchmarkNumLookups+ (0, 0)+ _bbloomQueries@(x2, y2) <-+ benchmark "benchBloomQueries"+ "Calculate batches of keys, and perform bloom queries for each batch. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches."+ (pure . benchBloomQueries blooms keyRng0) benchmarkNumLookups+ bgenKeyBatches+ _bindexSearches <-+ benchmark "benchIndexSearches"+ "Calculate batches of keys, perform bloom queries for each batch, and perform index searches for positively queried keys in each batch. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches and benchBloomQueries."+ (benchIndexSearches arenaManager blooms indexes handles keyRng0) benchmarkNumLookups+ (x1 + x2, y1 + y2)+ _bprepLookups <-+ benchmark "benchPrepLookups"+ "Calculate batches of keys, and prepare lookups for each batch. This is roughly doing the same amount of work as benchBloomQueries and benchIndexSearches. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches."+ (benchPrepLookups arenaManager blooms indexes handles keyRng0) benchmarkNumLookups+ bgenKeyBatches+ _blookupsIO <-+ benchmark "benchLookupsIO"+ "Calculate batches of keys, and perform disk lookups for each batch. This is roughly doing the same as benchPrepLookups, but also performing the disk I/O and resolving values. Net time/allocation is the result of subtracting the cost of benchGenKeyBatches."+ (\n -> do+ let wb_unused = WB.empty+ bracket (WBB.new hfs (FS.mkFsPath ["wbblobs_unused"])) releaseRef $ \wbblobs_unused ->+ benchLookupsIO hbio arenaManager benchmarkResolveSerialisedValue+ wb_unused wbblobs_unused runs blooms indexes handles+ keyRng0 n)+ benchmarkNumLookups+ bgenKeyBatches+ --TODO: consider adding benchmarks that also use the write buffer++ traceMarkerIO "Cleaning up"+ putStrLn "Cleaning up"+ V.mapM_ releaseRef runs++ traceMarkerIO "Computing statistics for prepLookups results"+ putStr "<Computing statistics for prepLookups>"+ let !x = classifyLookups blooms keyRng0 benchmarkNumLookups+ putStrLn "<finished>"+ putStrLn "Statistics for prepLookups:"+ putStrLn "(positives, fpr, tp, fn, fp, tn)"+ print x++type Alloc = Int++benchmark :: String+ -> String+ -> (Int -> IO ())+ -> Int+ -> (NominalDiffTime, Alloc)+ -> IO (NominalDiffTime, Alloc)+benchmark name description action n (subtractTime, subtractAlloc) = do+ traceMarkerIO ("Benchmarking " ++ name)+ putStrLn $ "Benchmarking " ++ name ++ " ... "+ putStrLn description+ performMajorGC+ allocBefore <- allocated_bytes <$> getRTSStats+ timeBefore <- getCurrentTime+ () <- action n+ timeAfter <- getCurrentTime+ performMajorGC+ allocAfter <- allocated_bytes <$> getRTSStats+ putStrLn "Finished."+ let allocTotal :: Alloc+ timeTotal = timeAfter `diffUTCTime` timeBefore+ allocTotal = fromIntegral allocAfter - fromIntegral allocBefore+ timeNet = timeTotal - subtractTime+ allocNet = allocTotal - subtractAlloc++ timePerKey, allocPerKey :: Double+ timePerKey = realToFrac timeNet / fromIntegral n+ allocPerKey = fromIntegral allocNet / fromIntegral n+ let printStat :: String -> Double -> String -> IO ()+ printStat label v unit =+ putStrLn $ label ++ showGFloat (Just 2) v (' ':unit)+ printStat "Time total: " (realToFrac timeTotal) "seconds"+ printStat "Alloc total: " (fromIntegral allocTotal) "bytes"+ printStat "Time net: " (realToFrac timeNet) "seconds"+ printStat "Alloc net: " (fromIntegral allocNet) "bytes"+ printStat "Time net per key: " timePerKey "seconds"+ printStat "Alloc net per key: " allocPerKey "bytes"++ putStrLn ""+ pure (timeNet, allocNet)++-- | (numEntries, sizeFactor)+type RunSizeInfo = (Int, Int)+type SizeBase = Int++-- | Calculate the sizes of a realistic LSM style set of runs. This uses base 4,+-- with 4 disk levels, using tiering for internal levels and leveling for the+-- final (biggest) level.+--+-- Due to the incremental merging, each level actually has (in the worst case)+-- 2x the number of runs, hence 8 per level for tiering levels.+--+lsmStyleRuns :: SizeBase -> [RunSizeInfo]+lsmStyleRuns l1 =+ replicate 8 (2^(l1+ 0), 1) -- 8 runs at level 1 (tiering)+ ++ replicate 8 (2^(l1+ 2), 4) -- 8 runs at level 2 (tiering)+ ++ replicate 8 (2^(l1+ 4), 16) -- 8 runs at level 3 (tiering)+ ++ replicate 8 (2^(l1+ 6), 64) -- 8 runs at level 4 (tiering)+ ++ replicate 8 (2^(l1+ 8), 256) -- 8 runs at level 5 (tiering)+ ++ [(2^(l1+12),4096)] -- 1 run at level 6 (leveling)++-- | The total number of entries.+--+-- This should be roughly @100_000_000@ when we pass in @benchmarkSizeBase@.+-- >>> totalNumEntries (lsmStyleRuns benchmarkSizeBase)+-- 111804416+totalNumEntries :: [RunSizeInfo] -> Int+totalNumEntries runSizes =+ sum [ numEntries | (numEntries, _) <- runSizes ]++totalNumEntriesSanityCheck :: SizeBase -> [RunSizeInfo] -> Bool+totalNumEntriesSanityCheck l1 runSizes =+ totalNumEntries runSizes+ ==+ sum [ 2^l1 * sizeFactor | (_, sizeFactor) <- runSizes ]++withFS ::+ (FS.HasFS IO FS.HandleIO -> FS.HasBlockIO IO FS.HandleIO -> IO a)+ -> IO a+withFS action =+ FS.withIOHasBlockIO (FS.MountPoint "_bench_lookups") FS.defaultIOCtxParams $ \hfs hbio -> do+ exists <- FS.doesDirectoryExist hfs (FS.mkFsPath [""])+ unless exists $ error ("_bench_lookups directory does not exist")+ action hfs hbio++-- | Input environment for benchmarking lookup functions.+--+-- The idea here is to have a collection of runs corresponding to the sizes used+-- in a largeish LSM. In particular, the sizes are in increasing powers of 4.+-- The keys in the runs are non-overlapping, and lookups will be true positives+-- in only one run. Thus most lookups will be true negatives.+--+-- The goal is to benchmark the critical path of performing asynchronous lookups+-- serially.+--+-- * where the same key is being looked up in many runs,+-- * with a true positive lookup in only one run,+-- * with true negative lookups in the other runs+-- * with false positives lookups in a fraction of the runs according to the+-- bloom filters' false positive rates+lookupsEnv ::+ [RunSizeInfo]+ -> StdGen -- ^ Key RNG+ -> FS.HasFS IO FS.HandleIO+ -> FS.HasBlockIO IO FS.HandleIO+ -> Run.RunDataCaching+ -> IO ( V.Vector (Ref (Run IO FS.HandleIO))+ , V.Vector (Bloom SerialisedKey)+ , V.Vector Index+ , V.Vector (FS.Handle FS.HandleIO)+ )+lookupsEnv runSizes keyRng0 hfs hbio caching = do+ -- create the vector of initial keys+ (mvec :: VUM.MVector RealWorld UTxOKey) <- VUM.unsafeNew (totalNumEntries runSizes)+ !keyRng1 <- vectorOfUniforms mvec keyRng0+ -- we reuse keyRng0 to generate batches of lookups, so by shuffling the+ -- vector we ensure that these batches of lookups will do random disk+ -- access.+ !_ <- shuffle mvec keyRng1++ -- create the runs+ rbs <- sequence+ [ RunBuilder.new hfs hbio benchSalt+ RunParams {+ runParamCaching = caching,+ runParamAlloc = RunAllocFixed benchmarkNumBitsPerEntry,+ runParamIndex = Index.Compact+ }+ (RunFsPaths (FS.mkFsPath []) (RunNumber i))+ (NumEntries numEntries)+ | ((numEntries, _), i) <- zip runSizes [0..] ]++ -- fill the runs+ putStr "addKeyOp"+ let zero = serialiseValue zeroUTxOValue+ foldM_+ (\ !i (!rb, !n) -> do+ let !mvecLocal = VUM.unsafeSlice i n mvec+ Merge.sort mvecLocal+ flip VUM.imapM_ mvecLocal $ \ !j !k -> do+ -- progress+ when (j .&. 0xFFFF == 0) (putStr ".")+ void $ RunBuilder.addKeyOp rb (serialiseKey k) (Insert zero)+ pure (i+n)+ )+ 0+ (zip rbs (fmap fst runSizes))+ putStr "DONE"++ -- return runs+ runs <- V.fromList <$> mapM Run.fromBuilder rbs+ let blooms = V.map (\(DeRef r) -> Run.runFilter r) runs+ indexes = V.map (\(DeRef r) -> Run.runIndex r) runs+ handles = V.map (\(DeRef r) -> Run.runKOpsFile r) runs+ pure $!! (runs, blooms, indexes, handles)++genLookupBatch :: StdGen -> Int -> (V.Vector SerialisedKey, StdGen)+genLookupBatch !rng0 !n0+ | n0 <= 0 = error "mkBatch: must be positive"+ | otherwise = runST $ do+ mres <- VM.unsafeNew n0+ go rng0 0 mres+ where+ go ::+ StdGen -> Int -> VM.MVector s SerialisedKey+ -> ST s (V.Vector SerialisedKey, StdGen)+ go !rng !i !mres+ | n0 == i = do+ !res <- V.unsafeFreeze mres+ pure (res, rng)+ | otherwise = do+ let (!k, !rng') = uniform @UTxOKey @StdGen rng+ !sk = serialiseKey k+ VM.write mres i $! sk+ go rng' (i+1) mres++-- | This gives us the baseline cost of calculating batches of keys.+benchGenKeyBatches ::+ V.Vector (Bloom SerialisedKey)+ -> StdGen+ -> Int+ -> ()+benchGenKeyBatches !bs !keyRng !n+ | n <= 0 = ()+ | otherwise =+ let (!_ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize+ in benchGenKeyBatches bs keyRng' (n-benchmarkGenBatchSize)++-- | This gives us the combined cost of calculating batches of keys, and+-- performing bloom queries for each batch.+benchBloomQueries ::+ V.Vector (Bloom SerialisedKey)+ -> StdGen+ -> Int+ -> ()+benchBloomQueries !bs !keyRng !n+ | n <= 0 = ()+ | otherwise =+ let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize+ in bloomQueries benchSalt bs ks `seq`+ benchBloomQueries bs keyRng' (n-benchmarkGenBatchSize)++-- | This gives us the combined cost of calculating batches of keys, performing+-- bloom queries for each batch, and performing index searches for each batch.+benchIndexSearches ::+ ArenaManager RealWorld+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector Index+ -> V.Vector (FS.Handle h)+ -> StdGen+ -> Int+ -> IO ()+benchIndexSearches !arenaManager !bs !ics !hs !keyRng !n+ | n <= 0 = pure ()+ | otherwise = do+ let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize+ !rkixs = bloomQueries benchSalt bs ks+ !_ioops <- withArena arenaManager $ \arena -> stToIO $ indexSearches arena ics hs ks rkixs+ benchIndexSearches arenaManager bs ics hs keyRng' (n-benchmarkGenBatchSize)++-- | This gives us the combined cost of calculating batches of keys, and+-- preparing lookups for each batch.+benchPrepLookups ::+ ArenaManager RealWorld+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector Index+ -> V.Vector (FS.Handle h)+ -> StdGen+ -> Int+ -> IO ()+benchPrepLookups !arenaManager !bs !ics !hs !keyRng !n+ | n <= 0 = pure ()+ | otherwise = do+ let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize+ (!_rkixs, !_ioops) <- withArena arenaManager $ \arena -> stToIO $ prepLookups arena benchSalt bs ics hs ks+ benchPrepLookups arenaManager bs ics hs keyRng' (n-benchmarkGenBatchSize)++-- | This gives us the combined cost of calculating batches of keys, and+-- performing disk lookups for each batch.+benchLookupsIO ::+ FS.HasBlockIO IO h+ -> ArenaManager RealWorld+ -> ResolveSerialisedValue+ -> WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs IO h)+ -> V.Vector (Ref (Run IO h))+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector Index+ -> V.Vector (FS.Handle h)+ -> StdGen+ -> Int+ -> IO ()+benchLookupsIO !hbio !arenaManager !resolve !wb !wbblobs !rs !bs !ics !hs =+ go+ where+ go !keyRng !n+ | n <= 0 = pure ()+ | otherwise = do+ let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize+ !_ <- lookupsIOWithWriteBuffer+ hbio arenaManager resolve benchSalt wb wbblobs rs bs ics hs ks+ go keyRng' (n-benchmarkGenBatchSize)++{-------------------------------------------------------------------------------+ Utilities+-------------------------------------------------------------------------------}++classifyLookups ::+ V.Vector (Bloom SerialisedKey)+ -> StdGen+ -> Int+ -> ( Int, Double -- (all) positives, fpr+ , Int, Int -- true positives, false negatives+ , Int, Int -- false positives, true negatives+ )+classifyLookups !bs !keyRng0 !n0 =+ let !positives = unsafePerformIO (putStr "classifyLookups")+ `seq` loop 0 keyRng0 n0+ !tp = n0+ !fn = 0+ !fp = positives - tp+ !tn = (V.length bs - 1) * n0+ !fpr = fromIntegral fp / (fromIntegral fp + fromIntegral tn)+ in positives `seq`+ ( positives, fpr+ , tp, fn+ , fp, tn)+ where+ loop !positives !keyRng !n+ | n <= 0 =+ unsafePerformIO (putStr "DONE") `seq`+ positives+ | otherwise =+ unsafePerformIO (putStr ".") `seq`+ let (!ks, !keyRng') = genLookupBatch keyRng benchmarkGenBatchSize+ !rkixs = bloomQueries benchSalt bs ks+ in loop (positives + VP.length rkixs) keyRng' (n-benchmarkGenBatchSize)++-- | Fill a mutable vector with uniformly random values.+vectorOfUniforms ::+ (PrimMonad m, VGM.MVector v a, RandomGen g, Uniform a)+ => v (PrimState m) a+ -> g+ -> m g+vectorOfUniforms !vec !g0 = do+ unsafeIOToPrim $ putStr "vectorOfUniforms"+ !g0' <- loop 0 g0+ unsafeIOToPrim $ putStr "DONE"+ pure g0'+ where+ !n = VGM.length vec+ loop !i !g+ | i == n-1 = pure g+ | otherwise = do+ when (i .&. 0xFFFF == 0) (unsafeIOToPrim $ putStr ".")+ let (!x, !g') = uniform g+ VGM.unsafeWrite vec i x+ loop (i+1) g'++-- https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle+shuffle ::+ (PrimMonad m, VGM.MVector v a, RandomGen g)+ => v (PrimState m) a+ -> g+ -> m g+shuffle !xs !g0 = do+ unsafeIOToPrim $ putStr "shuffle"+ !g0' <- loop 0 g0+ unsafeIOToPrim $ putStr "DONE"+ pure g0'+ where+ !n = VGM.length xs+ loop !i !g+ | i == n-1 = pure g+ | otherwise = do+ when (i .&. 0xFFFF == 0) (unsafeIOToPrim $ putStr ".")+ let (!j, !g') = randomR (i, n-1) g+ VGM.swap xs i j+ loop (i+1) g'
@@ -0,0 +1,1125 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DuplicateRecordFields #-}+{-# LANGUAGE OverloadedStrings #-}+{-# OPTIONS_GHC -Wno-orphans #-}++{- Benchmark requirements:++A. The benchmark should use the external interface of the disk backend,+ and no internal interfaces.+B. The benchmark should use a workload ratio of 1 insert, to 1 delete,+ to 1 lookup. This is the workload ratio for the UTxO. Thus the+ performance is to be evaluated on the combination of the operations,+ not on operations individually.+C. The benchmark should use 34 byte keys, and 60 byte values. This+ corresponds roughly to the UTxO.+D. The benchmark should use keys that are evenly spread through the+ key space, such as cryptographic hashes.+E. The benchmark should start with a table of 100 million entries.+ This corresponds to the stretch target for the UTxO size.+ This table may be pre-generated. The time to construct the table+ should not be included in the benchmark time.+F. The benchmark workload should ensure that all inserts are for+ `fresh' keys, and that all lookups and deletes are for keys that+ are present. This is the typical workload for the UTxO.+G. It is acceptable to pre-generate the sequence of operations for the+ benchmark, or to take any other measure to exclude the cost of+ generating or reading the sequence of operations.+H. The benchmark should use the external interface of the disk backend+ to present batches of operations: a first batch consisting of 256+ lookups, followed by a second batch consisting of 256 inserts plus+ 256 deletes. This corresponds to the UTxO workload using 64kb+ blocks, with 512 byte txs with 2 inputs and 2 outputs.+I. The benchmark should be able to run in two modes, using the+ external interface of the disk backend in two ways: serially (in+ batches), or fully pipelined (in batches).+-}+module Main (main) where++import Control.Applicative ((<**>))+import Control.Concurrent (getNumCapabilities)+import Control.Concurrent.Async+import Control.Concurrent.MVar+import Control.DeepSeq (force)+import Control.Exception+import Control.Monad (forM_, unless, void, when)+import Control.Monad.Trans.State.Strict (runState, state)+import Control.Tracer+import qualified Data.ByteString.Short as BS+import qualified Data.Foldable as Fold+import qualified Data.IntSet as IS+import Data.IORef+import qualified Data.List.NonEmpty as NE+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Monoid+import qualified Data.Primitive as P+import qualified Data.Vector as V+import Data.Void (Void)+import Data.Word (Word32, Word64)+import qualified GHC.Stats as GHC+import qualified MCG+import qualified Options.Applicative as O+import Prelude hiding (lookup)+import qualified System.Clock as Clock+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.IO as FsIO+import System.IO+import System.Mem (performMajorGC)+import qualified System.Random as Random+import Text.Printf (printf)+import Text.Show.Pretty++import Database.LSMTree.Extras (groupsOfN)++-- We should be able to write this benchmark+-- using only use public lsm-tree interface+import qualified Database.LSMTree as LSM++-------------------------------------------------------------------------------+-- Table configuration+-------------------------------------------------------------------------------++benchTableConfig :: LSM.TableConfig+benchTableConfig =+ LSM.defaultTableConfig {LSM.confFencePointerIndex = LSM.CompactIndex}++benchSalt :: LSM.Salt+benchSalt = 4++-------------------------------------------------------------------------------+-- Keys and values+-------------------------------------------------------------------------------++type K = BS.ShortByteString+type V = BS.ShortByteString+type B = Void++deriving via LSM.ResolveAsFirst V instance LSM.ResolveValue V++label :: LSM.SnapshotLabel+label = LSM.SnapshotLabel "K V B"++-- | We generate 34 byte keys by using a PRNG to extend a word64 to 32 bytes+-- and then appending two constant bytes. This corresponds relatively closely+-- to UTxO keys, which are 32 byte cryptographic hashes, followed by two bytes+-- which are typically the 16bit value 0 or 1 (a transaction output index).+--+makeKey :: Word64 -> K+makeKey seed =+ case P.runPrimArray $ do+ v <- P.newPrimArray 5+ let g0 = Random.mkStdGen (fromIntegral seed)+ let (!w0, !g1) = Random.uniform g0+ P.writePrimArray v 0 w0+ let (!w1, !g2) = Random.uniform g1+ P.writePrimArray v 1 w1+ let (!w2, !g3) = Random.uniform g2+ P.writePrimArray v 2 w2+ let (!w3, _g4) = Random.uniform g3+ P.writePrimArray v 3 w3+ P.writePrimArray v 4 0x3d3d3d3d3d3d3d3d -- ========+ case v of+ P.MutablePrimArray mba -> do+ _ <- P.resizeMutableByteArray (P.MutableByteArray mba) 34+ pure v++ of (P.PrimArray ba :: P.PrimArray Word64) ->+ byteArrayToSBS (P.ByteArray ba)++-- | \( O(1) \) conversion.+byteArrayToSBS :: P.ByteArray -> BS.ShortByteString+#if MIN_VERSION_bytestring(0,12,0)+byteArrayToSBS ba = BS.ShortByteString ba+#else+byteArrayToSBS (P.ByteArray ba) = BS.SBS ba+#endif++-- We use constant value. This shouldn't affect anything.+theValue :: V+theValue = BS.replicate 60 120 -- 'x'+{-# NOINLINE theValue #-}++-------------------------------------------------------------------------------+-- Options and commands+-------------------------------------------------------------------------------++data GlobalOpts = GlobalOpts+ { rootDir :: !FilePath -- ^ session directory.+ , initialSize :: !Int+ -- | The cache policy for the LSM table. This configuration option is used+ -- both during setup, and during a run (where it is used to override the+ -- config option of the snapshot).+ , diskCachePolicy :: !LSM.DiskCachePolicy+ -- | Enable trace output+ , trace :: !Bool+ }+ deriving stock Show++data SetupOpts = SetupOpts {+ bloomFilterAlloc :: !LSM.BloomFilterAlloc+ , writeBufferAlloc :: !LSM.WriteBufferAlloc+ }+ deriving stock Show++data RunOpts = RunOpts+ { batchCount :: !Int+ , batchSize :: !Int+ , check :: !Bool+ , seed :: !Word64+ , mode :: !RunMode+ }+ deriving stock Show++data Cmd+ -- | Setup benchmark: generate initial LSM tree etc.+ = CmdSetup SetupOpts++ -- | Make a dry run, measure the overhead.+ | CmdDryRun RunOpts++ -- | Run the actual benchmark+ | CmdRun RunOpts+ deriving stock Show++mkTableConfigSetup :: GlobalOpts -> SetupOpts -> LSM.TableConfig -> LSM.TableConfig+mkTableConfigSetup GlobalOpts{diskCachePolicy}+ SetupOpts {+ bloomFilterAlloc,+ writeBufferAlloc+ } conf =+ conf {+ LSM.confDiskCachePolicy = diskCachePolicy+ , LSM.confBloomFilterAlloc = bloomFilterAlloc+ , LSM.confWriteBufferAlloc = writeBufferAlloc+ }++mkTableConfigRun :: GlobalOpts -> RunOpts -> LSM.TableConfig -> LSM.TableConfig+mkTableConfigRun GlobalOpts{diskCachePolicy} RunOpts {mode} conf =+ conf {+ LSM.confDiskCachePolicy = diskCachePolicy,+ LSM.confMergeBatchSize =+ case mode of+ Sequential -> LSM.confMergeBatchSize conf+ Pipelined -> LSM.MergeBatchSize 1+ LookupOnly -> LSM.confMergeBatchSize conf+ UpdateOnly -> LSM.confMergeBatchSize conf+ }++mkTableConfigOverride :: GlobalOpts -> RunOpts -> LSM.TableConfigOverride+mkTableConfigOverride GlobalOpts{diskCachePolicy} RunOpts {mode} =+ LSM.noTableConfigOverride {+ LSM.overrideDiskCachePolicy = Just diskCachePolicy,+ LSM.overrideMergeBatchSize =+ case mode of+ Sequential -> Nothing+ Pipelined -> Just $ LSM.MergeBatchSize 1+ LookupOnly -> Nothing+ UpdateOnly -> Nothing+ }++mkTracer :: GlobalOpts -> Tracer IO LSM.LSMTreeTrace+mkTracer gopts+ | trace gopts =+ -- Don't trace update/lookup messages, because they are too noisy+ squelchUnless+ (\case+ LSM.TraceTable _ LSM.TraceUpdates{} -> False+ LSM.TraceTable _ LSM.TraceLookups{} -> False+ _ -> True )+ (show `contramap` stdoutTracer)+ | otherwise = nullTracer++-------------------------------------------------------------------------------+-- command line interface+-------------------------------------------------------------------------------++globalOptsP :: O.Parser GlobalOpts+globalOptsP = pure GlobalOpts+ <*> O.option O.str (O.long "bench-dir" <> O.value "_bench_utxo" <> O.showDefault <> O.help "Benchmark directory to put files in")+ <*> O.option O.auto (O.long "initial-size" <> O.value 100_000_000 <> O.showDefault <> O.help "Initial LSM tree size")+ <*> O.option O.auto (O.long "disk-cache-policy" <> O.value LSM.DiskCacheAll <> O.showDefault <> O.help "Disk cache policy [DiskCacheAll | DiskCacheLevelOneTo Int | DiskCacheNone]")+ <*> O.flag False True (O.long "trace" <> O.help "Enable trace messages (disabled by default)")++cmdP :: O.Parser Cmd+cmdP = O.subparser $ mconcat+ [ O.command "setup" $ O.info+ (CmdSetup <$> setupOptsP <**> O.helper)+ (O.progDesc "Setup benchmark")+ , O.command "dry-run" $ O.info+ (CmdDryRun <$> runOptsP <**> O.helper)+ (O.progDesc "Dry run, measure overhead")++ , O.command "run" $ O.info+ (CmdRun <$> runOptsP <**> O.helper)+ (O.progDesc "Proper run")+ ]++setupOptsP :: O.Parser SetupOpts+setupOptsP =+ pure SetupOpts+ <*> O.option+ O.auto+ (O.long "bloom-filter-alloc" <>+ O.value (LSM.confBloomFilterAlloc LSM.defaultTableConfig) <>+ O.showDefault <>+ O.help "Bloom filter allocation method [AllocFixed n | AllocRequestFPR d]")++ <*> LSM.AllocNumEntries `fmap` O.option+ (O.auto)+ (O.long "write-buffer-alloc" <>+ O.value defaultAllocNumEntries <>+ O.showDefault <>+ O.help "Write buffer size")+ where+ LSM.AllocNumEntries defaultAllocNumEntries =+ LSM.confWriteBufferAlloc LSM.defaultTableConfig++runOptsP :: O.Parser RunOpts+runOptsP = pure RunOpts+ <*> O.option O.auto (O.long "batch-count" <> O.value 200 <> O.showDefault <> O.help "Batch count")+ <*> O.option O.auto (O.long "batch-size" <> O.value 256 <> O.showDefault <> O.help "Batch size")+ <*> O.switch (O.long "check" <> O.help "Check generated key distribution")+ <*> O.option O.auto (O.long "seed" <> O.value 1337 <> O.showDefault <> O.help "Random seed")+ <*> O.option O.auto (O.long "mode" <> O.value Sequential <> O.showDefault+ <> O.help "Mode [Sequential | Pipelined | LookupOnly | UpdateOnly]")++-- | The run mode affects the method of performing+data RunMode =+ -- | Use sequential mode: wait for a batch of operations to complete before+ -- starting the next one.+ Sequential+ -- | Use pipelined mode: overlap batches of operations.+ | Pipelined+ -- | Use lookup-only sequential mode: like sequential mode, but only perform+ -- the lookups and not the updates.+ | LookupOnly+ -- | Use update-only sequential mode: like sequential mode, but only perform+ -- the updates and not the lookups.+ | UpdateOnly+ deriving stock (Show, Eq, Read)++deriving stock instance Read LSM.DiskCachePolicy+deriving stock instance Read LSM.BloomFilterAlloc++-------------------------------------------------------------------------------+-- measurements+-------------------------------------------------------------------------------++timed :: IO a -> IO (a, Double, RTSStatsDiff Triple, ProcIODiff)+timed action = do+ !p1 <- getProcIO+ performMajorGC+ s1 <- GHC.getRTSStats+ t1 <- Clock.getTime Clock.Monotonic+ x <- action+ t2 <- Clock.getTime Clock.Monotonic+ performMajorGC+ s2 <- GHC.getRTSStats+ !p2 <- getProcIO+ let !t = fromIntegral (Clock.toNanoSecs (Clock.diffTimeSpec t2 t1)) * 1e-9+ !s = s2 `diffRTSStats` s1+ !p = p2 `diffProcIO` p1+ printf "Running time: %.03f sec\n" t+ printf "/proc/self/io after vs. before: %s\n" (ppShow p)+ printf "RTSStats after vs. before: %s\n" (ppShow s)+ pure (x, t, s, p)++timed_ :: IO () -> IO (Double, RTSStatsDiff Triple, ProcIODiff)+timed_ action = do+ ((), t, sdiff, pdiff) <- timed action+ pure (t, sdiff, pdiff)++-- | This /should/ include the statistics of any child processes.+getProcIO :: IO ProcIO+getProcIO = do+ s <- readFile "/proc/self/io"+ let ss = concatMap words $ lines s+ pure $ parse ss+ where+ parse [+ "rchar:", rcharS+ , "wchar:", wcharS+ , "syscr:", syscrS+ , "syscw:", syscwS+ , "read_bytes:", read_bytesS+ , "write_bytes:", write_bytesS+ , "cancelled_write_bytes:", cancellled_write_bytesS+ ] = ProcIO {+ rchar = read rcharS+ , wchar = read wcharS+ , syscr = read syscrS+ , syscw = read syscwS+ , read_bytes = read read_bytesS+ , write_bytes = read write_bytesS+ , cancelled_write_bytes = read cancellled_write_bytesS+ }+ parse s = error $ "getProcIO: parse of /proc/self/io failed. Input is " <> show s++diffProcIO :: ProcIO -> ProcIO -> ProcIODiff+diffProcIO after before = ProcIODiff ProcIO {+ rchar = subtractOn rchar+ , wchar = subtractOn wchar+ , syscr = subtractOn syscr+ , syscw = subtractOn syscw+ , read_bytes = subtractOn read_bytes+ , write_bytes = subtractOn write_bytes+ , cancelled_write_bytes = subtractOn cancelled_write_bytes+ }+ where+ subtractOn f = f after - f before++newtype ProcIODiff = ProcIODiff ProcIO+ deriving stock Show++-- | See the @/proc/[pid]/io@ section in @man proc@+data ProcIO = ProcIO {+ rchar :: !Integer+ , wchar :: !Integer+ , syscr :: !Integer+ , syscw :: !Integer+ , read_bytes :: !Integer+ , write_bytes :: !Integer+ , cancelled_write_bytes :: !Integer+ }+ deriving stock Show++-- | 'diffRTSStats a b = b - a'+diffRTSStats :: GHC.RTSStats -> GHC.RTSStats -> RTSStatsDiff Triple+diffRTSStats after before = RTSStatsDiff {+ gcs = subtractOn GHC.gcs+ , major_gcs = subtractOn GHC.major_gcs+ , allocated_bytes = subtractOn GHC.allocated_bytes+ , max_live_bytes = subtractOn GHC.max_live_bytes+ , max_large_objects_bytes = subtractOn GHC.max_large_objects_bytes+ , max_compact_bytes = subtractOn GHC.max_compact_bytes+ , max_slop_bytes = subtractOn GHC.max_slop_bytes+ , max_mem_in_use_bytes = subtractOn GHC.max_mem_in_use_bytes+ , cumulative_live_bytes = subtractOn GHC.cumulative_live_bytes+ , copied_bytes = subtractOn GHC.copied_bytes+ , par_copied_bytes = subtractOn GHC.par_copied_bytes+ , cumulative_par_balanced_copied_bytes = subtractOn GHC.cumulative_par_balanced_copied_bytes+ , init_cpu_ns = subtractOn GHC.init_cpu_ns+ , init_elapsed_ns = subtractOn GHC.init_elapsed_ns+ , mutator_cpu_ns = subtractOn GHC.mutator_cpu_ns+ , mutator_elapsed_ns = subtractOn GHC.mutator_elapsed_ns+ , gc_cpu_ns = subtractOn GHC.gc_cpu_ns+ , gc_elapsed_ns = subtractOn GHC.gc_elapsed_ns+ , cpu_ns = subtractOn GHC.cpu_ns+ , elapsed_ns = subtractOn GHC.elapsed_ns+ }+ where+ subtractOn :: Num a => (GHC.RTSStats -> a) -> Triple a+ subtractOn f = Triple {before = x, after = y, difference = y - x}+ where x = f before+ y = f after++-- | A difference datatype for 'GHC.RTSStats'.+--+-- Most fields, like 'GHC.gcs' or 'GHC.cpu_ns', are an aggregate sum, and so a+-- diff can be computed by pointwise subtraction.+--+-- Others fields, like 'GHC.max_live_bytes' only record the maximum value thus+-- far seen. We report a triplet containing the maximum before and after, and+-- their difference.+data RTSStatsDiff f = RTSStatsDiff {+ gcs :: !(f Word32)+ , major_gcs :: !(f Word32)+ , allocated_bytes :: !(f Word64)+ , max_live_bytes :: !(f Word64)+ , max_large_objects_bytes :: !(f Word64)+ , max_compact_bytes :: !(f Word64)+ , max_slop_bytes :: !(f Word64)+ , max_mem_in_use_bytes :: !(f Word64)+ , cumulative_live_bytes :: !(f Word64)+ , copied_bytes :: !(f Word64)+ , par_copied_bytes :: !(f Word64)+ , cumulative_par_balanced_copied_bytes :: !(f Word64)+ , init_cpu_ns :: !(f GHC.RtsTime)+ , init_elapsed_ns :: !(f GHC.RtsTime)+ , mutator_cpu_ns :: !(f GHC.RtsTime)+ , mutator_elapsed_ns :: !(f GHC.RtsTime)+ , gc_cpu_ns :: !(f GHC.RtsTime)+ , gc_elapsed_ns :: !(f GHC.RtsTime)+ , cpu_ns :: !(f GHC.RtsTime)+ , elapsed_ns :: !(f GHC.RtsTime)+ }++deriving stock instance Show (RTSStatsDiff Triple)++data Triple a = Triple {+ before :: !a+ , after :: !a+ , difference :: !a+ }+ deriving stock Show++-------------------------------------------------------------------------------+-- setup+-------------------------------------------------------------------------------++doSetup :: GlobalOpts -> SetupOpts -> IO ()+doSetup gopts opts = do+ void $ timed_ $ doSetup' gopts opts++doSetup' :: GlobalOpts -> SetupOpts -> IO ()+doSetup' gopts opts =+ FsIO.withIOHasBlockIO mountPoint FsIO.defaultIOCtxParams $ \hasFS hasBlockIO ->+ LSM.withOpenSession (mkTracer gopts) hasFS hasBlockIO benchSalt (FS.mkFsPath []) $ \session -> do+ tbl <- LSM.newTableWith @IO @K @V @B (mkTableConfigSetup gopts opts benchTableConfig) session++ forM_ (groupsOfN 256 [ 0 .. initialSize gopts ]) $ \batch -> do+ -- TODO: this procedure simply inserts all the keys into initial lsm tree+ -- We might want to do deletes, so there would be delete-insert pairs+ -- Let's do that when we can actually test that benchmark works.+ LSM.inserts tbl $ V.fromList [+ (makeKey (fromIntegral i), theValue, Nothing)+ | i <- NE.toList batch+ ]++ LSM.saveSnapshot name label tbl+ where+ mountPoint :: FS.MountPoint+ mountPoint = FS.MountPoint (rootDir gopts)++ name = LSM.toSnapshotName ("bench_" ++ show (initialSize gopts))+++-------------------------------------------------------------------------------+-- dry-run+-------------------------------------------------------------------------------++doDryRun :: GlobalOpts -> RunOpts -> IO ()+doDryRun gopts opts = do+ void $ timed_ $ doDryRun' gopts opts++doDryRun' :: GlobalOpts -> RunOpts -> IO ()+doDryRun' gopts opts = do+ -- calculated some expected statistics for generated batches+ -- using nested do block to limit scope of intermediate bindings n, d, p, and q+ do+ -- we generate n random numbers in range of [ 1 .. d ]+ -- what is the chance they are all distinct+ let n = fromIntegral (batchCount opts * batchSize opts) :: Double+ let d = fromIntegral (initialSize gopts) :: Double+ -- this is birthday problem.+ let p = 1 - exp (negate $ (n * (n - 1)) / (2 * d))++ -- number of people with a shared birthday+ -- https://en.wikipedia.org/wiki/Birthday_problem#Number_of_people_with_a_shared_birthday+ let q = n * (1 - ((d - 1) / d) ** (n - 1))++ printf "Probability of a duplicate: %5f\n" p+ printf "Expected number of duplicates (extreme upper bound): %5f out of %f\n" q n++ let g0 = initGen (initialSize gopts) (batchSize opts) (batchCount opts) (seed opts)++ keysRef <- newIORef $+ if check opts+ then IS.fromList [ 0 .. (initialSize gopts) - 1 ]+ else IS.empty+ duplicateRef <- newIORef (0 :: Int)++ void $ forFoldM_ g0 [ 0 .. batchCount opts - 1 ] $ \b g -> do+ let lookups :: V.Vector Word64+ inserts :: V.Vector Word64+ (!g', lookups, inserts) = generateBatch' (initialSize gopts) (batchSize opts) g b++ when (check opts) $ do+ keys <- readIORef keysRef+ let new = intSetFromVector lookups+ let diff = IS.difference new keys+ -- when (IS.notNull diff) $ printf "missing in batch %d %s\n" b (show diff)+ modifyIORef' duplicateRef $ \n -> n + IS.size diff+ writeIORef keysRef $! IS.union (IS.difference keys new)+ (intSetFromVector inserts)++ let (batch1, batch2) = toOperations lookups inserts+ _ <- evaluate $ force (batch1, batch2)++ pure g'++ when (check opts) $ do+ duplicates <- readIORef duplicateRef+ printf "True duplicates: %d\n" duplicates++ -- See batchOverlaps for explanation of this check.+ when (check opts) $+ let anyOverlap = (not . null)+ (batchOverlaps (initialSize gopts) (batchSize opts)+ (batchCount opts) (seed opts))+ in putStrLn $ "Any adjacent batches with overlap: " ++ show anyOverlap++-------------------------------------------------------------------------------+-- PRNG initialisation+-------------------------------------------------------------------------------++initGen :: Int -> Int -> Int -> Word64 -> MCG.MCG+initGen initialSize batchSize batchCount seed =+ let period = initialSize + batchSize * batchCount+ in MCG.make (fromIntegral period) seed++-------------------------------------------------------------------------------+-- Batch generation+-------------------------------------------------------------------------------++generateBatch ::+ Int -- ^ initial size of the collection+ -> Int -- ^ batch size+ -> MCG.MCG -- ^ generator+ -> Int -- ^ batch number+ -> (MCG.MCG, V.Vector K, V.Vector (K, LSM.Update V B))+generateBatch initialSize batchSize g b =+ (g', lookups', inserts')+ where+ (lookups', inserts') = toOperations lookups inserts+ (!g', lookups, inserts) = generateBatch' initialSize batchSize g b++{- | Implement generation of unbounded sequence of insert\/delete operations++matching UTxO style from spec: interleaved batches insert and lookup+configurable batch sizes+1 insert, 1 delete, 1 lookup per key.++Current approach is probabilistic, but uses very little state.+We could also make it exact, but then we'll need to carry some state around+(at least the difference).++-}+{-# INLINE generateBatch' #-}+generateBatch' ::+ Int -- ^ initial size of the collection+ -> Int -- ^ batch size+ -> MCG.MCG -- ^ generator+ -> Int -- ^ batch number+ -> (MCG.MCG, V.Vector Word64, V.Vector Word64)+generateBatch' initialSize batchSize g b = (g'', lookups, inserts)+ where+ maxK :: Word64+ maxK = fromIntegral $ initialSize + batchSize * b++ lookups :: V.Vector Word64+ (lookups, !g'') =+ runState (V.replicateM batchSize (state (MCG.reject maxK))) g++ inserts :: V.Vector Word64+ inserts = V.enumFromTo maxK (maxK + fromIntegral batchSize - 1)++-- | Generate operation inputs+{-# INLINE toOperations #-}+toOperations :: V.Vector Word64 -> V.Vector Word64 -> (V.Vector K, V.Vector (K, LSM.Update V B))+toOperations lookups inserts = (batch1, batch2)+ where+ batch1 :: V.Vector K+ batch1 = V.map makeKey lookups++ batch2 :: V.Vector (K, LSM.Update V B)+ batch2 = V.map (\k -> (k, LSM.Delete)) batch1 V.+++ V.map (\k -> (makeKey k, LSM.Insert theValue Nothing)) inserts++-------------------------------------------------------------------------------+-- run+-------------------------------------------------------------------------------++doRun :: GlobalOpts -> RunOpts -> IO ()+doRun gopts opts =+ FsIO.withIOHasBlockIO mountPoint FsIO.defaultIOCtxParams $ \hasFS hasBlockIO ->+ LSM.withOpenSession (mkTracer gopts) hasFS hasBlockIO benchSalt (FS.mkFsPath []) $ \session ->+ withLatencyHandle $ \h -> do+ -- open snapshot+ -- In checking mode we start with an empty table, since our pure+ -- reference version starts with empty (as it's not practical or+ -- necessary for testing to load the whole snapshot).+ tbl <- if check opts+ then let conf = mkTableConfigRun gopts opts benchTableConfig+ in LSM.newTableWith @IO @K @V @B conf session+ else let conf = mkTableConfigOverride gopts opts+ in LSM.openTableFromSnapshotWith @IO @K @V @B conf session name label++ -- In checking mode, compare each output against a pure reference.+ checkvar <- newIORef $ pureReference+ (initialSize gopts) (batchSize opts)+ (batchCount opts) (seed opts)+ let fcheck | not (check opts) || mode opts == UpdateOnly = \_ _ -> pure ()+ | otherwise = \b y -> do+ (x:xs) <- readIORef checkvar+ unless (x == y) $+ fail $ "lookup result mismatch in batch " ++ show b+ writeIORef checkvar xs++ let benchmarkIterations = case mode opts of+ Sequential -> sequentialIterations h+ Pipelined -> pipelinedIterations h+ LookupOnly -> sequentialIterationsLO+ UpdateOnly -> sequentialIterationsUO+ !progressInterval = max 1 ((batchCount opts) `div` 100)+ madeProgress b = b `mod` progressInterval == 0+ (time, _, _) <- timed_ $ do+ benchmarkIterations+ (\b y -> fcheck b y >> when (madeProgress b) (putChar '.'))+ (initialSize gopts)+ (batchSize opts)+ (batchCount opts)+ (seed opts)+ tbl+ putStrLn ""++ let ops = batchCount opts * batchSize opts+ printf "Operations per second: %7.01f ops/sec\n" (fromIntegral ops / time)+ where+ mountPoint :: FS.MountPoint+ mountPoint = FS.MountPoint (rootDir gopts)++ name = LSM.toSnapshotName ("bench_" ++ show (initialSize gopts))++-------------------------------------------------------------------------------+-- sequential+-------------------------------------------------------------------------------++type LookupResults = V.Vector (K, LSM.LookupResult V ())++{-# INLINE sequentialIteration #-}+sequentialIteration :: LatencyHandle+ -> (Int -> LookupResults -> IO ())+ -> Int+ -> Int+ -> LSM.Table IO K V B+ -> Int+ -> MCG.MCG+ -> IO MCG.MCG+sequentialIteration h output !initialSize !batchSize !tbl !b !g =+ withTimedBatch h b $ \tref -> do+ let (!g', ls, is) = generateBatch initialSize batchSize g b++ -- lookups+ results <- timeLatency tref $ LSM.lookups tbl ls+ output b (V.zip ls (fmap (fmap (const ())) results))++ -- deletes and inserts+ _ <- timeLatency tref $ LSM.updates tbl is++ -- continue to the next batch+ pure g'++sequentialIterations :: LatencyHandle+ -> (Int -> LookupResults -> IO ())+ -> Int -> Int -> Int -> Word64+ -> LSM.Table IO K V B+ -> IO ()+sequentialIterations h output !initialSize !batchSize !batchCount !seed !tbl = do+ createGnuplotExampleFileSequential+ hPutHeaderSequential h+ void $ forFoldM_ g0 [ 0 .. batchCount - 1 ] $ \b g ->+ sequentialIteration h output initialSize batchSize tbl b g+ where+ g0 = initGen initialSize batchSize batchCount seed++{-# INLINE sequentialIterationLO #-}+sequentialIterationLO :: (Int -> LookupResults -> IO ())+ -> Int+ -> Int+ -> LSM.Table IO K V B+ -> Int+ -> MCG.MCG+ -> IO MCG.MCG+sequentialIterationLO output !initialSize !batchSize !tbl !b !g = do+ let (!g', ls, _is) = generateBatch initialSize batchSize g b++ -- lookups+ results <- LSM.lookups tbl ls+ output b (V.zip ls (fmap (fmap (const ())) results))++ -- continue to the next batch+ pure g'++sequentialIterationsLO :: (Int -> LookupResults -> IO ())+ -> Int -> Int -> Int -> Word64+ -> LSM.Table IO K V B+ -> IO ()+sequentialIterationsLO output !initialSize !batchSize !batchCount !seed !tbl =+ void $ forFoldM_ g0 [ 0 .. batchCount - 1 ] $ \b g ->+ sequentialIterationLO output initialSize batchSize tbl b g+ where+ g0 = initGen initialSize batchSize batchCount seed++{-# INLINE sequentialIterationUO #-}+sequentialIterationUO ::+ (Int -> LookupResults -> IO ())+ -> Int+ -> Int+ -> LSM.Table IO K V B+ -> Int+ -> MCG.MCG+ -> IO MCG.MCG+sequentialIterationUO output !initialSize !batchSize !tbl !b !g = do+ let (!g', _ls, is) = generateBatch initialSize batchSize g b++ -- lookups+ output b V.empty++ -- deletes and inserts+ _ <- LSM.updates tbl is++ -- continue to the next batch+ pure g'++sequentialIterationsUO ::+ (Int -> LookupResults -> IO ())+ -> Int -> Int -> Int -> Word64+ -> LSM.Table IO K V B+ -> IO ()+sequentialIterationsUO output !initialSize !batchSize !batchCount !seed !tbl = do+ void $ forFoldM_ g0 [ 0 .. batchCount - 1 ] $ \b g ->+ sequentialIterationUO output initialSize batchSize tbl b g+ where+ g0 = initGen initialSize batchSize batchCount seed++-------------------------------------------------------------------------------+-- pipelined+-------------------------------------------------------------------------------++{- One iteration of the protocol for one thread looks like this:++1. Lookups (db_n-1) tx_n+0+2. Sync ? (db_n+0, updates)+3. db_n+1 <- Dup (db_n+0)+ Updates (db_n+1) tx_n+0+4. Sync ! (db_n+1, updates)++Thus for two threads running iterations concurrently, it looks like this:++1. Lookups (db_n-1) tx_n+0 3. db_n+0 <- Dup (db_n-1)+ Updates (db_n+0) tx_n-1+2. Sync ? (db_n+0, updates) <- 4. Sync ! (db_n+0, updates)+3. db_n+1 <- Dup (db_n+0) 1. Lookups (db_n+0) tx_n+1+ Updates (db_n+1) tx_n+0+4. Sync ! (db_n+1, updates) -> 2. Sync ? (db_n+1, updates)+1. Lookups (db_n+1) tx_n+2 3. db_n+2 <- Dup (db_n+1)+ Updates (db_n+2) tx_n+1+2. Sync ? (db_n+2, updates) <- 4. Sync ! (db_n+2, updates)+3. db_n+3 <- Dup (db_n+2) 1. Lookups (db_n+2) tx_n+3+ Updates (db_n+3) tx_n+2+4. Sync ! (db_n+3, updates) -> 2. Sync ? (db_n+3, updates)+1. Lookups (db_n+3) tx_n+4 3. db_n+4 <- Dup (db_n+3)+ Updates (db_n+4) tx_n+3+2. Sync ? (db_n+4, updates) <- 4. Sync ! (db_n+4, updates)+3. db_n+5 <- Dup (db_n+4) 1. Lookups (db_n+4) tx_n+5+ Updates (db_n+5) tx_n+4+4. Sync ! (db_n+5, updates) -> 2. Sync ? (db_n+5, updates)++And the initial setup looks like this:++ db_1 <- Dup (db_0)+ Lookups (db_0) tx_0+ Updates (db_1, tx_0)+ Sync ! (db_1, updates) ->+ 1. Lookups (db_0) tx_1++ 2. Sync ? (db_1, updates)+1. Lookups (db_1) tx_2 3. db_2 <- Dup (db_1)+ Updates (db_2) tx_1+2. Sync ? (db_2, updates) <- 4. Sync ! (db_2, updates)+3. db_3 <- Dup (db_2) 1. Lookups (db_2) tx_3+ Updates (db_3) tx_2+4. Sync ! (db_3, updates) 2. Sync ? (db_3, updates)+-}+pipelinedIteration :: LatencyHandle+ -> (Int -> LookupResults -> IO ())+ -> Int+ -> Int+ -> MVar (LSM.Table IO K V B, Map K (LSM.Update V B))+ -> MVar (LSM.Table IO K V B, Map K (LSM.Update V B))+ -> MVar MCG.MCG+ -> MVar MCG.MCG+ -> LSM.Table IO K V B+ -> Int+ -> IO (LSM.Table IO K V B)+pipelinedIteration h output !initialSize !batchSize+ !syncTblIn !syncTblOut+ !syncRngIn !syncRngOut+ !tbl_n !b =+ withTimedBatch h b $ \tref -> do+ g <- takeMVar syncRngIn+ let (!g', !ls, !is) = generateBatch initialSize batchSize g b++ -- 1: perform the lookups+ lrs <- timeLatency tref $ LSM.lookups tbl_n ls++ -- 2. sync: receive updates and new table+ tbl_n1 <- timeLatency tref $ do+ putMVar syncRngOut g'+ (tbl_n1, delta) <- takeMVar syncTblIn++ -- At this point, after syncing, our peer is guaranteed to no longer be+ -- using tbl_n. They used it to generate tbl_n+1 (which they gave us).+ LSM.closeTable tbl_n+ output b $! applyUpdates delta (V.zip ls lrs)+ pure tbl_n1++ -- 3. perform the inserts and report outputs (in any order)+ tbl_n2 <- timeLatency tref $ do+ tbl_n2 <- LSM.duplicate tbl_n1+ LSM.updates tbl_n2 is+ pure tbl_n2++ -- 4. sync: send the updates and new table+ timeLatency tref $ do+ let delta' :: Map K (LSM.Update V B)+ !delta' = Map.fromList (V.toList is)+ putMVar syncTblOut (tbl_n2, delta')++ pure tbl_n2+ where+ applyUpdates :: Map K (LSM.Update V a)+ -> V.Vector (K, LSM.LookupResult V b)+ -> V.Vector (K, LSM.LookupResult V ())+ applyUpdates m lrs =+ flip V.map lrs $ \(k, lr) ->+ case Map.lookup k m of+ Nothing -> (k, fmap (const ()) lr)+ Just u -> (k, updateToLookupResult u)++pipelinedIterations :: LatencyHandle+ -> (Int -> LookupResults -> IO ())+ -> Int -> Int -> Int -> Word64+ -> LSM.Table IO K V B+ -> IO ()+pipelinedIterations h output !initialSize !batchSize !batchCount !seed tbl_0 = do+ createGnuplotExampleFilePipelined+ hPutHeaderPipelined h+ n <- getNumCapabilities+ printf "INFO: the pipelined benchmark is running with %d capabilities.\n" n++ syncTblA2B <- newEmptyMVar+ syncTblB2A <- newEmptyMVar+ syncRngA2B <- newEmptyMVar+ syncRngB2A <- newEmptyMVar++ let g0 = initGen initialSize batchSize batchCount seed++ tbl_1 <- LSM.duplicate tbl_0+ let prelude = do+ let (g1, ls0, is0) = generateBatch initialSize batchSize g0 0+ lrs0 <- LSM.lookups tbl_0 ls0+ output 0 $! V.zip ls0 (fmap (fmap (const ())) lrs0)+ LSM.updates tbl_1 is0+ let !delta = Map.fromList (V.toList is0)+ putMVar syncTblA2B (tbl_1, delta)+ putMVar syncRngA2B g1++ threadA =+ forFoldM_ tbl_1 [ 2, 4 .. batchCount - 1 ] $ \b tbl_n ->+ pipelinedIteration h output initialSize batchSize+ syncTblB2A syncTblA2B -- in, out+ syncRngB2A syncRngA2B -- in, out+ tbl_n b++ threadB =+ forFoldM_ tbl_0 [ 1, 3 .. batchCount - 1 ] $ \b tbl_n ->+ pipelinedIteration h output initialSize batchSize+ syncTblA2B syncTblB2A -- in, out+ syncRngA2B syncRngB2A -- in, out+ tbl_n b++ -- We do batch 0 as a special prelude to get the pipeline started...+ prelude+ -- Run the pipeline: batches 2,4,6... concurrently with batches 1,3,5...+ -- If run with +RTS -N2 then we'll put each thread on a separate core.+ withAsyncOn 0 threadA $ \ta ->+ withAsyncOn 1 threadB $ \tb ->+ waitBoth ta tb >> pure ()++-------------------------------------------------------------------------------+-- Testing+-------------------------------------------------------------------------------++pureReference :: Int -> Int -> Int -> Word64 -> [V.Vector (K, LSM.LookupResult V ())]+pureReference !initialSize !batchSize !batchCount !seed =+ generate g0 Map.empty 0+ where+ g0 = initGen initialSize batchSize batchCount seed++ generate !_ !_ !b | b == batchCount = []+ generate !g !m !b = results : generate g' m' (b+1)+ where+ (g', lookups, inserts) = generateBatch initialSize batchSize g b+ !results = V.map (lookup m) lookups+ !m' = Fold.foldl' (flip (uncurry Map.insert)) m inserts++ lookup m k =+ case Map.lookup k m of+ Nothing -> (,) k LSM.NotFound+ Just u -> (,) k $! updateToLookupResult u++updateToLookupResult :: LSM.Update v b -> LSM.LookupResult v ()+updateToLookupResult (LSM.Insert v Nothing) = LSM.Found v+updateToLookupResult (LSM.Insert v (Just _)) = LSM.FoundWithBlob v ()+updateToLookupResult LSM.Delete = LSM.NotFound+updateToLookupResult (LSM.Upsert _) = error $+ "Unexpected upsert encountered"++-- | Return the adjacent batches where there is overlap between one batch's+-- inserts and the next batch's lookups. Testing the pipelined version needs+-- some overlap to get proper coverage. So this function is used as a coverage+-- check.+--+batchOverlaps :: Int -> Int -> Int -> Word64 -> [IS.IntSet]+batchOverlaps initialSize batchSize batchCount seed =+ let xs = generate g0 0++ in filter (not . IS.null)+ $ map (\((_, is),(ls, _)) -> IS.intersection (intSetFromVector is)+ (intSetFromVector ls))+ $ zip xs (tail xs)+ where+ generate _ b | b == batchCount = []+ generate g b = (lookups, inserts) : generate g' (b+1)+ where+ (g', lookups, inserts) = generateBatch' initialSize batchSize g b++ g0 = initGen initialSize batchSize batchCount seed++-------------------------------------------------------------------------------+-- measure batch latency+-------------------------------------------------------------------------------++_mEASURE_BATCH_LATENCY :: Bool+#ifdef MEASURE_BATCH_LATENCY+_mEASURE_BATCH_LATENCY = True+#else+_mEASURE_BATCH_LATENCY = False+#endif++type LatencyHandle = Handle++type TimeRef = IORef [Integer]++withLatencyHandle :: (LatencyHandle -> IO a) -> IO a+withLatencyHandle k+ | _mEASURE_BATCH_LATENCY = withFile "latency.dat" WriteMode k+ | otherwise = k (error "LatencyHandle: do not use")++{-# INLINE hPutHeaderSequential #-}+hPutHeaderSequential :: LatencyHandle -> IO ()+hPutHeaderSequential h+ | _mEASURE_BATCH_LATENCY = do+ hPutStrLn h "# batch number \t lookup time (ns) \t update time (ns)"+ | otherwise = pure ()++{-# INLINE createGnuplotExampleFileSequential #-}+createGnuplotExampleFileSequential :: IO ()+createGnuplotExampleFileSequential+ | _mEASURE_BATCH_LATENCY = do+ withFile "latency.gp" WriteMode $ \h -> do+ mapM_ (hPutStrLn h) [+ "set title \"Latency (sequential)\""+ , ""+ , "set xlabel \"Batch number\""+ , ""+ , "set logscale y"+ , "set ylabel \"Time (nanoseconds)\""+ , ""+ , "plot \"latency.dat\" using 1:2 title 'lookups' axis x1y1, \\"+ , " \"latency.dat\" using 1:3 title 'updates' axis x1y1"+ ]+ | otherwise = pure ()++{-# INLINE hPutHeaderPipelined #-}+hPutHeaderPipelined :: LatencyHandle -> IO ()+hPutHeaderPipelined h+ | _mEASURE_BATCH_LATENCY = do+ hPutStr h "# batch number"+ hPutStr h "\t lookup time (ns) \t sync receive time (ns) "+ hPutStrLn h "\t update time (ns) \t sync send time (ns)"+ | otherwise = pure ()++{-# INLINE createGnuplotExampleFilePipelined #-}+createGnuplotExampleFilePipelined :: IO ()+createGnuplotExampleFilePipelined+ | _mEASURE_BATCH_LATENCY =+ withFile "latency.gp" WriteMode $ \h -> do+ mapM_ (hPutStrLn h) [+ "set title \"Latency (pipelined)\""+ , ""+ , "set xlabel \"Batch number\""+ , ""+ , "set logscale y"+ , "set ylabel \"Time (nanoseconds)\""+ , ""+ , "plot \"latency.dat\" using 1:2 title 'lookups' axis x1y1, \\"+ , " \"latency.dat\" using 1:3 title 'sync receive' axis x1y1, \\"+ , " \"latency.dat\" using 1:4 title 'updates' axis x1y1, \\"+ , " \"latency.dat\" using 1:5 title 'sync send' axis x1y1"+ ]+ | otherwise = pure ()++{-# INLINE withTimedBatch #-}+withTimedBatch :: LatencyHandle -> Int -> (TimeRef -> IO a) -> IO a+withTimedBatch h b k+ | _mEASURE_BATCH_LATENCY = do+ tref <- newIORef []+ x <- k tref+ ts <- readIORef tref+ let s = shows b+ . getDual (foldMap (\t -> Dual (showString "\t" <> shows t)) ts)+ hPutStrLn h (s "")+ pure x+ | otherwise = k (error "TimeRef: do not use")++{-# INLINE timeLatency #-}+timeLatency :: TimeRef -> IO a -> IO a+timeLatency tref k+ | _mEASURE_BATCH_LATENCY = do+ t1 <- Clock.getTime Clock.Monotonic+ x <- k+ t2 <- Clock.getTime Clock.Monotonic+ let !t = Clock.toNanoSecs (Clock.diffTimeSpec t2 t1)+ modifyIORef tref (t :)+ pure x+ | otherwise = k++-------------------------------------------------------------------------------+-- main+-------------------------------------------------------------------------------++main :: IO ()+main = do+ hSetBuffering stdout NoBuffering+#ifdef NO_IGNORE_ASSERTS+ putStrLn "WARNING: Benchmarking in debug mode."+ putStrLn " To benchmark in release mode, pass:"+ putStrLn " --project-file=cabal.project.release"+#endif+ (gopts, cmd) <- O.customExecParser prefs cliP+ print gopts+ print cmd+ case cmd of+ CmdSetup opts -> doSetup gopts opts+ CmdDryRun opts -> doDryRun gopts opts+ CmdRun opts -> doRun gopts opts+ where+ cliP = O.info ((,) <$> globalOptsP <*> cmdP <**> O.helper) O.fullDesc+ prefs = O.prefs $ O.showHelpOnEmpty <> O.helpShowGlobals <> O.subparserInline++-------------------------------------------------------------------------------+-- general utils+-------------------------------------------------------------------------------++forFoldM_ :: Monad m => s -> [a] -> (a -> s -> m s) -> m s+forFoldM_ !s0 xs0 f = go s0 xs0+ where+ go !s [] = pure s+ go !s (x:xs) = do+ !s' <- f x s+ go s' xs++intSetFromVector :: V.Vector Word64 -> IS.IntSet+intSetFromVector = V.foldl' (\acc x -> IS.insert (fromIntegral x) acc) IS.empty++-------------------------------------------------------------------------------+-- unused for now+-------------------------------------------------------------------------------++_unused :: ()+_unused = const ()+ timed
@@ -0,0 +1,423 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DuplicateRecordFields #-}+{-# LANGUAGE NondecreasingIndentation #-}+{-# LANGUAGE OverloadedStrings #-}++{- Benchmark requirements:++A. The benchmark should use the external interface of the disk backend,+ and no internal interfaces.+B. The benchmark should use a workload ratio of 1 insert, to 1 delete,+ to 1 lookup. This is the workload ratio for the UTxO. Thus the+ performance is to be evaluated on the combination of the operations,+ not on operations individually.+C. The benchmark should use 34 byte keys, and 60 byte values. This+ corresponds roughly to the UTxO.+D. The benchmark should use keys that are evenly spread through the+ key space, such as cryptographic hashes.+E. The benchmark should start with a table of 100 million entries.+ This corresponds to the stretch target for the UTxO size.+ This table may be pre-generated. The time to construct the table+ should not be included in the benchmark time.+F. The benchmark workload should ensure that all inserts are for+ `fresh' keys, and that all lookups and deletes are for keys that+ are present. This is the typical workload for the UTxO.+G. It is acceptable to pre-generate the sequence of operations for the+ benchmark, or to take any other measure to exclude the cost of+ generating or reading the sequence of operations.+H. The benchmark should use the external interface of the disk backend+ to present batches of operations: a first batch consisting of 256+ lookups, followed by a second batch consisting of 256 inserts plus+ 256 deletes. This corresponds to the UTxO workload using 64kb+ blocks, with 512 byte txs with 2 inputs and 2 outputs.+I. The benchmark should be able to run in two modes, using the+ external interface of the disk backend in two ways: serially (in+ batches), or fully pipelined (in batches).++TODO 2024-04-29 consider alternative methods of implementing key generation+TODO 2024-04-29 pipelined mode is not implemented.++-}+module Main (main) where++import Control.Applicative ((<**>))+import Control.DeepSeq (force)+import Control.Exception (evaluate)+import Control.Monad (forM, forM_, void, when)+import qualified Crypto.Hash.SHA256 as SHA256+import qualified Data.Binary as B+import qualified Data.ByteString as BS+import qualified Data.IntSet as IS+import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef)+import Data.List.Split (chunksOf)+import Data.Traversable (mapAccumL)+import Data.Tuple (swap)+import Data.Word (Word64)+import qualified MCG+import qualified Options.Applicative as O+import qualified System.Clock as Clock+import System.Directory (removePathForcibly)+import Text.Printf (printf)++import qualified RocksDB++-------------------------------------------------------------------------------+-- Keys and values+-------------------------------------------------------------------------------++type K = BS.ByteString+type V = BS.ByteString++-- We generate keys by hashing a word64 and adding two "random" bytes.+-- This way we can ensure that keys are distinct.+--+-- I think this approach of generating keys should match UTxO quite well.+-- This is purely CPU bound operation, and we should be able to push IO+-- when doing these in between.+makeKey :: Word64 -> K+makeKey w64 = SHA256.hashlazy (B.encode w64) <> "=="++makeValue :: K -> V+makeValue k = k <> BS.replicate (60 - 34) 120 -- 'x'++-------------------------------------------------------------------------------+-- Options and commands+-------------------------------------------------------------------------------++data GlobalOpts = GlobalOpts+ { rootDir :: !FilePath -- ^ session directory.+ , initialSize :: !Int+ }+ deriving stock Show++data SetupOpts = SetupOpts+ deriving stock Show++data RunOpts = RunOpts+ { batchCount :: !Int+ , batchSize :: !Int+ , check :: !Bool+ , seed :: !Word64+ }+ deriving stock Show++data Cmd+ -- | Setup benchmark: generate initial LSM tree etc.+ = CmdSetup SetupOpts++ -- | Make a dry run, measure the overhead.+ | CmdDryRun RunOpts++ -- | Run the actual benchmark+ | CmdRun RunOpts+ deriving stock Show++-------------------------------------------------------------------------------+-- command line interface+-------------------------------------------------------------------------------++globalOptsP :: O.Parser GlobalOpts+globalOptsP = pure GlobalOpts+ <*> pure "_bench_rocksdb"+ <*> O.option O.auto (O.long "initial-size" <> O.value 100_000_000 <> O.showDefault <> O.help "Initial LSM tree size")++cmdP :: O.Parser Cmd+cmdP = O.subparser $ mconcat+ [ O.command "setup" $ O.info+ (CmdSetup <$> setupOptsP <**> O.helper)+ (O.progDesc "Setup benchmark")++ , O.command "dry-run" $ O.info+ (CmdDryRun <$> runOptsP <**> O.helper)+ (O.progDesc "Dry run, measure overhead")++ , O.command "run" $ O.info+ (CmdRun <$> runOptsP <**> O.helper)+ (O.progDesc "Proper run")+ ]++setupOptsP :: O.Parser SetupOpts+setupOptsP = pure SetupOpts++runOptsP :: O.Parser RunOpts+runOptsP = pure RunOpts+ <*> O.option O.auto (O.long "batch-count" <> O.value 1000 <> O.showDefault <> O.help "Batch count")+ <*> O.option O.auto (O.long "batch-size" <> O.value 256 <> O.showDefault <> O.help "Batch size")+ <*> O.switch (O.long "check" <> O.help "Check generated key distribution")+ <*> O.option O.auto (O.long "seed" <> O.value 1337 <> O.showDefault <> O.help "Random seed")++-------------------------------------------------------------------------------+-- clock+-------------------------------------------------------------------------------++timed :: IO a -> IO (a, Double)+timed action = do+ t1 <- Clock.getTime Clock.Monotonic+ x <- action+ t2 <- Clock.getTime Clock.Monotonic+ let !t = fromIntegral (Clock.toNanoSecs (Clock.diffTimeSpec t2 t1)) * 1e-9+ pure (x, t)++timed_ :: IO () -> IO Double+timed_ action = do+ t1 <- Clock.getTime Clock.Monotonic+ action+ t2 <- Clock.getTime Clock.Monotonic+ pure $! fromIntegral (Clock.toNanoSecs (Clock.diffTimeSpec t2 t1)) * 1e-9++-------------------------------------------------------------------------------+-- setup+-------------------------------------------------------------------------------++{-+rocksdb::BlockBasedTableOptions table_options;+ table_options.filter_policy.reset(rocksdb::NewBloomFilterPolicy(10, false));+ my_cf_options.table_factory.reset(+ rocksdb::NewBlockBasedTableFactory(table_options));+-}++rocksDbWithOptions :: (RocksDB.Options -> IO r) -> IO r+rocksDbWithOptions kont =+ RocksDB.withOptions $ \options ->+ RocksDB.withFilterPolicyBloom 10 $ \filterPolicy ->+ RocksDB.withBlockTableOptions $ \topts -> do+ RocksDB.blockBasedOptionsSetFilterPolicy topts filterPolicy++ RocksDB.optionsSetCreateIfMissing options True+ RocksDB.optionsIncreaseParallelism options 4 -- the benchmark spec says to use one core though.+ RocksDB.optionsOptimizeLevelStyleCompaction options 0x200_00000 -- 32*1024*1024; default is 512*1024*1024?+ -- RocksDB.optionsSetCompression options 0 -- no compression+ RocksDB.optionsSetCompression options 4 -- lz4+ -- RocksDB.optionsSetCompression options 7 -- zstd+ RocksDB.optionsSetMaxOpenFiles options 512++ RocksDB.optionsSetBlockBasedTableFactory options topts++ kont options++rocksDbWithWriteOptions :: (RocksDB.WriteOptions -> IO r) -> IO r+rocksDbWithWriteOptions kont = RocksDB.withWriteOptions $ \options -> do+ -- lsm-tree doesn't have WAL, so we don't use with RocksDB to be fair.+ RocksDB.writeOptionsDisableWAL options True++ kont options++--+--+-- It's useful to run+--+-- sst_dump --file=_bench_rocksdb --command=identity --show_properties+--+-- after setup+doSetup :: GlobalOpts -> SetupOpts -> IO ()+doSetup gopts opts = do+ time <- timed_ $ doSetup' gopts opts+ printf "Setup %.03f sec\n" time++doSetup' :: GlobalOpts -> SetupOpts -> IO ()+doSetup' gopts _opts =+ rocksDbWithOptions $ \options ->+ RocksDB.withRocksDB options (rootDir gopts) $ \db ->+ rocksDbWithWriteOptions $ \wopts ->+ forM_ (chunksOf 256 [ 0 .. initialSize gopts ]) $ \chunk ->+ RocksDB.withWriteBatch $ \batch -> do+ forM_ chunk $ \ (fromIntegral -> i) -> do+ when (mod i (fromIntegral (div (initialSize gopts) 50)) == 0) $ do+ printf "%3.0f%%\n" (100 * fromIntegral i / fromIntegral (initialSize gopts) :: Double)+ let k = makeKey i+ let v = makeValue k+ RocksDB.writeBatchPut batch k v++ RocksDB.write db wopts batch++-------------------------------------------------------------------------------+-- dry-run+-------------------------------------------------------------------------------++doDryRun :: GlobalOpts -> RunOpts -> IO ()+doDryRun gopts opts = do+ time <- timed_ $ doDryRun' gopts opts+ printf "Batch generation: %.03f sec\n" time++doDryRun' :: GlobalOpts -> RunOpts -> IO ()+doDryRun' gopts opts = do+ keysRef <- newIORef $+ if (check opts)+ then IS.fromList [ 0 .. initialSize gopts - 1 ]+ else IS.empty+ duplicateRef <- newIORef (0 :: Int)++ void $ forFoldM_ initGen [ 0 .. batchCount opts - 1 ] $ \b g -> do+ let lookups :: [Word64]+ inserts :: [Word64]+ (!nextG, lookups, inserts) = generateBatch (initialSize gopts) (batchSize opts) g b++ when (check opts) $ do+ keys <- readIORef keysRef+ let new = IS.fromList $ map fromIntegral lookups+ let diff = IS.difference new keys+ -- when (IS.notNull diff) $ printf "missing in batch %d %s\n" b (show diff)+ modifyIORef' duplicateRef $ \n -> n + IS.size diff+ writeIORef keysRef $! IS.union+ (IS.difference keys new)+ (IS.fromList $ map fromIntegral inserts)++ -- lookups+ forM_ lookups $ \k -> evaluate (makeKey k)++ -- deletes & inserts; deletes done above.+ forM_ inserts $ \k' -> do+ let k = makeKey k'+ evaluate k >> evaluate (makeValue k)++ pure nextG++ when (check opts) $ do+ duplicates <- readIORef duplicateRef+ printf "True duplicates: %d\n" duplicates+ where+ initGen = MCG.make+ (fromIntegral $ initialSize gopts + batchSize opts * batchCount opts)+ (seed opts)++-------------------------------------------------------------------------------+-- Batch generation+-------------------------------------------------------------------------------++{- | Implement generation of unbounded sequence of insert/delete operations++matching UTxO style from spec: interleaved batches insert and lookup+configurable batch sizes+1 insert, 1 delete, 1 lookup per key.++Current approach is probabilistic, but uses very little state.+We could also make it exact, but then we'll need to carry some state around+(at least the difference).++-}+generateBatch ::+ Int -- ^ initial size of the collection+ -> Int -- ^ batch size+ -> MCG.MCG -- ^ generator+ -> Int -- ^ batch number+ -> (MCG.MCG, [Word64], [Word64])+generateBatch initialSize batchSize g b = (nextG, lookups, inserts)+ where+ maxK :: Word64+ maxK = fromIntegral $ initialSize + batchSize * b++ lookups :: [Word64]+ (!nextG, lookups) = mapAccumL (\g' _ -> swap (MCG.reject maxK g')) g [1 .. batchSize]++ inserts :: [Word64]+ inserts = [ maxK .. maxK + fromIntegral batchSize - 1 ]++-------------------------------------------------------------------------------+-- run+-------------------------------------------------------------------------------++doRun :: GlobalOpts -> RunOpts -> IO ()+doRun gopts opts = do+ removePathForcibly $ rootDir gopts ++ "_cp"+ makeCheckpoint gopts++ time <- timed_ $ doRun' gopts { rootDir = rootDir gopts ++ "_cp" } opts+ -- TODO: collect more statistic, save them in dry-run,+ -- TODO: make the results human comprehensible.+ printf "Proper run: %7.03f sec\n" time+ let ops = batchCount opts * batchSize opts+ printf "Operations per second: %7.01f ops/sec\n" (fromIntegral ops / time)++makeCheckpoint :: GlobalOpts -> IO ()+makeCheckpoint gopts =+ rocksDbWithOptions $ \options ->+ RocksDB.withRocksDB options (rootDir gopts) $ \db ->+ RocksDB.checkpoint db $ rootDir gopts ++ "_cp"++doRun' :: GlobalOpts -> RunOpts -> IO ()+doRun' gopts opts =+ rocksDbWithOptions $ \options ->+ RocksDB.withRocksDB options (rootDir gopts) $ \db ->+ rocksDbWithWriteOptions $ \wopts ->+ RocksDB.withReadOptions $ \ropts ->+ void $ forFoldM_ initGen [ 0 .. batchCount opts - 1 ] $ \b g -> do+ let lookups :: [Word64]+ inserts :: [Word64]+ (!nextG, lookups, inserts) = generateBatch (initialSize gopts) (batchSize opts) g b++ -- lookups+ let ks = makeKey <$> lookups++ vs' <-+ -- multi get or not.+ if True+ then RocksDB.multiGet db ropts ks+ else forM ks (RocksDB.get db ropts)++ vs <- evaluate (force vs')++ -- check that we get values we expect+ when (check opts) $ do+ let expected = map (Just . makeValue) ks+ when (vs /= expected) $ do+ printf "Value mismatch in batch %d\n" b+ print vs+ print expected++ -- deletes and inserts+ RocksDB.withWriteBatch $ \batch -> do+ forM_ ks $ \k -> do+ RocksDB.writeBatchDelete batch k++ forM_ inserts $ \k' -> do+ let k = makeKey k'+ let v = makeValue k+ RocksDB.writeBatchPut batch k v++ RocksDB.write db wopts batch++ pure nextG+ where+ initGen = MCG.make+ (fromIntegral $ initialSize gopts + batchSize opts * batchCount opts)+ (seed opts)++-------------------------------------------------------------------------------+-- main+-------------------------------------------------------------------------------++main :: IO ()+main = do+#ifdef NO_IGNORE_ASSERTS+ putStrLn "WARNING: Benchmarking in debug mode."+ putStrLn " To benchmark in release mode, pass:"+ putStrLn " --project-file=cabal.project.release"+#endif+ (gopts, cmd) <- O.customExecParser prefs cliP+ print gopts+ print cmd+ case cmd of+ CmdSetup opts -> doSetup gopts opts+ CmdDryRun opts -> doDryRun gopts opts+ CmdRun opts -> doRun gopts opts+ where+ cliP = O.info ((,) <$> globalOptsP <*> cmdP <**> O.helper) O.fullDesc+ prefs = O.prefs $ O.showHelpOnEmpty <> O.helpShowGlobals <> O.subparserInline++-------------------------------------------------------------------------------+-- general utils+-------------------------------------------------------------------------------++forFoldM_ :: Monad m => s -> [a] -> (a -> s -> m s) -> m s+forFoldM_ !s [] _ = pure s+forFoldM_ !s (x:xs) f = do+ !s' <- f x s+ forFoldM_ s' xs f++-------------------------------------------------------------------------------+-- unused for now+-------------------------------------------------------------------------------++_unused :: ()+_unused = const ()+ timed
@@ -0,0 +1,469 @@+{-# LANGUAGE OverloadedLists #-}++module Bench.Database.LSMTree (benchmarks) where++import Control.DeepSeq+import Control.Exception+import Control.Tracer+import Criterion.Main+import qualified Data.BloomFilter.Hash as Bloom+import Data.ByteString.Short (ShortByteString)+import qualified Data.ByteString.Short as SBS+import Data.Foldable+import Data.Functor.Compose+import qualified Data.Vector as V+import Data.Void+import Data.Word+import Database.LSMTree hiding (withTable)+import Database.LSMTree.Extras+import Database.LSMTree.Extras.Orphans ()+import Database.LSMTree.Internal.Assertions (fromIntegralChecked)+import qualified Database.LSMTree.Internal.RawBytes as RB+import GHC.Generics (Generic)+import Prelude hiding (getContents, take)+import System.Directory (removeDirectoryRecursive)+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.IO as FS+import qualified System.FS.IO as FS+import System.IO.Temp+import System.Random++benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree" [+ benchLargeValueVsSmallValueBlob+ , benchCursorScanVsRangeLookupScan+ , benchInsertBatches+ , benchInsertsVsUpserts+ , benchLookupsInsertsVsUpserts+ , benchLookupInsertsVsLookupUpserts+ ]++{-------------------------------------------------------------------------------+ Types+-------------------------------------------------------------------------------}++newtype K = K Word64+ deriving stock Generic+ deriving newtype (Show, Eq, Ord, Num, NFData, SerialiseKey)+ deriving anyclass Uniform++data V1 = V1 !Word64 !ShortByteString+ deriving stock (Show, Eq, Ord)+ deriving ResolveValue via ResolveAsFirst V1++instance NFData V1 where+ rnf (V1 x s) = rnf x `seq` rnf s++instance SerialiseValue V1 where+ serialiseValue (V1 x s) = serialiseValue x <> serialiseValue s+ deserialiseValue rb = V1 (deserialiseValue $ RB.take 8 rb) (deserialiseValue $ RB.drop 8 rb)++newtype B1 = B1 Void+ deriving newtype (Show, Eq, Ord, NFData, SerialiseValue)++newtype V2 = V2 Word64+ deriving newtype (Show, Eq, Ord, NFData, SerialiseValue)+ deriving ResolveValue via ResolveAsFirst V2++newtype B2 = B2 ShortByteString+ deriving newtype (Show, Eq, Ord, NFData, SerialiseValue)++newtype V3 = V3 Word64+ deriving newtype (Show, Eq, Ord, Num, NFData, SerialiseValue)++type B3 = Void++instance ResolveValue V3 where+ resolve = (+)++benchConfig :: TableConfig+benchConfig = defaultTableConfig+ { confWriteBufferAlloc = AllocNumEntries 1000+ , confFencePointerIndex = CompactIndex+ , confDiskCachePolicy = DiskCacheNone+ }++benchSalt :: Bloom.Salt+benchSalt = 4++{-------------------------------------------------------------------------------+ Large Value vs. Small Value Blob+-------------------------------------------------------------------------------}++benchLargeValueVsSmallValueBlob :: Benchmark+benchLargeValueVsSmallValueBlob =+ env mkEntries $ \es -> bgroup "large-value-vs-small-value-blob" [+ env (mkGrouped (mkV1 es)) $ \ ~(ess, kss) -> bgroup "V1" [+ withEnv ess $ \ ~(_, _, _, _, t :: Table IO K V1 B1) -> do+ bench "lookups-large-value" $ whnfIO $+ V.mapM_ (lookups t) kss+ ]+ , env (mkGrouped (mkV2 es)) $ \ ~(ess, kss) -> bgroup "V2" [+ withEnv ess $ \ ~(_, _, _, _, t :: Table IO K V2 B2) -> do+ bench "lookups-small-value" $ whnfIO $+ V.mapM_ (lookups t) kss+ , withEnv ess $ \ ~(_, _, _, s, t :: Table IO K V2 B2) -> do+ bench "lookups-small-value-blob" $ whnfIO $ do+ V.forM_ kss $ \ks -> do+ lrs <- lookups t ks+ retrieveBlobs s (V.fromList $ toList $ Compose lrs)+ ]+ ]+ where+ !initialSize = 80_000+ !batchSize = 250++ customRandomEntries :: Int -> V.Vector (K, Word64, ShortByteString)+ customRandomEntries n = V.unfoldrExactN n f (mkStdGen 17)+ where f !g = let (!k, !g') = uniform g+ in ((k, v, b), g')+ -- The exact value does not matter much, so we pick an arbitrary+ -- hardcoded one.+ !v = 138+ -- TODO: tweak size of blob+ !b = SBS.pack [0 | _ <- [1 :: Int .. 1500]]++ mkEntries :: IO (V.Vector (K, Word64, ShortByteString))+ mkEntries = pure $ customRandomEntries initialSize++ mkGrouped ::+ V.Vector (k, v, b)+ -> IO ( V.Vector (V.Vector (k, v, b))+ , V.Vector (V.Vector k) )+ mkGrouped es = pure $+ let ess = vgroupsOfN batchSize es+ kss = V.map (V.map fst3) ess+ in (ess, kss)++ withEnv inss = envWithCleanup (initialise inss) cleanup++ initialise inss = do+ (tmpDir, hfs, hbio) <- mkFiles+ s <- openSession nullTracer hfs hbio benchSalt (FS.mkFsPath [])+ t <- newTableWith benchConfig s+ V.mapM_ (inserts t) inss+ pure (tmpDir, hfs, hbio, s, t)++ cleanup (tmpDir, hfs, hbio, s, t) = do+ closeTable t+ closeSession s+ cleanupFiles (tmpDir, hfs, hbio)++ mkV1 :: V.Vector (K, Word64, ShortByteString) -> V.Vector (K, V1, Maybe B1)+ mkV1 = V.map (\(k, v, b) -> (k, V1 v b, Nothing))++ mkV2 :: V.Vector (K, Word64, ShortByteString) -> V.Vector (K, V2, Maybe B2)+ mkV2 = V.map (\(k, v, b) -> (k, V2 v, Just $ B2 b))++fst3 :: (a, b, c) -> a+fst3 (x, _, _) = x++{-------------------------------------------------------------------------------+ Cursor Scan vs. Range Lookup Scan+-------------------------------------------------------------------------------}++benchCursorScanVsRangeLookupScan :: Benchmark+benchCursorScanVsRangeLookupScan =+ env mkEntries $ \es ->+ env (mkGrouped es) $ \ ess ->+ withEnv ess $ \ ~(_, _, _, _, t :: Table IO K V2 B2) ->+ bgroup "cursor-scan-vs-range-lookup-scan" [+ bench "cursor-scan-full" $ whnfIO $ do+ withCursor t $ \c -> do+ take initialSize c+ , bench "cursor-scan-chunked" $ whnfIO $ do+ withCursor t $ \c -> do+ forM_ ([1 .. numChunks] :: [Int]) $ \_ -> do+ take readSize c+ , bench "range-scan-full" $ whnfIO $ do+ rangeLookup t (FromToIncluding (K minBound) (K maxBound))+ , bench "range-scan-chunked" $ whnfIO $ do+ forM_ ranges $ \r -> do+ rangeLookup t r+ ]+ where+ initialSize, batchSize, numChunks :: Int+ !initialSize = 80_000+ !batchSize = 250+ !numChunks = 100++ readSize :: Int+ !readSize = check $ initialSize `div` numChunks+ where+ check x = assert (x * numChunks == initialSize) $ x++ ranges :: [Range K]+ !ranges = check $ force $ [+ FromToExcluding (K $ c * i) (K $ c * (i + 1))+ | i <- [0 .. fromIntegralChecked numChunks - 2]+ ] <> [+ FromToIncluding (K $ c * (fromIntegralChecked numChunks - 1)) (K maxBound)+ ]+ where+ c = fromIntegralChecked (maxBound `div` numChunks)+ check rs = assert (length rs == numChunks) rs+++ customRandomEntries :: Int -> V.Vector (K, V2, Maybe B2)+ customRandomEntries n = V.unfoldrExactN n f (mkStdGen 17)+ where f !g = let (!k, !g') = uniform g+ in ((k, v, Nothing), g')+ -- The exact value does not matter much, so we pick an arbitrary+ -- hardcoded one.+ !v = V2 138++ mkEntries :: IO (V.Vector (K, V2, Maybe B2))+ mkEntries = pure $ customRandomEntries initialSize++ mkGrouped ::+ V.Vector (k, v, b)+ -> IO (V.Vector (V.Vector (k, v, b)))+ mkGrouped es = pure $ vgroupsOfN batchSize es++ withEnv inss = envWithCleanup (initialise inss) cleanup++ initialise inss = do+ (tmpDir, hfs, hbio) <- mkFiles+ s <- openSession nullTracer hfs hbio benchSalt (FS.mkFsPath [])+ t <- newTableWith benchConfig s+ V.mapM_ (inserts t) inss+ pure (tmpDir, hfs, hbio, s, t)++ cleanup (tmpDir, hfs, hbio, s, t) = do+ closeTable t+ closeSession s+ cleanupFiles (tmpDir, hfs, hbio)+++{-------------------------------------------------------------------------------+ Insert Batches+-------------------------------------------------------------------------------}++benchInsertBatches :: Benchmark+benchInsertBatches =+ env genInserts $ \iss ->+ withEnv $ \ ~(_, _, _, _, t :: Table IO K V2 Void) -> do+ bench "insert-batches" $ whnfIO $+ V.mapM_ (inserts t) iss+ where+ !initialSize = 100_000+ !batchSize = 256++ _benchConfig :: TableConfig+ _benchConfig = benchConfig {+ confWriteBufferAlloc = AllocNumEntries 1000+ }++ randomInserts :: Int -> V.Vector (K, V2, Maybe Void)+ randomInserts n = V.unfoldrExactN n f (mkStdGen 17)+ where f !g = let (!k, !g') = uniform g+ in ((k, v, Nothing), g')+ -- The exact value does not matter much, so we pick an arbitrary+ -- hardcoded one.+ !v = V2 17++ genInserts :: IO (V.Vector (V.Vector (K, V2, Maybe Void)))+ genInserts = pure $ vgroupsOfN batchSize $ randomInserts initialSize++ withEnv = envWithCleanup initialise cleanup++ initialise = do+ (tmpDir, hfs, hbio) <- mkFiles+ s <- openSession nullTracer hfs hbio benchSalt (FS.mkFsPath [])+ t <- newTableWith _benchConfig s+ pure (tmpDir, hfs, hbio, s, t)++ cleanup (tmpDir, hfs, hbio, s, t) = do+ closeTable t+ closeSession s+ cleanupFiles (tmpDir, hfs, hbio)++{-------------------------------------------------------------------------------+ Inserts vs. Upserts+-------------------------------------------------------------------------------}++-- | Compare inserts and upserts. The logical contents of the resulting+-- database are the same.+benchInsertsVsUpserts :: Benchmark+benchInsertsVsUpserts =+ env (pure $ snd $ randomEntriesGrouped 800_000 250) $ \ess ->+ env (pure $ V.map mkInserts ess) $ \inss ->+ bgroup "inserts-vs-upserts" [+ bench "inserts" $+ withEmptyTable $ \(_, _, _, _, t) ->+ V.mapM_ (inserts t) inss+ , bench "upserts" $+ withEmptyTable $ \(_, _, _, _, t) ->+ V.mapM_ (upserts t) ess+ ]+ where+ withEmptyTable =+ perRunEnvWithCleanup+ (do (tmpDir, hfs, hbio) <- mkFiles+ (s, t) <- mkTable hfs hbio benchConfig+ pure (tmpDir, hfs, hbio, s, t)+ )+ (\(tmpDir, hfs, hbio, s, t) -> do+ cleanupTable (s, t)+ cleanupFiles (tmpDir, hfs, hbio)+ )++{-------------------------------------------------------------------------------+ Lookups plus Inserts vs. Upserts+-------------------------------------------------------------------------------}++-- | Compare lookups+inserts to upserts. The former costs 2 LSMT operations,+-- while Upserts only cost 1 LSMT operation. The number of operations do not+-- directly translate to the number of I\/O operations, but one can assume that+-- lookup+insert is roughly twice as costly as upsert.+benchLookupsInsertsVsUpserts :: Benchmark+benchLookupsInsertsVsUpserts =+ env (pure $ snd $ randomEntriesGrouped 800_000 250) $ \ess ->+ env (pure $ V.map mkInserts ess) $ \inss ->+ bgroup "lookups-inserts-vs-upserts" [+ bench "lookups-inserts" $+ withTable inss $ \(_, _, _, _, t) ->+ -- Insert the same keys again, but we sum the existing values in+ -- the table with the values we are going to insert: first lookup+ -- the existing values, sum those with the insert values, then+ -- insert the updated values.+ V.forM_ ess $ \es -> do+ lrs <- lookups t (V.map fst es)+ let ins' = V.zipWith f es lrs+ inserts t ins'+ , bench "upserts" $+ withTable inss $ \(_, _, _, _, t) ->+ -- Insert the same keys again, but we sum the existing values in+ -- the table with the values we are going to insert: submit+ -- upserts with the insert values.+ V.forM_ ess $ \es -> upserts t es+ ]+ where+ f (k, v) = \case+ NotFound -> (k, v, Nothing)+ Found v' -> (k, v `resolve` v', Nothing)+ FoundWithBlob _ _ -> error "Unexpected blob found"++ withTable inss = perRunEnvWithCleanup+ -- Make a table and fill it up+ (do (tmpDir, hfs, hbio) <- mkFiles+ (s, t) <- mkTable hfs hbio benchConfig+ V.mapM_ (inserts t) inss+ pure (tmpDir, hfs, hbio, s, t)+ )+ (\(tmpDir, hfs, hbio, s, t) -> do+ cleanupTable (s, t)+ cleanupFiles (tmpDir, hfs, hbio)+ )++{-------------------------------------------------------------------------------+ Lookup Inserts vs. Lookup Upserts+-------------------------------------------------------------------------------}++-- | Compare lookups after inserts against lookups after upserts.+benchLookupInsertsVsLookupUpserts :: Benchmark+benchLookupInsertsVsLookupUpserts =+ env (pure $ snd $ randomEntriesGrouped 80_000 250) $ \ess ->+ env (pure $ V.map mkInserts ess) $ \inss ->+ bgroup "lookup-inserts-vs-lookup-upserts" [+ bench "lookup-inserts" $+ withInsertTable inss $ \(_, _, _, _, t) -> do+ V.forM_ ess $ \es -> lookups t (V.map fst es)+ , bench "lookup-upserts" $+ withUpsertTable ess $ \(_, _, _, _, t) -> do+ V.forM_ ess $ \es -> lookups t (V.map fst es)+ ]+ where+ withInsertTable inss =+ perRunEnvWithCleanup+ -- Insert the same keys 10 times, where each new insert behaves like+ -- a lookup+insert. This results in a logical database containing+ -- the original keys with the original value *10.+ (do (tmpDir, hfs, hbio) <- mkFiles+ (s, t) <- mkTable hfs hbio benchConfig+ V.forM_ [1..10] $ \(i::Int) -> do+ let inss' = (V.map . V.map) (\(k, v, b) -> (k, fromIntegral i * v, b)) inss+ V.mapM_ (inserts t) inss'+ pure (tmpDir, hfs, hbio, s, t)+ )+ (\(tmpDir, hfs, hbio, s, t) -> do+ cleanupTable (s, t)+ cleanupFiles (tmpDir, hfs, hbio)+ )++ withUpsertTable ess =+ perRunEnvWithCleanup+ -- Upsert the same key 10 times. The results in a logical database+ -- containing the original keys with the original value *10.+ (do (tmpDir, hfs, hbio) <- mkFiles+ (s, t) <- mkTable hfs hbio benchConfig+ V.forM_ [1..10] $ \(_::Int) ->+ V.mapM_ (upserts t) ess+ pure (tmpDir, hfs, hbio, s, t)+ )+ (\(tmpDir, hfs, hbio, s, t) -> do+ cleanupTable (s, t)+ cleanupFiles (tmpDir, hfs, hbio)+ )++{-------------------------------------------------------------------------------+ Setup+-------------------------------------------------------------------------------}++-- | Random keys, default values @1@+randomEntries :: Int -> V.Vector (K, V3)+randomEntries n = V.unfoldrExactN n f (mkStdGen 17)+ where f !g = let (!k, !g') = uniform g+ in ((k, 1), g')++-- | Like 'randomEntries', but also returns groups of size 'm'+randomEntriesGrouped :: Int -> Int -> (V.Vector (K, V3), V.Vector (V.Vector (K, V3)))+randomEntriesGrouped n m =+ let es = randomEntries n+ in (es, vgroupsOfN m es)++mkInserts :: V.Vector (K, V3) -> V.Vector (K, V3, Maybe B3)+mkInserts = V.map (\(k, v) -> (k, v, Nothing))++mkFiles ::+ IO ( FilePath -- ^ Temporary directory+ , FS.HasFS IO FS.HandleIO+ , FS.HasBlockIO IO FS.HandleIO+ )+mkFiles = do+ sysTmpDir <- getCanonicalTemporaryDirectory+ benchTmpDir <- createTempDirectory sysTmpDir "full"+ (hfs, hbio) <- FS.ioHasBlockIO (FS.MountPoint benchTmpDir) FS.defaultIOCtxParams+ pure (benchTmpDir, hfs, hbio)++cleanupFiles ::+ ( FilePath -- ^ Temporary directory+ , FS.HasFS IO FS.HandleIO+ , FS.HasBlockIO IO FS.HandleIO+ )+ -> IO ()+cleanupFiles (tmpDir, _hfs, hbio) = do+ FS.close hbio+ removeDirectoryRecursive tmpDir++mkTable ::+ FS.HasFS IO FS.HandleIO+ -> FS.HasBlockIO IO FS.HandleIO+ -> TableConfig+ -> IO ( Session IO+ , Table IO K V3 B3+ )+mkTable hfs hbio conf = do+ sesh <- openSession nullTracer hfs hbio benchSalt (FS.mkFsPath [])+ t <- newTableWith conf sesh+ pure (sesh, t)++cleanupTable ::+ ( Session IO+ , Table IO K V3 B3+ )+ -> IO ()+cleanupTable (s, t) = do+ closeTable t+ closeSession s
@@ -0,0 +1,93 @@+{-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE TypeApplications #-}+++module Bench.Database.LSMTree.Internal.BloomFilter (+ benchmarks+ -- * Benchmarked functions+ , elems+ ) where++import Criterion.Main+import qualified Data.Bifoldable as BiFold+import Data.BloomFilter (Bloom)+import qualified Data.BloomFilter as Bloom+import Data.BloomFilter.Hash (Hashable)+import qualified Data.Foldable as Fold+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Database.LSMTree.Extras.Random+import Database.LSMTree.Extras.UTxO (UTxOKey)+import Database.LSMTree.Internal.Serialise (SerialisedKey,+ serialiseKey)+import System.Random as R++-- See 'utxoNumPages'.+benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.BloomFilter" [+ bgroup "elems" [+ env (elemEnv 0.1 2_500_000 1_000_000 0) $ \ ~(b, xs) ->+ bench "onlyTruePositives 0.1" $ whnf (elems b) xs+ , env (elemEnv 0.9 2_500_000 1_000_000 0) $ \ ~(b, xs) ->+ bench "onlyTruePositives 0.9" $ whnf (elems b) xs+ , env (elemEnv 0.1 2_500_000 0 1_000_000) $ \ ~(b, xs) ->+ bench "onlyNegatives 0.1" $ whnf (elems b) xs+ , env (elemEnv 0.9 2_500_000 0 1_000_000) $ \ ~(b, xs) ->+ bench "onlyNegatives 0.9" $ whnf (elems b) xs+ ]+ , env (constructionEnv 2_500_000) $ \ m ->+ bgroup "construction" [+ bench "FPR = 0.1" $+ whnf (constructBloom 0.1) m++ , bench "FPR = 0.9" $+ whnf (constructBloom 0.9) m+ ]+ ]++benchSalt :: Bloom.Salt+benchSalt = 4++-- | Input environment for benchmarking 'Bloom.elem'.+elemEnv ::+ Double -- ^ False positive rate+ -> Int -- ^ Number of entries in the bloom filter+ -> Int -- ^ Number of positive lookups+ -> Int -- ^ Number of negative lookups+ -> IO (Bloom SerialisedKey, [SerialisedKey])+elemEnv fpr nbloom nelemsPositive nelemsNegative = do+ let g = mkStdGen 100+ (g1, g') = R.splitGen g+ (g2, g3) = R.splitGen g'++ let (xs, ys1) = splitAt nbloom+ $ uniformWithoutReplacement @UTxOKey g1 (nbloom + nelemsNegative)+ ys2 = sampleUniformWithReplacement @UTxOKey g2 nelemsPositive xs+ zs = shuffle (ys1 ++ ys2) g3+ pure ( Bloom.fromList (Bloom.policyForFPR fpr) benchSalt (fmap serialiseKey xs)+ , fmap serialiseKey zs+ )++-- | Used for benchmarking 'Bloom.elem'.+elems :: Hashable a => Bloom a -> [a] -> ()+elems b xs = Fold.foldl' (\acc x -> Bloom.elem x b `seq` acc) () xs++-- | Input environment for benchmarking 'constructBloom'.+constructionEnv :: Int -> IO (Map SerialisedKey SerialisedKey)+constructionEnv n = do+ stdgen <- newStdGen+ stdgen' <- newStdGen+ let ks = uniformWithoutReplacement @UTxOKey stdgen n+ vs = uniformWithReplacement @UTxOKey stdgen' n+ pure $ Map.fromList (zipWith (\k v -> (serialiseKey k, serialiseKey v)) ks vs)++-- | Used for benchmarking the construction of bloom filters from write buffers.+constructBloom ::+ Double+ -> Map SerialisedKey SerialisedKey+ -> Bloom SerialisedKey+constructBloom fpr m =+ -- For faster construction, avoid going via lists and use Bloom.create,+ -- traversing the map inserting the keys+ Bloom.create (Bloom.sizeForFPR fpr (Map.size m)) benchSalt $ \b ->+ BiFold.bifoldMap (\k -> Bloom.insert b k) (\_v -> pure ()) m
@@ -0,0 +1,121 @@+{-# LANGUAGE CPP #-}++module Bench.Database.LSMTree.Internal.Index (benchmarks) where++import Control.Category ((>>>))+import Control.DeepSeq (rnf)+import Control.Monad.ST.Strict (runST)+import Criterion.Main (Benchmark, Benchmarkable, bench, bgroup, env,+ whnf)+#if !MIN_VERSION_base(4,20,0)+import Data.List (foldl')+ -- foldl' is included in the Prelude from base 4.20 onwards+#endif+import Database.LSMTree.Extras.Generators (mkPages, toAppends)+ -- also for @Arbitrary@ instantiation of @SerialisedKey@+import Database.LSMTree.Extras.Index (Append, append)+import Database.LSMTree.Internal.Index (Index,+ IndexType (Compact, Ordinary), newWithDefaults, search,+ unsafeEnd)+import Database.LSMTree.Internal.Serialise (SerialisedKey)+import Test.QuickCheck (choose, vector)+import Test.QuickCheck.Gen (Gen (MkGen))+import Test.QuickCheck.Random (mkQCGen)++-- * Benchmarks++benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.Index" $+ map (uncurry benchmarksForSingleType) $+ [("Compact", Compact), ("Ordinary", Ordinary)]+ where++ benchmarksForSingleType :: String -> IndexType -> Benchmark+ benchmarksForSingleType indexTypeName indexType+ = bgroup (indexTypeName ++ " index") $+ [+ -- Search+ env (pure $ searchIndex indexType 10000) $ \ index ->+ env (pure $ searchKeys 1000) $ \ keys ->+ bench "Search" $+ searchBenchmarkable index keys,++ -- Incremental construction+ env (pure $ incrementalConstructionAppends 10000) $ \ appends ->+ bench "Incremental construction" $+ incrementalConstructionBenchmarkable indexType appends+ ]++-- * Utilities++-- | Deterministically constructs a value using a QuickCheck generator.+generated :: Gen a -> a+generated (MkGen exec) = exec (mkQCGen 411) 30++{-|+ Constructs serialised keys that conform to the key size constraint of+ compact indexes.+-}+keysForIndexCompact :: Int -- ^ Number of keys+ -> [SerialisedKey] -- ^ Constructed keys+keysForIndexCompact = vector >>>+ generated++{-|+ Constructs append operations whose serialised keys conform to the key size+ constraint of compact indexes.+-}+appendsForIndexCompact :: Int -- ^ Number of keys used in the construction+ -> [Append] -- ^ Constructed append operations+appendsForIndexCompact = keysForIndexCompact >>>+ mkPages 0.03 (choose (0, 16)) 0.01 >>>+ generated >>>+ toAppends+{-+ The arguments used for 'mkPages' are the same as the ones used for+ 'genPages' in the instantiation of 'Arbitrary' for 'LogicalPageSummaries' at+ the time of writing.+-}++{-|+ Constructs an index by applying append operations to an initially empty+ index.+-}+indexFromAppends :: IndexType -> [Append] -> Index+indexFromAppends indexType appends = runST $ do+ indexAcc <- newWithDefaults indexType+ mapM_ (flip append indexAcc) appends+ snd <$> unsafeEnd indexAcc++-- * Benchmark ingredients++-- ** Search++-- | Constructs an index to be searched.+searchIndex :: IndexType -- ^ Type of index to construct+ -> Int -- ^ Number of keys used in the construction+ -> Index -- ^ Constructed index+searchIndex indexType keyCount+ = indexFromAppends indexType (appendsForIndexCompact keyCount)++-- | Constructs a list of keys to search for.+searchKeys :: Int -- ^ Number of searches+ -> [SerialisedKey] -- ^ Constructed search keys+searchKeys = keysForIndexCompact++-- | The action to be performed by a search benchmark.+searchBenchmarkable :: Index -> [SerialisedKey] -> Benchmarkable+searchBenchmarkable index = whnf $ foldl' (\ _ key -> rnf (search key index)) ()++-- ** Incremental construction++-- | Constructs append operations to be used in index construction.+incrementalConstructionAppends ::+ Int -- ^ Number of keys used in the construction+ -> [Append] -- ^ Constructed append operations+incrementalConstructionAppends = appendsForIndexCompact++-- | The action to be performed by an incremental-construction benchmark.+incrementalConstructionBenchmarkable :: IndexType -> [Append] -> Benchmarkable+incrementalConstructionBenchmarkable indexType appends+ = whnf (indexFromAppends indexType) appends
@@ -0,0 +1,154 @@+{-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE TypeApplications #-}++module Bench.Database.LSMTree.Internal.Index.Compact (+ benchmarks+ -- * Benchmarked functions+ , searches+ , constructIndexCompact+ ) where++import Control.DeepSeq (deepseq)+import Control.Exception (assert)+import Control.Monad.ST.Strict+import Criterion.Main+import qualified Data.Foldable as Fold+import qualified Data.List as List+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word+import Database.LSMTree.Extras+import Database.LSMTree.Extras.Generators+import Database.LSMTree.Extras.Index+import Database.LSMTree.Extras.Random+import Database.LSMTree.Extras.UTxO+import Database.LSMTree.Internal.Index.Compact+import Database.LSMTree.Internal.Index.CompactAcc+import Database.LSMTree.Internal.Map.Range+import Database.LSMTree.Internal.Serialise (SerialisedKey,+ serialiseKey)+import System.Random+import Test.QuickCheck (generate)++benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.Index.Compact" [+ env (searchEnv 10000 1000) $ \ ~(ic, ks) ->+ bench "searches-10-1k" $ whnf (searches ic) ks+ , bgroup "construction" [+ env (constructionEnv 1000) $ \ pages ->+ bench "construction-1k-100" $ whnf (constructIndexCompact 100) pages+ , env (VUM.replicate 3000 (7 :: Word32)) $ \ mv ->+ bench "unsafeWriteRange-1k" $+ whnfAppIO (\x -> stToIO (unsafeWriteRange mv (BoundInclusive 1000) (BoundInclusive 2000) x)) 17+ , env (VUM.replicate 30000 (7 :: Word32)) $ \ mv ->+ bench "unsafeWriteRange-10k" $+ whnfAppIO (\x -> stToIO (unsafeWriteRange mv (BoundInclusive 10000) (BoundInclusive 20000) x)) 17+ ]+ , benchNonUniform+ ]++-- | Input environment for benchmarking 'searches'.+searchEnv ::+ Int -- ^ Number of pages+ -> Int -- ^ Number of searches+ -> IO (IndexCompact, [SerialisedKey])+searchEnv npages nsearches = do+ ic <- constructIndexCompact 100 <$> constructionEnv npages+ let stdgen = mkStdGen 17+ let ks = serialiseKey <$> uniformWithReplacement @UTxOKey stdgen nsearches+ pure (ic, ks)++-- | Used for benchmarking 'search'.+searches ::+ IndexCompact+ -> [SerialisedKey] -- ^ Keys to search for+ -> ()+searches ic ks = Fold.foldl' (\acc k -> search k ic `deepseq` acc) () ks++-- | Input environment for benchmarking 'constructIndexCompact'.+constructionEnv ::+ Int -- ^ Number of pages+ -> IO [Append]+constructionEnv n = do+ let stdgen = mkStdGen 17+ let ks = uniformWithoutReplacement @UTxOKey stdgen (2 * n)+ ps <- generate (mkPages 0 (error "unused in constructionEnv") 0 ks)+ pure (toAppends ps)++-- | Used for benchmarking the incremental construction of a 'IndexCompact'.+constructIndexCompact ::+ ChunkSize+ -> [Append] -- ^ Pages to add in succession+ -> IndexCompact+constructIndexCompact (ChunkSize csize) apps = runST $ do+ ica <- new csize+ mapM_ (`appendToCompact` ica) apps+ (_, index) <- unsafeEnd ica+ pure index++{-------------------------------------------------------------------------------+ Benchmarks for UTxO keys that are /almost/ uniformly distributed+-------------------------------------------------------------------------------}++-- | UTXO keys are not truly uniformly distributed. The 'txId' is a uniformly+-- distributed hash, but the same 'txId' can appear in multiple UTXO keys, but+-- with a different 'txIx'. In the worst case, this means that we have a clash+-- in the compact index for /every page/. The following benchmarks show+benchNonUniform :: Benchmark+benchNonUniform =+ bgroup "non-uniformity" [+ -- construction+ env (pure $ (appsWithNearDups (mkStdGen 17) 1000)) $ \as ->+ bench ("construct appsWithNearDups") $ whnf (constructIndexCompact 1000) as+ , env (pure $ (appsWithoutNearDups (mkStdGen 17) 1000)) $ \as ->+ bench ("construct appsWithoutNearDups") $ whnf (constructIndexCompact 1000) as+ -- search+ , env ( let ic = constructIndexCompact 100 (appsWithNearDups (mkStdGen 17) 1000)+ g = mkStdGen 42+ ks = serialiseKey <$> uniformWithReplacement @UTxOKey g 1000+ in pure (ic, ks) ) $ \ ~(ic, ks) ->+ bench "search appsWithNearDups" $ whnf (searches ic) ks+ , env ( let ic = constructIndexCompact 100 (appsWithoutNearDups (mkStdGen 17) 1000)+ g = mkStdGen 42+ ks = serialiseKey <$> uniformWithReplacement @UTxOKey g 1000+ in pure (ic, ks) ) $ \ ~(ic, ks) ->+ bench "search appsWithoutNearDups" $ whnf (searches ic) ks+ ]++-- | 'Append's with truly uniformly distributed UTXO keys.+appsWithoutNearDups ::+ StdGen+ -> Int -- ^ Number of pages+ -> [Append]+appsWithoutNearDups g n =+ let ks = uniformWithoutReplacement @UTxOKey g (n * 2)+ ks' = List.sort ks+ -- append a dummy UTXO key because appsWithNearDups does so too.+ ps = groupsOfN 2 (UTxOKey 0 0 : ks')+ in assert (length ps == n + 1) $+ fmap fromNE ps++-- | 'Append's with worst-case near-duplicates. Each page boundary splits UTXO+-- keys with the same 'txId' but different a 'txIx'.+appsWithNearDups ::+ StdGen+ -> Int -- ^ Number of pages+ -> [Append]+appsWithNearDups g n =+ let ks = uniformWithoutReplacement @UTxOKey g n+ ks' = flip concatMap (List.sort ks) $ \k -> [k {txIx = 0}, k {txIx = 1}]+ -- append a dummy UTXO key so that each pair of near-duplicate keys is+ -- split between pages. That is, the left element of the pair is the+ -- maximum key in a page, and the right element of the pair is the+ -- minimum key on the next page.+ ps = groupsOfN 2 (UTxOKey 0 0 : ks')+ in -- check that the number of pages+ assert (length ps == n + 1) $+ fmap fromNE ps++fromNE :: NonEmpty UTxOKey -> Append+fromNE xs =+ assert (NE.sort xs == xs) $+ assert (NE.nub xs == xs) $+ AppendSinglePage (serialiseKey $ NE.head xs) (serialiseKey $ NE.last xs)
@@ -0,0 +1,281 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Bench.Database.LSMTree.Internal.Lookup (benchmarks) where++import Control.Monad+import Control.Monad.ST.Strict (stToIO)+import Control.RefCount+import Criterion.Main (Benchmark, bench, bgroup, env, envWithCleanup,+ perRunEnv, perRunEnvWithCleanup, whnf, whnfAppIO)+import Data.Bifunctor (Bifunctor (..))+import Data.ByteString (ByteString)+import qualified Data.List as List+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromMaybe)+import qualified Data.Vector as V+import Database.LSMTree.Extras.Orphans ()+import Database.LSMTree.Extras.Random (frequency, randomByteStringR,+ sampleUniformWithReplacement, shuffle,+ uniformWithoutReplacement)+import Database.LSMTree.Extras.UTxO+import Database.LSMTree.Internal.Arena (ArenaManager, closeArena,+ newArena, newArenaManager, withArena)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))+import Database.LSMTree.Internal.Index as Index+import Database.LSMTree.Internal.Lookup (bloomQueries, indexSearches,+ intraPageLookupsWithWriteBuffer, lookupsIOWithWriteBuffer,+ prepLookups)+import Database.LSMTree.Internal.Page (getNumPages)+import Database.LSMTree.Internal.Paths (RunFsPaths (..))+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import GHC.Exts (RealWorld)+import GHC.Stack (HasCallStack)+import Prelude hiding (getContents)+import System.Directory (removeDirectoryRecursive)+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.IO as FS+import qualified System.FS.IO as FS+import System.IO.Temp+import System.Random as R++-- | TODO: add a separate micro-benchmark that includes multi-pages.+benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.Lookup" [+ benchLookups Config {+ name = "2_000_000 entries, 256 positive lookups"+ , nentries = 2_000_000+ , npos = 256+ , nneg = 0+ , ioctxps = Nothing+ , caching = Run.CacheRunData+ }+ , benchLookups Config {+ name = "2_000_000 entries, 256 negative lookups"+ , nentries = 2_000_000+ , npos = 0+ , nneg = 256+ , ioctxps = Nothing+ , caching = Run.CacheRunData+ }+ , benchLookups Config {+ name = "2_000_000 entries, 256 positive lookups, NoCache"+ , nentries = 2_000_000+ , npos = 256+ , nneg = 0+ , ioctxps = Nothing+ , caching = Run.NoCacheRunData+ }+ , benchLookups Config {+ name = "2_000_000 entries, 256 negative lookups, NoCache"+ , nentries = 2_000_000+ , npos = 0+ , nneg = 256+ , ioctxps = Nothing+ , caching = Run.NoCacheRunData+ }+ ]++benchSalt :: Bloom.Salt+benchSalt = 4++benchLookups :: Config -> Benchmark+benchLookups conf@Config{name} =+ withEnv $ \ ~(_dir, arenaManager, _hasFS, hasBlockIO, wbblobs, rs, ks) ->+ env ( pure ( V.map (\(DeRef r) -> Run.runFilter r) rs+ , V.map (\(DeRef r) -> Run.runIndex r) rs+ , V.map (\(DeRef r) -> Run.runKOpsFile r) rs+ )+ ) $ \ ~(blooms, indexes, kopsFiles) ->+ bgroup name [+ -- The bloomfilter is queried for all lookup keys. The result is an+ -- unboxed vector, so only use @whnf@.+ bench "Bloomfilter query" $+ whnf (\ks' -> bloomQueries benchSalt blooms ks') ks+ -- The compact index is only searched for (true and false) positive+ -- lookup keys. We use whnf here because the result is+ , env (pure $ bloomQueries benchSalt blooms ks) $ \rkixs ->+ bench "Compact index search" $+ whnfAppIO (\ks' -> withArena arenaManager $ \arena -> stToIO $ indexSearches arena indexes kopsFiles ks' rkixs) ks+ -- prepLookups combines bloom filter querying and index searching.+ -- The implementation forces the results to WHNF, so we use+ -- whnfAppIO here instead of nfAppIO.+ , bench "Lookup preparation in memory" $+ whnfAppIO (\ks' -> withArena arenaManager $ \arena -> stToIO $ prepLookups arena benchSalt blooms indexes kopsFiles ks') ks+ -- Submit the IOOps we get from prepLookups to HasBlockIO. We use+ -- perRunEnv because IOOps contain mutable buffers, so we want fresh+ -- ones for each run of the benchmark. We manually evaluate the+ -- result to WHNF since it is unboxed vector.+ , bench "Submit IOOps" $+ perRunEnv (withArena arenaManager $ \arena -> stToIO $ prepLookups arena benchSalt blooms indexes kopsFiles ks) $ \ ~(_rkixs, ioops) -> do+ !_ioress <- FS.submitIO hasBlockIO ioops+ pure ()+ -- When IO result have been collected, intra-page lookups searches+ -- through the raw bytes (representing a disk page) for the lookup+ -- key. Again, we use perRunEnv here because IOOps contain mutable+ -- buffers, so we want fresh ones for each run of the benchmark. The+ -- result is a boxed vector of Maybe Entry, but since the+ -- implementation takes care to evaluate each of the elements, we+ -- only compute WHNF.+ , bench "Perform intra-page lookups" $+ perRunEnvWithCleanup+ ( do arena <- newArena arenaManager+ (rkixs, ioops) <- stToIO (prepLookups arena benchSalt blooms indexes kopsFiles ks)+ ioress <- FS.submitIO hasBlockIO ioops+ pure (rkixs, ioops, ioress, arena)+ )+ (\(_, _, _, arena) -> closeArena arenaManager arena)+ (\ ~(rkixs, ioops, ioress, _) -> do+ !_ <- intraPageLookupsWithWriteBuffer+ resolveV WB.empty wbblobs rs ks rkixs ioops ioress+ pure ())+ -- The whole shebang: lookup preparation, doing the IO, and then+ -- performing intra-page-lookups. Again, we evaluate the result to+ -- WHNF because it is the same result that+ -- intraPageLookupsWithWriteBuffer produces (see above).+ , bench "Lookups in IO" $+ whnfAppIO (\ks' -> lookupsIOWithWriteBuffer+ hasBlockIO arenaManager resolveV+ benchSalt WB.empty wbblobs+ rs blooms indexes kopsFiles ks') ks+ ]+ -- TODO: consider adding benchmarks that also use the write buffer+ -- (then we can't just use 'WB.empty', but must take it from the env)+ where+ withEnv = envWithCleanup+ (lookupsInBatchesEnv conf)+ lookupsInBatchesCleanup+ resolveV = \v1 _v2 -> v1++{-------------------------------------------------------------------------------+ Environments+-------------------------------------------------------------------------------}++-- | Config options describing a benchmarking scenario+data Config = Config {+ -- | Name for the benchmark scenario described by this config.+ name :: !String+ -- | Number of key\/operation pairs in the run+ , nentries+ -- | Number of positive lookups+ , npos :: !Int+ -- | Number of negative lookups+ , nneg :: !Int+ -- | Optional parameters for the io-uring context+ , ioctxps :: !(Maybe FS.IOCtxParams)+ -- | Whether to use or bypass the OS page cache+ , caching :: !Run.RunDataCaching+ }++lookupsInBatchesEnv ::+ Config+ -> IO ( FilePath -- ^ Temporary directory+ , ArenaManager RealWorld+ , FS.HasFS IO FS.HandleIO+ , FS.HasBlockIO IO FS.HandleIO+ , Ref (WBB.WriteBufferBlobs IO FS.HandleIO)+ , V.Vector (Ref (Run IO FS.HandleIO))+ , V.Vector SerialisedKey+ )+lookupsInBatchesEnv Config {..} = do+ arenaManager <- newArenaManager+ sysTmpDir <- getCanonicalTemporaryDirectory+ benchTmpDir <- createTempDirectory sysTmpDir "lookupsInBatchesEnv"+ (storedKeys, lookupKeys) <- lookupsEnv (mkStdGen 17) nentries npos nneg+ (hasFS, hasBlockIO) <- FS.ioHasBlockIO (FS.MountPoint benchTmpDir) (fromMaybe FS.defaultIOCtxParams ioctxps)+ wbblobs <- WBB.new hasFS (FS.mkFsPath ["0.wbblobs"])+ wb <- WB.fromMap <$> traverse (traverse (WBB.addBlob hasFS wbblobs)) storedKeys+ let fsps = RunFsPaths (FS.mkFsPath []) (RunNumber 0)+ r <- Run.fromWriteBuffer hasFS hasBlockIO benchSalt runParams fsps wb wbblobs+ let NumEntries nentriesReal = Run.size r+ assertEqual nentriesReal nentries $ pure ()+ -- 42 to 43 entries per page+ assertEqual (nentriesReal `div` getNumPages (Run.sizeInPages r)) 42 $ pure ()+ pure ( benchTmpDir+ , arenaManager+ , hasFS+ , hasBlockIO+ , wbblobs+ , V.singleton r+ , lookupKeys+ )+ where+ runParams :: RunBuilder.RunParams+ runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Compact+ }++lookupsInBatchesCleanup ::+ ( FilePath -- ^ Temporary directory+ , ArenaManager RealWorld+ , FS.HasFS IO FS.HandleIO+ , FS.HasBlockIO IO FS.HandleIO+ , Ref (WBB.WriteBufferBlobs IO FS.HandleIO)+ , V.Vector (Ref (Run IO FS.HandleIO))+ , V.Vector SerialisedKey+ )+ -> IO ()+lookupsInBatchesCleanup (tmpDir, _arenaManager, _hasFS, hasBlockIO, wbblobs, rs, _) = do+ FS.close hasBlockIO+ forM_ rs releaseRef+ releaseRef wbblobs+ removeDirectoryRecursive tmpDir++-- | Generate keys to store and keys to lookup+lookupsEnv ::+ StdGen -- ^ RNG+ -> Int -- ^ Number of stored key\/operation pairs+ -> Int -- ^ Number of positive lookups+ -> Int -- ^ Number of negative lookups+ -> IO ( Map SerialisedKey (Entry SerialisedValue SerialisedBlob)+ , V.Vector SerialisedKey+ )+lookupsEnv g nentries npos nneg = do+ let (g1, g') = R.splitGen g+ (g2, g'') = R.splitGen g'+ (g3, g4) = R.splitGen g''+ let (keys, negLookups) = splitAt nentries+ $ uniformWithoutReplacement @UTxOKey g1 (nentries + nneg)+ posLookups = sampleUniformWithReplacement g2 npos keys+ let values = take nentries $ List.unfoldr (Just . randomEntry) g3+ entries = Map.fromList $ zip keys values++ let lookups = shuffle (negLookups ++ posLookups) g4+ entries' = Map.mapKeys serialiseKey+ $ Map.map (bimap serialiseValue serialiseBlob) entries+ lookups' = V.fromList $ fmap serialiseKey lookups+ assertEqual (Map.size entries') (nentries) $ pure ()+ assertEqual (length lookups') (npos + nneg) $ pure ()+ pure (entries', lookups')++randomEntry :: StdGen -> (Entry UTxOValue ByteString, StdGen)+randomEntry g = frequency [+ (20, \g' -> let (!v, !g'') = uniform g' in (Insert v, g''))+ , (1, \g' -> let (!v, !g'') = uniform g'+ -- The size of the blobs doesn't matter for the benchmark,+ -- as it only deals with the blob references. So we make+ -- them tiny to not slow down the setup.+ (!b, !g''') = randomByteStringR (0, 100) g''+ in (InsertWithBlob v b, g'''))+ , (2, \g' -> let (!v, !g'') = uniform g' in (Upsert v, g''))+ , (2, \g' -> (Delete, g'))+ ] g++-- | Assertions on the generated environment should also be checked for release+-- builds, so don't use 'Control.Exception.assert'.+assertEqual :: (HasCallStack, Eq a, Show a) => a -> a -> b -> b+assertEqual x y+ | x == y = id+ | otherwise = error $ show x ++ " /= " ++ show y
@@ -0,0 +1,453 @@+module Bench.Database.LSMTree.Internal.Merge (benchmarks) where++import Control.RefCount+import Criterion.Main (Benchmark, bench, bgroup)+import qualified Criterion.Main as Cr+import Data.Bifunctor (first)+import qualified Data.BloomFilter.Hash as Hash+import Data.Foldable (traverse_)+import Data.IORef+import qualified Data.List as List+import qualified Data.Map.Strict as Map+import Data.Maybe (fromMaybe)+import qualified Data.Vector as V+import Data.Word (Word64)+import Database.LSMTree.Extras.Orphans ()+import qualified Database.LSMTree.Extras.Random as R+import Database.LSMTree.Extras.RunData+import Database.LSMTree.Extras.UTxO+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Compact))+import Database.LSMTree.Internal.Merge (MergeType (..))+import qualified Database.LSMTree.Internal.Merge as Merge+import Database.LSMTree.Internal.Paths (RunFsPaths (..))+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import Prelude hiding (getContents)+import System.Directory (removeDirectoryRecursive)+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.IO as FS+import qualified System.FS.IO as FS+import System.IO.Temp+import qualified System.Random as R+import System.Random (StdGen, mkStdGen, uniform, uniformR)++benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.Merge" [+ -- various numbers of runs+ benchMerge configWord64+ { name = "word64-insert-x2"+ , nentries = totalEntries `splitInto` 2+ , finserts = 1+ }+ , benchMerge configWord64+ { name = "word64-insert-x4"+ , nentries = totalEntries `splitInto` 4+ , finserts = 1+ }+ , benchMerge configWord64+ { name = "word64-insert-x7"+ , nentries = totalEntries `splitInto` 7+ , finserts = 1+ }+ , benchMerge configWord64+ { name = "word64-insert-x13"+ , nentries = totalEntries `splitInto` 13+ , finserts = 1+ }+ -- different operations+ , benchMerge configWord64+ { name = "word64-delete-x4"+ , nentries = totalEntries `splitInto` 4+ , fdeletes = 1+ }+ , benchMerge configWord64+ { name = "word64-blob-x4"+ , nentries = totalEntries `splitInto` 4+ , fblobinserts = 1+ }+ , benchMerge configWord64+ { name = "word64-mupsert-x4" -- basically no collisions+ , nentries = totalEntries `splitInto` 4+ , fmupserts = 1+ , mergeResolve = Just (onDeserialisedValues ((+) @Word64))+ }+ , benchMerge configWord64+ { name = "word64-mupsert-collisions-x4"+ , nentries = totalEntries `splitInto` 4+ , fmupserts = 1+ , randomKey = -- each run uses half of the possible keys+ randomWord64OutOf (totalEntries `div` 2)+ , mergeResolve = Just (onDeserialisedValues ((+) @Word64))+ }+ -- different merge types+ , benchMerge configWord64+ { name = "word64-mix-collisions-x4-midlevel"+ , nentries = totalEntries `splitInto` 4+ , finserts = 1+ , fdeletes = 1+ , fmupserts = 1+ , randomKey = -- each run uses half of the possible keys+ randomWord64OutOf (totalEntries `div` 2)+ , mergeResolve = Just (onDeserialisedValues ((+) @Word64))+ , mergeType = MergeTypeMidLevel+ }+ , benchMerge configWord64+ { name = "word64-mix-collisions-x4-lastlevel"+ , nentries = totalEntries `splitInto` 4+ , finserts = 1+ , fdeletes = 1+ , fmupserts = 1+ , randomKey = -- each run uses half of the possible keys+ randomWord64OutOf (totalEntries `div` 2)+ , mergeResolve = Just (onDeserialisedValues ((+) @Word64))+ , mergeType = MergeTypeLastLevel+ }+ , benchMerge configWord64+ { name = "word64-mix-collisions-x4-union"+ , nentries = totalEntries `splitInto` 4+ , finserts = 1+ , fdeletes = 1+ , fmupserts = 1+ , randomKey = -- each run uses half of the possible keys+ randomWord64OutOf (totalEntries `div` 2)+ , mergeResolve = Just (onDeserialisedValues ((+) @Word64))+ , mergeType = MergeTypeUnion+ }+ -- not writing anything at all+ , benchMerge configWord64+ { name = "word64-delete-x4-lastlevel"+ , nentries = totalEntries `splitInto` 4+ , fdeletes = 1+ , mergeType = MergeTypeLastLevel+ }+ -- different key and value sizes+ , benchMerge configWord64+ { name = "insert-large-keys-x4" -- potentially long keys+ , nentries = (totalEntries `div` 10) `splitInto` 4+ , finserts = 1+ , randomKey = first serialiseKey . R.randomByteStringR (8, 4000)+ }+ , benchMerge configWord64+ { name = "insert-mixed-vals-x4" -- potentially long values+ , nentries = (totalEntries `div` 10) `splitInto` 4+ , finserts = 1+ , randomValue = first serialiseValue . R.randomByteStringR (0, 4000)+ }+ , benchMerge configWord64+ { name = "insert-page-x4" -- 1 page+ , nentries = (totalEntries `div` 10) `splitInto` 4+ , finserts = 1+ , randomValue = first serialiseValue . R.randomByteStringR (4056, 4056)+ }+ , benchMerge configWord64+ { name = "insert-page-plus-byte-x4" -- 1 page + 1 byte+ , nentries = (totalEntries `div` 10) `splitInto` 4+ , finserts = 1+ , randomValue = first serialiseValue . R.randomByteStringR (4057, 4057)+ }+ , benchMerge configWord64+ { name = "insert-huge-vals-x4" -- 1-5 pages+ , nentries = (totalEntries `div` 10) `splitInto` 4+ , finserts = 1+ , randomValue = first serialiseValue . R.randomByteStringR (10_000, 20_000)+ }+ -- common UTxO scenarios+ , benchMerge configUTxO+ { name = "utxo-x4" -- like tiering merge+ , nentries = totalEntries `splitInto` 4+ , finserts = 1+ , fdeletes = 1+ }+ , benchMerge configUTxO+ { name = "utxo-x4-uneven"+ , nentries = totalEntries `distributed` [1, 3, 1.5, 2.5]+ , finserts = 1+ , fdeletes = 1+ }+ , benchMerge configUTxO+ { name = "utxo-x4-lastlevel"+ , nentries = totalEntries `splitInto` 4+ , finserts = 1+ , fdeletes = 1+ , mergeType = MergeTypeLastLevel+ }+ , benchMerge configUTxO+ { name = "utxo-x4+1-min-skewed-lastlevel" -- live levelling merge+ , nentries = totalEntries `distributed` [1, 1, 1, 1, 4]+ , finserts = 1+ , fdeletes = 1+ , mergeType = MergeTypeLastLevel+ }+ , benchMerge configUTxO+ { name = "utxo-x4+1-max-skewed-lastlevel" -- live levelling merge+ , nentries = totalEntries `distributed` [1, 1, 1, 1, 16]+ , finserts = 1+ , fdeletes = 1+ , mergeType = MergeTypeLastLevel+ }+ , benchMerge configUTxOStaking+ { name = "utxo-x2-tree-union" -- binary union merge+ , nentries = totalEntries `distributed` [4, 1]+ , mergeType = MergeTypeUnion+ }+ , benchMerge configUTxOStaking+ { name = "utxo-x10-tree-level" -- merge a whole table (for union)+ , nentries = totalEntries `distributed` [ 1, 1, 1+ , 4, 4, 4+ , 16, 16, 16+ , 100 -- last level+ ]+ , mergeType = MergeTypeLastLevel+ }+ ]+ where+ totalEntries = 50_000++ splitInto :: Int -> Int -> [Int]+ n `splitInto` k = n `distributed` replicate k 1++ distributed :: Int -> [Double] -> [Int]+ n `distributed` weights =+ let total = sum weights+ in [ round (fromIntegral n * w / total)+ | w <- weights+ ]++benchSalt :: Bloom.Salt+benchSalt = 4++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Compact+ }++benchMerge :: Config -> Benchmark+benchMerge conf@Config{name} =+ withEnv $ \ ~(_dir, hasFS, hasBlockIO, runs) ->+ bgroup name [+ bench "merge" $+ -- We'd like to do: `whnfAppIO (runs' -> ...) runs`.+ -- However, we also need per-run cleanup to avoid running out of+ -- disk space. We use `perRunEnvWithCleanup`, which has two issues:+ -- 1. Just as `whnfAppIO` etc., it takes an IO action and returns+ -- `Benchmarkable`, which does not compose. As a workaround, we+ -- thread `runs` through the environment, too.+ -- 2. It forces the result to normal form, which would traverse the+ -- whole run, so we force to WHNF ourselves and just return `()`.+ -- 3. It doesn't have access to the run we created in the benchmark,+ -- but the cleanup should not be part of the measurement, as it+ -- includes deleting files. So we smuggle the reference out using+ -- an `IORef`.+ Cr.perRunEnvWithCleanup+ ((runs,) <$> newIORef Nothing)+ (releaseRun . snd) $ \(runs', ref) -> do+ !run <- merge hasFS hasBlockIO conf outputRunPaths runs'+ writeIORef ref $ Just $ releaseRef run+ ]+ where+ withEnv =+ Cr.envWithCleanup+ (mergeEnv conf)+ mergeEnvCleanup++ releaseRun :: IORef (Maybe (IO ())) -> IO ()+ releaseRun ref =+ readIORef ref >>= \case+ Nothing -> pure ()+ Just release -> release++merge ::+ FS.HasFS IO FS.HandleIO+ -> FS.HasBlockIO IO FS.HandleIO+ -> Config+ -> Run.RunFsPaths+ -> InputRuns+ -> IO (Ref (Run IO FS.HandleIO))+merge fs hbio Config {..} targetPaths runs = do+ let f = fromMaybe const mergeResolve+ m <- fromMaybe (error "empty inputs, no merge created") <$>+ Merge.new fs hbio benchSalt runParams mergeType f targetPaths runs+ Merge.stepsToCompletion m stepSize++fsPath :: FS.FsPath+fsPath = FS.mkFsPath []++outputRunPaths :: Run.RunFsPaths+outputRunPaths = RunFsPaths fsPath (RunNumber 0)++inputRunPathsCounter :: IO (UniqCounter IO)+inputRunPathsCounter = newUniqCounter 1 -- 0 is for output++type InputRuns = V.Vector (Ref (Run IO FS.HandleIO))++onDeserialisedValues ::+ SerialiseValue v => (v -> v -> v) -> ResolveSerialisedValue+onDeserialisedValues f x y =+ serialiseValue (f (deserialiseValue x) (deserialiseValue y))++{-------------------------------------------------------------------------------+ Environments+-------------------------------------------------------------------------------}++-- | Config options describing a benchmarking scenario+data Config = Config {+ -- | Name for the benchmark scenario described by this config.+ name :: !String+ -- | Number of key\/operation pairs, one for each run.+ , nentries :: ![Int]+ -- | Frequency of inserts within the key\/op pairs.+ , finserts :: !Int+ -- | Frequency of inserts with blobs within the key\/op pairs.+ , fblobinserts :: !Int+ -- | Frequency of deletes within the key\/op pairs.+ , fdeletes :: !Int+ -- | Frequency of mupserts within the key\/op pairs.+ , fmupserts :: !Int+ , randomKey :: Rnd SerialisedKey+ , randomValue :: Rnd SerialisedValue+ , randomBlob :: Rnd SerialisedBlob+ , mergeType :: !MergeType+ -- | Needs to be defined when generating mupserts.+ , mergeResolve :: !(Maybe ResolveSerialisedValue)+ -- | Merging is done in chunks of @stepSize@ entries.+ , stepSize :: !Int+ }++type Rnd a = StdGen -> (a, StdGen)++defaultConfig :: Config+defaultConfig = Config {+ name = "default"+ , nentries = []+ , finserts = 0+ , fblobinserts = 0+ , fdeletes = 0+ , fmupserts = 0+ , randomKey = error "randomKey not implemented"+ , randomValue = error "randomValue not implemented"+ , randomBlob = error "randomBlob not implemented"+ , mergeType = MergeTypeMidLevel+ , mergeResolve = Nothing+ , stepSize = maxBound -- by default, just do in one go+ }++configWord64 :: Config+configWord64 = defaultConfig {+ randomKey = first serialiseKey . uniform @Word64 @_+ , randomValue = first serialiseValue . uniform @Word64 @_+ , randomBlob = first serialiseBlob . R.randomByteStringR (0, 0x2000) -- up to 8 kB+ }++configUTxO :: Config+configUTxO = defaultConfig {+ randomKey = first serialiseKey . uniform @UTxOKey @_+ , randomValue = first serialiseValue . uniform @UTxOValue @_+ }++configUTxOStaking :: Config+configUTxOStaking = defaultConfig {+ fmupserts = 1+ , randomKey = first serialiseKey . uniform @UTxOKey @_+ , randomValue = first serialiseValue . uniform @Word64 @_+ , mergeResolve = Just (onDeserialisedValues ((+) @Word64))+ }++mergeEnv ::+ Config+ -> IO ( FilePath -- ^ Temporary directory+ , FS.HasFS IO FS.HandleIO+ , FS.HasBlockIO IO FS.HandleIO+ , InputRuns+ )+mergeEnv config = do+ sysTmpDir <- getCanonicalTemporaryDirectory+ benchTmpDir <- createTempDirectory sysTmpDir "mergeEnv"+ (hasFS, hasBlockIO) <- FS.ioHasBlockIO (FS.MountPoint benchTmpDir) FS.defaultIOCtxParams+ runs <- randomRuns hasFS hasBlockIO config (mkStdGen 17)+ pure (benchTmpDir, hasFS, hasBlockIO, runs)++mergeEnvCleanup ::+ ( FilePath -- ^ Temporary directory+ , FS.HasFS IO FS.HandleIO+ , FS.HasBlockIO IO FS.HandleIO+ , InputRuns+ )+ -> IO ()+mergeEnvCleanup (tmpDir, _hasFS, hasBlockIO, runs) = do+ traverse_ releaseRef runs+ removeDirectoryRecursive tmpDir+ FS.close hasBlockIO++-- | Generate keys and entries to insert into the write buffer.+-- They are already serialised to exclude the cost from the benchmark.+randomRuns ::+ FS.HasFS IO FS.HandleIO+ -> FS.HasBlockIO IO FS.HandleIO+ -> Config+ -> StdGen+ -> IO InputRuns+randomRuns hasFS hasBlockIO config@Config {..} rng0 = do+ counter <- inputRunPathsCounter+ fmap V.fromList $+ mapM (unsafeCreateRun hasFS hasBlockIO benchSalt runParams fsPath counter) $+ zipWith+ (randomRunData config)+ nentries+ (List.unfoldr (Just . R.splitGen) rng0)++-- | Generate keys and entries to insert into the write buffer.+-- They are already serialised to exclude the cost from the benchmark.+randomRunData ::+ Config+ -> Int -- ^ number of entries+ -> StdGen -- ^ RNG+ -> SerialisedRunData+randomRunData Config {..} runentries g0 =+ RunData . Map.fromList $+ zip+ (R.withoutReplacement g1 runentries randomKey)+ (R.withReplacement g2 runentries randomEntry)+ where+ (g1, g2) = R.splitGen g0++ randomEntry :: Rnd (Entry SerialisedValue SerialisedBlob)+ randomEntry = R.frequency+ [ ( finserts+ , \g -> let (!v, !g') = randomValue g+ in (Insert v, g')+ )+ , ( fblobinserts+ , \g -> let (!v, !g') = randomValue g+ (!b, !g'') = randomBlob g'+ in (InsertWithBlob v b, g'')+ )+ , ( fdeletes+ , \g -> (Delete, g)+ )+ , ( fmupserts+ , \g -> let (!v, !g') = randomValue g+ in (Upsert v, g')+ )+ ]++-- | @randomWord64OutOf n@ generates one out of @n@ distinct keys, which are+-- uniformly distributed.+-- This can be used to make collisions more likely.+--+-- Make sure to pick an @n@ that is significantly larger than the size of a run!+-- Each run entry needs a distinct key.+randomWord64OutOf :: Int -> Rnd SerialisedKey+randomWord64OutOf possibleKeys =+ first (serialiseKey . Hash.hashSalt64 benchSalt)+ . uniformR (0, fromIntegral possibleKeys :: Word64)
@@ -0,0 +1,68 @@+{-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE TypeApplications #-}++{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}++module Bench.Database.LSMTree.Internal.RawPage (+ benchmarks+ -- * Benchmarked functions+ ) where++import Control.DeepSeq (deepseq)+import qualified Data.ByteString as BS++import Database.LSMTree.Extras.ReferenceImpl+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Serialise++import Criterion.Main+import Test.QuickCheck+import Test.QuickCheck.Gen (Gen (..))+import Test.QuickCheck.Random (mkQCGen)++benchmarks :: Benchmark+benchmarks = rawpage `deepseq` bgroup "Bench.Database.LSMTree.Internal.RawPage"+ [ bRawPageFindKey+ , bRawPageLookup+ ]+ where+ bRawPageFindKey = bgroup "rawPageFindKey"+ [ bench "missing" $ whnf (rawPageFindKey rawpage) missing+ , bench "existing-head" $ whnf (rawPageFindKey rawpage) existingHead+ , bench "existing-last" $ whnf (rawPageFindKey rawpage) existingLast+ ]++ bRawPageLookup = bgroup "rawPageLookup"+ [ bench "missing" $ whnf (rawPageLookup rawpage) missing+ , bench "existing-head" $ whnf (rawPageLookup rawpage) existingHead+ , bench "existing-last" $ whnf (rawPageLookup rawpage) existingLast+ ]++ kops :: [(Key, Operation)]+ kops = unGen genPage (mkQCGen 42) 200+ where+ genPage = orderdKeyOps <$>+ genPageContentNearFull DiskPage4k genSmallKey genSmallValue++ rawpage :: RawPage+ rawpage = fst $ toRawPage (PageContentFits kops)++ genSmallKey :: Gen Key+ genSmallKey = Key . BS.pack <$> vectorOf 8 arbitrary++ genSmallValue :: Gen Value+ genSmallValue = Value . BS.pack <$> vectorOf 8 arbitrary++ missing :: SerialisedKey+ missing = SerialisedKey $ RB.pack [1, 2, 3]++ keys :: [Key]+ keys = map fst kops++ existingHead :: SerialisedKey+ existingHead = SerialisedKey $ RB.fromByteString $ unKey $ head keys++ existingLast :: SerialisedKey+ existingLast = SerialisedKey $ RB.fromByteString $ unKey $ last keys+
@@ -0,0 +1,20 @@+module Bench.Database.LSMTree.Internal.Serialise (+ benchmarks+ ) where++import Criterion.Main+import Database.LSMTree.Extras.UTxO+import Database.LSMTree.Internal.Serialise.Class+import System.Random++benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.Serialise" [+ env (pure $ fst $ uniform (mkStdGen 12)) $ \(k :: UTxOKey) ->+ bgroup "UTxOKey" [+ bench "serialiseKey" $ whnf serialiseKey k+ , bench "serialiseKeyRoundtrip" $ whnf serialiseKeyRoundtrip k+ ]+ ]++serialiseKeyRoundtrip :: SerialiseKey k => k -> k+serialiseKeyRoundtrip k = deserialiseKey (serialiseKey k)
@@ -0,0 +1,224 @@+module Bench.Database.LSMTree.Internal.WriteBuffer (benchmarks) where++import Control.DeepSeq (NFData (..), rwhnf)+import Control.Exception (assert)+import Criterion.Main (Benchmark, bench, bgroup)+import qualified Criterion.Main as Cr+import Data.Bifunctor (first)+import qualified Data.Foldable as Fold+import qualified Data.List as List+import Data.Maybe (fromMaybe, isJust, isNothing)+import Data.Word (Word64)+import Database.LSMTree.Extras.Orphans ()+import Database.LSMTree.Extras.Random (frequency, randomByteStringR)+import Database.LSMTree.Extras.UTxO+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.WriteBuffer (WriteBuffer)+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import System.Random (StdGen, mkStdGen, uniform)++benchmarks :: Benchmark+benchmarks = bgroup "Bench.Database.LSMTree.Internal.WriteBuffer" [+ benchWriteBuffer configWord64+ { name = "word64-insert-10k"+ , nentries = 10_000+ , finserts = 1+ }+ , benchWriteBuffer configWord64+ { name = "word64-delete-10k"+ , nentries = 10_000+ , fdeletes = 1+ }+ , benchWriteBuffer configWord64+ { name = "word64-blob-10k"+ , nentries = 10_000+ , fblobinserts = 1+ }+ , benchWriteBuffer configWord64+ { name = "word64-mupsert-10k"+ , nentries = 10_000+ , fmupserts = 1+ -- TODO: too few collisions to really measure resolution+ , resolveVal = Just (onDeserialisedValues ((+) @Word64))+ }+ -- different key and value sizes+ , benchWriteBuffer configWord64+ { name = "insert-large-keys-1k" -- large keys+ , nentries = 1_000+ , finserts = 1+ , randomKey = first serialiseKey . randomByteStringR (6, 4000)+ }+ , benchWriteBuffer configWord64+ { name = "insert-mixed-vals-1k" -- small and large values+ , nentries = 1_000+ , finserts = 1+ , randomValue = first serialiseValue . randomByteStringR (0, 8000)+ }+ , benchWriteBuffer configWord64+ { name = "insert-page-1k" -- 1 page+ , nentries = 1_000+ , finserts = 1+ , randomValue = first serialiseValue . randomByteStringR (4056, 4056)+ }+ , benchWriteBuffer configWord64+ { name = "insert-page-plus-byte-1k" -- 1 page + 1 byte+ , nentries = 1_000+ , finserts = 1+ , randomValue = first serialiseValue . randomByteStringR (4057, 4057)+ }+ , benchWriteBuffer configWord64+ { name = "insert-huge-vals-1k" -- 3-5 pages+ , nentries = 1_000+ , finserts = 1+ , randomValue = first serialiseValue . randomByteStringR (10_000, 20_000)+ }+ -- UTxO workload+ -- compare different buffer sizes to see superlinear cost of map insertion+ , benchWriteBuffer configUTxO+ { name = "utxo-2k"+ , nentries = 2_000+ , finserts = 1+ , fdeletes = 1+ }+ , benchWriteBuffer configUTxO+ { name = "utxo-10k"+ , nentries = 10_000+ , finserts = 1+ , fdeletes = 1+ }+ , benchWriteBuffer configUTxO+ { name = "utxo-50k"+ , nentries = 50_000+ , finserts = 1+ , fdeletes = 1+ }+ ]++benchWriteBuffer :: Config -> Benchmark+benchWriteBuffer conf@Config{name} =+ Cr.env (pure (envInputKOps conf)) $ \ kops ->+ bgroup name [+ bench "insert" $+ Cr.whnf (\kops' -> insert kops') kops+ --TODO: re-add I/O tests here:+ -- * writing out blobs during insert+ -- * flushing write buffer to run+ ]++insert :: InputKOps -> WriteBuffer+insert (InputKOps kops resolveVal) =+ Fold.foldl' (\wb (k, e) -> WB.addEntry resolveVal k e wb) WB.empty kops++data InputKOps =+ InputKOps+ [(SerialisedKey, Entry SerialisedValue BlobSpan)]+ ResolveSerialisedValue++instance NFData InputKOps where+ rnf (InputKOps kops resolveVal) = rnf kops `seq` rwhnf resolveVal++onDeserialisedValues ::+ SerialiseValue v => (v -> v -> v) -> ResolveSerialisedValue+onDeserialisedValues f x y =+ serialiseValue (f (deserialiseValue x) (deserialiseValue y))++type SerialisedKOp = (SerialisedKey, SerialisedEntry)+type SerialisedEntry = Entry SerialisedValue BlobSpan++{-------------------------------------------------------------------------------+ Environments+-------------------------------------------------------------------------------}++-- | Config options describing a benchmarking scenario+data Config = Config {+ -- | Name for the benchmark scenario described by this config.+ name :: !String+ -- | Number of key\/operation pairs in the run+ , nentries :: !Int+ -- | Frequency of inserts within the key\/op pairs.+ , finserts :: !Int+ -- | Frequency of inserts with blobs within the key\/op pairs.+ , fblobinserts :: !Int+ -- | Frequency of deletes within the key\/op pairs.+ , fdeletes :: !Int+ -- | Frequency of mupserts within the key\/op pairs.+ , fmupserts :: !Int+ , randomKey :: Rnd SerialisedKey+ , randomValue :: Rnd SerialisedValue+ -- | Needs to be defined when generating mupserts.+ , resolveVal :: !(Maybe ResolveSerialisedValue)+ }++type Rnd a = StdGen -> (a, StdGen)++defaultConfig :: Config+defaultConfig = Config {+ name = "default"+ , nentries = 0+ , finserts = 0+ , fblobinserts = 0+ , fdeletes = 0+ , fmupserts = 0+ , randomKey = error "randomKey not implemented"+ , randomValue = error "randomValue not implemented"+ , resolveVal = Nothing+ }++configWord64 :: Config+configWord64 = defaultConfig {+ randomKey = first serialiseKey . uniform @Word64 @_+ , randomValue = first serialiseValue . uniform @Word64 @_+ }++configUTxO :: Config+configUTxO = defaultConfig {+ randomKey = first serialiseKey . uniform @UTxOKey @_+ , randomValue = first serialiseValue . uniform @UTxOValue @_+ }++envInputKOps :: Config -> InputKOps+envInputKOps config = do+ let kops = randomKOps config (mkStdGen 17)+ in InputKOps kops (fromMaybe const (resolveVal config))++-- | Generate keys and entries to insert into the write buffer.+-- They are already serialised to exclude the cost from the benchmark.+randomKOps ::+ Config+ -> StdGen -- ^ RNG+ -> [SerialisedKOp]+randomKOps Config {..} = take nentries . List.unfoldr (Just . randomKOp) .+ assert (if fmupserts > 0 then isJust resolveVal else isNothing resolveVal)+ where+ randomKOp :: Rnd SerialisedKOp+ randomKOp g = let (!k, !g') = randomKey g+ (!e, !g'') = randomEntry g'+ in ((k, e), g'')++ randomEntry :: Rnd SerialisedEntry+ randomEntry = frequency+ [ ( finserts+ , \g -> let (!v, !g') = randomValue g+ in (Insert v, g')+ )+ , ( fblobinserts+ , \g -> let (!v, !g') = randomValue g+ (!b, !g'') = randomBlobSpan g'+ in (InsertWithBlob v b, g'')+ )+ , ( fdeletes+ , \g -> (Delete, g)+ )+ , ( fmupserts+ , \g -> let (!v, !g') = randomValue g+ in (Upsert v, g')+ )+ ]++randomBlobSpan :: Rnd BlobSpan+randomBlobSpan !g =+ let (off, !g') = uniform g+ (len, !g'') = uniform g'+ in (BlobSpan off len, g'')
@@ -0,0 +1,32 @@+{-# LANGUAGE CPP #-}++-- | Micro-benchmarks for the @lsm-tree@ library.+module Main (main) where++import qualified Bench.Database.LSMTree+import qualified Bench.Database.LSMTree.Internal.BloomFilter+import qualified Bench.Database.LSMTree.Internal.Index+import qualified Bench.Database.LSMTree.Internal.Index.Compact+import qualified Bench.Database.LSMTree.Internal.Lookup+import qualified Bench.Database.LSMTree.Internal.Merge+import qualified Bench.Database.LSMTree.Internal.RawPage+import qualified Bench.Database.LSMTree.Internal.Serialise+import qualified Bench.Database.LSMTree.Internal.WriteBuffer+import Criterion.Main (defaultMain)++main :: IO ()+main = do+#ifdef NO_IGNORE_ASSERTS+ putStrLn "WARNING: BENCHMARKING A BUILD IN DEBUG MODE"+#endif+ defaultMain [+ Bench.Database.LSMTree.Internal.BloomFilter.benchmarks+ , Bench.Database.LSMTree.Internal.Index.benchmarks+ , Bench.Database.LSMTree.Internal.Index.Compact.benchmarks+ , Bench.Database.LSMTree.Internal.Lookup.benchmarks+ , Bench.Database.LSMTree.Internal.Merge.benchmarks+ , Bench.Database.LSMTree.Internal.RawPage.benchmarks+ , Bench.Database.LSMTree.Internal.Serialise.benchmarks+ , Bench.Database.LSMTree.Internal.WriteBuffer.benchmarks+ , Bench.Database.LSMTree.benchmarks+ ]
@@ -0,0 +1,1180 @@+cabal-version: 3.4+name: lsm-tree+version: 1.0.0.0+synopsis: Log-structured merge-trees+description:+ This package contains an efficient implementation of on-disk key–value storage, implemented as a log-structured merge-tree, LSM-tree or LSMT.+ An LSM-tree is a data structure for key–value mappings, similar to "Data.Map", but optimized for large tables with a high insertion volume.+ It has support for:++ * Basic key–value operations, such as lookup, insert, and delete.+ * Range lookups, which efficiently retrieve the values for all keys in a given range.+ * Monoidal upserts which combine the stored and new values.+ * BLOB storage which associates a large auxiliary BLOB with a key.+ * Durable on-disk persistence and rollback via named snapshots.+ * Cheap table duplication where all duplicates can be independently accessed and modified.+ * High-performance lookups on SSDs using I\/O batching and parallelism.++ This package exports two modules:++ * "Database.LSMTree.Simple"++ This module exports a simplified API which picks sensible defaults for a number of configuration parameters.++ It does not support upserts or BLOBs, due to their unintuitive interaction, see [Upsert and BLOB](#upsertandblob).++ If you are looking at this package for the first time, it is strongly recommended that you start by reading this module.++ * "Database.LSMTree"++ This module exports the full API.++ == Upsert and BLOB #upsertandblob#++ The interaction between upserts and BLOBs is unintuitive.+ A upsert updates the value associated with the key by combining the new and old values with a user-specified function.+ However, any BLOB associated with the key is simply deleted.++ == Portability #portability#++ * This package only supports 64-bit, little-endian systems.+ * On Windows, the package has only been tested with NTFS filesystems.+ * On Linux, executables using this package, including test and benchmark suites, must be compiled with the [@-threaded@](https://downloads.haskell.org/ghc/latest/docs/users_guide/phases.html#ghc-flag-threaded) RTS option enabled.++ == Concurrency #concurrency#++ LSM-trees can be used concurrently, but with a few restrictions:++ * Each session locks its session directory.+ This means that a database cannot be accessed from different processes at the same time.+ * Tables can be used concurrently and concurrent use of read operations such as lookups is deterministic.+ However, concurrent use of write operations such as insert or delete with any other operation results in a race condition.++ == Performance #performance#++ The worst-case behaviour of the library is described using [big-O notation](http://en.wikipedia.org/wiki/Big_O_notation).+ The documentation provides two measures of complexity:++ * The time complexity of operations is described in terms of the number of disk I\/O operations and referred to as the disk I\/O complexity.+ In practice, the time of the operations on LSM-trees is dominated by the number of disk I\/O actions.+ * The space complexity of tables is described in terms of the in-memory size of an LSM-tree table.+ Both the in-memory and on-disk size of an LSM-tree table scale linearly with the number of physical entries.+ However, the in-memory size of an LSM-tree table is smaller than its on-disk size by a significant constant.+ This is discussed in detail below, under [In-memory size of tables](#performance_size).++ The complexities are described in terms of the following variables and constants:++ * The variable \(n\) refers to the number of /physical/ table entries.+ A /physical/ table entry is any key–operation pair, e.g., @Insert k v@ or @Delete k@, whereas a /logical/ table entry is determined by all physical entries with the same key.+ If the variable \(n\) is used to describe the complexity of an operation that involves multiple tables, it refers to the sum of all table entries.+ * The variable \(o\) refers to the number of open tables and cursors in the session.+ * The variable \(s\) refers to the number of snapshots in the session.+ * The variable \(b\) usually refers to the size of a batch of inputs\/outputs.+ Its precise meaning is explained for each occurrence.+ * The constant \(B\) refers to the size of the write buffer,+ which is determined by the @TableConfig@ parameter @confWriteBufferAlloc@.+ * The constant \(T\) refers to the size ratio of the table,+ which is determined by the @TableConfig@ parameter @confSizeRatio@.+ * The constant \(P\) refers to the average number of key–value pairs that fit in a page of memory.++ === Disk I\/O cost of operations #performance_time#++ The following table summarises the worst-case cost of the operations on LSM-trees measured in the number of disk I\/O operations.+ If the cost depends on the merge policy or merge schedule, then the table contains one entry for each relevant combination.+ Otherwise, the merge policy and\/or merge schedule is listed as N\/A.+ The merge policy and merge schedule are determined by the @TableConfig@ parameters @confMergePolicy@ and @confMergeSchedule@.++ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | Resource | Operation | Merge policy | Merge schedule | Worst-case disk I\/O complexity |+ +==========+========================+=================+=================+================================================++ | Session | Open | N\/A | N\/A | \(O(1)\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Close | @LazyLevelling@ | N\/A | \(O(o \: T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | Table | New | N\/A | N\/A | \(O(1)\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Close | @LazyLevelling@ | N\/A | \(O(T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Lookup | @LazyLevelling@ | N\/A | \(O(T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Range Lookup | N\/A | N\/A | \(O(T \: \log_T \frac{n}{B} + \frac{b}{P})\)* |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Insert\/Delete\/Update | @LazyLevelling@ | @Incremental@ | \(O(\frac{1}{P} \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | | @LazyLevelling@ | @OneShot@ | \(O(\frac{n}{P})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Duplicate | N\/A | N\/A | \(O(0)\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Union | N\/A | N\/A | \(O(\frac{n}{P})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | Snapshot | Save | @LazyLevelling@ | N\/A | \(O(T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Open | N\/A | N\/A | \(O(\frac{n}{P})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Delete | @LazyLevelling@ | N\/A | \(O(T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | List | N\/A | N\/A | \(O(s)\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | Cursor | New | @LazyLevelling@ | N\/A | \(O(T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Close | @LazyLevelling@ | N\/A | \(O(T \: \log_T \frac{n}{B})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------++ | | Next | N\/A | N\/A | \(O(\frac{1}{P})\) |+ +----------+------------------------+-----------------+-----------------+------------------------------------------------+++ (*The variable \(b\) refers to the number of entries retrieved by the range lookup.)++ === Table Size #performance_size#++ The in-memory and the on-disk size of an LSM-tree scale /linearly/ with the number of physical entries.+ However, the in-memory size is smaller by a significant factor.+ Let us look at a table that uses the default configuration and has 100 million entries with 34 byte keys and 60 byte values.+ The total size of 100 million key–value pairs is approximately 8.75GiB.+ Hence, the on-disk size would be at least 8.75GiB, not counting the overhead for metadata.++ The in-memory size would be approximately 265.39MiB:++ * The write buffer would store at most 20,000 entries, which is approximately 2.86MiB.+ * The fence-pointer indexes would store approximately 2.29 million keys, which is approximately 9.30MiB.+ * The Bloom filters would use 15.78 bits per entry, which is approximately 188.11MiB.++ For a discussion of how the sizes of these components are determined by the table configuration, see [Fine-tuning Table Configuration](#fine_tuning).++ The total size of an LSM-tree must not exceed \(2^{41}\) physical entries.+ Violation of this condition /is/ checked and will throw a 'TableTooLargeError'.++ === Fine-tuning Table Configuration #fine_tuning#++ [@confMergePolicy@]+ The /merge policy/ balances the performance of lookups against the performance of updates.+ Levelling favours lookups.+ Tiering favours updates.+ Lazy levelling strikes a middle ground between levelling and tiering, and moderately favours updates.+ This parameter is explicitly referenced in the documentation of those operations it affects.++ [@confSizeRatio@]+ The /size ratio/ pushes the effects of the merge policy to the extreme.+ If the size ratio is higher, levelling favours lookups more, and tiering and lazy levelling favour updates more.+ This parameter is referred to as \(T\) in the disk I\/O cost of operations.++ [@confWriteBufferAlloc@]+ The /write buffer capacity/ balances the performance of lookups and updates against the in-memory size of the table.+ If the write buffer is larger, it takes up more memory, but lookups and updates are more efficient.+ This parameter is referred to as \(B\) in the disk I\/O cost of operations.+ Irrespective of this parameter, the write buffer size cannot exceed 4GiB.++ [@confMergeSchedule@]+ The /merge schedule/ balances the performance of lookups and updates against the smooth performance of updates.+ The merge schedule does not affect the performance of table unions.+ With the one-shot merge schedule, lookups and updates are more efficient overall, but some updates may take much longer than others.+ With the incremental merge schedule, lookups and updates are less efficient overall, but each update does a similar amount of work.+ This parameter is explicitly referenced in the documentation of those operations it affects.++ [@confBloomFilterAlloc@]+ The Bloom filter size balances the performance of lookups against the in-memory size of the table.+ If the Bloom filters are larger, they take up more memory, but lookup operations are more efficient.++ [@confFencePointerIndex@]+ The /fence-pointer index type/ supports two types of indexes.+ The /ordinary/ indexes are designed to work with any key.+ The /compact/ indexes are optimised for the case where the keys in the database are uniformly distributed, e.g., when the keys are hashes.++ [@confDiskCachePolicy@]+ The /disk cache policy/ determines if lookup operations use the OS page cache.+ Caching may improve the performance of lookups and updates if database access follows certain patterns.++ [@confMergeBatchSize@]+ The merge batch size balances the maximum latency of individual update+ operations, versus the latency of a sequence of update operations. Bigger+ batches improves overall performance but some updates will take a lot+ longer than others. The default is to use a large batch size.++ ==== Fine-tuning: Merge Policy, Size Ratio, and Write Buffer Size #fine_tuning_data_layout#++ The configuration parameters @confMergePolicy@, @confSizeRatio@, and @confWriteBufferAlloc@ affect how the table organises its data.+ To understand what effect these parameters have, one must have a basic understanding of how an LSM-tree stores its data.+ The physical entries in an LSM-tree are key–operation pairs, which pair a key with an operation such as an @Insert@ with a value or a @Delete@.+ These key–operation pairs are organised into /runs/, which are sequences of key–operation pairs sorted by their key.+ Runs are organised into /levels/, which are unordered sequences or runs.+ Levels are organised hierarchically.+ Level 0 is kept in memory, and is referred to as the /write buffer/.+ All subsequent levels are stored on disk, with each run stored in its own file.+ The following shows an example LSM-tree layout, with each run as a boxed sequence of keys and each level as a row.++ \[+ \begin{array}{l:l}+ \text{Level}+ &+ \text{Data}+ \\+ 0+ &+ \fbox{\(\texttt{4}\,\_\)}+ \\+ 1+ &+ \fbox{\(\texttt{1}\,\texttt{3}\)}+ \quad+ \fbox{\(\texttt{2}\,\texttt{7}\)}+ \\+ 2+ &+ \fbox{\(\texttt{0}\,\texttt{2}\,\texttt{3}\,\texttt{4}\,\texttt{5}\,\texttt{6}\,\texttt{8}\,\texttt{9}\)}+ \end{array}+ \]++ The data in an LSM-tree is /partially sorted/: only the key–operation pairs within each run are sorted and deduplicated.+ As a rule of thumb, keeping more of the data sorted means lookup operations are faster but update operations are slower.++ The configuration parameters @confMergePolicy@, @confSizeRatio@, and @confWriteBufferAlloc@ directly affect a table's data layout.+ The parameter @confWriteBufferAlloc@ determines the capacity of the write buffer.++ [@AllocNumEntries maxEntries@]:+ The write buffer can contain at most @maxEntries@ entries.+ The constant \(B\) refers to the value of @maxEntries@.+ Irrespective of this parameter, the write buffer size cannot exceed 4GiB.++ The parameter @confSizeRatio@ determines the ratio between the capacities of successive levels.+ The constant \(T\) refers to the value of @confSizeRatio@.+ For instance, if \(B = 2\) and \(T = 2\), then++ \[+ \begin{array}{l:l}+ \text{Level} & \text{Capacity}+ \\+ 0 & B \cdot T^0 = 2+ \\+ 1 & B \cdot T^1 = 4+ \\+ 2 & B \cdot T^2 = 8+ \\+ \ell & B \cdot T^\ell+ \end{array}+ \]++ The merge policy @confMergePolicy@ determines the number of runs per level.+ In a /tiering/ LSM-tree, each level contains \(T\) runs.+ In a /levelling/ LSM-tree, each level contains one single run.+ The /lazy levelling/ policy uses levelling only for the last level and uses tiering for all preceding levels.+ The previous example used lazy levelling.+ The following examples illustrate the different merge policies using the same data, assuming \(B = 2\) and \(T = 2\).++ \[+ \begin{array}{l:l:l:l}+ \text{Level}+ &+ \text{Tiering}+ &+ \text{Levelling}+ &+ \text{Lazy Levelling}+ \\+ 0+ &+ \fbox{\(\texttt{4}\,\_\)}+ &+ \fbox{\(\texttt{4}\,\_\)}+ &+ \fbox{\(\texttt{4}\,\_\)}+ \\+ 1+ &+ \fbox{\(\texttt{1}\,\texttt{3}\)}+ \quad+ \fbox{\(\texttt{2}\,\texttt{7}\)}+ &+ \fbox{\(\texttt{1}\,\texttt{2}\,\texttt{3}\,\texttt{7}\)}+ &+ \fbox{\(\texttt{1}\,\texttt{3}\)}+ \quad+ \fbox{\(\texttt{2}\,\texttt{7}\)}+ \\+ 2+ &+ \fbox{\(\texttt{4}\,\texttt{5}\,\texttt{7}\,\texttt{8}\)}+ \quad+ \fbox{\(\texttt{0}\,\texttt{2}\,\texttt{3}\,\texttt{9}\)}+ &+ \fbox{\(\texttt{0}\,\texttt{2}\,\texttt{3}\,\texttt{4}\,\texttt{5}\,\texttt{6}\,\texttt{8}\,\texttt{9}\)}+ &+ \fbox{\(\texttt{0}\,\texttt{2}\,\texttt{3}\,\texttt{4}\,\texttt{5}\,\texttt{6}\,\texttt{8}\,\texttt{9}\)}+ \end{array}+ \]++ Tiering favours the performance of updates.+ Levelling favours the performance of lookups.+ Lazy levelling strikes a middle ground between tiering and levelling.+ It favours the performance of lookup operations for the oldest data and enables more deduplication,+ without the impact that full levelling has on update operations.++ ==== Fine-tuning: Merge Schedule #fine_tuning_merge_schedule#++ The configuration parameter @confMergeSchedule@ affects the worst-case performance of lookup and update operations and the structure of runs.+ Regardless of the merge schedule, the amortised disk I\/O complexity of lookups and updates is /logarithmic/ in the size of the table.+ When the write buffer fills up, its contents are flushed to disk as a run and added to level 1.+ When some level fills up, its contents are flushed down to the next level.+ Eventually, as data is flushed down, runs must be merged.+ This package supports two schedules for merging:++ * Using the @OneShot@ merge schedule, runs must always be kept fully sorted and deduplicated.+ However, flushing a run down to the next level may cause the next level to fill up,+ in which case it too must be flushed and merged futher down.+ In the worst case, this can cascade down the entire table.+ Consequently, the worst-case disk I\/O complexity of updates is /linear/ in the size of the table.+ This is unacceptable for real-time systems and other use cases where unresponsiveness is unacceptable.+ * Using the @Incremental@ merge schedule, runs can be /partially merged/, which allows the merging work to be spead out evenly across all update operations.+ This aligns the worst-case and average-case disk I\/O complexity of updates—both are /logarithmic/ in the size of the table.+ The cost is a small constant overhead for both lookup and update operations.++ The merge schedule does not affect the performance of table unions.+ The amortised disk I\/O complexity of one-shot unions is /linear/ in the size of the tables.+ Instead, there are separate operations for incremental and oneshot unions.+ For incremental unions, it is up to the user to spread the merging work out evenly over time.++ ==== Fine-tuning: Bloom Filter Size #fine_tuning_bloom_filter_size#++ The configuration parameter @confBloomFilterAlloc@ affects the size of the Bloom filters,+ which balances the performance of lookups against the in-memory size of the table.++ Tables maintain a [Bloom filter](https://en.wikipedia.org/wiki/Bloom_filter) in memory for each run on disk.+ These Bloom filters are probablilistic datastructures that are used to track which keys are present in their corresponding run.+ Querying a Bloom filter returns either \"maybe\" meaning the key is possibly in the run or \"no\" meaning the key is definitely not in the run.+ When a query returns \"maybe\" while the key is /not/ in the run, this is referred to as a /false positive/.+ While the database executes a lookup operation, any Bloom filter query that returns a false positive causes the database to unnecessarily read a page from disk.+ The probabliliy of these spurious reads follow a [binomial distribution](https://en.wikipedia.org/wiki/Binomial_distribution) \(\text{Binomial}(r,\text{FPR})\)+ where \(r\) refers to the number of runs and \(\text{FPR}\) refers to the false-positive rate of the Bloom filters.+ Hence, the expected number of spurious reads for each lookup operation is \(r\cdot\text{FPR}\).+ The number of runs \(r\) is proportional to the number of physical entries in the table. Its exact value depends on the merge policy of the table:++ [@LazyLevelling@]+ \(r = T (\log_T\frac{n}{B} - 1) + 1\).++ The false-positive rate scales exponentially with size of the Bloom filters in bits per entry.++ +---------------------------+----------------------++ | False-positive rate (FPR) | Bits per entry (BPE) |+ +===========================+======================++ | \(1\text{ in }10\) | \(\approx 4.77 \) |+ +---------------------------+----------------------++ | \(1\text{ in }100\) | \(\approx 9.85 \) |+ +---------------------------+----------------------++ | \(1\text{ in }1{,}000\) | \(\approx 15.78 \) |+ +---------------------------+----------------------++ | \(1\text{ in }10{,}000\) | \(\approx 22.57 \) |+ +---------------------------+----------------------++ | \(1\text{ in }100{,}000\) | \(\approx 30.22 \) |+ +---------------------------+----------------------+++ The configuration parameter @confBloomFilterAlloc@ can be specified in two ways:++ [@AllocFixed bitsPerEntry@]+ Allocate the requested number of bits per entry in the table.++ The value must strictly positive, but fractional values are permitted.+ The recommended range is \([2, 24]\).++ [@AllocRequestFPR falsePositiveRate@]+ Allocate the required number of bits per entry to get the requested false-positive rate.++ The value must be in the range \((0, 1)\).+ The recommended range is \([1\mathrm{e}{ -5 },1\mathrm{e}{ -2 }]\).++ The total in-memory size of all Bloom filters scales /linearly/ with the number of physical entries in the table and is determined by the number of physical entries multiplied by the number of bits per physical entry, i.e, \(n\cdot\text{BPE}\).+ Let us consider a table with 100 million physical entries which uses the default table configuration for every parameter other than the Bloom filter size.++ +---------------------------+----------------------+------------------------------------------------------------------++ | False-positive rate (FPR) | Bloom filter size | Expected spurious reads per lookup |+ +===========================+======================+==================================================================++ | \(1\text{ in }10\) | \( 56.86\text{MiB}\) | \( 2.56\text{ spurious reads every lookup }\) |+ +---------------------------+----------------------+------------------------------------------------------------------++ | \(1\text{ in }100\) | \(117.42\text{MiB}\) | \( 1 \text{ spurious read every } 3.91\text{ lookups }\) |+ +---------------------------+----------------------+------------------------------------------------------------------++ | \(1\text{ in }1{,}000\) | \(188.11\text{MiB}\) | \( 1 \text{ spurious read every } 39.10\text{ lookups }\) |+ +---------------------------+----------------------+------------------------------------------------------------------++ | \(1\text{ in }10{,}000\) | \(269.06\text{MiB}\) | \( 1 \text{ spurious read every } 391.01\text{ lookups }\) |+ +---------------------------+----------------------+------------------------------------------------------------------++ | \(1\text{ in }100{,}000\) | \(360.25\text{MiB}\) | \( 1 \text{ spurious read every } 3910.19\text{ lookups }\) |+ +---------------------------+----------------------+------------------------------------------------------------------+++ ==== Fine-tuning: Fence-Pointer Index Type #fine_tuning_fence_pointer_index_type#++ The configuration parameter @confFencePointerIndex@ affects the type and size of the fence-pointer indexes.+ Tables maintain a fence-pointer index in memory for each run on disk.+ These fence-pointer indexes store the keys at the boundaries of each page of memory to ensure that each lookup has to read at most one page of memory from each run.+ Tables support two types of fence-pointer indexes:++ [@OrdinaryIndex@]+ Ordinary indexes are designed for any use case.++ Ordinary indexes store one serialised key per page of memory.+ The average total in-memory size of all indexes is \(K \cdot \frac{n}{P}\) bits,+ where \(K\) refers to the average size of a serialised key in bits.++ [@CompactIndex@]+ Compact indexes are designed for the use case where the keys in the table are uniformly distributed, such as when using hashes.++ Compact indexes store the 64 most significant bits of the minimum serialised key of each page of memory.+ This requires that serialised keys are /at least/ 64 bits in size.+ Compact indexes store 1 additional bit per page of memory to resolve collisions, 1 additional bit per page of memory to mark entries that are larger than one page, and a negligible amount of memory for tie breakers.+ The average total in-memory size of all indexes is \(66 \cdot \frac{n}{P}\) bits.++ ==== Fine-tuning: Disk Cache Policy #fine_tuning_disk_cache_policy#++ The configuration parameter @confDiskCachePolicy@ determines how the database uses the OS page cache.+ This may improve performance if the database's /access pattern/ has good /temporal locality/ or good /spatial locality/.+ The database's access pattern refers to the pattern by which entries are accessed by lookup operations.+ An access pattern has good temporal locality if it is likely to access entries that were recently accessed or updated.+ An access pattern has good spatial locality if it is likely to access entries that have nearby keys.++ * Use the @DiskCacheAll@ policy if the database's access pattern has either good spatial locality or both good spatial and temporal locality.+ * Use the @DiskCacheLevelOneTo l@ policy if the database's access pattern has good temporal locality for updates only.+ The variable @l@ determines the number of levels that are cached.+ For a description of levels, see [Merge Policy, Size Ratio, and Write Buffer Size](#fine_tuning_data_layout).+ With this setting, the database can be expected to cache up to \(\frac{k}{P}\) pages of memory,+ where \(k\) refers to the number of entries that fit in levels \([1,l]\) and is defined as \(\sum_{i=1}^{l}BT^{i}\).+ * Use the @DiskCacheNone@ policy if the database's access pattern has does not have good spatial or temporal locality.+ For instance, if the access pattern is uniformly random.++ ==== Fine-tuning: Merge Batch Size #fine_tuning_merge_batch_size#++ The /merge batch size/ is a micro-tuning parameter, and in most cases you do+ need to think about it and can leave it at its default.++ When using the 'Incremental' merge schedule, merging is done in batches. This+ is a trade-off: larger batches tends to mean better overall performance but the+ downside is that while most updates (inserts, deletes, upserts) are fast, some+ are slower (when a batch of merging work has to be done).++ If you care most about the maximum latency of updates, then use a small batch+ size. If you don't care about latency of individual operations, just the+ latency of the overall sequence of operations then use a large batch size. The+ default is to use a large batch size, the same size as the write buffer itself.+ The minimum batch size is 1. The maximum batch size is the size of the write+ buffer 'confWriteBufferAlloc'.++ Note that the actual batch size is the minimum of this configuration+ parameter and the size of the batch of operations performed (e.g. 'inserts').+ So if you consistently use large batches, you can use a batch size of 1 and+ the merge batch size will always be determined by the operation batch size.++ A further reason why it may be preferable to use minimal batch sizes is to get+ good parallel work balance, when using parallelism.++ == References++ The implementation of LSM-trees in this package draws inspiration from:++ * Chris Okasaki.+ 1998.+ \"Purely Functional Data Structures\"+ [doi:10.1017/CBO9780511530104](https://doi.org/10.1017/CBO9780511530104)+ * Niv Dayan, Manos Athanassoulis, and Stratos Idreos.+ 2017.+ \"Monkey: Optimal Navigable Key-Value Store.\"+ [doi:10.1145/3035918.3064054](https://doi.org/10.1145/3035918.3064054)+ * Subhadeep Sarkar, Dimitris Staratzis, Ziehen Zhu, and Manos Athanassoulis.+ 2021.+ \"Constructing and analyzing the LSM compaction design space.\"+ [doi:10.14778/3476249.3476274](https://doi.org/10.14778/3476249.3476274)++license: Apache-2.0+license-files:+ LICENSE+ NOTICE++author:+ Duncan Coutts, Joris Dral, Matthias Heinzel, Wolfgang Jeltsch, Wen Kokke, and Alex Washburn++maintainer: oso@intersectmbo.org+copyright: (c) 2023-2025 Cardano Development Foundation+category: Database+build-type: Simple+tested-with: GHC ==9.2 || ==9.4 || ==9.6 || ==9.8 || ==9.10 || ==9.12+extra-doc-files: CHANGELOG.md+data-dir: test/golden-file-data/++source-repository head+ type: git+ location: https://github.com/IntersectMBO/lsm-tree+ subdir: lsm-tree++source-repository this+ type: git+ location: https://github.com/IntersectMBO/lsm-tree+ subdir: lsm-tree+ tag: lsm-tree-1.0.0.0++common warnings+ ghc-options:+ -Wall -Wcompat -Wincomplete-uni-patterns+ -Wincomplete-record-updates -Wpartial-fields -Widentities+ -Wredundant-constraints -Wmissing-export-lists+ -Wno-unticked-promoted-constructors -Wunused-packages++ ghc-options: -Werror=missing-deriving-strategies++common wno-x-partial+ if impl(ghc >=9.8)+ -- No errors for x-partial functions. We might remove this in the future if+ -- we decide to refactor code that uses partial functions.+ ghc-options: -Wno-x-partial++common language+ default-language: GHC2021+ default-extensions:+ DeriveAnyClass+ DerivingStrategies+ DerivingVia+ ExplicitNamespaces+ GADTs+ LambdaCase+ RecordWildCards+ RoleAnnotations+ ViewPatterns++library+ import: language, warnings, wno-x-partial+ hs-source-dirs: src+ exposed-modules:+ Database.LSMTree+ Database.LSMTree.Simple++ build-depends:+ , base >=4.16 && <4.22+ , blockio ^>=0.1+ , contra-tracer ^>=0.1 || ^>=0.2+ , deepseq ^>=1.4 || ^>=1.5+ , fs-api ^>=0.4+ , io-classes ^>=1.6 || ^>=1.7 || ^>=1.8.0.1+ , io-classes:strict-mvar+ , lsm-tree:control+ , lsm-tree:core+ , primitive ^>=0.9+ , random ^>=1.0 || ^>=1.1 || ^>=1.2 || ^>=1.3+ , text ^>=2.1.1+ , vector ^>=0.13++library core+ import: language, warnings, wno-x-partial+ visibility: private+ hs-source-dirs: src-core+ exposed-modules:+ Database.LSMTree.Internal.Arena+ Database.LSMTree.Internal.Assertions+ Database.LSMTree.Internal.BitMath+ Database.LSMTree.Internal.BlobFile+ Database.LSMTree.Internal.BlobRef+ Database.LSMTree.Internal.BloomFilter+ Database.LSMTree.Internal.ByteString+ Database.LSMTree.Internal.ChecksumHandle+ Database.LSMTree.Internal.Chunk+ Database.LSMTree.Internal.Config+ Database.LSMTree.Internal.Config.Override+ Database.LSMTree.Internal.CRC32C+ Database.LSMTree.Internal.Cursor+ Database.LSMTree.Internal.Entry+ Database.LSMTree.Internal.IncomingRun+ Database.LSMTree.Internal.Index+ Database.LSMTree.Internal.Index.Compact+ Database.LSMTree.Internal.Index.CompactAcc+ Database.LSMTree.Internal.Index.Ordinary+ Database.LSMTree.Internal.Index.OrdinaryAcc+ Database.LSMTree.Internal.Lookup+ Database.LSMTree.Internal.Map.Range+ Database.LSMTree.Internal.Merge+ Database.LSMTree.Internal.MergeSchedule+ Database.LSMTree.Internal.MergingRun+ Database.LSMTree.Internal.MergingTree+ Database.LSMTree.Internal.MergingTree.Lookup+ Database.LSMTree.Internal.Page+ Database.LSMTree.Internal.PageAcc+ Database.LSMTree.Internal.PageAcc1+ Database.LSMTree.Internal.Paths+ Database.LSMTree.Internal.Primitive+ Database.LSMTree.Internal.Range+ Database.LSMTree.Internal.RawBytes+ Database.LSMTree.Internal.RawOverflowPage+ Database.LSMTree.Internal.RawPage+ Database.LSMTree.Internal.Readers+ Database.LSMTree.Internal.Run+ Database.LSMTree.Internal.RunAcc+ Database.LSMTree.Internal.RunBuilder+ Database.LSMTree.Internal.RunNumber+ Database.LSMTree.Internal.RunReader+ Database.LSMTree.Internal.Serialise+ Database.LSMTree.Internal.Serialise.Class+ Database.LSMTree.Internal.Snapshot+ Database.LSMTree.Internal.Snapshot.Codec+ Database.LSMTree.Internal.Types+ Database.LSMTree.Internal.UniqCounter+ Database.LSMTree.Internal.Unsafe+ Database.LSMTree.Internal.Unsliced+ Database.LSMTree.Internal.Vector+ Database.LSMTree.Internal.Vector.Growing+ Database.LSMTree.Internal.WriteBuffer+ Database.LSMTree.Internal.WriteBufferBlobs+ Database.LSMTree.Internal.WriteBufferReader+ Database.LSMTree.Internal.WriteBufferWriter++ build-depends:+ , base >=4.16 && <4.22+ , bitvec ^>=1.1+ , blockio ^>=0.1+ , bloomfilter-blocked ^>=0.1+ , bytestring ^>=0.11.4.0 || ^>=0.12.1.0+ , cborg ^>=0.2.10.0+ , containers ^>=0.6 || ^>=0.7+ , contra-tracer ^>=0.1 || ^>=0.2+ , crc32c ^>=0.2.1+ , deepseq ^>=1.4 || ^>=1.5+ , filepath ^>=1.5+ , fs-api ^>=0.4+ , io-classes ^>=1.6 || ^>=1.7 || ^>=1.8.0.1+ , io-classes:strict-mvar+ , lsm-tree:control+ , lsm-tree:kmerge+ , primitive ^>=0.9+ , serialise ^>=0.2+ , text ^>=2.1.1+ , utf8-string ^>=1.0+ , vector ^>=0.13+ , vector-algorithms ^>=0.9++ if impl(ghc >=9.4)+ other-modules: Database.LSMTree.Internal.StrictArray+ build-depends: data-elevator ^>=0.1.0.2 || ^>=0.2+ cpp-options: -DHAVE_STRICT_ARRAY++library extras+ import: language, warnings+ visibility: private+ hs-source-dirs: src-extras+ exposed-modules:+ Database.LSMTree.Extras+ Database.LSMTree.Extras.Generators+ Database.LSMTree.Extras.Index+ Database.LSMTree.Extras.MergingRunData+ Database.LSMTree.Extras.MergingTreeData+ Database.LSMTree.Extras.NoThunks+ Database.LSMTree.Extras.Orphans+ Database.LSMTree.Extras.Random+ Database.LSMTree.Extras.ReferenceImpl+ Database.LSMTree.Extras.RunData+ Database.LSMTree.Extras.UTxO++ build-depends:+ , base >=4.16 && <4.22+ , bitvec+ , blockio+ , bytestring+ , containers+ , contra-tracer+ , deepseq+ , fs-api+ , fs-sim+ , io-classes:strict-mvar+ , io-classes:strict-stm+ , lsm-tree+ , lsm-tree:control+ , lsm-tree:core+ , lsm-tree:kmerge+ , lsm-tree:prototypes+ , nonempty-containers+ , nothunks+ , primitive+ , QuickCheck+ , quickcheck-instances+ , random+ , vector+ , wide-word++test-suite lsm-tree-test+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Main.hs+ other-modules:+ Database.LSMTree.Class+ Database.LSMTree.Class.Common+ Database.LSMTree.Model+ Database.LSMTree.Model.IO+ Database.LSMTree.Model.Session+ Database.LSMTree.Model.Table+ Paths_lsm_tree+ Test.Database.LSMTree+ Test.Database.LSMTree.Class+ Test.Database.LSMTree.Generators+ Test.Database.LSMTree.Internal+ Test.Database.LSMTree.Internal.Arena+ Test.Database.LSMTree.Internal.BlobFile.FS+ Test.Database.LSMTree.Internal.BloomFilter+ Test.Database.LSMTree.Internal.Chunk+ Test.Database.LSMTree.Internal.CRC32C+ Test.Database.LSMTree.Internal.Entry+ Test.Database.LSMTree.Internal.Index.Compact+ Test.Database.LSMTree.Internal.Index.Ordinary+ Test.Database.LSMTree.Internal.Lookup+ Test.Database.LSMTree.Internal.Merge+ Test.Database.LSMTree.Internal.MergingRun+ Test.Database.LSMTree.Internal.MergingTree+ Test.Database.LSMTree.Internal.PageAcc+ Test.Database.LSMTree.Internal.PageAcc1+ Test.Database.LSMTree.Internal.RawBytes+ Test.Database.LSMTree.Internal.RawOverflowPage+ Test.Database.LSMTree.Internal.RawPage+ Test.Database.LSMTree.Internal.Readers+ Test.Database.LSMTree.Internal.Run+ Test.Database.LSMTree.Internal.RunAcc+ Test.Database.LSMTree.Internal.RunBloomFilterAlloc+ Test.Database.LSMTree.Internal.RunBuilder+ Test.Database.LSMTree.Internal.RunReader+ Test.Database.LSMTree.Internal.Serialise+ Test.Database.LSMTree.Internal.Serialise.Class+ Test.Database.LSMTree.Internal.Snapshot.Codec+ Test.Database.LSMTree.Internal.Snapshot.Codec.Golden+ Test.Database.LSMTree.Internal.Snapshot.FS+ Test.Database.LSMTree.Internal.Unsliced+ Test.Database.LSMTree.Internal.Vector+ Test.Database.LSMTree.Internal.Vector.Growing+ Test.Database.LSMTree.Internal.WriteBufferBlobs.FS+ Test.Database.LSMTree.Internal.WriteBufferReader.FS+ Test.Database.LSMTree.Model.Table+ Test.Database.LSMTree.Resolve+ Test.Database.LSMTree.StateMachine+ Test.Database.LSMTree.StateMachine.DL+ Test.Database.LSMTree.StateMachine.Op+ Test.Database.LSMTree.Tracer.Golden+ Test.Database.LSMTree.UnitTests+ Test.FS+ Test.Util.Arbitrary+ Test.Util.FS+ Test.Util.FS.Error+ Test.Util.Orphans+ Test.Util.PrettyProxy+ Test.Util.QC+ Test.Util.QLS+ Test.Util.RawPage+ Test.Util.TypeFamilyWrappers++ autogen-modules: Paths_lsm_tree+ build-depends:+ , ansi-terminal+ , barbies+ , base <5+ , bitvec+ , blockio+ , blockio:sim+ , bloomfilter-blocked+ , bytestring+ , cborg+ , constraints+ , containers+ , contra-tracer+ , crc32c+ , cryptohash-sha256+ , deepseq+ , directory+ , filepath+ , fs-api+ , fs-sim+ , io-classes+ , io-classes:strict-mvar+ , io-classes:strict-stm+ , io-sim+ , lsm-tree+ , lsm-tree:control+ , lsm-tree:core+ , lsm-tree:extras+ , lsm-tree:prototypes+ , mtl+ , nothunks+ , primitive+ , QuickCheck+ , quickcheck-classes+ , quickcheck-dynamic+ , quickcheck-instances+ , quickcheck-lockstep >=0.8+ , random+ , safe-wild-cards+ , semialign+ , split+ , tasty+ , tasty-golden+ , tasty-hunit+ , tasty-quickcheck+ , temporary+ , text+ , these+ , transformers+ , vector+ , vector-algorithms+ , wide-word++ ghc-options: -threaded++benchmark lsm-tree-micro-bench+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: bench/micro+ main-is: Main.hs+ other-modules:+ Bench.Database.LSMTree+ Bench.Database.LSMTree.Internal.BloomFilter+ Bench.Database.LSMTree.Internal.Index+ Bench.Database.LSMTree.Internal.Index.Compact+ Bench.Database.LSMTree.Internal.Lookup+ Bench.Database.LSMTree.Internal.Merge+ Bench.Database.LSMTree.Internal.RawPage+ Bench.Database.LSMTree.Internal.Serialise+ Bench.Database.LSMTree.Internal.WriteBuffer++ build-depends:+ , base <5+ , blockio+ , bloomfilter-blocked+ , bytestring+ , containers+ , contra-tracer+ , criterion+ , deepseq+ , directory+ , fs-api+ , lsm-tree+ , lsm-tree:control+ , lsm-tree:core+ , lsm-tree:extras+ , QuickCheck+ , random+ , temporary+ , vector++ ghc-options: -rtsopts -with-rtsopts=-T -threaded++benchmark lsm-tree-bench-bloomfilter+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: bench/macro+ main-is: lsm-tree-bench-bloomfilter.hs+ build-depends:+ , base <5+ , bloomfilter-blocked+ , lsm-tree:core+ , lsm-tree:extras+ , random+ , time+ , vector+ , wide-word++ ghc-options: -rtsopts -with-rtsopts=-T -threaded++benchmark lsm-tree-bench-lookups+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: bench/macro+ main-is: lsm-tree-bench-lookups.hs+ build-depends:+ , base <5+ , blockio+ , bloomfilter-blocked+ , deepseq+ , fs-api+ , io-classes+ , lsm-tree:control+ , lsm-tree:core+ , lsm-tree:extras+ , primitive+ , random+ , time+ , vector+ , vector-algorithms++ ghc-options: -rtsopts -with-rtsopts=-T -threaded++library mcg+ import: language, warnings, wno-x-partial+ visibility: private+ hs-source-dirs: src-mcg+ exposed-modules: MCG+ build-depends:+ , base <5+ , primes++benchmark unions-bench+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: bench-unions+ main-is: Main.hs+ other-modules: Bench.Unions+ build-depends:+ , async+ , base+ , bytestring+ , clock+ , containers+ , directory+ , lsm-tree+ , lsm-tree:extras+ , mtl+ , optparse-applicative+ , primitive+ , random+ , vector++ ghc-options: -rtsopts -with-rtsopts=-T -threaded++flag measure-batch-latency+ description:+ Measure the latency for individual batches of updates and lookups++ default: False+ manual: True++common measure-batch-latency+ if flag(measure-batch-latency)+ cpp-options: -DMEASURE_BATCH_LATENCY++benchmark utxo-bench+ import: language, warnings, wno-x-partial, measure-batch-latency+ type: exitcode-stdio-1.0+ hs-source-dirs: bench/macro+ main-is: utxo-bench.hs+ build-depends:+ , async+ , base <5+ , blockio+ , bytestring+ , clock+ , containers+ , contra-tracer+ , deepseq+ , fs-api+ , lsm-tree+ , lsm-tree:extras+ , lsm-tree:mcg+ , optparse-applicative+ , pretty-show+ , primitive+ , random+ , transformers+ , vector++ ghc-options: -rtsopts -with-rtsopts=-T -threaded++flag rocksdb+ description: Build components that rely on RocksDB (only on Linux)+ default: True+ manual: False++benchmark utxo-rocksdb-bench+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: bench/macro+ main-is: utxo-rocksdb-bench.hs++ if !(os(linux) && flag(rocksdb))+ buildable: False++ build-depends:+ , base <5+ , binary+ , bytestring+ , clock+ , containers+ , cryptohash-sha256+ , deepseq+ , directory+ , lsm-tree:mcg+ , lsm-tree:rocksdb+ , optparse-applicative+ , split++ ghc-options: -rtsopts -with-rtsopts=-T -threaded++library rocksdb+ import: language, warnings+ visibility: private+ hs-source-dirs: src-rocksdb+ exposed-modules: RocksDB+ other-modules: RocksDB.FFI++ if !(os(linux) && flag(rocksdb))+ buildable: False++ -- Ubuntu 22.04 doesn't have pkgconfig files for rocksdb+ extra-libraries: rocksdb+ build-depends:+ , base <5+ , bytestring+ , indexed-traversable++library kmerge+ import: language, warnings, wno-x-partial+ visibility: private+ hs-source-dirs: src-kmerge+ exposed-modules:+ KMerge.Heap+ KMerge.LoserTree++ build-depends:+ , base <5+ , indexed-traversable ^>=0.1+ , primitive ^>=0.9++test-suite kmerge-test+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: kmerge-test.hs+ build-depends:+ , base >=4.16 && <4.22+ , deepseq+ , heaps+ , lsm-tree:kmerge+ , primitive+ , splitmix+ , tasty+ , tasty-bench+ , tasty-hunit+ , tasty-quickcheck+ , vector++benchmark kmerge-bench+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: kmerge-test.hs+ cpp-options: -DKMERGE_BENCHMARKS+ build-depends:+ , base >=4.16 && <4.22+ , deepseq+ , heaps+ , lsm-tree:kmerge+ , primitive+ , splitmix+ , tasty+ , tasty-bench+ , tasty-hunit+ , tasty-quickcheck+ , vector++test-suite map-range-test+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: map-range-test.hs+ build-depends:+ , base >=4.16 && <4.22+ , bytestring+ , containers+ , lsm-tree:core+ , QuickCheck+ , tasty+ , tasty-hunit+ , tasty-quickcheck++library prototypes+ import: language, warnings, wno-x-partial+ visibility: private+ hs-source-dirs: src-prototypes+ exposed-modules:+ FormatPage+ ScheduledMerges++ build-depends:+ , base <5+ , binary+ , bytestring+ , containers+ , contra-tracer+ , primitive+ , QuickCheck+ , transformers++test-suite prototypes-test+ import: language, warnings, wno-x-partial+ type: exitcode-stdio-1.0+ hs-source-dirs: test-prototypes+ main-is: Main.hs+ other-modules:+ Test.FormatPage+ Test.ScheduledMerges+ Test.ScheduledMerges.RunSizes+ Test.ScheduledMergesQLS++ build-depends:+ , base <5+ , bytestring+ , constraints+ , containers+ , contra-tracer+ , lsm-tree:prototypes+ , mtl+ , primitive+ , QuickCheck+ , quickcheck-dynamic+ , quickcheck-lockstep >=0.8+ , tasty+ , tasty-hunit+ , tasty-quickcheck++library control+ import: language, warnings+ visibility: private+ hs-source-dirs: src-control+ exposed-modules:+ Control.ActionRegistry+ Control.Concurrent.Class.MonadSTM.RWVar+ Control.RefCount++ build-depends:+ , base >=4.16 && <4.22+ , deepseq ^>=1.4 || ^>=1.5+ , io-classes ^>=1.6 || ^>=1.7 || ^>=1.8.0.1+ , io-classes:strict-stm+ , primitive ^>=0.9++test-suite control-test+ import: language, warnings+ type: exitcode-stdio-1.0+ hs-source-dirs: test-control+ main-is: Main.hs+ other-modules:+ Test.Control.ActionRegistry+ Test.Control.Concurrent.Class.MonadSTM.RWVar+ Test.Control.RefCount++ build-depends:+ , base <5+ , io-classes+ , io-sim+ , lsm-tree:control+ , primitive+ , QuickCheck+ , tasty+ , tasty-quickcheck++-- It's not really a test suite, but if we make it an executable then its+-- dependencies will be included for dependency resolution when building the+-- main library. As a test-suite, it's more accurately represented as an+-- internal component.+test-suite demo+ import: language, warnings+ type: exitcode-stdio-1.0+ hs-source-dirs: app+ main-is: Main.hs+ other-modules: Database.LSMTree.Demo+ build-depends:+ , base <5+ , blockio+ , blockio:sim+ , contra-tracer+ , directory+ , fs-api+ , fs-sim+ , io-classes+ , io-sim+ , lsm-tree+ , primitive+ , vector++ ghc-options: -threaded
@@ -0,0 +1,644 @@+{-# LANGUAGE CPP #-}++-- | Registry of monadic actions supporting rollback actions and delayed actions+-- in the presence of (a-)synchronous exceptions.+--+-- This module is heavily inspired by:+--+-- * [resource-registry](https://github.com/IntersectMBO/io-classes-extra/blob/main/resource-registry/src/Control/ResourceRegistry.hs)+--+-- * [resourcet](https://hackage.haskell.org/package/resourcet)+module Control.ActionRegistry (+ -- * Modify mutable state #modifyMutableState#+ -- $modify-mutable-state+ modifyWithActionRegistry+ , modifyWithActionRegistry_+ -- * Action registry #actionRegistry#+ -- $action-registry+ , ActionRegistry+ , ActionError+ , getActionError+ , mapActionError+ -- * Runners+ , withActionRegistry+ , unsafeNewActionRegistry+ , unsafeFinaliseActionRegistry+ , CommitActionRegistryError (..)+ , AbortActionRegistryError (..)+ , AbortActionRegistryReason (..)+ , getReasonExitCaseException+ , mapExceptionWithActionRegistry+ -- * Registering actions #registeringActions#+ -- $registering-actions+ , withRollback+ , withRollback_+ , withRollbackMaybe+ , withRollbackEither+ , withRollbackFun+ , delayedCommit+ ) where++import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive+import Data.Kind+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import Data.Maybe (fromMaybe)+import Data.Monoid (First (..))+import Data.Primitive.MutVar++#ifdef NO_IGNORE_ASSERTS+import GHC.Stack+#endif++-- TODO: add tests using fs-sim/io-sim to make sure exception safety is+-- guaranteed.++-- TODO: add assertions that allocated resources end up in the final state, and+-- that temporarily freed resources are removed from the final state.++-- TODO: could we statically disallow using a resource after it is freed using+-- @delayedCommit@, for example through data abstraction?++-- Call stack instrumentation is enabled if assertions are enabled.+#ifdef NO_IGNORE_ASSERTS+#define HasCallStackIfDebug HasCallStack+#else+#define HasCallStackIfDebug ()+#endif++{-------------------------------------------------------------------------------+ Printing utilities+-------------------------------------------------------------------------------}++tabLines1 :: String -> String+tabLines1 = tabLinesN 1++#ifdef NO_IGNORE_ASSERTS+tabLines2 :: String -> String+tabLines2 = tabLinesN 2+#endif++tabLinesN :: Int -> String -> String+tabLinesN n = unlines . fmap (ts++) . lines+ where+ ts = concat $ replicate n " "++{-------------------------------------------------------------------------------+ Modify mutable state+-------------------------------------------------------------------------------}++{- $modify-mutable-state++ When a piece of mutable state holding system resources is updated, then it is+ important to guarantee in the presence of (a-)synchronous exceptions that:++ 1. Allocated resources end up in the state+ 2. Freed resources are removed from the state++ Consider the example program below. We have some mutable @State@ that holds a+ file handle/descriptor. We want to mutate this state by closing the current+ handle, and replacing it by a newly opened handle. Using the tools at our+ disposal in "Control.ActionRegistry", we guarantee (1) and (2).++ @+ type State = MVar Handle++ example :: State -> IO ()+ example st =+ 'modifyWithActionRegistry_'+ (takeMVar st)+ (putMVar st)+ $ \\reg h -> do+ h' <- 'withRollback' reg+ (openFile "file.txt" ReadWriteMode)+ hClose+ 'delayedCommit' reg (hClose h)+ pure h'+ @++ What is also nice about this examples is that it is atomic: other threads will+ not be able to see the updated @State@ until 'modifyWithActionRegistry_' has+ exited and the necessary side effects have been performed. Of course, another+ thread *could* observe that the @file.txt@ was created before+ 'modifyWithActionRegistry_' has exited, but the assumption is that the threads+ in our program are cooperative. It is up to the user to ensure that actions+ that are performed as part of the state update do not conflict with other+ actions.+-}++{-# SPECIALISE modifyWithActionRegistry ::+ IO st+ -> (st -> IO ())+ -> (ActionRegistry IO -> st -> IO (st, a))+ -> IO a+ #-}+-- | Modify a piece piece of state given a fresh action registry.+modifyWithActionRegistry ::+ (PrimMonad m, MonadCatch m)+ => m st -- ^ Get the state+ -> (st -> m ()) -- ^ Store a state+ -> (ActionRegistry m -> st -> m (st, a)) -- ^ Modify the state+ -> m a+modifyWithActionRegistry getSt putSt action =+ snd . fst <$> generalBracket acquire release (uncurry action)+ where+ acquire = (,) <$> unsafeNewActionRegistry <*> getSt+ release (reg, oldSt) ec = do+ case ec of+ ExitCaseSuccess (newSt, _) -> putSt newSt+ ExitCaseException _ -> putSt oldSt+ ExitCaseAbort -> putSt oldSt+ unsafeFinaliseActionRegistry reg ec++{-# SPECIALISE modifyWithActionRegistry_ ::+ IO st+ -> (st -> IO ())+ -> (ActionRegistry IO -> st -> IO st)+ -> IO ()+ #-}+-- | Like 'modifyWithActionRegistry', but without a return value.+modifyWithActionRegistry_ ::+ (PrimMonad m, MonadCatch m)+ => m st -- ^ Get the state+ -> (st -> m ()) -- ^ Store a state+ -> (ActionRegistry m -> st -> m st)+ -> m ()+modifyWithActionRegistry_ getSt putSt action =+ modifyWithActionRegistry getSt putSt (\reg content -> (,()) <$> action reg content)++{-------------------------------------------------------------------------------+ Action registry+-------------------------------------------------------------------------------}++{- $action-registry++ An 'ActionRegistry' is a registry of monadic actions to support working with+ resources and mutable state in the presence of (a)synchronous exceptions. It+ works analogously to database transactions: within the \"transaction\" scope+ we can perform actions (such as resource allocations and state changes) and we+ can register delayed (commit) and rollback actions. The delayed actions are+ all executed at the end if the transaction scope is exited successfully, but+ if an exception is thrown (sync or async) then the rollback actions are+ executed instead, and the exception is propagated.++ * Rollback actions are executed in the reverse order in which they were+ registered, which is the natural nesting order when considered as bracketing.++ * Delayed actions are executed in the same order in which they are registered.+-}++-- | Registry of monadic actions supporting rollback actions and delayed actions+-- in the presence of (a-)synchronous exceptions.+--+-- See [Action registry](#g:actionRegistry) for more information.+--+-- An action registry should be short-lived, and it is not thread-safe.+data ActionRegistry m = ActionRegistry {+ -- | Registered rollback actions. Use 'consAction' when modifying this+ -- variable.+ --+ -- INVARIANT: actions are stored in LIFO order.+ --+ -- INVARIANT: the contents of this variable are in NF.+ registryRollback :: !(MutVar (PrimState m) [Action m])++ -- | Registered, delayed actions. Use 'consAction' when modifying this+ -- variable.+ --+ -- INVARIANT: actions are stored in LIFO order.+ --+ -- INVARIANT: the contents of this variable are in NF.+ , registryDelay :: !(MutVar (PrimState m) [Action m])+ }++{-# SPECIALISE consAction :: Action IO -> MutVar RealWorld [Action IO] -> IO () #-}+-- | Cons an action onto the contents of an actions variable.+--+-- Both the action and the resulting variable contents are evaluated to WHNF. If+-- the contents of the variable were already in NF, then the result will also be+-- in NF.+consAction :: PrimMonad m => Action m -> MutVar (PrimState m) [Action m] -> m ()+consAction !a var = modifyMutVar' var $ \as -> a `consStrict` as+ where consStrict !x xs = x : xs++-- | Monadic computations that (may) produce side effects+type Action :: (Type -> Type) -> Type++-- | An action failed with an exception+type ActionError :: Type++mkAction :: HasCallStackIfDebug => m () -> Action m+mkActionError :: SomeException -> Action m -> ActionError+getActionError :: ActionError -> SomeException+mapActionError :: (SomeException -> SomeException) -> ActionError -> ActionError++#ifdef NO_IGNORE_ASSERTS+data Action m = Action {+ runAction :: !(m ())+ , actionCallStack :: !CallStack+ }++data ActionError = ActionError SomeException CallStack+ deriving stock Show++instance Exception ActionError where+ displayException (ActionError err registerSite) = unlines [+ "A registered action threw an error: "+ , tabLines1 "The error:"+ , tabLines2 (displayException err)+ , tabLines1 "Registration site:"+ , tabLines2 (prettyCallStack registerSite)+ ]++mkAction a = Action a callStack++mkActionError e a = ActionError e (actionCallStack a)++getActionError (ActionError e _) = e++mapActionError f (ActionError e s) = ActionError (f e) s+#else+newtype Action m = Action {+ runAction :: m ()+ }++newtype ActionError = ActionError SomeException+ deriving stock Show+ deriving anyclass Exception++mkAction a = Action a++mkActionError e _ = ActionError e++getActionError (ActionError e) = e++mapActionError f (ActionError e) = ActionError (f e)+#endif++{-------------------------------------------------------------------------------+ Runners+-------------------------------------------------------------------------------}++{-# SPECIALISE withActionRegistry :: (ActionRegistry IO -> IO a) -> IO a #-}+-- | Run code with a new 'ActionRegistry'.+--+-- (A-)synchronous exception safety is only guaranteed within the scope of+-- 'withActionRegistry' (and only for properly registered actions). As soon as+-- we leave this scope, all bets are off. If, for example, a newly allocated+-- file handle escapes the scope, then that file handle can be leaked. If such+-- is the case, then it is highly likely that you should be using+-- 'modifyWithActionRegistry' instead.+--+-- If the code was interrupted due to an exception for example, then the+-- registry is aborted, which performs registered rollback actions. If the code+-- successfully terminated, then the registry is committed, in which case+-- registered, delayed actions will be performed.+--+-- Registered actions are run in LIFO order, whether they be rollback actions or+-- delayed actions.+withActionRegistry ::+ (PrimMonad m, MonadCatch m)+ => (ActionRegistry m -> m a)+ -> m a+withActionRegistry k = fst <$> generalBracket acquire release k+ where+ acquire = unsafeNewActionRegistry+ release reg ec = unsafeFinaliseActionRegistry reg ec++{-# SPECIALISE unsafeNewActionRegistry :: IO (ActionRegistry IO) #-}+-- | This function is considered unsafe. Preferably, use 'withActionRegistry'+-- instead.+--+-- If this function is used directly, use 'generalBracket' to pair+-- 'unsafeNewActionRegistry' with an 'unsafeFinaliseActionRegistry'.+unsafeNewActionRegistry :: PrimMonad m => m (ActionRegistry m)+unsafeNewActionRegistry = do+ registryRollback <- newMutVar $! []+ registryDelay <- newMutVar $! []+ pure $! ActionRegistry {..}++{-# SPECIALISE unsafeFinaliseActionRegistry :: ActionRegistry IO -> ExitCase a -> IO () #-}+-- | This function is considered unsafe. See 'unsafeNewActionRegistry'.+--+-- This commits the action registry on 'ExitCaseSuccess', and otherwise aborts+-- the action registry.+unsafeFinaliseActionRegistry ::+ (PrimMonad m, MonadCatch m)+ => ActionRegistry m+ -> ExitCase a+ -> m ()+unsafeFinaliseActionRegistry reg ec = case ec of+ ExitCaseSuccess{} -> unsafeCommitActionRegistry reg+ ExitCaseException e -> unsafeAbortActionRegistry reg (ReasonExitCaseException e)+ ExitCaseAbort -> unsafeAbortActionRegistry reg ReasonExitCaseAbort++{-# SPECIALISE unsafeCommitActionRegistry :: ActionRegistry IO -> IO () #-}+-- | Perform delayed actions, but not rollback actions.+unsafeCommitActionRegistry :: (PrimMonad m, MonadCatch m) => ActionRegistry m -> m ()+unsafeCommitActionRegistry reg = do+ as <- readMutVar (registryDelay reg)+ -- Run actions in FIFO order+ r <- runActions (reverse as)+ case NE.nonEmpty r of+ Nothing -> pure ()+ Just exceptions -> throwIO (CommitActionRegistryError exceptions)++data CommitActionRegistryError = CommitActionRegistryError (NonEmpty ActionError)+ deriving stock Show++instance Exception CommitActionRegistryError where+ displayException (CommitActionRegistryError es) = unlines $ [+ "Exceptions thrown while committing an action registry."+ ] <> NE.toList (fmap displayOne es)+ where+ displayOne e = tabLines1 (displayException e)++{-# SPECIALISE unsafeAbortActionRegistry ::+ ActionRegistry IO+ -> AbortActionRegistryReason+ -> IO () #-}+-- | Perform rollback actions, but not delayed actions+unsafeAbortActionRegistry ::+ (PrimMonad m, MonadCatch m)+ => ActionRegistry m+ -> AbortActionRegistryReason+ -> m ()+unsafeAbortActionRegistry reg reason = do+ as <- readMutVar (registryRollback reg)+ -- Run actions in LIFO order+ r <- runActions as+ case NE.nonEmpty r of+ Nothing -> pure ()+ Just exceptions -> throwIO (AbortActionRegistryError reason exceptions)++-- | Reasons why an action registry was aborted.+data AbortActionRegistryReason =+ -- | The action registry was aborted because the code that it scoped over+ -- threw an exception (see 'ExitCaseException').+ ReasonExitCaseException SomeException+ -- | The action registry was aborted because the code that it scoped over+ -- aborted (see 'ExitCaseAbort').+ | ReasonExitCaseAbort+ deriving stock Show++getReasonExitCaseException :: AbortActionRegistryReason -> Maybe SomeException+getReasonExitCaseException = \case+ ReasonExitCaseException e -> Just e+ ReasonExitCaseAbort -> Nothing++data AbortActionRegistryError =+ AbortActionRegistryError AbortActionRegistryReason (NonEmpty ActionError)+ deriving stock Show++instance Exception AbortActionRegistryError where+ displayException (AbortActionRegistryError reason es) = unlines $ [+ "Exceptions thrown while aborting an action registry."+ , ("Reason for aborting the registry: " ++ show reason)+ ] <> NE.toList (fmap displayOne es)+ where+ displayOne e = tabLines1 (displayException e)++{-# SPECIALISE runActions :: [Action IO] -> IO [ActionError] #-}+-- | Run all actions even if previous actions threw exceptions.+runActions :: MonadCatch m => [Action m] -> m [ActionError]+runActions = go []+ where+ go es [] = pure (reverse es)+ go es (a:as) = do+ eith <- try @_ @SomeException (runAction a)+ case eith of+ Left e -> go (mkActionError e a : es) as+ Right _ -> go es as++{-# SPECIALISE mapExceptionWithActionRegistry ::+ (Exception e1, Exception e2)+ => (e1 -> e2)+ -> IO a+ -> IO a #-}+-- | As 'Control.Exception.mapException', but aware of the structure of+-- 'AbortActionRegistryError' and 'CommitActionRegistryError'.+mapExceptionWithActionRegistry ::+ (Exception e1, Exception e2, MonadCatch m)+ => (e1 -> e2)+ -> m a+ -> m a+mapExceptionWithActionRegistry f action = action `catch` (throwIO . mapSomeException)+ where+ -- TODO: This erases the `ExceptionContext` of the underlying exception.+ -- Unfortunately, the API exposed by `io-classes` does not currently+ -- have the primitives to preserve the `ExceptionContext`.+ mapSomeException :: SomeException -> SomeException+ mapSomeException e =+ fromMaybe e . getFirst . mconcat . fmap First $+ [ toException . f <$> fromException e+ , toException . mapAbortActionRegistryError <$> fromException e+ , toException . mapCommitActionRegistryError <$> fromException e+ ]++ mapAbortActionRegistryError :: AbortActionRegistryError -> AbortActionRegistryError+ mapAbortActionRegistryError = \case+ AbortActionRegistryError reason es ->+ AbortActionRegistryError (mapAbortActionRegistryReason reason) (mapActionError mapSomeException <$> es)++ mapAbortActionRegistryReason :: AbortActionRegistryReason -> AbortActionRegistryReason+ mapAbortActionRegistryReason = \case+ ReasonExitCaseException e -> ReasonExitCaseException (mapSomeException e)+ ReasonExitCaseAbort -> ReasonExitCaseAbort++ mapCommitActionRegistryError :: CommitActionRegistryError -> CommitActionRegistryError+ mapCommitActionRegistryError = \case+ CommitActionRegistryError es ->+ CommitActionRegistryError (mapActionError mapSomeException <$> es)++{-------------------------------------------------------------------------------+ Registering actions+-------------------------------------------------------------------------------}++{- $registering-actions++ /Actions/ are monadic computations that (may) produce side effects. Such side+ effects can include opening or closing a file handle, but also modifying a+ mutable variable.++ We make a distinction between three types of actions:++ * An /immediate action/ is performed immediately, as the name suggests.++ * A /rollback action/ is an action that is registered in an action registry,+ and it is performed precisely when the corresponding action registry is+ aborted. See 'withRollback' for examples.++ * A /delayed action/ is an action that is registered in an action registry,+ and it is performed precisely when the corresponding action registry is+ committed. See 'delayedCommit' for examples.++ Immediate actions are run with asynchronous exceptions masked to guarantee+ that the rollback action is registered after the immediate action has returned+ successfully. This means that all the usual masking caveats apply for the+ immediate acion.++ Rollback actions and delayed actions are performed /precisely/ when aborting+ or committing an action registry respectively (see [Action+ registry](#g:actionRegistry)). To achieve this, finalisation of the action+ registry happens in the same masked state as running the registered actions.+ This means all the usual masking caveats apply for the registered actions.+-}++{-# SPECIALISE withRollback ::+ HasCallStackIfDebug+ => ActionRegistry IO+ -> IO a+ -> (a -> IO ())+ -> IO a #-}+-- | Perform an immediate action and register a rollback action.+--+-- See [Registering actions](#g:registeringActions) for more information about+-- the different types of actions.+--+-- A typical use case for 'withRollback' is to allocate a resource as the+-- immediate action, and to release said resource as the rollback action. In+-- that sense, 'withRollback' is similar to 'bracketOnError', but 'withRollback'+-- offers stronger guarantees.+--+-- Note that the following two expressions are /not/ equivalent. The former is+-- correct in the presence of asynchronous exceptions, while the latter is not!+--+-- @+-- withRollback reg acquire free+-- =/=+-- acquire >>= \x -> withRollback reg free (pure x)+-- @+withRollback ::+ (PrimMonad m, MonadMask m)+ => HasCallStackIfDebug+ => ActionRegistry m+ -> m a+ -> (a -> m ())+ -> m a+withRollback reg acquire release =+ withRollbackFun reg Just acquire release++{-# SPECIALISE withRollback_ ::+ HasCallStackIfDebug+ => ActionRegistry IO+ -> IO a+ -> IO ()+ -> IO a #-}+-- | Like 'withRollback', but the rollback action does not get access to the+-- result of the immediate action.+--+withRollback_ ::+ (PrimMonad m, MonadMask m)+ => HasCallStackIfDebug+ => ActionRegistry m+ -> m a+ -> m ()+ -> m a+withRollback_ reg acquire release =+ withRollbackFun reg Just acquire (\_ -> release)++{-# SPECIALISE withRollbackMaybe ::+ HasCallStackIfDebug+ => ActionRegistry IO+ -> IO (Maybe a)+ -> (a -> IO ())+ -> IO (Maybe a)+ #-}+-- | Like 'withRollback', but the immediate action may fail with a 'Nothing'.+-- The rollback action will only be registered if 'Just'.+--+withRollbackMaybe ::+ (PrimMonad m, MonadMask m)+ => HasCallStackIfDebug+ => ActionRegistry m+ -> m (Maybe a)+ -> (a -> m ())+ -> m (Maybe a)+withRollbackMaybe reg acquire release =+ withRollbackFun reg id acquire release++{-# SPECIALISE withRollbackEither ::+ HasCallStackIfDebug+ => ActionRegistry IO+ -> IO (Either e a)+ -> (a -> IO ())+ -> IO (Either e a)+ #-}+-- | Like 'withRollback', but the immediate action may fail with a 'Left'. The+-- rollback action will only be registered if 'Right'.+--+withRollbackEither ::+ (PrimMonad m, MonadMask m)+ => HasCallStackIfDebug+ => ActionRegistry m+ -> m (Either e a)+ -> (a -> m ())+ -> m (Either e a)+withRollbackEither reg acquire release =+ withRollbackFun reg fromEither acquire release+ where+ fromEither :: Either e a -> Maybe a+ fromEither (Left _) = Nothing+ fromEither (Right x) = Just x++{-# SPECIALISE withRollbackFun ::+ HasCallStackIfDebug+ => ActionRegistry IO+ -> (a -> Maybe b)+ -> IO a+ -> (b -> IO ())+ -> IO a+ #-}+-- | Like 'withRollback', but the immediate action may fail in some general+-- way. The rollback function will only be registered if the @(a -> Maybe b)@+-- function returned 'Just'.+--+-- 'withRollbackFun' is the most general form in the 'withRollback*' family of+-- functions. All 'withRollback*' functions can be defined in terms of+-- 'withRollBackFun'.+--+withRollbackFun ::+ (PrimMonad m, MonadMask m)+ => HasCallStackIfDebug+ => ActionRegistry m+ -> (a -> Maybe b)+ -> m a+ -> (b -> m ())+ -> m a+withRollbackFun reg extract acquire release = do+ mask_ $ do+ x <- acquire+ case extract x of+ Nothing -> pure x+ Just y -> do+ consAction (mkAction (release y)) (registryRollback reg)+ pure x++{-# SPECIALISE delayedCommit ::+ HasCallStackIfDebug+ => ActionRegistry IO+ -> IO ()+ -> IO () #-}+-- | Register a delayed action.+--+-- See [Registering actions](#g:registeringActions) for more information about+-- the different types of actions.+--+-- A typical use case for 'delayedCommit' is to delay destructive actions until+-- they are safe to be performed. For example, a destructive action such as+-- removing a file can often not be rolled back without jumping through+-- additional hoops.+--+-- If you can think of a sensible rollback action for the action you want to+-- delay then 'withRollback' might be a more suitable fit than 'delayedCommit'.+-- For example, incrementing a thread-safe mutable variable can easily be rolled+-- back by decrementing the same variable again.+--+delayedCommit ::+ PrimMonad m+ => HasCallStackIfDebug+ => ActionRegistry m+ -> m ()+ -> m ()+delayedCommit reg action = consAction (mkAction action) (registryDelay reg)
@@ -0,0 +1,179 @@+-- | A read-write-locked mutable variable with a bias towards write locks.+--+-- This module is intended to be imported qualified:+--+-- @+-- import Control.Concurrent.Class.MonadSTM.RWVar (RWVar)+-- import qualified Control.Concurrent.Class.MonadSTM.RWVar as RW+-- @+module Control.Concurrent.Class.MonadSTM.RWVar (+ RWVar (..)+ , RWState (..)+ , new+ , unsafeAcquireReadAccess+ , unsafeReleaseReadAccess+ , withReadAccess+ , unsafeAcquireWriteAccess+ , unsafeReleaseWriteAccess+ , withWriteAccess+ , withWriteAccess_+ ) where++import Control.Concurrent.Class.MonadSTM.Strict+import Control.DeepSeq+import Control.Monad.Class.MonadThrow+import Data.Word++-- | A read-write-locked mutable variable with a bias towards write-locks.+newtype RWVar m a = RWVar (StrictTVar m (RWState a))++-- | __NOTE__: Only strict in the reference and not the referenced value.+instance NFData (RWVar m a) where+ rnf = rwhnf++data RWState a =+ -- | @n@ concurrent readers and no writer.+ Reading !Word64 !a+ -- | @n@ concurrent readers and no writer, but no new readers can get+ -- access.+ | WaitingToWrite !Word64 !a+ -- | A single writer and no concurrent readers.+ | Writing++{-# SPECIALISE new :: a -> IO (RWVar IO a) #-}+new :: MonadSTM m => a -> m (RWVar m a)+new !x = RWVar <$> newTVarIO (Reading 0 x)++{-# SPECIALISE unsafeAcquireReadAccess :: RWVar IO a -> STM IO a #-}+unsafeAcquireReadAccess :: MonadSTM m => RWVar m a -> STM m a+unsafeAcquireReadAccess (RWVar !var) = do+ readTVar var >>= \case+ Reading n x -> do+ writeTVar var (Reading (n+1) x)+ pure x+ WaitingToWrite{} -> retry+ Writing -> retry++{-# SPECIALISE unsafeReleaseReadAccess :: RWVar IO a -> STM IO () #-}+unsafeReleaseReadAccess :: MonadSTM m => RWVar m a -> STM m ()+unsafeReleaseReadAccess (RWVar !var) = do+ readTVar var >>= \case+ Reading n x+ | n == 0 -> error "releasing a reader without read access (Reading)"+ | otherwise -> writeTVar var (Reading (n - 1) x)+ WaitingToWrite n x+ | n == 0 -> error "releasing a reader without read access (WaitingToWrite)"+ | otherwise -> writeTVar var (WaitingToWrite (n - 1) x)+ Writing -> error "releasing a reader without read access (Writing)"++{-# SPECIALISE withReadAccess :: RWVar IO a -> (a -> IO b) -> IO b #-}+withReadAccess :: (MonadSTM m, MonadThrow m) => RWVar m a -> (a -> m b) -> m b+withReadAccess rwvar k =+ bracket+ (atomically $ unsafeAcquireReadAccess rwvar)+ (\_ -> atomically $ unsafeReleaseReadAccess rwvar)+ k++{-# SPECIALISE unsafeAcquireWriteAccess :: RWVar IO a -> IO a #-}+-- | Acquire write access. This function assumes that it runs in a masked+-- context, and that is properly paired with an 'unsafeReleaseWriteAccess'!+--+-- If multiple threads try to acquire write access concurrently, then they will+-- race for access. However, if a thread has set RWState to WaitingToWrite, then+-- it is guaranteed that the same thread will acquire write access when all+-- readers have finished. That is, other writes can not "jump the queue". When+-- the writer finishes, then all other waiting threads will race for write+-- access again.+--+-- TODO: unsafeReleaseWriteAccess will set RWState to Reading 0. In case we have+-- readers *and* writers waiting for a writer to finish, once the writer is+-- finished there will be a race. In this race, readers and writers are just as+-- likely to acquire access first. However, if we wanted to make RWVar even more+-- biased towards writers, then we could ensure that all waiting writers get+-- access before the readers get a chance. This would probably require us to+-- change RWState to represent the case where writers are waiting for a writer+-- to finish.+unsafeAcquireWriteAccess :: (MonadSTM m, MonadCatch m) => RWVar m a -> m a+unsafeAcquireWriteAccess (RWVar !var) =+ -- trySetWriting is interruptible, but it is fine if it is interrupted+ -- because the RWState can not be changed before the interruption.+ --+ -- trySetWriting might update the RWState. There are interruptible+ -- operations in the body of the bracketOnError (in waitToWrite), so async+ -- exceptions can be delivered there. If an async exception happens because+ -- of an interrupt, we undo the RWState change using undoWaitingToWrite.+ --+ -- Note that if waitToWrite is interrupted, that it is impossible for the+ -- RWState to have changed from WaitingToWrite to either Reading or Writing.+ -- Therefore, undoWaitingToWrite can assume that it will find WaitingToWrite+ -- in the lock.+ bracketOnError trySetWriting undoWaitingToWrite $+ -- When Nothing is returned, it means that we set the RWState to+ -- WaitingToWrite, and so we wait to acquire the final write access.+ --+ -- When Just is returned, we already have write access.+ maybe waitToWrite pure+ where+ -- Try to acquire a write lock immediately, or otherwise set the internal+ -- state to WaitingToWrite as soon as possible.+ --+ -- Note: this is interruptible+ trySetWriting = atomically $ readTVar var >>= \case+ Reading n x+ | n == 0 -> do+ writeTVar var Writing+ pure (Just x)+ | otherwise -> do+ writeTVar var (WaitingToWrite n x)+ pure Nothing+ -- The following two branches are interruptible+ WaitingToWrite _n _x -> retry+ Writing -> retry++ -- Note: this is uninterruptible+ undoWaitingToWrite Nothing = atomically $ readTVar var >>= \case+ Reading _n _x -> error "undoWaitingToWrite: found Reading but expected WaitingToWrite"+ WaitingToWrite n x -> writeTVar var (Reading n x)+ Writing -> error "undoWaitingToWrite: found Writing but expected WaitingToWrite"+ undoWaitingToWrite (Just _) = error "undoWaitingToWrite: found Just but expected Nothing"++ -- Wait for the number of readers to go to 0, and then finally acquire write+ -- access.+ --+ -- Note: this is interruptible+ waitToWrite = atomically $ readTVar var >>= \case+ Reading _n _x -> error "waitToWrite: found Reading but expected WaitingToWrite"+ WaitingToWrite n x+ | n == 0 -> do+ writeTVar var Writing+ pure x+ -- This branch is interruptible+ | otherwise -> retry+ Writing -> error "waitToWrite: found Reading but expected Writing"++{-# SPECIALISE unsafeReleaseWriteAccess :: RWVar IO a -> a -> STM IO () #-}+-- | Release write access. This function assumes that it runs in a masked+-- context, and that is properly paired with an 'unsafeAcquireWriteAccess'!+unsafeReleaseWriteAccess :: MonadSTM m => RWVar m a -> a -> STM m ()+unsafeReleaseWriteAccess (RWVar !var) !x = do+ readTVar var >>= \case+ Reading _ _ -> error "releasing a writer without write access (Reading)"+ WaitingToWrite _ _ -> error "releasing a writer without write access (WaitingToWrite)"+ Writing -> writeTVar var (Reading 0 x)++{-# SPECIALISE withWriteAccess :: RWVar IO a -> (a -> IO (a, b)) -> IO b #-}+withWriteAccess :: (MonadSTM m, MonadCatch m) => RWVar m a -> (a -> m (a, b)) -> m b+withWriteAccess rwvar k = snd . fst <$>+ generalBracket+ (unsafeAcquireWriteAccess rwvar)+ (\x ec -> do+ atomically $ case ec of+ ExitCaseSuccess (x', _) -> unsafeReleaseWriteAccess rwvar x'+ ExitCaseException _ -> unsafeReleaseWriteAccess rwvar x+ ExitCaseAbort -> unsafeReleaseWriteAccess rwvar x+ )+ k++{-# SPECIALISE withWriteAccess_ :: RWVar IO a -> (a -> IO a) -> IO () #-}+withWriteAccess_ :: (MonadSTM m, MonadCatch m) => RWVar m a -> (a -> m a) -> m ()+withWriteAccess_ rwvar k = withWriteAccess rwvar (fmap (,()) . k)
@@ -0,0 +1,621 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE FunctionalDependencies #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE TypeFamilies #-}++module Control.RefCount (+ -- * Using references+ Ref(DeRef)+ , releaseRef+ , withRef+ , dupRef+ , RefException (..)+ -- ** Weak references+ , WeakRef (..)+ , mkWeakRef+ , mkWeakRefFromRaw+ , deRefWeak+ -- * Implementing objects with finalisers+ , RefCounted (..)+ , newRef+ -- ** Low level reference counts+ , RefCounter (RefCounter)+ , newRefCounter+ , incrementRefCounter+ , decrementRefCounter+ , tryIncrementRefCounter++ -- * Test API+ , checkForgottenRefs+ , ignoreForgottenRefs+ , enableForgottenRefChecks+ , disableForgottenRefChecks+ ) where++import Control.DeepSeq+import Control.Exception (assert)+import Control.Monad (void, when)+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive+import Data.Primitive.PrimVar+import GHC.Show (appPrec)+import GHC.Stack (CallStack, prettyCallStack)++#ifdef NO_IGNORE_ASSERTS+import Control.Concurrent (yield)+import Data.IORef+import GHC.Stack (HasCallStack, callStack)+import System.IO.Unsafe (unsafeDupablePerformIO, unsafePerformIO)+import System.Mem.Weak hiding (deRefWeak)+#if MIN_VERSION_base(4,20,0)+import System.Mem (performBlockingMajorGC)+#else+import System.Mem (performMajorGC)+#endif+#endif+++-------------------------------------------------------------------------------+-- Low level RefCounter API+--++-- | A reference counter with an optional finaliser action. Once the reference+-- count reaches @0@, the finaliser will be run.+data RefCounter m = RefCounter {+ countVar :: !(PrimVar (PrimState m) Int)+ , finaliser :: !(m ())+ }++instance Show (RefCounter m) where+ show _ = "<RefCounter>"++-- | NOTE: Only strict in the variable and not the referenced value.+instance NFData (RefCounter m) where+ rnf RefCounter{countVar, finaliser} =+ rwhnf countVar `seq` rwhnf finaliser++{-# SPECIALISE newRefCounter :: IO () -> IO (RefCounter IO) #-}+-- | Make a reference counter with initial value @1@.+--+-- The given finaliser is run when the reference counter reaches @0@. The+-- finaliser is run with async exceptions masked.+--+newRefCounter :: PrimMonad m => m () -> m (RefCounter m)+newRefCounter finaliser = do+ countVar <- newPrimVar 1+ pure $! RefCounter { countVar, finaliser }++{-# SPECIALISE incrementRefCounter :: RefCounter IO -> IO () #-}+-- | Increase the reference counter by one.+--+-- The count must be known (from context) to be non-zero already. Typically+-- this will be because the caller has a reference already and is handing out+-- another reference to some other code.+incrementRefCounter :: PrimMonad m => RefCounter m -> m ()+incrementRefCounter RefCounter{countVar} = do+ prevCount <- fetchAddInt countVar 1+ assert (prevCount > 0) $ pure ()++{-# SPECIALISE decrementRefCounter :: RefCounter IO -> IO () #-}+-- | Decrease the reference counter by one.+--+-- The count must be known (from context) to be non-zero. Typically this will+-- be because the caller has a reference already (that they took out themselves+-- or were given).+decrementRefCounter :: (PrimMonad m, MonadMask m) => RefCounter m -> m ()+decrementRefCounter RefCounter{countVar, finaliser} =+ --TODO: remove mask and require all uses to run with exceptions mask.+ mask_ $ do+ prevCount <- fetchSubInt countVar 1+ assert (prevCount > 0) $ pure ()+ when (prevCount == 1) finaliser++{-# SPECIALISE tryIncrementRefCounter :: RefCounter IO -> IO Bool #-}+-- | Try to turn a \"weak\" reference on something into a proper reference.+-- This is by analogy with @deRefWeak :: Weak v -> IO (Maybe v)@, but for+-- reference counts.+--+-- This amounts to trying to increase the reference count, but if it is already+-- zero then this will fail. And unlike with 'addReference' where such failure+-- would be a programmer error, this corresponds to the case when the thing the+-- reference count is tracking has been closed already.+--+-- The result is @True@ when a strong reference has been obtained and @False@+-- when upgrading fails.+--+tryIncrementRefCounter :: PrimMonad m => RefCounter m -> m Bool+tryIncrementRefCounter RefCounter{countVar} = do+ prevCount <- atomicReadInt countVar+ casLoop prevCount+ where+ -- A classic lock-free CAS loop.+ -- Check the value before is non-zero, return failure or continue.+ -- Atomically write the new (incremented) value if the old value is+ -- unchanged, and return the old value (either way).+ -- If no other thread changed the old value, we succeed.+ -- Otherwise we go round the loop again.+ casLoop prevCount+ | prevCount <= 0 = pure False+ | otherwise = do+ prevCount' <- casInt countVar prevCount (prevCount+1)+ if prevCount' == prevCount+ then pure True+ else casLoop prevCount'+++-------------------------------------------------------------------------------+-- Ref API+--++-- | A reference to an object of type @a@. Use references to support prompt+-- finalisation of object resources.+--+-- Rules of use:+--+-- * Each 'Ref' must eventually be released /exactly/ once with 'releaseRef'.+-- * Use 'withRef', or 'DeRef' to (temporarily) obtain the underlying+-- object.+-- * After calling 'releaseRef', the operations 'withRef' and pattern 'DeRef'+-- must /not/ be used.+-- * After calling 'releaseRef', an object obtained previously from+-- 'DeRef' must /not/ be used. For this reason, it is advisable to use+-- 'withRef' where possible, and be careful with use of 'DeRef'.+-- * A 'Ref' may be duplicated using 'dupRef' to produce an independent+-- reference (which must itself be released with 'releaseRef').+--+-- All of these operations are thread safe. They are not async-exception safe+-- however: the operations that allocate or deallocate must be called with+-- async exceptions masked. This includes 'newRef', 'dupRef' and 'releaseRef'.+--+-- Provided that all these rules are followed, this guarantees that the+-- object's finaliser will be run exactly once, promptly, when the final+-- reference is released.+--+-- In debug mode (when using CPP define @NO_IGNORE_ASSERTS@), adherence to+-- these rules are checked dynamically. These dynamic checks are however not+-- thread safe, so it is not guaranteed that all violations are always detected.+--+#ifndef NO_IGNORE_ASSERTS+newtype Ref obj = Ref { refobj :: obj }+#else+data Ref obj = Ref { refobj :: !obj, reftracker :: !RefTracker }+#endif++instance Show obj => Show (Ref obj) where+ showsPrec d Ref{refobj} =+ showParen (d > 10) $+ showString "Ref " . showsPrec 11 refobj++instance NFData obj => NFData (Ref obj) where+ rnf Ref{refobj} = rnf refobj++-- | Class of objects which support 'Ref'.+--+-- For objects in this class the guarantee is that (when the 'Ref' rules are+-- followed) the object's finaliser is called exactly once.+--+class RefCounted m obj | obj -> m where+ getRefCounter :: obj -> RefCounter m++#ifdef NO_IGNORE_ASSERTS+#define HasCallStackIfDebug HasCallStack+#else+#define HasCallStackIfDebug ()+#endif++{-# SPECIALISE+ newRef ::+ RefCounted IO obj+ => HasCallStackIfDebug+ => IO ()+ -> (RefCounter IO -> obj)+ -> IO (Ref obj)+ #-}+-- | Make a new reference.+--+-- The given finaliser is run when the last reference is released. The+-- finaliser is run with async exceptions masked.+--+newRef ::+ (RefCounted m obj, PrimMonad m)+ => HasCallStackIfDebug+ => m ()+ -> (RefCounter m -> obj)+ -> m (Ref obj)+newRef finaliser mkObject = do+ rc <- newRefCounter finaliser+ let !obj = mkObject rc+ assert (countVar (getRefCounter obj) == countVar rc) $+ newRefWithTracker obj++-- | Release a reference to an object that will no longer be used (via this+-- reference).+--+{-# SPECIALISE+ releaseRef ::+ RefCounted IO obj+ => HasCallStackIfDebug+ => Ref obj+ -> IO ()+ #-}+releaseRef ::+ (RefCounted m obj, PrimMonad m, MonadMask m)+ => HasCallStackIfDebug+ => Ref obj+ -> m ()+releaseRef ref@Ref{refobj} = do+ assertNoDoubleRelease ref+ assertNoForgottenRefs+ releaseRefTracker ref+ decrementRefCounter (getRefCounter refobj)++{-# COMPLETE DeRef #-}+{-# INLINE DeRef #-}+-- | Get the object in a 'Ref'. Be careful with retaining the object for too+-- long, since the object must not be used after 'releaseRef' is called.+--+pattern DeRef :: HasCallStackIfDebug => obj -> Ref obj+#ifndef NO_IGNORE_ASSERTS+pattern DeRef obj <- Ref obj+#else+pattern DeRef obj <- (deRef -> !obj) -- So we get assertion checking++deRef :: HasCallStack => Ref obj -> obj+deRef ref@Ref{refobj} =+ unsafeDupablePerformIO (assertNoUseAfterRelease ref)+ `seq` refobj+#endif++{-# SPECIALISE+ withRef ::+ HasCallStackIfDebug+ => Ref obj+ -> (obj -> IO a)+ -> IO a+ #-}+{-# INLINE withRef #-}+-- | Use the object in a 'Ref'. Do not retain the object after the scope of+-- the body. If you cannot use scoped \"with\" style, use pattern 'DeRef'.+--+withRef ::+ forall m obj a.+ (PrimMonad m, MonadThrow m)+ => HasCallStackIfDebug+ => Ref obj+ -> (obj -> m a)+ -> m a+withRef ref@Ref{refobj} f = do+ assertNoUseAfterRelease ref+ assertNoForgottenRefs+ f refobj+#ifndef NO_IGNORE_ASSERTS+ where+ _unused = throwIO @m @SomeException+#endif++{-# SPECIALISE+ dupRef ::+ RefCounted IO obj+ => HasCallStackIfDebug+ => Ref obj+ -> IO (Ref obj)+ #-}+-- | Duplicate an existing reference, to produce a new reference.+--+dupRef ::+ forall m obj. (RefCounted m obj, PrimMonad m, MonadThrow m)+ => HasCallStackIfDebug+ => Ref obj+ -> m (Ref obj)+dupRef ref@Ref{refobj} = do+ assertNoUseAfterRelease ref+ assertNoForgottenRefs+ incrementRefCounter (getRefCounter refobj)+ newRefWithTracker refobj+#ifndef NO_IGNORE_ASSERTS+ where+ _unused = throwIO @m @SomeException+#endif++-- | A \"weak\" reference to an object: that is, a reference that does not+-- guarantee to keep the object alive. If however the object is still alive+-- (due to other normal references still existing) then it can be converted+-- back into a normal reference with 'deRefWeak'.+--+-- Weak references do not themselves need to be released.+--+newtype WeakRef a = WeakRef a+ deriving stock Show++-- | Given an existing normal reference, create a new weak reference.+--+mkWeakRef :: Ref obj -> WeakRef obj+mkWeakRef Ref {refobj} = WeakRef refobj++-- | Given an existing raw reference, create a new weak reference.+--+mkWeakRefFromRaw :: obj -> WeakRef obj+mkWeakRefFromRaw obj = WeakRef obj++{-# SPECIALISE+ deRefWeak ::+ RefCounted IO obj+ => HasCallStackIfDebug+ => WeakRef obj+ -> IO (Maybe (Ref obj))+ #-}+-- | If the object is still alive, obtain a /new/ normal reference. The normal+-- rules for 'Ref' apply, including the need to eventually call 'releaseRef'.+--+deRefWeak ::+ (RefCounted m obj, PrimMonad m)+ => HasCallStackIfDebug+ => WeakRef obj+ -> m (Maybe (Ref obj))+deRefWeak (WeakRef obj) = do+ success <- tryIncrementRefCounter (getRefCounter obj)+ if success then Just <$> newRefWithTracker obj+ else pure Nothing++{-# INLINE newRefWithTracker #-}+#ifndef NO_IGNORE_ASSERTS+newRefWithTracker :: PrimMonad m => obj -> m (Ref obj)+newRefWithTracker obj =+ pure $! Ref obj+#else+newRefWithTracker :: (PrimMonad m, HasCallStack) => obj -> m (Ref obj)+newRefWithTracker obj = do+ reftracker' <- newRefTracker callStack+ pure $! Ref obj reftracker'+#endif++data RefException =+ RefUseAfterRelease RefId+ CallStack -- ^ Allocation site+ CallStack -- ^ Release site+ CallStack -- ^ Use site+ | RefDoubleRelease RefId+ CallStack -- ^ Allocation site+ CallStack -- ^ First release site+ CallStack -- ^ Second release site+ | RefNeverReleased RefId+ CallStack -- ^ Allocation site++newtype RefId = RefId Int+ deriving stock (Show, Eq, Ord)++instance Show RefException where+ --Sigh. QuickCheck still uses 'show' rather than 'displayException'.+ showsPrec d x = showParen (d > appPrec) $ showString (displayException x)++instance Exception RefException where+ displayException (RefUseAfterRelease refid allocSite releaseSite useSite) =+ "Reference is used after release: " ++ show refid+ ++ "\nAllocation site: " ++ prettyCallStack allocSite+ ++ "\nRelease site: " ++ prettyCallStack releaseSite+ ++ "\nUse site: " ++ prettyCallStack useSite+ displayException (RefDoubleRelease refid allocSite releaseSite1 releaseSite2) =+ "Reference is released twice: " ++ show refid+ ++ "\nAllocation site: " ++ prettyCallStack allocSite+ ++ "\nFirst release site: " ++ prettyCallStack releaseSite1+ ++ "\nSecond release site: " ++ prettyCallStack releaseSite2+ displayException (RefNeverReleased refid allocSite) =+ "Reference is never released: " ++ show refid+ ++ "\nAllocation site: " ++ prettyCallStack allocSite++#ifndef NO_IGNORE_ASSERTS++{-# INLINE releaseRefTracker #-}+releaseRefTracker :: PrimMonad m => Ref a -> m ()+releaseRefTracker _ = pure ()++{-# INLINE assertNoForgottenRefs #-}+assertNoForgottenRefs :: PrimMonad m => m ()+assertNoForgottenRefs = pure ()++{-# INLINE assertNoUseAfterRelease #-}+assertNoUseAfterRelease :: PrimMonad m => Ref a -> m ()+assertNoUseAfterRelease _ = pure ()++{-# INLINE assertNoDoubleRelease #-}+assertNoDoubleRelease :: PrimMonad m => Ref a -> m ()+assertNoDoubleRelease _ = pure ()++#else++-- | A weak pointer to an outer IORef, containing an inner IORef with a maybe to+-- indicate if the ref has been explicitly released.+--+-- The finaliser for the outer weak pointer is given access to the inner IORef+-- so that it can tell if the reference has become garbage without being+-- explicitly released.+--+-- The outer IORef is also stored directly. This ensures the weak pointer to+-- the same is not garbage collected until the RefTracker itself (and thus the+-- parent Ref) is itself garbage collected.+--+-- The inner IORef is mutated when explicitly released. The outer IORef is+-- never modified, but we use an IORef to ensure the weak pointer is reliable.+--+-- The inner IORef also tracks the call stack for the site where the reference+-- (tracker) is released. This call stack is used in exceptions for easier+-- debugging.+data RefTracker = RefTracker !RefId+ !(Weak (IORef (IORef (Maybe CallStack))))+ !(IORef (IORef (Maybe CallStack))) -- ^ Release site+ !CallStack -- ^ Allocation site++{-# NOINLINE globalRefIdSupply #-}+globalRefIdSupply :: PrimVar RealWorld Int+globalRefIdSupply = unsafePerformIO $ newPrimVar 0++data Enabled a = Enabled !a | Disabled++{-# NOINLINE globalForgottenRef #-}+globalForgottenRef :: IORef (Enabled (Maybe (RefId, CallStack)))+globalForgottenRef = unsafePerformIO $ newIORef (Enabled Nothing)++-- | This version of 'unsafeIOToPrim' is strict in the result of the argument+-- action.+--+-- Without strictness it seems that some IO side effects are not happening at+-- the right time, like clearing the @globalForgottenRef@ in+-- @assertNoForgottenRefs@.+unsafeIOToPrimStrict :: PrimMonad m => IO a -> m a+unsafeIOToPrimStrict k = do+ !x <- unsafeIOToPrim k+ pure x++newRefTracker :: PrimMonad m => CallStack -> m RefTracker+newRefTracker allocSite = unsafeIOToPrimStrict $ do+ inner <- newIORef Nothing+ outer <- newIORef inner+ refid <- fetchAddInt globalRefIdSupply 1+ weak <- mkWeakIORef outer $+ finaliserRefTracker inner (RefId refid) allocSite+ pure (RefTracker (RefId refid) weak outer allocSite)++releaseRefTracker :: (HasCallStack, PrimMonad m) => Ref a -> m ()+releaseRefTracker Ref { reftracker = RefTracker _refid _weak outer _ } =+ unsafeIOToPrimStrict $ do+ inner <- readIORef outer+ let releaseSite = callStack+ writeIORef inner (Just releaseSite)++finaliserRefTracker :: IORef (Maybe CallStack) -> RefId -> CallStack -> IO ()+finaliserRefTracker inner refid allocSite = do+ released <- readIORef inner+ case released of+ Just _releaseSite -> pure ()+ Nothing -> do+ -- Uh oh! Forgot a reference without releasing!+ -- Add it to a global var which we can poll elsewhere.+ mref <- readIORef globalForgottenRef+ case mref of+ Disabled -> pure ()+ -- Just keep one, but keep the last allocated one.+ -- The reason for last is that when there are nested structures with+ -- refs then the last allocated is likely to be the outermost, which+ -- is the best place to start hunting for ref leaks. Otherwise one can+ -- go on a wild goose chase tracking down inner refs that were only+ -- forgotten due to an outer ref being forgotten.+ Enabled (Just (refid', _)) | refid < refid' -> pure ()+ Enabled _ -> writeIORef globalForgottenRef (Enabled (Just (refid, allocSite)))++assertNoForgottenRefs :: (PrimMonad m, MonadThrow m) => m ()+assertNoForgottenRefs = do+ mrefs <- unsafeIOToPrimStrict $ readIORef globalForgottenRef+ case mrefs of+ Disabled -> pure ()+ Enabled Nothing -> pure ()+ Enabled (Just (refid, allocSite)) -> do+ -- Clear the var so we don't assert again.+ --+ -- Using the strict version is important here: if @m ~ IOSim s@, then+ -- using the non-strict version will lead to @RefNeverReleased@+ -- exceptions.+ unsafeIOToPrimStrict $ writeIORef globalForgottenRef (Enabled Nothing)+ throwIO (RefNeverReleased refid allocSite)+++assertNoUseAfterRelease :: (PrimMonad m, MonadThrow m, HasCallStack) => Ref a -> m ()+assertNoUseAfterRelease Ref { reftracker = RefTracker refid _weak outer allocSite } = do+ released <- unsafeIOToPrimStrict (readIORef =<< readIORef outer)+ case released of+ Nothing -> pure ()+ Just releaseSite -> do+ -- The site where the reference is used after release+ let useSite = callStack+ throwIO (RefUseAfterRelease refid allocSite releaseSite useSite)+#if !(MIN_VERSION_base(4,20,0))+ where+ _unused = callStack+#endif++assertNoDoubleRelease :: (PrimMonad m, MonadThrow m, HasCallStack) => Ref a -> m ()+assertNoDoubleRelease Ref { reftracker = RefTracker refid _weak outer allocSite } = do+ released <- unsafeIOToPrimStrict (readIORef =<< readIORef outer)+ case released of+ Nothing -> pure ()+ Just releaseSite1 -> do+ -- The second release site+ let releaseSite2 = callStack+ throwIO (RefDoubleRelease refid allocSite releaseSite1 releaseSite2)+#if !(MIN_VERSION_base(4,20,0))+ where+ _unused = callStack+#endif++#endif++-- | Run a GC to try and see if any refs have been forgotten without being+-- released. If so, this will throw a synchronous exception.+--+-- Note however that this is not the only place where 'RefNeverReleased'+-- exceptions can be thrown. All Ref operations poll for forgotten refs.+--+checkForgottenRefs :: forall m. (PrimMonad m, MonadThrow m) => m ()+checkForgottenRefs = do+#ifndef NO_IGNORE_ASSERTS+ pure ()+#else+ -- The hope is that by combining `performMajorGC` with `yield` that the+ -- former starts the finalizer threads for all dropped weak references and+ -- the latter suspends the current process and puts it at the end of the+ -- thread queue, such that when the current process resumes the finalizer+ -- threads for all dropped weak references have finished.+ -- Unfortunately, this relies on the implementation of the GHC scheduler,+ -- not on any Haskell specification, and is therefore both non-portable and+ -- presumably rather brittle. Therefore, for good measure, we do it twice.+ unsafeIOToPrimStrict $ do+ performMajorGCWithBlockingIfAvailable+ yield+ performMajorGCWithBlockingIfAvailable+ yield+ assertNoForgottenRefs+#endif+ where+ _unused = throwIO @m @SomeException++-- | Ignore and reset the state of forgotten reference tracking. This ensures+-- that any stale forgotten references are not reported later.+--+-- This is especillay important in QC tests with shrinking which otherwise+-- leads to confusion.+ignoreForgottenRefs :: (PrimMonad m, MonadCatch m) => m ()+ignoreForgottenRefs = void $ try @_ @SomeException $ checkForgottenRefs++#ifdef NO_IGNORE_ASSERTS+performMajorGCWithBlockingIfAvailable :: IO ()+performMajorGCWithBlockingIfAvailable = do+#if MIN_VERSION_base(4,20,0)+ performBlockingMajorGC+#else+ performMajorGC+#endif+#endif++-- | Enable forgotten reference checks.+enableForgottenRefChecks :: IO ()++-- | Disable forgotten reference checks. This will error if there are already+-- forgotten references while we are trying to disable the checks.+disableForgottenRefChecks :: IO ()++#ifdef NO_IGNORE_ASSERTS+enableForgottenRefChecks =+ modifyIORef globalForgottenRef $ \case+ Disabled -> Enabled Nothing+ Enabled _ -> error "enableForgottenRefChecks: already enabled"++disableForgottenRefChecks =+ modifyIORef globalForgottenRef $ \case+ Disabled -> error "disableForgottenRefChecks: already disabled"+ Enabled Nothing -> Disabled+ Enabled _ -> error "disableForgottenRefChecks: can not disable when there are forgotten references"+#else+enableForgottenRefChecks = pure ()+disableForgottenRefChecks = pure ()+#endif
@@ -0,0 +1,252 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE RecordWildCards #-}+{-# OPTIONS_HADDOCK not-home #-}+module Database.LSMTree.Internal.Arena (+ ArenaManager,+ newArenaManager,+ Arena,+ Size,+ Offset,+ Alignment,+ withArena,+ newArena,+ closeArena,++ allocateFromArena,+ -- * Test helpers+ withUnmanagedArena,+) where++import Control.DeepSeq (NFData (..))+import Control.Exception (assert)+import Control.Monad.Primitive+import Control.Monad.ST (ST)+import Data.Bits (complement, popCount, (.&.))+import Data.Primitive.ByteArray+import Data.Primitive.MutVar+import Data.Primitive.MVar+import Data.Primitive.PrimVar++#ifdef NO_IGNORE_ASSERTS+import Data.Word (Word8)+#endif++data ArenaManager s = ArenaManager (MutVar s [Arena s])++{-# SPECIALISE+ newArenaManager :: ST s (ArenaManager s)+ #-}+{-# SPECIALISE+ newArenaManager :: IO (ArenaManager RealWorld)+ #-}+newArenaManager :: PrimMonad m => m (ArenaManager (PrimState m))+newArenaManager = do+ m <- newMutVar []+ pure $ ArenaManager m++-- | For use in bencmark environments+instance NFData (ArenaManager s) where+ rnf (ArenaManager !_) = ()++data Arena s = Arena+ { curr :: !(MVar s (Block s)) -- current block, also acts as a lock+ , free :: !(MutVar s [Block s])+ , full :: !(MutVar s [Block s])+ }++data Block s = Block !(PrimVar s Int) !(MutableByteArray s)++instance NFData (Arena s) where+ rnf (Arena !_ !_ !_) = ()++type Size = Int+type Offset = Int+type Alignment = Int++blockSize :: Int+blockSize = 0x100000++{-# SPECIALISE+ newBlock :: ST s (Block s)+ #-}+{-# SPECIALISE+ newBlock :: IO (Block RealWorld)+ #-}+newBlock :: PrimMonad m => m (Block (PrimState m))+newBlock = do+ off <- newPrimVar 0+ mba <- newAlignedPinnedByteArray blockSize 4096+ pure (Block off mba)++{-# INLINE withArena #-}+withArena :: PrimMonad m => ArenaManager (PrimState m) -> (Arena (PrimState m) -> m a) -> m a+withArena am f = do+ a <- newArena am+ x <- f a+ closeArena am a+ pure x++{-# SPECIALISE+ newArena :: ArenaManager s -> ST s (Arena s)+ #-}+{-# SPECIALISE+ newArena :: ArenaManager RealWorld -> IO (Arena RealWorld)+ #-}+newArena :: PrimMonad m => ArenaManager (PrimState m) -> m (Arena (PrimState m))+newArena (ArenaManager arenas) = do+ marena <- atomicModifyMutVar' arenas $ \case+ [] -> ([], Nothing)+ (x:xs) -> (xs, Just x)++ case marena of+ Just arena -> pure arena+ Nothing -> do+ curr <- newBlock >>= newMVar+ free <- newMutVar []+ full <- newMutVar []+ pure Arena {..}++{-# SPECIALISE+ closeArena :: ArenaManager s -> Arena s -> ST s ()+ #-}+{-# SPECIALISE+ closeArena :: ArenaManager RealWorld -> Arena RealWorld -> IO ()+ #-}+closeArena :: PrimMonad m => ArenaManager (PrimState m) -> Arena (PrimState m) -> m ()+closeArena (ArenaManager arenas) arena = do+ scrambleArena arena++ -- reset the arena to clear state+ resetArena arena++ atomicModifyMutVar' arenas $ \xs -> (arena : xs, ())++{-# SPECIALISE+ scrambleArena :: Arena s -> ST s ()+ #-}+{-# SPECIALISE+ scrambleArena :: Arena RealWorld -> IO ()+ #-}+scrambleArena :: PrimMonad m => Arena (PrimState m) -> m ()+#ifndef NO_IGNORE_ASSERTS+scrambleArena _ = pure ()+#else+scrambleArena Arena {..} = do+ readMVar curr >>= scrambleBlock+ readMutVar full >>= mapM_ scrambleBlock+ readMutVar free >>= mapM_ scrambleBlock++{-# SPECIALISE+ scrambleBlock :: Block s -> ST s ()+ #-}+{-# SPECIALISE+ scrambleBlock :: Block RealWorld -> IO ()+ #-}+scrambleBlock :: PrimMonad m => Block (PrimState m) -> m ()+scrambleBlock (Block _ mba) = do+ size <- getSizeofMutableByteArray mba+ setByteArray mba 0 size (0x77 :: Word8)+#endif++{-# SPECIALISE+ resetArena :: Arena s -> ST s ()+ #-}+{-# SPECIALISE+ resetArena :: Arena RealWorld -> IO ()+ #-}+-- | Reset arena, i.e. return used blocks to free list.+resetArena :: PrimMonad m => Arena (PrimState m) -> m ()+resetArena Arena {..} = do+ Block off mba <- takeMVar curr++ -- reset current block+ writePrimVar off 0++ -- move full block to free blocks.+ -- block's offset will be reset in 'newBlockWithFree'+ full' <- atomicModifyMutVar' full $ \xs -> ([], xs)+ atomicModifyMutVar' free $ \xs -> (full' <> xs, ())++ putMVar curr $! Block off mba++-- | Create unmanaged arena.+--+-- Never use this in non-tests code.+{-# SPECIALISE+ withUnmanagedArena :: (Arena s -> ST s a) -> ST s a+ #-}+{-# SPECIALISE+ withUnmanagedArena :: (Arena RealWorld -> IO a) -> IO a+ #-}+withUnmanagedArena :: PrimMonad m => (Arena (PrimState m) -> m a) -> m a+withUnmanagedArena k = do+ mgr <- newArenaManager+ withArena mgr k++{-# SPECIALISE+ allocateFromArena :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s)+ #-}+{-# SPECIALISE+ allocateFromArena :: Arena RealWorld -> Size -> Alignment -> IO (Offset, MutableByteArray RealWorld)+ #-}+-- | Allocate a slice of mutable byte array from the arena.+allocateFromArena :: PrimMonad m => Arena (PrimState m)-> Size -> Alignment -> m (Offset, MutableByteArray (PrimState m))+allocateFromArena !arena !size !alignment =+ assert (popCount alignment == 1) $ -- powers of 2+ assert (size <= blockSize) $ -- not too large allocations+ allocateFromArena' arena size alignment++{-# SPECIALISE+ allocateFromArena' :: Arena s -> Size -> Alignment -> ST s (Offset, MutableByteArray s)+ #-}+{-# SPECIALISE+ allocateFromArena' :: Arena RealWorld -> Size -> Alignment -> IO (Offset, MutableByteArray RealWorld)+ #-}+-- TODO!? this is not async exception safe+allocateFromArena' :: PrimMonad m => Arena (PrimState m)-> Size -> Alignment -> m (Offset, MutableByteArray (PrimState m))+allocateFromArena' arena@Arena { .. } !size !alignment = do+ -- take current block, lock the arena+ curr'@(Block off mba) <- takeMVar curr++ off' <- readPrimVar off+ let !ali = alignment - 1+ let !off'' = (off' + ali) .&. complement ali -- ceil towards next alignment+ let !end = off'' + size+ if end <= blockSize+ then do+ -- fits into current block:+ -- * update offset+ writePrimVar off end+ -- * release lock+ putMVar curr curr'+ -- * return data+ pure (off'', mba)++ else do+ -- doesn't fit into current block:+ -- * move current block into full+ atomicModifyMutVar' full (\xs -> (curr' : xs, ()))+ -- * allocate new block+ new <- newBlockWithFree free+ -- * set new block as current, release the lock+ putMVar curr new+ -- * go again+ allocateFromArena' arena size alignment++{-# SPECIALISE+ newBlockWithFree :: MutVar s [Block s] -> ST s (Block s)+ #-}+{-# SPECIALISE+ newBlockWithFree :: MutVar RealWorld [Block RealWorld] -> IO (Block RealWorld)+ #-}+-- | Allocate new block, possibly taking it from a free list+newBlockWithFree :: PrimMonad m => MutVar (PrimState m) [Block (PrimState m)] -> m (Block (PrimState m))+newBlockWithFree free = do+ free' <- readMutVar free+ case free' of+ [] -> newBlock+ x@(Block off _):xs -> do+ writePrimVar off 0+ writeMutVar free xs+ pure x
@@ -0,0 +1,54 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Assertions (+ assert,+ isValidSlice,+ sameByteArray,+ fromIntegralChecked,+) where++#if MIN_VERSION_base(4,17,0)+import GHC.Exts (isTrue#, sameByteArray#)+#else+import GHC.Exts (ByteArray#, MutableByteArray#, isTrue#,+ sameMutableByteArray#, unsafeCoerce#)+#endif++import Control.Exception (assert)+import Data.Primitive.ByteArray (ByteArray (..), sizeofByteArray)+import GHC.Stack (HasCallStack)+import Text.Printf++isValidSlice :: Int -> Int -> ByteArray -> Bool+isValidSlice off len ba =+ off >= 0 &&+ len >= 0 &&+ (off + len) >= 0 && -- sum doesn't overflow+ (off + len) <= sizeofByteArray ba++sameByteArray :: ByteArray -> ByteArray -> Bool+sameByteArray (ByteArray ba1#) (ByteArray ba2#) =+#if MIN_VERSION_base(4,17,0)+ isTrue# (sameByteArray# ba1# ba2#)+#else+ isTrue# (sameMutableByteArray# (unsafeCoerceByteArray# ba1#)+ (unsafeCoerceByteArray# ba2#))+ where+ unsafeCoerceByteArray# :: ByteArray# -> MutableByteArray# s+ unsafeCoerceByteArray# = unsafeCoerce#+#endif++{-# INLINABLE fromIntegralChecked #-}+-- | Like 'fromIntegral', but throws an error when @(x :: a) /= fromIntegral+-- (fromIntegral x :: b)@.+fromIntegralChecked :: (HasCallStack, Integral a, Integral b, Show a) => a -> b+fromIntegralChecked x+ | x'' == x+ = x'+ | otherwise+ = error $ printf "fromIntegralChecked: conversion failed, %s /= %s" (show x) (show x'')+ where+ x' = fromIntegral x+ x'' = fromIntegral x'
@@ -0,0 +1,189 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Some arithmetic implemented using bit operations+--+-- Valid for non-negative arguments.+module Database.LSMTree.Internal.BitMath (+ -- * Two+ div2,+ mod2,+ mul2,+ ceilDiv2,+ -- * Four+ div4,+ mod4,+ mul4,+ -- * Eight+ div8,+ mod8,+ mul8,+ ceilDiv8,+ -- * Sixteen+ div16,+ mod16,+ mul16,+ -- * 32+ div32,+ mod32,+ -- * 64+ div64,+ mod64,+ mul64,+ ceilDiv64,+ -- * 4096, page size+ divPageSize,+ modPageSize,+ mulPageSize,+ ceilDivPageSize,+ roundUpToPageSize,+) where++import Data.Bits++-------------------------------------------------------------------------------+-- 2+-------------------------------------------------------------------------------++div2 :: Bits a => a -> a+div2 x = unsafeShiftR x 1+{-# INLINE div2 #-}++mod2 :: (Bits a, Num a) => a -> a+mod2 x = x .&. 1+{-# INLINE mod2 #-}++mul2 :: Bits a => a -> a+mul2 x = unsafeShiftL x 1+{-# INLINE mul2 #-}++ceilDiv2 :: (Bits a, Num a) => a -> a+ceilDiv2 i = unsafeShiftR (i + 1) 1+{-# INLINE ceilDiv2 #-}++-------------------------------------------------------------------------------+-- 4+-------------------------------------------------------------------------------++div4 :: Bits a => a -> a+div4 x = unsafeShiftR x 2+{-# INLINE div4 #-}++mod4 :: (Bits a, Num a) => a -> a+mod4 x = x .&. 3+{-# INLINE mod4 #-}++mul4 :: Bits a => a -> a+mul4 x = unsafeShiftL x 2+{-# INLINE mul4 #-}++-------------------------------------------------------------------------------+-- 8+-------------------------------------------------------------------------------++div8 :: Bits a => a -> a+div8 x = unsafeShiftR x 3+{-# INLINE div8 #-}++mod8 :: (Bits a, Num a) => a -> a+mod8 x = x .&. 7+{-# INLINE mod8 #-}++mul8 :: Bits a => a -> a+mul8 x = unsafeShiftL x 3+{-# INLINE mul8 #-}++ceilDiv8 :: (Bits a, Num a) => a -> a+ceilDiv8 i = unsafeShiftR (i + 7) 3+{-# INLINE ceilDiv8 #-}++-------------------------------------------------------------------------------+-- 16+-------------------------------------------------------------------------------++div16 :: Bits a => a -> a+div16 x = unsafeShiftR x 4+{-# INLINE div16 #-}++mod16 :: (Bits a, Num a) => a -> a+mod16 x = x .&. 15+{-# INLINE mod16 #-}++mul16 :: Bits a => a -> a+mul16 x = unsafeShiftL x 4+{-# INLINE mul16 #-}++-------------------------------------------------------------------------------+-- 32+-------------------------------------------------------------------------------++div32 :: Bits a => a -> a+div32 x = unsafeShiftR x 5+{-# INLINE div32 #-}++mod32 :: (Bits a, Num a) => a -> a+mod32 x = x .&. 31+{-# INLINE mod32 #-}++-------------------------------------------------------------------------------+-- 64+-------------------------------------------------------------------------------++-- |+--+-- >>> map div64 [0 :: Word, 1, 63, 64, 65]+-- [0,0,0,1,1]+--+div64 :: Bits a => a -> a+div64 x = unsafeShiftR x 6+{-# INLINE div64 #-}++-- |+--+-- >>> map mod64 [0 :: Word, 1, 63, 64, 65]+-- [0,1,63,0,1]+--+mod64 :: (Bits a, Num a) => a -> a+mod64 x = x .&. 63+{-# INLINE mod64 #-}++mul64 :: Bits a => a -> a+mul64 x = unsafeShiftL x 6+{-# INLINE mul64 #-}++-- | rounding up division+--+-- >>> map ceilDiv64 [0 :: Word, 1, 63, 64, 65]+-- [0,1,1,1,2]+--+ceilDiv64 :: (Bits a, Num a) => a -> a+ceilDiv64 i = unsafeShiftR (i + 63) 6+{-# INLINE ceilDiv64 #-}++-------------------------------------------------------------------------------+-- 4096+-------------------------------------------------------------------------------++-- | Assumes @pageSize = 4096@.+divPageSize :: Bits a => a -> a+divPageSize x = unsafeShiftR x 12+{-# INLINE divPageSize #-}++-- | Assumes @pageSize = 4096@.+modPageSize :: (Bits a, Num a) => a -> a+modPageSize x = x .&. 4095+{-# INLINE modPageSize #-}++-- | Assumes @pageSize = 4096@.+mulPageSize :: Bits a => a -> a+mulPageSize x = unsafeShiftL x 12+{-# INLINE mulPageSize #-}++-- | Assumes @pageSize = 4096@.+ceilDivPageSize :: (Bits a, Num a) => a -> a+ceilDivPageSize x = unsafeShiftR (x + 4095) 12+{-# INLINE ceilDivPageSize #-}++-- | Assumes @pageSize = 4096@.+roundUpToPageSize :: (Bits a, Num a) => a -> a+roundUpToPageSize x = (x + 4095) .&. complement 4095+{-# INLINE roundUpToPageSize #-}
@@ -0,0 +1,129 @@+{-# LANGUAGE TypeFamilies #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.BlobFile (+ BlobFile (..)+ , BlobSpan (..)+ , openBlobFile+ , readBlob+ , readBlobRaw+ , writeBlob+ ) where++import Control.DeepSeq (NFData (..))+import Control.Monad.Class.MonadThrow (MonadCatch (bracketOnError),+ MonadThrow (..))+import Control.Monad.Primitive (PrimMonad)+import Control.RefCount+import qualified Data.Primitive.ByteArray as P+import qualified Data.Vector.Primitive as VP+import Data.Word (Word32, Word64)+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise (SerialisedBlob (..))+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.CallStack (HasCallStack)++-- | A handle to a file containing blobs.+--+-- This is a reference counted object. Upon finalisation, the file is closed+-- and deleted.+--+data BlobFile m h = BlobFile {+ blobFileHandle :: {-# UNPACK #-} !(FS.Handle h),+ blobFileRefCounter :: {-# UNPACK #-} !(RefCounter m)+ }+ deriving stock (Show)++instance RefCounted m (BlobFile m h) where+ getRefCounter = blobFileRefCounter++instance NFData h => NFData (BlobFile m h) where+ rnf (BlobFile a b) = rnf a `seq` rnf b++-- | The location of a blob inside a blob file.+data BlobSpan = BlobSpan {+ blobSpanOffset :: {-# UNPACK #-} !Word64+ , blobSpanSize :: {-# UNPACK #-} !Word32+ }+ deriving stock (Show, Eq)++instance NFData BlobSpan where+ rnf (BlobSpan a b) = rnf a `seq` rnf b++-- | Open the given file to make a 'BlobFile'. The finaliser will close and+-- delete the file.+--+-- REF: the resulting reference must be released once it is no longer used.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates/creates resources.+{-# SPECIALISE openBlobFile :: HasCallStack => HasFS IO h -> FS.FsPath -> FS.OpenMode -> IO (Ref (BlobFile IO h)) #-}+openBlobFile ::+ (PrimMonad m, MonadCatch m)+ => HasCallStack+ => HasFS m h+ -> FS.FsPath+ -> FS.OpenMode+ -> m (Ref (BlobFile m h))+openBlobFile fs path mode =+ bracketOnError (FS.hOpen fs path mode) (FS.hClose fs) $ \blobFileHandle -> do+ let finaliser =+ FS.hClose fs blobFileHandle `finally`+ -- TODO: this function takes ownership of the file path. The file is+ -- removed when the blob file is finalised, which may lead to+ -- surprise errors when the file is also deleted elsewhere. Maybe+ -- file paths should be guarded by 'Ref's as well?+ FS.removeFile fs (FS.handlePath blobFileHandle)+ newRef finaliser $ \blobFileRefCounter ->+ BlobFile {+ blobFileHandle,+ blobFileRefCounter+ }++{-# INLINE readBlob #-}+readBlob ::+ (MonadThrow m, PrimMonad m)+ => HasFS m h+ -> Ref (BlobFile m h)+ -> BlobSpan+ -> m SerialisedBlob+readBlob fs (DeRef blobfile) blobspan = readBlobRaw fs blobfile blobspan++{-# SPECIALISE readBlobRaw :: HasFS IO h -> BlobFile IO h -> BlobSpan -> IO SerialisedBlob #-}+readBlobRaw ::+ (MonadThrow m, PrimMonad m)+ => HasFS m h+ -> BlobFile m h+ -> BlobSpan+ -> m SerialisedBlob+readBlobRaw fs BlobFile {blobFileHandle}+ BlobSpan {blobSpanOffset, blobSpanSize} = do+ let off = FS.AbsOffset blobSpanOffset+ len :: Int+ len = fromIntegral blobSpanSize+ mba <- P.newPinnedByteArray len+ _ <- FS.hGetBufExactlyAt fs blobFileHandle mba 0+ (fromIntegral len :: FS.ByteCount) off+ ba <- P.unsafeFreezeByteArray mba+ let !rb = RB.fromByteArray 0 len ba+ pure (SerialisedBlob rb)++{-# SPECIALISE writeBlob :: HasFS IO h -> Ref (BlobFile IO h) -> SerialisedBlob -> Word64 -> IO () #-}+writeBlob ::+ (MonadThrow m, PrimMonad m)+ => HasFS m h+ -> Ref (BlobFile m h)+ -> SerialisedBlob+ -> Word64+ -> m ()+writeBlob fs (DeRef BlobFile {blobFileHandle})+ (SerialisedBlob' (VP.Vector boff blen ba)) off = do+ mba <- P.unsafeThawByteArray ba+ _ <- FS.hPutBufExactlyAt+ fs blobFileHandle mba+ (FS.BufferOffset boff)+ (fromIntegral blen :: FS.ByteCount)+ (FS.AbsOffset off)+ pure ()
@@ -0,0 +1,221 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE TypeFamilies #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.BlobRef (+ BlobSpan (..)+ , RawBlobRef (..)+ , WeakBlobRef (..)+ , WeakBlobRefInvalid (..)+ , mkRawBlobRef+ , mkWeakBlobRef+ , rawToWeakBlobRef+ , readRawBlobRef+ , readWeakBlobRef+ , readWeakBlobRefs+ ) where++import Control.Monad.Class.MonadThrow (Exception, MonadMask,+ MonadThrow (..), bracket, throwIO)+import Control.Monad.Primitive+import Control.RefCount+import qualified Data.Primitive.ByteArray as P (newPinnedByteArray,+ unsafeFreezeByteArray)+import qualified Data.Vector as V+import qualified Data.Vector.Mutable as VM+import Database.LSMTree.Internal.BlobFile (BlobFile (..),+ BlobSpan (..))+import qualified Database.LSMTree.Internal.BlobFile as BlobFile+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise (SerialisedBlob (..))+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)+++-- | A raw blob reference is a reference to a blob within a blob file.+--+-- The \"raw\" means that it does not maintain ownership of the 'BlobFile' to+-- keep it open. Thus these are only safe to use in the context of code that+-- already (directly or indirectly) owns the blob file that the blob ref uses+-- (such as within run merging).+--+-- Thus these cannot be handed out via the API. Use 'WeakBlobRef' for that.+--+data RawBlobRef m h = RawBlobRef {+ rawBlobRefFile :: {-# NOUNPACK #-} !(BlobFile m h)+ , rawBlobRefSpan :: {-# UNPACK #-} !BlobSpan+ }+ deriving stock (Show)++-- | A \"weak\" reference to a blob within a blob file. These are the ones we+-- can return in the public API and can outlive their parent table.+--+-- They are weak references in that they do not keep the file open using a+-- reference. So when we want to use our weak reference we have to dereference+-- them to obtain a normal strong reference while we do the I\/O to read the+-- blob. This ensures the file is not closed under our feet.+--+-- See 'Database.LSMTree.BlobRef' for more info.+--+data WeakBlobRef m h = WeakBlobRef {+ weakBlobRefFile :: {-# NOUNPACK #-} !(WeakRef (BlobFile m h))+ , weakBlobRefSpan :: {-# UNPACK #-} !BlobSpan+ }+ deriving stock (Show)++-- | A \"strong\" reference to a blob within a blob file. The blob file remains+-- open while the strong reference is live. Thus it is safe to do I\/O to+-- retrieve the blob based on the reference. Strong references must be released+-- using 'releaseBlobRef' when no longer in use (e.g. after completing I\/O).+--+data StrongBlobRef m h = StrongBlobRef {+ strongBlobRefFile :: {-# NOUNPACK #-} !(Ref (BlobFile m h))+ , strongBlobRefSpan :: {-# UNPACK #-} !BlobSpan+ }+ deriving stock (Show)++-- | Convert a 'RawBlobRef' to a 'WeakBlobRef'.+rawToWeakBlobRef :: RawBlobRef m h -> WeakBlobRef m h+rawToWeakBlobRef RawBlobRef {rawBlobRefFile, rawBlobRefSpan} =+ -- This doesn't need to really do anything, because the raw version+ -- does not maintain an independent ref count, and the weak one does+ -- not either.+ WeakBlobRef {+ weakBlobRefFile = mkWeakRefFromRaw rawBlobRefFile,+ weakBlobRefSpan = rawBlobRefSpan+ }++mkRawBlobRef :: Ref (BlobFile m h) -> BlobSpan -> RawBlobRef m h+mkRawBlobRef (DeRef blobfile) blobspan =+ RawBlobRef {+ rawBlobRefFile = blobfile,+ rawBlobRefSpan = blobspan+ }++mkWeakBlobRef :: Ref (BlobFile m h) -> BlobSpan -> WeakBlobRef m h+mkWeakBlobRef blobfile blobspan =+ WeakBlobRef {+ weakBlobRefFile = mkWeakRef blobfile,+ weakBlobRefSpan = blobspan+ }++-- | The 'WeakBlobRef' now points to a blob that is no longer available.+newtype WeakBlobRefInvalid = WeakBlobRefInvalid Int+ deriving stock (Show)+ deriving anyclass (Exception)++{-# SPECIALISE deRefWeakBlobRef ::+ WeakBlobRef IO h+ -> IO (StrongBlobRef IO h) #-}+deRefWeakBlobRef ::+ (MonadThrow m, PrimMonad m)+ => WeakBlobRef m h+ -> m (StrongBlobRef m h)+deRefWeakBlobRef WeakBlobRef{weakBlobRefFile, weakBlobRefSpan} = do+ mstrongBlobRefFile <- deRefWeak weakBlobRefFile+ case mstrongBlobRefFile of+ Just strongBlobRefFile ->+ pure StrongBlobRef {+ strongBlobRefFile,+ strongBlobRefSpan = weakBlobRefSpan+ }+ Nothing -> throwIO (WeakBlobRefInvalid 0)++{-# SPECIALISE deRefWeakBlobRefs ::+ V.Vector (WeakBlobRef IO h)+ -> IO (V.Vector (StrongBlobRef IO h)) #-}+deRefWeakBlobRefs ::+ forall m h.+ (MonadMask m, PrimMonad m)+ => V.Vector (WeakBlobRef m h)+ -> m (V.Vector (StrongBlobRef m h))+deRefWeakBlobRefs wrefs = do+ refs <- VM.new (V.length wrefs)+ V.iforM_ wrefs $ \i WeakBlobRef {weakBlobRefFile, weakBlobRefSpan} -> do+ mstrongBlobRefFile <- deRefWeak weakBlobRefFile+ case mstrongBlobRefFile of+ Just strongBlobRefFile ->+ VM.write refs i StrongBlobRef {+ strongBlobRefFile,+ strongBlobRefSpan = weakBlobRefSpan+ }+ Nothing -> do+ -- drop refs on the previous ones taken successfully so far+ VM.mapM_ releaseBlobRef (VM.take i refs)+ throwIO (WeakBlobRefInvalid i)+ V.unsafeFreeze refs++{-# INLINE releaseBlobRef #-}+releaseBlobRef :: (MonadMask m, PrimMonad m) => StrongBlobRef m h -> m ()+releaseBlobRef = releaseRef . strongBlobRefFile++{-# INLINE readRawBlobRef #-}+readRawBlobRef ::+ (MonadThrow m, PrimMonad m)+ => HasFS m h+ -> RawBlobRef m h+ -> m SerialisedBlob+readRawBlobRef fs RawBlobRef {rawBlobRefFile, rawBlobRefSpan} =+ BlobFile.readBlobRaw fs rawBlobRefFile rawBlobRefSpan++{-# SPECIALISE readWeakBlobRef :: HasFS IO h -> WeakBlobRef IO h -> IO SerialisedBlob #-}+readWeakBlobRef ::+ (MonadMask m, PrimMonad m)+ => HasFS m h+ -> WeakBlobRef m h+ -> m SerialisedBlob+readWeakBlobRef fs wref =+ bracket (deRefWeakBlobRef wref) releaseBlobRef $+ \StrongBlobRef {strongBlobRefFile, strongBlobRefSpan} ->+ BlobFile.readBlob fs strongBlobRefFile strongBlobRefSpan++{-# SPECIALISE readWeakBlobRefs :: HasBlockIO IO h -> V.Vector (WeakBlobRef IO h) -> IO (V.Vector SerialisedBlob) #-}+readWeakBlobRefs ::+ (MonadMask m, PrimMonad m)+ => HasBlockIO m h+ -> V.Vector (WeakBlobRef m h)+ -> m (V.Vector SerialisedBlob)+readWeakBlobRefs hbio wrefs =+ bracket (deRefWeakBlobRefs wrefs) (V.mapM_ releaseBlobRef) $ \refs -> do+ -- Prepare the IOOps:+ -- We use a single large memory buffer, with appropriate offsets within+ -- the buffer.+ let bufSize :: Int+ !bufSize = V.sum (V.map blobRefSpanSize refs)++ {-# INLINE bufOffs #-}+ bufOffs :: V.Vector Int+ bufOffs = V.scanl (+) 0 (V.map blobRefSpanSize refs)+ buf <- P.newPinnedByteArray bufSize++ -- Submit the IOOps all in one go:+ _ <- FS.submitIO hbio $+ V.zipWith+ (\bufoff+ StrongBlobRef {+ strongBlobRefFile = DeRef BlobFile {blobFileHandle},+ strongBlobRefSpan = BlobSpan {blobSpanOffset, blobSpanSize}+ } ->+ FS.IOOpRead+ blobFileHandle+ (fromIntegral blobSpanOffset :: FS.FileOffset)+ buf (FS.BufferOffset bufoff)+ (fromIntegral blobSpanSize :: FS.ByteCount))+ bufOffs refs+ -- We do not need to inspect the results because IO errors are+ -- thrown as exceptions, and the result is just the read length+ -- which is already known. Short reads can't happen here.++ -- Construct the SerialisedBlobs results:+ -- This is just the different offsets within the shared buffer.+ ba <- P.unsafeFreezeByteArray buf+ pure $! V.zipWith+ (\off len -> SerialisedBlob (RB.fromByteArray off len ba))+ bufOffs+ (V.map blobRefSpanSize refs)+ where+ blobRefSpanSize = fromIntegral . blobSpanSize . strongBlobRefSpan
@@ -0,0 +1,326 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.BloomFilter (+ -- * Types+ Bloom.Bloom,+ Bloom.MBloom,+ Bloom.Salt,++ -- * Bulk query+ bloomQueries,+ RunIxKeyIx(RunIxKeyIx),+ RunIx, KeyIx,++ -- * Serialisation+ bloomFilterVersion,+ bloomFilterToLBS,+ bloomFilterFromFile,+) where++import Data.Bits+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder.Extra as B+import qualified Data.ByteString.Lazy as LBS+import qualified Data.Primitive as P+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import Data.Word (Word32, Word64, byteSwap32)+import Text.Printf (printf)++import Control.Exception (assert)+import Control.Monad (void, when)+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive (PrimMonad)+import Control.Monad.ST (ST, runST)+import System.FS.API++import Data.BloomFilter.Blocked (Bloom)+import qualified Data.BloomFilter.Blocked as Bloom+import Database.LSMTree.Internal.ByteString (byteArrayToByteString)+import Database.LSMTree.Internal.CRC32C (FileCorruptedError (..),+ FileFormat (..))+import Database.LSMTree.Internal.Serialise (SerialisedKey)+import qualified Database.LSMTree.Internal.Vector as P++#ifdef HAVE_STRICT_ARRAY+import qualified Database.LSMTree.Internal.StrictArray as P+#endif++import Prelude hiding (filter)++-- Bulk query+-----------------------------------------------------------++type KeyIx = Int+type RunIx = Int++-- | A 'RunIxKeyIx' is a (compact) pair of a 'RunIx' and a 'KeyIx'.+--+-- We represent it as a 32bit word, using:+--+-- * 16 bits for the run\/filter index (MSB)+-- * 16 bits for the key index (LSB)+--+newtype RunIxKeyIx = MkRunIxKeyIx Word32+ deriving stock Eq+ deriving newtype P.Prim++pattern RunIxKeyIx :: RunIx -> KeyIx -> RunIxKeyIx+pattern RunIxKeyIx r k <- (unpackRunIxKeyIx -> (r, k))+ where+ RunIxKeyIx r k = packRunIxKeyIx r k+{-# INLINE RunIxKeyIx #-}+{-# COMPLETE RunIxKeyIx #-}++packRunIxKeyIx :: Int -> Int -> RunIxKeyIx+packRunIxKeyIx r k =+ assert (r >= 0 && r <= 0xffff+ && k >= 0 && k <= 0xffff) $+ MkRunIxKeyIx $+ (fromIntegral :: Word -> Word32) $+ (fromIntegral r `unsafeShiftL` 16)+ .|. fromIntegral k+{-# INLINE packRunIxKeyIx #-}++unpackRunIxKeyIx :: RunIxKeyIx -> (Int, Int)+unpackRunIxKeyIx (MkRunIxKeyIx c) =+ ( fromIntegral (c `unsafeShiftR` 16)+ , fromIntegral (c .&. 0xfff)+ )+{-# INLINE unpackRunIxKeyIx #-}++instance Show RunIxKeyIx where+ showsPrec _ (RunIxKeyIx r k) =+ showString "RunIxKeyIx " . showsPrec 11 r+ . showChar ' ' . showsPrec 11 k++type ResIx = Int -- Result index++{-# NOINLINE bloomQueries #-}+-- | Perform a batch of bloom queries. The result is a tuple of indexes into the+-- vector of runs and vector of keys respectively. The order of keys and+-- runs\/filters in the input is maintained in the output. This implementation+-- produces results in key-major order.+--+-- The result vector can be of variable length. The initial estimate is 2x the+-- number of keys but this is grown if needed (using a doubling strategy).+--+bloomQueries ::+ Bloom.Salt+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector SerialisedKey+ -> VP.Vector RunIxKeyIx+bloomQueries !_salt !filters !keys+ | V.null filters || V.null keys = VP.empty+bloomQueries !salt !filters !keys =+ runST (bloomQueries_loop1 filters' keyhashes)+ where+ filters' = toFiltersArray filters+ keyhashes = P.generatePrimArray (V.length keys) $ \i ->+ Bloom.hashesWithSalt salt (V.unsafeIndex keys i)++-- loop over all keys+bloomQueries_loop1 ::+ BloomFilters+ -> P.PrimArray (Bloom.Hashes SerialisedKey)+ -> ST s (VP.Vector RunIxKeyIx)+bloomQueries_loop1 !filters !keyhashes = do+ res <- P.newPrimArray (P.sizeofPrimArray keyhashes * 2)+ (res', resix') <- go res 0 0+ parr <- P.unsafeFreezePrimArray =<< P.resizeMutablePrimArray res' resix'+ pure $! P.primArrayToPrimVector parr+ where+ go !res !resix !kix+ | kix == P.sizeofPrimArray keyhashes = pure (res, resix)+ | otherwise = do+ let !keyhash = P.indexPrimArray keyhashes kix+ bloomQueries_loop2_prefetch filters keyhash 0+ (res', resix') <- bloomQueries_loop2 filters keyhash kix res resix 0+ go res' resix' (kix+1)++-- loop over all filters+bloomQueries_loop2 ::+ BloomFilters+ -> Bloom.Hashes SerialisedKey+ -> KeyIx+ -> P.MutablePrimArray s RunIxKeyIx+ -> ResIx+ -> RunIx+ -> ST s (P.MutablePrimArray s RunIxKeyIx, ResIx)+bloomQueries_loop2 !filters !keyhash !kix = go+ where+ go res2 resix2 rix+ | rix == lengthFiltersArray filters = pure (res2, resix2)+ | let !filter = indexFiltersArray filters rix+ , Bloom.elemHashes filter keyhash = do+ P.writePrimArray res2 resix2 (RunIxKeyIx rix kix)+ ressz2 <- P.getSizeofMutablePrimArray res2+ res2' <- if resix2+1 < ressz2+ then pure res2+ else P.resizeMutablePrimArray res2 (ressz2 * 2)+ go res2' (resix2+1) (rix+1)++ | otherwise =+ go res2 resix2 (rix+1)++bloomQueries_loop2_prefetch ::+ BloomFilters+ -> Bloom.Hashes SerialisedKey+ -> RunIx+ -> ST s ()+bloomQueries_loop2_prefetch !filters !keyhash = go+ where+ go !rix+ | rix == lengthFiltersArray filters = pure ()+ | otherwise = do+ let !filter = indexFiltersArray filters rix+ Bloom.prefetchElem filter keyhash+ go (rix+1)++type BloomFilters =+#ifdef HAVE_STRICT_ARRAY+ P.StrictArray (Bloom SerialisedKey)+#else+ V.Vector (Bloom SerialisedKey)+#endif++{-# INLINE toFiltersArray #-}+toFiltersArray :: V.Vector (Bloom SerialisedKey) -> BloomFilters++{-# INLINE indexFiltersArray #-}+indexFiltersArray :: BloomFilters -> Int -> Bloom SerialisedKey++{-# INLINE lengthFiltersArray #-}+lengthFiltersArray :: BloomFilters -> Int++#ifdef HAVE_STRICT_ARRAY+toFiltersArray = P.vectorToStrictArray+indexFiltersArray = P.indexStrictArray+lengthFiltersArray = P.sizeofStrictArray+#else+toFiltersArray = id+indexFiltersArray = V.unsafeIndex+lengthFiltersArray = V.length+#endif+++-- serialising+-----------------------------------------------------------++-- | By writing out the version in host endianness, we also indicate endianness.+-- During deserialisation, we would discover an endianness mismatch.+--+-- We base our version number on the 'Bloom.formatVersion' from the @bloomfilter@+-- library, plus our own version here. This accounts both for changes in the+-- format code here, and changes in the library.+--+bloomFilterVersion :: Word32+bloomFilterVersion = 1 + fromIntegral Bloom.formatVersion++bloomFilterToLBS :: Bloom a -> LBS.ByteString+bloomFilterToLBS bf =+ let (size, salt, ba, off, len) = Bloom.serialise bf+ in header size salt <> byteArrayToLBS ba off len+ where+ header Bloom.BloomSize { sizeBits, sizeHashes } salt =+ -- creates a single 24 byte chunk+ B.toLazyByteStringWith (B.safeStrategy 24 B.smallChunkSize) mempty $+ B.word32Host bloomFilterVersion+ <> B.word32Host (fromIntegral sizeHashes)+ <> B.word64Host (fromIntegral sizeBits)+ <> B.word64Host salt++ byteArrayToLBS :: P.ByteArray -> Int -> Int -> LBS.ByteString+ byteArrayToLBS ba off len =+ LBS.fromStrict (byteArrayToByteString off len ba)++-- deserialising+-----------------------------------------------------------++{-# SPECIALISE bloomFilterFromFile ::+ HasFS IO h+ -> Bloom.Salt+ -> Handle h+ -> IO (Bloom a) #-}+-- | Read a 'Bloom' from a file.+--+bloomFilterFromFile ::+ (PrimMonad m, MonadCatch m)+ => HasFS m h+ -> Bloom.Salt -- ^ Expected salt+ -> Handle h -- ^ The open file, in read mode+ -> m (Bloom a)+bloomFilterFromFile hfs expectedSalt h = do+ header <- rethrowEOFError "Doesn't contain a header" $+ hGetByteArrayExactly hfs h 24++ let !version = P.indexByteArray header 0 :: Word32+ !nhashes = P.indexByteArray header 1 :: Word32+ !nbits = P.indexByteArray header 1 :: Word64+ !salt = P.indexByteArray header 2 :: Bloom.Salt++ when (version /= bloomFilterVersion) $ throwFormatError $+ if byteSwap32 version == bloomFilterVersion+ then "Different byte order"+ else "Unsupported version"++ when (nbits <= 0) $ throwFormatError "Length is zero"++ -- limit to 2^48 bits+ when (nbits >= fromIntegral Bloom.maxSizeBits) $ throwFormatError "Too large bloomfilter"++ when (expectedSalt /= salt) $ throwFormatError $+ printf "Expected salt does not match actual salt: %d /= %d"+ expectedSalt+ salt++ -- read the filter data from the file directly into the bloom filter+ bloom <-+ Bloom.deserialise+ Bloom.BloomSize {+ Bloom.sizeBits = fromIntegral nbits,+ Bloom.sizeHashes = fromIntegral nhashes+ }+ salt+ (\buf off len ->+ rethrowEOFError "bloom filter file too short" $+ void $ hGetBufExactly hfs+ h buf (BufferOffset off) (fromIntegral len))++ -- check we're now at EOF+ trailing <- hGetSome hfs h 1+ when (not (BS.null trailing)) $+ throwFormatError "Byte array is too large for components"+ pure bloom+ where+ throwFormatError = throwIO+ . ErrFileFormatInvalid+ (mkFsErrorPath hfs (handlePath h))+ FormatBloomFilterFile+ rethrowEOFError msg =+ handleJust+ (\e -> if isFsErrorType FsReachedEOF e then Just e else Nothing)+ (\e -> throwIO $+ ErrFileFormatInvalid+ (fsErrorPath e) FormatBloomFilterFile msg)++{-# SPECIALISE hGetByteArrayExactly ::+ HasFS IO h+ -> Handle h+ -> Int+ -> IO P.ByteArray #-}+hGetByteArrayExactly ::+ (PrimMonad m, MonadThrow m)+ => HasFS m h+ -> Handle h+ -> Int+ -> m P.ByteArray+hGetByteArrayExactly hfs h len = do+ buf <- P.newByteArray len+ _ <- hGetBufExactly hfs h buf 0 (fromIntegral len)+ P.unsafeFreezeByteArray buf+
@@ -0,0 +1,138 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | @bytestring@ extras+--+module Database.LSMTree.Internal.ByteString (+ tryCheapToShort,+ tryGetByteArray,+ shortByteStringFromTo,+ byteArrayFromTo,+ byteArrayToByteString,+ unsafePinnedByteArrayToByteString,+ byteArrayToSBS,+) where++import Control.Exception (assert)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BB+import qualified Data.ByteString.Builder.Internal as BB+import qualified Data.ByteString.Internal as BS.Internal+import Data.ByteString.Short (ShortByteString (SBS))+import qualified Data.ByteString.Short.Internal as SBS+import Data.Primitive.ByteArray+import Database.LSMTree.Internal.Assertions (isValidSlice)+import Foreign.Ptr (minusPtr, plusPtr)+import GHC.Exts+import qualified GHC.ForeignPtr as Foreign+import GHC.Stack (HasCallStack)++-- | \( O(1) \) conversion, if possible.+--+-- In addition to the conditions explained for 'tryGetByteArray', the+-- bytestring must use the full length of the underlying byte array.+tryCheapToShort :: BS.ByteString -> Either String ShortByteString+tryCheapToShort bs =+ tryGetByteArray bs >>= \(ba , n) ->+ if n /= sizeofByteArray ba then+ Left "ByteString does not use full ByteArray"+ else+ let !(ByteArray ba#) = ba in Right (SBS ba#)+++-- | \( O(1) \) conversion from a strict 'BS.ByteString' to its underlying+-- pinned 'ByteArray', if possible. Also returns the length (in bytes) of the+-- byte array prefix that was used by the bytestring.+--+-- Strict bytestrings are allocated using 'mallocPlainForeignPtrBytes', so we+-- are expecting a 'PlainPtr' (or 'FinalPtr' when the length is 0).+-- We also require that bytestrings referencing a byte array point point at the+-- beginning, without any offset.+tryGetByteArray :: BS.ByteString -> Either String (ByteArray, Int)+tryGetByteArray (BS.Internal.BS (Foreign.ForeignPtr addr# contents) n) =+ case contents of+ Foreign.PlainPtr mba# ->+ case mutableByteArrayContents# mba# `eqAddr#` addr# of+ 0# -> Left "non-zero offset into ByteArray"+ _ -> -- safe, ByteString's content is considered immutable+ Right $ case unsafeFreezeByteArray# mba# realWorld# of+ (# _, ba# #) -> (ByteArray ba#, n)+ Foreign.MallocPtr {} ->+ Left ("unsupported MallocPtr (length " <> show n <> ")")+ Foreign.PlainForeignPtr {} ->+ Left ("unsupported PlainForeignPtr (length " <> show n <> ")")+ Foreign.FinalPtr | n == 0 ->+ -- We can also handle empty bytestrings ('BS.empty' uses 'FinalPtr').+ Right (emptyByteArray, 0)+ Foreign.FinalPtr ->+ Left ("unsupported FinalPtr (length " <> show n <> ")")++-- | Copy of 'SBS.shortByteString', but with bounds (unchecked).+--+-- https://github.com/haskell/bytestring/issues/664+{-# INLINE shortByteStringFromTo #-}+shortByteStringFromTo :: Int -> Int -> ShortByteString -> BB.Builder+shortByteStringFromTo = \i j sbs -> BB.builder $ shortByteStringCopyStepFromTo i j sbs++-- | Like 'shortByteStringFromTo' but for 'ByteArray'+--+-- https://github.com/haskell/bytestring/issues/664+byteArrayFromTo :: Int -> Int -> ByteArray -> BB.Builder+byteArrayFromTo = \i j (ByteArray ba) -> BB.builder $ shortByteStringCopyStepFromTo i j (SBS ba)++-- | Copy of 'SBS.shortByteStringCopyStep' but with bounds (unchecked)+{-# INLINE shortByteStringCopyStepFromTo #-}+shortByteStringCopyStepFromTo ::+ Int -> Int -> ShortByteString -> BB.BuildStep a -> BB.BuildStep a+shortByteStringCopyStepFromTo !ip0 !ipe0 !sbs k =+ go ip0 ipe0+ where+ go !ip !ipe (BB.BufferRange op ope)+ | inpRemaining <= outRemaining = do+ SBS.copyToPtr sbs ip op inpRemaining+ let !br' = BB.BufferRange (op `plusPtr` inpRemaining) ope+ k br'+ | otherwise = do+ SBS.copyToPtr sbs ip op outRemaining+ let !ip' = ip + outRemaining+ pure $ BB.bufferFull 1 ope (go ip' ipe)+ where+ outRemaining = ope `minusPtr` op+ inpRemaining = ipe - ip++-- | \( O(1) \) conversion if the byte array is pinned, \( O(n) \) otherwise.+-- Takes offset and length of the slice to be used.+byteArrayToByteString :: Int -> Int -> ByteArray -> BS.ByteString+byteArrayToByteString off len ba =+ assert (isValidSlice off len ba) $+ if isByteArrayPinned ba+ then unsafePinnedByteArrayToByteString off len ba+ else unsafePinnedByteArrayToByteString 0 len $ runByteArray $ do+ mba <- newPinnedByteArray len+ copyByteArray mba 0 ba off len+ pure mba++-- | \( O(1) \) conversion. Takes offset and length of the slice to be used.+-- Fails if the byte array is not pinned.+--+-- Based on 'SBS.fromShort'.+unsafePinnedByteArrayToByteString :: HasCallStack => Int -> Int -> ByteArray -> BS.ByteString+unsafePinnedByteArrayToByteString off@(I# off#) len ba@(ByteArray ba#) =+ assert (isValidSlice off len ba) $+ if isByteArrayPinned ba+ then BS.Internal.BS fp len+ else error $ "unsafePinnedByteArrayToByteString: not pinned, length "+ <> show (sizeofByteArray ba)+ where+ addr# = plusAddr# (byteArrayContents# ba#) off#+ fp = Foreign.ForeignPtr addr# (Foreign.PlainPtr (unsafeCoerce# ba#))++-- | \( O(1) \) conversion.+byteArrayToSBS :: ByteArray -> ShortByteString+#if MIN_VERSION_bytestring(0,12,0)+byteArrayToSBS ba = SBS.ShortByteString ba+#else+byteArrayToSBS (ByteArray ba) = SBS.SBS ba+#endif
@@ -0,0 +1,501 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_HADDOCK not-home #-}++-- Needed by cabal repl as Data.Digest.CRC32C is ambiguous+{-# LANGUAGE PackageImports #-}++-- Needed by GHC <= 9.2 for newtype deriving Prim below+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE UnboxedTuples #-}++-- | Functionality related to CRC-32C (Castagnoli) checksums:+--+-- * Support for calculating checksums while incrementally writing files.+-- * Support for verifying checksums of files.+-- * Support for a text file format listing file checksums.+--+module Database.LSMTree.Internal.CRC32C (++ CRC32C(..),++ -- * Pure incremental checksum calculation+ initialCRC32C,+ updateCRC32C,++ -- * I\/O with checksum calculation+ hGetSomeCRC32C,+ hGetExactlyCRC32C,+ hPutSomeCRC32C,+ hPutAllCRC32C,+ hPutAllChunksCRC32C,+ readFileCRC32C,++ ChunkSize (..),+ defaultChunkSize,+ hGetExactlyCRC32C_SBS,+ hGetAllCRC32C',++ -- * Checksum files+ -- $checksum-files+ ChecksumsFile,+ ChecksumsFileName(..),+ getChecksum,+ readChecksumsFile,+ writeChecksumsFile,+ writeChecksumsFile',++ -- * Checksum checking+ checkCRC,+ expectChecksum,++ -- * File format errors+ FileFormat (..),+ FileCorruptedError (..),+ expectValidFile,+ ) where++import Control.Monad+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive+import Data.Bits+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BS+import qualified Data.ByteString.Char8 as BSC+import qualified Data.ByteString.Internal as BS.Internal+import qualified Data.ByteString.Lazy as BSL+import qualified Data.ByteString.Short as SBS+import Data.Char (ord)+import "crc32c" Data.Digest.CRC32C as CRC+import Data.Either (partitionEithers)+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Primitive+import Data.Word+import GHC.Exts+import qualified GHC.ForeignPtr as Foreign+import System.FS.API+import System.FS.API.Lazy+import System.FS.BlockIO.API (Advice (..), ByteCount, HasBlockIO,+ hAdviseAll, hDropCacheAll)++newtype CRC32C = CRC32C {unCRC32C :: Word32}+ deriving stock (Eq, Ord, Show)+ deriving newtype (Prim)+++initialCRC32C :: CRC32C+initialCRC32C = CRC32C 0 -- same as crc32c BS.empty+++updateCRC32C :: BS.ByteString -> CRC32C -> CRC32C+updateCRC32C bs (CRC32C crc) = CRC32C (CRC.crc32c_update crc bs)++{-# SPECIALISE hGetSomeCRC32C :: HasFS IO h -> Handle h -> Word64 -> CRC32C -> IO (BS.ByteString, CRC32C) #-}+hGetSomeCRC32C :: Monad m+ => HasFS m h+ -> Handle h+ -> Word64+ -> CRC32C -> m (BS.ByteString, CRC32C)+hGetSomeCRC32C fs h n crc = do+ bs <- hGetSome fs h n+ let !crc' = updateCRC32C bs crc+ pure (bs, crc')+++-- | This function ensures that exactly the requested number of bytes is read.+-- If the file is too short, an 'FsError' of type 'FsReachedEOF' is thrown.+--+-- It attempts to read everything into a single strict chunk, which should+-- almost always succeed. If it doesn't, multiple chunks are produced.+--+-- TODO: To reliably return a strict bytestring without additional copying,+-- @fs-api@ needs to support directly reading into a buffer, which is currently+-- work in progress: <https://github.com/input-output-hk/fs-sim/pull/46>+{-# SPECIALISE hGetExactlyCRC32C :: HasFS IO h -> Handle h -> Word64 -> CRC32C -> IO (BSL.ByteString, CRC32C) #-}+hGetExactlyCRC32C :: MonadThrow m+ => HasFS m h+ -> Handle h+ -> Word64+ -> CRC32C -> m (BSL.ByteString, CRC32C)+hGetExactlyCRC32C fs h n crc = do+ lbs <- hGetExactly fs h n+ let !crc' = BSL.foldlChunks (flip updateCRC32C) crc lbs+ pure (lbs, crc')+++{-# SPECIALISE hPutSomeCRC32C :: HasFS IO h -> Handle h -> BS.ByteString -> CRC32C -> IO (Word64, CRC32C) #-}+hPutSomeCRC32C :: Monad m+ => HasFS m h+ -> Handle h+ -> BS.ByteString+ -> CRC32C -> m (Word64, CRC32C)+hPutSomeCRC32C fs h bs crc = do+ !n <- hPutSome fs h bs+ let !crc' = updateCRC32C (BS.take (fromIntegral n) bs) crc+ pure (n, crc')+++-- | This function makes sure that the whole 'BS.ByteString' is written.+{-# SPECIALISE hPutAllCRC32C :: HasFS IO h -> Handle h -> BS.ByteString -> CRC32C -> IO (Word64, CRC32C) #-}+hPutAllCRC32C :: forall m h+ . Monad m+ => HasFS m h+ -> Handle h+ -> BS.ByteString+ -> CRC32C -> m (Word64, CRC32C)+hPutAllCRC32C fs h = go 0+ where+ go :: Word64 -> BS.ByteString -> CRC32C -> m (Word64, CRC32C)+ go !written !bs !crc = do+ (n, crc') <- hPutSomeCRC32C fs h bs crc+ let bs' = BS.drop (fromIntegral n) bs+ written' = written + n+ if BS.null bs'+ then pure (written', crc')+ else go written' bs' crc'++-- | This function makes sure that the whole /lazy/ 'BSL.ByteString' is written.+{-# SPECIALISE hPutAllChunksCRC32C :: HasFS IO h -> Handle h -> BSL.ByteString -> CRC32C -> IO (Word64, CRC32C) #-}+hPutAllChunksCRC32C :: forall m h+ . Monad m+ => HasFS m h+ -> Handle h+ -> BSL.ByteString+ -> CRC32C -> m (Word64, CRC32C)+hPutAllChunksCRC32C fs h = \lbs crc ->+ foldM (uncurry putChunk) (0, crc) (BSL.toChunks lbs)+ where+ putChunk :: Word64 -> CRC32C -> BS.ByteString -> m (Word64, CRC32C)+ putChunk !written !crc !bs = do+ (n, crc') <- hPutAllCRC32C fs h bs crc+ pure (written + n, crc')++{-# SPECIALISE readFileCRC32C :: HasFS IO h -> FsPath -> IO CRC32C #-}+readFileCRC32C :: forall m h. MonadThrow m => HasFS m h -> FsPath -> m CRC32C+readFileCRC32C fs file =+ withFile fs file ReadMode (\h -> go h initialCRC32C)+ where+ go :: Handle h -> CRC32C -> m CRC32C+ go !h !crc = do+ bs <- hGetSome fs h 65504 -- 2^16 - 4 words overhead+ if BS.null bs+ then pure crc+ else go h (updateCRC32C bs crc)++newtype ChunkSize = ChunkSize ByteCount++defaultChunkSize :: ChunkSize+defaultChunkSize = ChunkSize 65504 -- 2^16 - 4 words overhead++{-# SPECIALISE hGetExactlyCRC32C_SBS :: HasFS IO h -> Handle h -> ByteCount -> CRC32C -> IO (SBS.ShortByteString, CRC32C) #-}+-- | Reads exactly as many bytes as requested, returning a 'ShortByteString' and+-- updating a given 'CRC32C' value.+--+-- If EOF is found before the requested number of bytes is read, an FsError+-- exception is thrown.+--+-- The returned 'ShortByteString' is backed by pinned memory.+hGetExactlyCRC32C_SBS ::+ forall m h. (MonadThrow m, PrimMonad m)+ => HasFS m h+ -> Handle h+ -> ByteCount -- ^ Number of bytes to read+ -> CRC32C+ -> m (SBS.ShortByteString, CRC32C)+hGetExactlyCRC32C_SBS hfs h !c !crc = do+ buf@(MutableByteArray !mba#) <- newPinnedByteArray (fromIntegral c)+ void $ hGetBufExactly hfs h buf 0 c+ (ByteArray !ba#) <- unsafeFreezeByteArray buf+ let fp = Foreign.ForeignPtr (byteArrayContents# ba#) (Foreign.PlainPtr (unsafeCoerce# mba#))+ !bs = BS.Internal.BS fp (fromIntegral c)+ !crc' = updateCRC32C bs crc+ pure (SBS.SBS ba#, crc')++{-# SPECIALISE hGetAllCRC32C' :: HasFS IO h -> Handle h -> ChunkSize -> CRC32C -> IO CRC32C #-}+-- | Reads all bytes, updating a given 'CRC32C' value without returning the+-- bytes.+hGetAllCRC32C' ::+ forall m h. PrimMonad m+ => HasFS m h+ -> Handle h+ -> ChunkSize -- ^ Chunk size, must be larger than 0+ -> CRC32C+ -> m CRC32C+hGetAllCRC32C' hfs h (ChunkSize !chunkSize) !crc0+ | chunkSize <= 0+ = error "hGetAllCRC32C': chunkSize must be >0"+ | otherwise+ = do+ buf@(MutableByteArray !mba#) <- newPinnedByteArray (fromIntegral chunkSize)+ (ByteArray !ba#) <- unsafeFreezeByteArray buf+ let fp = Foreign.ForeignPtr (byteArrayContents# ba#) (Foreign.PlainPtr (unsafeCoerce# mba#))+ !bs = BS.Internal.BS fp (fromIntegral chunkSize)+ go bs buf crc0+ where+ -- In particular, note that the "immutable" bs :: BS.ByteString aliases the+ -- mutable buf :: MutableByteArray. This is a bit hairy but we need to do+ -- something like this because the CRC code only works with ByteString.+ -- We thus have to be very careful about when bs is used.+ go :: BS.ByteString -> MutableByteArray (PrimState m) -> CRC32C -> m CRC32C+ go !bs buf !crc = do+ !n <- hGetBufSome hfs h buf 0 chunkSize+ if n == 0+ then pure crc+ else do+ -- compute the update CRC value before reading the next bytes+ let !crc' = updateCRC32C (BS.take (fromIntegral n) bs) crc+ go bs buf crc'+++{- $checksum-files+We use @.checksum@ files to help verify the integrity of on disk snapshots.+Each .checksum file lists the CRC-32C (Castagnoli) of other files. For further+details see @doc/format-directory.md@.++The file uses the BSD-style checksum format (e.g. as produced by tools like+@md5sum --tag@), with the algorithm name \"CRC32C\". This format is text,+one line per file, using hexedecimal for the 32bit output.++Checksum files are used for each LSM run, and for the snapshot metadata.++Typical examples are:++> CRC32C (keyops) = fd040004+> CRC32C (blobs) = 5a3b820c+> CRC32C (filter) = 6653e178+> CRC32C (index) = f4ec6724++Or++> CRC32C (snapshot) = 87972d7f+-}++type ChecksumsFile = Map ChecksumsFileName CRC32C++-- | File names must not include characters @'('@, @')'@ or @'\n'@.+--+newtype ChecksumsFileName = ChecksumsFileName {unChecksumsFileName :: BSC.ByteString}+ deriving stock (Eq, Ord, Show)++{-# SPECIALISE+ getChecksum ::+ HasFS IO h+ -> FsPath+ -> ChecksumsFile+ -> ChecksumsFileName+ -> IO CRC32C+ #-}+getChecksum ::+ MonadThrow m+ => HasFS m h+ -> FsPath+ -> ChecksumsFile+ -> ChecksumsFileName+ -> m CRC32C+getChecksum hfs fsPath checksumsFile checksumsFileName =+ case Map.lookup checksumsFileName checksumsFile of+ Just checksum ->+ pure checksum+ Nothing ->+ throwIO $+ ErrFileFormatInvalid+ (mkFsErrorPath hfs fsPath)+ FormatChecksumsFile+ ("could not find checksum for " <>+ show (unChecksumsFileName checksumsFileName))++{-# SPECIALISE+ readChecksumsFile ::+ HasFS IO h+ -> FsPath+ -> IO ChecksumsFile+ #-}+readChecksumsFile ::+ (MonadThrow m)+ => HasFS m h+ -> FsPath+ -> m ChecksumsFile+readChecksumsFile fs path = do+ str <- withFile fs path ReadMode (\h -> hGetAll fs h)+ expectValidFile fs path FormatChecksumsFile (parseChecksumsFile (BSL.toStrict str))++{-# SPECIALISE writeChecksumsFile :: HasFS IO h -> FsPath -> ChecksumsFile -> IO () #-}+writeChecksumsFile :: MonadThrow m+ => HasFS m h -> FsPath -> ChecksumsFile -> m ()+writeChecksumsFile fs path checksums =+ withFile fs path (WriteMode MustBeNew) $ \h -> do+ _ <- hPutAll fs h (formatChecksumsFile checksums)+ pure ()++{-# SPECIALISE writeChecksumsFile' :: HasFS IO h -> Handle h -> ChecksumsFile -> IO () #-}+writeChecksumsFile' :: MonadThrow m+ => HasFS m h -> Handle h -> ChecksumsFile -> m ()+writeChecksumsFile' fs h checksums = void $ hPutAll fs h (formatChecksumsFile checksums)++parseChecksumsFile :: BSC.ByteString -> Either String ChecksumsFile+parseChecksumsFile content =+ case partitionEithers (parseLines content) of+ ([], entries) -> Right $! Map.fromList entries+ ((badline:_), _) -> Left $! "could not parse '" <> BSC.unpack badline <> "'"+ where+ parseLines = map (\l -> maybe (Left l) Right (parseChecksumFileLine l))+ . filter (not . BSC.null)+ . BSC.lines++parseChecksumFileLine :: BSC.ByteString -> Maybe (ChecksumsFileName, CRC32C)+parseChecksumFileLine str0 = do+ guard (BSC.take 8 str0 == "CRC32C (")+ let str1 = BSC.drop 8 str0+ let (name, str2) = BSC.break (==')') str1+ guard (BSC.take 4 str2 == ") = ")+ let str3 = BSC.drop 4 str2+ guard (BSC.length str3 == 8 && BSC.all isHexDigit str3)+ let !crc = fromIntegral (hexdigitsToInt str3)+ pure (ChecksumsFileName name, CRC32C crc)++isHexDigit :: Char -> Bool+isHexDigit c = (c >= '0' && c <= '9')+ || (c >= 'a' && c <= 'f') --lower case only++-- Precondition: BSC.all isHexDigit+hexdigitsToInt :: BSC.ByteString -> Word+hexdigitsToInt =+ BSC.foldl' accumdigit 0+ where+ accumdigit :: Word -> Char -> Word+ accumdigit !a !c =+ (a `shiftL` 4) .|. hexdigitToWord c++-- Precondition: isHexDigit+hexdigitToWord :: Char -> Word+hexdigitToWord c+ | let !dec = fromIntegral (ord c - ord '0')+ , dec <= 9 = dec++ | let !hex = fromIntegral (ord c - ord 'a' + 10)+ , otherwise = hex++formatChecksumsFile :: ChecksumsFile -> BSL.ByteString+formatChecksumsFile checksums =+ BS.toLazyByteString $+ mconcat+ [ BS.byteString "CRC32C ("+ <> BS.byteString name+ <> BS.byteString ") = "+ <> BS.word32HexFixed crc+ <> BS.char8 '\n'+ | (ChecksumsFileName name, CRC32C crc) <- Map.toList checksums ]++{-------------------------------------------------------------------------------+ Checksum errors+-------------------------------------------------------------------------------}++-- | Check the CRC32C checksum for a file.+--+-- If the boolean argument is @True@, all file data for this path is evicted+-- from the page cache.+{-# SPECIALISE+ checkCRC ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bool+ -> CRC32C+ -> FsPath+ -> IO ()+ #-}+checkCRC ::+ forall m h.+ (MonadMask m, PrimMonad m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bool+ -> CRC32C+ -> FsPath+ -> m ()+checkCRC fs hbio dropCache expected fp = withFile fs fp ReadMode $ \h -> do+ -- double the file readahead window (only applies to this file descriptor)+ hAdviseAll hbio h AdviceSequential+ !checksum <- hGetAllCRC32C' fs h defaultChunkSize initialCRC32C+ when dropCache $ hDropCacheAll hbio h+ expectChecksum fs fp expected checksum++{-# SPECIALISE+ expectChecksum ::+ HasFS IO h+ -> FsPath+ -> CRC32C+ -> CRC32C+ -> IO ()+ #-}+expectChecksum ::+ MonadThrow m+ => HasFS m h+ -> FsPath+ -> CRC32C+ -> CRC32C+ -> m ()+expectChecksum hfs fp expected checksum =+ when (expected /= checksum) $+ throwIO $+ ErrFileChecksumMismatch+ (mkFsErrorPath hfs fp)+ (unCRC32C expected)+ (unCRC32C checksum)+++{-------------------------------------------------------------------------------+ File Format Errors+-------------------------------------------------------------------------------}++data FileFormat+ = FormatChecksumsFile+ | FormatBloomFilterFile+ | FormatIndexFile+ | FormatWriteBufferFile+ | FormatSnapshotMetaData+ deriving stock (Show, Eq)++-- | The file is corrupted.+data FileCorruptedError+ = -- | The file fails to parse.+ ErrFileFormatInvalid+ -- | File.+ !FsErrorPath+ -- | File format.+ !FileFormat+ -- | Error message.+ !String+ | -- | The file CRC32 checksum is invalid.+ ErrFileChecksumMismatch+ -- | File.+ !FsErrorPath+ -- | Expected checksum.+ !Word32+ -- | Actual checksum.+ !Word32+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE+ expectValidFile ::+ HasFS IO h+ -> FsPath+ -> FileFormat+ -> Either String a+ -> IO a+ #-}+expectValidFile ::+ MonadThrow m+ => HasFS f h+ -> FsPath+ -> FileFormat+ -> Either String a+ -> m a+expectValidFile _hfs _file _format (Right x) =+ pure x+expectValidFile hfs file format (Left msg) =+ throwIO $ ErrFileFormatInvalid (mkFsErrorPath hfs file) format msg+
@@ -0,0 +1,250 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.ChecksumHandle+ (+ -- * Checksum handles+ -- $checksum-handles+ ChecksumHandle (..),+ makeHandle,+ readChecksum,+ dropCache,+ closeHandle,+ writeToHandle,+ -- * Specialised writers+ writeRawPage,+ writeRawOverflowPages,+ writeBlob,+ copyBlob,+ writeFilter,+ writeIndexHeader,+ writeIndexChunk,+ writeIndexFinal,+ ) where++import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadThrow)+import Control.Monad.Primitive+import qualified Data.ByteString.Lazy as BSL+import Data.Primitive.PrimVar+import Data.Word (Word64)+import Database.LSMTree.Internal.BlobRef (BlobSpan (..), RawBlobRef)+import qualified Database.LSMTree.Internal.BlobRef as BlobRef+import Database.LSMTree.Internal.BloomFilter (Bloom, bloomFilterToLBS)+import Database.LSMTree.Internal.Chunk (Chunk)+import qualified Database.LSMTree.Internal.Chunk as Chunk (toByteString)+import Database.LSMTree.Internal.CRC32C (CRC32C)+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Index (Index, IndexType)+import qualified Database.LSMTree.Internal.Index as Index (finalLBS, headerLBS)+import Database.LSMTree.Internal.Paths (ForBlob (..), ForFilter (..),+ ForIndex (..), ForKOps (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)+import qualified Database.LSMTree.Internal.RawOverflowPage as RawOverflowPage+import Database.LSMTree.Internal.RawPage (RawPage)+import qualified Database.LSMTree.Internal.RawPage as RawPage+import Database.LSMTree.Internal.Serialise+import qualified System.FS.API as FS+import System.FS.API+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++{-------------------------------------------------------------------------------+ ChecksumHandle+-------------------------------------------------------------------------------}++{- $checksum-handles+ A handle ('ChecksumHandle') that maintains a running CRC32 checksum.+-}++-- | Tracks the checksum of a (write mode) file handle.+data ChecksumHandle s h = ChecksumHandle !(FS.Handle h) !(PrimVar s CRC32C)++{-# SPECIALISE makeHandle ::+ HasFS IO h+ -> FS.FsPath+ -> IO (ChecksumHandle RealWorld h) #-}+makeHandle ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> FS.FsPath+ -> m (ChecksumHandle (PrimState m) h)+makeHandle fs path =+ ChecksumHandle+ <$> FS.hOpen fs path (FS.WriteMode FS.MustBeNew)+ <*> newPrimVar CRC.initialCRC32C++{-# SPECIALISE readChecksum ::+ ChecksumHandle RealWorld h+ -> IO CRC32C #-}+readChecksum ::+ PrimMonad m+ => ChecksumHandle (PrimState m) h+ -> m CRC32C+readChecksum (ChecksumHandle _h checksum) = readPrimVar checksum++dropCache :: HasBlockIO m h -> ChecksumHandle (PrimState m) h -> m ()+dropCache hbio (ChecksumHandle h _) = FS.hDropCacheAll hbio h++closeHandle :: HasFS m h -> ChecksumHandle (PrimState m) h -> m ()+closeHandle fs (ChecksumHandle h _checksum) = FS.hClose fs h++{-# SPECIALISE writeToHandle ::+ HasFS IO h+ -> ChecksumHandle RealWorld h+ -> BSL.ByteString+ -> IO () #-}+writeToHandle ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ChecksumHandle (PrimState m) h+ -> BSL.ByteString+ -> m ()+writeToHandle fs (ChecksumHandle h checksum) lbs = do+ crc <- readPrimVar checksum+ (_, crc') <- CRC.hPutAllChunksCRC32C fs h lbs crc+ writePrimVar checksum crc'++{-------------------------------------------------------------------------------+ Specialised Writers for ChecksumHandle+-------------------------------------------------------------------------------}++{-# SPECIALISE writeRawPage ::+ HasFS IO h+ -> ForKOps (ChecksumHandle RealWorld h)+ -> RawPage+ -> IO () #-}+writeRawPage ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ForKOps (ChecksumHandle (PrimState m) h)+ -> RawPage+ -> m ()+writeRawPage hfs kOpsHandle =+ writeToHandle hfs (unForKOps kOpsHandle)+ . BSL.fromStrict+ . RB.unsafePinnedToByteString -- 'RawPage' is guaranteed to be pinned+ . RawPage.rawPageRawBytes++{-# SPECIALISE writeRawOverflowPages ::+ HasFS IO h+ -> ForKOps (ChecksumHandle RealWorld h)+ -> [RawOverflowPage]+ -> IO () #-}+writeRawOverflowPages ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ForKOps (ChecksumHandle (PrimState m) h)+ -> [RawOverflowPage]+ -> m ()+writeRawOverflowPages hfs kOpsHandle =+ writeToHandle hfs (unForKOps kOpsHandle)+ . BSL.fromChunks+ . map (RawOverflowPage.rawOverflowPageToByteString)++{-# SPECIALISE writeBlob ::+ HasFS IO h+ -> PrimVar RealWorld Word64+ -> ForBlob (ChecksumHandle RealWorld h)+ -> SerialisedBlob+ -> IO BlobSpan #-}+writeBlob ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> PrimVar (PrimState m) Word64+ -> ForBlob (ChecksumHandle (PrimState m) h)+ -> SerialisedBlob+ -> m BlobSpan+writeBlob hfs blobOffset blobHandle blob = do+ -- NOTE: This is different from BlobFile.writeBlob. This is because BlobFile+ -- internalises a regular Handle, rather than a ChecksumHandle. These two+ -- functions cannot be easily unified, because BlobFile.writeBlob permits+ -- writing blobs to arbitrary positions in the blob file, whereas, by the+ -- very nature of CRC32 checksums, ChecksumHandle.writeBlob only supports+ -- sequential writes.+ let size = sizeofBlob64 blob+ offset <- readPrimVar blobOffset+ modifyPrimVar blobOffset (+size)+ let SerialisedBlob rb = blob+ let lbs = BSL.fromStrict $ RB.toByteString rb+ writeToHandle hfs (unForBlob blobHandle) lbs+ pure (BlobSpan offset (fromIntegral size))++{-# SPECIALISE copyBlob ::+ HasFS IO h+ -> PrimVar RealWorld Word64+ -> ForBlob (ChecksumHandle RealWorld h)+ -> RawBlobRef IO h+ -> IO BlobSpan #-}+copyBlob ::+ (MonadSTM m, MonadThrow m, PrimMonad m)+ => HasFS m h+ -> PrimVar (PrimState m) Word64+ -> ForBlob (ChecksumHandle (PrimState m) h)+ -> RawBlobRef m h+ -> m BlobSpan+copyBlob hfs blobOffset blobHandle blobref = do+ blob <- BlobRef.readRawBlobRef hfs blobref+ writeBlob hfs blobOffset blobHandle blob++{-# SPECIALISE writeFilter ::+ HasFS IO h+ -> ForFilter (ChecksumHandle RealWorld h)+ -> Bloom SerialisedKey+ -> IO () #-}+writeFilter ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ForFilter (ChecksumHandle (PrimState m) h)+ -> Bloom SerialisedKey+ -> m ()+writeFilter hfs filterHandle bf =+ writeToHandle hfs (unForFilter filterHandle) (bloomFilterToLBS bf)++{-# SPECIALISE writeIndexHeader ::+ HasFS IO h+ -> ForIndex (ChecksumHandle RealWorld h)+ -> IndexType+ -> IO () #-}+writeIndexHeader ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ForIndex (ChecksumHandle (PrimState m) h)+ -> IndexType+ -> m ()+writeIndexHeader hfs indexHandle indexType =+ writeToHandle hfs (unForIndex indexHandle) $+ Index.headerLBS indexType++{-# SPECIALISE writeIndexChunk ::+ HasFS IO h+ -> ForIndex (ChecksumHandle RealWorld h)+ -> Chunk+ -> IO () #-}+writeIndexChunk ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ForIndex (ChecksumHandle (PrimState m) h)+ -> Chunk+ -> m ()+writeIndexChunk hfs indexHandle chunk =+ writeToHandle hfs (unForIndex indexHandle) $+ BSL.fromStrict $ Chunk.toByteString chunk++{-# SPECIALISE writeIndexFinal ::+ HasFS IO h+ -> ForIndex (ChecksumHandle RealWorld h)+ -> NumEntries+ -> Index+ -> IO () #-}+writeIndexFinal ::+ (MonadSTM m, PrimMonad m)+ => HasFS m h+ -> ForIndex (ChecksumHandle (PrimState m) h)+ -> NumEntries+ -> Index+ -> m ()+writeIndexFinal hfs indexHandle numEntries index =+ writeToHandle hfs (unForIndex indexHandle) $+ Index.finalLBS numEntries index
@@ -0,0 +1,133 @@+{-# OPTIONS_HADDOCK not-home #-}+{- HLINT ignore "Avoid restricted alias" -}++-- | Chunks of bytes, typically output during incremental index serialisation.+module Database.LSMTree.Internal.Chunk+(+ -- * Chunks+ Chunk (Chunk),+ toByteVector,+ toByteString,++ -- * Balers+ Baler (Baler),+ createBaler,+ feedBaler,+ unsafeEndBaler+)+where++import Prelude hiding (length)++import Control.Exception (assert)+import Control.Monad.ST.Strict (ST)+import Data.ByteString (ByteString)+import Data.List (scanl')+import Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,+ writePrimVar)+import Data.Vector.Primitive (Vector (Vector), length, unsafeCopy,+ unsafeFreeze)+import Data.Vector.Primitive.Mutable (MVector)+import qualified Data.Vector.Primitive.Mutable as Mutable (drop, length, slice,+ take, unsafeCopy, unsafeNew)+import Data.Word (Word8)+import Database.LSMTree.Internal.ByteString (byteArrayToByteString)++-- * Chunks++-- | A chunk of bytes, typically output during incremental index serialisation.+newtype Chunk = Chunk (Vector Word8) deriving stock (Eq, Show)++-- | Yields the contents of a chunk as a byte vector.+toByteVector :: Chunk -> Vector Word8+toByteVector (Chunk byteVector) = byteVector++-- | Yields the contents of a chunk as a (strict) byte string.+toByteString :: Chunk -> ByteString+toByteString (Chunk (Vector vecOffset vecLength byteArray))+ = byteArrayToByteString vecOffset vecLength byteArray++-- * Balers++{-|+ An object that receives blocks of bytes and repackages them into chunks such+ that all chunks except for a possible remnant chunk at the end are of at+ least a given minimum size.+-}+data Baler s = Baler+ !(MVector s Word8) -- Buffer storing queued bytes+ !(PrimVar s Int) -- Reference to the number of queued bytes++-- | Creates a new baler.+createBaler :: Int -- ^ Minimum chunk size in bytes+ -> ST s (Baler s) -- ^ Creation of the baler+createBaler minChunkSize = assert (minChunkSize > 0) $+ Baler <$>+ Mutable.unsafeNew (pred minChunkSize) <*>+ newPrimVar 0++{-|+ Feeds a baler blocks of bytes.++ Bytes received by a baler are generally queued for later output, but if+ feeding new bytes makes the accumulated content exceed the minimum chunk+ size then a chunk containing all the accumulated content is output.+-}+feedBaler :: forall s . [Vector Word8] -> Baler s -> ST s (Maybe Chunk)+feedBaler blocks (Baler buffer remnantSizeRef) = do+ remnantSize <- readPrimVar remnantSizeRef+ let++ inputSize :: Int+ !inputSize = sum (map length blocks)++ totalSize :: Int+ !totalSize = remnantSize + inputSize++ if totalSize <= Mutable.length buffer+ then do+ unsafeCopyBlocks (Mutable.drop remnantSize buffer)+ writePrimVar remnantSizeRef totalSize+ pure Nothing+ else do+ protoChunk <- Mutable.unsafeNew totalSize+ Mutable.unsafeCopy (Mutable.take remnantSize protoChunk)+ (Mutable.take remnantSize buffer)+ unsafeCopyBlocks (Mutable.drop remnantSize protoChunk)+ writePrimVar remnantSizeRef 0+ chunk <- Chunk <$> unsafeFreeze protoChunk+ pure (Just chunk)+ where++ unsafeCopyBlocks :: MVector s Word8 -> ST s ()+ unsafeCopyBlocks vec+ = sequence_ $+ zipWith3 (\ start size block -> unsafeCopy+ (Mutable.slice start size vec)+ block)+ (scanl' (+) 0 blockSizes)+ blockSizes+ blocks+ where++ blockSizes :: [Int]+ blockSizes = map length blocks++{-|+ Returns the bytes still queued in a baler, if any, thereby invalidating the+ baler. Executing @unsafeEndBaler baler@ is only safe when @baler@ is not+ used afterwards.+-}+unsafeEndBaler :: forall s . Baler s -> ST s (Maybe Chunk)+unsafeEndBaler (Baler buffer remnantSizeRef) = do+ remnantSize <- readPrimVar remnantSizeRef+ if remnantSize == 0+ then pure Nothing+ else do+ let++ protoChunk :: MVector s Word8+ !protoChunk = Mutable.take remnantSize buffer++ chunk <- Chunk <$> unsafeFreeze protoChunk+ pure (Just chunk)
@@ -0,0 +1,443 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Config (+ LevelNo (..)+ -- * Table configuration+ , TableConfig (..)+ , defaultTableConfig+ , RunLevelNo (..)+ , runParamsForLevel+ -- * Merge policy+ , MergePolicy (..)+ -- * Size ratio+ , SizeRatio (..)+ , sizeRatioInt+ -- * Write buffer allocation+ , WriteBufferAlloc (..)+ -- * Bloom filter allocation+ , BloomFilterAlloc (..)+ , bloomFilterAllocForLevel+ -- * Fence pointer index+ , FencePointerIndexType (..)+ , indexTypeForRun+ -- * Disk cache policy+ , DiskCachePolicy (..)+ , diskCachePolicyForLevel+ -- * Merge schedule+ , MergeSchedule (..)+ -- * Merge batch size+ , MergeBatchSize (..)+ , creditThresholdForLevel+ ) where++import Control.DeepSeq (NFData (..))+import Database.LSMTree.Internal.Index (IndexType)+import qualified Database.LSMTree.Internal.Index as Index+ (IndexType (Compact, Ordinary))+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.Run (RunDataCaching (..))+import Database.LSMTree.Internal.RunAcc (RunBloomFilterAlloc (..))+import Database.LSMTree.Internal.RunBuilder (RunParams (..))++newtype LevelNo = LevelNo Int+ deriving stock (Show, Eq, Ord)+ deriving newtype (Enum, NFData)++{-------------------------------------------------------------------------------+ Table configuration+-------------------------------------------------------------------------------}++{- |+A collection of configuration parameters for tables, which can be used to tune the performance of the table.+To construct a 'TableConfig', modify the 'defaultTableConfig', which defines reasonable defaults for all parameters.++For a detailed discussion of fine-tuning the table configuration, see [Fine-tuning Table Configuration](../#fine_tuning).++[@confMergePolicy :: t'MergePolicy'@]+ The /merge policy/ balances the performance of lookups against the performance of updates.+ Levelling favours lookups.+ Tiering favours updates.+ Lazy levelling strikes a middle ground between levelling and tiering, and moderately favours updates.+ This parameter is explicitly referenced in the documentation of those operations it affects.++[@confSizeRatio :: t'SizeRatio'@]+ The /size ratio/ pushes the effects of the merge policy to the extreme.+ If the size ratio is higher, levelling favours lookups more, and tiering and lazy levelling favour updates more.+ This parameter is referred to as \(T\) in the disk I\/O cost of operations.++[@confWriteBufferAlloc :: t'WriteBufferAlloc'@]+ The /write buffer capacity/ balances the performance of lookups and updates against the in-memory size of the database.+ If the write buffer is larger, it takes up more memory, but lookups and updates are more efficient.+ This parameter is referred to as \(B\) in the disk I\/O cost of operations.+ Irrespective of this parameter, the write buffer size cannot exceed 4GiB.++[@confMergeSchedule :: t'MergeSchedule'@]+ The /merge schedule/ balances the performance of lookups and updates against the consistency of updates.+ With the one-shot merge schedule, lookups and updates are more efficient overall, but some updates may take much longer than others.+ With the incremental merge schedule, lookups and updates are less efficient overall, but each update does a similar amount of work.+ This parameter is explicitly referenced in the documentation of those operations it affects.+ The merge schedule does not affect the way that table unions are computed.+ However, any table union must complete all outstanding incremental updates.++[@confBloomFilterAlloc :: t'BloomFilterAlloc'@]+ The Bloom filter size balances the performance of lookups against the in-memory size of the database.+ If the Bloom filters are larger, they take up more memory, but lookup operations are more efficient.++[@confFencePointerIndex :: t'FencePointerIndexType'@]+ The /fence-pointer index type/ supports two types of indexes.+ The /ordinary/ indexes are designed to work with any key.+ The /compact/ indexes are optimised for the case where the keys in the database are uniformly distributed, e.g., when the keys are hashes.++[@confDiskCachePolicy :: t'DiskCachePolicy'@]+ The /disk cache policy/ supports caching lookup operations using the OS page cache.+ Caching may improve the performance of lookups and updates if database access follows certain patterns.++[@confMergeBatchSize :: t'MergeBatchSize'@]+ The merge batch size balances the maximum latency of individual update+ operations, versus the latency of a sequence of update operations. Bigger+ batches improves overall performance but some updates will take a lot+ longer than others. The default is to use a large batch size.+-}+data TableConfig = TableConfig {+ confMergePolicy :: !MergePolicy+ , confMergeSchedule :: !MergeSchedule+ , confSizeRatio :: !SizeRatio+ , confWriteBufferAlloc :: !WriteBufferAlloc+ , confBloomFilterAlloc :: !BloomFilterAlloc+ , confFencePointerIndex :: !FencePointerIndexType+ , confDiskCachePolicy :: !DiskCachePolicy+ , confMergeBatchSize :: !MergeBatchSize+ }+ deriving stock (Show, Eq)++instance NFData TableConfig where+ rnf (TableConfig a b c d e f g h) =+ rnf a `seq` rnf b `seq` rnf c `seq` rnf d `seq`+ rnf e `seq` rnf f `seq` rnf g `seq` rnf h++-- | The 'defaultTableConfig' defines reasonable defaults for all 'TableConfig' parameters.+--+-- >>> confMergePolicy defaultTableConfig+-- LazyLevelling+-- >>> confMergeSchedule defaultTableConfig+-- Incremental+-- >>> confSizeRatio defaultTableConfig+-- Four+-- >>> confWriteBufferAlloc defaultTableConfig+-- AllocNumEntries 20000+-- >>> confBloomFilterAlloc defaultTableConfig+-- AllocRequestFPR 1.0e-3+-- >>> confFencePointerIndex defaultTableConfig+-- OrdinaryIndex+-- >>> confDiskCachePolicy defaultTableConfig+-- DiskCacheAll+-- >>> confMergeBatchSize defaultTableConfig+-- MergeBatchSize 20000+--+defaultTableConfig :: TableConfig+defaultTableConfig =+ TableConfig+ { confMergePolicy = LazyLevelling+ , confMergeSchedule = Incremental+ , confSizeRatio = Four+ , confWriteBufferAlloc = AllocNumEntries 20_000+ , confBloomFilterAlloc = AllocRequestFPR 1.0e-3+ , confFencePointerIndex = OrdinaryIndex+ , confDiskCachePolicy = DiskCacheAll+ , confMergeBatchSize = MergeBatchSize 20_000 -- same as write buffer+ }++data RunLevelNo = RegularLevel LevelNo | UnionLevel++runParamsForLevel :: TableConfig -> RunLevelNo -> RunParams+runParamsForLevel conf@TableConfig {..} levelNo =+ RunParams+ { runParamCaching = diskCachePolicyForLevel confDiskCachePolicy levelNo+ , runParamAlloc = bloomFilterAllocForLevel conf levelNo+ , runParamIndex = indexTypeForRun confFencePointerIndex+ }++{-------------------------------------------------------------------------------+ Merge policy+-------------------------------------------------------------------------------}++{- |+The /merge policy/ balances the performance of lookups against the performance of updates.+Levelling favours lookups.+Tiering favours updates.+Lazy levelling strikes a middle ground between levelling and tiering, and moderately favours updates.+This parameter is explicitly referenced in the documentation of those operations it affects.++__NOTE:__ This package only supports lazy levelling.++For a detailed discussion of the merge policy, see [Fine-tuning: Merge Policy, Size Ratio, and Write Buffer Size](../#fine_tuning_data_layout).+-}+data MergePolicy =+ LazyLevelling+ deriving stock (Eq, Show)++instance NFData MergePolicy where+ rnf LazyLevelling = ()++{-------------------------------------------------------------------------------+ Size ratio+-------------------------------------------------------------------------------}++{- |+The /size ratio/ pushes the effects of the merge policy to the extreme.+If the size ratio is higher, levelling favours lookups more, and tiering and lazy levelling favour updates more.+This parameter is referred to as \(T\) in the disk I\/O cost of operations.++__NOTE:__ This package only supports a size ratio of four.++For a detailed discussion of the size ratio, see [Fine-tuning: Merge Policy, Size Ratio, and Write Buffer Size](../#fine_tuning_data_layout).+-}+data SizeRatio = Four+ deriving stock (Eq, Show)++instance NFData SizeRatio where+ rnf Four = ()++sizeRatioInt :: SizeRatio -> Int+sizeRatioInt = \case Four -> 4++{-------------------------------------------------------------------------------+ Write buffer allocation+-------------------------------------------------------------------------------}++-- TODO: "If the sizes of values vary greatly, this can lead to unevenly sized runs on disk and unpredictable performance."++{- |+The /write buffer capacity/ balances the performance of lookups and updates against the in-memory size of the table.+If the write buffer is larger, it takes up more memory, but lookups and updates are more efficient.+Irrespective of this parameter, the write buffer size cannot exceed 4GiB.++For a detailed discussion of the size ratio, see [Fine-tuning: Merge Policy, Size Ratio, and Write Buffer Size](../#fine_tuning_data_layout).+-}+data WriteBufferAlloc =+ {- |+ Allocate space for the in-memory write buffer to fit the requested number of entries.+ This parameter is referred to as \(B\) in the disk I\/O cost of operations.+ -}+ AllocNumEntries !Int+ deriving stock (Show, Eq)++instance NFData WriteBufferAlloc where+ rnf (AllocNumEntries n) = rnf n++{-------------------------------------------------------------------------------+ Merge schedule+-------------------------------------------------------------------------------}++{- |+The /merge schedule/ balances the performance of lookups and updates against the consistency of updates.+The merge schedule does not affect the performance of table unions.+With the one-shot merge schedule, lookups and updates are more efficient overall, but some updates may take much longer than others.+With the incremental merge schedule, lookups and updates are less efficient overall, but each update does a similar amount of work.+This parameter is explicitly referenced in the documentation of those operations it affects.++For a detailed discussion of the effect of the merge schedule, see [Fine-tuning: Merge Schedule](../#fine_tuning_merge_schedule).+-}+data MergeSchedule =+ {- |+ The 'OneShot' merge schedule causes the merging algorithm to complete merges immediately.+ This is more efficient than the 'Incremental' merge schedule, but has an inconsistent workload.+ Using the 'OneShot' merge schedule, the worst-case disk I\/O complexity of the update operations is /linear/ in the size of the table.+ For real-time systems and other use cases where unresponsiveness is unacceptable, use the 'Incremental' merge schedule.+ -}+ OneShot+ {- |+ The 'Incremental' merge schedule spreads out the merging work over time.+ This is less efficient than the 'OneShot' merge schedule, but has a consistent workload.+ Using the 'Incremental' merge schedule, the worst-case disk I\/O complexity of the update operations is /logarithmic/ in the size of the table.+ This 'Incremental' merge schedule still uses batching to improve performance.+ The batch size can be controlled using the 'MergeBatchSize'.+ -}+ | Incremental+ deriving stock (Eq, Show)++instance NFData MergeSchedule where+ rnf OneShot = ()+ rnf Incremental = ()++{-------------------------------------------------------------------------------+ Bloom filter allocation+-------------------------------------------------------------------------------}++{- |+The Bloom filter size balances the performance of lookups against the in-memory size of the table.+If the Bloom filters are larger, they take up more memory, but lookup operations are more efficient.++For a detailed discussion of the Bloom filter size, see [Fine-tuning: Bloom Filter Size](../#fine_tuning_bloom_filter_size).+-}+data BloomFilterAlloc =+ {- |+ Allocate the requested number of bits per entry in the table.++ The value must strictly positive, but fractional values are permitted.+ The recommended range is \([2, 24]\).+ -}+ AllocFixed !Double+ | {- |+ Allocate the required number of bits per entry to get the requested false-positive rate.++ The value must be in the range \((0, 1)\).+ The recommended range is \([1\mathrm{e}{ -5 },1\mathrm{e}{ -2 }]\).+ -}+ AllocRequestFPR !Double+ deriving stock (Show, Eq)++instance NFData BloomFilterAlloc where+ rnf (AllocFixed n) = rnf n+ rnf (AllocRequestFPR fpr) = rnf fpr++bloomFilterAllocForLevel :: TableConfig -> RunLevelNo -> RunBloomFilterAlloc+bloomFilterAllocForLevel conf _levelNo =+ case confBloomFilterAlloc conf of+ AllocFixed n -> RunAllocFixed n+ AllocRequestFPR fpr -> RunAllocRequestFPR fpr++{-------------------------------------------------------------------------------+ Fence pointer index+-------------------------------------------------------------------------------}++{- |+The /fence-pointer index type/ supports two types of indexes.+The /ordinary/ indexes are designed to work with any key.+The /compact/ indexes are optimised for the case where the keys in the database are uniformly distributed, e.g., when the keys are hashes.++For a detailed discussion the fence-pointer index types, see [Fine-tuning: Fence-Pointer Index Type](../#fine_tuning_fence_pointer_index_type).+-}+data FencePointerIndexType =+ {- |+ Ordinary indexes are designed to work with any key.++ When using an ordinary index, the 'Database.LSMTree.Internal.Serialise.Class.serialiseKey' function cannot produce output larger than 64KiB.+ -}+ OrdinaryIndex+ | {- |+ Compact indexes are designed for the case where the keys in the database are uniformly distributed, e.g., when the keys are hashes.++ When using a compact index, some requirements apply to serialised keys:++ * keys must be uniformly distributed;+ * keys can be of variable length;+ * keys less than 8 bytes (64bits) are padded with zeros (in LSB position).+ -}+ CompactIndex+ deriving stock (Eq, Show)++instance NFData FencePointerIndexType where+ rnf CompactIndex = ()+ rnf OrdinaryIndex = ()++indexTypeForRun :: FencePointerIndexType -> IndexType+indexTypeForRun CompactIndex = Index.Compact+indexTypeForRun OrdinaryIndex = Index.Ordinary++{-------------------------------------------------------------------------------+ Disk cache policy+-------------------------------------------------------------------------------}++{- |+The /disk cache policy/ determines if lookup operations use the OS page cache.+Caching may improve the performance of lookups if database access follows certain patterns.++For a detailed discussion the disk cache policy, see [Fine-tuning: Disk Cache Policy](../#fine_tuning_disk_cache_policy).+-}+data DiskCachePolicy =+ {- |+ Cache all data in the table.++ Use this policy if the database's access pattern has either good spatial locality or both good spatial and temporal locality.+ -}+ DiskCacheAll++ | {- |+ Cache the data in the freshest @l@ levels.++ Use this policy if the database's access pattern only has good temporal locality.++ The variable @l@ determines the number of levels that are cached.+ For a description of levels, see [Merge Policy, Size Ratio, and Write Buffer Size](#fine_tuning_data_layout).+ With this setting, the database can be expected to cache up to \(\frac{k}{P}\) pages of memory,+ where \(k\) refers to the number of entries that fit in levels \([1,l]\) and is defined as \(\sum_{i=1}^{l}BT^{i}\).+ -}+ -- TODO: Add a policy for caching based on size in bytes, rather than exposing internal details such as levels.+ -- For instance, a policy that states "cache the freshest 10MiB"+ DiskCacheLevelOneTo !Int++ | {- |+ Do not cache any table data.++ Use this policy if the database's access pattern has does not have good spatial or temporal locality.+ For instance, if the access pattern is uniformly random.+ -}+ DiskCacheNone+ deriving stock (Show, Eq)++instance NFData DiskCachePolicy where+ rnf DiskCacheAll = ()+ rnf (DiskCacheLevelOneTo l) = rnf l+ rnf DiskCacheNone = ()++-- | Interpret the 'DiskCachePolicy' for a level: should we cache data in runs+-- at this level.+--+diskCachePolicyForLevel :: DiskCachePolicy -> RunLevelNo -> RunDataCaching+diskCachePolicyForLevel policy levelNo =+ case policy of+ DiskCacheAll -> CacheRunData+ DiskCacheNone -> NoCacheRunData+ DiskCacheLevelOneTo n ->+ case levelNo of+ RegularLevel l | l <= LevelNo n -> CacheRunData+ | otherwise -> NoCacheRunData+ UnionLevel -> NoCacheRunData++{-------------------------------------------------------------------------------+ Merge batch size+-------------------------------------------------------------------------------}++{- |+The /merge batch size/ is a micro-tuning parameter, and in most cases you do+need to think about it and can leave it at its default.++When using the 'Incremental' merge schedule, merging is done in batches. This+is a trade-off: larger batches tends to mean better overall performance but the+downside is that while most updates (inserts, deletes, upserts) are fast, some+are slower (when a batch of merging work has to be done).++If you care most about the maximum latency of updates, then use a small batch+size. If you don't care about latency of individual operations, just the+latency of the overall sequence of operations then use a large batch size. The+default is to use a large batch size, the same size as the write buffer itself.+The minimum batch size is 1. The maximum batch size is the size of the write+buffer 'confWriteBufferAlloc'.++Note that the actual batch size is the minimum of this configuration+parameter and the size of the batch of operations performed (e.g. 'inserts').+So if you consistently use large batches, you can use a batch size of 1 and+the merge batch size will always be determined by the operation batch size.++A further reason why it may be preferable to use minimal batch sizes is to get+good parallel work balance, when using parallelism.+-}+newtype MergeBatchSize = MergeBatchSize Int+ deriving stock (Show, Eq, Ord)+ deriving newtype (NFData)++-- TODO: the thresholds for doing merge work should be different for each level,+-- and ideally all-pairs co-prime.+creditThresholdForLevel :: TableConfig -> LevelNo -> MR.CreditThreshold+creditThresholdForLevel TableConfig {+ confMergeBatchSize = MergeBatchSize mergeBatchSz,+ confWriteBufferAlloc = AllocNumEntries writeBufferSz+ }+ (LevelNo _i) =+ MR.CreditThreshold+ . MR.UnspentCredits+ . MR.MergeCredits+ . max 1+ . min writeBufferSz+ $ mergeBatchSz
@@ -0,0 +1,205 @@+{-# LANGUAGE NamedFieldPuns #-}+{-# OPTIONS_HADDOCK not-home #-}++-- Definitions for override table config options.+module Database.LSMTree.Internal.Config.Override (+ -- $override-policy++ -- * Override table config+ TableConfigOverride (..)+ , noTableConfigOverride+ , overrideTableConfig+ ) where++import qualified Data.Vector as V+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.MergeSchedule (MergePolicyForLevel,+ NominalCredits, NominalDebt)+import Database.LSMTree.Internal.MergingRun (LevelMergeType,+ MergeCredits, MergeDebt, TreeMergeType)+import Database.LSMTree.Internal.Run+import Database.LSMTree.Internal.RunAcc (RunBloomFilterAlloc)+import Database.LSMTree.Internal.RunNumber (RunNumber)+import Database.LSMTree.Internal.Snapshot++-- $override-policy+--+-- === Limitations+--+-- Overriding config options should in many cases be possible as long as there+-- is a mitigation strategy to ensure that a change in options does not cause+-- the database state of a table to become inconsistent. But what does this+-- strategy look like? And if we allow a config option to be overridden on a+-- live table (one that a user has access to), how should we apply these new+-- options to shared data like merging runs? Moreover, would we answer these+-- questions differently for each type of config option?+--+-- For now, it seems to be the most straightforward to limit the config options+-- we allow to be overridden, and that we only change the config options+-- offline. That is, we override the config option just before opening a+-- snapshot from disk. At that point, there is no sharing because the table is+-- not live yet, which simplifies how changing a config option is handled.+--+-- Another complicating factor is that we have thought about the possibility of+-- restoring sharing of ongoing merges between live tables and newly opened+-- snapshots. At that point, we run into the same challenges again... But for+-- now, changing only the disk cache policy and merge batch size offline should+-- work fine.++{-------------------------------------------------------------------------------+ Helper class+-------------------------------------------------------------------------------}++-- | This class is only here so that we can recursively call 'override' on all+-- fields of a datatype, instead of having to invent a new name for each type+-- that the function is called on such as @overrideTableConfig@,+-- @overrideSnapshotRun@, etc.+class Override o a where+ override :: o -> a -> a++instance Override a c => Override (Maybe a) c where+ override = maybe id override++{-------------------------------------------------------------------------------+ Override table config+-------------------------------------------------------------------------------}++{- |+The 'TableConfigOverride' can be used to override the 'TableConfig'+when opening a table from a snapshot.+-}+data TableConfigOverride = TableConfigOverride {+ overrideDiskCachePolicy :: Maybe DiskCachePolicy,+ overrideMergeBatchSize :: Maybe MergeBatchSize+ }+ deriving stock (Show, Eq)++-- | No override of the 'TableConfig'. You can use this as a default value and+-- record update to override some parameters, while being future-proof to new+-- parameters, e.g.+--+-- > noTableConfigOverride { overrideDiskCachePolicy = DiskCacheNone }+--+noTableConfigOverride :: TableConfigOverride+noTableConfigOverride = TableConfigOverride Nothing Nothing++-- | Override the a subset of the table configuration parameters that are+-- stored in snapshot metadata.+--+-- Tables opened from the new 'SnapshotMetaData' will use the new value for the+-- table configuration.+overrideTableConfig :: TableConfigOverride+ -> SnapshotMetaData -> SnapshotMetaData+overrideTableConfig = override++instance Override TableConfigOverride SnapshotMetaData where+ override TableConfigOverride {..} =+ override overrideMergeBatchSize+ . override overrideDiskCachePolicy++{-------------------------------------------------------------------------------+ Override merge batch size+-------------------------------------------------------------------------------}++instance Override MergeBatchSize SnapshotMetaData where+ override mbs smd =+ smd { snapMetaConfig = override mbs (snapMetaConfig smd) }++instance Override MergeBatchSize TableConfig where+ override confMergeBatchSize' tc =+ tc { confMergeBatchSize = confMergeBatchSize' }++{-------------------------------------------------------------------------------+ Override disk cache policy+-------------------------------------------------------------------------------}++-- NOTE: the instances below explicitly pattern match on the types of+-- constructor fields. This makes the code more verbose, but it also makes the+-- code a little more future proof. It should help us not to forget to update+-- the instances when new fields are added or existing fields change. In+-- particular, if anything changes about the constructor or its fields (and+-- their types), then we will see a compiler error, and then we are forced to+-- look at the code and make adjustments.++instance Override DiskCachePolicy SnapshotMetaData where+ override dcp+ (SnapshotMetaData (sl :: SnapshotLabel)+ (tc :: TableConfig) (rn :: RunNumber)+ (sls :: (SnapLevels SnapshotRun))+ (smt :: (Maybe (SnapMergingTree SnapshotRun))))+ = SnapshotMetaData sl (override dcp tc) rn (override dcp sls) $+ let rdc = diskCachePolicyForLevel dcp UnionLevel+ in fmap (override rdc) smt++instance Override DiskCachePolicy TableConfig where+ override confDiskCachePolicy' tc =+ tc { confDiskCachePolicy = confDiskCachePolicy' }++instance Override DiskCachePolicy (SnapLevels SnapshotRun) where+ override dcp (SnapLevels (vec :: V.Vector (SnapLevel SnapshotRun))) =+ SnapLevels $ V.imap go vec+ where+ go (LevelNo . (+1) -> ln) (x :: SnapLevel SnapshotRun) =+ let rdc = diskCachePolicyForLevel dcp (RegularLevel ln)+ in override rdc x++instance Override RunDataCaching (SnapLevel SnapshotRun) where+ override rdc+ (SnapLevel (sir :: SnapIncomingRun SnapshotRun) (srs :: V.Vector SnapshotRun))+ = SnapLevel (override rdc sir) (V.map (override rdc) srs)++instance Override RunDataCaching (SnapIncomingRun SnapshotRun) where+ override rdc = \case+ SnapIncomingSingleRun (sr :: SnapshotRun) ->+ SnapIncomingSingleRun $ override rdc sr+ SnapIncomingMergingRun+ (mpfl :: MergePolicyForLevel) (nd :: NominalDebt)+ (nc :: NominalCredits) (smr :: SnapMergingRun LevelMergeType SnapshotRun) ->+ SnapIncomingMergingRun mpfl nd nc (override rdc smr)++instance Override RunDataCaching (SnapMergingRun t SnapshotRun) where+ override rdc = \case+ SnapCompletedMerge (md :: MergeDebt) (sr :: SnapshotRun) ->+ SnapCompletedMerge md (override rdc sr)+ SnapOngoingMerge+ (rp :: RunParams) (mc :: MergeCredits)+ (srs :: V.Vector SnapshotRun) (t :: t) ->+ SnapOngoingMerge (override rdc rp) mc (V.map (override rdc) srs) t++instance Override RunDataCaching RunParams where+ override rdc+ (RunParams (_rdc :: RunDataCaching) (rbfa :: RunBloomFilterAlloc) (it :: IndexType))+ = RunParams rdc rbfa it++instance Override RunDataCaching SnapshotRun where+ override rdc+ (SnapshotRun (rn :: RunNumber) (_rdc :: RunDataCaching) (it ::IndexType))+ = SnapshotRun rn rdc it++instance Override RunDataCaching (SnapMergingTree SnapshotRun) where+ override rdc (SnapMergingTree (smts :: SnapMergingTreeState SnapshotRun))+ = SnapMergingTree (override rdc smts)++instance Override RunDataCaching (SnapMergingTreeState SnapshotRun) where+ override rdc = \case+ SnapCompletedTreeMerge (sr :: SnapshotRun) ->+ SnapCompletedTreeMerge (override rdc sr)+ SnapPendingTreeMerge (spm :: SnapPendingMerge SnapshotRun) ->+ SnapPendingTreeMerge (override rdc spm)+ SnapOngoingTreeMerge (smr :: SnapMergingRun TreeMergeType SnapshotRun) ->+ SnapOngoingTreeMerge (override rdc smr)++instance Override RunDataCaching (SnapPendingMerge SnapshotRun) where+ override rdc = \case+ SnapPendingLevelMerge+ (spers :: [SnapPreExistingRun SnapshotRun])+ (msmt :: Maybe (SnapMergingTree SnapshotRun)) ->+ SnapPendingLevelMerge (fmap (override rdc) spers) (fmap (override rdc) msmt)+ SnapPendingUnionMerge (smts :: [SnapMergingTree SnapshotRun]) ->+ SnapPendingUnionMerge (fmap (override rdc) smts)++instance Override RunDataCaching (SnapPreExistingRun SnapshotRun) where+ override rdc = \case+ SnapPreExistingRun (sr :: SnapshotRun) -> SnapPreExistingRun (override rdc sr)+ SnapPreExistingMergingRun (smr :: SnapMergingRun LevelMergeType SnapshotRun) ->+ SnapPreExistingMergingRun (override rdc smr)
@@ -0,0 +1,131 @@+{-# LANGUAGE DataKinds #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Cursor (+ readEntriesWhile+ ) where++import Control.Concurrent.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadST (MonadST (..))+import Control.Monad.Class.MonadThrow+import qualified Data.Vector as V+import Database.LSMTree.Internal.BlobRef (RawBlobRef,+ WeakBlobRef (..))+import qualified Database.LSMTree.Internal.BlobRef as BlobRef+import Database.LSMTree.Internal.Entry (Entry)+import qualified Database.LSMTree.Internal.Entry as Entry+import qualified Database.LSMTree.Internal.Readers as Readers+import qualified Database.LSMTree.Internal.RunReader as Reader+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue,+ SerialisedKey, SerialisedValue)+import qualified Database.LSMTree.Internal.Vector as V++{-# INLINE readEntriesWhile #-}+{-# SPECIALISE readEntriesWhile :: forall h res.+ ResolveSerialisedValue+ -> (SerialisedKey -> Bool)+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef IO h) -> res)+ -> Readers.Readers IO h+ -> Int+ -> IO (V.Vector res, Readers.HasMore) #-}+-- | General notes on the code below:+-- * it is quite similar to the one in Internal.Merge, but different enough+-- that it's probably easier to keep them separate+-- * any function that doesn't take a 'hasMore' argument assumes that the+-- readers have not been drained yet, so we must check before calling them+-- * there is probably opportunity for optimisations+readEntriesWhile :: forall h m res.+ (MonadMask m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> (SerialisedKey -> Bool)+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef m h) -> res)+ -> Readers.Readers m h+ -> Int+ -> m (V.Vector res, Readers.HasMore)+readEntriesWhile resolve keyIsWanted fromEntry readers n =+ flip (V.unfoldrNM' n) Readers.HasMore $ \case+ Readers.Drained -> pure (Nothing, Readers.Drained)+ Readers.HasMore -> readEntryIfWanted+ where+ -- Produces a result unless the readers have been drained or 'keyIsWanted'+ -- returned False.+ readEntryIfWanted :: m (Maybe res, Readers.HasMore)+ readEntryIfWanted = do+ key <- Readers.peekKey readers+ if keyIsWanted key then readEntry+ else pure (Nothing, Readers.HasMore)++ readEntry :: m (Maybe res, Readers.HasMore)+ readEntry = do+ (key, readerEntry, hasMore) <- Readers.pop resolve readers+ let !entry = Reader.toFullEntry readerEntry+ case hasMore of+ Readers.Drained -> do+ handleResolved key entry Readers.Drained+ Readers.HasMore -> do+ case entry of+ Entry.Upsert v ->+ handleMupdate key v+ _ -> do+ -- Anything but Upsert supersedes all previous entries of+ -- the same key, so we can simply drop them and are done.+ hasMore' <- dropRemaining key+ handleResolved key entry hasMore'++ dropRemaining :: SerialisedKey -> m Readers.HasMore+ dropRemaining key = do+ (_, hasMore) <- Readers.dropWhileKey resolve readers key+ pure hasMore++ -- Resolve a 'Mupsert' value with the other entries of the same key.+ handleMupdate :: SerialisedKey+ -> SerialisedValue+ -> m (Maybe res, Readers.HasMore)+ handleMupdate key v = do+ nextKey <- Readers.peekKey readers+ if nextKey /= key+ then+ -- No more entries for same key, done.+ handleResolved key (Entry.Upsert v) Readers.HasMore+ else do+ (_, nextEntry, hasMore) <- Readers.pop resolve readers+ let resolved = Entry.combine resolve (Entry.Upsert v)+ (Reader.toFullEntry nextEntry)+ case hasMore of+ Readers.HasMore -> case resolved of+ Entry.Upsert v' ->+ -- Still a mupsert, keep resolving!+ handleMupdate key v'+ _ -> do+ -- Done with this key, remaining entries are obsolete.+ hasMore' <- dropRemaining key+ handleResolved key resolved hasMore'+ Readers.Drained -> do+ handleResolved key resolved Readers.Drained++ -- Once we have a resolved entry, we still have to make sure it's not+ -- a 'Delete', since we only want to write values to the result vector.+ handleResolved :: SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> Readers.HasMore+ -> m (Maybe res, Readers.HasMore)+ handleResolved key entry hasMore =+ case toResult key entry of+ Just !res ->+ -- Found one resolved value, done.+ pure (Just res, hasMore)+ Nothing ->+ -- Resolved value was a Delete, which we don't want to include.+ -- So look for another one (unless there are no more entries!).+ case hasMore of+ Readers.HasMore -> readEntryIfWanted+ Readers.Drained -> pure (Nothing, Readers.Drained)++ toResult :: SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> Maybe res+ toResult key = \case+ Entry.Insert v -> Just $ fromEntry key v Nothing+ Entry.InsertWithBlob v b -> Just $ fromEntry key v (Just (BlobRef.rawToWeakBlobRef b))+ Entry.Upsert v -> Just $ fromEntry key v Nothing+ Entry.Delete -> Nothing
@@ -0,0 +1,138 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Entry (+ Entry (..)+ , hasBlob+ , onValue+ , onBlobRef+ , NumEntries (..)+ , unNumEntries+ -- * Value resolution/merging+ , combine+ , combineUnion+ , combineMaybe+ ) where++import Control.DeepSeq (NFData (..))+import Data.Bifoldable (Bifoldable (..))+import Data.Bifunctor (Bifunctor (..))++data Entry v b+ = Insert !v+ | InsertWithBlob !v !b+ | Upsert !v+ | Delete+ deriving stock (Eq, Show, Functor, Foldable, Traversable)++hasBlob :: Entry v b -> Bool+hasBlob Insert{} = False+hasBlob InsertWithBlob{} = True+hasBlob Upsert{} = False+hasBlob Delete{} = False++instance (NFData v, NFData b) => NFData (Entry v b) where+ rnf (Insert v) = rnf v+ rnf (InsertWithBlob v br) = rnf v `seq` rnf br+ rnf (Upsert v) = rnf v+ rnf Delete = ()++onValue :: v' -> (v -> v') -> Entry v b -> v'+onValue def f = \case+ Insert v -> f v+ InsertWithBlob v _ -> f v+ Upsert v -> f v+ Delete -> def++onBlobRef :: b' -> (b -> b') -> Entry v b -> b'+onBlobRef def g = \case+ Insert{} -> def+ InsertWithBlob _ br -> g br+ Upsert{} -> def+ Delete -> def++instance Bifunctor Entry where+ first f = \case+ Insert v -> Insert (f v)+ InsertWithBlob v br -> InsertWithBlob (f v) br+ Upsert v -> Upsert (f v)+ Delete -> Delete++ second g = \case+ Insert v -> Insert v+ InsertWithBlob v br -> InsertWithBlob v (g br)+ Upsert v -> Upsert v+ Delete -> Delete++instance Bifoldable Entry where+ bifoldMap f g = \case+ Insert v -> f v+ InsertWithBlob v br -> f v <> g br+ Upsert v -> f v+ Delete -> mempty++-- | A count of entries, for example the number of entries in a run.+--+-- This number is limited by the machine's word size. On 32-bit systems, the+-- maximum number we can represent is @2^31@ which is roughly 2 billion. This+-- should be a sufficiently large limit that we never reach it in practice. By+-- extension for 64-bit and higher-bit systems this limit is also sufficiently+-- large.+newtype NumEntries = NumEntries Int+ deriving stock (Eq, Ord, Show)+ deriving newtype NFData++instance Semigroup NumEntries where+ NumEntries a <> NumEntries b = NumEntries (a + b)++instance Monoid NumEntries where+ mempty = NumEntries 0++unNumEntries :: NumEntries -> Int+unNumEntries (NumEntries x) = x++{-------------------------------------------------------------------------------+ Value resolution/merging+-------------------------------------------------------------------------------}++-- | Given a value-merge function, combine entries. Only take a blob from the+-- left entry.+--+-- Note: 'Entry' is a semigroup with 'combine' if the @(v -> v -> v)@ argument+-- is associative.+combine :: (v -> v -> v) -> Entry v b -> Entry v b -> Entry v b+combine _ e@Delete _ = e+combine _ e@Insert {} _ = e+combine _ e@InsertWithBlob {} _ = e+combine _ (Upsert u) Delete = Insert u+combine f (Upsert u) (Insert v) = Insert (f u v)+combine f (Upsert u) (InsertWithBlob v _) = Insert (f u v)+combine f (Upsert u) (Upsert v) = Upsert (f u v)++-- | Combine two entries of runs that have been 'union'ed together. If any one+-- has a value, the result should have a value (represented by 'Insert'). If+-- both have a value, these values get combined monoidally. Only take a blob+-- from the left entry.+--+-- Note: 'Entry' is a semigroup with 'combineUnion' if the @(v -> v -> v)@+-- argument is associative.+combineUnion :: (v -> v -> v) -> Entry v b -> Entry v b -> Entry v b+combineUnion f = go+ where+ go Delete (Upsert v) = Insert v+ go Delete e = e+ go (Upsert u) Delete = Insert u+ go e Delete = e+ go (Insert u) (Insert v) = Insert (f u v)+ go (Insert u) (InsertWithBlob v _) = Insert (f u v)+ go (Insert u) (Upsert v) = Insert (f u v)+ go (InsertWithBlob u b) (Insert v) = InsertWithBlob (f u v) b+ go (InsertWithBlob u b) (InsertWithBlob v _) = InsertWithBlob (f u v) b+ go (InsertWithBlob u b) (Upsert v) = InsertWithBlob (f u v) b+ go (Upsert u) (Insert v) = Insert (f u v)+ go (Upsert u) (InsertWithBlob v _) = Insert (f u v)+ go (Upsert u) (Upsert v) = Insert (f u v)++combineMaybe :: (v -> v -> v) -> Maybe (Entry v b) -> Maybe (Entry v b) -> Maybe (Entry v b)+combineMaybe _ e1 Nothing = e1+combineMaybe _ Nothing e2 = e2+combineMaybe f (Just e1) (Just e2) = Just $! combine f e1 e2
@@ -0,0 +1,369 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.IncomingRun (+ IncomingRun (..)+ , MergePolicyForLevel (..)+ , duplicateIncomingRun+ , releaseIncomingRun+ , newIncomingSingleRun+ , newIncomingMergingRun+ , snapshotIncomingRun++ -- * Credits and credit tracking+ -- $credittracking+ , NominalDebt (..)+ , NominalCredits (..)+ , nominalDebtAsCredits+ , supplyCreditsIncomingRun+ , immediatelyCompleteIncomingRun+ ) where++import Control.Concurrent.Class.MonadMVar.Strict+import Control.DeepSeq (NFData (..))+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..))+import Control.Monad.Primitive+import Control.RefCount+import Data.Primitive (Prim)+import Data.Primitive.PrimVar+import Database.LSMTree.Internal.Assertions (assert)+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.MergingRun (MergeCredits (..),+ MergeDebt (..), MergingRun)+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.Run (Run)++import GHC.Exts (Word (W#), quotRemWord2#, timesWord2#)++{-------------------------------------------------------------------------------+ Incoming runs+-------------------------------------------------------------------------------}++-- | An incoming run is either a single run, or a merge.+data IncomingRun m h =+ Single !(Ref (Run m h))+ | Merging !MergePolicyForLevel+ !NominalDebt+ !(PrimVar (PrimState m) NominalCredits)+ !(Ref (MergingRun MR.LevelMergeType m h))++data MergePolicyForLevel = LevelTiering | LevelLevelling+ deriving stock (Show, Eq)++instance NFData MergePolicyForLevel where+ rnf LevelTiering = ()+ rnf LevelLevelling = ()++{-# SPECIALISE duplicateIncomingRun :: IncomingRun IO h -> IO (IncomingRun IO h) #-}+duplicateIncomingRun ::+ (PrimMonad m, MonadMask m)+ => IncomingRun m h+ -> m (IncomingRun m h)+duplicateIncomingRun (Single r) =+ Single <$> dupRef r++duplicateIncomingRun (Merging mp md mcv mr) =+ Merging mp md <$> (newPrimVar =<< readPrimVar mcv)+ <*> dupRef mr++{-# SPECIALISE releaseIncomingRun :: IncomingRun IO h -> IO () #-}+releaseIncomingRun ::+ (PrimMonad m, MonadMask m)+ => IncomingRun m h -> m ()+releaseIncomingRun (Single r) = releaseRef r+releaseIncomingRun (Merging _ _ _ mr) = releaseRef mr++{-# INLINE newIncomingSingleRun #-}+newIncomingSingleRun ::+ (PrimMonad m, MonadThrow m)+ => Ref (Run m h)+ -> m (IncomingRun m h)+newIncomingSingleRun r = Single <$> dupRef r++{-# INLINE newIncomingMergingRun #-}+newIncomingMergingRun ::+ (PrimMonad m, MonadThrow m)+ => MergePolicyForLevel+ -> NominalDebt+ -> Ref (MergingRun MR.LevelMergeType m h)+ -> m (IncomingRun m h)+newIncomingMergingRun mergePolicy nominalDebt mr = do+ nominalCreditsVar <- newPrimVar (NominalCredits 0)+ Merging mergePolicy nominalDebt nominalCreditsVar <$> dupRef mr++{-# SPECIALISE snapshotIncomingRun ::+ IncomingRun IO h+ -> IO (Either (Ref (Run IO h))+ (MergePolicyForLevel,+ NominalDebt,+ NominalCredits,+ Ref (MergingRun MR.LevelMergeType IO h))) #-}+snapshotIncomingRun ::+ PrimMonad m+ => IncomingRun m h+ -> m (Either (Ref (Run m h))+ (MergePolicyForLevel,+ NominalDebt,+ NominalCredits,+ Ref (MergingRun MR.LevelMergeType m h)))+snapshotIncomingRun (Single r) = pure (Left r)+snapshotIncomingRun (Merging mergePolicy nominalDebt nominalCreditsVar mr) = do+ nominalCredits <- readPrimVar nominalCreditsVar+ pure (Right (mergePolicy, nominalDebt, nominalCredits, mr))++{-------------------------------------------------------------------------------+ Credits+-------------------------------------------------------------------------------}++{- $credittracking++With scheduled merges, each update (e.g., insert) on a table contributes to the+progression of ongoing merges in the levels structure. This ensures that merges+are finished in time before a new merge has to be started. The points in the+evolution of the levels structure where new merges are started are known: a+flush of a full write buffer will create a new run on the first level, and+after sufficient flushes (e.g., 4) we will start at least one new merge on the+second level. This may cascade down to lower levels depending on how full the+levels are. As such, we have a well-defined measure to determine when merges+should be finished: it only depends on the maximum size of the write buffer!++The simplest solution to making sure merges are done in time is to step them to+completion immediately when started. This does not, however, spread out work+over time nicely. Instead, we schedule merge work based on how many updates are+made on the table, taking care to ensure that the merge is finished /just/ in+time before the next flush comes around, and not too early.++The progression is tracked using nominal credits. Each individual update+contributes a single credit to each level, since each level contains precisely+one ongoing merge. Contributing a credit does not, however, translate directly+to performing one /unit/ of merging work:++* The amount of work to do for one credit is adjusted depending on the actual+ size of the merge we are doing. Last-level merges, for example, can have+ larger inputs, and therefore we have to do a little more work for each+ credit. Or input runs involved in a merge can be less than maximal size for+ the level, and so there may be less merging work to do. As such, we /scale/+ 'NominalCredits' to 'MergeCredits', and then supply the 'MergeCredits' to+ the 'MergingRun'.++* Supplying 'MergeCredits' to a 'MergingRun' does not necessarily directly+ translate into performing merging work. Merge credits are accumulated until+ they go over a threshold, after which a batch of merge work will be performed.+ Configuring this threshold should allow a good balance between spreading out+ I\/O and achieving good (concurrent) performance.++Merging runs can be shared across tables, which means that multiple threads+can contribute to the same merge concurrently. Incoming runs however are /not/+shared between tables. As such the tracking of 'NominalCredits' does not need+to use any concurrency precautions.+-}++-- | Total merge debt to complete the merge in an incoming run.+--+-- This corresponds to the number (worst case, minimum number) of update+-- operations inserted into the table, before we will expect the merge to+-- complete.+newtype NominalDebt = NominalDebt Int+ deriving stock Eq+ deriving newtype (NFData)++-- | Merge credits that get supplied to a table's levels.+--+-- This corresponds to the number of update operations inserted into the table.+newtype NominalCredits = NominalCredits Int+ deriving stock Eq+ deriving newtype (Prim, NFData)++nominalDebtAsCredits :: NominalDebt -> NominalCredits+nominalDebtAsCredits (NominalDebt c) = NominalCredits c++{-# SPECIALISE supplyCreditsIncomingRun ::+ TableConfig+ -> LevelNo+ -> IncomingRun IO h+ -> NominalCredits+ -> IO () #-}+-- | Supply a given number of nominal credits to the merge in an incoming run.+-- This is a relative addition of credits, not a new absolute total value.+supplyCreditsIncomingRun ::+ (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => TableConfig+ -> LevelNo+ -> IncomingRun m h+ -> NominalCredits+ -> m ()+supplyCreditsIncomingRun _ _ (Single _r) _ = pure ()+supplyCreditsIncomingRun conf ln (Merging _ nominalDebt nominalCreditsVar mr)+ deposit = do+ (_nominalCredits,+ nominalCredits') <- depositNominalCredits nominalDebt nominalCreditsVar+ deposit+ let !mergeDebt = MR.totalMergeDebt mr+ !mergeCredits' = scaleNominalToMergeCredit nominalDebt mergeDebt+ nominalCredits'+ !thresh = creditThresholdForLevel conf ln+ (_suppliedCredits,+ _suppliedCredits') <- MR.supplyCreditsAbsolute mr thresh mergeCredits'+ pure ()+ --TODO: currently each supplying credits action results in contributing+ -- credits to the underlying merge, but this need not be the case. We+ -- _could_ do threshold based batching at the level of the IncomingRun.+ -- The IncomingRun does not need to worry about concurrency, so does not+ -- pay the cost of atomic operations on the counters. Then when we+ -- accumulate a batch we could supply that to the MergingRun (which must+ -- use atomic operations for its counters). We could potentially simplify+ -- MergingRun by dispensing with batching for the MergeCredits counters.++-- | Deposit nominal credits in the local credits var, ensuring the total+-- credits does not exceed the total debt.+--+-- Depositing /could/ leave the credit higher than the total debt. It is not+-- avoided by construction. The scenario is this: when a completed merge is+-- underfull, we combine it with the incoming run, so it means we have one run+-- fewer on the level then we'd normally have. This means that the level+-- becomes full at a later time, so more time passes before we call+-- 'MR.expectCompleted' on any levels further down the tree. This means we keep+-- supplying nominal credits to levels further down past the point their+-- nominal debt is paid off. So the solution here is just to drop any nominal+-- credits that are in excess of the nominal debt.+--+-- This is /not/ itself thread safe. All 'TableContent' update operations are+-- expected to be serialised by the caller. See concurrency comments for+-- 'TableContent' for detail.+{-# SPECIALISE depositNominalCredits ::+ NominalDebt+ -> PrimVar RealWorld NominalCredits+ -> NominalCredits+ -> IO (NominalCredits, NominalCredits) #-}+depositNominalCredits ::+ PrimMonad m+ => NominalDebt+ -> PrimVar (PrimState m) NominalCredits+ -> NominalCredits+ -> m (NominalCredits, NominalCredits)+depositNominalCredits (NominalDebt nominalDebt)+ nominalCreditsVar+ (NominalCredits deposit) = do+ NominalCredits before <- readPrimVar nominalCreditsVar+ let !after = NominalCredits (min (before + deposit) nominalDebt)+ writePrimVar nominalCreditsVar after+ pure (NominalCredits before, after)++-- | Linearly scale a nominal credit (between 0 and the nominal debt) into an+-- equivalent merge credit (between 0 and the total merge debt).+--+-- Crucially, @100% nominal credit ~~ 100% merge credit@, so when we pay off+-- the nominal debt, we also exactly pay off the merge debt. That is:+--+-- > scaleNominalToMergeCredit nominalDebt mergeDebt nominalDebt == mergeDebt+--+-- (modulo some newtype conversions)+--+scaleNominalToMergeCredit ::+ NominalDebt+ -> MergeDebt+ -> NominalCredits+ -> MergeCredits+scaleNominalToMergeCredit (NominalDebt nominalDebt)+ (MergeDebt (MergeCredits mergeDebt))+ (NominalCredits nominalCredits) =+ -- The scaling involves an operation: (a * b) `div` c+ -- but where potentially the variables a,b,c may be bigger than a 32bit+ -- integer can hold. This would be the case for runs that have more than+ -- 4 billion entries.+ --+ -- (This is assuming 64bit Int, the problem would be even worse for 32bit+ -- systems. The solution here would also work for 32bit systems, allowing+ -- up to, 2^31, 2 billion entries per run.)+ --+ -- To work correctly in this case we need higher range for the intermediate+ -- result a*b which could be bigger than 64bits can hold. A correct+ -- implementation can use Rational, but a fast implementation should use+ -- only integer operations. This is relevant because this is on the fast+ -- path for small insertions into the table that often do no merging work+ -- and just update credit counters.++ -- The fast implementation uses integer operations that produce a 128bit+ -- intermediate result for the a*b result, and use a 128bit numerator in+ -- the division operation (but 64bit denominator). These are known as+ -- "widening multiplication" and "narrowing division". GHC has direct+ -- support for these operations as primops: timesWord2# and quotRemWord2#,+ -- but they are not exposed through any high level API shipped with GHC.++ -- The specification using Rational is:+ let mergeCredits_spec = floor $ toRational nominalCredits+ * toRational mergeDebt+ / toRational nominalDebt+ -- Note that it doesn't matter if we use floor or ceiling here.+ -- Rounding errors will not compound because we sum nominal debt and+ -- convert absolute nominal to absolute merging credit. We don't+ -- convert each deposit and sum all the rounding errors.+ -- When nominalCredits == nominalDebt then the result is exact anyway+ -- (being mergeDebt) so the rounding mode makes no difference when we+ -- get to the end of the merge. Using floor makes things simpler for+ -- the fast integer implementation below, so we take that as the spec.++ -- If the nominalCredits is between 0 and nominalDebt then it's+ -- guaranteed that the mergeCredit is between 0 and mergeDebt.+ -- The mergeDebt fits in an Int, therefore the result does too.+ -- Therefore the undefined behaviour case of timesDivABC_fast is+ -- avoided and the w2i cannot overflow.+ mergeCredits_fast = w2i $ timesDivABC_fast (i2w nominalCredits)+ (i2w mergeDebt)+ (i2w nominalDebt)+ in assert (0 < nominalDebt) $+ assert (0 <= nominalCredits && nominalCredits <= nominalDebt) $+ assert (mergeCredits_spec == mergeCredits_fast) $+ MergeCredits mergeCredits_fast+ where+ {-# INLINE i2w #-}+ {-# INLINE w2i #-}+ i2w :: Int -> Word+ w2i :: Word -> Int+ i2w = fromIntegral+ w2i = fromIntegral++-- | Compute @(a * b) `div` c@ for unsigned integers for the full range of+-- 64bit unsigned integers, provided that @a <= c@ and thus the result will+-- fit in 64bits.+--+-- The @a * b@ intermediate result is computed using 128bit precision.+--+-- Note: the behaviour is undefined if the result will not fit in 64bits.+-- It will probably result in immediate termination with SIGFPE.+--+timesDivABC_fast :: Word -> Word -> Word -> Word+timesDivABC_fast (W# a) (W# b) (W# c) =+ case timesWord2# a b of+ (# ph, pl #) ->+ case quotRemWord2# ph pl c of+ (# q, _r #) -> W# q++{-# SPECIALISE immediatelyCompleteIncomingRun ::+ TableConfig+ -> LevelNo+ -> IncomingRun IO h+ -> IO (Ref (Run IO h)) #-}+-- | Supply enough credits to complete the merge now.+immediatelyCompleteIncomingRun ::+ (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => TableConfig+ -> LevelNo+ -> IncomingRun m h+ -> m (Ref (Run m h))+immediatelyCompleteIncomingRun conf ln ir =+ case ir of+ Single r -> dupRef r+ Merging _ (NominalDebt nominalDebt) nominalCreditsVar mr -> do++ NominalCredits nominalCredits <- readPrimVar nominalCreditsVar+ let !deposit = NominalCredits (nominalDebt - nominalCredits)+ supplyCreditsIncomingRun conf ln ir deposit++ -- This ensures the merge is really completed. However, we don't+ -- release the merge yet, but we do return a new reference to the run.+ MR.expectCompleted mr
@@ -0,0 +1,211 @@+{-# OPTIONS_HADDOCK not-home #-}++-- |+-- Provides support for working with fence pointer indexes of different types+-- and their accumulators.+--+-- Keys used with an index are subject to the key size constraints of the+-- concrete type of the index. These constraints are stated in the descriptions+-- of the modules "Database.LSMTree.Internal.Index.Compact" and+-- "Database.LSMTree.Internal.Index.Ordinary", respectively.+--+-- Part of the functionality that this module provides is the construction of+-- serialised indexes in a mostly incremental fashion. The incremental part of+-- serialisation is provided through index accumulators, while the+-- non-incremental bits are provided through the index operations 'headerLBS'+-- and 'finalLBS'. To completely serialise an index interleaved with its+-- construction, proceed as follows:+--+-- 1. Use 'headerLBS' to generate the header of the serialised index.+--+-- 2. Incrementally construct the index using the operations of 'IndexAcc',+-- and assemble the body of the serialised index from the generated chunks.+--+-- 3. Use 'finalLBS' to generate the footer of the serialised index.+--+module Database.LSMTree.Internal.Index+(+ -- * Index types+ IndexType (Compact, Ordinary),+ indexToIndexType,++ -- * Indexes+ Index (CompactIndex, OrdinaryIndex),+ search,+ sizeInPages,+ headerLBS,+ finalLBS,+ fromSBS,++ -- * Index accumulators+ IndexAcc (CompactIndexAcc, OrdinaryIndexAcc),+ newWithDefaults,+ appendSingle,+ appendMulti,+ unsafeEnd+)+where++import Control.Arrow (second)+import Control.DeepSeq (NFData (rnf))+import Control.Monad.ST.Strict (ST)+import Data.ByteString.Lazy (LazyByteString)+import Data.ByteString.Short (ShortByteString)+import Data.Word (Word32)+import Database.LSMTree.Internal.Chunk (Chunk)+import Database.LSMTree.Internal.Entry (NumEntries)+import Database.LSMTree.Internal.Index.Compact (IndexCompact)+import qualified Database.LSMTree.Internal.Index.Compact as Compact (finalLBS,+ fromSBS, headerLBS, search, sizeInPages)+import Database.LSMTree.Internal.Index.CompactAcc (IndexCompactAcc)+import qualified Database.LSMTree.Internal.Index.CompactAcc as Compact+ (appendMulti, appendSingle, newWithDefaults, unsafeEnd)+import Database.LSMTree.Internal.Index.Ordinary (IndexOrdinary)+import qualified Database.LSMTree.Internal.Index.Ordinary as Ordinary (finalLBS,+ fromSBS, headerLBS, search, sizeInPages)+import Database.LSMTree.Internal.Index.OrdinaryAcc (IndexOrdinaryAcc)+import qualified Database.LSMTree.Internal.Index.OrdinaryAcc as Ordinary+ (appendMulti, appendSingle, newWithDefaults, unsafeEnd)+import Database.LSMTree.Internal.Page (NumPages, PageSpan)+import Database.LSMTree.Internal.Serialise (SerialisedKey)++-- * Index types++-- | The type of supported index types.+data IndexType = Compact | Ordinary+ deriving stock (Eq, Show)++instance NFData IndexType where+ rnf Compact = ()+ rnf Ordinary = ()++-- * Indexes++-- | The type of supported indexes.+data Index+ = CompactIndex !IndexCompact+ | OrdinaryIndex !IndexOrdinary+ deriving stock (Eq, Show)++instance NFData Index where+ rnf (CompactIndex index) = rnf index+ rnf (OrdinaryIndex index) = rnf index++indexToIndexType :: Index -> IndexType+indexToIndexType CompactIndex{} = Compact+indexToIndexType OrdinaryIndex{} = Ordinary++{-|+ Searches for a page span that contains a key–value pair with the given key.+ If there is indeed such a pair, the result is the corresponding page span;+ if there is no such pair, the result is an arbitrary but valid page span.+-}+search :: SerialisedKey -> Index -> PageSpan+search key (CompactIndex index) = Compact.search key index+search key (OrdinaryIndex index) = Ordinary.search key index++-- | Yields the number of pages covered by an index.+sizeInPages :: Index -> NumPages+sizeInPages (CompactIndex index) = Compact.sizeInPages index+sizeInPages (OrdinaryIndex index) = Ordinary.sizeInPages index++{-|+ Yields the header of the serialised form of an index.++ See [the module documentation]("Database.LSMTree.Internal.Index") for how to+ generate a complete serialised index.+-}+headerLBS :: IndexType -> LazyByteString+headerLBS Compact = Compact.headerLBS+headerLBS Ordinary = Ordinary.headerLBS++{-|+ Yields the footer of the serialised form of an index.++ See [the module documentation]("Database.LSMTree.Internal.Index") for how to+ generate a complete serialised index.+-}+finalLBS :: NumEntries -> Index -> LazyByteString+finalLBS entryCount (CompactIndex index) = Compact.finalLBS entryCount index+finalLBS entryCount (OrdinaryIndex index) = Ordinary.finalLBS entryCount index++{-|+ Reads an index along with the number of entries of the respective run from a+ byte string.++ The byte string must contain the serialised index exactly, with no leading+ or trailing space. Furthermore, its contents must be stored 64-bit-aligned.++ The contents of the byte string may be directly used as the backing memory+ for the constructed index. Currently, this is done for compact indexes.++ For deserialising numbers, the endianness of the host system is used. If+ serialisation has been done with a different endianness, this mismatch is+ detected by looking at the type–version indicator.+-}+fromSBS :: IndexType -> ShortByteString -> Either String (NumEntries, Index)+fromSBS Compact input = second CompactIndex <$> Compact.fromSBS input+fromSBS Ordinary input = second OrdinaryIndex <$> Ordinary.fromSBS input++-- * Index accumulators++{-|+ The type of supported index accumulators, where an index accumulator denotes+ an index under incremental construction.++ Incremental index construction is only guaranteed to work correctly when the+ supplied key ranges do not overlap and are given in ascending order.+-}+data IndexAcc s = CompactIndexAcc !(IndexCompactAcc s)+ | OrdinaryIndexAcc !(IndexOrdinaryAcc s)++-- | Create a new index accumulator, using a default configuration.+newWithDefaults :: IndexType -> ST s (IndexAcc s)+newWithDefaults Compact = CompactIndexAcc <$> Compact.newWithDefaults+newWithDefaults Ordinary = OrdinaryIndexAcc <$> Ordinary.newWithDefaults++{-|+ Adds information about a single page that fully comprises one or more+ key–value pairs to an index and outputs newly available chunks.++ See the documentation of the 'IndexAcc' type for constraints to adhere to.+-}+appendSingle :: (SerialisedKey, SerialisedKey)+ -> IndexAcc s+ -> ST s (Maybe Chunk)+appendSingle pageInfo (CompactIndexAcc indexAcc) = Compact.appendSingle+ pageInfo+ indexAcc+appendSingle pageInfo (OrdinaryIndexAcc indexAcc) = Ordinary.appendSingle+ pageInfo+ indexAcc++{-|+ Adds information about multiple pages that together comprise a single+ key–value pair to an index and outputs newly available chunks.++ The provided 'Word32' value denotes the number of /overflow/ pages, so that+ the number of pages that comprise the key–value pair is the successor of+ that number.++ See the documentation of the 'IndexAcc' type for constraints to adhere to.+-}+appendMulti :: (SerialisedKey, Word32) -> IndexAcc s -> ST s [Chunk]+appendMulti pagesInfo (CompactIndexAcc indexAcc) = Compact.appendMulti+ pagesInfo+ indexAcc+appendMulti pagesInfo (OrdinaryIndexAcc indexAcc) = Ordinary.appendMulti+ pagesInfo+ indexAcc++{-|+ Returns the constructed index, along with a final chunk in case the+ serialised key list has not been fully output yet, thereby invalidating the+ index under construction. Executing @unsafeEnd index@ is only safe when+ @index@ is not used afterwards.+-}+unsafeEnd :: IndexAcc s -> ST s (Maybe Chunk, Index)+unsafeEnd (CompactIndexAcc indexAcc) = second CompactIndex <$>+ Compact.unsafeEnd indexAcc+unsafeEnd (OrdinaryIndexAcc indexAcc) = second OrdinaryIndex <$>+ Ordinary.unsafeEnd indexAcc
@@ -0,0 +1,743 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | A compact fence-pointer index for uniformly distributed keys.+--+-- Keys used with a compact index must be at least 8 bytes long.+--+-- TODO: add utility functions for clash probability calculations+--+module Database.LSMTree.Internal.Index.Compact (+ -- $compact+ IndexCompact (..)+ -- * Queries+ , search+ , sizeInPages+ , countClashes+ , hasClashes+ -- * Non-incremental serialisation+ , toLBS+ -- * Incremental serialisation+ , headerLBS+ , finalLBS+ , word64VectorToChunk+ -- * Deserialisation+ , fromSBS+ ) where++import Control.DeepSeq (NFData (..))+import Control.Monad (when)+import Control.Monad.ST+import Data.Bit hiding (flipBit)+import Data.Bits (unsafeShiftR, (.&.))+import qualified Data.ByteString.Builder as BB+import qualified Data.ByteString.Builder.Extra as BB+import qualified Data.ByteString.Lazy as LBS+import Data.ByteString.Short (ShortByteString (..))+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromMaybe)+import Data.Primitive.ByteArray (ByteArray (..), indexByteArray,+ sizeofByteArray)+import Data.Primitive.Types (sizeOf)+import qualified Data.Vector.Algorithms.Search as VA+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Base as VU+import Data.Word+import Database.LSMTree.Internal.BitMath+import Database.LSMTree.Internal.ByteString (byteArrayFromTo)+import Database.LSMTree.Internal.Chunk (Chunk (Chunk))+import qualified Database.LSMTree.Internal.Chunk as Chunk (toByteString)+import Database.LSMTree.Internal.Entry (NumEntries (..))+import Database.LSMTree.Internal.Map.Range (Bound (..))+import Database.LSMTree.Internal.Page+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Unsliced+import Database.LSMTree.Internal.Vector++{- $compact++ A fence-pointer index is a mapping of disk pages, identified by some number+ @i@, to min-max information for keys on that page.++ Fence-pointer indexes can be constructed and serialised incrementally, see+ module "Database.LSMTree.Internal.Index.CompactAcc".++ Given a serialised target key @k@, an index can be 'search'ed to find a disk+ page @i@ that /might/ contain @k@. Fence-pointer indices offer no guarantee of+ whether the page contains the key, but the indices do guarantee that no page+ other than page @i@ could store @k@.++ === Intro++ A run is a file that stores a collection of key-value pairs. The data is+ sorted by key, and the data is divided into disk pages. A fence pointer index+ is a mapping of each disk page to min\/max information for keys on that page.+ As such, the index stores the keys that mark the /boundaries/ for each page. A+ lookup of a key in a run first searches the fence pointer index to find the+ page that contains the relevant key range, before doing a single I/O to+ retrieve the page. The guarantee of an index search is this:++ GUARANTEE: /if/ the key is somewhere in the run, then it can only be in the+ page that is returned by the index search.++ This module presents a /compact/ implementation of a fence pointer index.+ However, we first consider a simple implementation of a fence pointer index.+ After that, we show how we can save on the memory size of the index+ representation if the keys are uniformly distributed, like hashes. Let's start+ with some basic definitions:++ > type Run k v = -- elided+ > type Page k v = -- elided+ > minKey :: Page k v -> k+ > maxKey :: Page k v - k+ > type PageNo = Int++ The first implementation one may come up with is a vector containing the+ minimum key and maximum key on each page. As such, the index stores the+ key-interval @[minKey p, maxKey p]@ for each page @p@. @search1@ searches the+ vector for the key-interval that contains the search key, if it exists, and+ returns the corresponding vector index as a page number. A design choice of+ ours is that the search will __always__ return a page number, even if the+ index could answer that the key is definitely not in a page (see+ 'indexSearches').++ > type Index1 k = V.Vector (k, k)+ >+ > mkIndex1 :: Run k v -> Index1 k+ > mkIndex1 = V.fromList . fmap (\p -> (minKey p, maxKey p))+ >+ > search1 :: k -> Index1 k -> PageNo+ > search1 = -- elided++ We can reduce the memory size of `Index1` by half if we store only the minimum+ keys on each page. As such, the index now stores the key-interval @[minKey p,+ minKey p')@ for each page @p@ and successor page @p'@. GUARANTEE is still+ guaranteed, because the old intervals are strictly contained in the new ones.+ @search2@ searches the vector for the largest key smaller or equal to the+ given one, if it exists, and returns the corresponding vector index as a page+ number.++ > type Index2 k = V.Vector k+ >+ > mkIndex2 :: Run k v -> Index2 k+ > mkIndex2 = V.fromList . fmap minKey+ >+ > search2 :: k -> Index k -> PageNo+ > search2 = -- elided++ Now on to creating a more compact representation, which relies on a property+ of the keys that are used: the keys must be uniformly distributed values, like+ hashes.++ === Compact representation++ As mentioned before, we can save on the memory size of the index+ representation if the keys are uniformly distributed, like hashes. From now+ on, we will just assume that our keys are hashes, which can be viewed as+ strings of bits.++ The intuition behind the compact index is this: often, we don't need to look+ at full bit-strings to compare independently generated hashes against+ each other. The probability of the @n@ most significant bits of two+ independently generated hashes matching is @(1/(2^n))@, and if we pick @n@+ large enough then we can expect a very small number of collisions. More+ generally, the expected number of @n@-bit hashes that have to be generated+ before a collision is observed is @2^(n/2)@, see [the birthday+ problem](https://en.wikipedia.org/wiki/Birthday_problem#Probability_of_a_shared_birthday_(collision).+ Or, phrased differently, if we store only the @n@ most significant bits of+ independently generated hashes in our index, then we can store up to @2^(n/2)@+ of those hashes before the expected number of collisions becomes one. Still,+ we need a way to break ties if there are collisions, because the probability+ of collisions is non-zero.++ ==== Clashes++ In our previous incarnations of the "simple" index, we relied on the minimum+ keys to define the boundaries between pages, or more precisely, the pages'+ key-intervals. If we look only at a subset of the most significant bits, that+ is no longer true in the presence of collisions. Let's illustrate this using+ an example of two contiguous pages @pi@ and @pj@:++ > minKey pi == "00000000"+ > maxKey pi == "10000000"+ > minKey pj == "10000001"+ > maxKey pj == "11111111"++ Say we store only the 4 most significant bits (left-to-right) in our compact+ index. This means that those 4 bits of @maxKey pi@ and @minKey pj@ collide.+ The problem is that the interval @[minKey pi, maxKey pi]@ strictly contains+ the interval @[minKey pi, minKey pj)@! The first 4 bits of @minKey pj@ do not+ actually define the boundary between @pi@ and @pj@, and so there would be no+ way to answer if a search key with first 4 bits "1000" could be found in @pi@+ or @pj@. We call such a situation a /clash/,++ The solution in these cases is to: (i) record that a clash occurred between+ pages @pi@ and @pj@, and (ii) store the full key @maxKey pi@ separately. The+ record of clashes can be implemented simply as a single bit per page: with+ a @True@ bit meaning a clash with the previous page. Note therefore that the+ first page's bit is always going to be @False@. We store the full keys using+ a 'Map'. It is ok that this is not a compact representation, because we+ expect to store full keys for only a very small number of pages.++ The example below shows a simplified view of the compact index implementation+ so far. As an example, we store the @64@ most significant bits of each minimum+ key in the @primary@ index, the record of clashes is called @clashes@, and the+ @IntMap@ is named the @tieBreaker@ map. @search3@ can at any point during the+ search, consult @clashes@ and @tieBreaker@ to /break ties/.++ > -- (primary , clashes , tieBreaker)+ > type Index3 k = (VU.Vector Word64, VU.Vector Bit, IntMap k )+ >+ > mkIndex3 :: Run k v -> Index3+ > mkIndex3 = -- elided+ >+ > search3 :: k -> Index3 -> PageNo+ > search3 = -- elided++ Let's compare the memory size of @Index2@ with @Index3@. Say we have \(n\)+ pages, and keys that are \(k\) bits each, then the memory size of @Index2@ is+ \(O(n~k)\) bits. Alternatively, for the same \(n\) pages, if we store only the+ \(c\) most significant bits, then the memory size of @Index3@ is+ \(O\left(n~c + n + k E\left[\text{collision}~n~c\right]\right)\) bits, where+ the last summand is the expected number of collisions for \(n\) independently+ generated hashes of bit-size \(c\). Precisely, the expected number of+ collisions is \(\frac{n^2 - n}{2 \times 2^c}\), so we can simplify the memory+ size to \(O\left(n~c + n + k \frac{n^2}{2^c}\right)\).++ So, @Index3@ is not strictly an improvement over @Index2@ if we look at memory+ complexity alone. However, /in practice/ @Index3@ is an improvement over+ @Index2@ if we can pick an \(c\) that is (i) much smaller than \(k\) and (ii)+ keeps \( \text{E}\left[\text{collision}~n~c\right] \) small (e.g. close to 1).+ For example, storing the first 64 bits of 100 million SHA256 hashes reduces+ memory size from \(256 n\) bits to \(64 n + n\) bits, because the expected+ number of collisions is smaller than 1.++ === Representing clashes and larger-than-page entries++ A complicating factor is the need to represent larger-than-page entries in+ the compact index. This need arises from the fact that we can have single+ key/value pairs that need more than one disk page to represent.++ One strategy would be to read the first page, discover that it is a+ multi-page entry and then read the subsequent pages. This would however+ double the disk I\/O latency and complicate the I\/O handling logic (which is+ non-trivial due to it being asynchronous).++ The strategy we use is to have the index be able to return the range of pages+ that need to be read, and then all pages can be read in one batch. This means+ the index must return not just individual page numbers but intervals, and+ with a representation capable of doing so. This is not implausible since we+ have an entry in the primary index for every disk page, and so we have+ entries both for the first and subsequent pages of a larger-than-page entry.+ The natural thing to do is have each of these subsequent primary index+ entries contain the same key prefix value. This means a binary search will+ find the /last/ entry in a run of equal prefix values.++ What leads to complexity is that we will /also/ get runs of equal values if+ we have clashes between pages (as discussed above). So in the general case+ we may have a run of equal values made up of a mixture of clashes and+ larger-than-page entries.++ So the general situation is that after a binary search we have found the+ end of what may turn out to be a run of clashes and larger-than-page values+ and we must disambigutate and return the appropriate single page (for the+ ordinary case) or an interval of pages (for the LTP case).++ To disambigutate we make use of the clash bits, and we make the choice to+ say that /all/ the entries for a LTP have their clash bit set, irrespective+ of whether the LTP is in fact involved in a clash. This may seem+ counter-intuitive but it leads to a simpler mathematical definition (below).++ The search algorithm involves searching backwards in the clash bits to find+ the beginning of the run of entries that are involved. To establish which+ entry within the run is the right one to return, we can consult the tie+ breaker map by looking for the biggest entry that is less than or equal to+ the full search key. This may then point to an index within the run of+ clashing entries, in which case this is the right entry, but it may also+ point to an earlier and thus irrelevant entry, in which case the first entry+ in the run is the right one.++ Note that this second case also covers the case of a single non-clashing LTP.++ Finally, to determine if the selected entry is an LTP and if so what interval+ of pages to return, we make use of a second bit vector of LTP \"overflow\"+ pages. This bit vector has 1 values for LTP overflow pages (i.e. the 2nd and+ subsequent pages) and 0 otherwise. We can then simply search forwards to find+ the length of the LTP (or 1 if it is not an LTP).++ === A semi-formal description of the compact index #rep-descr#++ * \(n\) is the number of pages+ * \(ps = \{p_i \mid 0 \leq i < n \}\) is a sorted set of pages+ * \(p^{min}_i\) is the full minimum key on a page \(p_i \in ps\).+ * \(p^{max}_i\) is the full maximum key on a page \(p_i \in ps\).+ * \(\texttt{topBits64}(k)\) extracts the \(64\) most significant bits from+ \(k\). We call these \(64\) bits the primary bits.+ * \(i \in \left[0, n \right)\), unless stated otherwise++ \[+ \begin{align*}+ P :&~ \texttt{Array PageNo Word64} \\+ P[i] =&~ \texttt{topBits64}(p^{min}_i) \\+ \\+ C :&~ \texttt{Array PageNo Bit} \\+ C[0] =&~ \texttt{false} \\+ C[i] =&~ \texttt{topBits64}(p^{max}_{i-1}) ~\texttt{==}~ \texttt{topBits64}(p^{min}_i) \\+ \\+ TB :&~ \texttt{Map Key PageNo} \\+ TB(p^{min}_i) =&~+ \begin{cases}+ p^{min}_i &, \text{if}~ C[i] \land \neg LTP[i] \\+ \text{undefined} &, \text{otherwise} \\+ \end{cases} \\+ \\+ LTP :&~ \texttt{Array PageNo Bit} \\+ LTP[0] =&~ \texttt{false} \\+ LTP[i] =&~ p^{min}_{i-1} ~\texttt{==}~ p^{min}_i \\+ \end{align*}+ \]++ === An informal description of the search algorithm #search-descr#++ The easiest way to think about the search algorithm is that we start with the+ full interval of page numbers, and shrink it until we get to an interval that+ contains only a single page (or in case of a larger-than-page value, multiple+ pages that have the same minimum key). Assume @k@ is our search key.++ * Search \(P\) for the vector index @i@ that maps to the largest prim-bits+ value that is smaller or equal to the primary bits of @k@.+ * Check \(C\) if the page corresponding to @i@ is part of a+ clash, if it is not, then we are done!+ * If there is a clash, we go into clash recovery mode. This means we have to+ resolve the ambiguous page boundary using \(TB\). Note, however, that if+ there are multiple clashes in a row, there could be multiple ambiguous page+ boundaries that have to be resolved. We can achieve this using a three-step+ process:++ * Search \(TB\) for the vector index @j1@ that maps to the largest full+ key that is smaller than or equal to @k@.+ * Do a linear search backwards through \(C\) starting from @i@ to find the+ first page @j2@ that is not part of a clash.+ * Take the maximum of @j1@ or @j2@. Consider the two cases where either+ @j1@ or @j2@ is largest (@j1@ can not be equal to @j2@):++ * @j1@ is largest: we resolved ambiguous page boundaries "to the left"+ (toward lower vector indexes) until we resolved an ambiguous page+ boundary "to the right" (toward the current vector index).+ * @j2@ is largest: we resolved ambiguous page boundaries only "to the+ left", and ended up on a page that doesn't have a clash.++ * For larger-than-page values, the steps up until now would only provide us+ with the page where the larger-than-page value starts. We use \(LTP\) to do+ a linear search to find the page where the larger-than-page value ends.++ Convince yourself that clash recovery works without any larger-than-page+ values, and then consider the case where the index does contain+ larger-than-page values. Hints:++ \[+ \begin{align*}+ LTP[i] &~ \implies C[i] \\+ LTP[i] &~ \implies TB(p^{min}_i) = \text{undefined} \\+ \end{align*}+ \]+-}++-- | A compact fence-pointer index for uniformly distributed keys.+--+-- Compact indexes save space by storing bit-prefixes of keys.+--+-- See [a semi-formal description of the compact index](#rep-descr) for more+-- details about the representation.+--+-- While the semi-formal description mentions the number of pages \(n\),+-- we do not store it, as it can be inferred from the length of 'icPrimary'.+data IndexCompact = IndexCompact {+ -- | \(P\): Maps a page @i@ to the 64-bit slice of primary bits of its+ -- minimum key.+ icPrimary :: {-# UNPACK #-} !(VU.Vector Word64)+ -- | \(C\): A clash on page @i@ means that the primary bits of the minimum+ -- key on that page aren't sufficient to decide whether a search for a key+ -- should continue left or right of the page.+ , icClashes :: {-# UNPACK #-} !(VU.Vector Bit)+ -- | \(TB\): Maps a full minimum key to the page @i@ that contains it, but+ -- only if there is a clash on page @i@.+ , icTieBreaker :: !(Map (Unsliced SerialisedKey) PageNo)+ -- | \(LTP\): Record of larger-than-page values. Given a span of pages for+ -- the larger-than-page value, the first page will map to 'False', and the+ -- remainder of the pages will be set to 'True'. Regular pages default to+ -- 'False'.+ , icLargerThanPage :: {-# UNPACK #-} !(VU.Vector Bit)+ }+ deriving stock (Show, Eq)++instance NFData IndexCompact where+ rnf ic = rnf a `seq` rnf b `seq` rnf c `seq` rnf d+ where IndexCompact a b c d = ic++{-------------------------------------------------------------------------------+ Queries+-------------------------------------------------------------------------------}++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.search').++ See [an informal description of the search algorithm](#search-descr) for+ more details about the search algorithm.+-}+search :: SerialisedKey -> IndexCompact -> PageSpan+-- One way to think of the search algorithm is that it starts with the full page+-- number interval, and shrinks it to a minimal interval that contains the+-- search key. The code below is annotated with [x,y] or [x, y) comments that+-- describe the known page number interval at that point in the search+-- algorithm.+search k IndexCompact{..} =+ let !primbits = keyTopBits64 k in+ -- [0, n), where n is the length of the P array+ case unsafeSearchLE primbits icPrimary of+ Nothing ->+ -- TODO: if the P array is indeed empty, then this violates the+ -- guarantee that we return a valid page span! We should specify that a+ -- compact index should be non-empty.+ if VU.length icLargerThanPage == 0 then singlePage (PageNo 0) else+ -- [0, n), our page span definitely starts at 0, but we still have to+ -- consult the LTP array to check whether the value on page 0 overflows+ -- into subsequent pages.+ let !i = bitLongestPrefixFromTo (BoundExclusive 0) NoBound (Bit True) icLargerThanPage+ -- [0, i]+ in multiPage (PageNo 0) (PageNo i)+ Just !i ->+ -- [0, i]+ if unBit $ icClashes VU.! i then+ -- [0, i], now in clash recovery mode.+ let -- i is the *last* index in a range of contiguous pages that all+ -- clash. Since i is the end of the range, we search backwards+ -- through the C array to find the start of this range.+ !i1 = PageNo $ fromMaybe 0 $+ bitIndexFromToRev (BoundInclusive 0) (BoundInclusive i) (Bit False) icClashes+ -- The TB map is consulted to find the closest key smaller than k.+ !i2 = maybe (PageNo 0) snd $+ Map.lookupLE (makeUnslicedKey k) icTieBreaker+ -- If i2 < i1, then it means the clashing pages were all just part+ -- of the same larger-than-page value. Entries are only included+ -- in the TB map if the clash was a *proper* clash.+ --+ -- If i1 <= i2, then there was a proper clash in [i1, i] that+ -- required a comparison with a tiebreaker key.+ PageNo !i3 = max i1 i2+ -- [max i1 i2, i], this is equivalent to taking the intersection+ -- of [i1, i] and [i2, i]+ !i4 = bitLongestPrefixFromTo (BoundExclusive i3) (BoundInclusive i) (Bit True) icLargerThanPage+ in multiPage (PageNo i3) (PageNo i4)+ -- [i3, i4], we consulted the LTP array to check whether the value+ -- on page i3 overflows into subsequent pages+ else+ -- [i, i], there is no clash with the previous page and so this page+ -- is also not part of a large value that spans multiple pages.+ singlePage (PageNo i)+++countClashes :: IndexCompact -> Int+countClashes = Map.size . icTieBreaker++hasClashes :: IndexCompact -> Bool+hasClashes = not . Map.null . icTieBreaker++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.sizeInPages').+-}+sizeInPages :: IndexCompact -> NumPages+sizeInPages = NumPages . toEnum . VU.length . icPrimary++{-------------------------------------------------------------------------------+ Non-incremental serialisation+-------------------------------------------------------------------------------}++-- | Serialises a compact index in one go.+toLBS :: NumEntries -> IndexCompact -> LBS.ByteString+toLBS numEntries index =+ headerLBS+ <> LBS.fromStrict (Chunk.toByteString (word64VectorToChunk (icPrimary index)))+ <> finalLBS numEntries index++{-------------------------------------------------------------------------------+ Incremental serialisation+-------------------------------------------------------------------------------}++-- | By writing out the type–version indicator in host endianness, we also+-- indicate endianness. During deserialisation, we would discover an endianness+-- mismatch.+supportedTypeAndVersion :: Word32+supportedTypeAndVersion = 0x0001++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.headerLBS').+-}+headerLBS :: LBS.ByteString+headerLBS =+ -- create a single 4 byte chunk+ BB.toLazyByteStringWith (BB.safeStrategy 4 BB.smallChunkSize) mempty $+ BB.word32Host supportedTypeAndVersion <> BB.word32Host 0++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.finalLBS').+-}+finalLBS :: NumEntries -> IndexCompact -> LBS.ByteString+finalLBS (NumEntries numEntries) IndexCompact {..} =+ -- use a builder, since it is all relatively small+ BB.toLazyByteString $+ putBitVec icClashes+ <> putBitVec icLargerThanPage+ <> putTieBreaker icTieBreaker+ <> BB.word64Host (fromIntegral numPages)+ <> BB.word64Host (fromIntegral numEntries)+ where+ numPages = VU.length icPrimary++-- | Constructs a chunk containing the contents of a vector of 64-bit words.+word64VectorToChunk :: VU.Vector Word64 -> Chunk+word64VectorToChunk (VU.V_Word64 (VP.Vector off len ba)) =+ Chunk (mkPrimVector (mul8 off) (mul8 len) ba)++-- | Padded to 64 bit.+--+-- Assumes that the bitvector has a byte-aligned offset.+putBitVec :: VU.Vector Bit -> BB.Builder+putBitVec (BitVec offsetBits lenBits ba)+ | mod8 offsetBits /= 0 = error "putBitVec: not byte aligned"+ | otherwise =+ -- first only write the bytes that are fully part of the bit vec+ byteArrayFromTo offsetBytes offsetLastByte ba+ -- then carefully write the last byte, might be partially uninitialised+ <> (if remainingBits /= 0 then+ BB.word8 (indexByteArray ba offsetLastByte .&. bitmaskLastByte)+ else+ mempty)+ <> putPaddingTo64 totalBytesWritten+ where+ offsetBytes = div8 offsetBits+ offsetLastByte = offsetBytes + div8 lenBits+ totalBytesWritten = ceilDiv8 lenBits++ bitmaskLastByte = unsafeShiftR 0xFF (8 - remainingBits)+ remainingBits = mod8 lenBits++-- | Padded to 64 bit.+putTieBreaker :: Map (Unsliced SerialisedKey) PageNo -> BB.Builder+putTieBreaker m =+ BB.word64Host (fromIntegral (Map.size m))+ <> foldMap putEntry (Map.assocs m)+ where+ putEntry :: (Unsliced SerialisedKey, PageNo) -> BB.Builder+ putEntry (fromUnslicedKey -> k, PageNo pageNo) =+ BB.word32Host (fromIntegral pageNo)+ <> BB.word32Host (fromIntegral (sizeofKey k))+ <> serialisedKey k+ <> putPaddingTo64 (sizeofKey k)++putPaddingTo64 :: Int -> BB.Builder+putPaddingTo64 written+ | mod8 written == 0 = mempty+ | otherwise = foldMap BB.word8 (replicate (8 - mod8 written) 0)++{-------------------------------------------------------------------------------+ Deserialisation+-------------------------------------------------------------------------------}++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.fromSBS').+-}+fromSBS :: ShortByteString -> Either String (NumEntries, IndexCompact)+fromSBS (SBS ba') = do+ let ba = ByteArray ba'+ let len8 = sizeofByteArray ba+ when (mod8 len8 /= 0) $ Left "Length is not multiple of 64 bit"+ when (len8 < 24) $ Left "Doesn't contain header and footer"++ -- check type and version+ let typeAndVersion = indexByteArray ba 0 :: Word32+ when (typeAndVersion == byteSwap32 supportedTypeAndVersion)+ (Left "Non-matching endianness")+ when (typeAndVersion /= supportedTypeAndVersion)+ (Left "Unsupported type or version")++ -- read footer+ let len64 = div8 len8+ let getPositive off64 = do+ let w = indexByteArray ba off64 :: Word64+ when (w > fromIntegral (maxBound :: Int)) $+ Left "Size information is too large for Int"+ pure (fromIntegral w)++ numPages <- getPositive (len64 - 2)+ numEntries <- getPositive (len64 - 1)++ -- read vectors+ -- offsets in 64 bits+ let off1_64 = 1 -- after type–version indicator+ (!off2_64, icPrimary) <- getVec64 "Primary array" ba off1_64 numPages+ -- offsets in 64 bits+ let !off3 = off2_64+ (!off4, icClashes) <- getBitVec "Clash bit vector" ba off3 numPages+ (!off5, icLargerThanPage) <- getBitVec "LTP bit vector" ba off4 numPages+ (!off6, icTieBreaker) <- getTieBreaker ba off5++ let bytesUsed = mul8 (off6 + 2)+ when (bytesUsed > sizeofByteArray ba) $+ Left "Byte array is too small for components"+ when (bytesUsed < sizeofByteArray ba) $+ Left "Byte array is too large for components"++ pure (NumEntries numEntries, IndexCompact {..})++type Offset32 = Int+type Offset64 = Int++getVec64 ::+ String -> ByteArray -> Offset32 -> Int+ -> Either String (Offset64, VU.Vector Word64)+getVec64 name ba off64 numEntries =+ case checkedPrimVec off64 numEntries ba of+ Nothing -> Left (name <> " is out of bounds")+ Just vec -> Right (off64 + numEntries, VU.V_Word64 vec)++getBitVec ::+ String -> ByteArray -> Offset64 -> Int+ -> Either String (Offset64, VU.Vector Bit)+getBitVec name ba off numEntries =+ case checkedBitVec (mul64 off) numEntries ba of+ Nothing -> Left (name <> " is out of bounds")+ Just vec -> Right (off + ceilDiv64 numEntries, vec)++-- | Checks bounds.+--+-- Inefficient, but okay for a small number of entries.+getTieBreaker ::+ ByteArray -> Offset64+ -> Either String (Offset64, Map (Unsliced SerialisedKey) PageNo)+getTieBreaker ba = \off -> do+ when (mul8 off >= sizeofByteArray ba) $+ Left "Tie breaker is out of bounds"+ let size = fromIntegral (indexByteArray ba off :: Word64)+ (off', pairs) <- go size (off + 1) []+ pure (off', Map.fromList pairs)+ where+ go :: Int -> Offset64 -> [(Unsliced SerialisedKey, PageNo)]+ -> Either String (Offset64, [(Unsliced SerialisedKey, PageNo)])+ go 0 off pairs = pure (off, pairs)+ go n off pairs = do+ when (mul8 off >= sizeofByteArray ba) $+ Left "Clash map entry is out of bounds"+ let off32 = mul2 off+ let !pageNo = fromIntegral (indexByteArray ba off32 :: Word32)+ let keyLen8 = fromIntegral (indexByteArray ba (off32 + 1) :: Word32)++ (off', key) <- getKey (off + 1) keyLen8+ go (n - 1) off' ((key, PageNo pageNo) : pairs)++ getKey :: Offset64 -> Int -> Either String (Offset64, Unsliced SerialisedKey)+ getKey off len8 = do+ let off8 = mul8 off+ -- We avoid retaining references to the bytearray.+ -- Probably not needed, since the bytearray will stay alive as long as+ -- the compact index anyway, but we expect very few keys in the tie+ -- breaker, so it is cheap and we don't have to worry about it any more.+ !key <- case checkedPrimVec off8 len8 ba of+ Nothing -> Left ("Clash map key is out of bounds")+ Just vec -> Right (SerialisedKey' vec)+ pure (off + ceilDiv8 len8, makeUnslicedKey key)++-- | Offset and length are in number of elements.+checkedPrimVec :: forall a.+ VP.Prim a => Int -> Int -> ByteArray -> Maybe (VP.Vector a)+checkedPrimVec off len ba+ | off >= 0, sizeOf (undefined :: a) * (off + len) <= sizeofByteArray ba =+ Just (mkPrimVector off len ba)+ | otherwise =+ Nothing++-- | Offset and length are in number of bits.+--+-- We can't use 'checkedPrimVec' here, since 'Bool' and 'Bit' are not 'VP.Prim'+-- (so the bit vector type doesn't use 'VP.Vector' under the hood).+checkedBitVec :: Int -> Int -> ByteArray -> Maybe (VU.Vector Bit)+checkedBitVec off len ba+ | off >= 0, off + len <= mul8 (sizeofByteArray ba) =+ Just (BitVec off len ba)+ | otherwise =+ Nothing++{-------------------------------------------------------------------------------+ Vector extras+-------------------------------------------------------------------------------}++-- | Find the largest vector element that is smaller or equal to to the given+-- one, and return its vector index.+--+-- Note: this function uses 'unsafeThaw', so all considerations for using+-- 'unsafeThaw' apply to using 'unsafeSearchLE' too.+--+-- PRECONDITION: the vector is sorted in ascending order.+unsafeSearchLE ::+ (VG.Vector v e, Ord e)+ => e -> v e -> Maybe Int -- TODO: return -1?+unsafeSearchLE e vec = runST $ do+ -- Vector search algorithms work on mutable vectors only.+ vec' <- VG.unsafeThaw vec+ -- @i@ is the first index where @e@ is strictly smaller than the element at+ -- @i@.+ i <- VA.gallopingSearchLeftP (> e) vec'+ -- The last (and therefore largest) element that is lesser-equal @e@ is+ -- @i-1@. However, if @i==lb@, then the interval @[lb, ub)@ doesn't contain+ -- any elements that are lesser-equal @e@.+ pure $ if i == 0 then Nothing else Just (i - 1)++-- | Return the index of the last bit in the vector with the specified value, if+-- any.+--+-- TODO(optimise): optimise by implementing this function similarly to how+-- 'bitIndex' is implemented internally. Another alternative I tried is using+-- the @vector-rotvec@ package and 'V.elemIndex', but 'V.elemIndex' is up to 64x+-- slower than bitIndex.+bitIndexFromToRev :: Bound Int -> Bound Int -> Bit -> VU.Vector Bit -> Maybe Int+bitIndexFromToRev lb ub b v = reverseIx <$> bitIndex b (VU.reverse $ VU.slice lb' (ub' - lb' + 1) v)+ where+ reverseIx x = ub' - x+ lb' = vectorLowerBound lb+ ub' = vectorUpperBound v ub++-- | Like 'bitIndex', but only searches the vector within the give lower and+-- upper bound. Returns the index into the original vector, not the slice.+bitIndexFromTo :: Bound Int -> Bound Int -> Bit -> VU.Vector Bit -> Maybe Int+bitIndexFromTo lb ub b v = shiftIx <$> bitIndex b (VU.slice lb' (ub' - lb' + 1) v)+ where+ shiftIx = (lb'+)+ lb' = vectorLowerBound lb+ ub' = vectorUpperBound v ub++-- | Find the longest prefix of the vector that consists of only bits matching+-- the given value. The return value is the index of the last bit in the prefix.+bitLongestPrefixFromTo :: Bound Int -> Bound Int -> Bit -> Vector Bit -> Int+bitLongestPrefixFromTo lb ub b v = maybe ub' (subtract 1) $ bitIndexFromTo lb ub (toggle b) v+ where+ toggle (Bit x) = Bit (not x)+ ub' = vectorUpperBound v ub++vectorLowerBound :: Bound Int -> Int+vectorLowerBound = \case+ NoBound -> 0+ BoundExclusive i -> i + 1+ BoundInclusive i -> i++vectorUpperBound :: VG.Vector v a => v a -> Bound Int -> Int+vectorUpperBound v = \case+ NoBound -> VG.length v - 1+ BoundExclusive i -> i - 1+ BoundInclusive i -> i
@@ -0,0 +1,295 @@+{-# LANGUAGE CPP #-}+{-# OPTIONS_HADDOCK not-home #-}++-- |+-- Incremental construction of a compact index yields chunks of the primary array+-- that can be serialised incrementally.+--+-- Incremental construction is an 'ST' computation that can be started using+-- 'new', returning an 'IndexCompactAcc' structure that accumulates internal+-- state. 'append'ing new pages to the 'IndexCompactAcc' /might/ yield 'Chunk's.+-- Incremental construction can be finalised with 'unsafeEnd', which yields both+-- a 'Chunk' (possibly) and the `IndexCompact'.+--+module Database.LSMTree.Internal.Index.CompactAcc (+ -- * Construction+ IndexCompactAcc (..)+ , new+ , newWithDefaults+ , appendSingle+ , appendMulti+ , unsafeEnd+ -- * Utility+ , SMaybe (..)+ , smaybe+ -- * Internal: exported for testing and benchmarking+ , unsafeWriteRange+ , vectorLowerBound+ , mvectorUpperBound+ ) where++#ifdef NO_IGNORE_ASSERTS+import Control.Exception (assert)+#endif++import Control.Monad (when)+import Control.Monad.ST.Strict+import Data.Bit hiding (flipBit)+import Data.List.NonEmpty (NonEmpty)+import qualified Data.List.NonEmpty as NE+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Primitive.ByteArray (newPinnedByteArray, setByteArray)+import Data.STRef.Strict+import qualified Data.Vector.Generic.Mutable as VGM+import qualified Data.Vector.Primitive.Mutable as VPM+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word+import Database.LSMTree.Internal.BitMath+import Database.LSMTree.Internal.Chunk (Chunk)+import Database.LSMTree.Internal.Index.Compact+import Database.LSMTree.Internal.Map.Range (Bound (..))+import Database.LSMTree.Internal.Page+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Unsliced++{-------------------------------------------------------------------------------+ Construction+-------------------------------------------------------------------------------}++-- | A mutable version of 'IndexCompact'. See [incremental+-- construction](#incremental).+data IndexCompactAcc s = IndexCompactAcc {+ -- * Core index structure+ -- | Accumulates pinned chunks of 'ciPrimary'.+ icaPrimary :: !(STRef s (NonEmpty (VU.MVector s Word64)))+ -- | Accumulates chunks of 'ciClashes'.+ , icaClashes :: !(STRef s (NonEmpty (VU.MVector s Bit)))+ -- | Accumulates the 'ciTieBreaker'.+ , icaTieBreaker :: !(STRef s (Map (Unsliced SerialisedKey) PageNo))+ -- | Accumulates chunks of 'ciLargerThanPage'.+ , icaLargerThanPage :: !(STRef s (NonEmpty (VU.MVector s Bit)))++ -- * Aux information required for incremental construction+ -- | Maximum size of a chunk+ , icaMaxChunkSize :: !Int+ -- | The number of the current disk page we are constructing the index for.+ , icaCurrentPageNumber :: !(STRef s Int)+ -- | The primary bits of the page-maximum key that we saw last.+ --+ -- This should be 'SNothing' if we haven't seen any keys/pages yet.+ , icaLastMaxPrimbits :: !(STRef s (SMaybe Word64))+ -- | The ful minimum key of the page that we saw last.+ --+ -- This should be 'SNothing' if we haven't seen any keys/pages yet.+ , icaLastMinKey :: !(STRef s (SMaybe SerialisedKey))+ }++-- | @'new' maxcsize@ creates a new mutable index with a maximum chunk size of+-- @maxcsize@.+--+-- PRECONDITION: maxcsize > 0+--+-- Note: after initialisation, @maxcsize@ can no longer be changed.+new ::Int -> ST s (IndexCompactAcc s)+new maxcsize = IndexCompactAcc+ -- Core index structure+ <$> (newSTRef . pure =<< newPinnedMVec64 maxcsize)+ <*> (newSTRef . pure =<< VUM.new maxcsize)+ <*> newSTRef Map.empty+ <*> (newSTRef . pure =<< VUM.new maxcsize)+ -- Aux information required for incremental construction+ <*> pure maxcsize+ <*> newSTRef 0+ <*> newSTRef SNothing+ <*> newSTRef SNothing++-- | We explicitly pin the byte arrays, since that allows for more efficient+-- serialisation, as the definition of 'isByteArrayPinned' changed in GHC 9.6,+-- see <https://gitlab.haskell.org/ghc/ghc/-/issues/22255>.+--+-- TODO: remove this workaround once a solution exists, e.g. a new primop that+-- allows checking for implicit pinning.+newPinnedMVec64 :: Int -> ST s (VUM.MVector s Word64)+newPinnedMVec64 lenWords = do+ mba <- newPinnedByteArray (mul8 lenWords)+ setByteArray mba 0 lenWords (0 :: Word64)+ pure (VUM.MV_Word64 (VPM.MVector 0 lenWords mba))++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.newWithDefaults').+-}+newWithDefaults :: ST s (IndexCompactAcc s)+newWithDefaults = new 1024++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.appendSingle').+-}+appendSingle :: forall s. (SerialisedKey, SerialisedKey) -> IndexCompactAcc s -> ST s (Maybe Chunk)+appendSingle (minKey, maxKey) ica@IndexCompactAcc{..} = do+#ifdef NO_IGNORE_ASSERTS+ lastMinKey <- readSTRef icaLastMinKey+ assert (minKey <= maxKey && smaybe True (<= minKey) lastMinKey) $ pure () -- sorted+#endif+ pageNo <- readSTRef icaCurrentPageNumber+ let ix = pageNo `mod` icaMaxChunkSize+ goAppend pageNo ix+ writeSTRef icaCurrentPageNumber $! pageNo + 1+ yield ica+ where+ minPrimbits, maxPrimbits :: Word64+ minPrimbits = keyTopBits64 minKey+ maxPrimbits = keyTopBits64 maxKey++ -- | Meat of the function+ goAppend ::+ Int -- ^ Current /global/ page number+ -> Int -- ^ Current /local/ page number (inside the current chunk)+ -> ST s ()+ goAppend pageNo ix = do+ writePrimary+ writeClashesAndLTP+ where+ -- | Set value in primary vector+ writePrimary :: ST s ()+ writePrimary =+ readSTRef icaPrimary >>= \cs -> VUM.write (NE.head cs) ix minPrimbits++ -- | Set value in clash vector, tie-breaker map and larger-than-page+ -- vector+ writeClashesAndLTP :: ST s ()+ writeClashesAndLTP = do+ lastMaxPrimbits <- readSTRef icaLastMaxPrimbits+ let clash = lastMaxPrimbits == SJust minPrimbits+ writeSTRef icaLastMaxPrimbits $! SJust maxPrimbits++ lastMinKey <- readSTRef icaLastMinKey+ let ltp = SJust minKey == lastMinKey+ writeSTRef icaLastMinKey $! SJust minKey++ readSTRef icaClashes >>= \cs -> VUM.write (NE.head cs) ix (Bit clash)+ readSTRef icaLargerThanPage >>= \cs -> VUM.write (NE.head cs) ix (Bit ltp)+ when (clash && not ltp) $+ modifySTRef' icaTieBreaker (Map.insert (makeUnslicedKey minKey) (PageNo pageNo))++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.appendMulti').+-}+appendMulti :: forall s. (SerialisedKey, Word32) -> IndexCompactAcc s -> ST s [Chunk]+appendMulti (k, n0) ica@IndexCompactAcc{..} =+ maybe id (:) <$> appendSingle (k, k) ica <*> overflows (fromIntegral n0)+ where+ minPrimbits :: Word64+ minPrimbits = keyTopBits64 k++ -- | Fill primary, clash and LTP vectors for a larger-than-page value. Yields+ -- chunks if necessary+ overflows :: Int -> ST s [Chunk]+ overflows n+ | n <= 0 = pure []+ | otherwise = do+ pageNo <- readSTRef icaCurrentPageNumber+ let ix = pageNo `mod` icaMaxChunkSize -- will be 0 in recursive calls+ remInChunk = min n (icaMaxChunkSize - ix)+ readSTRef icaPrimary >>= \cs ->+ unsafeWriteRange (NE.head cs) (BoundInclusive ix) (BoundExclusive $ ix + remInChunk) minPrimbits+ readSTRef icaClashes >>= \cs ->+ unsafeWriteRange (NE.head cs) (BoundInclusive ix) (BoundExclusive $ ix + remInChunk) (Bit True)+ readSTRef icaLargerThanPage >>= \cs ->+ unsafeWriteRange (NE.head cs) (BoundInclusive ix) (BoundExclusive $ ix + remInChunk) (Bit True)+ writeSTRef icaCurrentPageNumber $! pageNo + remInChunk+ res <- yield ica+ maybe id (:) res <$> overflows (n - remInChunk)++-- | Yield a chunk and start a new one if the current chunk is already full.+--+-- TODO(optimisation): yield will eagerly allocate new mutable vectors, but+-- maybe that should be done lazily.+--+-- INVARIANTS: see [construction invariants](#construction-invariants).+yield :: IndexCompactAcc s -> ST s (Maybe Chunk)+yield IndexCompactAcc{..} = do+ pageNo <- readSTRef icaCurrentPageNumber+ if pageNo `mod` icaMaxChunkSize == 0 then do -- The current chunk is full+ primaryChunk <- VU.unsafeFreeze . NE.head =<< readSTRef icaPrimary+ modifySTRef' icaPrimary . NE.cons =<< newPinnedMVec64 icaMaxChunkSize+ modifySTRef' icaClashes . NE.cons =<< VUM.new icaMaxChunkSize+ modifySTRef' icaLargerThanPage . NE.cons =<< VUM.new icaMaxChunkSize+ pure $ Just (word64VectorToChunk primaryChunk)+ else -- the current chunk is not yet full+ pure Nothing++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.unsafeEnd').+-}+unsafeEnd :: IndexCompactAcc s -> ST s (Maybe Chunk, IndexCompact)+unsafeEnd IndexCompactAcc{..} = do+ pageNo <- readSTRef icaCurrentPageNumber+ let ix = pageNo `mod` icaMaxChunkSize++ chunksPrimary <-+ traverse VU.unsafeFreeze . sliceCurrent ix =<< readSTRef icaPrimary+ chunksClashes <-+ traverse VU.unsafeFreeze . sliceCurrent ix =<< readSTRef icaClashes+ chunksLargerThanPage <-+ traverse VU.unsafeFreeze . sliceCurrent ix =<< readSTRef icaLargerThanPage++ -- Only slice out a chunk if there are entries in the chunk+ let mchunk = if ix == 0+ then Nothing+ else Just (word64VectorToChunk (head chunksPrimary))++ let icPrimary = VU.concat . reverse $ chunksPrimary+ let icClashes = VU.concat . reverse $ chunksClashes+ let icLargerThanPage = VU.concat . reverse $ chunksLargerThanPage+ icTieBreaker <- readSTRef icaTieBreaker++ pure (mchunk, IndexCompact {..})+ where+ -- The current (most recent) chunk of the bitvectors is only partially+ -- constructed, so we need to only use the part that is already filled.+ sliceCurrent ix (c NE.:| cs)+ | ix == 0 = cs -- current chunk is completely empty, just ignore it+ | otherwise = VUM.slice 0 ix c : cs++{-------------------------------------------------------------------------------+ Strict 'Maybe'+-------------------------------------------------------------------------------}++data SMaybe a = SNothing | SJust !a+ deriving stock (Eq, Show)++smaybe :: b -> (a -> b) -> SMaybe a -> b+smaybe snothing sjust = \case+ SNothing -> snothing+ SJust x -> sjust x++{-------------------------------------------------------------------------------+ Vector extras+-------------------------------------------------------------------------------}++unsafeWriteRange :: VU.Unbox a => VU.MVector s a -> Bound Int -> Bound Int -> a -> ST s ()+unsafeWriteRange !v !lb !ub !x = VUM.set (VUM.unsafeSlice lb' len v) x+ where+ !lb' = vectorLowerBound lb+ !ub' = mvectorUpperBound v ub+ !len = ub' - lb' + 1++-- | Map a 'Bound' to the equivalent inclusive lower bound.+vectorLowerBound :: Bound Int -> Int+vectorLowerBound = \case+ NoBound -> 0+ BoundExclusive i -> i + 1+ BoundInclusive i -> i++-- | Map a 'Bound' to the equivalent inclusive upper bound.+mvectorUpperBound :: VGM.MVector v a => v s a -> Bound Int -> Int+mvectorUpperBound v = \case+ NoBound -> VGM.length v - 1+ BoundExclusive i -> i - 1+ BoundInclusive i -> i
@@ -0,0 +1,244 @@+{-# OPTIONS_HADDOCK not-home #-}+{- HLINT ignore "Avoid restricted alias" -}++{-|+ A general-purpose fence pointer index.++ Keys used with an ordinary index must be smaller than 64 KiB.+-}+module Database.LSMTree.Internal.Index.Ordinary+(+ IndexOrdinary (IndexOrdinary),+ toUnslicedLastKeys,+ search,+ sizeInPages,+ headerLBS,+ finalLBS,+ fromSBS+)+where++import Prelude hiding (drop, last, length)++import Control.DeepSeq (NFData (rnf))+import Control.Exception (assert)+import Control.Monad (when)+import Data.ByteString.Builder (toLazyByteString)+import Data.ByteString.Builder.Extra (word32Host, word64Host)+import Data.ByteString.Lazy (LazyByteString)+import Data.ByteString.Short (ShortByteString (SBS))+import qualified Data.ByteString.Short as ShortByteString (length)+import Data.Primitive.ByteArray (ByteArray (ByteArray),+ indexByteArray)+import Data.Vector (Vector, drop, findIndex, findIndexR, fromList,+ last, length, (!))+import qualified Data.Vector.Primitive as Primitive (Vector (Vector), drop,+ force, length, null, splitAt, take)+import Data.Word (Word16, Word32, Word64, Word8, byteSwap32)+import Database.LSMTree.Internal.Entry (NumEntries (NumEntries),+ unNumEntries)+import Database.LSMTree.Internal.Page (NumPages (NumPages),+ PageNo (PageNo), PageSpan (PageSpan))+import Database.LSMTree.Internal.Serialise+ (SerialisedKey (SerialisedKey'))+import Database.LSMTree.Internal.Unsliced (Unsliced, makeUnslicedKey)+import Database.LSMTree.Internal.Vector (binarySearchL, mkPrimVector)++{-|+ The type–version indicator for the ordinary index and its serialisation+ format as supported by this module.+-}+supportedTypeAndVersion :: Word32+supportedTypeAndVersion = 0x0101++{-|+ A general-purpose fence pointer index.++ An index is represented by a vector that maps the number of each page to the+ key stored last in this page or, if the page is an overflow page, to the key+ of the corresponding key–value pair. The vector must have the following+ properties:++ * It is non-empty.++ * Its elements are non-decreasing.++ This restriction follows from the fact that a run must contain keys in+ ascending order and must comprise at least one page for 'search' to be able+ to return a valid page span.+-}+newtype IndexOrdinary = IndexOrdinary (Vector (Unsliced SerialisedKey))+ deriving stock (Eq, Show)++instance NFData IndexOrdinary where++ rnf (IndexOrdinary unslicedLastKeys) = rnf unslicedLastKeys++toUnslicedLastKeys :: IndexOrdinary -> Vector (Unsliced SerialisedKey)+toUnslicedLastKeys (IndexOrdinary unslicedLastKeys) = unslicedLastKeys++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.search').+-}+search :: SerialisedKey -> IndexOrdinary -> PageSpan+search key (IndexOrdinary unslicedLastKeys)+ -- TODO: ideally, we could assert that an index is never empty, but+ -- unfortunately we can not currently do this. Runs (and thefeore indexes)+ -- /can/ be empty if they were created by a last-level merge where all input+ -- entries were deletes. Other parts of the @lsm-tree@ code won't fail as long+ -- as we return @PageSpan 0 0@ when we search an empty ordinary index. The+ -- ideal fix would be to remove empty runs from the levels entirely, but this+ -- requires more involved changes to the merge schedule and until then we'll+ -- just hack the @pageCount <= 0@ case in.+ | pageCount <= 0 = PageSpan (PageNo 0) (PageNo 0)+ | otherwise = assert (pageCount > 0) result where++ protoStart :: Int+ !protoStart = binarySearchL unslicedLastKeys (makeUnslicedKey key)++ pageCount :: Int+ !pageCount = length unslicedLastKeys++ result :: PageSpan+ result | protoStart < pageCount+ = let++ unslicedResultKey :: Unsliced SerialisedKey+ !unslicedResultKey = unslicedLastKeys ! protoStart++ end :: Int+ !end = maybe (pred pageCount) (+ protoStart) $+ findIndex (/= unslicedResultKey) $+ drop (succ protoStart) unslicedLastKeys++ in PageSpan (PageNo $ protoStart)+ (PageNo $ end)+ | otherwise+ = let++ unslicedResultKey :: Unsliced SerialisedKey+ !unslicedResultKey = last unslicedLastKeys++ start :: Int+ !start = maybe 0 succ $+ findIndexR (/= unslicedResultKey) $+ unslicedLastKeys++ in PageSpan (PageNo $ start)+ (PageNo $ pred pageCount)++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.sizeInPages').+-}+sizeInPages :: IndexOrdinary -> NumPages+sizeInPages (IndexOrdinary unslicedLastKeys)+ = NumPages $ fromIntegral (length unslicedLastKeys)++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.headerLBS').+-}+headerLBS :: LazyByteString+headerLBS = toLazyByteString $+ word32Host $+ supportedTypeAndVersion++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.finalLBS').+-}+finalLBS :: NumEntries -> IndexOrdinary -> LazyByteString+finalLBS entryCount _ = toLazyByteString $+ word64Host $+ fromIntegral $+ unNumEntries $+ entryCount++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.fromSBS').+-}+fromSBS :: ShortByteString -> Either String (NumEntries, IndexOrdinary)+fromSBS shortByteString@(SBS unliftedByteArray)+ | fullSize < 12+ = Left "Doesn't contain header and footer"+ | typeAndVersion == byteSwap32 supportedTypeAndVersion+ = Left "Non-matching endianness"+ | typeAndVersion /= supportedTypeAndVersion+ = Left "Unsupported type or version"+ | otherwise+ = (,) <$> entryCount <*> index+ where++ fullSize :: Int+ fullSize = ShortByteString.length shortByteString++ byteArray :: ByteArray+ byteArray = ByteArray unliftedByteArray++ fullBytes :: Primitive.Vector Word8+ fullBytes = mkPrimVector 0 fullSize byteArray++ typeAndVersion :: Word32+ typeAndVersion = indexByteArray byteArray 0++ postTypeAndVersionBytes :: Primitive.Vector Word8+ postTypeAndVersionBytes = Primitive.drop 4 fullBytes++ lastKeysBytes, entryCountBytes :: Primitive.Vector Word8+ (lastKeysBytes, entryCountBytes)+ = Primitive.splitAt (fullSize - 12) postTypeAndVersionBytes++ entryCount :: Either String NumEntries+ entryCount | toInteger asWord64 > toInteger (maxBound :: Int)+ = Left "Number of entries not representable as Int"+ | otherwise+ = Right (NumEntries (fromIntegral asWord64))+ where++ asWord64 :: Word64+ asWord64 = indexByteArray entryCountRep 0++ entryCountRep :: ByteArray+ Primitive.Vector _ _ entryCountRep = Primitive.force entryCountBytes++ index :: Either String IndexOrdinary+ index = IndexOrdinary <$> fromList <$> unslicedLastKeys lastKeysBytes++ unslicedLastKeys :: Primitive.Vector Word8+ -> Either String [Unsliced SerialisedKey]+ unslicedLastKeys bytes+ | Primitive.null bytes+ = Right []+ | otherwise+ = do+ when (Primitive.length bytes < 2)+ (Left "Too few bytes for key size")+ let++ firstSizeRep :: ByteArray+ Primitive.Vector _ _ firstSizeRep+ = Primitive.force (Primitive.take 2 bytes)++ firstSize :: Int+ firstSize = fromIntegral $+ (indexByteArray firstSizeRep 0 :: Word16)++ postFirstSizeBytes :: Primitive.Vector Word8+ postFirstSizeBytes = Primitive.drop 2 bytes++ when (Primitive.length postFirstSizeBytes < firstSize)+ (Left "Too few bytes for key")+ let++ firstBytes, othersBytes :: Primitive.Vector Word8+ (firstBytes, othersBytes)+ = Primitive.splitAt firstSize postFirstSizeBytes++ first :: Unsliced SerialisedKey+ !first = makeUnslicedKey (SerialisedKey' firstBytes)++ others <- unslicedLastKeys othersBytes+ pure (first : others)
@@ -0,0 +1,135 @@+{-# LANGUAGE CPP #-}+{-# OPTIONS_HADDOCK not-home #-}+{- HLINT ignore "Avoid restricted alias" -}++{-|+ Incremental construction functionality for the general-purpose fence pointer+ index.+-}+module Database.LSMTree.Internal.Index.OrdinaryAcc+(+ IndexOrdinaryAcc (IndexOrdinaryAcc),+ new,+ newWithDefaults,+ appendSingle,+ appendMulti,+ unsafeEnd+)+where++import Prelude hiding (take)++import Control.Exception (assert)+import Control.Monad.ST.Strict (ST)+import Data.Maybe (maybeToList)+import qualified Data.Vector.Primitive as Primitive (Vector, length)+import Data.Word (Word16, Word32, Word8)+import Database.LSMTree.Internal.Chunk (Baler, Chunk, createBaler,+ feedBaler, unsafeEndBaler)+import Database.LSMTree.Internal.Index.Ordinary+ (IndexOrdinary (IndexOrdinary))+import Database.LSMTree.Internal.Serialise+ (SerialisedKey (SerialisedKey'))+import Database.LSMTree.Internal.Unsliced (Unsliced, makeUnslicedKey)+import Database.LSMTree.Internal.Vector (byteVectorFromPrim)+import Database.LSMTree.Internal.Vector.Growing (GrowingVector)+import qualified Database.LSMTree.Internal.Vector.Growing as Growing (append,+ freeze, new)+#ifdef NO_IGNORE_ASSERTS+import Database.LSMTree.Internal.Unsliced (fromUnslicedKey)+import qualified Database.LSMTree.Internal.Vector.Growing as Growing+ (readMaybeLast)+#endif++{-|+ A general-purpose fence pointer index under incremental construction.++ A value @IndexOrdinaryAcc unslicedLastKeys baler@ denotes a partially+ constructed index that assigns keys to pages according to @unslicedLastKeys@+ and uses @baler@ for incremental output of the serialised key list.+-}+data IndexOrdinaryAcc s = IndexOrdinaryAcc+ !(GrowingVector s (Unsliced SerialisedKey))+ !(Baler s)++-- | Creates a new, initially empty, index.+new :: Int -- ^ Initial size of the key buffer+ -> Int -- ^ Minimum chunk size in bytes+ -> ST s (IndexOrdinaryAcc s) -- ^ Construction of the index+new initialKeyBufferSize minChunkSize = IndexOrdinaryAcc <$>+ Growing.new initialKeyBufferSize <*>+ createBaler minChunkSize++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.newWithDefaults').+-}+newWithDefaults :: ST s (IndexOrdinaryAcc s)+newWithDefaults = new 1024 4096++-- | Yields the serialisation of an element of a key list.+keyListElem :: SerialisedKey -> [Primitive.Vector Word8]+keyListElem (SerialisedKey' keyBytes) = [keySizeBytes, keyBytes] where++ keySize :: Int+ !keySize = Primitive.length keyBytes++ keySizeAsWord16 :: Word16+ !keySizeAsWord16 = assert (keySize <= fromIntegral (maxBound :: Word16)) $+ fromIntegral keySize++ keySizeBytes :: Primitive.Vector Word8+ !keySizeBytes = byteVectorFromPrim keySizeAsWord16++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.appendSingle').+-}+appendSingle :: (SerialisedKey, SerialisedKey)+ -> IndexOrdinaryAcc s+ -> ST s (Maybe Chunk)+appendSingle (firstKey, lastKey) (IndexOrdinaryAcc unslicedLastKeys baler)+ = assert (firstKey <= lastKey) $+ do+#ifdef NO_IGNORE_ASSERTS+ maybeLastUnslicedLastKey <- Growing.readMaybeLast unslicedLastKeys+ assert+ (all (< firstKey) (fromUnslicedKey <$> maybeLastUnslicedLastKey))+ (pure ())+#endif+ Growing.append unslicedLastKeys 1 (makeUnslicedKey lastKey)+ feedBaler (keyListElem lastKey) baler++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.appendMulti').+-}+appendMulti :: (SerialisedKey, Word32)+ -> IndexOrdinaryAcc s+ -> ST s [Chunk]+appendMulti (key, overflowPageCount) (IndexOrdinaryAcc unslicedLastKeys baler)+ = do+#ifdef NO_IGNORE_ASSERTS+ maybeLastUnslicedLastKey <- Growing.readMaybeLast unslicedLastKeys+ assert (all (< key) (fromUnslicedKey <$> maybeLastUnslicedLastKey))+ (pure ())+#endif+ Growing.append unslicedLastKeys pageCount (makeUnslicedKey key)+ maybeToList <$> feedBaler keyListElems baler+ where++ pageCount :: Int+ !pageCount = succ (fromIntegral overflowPageCount)++ keyListElems :: [Primitive.Vector Word8]+ keyListElems = concat (replicate pageCount (keyListElem key))++{-|+ For a specification of this operation, see the documentation of [its+ type-agnostic version]('Database.LSMTree.Internal.Index.unsafeEnd').+-}+unsafeEnd :: IndexOrdinaryAcc s -> ST s (Maybe Chunk, IndexOrdinary)+unsafeEnd (IndexOrdinaryAcc unslicedLastKeys baler) = do+ frozenUnslicedLastKeys <- Growing.freeze unslicedLastKeys+ remnant <- unsafeEndBaler baler+ pure (remnant, IndexOrdinary frozenUnslicedLastKeys)
@@ -0,0 +1,344 @@+{-# LANGUAGE CPP #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Lookup (+ LookupAcc+ , lookupsIOWithWriteBuffer+ , lookupsIO+ -- * Errors+ , TableCorruptedError (..)+ -- * Internal: exposed for tests and benchmarks+ , RunIx+ , KeyIx+ , RunIxKeyIx(..)+ , prepLookups+ , bloomQueries+ , indexSearches+ , intraPageLookupsWithWriteBuffer+ , intraPageLookupsOn+ ) where++import Data.Bifunctor+import Data.Primitive.ByteArray+import qualified Data.Vector as V+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed as VU+import Database.LSMTree.Internal.Arena (Arena, ArenaManager,+ allocateFromArena, withArena)++import Control.Exception (assert)+import Control.Monad+import Control.Monad.Class.MonadST as ST+import Control.Monad.Class.MonadThrow (Exception, MonadThrow (..))+import Control.Monad.Primitive+import Control.Monad.ST.Strict+import Control.RefCount++import Database.LSMTree.Internal.BlobRef (WeakBlobRef (..))+import Database.LSMTree.Internal.BloomFilter (Bloom, RunIxKeyIx (..),+ bloomQueries)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Index (Index)+import qualified Database.LSMTree.Internal.Index as Index (search)+import Database.LSMTree.Internal.Page (PageSpan (..), getNumPages,+ pageSpanSize, unPageNo)+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.Serialise+import qualified Database.LSMTree.Internal.Vector as V+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import System.FS.API (BufferOffset (..), Handle)+import System.FS.BlockIO.API++-- | Prepare disk lookups by doing bloom filter queries, index searches and+-- creating 'IOOp's. The result is a vector of 'IOOp's and a vector of indexes,+-- both of which are the same length. The indexes record the run and key+-- associated with each 'IOOp'.+prepLookups ::+ Arena s+ -> Bloom.Salt+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector Index+ -> V.Vector (Handle h)+ -> V.Vector SerialisedKey+ -> ST s (VP.Vector RunIxKeyIx, V.Vector (IOOp s h))+prepLookups arena salt blooms indexes kopsFiles ks = do+ let !rkixs = bloomQueries salt blooms ks+ !ioops <- indexSearches arena indexes kopsFiles ks rkixs+ pure (rkixs, ioops)++type KeyIx = Int+type RunIx = Int++-- | Perform a batch of fence pointer index searches, and create an 'IOOp' for+-- each search result. The resulting vector has the same length as the+-- @VP.Vector RunIxKeyIx@ argument, because index searching always returns a+-- positive search result.+indexSearches ::+ Arena s+ -> V.Vector Index+ -> V.Vector (Handle h)+ -> V.Vector SerialisedKey+ -> VP.Vector RunIxKeyIx -- ^ Result of 'bloomQueries'+ -> ST s (V.Vector (IOOp s h))+indexSearches !arena !indexes !kopsFiles !ks !rkixs = V.generateM n $ \i -> do+ let (RunIxKeyIx !rix !kix)+ = rkixs `VP.unsafeIndex` i+ !c = indexes `V.unsafeIndex` rix+ !h = kopsFiles `V.unsafeIndex` rix+ !k = ks `V.unsafeIndex` kix+ !pspan = Index.search k c+ !size = pageSpanSize pspan+ -- Acquire a reusable buffer+ (!off, !buf) <- allocateFromArena arena (getNumPages size * 4096) 4096+ pure $! IOOpRead+ h+ (fromIntegral $ unPageNo (pageSpanStart pspan) * 4096)+ buf+ (fromIntegral off)+ (getNumPages size * 4096)+ where+ !n = VP.length rkixs++type LookupAcc m h = V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h)))++{-# SPECIALISE lookupsIOWithWriteBuffer ::+ HasBlockIO IO h+ -> ArenaManager RealWorld+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs IO h)+ -> V.Vector (Ref (Run IO h))+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector Index+ -> V.Vector (Handle h)+ -> V.Vector SerialisedKey+ -> IO (LookupAcc IO h)+ #-}+-- | Like 'lookupsIO', but takes a write buffer into account.+lookupsIOWithWriteBuffer ::+ forall m h. (MonadThrow m, MonadST m)+ => HasBlockIO m h+ -> ArenaManager (PrimState m)+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs m h)+ -> V.Vector (Ref (Run m h)) -- ^ Runs @rs@+ -> V.Vector (Bloom SerialisedKey) -- ^ The bloom filters inside @rs@+ -> V.Vector Index -- ^ The indexes inside @rs@+ -> V.Vector (Handle h) -- ^ The file handles to the key\/value files inside @rs@+ -> V.Vector SerialisedKey+ -> m (LookupAcc m h)+lookupsIOWithWriteBuffer !hbio !mgr !resolveV !salt !wb !wbblobs !rs !blooms !indexes !kopsFiles !ks =+ assert precondition $+ withArena mgr $ \arena -> do+ (rkixs, ioops) <- ST.stToIO $ prepLookups arena salt blooms indexes kopsFiles ks+ ioress <- submitIO hbio ioops+ intraPageLookupsWithWriteBuffer resolveV wb wbblobs rs ks rkixs ioops ioress+ where+ -- we check only that the lengths match, because checking the contents is+ -- too expensive.+ precondition =+ assert (V.length rs == V.length blooms) $+ assert (V.length rs == V.length indexes) $+ assert (V.length rs == V.length kopsFiles) $+ True++{-# SPECIALISE lookupsIO ::+ HasBlockIO IO h+ -> ArenaManager RealWorld+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> V.Vector (Ref (Run IO h))+ -> V.Vector (Bloom SerialisedKey)+ -> V.Vector Index+ -> V.Vector (Handle h)+ -> V.Vector SerialisedKey+ -> IO (LookupAcc IO h)+ #-}+-- | Batched lookups in I\/O.+--+-- PRECONDITION: the vectors of bloom filters, indexes and file handles+-- should pointwise match with the vectors of runs.+lookupsIO ::+ forall m h. (MonadThrow m, MonadST m)+ => HasBlockIO m h+ -> ArenaManager (PrimState m)+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> V.Vector (Ref (Run m h)) -- ^ Runs @rs@+ -> V.Vector (Bloom SerialisedKey) -- ^ The bloom filters inside @rs@+ -> V.Vector Index -- ^ The indexes inside @rs@+ -> V.Vector (Handle h) -- ^ The file handles to the key\/value files inside @rs@+ -> V.Vector SerialisedKey+ -> m (LookupAcc m h)+lookupsIO !hbio !mgr !resolveV !salt !rs !blooms !indexes !kopsFiles !ks =+ assert precondition $+ withArena mgr $ \arena -> do+ (rkixs, ioops) <- ST.stToIO $ prepLookups arena salt blooms indexes kopsFiles ks+ ioress <- submitIO hbio ioops+ intraPageLookupsOn resolveV (V.map (const Nothing) ks) rs ks rkixs ioops ioress+ where+ -- we check only that the lengths match, because checking the contents is+ -- too expensive.+ precondition =+ assert (V.length rs == V.length blooms) $+ assert (V.length rs == V.length indexes) $+ assert (V.length rs == V.length kopsFiles) $+ True++{-# SPECIALISE intraPageLookupsWithWriteBuffer ::+ ResolveSerialisedValue+ -> WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs IO h)+ -> V.Vector (Ref (Run IO h))+ -> V.Vector SerialisedKey+ -> VP.Vector RunIxKeyIx+ -> V.Vector (IOOp RealWorld h)+ -> VU.Vector IOResult+ -> IO (LookupAcc IO h)+ #-}+-- | Like 'intraPageLookupsOn', but uses the write buffer as the initial+-- accumulator.+--+intraPageLookupsWithWriteBuffer ::+ forall m h. (PrimMonad m, MonadThrow m)+ => ResolveSerialisedValue+ -> WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs m h)+ -> V.Vector (Ref (Run m h))+ -> V.Vector SerialisedKey+ -> VP.Vector RunIxKeyIx+ -> V.Vector (IOOp (PrimState m) h)+ -> VU.Vector IOResult+ -> m (LookupAcc m h)+intraPageLookupsWithWriteBuffer !resolveV !wb !wbblobs !rs !ks !rkixs !ioops !ioress = do+ -- The most recent values are in the write buffer, so we use it to+ -- initialise the accumulator.+ acc0 <-+ V.generateM (V.length ks) $ \ki ->+ case WB.lookup wb (V.unsafeIndex ks ki) of+ Nothing -> pure Nothing+ Just e -> pure $! Just $! fmap (WBB.mkWeakBlobRef wbblobs) e+ -- TODO: ^^ we should be able to avoid this allocation by+ -- combining the conversion with other later conversions.+ intraPageLookupsOn resolveV acc0 rs ks rkixs ioops ioress++-- | The table data is corrupted.+data TableCorruptedError+ = ErrLookupByteCountDiscrepancy+ -- | Expected byte count.+ !ByteCount+ -- | Actual byte count.+ !ByteCount+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE intraPageLookupsOn ::+ ResolveSerialisedValue+ -> LookupAcc IO h+ -> V.Vector (Ref (Run IO h))+ -> V.Vector SerialisedKey+ -> VP.Vector RunIxKeyIx+ -> V.Vector (IOOp RealWorld h)+ -> VU.Vector IOResult+ -> IO (LookupAcc IO h)+ #-}+-- | Intra-page lookups, and combining lookup results from multiple runs and+-- a potential initial accumulator (e.g. from the write buffer).+--+-- This function assumes that @rkixs@ is ordered such that newer runs are+-- handled first. The order matters for resolving cases where we find the same+-- key in multiple runs.+--+intraPageLookupsOn ::+ forall m h. (PrimMonad m, MonadThrow m)+ => ResolveSerialisedValue+ -> LookupAcc m h -- initial acc+ -> V.Vector (Ref (Run m h))+ -> V.Vector SerialisedKey+ -> VP.Vector RunIxKeyIx+ -> V.Vector (IOOp (PrimState m) h)+ -> VU.Vector IOResult+ -> m (LookupAcc m h)+intraPageLookupsOn !resolveV !acc0 !rs !ks !rkixs !ioops !ioress =+ assert (V.length acc0 == V.length ks) $ do+ -- We accumulate results into the 'res' vector. When there are several+ -- lookup hits for the same key then we combine the results. The combining+ -- operator is associative but not commutative, so we must do this in the+ -- right order. We start with the write buffer lookup results and then go+ -- through the run lookup results in rkixs, which must be ordered by run.+ --+ -- TODO: reassess the representation of the result vector to try to reduce+ -- intermediate allocations. For example use a less convenient+ -- representation with several vectors (e.g. separate blob info) and+ -- convert to the final convenient representation in a single pass near+ -- the surface API so that all the conversions can be done in one pass+ -- without intermediate allocations.+ --+ res <- V.unsafeThaw acc0+ loop res 0+ V.unsafeFreeze res+ where+ !n = V.length ioops++ loop ::+ VM.MVector (PrimState m)+ (Maybe (Entry SerialisedValue (WeakBlobRef m h)))+ -> Int+ -> m ()+ loop !res !ioopix+ | ioopix == n = pure ()+ | otherwise = do+ let ioop = ioops `V.unsafeIndex` ioopix+ iores = ioress `VU.unsafeIndex` ioopix+ checkIOResult ioop iores+ let (RunIxKeyIx !rix !kix) = rkixs `VP.unsafeIndex` ioopix+ r = rs `V.unsafeIndex` rix+ k = ks `V.unsafeIndex` kix+ buf <- unsafeFreezeByteArray (ioopBuffer ioop)+ let boff = unBufferOffset $ ioopBufferOffset ioop+ case rawPageLookup (makeRawPage buf boff) k of+ LookupEntryNotPresent -> pure ()+ -- Laziness ensures that we only compute the forcing of the value in+ -- the entry when the result is needed.+ LookupEntry e -> do+ let e' = bimap copySerialisedValue (Run.mkWeakBlobRef r) e+ -- TODO: ^^ we should be able to avoid this allocation by+ -- combining the conversion with other later conversions.+ V.unsafeInsertWithMStrict res (combine resolveV) kix e'+ -- Laziness ensures that we only compute the appending of the prefix+ -- and suffix when the result is needed. We do not use 'force' here,+ -- since appending already creates a new primary vector.+ LookupEntryOverflow e m -> do+ let v' (SerialisedValue v) = SerialisedValue $ v <>+ RawBytes (V.mkPrimVector+ (unBufferOffset (ioopBufferOffset ioop) + 4096)+ (fromIntegral m)+ buf)+ e' = bimap v' (Run.mkWeakBlobRef r) e+ V.unsafeInsertWithMStrict res (combine resolveV) kix e'+ loop res (ioopix + 1)++ -- Check that the IOOp was performed successfully, and that it wrote/read+ -- exactly as many bytes as we expected. If not, then the buffer won't+ -- contain the correct number of disk-page bytes, so we throw an exception.+ checkIOResult :: IOOp (PrimState m) h -> IOResult -> m ()+ checkIOResult ioop (IOResult m) =+ unless (expected == m) . throwIO $+ ErrLookupByteCountDiscrepancy expected m+ where expected = ioopByteCount ioop++ -- Force a serialised value to not retain any memory by copying the+ -- underlying raw bytes.+ copySerialisedValue :: SerialisedValue -> SerialisedValue+ copySerialisedValue (SerialisedValue rb) =+ SerialisedValue (RB.copy rb)
@@ -0,0 +1,91 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_HADDOCK not-home #-}+module Database.LSMTree.Internal.Map.Range (+ Bound (.., BoundExclusive, BoundInclusive)+ , Clusive (..)+ , rangeLookup+ ) where++import Data.Map (Map)+import qualified Data.Map as Map+import Data.Map.Internal (Map (..))++data Clusive = Exclusive | Inclusive deriving stock Show+data Bound k = NoBound | Bound !k !Clusive deriving stock Show++{-# COMPLETE BoundExclusive, BoundInclusive #-}++pattern BoundExclusive :: k -> Bound k+pattern BoundExclusive k = Bound k Exclusive++pattern BoundInclusive :: k -> Bound k+pattern BoundInclusive k = Bound k Inclusive++-- | Find all the keys in the given range and return the corresponding+-- (key, value) pairs (in ascending order).+--+rangeLookup ::+ forall k v. Ord k+ => Bound k -- ^ lower bound+ -> Bound k -- ^ upper bound+ -> Map k v+ -> [(k, v)]+rangeLookup NoBound NoBound m = Map.toList m+rangeLookup (Bound lb lc) NoBound m = rangeLookupLo lb lc m []+rangeLookup NoBound (Bound ub uc) m = rangeLookupHi ub uc m []+rangeLookup (Bound lb lc) (Bound ub uc) m = rangeLookupBoth lb lc ub uc m []++toDList :: Map k v -> [(k, v)] -> [(k, v)]+toDList Tip = id+toDList (Bin _ k v l r) = toDList l . ((k,v):) . toDList r++rangeLookupLo :: Ord k => k -> Clusive -> Map k v -> [(k, v)] -> [(k, v)]+rangeLookupLo !_ !_ Tip = id+rangeLookupLo lb lc (Bin _ k v l r)+ -- ... | --- k -----+ | evalLowerBound lb lc k+ = rangeLookupLo lb lc l . ((k, v) :) . toDList r++ -- ... k ... |--------+ | otherwise+ = rangeLookupLo lb lc r++rangeLookupHi :: Ord k => k -> Clusive -> Map k v -> [(k, v)] -> [(k, v)]+rangeLookupHi !_ !_ Tip = id+rangeLookupHi ub uc (Bin _ k v l r)+ -- --- k --- | ...+ | evalUpperBound ub uc k+ = toDList l . ((k, v) :) . rangeLookupHi ub uc r++ -- --------- | ... k ...+ | otherwise+ = rangeLookupHi ub uc l++rangeLookupBoth :: Ord k => k -> Clusive -> k -> Clusive -> Map k v -> [(k, v)] -> [(k, v)]+rangeLookupBoth !_ !_ !_ !_ Tip = id+rangeLookupBoth lb lc ub uc (Bin _ k v l r)+ -- ... |--- k ---| ...+ | evalLowerBound lb lc k+ , evalUpperBound ub uc k+ = rangeLookupLo lb lc l . ((k,v):) . rangeLookupHi ub uc r++ -- ... |-------| ... k ...+ | evalLowerBound lb lc k+ = rangeLookupBoth lb lc ub uc l++ -- ... k ... |-------| ...+ | evalUpperBound ub uc k+ = rangeLookupBoth lb lc ub uc r++ | otherwise+ = id++evalLowerBound :: Ord k => k -> Clusive -> k -> Bool+evalLowerBound b Exclusive k = b < k+evalLowerBound b Inclusive k = b <= k++evalUpperBound :: Ord k => k -> Clusive -> k -> Bool+evalUpperBound b Exclusive k = k < b+evalUpperBound b Inclusive k = k <= b
@@ -0,0 +1,500 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | The 'Merge' type and its functions are not intended for concurrent use.+-- Concurrent access should therefore be sequentialised using a suitable+-- concurrency primitive, such as an 'MVar'.+--+module Database.LSMTree.Internal.Merge (+ Merge (..)+ , MergeType (..)+ , IsMergeType (..)+ , LevelMergeType (..)+ , TreeMergeType (..)+ , MergeState (..)+ , RunParams (..)+ , new+ , abort+ , complete+ , stepsToCompletion+ , stepsToCompletionCounted+ , StepResult (..)+ , steps+ , mergeRunParams+ ) where++import Control.DeepSeq (NFData (..))+import Control.Exception (assert)+import Control.Monad (when)+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadMask, MonadThrow)+import Control.Monad.Primitive (PrimState)+import Control.RefCount+import Data.Primitive.MutVar+import Data.Traversable (for)+import qualified Data.Vector as V+import Database.LSMTree.Internal.BlobRef (RawBlobRef)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Readers (Readers)+import qualified Database.LSMTree.Internal.Readers as Readers+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunBuilder (RunBuilder, RunParams)+import qualified Database.LSMTree.Internal.RunBuilder as Builder+import qualified Database.LSMTree.Internal.RunReader as Reader+import Database.LSMTree.Internal.Serialise+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import System.FS.BlockIO.API (HasBlockIO)++-- | An in-progress incremental k-way merge of 'Run's.+--+-- Since we always resolve all entries of the same key in one go, there is no+-- need to store incompletely-resolved entries.+data Merge t m h = Merge {+ mergeType :: !t+ -- | We also store @isLastLevel mergeType@ and @isUnion mergeType@ here to+ -- avoid recomputing them it again and again for each call to 'steps',+ -- which would also add an 'IsMergeType' constraint to large parts of the+ -- interface.+ , mergeIsLastLevel :: !Bool+ , mergeIsUnion :: !Bool+ , mergeResolve :: !ResolveSerialisedValue+ , mergeReaders :: {-# UNPACK #-} !(Readers m h)+ , mergeBuilder :: !(RunBuilder m h)+ -- | The result of the latest call to 'steps'. This is used to determine+ -- whether a merge can be 'complete'd.+ , mergeState :: !(MutVar (PrimState m) MergeState)+ , mergeHasFS :: !(HasFS m h)+ , mergeHasBlockIO :: !(HasBlockIO m h)+ }++mergeRunParams :: Merge t m h -> RunParams+mergeRunParams = Builder.runBuilderParams . mergeBuilder++-- | The current state of the merge.+data MergeState =+ -- | There is still merging work to be done+ Merging+ -- | There is no more merging work to be done, but the merge still has to be+ -- completed to yield a new run.+ | MergingDone+ -- | A run was yielded as the result of a merge. The merge is implicitly+ -- closed.+ | Completed+ -- | The merge was closed before it was completed.+ | Closed++-- | Merges can either exist on a level of the LSM, or be a union merge of two+-- tables.+class IsMergeType t where+ -- | A last level merge behaves differently from a mid-level merge: last level+ -- merges can actually remove delete operations, whereas mid-level merges must+ -- preserve them.+ isLastLevel :: t -> Bool+ -- | Union merges follow the semantics of @Data.Map.unionWith (<>)@. Since the+ -- input runs are semantically treated like @Data.Map@s, deletes are ignored+ -- and inserts act like mupserts, so they need to be merged monoidally using+ -- 'resolveValue'.+ isUnion :: t -> Bool++-- | Fully general merge type, mainly useful for testing.+data MergeType = MergeTypeMidLevel | MergeTypeLastLevel | MergeTypeUnion+ deriving stock (Eq, Show)++instance NFData MergeType where+ rnf MergeTypeMidLevel = ()+ rnf MergeTypeLastLevel = ()+ rnf MergeTypeUnion = ()++instance IsMergeType MergeType where+ isLastLevel = \case+ MergeTypeMidLevel -> False+ MergeTypeLastLevel -> True+ MergeTypeUnion -> True+ isUnion = \case+ MergeTypeMidLevel -> False+ MergeTypeLastLevel -> False+ MergeTypeUnion -> True++-- | Different types of merges created as part of a regular (non-union) level.+--+-- A last level merge behaves differently from a mid-level merge: last level+-- merges can actually remove delete operations, whereas mid-level merges must+-- preserve them. This is orthogonal to the 'MergePolicy'.+data LevelMergeType = MergeMidLevel | MergeLastLevel+ deriving stock (Eq, Show)++instance NFData LevelMergeType where+ rnf MergeMidLevel = ()+ rnf MergeLastLevel = ()++instance IsMergeType LevelMergeType where+ isLastLevel = \case+ MergeMidLevel -> False+ MergeLastLevel -> True+ isUnion = const False++-- | Different types of merges created as part of the merging tree.+data TreeMergeType = MergeLevel | MergeUnion+ deriving stock (Eq, Show)++instance NFData TreeMergeType where+ rnf MergeLevel = ()+ rnf MergeUnion = ()++instance IsMergeType TreeMergeType where+ isLastLevel = const True+ isUnion = \case+ MergeLevel -> False+ MergeUnion -> True++{-# SPECIALISE new ::+ IsMergeType t+ => HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> t+ -> ResolveSerialisedValue+ -> Run.RunFsPaths+ -> V.Vector (Ref (Run IO h))+ -> IO (Maybe (Merge t IO h)) #-}+-- | Returns 'Nothing' if no input 'Run' contains any entries.+-- The list of runs should be sorted from new to old.+new ::+ (IsMergeType t, MonadMask m, MonadSTM m, MonadST m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> RunParams+ -> t+ -> ResolveSerialisedValue+ -> Run.RunFsPaths+ -> V.Vector (Ref (Run m h))+ -> m (Maybe (Merge t m h))+new hfs hbio salt runParams mergeType mergeResolve targetPaths runs = do+ let sources = Readers.FromRun <$> V.toList runs+ mreaders <- Readers.new mergeResolve Readers.NoOffsetKey sources+ -- TODO: Exception safety! If Readers.new fails after already creating some+ -- run readers, or Builder.new fails, the run readers will stay open,+ -- holding handles of the input runs' files.+ for mreaders $ \mergeReaders -> do+ -- calculate upper bounds based on input runs+ let numEntries = V.foldMap' Run.size runs+ mergeBuilder <- Builder.new hfs hbio salt runParams targetPaths numEntries+ mergeState <- newMutVar $! Merging+ pure Merge {+ mergeIsLastLevel = isLastLevel mergeType+ , mergeIsUnion = isUnion mergeType+ , mergeHasFS = hfs+ , mergeHasBlockIO = hbio+ , ..+ }++{-# SPECIALISE abort :: Merge t IO (FS.Handle h) -> IO () #-}+-- | This function should be called when discarding a 'Merge' before it+-- was done (i.e. returned 'MergeComplete'). This removes the incomplete files+-- created for the new run so far and avoids leaking file handles.+--+-- Once it has been called, do not use the 'Merge' any more!+abort :: (MonadMask m, MonadSTM m, MonadST m) => Merge t m h -> m ()+abort Merge {..} = do+ readMutVar mergeState >>= \case+ Merging -> do+ Readers.close mergeReaders+ Builder.close mergeBuilder+ MergingDone -> do+ -- the readers are already drained, therefore closed+ Builder.close mergeBuilder+ Completed ->+ assert False $ pure ()+ Closed ->+ assert False $ pure ()+ writeMutVar mergeState $! Closed++{-# SPECIALISE complete ::+ Merge t IO h+ -> IO (Ref (Run IO h)) #-}+-- | Complete a 'Merge', returning a new 'Run' as the result of merging the+-- input runs.+--+-- All resources held by the merge are released, so do not use the it any more!+--+-- This function will /not/ do any merging work if there is any remaining. That+-- is, if not enough 'steps' were performed to exhaust the input 'Readers', this+-- function will throw an error.+--+-- Returns an error if the merge was not yet done, if it was already completed+-- before, or if it was already closed.+--+-- Note: this function creates new 'Run' resources, so it is recommended to run+-- this function with async exceptions masked. Otherwise, these resources can+-- leak. And it must eventually be released with 'releaseRef'.+--+complete ::+ (MonadSTM m, MonadST m, MonadMask m)+ => Merge t m h+ -> m (Ref (Run m h))+complete Merge{..} = do+ readMutVar mergeState >>= \case+ Merging -> error "complete: Merge is not done"+ MergingDone -> do+ -- the readers are already drained, therefore closed+ r <- Run.fromBuilder mergeBuilder+ writeMutVar mergeState $! Completed+ pure r+ Completed -> error "complete: Merge is already completed"+ Closed -> error "complete: Merge is closed"++{-# SPECIALISE stepsToCompletion ::+ Merge t IO h+ -> Int+ -> IO (Ref (Run IO h)) #-}+-- | Like 'steps', but calling 'complete' once the merge is finished.+--+-- Note: run with async exceptions masked. See 'complete'.+stepsToCompletion ::+ (MonadMask m, MonadSTM m, MonadST m)+ => Merge t m h+ -> Int+ -> m (Ref (Run m h))+stepsToCompletion m stepBatchSize = go+ where+ go = do+ steps m stepBatchSize >>= \case+ (_, MergeInProgress) -> go+ (_, MergeDone) -> complete m++{-# SPECIALISE stepsToCompletionCounted ::+ Merge t IO h+ -> Int+ -> IO (Int, Ref (Run IO h)) #-}+-- | Like 'steps', but calling 'complete' once the merge is finished.+--+-- Note: run with async exceptions masked. See 'complete'.+stepsToCompletionCounted ::+ (MonadMask m, MonadSTM m, MonadST m)+ => Merge t m h+ -> Int+ -> m (Int, Ref (Run m h))+stepsToCompletionCounted m stepBatchSize = go 0+ where+ go !stepsSum = do+ steps m stepBatchSize >>= \case+ (n, MergeInProgress) -> go (stepsSum + n)+ (n, MergeDone) -> let !stepsSum' = stepsSum + n+ in (stepsSum',) <$> complete m++data StepResult = MergeInProgress | MergeDone+ deriving stock Eq++stepsInvariant :: Int -> (Int, StepResult) -> Bool+stepsInvariant requestedSteps = \case+ (n, MergeInProgress) -> n >= requestedSteps+ _ -> True++{-# SPECIALISE steps ::+ Merge t IO h+ -> Int+ -> IO (Int, StepResult) #-}+-- | Do at least a given number of steps of merging. Each step reads a single+-- entry, then either resolves the previous entry with the new one or writes it+-- out to the run being created. Since we always finish resolving a key we+-- started, we might do slightly more work than requested.+--+-- Returns the number of input entries read, which is guaranteed to be at least+-- as many as requested (unless the merge is complete).+--+-- Returns an error if the merge was already completed or closed.+steps ::+ (MonadMask m, MonadSTM m, MonadST m)+ => Merge t m h+ -> Int -- ^ How many input entries to consume (at least)+ -> m (Int, StepResult)+steps m@Merge {..} requestedSteps = assertStepsInvariant <$> do+ -- TODO: ideally, we would not check whether the merge was already done on+ -- every call to @steps@. It is important for correctness, however, that we+ -- do not call @steps@ on a merge when it was already done. It is not yet+ -- clear whether our (upcoming) implementation of scheduled merges is going+ -- to satisfy this precondition when it calls @steps@, so for now we do the+ -- check.+ readMutVar mergeState >>= \case+ Merging -> if mergeIsUnion then doStepsUnion m requestedSteps+ else doStepsLevel m requestedSteps+ MergingDone -> pure (0, MergeDone)+ Completed -> error "steps: Merge is completed"+ Closed -> error "steps: Merge is closed"+ where+ assertStepsInvariant res = assert (stepsInvariant requestedSteps res) res++{-# SPECIALISE doStepsLevel ::+ Merge t IO h+ -> Int+ -> IO (Int, StepResult) #-}+doStepsLevel ::+ (MonadMask m, MonadSTM m, MonadST m)+ => Merge t m h+ -> Int -- ^ How many input entries to consume (at least)+ -> m (Int, StepResult)+doStepsLevel m@Merge {..} requestedSteps = go 0+ where+ go !n+ | n >= requestedSteps =+ pure (n, MergeInProgress)+ | otherwise = do+ (key, entry, hasMore) <- Readers.pop mergeResolve mergeReaders+ case hasMore of+ Readers.HasMore ->+ handleEntry (n + 1) key entry+ Readers.Drained -> do+ -- no future entries, no previous entry to resolve, just write!+ writeReaderEntry m key entry+ writeMutVar mergeState $! MergingDone+ pure (n + 1, MergeDone)++ handleEntry !n !key (Reader.Entry (Upsert v)) =+ -- resolve small mupsert vals with the following entries of the same key+ handleMupdate n key v+ handleEntry !n !key (Reader.EntryOverflow (Upsert v) _ len overflowPages) =+ -- resolve large mupsert vals with following entries of the same key+ handleMupdate n key (Reader.appendOverflow len overflowPages v)+ handleEntry !n !key entry = do+ -- otherwise, we can just drop all following entries of same key+ writeReaderEntry m key entry+ dropRemaining n key++ -- the value is from a mupsert, complete (not just a prefix)+ handleMupdate !n !key !v = do+ nextKey <- Readers.peekKey mergeReaders+ if nextKey /= key+ then do+ -- resolved all entries for this key, write it+ writeSerialisedEntry m key (Upsert v)+ go n+ else do+ (_, nextEntry, hasMore) <- Readers.pop mergeResolve mergeReaders+ -- for resolution, we need the full second value to be present+ let resolved = combine mergeResolve+ (Upsert v)+ (Reader.toFullEntry nextEntry)+ case hasMore of+ Readers.HasMore -> case resolved of+ Upsert v' ->+ -- still a mupsert, keep resolving+ handleMupdate (n + 1) key v'+ _ -> do+ -- done with this key, now the remaining entries are obsolete+ writeSerialisedEntry m key resolved+ dropRemaining (n + 1) key+ Readers.Drained -> do+ writeSerialisedEntry m key resolved+ writeMutVar mergeState $! MergingDone+ pure (n + 1, MergeDone)++ dropRemaining !n !key = do+ (dropped, hasMore) <- Readers.dropWhileKey mergeResolve mergeReaders key+ case hasMore of+ Readers.HasMore -> go (n + dropped)+ Readers.Drained -> do+ writeMutVar mergeState $! MergingDone+ pure (n + dropped, MergeDone)++{-# SPECIALISE doStepsUnion ::+ Merge t IO h+ -> Int+ -> IO (Int, StepResult) #-}+doStepsUnion ::+ (MonadMask m, MonadSTM m, MonadST m)+ => Merge t m h+ -> Int -- ^ How many input entries to consume (at least)+ -> m (Int, StepResult)+doStepsUnion m@Merge {..} requestedSteps = go 0+ where+ go !n+ | n >= requestedSteps =+ pure (n, MergeInProgress)+ | otherwise = do+ (key, entry, hasMore) <- Readers.pop mergeResolve mergeReaders+ handleEntry (n + 1) key entry hasMore++ -- Similar to 'handleMupdate' in 'stepsLevel', but here we have to combine+ -- all entries monoidally, so there are no obsolete/overwritten entries+ -- that we could skip.+ --+ -- TODO(optimisation): If mergeResolve is const, we could skip all remaining+ -- entries for the key. Unfortunately, we can't inspect the function. This+ -- would require encoding it as something like `Const | Resolve (_ -> _ ->+ -- _)`.+ handleEntry !n !key !entry Readers.Drained = do+ -- no future entries, no previous entry to resolve, just write!+ writeReaderEntry m key entry+ writeMutVar mergeState $! MergingDone+ pure (n, MergeDone)++ handleEntry !n !key !entry Readers.HasMore = do+ nextKey <- Readers.peekKey mergeReaders+ if nextKey /= key+ then do+ -- resolved all entries for this key, write it+ writeReaderEntry m key entry+ go n+ else do+ (_, nextEntry, hasMore) <- Readers.pop mergeResolve mergeReaders+ -- for resolution, we need the full second value to be present+ let resolved = combineUnion mergeResolve+ (Reader.toFullEntry entry)+ (Reader.toFullEntry nextEntry)+ handleEntry (n + 1) key (Reader.Entry resolved) hasMore++{-# INLINE writeReaderEntry #-}+writeReaderEntry ::+ (MonadSTM m, MonadST m, MonadThrow m)+ => Merge t m h+ -> SerialisedKey+ -> Reader.Entry m h+ -> m ()+writeReaderEntry m key (Reader.Entry entryFull) =+ -- Small entry.+ -- Note that this small entry could be the only one on the page. We only+ -- care about it being small, not single-entry, since it could still end+ -- up sharing a page with other entries in the merged run.+ -- TODO(optimise): This doesn't fully exploit the case where there is a+ -- single page small entry on the page which again ends up as the only+ -- entry of a page (which would for example happen a lot if most entries+ -- have 2k-4k bytes). In that case we could have copied the RawPage+ -- (but we find out too late to easily exploit it).+ writeSerialisedEntry m key entryFull+writeReaderEntry m key entry@(Reader.EntryOverflow prefix page _ overflowPages)+ | InsertWithBlob {} <- prefix =+ assert (shouldWriteEntry m prefix) $ do -- large, can't be delete+ -- has blob, we can't just copy the first page, fall back+ -- we simply append the overflow pages to the value+ Builder.addKeyOp (mergeBuilder m) key (Reader.toFullEntry entry)+ -- TODO(optimise): This copies the overflow pages unnecessarily.+ -- We could extend the RunBuilder API to allow to either:+ -- 1. write an Entry (containing the value prefix) + [RawOverflowPage]+ -- 2. write a RawPage + SerialisedBlob + [RawOverflowPage], rewriting+ -- the raw page's blob offset (slightly faster, but a bit hacky)+ | otherwise =+ assert (shouldWriteEntry m prefix) $ -- large, can't be delete+ -- no blob, directly copy all pages as they are+ Builder.addLargeSerialisedKeyOp (mergeBuilder m) key page overflowPages++{-# INLINE writeSerialisedEntry #-}+writeSerialisedEntry ::+ (MonadSTM m, MonadST m, MonadThrow m)+ => Merge t m h+ -> SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> m ()+writeSerialisedEntry m key entry =+ when (shouldWriteEntry m entry) $+ Builder.addKeyOp (mergeBuilder m) key entry++-- On the last level we could also turn Upsert into Insert, but no need to+-- complicate things.+shouldWriteEntry :: Merge t m h -> Entry v b -> Bool+shouldWriteEntry m Delete = not (mergeIsLastLevel m)+shouldWriteEntry _ _ = True
@@ -0,0 +1,998 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_HADDOCK not-home #-}++-- TODO: establish that this implementation matches up with the ScheduledMerges+-- prototype. See lsm-tree#445.+module Database.LSMTree.Internal.MergeSchedule (+ -- * Traces+ AtLevel (..)+ , MergeTrace (..)+ -- * Table content+ , TableContent (..)+ , duplicateTableContent+ , releaseTableContent+ -- * Levels cache+ , LevelsCache (..)+ , mkLevelsCache+ -- * Levels, runs and ongoing merges+ , Levels+ , Level (..)+ , MergePolicyForLevel (..)+ , mergingRunParamsForLevel+ -- * Union level+ , UnionLevel (..)+ -- * Union cache+ , UnionCache (..)+ , mkUnionCache+ , duplicateUnionCache+ , releaseUnionCache+ -- * Flushes and scheduled merges+ , updatesWithInterleavedFlushes+ , flushWriteBuffer+ -- * Exported for cabal-docspec+ , maxRunSize+ -- * Credits+ , MergeDebt (..)+ , MergeCredits (..)+ , supplyCredits+ , NominalDebt (..)+ , NominalCredits (..)+ , nominalDebtAsCredits+ , nominalDebtForLevel+ -- * Exported for testing+ , addWriteBufferEntries+ ) where++import Control.ActionRegistry+import Control.Concurrent.Class.MonadMVar.Strict+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..))+import Control.Monad.Primitive+import Control.RefCount+import Control.Tracer+import Data.Foldable (fold, traverse_)+import qualified Data.Vector as V+import Database.LSMTree.Internal.Assertions (assert)+import Database.LSMTree.Internal.BloomFilter (Bloom)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.Entry (Entry, NumEntries (..),+ unNumEntries)+import Database.LSMTree.Internal.IncomingRun+import Database.LSMTree.Internal.Index (Index)+import Database.LSMTree.Internal.MergingRun (MergeCredits (..),+ MergeDebt (..), MergingRun, RunParams (..))+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.MergingTree (MergingTree)+import qualified Database.LSMTree.Internal.MergingTree.Lookup as MT+import Database.LSMTree.Internal.Paths (ActiveDir, RunFsPaths (..),+ SessionRoot)+import qualified Database.LSMTree.Internal.Paths as Paths+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue,+ SerialisedBlob, SerialisedKey, SerialisedValue)+import Database.LSMTree.Internal.UniqCounter+import Database.LSMTree.Internal.Vector (forMStrict, mapMStrict,+ mapStrict)+import Database.LSMTree.Internal.WriteBuffer (WriteBuffer)+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import Database.LSMTree.Internal.WriteBufferBlobs (WriteBufferBlobs)+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import System.FS.BlockIO.API (HasBlockIO)++{-------------------------------------------------------------------------------+ Traces+-------------------------------------------------------------------------------}++data AtLevel a = AtLevel LevelNo a+ deriving stock (Show, Eq)++data MergeTrace =+ TraceFlushWriteBuffer+ NumEntries -- ^ Size of the write buffer+ RunNumber+ RunParams+ | TraceAddLevel+ | TraceAddRun+ RunNumber -- ^ newly added run+ (V.Vector RunNumber) -- ^ resident runs+ | TraceNewMerge+ (V.Vector NumEntries) -- ^ Sizes of input runs+ RunNumber+ RunParams+ MergePolicyForLevel+ MR.LevelMergeType+ | TraceNewMergeSingleRun+ NumEntries -- ^ Size of run+ RunNumber+ | TraceCompletedMerge -- TODO: currently not traced for Incremental merges+ NumEntries -- ^ Size of output run+ RunNumber+ -- | This is traced at the latest point the merge could complete.+ | TraceExpectCompletedMerge+ RunNumber+ deriving stock (Show, Eq)++{-------------------------------------------------------------------------------+ Table content+-------------------------------------------------------------------------------}++-- | The levels of the table, from most to least recently inserted.+--+-- Concurrency: read-only operations are allowed to be concurrent with each+-- other, but update operations must not be concurrent with each other or read+-- operations. For example, inspecting the levels cache can be done+-- concurrently, but 'updatesWithInterleavedFlushes' must be serialised.+--+data TableContent m h = TableContent {+ -- | The in-memory level 0 of the table+ --+ -- TODO: probably less allocation to make this a MutVar+ tableWriteBuffer :: !WriteBuffer+ -- | The blob storage for entries in the write buffer+ , tableWriteBufferBlobs :: !(Ref (WriteBufferBlobs m h))+ -- | A hierarchy of \"regular\" on-disk levels numbered 1 and up. Note that+ -- vector index @n@ refers to level @n+1@.+ , tableLevels :: !(Levels m h)+ -- | Cache of flattened regular 'levels'.+ , tableCache :: !(LevelsCache m h)+ -- | An optional final union level, containing its own cache.+ , tableUnionLevel :: !(UnionLevel m h)+ }++{-# SPECIALISE duplicateTableContent :: ActionRegistry IO -> TableContent IO h -> IO (TableContent IO h) #-}+duplicateTableContent ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> TableContent m h+ -> m (TableContent m h)+duplicateTableContent reg (TableContent wb wbb levels cache ul) = do+ wbb' <- withRollback reg (dupRef wbb) releaseRef+ levels' <- duplicateLevels reg levels+ cache' <- duplicateLevelsCache reg cache+ ul' <- duplicateUnionLevel reg ul+ pure $! TableContent wb wbb' levels' cache' ul'++{-# SPECIALISE releaseTableContent :: ActionRegistry IO -> TableContent IO h -> IO () #-}+releaseTableContent ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> TableContent m h+ -> m ()+releaseTableContent reg (TableContent _wb wbb levels cache ul) = do+ delayedCommit reg (releaseRef wbb)+ releaseLevels reg levels+ releaseLevelsCache reg cache+ releaseUnionLevel reg ul++{-------------------------------------------------------------------------------+ Levels cache+-------------------------------------------------------------------------------}++-- | Flattened cache of the runs that referenced by a table.+--+-- This cache includes a vector of runs, but also vectors of the runs broken+-- down into components, like bloom filters, fence pointer indexes and file+-- handles. This allows for quick access in the lookup code. Recomputing this+-- cache should be relatively rare.+--+-- Caches keep references to its runs on construction, and they release each+-- reference when the cache is invalidated. This is done so that incremental+-- merges can remove references for their input runs when a merge completes,+-- without closing runs that might be in use for other operations such as+-- lookups. This does mean that a cache can keep runs open for longer than+-- necessary, so caches should be rebuilt using, e.g., 'rebuildCache', in a+-- timely manner.+data LevelsCache m h = LevelsCache_ {+ cachedRuns :: !(V.Vector (Ref (Run m h)))+ , cachedFilters :: !(V.Vector (Bloom SerialisedKey))+ , cachedIndexes :: !(V.Vector Index)+ , cachedKOpsFiles :: !(V.Vector (FS.Handle h))+ }++{-# SPECIALISE mkLevelsCache ::+ ActionRegistry IO+ -> Levels IO h+ -> IO (LevelsCache IO h) #-}+-- | Flatten the argument 'Level's into a single vector of runs, including all+-- runs that are inputs to an ongoing merge. Use that to populate the+-- 'LevelsCache'. The cache will take a reference for each of its runs.+mkLevelsCache ::+ forall m h. (PrimMonad m, MonadMVar m, MonadMask m)+ => ActionRegistry m+ -> Levels m h+ -> m (LevelsCache m h)+mkLevelsCache reg lvls = do+ rs <- foldRunAndMergeM+ (fmap V.singleton . dupRun)+ (\mr -> withRollback reg (MR.duplicateRuns mr) (V.mapM_ releaseRef))+ lvls+ pure $! LevelsCache_ {+ cachedRuns = rs+ , cachedFilters = mapStrict (\(DeRef r) -> Run.runFilter r) rs+ , cachedIndexes = mapStrict (\(DeRef r) -> Run.runIndex r) rs+ , cachedKOpsFiles = mapStrict (\(DeRef r) -> Run.runKOpsFile r) rs+ }+ where+ dupRun r = withRollback reg (dupRef r) releaseRef++ -- TODO: this is not terribly performant, but it is also not sure if we are+ -- going to need this in the end. We might get rid of the LevelsCache.+ foldRunAndMergeM ::+ Monoid a+ => (Ref (Run m h) -> m a)+ -> (Ref (MergingRun MR.LevelMergeType m h) -> m a)+ -> Levels m h+ -> m a+ foldRunAndMergeM k1 k2 ls =+ fmap fold $ forMStrict ls $ \(Level ir rs) -> do+ incoming <- case ir of+ Single r -> k1 r+ Merging _ _ _ mr -> k2 mr+ (incoming <>) . fold <$> V.forM rs k1++{-# SPECIALISE rebuildCache ::+ ActionRegistry IO+ -> LevelsCache IO h+ -> Levels IO h+ -> IO (LevelsCache IO h) #-}+-- | Remove references to runs in the old cache, and create a new cache with+-- fresh references taken for the runs in the new levels.+--+-- TODO: caches are currently only rebuilt in flushWriteBuffer. If an+-- OngoingMerge is completed, then tables will only rebuild the cache, and+-- therefore release "old" runs, when a flush is initiated. This is sub-optimal,+-- and there are at least two solutions, but it is unclear which is faster or+-- more convenient.+--+-- * Get rid of the cache entirely, and have each batch of lookups take+-- references for runs in the levels structure.+--+-- * Keep the cache feature, but force a rebuild every once in a while, e.g.,+-- once in every 100 lookups.+--+-- TODO: rebuilding the cache can invalidate blob references if the cache was+-- holding the last reference to a run. This is not really a problem of just the+-- caching approach, but allowing merges to finish early. We should come up with+-- a solution to keep blob references valid until the next /update/ comes along.+-- Lookups should no invalidate blob erferences.+rebuildCache ::+ (PrimMonad m, MonadMVar m, MonadMask m)+ => ActionRegistry m+ -> LevelsCache m h -- ^ old cache+ -> Levels m h -- ^ new levels+ -> m (LevelsCache m h) -- ^ new cache+rebuildCache reg oldCache newLevels = do+ releaseLevelsCache reg oldCache+ mkLevelsCache reg newLevels++{-# SPECIALISE duplicateLevelsCache ::+ ActionRegistry IO+ -> LevelsCache IO h+ -> IO (LevelsCache IO h) #-}+duplicateLevelsCache ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> LevelsCache m h+ -> m (LevelsCache m h)+duplicateLevelsCache reg cache = do+ rs' <- forMStrict (cachedRuns cache) $ \r ->+ withRollback reg (dupRef r) releaseRef+ pure cache { cachedRuns = rs' }++{-# SPECIALISE releaseLevelsCache ::+ ActionRegistry IO+ -> LevelsCache IO h+ -> IO () #-}+releaseLevelsCache ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> LevelsCache m h+ -> m ()+releaseLevelsCache reg cache =+ V.forM_ (cachedRuns cache) $ \r ->+ delayedCommit reg (releaseRef r)++{-------------------------------------------------------------------------------+ Levels+-------------------------------------------------------------------------------}++type Levels m h = V.Vector (Level m h)++-- | A level is a sequence of resident runs at this level, prefixed by an+-- incoming run, which is usually multiple runs that are being merged. Once+-- completed, the resulting run will become a resident run at this level.+data Level m h = Level {+ incomingRun :: !(IncomingRun m h)+ , residentRuns :: !(V.Vector (Ref (Run m h)))+ }++{-# SPECIALISE duplicateLevels :: ActionRegistry IO -> Levels IO h -> IO (Levels IO h) #-}+duplicateLevels ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> Levels m h+ -> m (Levels m h)+duplicateLevels reg levels =+ forMStrict levels $ \Level {incomingRun, residentRuns} -> do+ incomingRun' <- withRollback reg (duplicateIncomingRun incomingRun) releaseIncomingRun+ residentRuns' <- forMStrict residentRuns $ \r ->+ withRollback reg (dupRef r) releaseRef+ pure $! Level {+ incomingRun = incomingRun',+ residentRuns = residentRuns'+ }++{-# SPECIALISE releaseLevels :: ActionRegistry IO -> Levels IO h -> IO () #-}+releaseLevels ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> Levels m h+ -> m ()+releaseLevels reg levels =+ V.forM_ levels $ \Level {incomingRun, residentRuns} -> do+ delayedCommit reg (releaseIncomingRun incomingRun)+ V.mapM_ (delayedCommit reg . releaseRef) residentRuns++{-# SPECIALISE iforLevelM_ :: Levels IO h -> (LevelNo -> Level IO h -> IO ()) -> IO () #-}+iforLevelM_ :: Monad m => Levels m h -> (LevelNo -> Level m h -> m ()) -> m ()+iforLevelM_ lvls k = V.iforM_ lvls $ \i lvl -> k (LevelNo (i + 1)) lvl++{-------------------------------------------------------------------------------+ Union level+-------------------------------------------------------------------------------}++-- | An additional optional last level, created as a result of+-- 'Database.LSMTree.union'. It can not only contain an ongoing merge of+-- multiple runs, but a nested tree of merges.+--+-- TODO: So far, this is+-- * not considered when creating cursors (also used for range lookups)+-- * never merged into the regular levels+data UnionLevel m h =+ NoUnion+ | Union !(Ref (MergingTree m h)) !(UnionCache m h)++{-# SPECIALISE duplicateUnionLevel ::+ ActionRegistry IO+ -> UnionLevel IO h+ -> IO (UnionLevel IO h) #-}+duplicateUnionLevel ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> UnionLevel m h+ -> m (UnionLevel m h)+duplicateUnionLevel reg ul =+ case ul of+ NoUnion -> pure ul+ Union tree cache -> Union <$> withRollback reg (dupRef tree) releaseRef+ <*> duplicateUnionCache reg cache++{-# SPECIALISE releaseUnionLevel ::+ ActionRegistry IO+ -> UnionLevel IO h+ -> IO () #-}+releaseUnionLevel ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> UnionLevel m h+ -> m ()+releaseUnionLevel _ NoUnion = pure ()+releaseUnionLevel reg (Union tree cache) = delayedCommit reg (releaseRef tree)+ >> releaseUnionCache reg cache++{-------------------------------------------------------------------------------+ Union cache+-------------------------------------------------------------------------------}++-- | Similar to the 'LevelsCache', but in a tree shape, since this structure is+-- required to combine the individual lookup results.+newtype UnionCache m h = UnionCache {+ cachedTree :: MT.LookupTree (V.Vector (Ref (Run m h)))+ }++{-# SPECIALISE mkUnionCache ::+ ActionRegistry IO+ -> Ref (MergingTree IO h)+ -> IO (UnionCache IO h) #-}+mkUnionCache ::+ (PrimMonad m, MonadMVar m, MonadMask m)+ => ActionRegistry m+ -> Ref (MergingTree m h)+ -> m (UnionCache m h)+mkUnionCache reg mt =+ UnionCache <$> MT.buildLookupTree reg mt++{-# SPECIALISE duplicateUnionCache ::+ ActionRegistry IO+ -> UnionCache IO h+ -> IO (UnionCache IO h) #-}+duplicateUnionCache ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> UnionCache m h+ -> m (UnionCache m h)+duplicateUnionCache reg (UnionCache mt) =+ UnionCache <$>+ MT.mapMStrict (mapMStrict (\r -> withRollback reg (dupRef r) releaseRef)) mt++{-# SPECIALISE releaseUnionCache ::+ ActionRegistry IO+ -> UnionCache IO h+ -> IO () #-}+releaseUnionCache ::+ (PrimMonad m, MonadMask m)+ => ActionRegistry m+ -> UnionCache m h+ -> m ()+releaseUnionCache reg (UnionCache mt) =+ traverse_ (traverse_ (\r -> delayedCommit reg (releaseRef r))) mt++{-------------------------------------------------------------------------------+ Flushes and scheduled merges+-------------------------------------------------------------------------------}++{-# SPECIALISE updatesWithInterleavedFlushes ::+ Tracer IO (AtLevel MergeTrace)+ -> TableConfig+ -> ResolveSerialisedValue+ -> HasFS IO h+ -> HasBlockIO IO h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter IO+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> ActionRegistry IO+ -> TableContent IO h+ -> IO (TableContent IO h) #-}+-- | A single batch of updates can fill up the write buffer multiple times. We+-- flush the write buffer each time it fills up before trying to fill it up+-- again.+--+-- TODO: in practice the size of a batch will be much smaller than the maximum+-- size of the write buffer, so we should optimise for the case that small+-- batches are inserted. Ideas:+--+-- * we can allow a range of sizes to flush to disk rather than just the max size+--+-- * could do a map bulk merge rather than sequential insert, on the prefix of+-- the batch that's guaranteed to fit+--+-- * or flush the existing buffer if we would expect the next batch to cause the+-- buffer to become too large+--+-- TODO: we could also optimise for the case where the write buffer is small+-- compared to the size of the batch, but it is less critical. In particular, in+-- case the write buffer is empty, or if it fills up multiple times for a single+-- batch of updates, we might be able to skip adding entries to the write buffer+-- for a large part. When the write buffer is empty, we can sort and deduplicate+-- the vector of updates directly, slice it up into properly sized sub-vectors,+-- and write those to disk. Of course, any remainder that did not fit into a+-- whole run should then end up in a fresh write buffer.+updatesWithInterleavedFlushes ::+ forall m h.+ (MonadMask m, MonadMVar m, MonadSTM m, MonadST m)+ => Tracer m (AtLevel MergeTrace)+ -> TableConfig+ -> ResolveSerialisedValue+ -> HasFS m h+ -> HasBlockIO m h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter m+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> ActionRegistry m+ -> TableContent m h+ -> m (TableContent m h)+updatesWithInterleavedFlushes tr conf resolve hfs hbio root salt uc es reg tc = do+ let wb = tableWriteBuffer tc+ wbblobs = tableWriteBufferBlobs tc+ (wb', es') <- addWriteBufferEntries hfs resolve wbblobs maxn wb es+ -- Supply credits before flushing, so that we complete merges in time. The+ -- number of supplied credits is based on the size increase of the write+ -- buffer, not the number of processed entries @length es' - length es@.+ let numAdded = unNumEntries (WB.numEntries wb') - unNumEntries (WB.numEntries wb)+ supplyCredits conf (NominalCredits numAdded) (tableLevels tc)+ let tc' = tc { tableWriteBuffer = wb' }+ if WB.numEntries wb' < maxn then do+ pure $! tc'+ -- If the write buffer did reach capacity, then we flush.+ else do+ tc'' <- flushWriteBuffer tr conf resolve hfs hbio root salt uc reg tc'+ -- In the fortunate case where we have already performed all the updates,+ -- return,+ if V.null es' then+ pure $! tc''+ -- otherwise, keep going+ else+ updatesWithInterleavedFlushes tr conf resolve hfs hbio root salt uc es' reg tc''+ where+ AllocNumEntries (NumEntries -> maxn) = confWriteBufferAlloc conf++{-# SPECIALISE addWriteBufferEntries ::+ HasFS IO h+ -> ResolveSerialisedValue+ -> Ref (WriteBufferBlobs IO h)+ -> NumEntries+ -> WriteBuffer+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> IO (WriteBuffer, V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)) #-}+-- | Add entries to the write buffer up until a certain write buffer size @n@.+--+-- NOTE: if the write buffer is larger @n@ already, this is a no-op.+addWriteBufferEntries ::+ (MonadSTM m, MonadThrow m, PrimMonad m)+ => HasFS m h+ -> ResolveSerialisedValue+ -> Ref (WriteBufferBlobs m h)+ -> NumEntries+ -> WriteBuffer+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> m (WriteBuffer, V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob))+addWriteBufferEntries hfs f wbblobs maxn =+ \wb es ->+ (\ r@(wb', es') ->+ -- never exceed the write buffer capacity+ assert (WB.numEntries wb' <= maxn) $+ -- If the new write buffer has not reached capacity yet, then it must+ -- be the case that we have performed all the updates.+ assert ((WB.numEntries wb' < maxn && V.null es')+ || (WB.numEntries wb' == maxn)) $+ r)+ <$> go wb es+ where+ --TODO: exception safety for async exceptions or I/O errors from writing blobs+ go !wb !es+ | WB.numEntries wb >= maxn = pure (wb, es)++ | Just ((k, e), es') <- V.uncons es = do+ e' <- traverse (WBB.addBlob hfs wbblobs) e+ go (WB.addEntry f k e' wb) es'++ | otherwise = pure (wb, es)+++{-# SPECIALISE flushWriteBuffer ::+ Tracer IO (AtLevel MergeTrace)+ -> TableConfig+ -> ResolveSerialisedValue+ -> HasFS IO h+ -> HasBlockIO IO h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter IO+ -> ActionRegistry IO+ -> TableContent IO h+ -> IO (TableContent IO h) #-}+-- | Flush the write buffer to disk, regardless of whether it is full or not.+--+-- The returned table content contains an updated set of levels, where the write+-- buffer is inserted into level 1.+flushWriteBuffer ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => Tracer m (AtLevel MergeTrace)+ -> TableConfig+ -> ResolveSerialisedValue+ -> HasFS m h+ -> HasBlockIO m h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter m+ -> ActionRegistry m+ -> TableContent m h+ -> m (TableContent m h)+flushWriteBuffer tr conf resolve hfs hbio root salt uc reg tc+ | WB.null (tableWriteBuffer tc) = pure tc+ | otherwise = do+ !uniq <- incrUniqCounter uc+ let !size = WB.numEntries (tableWriteBuffer tc)+ !ln = LevelNo 1+ (!runParams,+ runPaths) = mergingRunParamsForLevel+ (Paths.activeDir root) conf uniq ln++ traceWith tr $ AtLevel ln $+ TraceFlushWriteBuffer size (runNumber runPaths) runParams+ r <- withRollback reg+ (Run.fromWriteBuffer+ hfs hbio salt+ runParams runPaths+ (tableWriteBuffer tc)+ (tableWriteBufferBlobs tc))+ releaseRef+ delayedCommit reg (releaseRef (tableWriteBufferBlobs tc))+ wbblobs' <- withRollback reg (WBB.new hfs (Paths.tableBlobPath root uniq))+ releaseRef+ levels' <- addRunToLevels tr conf resolve hfs hbio root salt uc r reg+ (tableLevels tc)+ (tableUnionLevel tc)+ tableCache' <- rebuildCache reg (tableCache tc) levels'+ pure $! TableContent {+ tableWriteBuffer = WB.empty+ , tableWriteBufferBlobs = wbblobs'+ , tableLevels = levels'+ , tableCache = tableCache'+ -- TODO: move into regular levels if merge completed and size fits+ , tableUnionLevel = tableUnionLevel tc+ }++{-# SPECIALISE addRunToLevels ::+ Tracer IO (AtLevel MergeTrace)+ -> TableConfig+ -> ResolveSerialisedValue+ -> HasFS IO h+ -> HasBlockIO IO h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter IO+ -> Ref (Run IO h)+ -> ActionRegistry IO+ -> Levels IO h+ -> UnionLevel IO h+ -> IO (Levels IO h) #-}+-- | Add a run to the levels, and propagate merges.+--+-- NOTE: @go@ is based on the @ScheduledMerges.increment@ prototype.+-- See @ScheduledMerges.increment@ for documentation about the merge algorithm.+addRunToLevels ::+ forall m h.+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => Tracer m (AtLevel MergeTrace)+ -> TableConfig+ -> ResolveSerialisedValue+ -> HasFS m h+ -> HasBlockIO m h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter m+ -> Ref (Run m h)+ -> ActionRegistry m+ -> Levels m h+ -> UnionLevel m h+ -> m (Levels m h)+addRunToLevels tr conf@TableConfig{..} resolve hfs hbio root salt uc r0 reg levels ul = do+ go (LevelNo 1) (V.singleton r0) levels+ where+ -- NOTE: @go@ is based on the @increment@ function from the+ -- @ScheduledMerges@ prototype.+ --+ -- Releases the vector of runs.+ go ::+ LevelNo+ -> V.Vector (Ref (Run m h))+ -> V.Vector (Level m h )+ -> m (V.Vector (Level m h))+ go !ln rs (V.uncons -> Nothing) = do+ traceWith tr $ AtLevel ln TraceAddLevel+ -- Make a new level+ let policyForLevel = mergePolicyForLevel confMergePolicy ln V.empty ul+ ir <- newMerge policyForLevel (mergeTypeForLevel V.empty ul) ln rs+ pure $! V.singleton $ Level ir V.empty+ go !ln rs' (V.uncons -> Just (Level ir rs, ls)) = do+ r <- expectCompletedMerge ln ir+ case mergePolicyForLevel confMergePolicy ln ls ul of+ -- If r is still too small for this level then keep it and merge again+ -- with the incoming runs.+ LevelTiering | Run.size r <= maxRunSize' conf LevelTiering (pred ln) -> do+ let mergeType = mergeTypeForLevel ls ul+ ir' <- newMerge LevelTiering mergeType ln (rs' `V.snoc` r)+ pure $! Level ir' rs `V.cons` ls+ -- This tiering level is now full. We take the completed merged run+ -- (the previous incoming runs), plus all the other runs on this level+ -- as a bundle and move them down to the level below. We start a merge+ -- for the new incoming runs. This level is otherwise empty.+ LevelTiering | levelIsFull confSizeRatio rs -> do+ ir' <- newMerge LevelTiering MR.MergeMidLevel ln rs'+ ls' <- go (succ ln) (r `V.cons` rs) ls+ pure $! Level ir' V.empty `V.cons` ls'+ -- This tiering level is not yet full. We move the completed merged run+ -- into the level proper, and start the new merge for the incoming runs.+ LevelTiering -> do+ let mergeType = mergeTypeForLevel ls ul+ ir' <- newMerge LevelTiering mergeType ln rs'+ traceWith tr $ AtLevel ln+ $ TraceAddRun+ (Run.runFsPathsNumber r)+ (V.map Run.runFsPathsNumber rs)+ pure $! Level ir' (r `V.cons` rs) `V.cons` ls+ -- The final level is using levelling. If the existing completed merge+ -- run is too large for this level, we promote the run to the next+ -- level and start merging the incoming runs into this (otherwise+ -- empty) level .+ LevelLevelling | Run.size r > maxRunSize' conf LevelLevelling ln -> do+ assert (V.null rs && V.null ls) $ pure ()+ ir' <- newMerge LevelTiering MR.MergeMidLevel ln rs'+ ls' <- go (succ ln) (V.singleton r) V.empty+ pure $! Level ir' V.empty `V.cons` ls'+ -- Otherwise we start merging the incoming runs into the run.+ LevelLevelling -> do+ assert (V.null rs && V.null ls) $ pure ()+ ir' <- newMerge LevelLevelling (mergeTypeForLevel ls ul) ln (rs' `V.snoc` r)+ pure $! Level ir' V.empty `V.cons` V.empty++ -- Releases the incoming run.+ expectCompletedMerge :: LevelNo -> IncomingRun m h -> m (Ref (Run m h))+ expectCompletedMerge ln ir = do+ r <- case ir of+ Single r -> pure r+ Merging _ _ _ mr -> do+ r <- withRollback reg (MR.expectCompleted mr) releaseRef+ delayedCommit reg (releaseRef mr)+ pure r+ traceWith tr $ AtLevel ln $+ TraceExpectCompletedMerge (Run.runFsPathsNumber r)+ pure r++ -- Consumes and releases the runs.+ newMerge :: MergePolicyForLevel+ -> MR.LevelMergeType+ -> LevelNo+ -> V.Vector (Ref (Run m h))+ -> m (IncomingRun m h)+ newMerge mergePolicy mergeType ln rs = do+ ir <- withRollback reg+ (newIncomingRunAtLevel tr hfs hbio+ root salt uc conf resolve+ mergePolicy mergeType ln rs)+ releaseIncomingRun+ -- The runs will end up inside the incoming/merging run, with fresh+ -- references (since newIncoming* will make duplicates).+ -- The original references must be released (but only on the happy path).+ V.forM_ rs $ \r -> delayedCommit reg (releaseRef r)+ case confMergeSchedule of+ Incremental -> pure ()+ OneShot ->+ bracket+ (immediatelyCompleteIncomingRun conf ln ir)+ releaseRef $ \r ->++ traceWith tr $ AtLevel ln $+ TraceCompletedMerge (Run.size r) (Run.runFsPathsNumber r)++ pure ir++{-# SPECIALISE newIncomingRunAtLevel ::+ Tracer IO (AtLevel MergeTrace)+ -> HasFS IO h+ -> HasBlockIO IO h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter IO+ -> TableConfig+ -> ResolveSerialisedValue+ -> MergePolicyForLevel+ -> MR.LevelMergeType+ -> LevelNo+ -> V.Vector (Ref (Run IO h))+ -> IO (IncomingRun IO h) #-}+newIncomingRunAtLevel ::+ (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => Tracer m (AtLevel MergeTrace)+ -> HasFS m h+ -> HasBlockIO m h+ -> SessionRoot+ -> Bloom.Salt+ -> UniqCounter m+ -> TableConfig+ -> ResolveSerialisedValue+ -> MergePolicyForLevel+ -> MR.LevelMergeType+ -> LevelNo+ -> V.Vector (Ref (Run m h))+ -> m (IncomingRun m h)+newIncomingRunAtLevel tr hfs hbio+ root salt uc conf resolve+ mergePolicy mergeType ln rs+ | Just (r, rest) <- V.uncons rs, V.null rest = do++ traceWith tr $ AtLevel ln $+ TraceNewMergeSingleRun (Run.size r) (Run.runFsPathsNumber r)++ newIncomingSingleRun r++ | otherwise = do++ uniq <- incrUniqCounter uc+ let (!runParams, !runPaths) =+ mergingRunParamsForLevel (Paths.activeDir root) conf uniq ln++ traceWith tr $ AtLevel ln $+ TraceNewMerge (V.map Run.size rs) (runNumber runPaths)+ runParams mergePolicy mergeType++ bracket+ (MR.new hfs hbio resolve salt runParams mergeType runPaths rs)+ releaseRef $ \mr ->+ assert (MR.totalMergeDebt mr <= maxMergeDebt conf mergePolicy ln) $+ let nominalDebt = nominalDebtForLevel conf ln in+ newIncomingMergingRun mergePolicy nominalDebt mr++mergingRunParamsForLevel ::+ ActiveDir+ -> TableConfig+ -> Unique+ -> LevelNo+ -> (RunParams, RunFsPaths)+mergingRunParamsForLevel dir conf unique ln =+ (runParamsForLevel conf (RegularLevel ln), runPaths)+ where+ !runPaths = Paths.runPath dir (uniqueToRunNumber unique)++-- | We use levelling on the last level, unless that is also the first level.+mergePolicyForLevel ::+ MergePolicy+ -> LevelNo+ -> Levels m h+ -> UnionLevel m h+ -> MergePolicyForLevel+mergePolicyForLevel LazyLevelling (LevelNo n) nextLevels unionLevel+ | n == 1+ = LevelTiering -- always use tiering on first level++ | V.null nextLevels+ , NoUnion <- unionLevel+ = LevelLevelling -- levelling on last level++ | otherwise+ = LevelTiering++-- $setup+-- >>> import Database.LSMTree.Internal.Entry+-- >>> import Database.LSMTree.Internal.Config++-- | Compute the maximum size of a run for a given level.+--+-- The @size@ of a tiering run at each level is allowed to be+-- @bufferSize*sizeRatio^(level-1) < size <= bufferSize*sizeRatio^level@.+--+-- >>> unNumEntries . maxRunSize Four (AllocNumEntries 2) LevelTiering . LevelNo <$> [0, 1, 2, 3, 4]+-- [0,2,8,32,128]+--+-- The @size@ of a levelling run at each level is allowed to be+-- @bufferSize*sizeRatio^(level-1) < size <= bufferSize*sizeRatio^(level+1)@. A+-- levelling run can take take up a whole level, so the maximum size of a run is+-- @sizeRatio*@ larger than the maximum size of a tiering run on the same level.+--+-- >>> unNumEntries . maxRunSize Four (AllocNumEntries 2) LevelLevelling . LevelNo <$> [0, 1, 2, 3, 4]+-- [0,8,32,128,512]+maxRunSize ::+ SizeRatio+ -> WriteBufferAlloc+ -> MergePolicyForLevel+ -> LevelNo+ -> NumEntries+maxRunSize _ _ _ (LevelNo ln)+ | ln < 0 = error "maxRunSize: non-positive level number"+ | ln == 0 = NumEntries 0++maxRunSize sizeRatio (AllocNumEntries (NumEntries -> bufferSize)) LevelTiering ln =+ NumEntries $ maxRunSizeTiering (sizeRatioInt sizeRatio) bufferSize ln++maxRunSize sizeRatio (AllocNumEntries (NumEntries -> bufferSize)) LevelLevelling ln =+ NumEntries $ maxRunSizeLevelling (sizeRatioInt sizeRatio) bufferSize ln++maxRunSizeTiering, maxRunSizeLevelling :: Int -> NumEntries -> LevelNo -> Int+maxRunSizeTiering sizeRatio (NumEntries bufferSize) (LevelNo ln) =+ bufferSize * sizeRatio ^ pred ln++maxRunSizeLevelling sizeRatio bufferSize ln =+ maxRunSizeTiering sizeRatio bufferSize (succ ln)++maxRunSize' :: TableConfig -> MergePolicyForLevel -> LevelNo -> NumEntries+maxRunSize' config policy ln =+ maxRunSize (confSizeRatio config) (confWriteBufferAlloc config) policy ln++-- | If there are no further levels provided, this level is the last one.+-- However, if a 'Union' is present, it acts as another (last) level.+mergeTypeForLevel :: Levels m h -> UnionLevel m h -> MR.LevelMergeType+mergeTypeForLevel levels unionLevel+ | V.null levels, NoUnion <- unionLevel = MR.MergeLastLevel+ | otherwise = MR.MergeMidLevel++levelIsFull :: SizeRatio -> V.Vector run -> Bool+levelIsFull sr rs = V.length rs + 1 >= (sizeRatioInt sr)++{-------------------------------------------------------------------------------+ Credits+-------------------------------------------------------------------------------}++{-+ Note [Credits]+~~~~~~~~~~~~~~++ With scheduled merges, each update (e.g., insert) on a table contributes to+ the progression of ongoing merges in the levels structure. This ensures that+ merges are finished in time before a new merge has to be started. The points+ in the evolution of the levels structure where new merges are started are+ known: a flush of a full write buffer will create a new run on the first+ level, and after sufficient flushes (e.g., 4) we will start at least one new+ merge on the second level. This may cascade down to lower levels depending on+ how full the levels are. As such, we have a well-defined measure to determine+ when merges should be finished: it only depends on the maximum size of the+ write buffer!++ The simplest solution to making sure merges are done in time is to step them+ to completion immediately when started. This does not, however, spread out+ work over time nicely. Instead, we schedule merge work based on how many+ updates are made on the table, taking care to ensure that the merge is+ finished /just/ in time before the next flush comes around, and not too early.++ TODO: we can still spread out work more evenly over time. We are finishing+ some merges too early, for example. See 'creditsForMerge'.++ The progression is tracked using merge credits, where each single update+ contributes a single credit to each ongoing merge. This is equivalent to+ saying we contribute a credit to each level, since each level contains+ precisely one ongoing merge. Contributing a credit does not, however,+ translate directly to doing one /unit/ of merging work:++ * The amount of work to do for one credit is adjusted depending on the type of+ merge we are doing. Last-level merges, for example, can have larger inputs,+ and therefore we have to do a little more work for each credit. As such, we+ /scale/ credits for the specific type of merge.++ * Unspent credits are accumulated until they go over a threshold, after which+ a batch of merge work will be performed. Configuring this threshold should+ allow to achieve a nice balance between spreading out I/O and achieving good+ (concurrent) performance.++ Merging runs can be shared across tables, which means that multiple threads+ can contribute to the same merge concurrently.+-}++{-# SPECIALISE supplyCredits ::+ TableConfig+ -> NominalCredits+ -> Levels IO h+ -> IO ()+ #-}+-- | Supply the given amount of credits to each merge in the levels structure.+-- This /may/ cause some merges to progress.+supplyCredits ::+ (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => TableConfig+ -> NominalCredits+ -> Levels m h+ -> m ()+supplyCredits conf deposit levels =+ iforLevelM_ levels $ \ln (Level ir _rs) ->+ supplyCreditsIncomingRun conf ln ir deposit+ --TODO: consider tracing supply of credits,+ -- supplyCreditsIncomingRun could easily return the supplied credits+ -- before & after, which may be useful for tracing.++-- | See 'maxPhysicalDebt' in 'newLevelMerge' in the 'ScheduledMerges'+-- prototype.+maxMergeDebt :: TableConfig -> MergePolicyForLevel -> LevelNo -> MergeDebt+maxMergeDebt TableConfig {+ confWriteBufferAlloc = AllocNumEntries (NumEntries -> bufferSize),+ confSizeRatio+ } mergePolicy ln =+ let !sizeRatio = sizeRatioInt confSizeRatio in+ case mergePolicy of+ LevelLevelling ->+ MergeDebt . MergeCredits $+ sizeRatio * maxRunSizeTiering sizeRatio bufferSize ln+ + maxRunSizeLevelling sizeRatio bufferSize ln++ LevelTiering ->+ MergeDebt . MergeCredits $+ sizeRatio * maxRunSizeTiering sizeRatio bufferSize ln+ + maxRunSizeTiering sizeRatio bufferSize (pred ln)++-- | The nominal debt equals the minimum of credits we will supply before we+-- expect the merge to complete. This is the same as the number of updates+-- in a run that gets moved to this level.+nominalDebtForLevel :: TableConfig -> LevelNo -> NominalDebt+nominalDebtForLevel TableConfig {+ confWriteBufferAlloc = AllocNumEntries (NumEntries -> !bufferSize),+ confSizeRatio+ } ln =+ NominalDebt (maxRunSizeTiering (sizeRatioInt confSizeRatio) bufferSize ln)
@@ -0,0 +1,1031 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE PatternSynonyms #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | An incremental merge of multiple runs.+module Database.LSMTree.Internal.MergingRun (+ -- * Merging run+ MergingRun+ , RunParams (..)+ , new+ , newCompleted+ , duplicateRuns+ , remainingMergeDebt+ , supplyChecked+ , supplyCreditsRelative+ , supplyCreditsAbsolute+ , expectCompleted+ , snapshot+ , totalMergeDebt+ , mergeType++ -- * Merge types+ , IsMergeType (..)+ , LevelMergeType (..)+ , TreeMergeType (..)++ -- * Credit tracking+ -- $credittracking+ , MergeDebt (..)+ , numEntriesToMergeDebt+ , MergeCredits (..)+ , CreditThreshold (..)+ , SpentCredits (..)+ , UnspentCredits (..)++ -- * Concurrency+ -- $concurrency++ -- * Internal state+ , pattern MergingRun+ , mergeState+ , MergingRunState (..)+ , MergeKnownCompleted (..)+ , CreditsVar (..)+ , pattern CreditsPair++ -- * Errors+ , TableTooLargeError (..)+ ) where++import Control.ActionRegistry+import Control.Concurrent.Class.MonadMVar.Strict+import Control.DeepSeq (NFData (..))+import Control.Monad (when)+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (Exception,+ MonadCatch (bracketOnError), MonadMask,+ MonadThrow (throwIO))+import Control.Monad.Primitive+import Control.RefCount+import Data.Bits+import Data.Maybe (fromMaybe)+import Data.Primitive.MutVar+import Data.Primitive.PrimVar+import qualified Data.Vector as V+import Database.LSMTree.Internal.Assertions (assert)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry (NumEntries (..))+import Database.LSMTree.Internal.Merge (IsMergeType (..),+ LevelMergeType (..), Merge, RunParams (..),+ StepResult (..), TreeMergeType (..))+import qualified Database.LSMTree.Internal.Merge as Merge+import Database.LSMTree.Internal.Paths (RunFsPaths (..))+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue)+import qualified Database.LSMTree.Internal.Vector as V+import GHC.Stack (HasCallStack, callStack)+import System.FS.API (HasFS)+import System.FS.BlockIO.API (HasBlockIO)++data MergingRun t m h = MergingRun {++ -- | The total merge debt.+ --+ -- This corresponds to the sum of the number of entries in the input runs.+ mergeDebt :: !MergeDebt++ -- See $credittracking++ -- | A pair of counters for tracking supplied credits:+ --+ -- 1. The supplied credits that have been spent on merging steps plus the+ -- supplied credits that are in the process of being spent.+ -- 2. The supplied credits that have not been spent and are available to+ -- spend.+ --+ -- The counters are always read & modified together atomically.+ --+ , mergeCreditsVar :: !(CreditsVar (PrimState m))++ -- | A variable that caches knowledge about whether the merge has been+ -- completed. If 'MergeKnownCompleted', then we are sure the merge has+ -- been completed, otherwise if 'MergeMaybeCompleted' we have to check the+ -- 'MergingRunState'.+ , mergeKnownCompleted :: !(MutVar (PrimState m) MergeKnownCompleted)+ , mergeState :: !(StrictMVar m (MergingRunState t m h))+ , mergeRefCounter :: !(RefCounter m)+ }++instance RefCounted m (MergingRun t m h) where+ getRefCounter = mergeRefCounter++data MergingRunState t m h =+ CompletedMerge+ !(Ref (Run m h))+ -- ^ Output run+ | OngoingMerge+ !(V.Vector (Ref (Run m h)))+ -- ^ Input runs+ !(Merge t m h)++data MergeKnownCompleted = MergeKnownCompleted | MergeMaybeCompleted+ deriving stock Eq++instance NFData MergeKnownCompleted where+ rnf MergeKnownCompleted = ()+ rnf MergeMaybeCompleted = ()++{-# SPECIALISE new ::+ Merge.IsMergeType t+ => HasFS IO h+ -> HasBlockIO IO h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> RunParams+ -> t+ -> RunFsPaths+ -> V.Vector (Ref (Run IO h))+ -> IO (Ref (MergingRun t IO h)) #-}+-- | Create a new merging run, returning a reference to it that must ultimately+-- be released via 'releaseRef'.+--+-- Duplicates the supplied references to the runs.+--+-- This function should be run with asynchronous exceptions masked to prevent+-- failing after internal resources have already been created.+new ::+ (Merge.IsMergeType t, MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => HasFS m h+ -> HasBlockIO m h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> RunParams+ -> t+ -> RunFsPaths+ -> V.Vector (Ref (Run m h))+ -> m (Ref (MergingRun t m h))+new hfs hbio resolve salt runParams ty runPaths inputRuns =+ assert (V.length inputRuns > 0) $ do+ -- there can be empty runs, which we don't want to include in the merge+ -- TODO: making runs non-empty would involve introducing a constructor+ -- @CompletedMergeEmpty@, but would simplify things and should be possible.+ let nonEmptyRuns = V.filter (\r -> Run.size r > NumEntries 0) inputRuns+ -- If creating the Merge fails, we must release the references again.+ withActionRegistry $ \reg -> do+ let dupRun r = withRollback reg (dupRef r) releaseRef+ case V.length nonEmptyRuns of+ 0 -> do+ -- we can't have an empty merge, but create a new empty run+ --+ -- potentially, we could have re-used one of the empty input runs (or+ -- even re-used a single non-empty input run if there are no others),+ -- as we do in the prototype. but that would mean that the result+ -- doesn't follow the supplied @runParams@.+ -- TODO: decide whether that optimisation is okay+ r <- Run.newEmpty hfs hbio salt runParams runPaths+ unsafeNew+ (MergeDebt 0)+ (SpentCredits 0)+ MergeKnownCompleted+ (CompletedMerge r)+ _ -> do+ rs <- V.mapM dupRun nonEmptyRuns+ merge <- fromMaybe (error "newMerge: merges can not be empty")+ <$> Merge.new hfs hbio salt runParams ty resolve runPaths rs+ unsafeNew+ (numEntriesToMergeDebt (V.foldMap' Run.size rs))+ (SpentCredits 0)+ MergeMaybeCompleted+ (OngoingMerge rs merge)++{-# SPECIALISE newCompleted ::+ MergeDebt+ -> Ref (Run IO h)+ -> IO (Ref (MergingRun t IO h)) #-}+-- | Create a merging run that is already in the completed state, returning a+-- reference that must ultimately be released via 'releaseRef'.+--+-- Duplicates the supplied reference to the run.+--+-- This function should be run with asynchronous exceptions masked to prevent+-- failing after internal resources have already been created.+newCompleted ::+ (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => MergeDebt -- ^ Since there are no longer any input runs, we need to be+ -- told what the merge debt was.+ -> Ref (Run m h)+ -> m (Ref (MergingRun t m h))+newCompleted mergeDebt inputRun = do+ bracketOnError (dupRef inputRun) releaseRef $ \run ->+ unsafeNew+ mergeDebt+ (SpentCredits (mergeDebtAsCredits mergeDebt)) -- since it is completed+ MergeKnownCompleted+ (CompletedMerge run)++-- | The table contains a run that has more than \(2^{40}\) physical entries.+data TableTooLargeError+ = ErrTableTooLarge+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# INLINE unsafeNew #-}+unsafeNew ::+ (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => MergeDebt+ -> SpentCredits+ -> MergeKnownCompleted+ -> MergingRunState t m h+ -> m (Ref (MergingRun t m h))+unsafeNew (MergeDebt mergeDebt) _ _ _+ | SpentCredits mergeDebt > maxBound+ = throwIO ErrTableTooLarge++unsafeNew mergeDebt (SpentCredits spentCredits)+ knownCompleted state = do+ let !credits = CreditsPair (SpentCredits spentCredits) (UnspentCredits 0)+ mergeCreditsVar <- CreditsVar <$> newPrimVar credits+ mergeKnownCompleted <- newMutVar knownCompleted+ mergeState <- newMVar $! state+ newRef (finalise mergeState) $ \mergeRefCounter ->+ MergingRun {+ mergeDebt+ , mergeCreditsVar+ , mergeKnownCompleted+ , mergeState+ , mergeRefCounter+ }+ where+ finalise var = withMVar var $ \case+ CompletedMerge r ->+ releaseRef r+ OngoingMerge rs m -> do+ -- The RunReaders in the Merge keep their own file handles to the+ -- run kopsFile open. We must close these handles *before* we release+ -- the runs themselves, which will close and delete the files.+ -- Otherwise we would be removing files that still have open handles+ -- (which does not work on Windows, and is caught by the MockFS).+ Merge.abort m+ V.forM_ rs releaseRef++-- | Create references to the runs that should be queried for lookups.+-- In particular, if the merge is not complete, these are the input runs.+--+-- TODO: This interface doesn't work well with the action registry. Just doing+-- @withRollback reg (duplicateRuns mr) (mapM_ releaseRef)@ isn't exception-safe+-- since if one of the @releaseRef@ calls fails, the following ones aren't run.+{-# SPECIALISE duplicateRuns ::+ Ref (MergingRun t IO h) -> IO (V.Vector (Ref (Run IO h))) #-}+duplicateRuns ::+ (PrimMonad m, MonadMVar m, MonadMask m)+ => Ref (MergingRun t m h)+ -> m (V.Vector (Ref (Run m h)))+duplicateRuns (DeRef mr) =+ -- We take the references while holding the MVar to make sure the MergingRun+ -- does not get completed concurrently before we are done.+ withMVar (mergeState mr) $ \case+ CompletedMerge r -> V.singleton <$> dupRef r+ OngoingMerge rs _ -> withActionRegistry $ \reg ->+ V.mapMStrict (\r -> withRollback reg (dupRef r) releaseRef) rs++-- | Take a snapshot of the state of a merging run.+--+-- TODO: this is not concurrency safe! The inputs runs to the merging run could+-- be released concurrently by another thread that completes the merge, while+-- the snapshot is taking place. The solution is for snapshot here to duplicate+-- the runs it returns _while_ holding the mergeState MVar (to exclude threads+-- that might concurrently complete the merge). And then the caller of course+-- must be updated to release the extra references.+--+{-# SPECIALISE snapshot ::+ Ref (MergingRun t IO h)+ -> IO (MergeDebt,+ MergeCredits,+ MergingRunState t IO h) #-}+snapshot ::+ (PrimMonad m, MonadMVar m)+ => Ref (MergingRun t m h)+ -> m (MergeDebt,+ MergeCredits,+ MergingRunState t m h)+snapshot (DeRef MergingRun {..}) = do+ state <- readMVar mergeState+ (SpentCredits spent,+ UnspentCredits unspent) <- atomicReadCredits mergeCreditsVar+ let supplied = spent + unspent+ pure (mergeDebt, supplied, state)++totalMergeDebt :: Ref (MergingRun t m h) -> MergeDebt+totalMergeDebt (DeRef MergingRun {mergeDebt}) = mergeDebt++{-# INLINE mergeType #-}+mergeType :: MonadMVar m => Ref (MergingRun t m h) -> m (Maybe t)+mergeType (DeRef mr) = do+ s <- readMVar (mergeState mr)+ pure $ case s of+ CompletedMerge _ -> Nothing+ OngoingMerge _ m -> Just (Merge.mergeType m)++{-------------------------------------------------------------------------------+ Credits+-------------------------------------------------------------------------------}++{- $credittracking++The credits and debt concept we use here comes from amortised analysis of data+structures (see the Bankers Method from Okasaki). Though here we use it not as+an analysis method but within the code itself for tracking the state of the+scheduled (i.e. incremental) merge.++There are two notions of credits (and corresponding debt) in this LSM+implementation: nominal credits and merge credits. The merging run deals+exclusively with merge credits. See 'IncomingRun' for nominal credits.++A single merge credit corresponds with a merge step performed. Merge steps are+measured by the number of entries in the input runs that are consumed. We+measure the merge in terms of inputs, not outputs, because the number of inputs+is known beforehand, whereas the number of outputs is not. The total merge debt+is therefore defined to be the sum of the number of entries across the input+runs to the merge. Once the merge credits spent equals the merge debt then the+merge is (or rather must be) complete.++In both the prototype and implementation we accumulate unspent credits until+they reach a threshold at which point we do a batch of merging work. We track+both credits spent and credits as yet unspent.++In the prototype, the credits spent equals the merge steps performed. The+same holds in the real implementation, but making it so is more complicated.+When we spend credits on merging work, the number of steps we perform is not+guaranteed to be the same as the credits supplied. For example we may ask to do+100 credits of merging work, but the merge code (for perfectly sensible+efficiency reasons) will decide to do 102 units of merging work. The rule is+that we may do (slightly) more work than the credits supplied but not less.+To account for this we spend more credits, corresponding to the excess merging+work performed. We spend them by borrowing them from the unspent credits, which+may leave the unspent credits with a negative balance.++Furthermore, the real implementation has to cope with concurrency: multiple+threads sharing the same 'MergingRun' and calling 'supplyCredits' concurrently.+The credit accounting thus needs to define the state of the credits while+merging work is in progress by some thread. The approach we take is to define+spent credit to include credits that are in the process of being spent,+leaving unspent credit as credits that are available for a thread to spend on+merging work.++Thus we track three things:++ * spent credits ('SpentCredits'): credits supplied that have been or are in+ the process of being spent on performing merging steps;++ * unspent credits ('UnspentCredits'): credits supplied that are not yet spent+ and are thus available to spend; and++ * merge debt ('MergeDebt'): the sum of the sizes of the input runs, and thus+ the total merge credits that have to be spent for the merge to be complete.++And define a derived measure:++ * supplied credits: the sum of the spent and unspent credits. This is+ therefore also the sum of all the credits that have been (successfully)+ supplied to a merging run via 'supplyCredits'.++ The supplied credits increases monotonically, even in the presence of+ (a)synchronous exceptions.++ We guarantee that the supplied credits never exceeds the total debt.++When the supplied credits equals the merge debt then we may not have actually+completed the merge (since that requires spending the credits) but we have the+potential to complete the merge whenever needed without supplying any more+credits.++The credits spent and the steps performed (or in the process of being+performed) will typically be equal. They are not guaranteed to be equal in the+presence of exceptions (synchronous or asynchronous). In this case we offer a+weaker guarantee: : a merge /may/ progress more steps than the number of+credits that were spent. If an exception happens at some point during merging+work, we will \"unspend\" all the credits we intended to spend, but we will not+revert all merging steps that we already successfully performed before the+exception. Thus we may do more merging steps than the credits we accounted as+spent. This makes the implementation simple, and merges will still finish in+time. It would be bad if we did not put back credits, because then a merge+might not finish in time, which will mess up the shape of the levels tree.+-}++newtype MergeCredits = MergeCredits Int+ deriving stock (Eq, Ord, Show)+ deriving newtype (Num, Real, Enum, Integral, NFData)++newtype MergeDebt = MergeDebt MergeCredits+ deriving stock (Eq, Ord)+ deriving newtype (NFData)++mergeDebtAsCredits :: MergeDebt -> MergeCredits+mergeDebtAsCredits (MergeDebt c) = c++{-# INLINE numEntriesToMergeDebt #-}+-- | The total debt of the merging run is exactly the sum total number of+-- entries across all the input runs to be merged.+--+numEntriesToMergeDebt :: NumEntries -> MergeDebt+numEntriesToMergeDebt (NumEntries n) = MergeDebt (MergeCredits n)++-- | Unspent credits are accumulated until they go over the 'CreditThreshold',+-- after which a batch of merge work will be performed. Configuring this+-- threshold should allow to achieve a nice balance between spreading out+-- I\/O and achieving good (concurrent) performance.+--+-- Note that ideally the batch size for different LSM levels should be+-- co-prime so that merge work at different levels is not synchronised.+--+newtype CreditThreshold = CreditThreshold UnspentCredits+ deriving stock Show++-- | The spent credits are supplied credits that have been spent on performing+-- merging steps plus the supplied credits that are in the process of being+-- spent (by some thread calling 'supplyCredits').+--+newtype SpentCredits = SpentCredits MergeCredits+ deriving newtype (Eq, Ord, Show)++-- | 40 bit unsigned number+instance Bounded SpentCredits where+ minBound = SpentCredits 0+ maxBound = SpentCredits (MergeCredits (1 `unsafeShiftL` 40 - 1))++-- | The unspent credits are supplied credits that have not yet been spent on+-- performing merging steps and are available to spend.+--+-- Note: unspent credits may be negative! This can occur when more merge+-- steps were performed than there were credits to cover. In this case the+-- credits are borrowed from the unspent credits, which may result in the+-- current unspent credits being negative for a time.+--+newtype UnspentCredits = UnspentCredits MergeCredits+ deriving newtype (Eq, Ord, Show)++-- | 24 bit signed number+instance Bounded UnspentCredits where+ minBound = UnspentCredits (MergeCredits ((-1) `unsafeShiftL` 23))+ maxBound = UnspentCredits (MergeCredits ( 1 `unsafeShiftL` 23 - 1))++-- | This holds the pair of the 'SpentCredits' and the 'UnspentCredits'. All+-- operations on this pair are atomic.+--+-- The model to think about is a @TVar (SpentCredits, UnspentCredits)@ but the+-- physical representation is a single mutable unboxed 64bit signed @Int@,+-- using 40 bits for the spent credits and 24 for the unspent credits. The+-- spent credits are unsigned, while the unspent credits are signed, so 40 bits+-- and 23+1 bits respectively. This imposes a limit of just over 1 trillion for+-- the spent credits and thus run size, and 8.3 million for the unspent credits+-- (23 + sign bit).+--+-- If these limits ever become restrictive, then the implementation could be+-- changed to use a TVar or a double-word CAS (DWCAS, i.e. 128bit).+--+newtype CreditsVar s = CreditsVar (PrimVar s Int)++pattern CreditsPair :: SpentCredits -> UnspentCredits -> Int+pattern CreditsPair sc uc <- (unpackCreditsPair -> (sc, uc))+ where+ CreditsPair sc uc = packCreditsPair sc uc+{-# INLINE CreditsPair #-}+{-# COMPLETE CreditsPair #-}++{-# INLINE packCreditsPair #-}+packCreditsPair :: SpentCredits -> UnspentCredits -> Int+packCreditsPair spent@(SpentCredits (MergeCredits sc))+ unspent@(UnspentCredits (MergeCredits uc)) =+ assert (spent >= minBound && spent <= maxBound) $+ assert (unspent >= minBound && unspent <= maxBound) $++ sc `unsafeShiftL` 24+ .|. (uc .&. 0xffffff)++{-# INLINE unpackCreditsPair #-}+unpackCreditsPair :: Int -> (SpentCredits, UnspentCredits)+unpackCreditsPair cp =+ -- we use unsigned shift for spent, and sign extending shift for unspent+ ( SpentCredits (MergeCredits (w2i (i2w cp `unsafeShiftR` 24)))+ , UnspentCredits (MergeCredits ((cp `unsafeShiftL` 40) `unsafeShiftR` 40))+ )+ where+ i2w :: Int -> Word+ w2i :: Word -> Int+ i2w = fromIntegral+ w2i = fromIntegral++{-------------------------------------------------------------------------------+ Credit transactions+-------------------------------------------------------------------------------}++{- $concurrency++Merging runs can be shared across tables, which means that multiple threads can+contribute to the same merge concurrently. The design to contribute credits to+the same merging run is largely lock-free. It ensures consistency of the+unspent credits and the merge state, while allowing threads to progress without+waiting on other threads.++The entry point for merging is 'supplyCredits'. This may be called by+concurrent threads that share the same merging run. No locks are held+initially.++The credits to supply can be specified as either an absolute or relative value.+That is, we can ask that the number of supplied credits be set to a value N, or+we can specify an additional N credits.+increasing so in the absolute case, there is no change if the requested new+supplied credit value is less than the current value. Supplying credits from+the levels (via incoming runs) uses absolute credits, while supplying credits+from merging trees using relative credits.++The main lock we will discuss is the 'mergeState' 'StrictMVar', and we will+refer to it as the merge lock.++We get the easy things out of the way first: the 'mergeKnownCompleted'+variable is purely an optimisation. It starts out as 'MergeMaybeCompleted'+and is only ever modified once to 'MergeKnownCompleted'. It is modified with+the merge lock held, but read without the lock. It does not matter if a thread+reads a stale value of 'MergeMaybeCompleted'. We can analyse the remainder of+the algorithm as if we were always in the 'MergeMaybeCompleted' state.++Variable access and locks:++* 'CreditsVar' contains the pair of the current 'SpentCredits' and+ 'UnspentCredits'. Is only operated upon using transactions (atomic CAS),+ and most of these transactions are done without the merge lock held.+ The two constituent components can increase and decrease, but the total+ supplied credits (sum of spent and unspent) can only increase.++* 'MergeState' contains the state of the merge itself. It is protected by the+ merge lock.++First, we do a moderately complex transaction 'atomicDepositAndSpendCredits',+which does the following:++ * Deposit credits to the unspent pot, while guaranteeing that the total+ supplied credits does not exceed the total debt for the merging run.+ * Compute any leftover credits (that would have exceeded the total debt).+ * Compute the credits to spend on performing merge steps, depending on which+ of three cases we are in:++ 1. we have supplied enough credits to complete the merge;+ 2. not case 1, but enough unspent credits have accumulated to do a batch of+ merge work;+ 3. not case 1 or 2, not enough credits to do any merge work.++ * Update the spent and unspent pots+ * Return the credits to spend now and any leftover credits.++If there are now credits to spend, then we attempt to perform that number of+merging steps. While doing the merging work, the (more expensive) merge lock is+taken to ensure that the merging work itself is performed only sequentially.++Note that it is not guaranteed that the merge gets completed, even if the+credits supplied has reached the total debt. It may be interrupted during the+merge (by an async exception). This does not matter because the merge will be+completed in 'expectCompleted'. Completing early is an optimisation.++If an exception occurs during the merge then the credits that were in the+process of being spent are transferred back from the spent to the unspent pot+using 'atomicSpendCredits' (with a negative amount). It is this case that+implies that the spent credits may not increase monotonically, even though the+supplied credits do increase monotonically.++Once performing merge steps is done, if it turns out that excess merge steps+were performed then we must do a further accounting transaction:+'atomicSpendCredits' to spend the excess credits. This is done without respect+to the balance of the unspent credits, which may result in the unspent credit+balance becoming negative. This is ok, and will result in more credits having+to be supplied next time before reaching the credit batch threshold. The+unspent credits can not be negative by the time the merge is complete because+the performing of merge steps cannot do excess steps when it reaches the end of+the merge.++-}++{-# INLINE atomicReadCredits #-}+atomicReadCredits ::+ PrimMonad m+ => CreditsVar (PrimState m)+ -> m (SpentCredits, UnspentCredits)+atomicReadCredits (CreditsVar v) =+ unpackCreditsPair <$> atomicReadInt v++{-# INLINE atomicModifyInt #-}+-- | Atomically modify a single mutable integer variable, using a CAS loop.+atomicModifyInt ::+ PrimMonad m+ => PrimVar (PrimState m) Int+ -> (Int -> (Int, a))+ -> m a+atomicModifyInt var f =+ readPrimVar var >>= casLoop+ where+ casLoop !before = do+ let (!after, !result) = f before+ before' <- casInt var before after+ if before' == before+ then pure result+ else casLoop before'++-- | Credits supplied using a relative value or an absolute value.+data SupplyMergeCredits = SupplyMergeCredits+ !SupplyRelativeOrAbsolute+ !MergeCredits+-- Note this is deliberately represented as a product type, not a sum type, to+-- get better unboxing in function args.++-- | Should we supply credits using a relative value or an absolute value.+data SupplyRelativeOrAbsolute = SupplyRelative | SupplyAbsolute++{-# SPECIALISE atomicDepositAndSpendCredits ::+ CreditsVar RealWorld+ -> MergeDebt+ -> CreditThreshold+ -> SupplyMergeCredits+ -> IO (MergeCredits, MergeCredits, MergeCredits, MergeCredits) #-}+-- | Atomically: add to the unspent credits pot, subject to the supplied+-- credits not exceeding the total debt. Return the new spent and unspent+-- credits, plus any leftover credits in excess of the total debt.+--+-- This is the only operation that changes the total supplied credits, and in+-- a non-decreasing way. Hence overall the supplied credits is monotonically+-- non-decreasing.+--+atomicDepositAndSpendCredits ::+ PrimMonad m+ => CreditsVar (PrimState m)+ -> MergeDebt -- ^ total debt+ -> CreditThreshold+ -> SupplyMergeCredits+ -> m (MergeCredits, MergeCredits, MergeCredits, MergeCredits)+ -- ^ (suppliedBefore, suppliedAfter, spendCredits, leftoverCredits)+atomicDepositAndSpendCredits (CreditsVar !var) (MergeDebt !totalDebt)+ (CreditThreshold !batchThreshold)+ (SupplyMergeCredits !supplyRelOrAbs !credits) =+ assert (credits >= 0) $+ atomicModifyInt var $ \(CreditsPair !spent !unspent) ->+ let (supplied, supplied', unspent', leftover)+ = depositCredits spent unspent supplyRelOrAbs credits++ (spend, spent'', unspent'')+ -- 1. supplied enough credits to complete the merge;+ | supplied' == totalDebt+ = spendAllCredits spent unspent'++ -- 2. not case 1, but enough unspent credits have accumulated to do+ -- a batch of merge work;+ | unspent' >= batchThreshold+ = spendBatchCredits spent unspent' batchThreshold++ -- 3. not case 1 or 2, not enough credits to do any merge work.+ | otherwise+ = (0, spent, unspent')++ in txResultFor supplied spent'' unspent'' spend leftover+ where+ txResultFor !supplied (SpentCredits !spent) (UnspentCredits !unspent)+ !spend !leftover =+ let !after = CreditsPair (SpentCredits spent) (UnspentCredits unspent)+ !supplied' = spent + unspent+ !result = (supplied, supplied', spend, leftover)++ in assert (supplied <= supplied') $+ assert (supplied' <= totalDebt) $+ (after, result)++ depositCredits (SpentCredits !spent) (UnspentCredits !unspent)+ SupplyRelative !deposit =+ let !supplied = spent + unspent+ !leftover = max 0 (supplied + deposit - totalDebt)+ !deposit' = deposit - leftover+ !unspent' = unspent + deposit'+ !supplied' = spent + unspent'+ in assert (unspent' >= unspent) $+ assert (deposit' >= 0) $+ assert (leftover >= 0 && leftover <= deposit) $+ (supplied, supplied', UnspentCredits unspent', leftover)++ depositCredits (SpentCredits !spent) (UnspentCredits !unspent)+ SupplyAbsolute !targetSupplied =+ let !supplied = spent + unspent+ !supplied' = min totalDebt (max supplied targetSupplied)+ !deposit = supplied' - supplied+ !leftover = 0 -- meaningless concept for absolute case+ !unspent' = unspent + deposit+ in assert (unspent' >= unspent) $+ assert (supplied' == spent + unspent') $+ assert (deposit >= 0) $+ (supplied, supplied', UnspentCredits unspent', leftover)++ spendBatchCredits (SpentCredits !spent) (UnspentCredits !unspent)+ (UnspentCredits !unspentBatchThreshold) =+ -- numBatches may be zero, in which case the result will be zero+ let !nBatches = unspent `div` unspentBatchThreshold+ !spend = nBatches * unspentBatchThreshold+ !spent' = spent + spend+ !unspent' = unspent - spend+ in assert (spend >= 0) $+ assert (unspent' < unspentBatchThreshold) $+ assert (spent' + unspent' == spent + unspent) $+ (spend, SpentCredits spent', UnspentCredits unspent')++ spendAllCredits (SpentCredits !spent) (UnspentCredits !unspent) =+ let spend = unspent+ spent' = spent + spend+ unspent' = 0+ in assert (spent' + unspent' == spent + unspent) $+ (spend, SpentCredits spent', UnspentCredits unspent')+++{-# SPECIALISE atomicSpendCredits ::+ CreditsVar RealWorld+ -> MergeCredits+ -> IO () #-}+-- | Atomically: transfer the given number of credits from the unspent pot to+-- the spent pot. The new unspent credits balance may be negative.+--+-- The amount to spend can also be negative to reverse a spending transaction.+--+-- The total supplied credits is unchanged.+--+atomicSpendCredits ::+ PrimMonad m+ => CreditsVar (PrimState m)+ -> MergeCredits -- ^ Can be positive to spend, or negative to unspend.+ -> m ()+atomicSpendCredits (CreditsVar var) spend =+ atomicModifyInt var $ \(CreditsPair (SpentCredits !spent)+ (UnspentCredits !unspent)) ->+ let spent' = spent + spend+ unspent' = unspent - spend+ after = CreditsPair (SpentCredits spent')+ (UnspentCredits unspent')+ in assert (spent' + unspent' == spent + unspent) $+ (after, ())++{-------------------------------------------------------------------------------+ The main algorithms+-------------------------------------------------------------------------------}++{-# SPECIALISE remainingMergeDebt ::+ Ref (MergingRun t IO h) -> IO (MergeDebt, NumEntries) #-}+-- | Calculate the merge credits required to complete the merge, as well as an+-- upper bound on the size of the resulting run.+remainingMergeDebt ::+ (MonadMVar m, PrimMonad m)+ => Ref (MergingRun t m h) -> m (MergeDebt, NumEntries)+remainingMergeDebt (DeRef mr) = do+ readMVar (mergeState mr) >>= \case+ CompletedMerge r -> do+ pure (MergeDebt 0, Run.size r)+ OngoingMerge _ _ -> do+ let MergeDebt totalDebt = mergeDebt mr+ let size = let MergeCredits n = totalDebt in NumEntries n+ (SpentCredits spent, UnspentCredits unspent) <-+ atomicReadCredits (mergeCreditsVar mr)+ let debt = totalDebt - (spent + unspent)+ assert (debt >= 0) $ pure ()+ pure (MergeDebt debt, size)++{-# INLINE supplyChecked #-}+-- | Helper function to assert common invariants for functions that supply+-- credits.+supplyChecked ::+ forall m r s. (HasCallStack, Monad m)+ => (r -> m (MergeDebt, s)) -- how to query current debt+ -> (r -> MergeCredits -> m MergeCredits) -- how to supply+ -> (r -> MergeCredits -> m MergeCredits)+supplyChecked _query supply x credits = do+ assertM $ credits > 0 -- only call them when there are credits to supply+#ifdef NO_IGNORE_ASSERTS+ debt <- fst <$> _query x+ assertM $ debt >= MergeDebt 0 -- debt can't be negative+ leftovers <- supply x credits+ assertM $ leftovers <= credits -- can't have more left than we started with+ assertM $ leftovers >= 0 -- leftovers can't be negative+ debt' <- fst <$> _query x+ assertM $ debt' >= MergeDebt 0+ -- the debt was reduced sufficiently (amount of credits spent)+ assertM $ debt' <= let MergeDebt d = debt+ in MergeDebt (d - (credits - leftovers))+ pure leftovers+#else+ supply x credits+#endif+ where+ assertM :: HasCallStack => Bool -> m ()+ assertM p = let _ = callStack in assert p (pure ())+ -- just uses callStack so the constraint is not redundant in release builds++{-# INLINE supplyCreditsRelative #-}+-- | Supply the given amount of credits to a merging run. This /may/ cause an+-- ongoing merge to progress.+--+-- The credits are given in relative terms: as an addition to the current+-- supplied credits. See 'supplyCreditsAbsolute' to set the supplied credits+-- to an absolute value.+--+-- The result is the number of credits left over. This will be non-zero if the+-- credits supplied would take the total supplied credits over the total merge+-- debt.+--+supplyCreditsRelative ::+ forall t m h. (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => Ref (MergingRun t m h)+ -> CreditThreshold+ -> MergeCredits+ -> m MergeCredits+supplyCreditsRelative = flip $ \th ->+ supplyChecked remainingMergeDebt $ \mr c -> do+ (_suppliedCredits, suppliedCredits', leftoverCredits)+ <- supplyCredits mr th (SupplyMergeCredits SupplyRelative c)++ assert (suppliedCredits' == mergeDebtAsCredits (totalMergeDebt mr)+ || leftoverCredits == 0) $+ pure leftoverCredits++{-# INLINE supplyCreditsAbsolute #-}+-- | Set the supplied credits to the given value, unless the current value is+-- already greater. This /may/ cause an ongoing merge to progress.+--+-- The credits are given in absolute terms: as the new value for the current+-- supplied credits. See 'supplyCreditsRelative' to set the supplied credits+-- as a relative addition to the current value.+--+-- The given credit value must be no greater than the 'totalMergeDebt'.+--+-- The result is the new value of the total supplied credits, which may be more+-- than the specified value if the current value was already greater than the+-- specified value.+--+-- The result is:+--+-- 1. The (absolute value of the) supplied credits beforehand.+-- 2. The (absolute value of the) supplied credits afterwards. This will be+-- equal to the given value or to the supplied credits beforehand,+-- whichever is the greater.+--+supplyCreditsAbsolute ::+ forall t m h. (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => Ref (MergingRun t m h)+ -> CreditThreshold+ -> MergeCredits+ -> m (MergeCredits, MergeCredits)+ -- ^ (suppliedCredits, suppliedCredits')+supplyCreditsAbsolute mr th c =+ assert (0 <= c && c <= mergeDebtAsCredits (totalMergeDebt mr)) $ do+ (suppliedCredits, suppliedCredits', _leftoverCredits)+ <- supplyCredits mr th (SupplyMergeCredits SupplyAbsolute c)+ assert (suppliedCredits' == max c suppliedCredits) $+ pure (suppliedCredits, suppliedCredits')++{-# SPECIALISE supplyCredits ::+ Ref (MergingRun t IO h)+ -> CreditThreshold+ -> SupplyMergeCredits+ -> IO (MergeCredits, MergeCredits, MergeCredits) #-}+-- | Supply the given amount of credits to a merging run. This /may/ cause an+-- ongoing merge to progress.+supplyCredits ::+ forall t m h. (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => Ref (MergingRun t m h)+ -> CreditThreshold+ -> SupplyMergeCredits+ -> m (MergeCredits, MergeCredits, MergeCredits)+ -- ^ (suppliedCredits, suppliedCredits', leftoverCredits)+supplyCredits (DeRef MergingRun {+ mergeKnownCompleted,+ mergeDebt,+ mergeCreditsVar,+ mergeState+ })+ !creditBatchThreshold+ (SupplyMergeCredits !supplyRelOrAbs !credits) =+ assert (credits >= 0) $ do+ mergeCompleted <- readMutVar mergeKnownCompleted+ case mergeCompleted of+ MergeKnownCompleted ->+ let suppliedCredits = mergeDebtAsCredits mergeDebt -- we're completed!+ suppliedCredits' = suppliedCredits -- we can't supply more now+ leftoverCredits = credits -- but meaningless for SupplyAbsolute+ in pure (suppliedCredits, suppliedCredits', leftoverCredits)+ MergeMaybeCompleted ->+ bracketOnError+ -- Atomically add credits to the unspent credits (but not allowing+ -- supplied credits to exceed the total debt), determine which case+ -- we're in and thus how many credits we should try to spend now on+ -- performing merge steps. Return the credits to spend now and any+ -- leftover credits that would exceed the debt limit.+ (atomicDepositAndSpendCredits+ mergeCreditsVar mergeDebt+ creditBatchThreshold+ (SupplyMergeCredits supplyRelOrAbs credits))++ -- If an exception occurs while merging (sync or async) then we+ -- reverse the spending of the credits (but not the deposit).+ (\(_, _, spendCredits, _) ->+ atomicSpendCredits mergeCreditsVar (-spendCredits))++ (\(suppliedCredits, suppliedCredits',+ spendCredits, leftoverCredits) -> do+ when (spendCredits > 0) $ do+ weFinishedMerge <-+ performMergeSteps mergeState mergeCreditsVar spendCredits++ -- If an async exception happens before we get to perform the+ -- completion, then that is fine. The next supplyCredits will+ -- complete the merge.+ when weFinishedMerge $+ completeMerge mergeState mergeKnownCompleted++ assert ( 0 <= suppliedCredits) $+ assert (suppliedCredits <= suppliedCredits') $+ assert (suppliedCredits' <= mergeDebtAsCredits mergeDebt) $+ pure (suppliedCredits, suppliedCredits', leftoverCredits))++{-# SPECIALISE performMergeSteps ::+ StrictMVar IO (MergingRunState t IO h)+ -> CreditsVar RealWorld+ -> MergeCredits+ -> IO Bool #-}+performMergeSteps ::+ (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => StrictMVar m (MergingRunState t m h)+ -> CreditsVar (PrimState m)+ -> MergeCredits+ -> m Bool+performMergeSteps mergeVar creditsVar credits =+ assert (credits >= 0) $+ withMVar mergeVar $ \case+ CompletedMerge{} -> pure False+ OngoingMerge _rs m -> do+ let MergeCredits stepsToDo = credits+ (stepsDone, stepResult) <- Merge.steps m stepsToDo+ assert (stepResult == MergeDone || stepsDone >= stepsToDo) (pure ())+ -- Merge.steps guarantees that @stepsDone >= stepsToDo@ /unless/ the+ -- merge was just now finished and excess credit was supplied.+ -- The latter is possible. As noted elsewhere, exceptions can result in+ -- us having done more merge steps than we accounted for with spent+ -- credits, hence it is possible when getting to the end of the merge+ -- for us to try to do more steps than there are steps possible to do.++ -- If excess merging steps were done then we must account for that.+ -- We do so by borrowing the excess from the unspent credits pot and+ -- spending them, i.e. doing a transfer from unspent to spent. This+ -- may result in the unspent credits pot becoming negative.+ let stepsExcess = MergeCredits (stepsDone - stepsToDo)+ when (stepsExcess > 0) $+ atomicSpendCredits creditsVar stepsExcess++ pure $ stepResult == MergeDone++{-# SPECIALISE completeMerge ::+ StrictMVar IO (MergingRunState t IO h)+ -> MutVar RealWorld MergeKnownCompleted+ -> IO () #-}+-- | Convert an 'OngoingMerge' to a 'CompletedMerge'.+completeMerge ::+ (MonadSTM m, MonadST m, MonadMVar m, MonadMask m)+ => StrictMVar m (MergingRunState t m h)+ -> MutVar (PrimState m) MergeKnownCompleted+ -> m ()+completeMerge mergeVar mergeKnownCompletedVar = do+ modifyMVarMasked_ mergeVar $ \case+ mrs@CompletedMerge{} -> pure $! mrs+ (OngoingMerge rs m) -> do+ -- first try to complete the merge before performing other side effects,+ -- in case the completion fails+ --TODO: Run.fromBuilder (used in Merge.complete) claims not to be+ -- exception safe so we should probably be using the resource registry+ -- and test for exception safety.+ r <- Merge.complete m+ V.forM_ rs releaseRef+ -- Cache the knowledge that we completed the merge+ writeMutVar mergeKnownCompletedVar MergeKnownCompleted+ pure $! CompletedMerge r++{-# SPECIALISE expectCompleted ::+ Ref (MergingRun t IO h)+ -> IO (Ref (Run IO h)) #-}+-- | This does /not/ release the reference, but allocates a new reference for+-- the returned run, which must be released at some point.+expectCompleted ::+ (MonadMVar m, MonadSTM m, MonadST m, MonadMask m)+ => Ref (MergingRun t m h) -> m (Ref (Run m h))+expectCompleted (DeRef MergingRun {..}) = do+ knownCompleted <- readMutVar mergeKnownCompleted+ -- The merge is not guaranteed to be complete, so we do the remaining steps+ when (knownCompleted == MergeMaybeCompleted) $ do+ (SpentCredits spentCredits,+ UnspentCredits unspentCredits) <- atomicReadCredits mergeCreditsVar+ let suppliedCredits = spentCredits + unspentCredits+ !credits = assert (MergeDebt suppliedCredits == mergeDebt) $+ assert (unspentCredits >= 0) $+ unspentCredits++ weFinishedMerge <- performMergeSteps mergeState mergeCreditsVar credits+ -- If an async exception happens before we get to perform the+ -- completion, then that is fine. The next 'expectCompleted' will+ -- complete the merge.+ when weFinishedMerge $ completeMerge mergeState mergeKnownCompleted+ withMVar mergeState $ \case+ CompletedMerge r -> dupRef r -- return a fresh reference to the run+ OngoingMerge{} -> do+ -- If the algorithm finds an ongoing merge here, then it is a bug in+ -- our merge scheduling algorithm. As such, we throw a pure error.+ error "expectCompleted: expected a completed merge, but found an ongoing merge"
@@ -0,0 +1,549 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE PatternSynonyms #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | An incremental merge of multiple runs, preserving a bracketing structure.+--+module Database.LSMTree.Internal.MergingTree (+ -- $mergingtrees+ MergingTree (..)+ , PreExistingRun (..)+ , newCompletedMerge+ , newOngoingMerge+ , newPendingLevelMerge+ , newPendingUnionMerge+ , isStructurallyEmpty+ , remainingMergeDebt+ , supplyCredits+ -- * Internal state+ , MergingTreeState (..)+ , PendingMerge (..)+ ) where++import Control.ActionRegistry+import Control.Concurrent.Class.MonadMVar.Strict+import Control.Exception (assert)+import Control.Monad (foldM, (<$!>))+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadMask)+import Control.Monad.Primitive+import Control.RefCount+import Data.Foldable (toList, traverse_)+#if !MIN_VERSION_base(4,20,0)+import Data.List (foldl')+ -- foldl' is included in the Prelude from base 4.20 onwards+#endif+import Data.Vector (Vector)+import qualified Data.Vector as V+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry (NumEntries (..))+import Database.LSMTree.Internal.MergingRun (MergeDebt (..),+ MergingRun)+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.Paths (SessionRoot)+import qualified Database.LSMTree.Internal.Paths as Paths+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue)+import Database.LSMTree.Internal.UniqCounter+import System.FS.API (HasFS)+import System.FS.BlockIO.API (HasBlockIO)++-- $mergingtrees Semantically, tables are key-value stores like Haskell's+-- @Map@. Table unions then behave like @Map.unionWith (<>)@. If one of the+-- input tables contains a value at a particular key, the result will also+-- contain it. If multiple tables share that key, the values will be combined+-- monoidally.+--+-- Looking at the implementation, tables are not just key-value pairs, but+-- consist of runs. If each table was just a single run, unioning would involve+-- a run merge similar to the one used for compaction (when a level is full),+-- but with a different merge type 'MR.MergeUnion' that differs semantically:+-- Here, runs don't represent updates (overwriting each other), but they each+-- represent the full state of a table. There is no distinction between no+-- entry and a 'Delete', between an 'Insert' and a 'Mupsert'.+--+-- To union two tables, we can therefore first merge down each table into a+-- single run (using regular level merges) and then union merge these.+--+-- However, we want to spread out the work required and perform these merges+-- incrementally. At first, we only create a new table that is empty except for+-- a data structure 'MergingTree', representing the merges that need to be+-- done. The usual operations can then be performed on the table while the+-- merge is in progress: Inserts go into the table as usual, not affecting its+-- last level ('UnionLevel'), lookups need to consider the tree (requiring some+-- complexity and runtime overhead), further unions incorporate the in-progress+-- tree into the resulting one, which also shares future merging work.+--+-- It seems necessary to represent the suspended merges using a tree. Other+-- approaches don't allow for full sharing of the incremental work (e.g.+-- because they effectively \"re-bracket\" nested unions). It also seems+-- necessary to first merge each input table into a single run, as there is no+-- practical distributive property between level and union merges.+++-- | A \"merging tree\" is a mutable representation of an incremental+-- tree-shaped nested merge. This allows to represent union merges of entire+-- tables, each of which itself first need to be merged to become a single run.+--+-- Trees have to support arbitrarily deep nesting, since each input to 'union'+-- might already contain an in-progress merging tree (which then becomes shared+-- between multiple tables).+--+data MergingTree m h = MergingTree {+ mergeState :: !(StrictMVar m (MergingTreeState m h))+ , mergeRefCounter :: !(RefCounter m)+ }++instance RefCounted m (MergingTree m h) where+ getRefCounter = mergeRefCounter++data MergingTreeState m h =+ CompletedTreeMerge+ !(Ref (Run m h))+ -- ^ Output run++ -- | Reuses MergingRun to allow sharing existing merges.+ | OngoingTreeMerge+ !(Ref (MergingRun MR.TreeMergeType m h))++ | PendingTreeMerge+ !(PendingMerge m h)++-- | A merge that is waiting for its inputs to complete.+data PendingMerge m h =+ -- | The collection of inputs is the entire contents of a table,+ -- i.e. its (merging) runs and finally a union merge (if that table+ -- already contained a union).+ PendingLevelMerge+ !(Vector (PreExistingRun m h))+ !(Maybe (Ref (MergingTree m h)))++ -- | Each input is the entire content of a table (as a merging tree).+ | PendingUnionMerge+ !(Vector (Ref (MergingTree m h)))++pendingContent ::+ PendingMerge m h+ -> ( MR.TreeMergeType+ , Vector (PreExistingRun m h)+ , Vector (Ref (MergingTree m h))+ )+pendingContent = \case+ PendingLevelMerge prs t -> (MR.MergeLevel, prs, maybe V.empty V.singleton t)+ PendingUnionMerge ts -> (MR.MergeUnion, V.empty, ts)++{-# COMPLETE PendingMerge #-}+pattern PendingMerge ::+ MR.TreeMergeType+ -> Vector (PreExistingRun m h)+ -> Vector (Ref (MergingTree m h))+ -> PendingMerge m h+pattern PendingMerge mt prs ts <- (pendingContent -> (mt, prs, ts))++data PreExistingRun m h =+ PreExistingRun !(Ref (Run m h))+ | PreExistingMergingRun !(Ref (MergingRun MR.LevelMergeType m h))++{-# SPECIALISE newCompletedMerge ::+ Ref (Run IO h)+ -> IO (Ref (MergingTree IO h)) #-}+newCompletedMerge ::+ (MonadMVar m, PrimMonad m, MonadMask m)+ => Ref (Run m h)+ -> m (Ref (MergingTree m h))+newCompletedMerge run = mkMergingTree . CompletedTreeMerge =<< dupRef run++{-# SPECIALISE newOngoingMerge ::+ Ref (MergingRun MR.TreeMergeType IO h)+ -> IO (Ref (MergingTree IO h)) #-}+-- | Create a new 'MergingTree' representing the merge of an ongoing run.+-- The usage of this function is primarily to facilitate the reloading of an+-- ongoing merge from a persistent snapshot.+newOngoingMerge ::+ (MonadMVar m, PrimMonad m, MonadMask m)+ => Ref (MergingRun MR.TreeMergeType m h)+ -> m (Ref (MergingTree m h))+newOngoingMerge mr = mkMergingTree . OngoingTreeMerge =<< dupRef mr++{-# SPECIALISE newPendingLevelMerge ::+ [PreExistingRun IO h]+ -> Maybe (Ref (MergingTree IO h))+ -> IO (Ref (MergingTree IO h)) #-}+-- | Create a new 'MergingTree' representing the merge of a sequence of+-- pre-existing runs (completed or ongoing, plus a optional final tree).+-- This is for merging the entire contents of a table down to a single run+-- (while sharing existing ongoing merges).+--+-- Shape: if the list of runs is empty and the optional input tree is+-- structurally empty, the result will also be structurally empty. See+-- 'isStructurallyEmpty'.+--+-- Resource tracking:+-- * This allocates a new 'Ref' which the caller is responsible for releasing+-- eventually.+-- * The ownership of all input 'Ref's remains with the caller. This action+-- will create duplicate references, not adopt the given ones.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates\/creates resources.+newPendingLevelMerge ::+ forall m h.+ (MonadMVar m, MonadMask m, PrimMonad m)+ => [PreExistingRun m h]+ -> Maybe (Ref (MergingTree m h))+ -> m (Ref (MergingTree m h))+newPendingLevelMerge [] (Just t) = dupRef t+newPendingLevelMerge [PreExistingRun r] Nothing = do+ -- No need to create a pending merge here.+ --+ -- We could do something similar for PreExistingMergingRun, but it's:+ -- * complicated, because of the LevelMergeType\/TreeMergeType mismatch.+ -- * unneeded, since that case should never occur. If there is only a+ -- single entry in the list, there can only be one level in the input+ -- table. At level 1 there are no merging runs, so it must be a+ -- PreExistingRun.+ r' <- dupRef r+ -- There are no interruption points here, and thus provided async+ -- exceptions are masked then there can be no async exceptions here at all.+ mkMergingTree (CompletedTreeMerge r')++newPendingLevelMerge prs mmt = do+ -- isStructurallyEmpty is an interruption point, and can receive async+ -- exceptions even when masked. So we use it first, *before* allocating+ -- new references.+ mmt' <- dupMaybeMergingTree mmt+ prs' <- traverse dupPreExistingRun (V.fromList prs)+ mkMergingTree (PendingTreeMerge (PendingLevelMerge prs' mmt'))+ where+ dupPreExistingRun (PreExistingRun r) =+ PreExistingRun <$!> dupRef r+ dupPreExistingRun (PreExistingMergingRun mr) =+ PreExistingMergingRun <$!> dupRef mr++ dupMaybeMergingTree :: Maybe (Ref (MergingTree m h))+ -> m (Maybe (Ref (MergingTree m h)))+ dupMaybeMergingTree Nothing = pure Nothing+ dupMaybeMergingTree (Just mt) = do+ isempty <- isStructurallyEmpty mt+ if isempty+ then pure Nothing+ else Just <$!> dupRef mt++{-# SPECIALISE newPendingUnionMerge ::+ [Ref (MergingTree IO h)]+ -> IO (Ref (MergingTree IO h)) #-}+-- | Create a new 'MergingTree' representing the union of one or more merging+-- trees. This is for unioning the content of multiple tables (represented+-- themselves as merging trees).+--+-- Shape: if all of the input trees are structurally empty, the result will+-- also be structurally empty. See 'isStructurallyEmpty'.+--+-- Resource tracking:+-- * This allocates a new 'Ref' which the caller is responsible for releasing+-- eventually.+-- * The ownership of all input 'Ref's remains with the caller. This action+-- will create duplicate references, not adopt the given ones.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates\/creates resources.+newPendingUnionMerge ::+ (MonadMVar m, MonadMask m, PrimMonad m)+ => [Ref (MergingTree m h)]+ -> m (Ref (MergingTree m h))+newPendingUnionMerge mts = do+ mts' <- V.filterM (fmap not . isStructurallyEmpty) (V.fromList mts)+ -- isStructurallyEmpty is interruptible even with async exceptions masked,+ -- but we use it before allocating new references.+ mts'' <- V.mapM dupRef mts'+ case V.uncons mts'' of+ Just (mt, x) | V.null x+ -> pure mt+ _ -> mkMergingTree (PendingTreeMerge (PendingUnionMerge mts''))++{-# SPECIALISE isStructurallyEmpty :: Ref (MergingTree IO h) -> IO Bool #-}+-- | Test if a 'MergingTree' is \"obviously\" empty by virtue of its structure.+-- This is not the same as being empty due to a pending or ongoing merge+-- happening to produce an empty run.+--+isStructurallyEmpty :: MonadMVar m => Ref (MergingTree m h) -> m Bool+isStructurallyEmpty (DeRef MergingTree {mergeState}) =+ isStructurallyEmptyState <$> readMVar mergeState++isStructurallyEmptyState :: MergingTreeState m h -> Bool+isStructurallyEmptyState = \case+ PendingTreeMerge (PendingLevelMerge prs Nothing) -> V.null prs+ PendingTreeMerge (PendingUnionMerge mts) -> V.null mts+ _ -> False+ -- It may also turn out to be useful to consider CompletedTreeMerge with+ -- a zero length runs as empty.++{-# SPECIALISE mkMergingTree ::+ MergingTreeState IO h+ -> IO (Ref (MergingTree IO h)) #-}+-- | Constructor helper.+--+-- This adopts the references in the MergingTreeState, so callers should+-- duplicate first. This is not the normal pattern, but this is an internal+-- helper only.+--+mkMergingTree ::+ (MonadMVar m, PrimMonad m, MonadMask m)+ => MergingTreeState m h+ -> m (Ref (MergingTree m h))+mkMergingTree mergeTreeState = do+ mergeState <- newMVar mergeTreeState+ newRef (finalise mergeState) $ \mergeRefCounter ->+ MergingTree {+ mergeState+ , mergeRefCounter+ }++{-# SPECIALISE finalise :: StrictMVar IO (MergingTreeState IO h) -> IO () #-}+finalise :: (MonadMVar m, PrimMonad m, MonadMask m)+ => StrictMVar m (MergingTreeState m h) -> m ()+finalise mergeState = releaseMTS =<< readMVar mergeState+ where+ releaseMTS (CompletedTreeMerge r) = releaseRef r+ releaseMTS (OngoingTreeMerge mr) = releaseRef mr+ releaseMTS (PendingTreeMerge ptm) =+ case ptm of+ PendingUnionMerge mts -> traverse_ releaseRef mts+ PendingLevelMerge prs mmt -> traverse_ releasePER prs+ >> traverse_ releaseRef mmt++ releasePER (PreExistingRun r) = releaseRef r+ releasePER (PreExistingMergingRun mr) = releaseRef mr++{-# SPECIALISE remainingMergeDebt ::+ Ref (MergingTree IO h) -> IO (MergeDebt, NumEntries) #-}+-- | Calculate an upper bound on the merge credits required to complete the+-- merge, i.e. turn the tree into a 'CompletedTreeMerge'. For the recursive+-- calculation, we also return an upper bound on the size of the resulting run.+remainingMergeDebt ::+ (MonadMVar m, PrimMonad m)+ => Ref (MergingTree m h) -> m (MergeDebt, NumEntries)+remainingMergeDebt (DeRef mt) = do+ readMVar (mergeState mt) >>= \case+ CompletedTreeMerge r -> pure (MergeDebt 0, Run.size r)+ OngoingTreeMerge mr -> addDebtOne <$> MR.remainingMergeDebt mr+ PendingTreeMerge ptm -> addDebtOne <$> remainingMergeDebtPendingMerge ptm+ where+ -- An ongoing merge should never have 0 debt, even if the 'MergingRun' in it+ -- says it is completed. We still need to update it to 'CompletedTreeMerge'.+ -- Similarly, a pending merge needs some work to complete it, even if all+ -- its inputs are empty.+ --+ -- Note that we can't use @max 1@, as this would violate the property that+ -- supplying N credits reduces the remaining debt by at least N.+ addDebtOne (MergeDebt !debt, !size) = (MergeDebt (debt + 1), size)++{-# SPECIALISE remainingMergeDebtPendingMerge ::+ PendingMerge IO h -> IO (MergeDebt, NumEntries) #-}+remainingMergeDebtPendingMerge ::+ (MonadMVar m, PrimMonad m)+ => PendingMerge m h -> m (MergeDebt, NumEntries)+remainingMergeDebtPendingMerge (PendingMerge _ prs mts) = do+ -- TODO: optimise to reduce allocations+ debtsPre <- traverse remainingMergeDebtPreExistingRun prs+ debtsChild <- traverse remainingMergeDebt mts+ pure (debtOfNestedMerge (debtsPre <> debtsChild))+ where+ remainingMergeDebtPreExistingRun = \case+ PreExistingRun r -> pure (MergeDebt 0, Run.size r)+ PreExistingMergingRun mr -> MR.remainingMergeDebt mr++debtOfNestedMerge :: Vector (MergeDebt, NumEntries) -> (MergeDebt, NumEntries)+debtOfNestedMerge debts =+ -- complete all children, then one merge of them all (so debt is their size)+ (MergeDebt (c + MR.MergeCredits n), NumEntries n)+ where+ (MergeDebt c, NumEntries n) = foldl' add (MergeDebt 0, NumEntries 0) debts++ add (MergeDebt !d1, NumEntries !n1) (MergeDebt !d2, NumEntries !n2) =+ (MergeDebt (d1 + d2), NumEntries (n1 + n2))++{-# SPECIALISE supplyCredits ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> Run.RunParams+ -> MR.CreditThreshold+ -> SessionRoot+ -> UniqCounter IO+ -> Ref (MergingTree IO h)+ -> MR.MergeCredits+ -> IO MR.MergeCredits #-}+supplyCredits ::+ forall m h.+ (MonadMVar m, MonadST m, MonadSTM m, MonadMask m)+ => HasFS m h+ -> HasBlockIO m h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> Run.RunParams+ -> MR.CreditThreshold+ -> SessionRoot+ -> UniqCounter m+ -> Ref (MergingTree m h)+ -> MR.MergeCredits+ -> m MR.MergeCredits+supplyCredits hfs hbio resolve salt runParams threshold root uc = \mt0 c0 -> do+ if c0 <= 0+ then pure 0+ else supplyTree mt0 c0+ where+ mkFreshRunPaths = do+ runNumber <- uniqueToRunNumber <$> incrUniqCounter uc+ pure (Paths.runPath (Paths.activeDir root) runNumber)++ supplyTree =+ MR.supplyChecked remainingMergeDebt $ \(DeRef mt) credits ->+ -- TODO: This locks the tree for everyone, for the entire call.+ -- Lookups have to wait until supplyCredits is done.+ -- It should be enough to take the lock only to turn a pending into+ -- an ongoing or ongoing into completed tree, very briefly.+ modifyWithActionRegistry+ (takeMVar (mergeState mt))+ (putMVar (mergeState mt))+ (\reg state -> supplyState reg state credits)++ supplyState reg state credits =+ case state of+ CompletedTreeMerge _ ->+ pure (state, credits)++ OngoingTreeMerge mr -> do+ leftovers <- MR.supplyCreditsRelative mr threshold credits+ if leftovers <= 0+ then+ pure (state, 0)+ else do+ -- complete ongoing merge+ r <- withRollback reg (MR.expectCompleted mr) releaseRef+ delayedCommit reg (releaseRef mr)+ -- all work is done, we can't spend any more credits+ pure (CompletedTreeMerge r, leftovers)++ PendingTreeMerge _+ | isStructurallyEmptyState state -> do+ -- make a completely fresh empty run. this can only happen at the+ -- root. the structurally empty tree still has debt 1, so we want to+ -- merge it into a single run.+ -- we handle this as a special case here since in several places+ -- below we require the list of children to be non-empty.+ runPaths <- mkFreshRunPaths+ run <-+ withRollback reg+ -- TODO: the builder's handles aren't cleaned up if we fail+ -- before fromBuilder closes them+ (Run.newEmpty hfs hbio salt runParams runPaths)+ releaseRef+ pure (CompletedTreeMerge run, credits)++ PendingTreeMerge pm -> do+ leftovers <- supplyPending pm credits+ if leftovers <= 0+ then+ -- still remaining work in children, we can't do more for now+ pure (state, leftovers)+ else do+ -- all children must be done, create new merge!+ state' <- startPendingMerge reg pm+ -- use any remaining credits to progress the new merge+ supplyState reg state' leftovers++ supplyPending ::+ PendingMerge m h -> MR.MergeCredits -> m MR.MergeCredits+ supplyPending =+ MR.supplyChecked remainingMergeDebtPendingMerge $ \pm credits -> do+ case pm of+ PendingLevelMerge prs mt ->+ leftToRight supplyPreExisting (V.toList prs) credits+ >>= leftToRight (flip supplyTree) (toList mt)+ PendingUnionMerge mts ->+ splitEqually (flip supplyTree) (V.toList mts) credits++ supplyPreExisting c = \case+ PreExistingRun _r -> pure c -- no work to do, all leftovers+ PreExistingMergingRun mr -> MR.supplyCreditsRelative mr threshold c++ -- supply credits left to right until they are used up+ leftToRight ::+ (MR.MergeCredits -> a -> m MR.MergeCredits)+ -> [a] -> MR.MergeCredits -> m MR.MergeCredits+ leftToRight _ _ 0 = pure 0+ leftToRight _ [] c = pure c+ leftToRight f (x:xs) c = f c x >>= leftToRight f xs++ -- approximately equal, being more precise would require more iterations+ splitEqually ::+ (MR.MergeCredits -> a -> m MR.MergeCredits)+ -> [a] -> MR.MergeCredits -> m MR.MergeCredits+ splitEqually f xs (MR.MergeCredits credits) =+ -- first give each tree k = ceil(1/n) credits (last ones might get less).+ -- it's important we fold here to collect leftovers.+ -- any remainders go left to right.+ foldM supplyNth (MR.MergeCredits credits) xs >>= leftToRight f xs+ where+ !n = length xs+ !k = MR.MergeCredits ((credits + (n - 1)) `div` n)++ supplyNth 0 _ = pure 0+ supplyNth c t = do+ let creditsToSpend = min k c+ leftovers <- f creditsToSpend t+ pure (c - creditsToSpend + leftovers)++ startPendingMerge reg pm = do+ (mergeType, rs) <- expectCompletedChildren reg pm+ assert (V.length rs > 0) $ pure ()+ runPaths <- mkFreshRunPaths+ mr <-+ withRollback reg+ (MR.new hfs hbio resolve salt runParams mergeType runPaths rs)+ releaseRef+ -- no need for the runs anymore, 'MR.new' made duplicates+ traverse_ (\r -> delayedCommit reg (releaseRef r)) rs+ pure (OngoingTreeMerge mr)++ -- Child references are released using 'delayedCommit', so they get released+ -- if the whole supply operation runs successfully (so the pending merge+ -- is replaced).+ --+ -- Returned references are registered in the ActionRegistry, so they will+ -- get released in case of an exception.+ expectCompletedChildren ::+ ActionRegistry m+ -> PendingMerge m h+ -> m (MR.TreeMergeType, Vector (Ref (Run m h)))+ expectCompletedChildren reg (PendingMerge ty prs mts) = do+ rs1 <- V.forM prs $ \case+ PreExistingRun r -> do+ delayedCommit reg (releaseRef r) -- only released at the end+ withRollback reg (dupRef r) releaseRef+ PreExistingMergingRun mr -> do+ delayedCommit reg (releaseRef mr) -- only released at the end+ withRollback reg (MR.expectCompleted mr) releaseRef+ rs2 <- V.forM mts $ \mt -> do+ delayedCommit reg (releaseRef mt) -- only released at the end+ withRollback reg (expectCompleted mt) releaseRef+ pure (ty, rs1 <> rs2)++-- | This does /not/ release the reference, but allocates a new reference for+-- the returned run, which must be released at some point.+{-# SPECIALISE expectCompleted ::+ Ref (MergingTree IO h)+ -> IO (Ref (Run IO h)) #-}+expectCompleted ::+ (MonadMVar m, MonadSTM m, MonadST m, MonadMask m)+ => Ref (MergingTree m h) -> m (Ref (Run m h))+expectCompleted (DeRef MergingTree {..}) = do+ withMVar mergeState $ \case+ CompletedTreeMerge r -> dupRef r -- return a fresh reference to the run+ OngoingTreeMerge mr -> MR.expectCompleted mr+ PendingTreeMerge{} ->+ error "expectCompleted: expected a completed merging tree, but found a pending one"
@@ -0,0 +1,141 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.MergingTree.Lookup (+ LookupTree (..)+ , mapMStrict+ , mkLookupNode+ , buildLookupTree+ , foldLookupTree+ ) where++import Control.ActionRegistry+import Control.Concurrent.Class.MonadMVar.Strict+import Control.Exception (assert)+import Control.Monad+import Control.Monad.Class.MonadAsync (Async, MonadAsync)+import qualified Control.Monad.Class.MonadAsync as Async+import Control.Monad.Class.MonadThrow (MonadMask)+import Control.Monad.Primitive+import Control.RefCount+import qualified Data.Vector as V+import qualified Database.LSMTree.Internal.Entry as Entry+import Database.LSMTree.Internal.Lookup (LookupAcc)+import qualified Database.LSMTree.Internal.MergingRun as MR+import qualified Database.LSMTree.Internal.MergingTree as MT+import Database.LSMTree.Internal.Run (Run)+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue)+import qualified Database.LSMTree.Internal.Vector as V++-- | A simplified representation of the shape of a 'MT.MergingTree'.+data LookupTree a =+ LookupBatch !a+ -- | Use 'mkLookupNode' to construct this.+ | LookupNode !MR.TreeMergeType !(V.Vector (LookupTree a)) -- ^ length 2 or more+ deriving stock (Foldable)++-- | Deriving 'Traversable' leads to functions that are not strict in the+-- elements of the vector of children. This function avoids that issue.+{-# SPECIALISE mapMStrict ::+ (a -> IO b)+ -> LookupTree a+ -> IO (LookupTree b) #-}+mapMStrict :: Monad m => (a -> m b) -> LookupTree a -> m (LookupTree b)+mapMStrict f = \case+ LookupBatch a -> LookupBatch <$!> f a+ LookupNode t v -> LookupNode t <$!> V.mapMStrict (mapMStrict f) v++-- | Asserts that the vector is non-empty. Collapses singleton nodes.+mkLookupNode :: MR.TreeMergeType -> V.Vector (LookupTree a) -> LookupTree a+mkLookupNode ty ts+ | assert (not (null ts)) (V.length ts == 1) = V.head ts+ | otherwise = LookupNode ty ts++-- | Combine a tree of accs into a single one, using the 'MR.TreeMergeType' of+-- each node.+{-# SPECIALISE foldLookupTree ::+ ResolveSerialisedValue+ -> LookupTree (Async IO (LookupAcc IO h))+ -> IO (LookupAcc IO h) #-}+foldLookupTree ::+ MonadAsync m+ => ResolveSerialisedValue+ -> LookupTree (Async m (LookupAcc m h))+ -> m (LookupAcc m h)+foldLookupTree resolve = \case+ LookupBatch batch ->+ Async.wait batch+ LookupNode mt children ->+ mergeLookupAcc resolve mt <$> traverse (foldLookupTree resolve) children++-- | Requires multiple inputs, all of the same length.+--+-- TODO: do more efficiently on mutable vectors?+mergeLookupAcc ::+ ResolveSerialisedValue+ -> MR.TreeMergeType+ -> V.Vector (LookupAcc m h)+ -> LookupAcc m h+mergeLookupAcc resolve mt accs =+ assert (V.length accs > 1) $+ assert (V.all ((== V.length (V.head accs)) . V.length) accs) $+ foldl1 (V.zipWith updateEntry) accs+ where+ updateEntry Nothing old = old+ updateEntry new Nothing = new+ updateEntry (Just new) (Just old) = Just (combine new old)++ combine = case mt of+ MR.MergeLevel -> Entry.combine resolve+ MR.MergeUnion -> Entry.combineUnion resolve++-- | Create a 'LookupTree' of batches of runs, e.g. to do lookups on. The+-- entries within each batch are to be combined using 'MR.MergeLevel'.+--+-- Assumes that the merging tree is not 'MT.isStructurallyEmpty'.+--+-- This function duplicates the references to all the tree's runs. These+-- references later need to be released.+{-# SPECIALISE buildLookupTree ::+ ActionRegistry IO+ -> Ref (MT.MergingTree IO h)+ -> IO (LookupTree (V.Vector (Ref (Run IO h)))) #-}+buildLookupTree ::+ (PrimMonad m, MonadMVar m, MonadMask m)+ => ActionRegistry m+ -> Ref (MT.MergingTree m h)+ -> m (LookupTree (V.Vector (Ref (Run m h))))+buildLookupTree reg (DeRef mt) =+ -- we make sure the state is not updated while we look at it, so no runs get+ -- dropped before we duplicated the reference.+ withMVar (MT.mergeState mt) $ \case+ MT.CompletedTreeMerge r ->+ LookupBatch . V.singleton <$!> dupRun r+ MT.OngoingTreeMerge mr -> do+ !rs <- withRollback reg (MR.duplicateRuns mr) (V.mapM_ releaseRef)+ ty <- MR.mergeType mr+ pure $ case ty of+ Nothing -> LookupBatch rs -- just one run+ Just MR.MergeLevel -> LookupBatch rs -- combine runs+ Just MR.MergeUnion -> mkLookupNode MR.MergeUnion -- separate+ (LookupBatch . V.singleton <$!> rs)+ MT.PendingTreeMerge (MT.PendingLevelMerge prs Nothing) -> do+ LookupBatch . V.concatMap id <$!> -- combine runs+ V.mapMStrict duplicatePreExistingRun prs+ MT.PendingTreeMerge (MT.PendingLevelMerge prs (Just tree)) -> do+ !child <- buildLookupTree reg tree+ if V.null prs+ then pure child+ else do+ preExisting <- do+ LookupBatch . V.concatMap id <$!> -- combine runs+ V.mapMStrict duplicatePreExistingRun prs+ pure $ mkLookupNode MR.MergeLevel $ V.fromList [preExisting, child]+ MT.PendingTreeMerge (MT.PendingUnionMerge trees) ->+ mkLookupNode MR.MergeUnion <$!> V.mapMStrict (buildLookupTree reg) trees+ where+ dupRun r = withRollback reg (dupRef r) releaseRef++ duplicatePreExistingRun (MT.PreExistingRun r) =+ V.singleton <$!> dupRun r+ duplicatePreExistingRun (MT.PreExistingMergingRun mr) =+ withRollback reg (MR.duplicateRuns mr) (V.mapM_ releaseRef)
@@ -0,0 +1,67 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Utilities related to pages.+--+module Database.LSMTree.Internal.Page (+ PageNo (..)+ , nextPageNo+ , NumPages (..)+ , getNumPages+ , PageSpan (..)+ , singlePage+ , multiPage+ , pageSpanSize+ ) where++import Control.DeepSeq (NFData (..))++-- | A 0-based number identifying a disk page.+newtype PageNo = PageNo { unPageNo :: Int }+ deriving stock (Show, Eq, Ord)+ deriving newtype NFData++-- | Increment the page number.+--+-- Note: This does not ensure that the incremented page number exists within a given page span.+{-# INLINE nextPageNo #-}+nextPageNo :: PageNo -> PageNo+nextPageNo = PageNo . succ . unPageNo++-- | The number of pages contained by an index or other paging data-structure.+--+-- Note: This is a 0-based number; take care to ensure arithmetic underflow+-- does not occur during subtraction operations!+newtype NumPages = NumPages Word+ deriving stock (Eq, Ord, Show)+ deriving newtype (NFData)++-- | A type-safe "unwrapper" for 'NumPages'. Use this accessor whenever you want+-- to convert 'NumPages' to a more versatile number type.+{-# INLINE getNumPages #-}+getNumPages :: Integral i => NumPages -> i+getNumPages (NumPages w) = fromIntegral w++-- | A span of pages, representing an inclusive interval of page numbers.+--+-- Typlically used to denote the contiguous page span for a database entry.+data PageSpan = PageSpan {+ pageSpanStart :: {-# UNPACK #-} !PageNo+ , pageSpanEnd :: {-# UNPACK #-} !PageNo+ }+ deriving stock (Show, Eq)++instance NFData PageSpan where+ rnf (PageSpan x y) = rnf x `seq` rnf y++{-# INLINE singlePage #-}+singlePage :: PageNo -> PageSpan+singlePage i = PageSpan i i++{-# INLINE multiPage #-}+multiPage :: PageNo -> PageNo -> PageSpan+multiPage i j = PageSpan i j++{-# INLINE pageSpanSize #-}+pageSpanSize :: PageSpan -> NumPages+pageSpanSize pspan = NumPages . toEnum $+ unPageNo (pageSpanEnd pspan) - unPageNo (pageSpanStart pspan) + 1
@@ -0,0 +1,404 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Page accumulator.+--+module Database.LSMTree.Internal.PageAcc (+ -- * Incrementally accumulating a single page.+ PageAcc (..),+ newPageAcc,+ resetPageAcc,+ pageAccAddElem,+ serialisePageAcc,+ -- ** Inspection+ keysCountPageAcc,+ indexKeyPageAcc,+ -- ** Entry sizes+ entryWouldFitInPage,+ sizeofEntry,+) where++import Control.Monad.ST.Strict (ST)+import Data.Bits (unsafeShiftL, (.|.))+import qualified Data.Primitive as P+import qualified Data.Vector as V+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Primitive as VP+import Data.Word (Word16, Word32, Word64)+import Database.LSMTree.Internal.BitMath+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry (Entry (..))+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Serialise++-- |+--+-- The delete operations take the least amount of space, as there's only the key.+--+-- A smallest page is with empty key:+--+-- >>> import FormatPage+-- >>> let Just page0 = pageSizeAddElem (Key "", Delete) (pageSizeEmpty DiskPage4k)+-- >>> page0+-- PageSize {pageSizeElems = 1, pageSizeBlobs = 0, pageSizeBytes = 32, pageSizeDisk = DiskPage4k}+--+-- Then we can add pages with a single byte key, e.g.+--+-- >>> pageSizeAddElem (Key "a", Delete) page0+-- Just (PageSize {pageSizeElems = 2, pageSizeBlobs = 0, pageSizeBytes = 35, pageSizeDisk = DiskPage4k})+--+-- i.e. roughly 3-4 bytes (when we get to 32/64 elements we add more bytes for bitmaps).+-- (key and value offset is together 4 bytes: so it's at least 4, the encoding of single element page takes more space).+--+-- If we write as small program, adding single byte keys to a page size:+--+-- >>> let calc s ps = case pageSizeAddElem (Key "x", Delete) ps of { Nothing -> s; Just ps' -> calc (s + 1) ps' }+-- >>> calc 1 page0+-- 759+--+-- I.e. we can have a 4096 byte page with at most 759 keys, assuming keys are+-- length 1 or longer, but without assuming that there are no duplicate keys.+--+-- And 759 rounded to the next multiple of 64 (for the bitmaps) is 768.+--+-- 'PageAcc' can hold up to 759 elements, but most likely 'pageAccAddElem' will make it overflow sooner.+-- Having an upper bound allows us to allocate all memory for the accumulator in advance.+--+-- We don't store or calculate individual key nor value offsets in 'PageAcc', as these will be naturally calculated during serialisation ('serialisePageAcc').+--+data PageAcc s = PageAcc+ { paDir :: !(P.MutablePrimArray s Int) -- ^ various counters (directory + extra counters). It is convenient to have counters as 'Int', as all indexing uses 'Int's.+ , paOpMap :: !(P.MutablePrimArray s Word64) -- ^ operations crumb map+ , paBlobRefsMap :: !(P.MutablePrimArray s Word64) -- ^ blob reference bitmap+ , paBlobRefs1 :: !(P.MutablePrimArray s Word64) -- ^ blob spans 64 bit part (offset)+ , paBlobRefs2 :: !(P.MutablePrimArray s Word32) -- ^ blob spans 32 bit part (size)+ , paKeys :: !(V.MVector s SerialisedKey) -- ^ keys+ , paValues :: !(V.MVector s SerialisedValue) -- ^ values+ }++-------------------------------------------------------------------------------+-- Constants+-------------------------------------------------------------------------------++keysCountIdx :: Int+keysCountIdx = 0+{-# INLINE keysCountIdx #-}++blobRefCountIdx :: Int+blobRefCountIdx = 1+{-# INLINE blobRefCountIdx #-}++byteSizeIdx :: Int+byteSizeIdx = 2+{-# INLINE byteSizeIdx #-}++keysSizeIdx :: Int+keysSizeIdx = 3+{-# INLINE keysSizeIdx #-}++pageSize :: Int+pageSize = 4096+{-# INLINE pageSize #-}++-- | See calculation in 'PageAcc' comments.+maxKeys :: Int+maxKeys = 759+{-# INLINE maxKeys #-}++-- | See calculation in 'PageAcc' comments.+maxOpMap :: Int+maxOpMap = 24 -- 768 / 32+{-# INLINE maxOpMap #-}++-- | See calculation in 'PageAcc' comments.+maxBlobRefsMap :: Int+maxBlobRefsMap = 12 -- 768 / 64+{-# INLINE maxBlobRefsMap #-}++-------------------------------------------------------------------------------+-- Entry operations+-------------------------------------------------------------------------------++-- | Calculate the total byte size of key, value and optional blobspan.+--+-- To fit into single page this has to be at most 4064. If the entry is larger,+-- the 'pageAccAddElem' is guaranteed to return 'False'.+--+-- In other words, if you have a large entry (i.e. Insert with big value),+-- don't use the 'PageAcc', but construct the single value page directly,+-- using 'Database.LSMTree.Internal.PageAcc1.singletonPage'.+--+-- If it's not known from context, use 'entryWouldFitInPage' to determine if+-- you're in the small or large case.+--+-- Checking entry size allows us to use 'Word16' arithmetic, we don't need to+-- worry about overflows.+--+sizeofEntry :: SerialisedKey -> Entry SerialisedValue b -> Int+sizeofEntry k Delete = sizeofKey k+sizeofEntry k (Upsert v) = sizeofKey k + sizeofValue v+sizeofEntry k (Insert v) = sizeofKey k + sizeofValue v+sizeofEntry k (InsertWithBlob v _) = sizeofKey k + sizeofValue v + 12++-- | Determine if the key, value and optional blobspan would fit within a+-- single page. If it does, then it's appropriate to use 'pageAccAddElem', but+-- otherwise use 'Database.LSMTree.Internal.PageAcc1.singletonPage'.+--+-- If 'entryWouldFitInPage' is @True@ and the 'PageAcc' is empty (i.e. using+--'resetPageAcc') then 'pageAccAddElem' is guaranteed to succeed.+--+entryWouldFitInPage :: SerialisedKey -> Entry SerialisedValue b -> Bool+entryWouldFitInPage k e = sizeofEntry k e + 32 <= pageSize++-- | Whether 'Entry' adds a blob reference+hasBlobRef :: Entry a b -> Bool+hasBlobRef (InsertWithBlob _ _) = True+hasBlobRef _ = False++-- | Entry's operations crumb. We return 'Word64' as we'd write that.+entryCrumb :: Entry SerialisedValue BlobSpan -> Word64+entryCrumb Insert {} = 0+entryCrumb InsertWithBlob {} = 0+entryCrumb Upsert {} = 1+entryCrumb Delete {} = 2++-- | Entry value. Return 'emptyValue' for 'Delete'+-- (the empty value is in the page even for 'Delete's)+entryValue :: Entry SerialisedValue BlobSpan -> SerialisedValue+entryValue (Insert v) = v+entryValue (InsertWithBlob v _) = v+entryValue (Upsert v) = v+entryValue Delete = emptyValue++emptyValue :: SerialisedValue+emptyValue = SerialisedValue (RB.pack [])++-------------------------------------------------------------------------------+-- PageAcc functions+-------------------------------------------------------------------------------++-- | Create new 'PageAcc'.+--+-- The create 'PageAcc' will be empty.+newPageAcc :: ST s (PageAcc s)+newPageAcc = do+ paDir <- P.newPrimArray 4+ paOpMap <- P.newPrimArray maxOpMap+ paBlobRefsMap <- P.newPrimArray maxBlobRefsMap+ paBlobRefs1 <- P.newPrimArray maxKeys+ paBlobRefs2 <- P.newPrimArray maxKeys+ paKeys <- VM.new maxKeys+ paValues <- VM.new maxKeys++ -- reset the memory, as it's not initialized+ let page = PageAcc {..}+ resetPageAccN page maxKeys+ pure page++dummyKey :: SerialisedKey+dummyKey = SerialisedKey mempty++dummyValue :: SerialisedValue+dummyValue = SerialisedValue mempty++-- | Reuse 'PageAcc' for constructing new page, the old data will be reset.+resetPageAcc :: PageAcc s+ -> ST s ()+resetPageAcc pa = do+ !n <- keysCountPageAcc pa+ resetPageAccN pa n++-- | Reset the page for the given number of key\/value pairs. This does not+-- check whether the given number exceeds 'maxKeys', in which case behaviour is+-- undefined.+resetPageAccN :: PageAcc s -> Int -> ST s ()+resetPageAccN PageAcc {..} !n = do+ P.setPrimArray paDir 0 4 0+ P.setPrimArray paOpMap 0 maxOpMap 0+ P.setPrimArray paBlobRefsMap 0 maxBlobRefsMap 0++ -- we don't need to clear these, we set what we need.+ -- P.setPrimArray paBlobRefs1 0 maxKeys 0+ -- P.setPrimArray paBlobRefs1 0 maxKeys 0++ VM.set (VM.slice 0 n paKeys) $! dummyKey+ VM.set (VM.slice 0 n paValues) $! dummyValue++ -- initial size is 8 bytes for directory and 2 bytes for last value offset.+ P.writePrimArray paDir byteSizeIdx 10++-- | Add an entry to 'PageAcc'.+--+pageAccAddElem ::+ PageAcc s+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan+ -> ST s Bool -- ^ 'True' if value was successfully added.+pageAccAddElem PageAcc {..} k e+ -- quick short circuit: if the entry is bigger than page: no luck.+ | sizeofEntry k e >= pageSize = pure False+ | otherwise = do+ n <- P.readPrimArray paDir keysCountIdx+ b <- P.readPrimArray paDir blobRefCountIdx+ s <- P.readPrimArray paDir byteSizeIdx++ let !n' = n + 1+ let !b' = if hasBlobRef e then b + 1 else b++ let !s' = s+ + (if mod64 n == 0 then 8 else 0) -- blobrefs bitmap+ + (if mod32 n == 0 then 8 else 0) -- operations bitmap+ + (case n of { 0 -> 6; 1 -> 2; _ -> 4 }) -- key and value offsets+ + sizeofEntry k e++ if s' > pageSize || -- check for size overflow+ n >= maxKeys -- check for buffer overflow+ then pure False+ else do+ -- key sizes+ ks <- P.readPrimArray paDir keysSizeIdx+ let !ks' = ks + sizeofKey k++ P.writePrimArray paDir keysCountIdx n'+ P.writePrimArray paDir blobRefCountIdx b'+ P.writePrimArray paDir keysSizeIdx ks'+ P.writePrimArray paDir byteSizeIdx s'++ -- blob reference+ case e of+ InsertWithBlob _ (BlobSpan w64 w32) -> do+ setBlobRef paBlobRefsMap n+ P.writePrimArray paBlobRefs1 b w64+ P.writePrimArray paBlobRefs2 b w32+ _ -> pure ()++ -- operation+ setOperation paOpMap n (entryCrumb e)++ -- keys and values+ VM.write paKeys n k+ VM.write paValues n $! entryValue e++ pure True++setBlobRef :: P.MutablePrimArray s Word64 -> Int -> ST s ()+setBlobRef arr i = do+ let !j = div64 i+ let !k = mod64 i+ w64 <- P.readPrimArray arr j+ let w64' = w64 .|. unsafeShiftL 1 k+ P.writePrimArray arr j $! w64'++setOperation :: P.MutablePrimArray s Word64 -> Int -> Word64 -> ST s ()+setOperation arr i crumb = do+ let !j = div32 i+ let !k = mod32 i+ w64 <- P.readPrimArray arr j+ let w64' = w64 .|. unsafeShiftL crumb (mul2 k)+ P.writePrimArray arr j $! w64'++-- | Convert mutable 'PageAcc' accumulator to concrete 'RawPage'.+--+-- After this operation 'PageAcc' argument can be reset with 'resetPageAcc',+-- and reused.+serialisePageAcc :: PageAcc s -> ST s RawPage+serialisePageAcc page@PageAcc {..} = do+ size <- P.readPrimArray paDir keysCountIdx+ case size of+ 0 -> pure emptyRawPage+ _ -> serialisePageAccN size page++-- | Serialise non-empty page.+serialisePageAccN :: forall s. Int -> PageAcc s -> ST s RawPage+serialisePageAccN size PageAcc {..} = do+ b <- P.readPrimArray paDir blobRefCountIdx+ ks <- P.readPrimArray paDir keysSizeIdx++ -- keys offsets offset+ let ko0 :: Int+ !ko0 = 8+ + mul8 (ceilDiv64 size + ceilDiv64 (mul2 size)) -- blob and operations bitmaps+ + mul8 b + mul4 b -- blob refs++ -- values offset offset+ let vo0 :: Int+ !vo0 = ko0 + mul2 size++ -- keys data offset+ let kd0 :: Int+ !kd0 = vo0 + if size == 1 then 6 else mul2 (size + 1)++ -- values data offset+ let vd0 :: Int+ !vd0 = kd0 + ks++ -- allocate bytearray+ ba <- P.newByteArray pageSize :: ST s (P.MutableByteArray s)+ P.fillByteArray ba 0 pageSize 0++ -- directory: 64 bytes+ P.writeByteArray ba 0 (fromIntegral size :: Word16)+ P.writeByteArray ba 1 (fromIntegral b :: Word16)+ P.writeByteArray ba 2 (fromIntegral ko0 :: Word16)++ -- blobref and operation bitmap sizes in bytes+ let !blobRefMapSize = mul8 (ceilDiv64 size)+ let !opMapSize = mul8 (ceilDiv64 (mul2 size))++ -- traceM $ "sizes " ++ show (blobRefMapSize, opMapSize)++ P.copyMutableByteArray ba 8 (primToByte paBlobRefsMap) 0 blobRefMapSize+ P.copyMutableByteArray ba (8 + blobRefMapSize) (primToByte paOpMap) 0 opMapSize++ -- blob references+ P.copyMutableByteArray ba (8 + blobRefMapSize + opMapSize) (primToByte paBlobRefs1) 0 (mul8 b)+ P.copyMutableByteArray ba (8 + blobRefMapSize + opMapSize + mul8 b) (primToByte paBlobRefs2) 0 (mul4 b)++ -- traceM $ "preloop " ++ show (ko0, kd0, vd0)++ let loop :: Int -> Int -> Int -> Int -> Int -> ST s ()+ loop !i !ko !vo !kd !vd+ | i >= size = do+ -- past last value offset.+ -- Note: we don't need to write this offset as Word32+ -- even in size == 1 case, as preconditions require the+ -- entries to fit into page, i.e. fit into Word16.+ P.writeByteArray ba (div2 vo) (fromIntegral vd :: Word16)++ | otherwise = do+ -- traceM $ "loop " ++ show (i, ko, kd, vo, vd)++ SerialisedKey (RawBytes (VP.Vector koff klen kba)) <- VM.read paKeys i+ SerialisedValue (RawBytes (VP.Vector voff vlen vba)) <- VM.read paValues i++ -- key and value offsets+ P.writeByteArray ba (div2 ko) (fromIntegral kd :: Word16)+ P.writeByteArray ba (div2 vo) (fromIntegral vd :: Word16)++ -- key and value data+ P.copyByteArray ba kd kba koff klen+ P.copyByteArray ba vd vba voff vlen++ loop (i + 1) (ko + 2) (vo + 2) (kd + klen) (vd + vlen)++ loop 0 ko0 vo0 kd0 vd0++ ba' <- P.unsafeFreezeByteArray ba+ pure (makeRawPage ba' 0)+++keysCountPageAcc :: PageAcc s -> ST s Int+keysCountPageAcc PageAcc {paDir} = P.readPrimArray paDir keysCountIdx++indexKeyPageAcc :: PageAcc s -> Int -> ST s SerialisedKey+indexKeyPageAcc PageAcc {paKeys} ix = VM.read paKeys ix++-------------------------------------------------------------------------------+-- Utils+-------------------------------------------------------------------------------++-- | Extract underlying bytearray fromn 'P.MutableByteArray',+-- so we can copy its contents.+primToByte :: P.MutablePrimArray s a -> P.MutableByteArray s+primToByte (P.MutablePrimArray ba) = P.MutableByteArray ba
@@ -0,0 +1,137 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Create a single value page+--+module Database.LSMTree.Internal.PageAcc1 (+ singletonPage,+) where++import Control.Monad.ST.Strict (ST, runST)+import qualified Data.Primitive as P+import qualified Data.Vector.Primitive as VP+import Data.Word (Word16, Word32, Word64)+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry (Entry (..))+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Serialise++pageSize :: Int+pageSize = 4096+{-# INLINE pageSize #-}++-- | Create a singleton page, also returning the overflow value bytes.+singletonPage ::+ SerialisedKey+ -> Entry SerialisedValue BlobSpan+ -> (RawPage, [RawOverflowPage])+singletonPage k (Insert v) = runST $ do+ -- allocate bytearray+ ba <- P.newPinnedByteArray pageSize :: ST s (P.MutableByteArray s)+ P.fillByteArray ba 0 pageSize 0++ -- directory: 64 bytes+ P.writeByteArray ba 0 (1 :: Word16)+ P.writeByteArray ba 1 (0 :: Word16)+ P.writeByteArray ba 2 (24 :: Word16)++ -- blobref and operation bitmap sizes in bytes+ -- P.writeByteArray ba 1 (0 :: Word64)+ -- P.writeByteArray ba 2 (0 :: Word64)++ -- no blob references++ P.writeByteArray ba 12 (32 :: Word16) -- key offset+ P.writeByteArray ba 13 (32 + sizeofKey16 k :: Word16) -- value offset+ P.writeByteArray ba 7 (32 + sizeofKey32 k + sizeofValue32 v :: Word32) -- post value offset++ -- copy key+ P.copyByteArray ba 32 kba koff klen++ -- copy value prefix+ let vlen' = min vlen (pageSize - klen - 32)+ P.copyByteArray ba (32 + klen) vba voff vlen'++ ba' <- P.unsafeFreezeByteArray ba+ let !page = unsafeMakeRawPage ba' 0+ !overflowPages = rawBytesToOverflowPages (RB.drop vlen' v')+ pure (page, overflowPages)+ where+ SerialisedKey (RawBytes (VP.Vector koff klen kba)) = k+ SerialisedValue v'@(RawBytes (VP.Vector voff vlen vba)) = v++singletonPage k (InsertWithBlob v (BlobSpan w64 w32)) = runST $ do+ -- allocate bytearray+ ba <- P.newPinnedByteArray pageSize :: ST s (P.MutableByteArray s)+ P.fillByteArray ba 0 pageSize 0++ -- directory: 64 bytes+ P.writeByteArray ba 0 (1 :: Word16)+ P.writeByteArray ba 1 (1 :: Word16)+ P.writeByteArray ba 2 (36 :: Word16)++ -- blobref and operation bitmap sizes in bytes+ P.writeByteArray ba 1 (1 :: Word64)+ -- P.writeByteArray ba 2 (0 :: Word64)++ -- blob references+ P.writeByteArray ba 3 w64+ P.writeByteArray ba 8 w32++ P.writeByteArray ba 18 (44 :: Word16) -- key offset+ P.writeByteArray ba 19 (44 + sizeofKey16 k :: Word16) -- value offset+ P.writeByteArray ba 10 (44 + sizeofKey32 k + sizeofValue32 v :: Word32) -- post value offset++ -- copy key+ P.copyByteArray ba 44 kba koff klen++ -- copy value prefix+ let vlen' = min vlen (pageSize - klen - 44)+ P.copyByteArray ba (44 + klen) vba voff vlen'++ ba' <- P.unsafeFreezeByteArray ba+ let !page = unsafeMakeRawPage ba' 0+ !overflowPages = rawBytesToOverflowPages (RB.drop vlen' v')+ pure (page, overflowPages)+ where+ SerialisedKey (RawBytes (VP.Vector koff klen kba)) = k+ SerialisedValue v'@(RawBytes (VP.Vector voff vlen vba)) = v++singletonPage k (Upsert v) = runST $ do+ -- allocate bytearray+ ba <- P.newPinnedByteArray pageSize :: ST s (P.MutableByteArray s)+ P.fillByteArray ba 0 pageSize 0++ -- directory: 64 bytes+ P.writeByteArray ba 0 (1 :: Word16)+ P.writeByteArray ba 1 (0 :: Word16)+ P.writeByteArray ba 2 (24 :: Word16)++ -- blobref and operation bitmap sizes in bytes+ -- P.writeByteArray ba 1 (0 :: Word64)+ P.writeByteArray ba 2 (1 :: Word64)++ -- no blob references++ P.writeByteArray ba 12 (32 :: Word16) -- key offset+ P.writeByteArray ba 13 (32 + sizeofKey16 k :: Word16) -- value offset+ P.writeByteArray ba 7 (32 + sizeofKey32 k + sizeofValue32 v :: Word32) -- post value offset++ -- copy key+ P.copyByteArray ba 32 kba koff klen++ -- copy value prefix+ let vlen' = min vlen (pageSize - klen - 32)+ P.copyByteArray ba (32 + klen) vba voff vlen'++ ba' <- P.unsafeFreezeByteArray ba+ let !page = unsafeMakeRawPage ba' 0+ !overflowPages = rawBytesToOverflowPages (RB.drop vlen' v')+ pure (page, overflowPages)+ where+ SerialisedKey (RawBytes (VP.Vector koff klen kba)) = k+ SerialisedValue v'@(RawBytes (VP.Vector voff vlen vba)) = v++singletonPage _ Delete = error "singletonPage: unexpected Delete entry"
@@ -0,0 +1,392 @@+{-# LANGUAGE PatternSynonyms #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Paths (+ SessionRoot (..)+ , lockFile+ , lockFileName+ , metadataFile+ , ActiveDir (..)+ , activeDir+ , runPath+ , SnapshotName+ , toSnapshotName+ , isValidSnapshotName+ , InvalidSnapshotNameError (..)+ , snapshotsDir+ , NamedSnapshotDir (..)+ , namedSnapshotDir+ , SnapshotMetaDataFile (..)+ , snapshotMetaDataFile+ , SnapshotMetaDataChecksumFile (..)+ , snapshotMetaDataChecksumFile+ -- * Table paths+ , tableBlobPath+ -- * Run paths+ , RunFsPaths (..)+ , pathsForRunFiles+ , runKOpsPath+ , runBlobPath+ , runFilterPath+ , runIndexPath+ , runChecksumsPath+ -- * Checksums for Run files+ , checksumFileNamesForRunFiles+ , toChecksumsFile+ , fromChecksumsFile+ -- * Checksums for WriteBuffer files+ , toChecksumsFileForWriteBufferFiles+ , fromChecksumsFileForWriteBufferFiles+ -- * ForRunFiles abstraction+ , ForKOps (..)+ , ForBlob (..)+ , ForFilter (..)+ , ForIndex (..)+ , ForRunFiles (..)+ , forRunKOpsRaw+ , forRunBlobRaw+ , forRunFilterRaw+ , forRunIndexRaw+ -- * WriteBuffer paths+ , WriteBufferFsPaths (WrapRunFsPaths, WriteBufferFsPaths, writeBufferDir, writeBufferNumber)+ , writeBufferKOpsPath+ , writeBufferBlobPath+ , writeBufferChecksumsPath+ , writeBufferFilePathWithExt+ ) where++import Control.Applicative (Applicative (..))+import Control.DeepSeq (NFData (..))+import Control.Exception.Base (throw)+import Control.Monad.Class.MonadThrow (Exception)+import qualified Data.ByteString.Char8 as BS+import Data.Foldable (toList)+import qualified Data.Map as Map+import Data.String (IsString (..))+import Data.Traversable (for)+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.UniqCounter+import Prelude hiding (Applicative (..))+import qualified System.FilePath.Posix+import qualified System.FilePath.Windows+import System.FS.API+++newtype SessionRoot = SessionRoot { getSessionRoot :: FsPath }+ deriving stock Eq++lockFile :: SessionRoot -> FsPath+lockFile (SessionRoot dir) = dir </> mkFsPath [lockFileName]++lockFileName :: String+lockFileName = "lock"++metadataFile :: SessionRoot -> FsPath+metadataFile (SessionRoot dir) = dir </> mkFsPath ["metadata"]++newtype ActiveDir = ActiveDir { getActiveDir :: FsPath }++activeDir :: SessionRoot -> ActiveDir+activeDir (SessionRoot dir) = ActiveDir (dir </> mkFsPath ["active"])++runPath :: ActiveDir -> RunNumber -> RunFsPaths+runPath (ActiveDir dir) = RunFsPaths dir+++{-------------------------------------------------------------------------------+ Snapshot name+-------------------------------------------------------------------------------}++newtype SnapshotName = SnapshotName FilePath+ deriving stock (Eq, Ord)++instance Show SnapshotName where+ showsPrec d (SnapshotName p) = showsPrec d p++-- | The given string must satisfy 'isValidSnapshotName'.+-- Otherwise, 'fromString' throws an 'InvalidSnapshotNameError'.+instance IsString SnapshotName where+ fromString :: String -> SnapshotName+ fromString = toSnapshotName++data InvalidSnapshotNameError+ = ErrInvalidSnapshotName !String+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | Check if a 'String' would be a valid snapshot name.+--+-- Snapshot names consist of lowercase characters, digits, dashes @-@,+-- and underscores @_@, and must be between 1 and 64 characters long.+-- >>> isValidSnapshotName "main"+-- True+--+-- >>> isValidSnapshotName "temporary-123-test_"+-- True+--+-- >>> isValidSnapshotName "UPPER"+-- False+-- >>> isValidSnapshotName "dir/dot.exe"+-- False+-- >>> isValidSnapshotName ".."+-- False+-- >>> isValidSnapshotName "\\"+-- False+-- >>> isValidSnapshotName ""+-- False+-- >>> isValidSnapshotName (replicate 100 'a')+-- False+--+-- Snapshot names must be valid directory on both POSIX and Windows.+-- This rules out the following reserved file and directory names on Windows:+--+-- >>> isValidSnapshotName "con"+-- False+-- >>> isValidSnapshotName "prn"+-- False+-- >>> isValidSnapshotName "aux"+-- False+-- >>> isValidSnapshotName "nul"+-- False+-- >>> isValidSnapshotName "com1" -- "com2", "com3", etc.+-- False+-- >>> isValidSnapshotName "lpt1" -- "lpt2", "lpt3", etc.+-- False+--+-- See, e.g., [the VBA docs for the "Bad file name or number" error](https://learn.microsoft.com/en-us/office/vba/language/reference/user-interface-help/bad-file-name-or-number-error-52).+isValidSnapshotName :: String -> Bool+isValidSnapshotName str =+ and [ all isValidChar str+ , strLength >= 1+ , strLength <= 64+ , System.FilePath.Posix.isValid str+ , System.FilePath.Windows.isValid str+ ]+ where+ strLength :: Int+ strLength = length str+ isValidChar :: Char -> Bool+ isValidChar c = ('a' <= c && c <= 'z') || ('0' <= c && c <= '9' ) || c `elem` "-_"++-- | Create snapshot name.+--+-- The given string must satisfy 'isValidSnapshotName'.+--+-- Throws the following exceptions:+--+-- ['InvalidSnapshotNameError']:+-- If the given string is not a valid snapshot name.+--+toSnapshotName :: String -> SnapshotName+toSnapshotName str+ | isValidSnapshotName str = SnapshotName str+ | otherwise = throw (ErrInvalidSnapshotName str)++snapshotsDir :: SessionRoot -> FsPath+snapshotsDir (SessionRoot dir) = dir </> mkFsPath ["snapshots"]++-- | The directory for a specific, /named/ snapshot.+--+-- Not to be confused with the snapshot/s/ directory, which holds all named+-- snapshot directories.+newtype NamedSnapshotDir = NamedSnapshotDir { getNamedSnapshotDir :: FsPath }++namedSnapshotDir :: SessionRoot -> SnapshotName -> NamedSnapshotDir+namedSnapshotDir root (SnapshotName name) =+ NamedSnapshotDir (snapshotsDir root </> mkFsPath [name])++newtype SnapshotMetaDataFile = SnapshotMetaDataFile FsPath++snapshotMetaDataFile :: NamedSnapshotDir -> SnapshotMetaDataFile+snapshotMetaDataFile (NamedSnapshotDir dir) =+ SnapshotMetaDataFile (dir </> mkFsPath ["metadata"])++newtype SnapshotMetaDataChecksumFile = SnapshotMetaDataChecksumFile FsPath++snapshotMetaDataChecksumFile :: NamedSnapshotDir -> SnapshotMetaDataChecksumFile+snapshotMetaDataChecksumFile (NamedSnapshotDir dir) =+ SnapshotMetaDataChecksumFile (dir </> mkFsPath ["metadata.checksum"])++{-------------------------------------------------------------------------------+ Table paths+-------------------------------------------------------------------------------}++-- | The file name for a table's write buffer blob file+tableBlobPath :: SessionRoot -> Unique -> FsPath+tableBlobPath session n =+ getActiveDir (activeDir session) </> mkFsPath [show (uniqueToInt n)] <.> "wbblobs"++{-------------------------------------------------------------------------------+ Run paths+-------------------------------------------------------------------------------}++-- | The (relative) file path locations of all the files used by the run:+--+-- The following files exist for a run:+--+-- 1. @${n}.keyops@: the sorted run of key\/operation pairs+-- 2. @${n}.blobs@: the blob values associated with the key\/operations+-- 3. @${n}.filter@: a Bloom filter of all the keys in the run+-- 4. @${n}.index@: an index from keys to disk page numbers+-- 5. @${n}.checksums@: a file listing the crc32c checksums of the other+-- files+--+-- The representation doesn't store the full, name, just the number @n@. Use+-- the accessor functions to get the actual names.+--+data RunFsPaths = RunFsPaths {+ -- | The directory that run files live in.+ runDir :: !FsPath+ , runNumber :: !RunNumber }+ deriving stock Show++instance NFData RunFsPaths where+ rnf (RunFsPaths x y) = rnf x `seq` rnf y++-- | Paths to all files associated with this run, except 'runChecksumsPath'.+pathsForRunFiles :: RunFsPaths -> ForRunFiles FsPath+pathsForRunFiles fsPaths = fmap (runFilePathWithExt fsPaths) runFileExts++runKOpsPath :: RunFsPaths -> FsPath+runKOpsPath = unForKOps . forRunKOps . pathsForRunFiles++runBlobPath :: RunFsPaths -> FsPath+runBlobPath = unForBlob . forRunBlob . pathsForRunFiles++runFilterPath :: RunFsPaths -> FsPath+runFilterPath = unForFilter . forRunFilter . pathsForRunFiles++runIndexPath :: RunFsPaths -> FsPath+runIndexPath = unForIndex . forRunIndex . pathsForRunFiles++runChecksumsPath :: RunFsPaths -> FsPath+runChecksumsPath = flip runFilePathWithExt "checksums"++runFilePathWithExt :: RunFsPaths -> String -> FsPath+runFilePathWithExt (RunFsPaths dir (RunNumber n)) ext =+ dir </> mkFsPath [show n] <.> ext++runFileExts :: ForRunFiles String+runFileExts = ForRunFiles {+ forRunKOps = ForKOps "keyops"+ , forRunBlob = ForBlob "blobs"+ , forRunFilter = ForFilter "filter"+ , forRunIndex = ForIndex "index"+ }++{-------------------------------------------------------------------------------+ Checksums For Run Files+-------------------------------------------------------------------------------}++checksumFileNamesForRunFiles :: ForRunFiles CRC.ChecksumsFileName+checksumFileNamesForRunFiles = fmap (CRC.ChecksumsFileName . BS.pack) runFileExts++toChecksumsFile :: ForRunFiles CRC.CRC32C -> CRC.ChecksumsFile+toChecksumsFile = Map.fromList . toList . liftA2 (,) checksumFileNamesForRunFiles++fromChecksumsFile :: CRC.ChecksumsFile -> Either String (ForRunFiles CRC.CRC32C)+fromChecksumsFile file = for checksumFileNamesForRunFiles $ \name ->+ case Map.lookup name file of+ Just crc -> Right crc+ Nothing -> Left ("key not found: " <> show name)++{-------------------------------------------------------------------------------+ Marker newtypes for individual elements of the ForRunFiles and the+ ForWriteBufferFiles abstractions+-------------------------------------------------------------------------------}++newtype ForKOps a = ForKOps {unForKOps :: a}+ deriving stock (Show, Foldable, Functor, Traversable)++newtype ForBlob a = ForBlob {unForBlob :: a}+ deriving stock (Show, Foldable, Functor, Traversable)++newtype ForFilter a = ForFilter {unForFilter :: a}+ deriving stock (Show, Foldable, Functor, Traversable)++newtype ForIndex a = ForIndex {unForIndex :: a}+ deriving stock (Show, Foldable, Functor, Traversable)++{-------------------------------------------------------------------------------+ ForRunFiles abstraction+-------------------------------------------------------------------------------}++-- | Stores something for each run file (except the checksums file), allowing to+-- easily do something for all of them without mixing them up.+data ForRunFiles a = ForRunFiles {+ forRunKOps :: !(ForKOps a)+ , forRunBlob :: !(ForBlob a)+ , forRunFilter :: !(ForFilter a)+ , forRunIndex :: !(ForIndex a)+ }+ deriving stock (Show, Foldable, Functor, Traversable)++forRunKOpsRaw :: ForRunFiles a -> a+forRunKOpsRaw = unForKOps . forRunKOps++forRunBlobRaw :: ForRunFiles a -> a+forRunBlobRaw = unForBlob . forRunBlob++forRunFilterRaw :: ForRunFiles a -> a+forRunFilterRaw = unForFilter . forRunFilter++forRunIndexRaw :: ForRunFiles a -> a+forRunIndexRaw = unForIndex . forRunIndex++instance Applicative ForRunFiles where+ pure x = ForRunFiles (ForKOps x) (ForBlob x) (ForFilter x) (ForIndex x)+ ForRunFiles (ForKOps f1) (ForBlob f2) (ForFilter f3) (ForIndex f4) <*> ForRunFiles (ForKOps x1) (ForBlob x2) (ForFilter x3) (ForIndex x4) =+ ForRunFiles (ForKOps $ f1 x1) (ForBlob $ f2 x2) (ForFilter $ f3 x3) (ForIndex $ f4 x4)++{-------------------------------------------------------------------------------+ WriteBuffer paths+-------------------------------------------------------------------------------}++newtype WriteBufferFsPaths = WrapRunFsPaths RunFsPaths++pattern WriteBufferFsPaths :: FsPath -> RunNumber -> WriteBufferFsPaths+pattern WriteBufferFsPaths {writeBufferDir, writeBufferNumber} = WrapRunFsPaths (RunFsPaths writeBufferDir writeBufferNumber)++{-# COMPLETE WriteBufferFsPaths #-}++writeBufferKOpsExt :: String+writeBufferKOpsExt = unForKOps . forRunKOps $ runFileExts++writeBufferBlobExt :: String+writeBufferBlobExt = unForBlob . forRunBlob $ runFileExts++writeBufferKOpsPath :: WriteBufferFsPaths -> FsPath+writeBufferKOpsPath = flip writeBufferFilePathWithExt writeBufferKOpsExt++writeBufferBlobPath :: WriteBufferFsPaths -> FsPath+writeBufferBlobPath = flip writeBufferFilePathWithExt writeBufferBlobExt++writeBufferChecksumsPath :: WriteBufferFsPaths -> FsPath+writeBufferChecksumsPath = flip writeBufferFilePathWithExt "checksums"++writeBufferFilePathWithExt :: WriteBufferFsPaths -> String -> FsPath+writeBufferFilePathWithExt (WriteBufferFsPaths dir (RunNumber n)) ext =+ dir </> mkFsPath [show n] <.> ext+++{-------------------------------------------------------------------------------+ Checksums For Run Files+-------------------------------------------------------------------------------}++toChecksumsFileForWriteBufferFiles :: (ForKOps CRC.CRC32C, ForBlob CRC.CRC32C) -> CRC.ChecksumsFile+toChecksumsFileForWriteBufferFiles (ForKOps kOpsChecksum, ForBlob blobChecksum) =+ Map.fromList+ [ (toChecksumsFileName writeBufferKOpsExt, kOpsChecksum)+ , (toChecksumsFileName writeBufferBlobExt, blobChecksum)+ ]+ where+ toChecksumsFileName = CRC.ChecksumsFileName . BS.pack++fromChecksumsFileForWriteBufferFiles :: CRC.ChecksumsFile -> Either String (ForKOps CRC.CRC32C, ForBlob CRC.CRC32C)+fromChecksumsFileForWriteBufferFiles file = do+ (,) <$> (ForKOps <$> fromChecksumFile writeBufferKOpsExt) <*> (ForBlob <$> fromChecksumFile writeBufferBlobExt)+ where+ fromChecksumFile key =+ maybe (Left $ "key not found: " <> key) Right $+ Map.lookup (CRC.ChecksumsFileName . fromString $ key) file
@@ -0,0 +1,180 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MagicHash #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | Primitive operations lifted into boxed types+module Database.LSMTree.Internal.Primitive (+ -- * Byte swaps+ byteSwapInt16+ , byteSwapInt32+ , byteSwapInt64+ , byteSwapInt+ , byteSwapWord16+ , byteSwapWord32+ , byteSwapWord64+ , byteSwapWord+ -- * Indexing byte arrays+ , indexInt8Array+ , indexWord8ArrayAsInt16+ , indexWord8ArrayAsInt32+ , indexWord8ArrayAsInt64+ , indexWord8ArrayAsInt+ , indexWord8Array+ , indexWord8ArrayAsWord16+ , indexWord8ArrayAsWord32+ , indexWord8ArrayAsWord64+ , indexWord8ArrayAsWord+ ) where++import Data.Primitive.ByteArray (ByteArray (..))+import GHC.Exts+import GHC.Int+import GHC.Word++{-------------------------------------------------------------------------------+ Conversions+-------------------------------------------------------------------------------}++{-# INLINE wordToInt16# #-}+wordToInt16# :: Word# -> Int16#+wordToInt16# w# = intToInt16# (word2Int# w#)++{-# INLINE int16ToWord# #-}+int16ToWord# :: Int16# -> Word#+int16ToWord# i# = int2Word# (int16ToInt# i#)++{-# INLINE wordToInt32# #-}+wordToInt32# :: Word# -> Int32#+wordToInt32# w# = intToInt32# (word2Int# w#)++{-# INLINE int32ToWord# #-}+int32ToWord# :: Int32# -> Word#+int32ToWord# i# = int2Word# (int32ToInt# i#)++#if MIN_VERSION_base(4,17,0)++{-# INLINE wordToInt64# #-}+wordToInt64# :: Word# -> Int64#+wordToInt64# w# = intToInt64# (word2Int# w#)++{-# INLINE int64ToWord# #-}+int64ToWord# :: Int64# -> Word#+int64ToWord# i# = int2Word# (int64ToInt# i#)++#endif++{-------------------------------------------------------------------------------+ Conversions: shims+-------------------------------------------------------------------------------}++{-# INLINE wordToInt64Shim# #-}+{-# INLINE int64ToWordShim# #-}++#if MIN_VERSION_base(4,17,0)++wordToInt64Shim# :: Word# -> Int64#+wordToInt64Shim# = wordToInt64#++int64ToWordShim# :: Int64# -> Word#+int64ToWordShim# = int64ToWord#++#else++wordToInt64Shim# :: Word# -> Int#+wordToInt64Shim# = word2Int#++int64ToWordShim# :: Int# -> Word#+int64ToWordShim# = int2Word#++#endif++{-------------------------------------------------------------------------------+ Byte swaps+-------------------------------------------------------------------------------}++{-# INLINE byteSwapInt16 #-}+byteSwapInt16 :: Int16 -> Int16+byteSwapInt16 (I16# i#) = I16# (wordToInt16# (byteSwap16# (int16ToWord# i#)))++{-# INLINE byteSwapInt32 #-}+byteSwapInt32 :: Int32 -> Int32+byteSwapInt32 (I32# i#) = I32# (wordToInt32# (byteSwap32# (int32ToWord# i#)))++{-# INLINE byteSwapInt64 #-}+byteSwapInt64 :: Int64 -> Int64+byteSwapInt64 (I64# i#) = I64# (wordToInt64Shim# (byteSwap# (int64ToWordShim# i#)))++{-# INLINE byteSwapInt #-}+byteSwapInt :: Int -> Int+byteSwapInt (I# i#) = I# (word2Int# (byteSwap# (int2Word# i#)))++{-# INLINE byteSwapWord16 #-}+byteSwapWord16 :: Word16 -> Word16+byteSwapWord16 = byteSwap16++{-# INLINE byteSwapWord32 #-}+byteSwapWord32 :: Word32 -> Word32+byteSwapWord32 = byteSwap32++{-# INLINE byteSwapWord64 #-}+byteSwapWord64 :: Word64 -> Word64+byteSwapWord64 = byteSwap64++{-# INLINE byteSwapWord #-}+byteSwapWord :: Word -> Word+byteSwapWord (W# w#) = W# (byteSwap# w#)++{-------------------------------------------------------------------------------+ Indexing byte arrays+-------------------------------------------------------------------------------}++{-# INLINE indexInt8Array #-}+indexInt8Array :: ByteArray -> Int -> Int8+indexInt8Array (ByteArray !ba#) (I# !off#) =+ I8# (indexInt8Array# ba# off#)++{-# INLINE indexWord8ArrayAsInt16 #-}+indexWord8ArrayAsInt16 :: ByteArray -> Int -> Int16+indexWord8ArrayAsInt16 (ByteArray !ba#) (I# !off#) =+ I16# (indexWord8ArrayAsInt16# ba# off#)++{-# INLINE indexWord8ArrayAsInt32 #-}+indexWord8ArrayAsInt32 :: ByteArray -> Int -> Int32+indexWord8ArrayAsInt32 (ByteArray !ba#) (I# !off#) =+ I32# (indexWord8ArrayAsInt32# ba# off#)++{-# INLINE indexWord8ArrayAsInt64 #-}+indexWord8ArrayAsInt64 :: ByteArray -> Int -> Int64+indexWord8ArrayAsInt64 (ByteArray !ba#) (I# !off#) =+ I64# (indexWord8ArrayAsInt64# ba# off#)++{-# INLINE indexWord8ArrayAsInt #-}+indexWord8ArrayAsInt :: ByteArray -> Int -> Int+indexWord8ArrayAsInt (ByteArray !ba#) (I# !off#) =+ I# (indexWord8ArrayAsInt# ba# off#)++{-# INLINE indexWord8Array #-}+indexWord8Array :: ByteArray -> Int -> Word8+indexWord8Array (ByteArray !ba#) (I# !off#) =+ W8# (indexWord8Array# ba# off#)++{-# INLINE indexWord8ArrayAsWord16 #-}+indexWord8ArrayAsWord16 :: ByteArray -> Int -> Word16+indexWord8ArrayAsWord16 (ByteArray !ba#) (I# !off#) =+ W16# (indexWord8ArrayAsWord16# ba# off#)++{-# INLINE indexWord8ArrayAsWord32 #-}+indexWord8ArrayAsWord32 :: ByteArray -> Int -> Word32+indexWord8ArrayAsWord32 (ByteArray !ba#) (I# !off#) =+ W32# (indexWord8ArrayAsWord32# ba# off#)++{-# INLINE indexWord8ArrayAsWord64 #-}+indexWord8ArrayAsWord64 :: ByteArray -> Int -> Word64+indexWord8ArrayAsWord64 (ByteArray !ba#) (I# !off#) =+ W64# (indexWord8ArrayAsWord64# ba# off#)++{-# INLINE indexWord8ArrayAsWord #-}+indexWord8ArrayAsWord :: ByteArray -> Int -> Word+indexWord8ArrayAsWord (ByteArray !ba#) (I# !off#) =+ W# (indexWord8ArrayAsWord# ba# off#)
@@ -0,0 +1,28 @@+{-# LANGUAGE DeriveFunctor #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Range (+ Range (..)+ ) where++import Control.DeepSeq (NFData (..))++{-------------------------------------------------------------------------------+ Small auxiliary types+-------------------------------------------------------------------------------}++-- | A range of keys.+data Range k =+ {- |+ @'FromToExcluding' i j@ is the range from @i@ (inclusive) to @j@ (exclusive).+ -}+ FromToExcluding k k+ {- |+ @'FromToIncluding' i j@ is the range from @i@ (inclusive) to @j@ (inclusive).+ -}+ | FromToIncluding k k+ deriving stock (Show, Eq, Functor)++instance NFData k => NFData (Range k) where+ rnf (FromToExcluding k1 k2) = rnf k1 `seq` rnf k2+ rnf (FromToIncluding k1 k2) = rnf k1 `seq` rnf k2
@@ -0,0 +1,312 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_HADDOCK not-home #-}++#ifndef __HLINT__+-- HLint would fail to find this file and emit a warning+#include <MachDeps.h>+#endif++-- |+--+-- This module is intended to be imported qualified, to avoid name clashes with+-- "Prelude" functions:+--+-- @+-- import Database.LSMTree.Internal.RawBytes (RawBytes (..))+-- import qualified Database.LSMTree.Internal.RawBytes as RB+-- @+module Database.LSMTree.Internal.RawBytes (+ -- See Note: [Export structure]+ -- * Raw bytes+ RawBytes (..)+ -- * Accessors+ -- ** Length information+ , size+ -- ** Extracting subvectors (slicing)+ , take+ , drop+ , topBits64+ -- * Construction+ -- | Use 'Semigroup' and 'Monoid' operations+ -- ** Restricting memory usage+ , copy+ , force+ -- * Conversions+ , fromVector+ , fromByteArray+ -- ** Lists+ , pack+ , unpack+ -- * @bytestring@ utils+ , fromByteString+ , unsafeFromByteString+ , toByteString+ , unsafePinnedToByteString+ , fromShortByteString+ , builder+ ) where++import Control.DeepSeq (NFData)+import Data.Bits (Bits (..))+import Data.BloomFilter.Hash (Hashable (..), hashByteArray)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BB+import Data.ByteString.Short (ShortByteString (SBS))+import qualified Data.ByteString.Short as SBS+import Data.Primitive.ByteArray (ByteArray (..), compareByteArrays)+import qualified Data.Primitive.ByteArray as BA+import qualified Data.Vector.Primitive as VP+import Database.LSMTree.Internal.ByteString (byteArrayToByteString,+ shortByteStringFromTo, tryGetByteArray,+ unsafePinnedByteArrayToByteString)+import Database.LSMTree.Internal.Vector+import Prelude hiding (drop, take)++import GHC.Exts+import GHC.Stack+import GHC.Word+import Text.Printf (printf)++-- $setup+-- >>> import Numeric++{- Note: [Export structure]+ ~~~~~~~~~~~~~~~~~~~~~~~+ Since RawBytes are very similar to Primitive Vectors, the code is sectioned+ and structured much like the "Data.Vector.Primitive" module.+-}++{-------------------------------------------------------------------------------+ Raw bytes+-------------------------------------------------------------------------------}++{- |+Raw bytes.++This type imposes no alignment constraint and provides no guarantee of whether the memory is pinned or unpinned.+-}+newtype RawBytes = RawBytes (VP.Vector Word8)+ deriving newtype (NFData)++-- TODO: Should we have a more well-behaved instance for 'Show'?+-- For instance, an instance that prints the bytes as a hexadecimal string?+deriving newtype instance Show RawBytes++_showBytesAsHex :: RawBytes -> ShowS+_showBytesAsHex (RawBytes bytes) = VP.foldr ((.) . showByte) id bytes+ where+ showByte :: Word8 -> ShowS+ showByte = showString . printf "%02x"++instance Eq RawBytes where+ bs1 == bs2 = compareBytes bs1 bs2 == EQ++{- |+This instance uses lexicographic ordering.+-}+instance Ord RawBytes where+ compare = compareBytes++-- | Based on @Ord 'ShortByteString'@.+compareBytes :: RawBytes -> RawBytes -> Ordering+compareBytes rb1@(RawBytes vec1) rb2@(RawBytes vec2) =+ let !len1 = size rb1+ !len2 = size rb2+ !len = min len1 len2+ in case compareByteArrays ba1 off1 ba2 off2 len of+ EQ | len1 < len2 -> LT+ | len1 > len2 -> GT+ o -> o+ where+ VP.Vector off1 _size1 ba1 = vec1+ VP.Vector off2 _size2 ba2 = vec2++instance Hashable RawBytes where+ hashSalt64 :: Word64 -> RawBytes -> Word64+ hashSalt64 = hash++hash :: Word64 -> RawBytes -> Word64+hash salt (RawBytes (VP.Vector off len ba)) = hashByteArray ba off len salt++{- |+@'fromList'@: \(O(n)\).++@'toList'@: \(O(n)\).+-}+instance IsList RawBytes where+ type Item RawBytes = Word8++ fromList :: [Item RawBytes] -> RawBytes+ fromList = pack++ toList :: RawBytes -> [Item RawBytes]+ toList = unpack++{- |+@'fromString'@: \(O(n)\).++__Warning:__ 'fromString' truncates multi-byte characters to octets. e.g. \"枯朶に烏のとまりけり秋の暮\" becomes \"�6k�nh~�Q��n�\".+-}+instance IsString RawBytes where+ fromString = fromByteString . fromString++{-------------------------------------------------------------------------------+ Accessors+-------------------------------------------------------------------------------}++-- | \( O(1) \)+size :: RawBytes -> Int+size = coerce VP.length++-- | \( O(1) \)+take :: Int -> RawBytes -> RawBytes+take = coerce VP.take++-- | \( O(1) \)+drop :: Int -> RawBytes -> RawBytes+drop = coerce VP.drop++-- | @'topBits64' rb@ slices the first @64@ bits from the /top/ of the raw bytes+-- @rb@. Returns the string of bits as a 'Word64'.+--+-- The /top/ corresponds to the most significant bit (big-endian).+--+-- If the number of bits is smaller than @64@, then any missing bits default to+-- @0@s.+--+-- >>> showHex (topBits64 (pack [1,0,0,0,0,0,0,0])) ""+-- "100000000000000"+--+-- >>> showHex (topBits64 (pack [1,0,0])) ""+-- "100000000000000"+--+-- TODO: optimisation ideas: use unsafe shift/byteswap primops, look at GHC+-- core, find other opportunities for using primops.+--+topBits64 :: RawBytes -> Word64+topBits64 rb@(RawBytes v@(VP.Vector (I# off#) _size (ByteArray k#)))+ | n >= 8+ = toWord64 (indexWord8ArrayAsWord64# k# off#)+ | otherwise+ = VP.foldl' f 0 v `unsafeShiftL` ((8 - n) * 8)+ where+ !n = size rb++ f :: Word64 -> Word8 -> Word64+ f acc w = acc `unsafeShiftL` 8 + fromIntegral w++#if (MIN_VERSION_GLASGOW_HASKELL(9, 4, 0, 0))+toWord64 :: Word64# -> Word64+#else+toWord64 :: Word# -> Word64+#endif+#if WORDS_BIGENDIAN+toWord64 = W64#+#else+toWord64 x# = byteSwap64 (W64# x#)+#endif++{-------------------------------------------------------------------------------+ Construction+-------------------------------------------------------------------------------}++{- |+@('<>')@: \(O(n)\).++@'Data.Semigroup.sconcat'@: \(O(n)\).+-}+instance Semigroup RawBytes where+ (<>) = coerce (VP.++)++{- |+@'mempty'@: \(O(1)\).++@'mconcat'@: \(O(n)\).+-}+instance Monoid RawBytes where+ mempty = coerce VP.empty+ mconcat = coerce VP.concat++-- | O(n) Yield the argument, but force it not to retain any extra memory by+-- copying it.+--+-- Useful when dealing with slices. Also, see+-- "Database.LSMTree.Internal.Unsliced"+copy :: RawBytes -> RawBytes+copy (RawBytes pvec) = RawBytes (VP.force pvec)++-- | Force 'RawBytes' to not retain any extra memory. This may copy the contents.+force :: RawBytes -> ByteArray+force (RawBytes (VP.Vector off len ba))+ | off == 0+ , BA.sizeofByteArray ba == len+ = ba++ | otherwise+ = BA.runByteArray $ do+ mba <- BA.newByteArray len+ BA.copyByteArray mba 0 ba off len+ pure mba++{-------------------------------------------------------------------------------+ Conversions+-------------------------------------------------------------------------------}++-- | \( O(1) \)+fromVector :: VP.Vector Word8 -> RawBytes+fromVector v = RawBytes v++-- | \( O(1) \)+fromByteArray :: Int -> Int -> ByteArray -> RawBytes+fromByteArray off len ba = RawBytes (mkPrimVector off len ba)++pack :: [Word8] -> RawBytes+pack = coerce VP.fromList++unpack :: RawBytes -> [Word8]+unpack = coerce VP.toList++{-------------------------------------------------------------------------------+ @bytestring@ utils+-------------------------------------------------------------------------------}++-- | \( O(n) \) conversion from a strict bytestring to raw bytes.+fromByteString :: BS.ByteString -> RawBytes+fromByteString = fromShortByteString . SBS.toShort++-- | \( O(1) \) conversion from a strict bytestring to raw bytes.+--+-- Strict bytestrings are allocated using 'mallocPlainForeignPtrBytes', so we+-- are expecting a 'PlainPtr' (or 'FinalPtr' with length 0).+-- For other variants, this function will fail.+unsafeFromByteString :: HasCallStack => BS.ByteString -> RawBytes+unsafeFromByteString bs =+ case tryGetByteArray bs of+ Right (ba, n) -> RawBytes (mkPrimVector 0 n ba)+ Left err -> error $ "unsafeFromByteString: " <> err++-- | \( O(1) \) conversion from raw bytes to a bytestring if pinned,+-- \( O(n) \) if unpinned.+toByteString :: RawBytes -> BS.ByteString+toByteString (RawBytes (VP.Vector off len ba)) =+ byteArrayToByteString off len ba++-- | \( O(1) \) conversion from raw bytes to a bytestring.+-- Fails if the underlying byte array is not pinned.+unsafePinnedToByteString :: HasCallStack => RawBytes -> BS.ByteString+unsafePinnedToByteString (RawBytes (VP.Vector off len ba)) =+ unsafePinnedByteArrayToByteString off len ba++-- | \( O(1) \) conversion from a short bytestring to raw bytes.+fromShortByteString :: ShortByteString -> RawBytes+fromShortByteString sbs@(SBS ba#) =+ RawBytes (mkPrimVector 0 (SBS.length sbs) (ByteArray ba#))++{-# INLINE builder #-}+builder :: RawBytes -> BB.Builder+builder (RawBytes (VP.Vector off sz (ByteArray ba#))) =+ shortByteStringFromTo off (off + sz) (SBS ba#)
@@ -0,0 +1,173 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.RawOverflowPage (+ RawOverflowPage (..),+ makeRawOverflowPage,+ unsafeMakeRawOverflowPage,+ rawOverflowPageRawBytes,+ rawOverflowPageToByteString,+ rawBytesToOverflowPages,+ pinnedByteArrayToOverflowPages,+ unpinnedByteArrayToOverflowPages,+) where++import Control.DeepSeq (NFData (rnf))+import Control.Monad (when)+import Data.ByteString (ByteString)+import Data.Primitive.ByteArray (ByteArray (..), copyByteArray,+ fillByteArray, isByteArrayPinned, newPinnedByteArray,+ runByteArray)+import qualified Data.Vector.Primitive as VP+import Database.LSMTree.Internal.Assertions+import Database.LSMTree.Internal.BitMath (roundUpToPageSize)+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB++-------------------------------------------------------------------------------+-- RawOverflowPage type+-------------------------------------------------------------------------------++-- | When a key\/op pair is too large to fit in a single disk page, the+-- representation is split into a normal page, and one or more overflow pages.+-- The normal 'RawPage' follows the run page format, and contains the key,+-- optional blob reference and a prefix of the value, while the overflow pages+-- contain the suffix of a large value that didn't fit in the normal page.+--+-- Each overflow page is the same size as normal pages (currently 4096 only).+--+data RawOverflowPage = RawOverflowPage+ !Int -- ^ offset in Word8s.+ !ByteArray+ deriving stock (Show)++-- | This invariant is the same as for 'RawPage', but there is no alignment+-- constraint. This is for two reasons: 1. we don't need alignment, because+-- the page has no structure that we need to read with aligned memory+-- operations; 2. we don't want alignment because we want to convert the suffix+-- of serialised values (as 'RawBytes') into a 'RawOverflowPage' without+-- copying, and a suffix can start at an arbitrary offset.+--+invariant :: RawOverflowPage -> Bool+invariant (RawOverflowPage off ba) = and+ [ isValidSlice off 4096 ba -- valid bytearray slice for length 4096+ , isByteArrayPinned ba -- bytearray must be pinned (for I/O)+ ]++instance NFData RawOverflowPage where+ rnf (RawOverflowPage _ _) = ()++-- | This instance assumes pages are 4096 bytes in size+instance Eq RawOverflowPage where+ r1 == r2 = rawOverflowPageRawBytes r1 == rawOverflowPageRawBytes r2++rawOverflowPageRawBytes :: RawOverflowPage -> RawBytes+rawOverflowPageRawBytes (RawOverflowPage off ba) =+ RB.fromByteArray off 4096 ba++-- | \( O(1) \) since we can avoid copying the pinned byte array.+rawOverflowPageToByteString :: RawOverflowPage -> ByteString+rawOverflowPageToByteString =+ RB.unsafePinnedToByteString . rawOverflowPageRawBytes++-- | Create a 'RawOverflowPage'.+--+-- The length must be 4096 or less.+--+-- This function will copy data if the byte array is not pinned, or the length+-- is strictly less than 4096.+--+makeRawOverflowPage ::+ ByteArray -- ^ bytearray+ -> Int -- ^ offset in bytes into the bytearray+ -> Int -- ^ length in bytes, must be @>= 0 && <= 4096@+ -> RawOverflowPage+makeRawOverflowPage ba off len+ | len == 4096+ , let page = RawOverflowPage off ba+ , invariant page+ = page++ | otherwise+ = makeRawOverflowPageCopy ba off len++makeRawOverflowPageCopy ::+ ByteArray -- ^ bytearray+ -> Int -- ^ offset in bytes into the bytearray+ -> Int -- ^ length in bytes+ -> RawOverflowPage+makeRawOverflowPageCopy ba off len =+ assert (isValidSlice off len ba && len <= 4096) $+ (\page -> assert (invariant page) page) $+ RawOverflowPage 0 $ runByteArray $ do+ mba <- newPinnedByteArray 4096+ let suffixlen = min 4096 len -- would only do anything with assertions off+ copyByteArray mba 0 ba off suffixlen+ when (suffixlen < 4096) $ fillByteArray mba suffixlen (4096-suffixlen) 0+ pure mba++-- | Create a 'RawOverflowPage' without copying. The byte array and offset must+-- satisfy the invariant for 'RawOverflowPage'.+--+unsafeMakeRawOverflowPage ::+ ByteArray -- ^ bytearray, must be pinned and contain 4096 bytes (after offset)+ -> Int -- ^ offset in bytes+ -> RawOverflowPage+unsafeMakeRawOverflowPage ba off =+ assert (invariant page) page+ where+ page = RawOverflowPage off ba++-- | Convert 'RawBytes' representing the \"overflow\" part of a value into one+-- or more 'RawOverflowPage's.+--+-- This will avoid copying where possible.+--+rawBytesToOverflowPages :: RawBytes -> [RawOverflowPage]+rawBytesToOverflowPages (RawBytes (VP.Vector off len ba))+ | isByteArrayPinned ba+ = pinnedByteArrayToOverflowPages off len ba+ | otherwise+ = unpinnedByteArrayToOverflowPages off len ba++pinnedByteArrayToOverflowPages :: Int -> Int -> ByteArray -> [RawOverflowPage]+pinnedByteArrayToOverflowPages !off !len !ba+ | len >= 4096+ , let !page = unsafeMakeRawOverflowPage ba off+ = page : pinnedByteArrayToOverflowPages (off+4096) (len-4096) ba++ | len <= 0+ = []++ | otherwise -- > 0 && < 4096+ -- have to copy the partial last page+ , let !page = makeRawOverflowPageCopy ba off len+ = page : []++-- | Not pinned, in principle shouldn't happen much because if the value+-- is big enough to overflow then it's big enough to be pinned.+-- It is possible however if a page has a huge key and a small value.+--+-- Unfortunately, with GHC versions 9.6.x we also get this because the meaning+-- of pinned has changed. Sigh.+-- See <https://gitlab.haskell.org/ghc/ghc/-/issues/22255>+--+unpinnedByteArrayToOverflowPages :: Int -> Int -> ByteArray -> [RawOverflowPage]+unpinnedByteArrayToOverflowPages !off !len !ba =+ let !lenPages = roundUpToPageSize len+ ba' = runByteArray $ do+ mba <- newPinnedByteArray lenPages+ copyByteArray mba 0 ba off len+ fillByteArray mba len (lenPages-len) 0+ pure mba+ pages = pinnedByteArrayToOverflowPages 0 lenPages ba'+ -- We've arranged to do the conversion without any extra copying,+ -- so assert that we got that right:+ in assert (all (isSliceNotCopy ba') pages)+ pages+ where+ isSliceNotCopy :: ByteArray -> RawOverflowPage -> Bool+ isSliceNotCopy ba1 (RawOverflowPage _ ba2) = sameByteArray ba1 ba2+
@@ -0,0 +1,474 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.RawPage (+ RawPage (..),+ emptyRawPage,+ makeRawPage,+ unsafeMakeRawPage,+ rawPageRawBytes,+ rawPageNumKeys,+ rawPageNumBlobs,+ rawPageLookup,+ RawPageLookup(..),+ rawPageOverflowPages,+ rawPageFindKey,+ rawPageIndex,+ RawPageIndex(..),+ getRawPageIndexKey,+ -- * Test and debug+ rawPageKeyOffsets,+ rawPageValueOffsets,+ rawPageValueOffsets1,+ rawPageHasBlobSpanAt,+ rawPageBlobSpanIndex,+ rawPageOpAt,+ rawPageKeys,+ rawPageValues,+) where++import Control.DeepSeq (NFData (rnf))+import Control.Exception (assert)+import Data.Bits (complement, popCount, unsafeShiftL, unsafeShiftR,+ (.&.))+import Data.Primitive.ByteArray (ByteArray (..), byteArrayFromList,+ copyByteArray, fillByteArray, indexByteArray,+ isByteArrayPinned, newAlignedPinnedByteArray, runByteArray,+ sizeofByteArray)+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import Data.Word (Word16, Word32, Word64, Word8)+import Database.LSMTree.Internal.BitMath+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry (Entry (..))+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise (SerialisedKey (..),+ SerialisedValue (..))+import Database.LSMTree.Internal.Vector+import qualified GHC.List as List++-------------------------------------------------------------------------------+-- RawPage type+-------------------------------------------------------------------------------++data RawPage = RawPage+ !Int -- ^ offset in Word16s.+ !ByteArray+ deriving stock (Show)++emptyRawPage :: RawPage+emptyRawPage = flip makeRawPage 0 $ byteArrayFromList+ [ 0, 0, 0, 0+ , 8, 0, 0, 0+ , 10 :: Word8+ ]++invariant :: RawPage -> Bool+invariant (RawPage off ba) = and+ [ off >= 0 -- offset is positive+ , mod8 offsetInBytes == 0 -- 64 bit/8 byte alignment+ , (sizeofByteArray ba - offsetInBytes) >= 4096 -- there is always 4096 bytes+ , isByteArrayPinned ba -- bytearray should be pinned+ ]+ where+ offsetInBytes = mul2 off++instance NFData RawPage where+ rnf (RawPage _ _) = ()++-- | This instance assumes pages are 4096 bytes in size+instance Eq RawPage where+ RawPage off1 ba1 == RawPage off2 ba2 = v1 == v2+ where+ v1, v2 :: VP.Vector Word16+ v1 = mkPrimVector off1 2048 ba1+ v2 = mkPrimVector off2 2048 ba2++-- | Create 'RawPage'.+--+-- This function may copy data to satisfy internal 'RawPage' invariants.+-- Use 'unsafeMakeRawPage' if you don't want copy.+makeRawPage ::+ ByteArray -- ^ bytearray, must contain 4096 bytes (after offset)+ -> Int -- ^ offset in bytes, must be 8 byte aligned.+ -> RawPage+makeRawPage ba off+ | invariant page = page+ | otherwise = RawPage 0 $ runByteArray $ do+ mba <- newAlignedPinnedByteArray 4096 8+ fillByteArray mba 0 4096 0+ copyByteArray mba 0 ba off (clamp 0 4096 (sizeofByteArray ba - off))+ pure mba+ where+ page = RawPage (div2 off) ba+ clamp l u x = max l (min u x)++unsafeMakeRawPage ::+ ByteArray -- ^ bytearray, must be pinned and contain 4096 bytes (after offset)+ -> Int -- ^ offset in bytes, must be 8 byte aligned.+ -> RawPage+unsafeMakeRawPage ba off = assert (invariant page) page+ where+ page = RawPage (div2 off) ba++rawPageRawBytes :: RawPage -> RawBytes+rawPageRawBytes (RawPage off ba) =+ RB.fromByteArray (mul2 off) 4096 ba++-------------------------------------------------------------------------------+-- Lookup function+-------------------------------------------------------------------------------++data RawPageLookup entry =+ -- | The key is not present on the page.+ LookupEntryNotPresent++ -- | The key is present and corresponds to a normal entry that fits+ -- fully within the page.+ | LookupEntry !entry++ -- | The key is present and corresponds to an entry where the value+ -- may have overflowed onto subsequent pages. In this case the entry+ -- contains the /prefix/ of the value (that did fit on the page). The+ -- length of the suffix is returned separately.+ | LookupEntryOverflow !entry !Word32+ deriving stock (Eq, Functor, Show)++-- |+-- __Time:__ \( \mathcal{O}\left( \log_2 n \right) \)+-- where \( n \) is the number of entries in the 'RawPage'.+--+-- Return the 'Entry' corresponding to the supplied 'SerialisedKey' if it exists+-- within the 'RawPage'.+rawPageLookup ::+ RawPage+ -> SerialisedKey+ -> RawPageLookup (Entry SerialisedValue BlobSpan)+rawPageLookup !page !key+ | dirNumKeys == 1 = lookup1+ | otherwise = case bisectPageToKey (fromIntegral dirNumKeys) page key of+ KeyWasFoundAt i -> LookupEntry $ rawPageEntryAt page i+ _ -> LookupEntryNotPresent+ where+ !dirNumKeys = rawPageNumKeys page++ lookup1+ | key == rawPageKeyAt page 0+ = let !entry = rawPageEntry1 page+ !suffix = rawPageSingleValueSuffix page+ in if suffix > 0+ then LookupEntryOverflow entry suffix+ else LookupEntry entry+ | otherwise+ = LookupEntryNotPresent++-- |+-- __Time:__ \( \mathcal{O}\left( \log_2 n \right) \)+-- where \( n \) is the number of entries in the 'RawPage'.+--+-- Return the least entry number in the 'RawPage' (if it exists)+-- which is greater than or equal to the supplied 'SerialisedKey'.+--+-- The following law always holds \( \forall \mathtt{key} \mathtt{page} \):+--+-- * maybe True (key <=) (getRawPageIndexKey . rawPageIndex page . id =<< rawPageFindKey page key)+--+-- * maybe True (key > ) (getRawPageIndexKey . rawPageIndex page . pred =<< rawPageFindKey page key)+--+-- * maybe (maximum (rawPageKeys page) < key) (rawPageFindKey page key)+rawPageFindKey ::+ RawPage+ -> SerialisedKey+ -> Maybe Word16 -- ^ entry number of first entry greater or equal to the key+rawPageFindKey !page !key+ | dirNumKeys == 1 = lookup1+ | otherwise = case bisectPageToKey (fromIntegral dirNumKeys) page key of+ KeyNotFoundExceedsPage -> Nothing+ KeyNotFoundNearestIsAt i -> Just $ fromIntegral i+ KeyWasFoundAt i -> Just $ fromIntegral i++ where+ !dirNumKeys = rawPageNumKeys page++ lookup1+ | key <= rawPageKeyAt page 0 = Just 0+ | otherwise = Nothing++-- | This data-type facilitates the code reuse of 'bisectPageToKey' between the+-- and 'rawPageLookup' and 'rawPageFindKey'.+data BinarySearchResult+ = KeyNotFoundExceedsPage+ | KeyNotFoundNearestIsAt {-# UNPACK #-} !Int+ | KeyWasFoundAt {-# UNPACK #-} !Int++-- | Binary search procedure shared between 'rawPageLookup' and 'rawPageFindKey'.+{-# INLINE bisectPageToKey #-}+bisectPageToKey ::+ Int+ -> RawPage+ -> SerialisedKey+ -> BinarySearchResult+bisectPageToKey !numKeys !page !key = go 0 numKeys+ where+ -- Binary search for within the range [i, j)+ -- for the index k corresponding to key.+ --+ -- If k > key, search [ k + 1, j )+ -- If k < key, search [ i, j )+ go :: Int -> Int -> BinarySearchResult+ go !i !j+ | i >= j =+ if i == numKeys+ then KeyNotFoundExceedsPage+ else KeyNotFoundNearestIsAt j+ | otherwise =+ let !k = i + div2 (j - i)+ in case key `compare` rawPageKeyAt page k of+ GT -> go (k + 1) j+ EQ -> KeyWasFoundAt k+ LT -> go i k++data RawPageIndex entry =+ IndexNotPresent+ -- | The index is present and corresponds to a normal entry that fits+ -- fully within the page (but might be the only entry on the page).+ | IndexEntry !SerialisedKey !entry+ -- | The index is present and corresponds to an entry where the value+ -- has overflowed onto subsequent pages. In this case only prefix entry+ -- and the length of the suffix are returned.+ -- The caller can copy the full serialised pages themselves.+ | IndexEntryOverflow !SerialisedKey !entry !Word32+ deriving stock (Eq, Functor, Show)++-- |+-- __Time:__ \( \mathcal{O}\left( 1 \right) \)+--+-- Conveniently access the 'SerialisedKey' of a 'RawPageIndex'.+{-# INLINE getRawPageIndexKey #-}+getRawPageIndexKey :: RawPageIndex e -> Maybe SerialisedKey+getRawPageIndexKey = \case+ IndexEntry k _ -> Just k+ IndexEntryOverflow k _ _ -> Just k+ IndexNotPresent -> Nothing+++{-# INLINE rawPageIndex #-}+rawPageIndex ::+ RawPage+ -> Word16+ -> RawPageIndex (Entry SerialisedValue BlobSpan)+rawPageIndex !page !ix+ | ix >= dirNumKeys =+ IndexNotPresent+ | dirNumKeys > 1 =+ let ix' = fromIntegral ix+ in IndexEntry (rawPageKeyAt page ix') (rawPageEntryAt page ix')+ | otherwise =+ let key = rawPageKeyAt page 0+ entry = rawPageEntry1 page+ !suffix = rawPageSingleValueSuffix page+ in if suffix <= 0+ then IndexEntry key entry+ else IndexEntryOverflow key entry suffix+ where+ !dirNumKeys = rawPageNumKeys page++-- | for non-single key page case+rawPageEntryAt :: RawPage -> Int -> Entry SerialisedValue BlobSpan+rawPageEntryAt page i =+ assert (i < fromIntegral (rawPageNumKeys page)) $+ assert (rawPageNumKeys page > 1) $+ case rawPageOpAt page i of+ 0 -> if rawPageHasBlobSpanAt page i == 0+ then Insert (rawPageValueAt page i)+ else InsertWithBlob (rawPageValueAt page i)+ (rawPageBlobSpanIndex page+ (rawPageCalculateBlobIndex page i))+ 1 -> Upsert (rawPageValueAt page i)+ _ -> Delete++-- | Single key page case+rawPageEntry1 :: RawPage -> Entry SerialisedValue BlobSpan+rawPageEntry1 page =+ case rawPageOpAt page 0 of+ 0 -> if rawPageHasBlobSpanAt page 0 == 0+ then Insert (rawPageSingleValuePrefix page)+ else InsertWithBlob (rawPageSingleValuePrefix page)+ (rawPageBlobSpanIndex page+ (rawPageCalculateBlobIndex page 0))+ 1 -> Upsert (rawPageSingleValuePrefix page)+ _ -> Delete++-- | Calculate the number of overflow pages that are expected to follow this+-- page.+--+-- This will be non-zero when the page contains a single key\/op entry that is+-- itself too large to fit within the page.+--+rawPageOverflowPages :: RawPage -> Int+rawPageOverflowPages page+ | rawPageNumKeys page == 1+ , let (_, end) = rawPageValueOffsets1 page+ = fromIntegral (ceilDivPageSize end - 1) -- don't count the first page++ | otherwise = 0++-------------------------------------------------------------------------------+-- Accessors+-------------------------------------------------------------------------------++-- Note: indexByteArray's offset is given in elements of type a rather than in bytes.++rawPageNumKeys :: RawPage -> Word16+rawPageNumKeys (RawPage off ba) = indexByteArray ba off++rawPageNumBlobs :: RawPage -> Word16+rawPageNumBlobs (RawPage off ba) = indexByteArray ba (off + 1)++rawPageKeysOffset :: RawPage -> Word16+rawPageKeysOffset (RawPage off ba) = indexByteArray ba (off + 2)++type KeyOffset = Word16+type ValueOffset = Word16++rawPageKeyOffsets :: RawPage -> VP.Vector KeyOffset+rawPageKeyOffsets page@(RawPage off ba) =+ mkPrimVector+ (off + fromIntegral (div2 dirOffset))+ (fromIntegral dirNumKeys + 1)+ ba+ where+ !dirNumKeys = rawPageNumKeys page+ !dirOffset = rawPageKeysOffset page++-- | for non-single key page case+rawPageValueOffsets :: RawPage -> VP.Vector ValueOffset+rawPageValueOffsets page@(RawPage off ba) =+ assert (dirNumKeys /= 1) $+ mkPrimVector+ (off + fromIntegral (div2 dirOffset) + fromIntegral dirNumKeys)+ (fromIntegral dirNumKeys + 1)+ ba+ where+ !dirNumKeys = rawPageNumKeys page+ !dirOffset = rawPageKeysOffset page++-- | single key page case+{-# INLINE rawPageValueOffsets1 #-}+rawPageValueOffsets1 :: RawPage -> (Word16, Word32)+rawPageValueOffsets1 page@(RawPage off ba) =+ assert (rawPageNumKeys page == 1) $+ ( indexByteArray ba (off + fromIntegral (div2 dirOffset) + 1)+ , indexByteArray ba (div2 (off + fromIntegral (div2 dirOffset)) + 1)+ )+ where+ !dirOffset = rawPageKeysOffset page++rawPageHasBlobSpanAt :: RawPage -> Int -> Word64+rawPageHasBlobSpanAt _page@(RawPage off ba) i = do+ let j = div64 i+ let k = mod64 i+ let word = indexByteArray ba (div4 off + 1 + j)+ unsafeShiftR word k .&. 1++rawPageOpAt :: RawPage -> Int -> Word64+rawPageOpAt page@(RawPage off ba) i = do+ let j = div32 i+ let k = mod32 i+ let word = indexByteArray ba (div4 off + 1 + ceilDiv64 (fromIntegral dirNumKeys) + j)+ unsafeShiftR word (mul2 k) .&. 3+ where+ !dirNumKeys = rawPageNumKeys page++rawPageKeys :: RawPage -> V.Vector SerialisedKey+rawPageKeys page@(RawPage off ba) = do+ let offs = rawPageKeyOffsets page+ V.fromList+ [ SerialisedKey (RB.fromByteArray (mul2 off + start) (end - start) ba)+ | i <- [ 0 .. fromIntegral dirNumKeys - 1 ] :: [Int]+ , let start = fromIntegral (VP.unsafeIndex offs i) :: Int+ , let end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int+ ]+ where+ !dirNumKeys = rawPageNumKeys page++rawPageKeyAt :: RawPage -> Int -> SerialisedKey+rawPageKeyAt page@(RawPage off ba) i = do+ SerialisedKey (RB.fromByteArray (mul2 off + start) (end - start) ba)+ where+ offs = rawPageKeyOffsets page+ start = fromIntegral (VP.unsafeIndex offs i) :: Int+ end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int++-- | Non-single page case+rawPageValues :: RawPage -> V.Vector SerialisedValue+rawPageValues page@(RawPage off ba) =+ let offs = rawPageValueOffsets page in+ V.fromList+ [ SerialisedValue $ RB.fromByteArray (mul2 off + start) (end - start) ba+ | i <- [ 0 .. fromIntegral dirNumKeys - 1 ] :: [Int]+ , let start = fromIntegral (VP.unsafeIndex offs i) :: Int+ , let end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int+ ]+ where+ !dirNumKeys = rawPageNumKeys page++-- | Non-single page case+rawPageValueAt :: RawPage -> Int -> SerialisedValue+rawPageValueAt page@(RawPage off ba) i =+ SerialisedValue (RB.fromByteArray (mul2 off + start) (end - start) ba)+ where+ offs = rawPageValueOffsets page+ start = fromIntegral (VP.unsafeIndex offs i) :: Int+ end = fromIntegral (VP.unsafeIndex offs (i + 1)) :: Int++-- | Single key page case+rawPageSingleValuePrefix :: RawPage -> SerialisedValue+rawPageSingleValuePrefix page@(RawPage off ba) =+ SerialisedValue $+ RB.fromByteArray+ (mul2 off + fromIntegral start)+ (fromIntegral prefix_end - fromIntegral start)+ ba+ where+ (start, end) = rawPageValueOffsets1 page+ prefix_end = min 4096 end++-- | Single key page case+rawPageSingleValueSuffix :: RawPage -> Word32+rawPageSingleValueSuffix page+ | end > 4096 = end - 4096+ | otherwise = 0+ where+ (_, end) = rawPageValueOffsets1 page++-- we could create unboxed array.+{-# INLINE rawPageBlobSpanIndex #-}+rawPageBlobSpanIndex :: RawPage+ -> Int -- ^ blobspan index. Calculate with 'rawPageCalculateBlobIndex'+ -> BlobSpan+rawPageBlobSpanIndex page@(RawPage off ba) i = BlobSpan+ ( indexByteArray ba (off1 + i) )+ ( indexByteArray ba (off2 + i) )+ where+ !dirNumKeys = rawPageNumKeys page+ !dirNumBlobs = rawPageNumBlobs page++ -- offset to start of blobspan arr+ off1 = div4 off + 1 + ceilDiv64 (fromIntegral dirNumKeys) + ceilDiv64 (fromIntegral (mul2 dirNumKeys))+ off2 = mul2 (off1 + fromIntegral dirNumBlobs)++rawPageCalculateBlobIndex ::+ RawPage+ -> Int -- ^ key index+ -> Int -- ^ blobspan index+rawPageCalculateBlobIndex (RawPage off ba) i = do+ let j = unsafeShiftR i 6 -- `div` 64+ let k = i .&. 63 -- `mod` 64+ -- generic sum isn't too great+ let s = List.foldl' (+) 0 [ popCount (indexByteArray ba (div4 off + 1 + jj) :: Word64) | jj <- [0 .. j-1 ] ]+ let word = indexByteArray ba (div4 off + 1 + j) :: Word64+ s + popCount (word .&. complement (unsafeShiftL 0xffffffffffffffff k))
@@ -0,0 +1,441 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Multiple inputs (write buffers, runs) that are being read incrementally.+module Database.LSMTree.Internal.Readers (+ Readers (..)+ , OffsetKey (..)+ , ReaderSource (..)+ , ReadersMergeType (..)+ , new+ , close+ , peekKey+ , HasMore (..)+ , pop+ , dropWhileKey+ -- * Internals+ , Reader (..)+ , ReaderNumber (..)+ , ReadCtx (..)+ ) where++import Control.Monad (zipWithM)+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadMask)+import Control.Monad.Primitive+import Control.RefCount+import Data.Function (on)+import Data.Functor ((<&>))+import Data.List.NonEmpty (nonEmpty)+import qualified Data.Map.Strict as Map+import Data.Maybe (catMaybes)+import Data.Primitive.MutVar+import Data.Traversable (for)+import Database.LSMTree.Internal.BlobRef (BlobSpan, RawBlobRef)+import Database.LSMTree.Internal.Entry (Entry (..))+import qualified Database.LSMTree.Internal.Entry as Entry+import Database.LSMTree.Internal.Index.CompactAcc (SMaybe (..),+ smaybe)+import Database.LSMTree.Internal.Run (Run)+import Database.LSMTree.Internal.RunReader (OffsetKey (..),+ RunReader (..))+import qualified Database.LSMTree.Internal.RunReader as RunReader+import Database.LSMTree.Internal.Serialise+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WB+import qualified KMerge.Heap as Heap+import qualified System.FS.API as FS++-- | A collection of runs and write buffers being read from, yielding elements+-- in order. More precisely, that means first ordered by their key, then by the+-- input run they came from. This is important for resolving multiple entries+-- with the same key into one.+--+-- Construct with 'new', then keep calling 'pop'.+-- If aborting early, remember to call 'close'!+--+-- Creating a 'Readers' does not retain a reference to the input 'Run's or the+-- 'WriteBufferBlobs', but does retain an independent reference on their blob+-- files. It is not necessary to separately retain the 'Run's or the+-- 'WriteBufferBlobs' for correct use of the 'Readers'. There is one important+-- caveat however: to preserve the validity of 'BlobRef's then it is necessary+-- to separately retain a reference to the 'Run' or its 'BlobFile' to preserve+-- the validity of 'BlobRefs'.+--+-- TODO: do this more nicely by changing 'Reader' to preserve the 'BlobFile'+-- ref until it is explicitly closed, and also retain the 'BlobFile' from the+-- WBB and release all of these 'BlobFiles' once the 'Readers' is itself closed.+--+data Readers m h = Readers {+ readersHeap :: !(Heap.MutableHeap (PrimState m) (ReadCtx m h))+ -- | Since there is always one reader outside of the heap, we need to+ -- store it separately. This also contains the next k\/op to yield, unless+ -- all readers are drained, i.e. both:+ -- 1. the reader inside the 'ReadCtx' is empty+ -- 2. the heap is empty+ , readersNext :: !(MutVar (PrimState m) (ReadCtx m h))+ }++newtype ReaderNumber = ReaderNumber Int+ deriving stock (Eq, Ord)++-- | Each heap element needs some more context than just the reader.+-- E.g. the 'Eq' instance we need to be able to access the first key to be read+-- in a pure way.+--+-- TODO(optimisation): We allocate this record for each k/op. This might be+-- avoidable, see ideas below.+data ReadCtx m h = ReadCtx {+ -- We could avoid this using a more specialised mutable heap with separate+ -- arrays for keys and values (or even each of their components).+ -- Using an 'STRef' could avoid reallocating the record for every entry,+ -- but that might not be straightforward to integrate with the heap.+ readCtxHeadKey :: !SerialisedKey+ , readCtxHeadEntry :: !(RunReader.Entry m h)+ -- We could get rid of this by making 'LoserTree' stable (for which there+ -- is a prototype already).+ -- Alternatively, if we decide to have an invariant that the number in+ -- 'RunFsPaths' is always higher for newer runs, then we could use that+ -- in the 'Ord' instance.+ , readCtxNumber :: !ReaderNumber+ , readCtxReader :: !(Reader m h)+ }++instance Eq (ReadCtx m h) where+ (==) = (==) `on` (\r -> (readCtxHeadKey r, readCtxNumber r))++-- | Makes sure we resolve entries in the right order.+instance Ord (ReadCtx m h) where+ compare = compare `on` (\r -> (readCtxHeadKey r, readCtxNumber r))++-- | An individual reader must be able to produce a sequence of pairs of+-- 'SerialisedKey' and 'RunReader.Entry', with ordered und unique keys.+--+-- TODO: This is slightly inelegant. This module could work generally for+-- anything that can produce elements, but currently is very specific to having+-- write buffer and run readers. Also, for run merging, no write buffer is+-- involved, but we still need to branch on this sum type.+-- A more general version is possible, but despite SPECIALISE-ing everything+-- showed ~100 bytes of extra allocations per entry that is read (which might be+-- avoidable with some tinkering).+data Reader m h =+ -- | The list allows to incrementally read from the write buffer without+ -- having to find the next entry in the Map again (requiring key+ -- comparisons) or having to copy out all entries.+ --+ -- TODO: more efficient representation? benchmark!+ ReadBuffer !(MutVar (PrimState m) [KOp m h])+ | ReadRun !(RunReader m h)+ -- | Recursively read from another reader. This requires keeping track of+ -- its 'HasMore' status, since we should not try to read another entry from+ -- it once it is drained.+ --+ -- We represent the recursive reader and 'HasMore' status together as a+ -- 'Maybe' 'Readers'. The reason is subtle: once a 'Readers' becomes drained+ -- it is immediately closed, after which the structure should not be used+ -- anymore or you'd be using resources after they have been closed already.+ --+ -- TODO: maybe it's a slightly more ergonomic alternative to no close the+ -- 'Readers' automatically.+ | ReadReaders !ReadersMergeType !(SMaybe (Readers m h))++type KOp m h = (SerialisedKey, Entry SerialisedValue (RawBlobRef m h))++data ReaderSource m h =+ FromWriteBuffer !WB.WriteBuffer !(Ref (WB.WriteBufferBlobs m h))+ | FromRun !(Ref (Run m h))+ -- | Recursive case, allowing to build a tree of readers for a merging tree.+ | FromReaders !ReadersMergeType ![ReaderSource m h]++{-# SPECIALISE new ::+ ResolveSerialisedValue+ -> OffsetKey+ -> [ReaderSource IO h]+ -> IO (Maybe (Readers IO h)) #-}+new :: forall m h.+ (MonadMask m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> OffsetKey+ -> [ReaderSource m h]+ -> m (Maybe (Readers m h))+new resolve !offsetKey sources = do+ readers <- zipWithM (fromSource . ReaderNumber) [1..] sources+ for (nonEmpty (catMaybes readers)) $ \xs -> do+ (readersHeap, readCtx) <- Heap.newMutableHeap xs+ readersNext <- newMutVar readCtx+ pure Readers {..}+ where+ fromSource :: ReaderNumber -> ReaderSource m h -> m (Maybe (ReadCtx m h))+ fromSource n src =+ case src of+ FromWriteBuffer wb wbblobs -> do+ rs <- fromWB wb wbblobs+ nextReadCtx resolve n rs+ FromRun r -> do+ rs <- ReadRun <$> RunReader.new offsetKey r+ nextReadCtx resolve n rs+ FromReaders mergeType nestedSources -> do+ new resolve offsetKey nestedSources >>= \case+ Nothing -> pure Nothing+ Just rs -> nextReadCtx resolve n (ReadReaders mergeType (SJust rs))++ fromWB :: WB.WriteBuffer -> Ref (WB.WriteBufferBlobs m h) -> m (Reader m h)+ fromWB wb wbblobs = do+ let kops = Map.toList $ filterWB $ WB.toMap wb+ ReadBuffer <$> newMutVar (map convertBlobs kops)+ where+ -- TODO: this conversion involves quite a lot of allocation+ convertBlobs :: (k, Entry v BlobSpan) -> (k, Entry v (RawBlobRef m h))+ convertBlobs = fmap (fmap (WB.mkRawBlobRef wbblobs))++ filterWB = case offsetKey of+ NoOffsetKey -> id+ OffsetKey k -> Map.dropWhileAntitone (< k)++{-# SPECIALISE close :: Readers IO (FS.Handle h) -> IO () #-}+-- | Clean up the resources held by the readers.+--+-- Only call this function when aborting before all readers have been drained!+close ::+ (MonadMask m, MonadSTM m, PrimMonad m)+ => Readers m h+ -> m ()+close Readers {..} = do+ ReadCtx {readCtxReader} <- readMutVar readersNext+ closeReader readCtxReader+ closeHeap+ where+ closeReader = \case+ ReadBuffer _ -> pure ()+ ReadRun r -> RunReader.close r+ ReadReaders _ readersMay -> smaybe (pure ()) close readersMay+ closeHeap =+ Heap.extract readersHeap >>= \case+ Nothing -> pure ()+ Just ReadCtx {readCtxReader} -> do+ closeReader readCtxReader+ closeHeap++{-# SPECIALISE peekKey :: Readers IO h -> IO SerialisedKey #-}+-- | Return the smallest key present in the readers, without consuming any+-- entries.+peekKey ::+ PrimMonad m+ => Readers m h+ -> m SerialisedKey+peekKey Readers {..} = do+ readCtxHeadKey <$> readMutVar readersNext++-- | Once a function returned 'Drained', do not use the 'Readers' any more!+data HasMore = HasMore | Drained+ deriving stock (Eq, Show)++{-# SPECIALISE pop ::+ ResolveSerialisedValue+ -> Readers IO h+ -> IO (SerialisedKey, RunReader.Entry IO h, HasMore) #-}+-- | Remove the entry with the smallest key and return it. If there are multiple+-- entries with that key, it removes the one from the source that came first+-- in list supplied to 'new'. No resolution of multiple entries takes place.+pop ::+ (MonadMask m, MonadSTM m, MonadST m)+ => ResolveSerialisedValue+ -> Readers m h+ -> m (SerialisedKey, RunReader.Entry m h, HasMore)+pop resolve r@Readers {..} = do+ ReadCtx {..} <- readMutVar readersNext+ hasMore <- dropOne resolve r readCtxNumber readCtxReader+ pure (readCtxHeadKey, readCtxHeadEntry, hasMore)++-- TODO: avoid duplication with Merge.TreeMergeType?+data ReadersMergeType = MergeLevel | MergeUnion+ deriving stock (Eq, Show)++{-# SPECIALISE popResolved ::+ ResolveSerialisedValue+ -> ReadersMergeType+ -> Readers IO h+ -> IO (SerialisedKey, RunReader.Entry IO h, HasMore) #-}+-- | Produces an entry with the smallest key, resolving all input entries if+-- there are multiple. Therefore, the next call to 'peekKey' will return a+-- larger key than the one returned here.+--+-- General notes on the code below:+-- * It is quite similar to the one in Internal.Cursor and Internal.Merge. Maybe+-- we can avoid some duplication.+-- * Any function that doesn't take a 'hasMore' argument assumes that the+-- readers have not been drained yet, so we must check before calling them.+-- * There is probably opportunity for optimisations.+--+-- TODO: use this function in Internal.Cursor? Measure performance impact.+popResolved ::+ forall h m.+ (MonadMask m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> ReadersMergeType+ -> Readers m h+ -> m (SerialisedKey, RunReader.Entry m h, HasMore)+popResolved resolve mergeType readers = readEntry+ where+ readEntry :: m (SerialisedKey, RunReader.Entry m h, HasMore)+ readEntry = do+ (key, entry, hasMore) <- pop resolve readers+ case hasMore of+ Drained -> do+ pure (key, entry, Drained)+ HasMore -> do+ case mergeType of+ MergeLevel -> handleLevel key (RunReader.toFullEntry entry)+ MergeUnion -> handleUnion key (RunReader.toFullEntry entry)++ handleUnion :: SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> m (SerialisedKey, RunReader.Entry m h, HasMore)+ handleUnion key entry = do+ nextKey <- peekKey readers+ if nextKey /= key+ then+ -- No more entries for same key, done.+ pure (key, RunReader.Entry entry, HasMore)+ else do+ (_, nextEntry, hasMore) <- pop resolve readers+ let resolved = Entry.combineUnion resolve entry+ (RunReader.toFullEntry nextEntry)+ case hasMore of+ HasMore -> handleUnion key resolved+ Drained -> pure (key, RunReader.Entry resolved, Drained)++ handleLevel :: SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> m (SerialisedKey, RunReader.Entry m h, HasMore)+ handleLevel key entry =+ case entry of+ Upsert v ->+ handleMupdate key v+ _ -> do+ -- Anything but Upsert supersedes all previous entries of+ -- the same key, so we can simply drop them and are done.+ hasMore' <- dropRemaining key+ pure (key, RunReader.Entry entry, hasMore')++ -- Resolve a 'Mupsert' value with the other entries of the same key.+ handleMupdate :: SerialisedKey+ -> SerialisedValue+ -> m (SerialisedKey, RunReader.Entry m h, HasMore)+ handleMupdate key v = do+ nextKey <- peekKey readers+ if nextKey /= key+ then+ -- No more entries for same key, done.+ pure (key, RunReader.Entry (Upsert v), HasMore)+ else do+ (_, nextEntry, hasMore) <- pop resolve readers+ let resolved = Entry.combine resolve (Upsert v)+ (RunReader.toFullEntry nextEntry)+ case hasMore of+ HasMore -> handleLevel key resolved+ Drained -> pure (key, RunReader.Entry resolved, Drained)++ dropRemaining :: SerialisedKey -> m HasMore+ dropRemaining key = do+ (_, hasMore) <- dropWhileKey resolve readers key+ pure hasMore++{-# SPECIALISE dropWhileKey ::+ ResolveSerialisedValue+ -> Readers IO h+ -> SerialisedKey+ -> IO (Int, HasMore) #-}+-- | Drop all entries with a key that is smaller or equal to the supplied one.+dropWhileKey ::+ (MonadMask m, MonadSTM m, MonadST m)+ => ResolveSerialisedValue+ -> Readers m h+ -> SerialisedKey+ -> m (Int, HasMore) -- ^ How many were dropped?+dropWhileKey resolve Readers {..} key = do+ cur <- readMutVar readersNext+ if readCtxHeadKey cur <= key+ then go 0 cur+ else pure (0, HasMore) -- nothing to do+ where+ -- invariant: @readCtxHeadKey <= key@+ go !n ReadCtx {readCtxNumber, readCtxReader} = do+ mNext <- nextReadCtx resolve readCtxNumber readCtxReader >>= \case+ Nothing -> Heap.extract readersHeap+ Just ctx -> Just <$> Heap.replaceRoot readersHeap ctx+ let !n' = n + 1+ case mNext of+ Nothing -> do+ pure (n', Drained)+ Just next -> do+ -- hasMore+ if readCtxHeadKey next <= key+ then+ go n' next+ else do+ writeMutVar readersNext next+ pure (n', HasMore)++{-# SPECIALISE dropOne ::+ ResolveSerialisedValue+ -> Readers IO h+ -> ReaderNumber+ -> Reader IO h+ -> IO HasMore #-}+dropOne ::+ (MonadMask m, MonadSTM m, MonadST m)+ => ResolveSerialisedValue+ -> Readers m h+ -> ReaderNumber+ -> Reader m h+ -> m HasMore+dropOne resolve Readers {..} number reader = do+ mNext <- nextReadCtx resolve number reader >>= \case+ Nothing -> Heap.extract readersHeap+ Just ctx -> Just <$> Heap.replaceRoot readersHeap ctx+ case mNext of+ Nothing ->+ pure Drained+ Just next -> do+ writeMutVar readersNext next+ pure HasMore++{-# SPECIALISE nextReadCtx ::+ ResolveSerialisedValue+ -> ReaderNumber+ -> Reader IO h+ -> IO (Maybe (ReadCtx IO h)) #-}+nextReadCtx ::+ (MonadMask m, MonadSTM m, MonadST m)+ => ResolveSerialisedValue+ -> ReaderNumber+ -> Reader m h+ -> m (Maybe (ReadCtx m h))+nextReadCtx resolve readCtxNumber readCtxReader =+ case readCtxReader of+ ReadBuffer r -> atomicModifyMutVar r $ \case+ [] ->+ ([], Nothing)+ ((readCtxHeadKey, e) : rest) ->+ let readCtxHeadEntry = RunReader.Entry e+ in (rest, Just ReadCtx {..})+ ReadRun r -> RunReader.next r <&> \case+ RunReader.Empty ->+ Nothing+ RunReader.ReadEntry readCtxHeadKey readCtxHeadEntry ->+ Just ReadCtx {..}+ ReadReaders mergeType readersMay -> case readersMay of+ SNothing ->+ pure Nothing+ SJust readers -> do+ (readCtxHeadKey, readCtxHeadEntry, hasMore) <-+ popResolved resolve mergeType readers+ let readersMay' = case hasMore of+ Drained -> SNothing+ HasMore -> SJust readers+ pure $ Just ReadCtx {+ -- TODO: reduce allocations?+ readCtxReader = ReadReaders mergeType readersMay'+ , ..+ }
@@ -0,0 +1,351 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE RecordWildCards #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | Runs of sorted key\/value data.+module Database.LSMTree.Internal.Run (+ -- * Run+ Run (Run, runIndex, runHasFS, runHasBlockIO, runRunDataCaching,+ runBlobFile, runFilter, runKOpsFile)+ , RunFsPaths+ , size+ , sizeInPages+ , runFsPaths+ , runFsPathsNumber+ , runDataCaching+ , runIndexType+ , mkRawBlobRef+ , mkWeakBlobRef+ -- ** Run creation+ , newEmpty+ , fromBuilder+ , fromWriteBuffer+ , RunParams (..)+ -- * Snapshot+ , openFromDisk+ , RunDataCaching (..)+ , IndexType (..)+ ) where++import Control.DeepSeq (NFData (..), rwhnf)+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive+import Control.RefCount+import qualified Data.ByteString.Short as SBS+import Data.Foldable (for_)+import Database.LSMTree.Internal.BlobFile+import Database.LSMTree.Internal.BlobRef hiding (mkRawBlobRef,+ mkWeakBlobRef)+import qualified Database.LSMTree.Internal.BlobRef as BlobRef+import Database.LSMTree.Internal.BloomFilter (Bloom,+ bloomFilterFromFile)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.Entry (NumEntries (..))+import Database.LSMTree.Internal.Index (Index, IndexType (..))+import qualified Database.LSMTree.Internal.Index as Index+import Database.LSMTree.Internal.Page (NumPages)+import Database.LSMTree.Internal.Paths as Paths+import Database.LSMTree.Internal.RunBuilder (RunBuilder,+ RunDataCaching (..), RunParams (..))+import qualified Database.LSMTree.Internal.RunBuilder as Builder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.WriteBuffer (WriteBuffer)+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import Database.LSMTree.Internal.WriteBufferBlobs (WriteBufferBlobs)+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++-- | The in-memory representation of a completed LSM run.+--+data Run m h = Run {+ runNumEntries :: !NumEntries+ -- | The reference count for the LSM run. This counts the+ -- number of references from LSM handles to this run. When+ -- this drops to zero the open files will be closed.+ , runRefCounter :: !(RefCounter m)+ -- | The file system paths for all the files used by the run.+ , runRunFsPaths :: !RunFsPaths+ -- | The bloom filter for the set of keys in this run.+ , runFilter :: !(Bloom SerialisedKey)+ -- | The in-memory index mapping keys to page numbers in the+ -- Key\/Ops file. In future we may support alternative index+ -- representations.+ , runIndex :: !Index+ -- | The file handle for the Key\/Ops file. This file is opened+ -- read-only and is accessed in a page-oriented way, i.e. only+ -- reading whole pages, at page offsets. It will be opened with+ -- 'O_DIRECT' on supported platforms.+ , runKOpsFile :: !(FS.Handle h)+ -- | The file handle for the BLOBs file. This file is opened+ -- read-only and is accessed in a normal style using buffered+ -- I\/O, reading arbitrary file offset and length spans.+ , runBlobFile :: !(Ref (BlobFile m h))+ , runRunDataCaching :: !RunDataCaching+ , runHasFS :: !(HasFS m h)+ , runHasBlockIO :: !(HasBlockIO m h)+ }++-- | Shows only the 'runRunFsPaths' field.+instance Show (Run m h) where+ showsPrec _ run = showString "Run { runRunFsPaths = " . showsPrec 0 (runRunFsPaths run) . showString " }"++instance NFData h => NFData (Run m h) where+ rnf (Run a b c d e f g h i j) =+ rnf a `seq` rwhnf b `seq` rnf c `seq` rnf d `seq` rnf e `seq`+ rnf f `seq` rnf g `seq` rnf h `seq` rwhnf i `seq` rwhnf j++instance RefCounted m (Run m h) where+ getRefCounter = runRefCounter++size :: Ref (Run m h) -> NumEntries+size (DeRef run) = runNumEntries run++sizeInPages :: Ref (Run m h) -> NumPages+sizeInPages (DeRef run) = Index.sizeInPages (runIndex run)++runFsPaths :: Ref (Run m h) -> RunFsPaths+runFsPaths (DeRef r) = runRunFsPaths r++runFsPathsNumber :: Ref (Run m h) -> RunNumber+runFsPathsNumber = Paths.runNumber . runFsPaths++-- | See 'openFromDisk'+runIndexType :: Ref (Run m h) -> IndexType+runIndexType (DeRef r) = Index.indexToIndexType (runIndex r)++-- | See 'openFromDisk'+runDataCaching :: Ref (Run m h) -> RunDataCaching+runDataCaching (DeRef r) = runRunDataCaching r+++-- | Helper function to make a 'WeakBlobRef' that points into a 'Run'.+mkRawBlobRef :: Run m h -> BlobSpan -> RawBlobRef m h+mkRawBlobRef Run{runBlobFile} blobspan =+ BlobRef.mkRawBlobRef runBlobFile blobspan++-- | Helper function to make a 'WeakBlobRef' that points into a 'Run'.+mkWeakBlobRef :: Ref (Run m h) -> BlobSpan -> WeakBlobRef m h+mkWeakBlobRef (DeRef Run{runBlobFile}) blobspan =+ BlobRef.mkWeakBlobRef runBlobFile blobspan++{-# SPECIALISE finaliser ::+ HasFS IO h+ -> FS.Handle h+ -> Ref (BlobFile IO h)+ -> RunFsPaths+ -> IO () #-}+-- | Close the files used in the run and remove them from disk. After calling+-- this operation, the run must not be used anymore.+--+-- TODO: exception safety+finaliser ::+ (MonadSTM m, MonadMask m, PrimMonad m)+ => HasFS m h+ -> FS.Handle h+ -> Ref (BlobFile m h)+ -> RunFsPaths+ -> m ()+finaliser hfs kopsFile blobFile fsPaths = do+ FS.hClose hfs kopsFile+ releaseRef blobFile+ FS.removeFile hfs (runKOpsPath fsPaths)+ FS.removeFile hfs (runFilterPath fsPaths)+ FS.removeFile hfs (runIndexPath fsPaths)+ FS.removeFile hfs (runChecksumsPath fsPaths)++{-# SPECIALISE setRunDataCaching ::+ HasBlockIO IO h+ -> FS.Handle h+ -> RunDataCaching+ -> IO () #-}+setRunDataCaching ::+ MonadSTM m+ => HasBlockIO m h+ -> FS.Handle h+ -> RunDataCaching+ -> m ()+setRunDataCaching hbio runKOpsFile CacheRunData = do+ -- disable file readahead (only applies to this file descriptor)+ FS.hAdviseAll hbio runKOpsFile FS.AdviceRandom+ -- use the page cache for disk I/O reads+ FS.hSetNoCache hbio runKOpsFile False+setRunDataCaching hbio runKOpsFile NoCacheRunData = do+ -- do not use the page cache for disk I/O reads+ FS.hSetNoCache hbio runKOpsFile True++{-# SPECIALISE newEmpty ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> IO (Ref (Run IO h)) #-}+-- | This function should be run with asynchronous exceptions masked to prevent+-- failing after internal resources have already been created.+newEmpty ::+ (MonadST m, MonadSTM m, MonadMask m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> m (Ref (Run m h))+newEmpty hfs hbio salt runParams runPaths = do+ builder <- Builder.new hfs hbio salt runParams runPaths (NumEntries 0)+ fromBuilder builder++{-# SPECIALISE fromBuilder ::+ RunBuilder IO h+ -> IO (Ref (Run IO h)) #-}+-- TODO: make exception safe+fromBuilder ::+ (MonadST m, MonadSTM m, MonadMask m)+ => RunBuilder m h+ -> m (Ref (Run m h))+fromBuilder builder = do+ (runHasFS, runHasBlockIO,+ runRunFsPaths, runFilter, runIndex,+ RunParams {runParamCaching = runRunDataCaching}, runNumEntries) <-+ Builder.unsafeFinalise builder+ runKOpsFile <- FS.hOpen runHasFS (runKOpsPath runRunFsPaths) FS.ReadMode+ -- TODO: openBlobFile should be called with exceptions masked+ runBlobFile <- openBlobFile runHasFS (runBlobPath runRunFsPaths) FS.ReadMode+ setRunDataCaching runHasBlockIO runKOpsFile runRunDataCaching+ newRef (finaliser runHasFS runKOpsFile runBlobFile runRunFsPaths)+ (\runRefCounter -> Run { .. })++{-# SPECIALISE fromWriteBuffer ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> WriteBuffer+ -> Ref (WriteBufferBlobs IO h)+ -> IO (Ref (Run IO h)) #-}+-- | Write a write buffer to disk, including the blobs it contains.+--+-- This creates a new 'Run' which must eventually be released with 'releaseRef'.+--+-- TODO: As a possible optimisation, blobs could be written to a blob file+-- immediately when they are added to the write buffer, avoiding the need to do+-- it here.+--+-- This function should be run with asynchronous exceptions masked to prevent+-- failing after internal resources have already been created.+fromWriteBuffer ::+ (MonadST m, MonadSTM m, MonadMask m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> WriteBuffer+ -> Ref (WriteBufferBlobs m h)+ -> m (Ref (Run m h))+fromWriteBuffer fs hbio salt params fsPaths buffer blobs = do+ builder <- Builder.new fs hbio salt params fsPaths (WB.numEntries buffer)+ for_ (WB.toList buffer) $ \(k, e) ->+ Builder.addKeyOp builder k (fmap (WBB.mkRawBlobRef blobs) e)+ --TODO: the fmap entry here reallocates even when there are no blobs+ fromBuilder builder++{-------------------------------------------------------------------------------+ Snapshot+-------------------------------------------------------------------------------}++{-# SPECIALISE openFromDisk ::+ HasFS IO h+ -> HasBlockIO IO h+ -> RunDataCaching+ -> IndexType+ -> Bloom.Salt+ -> RunFsPaths+ -> IO (Ref (Run IO h)) #-}+-- | Load a previously written run from disk, checking each file's checksum+-- against the checksum file.+--+-- This creates a new 'Run' which must eventually be released with 'releaseRef'.+--+-- Exceptions will be raised when any of the file's contents don't match their+-- checksum ('ChecksumError') or can't be parsed ('FileFormatError').+--+-- The 'RunDataCaching' and 'IndexType' parameters need to be saved and+-- restored separately because these are not stored in the on-disk+-- representation. Use 'runDataCaching' and 'runIndexType' to obtain these+-- parameters from the open run before persisting to disk.+--+-- TODO: it may make more sense to persist these parameters with the run's+-- on-disk representation.+--+openFromDisk ::+ forall m h.+ (MonadSTM m, MonadMask m, PrimMonad m)+ => HasFS m h+ -> HasBlockIO m h+ -> RunDataCaching+ -> IndexType+ -> Bloom.Salt -- ^ Expected salt+ -> RunFsPaths+ -> m (Ref (Run m h))+-- TODO: make exception safe+openFromDisk fs hbio runRunDataCaching indexType expectedSalt runRunFsPaths = do+ expectedChecksums <-+ CRC.expectValidFile fs (runChecksumsPath runRunFsPaths) CRC.FormatChecksumsFile+ . fromChecksumsFile+ =<< CRC.readChecksumsFile fs (runChecksumsPath runRunFsPaths)++ -- verify checksums of files we don't read yet+ let paths = pathsForRunFiles runRunFsPaths+ checkCRC runRunDataCaching (forRunKOpsRaw expectedChecksums) (forRunKOpsRaw paths)+ checkCRC runRunDataCaching (forRunBlobRaw expectedChecksums) (forRunBlobRaw paths)++ -- read and try parsing files+ let filterPath = forRunFilterRaw paths+ checkCRC CacheRunData (forRunFilterRaw expectedChecksums) filterPath+ runFilter <- FS.withFile fs filterPath FS.ReadMode $+ bloomFilterFromFile fs expectedSalt++ (runNumEntries, runIndex) <-+ CRC.expectValidFile fs (forRunIndexRaw paths) CRC.FormatIndexFile+ . Index.fromSBS indexType+ =<< readCRC (forRunIndexRaw expectedChecksums) (forRunIndexRaw paths)++ runKOpsFile <- FS.hOpen fs (runKOpsPath runRunFsPaths) FS.ReadMode+ -- TODO: openBlobFile should be called with exceptions masked+ runBlobFile <- openBlobFile fs (runBlobPath runRunFsPaths) FS.ReadMode+ setRunDataCaching hbio runKOpsFile runRunDataCaching+ newRef (finaliser fs runKOpsFile runBlobFile runRunFsPaths) $ \runRefCounter ->+ Run {+ runHasFS = fs+ , runHasBlockIO = hbio+ , ..+ }+ where+ -- Note: all file data for this path is evicted from the page cache /if/ the+ -- caching argument is 'NoCacheRunData'.+ checkCRC :: RunDataCaching -> CRC.CRC32C -> FS.FsPath -> m ()+ checkCRC cache expected fp =+ CRC.checkCRC fs hbio (cache == NoCacheRunData) expected fp++ -- Note: all file data for this path is evicted from the page cache+ readCRC :: CRC.CRC32C -> FS.FsPath -> m SBS.ShortByteString+ readCRC expected fp = FS.withFile fs fp FS.ReadMode $ \h -> do+ n <- FS.hGetSize fs h+ -- double the file readahead window (only applies to this file descriptor)+ FS.hAdviseAll hbio h FS.AdviceSequential+ (sbs, !checksum) <- CRC.hGetExactlyCRC32C_SBS fs h (fromIntegral n) CRC.initialCRC32C+ -- drop the file from the OS page cache+ FS.hAdviseAll hbio h FS.AdviceDontNeed+ CRC.expectChecksum fs fp expected checksum+ pure sbs
@@ -0,0 +1,355 @@+{-# LANGUAGE RecordWildCards #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | Incremental (in-memory portion of) run construction+--+module Database.LSMTree.Internal.RunAcc (+ RunAcc (..)+ , new+ , unsafeFinalise+ -- * Adding key\/op pairs+ -- | There are a few variants of actions to add key\/op pairs to the run+ -- accumulator. Which one to use depends on a couple questions:+ --+ -- * Is it fully in memory or is it pre-serialised and only partly in+ -- memory?+ -- * Is the key\/op pair known to be \"small\" or \"large\"?+ --+ -- If it's in memory but it's not known whether it's small or large then+ -- use 'addKeyOp'. One can use 'entryWouldFitInPage' to find out if it's+ -- small or large. If it's in memory and known to be small or large then+ -- use 'addSmallKeyOp' or 'addLargeKeyOp' as appropriate. If it's large+ -- and pre-serialised, use 'addLargeSerialisedKeyOp' but note its+ -- constraints carefully.+ --+ , addKeyOp+ , addSmallKeyOp+ , addLargeKeyOp+ , addLargeSerialisedKeyOp+ , PageAcc.entryWouldFitInPage+ -- * Bloom filter allocation+ , RunBloomFilterAlloc (..)+ -- ** Exposed for testing+ , newMBloom+ ) where++import Control.DeepSeq (NFData (..))+import Control.Exception (assert)+import Control.Monad.ST.Strict+import qualified Data.BloomFilter.Blocked as Bloom+import Data.Primitive.PrimVar (PrimVar, modifyPrimVar, newPrimVar,+ readPrimVar)+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.BloomFilter (Bloom, MBloom)+import Database.LSMTree.Internal.Chunk (Chunk)+import Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))+import Database.LSMTree.Internal.Index (Index, IndexAcc, IndexType)+import qualified Database.LSMTree.Internal.Index as Index (appendMulti,+ appendSingle, newWithDefaults, unsafeEnd)+import Database.LSMTree.Internal.PageAcc (PageAcc)+import qualified Database.LSMTree.Internal.PageAcc as PageAcc+import qualified Database.LSMTree.Internal.PageAcc1 as PageAcc+import Database.LSMTree.Internal.RawOverflowPage+import Database.LSMTree.Internal.RawPage (RawPage)+import qualified Database.LSMTree.Internal.RawPage as RawPage+import Database.LSMTree.Internal.Serialise (SerialisedKey,+ SerialisedValue)++{-------------------------------------------------------------------------------+ Incremental, in-memory run construction+-------------------------------------------------------------------------------}++-- | The run accumulator is a mutable structure that accumulates key\/op pairs.+-- It yields pages and chunks of the index incrementally, and returns the+-- Bloom filter and complete index at the end.+--+-- Use 'new' to start run construction, add new key\/operation pairs to the run+-- by using 'addKeyOp' and co, and complete run construction using+-- 'unsafeFinalise'.+data RunAcc s = RunAcc {+ mbloom :: !(MBloom s SerialisedKey)+ , mindex :: !(IndexAcc s)+ , mpageacc :: !(PageAcc s)+ , entryCount :: !(PrimVar s Int)+ }++-- | @'new' nentries@ starts an incremental run construction.+--+-- @nentries@ should be an upper bound on the expected number of entries in the+-- output run.+new ::+ NumEntries+ -> RunBloomFilterAlloc+ -> Bloom.Salt+ -> IndexType+ -> ST s (RunAcc s)+new nentries alloc salt indexType = do+ mbloom <- newMBloom nentries alloc salt+ mindex <- Index.newWithDefaults indexType+ mpageacc <- PageAcc.newPageAcc+ entryCount <- newPrimVar 0+ pure RunAcc{..}++-- | Finalise an incremental run construction. Do /not/ use a 'RunAcc' after+-- finalising it.+--+-- The frozen bloom filter and compact index will be returned, along with the+-- final page of the run (if necessary), and the remaining chunks of the+-- incrementally constructed compact index.+unsafeFinalise ::+ RunAcc s+ -> ST s ( Maybe RawPage+ , Maybe Chunk+ , Bloom SerialisedKey+ , Index+ , NumEntries+ )+unsafeFinalise racc@RunAcc {..} = do+ mpagemchunk <- flushPageIfNonEmpty racc+ (mchunk', index) <- Index.unsafeEnd mindex+ bloom <- Bloom.unsafeFreeze mbloom+ numEntries <- readPrimVar entryCount+ let !mpage = fst <$> mpagemchunk+ !mchunk = selectChunk mpagemchunk mchunk'+ pure (mpage, mchunk, bloom, index, NumEntries numEntries)+ where+ selectChunk :: Maybe (RawPage, Maybe Chunk)+ -> Maybe Chunk+ -> Maybe Chunk+ selectChunk (Just (_page, Just _chunk)) (Just _chunk') =+ -- If flushing the page accumulator gives us an index chunk then+ -- the index can't have any more chunks when we finalise the index.+ error "unsafeFinalise: impossible double final chunk"+ selectChunk (Just (_page, Just chunk)) _ = Just chunk+ selectChunk _ (Just chunk) = Just chunk+ selectChunk _ _ = Nothing++-- | Add a key\/op pair with an optional blob span to the run accumulator.+--+-- Note that this version expects the full value to be in the given'Entry', not+-- just a prefix of the value that fits into a single page.+--+-- If the key\/op pair is known to be \"small\" or \"large\" then you can use+-- the special versions 'addSmallKeyOp' or 'addLargeKeyOp'. If it is+-- pre-serialised, use 'addLargeSerialisedKeyOp'.+--+addKeyOp ::+ RunAcc s+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan -- ^ the full value, not just a prefix+ -> ST s ([RawPage], [RawOverflowPage], [Chunk])+addKeyOp racc k e+ | PageAcc.entryWouldFitInPage k e = smallToLarge <$> addSmallKeyOp racc k e+ | otherwise = addLargeKeyOp racc k e+ where+ smallToLarge :: Maybe (RawPage, Maybe Chunk)+ -> ([RawPage], [RawOverflowPage], [Chunk])+ smallToLarge Nothing = ([], [], [])+ smallToLarge (Just (page, Nothing)) = ([page], [], [])+ smallToLarge (Just (page, Just chunk)) = ([page], [], [chunk])++-- | Add a \"small\" key\/op pair with an optional blob span to the run+-- accumulator.+--+-- This version is /only/ for small entries that can fit within a single page.+-- Use 'addLargeKeyOp' if the entry is bigger than a page. If this distinction+-- is not known at the use site, use 'PageAcc.entryWouldFitInPage' to determine+-- which case applies, or use 'addKeyOp'.+--+-- This is guaranteed to add the key\/op, and it may yield (at most one) page.+--+addSmallKeyOp ::+ RunAcc s+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan+ -> ST s (Maybe (RawPage, Maybe Chunk))+addSmallKeyOp racc@RunAcc{entryCount , mpageacc} k e =+ assert (PageAcc.entryWouldFitInPage k e) $ do+ modifyPrimVar entryCount (+1)+ -- We do not use Bloom.insert here. We accumulate keys (in the mpageacc)+ -- and add them all to the mbloom in a batch in flushPageIfNonEmpty.++ pageBoundaryNeeded <-+ -- Try adding the key/op to the page accumulator to see if it fits. If+ -- it does not fit, a page boundary is needed.+ not <$> PageAcc.pageAccAddElem mpageacc k e++ if pageBoundaryNeeded+ then do+ -- We need a page boundary. If the current page is empty then we have+ -- a boundary already, otherwise we need to flush the current page.+ mpagemchunk <- flushPageIfNonEmpty racc+ -- The current page is now empty, either because it was already empty+ -- or because we just flushed it. Adding the new key/op to an empty+ -- page must now succeed, because we know it fits in a page.+ added <- PageAcc.pageAccAddElem mpageacc k e+ assert added $ pure mpagemchunk++ else pure Nothing++-- | Add a \"large\" key\/op pair with an optional blob span to the run+-- accumulator.+--+-- This version is /only/ for large entries that span multiple pages. Use+-- 'addSmallKeyOp' if the entry is smaller than a page. If this distinction+-- is not known at the use site, use 'PageAcc.entryWouldFitInPage' to determine+-- which case applies.+--+-- Note that this version expects the full large value to be in the given+-- 'Entry', not just the prefix of the value that fits into a single page.+-- For large multi-page values that are represented by a pre-serialised+-- 'RawPage' (as occurs when merging runs), use 'addLargeSerialisedKeyOp'.+--+-- This is guaranteed to add the key\/op. It will yield one or two 'RawPage's,+-- and one or more 'RawOverflowPage's. These pages should be written out to+-- the run's page file in that order, the 'RawPage's followed by the+-- 'RawOverflowPage's.+--+addLargeKeyOp ::+ RunAcc s+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan -- ^ the full value, not just a prefix+ -> ST s ([RawPage], [RawOverflowPage], [Chunk])+addLargeKeyOp racc@RunAcc{entryCount, mindex, mbloom} k e =+ assert (not (PageAcc.entryWouldFitInPage k e)) $ do+ modifyPrimVar entryCount (+1)+ -- Large key/op pairs don't get added to the mpageacc, and thus bypass the+ -- bulk bloom insert in flushPageIfNonEmpty, so we add them directly here.+ Bloom.insert mbloom k++ -- If the existing page accumulator is non-empty, we flush it, since the+ -- new large key/op will need more than one page to itself.+ mpagemchunkPre <- flushPageIfNonEmpty racc++ -- Make the new page and overflow pages. Add the span of pages to the index.+ let (page, overflowPages) = PageAcc.singletonPage k e+ chunks <- Index.appendMulti (k, fromIntegral (length overflowPages)) mindex++ -- Combine the results with anything we flushed before+ let (!pages, !chunks') = selectPagesAndChunks mpagemchunkPre page chunks+ pure (pages, overflowPages, chunks')++-- | Add a \"large\" pre-serialised key\/op entry to the run accumulator.+--+-- This version is for large entries that span multiple pages and are+-- represented by already serialised 'RawPage' and one or more+-- 'RawOverflowPage's.+--+-- For this case, the caller provides the key, the raw page it is from and the+-- overflow pages. The raw page and overflow pages are returned along with any+-- other pages that need to be yielded (in order). The caller should write out+-- the pages to the run's page file in order: the returned 'RawPage's followed+-- by the 'RawOverflowPage's (the same as for 'addLargeKeyOp').+--+-- Note that this action is not appropriate for key\/op entries that would fit+-- within a page ('PageAcc.entryWouldFitInPage') but just /happen/ to have+-- ended up in a page on their own in an input to a merge. A page can end up+-- with a single entry because a page boundary was needed rather than because+-- the entry itself was too big. Furthermore, pre-serialised pages can only be+-- used unaltered if the entry does /not/ use a 'BlobSpan', since the 'BlobSpan'+-- typically needs to be modified. Thus the caller should use the following+-- tests to decide if 'addLargeSerialisedKeyOp' should be used:+--+-- 1. The entry does not use a 'BlobSpan'.+-- 2. The entry definitely overflows onto one or more overflow pages.+--+-- Otherwise, use 'addLargeKeyOp' or 'addSmallKeyOp' as appropriate.+--+addLargeSerialisedKeyOp ::+ RunAcc s+ -> SerialisedKey -- ^ The key+ -> RawPage -- ^ The page that this key\/op is in, which must be the+ -- first page of a multi-page representation of a single+ -- key\/op /without/ a 'BlobSpan'.+ -> [RawOverflowPage] -- ^ The overflow pages for this key\/op+ -> ST s ([RawPage], [RawOverflowPage], [Chunk])+addLargeSerialisedKeyOp racc@RunAcc{entryCount, mindex, mbloom}+ k page overflowPages =+ assert (RawPage.rawPageNumKeys page == 1) $+ assert (RawPage.rawPageHasBlobSpanAt page 0 == 0) $+ assert (RawPage.rawPageOverflowPages page > 0) $+ assert (RawPage.rawPageOverflowPages page == length overflowPages) $ do+ modifyPrimVar entryCount (+1)+ -- Large key/op pairs don't get added to the mpageacc, and thus bypass the+ -- bulk bloom insert in flushPageIfNonEmpty, so we add them directly here.+ Bloom.insert mbloom k++ -- If the existing page accumulator is non-empty, we flush it, since the+ -- new large key/op will need more than one page to itself.+ mpagemchunkPre <- flushPageIfNonEmpty racc+ let nOverflowPages = length overflowPages --TODO: consider using vector+ chunks <- Index.appendMulti (k, fromIntegral nOverflowPages) mindex+ let (!pages, !chunks') = selectPagesAndChunks mpagemchunkPre page chunks+ pure (pages, overflowPages, chunks')++-- | Internal helper: finalise the current page, add the page to the index,+-- reset the page accumulator and return the serialised 'RawPage' along with+-- any index chunk.+--+-- Returns @Nothing@ if the page accumulator was empty.+--+flushPageIfNonEmpty :: RunAcc s -> ST s (Maybe (RawPage, Maybe Chunk))+flushPageIfNonEmpty RunAcc{mpageacc, mindex, mbloom} = do+ nkeys <- PageAcc.keysCountPageAcc mpageacc+ if nkeys > 0+ then do+ bloomInserts mbloom mpageacc nkeys++ -- Grab the min and max keys, and add the page to the index.+ minKey <- PageAcc.indexKeyPageAcc mpageacc 0+ maxKey <- PageAcc.indexKeyPageAcc mpageacc (nkeys-1)+ mchunk <- Index.appendSingle (minKey, maxKey) mindex++ -- Now serialise the page and reset the accumulator+ page <- PageAcc.serialisePageAcc mpageacc+ PageAcc.resetPageAcc mpageacc+ pure (Just (page, mchunk))++ else pure Nothing++-- An instance of insertMany specialised to SerialisedKey and indexKeyPageAcc.+-- This is a performance-sensitive function. It is marked NOINLINE so we can+-- easily inspect the core and check all the specialisations worked as expected.+{-# NOINLINE bloomInserts #-}+bloomInserts :: MBloom s SerialisedKey -> PageAcc s -> Int -> ST s ()+bloomInserts !mbloom !mpageacc !nkeys =+ Bloom.insertMany mbloom (PageAcc.indexKeyPageAcc mpageacc) nkeys++-- | Internal helper for 'addLargeKeyOp' and 'addLargeSerialisedKeyOp'.+-- Combine the result of 'flushPageIfNonEmpty' with extra pages and index+-- chunks.+--+selectPagesAndChunks :: Maybe (RawPage, Maybe Chunk)+ -> RawPage+ -> [Chunk]+ -> ([RawPage], [Chunk])+selectPagesAndChunks mpagemchunkPre page chunks =+ case mpagemchunkPre of+ Nothing -> ( [page], chunks)+ Just (pagePre, Nothing) -> ([pagePre, page], chunks)+ Just (pagePre, Just chunkPre) -> ([pagePre, page], chunkPre:chunks)++{-------------------------------------------------------------------------------+ Bloom filter allocation+-------------------------------------------------------------------------------}++-- | See 'Database.LSMTree.Internal.Config.BloomFilterAlloc'+data RunBloomFilterAlloc =+ -- | Bits per element in a filter+ RunAllocFixed !Double+ | RunAllocRequestFPR !Double+ deriving stock (Show, Eq)++instance NFData RunBloomFilterAlloc where+ rnf (RunAllocFixed a) = rnf a+ rnf (RunAllocRequestFPR a) = rnf a++newMBloom :: NumEntries -> RunBloomFilterAlloc -> Bloom.Salt -> ST s (MBloom s a)+newMBloom (NumEntries nentries) alloc salt =+ Bloom.new (Bloom.sizeForPolicy (policy alloc) nentries) salt+ where+ --TODO: it'd be possible to turn the RunBloomFilterAlloc into a BloomPolicy+ -- without the NumEntries, and cache the policy, avoiding recalculating the+ -- policy every time.+ policy (RunAllocFixed bitsPerEntry) = Bloom.policyForBits bitsPerEntry+ policy (RunAllocRequestFPR fpr) = Bloom.policyForFPR fpr
@@ -0,0 +1,250 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | A mutable run ('RunBuilder') that is under construction.+--+module Database.LSMTree.Internal.RunBuilder (+ RunBuilder (..)+ , RunParams (..)+ , RunDataCaching (..)+ , RunBloomFilterAlloc (..)+ , IndexType (..)+ , new+ , addKeyOp+ , addLargeSerialisedKeyOp+ , unsafeFinalise+ , close+ ) where++import Control.DeepSeq (NFData (..))+import Control.Monad (when)+import Control.Monad.Class.MonadST (MonadST (..))+import qualified Control.Monad.Class.MonadST as ST+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadThrow)+import Control.Monad.Primitive+import Data.Foldable (for_, traverse_)+import Data.Primitive.PrimVar+import Data.Word (Word64)+import Database.LSMTree.Internal.BlobRef (RawBlobRef)+import Database.LSMTree.Internal.BloomFilter (Bloom)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.ChecksumHandle+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Index (Index, IndexType (..))+import Database.LSMTree.Internal.Paths+import Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)+import Database.LSMTree.Internal.RawPage (RawPage)+import Database.LSMTree.Internal.RunAcc (RunAcc,+ RunBloomFilterAlloc (..))+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import Database.LSMTree.Internal.Serialise+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++-- | The in-memory representation of an LSM run that is under construction.+-- (The \"M\" stands for mutable.) This is the output sink for two key+-- algorithms: 1. writing out the write buffer, and 2. incrementally merging+-- two or more runs.+--+-- It contains open file handles for all four files used in the disk+-- representation of a run. Each file handle is opened write-only and should be+-- written to using normal buffered I\/O.+--+-- __Not suitable for concurrent construction from multiple threads!__+--+data RunBuilder m h = RunBuilder {+ runBuilderParams :: !RunParams++ -- | The file system paths for all the files used by the run.+ , runBuilderFsPaths :: !RunFsPaths++ -- | The run accumulator. This is the representation used for the+ -- morally pure subset of the run cnstruction functionality. In+ -- particular it contains the (mutable) index, bloom filter and buffered+ -- pending output for the key\/ops file.+ , runBuilderAcc :: !(RunAcc (PrimState m))++ -- | The byte offset within the blob file for the next blob to be written.+ , runBuilderBlobOffset :: !(PrimVar (PrimState m) Word64)++ -- | The (write mode) file handles.+ , runBuilderHandles :: {-# UNPACK #-} !(ForRunFiles (ChecksumHandle (PrimState m) h))+ , runBuilderHasFS :: !(HasFS m h)+ , runBuilderHasBlockIO :: !(HasBlockIO m h)+ }++data RunParams = RunParams {+ runParamCaching :: !RunDataCaching,+ runParamAlloc :: !RunBloomFilterAlloc,+ runParamIndex :: !IndexType+ }+ deriving stock (Eq, Show)++instance NFData RunParams where+ rnf (RunParams a b c) = rnf a `seq` rnf b `seq` rnf c++-- | Should this run cache key\/ops data in memory?+data RunDataCaching = CacheRunData | NoCacheRunData+ deriving stock (Show, Eq)++instance NFData RunDataCaching where+ rnf CacheRunData = ()+ rnf NoCacheRunData = ()++{-# SPECIALISE new ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> NumEntries+ -> IO (RunBuilder IO h) #-}+-- | Create an 'RunBuilder' to start building a run.+--+-- NOTE: 'new' assumes that 'runDir' that the run is created in exists.+new ::+ (MonadST m, MonadSTM m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> NumEntries -- ^ an upper bound of the number of entries to be added+ -> m (RunBuilder m h)+new hfs hbio salt runBuilderParams@RunParams{..} runBuilderFsPaths numEntries = do+ runBuilderAcc <- ST.stToIO $+ RunAcc.new numEntries runParamAlloc salt runParamIndex+ runBuilderBlobOffset <- newPrimVar 0++ runBuilderHandles <- traverse (makeHandle hfs) (pathsForRunFiles runBuilderFsPaths)++ let builder = RunBuilder { runBuilderHasFS = hfs, runBuilderHasBlockIO = hbio, .. }+ writeIndexHeader hfs (forRunIndex runBuilderHandles) runParamIndex+ pure builder++{-# SPECIALISE addKeyOp ::+ RunBuilder IO h+ -> SerialisedKey+ -> Entry SerialisedValue (RawBlobRef IO h)+ -> IO () #-}+-- | Add a key\/op pair.+--+-- In the 'InsertWithBlob' case, the 'RawBlobRef' identifies where the blob can be+-- found (which is either from a write buffer or another run). The blobs will+-- be copied from their existing blob file into the new run's blob file.+--+-- Use only for entries that are fully in-memory (other than any blob).+-- To handle larger-than-page values in a chunked style during run merging,+-- use 'addLargeSerialisedKeyOp'.+--+-- The k\/ops and the primary array of the index get written incrementally,+-- everything else only at the end when 'unsafeFinalise' is called.+--+addKeyOp ::+ (MonadST m, MonadSTM m, MonadThrow m)+ => RunBuilder m h+ -> SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> m ()+addKeyOp RunBuilder{..} key op = do+ -- TODO: the fmap entry here reallocates even when there are no blobs.+ -- We need the Entry _ BlobSpan for RunAcc.add{Small,Large}KeyOp+ -- Perhaps pass the optional blob span separately from the Entry.+ op' <- traverse (copyBlob runBuilderHasFS runBuilderBlobOffset (forRunBlob runBuilderHandles)) op+ if RunAcc.entryWouldFitInPage key op'+ then do+ mpagemchunk <- ST.stToIO $ RunAcc.addSmallKeyOp runBuilderAcc key op'+ case mpagemchunk of+ Nothing -> pure ()+ Just (page, mchunk) -> do+ writeRawPage runBuilderHasFS (forRunKOps runBuilderHandles) page+ for_ mchunk $ writeIndexChunk runBuilderHasFS (forRunIndex runBuilderHandles)++ else do+ (pages, overflowPages, chunks)+ <- ST.stToIO $ RunAcc.addLargeKeyOp runBuilderAcc key op'+ --TODO: consider optimisation: use writev to write all pages in one go+ for_ pages $ writeRawPage runBuilderHasFS (forRunKOps runBuilderHandles)+ writeRawOverflowPages runBuilderHasFS (forRunKOps runBuilderHandles) overflowPages+ for_ chunks $ writeIndexChunk runBuilderHasFS (forRunIndex runBuilderHandles)++{-# SPECIALISE addLargeSerialisedKeyOp ::+ RunBuilder IO h+ -> SerialisedKey+ -> RawPage+ -> [RawOverflowPage]+ -> IO () #-}+-- | See 'RunAcc.addLargeSerialisedKeyOp' for details.+--+addLargeSerialisedKeyOp ::+ (MonadST m, MonadSTM m)+ => RunBuilder m h+ -> SerialisedKey+ -> RawPage+ -> [RawOverflowPage]+ -> m ()+addLargeSerialisedKeyOp RunBuilder{..} key page overflowPages = do+ (pages, overflowPages', chunks)+ <- ST.stToIO $+ RunAcc.addLargeSerialisedKeyOp runBuilderAcc key page overflowPages+ for_ pages $ writeRawPage runBuilderHasFS (forRunKOps runBuilderHandles)+ writeRawOverflowPages runBuilderHasFS (forRunKOps runBuilderHandles) overflowPages'+ for_ chunks $ writeIndexChunk runBuilderHasFS (forRunIndex runBuilderHandles)++{-# SPECIALISE unsafeFinalise ::+ RunBuilder IO h+ -> IO (HasFS IO h, HasBlockIO IO h,+ RunFsPaths, Bloom SerialisedKey, Index,+ RunParams, NumEntries) #-}+-- | Finish construction of the run.+-- Writes the filter and index to file and leaves all written files on disk.+--+-- __Do not use the 'RunBuilder' after calling this function!__+--+-- TODO: Ensure proper cleanup even in presence of exceptions.+unsafeFinalise ::+ (MonadST m, MonadSTM m, MonadThrow m)+ => RunBuilder m h+ -> m (HasFS m h, HasBlockIO m h,+ RunFsPaths, Bloom SerialisedKey, Index,+ RunParams, NumEntries)+-- TODO: consider introducing a type for this big tuple+unsafeFinalise RunBuilder {..} = do+ -- write final bits+ (mPage, mChunk, runFilter, runIndex, numEntries) <-+ ST.stToIO (RunAcc.unsafeFinalise runBuilderAcc)+ for_ mPage $ writeRawPage runBuilderHasFS (forRunKOps runBuilderHandles)+ for_ mChunk $ writeIndexChunk runBuilderHasFS (forRunIndex runBuilderHandles)+ writeIndexFinal runBuilderHasFS (forRunIndex runBuilderHandles) numEntries runIndex+ writeFilter runBuilderHasFS (forRunFilter runBuilderHandles) runFilter+ -- write checksums+ checksums <- toChecksumsFile <$> traverse readChecksum runBuilderHandles+ FS.withFile runBuilderHasFS (runChecksumsPath runBuilderFsPaths) (FS.WriteMode FS.MustBeNew) $ \h -> do+ CRC.writeChecksumsFile' runBuilderHasFS h checksums+ -- always drop the checksum file from the cache+ FS.hDropCacheAll runBuilderHasBlockIO h+ -- always drop filter and index files from the cache+ dropCache runBuilderHasBlockIO (forRunFilterRaw runBuilderHandles)+ dropCache runBuilderHasBlockIO (forRunIndexRaw runBuilderHandles)+ -- drop the KOps and blobs files from the cache if asked for+ when (runParamCaching runBuilderParams == NoCacheRunData) $ do+ dropCache runBuilderHasBlockIO (forRunKOpsRaw runBuilderHandles)+ dropCache runBuilderHasBlockIO (forRunBlobRaw runBuilderHandles)+ mapM_ (closeHandle runBuilderHasFS) runBuilderHandles+ pure (runBuilderHasFS, runBuilderHasBlockIO,+ runBuilderFsPaths, runFilter, runIndex,+ runBuilderParams, numEntries)++{-# SPECIALISE close :: RunBuilder IO h -> IO () #-}+-- | Close a run that is being constructed (has not been finalised yet),+-- removing all files associated with it from disk.+-- After calling this operation, the run must not be used anymore.+--+-- TODO: Ensure proper cleanup even in presence of exceptions.+close :: MonadSTM m => RunBuilder m h -> m ()+close RunBuilder {..} = do+ traverse_ (closeHandle runBuilderHasFS) runBuilderHandles+ traverse_ (FS.removeFile runBuilderHasFS) (pathsForRunFiles runBuilderFsPaths)
@@ -0,0 +1,21 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.RunNumber (+ RunNumber (..),+ TableId (..),+ CursorId (..),+) where++import Control.DeepSeq (NFData)++newtype RunNumber = RunNumber Int+ deriving stock (Eq, Ord, Show)+ deriving newtype (NFData)++newtype TableId = TableId Int+ deriving stock (Eq, Ord, Show)+ deriving newtype (NFData)++newtype CursorId = CursorId Int+ deriving stock (Eq, Ord, Show)+ deriving newtype (NFData)
@@ -0,0 +1,345 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | A run that is being read incrementally.+--+module Database.LSMTree.Internal.RunReader (+ RunReader (..)+ , OffsetKey (..)+ , new+ , next+ , close+ , Result (..)+ , Entry (..)+ , toFullEntry+ , appendOverflow+ -- * Exported for WriteBufferReader+ , mkEntryOverflow+ , readDiskPage+ , readOverflowPages+ ) where++import Control.Exception (assert)+import Control.Monad (guard, when)+import Control.Monad.Class.MonadST (MonadST (..))+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadCatch (..),+ MonadMask (..), MonadThrow (..))+import Control.Monad.Primitive (PrimMonad (..))+import Control.RefCount+import Data.Bifunctor (first)+import Data.Maybe (isNothing)+import Data.Primitive.ByteArray (newPinnedByteArray,+ unsafeFreezeByteArray)+import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar,+ writeMutVar)+import Data.Primitive.PrimVar+import Data.Word (Word16, Word32)+import Database.LSMTree.Internal.BitMath (ceilDivPageSize,+ mulPageSize, roundUpToPageSize)+import Database.LSMTree.Internal.BlobFile as BlobFile+import Database.LSMTree.Internal.BlobRef as BlobRef+import qualified Database.LSMTree.Internal.Entry as E+import qualified Database.LSMTree.Internal.Index as Index (search)+import Database.LSMTree.Internal.Page (PageNo (..), PageSpan (..),+ getNumPages, nextPageNo)+import Database.LSMTree.Internal.Paths+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage,+ pinnedByteArrayToOverflowPages, rawOverflowPageRawBytes)+import Database.LSMTree.Internal.RawPage+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.Serialise+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++-- | Allows reading the k\/ops of a run incrementally, using its own read-only+-- file handle and in-memory cache of the current disk page.+--+-- Creating a 'RunReader' does not retain a reference to the 'Run', but does+-- retain an independent reference on the run's blob file. It is not necessary+-- to separately retain the 'Run' for correct use of the 'RunReader'. There is+-- one important caveat however: the 'RunReader' maintains the validity of+-- 'BlobRef's only up until the point where the reader is drained (or+-- explicitly closed). In particular this means 'BlobRefs' can be invalidated+-- as soon as the 'next' returns 'Empty'. If this is not sufficient then it is+-- necessary to separately retain a reference to the 'Run' or its 'BlobFile' to+-- preserve the validity of 'BlobRefs'.+--+-- New pages are loaded when trying to read their first entry.+--+-- TODO(optimise): Reuse page buffers using some kind of allocator. However,+-- deciding how long a page needs to stay around is not trivial.+data RunReader m h = RunReader {+ -- | The disk page currently being read. If it is 'Nothing', the reader+ -- is considered closed.+ readerCurrentPage :: !(MutVar (PrimState m) (Maybe RawPage))+ -- | The index of the entry to be returned by the next call to 'next'.+ , readerCurrentEntryNo :: !(PrimVar (PrimState m) Word16)+ -- | Read mode file handle into the run's k\/ops file. We rely on it to+ -- track the position of the next disk page to read, instead of keeping+ -- a counter ourselves. Also, the run's handle is supposed to be opened+ -- with @O_DIRECT@, which is counterproductive here.+ , readerKOpsHandle :: !(FS.Handle h)+ -- | The blob file from the run this reader is reading from.+ , readerBlobFile :: !(Ref (BlobFile m h))+ , readerRunDataCaching :: !Run.RunDataCaching+ , readerHasFS :: !(HasFS m h)+ , readerHasBlockIO :: !(HasBlockIO m h)+ }++data OffsetKey = NoOffsetKey | OffsetKey !SerialisedKey+ deriving stock Show++{-# SPECIALISE new ::+ OffsetKey+ -> Ref (Run.Run IO h)+ -> IO (RunReader IO h) #-}+new :: forall m h.+ (MonadMask m, MonadSTM m, PrimMonad m)+ => OffsetKey+ -> Ref (Run.Run m h)+ -> m (RunReader m h)+new !offsetKey+ readerRun@(DeRef Run.Run {+ runBlobFile,+ runRunDataCaching = readerRunDataCaching,+ runHasFS = readerHasFS,+ runHasBlockIO = readerHasBlockIO,+ runIndex = index+ }) = do+ (readerKOpsHandle :: FS.Handle h) <-+ FS.hOpen readerHasFS (runKOpsPath (Run.runFsPaths readerRun)) FS.ReadMode >>= \h -> do+ fileSize <- FS.hGetSize readerHasFS h+ let fileSizeInPages = fileSize `div` toEnum pageSize+ let indexedPages = getNumPages $ Run.sizeInPages readerRun+ assert (indexedPages == fileSizeInPages) $ pure h+ -- Advise the OS that this file is being read sequentially, which will+ -- double the readahead window in response (only for this file descriptor)+ FS.hAdviseAll readerHasBlockIO readerKOpsHandle FS.AdviceSequential++ (page, entryNo) <- seekFirstEntry readerKOpsHandle++ readerBlobFile <- dupRef runBlobFile+ readerCurrentEntryNo <- newPrimVar entryNo+ readerCurrentPage <- newMutVar page+ let reader = RunReader {..}++ when (isNothing page) $+ close reader+ pure reader+ where+ seekFirstEntry readerKOpsHandle =+ case offsetKey of+ NoOffsetKey -> do+ -- Load first page from disk, if it exists.+ firstPage <- readDiskPage readerHasFS readerKOpsHandle+ pure (firstPage, 0)+ OffsetKey offset -> do+ -- Use the index to find the page number for the key (if it exists).+ let PageSpan pageNo pageEnd = Index.search offset index+ seekToDiskPage readerHasFS pageNo readerKOpsHandle+ readDiskPage readerHasFS readerKOpsHandle >>= \case+ Nothing ->+ pure (Nothing, 0)+ Just foundPage -> do+ case rawPageFindKey foundPage offset of+ Just n ->+ -- Found an appropriate index within the index's page.+ pure (Just foundPage, n)++ _ -> do+ -- The index said that the key, if it were to exist, would+ -- live on pageNo, but then rawPageFindKey tells us that in+ -- fact there is no key greater than or equal to the given+ -- offset on this page.+ -- This tells us that the key does not exist, but that if it+ -- were to exist, it would be between the last key in this+ -- page and the first key in the next page.+ -- Thus the reader should be initialised to return keys+ -- starting from the next (non-overflow) page.+ seekToDiskPage readerHasFS (nextPageNo pageEnd) readerKOpsHandle+ nextPage <- readDiskPage readerHasFS readerKOpsHandle+ pure (nextPage, 0)++{-# SPECIALISE close ::+ RunReader IO h+ -> IO () #-}+-- | This function should be called when discarding a 'RunReader' before it+-- was done (i.e. returned 'Empty'). This avoids leaking file handles.+-- Once it has been called, do not use the reader any more!+close ::+ (MonadSTM m, MonadMask m, PrimMonad m)+ => RunReader m h+ -> m ()+close RunReader{..} = do+ when (readerRunDataCaching == Run.NoCacheRunData) $+ -- drop the file from the OS page cache+ FS.hDropCacheAll readerHasBlockIO readerKOpsHandle+ FS.hClose readerHasFS readerKOpsHandle+ releaseRef readerBlobFile+ --TODO: arguably we should have distinct finish and close and require that+ -- readers are _always_ closed, even after they have been drained.+ -- This would allow BlobRefs to remain valid until the reader is closed.+ -- Currently they are invalidated as soon as the cursor is drained.++-- | The 'SerialisedKey' and 'SerialisedValue' point into the in-memory disk+-- page. Keeping them alive will also prevent garbage collection of the 4k byte+-- array, so if they're long-lived, consider making a copy!+data Result m h+ = Empty+ | ReadEntry !SerialisedKey !(Entry m h)++data Entry m h =+ Entry+ !(E.Entry SerialisedValue (RawBlobRef m h))+ | -- | A large entry. The caller might be interested in various different+ -- (redundant) representation, so we return all of them.+ EntryOverflow+ -- | The value is just a prefix, with the remainder in the overflow pages.+ !(E.Entry SerialisedValue (RawBlobRef m h))+ -- | A page containing the single entry (or rather its prefix).+ !RawPage+ -- | Non-zero length of the overflow in bytes.+ !Word32+ -- | The overflow pages containing the suffix of the value (so at least+ -- the number of bytes specified above).+ --+ -- TODO(optimise): Sometimes, reading the overflow pages is not necessary.+ -- We could just return the page index and offer a separate function to do+ -- the disk I/O once needed.+ ![RawOverflowPage]++mkEntryOverflow ::+ E.Entry SerialisedValue (RawBlobRef m h)+ -> RawPage+ -> Word32+ -> [RawOverflowPage]+ -> Entry m h+mkEntryOverflow entryPrefix page len overflowPages =+ assert (len > 0) $+ assert (rawPageOverflowPages page == ceilDivPageSize (fromIntegral len)) $+ assert (rawPageOverflowPages page == length overflowPages) $+ EntryOverflow entryPrefix page len overflowPages++{-# INLINE toFullEntry #-}+toFullEntry :: Entry m h -> E.Entry SerialisedValue (RawBlobRef m h)+toFullEntry = \case+ Entry e ->+ e+ EntryOverflow prefix _ len overflowPages ->+ first (appendOverflow len overflowPages) prefix++{-# INLINE appendOverflow #-}+appendOverflow :: Word32 -> [RawOverflowPage] -> SerialisedValue -> SerialisedValue+appendOverflow len overflowPages (SerialisedValue prefix) =+ SerialisedValue $+ RB.take (RB.size prefix + fromIntegral len) $+ mconcat (prefix : map rawOverflowPageRawBytes overflowPages)++{-# SPECIALISE next ::+ RunReader IO h+ -> IO (Result IO h) #-}+-- | Stop using the 'RunReader' after getting 'Empty', because the 'Reader' is+-- automatically closed!+next :: forall m h.+ (MonadMask m, MonadSTM m, MonadST m)+ => RunReader m h+ -> m (Result m h)+next reader@RunReader {..} = do+ readMutVar readerCurrentPage >>= \case+ Nothing ->+ pure Empty+ Just page -> do+ entryNo <- readPrimVar readerCurrentEntryNo+ go entryNo page+ where+ go :: Word16 -> RawPage -> m (Result m h)+ go !entryNo !page =+ -- take entry from current page (resolve blob if necessary)+ case rawPageIndex page entryNo of+ IndexNotPresent -> do+ -- if it is past the last one, load a new page from disk, try again+ newPage <- readDiskPage readerHasFS readerKOpsHandle+ stToIO $ writeMutVar readerCurrentPage newPage+ case newPage of+ Nothing -> do+ close reader+ pure Empty+ Just p -> do+ writePrimVar readerCurrentEntryNo 0+ go 0 p -- try again on the new page+ IndexEntry key entry -> do+ modifyPrimVar readerCurrentEntryNo (+1)+ let entry' = fmap (BlobRef.mkRawBlobRef readerBlobFile) entry+ let rawEntry = Entry entry'+ pure (ReadEntry key rawEntry)+ IndexEntryOverflow key entry lenSuffix -> do+ -- TODO: we know that we need the next page, could already load?+ modifyPrimVar readerCurrentEntryNo (+1)+ let entry' :: E.Entry SerialisedValue (RawBlobRef m h)+ entry' = fmap (BlobRef.mkRawBlobRef readerBlobFile) entry+ overflowPages <- readOverflowPages readerHasFS readerKOpsHandle lenSuffix+ let rawEntry = mkEntryOverflow entry' page lenSuffix overflowPages+ pure (ReadEntry key rawEntry)++{-------------------------------------------------------------------------------+ Utilities+-------------------------------------------------------------------------------}++seekToDiskPage :: HasFS m h -> PageNo -> FS.Handle h -> m ()+seekToDiskPage fs pageNo h = do+ FS.hSeek fs h FS.AbsoluteSeek (pageNoToByteOffset pageNo)+ where+ pageNoToByteOffset (PageNo n) =+ assert (n >= 0) $+ mulPageSize (fromIntegral n)++{-# SPECIALISE readDiskPage ::+ HasFS IO h+ -> FS.Handle h+ -> IO (Maybe RawPage) #-}+-- | Returns 'Nothing' on EOF.+readDiskPage ::+ (MonadCatch m, PrimMonad m)+ => HasFS m h+ -> FS.Handle h+ -> m (Maybe RawPage)+readDiskPage fs h = do+ mba <- newPinnedByteArray pageSize+ -- TODO: make sure no other exception type can be thrown+ --+ -- TODO: if FS.FsReachEOF is thrown as an injected disk fault, then we+ -- incorrectly deduce that the file has no more contents. We should probably+ -- use an explicit file pointer instead in the style of 'FilePointer'.+ handleJust (guard . FS.isFsErrorType FS.FsReachedEOF) (\_ -> pure Nothing) $ do+ bytesRead <- FS.hGetBufExactly fs h mba 0 (fromIntegral pageSize)+ assert (fromIntegral bytesRead == pageSize) $ pure ()+ ba <- unsafeFreezeByteArray mba+ let !rawPage = unsafeMakeRawPage ba 0+ pure (Just rawPage)++pageSize :: Int+pageSize = 4096++{-# SPECIALISE readOverflowPages ::+ HasFS IO h+ -> FS.Handle h+ -> Word32+ -> IO [RawOverflowPage] #-}+-- | Throws exception on EOF. If a suffix was expected, the file should have it.+-- Reads full pages, despite the suffix only using part of the last page.+readOverflowPages ::+ (MonadSTM m, MonadThrow m, PrimMonad m)+ => HasFS m h+ -> FS.Handle h+ -> Word32+ -> m [RawOverflowPage]+readOverflowPages fs h len = do+ let lenPages = fromIntegral (roundUpToPageSize len) -- always read whole pages+ mba <- newPinnedByteArray lenPages+ _ <- FS.hGetBufExactly fs h mba 0 (fromIntegral lenPages)+ ba <- unsafeFreezeByteArray mba+ -- should not copy since 'ba' is pinned and its length is a multiple of 4k.+ pure $ pinnedByteArrayToOverflowPages 0 lenPages ba
@@ -0,0 +1,184 @@+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE PatternSynonyms #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | Newtype wrappers and utilities for serialised keys, values and blobs.+--+module Database.LSMTree.Internal.Serialise (+ -- * Re-exports+ SerialiseKey+ , SerialiseValue+ -- * Keys+ , SerialisedKey (SerialisedKey, SerialisedKey')+ , serialiseKey+ , deserialiseKey+ , sizeofKey+ , sizeofKey16+ , sizeofKey32+ , sizeofKey64+ , serialisedKey+ , keyTopBits64+ -- * Values+ , SerialisedValue (SerialisedValue, SerialisedValue')+ , serialiseValue+ , deserialiseValue+ , sizeofValue+ , sizeofValue16+ , sizeofValue32+ , sizeofValue64+ , serialisedValue+ , ResolveSerialisedValue+ -- * Blobs+ , SerialisedBlob (SerialisedBlob, SerialisedBlob')+ , serialiseBlob+ , deserialiseBlob+ , sizeofBlob+ , sizeofBlob64+ , serialisedBlob+ ) where++import Control.DeepSeq (NFData)+import Data.BloomFilter.Hash (Hashable (..))+import qualified Data.ByteString.Builder as BB+import qualified Data.Vector.Primitive as VP+import Data.Word+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise.Class (SerialiseKey,+ SerialiseValue)+import qualified Database.LSMTree.Internal.Serialise.Class as Class++{-------------------------------------------------------------------------------+ Keys+-------------------------------------------------------------------------------}++-- | Representation of a serialised key.+--+-- Serialisation should preserve equality and ordering. The 'Ord' instance for+-- 'SerialisedKey' uses lexicographical ordering.+newtype SerialisedKey = SerialisedKey RawBytes+ deriving stock Show+ deriving newtype (Eq, Ord, Hashable, NFData)++{-# COMPLETE SerialisedKey' #-}+pattern SerialisedKey' :: VP.Vector Word8 -> SerialisedKey+pattern SerialisedKey' pvec = SerialisedKey (RawBytes pvec)++{-# INLINE serialiseKey #-}+serialiseKey :: SerialiseKey k => k -> SerialisedKey+serialiseKey k = SerialisedKey (Class.serialiseKey k)++{-# INLINE deserialiseKey #-}+deserialiseKey :: SerialiseKey k => SerialisedKey -> k+deserialiseKey (SerialisedKey bytes) = Class.deserialiseKey bytes++{-# INLINE sizeofKey #-}+-- | Size of key in number of bytes.+sizeofKey :: SerialisedKey -> Int+sizeofKey (SerialisedKey rb) = RB.size rb++{-# INLINE sizeofKey16 #-}+-- | Size of key in number of bytes.+sizeofKey16 :: SerialisedKey -> Word16+sizeofKey16 = fromIntegral . sizeofKey++{-# INLINE sizeofKey32 #-}+-- | Size of key in number of bytes.+sizeofKey32 :: SerialisedKey -> Word32+sizeofKey32 = fromIntegral . sizeofKey++{-# INLINE sizeofKey64 #-}+-- | Size of key in number of bytes.+sizeofKey64 :: SerialisedKey -> Word64+sizeofKey64 = fromIntegral . sizeofKey++{-# INLINE serialisedKey #-}+serialisedKey :: SerialisedKey -> BB.Builder+serialisedKey (SerialisedKey rb) = RB.builder rb++{-# INLINE keyTopBits64 #-}+-- | See 'RB.topBits64'+keyTopBits64 :: SerialisedKey -> Word64+keyTopBits64 (SerialisedKey rb) = RB.topBits64 rb++{-------------------------------------------------------------------------------+ Values+-------------------------------------------------------------------------------}++-- | Representation of a serialised value.+newtype SerialisedValue = SerialisedValue RawBytes+ deriving stock Show+ deriving newtype (Eq, Ord, NFData)++{-# COMPLETE SerialisedValue' #-}+pattern SerialisedValue' :: VP.Vector Word8 -> SerialisedValue+pattern SerialisedValue' pvec = (SerialisedValue (RawBytes pvec))++{-# INLINE serialiseValue #-}+serialiseValue :: SerialiseValue v => v -> SerialisedValue+serialiseValue v = SerialisedValue (Class.serialiseValue v)++{-# INLINE deserialiseValue #-}+deserialiseValue :: SerialiseValue v => SerialisedValue -> v+deserialiseValue (SerialisedValue bytes) = Class.deserialiseValue bytes++{-# INLINE sizeofValue #-}+sizeofValue :: SerialisedValue -> Int+sizeofValue (SerialisedValue rb) = RB.size rb++{-# INLINE sizeofValue16 #-}+-- | Size of value in number of bytes.+sizeofValue16 :: SerialisedValue -> Word16+sizeofValue16 = fromIntegral . sizeofValue++{-# INLINE sizeofValue32 #-}+-- | Size of value in number of bytes.+sizeofValue32 :: SerialisedValue -> Word32+sizeofValue32 = fromIntegral . sizeofValue++{-# INLINE sizeofValue64 #-}+-- | Size of value in number of bytes.+sizeofValue64 :: SerialisedValue -> Word64+sizeofValue64 = fromIntegral . sizeofValue++{-# LANGUAGE serialisedValue #-}+serialisedValue :: SerialisedValue -> BB.Builder+serialisedValue (SerialisedValue rb) = RB.builder rb++type ResolveSerialisedValue =+ SerialisedValue -> SerialisedValue -> SerialisedValue++{-------------------------------------------------------------------------------+ Blobs+-------------------------------------------------------------------------------}++-- | Representation of a serialised blob.+newtype SerialisedBlob = SerialisedBlob RawBytes+ deriving stock Show+ deriving newtype (Eq, Ord, NFData)++{-# COMPLETE SerialisedBlob' #-}+pattern SerialisedBlob' :: VP.Vector Word8 -> SerialisedBlob+pattern SerialisedBlob' pvec = (SerialisedBlob (RawBytes pvec))++{-# INLINE serialiseBlob #-}+serialiseBlob :: SerialiseValue v => v -> SerialisedBlob+serialiseBlob v = SerialisedBlob (Class.serialiseValue v)++{-# INLINE deserialiseBlob #-}+deserialiseBlob :: SerialiseValue v => SerialisedBlob -> v+deserialiseBlob (SerialisedBlob bytes) = Class.deserialiseValue bytes++{-# INLINE sizeofBlob #-}+-- | Size of blob in number of bytes.+sizeofBlob :: SerialisedBlob -> Int+sizeofBlob (SerialisedBlob rb) = RB.size rb++{-# INLINE sizeofBlob64 #-}+sizeofBlob64 :: SerialisedBlob -> Word64+sizeofBlob64 = fromIntegral . sizeofBlob++{-# INLINE serialisedBlob #-}+serialisedBlob :: SerialisedBlob -> BB.Builder+serialisedBlob (SerialisedBlob rb) = RB.builder rb
@@ -0,0 +1,536 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Public API for serialisation of keys, blobs and values+--+module Database.LSMTree.Internal.Serialise.Class (+ -- * SerialiseKey+ SerialiseKey (..)+ , serialiseKeyIdentity+ , serialiseKeyIdentityUpToSlicing+ , SerialiseKeyOrderPreserving+ , serialiseKeyPreservesOrdering+ -- * SerialiseValue+ , SerialiseValue (..)+ , serialiseValueIdentity+ , serialiseValueIdentityUpToSlicing+ -- * RawBytes+ , RawBytes (..)+ , packSlice+ -- * Errors+ , requireBytesExactly+ ) where++import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as B+import qualified Data.ByteString.Lazy as LBS+import qualified Data.ByteString.Short.Internal as SBS+import qualified Data.ByteString.UTF8 as UTF8+import Data.Int (Int16, Int32, Int64, Int8)+import Data.Monoid (Sum (..))+import qualified Data.Primitive as P+import qualified Data.Vector.Primitive as VP+import Data.Void (Void, absurd)+import Data.Word (Word16, Word32, Word64, Word8)+import Database.LSMTree.Internal.ByteString (byteArrayToSBS)+import Database.LSMTree.Internal.Primitive+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Vector+import Numeric (showInt)++{-------------------------------------------------------------------------------+ SerialiseKey+-------------------------------------------------------------------------------}++{- | Serialisation of keys.++Instances should satisfy the following laws:++[Identity]+ @'deserialiseKey' ('serialiseKey' x) == x@+[Identity up to slicing]+ @'deserialiseKey' ('packSlice' prefix ('serialiseKey' x) suffix) == x@+-}+class SerialiseKey k where+ serialiseKey :: k -> RawBytes+ -- TODO: 'deserialiseKey' is only strictly necessary for range queries.+ -- It might make sense to move it to a separate class, which could also+ -- require total deserialisation (potentially using 'Either').+ deserialiseKey :: RawBytes -> k++-- | Test the __Identity__ law for the 'SerialiseKey' class+serialiseKeyIdentity :: (Eq k, SerialiseKey k) => k -> Bool+serialiseKeyIdentity x = deserialiseKey (serialiseKey x) == x++-- | Test the __Identity up to slicing__ law for the 'SerialiseKey' class+serialiseKeyIdentityUpToSlicing ::+ (Eq k, SerialiseKey k)+ => RawBytes -> k -> RawBytes -> Bool+serialiseKeyIdentityUpToSlicing prefix x suffix =+ deserialiseKey (packSlice prefix (serialiseKey x) suffix) == x++{- |+Order-preserving serialisation of keys.++Table data is sorted by /serialised/ keys.+Range lookups and cursors return entries in this order.+If serialisation does not preserve the ordering of /unserialised/ keys,+then range lookups and cursors return entries out of order.++If the 'SerialiseKey' instance for a type preserves the ordering,+then it can safely be given an instance of 'SerialiseKeyOrderPreserving'.+These should satisfy the following law:++[Order-preserving]+ @x \`'compare'\` y == 'serialiseKey' x \`'compare'\` 'serialiseKey' y@++Serialised keys are lexicographically ordered.+To satisfy the __Order-preserving__ law, keys should be serialised into a big-endian format.+-}+class SerialiseKey k => SerialiseKeyOrderPreserving k where++-- | Test the __Order-preserving__ law for the 'SerialiseKeyOrderPreserving' class+serialiseKeyPreservesOrdering :: (Ord k, SerialiseKey k) => k -> k -> Bool+serialiseKeyPreservesOrdering x y = x `compare` y == serialiseKey x `compare` serialiseKey y++{-------------------------------------------------------------------------------+ SerialiseValue+-------------------------------------------------------------------------------}++{- | Serialisation of values and blobs.++Instances should satisfy the following laws:++[Identity]+ @'deserialiseValue' ('serialiseValue' x) == x@++[Identity up to slicing]+ @'deserialiseValue' ('packSlice' prefix ('serialiseValue' x) suffix) == x@+-}+class SerialiseValue v where+ serialiseValue :: v -> RawBytes+ deserialiseValue :: RawBytes -> v++-- | Test the __Identity__ law for the 'SerialiseValue' class+serialiseValueIdentity :: (Eq v, SerialiseValue v) => v -> Bool+serialiseValueIdentity x = deserialiseValue (serialiseValue x) == x++-- | Test the __Identity up to slicing__ law for the 'SerialiseValue' class+serialiseValueIdentityUpToSlicing ::+ (Eq v, SerialiseValue v)+ => RawBytes -> v -> RawBytes -> Bool+serialiseValueIdentityUpToSlicing prefix x suffix =+ deserialiseValue (packSlice prefix (serialiseValue x) suffix) == x++{-------------------------------------------------------------------------------+ RawBytes+-------------------------------------------------------------------------------}++-- | @'packSlice' prefix x suffix@ makes @x@ into a slice with @prefix@ bytes on+-- the left and @suffix@ bytes on the right.+packSlice :: RawBytes -> RawBytes -> RawBytes -> RawBytes+packSlice prefix x suffix =+ RB.take (RB.size x) (RB.drop (RB.size prefix) (prefix <> x <> suffix))++{-------------------------------------------------------------------------------+ Errors+-------------------------------------------------------------------------------}++-- | @'requireBytesExactly' tyName expected actual x@+requireBytesExactly :: String -> Int -> Int -> a -> a+requireBytesExactly tyName expected actual x+ | expected == actual = x+ | otherwise =+ error+ $ showString "deserialise "+ . showString tyName+ . showString ": expected "+ . showInt expected+ . showString " bytes, but got "+ . showInt actual+ $ ""++{-------------------------------------------------------------------------------+ Int+-------------------------------------------------------------------------------}++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Int8 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int8" 1 len $ indexInt8Array ba off++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Int8 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int8" 1 len $ indexInt8Array ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Int16 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapInt16 x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int16" 2 len $ byteSwapInt16 (indexWord8ArrayAsInt16 ba off)++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Int16 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int16" 2 len $ indexWord8ArrayAsInt16 ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Int32 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapInt32 x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int32" 4 len $ byteSwapInt32 (indexWord8ArrayAsInt32 ba off)++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Int32 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int32" 4 len $ indexWord8ArrayAsInt32 ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Int64 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapInt64 x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int64" 8 len $ byteSwapInt64 (indexWord8ArrayAsInt64 ba off)++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Int64 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int64" 8 len $ indexWord8ArrayAsInt64 ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Int where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapInt x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int" 8 len $ byteSwapInt (indexWord8ArrayAsInt ba off)++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Int where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Int" 8 len $ indexWord8ArrayAsInt ba off++{-------------------------------------------------------------------------------+ Word+-------------------------------------------------------------------------------}++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Word8 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word8" 1 len (indexWord8Array ba off)++instance SerialiseKeyOrderPreserving Word8++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Word8 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word8" 1 len $ indexWord8Array ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Word16 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapWord16 x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word16" 2 len $ byteSwapWord16 (indexWord8ArrayAsWord16 ba off)++instance SerialiseKeyOrderPreserving Word16++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Word16 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word16" 2 len $ indexWord8ArrayAsWord16 ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Word32 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapWord32 x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word32" 4 len $ byteSwapWord32 (indexWord8ArrayAsWord32 ba off)++instance SerialiseKeyOrderPreserving Word32++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Word32 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word32" 4 len $ indexWord8ArrayAsWord32 ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Word64 where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapWord64 x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word64" 8 len $ byteSwapWord64 (indexWord8ArrayAsWord64 ba off)++instance SerialiseKeyOrderPreserving Word64++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Word64 where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word64" 8 len $ indexWord8ArrayAsWord64 ba off++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(1)\).+-}+instance SerialiseKey Word where+ serialiseKey x = RB.RawBytes $ byteVectorFromPrim $ byteSwapWord x++ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word" 8 len $ byteSwapWord (indexWord8ArrayAsWord ba off)++instance SerialiseKeyOrderPreserving Word++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(1)\).+-}+instance SerialiseValue Word where+ serialiseValue x = RB.RawBytes $ byteVectorFromPrim $ x++ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word" 8 len $ indexWord8ArrayAsWord ba off++{-------------------------------------------------------------------------------+ String+-------------------------------------------------------------------------------}++{- |+@'serialiseKey'@: \(O(n)\).++@'deserialiseKey'@: \(O(n)\).++The 'String' is (de)serialised as UTF-8.+-}+instance SerialiseKey String where+ -- TODO: Optimise. The performance is \(O(n) + O(n)\) but it could be \(O(n)\).+ serialiseKey = serialiseKey . UTF8.fromString+ deserialiseKey = UTF8.toString . deserialiseKey++instance SerialiseKeyOrderPreserving String++{- |+@'serialiseKey'@: \(O(n)\).++@'deserialiseKey'@: \(O(n)\).++The 'String' is (de)serialised as UTF-8.+-}+instance SerialiseValue String where+ -- TODO: Optimise. The performance is \(O(n) + O(n)\) but it could be \(O(n)\).+ serialiseValue = serialiseValue . UTF8.fromString+ deserialiseValue = UTF8.toString . deserialiseValue++{-------------------------------------------------------------------------------+ ByteString+-------------------------------------------------------------------------------}++{- |+@'serialiseKey'@: \(O(n)\).++@'deserialiseKey'@: \(O(n)\).+-}+instance SerialiseKey LBS.ByteString where+ -- TODO: Optimise. The performance is \(O(n) + O(n)\) but it could be \(O(n)\).+ serialiseKey = serialiseKey . LBS.toStrict+ deserialiseKey = B.toLazyByteString . RB.builder++instance SerialiseKeyOrderPreserving LBS.ByteString++{- |+@'serialiseValue'@: \(O(n)\).++@'deserialiseValue'@: \(O(n)\).+-}+instance SerialiseValue LBS.ByteString where+ -- TODO: Optimise. The performance is \(O(n) + O(n)\) but it could be \(O(n)\).+ serialiseValue = serialiseValue . LBS.toStrict+ deserialiseValue = B.toLazyByteString . RB.builder++{- |+@'serialiseKey'@: \(O(n)\).++@'deserialiseKey'@: \(O(n)\).+-}+instance SerialiseKey BS.ByteString where+ serialiseKey = serialiseKey . SBS.toShort+ deserialiseKey = SBS.fromShort . deserialiseKey++instance SerialiseKeyOrderPreserving BS.ByteString++{- |+@'serialiseValue'@: \(O(n)\).++@'deserialiseValue'@: \(O(n)\).+-}+instance SerialiseValue BS.ByteString where+ serialiseValue = serialiseValue . SBS.toShort+ deserialiseValue = SBS.fromShort . deserialiseValue++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(n)\).+-}+instance SerialiseKey SBS.ShortByteString where+ serialiseKey = RB.fromShortByteString+ deserialiseKey = byteArrayToSBS . RB.force++instance SerialiseKeyOrderPreserving SBS.ShortByteString++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(n)\).+-}+instance SerialiseValue SBS.ShortByteString where+ serialiseValue = RB.fromShortByteString+ deserialiseValue = byteArrayToSBS . RB.force++{-------------------------------------------------------------------------------+ ByteArray+-------------------------------------------------------------------------------}++{- |+@'serialiseKey'@: \(O(1)\).++@'deserialiseKey'@: \(O(n)\).+-}+instance SerialiseKey P.ByteArray where+ serialiseKey ba = RB.fromByteArray 0 (P.sizeofByteArray ba) ba+ deserialiseKey = RB.force++{- |+@'serialiseValue'@: \(O(1)\).++@'deserialiseValue'@: \(O(n)\).+-}+instance SerialiseValue P.ByteArray where+ serialiseValue ba = RB.fromByteArray 0 (P.sizeofByteArray ba) ba+ deserialiseValue = RB.force++{-------------------------------------------------------------------------------+ Void+-------------------------------------------------------------------------------}++-- | The implementation of 'deserialiseKey' throws an exception.+instance SerialiseKey Void where+ serialiseKey = absurd+ deserialiseKey = error "deserialiseKey: cannot deserialise into Void"+++-- | The implementation of 'deserialiseValue' throws an exception.+instance SerialiseValue Void where+ serialiseValue = absurd+ deserialiseValue = error "deserialiseValue: cannot deserialise into Void"++{-------------------------------------------------------------------------------+ Sum+-------------------------------------------------------------------------------}++{- |+An instance for 'Sum' which is transparent to the serialisation of the value type.++__NOTE:__ If you want to serialise @'Sum' a@ differently from @a@, you must use another newtype wrapper.+-}+instance SerialiseValue a => SerialiseValue (Sum a) where+ serialiseValue (Sum v) = serialiseValue v++ deserialiseValue = Sum . deserialiseValue
@@ -0,0 +1,809 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Snapshot (+ -- * Snapshot metadata+ SnapshotLabel (..)+ , SnapshotMetaData (..)+ -- * Levels snapshot format+ , SnapLevels (..)+ , SnapLevel (..)+ , SnapIncomingRun (..)+ , SnapMergingRun (..)+ -- * MergeTree snapshot format+ , SnapMergingTree(..)+ , SnapMergingTreeState(..)+ , SnapPendingMerge(..)+ , SnapPreExistingRun(..)+ -- * Conversion to levels snapshot format+ , toSnapLevels+ -- * Conversion to merging tree snapshot format+ , toSnapMergingTree+ -- * Write buffer+ , snapshotWriteBuffer+ , openWriteBuffer+ -- * Run+ , SnapshotRun (..)+ , snapshotRun+ , openRun+ -- * Opening snapshot formats+ -- ** Levels format+ , fromSnapLevels+ -- ** Merging Tree format+ , fromSnapMergingTree+ -- * Hard links+ , hardLinkRunFiles+ ) where++import Control.ActionRegistry+import Control.Concurrent.Class.MonadMVar.Strict+import Control.Concurrent.Class.MonadSTM (MonadSTM)+import Control.DeepSeq (NFData (..))+import Control.Monad (void)+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadThrow (MonadMask, bracket,+ bracketOnError)+import Control.Monad.Primitive (PrimMonad)+import Control.RefCount+import Data.Foldable (sequenceA_, traverse_)+import Data.String (IsString)+import Data.Text (Text)+import qualified Data.Vector as V+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.CRC32C (checkCRC)+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.IncomingRun+import qualified Database.LSMTree.Internal.Merge as Merge+import Database.LSMTree.Internal.MergeSchedule+import qualified Database.LSMTree.Internal.MergingRun as MR+import qualified Database.LSMTree.Internal.MergingTree as MT+import Database.LSMTree.Internal.Paths (ActiveDir (..), ForBlob (..),+ ForKOps (..), NamedSnapshotDir (..), RunFsPaths (..),+ WriteBufferFsPaths (..),+ fromChecksumsFileForWriteBufferFiles, pathsForRunFiles,+ runChecksumsPath, runPath, writeBufferBlobPath,+ writeBufferChecksumsPath, writeBufferKOpsPath)+import Database.LSMTree.Internal.Run (Run, RunParams)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue)+import Database.LSMTree.Internal.UniqCounter (UniqCounter,+ incrUniqCounter, uniqueToInt, uniqueToRunNumber)+import Database.LSMTree.Internal.WriteBuffer (WriteBuffer)+import Database.LSMTree.Internal.WriteBufferBlobs (WriteBufferBlobs)+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import qualified Database.LSMTree.Internal.WriteBufferReader as WBR+import qualified Database.LSMTree.Internal.WriteBufferWriter as WBW+import qualified System.FS.API as FS+import System.FS.API (HasFS, (<.>), (</>))+import qualified System.FS.API.Lazy as FSL+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++{-------------------------------------------------------------------------------+ Snapshot metadata+-------------------------------------------------------------------------------}++-- | Custom, user-supplied text that is included in the metadata.+--+-- The main use case for a 'SnapshotLabel' is for the user to supply textual+-- information about the key\/value\/blob type for the table that corresponds to+-- the snapshot. This information is used to dynamically check that a snapshot+-- is opened at the correct key\/value\/blob type.+newtype SnapshotLabel = SnapshotLabel Text+ deriving stock (Show, Eq)+ deriving newtype (NFData, IsString)++data SnapshotMetaData = SnapshotMetaData {+ -- | See 'SnapshotLabel'.+ --+ -- One could argue that the 'SnapshotName' could be used to to hold this+ -- type information, but the file name of snapshot metadata is not guarded+ -- by a checksum, whereas the contents of the file are. Therefore using the+ -- 'SnapshotLabel' is safer.+ snapMetaLabel :: !SnapshotLabel+ -- | The 'TableConfig' for the snapshotted table.+ , snapMetaConfig :: !TableConfig+ -- | The write buffer.+ , snapWriteBuffer :: !RunNumber+ -- | The shape of the levels of the LSM tree.+ , snapMetaLevels :: !(SnapLevels SnapshotRun)+ -- | The state of tree merging of the LSM tree.+ , snapMergingTree :: !(Maybe (SnapMergingTree SnapshotRun))+ }+ deriving stock Eq++instance NFData SnapshotMetaData where+ rnf (SnapshotMetaData a b c d e) =+ rnf a `seq` rnf b `seq` rnf c `seq`+ rnf d `seq` rnf e++{-------------------------------------------------------------------------------+ Levels snapshot format+-------------------------------------------------------------------------------}++newtype SnapLevels r = SnapLevels { getSnapLevels :: V.Vector (SnapLevel r) }+ deriving stock (Eq, Functor, Foldable, Traversable)+ deriving newtype NFData++data SnapLevel r = SnapLevel {+ snapIncoming :: !(SnapIncomingRun r)+ , snapResidentRuns :: !(V.Vector r)+ }+ deriving stock (Eq, Functor, Foldable, Traversable)++instance NFData r => NFData (SnapLevel r) where+ rnf (SnapLevel a b) = rnf a `seq` rnf b++-- | Note that for snapshots of incoming runs, we store only the merge debt and+-- nominal credits, not the nominal debt or the merge credits. The rationale is+-- a bit subtle.+--+-- The nominal debt does not need to be stored because it can be derived based+-- on the table's write buffer capacity (which is stored in the snapshot's+-- TableConfig), and on the level number that the merge is at (which also known+-- from the snapshot structure).+--+-- The merge credits can be recalculated from the combination of the nominal debt,+-- nominal credits and merge debt.+--+-- The merge debt is always the sum of the size of the input runs, so at first+-- glance this seems redundant. However for completed merges we no longer have+-- the input runs, so we must store the merge debt if we are to perfectly round+-- trip the snapshot. This is a nice simple property to have though it is+-- probably not 100% essential. We could weaken the round trip property to+-- allow forgetting the merge debt and credit of completed merges (and set them+-- both to zero).+--+data SnapIncomingRun r =+ SnapIncomingMergingRun+ !MergePolicyForLevel+ !NominalDebt+ !NominalCredits -- ^ The nominal credits supplied, and that+ -- need to be supplied on snapshot open.+ !(SnapMergingRun MR.LevelMergeType r)+ | SnapIncomingSingleRun !r+ deriving stock (Eq, Functor, Foldable, Traversable)++instance NFData r => NFData (SnapIncomingRun r) where+ rnf (SnapIncomingMergingRun a b c d) =+ rnf a `seq` rnf b `seq` rnf c `seq` rnf d+ rnf (SnapIncomingSingleRun a) = rnf a++-- | The total number of supplied credits. This total is used on snapshot load+-- to restore merging work that was lost when the snapshot was created.+newtype SuppliedCredits = SuppliedCredits { getSuppliedCredits :: Int }+ deriving stock (Eq, Read)+ deriving newtype NFData++data SnapMergingRun t r =+ SnapCompletedMerge !MergeDebt !r+ | SnapOngoingMerge !RunParams !MergeCredits !(V.Vector r) !t+ deriving stock (Eq, Functor, Foldable, Traversable)++instance (NFData t, NFData r) => NFData (SnapMergingRun t r) where+ rnf (SnapCompletedMerge a b) = rnf a `seq` rnf b+ rnf (SnapOngoingMerge a b c d) = rnf a `seq` rnf b `seq` rnf c `seq` rnf d++{-------------------------------------------------------------------------------+ Snapshot MergingTree+-------------------------------------------------------------------------------}++newtype SnapMergingTree r = SnapMergingTree (SnapMergingTreeState r)+ deriving stock (Eq, Functor, Foldable, Traversable)+ deriving newtype NFData++data SnapMergingTreeState r =+ SnapCompletedTreeMerge !r+ | SnapPendingTreeMerge !(SnapPendingMerge r)+ | SnapOngoingTreeMerge !(SnapMergingRun MR.TreeMergeType r)+ deriving stock (Eq, Functor, Foldable, Traversable)++instance NFData r => NFData (SnapMergingTreeState r) where+ rnf (SnapCompletedTreeMerge a) = rnf a+ rnf (SnapPendingTreeMerge a) = rnf a+ rnf (SnapOngoingTreeMerge a) = rnf a++data SnapPendingMerge r =+ SnapPendingLevelMerge+ ![SnapPreExistingRun r]+ !(Maybe (SnapMergingTree r))+ | SnapPendingUnionMerge+ ![SnapMergingTree r]+ deriving stock (Eq, Functor, Foldable, Traversable)++instance NFData r => NFData (SnapPendingMerge r) where+ rnf (SnapPendingLevelMerge a b) = rnf a `seq` rnf b+ rnf (SnapPendingUnionMerge a) = rnf a++data SnapPreExistingRun r =+ SnapPreExistingRun !r+ | SnapPreExistingMergingRun !(SnapMergingRun MR.LevelMergeType r)+ deriving stock (Eq, Functor, Foldable, Traversable)++instance NFData r => NFData (SnapPreExistingRun r) where+ rnf (SnapPreExistingRun a) = rnf a+ rnf (SnapPreExistingMergingRun a) = rnf a++{-------------------------------------------------------------------------------+ Opening from merging tree snapshot format+-------------------------------------------------------------------------------}++{-# SPECIALISE fromSnapMergingTree ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> UniqCounter IO+ -> ResolveSerialisedValue+ -> ActiveDir+ -> ActionRegistry IO+ -> SnapMergingTree (Ref (Run IO h))+ -> IO (Ref (MT.MergingTree IO h))+ #-}+-- | Converts a snapshot of a merging tree of runs to a real merging tree.+--+-- Returns a new reference. Input runs remain owned by the caller.+fromSnapMergingTree ::+ forall m h. (MonadMask m, MonadMVar m, MonadSTM m, MonadST m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> UniqCounter m+ -> ResolveSerialisedValue+ -> ActiveDir+ -> ActionRegistry m+ -> SnapMergingTree (Ref (Run m h))+ -> m (Ref (MT.MergingTree m h))+fromSnapMergingTree hfs hbio salt uc resolve dir =+ go+ where+ -- Reference strategy:+ -- * go returns a fresh reference+ -- * go ensures the returned reference will be cleaned up on failure,+ -- using withRollback+ -- * All results from recursive calls must be released locally on the+ -- happy path.+ go :: ActionRegistry m+ -> SnapMergingTree (Ref (Run m h))+ -> m (Ref (MT.MergingTree m h))++ go reg (SnapMergingTree (SnapCompletedTreeMerge run)) =+ withRollback reg+ (MT.newCompletedMerge run)+ releaseRef++ go reg (SnapMergingTree (SnapPendingTreeMerge+ (SnapPendingLevelMerge prs mmt))) = do+ prs' <- traverse (fromSnapPreExistingRun reg) prs+ mmt' <- traverse (go reg) mmt+ mt <- withRollback reg+ (MT.newPendingLevelMerge prs' mmt')+ releaseRef+ traverse_ (delayedCommit reg . releasePER) prs'+ traverse_ (delayedCommit reg . releaseRef) mmt'+ pure mt++ go reg (SnapMergingTree (SnapPendingTreeMerge+ (SnapPendingUnionMerge mts))) = do+ mts' <- traverse (go reg) mts+ mt <- withRollback reg+ (MT.newPendingUnionMerge mts')+ releaseRef+ traverse_ (delayedCommit reg . releaseRef) mts'+ pure mt++ go reg (SnapMergingTree (SnapOngoingTreeMerge smrs)) = do+ mr <- withRollback reg+ (fromSnapMergingRun hfs hbio salt uc resolve dir smrs)+ releaseRef+ mt <- withRollback reg+ (MT.newOngoingMerge mr)+ releaseRef+ delayedCommit reg (releaseRef mr)+ pure mt++ -- Returns fresh refs, which must be released locally.+ fromSnapPreExistingRun :: ActionRegistry m+ -> SnapPreExistingRun (Ref (Run m h))+ -> m (MT.PreExistingRun m h)+ fromSnapPreExistingRun reg (SnapPreExistingRun run) =+ MT.PreExistingRun <$>+ withRollback reg (dupRef run) releaseRef+ fromSnapPreExistingRun reg (SnapPreExistingMergingRun smrs) =+ MT.PreExistingMergingRun <$>+ withRollback reg+ (fromSnapMergingRun hfs hbio salt uc resolve dir smrs)+ releaseRef++ releasePER (MT.PreExistingRun r) = releaseRef r+ releasePER (MT.PreExistingMergingRun mr) = releaseRef mr++{-------------------------------------------------------------------------------+ Conversion to merge tree snapshot format+-------------------------------------------------------------------------------}++{-# SPECIALISE toSnapMergingTree :: Ref (MT.MergingTree IO h) -> IO (SnapMergingTree (Ref (Run IO h))) #-}+toSnapMergingTree ::+ (PrimMonad m, MonadMVar m)+ => Ref (MT.MergingTree m h)+ -> m (SnapMergingTree (Ref (Run m h)))+toSnapMergingTree (DeRef (MT.MergingTree mStateVar _mCounter)) =+ withMVar mStateVar $ \mState -> SnapMergingTree <$> toSnapMergingTreeState mState++{-# SPECIALISE toSnapMergingTreeState :: MT.MergingTreeState IO h -> IO (SnapMergingTreeState (Ref (Run IO h))) #-}+toSnapMergingTreeState ::+ (PrimMonad m, MonadMVar m)+ => MT.MergingTreeState m h+ -> m (SnapMergingTreeState (Ref (Run m h)))+toSnapMergingTreeState (MT.CompletedTreeMerge r) = pure $ SnapCompletedTreeMerge r+toSnapMergingTreeState (MT.PendingTreeMerge p) = SnapPendingTreeMerge <$> toSnapPendingMerge p+toSnapMergingTreeState (MT.OngoingTreeMerge mergingRun) =+ SnapOngoingTreeMerge <$> toSnapMergingRun mergingRun++{-# SPECIALISE toSnapPendingMerge :: MT.PendingMerge IO h -> IO (SnapPendingMerge (Ref (Run IO h))) #-}+toSnapPendingMerge ::+ (PrimMonad m, MonadMVar m)+ => MT.PendingMerge m h+ -> m (SnapPendingMerge (Ref (Run m h)))+toSnapPendingMerge (MT.PendingUnionMerge mts) =+ SnapPendingUnionMerge <$> traverse toSnapMergingTree (V.toList mts)+toSnapPendingMerge (MT.PendingLevelMerge pes mmt) = do+ pes' <- traverse toSnapPreExistingRun pes+ mmt' <- traverse toSnapMergingTree mmt+ pure $ SnapPendingLevelMerge (V.toList pes') mmt'++{-# SPECIALISE toSnapPreExistingRun :: MT.PreExistingRun IO h -> IO (SnapPreExistingRun (Ref (Run IO h))) #-}+toSnapPreExistingRun ::+ (PrimMonad m, MonadMVar m)+ => MT.PreExistingRun m h+ -> m (SnapPreExistingRun (Ref (Run m h)))+toSnapPreExistingRun (MT.PreExistingRun run) = pure $ SnapPreExistingRun run+toSnapPreExistingRun (MT.PreExistingMergingRun peMergingRun) =+ SnapPreExistingMergingRun <$> toSnapMergingRun peMergingRun++{-------------------------------------------------------------------------------+ Conversion to levels snapshot format+-------------------------------------------------------------------------------}++--TODO: probably generally all the Ref (Run _) here ought to be fresh+-- references, created as we snapshot the levels, so that the runs don't+-- disappear under our feet during the process of making the snapshot durable.+-- At minimum the volatile runs are the inputs to merging runs, but it may be+-- simpler to duplicate them all, and release them all at the end.++{-# SPECIALISE toSnapLevels :: Levels IO h -> IO (SnapLevels (Ref (Run IO h))) #-}+toSnapLevels ::+ (PrimMonad m, MonadMVar m)+ => Levels m h+ -> m (SnapLevels (Ref (Run m h)))+toSnapLevels levels = SnapLevels <$> V.mapM toSnapLevel levels++{-# SPECIALISE toSnapLevel :: Level IO h -> IO (SnapLevel (Ref (Run IO h))) #-}+toSnapLevel ::+ (PrimMonad m, MonadMVar m)+ => Level m h+ -> m (SnapLevel (Ref (Run m h)))+toSnapLevel Level{..} = do+ sir <- toSnapIncomingRun incomingRun+ pure (SnapLevel sir residentRuns)++{-# SPECIALISE toSnapIncomingRun :: IncomingRun IO h -> IO (SnapIncomingRun (Ref (Run IO h))) #-}+toSnapIncomingRun ::+ (PrimMonad m, MonadMVar m)+ => IncomingRun m h+ -> m (SnapIncomingRun (Ref (Run m h)))+toSnapIncomingRun ir = do+ s <- snapshotIncomingRun ir+ case s of+ Left r -> pure $! SnapIncomingSingleRun r+ Right (mergePolicy,+ nominalDebt,+ nominalCredits,+ mergingRun) -> do+ -- We need to know how many credits were supplied so we can restore merge+ -- work on snapshot load.+ smrs <- toSnapMergingRun mergingRun+ pure $! SnapIncomingMergingRun mergePolicy nominalDebt nominalCredits smrs++{-# SPECIALISE toSnapMergingRun ::+ Ref (MR.MergingRun t IO h)+ -> IO (SnapMergingRun t (Ref (Run IO h))) #-}+toSnapMergingRun ::+ (PrimMonad m, MonadMVar m)+ => Ref (MR.MergingRun t m h)+ -> m (SnapMergingRun t (Ref (Run m h)))+toSnapMergingRun !mr = do+ -- TODO: MR.snapshot needs to return duplicated run references, and we+ -- need to arrange to release them when the snapshotting is done.+ ( mergeDebt, mergeCredits, state) <- MR.snapshot mr+ case state of+ MR.CompletedMerge r ->+ pure $! SnapCompletedMerge mergeDebt r++ MR.OngoingMerge rs m ->+ pure $! SnapOngoingMerge runParams mergeCredits rs mergeType+ where+ runParams = Merge.mergeRunParams m+ mergeType = Merge.mergeType m++{-------------------------------------------------------------------------------+ Write Buffer+-------------------------------------------------------------------------------}++{-# SPECIALISE+ snapshotWriteBuffer ::+ HasFS IO h+ -> HasBlockIO IO h+ -> UniqCounter IO+ -> UniqCounter IO+ -> ActionRegistry IO+ -> ActiveDir+ -> NamedSnapshotDir+ -> WriteBuffer+ -> Ref (WriteBufferBlobs IO h)+ -> IO WriteBufferFsPaths+ #-}+snapshotWriteBuffer ::+ (MonadMVar m, MonadSTM m, MonadST m, MonadMask m)+ => HasFS m h+ -> HasBlockIO m h+ -> UniqCounter m+ -> UniqCounter m+ -> ActionRegistry m+ -> ActiveDir+ -> NamedSnapshotDir+ -> WriteBuffer+ -> Ref (WriteBufferBlobs m h)+ -> m WriteBufferFsPaths+snapshotWriteBuffer hfs hbio activeUc snapUc reg activeDir snapDir wb wbb = do+ -- Write the write buffer and write buffer blobs to the active directory.+ activeWriteBufferNumber <- uniqueToRunNumber <$> incrUniqCounter activeUc+ let activeWriteBufferPaths = WriteBufferFsPaths (getActiveDir activeDir) activeWriteBufferNumber+ withRollback_ reg+ (WBW.writeWriteBuffer hfs hbio activeWriteBufferPaths wb wbb)+ -- TODO: it should probably be the responsibility of writeWriteBuffer to do+ -- cleanup+ $ do+ -- TODO: check files exist before removing them+ FS.removeFile hfs (writeBufferKOpsPath activeWriteBufferPaths)+ FS.removeFile hfs (writeBufferBlobPath activeWriteBufferPaths)+ -- Hard link the write buffer and write buffer blobs to the snapshot directory.+ snapWriteBufferNumber <- uniqueToRunNumber <$> incrUniqCounter snapUc+ let snapWriteBufferPaths = WriteBufferFsPaths (getNamedSnapshotDir snapDir) snapWriteBufferNumber+ hardLink hfs hbio reg+ (writeBufferKOpsPath activeWriteBufferPaths)+ (writeBufferKOpsPath snapWriteBufferPaths)+ hardLink hfs hbio reg+ (writeBufferBlobPath activeWriteBufferPaths)+ (writeBufferBlobPath snapWriteBufferPaths)+ hardLink hfs hbio reg+ (writeBufferChecksumsPath activeWriteBufferPaths)+ (writeBufferChecksumsPath snapWriteBufferPaths)+ pure snapWriteBufferPaths++{-# SPECIALISE+ openWriteBuffer ::+ ActionRegistry IO+ -> ResolveSerialisedValue+ -> HasFS IO h+ -> HasBlockIO IO h+ -> UniqCounter IO+ -> ActiveDir+ -> WriteBufferFsPaths+ -> IO (WriteBuffer, Ref (WriteBufferBlobs IO h))+ #-}+openWriteBuffer ::+ (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => ActionRegistry m+ -> ResolveSerialisedValue+ -> HasFS m h+ -> HasBlockIO m h+ -> UniqCounter m+ -> ActiveDir+ -> WriteBufferFsPaths+ -> m (WriteBuffer, Ref (WriteBufferBlobs m h))+openWriteBuffer reg resolve hfs hbio uc activeDir snapWriteBufferPaths = do+ -- Check the checksums+ -- TODO: This reads the blobfile twice: once to check the CRC and once more+ -- to copy it from the snapshot directory to the active directory.+ (expectedChecksumForKOps, expectedChecksumForBlob) <-+ CRC.expectValidFile hfs (writeBufferChecksumsPath snapWriteBufferPaths) CRC.FormatWriteBufferFile+ . fromChecksumsFileForWriteBufferFiles+ =<< CRC.readChecksumsFile hfs (writeBufferChecksumsPath snapWriteBufferPaths)+ checkCRC hfs hbio False (unForKOps expectedChecksumForKOps) (writeBufferKOpsPath snapWriteBufferPaths)+ checkCRC hfs hbio False (unForBlob expectedChecksumForBlob) (writeBufferBlobPath snapWriteBufferPaths)+ -- Copy the write buffer blobs file to the active directory and open it.+ activeWriteBufferNumber <- uniqueToInt <$> incrUniqCounter uc+ let activeWriteBufferBlobPath =+ getActiveDir activeDir </> FS.mkFsPath [show activeWriteBufferNumber] <.> "wbblobs"+ copyFile hfs reg (writeBufferBlobPath snapWriteBufferPaths) activeWriteBufferBlobPath+ writeBufferBlobs <-+ withRollback reg+ (WBB.open hfs activeWriteBufferBlobPath FS.AllowExisting)+ releaseRef+ -- Read write buffer key/ops+ let kOpsPath = ForKOps (writeBufferKOpsPath snapWriteBufferPaths)+ writeBuffer <-+ withRef writeBufferBlobs $ \wbb ->+ WBR.readWriteBuffer resolve hfs hbio kOpsPath (WBB.blobFile wbb)+ pure (writeBuffer, writeBufferBlobs)++{-------------------------------------------------------------------------------+ Runs+-------------------------------------------------------------------------------}++-- | Information needed to open a 'Run' from disk using 'snapshotRun' and+-- 'openRun'.+--+-- TODO: one could imagine needing only the 'RunNumber' to identify the files+-- on disk, and the other parameters being stored with the run itself, rather+-- than needing to be supplied.+data SnapshotRun = SnapshotRun {+ snapRunNumber :: !RunNumber,+ snapRunCaching :: !Run.RunDataCaching,+ snapRunIndex :: !Run.IndexType+ }+ deriving stock Eq++instance NFData SnapshotRun where+ rnf (SnapshotRun a b c) = rnf a `seq` rnf b `seq` rnf c++{-# SPECIALISE snapshotRun ::+ HasFS IO h+ -> HasBlockIO IO h+ -> UniqCounter IO+ -> ActionRegistry IO+ -> NamedSnapshotDir+ -> Ref (Run IO h)+ -> IO SnapshotRun #-}+-- | @'snapshotRun' _ _ snapUc _ targetDir run@ creates hard links for all files+-- associated with the @run@, and puts the new directory entries in the+-- @targetDir@ directory. The entries are renamed using @snapUc@.+snapshotRun ::+ (MonadMask m, PrimMonad m)+ => HasFS m h+ -> HasBlockIO m h+ -> UniqCounter m+ -> ActionRegistry m+ -> NamedSnapshotDir+ -> Ref (Run m h)+ -> m SnapshotRun+snapshotRun hfs hbio snapUc reg (NamedSnapshotDir targetDir) run = do+ rn <- uniqueToRunNumber <$> incrUniqCounter snapUc+ let sourcePaths = Run.runFsPaths run+ let targetPaths = sourcePaths { runDir = targetDir , runNumber = rn}+ hardLinkRunFiles hfs hbio reg sourcePaths targetPaths+ pure SnapshotRun {+ snapRunNumber = runNumber targetPaths,+ snapRunCaching = Run.runDataCaching run,+ snapRunIndex = Run.runIndexType run+ }++{-# SPECIALISE openRun ::+ HasFS IO h+ -> HasBlockIO IO h+ -> UniqCounter IO+ -> ActionRegistry IO+ -> NamedSnapshotDir+ -> ActiveDir+ -> Bloom.Salt+ -> SnapshotRun+ -> IO (Ref (Run IO h)) #-}+-- | @'openRun' _ _ uniqCounter _ sourceDir targetDir _ snaprun@ takes all run+-- files that are referenced by @snaprun@, and hard links them from @sourceDir@+-- into @targetDir@ with new, unique names (using @uniqCounter@). Each set of+-- (hard linked) files that represents a run is opened and verified, returning+-- 'Run' as a result.+--+-- The result must ultimately be released using 'releaseRef'.+openRun ::+ (MonadMask m, MonadSTM m, MonadST m)+ => HasFS m h+ -> HasBlockIO m h+ -> UniqCounter m+ -> ActionRegistry m+ -> NamedSnapshotDir+ -> ActiveDir+ -> Bloom.Salt+ -> SnapshotRun+ -> m (Ref (Run m h))+openRun hfs hbio uc reg+ (NamedSnapshotDir sourceDir) (ActiveDir targetDir)+ expectedSalt+ SnapshotRun {+ snapRunNumber = runNum,+ snapRunCaching = caching,+ snapRunIndex = indexType+ } = do+ let sourcePaths = RunFsPaths sourceDir runNum+ runNum' <- uniqueToRunNumber <$> incrUniqCounter uc+ let targetPaths = RunFsPaths targetDir runNum'+ hardLinkRunFiles hfs hbio reg sourcePaths targetPaths++ withRollback reg+ (Run.openFromDisk hfs hbio caching indexType expectedSalt targetPaths)+ releaseRef++{-------------------------------------------------------------------------------+ Opening from levels snapshot format+-------------------------------------------------------------------------------}++{-# SPECIALISE fromSnapLevels ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> UniqCounter IO+ -> TableConfig+ -> ResolveSerialisedValue+ -> ActionRegistry IO+ -> ActiveDir+ -> SnapLevels (Ref (Run IO h))+ -> IO (Levels IO h)+ #-}+-- | Duplicates runs and re-creates merging runs.+fromSnapLevels ::+ forall m h. (MonadMask m, MonadMVar m, MonadSTM m, MonadST m)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> UniqCounter m+ -> TableConfig+ -> ResolveSerialisedValue+ -> ActionRegistry m+ -> ActiveDir+ -> SnapLevels (Ref (Run m h))+ -> m (Levels m h)+fromSnapLevels hfs hbio salt uc conf resolve reg dir (SnapLevels levels) =+ V.iforM levels $ \i -> fromSnapLevel (LevelNo (i+1))+ where+ -- TODO: we may wish to trace the merges created during snapshot restore:++ fromSnapLevel :: LevelNo -> SnapLevel (Ref (Run m h)) -> m (Level m h)+ fromSnapLevel ln SnapLevel{snapIncoming, snapResidentRuns} = do+ incomingRun <- withRollback reg+ (fromSnapIncomingRun ln snapIncoming)+ releaseIncomingRun+ residentRuns <- V.forM snapResidentRuns $ \r ->+ withRollback reg+ (dupRef r)+ releaseRef+ pure Level {incomingRun , residentRuns}++ fromSnapIncomingRun ::+ LevelNo+ -> SnapIncomingRun (Ref (Run m h))+ -> m (IncomingRun m h)+ fromSnapIncomingRun _ln (SnapIncomingSingleRun run) =+ newIncomingSingleRun run++ fromSnapIncomingRun ln (SnapIncomingMergingRun mergePolicy nominalDebt+ nominalCredits smrs) =+ bracket+ (fromSnapMergingRun hfs hbio salt uc resolve dir smrs)+ releaseRef $ \mr -> do++ ir <- newIncomingMergingRun mergePolicy nominalDebt mr+ -- This will set the correct nominal credits, but it will not do any+ -- more merging work because fromSnapMergingRun already supplies+ -- all the merging credits already.+ supplyCreditsIncomingRun conf ln ir nominalCredits+ pure ir++{-# SPECIALISE fromSnapMergingRun ::+ MR.IsMergeType t+ => HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> UniqCounter IO+ -> ResolveSerialisedValue+ -> ActiveDir+ -> SnapMergingRun t (Ref (Run IO h))+ -> IO (Ref (MR.MergingRun t IO h)) #-}+fromSnapMergingRun ::+ (MonadMask m, MonadMVar m, MonadSTM m, MonadST m, MR.IsMergeType t)+ => HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> UniqCounter m+ -> ResolveSerialisedValue+ -> ActiveDir+ -> SnapMergingRun t (Ref (Run m h))+ -> m (Ref (MR.MergingRun t m h))+fromSnapMergingRun _ _ _ _ _ _ (SnapCompletedMerge mergeDebt r) =+ MR.newCompleted mergeDebt r++fromSnapMergingRun hfs hbio salt uc resolve dir+ (SnapOngoingMerge runParams mergeCredits rs mergeType) = do+ bracketOnError+ (do uniq <- incrUniqCounter uc+ let runPaths = runPath dir (uniqueToRunNumber uniq)+ MR.new hfs hbio resolve salt runParams mergeType runPaths rs)+ releaseRef $ \mr -> do+ -- When a snapshot is created, merge progress is lost, so we have to+ -- redo merging work here. The MergeCredits in SnapMergingRun tracks+ -- how many credits were supplied before the snapshot was taken.++ --TODO: the threshold should be stored with the MergingRun+ -- here we want to supply the credits now, so we can use a threshold of 1+ let thresh = MR.CreditThreshold (MR.UnspentCredits 1)+ _ <- MR.supplyCreditsAbsolute mr thresh mergeCredits+ pure mr++{-------------------------------------------------------------------------------+ Hard links+-------------------------------------------------------------------------------}++{-# SPECIALISE hardLinkRunFiles ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ActionRegistry IO+ -> RunFsPaths+ -> RunFsPaths+ -> IO () #-}+-- | @'hardLinkRunFiles' _ _ _ sourcePaths targetPaths@ creates a hard link for+-- each @sourcePaths@ path using the corresponding @targetPaths@ path as the+-- name for the new directory entry.+hardLinkRunFiles ::+ (MonadMask m, PrimMonad m)+ => HasFS m h+ -> HasBlockIO m h+ -> ActionRegistry m+ -> RunFsPaths+ -> RunFsPaths+ -> m ()+hardLinkRunFiles hfs hbio reg sourceRunFsPaths targetRunFsPaths = do+ let sourcePaths = pathsForRunFiles sourceRunFsPaths+ targetPaths = pathsForRunFiles targetRunFsPaths+ sequenceA_ (hardLink hfs hbio reg <$> sourcePaths <*> targetPaths)+ hardLink hfs hbio reg (runChecksumsPath sourceRunFsPaths) (runChecksumsPath targetRunFsPaths)++{-# SPECIALISE+ hardLink ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ActionRegistry IO+ -> FS.FsPath+ -> FS.FsPath+ -> IO ()+ #-}+-- | @'hardLink' hfs hbio reg sourcePath targetPath@ creates a hard link from+-- @sourcePath@ to @targetPath@.+hardLink ::+ (MonadMask m, PrimMonad m)+ => HasFS m h+ -> HasBlockIO m h+ -> ActionRegistry m+ -> FS.FsPath+ -> FS.FsPath+ -> m ()+hardLink hfs hbio reg sourcePath targetPath = do+ withRollback_ reg+ (FS.createHardLink hbio sourcePath targetPath)+ (FS.removeFile hfs targetPath)++{-------------------------------------------------------------------------------+ Copy file+-------------------------------------------------------------------------------}++{-# SPECIALISE+ copyFile ::+ HasFS IO h+ -> ActionRegistry IO+ -> FS.FsPath+ -> FS.FsPath+ -> IO ()+ #-}+-- | @'copyFile' hfs reg source target@ copies the @source@ path to the @target@ path.+copyFile ::+ (MonadMask m, PrimMonad m)+ => HasFS m h+ -> ActionRegistry m+ -> FS.FsPath+ -> FS.FsPath+ -> m ()+copyFile hfs reg sourcePath targetPath =+ flip (withRollback_ reg) (FS.removeFile hfs targetPath) $+ FS.withFile hfs sourcePath FS.ReadMode $ \sourceHandle ->+ FS.withFile hfs targetPath (FS.WriteMode FS.MustBeNew) $ \targetHandle -> do+ bs <- FSL.hGetAll hfs sourceHandle+ void $ FSL.hPutAll hfs targetHandle bs
@@ -0,0 +1,809 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | Encoders and decoders for snapshot metadata+--+module Database.LSMTree.Internal.Snapshot.Codec (+ -- * Versioning+ SnapshotVersion (..)+ , prettySnapshotVersion+ , currentSnapshotVersion+ , allCompatibleSnapshotVersions+ -- * Writing and reading files+ , writeFileSnapshotMetaData+ , readFileSnapshotMetaData+ , encodeSnapshotMetaData+ -- * Encoding and decoding+ , Encode (..)+ , Decode (..)+ , DecodeVersioned (..)+ , Versioned (..)+ ) where++import Codec.CBOR.Decoding+import Codec.CBOR.Encoding+import Codec.CBOR.Read+import Codec.CBOR.Write+import Control.Monad.Class.MonadThrow (Exception (displayException),+ MonadThrow (..))+import Data.Bifunctor (Bifunctor (..))+import qualified Data.ByteString.Char8 as BSC+import Data.ByteString.Lazy (ByteString)+import qualified Data.Map.Strict as Map+import qualified Data.Vector as V+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.CRC32C+import qualified Database.LSMTree.Internal.CRC32C as FS+import Database.LSMTree.Internal.MergeSchedule+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.RunBuilder (IndexType (..),+ RunBloomFilterAlloc (..), RunDataCaching (..),+ RunParams (..))+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Snapshot+import qualified System.FS.API as FS+import System.FS.API (FsPath, HasFS)+import Text.Printf++{-------------------------------------------------------------------------------+ Versioning+-------------------------------------------------------------------------------}++-- | The version of a snapshot.+--+-- A snapshot format version is a number. Version numbers are consecutive and+-- increasing. A single release of the library may support several older+-- snapshot format versions, and thereby provide backwards compatibility.+-- Support for old versions is not guaranteed indefinitely, but backwards+-- compatibility is guaranteed for at least the previous version, and preferably+-- for more. Forwards compatibility is not provided at all: snapshots with a+-- later version than the current version for the library release will always+-- fail.+data SnapshotVersion = V0 | V1+ deriving stock (Show, Eq, Ord)++-- | Pretty-print a snapshot version+--+-- >>> prettySnapshotVersion currentSnapshotVersion+-- "v1"+prettySnapshotVersion :: SnapshotVersion -> String+prettySnapshotVersion V0 = "v0"+prettySnapshotVersion V1 = "v1"++-- | The current snapshot version+--+-- >>> currentSnapshotVersion+-- V1+currentSnapshotVersion :: SnapshotVersion+currentSnapshotVersion = V1++-- | All snapshot versions that the current snapshpt version is compatible with.+--+-- >>> allCompatibleSnapshotVersions+-- [V0,V1]+--+-- >>> last allCompatibleSnapshotVersions == currentSnapshotVersion+-- True+allCompatibleSnapshotVersions :: [SnapshotVersion]+allCompatibleSnapshotVersions = [V0, V1]++isCompatible :: SnapshotVersion -> Either String ()+isCompatible otherVersion+ -- for the moment, all versions are backwards compatible:+ | otherVersion `elem` allCompatibleSnapshotVersions+ = Right ()+ | otherwise = Left "forward compatibility not supported"++{-------------------------------------------------------------------------------+ Writing and reading files+-------------------------------------------------------------------------------}++{-# SPECIALISE+ writeFileSnapshotMetaData ::+ HasFS IO h+ -> FsPath+ -> FsPath+ -> SnapshotMetaData+ -> IO ()+ #-}+-- | Encode 'SnapshotMetaData' and write it to 'SnapshotMetaDataFile'.+--+-- In the presence of exceptions, newly created files will not be removed. It is+-- up to the user of this function to clean up the files.+writeFileSnapshotMetaData ::+ MonadThrow m+ => HasFS m h+ -> FsPath -- ^ Target file for snapshot metadata+ -> FsPath -- ^ Target file for checksum+ -> SnapshotMetaData+ -> m ()+writeFileSnapshotMetaData hfs contentPath checksumPath snapMetaData = do+ (_, checksum) <-+ FS.withFile hfs contentPath (FS.WriteMode FS.MustBeNew) $ \h ->+ hPutAllChunksCRC32C hfs h (encodeSnapshotMetaData snapMetaData) initialCRC32C++ let checksumFileName = ChecksumsFileName (BSC.pack "metadata")+ checksumFile = Map.singleton checksumFileName checksum+ writeChecksumsFile hfs checksumPath checksumFile++encodeSnapshotMetaData :: SnapshotMetaData -> ByteString+encodeSnapshotMetaData = toLazyByteString . encode . Versioned++{-# SPECIALISE+ readFileSnapshotMetaData ::+ HasFS IO h+ -> FsPath+ -> FsPath+ -> IO SnapshotMetaData+ #-}+-- | Read from 'SnapshotMetaDataFile' and attempt to decode it to+-- 'SnapshotMetaData'.+readFileSnapshotMetaData ::+ (MonadThrow m)+ => HasFS m h+ -> FsPath -- ^ Source file for snapshot metadata+ -> FsPath -- ^ Source file for checksum+ -> m SnapshotMetaData+readFileSnapshotMetaData hfs contentPath checksumPath = do+ checksumFile <- readChecksumsFile hfs checksumPath+ let checksumFileName = ChecksumsFileName (BSC.pack "metadata")++ expectedChecksum <- getChecksum hfs checksumPath checksumFile checksumFileName++ (lbs, actualChecksum) <- FS.withFile hfs contentPath FS.ReadMode $ \h -> do+ n <- FS.hGetSize hfs h+ FS.hGetExactlyCRC32C hfs h n initialCRC32C++ expectChecksum hfs contentPath expectedChecksum actualChecksum++ expectValidFile hfs contentPath FormatSnapshotMetaData (decodeSnapshotMetaData lbs)++decodeSnapshotMetaData :: ByteString -> Either String SnapshotMetaData+decodeSnapshotMetaData lbs = bimap displayException (getVersioned . snd) (deserialiseFromBytes decode lbs)++{-------------------------------------------------------------------------------+ Encoding and decoding+-------------------------------------------------------------------------------}++class Encode a where+ encode :: a -> Encoding++-- | Decoder that is not parameterised by a 'SnapshotVersion'.+--+-- Used only for 'SnapshotVersion' and 'Versioned', which live outside the+-- 'SnapshotMetaData' type hierarchy.+class Decode a where+ decode :: Decoder s a++-- | Decoder parameterised by a 'SnapshotVersion'.+--+-- Used for every type in the 'SnapshotMetaData' type hierarchy.+class DecodeVersioned a where+ decodeVersioned :: SnapshotVersion -> Decoder s a++newtype Versioned a = Versioned { getVersioned :: a }+ deriving stock (Show, Eq)++instance Encode a => Encode (Versioned a) where+ encode (Versioned x) =+ encodeListLen 2+ <> encode currentSnapshotVersion+ <> encode x++-- | Decodes a 'SnapshotVersion' first, and then passes that into the versioned+-- decoder for @a@.+instance DecodeVersioned a => Decode (Versioned a) where+ decode = do+ _ <- decodeListLenOf 2+ version <- decode+ case isCompatible version of+ Right () -> pure ()+ Left errMsg ->+ fail $+ printf "Incompatible snapshot format version found. Version %s \+ \is not backwards compatible with version %s : %s"+ (prettySnapshotVersion currentSnapshotVersion)+ (prettySnapshotVersion version)+ errMsg+ Versioned <$> decodeVersioned version++{-------------------------------------------------------------------------------+ Encoding and decoding: Versioning+-------------------------------------------------------------------------------}++instance Encode SnapshotVersion where+ encode ver =+ encodeListLen 1+ <> case ver of+ V0 -> encodeWord 0+ V1 -> encodeWord 1++instance Decode SnapshotVersion where+ decode = do+ _ <- decodeListLenOf 1+ ver <- decodeWord+ case ver of+ 0 -> pure V0+ 1 -> pure V1+ _ -> fail ("Unknown snapshot format version number: " <> show ver)++{-------------------------------------------------------------------------------+ Encoding and decoding: SnapshotMetaData+-------------------------------------------------------------------------------}++-- SnapshotMetaData++instance Encode SnapshotMetaData where+ encode (SnapshotMetaData label config writeBuffer levels mergingTree) =+ encodeListLen 5+ <> encode label+ <> encode config+ <> encode writeBuffer+ <> encode levels+ <> encodeMaybe mergingTree++instance DecodeVersioned SnapshotMetaData where+ decodeVersioned ver = do+ _ <- decodeListLenOf 5+ SnapshotMetaData+ <$> decodeVersioned ver+ <*> decodeVersioned ver+ <*> decodeVersioned ver+ <*> decodeVersioned ver+ <*> decodeMaybe ver++-- SnapshotLabel++instance Encode SnapshotLabel where+ encode (SnapshotLabel s) = encodeString s++instance DecodeVersioned SnapshotLabel where+ decodeVersioned _v = SnapshotLabel <$> decodeString++instance Encode SnapshotRun where+ encode SnapshotRun { snapRunNumber, snapRunCaching, snapRunIndex } =+ encodeListLen 4+ <> encodeWord 0+ <> encode snapRunNumber+ <> encode snapRunCaching+ <> encode snapRunIndex++instance DecodeVersioned SnapshotRun where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (4, 0) -> do snapRunNumber <- decodeVersioned v+ snapRunCaching <- decodeVersioned v+ snapRunIndex <- decodeVersioned v+ pure SnapshotRun{..}+ _ -> fail ("[SnapshotRun] Unexpected combination of list length and tag: " <> show (n, tag))++{-------------------------------------------------------------------------------+ Encoding and decoding: TableConfig+-------------------------------------------------------------------------------}++-- TableConfig++instance Encode TableConfig where+ encode+ ( TableConfig+ { confMergePolicy = mergePolicy+ , confMergeSchedule = mergeSchedule+ , confSizeRatio = sizeRatio+ , confWriteBufferAlloc = writeBufferAlloc+ , confBloomFilterAlloc = bloomFilterAlloc+ , confFencePointerIndex = fencePointerIndex+ , confDiskCachePolicy = diskCachePolicy+ , confMergeBatchSize = mergeBatchSize+ }+ ) =+ encodeListLen 8+ <> encode mergePolicy+ <> encode mergeSchedule+ <> encode sizeRatio+ <> encode writeBufferAlloc+ <> encode bloomFilterAlloc+ <> encode fencePointerIndex+ <> encode diskCachePolicy+ <> encode mergeBatchSize++instance DecodeVersioned TableConfig where+ decodeVersioned v@V0 = do+ decodeListLenOf 7+ confMergePolicy <- decodeVersioned v+ confMergeSchedule <- decodeVersioned v+ confSizeRatio <- decodeVersioned v+ confWriteBufferAlloc <- decodeVersioned v+ confBloomFilterAlloc <- decodeVersioned v+ confFencePointerIndex <- decodeVersioned v+ confDiskCachePolicy <- decodeVersioned v+ let confMergeBatchSize = case confWriteBufferAlloc of+ AllocNumEntries n -> MergeBatchSize n+ pure TableConfig {..}++ -- We introduced the confMergeBatchSize in V1+ decodeVersioned v@V1 = do+ decodeListLenOf 8+ confMergePolicy <- decodeVersioned v+ confMergeSchedule <- decodeVersioned v+ confSizeRatio <- decodeVersioned v+ confWriteBufferAlloc <- decodeVersioned v+ confBloomFilterAlloc <- decodeVersioned v+ confFencePointerIndex <- decodeVersioned v+ confDiskCachePolicy <- decodeVersioned v+ confMergeBatchSize <- decodeVersioned v+ pure TableConfig {..}++-- MergePolicy++instance Encode MergePolicy where+ encode LazyLevelling = encodeWord 0++instance DecodeVersioned MergePolicy where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure LazyLevelling+ _ -> fail ("[MergePolicy] Unexpected tag: " <> show tag)++-- SizeRatio++instance Encode SizeRatio where+ encode Four = encodeInt 4++instance DecodeVersioned SizeRatio where+ decodeVersioned _v = do+ x <- decodeWord64+ case x of+ 4 -> pure Four+ _ -> fail ("Expected 4, but found " <> show x)++-- WriteBufferAlloc++instance Encode WriteBufferAlloc where+ encode (AllocNumEntries numEntries) =+ encodeListLen 2+ <> encodeWord 0+ <> encodeInt numEntries++instance DecodeVersioned WriteBufferAlloc where+ decodeVersioned _v = do+ _ <- decodeListLenOf 2+ tag <- decodeWord+ case tag of+ 0 -> AllocNumEntries <$> decodeInt+ _ -> fail ("[WriteBufferAlloc] Unexpected tag: " <> show tag)++-- RunParams and friends++instance Encode RunParams where+ encode RunParams { runParamCaching, runParamAlloc, runParamIndex } =+ encodeListLen 4+ <> encodeWord 0+ <> encode runParamCaching+ <> encode runParamAlloc+ <> encode runParamIndex++instance DecodeVersioned RunParams where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (4, 0) -> do runParamCaching <- decodeVersioned v+ runParamAlloc <- decodeVersioned v+ runParamIndex <- decodeVersioned v+ pure RunParams{..}+ _ -> fail ("[RunParams] Unexpected combination of list length and tag: " <> show (n, tag))++instance Encode RunDataCaching where+ encode CacheRunData = encodeWord 0+ encode NoCacheRunData = encodeWord 1++instance DecodeVersioned RunDataCaching where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure CacheRunData+ 1 -> pure NoCacheRunData+ _ -> fail ("[RunDataCaching] Unexpected tag: " <> show tag)++instance Encode IndexType where+ encode Ordinary = encodeWord 0+ encode Compact = encodeWord 1++instance DecodeVersioned IndexType where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure Ordinary+ 1 -> pure Compact+ _ -> fail ("[IndexType] Unexpected tag: " <> show tag)++instance Encode RunBloomFilterAlloc where+ encode (RunAllocFixed bits) =+ encodeListLen 2+ <> encodeWord 0+ <> encodeDouble bits+ encode (RunAllocRequestFPR fpr) =+ encodeListLen 2+ <> encodeWord 1+ <> encodeDouble fpr++instance DecodeVersioned RunBloomFilterAlloc where+ decodeVersioned _v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (2, 0) -> RunAllocFixed <$> decodeDouble+ (2, 1) -> RunAllocRequestFPR <$> decodeDouble+ _ -> fail ("[RunBloomFilterAlloc] Unexpected combination of list length and tag: " <> show (n, tag))++-- BloomFilterAlloc++instance Encode BloomFilterAlloc where+ encode (AllocFixed x) =+ encodeListLen 2+ <> encodeWord 0+ <> encodeDouble x+ encode (AllocRequestFPR x) =+ encodeListLen 2+ <> encodeWord 1+ <> encodeDouble x++instance DecodeVersioned BloomFilterAlloc where+ decodeVersioned _v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (2, 0) -> AllocFixed <$> decodeDouble+ (2, 1) -> AllocRequestFPR <$> decodeDouble+ _ -> fail ("[BloomFilterAlloc] Unexpected combination of list length and tag: " <> show (n, tag))++-- FencePointerIndexType++instance Encode FencePointerIndexType where+ encode CompactIndex = encodeWord 0+ encode OrdinaryIndex = encodeWord 1++instance DecodeVersioned FencePointerIndexType where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure CompactIndex+ 1 -> pure OrdinaryIndex+ _ -> fail ("[FencePointerIndexType] Unexpected tag: " <> show tag)++-- DiskCachePolicy++instance Encode DiskCachePolicy where+ encode DiskCacheAll =+ encodeListLen 1+ <> encodeWord 0+ encode (DiskCacheLevelOneTo x) =+ encodeListLen 2+ <> encodeWord 1+ <> encodeInt x+ encode DiskCacheNone =+ encodeListLen 1+ <> encodeWord 2++instance DecodeVersioned DiskCachePolicy where+ decodeVersioned _v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (1, 0) -> pure DiskCacheAll+ (2, 1) -> DiskCacheLevelOneTo <$> decodeInt+ (1, 2) -> pure DiskCacheNone+ _ -> fail ("[DiskCachePolicy] Unexpected combination of list length and tag: " <> show (n, tag))++-- MergeSchedule++instance Encode MergeSchedule where+ encode OneShot = encodeWord 0+ encode Incremental = encodeWord 1++instance DecodeVersioned MergeSchedule where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure OneShot+ 1 -> pure Incremental+ _ -> fail ("[MergeSchedule] Unexpected tag: " <> show tag)++-- MergeBatchSize++instance Encode MergeBatchSize where+ encode (MergeBatchSize n) = encodeInt n++instance DecodeVersioned MergeBatchSize where+ decodeVersioned _v = MergeBatchSize <$> decodeInt++{-------------------------------------------------------------------------------+ Encoding and decoding: SnapLevels+-------------------------------------------------------------------------------}++-- SnapLevels++instance Encode r => Encode (SnapLevels r) where+ encode (SnapLevels levels) = encode levels++instance DecodeVersioned r => DecodeVersioned (SnapLevels r) where+ decodeVersioned v = SnapLevels <$> decodeVersioned v++-- SnapLevel++instance Encode r => Encode (SnapLevel r) where+ encode (SnapLevel incomingRuns residentRuns) =+ encodeListLen 2+ <> encode incomingRuns+ <> encode residentRuns+++instance DecodeVersioned r => DecodeVersioned (SnapLevel r) where+ decodeVersioned v = do+ _ <- decodeListLenOf 2+ SnapLevel <$> decodeVersioned v <*> decodeVersioned v++-- Vector++instance Encode r => Encode (V.Vector r) where+ encode = encodeVector++instance DecodeVersioned r => DecodeVersioned (V.Vector r) where+ decodeVersioned = decodeVector++-- RunNumber++instance Encode RunNumber where+ encode (RunNumber x) = encodeInt x++instance DecodeVersioned RunNumber where+ decodeVersioned _v = RunNumber <$> decodeInt++-- SnapIncomingRun++instance Encode r => Encode (SnapIncomingRun r) where+ encode (SnapIncomingMergingRun mpfl nd nc smrs) =+ encodeListLen 5+ <> encodeWord 0+ <> encode mpfl+ <> encode nd+ <> encode nc+ <> encode smrs+ encode (SnapIncomingSingleRun x) =+ encodeListLen 2+ <> encodeWord 1+ <> encode x++instance DecodeVersioned r => DecodeVersioned (SnapIncomingRun r) where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (5, 0) -> SnapIncomingMergingRun+ <$> decodeVersioned v <*> decodeVersioned v+ <*> decodeVersioned v <*> decodeVersioned v+ (2, 1) -> SnapIncomingSingleRun <$> decodeVersioned v+ _ -> fail ("[SnapIncomingRun] Unexpected combination of list length and tag: " <> show (n, tag))++-- MergePolicyForLevel++instance Encode MergePolicyForLevel where+ encode LevelTiering = encodeWord 0+ encode LevelLevelling = encodeWord 1++instance DecodeVersioned MergePolicyForLevel where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure LevelTiering+ 1 -> pure LevelLevelling+ _ -> fail ("[MergePolicyForLevel] Unexpected tag: " <> show tag)++-- SnapMergingRun++instance (Encode t, Encode r) => Encode (SnapMergingRun t r) where+ encode (SnapCompletedMerge md r) =+ encodeListLen 3+ <> encodeWord 0+ <> encode md+ <> encode r+ encode (SnapOngoingMerge rp mc rs mt) =+ encodeListLen 5+ <> encodeWord 1+ <> encode rp+ <> encode mc+ <> encode rs+ <> encode mt++instance (DecodeVersioned t, DecodeVersioned r) => DecodeVersioned (SnapMergingRun t r) where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (3, 0) -> SnapCompletedMerge <$> decodeVersioned v+ <*> decodeVersioned v+ (5, 1) -> SnapOngoingMerge <$> decodeVersioned v <*> decodeVersioned v+ <*> decodeVersioned v <*> decodeVersioned v+ _ -> fail ("[SnapMergingRun] Unexpected combination of list length and tag: " <> show (n, tag))++-- NominalDebt, NominalCredits, MergeDebt and MergeCredits++instance Encode NominalDebt where+ encode (NominalDebt x) = encodeInt x++instance DecodeVersioned NominalDebt where+ decodeVersioned _v = NominalDebt <$> decodeInt++instance Encode NominalCredits where+ encode (NominalCredits x) = encodeInt x++instance DecodeVersioned NominalCredits where+ decodeVersioned _v = NominalCredits <$> decodeInt++instance Encode MergeDebt where+ encode (MergeDebt (MergeCredits x)) = encodeInt x++instance DecodeVersioned MergeDebt where+ decodeVersioned _v = (MergeDebt . MergeCredits) <$> decodeInt++instance Encode MergeCredits where+ encode (MergeCredits x) = encodeInt x++instance DecodeVersioned MergeCredits where+ decodeVersioned _v = MergeCredits <$> decodeInt++-- MergeType++instance Encode MR.LevelMergeType where+ encode MR.MergeMidLevel = encodeWord 0+ encode MR.MergeLastLevel = encodeWord 1++instance DecodeVersioned MR.LevelMergeType where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 0 -> pure MR.MergeMidLevel+ 1 -> pure MR.MergeLastLevel+ _ -> fail ("[LevelMergeType] Unexpected tag: " <> show tag)++-- | We start the tags for these merge types at an offset. This way, if we+-- serialise @MR.MergeMidLevel :: MR.LevelMergeType@ as 0 and then accidentally+-- try deserialising it as a @MR.TreeMergeType@, that will fail.+--+-- However, 'MR.LevelMergeType' and 'MR.TreeMergeType' are only different+-- (overlapping) subsets of 'MR.MergeType'. In particular, 'MR.MergeLastLevel'+-- and 'MR.MergeLevel' are semantically the same. Encoding them as the same+-- number leaves the door open to relaxing the restrictions on which merge types+-- can occur where, e.g. decoding them as a general 'MR.MergeType', without+-- having to change the file format.+instance Encode MR.TreeMergeType where+ encode MR.MergeLevel = encodeWord 1+ encode MR.MergeUnion = encodeWord 2++instance DecodeVersioned MR.TreeMergeType where+ decodeVersioned _v = do+ tag <- decodeWord+ case tag of+ 1 -> pure MR.MergeLevel+ 2 -> pure MR.MergeUnion+ _ -> fail ("[TreeMergeType] Unexpected tag: " <> show tag)++{-------------------------------------------------------------------------------+ Encoding and decoding: SnapMergingTree+-------------------------------------------------------------------------------}++-- SnapMergingTree++instance Encode r => Encode (SnapMergingTree r) where+ encode (SnapMergingTree tState) = encode tState++instance DecodeVersioned r => DecodeVersioned (SnapMergingTree r) where+ decodeVersioned ver = SnapMergingTree <$> decodeVersioned ver++-- SnapMergingTreeState++instance Encode r => Encode (SnapMergingTreeState r) where+ encode (SnapCompletedTreeMerge x) =+ encodeListLen 2+ <> encodeWord 0+ <> encode x+ encode (SnapPendingTreeMerge x) =+ encodeListLen 2+ <> encodeWord 1+ <> encode x+ encode (SnapOngoingTreeMerge smrs) =+ encodeListLen 2+ <> encodeWord 2+ <> encode smrs++instance DecodeVersioned r => DecodeVersioned (SnapMergingTreeState r) where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (2, 0) -> SnapCompletedTreeMerge <$> decodeVersioned v+ (2, 1) -> SnapPendingTreeMerge <$> decodeVersioned v+ (2, 2) -> SnapOngoingTreeMerge <$> decodeVersioned v+ _ -> fail ("[SnapMergingTreeState] Unexpected combination of list length and tag: " <> show (n, tag))++-- SnapPendingMerge++instance Encode r => Encode (SnapPendingMerge r) where+ encode (SnapPendingLevelMerge pe mt) =+ encodeListLen 3+ <> encodeWord 0+ <> encodeList pe+ <> encodeMaybe mt+ encode (SnapPendingUnionMerge mts) =+ encodeListLen 2+ <> encodeWord 1+ <> encodeList mts++instance DecodeVersioned r => DecodeVersioned (SnapPendingMerge r) where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (3, 0) -> SnapPendingLevelMerge <$> decodeList v <*> decodeMaybe v+ (2, 1) -> SnapPendingUnionMerge <$> decodeList v+ _ -> fail ("[SnapPendingMerge] Unexpected combination of list length and tag: " <> show (n, tag))++-- SnapPreExistingRun++instance Encode r => Encode (SnapPreExistingRun r) where+ encode (SnapPreExistingRun x) =+ encodeListLen 2+ <> encodeWord 0+ <> encode x+ encode (SnapPreExistingMergingRun smrs) =+ encodeListLen 2+ <> encodeWord 1+ <> encode smrs++instance DecodeVersioned r => DecodeVersioned (SnapPreExistingRun r) where+ decodeVersioned v = do+ n <- decodeListLen+ tag <- decodeWord+ case (n, tag) of+ (2, 0) -> SnapPreExistingRun <$> decodeVersioned v+ (2, 1) -> SnapPreExistingMergingRun <$> decodeVersioned v+ _ -> fail ("[SnapPreExistingRun] Unexpected combination of list length and tag: " <> show (n, tag))++-- Utilities for encoding/decoding Maybe values++--Note: the Maybe encoding cannot be nested like (Maybe (Maybe a)), but it is+-- ok for fields of records.+encodeMaybe :: Encode a => Maybe a -> Encoding+encodeMaybe = \case+ Nothing -> encodeNull+ Just en -> encode en++decodeMaybe :: DecodeVersioned a => SnapshotVersion -> Decoder s (Maybe a)+decodeMaybe v = do+ tok <- peekTokenType+ case tok of+ TypeNull -> Nothing <$ decodeNull+ _ -> Just <$> decodeVersioned v++encodeList :: Encode a => [a] -> Encoding+encodeList xs =+ encodeListLen (fromIntegral (length xs))+ <> foldr (\x r -> encode x <> r) mempty xs++decodeList :: DecodeVersioned a => SnapshotVersion -> Decoder s [a]+decodeList v = do+ n <- decodeListLen+ decodeSequenceLenN (flip (:)) [] reverse n (decodeVersioned v)++encodeVector :: Encode a => V.Vector a -> Encoding+encodeVector xs =+ encodeListLen (fromIntegral (V.length xs))+ <> foldr (\x r -> encode x <> r) mempty xs++decodeVector :: DecodeVersioned a => SnapshotVersion -> Decoder s (V.Vector a)+decodeVector v = do+ n <- decodeListLen+ decodeSequenceLenN (flip (:)) [] (V.reverse . V.fromList)+ n (decodeVersioned v)
@@ -0,0 +1,60 @@+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE UnboxedTuples #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.StrictArray (+ StrictArray,+ vectorToStrictArray,+ indexStrictArray,+ sizeofStrictArray,+) where++import Data.Elevator (Strict (Strict))+import qualified Data.Vector as V+import GHC.Exts ((+#))+import qualified GHC.Exts as GHC+import GHC.ST (ST (ST), runST)+import Unsafe.Coerce (unsafeCoerce)++data StrictArray a = StrictArray !(GHC.Array# (Strict a))++vectorToStrictArray :: forall a. V.Vector a -> StrictArray a+vectorToStrictArray v =+ runST $ ST $ \s0 ->+ case GHC.newArray# (case V.length v of GHC.I# l# -> l#)+ (Strict (unsafeCoerce ())) s0 of+ -- initialise with (), will be overwritten.+ (# s1, a# #) ->+ case go a# 0# s1 of+ s2 -> case GHC.unsafeFreezeArray# a# s2 of+ (# s3, a'# #) -> (# s3, StrictArray a'# #)+ where+ go :: forall s.+ GHC.MutableArray# s (Strict a)+ -> GHC.Int#+ -> GHC.State# s+ -> GHC.State# s+ go a# i# s+ | GHC.I# i# < V.length v+ = let x = V.unsafeIndex v (GHC.I# i#)+ -- We have to use seq# here to force the array element to WHNF+ -- before putting it into the strict array. This should not be+ -- necessary. https://github.com/sgraf812/data-elevator/issues/4+ in case GHC.seq# x s of+ (# s', x' #) ->+ case GHC.writeArray# a# i# (Strict x') s' of+ s'' -> go a# (i# +# 1#) s''+ | otherwise = s++{-# INLINE indexStrictArray #-}+indexStrictArray :: StrictArray a -> Int -> a+indexStrictArray (StrictArray a#) (GHC.I# i#) =+ case GHC.indexArray# a# i# of+ (# Strict x #) -> x++{-# INLINE sizeofStrictArray #-}+sizeofStrictArray :: StrictArray a -> Int+sizeofStrictArray (StrictArray a#) =+ GHC.I# (GHC.sizeofArray# a#)
@@ -0,0 +1,231 @@+{-# OPTIONS_HADDOCK not-home #-}+{-# LANGUAGE DefaultSignatures #-}++module Database.LSMTree.Internal.Types (+ Salt,+ Session (..),+ Table (..),+ BlobRef (..),+ Cursor (..),+ ResolveValue (..),+ resolveCompatibility,+ resolveAssociativity,+ resolveValidOutput,+ ResolveViaSemigroup (..),+ ResolveAsFirst (..),+ ) where++import Control.DeepSeq (NFData (..), deepseq)+import Data.Kind (Type)+import Data.Semigroup (Sum)+import Data.Typeable+import Data.Void (Void)+import Data.Word (Word64)+import qualified Database.LSMTree.Internal.BlobRef as Unsafe+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import Database.LSMTree.Internal.Serialise.Class (SerialiseValue (..))+import qualified Database.LSMTree.Internal.Unsafe as Unsafe++{- |+The session salt is used to secure the hash operations in the Bloom filters.++The value of the salt must be kept secret.+Otherwise, there are no restrictions on the value.+-}+type Salt = Word64++{- |+A session stores context that is shared by multiple tables.++Each session is associated with one session directory where the files+containing table data are stored. Each session locks its session directory.+There can only be one active session for each session directory at a time.+If a database is must be accessed from multiple parts of a program,+one session should be opened and shared between those parts of the program.+Session directories cannot be shared between OS processes.++A session may contain multiple tables, which may each have a different configuration and different key, value, and BLOB types.+Furthermore, sessions may contain both [simple]("Database.LSMTree.Simple") and [full-featured]("Database.LSMTree") tables.+-}+type Session :: (Type -> Type) -> Type+data Session m+ = forall h.+ (Typeable h) =>+ Session !(Unsafe.Session m h)++instance NFData (Session m) where+ rnf :: Session m -> ()+ rnf (Session session) = rnf session++{- |+A table is a handle to an individual LSM-tree key\/value store with both in-memory and on-disk parts.++__Warning:__+Tables are ephemeral. Once you close a table, its data is lost forever.+To persist tables, use [snapshots]("Database.LSMTree#g:snapshots").+-}+type role Table nominal nominal nominal nominal++type Table :: (Type -> Type) -> Type -> Type -> Type -> Type+data Table m k v b+ = forall h.+ (Typeable h) =>+ Table !(Unsafe.Table m h)++instance NFData (Table m k v b) where+ rnf :: Table m k v b -> ()+ rnf (Table table) = rnf table++{- |+A blob reference is a reference to an on-disk blob.++__Warning:__ A blob reference is /not stable/. Any operation that modifies the table,+cursor, or session that corresponds to a blob reference may cause it to be invalidated.++The word \"blob\" in this type comes from the acronym Binary Large Object.+-}+type role BlobRef nominal nominal++type BlobRef :: (Type -> Type) -> Type -> Type+data BlobRef m b+ = forall h.+ (Typeable h) =>+ BlobRef !(Unsafe.WeakBlobRef m h)++instance Show (BlobRef m b) where+ showsPrec :: Int -> BlobRef m b -> ShowS+ showsPrec d (BlobRef b) = showsPrec d b++{- |+A cursor is a stable read-only iterator for a table.++A cursor iterates over the entries in a table following the order of the+serialised keys. After the cursor is created, updates to the referenced table+do not affect the cursor.++The name of this type references [database cursors](https://en.wikipedia.org/wiki/Cursor_(databases\)), not, e.g., text editor cursors.+-}+type role Cursor nominal nominal nominal nominal++type Cursor :: (Type -> Type) -> Type -> Type -> Type -> Type+data Cursor m k v b+ = forall h.+ (Typeable h) =>+ Cursor !(Unsafe.Cursor m h)++instance NFData (Cursor m k v b) where+ rnf :: Cursor m k v b -> ()+ rnf (Cursor cursor) = rnf cursor++--------------------------------------------------------------------------------+-- Monoidal value resolution+--------------------------------------------------------------------------------++{- |+An instance of @'ResolveValue' v@ specifies how to merge values when using+monoidal upsert.++The class has two functions.+The function 'resolve' acts on values, whereas the function 'resolveSerialised' acts on serialised values.+Each function has a default implementation in terms of the other and serialisation\/deserialisation.+The user is encouraged to implement 'resolveSerialised'.++Instances should satisfy the following:++[Compatibility]:+ The functions 'resolve' and 'resolveSerialised' should be compatible:++ prop> serialiseValue (resolve v1 v2) == resolveSerialised (Proxy @v) (serialiseValue v1) (serialiseValue v2)++[Associativity]:+ The function 'resolve' should be associative:++ prop> serialiseValue (v1 `resolve` (v2 `resolve` v3)) == serialiseValue ((v1 `resolve` v2) `resolve` v3)++[Valid Output]:+ The function 'resolveSerialised' should only return deserialisable values:++ prop> deserialiseValue (resolveSerialised (Proxy @v) rb1 rb2) `deepseq` True+-}+-- TODO(optimisation): Include a function that determines whether or not it is safe to remove and Update from the last level of an LSM-tree.+-- TODO(optimisation): Include a function @v -> RawBytes -> RawBytes@ that can be used to merged deserialised and serialised values.+-- This can be used when inserting values into the write buffer.+class ResolveValue v where+ {-# MINIMAL resolve | resolveSerialised #-}++ {- |+ Combine two values.+ -}+ resolve :: v -> v -> v++ default resolve :: SerialiseValue v => v -> v -> v+ resolve v1 v2 =+ deserialiseValue $+ resolveSerialised (Proxy :: Proxy v) (serialiseValue v1) (serialiseValue v2)++ {- |+ Combine two serialised values.++ The user may assume that the input bytes are valid and can be deserialised using 'deserialiseValue'.+ The inputs are only ever produced by 'serialiseValue' and 'resolveSerialised'.+ -}+ resolveSerialised :: Proxy v -> RawBytes -> RawBytes -> RawBytes+ default resolveSerialised :: SerialiseValue v => Proxy v -> RawBytes -> RawBytes -> RawBytes+ resolveSerialised Proxy rb1 rb2 =+ serialiseValue (resolve (deserialiseValue @v rb1) (deserialiseValue @v rb2))++{- |+Test the __Compatibility__ law for the 'ResolveValue' class.+-}+resolveCompatibility :: (SerialiseValue v, ResolveValue v) => v -> v -> Bool+resolveCompatibility (v1 :: v) v2 =+ serialiseValue (resolve v1 v2) == resolveSerialised (Proxy @v) (serialiseValue v1) (serialiseValue v2)++{- |+Test the __Associativity__ law for the 'ResolveValue' class.+-}+resolveAssociativity :: (SerialiseValue v, ResolveValue v) => v -> v -> v -> Bool+resolveAssociativity v1 v2 v3 =+ serialiseValue (v1 `resolve` (v2 `resolve` v3)) == serialiseValue ((v1 `resolve` v2) `resolve` v3)++{- |+Test the __Valid Output__ law for the 'ResolveValue' class.+-}+resolveValidOutput :: (SerialiseValue v, ResolveValue v, NFData v) => v -> v -> Bool+resolveValidOutput (v1 :: v) (v2 :: v) =+ deserialiseValue @v (resolveSerialised (Proxy @v) (serialiseValue v1) (serialiseValue v2)) `deepseq` True++{- |+Wrapper that provides an instance of 'ResolveValue' via the 'Semigroup'+instance of the underlying type.++prop> resolve (ResolveViaSemigroup v1) (ResolveViaSemigroup v2) = ResolveViaSemigroup (v1 <> v2)+-}+newtype ResolveViaSemigroup v = ResolveViaSemigroup v+ deriving stock (Eq, Show)+ deriving newtype (SerialiseValue)++instance (SerialiseValue v, Semigroup v) => ResolveValue (ResolveViaSemigroup v) where+ resolve :: ResolveViaSemigroup v -> ResolveViaSemigroup v -> ResolveViaSemigroup v+ resolve (ResolveViaSemigroup v1) (ResolveViaSemigroup v2) = ResolveViaSemigroup (v1 <> v2)++{- |+Wrapper that provides an instance of 'ResolveValue' such that 'Database.LSMTree.upsert' behaves as 'Database.LSMTree.insert'.++The name 'ResolveAsFirst' is in reference to the wrapper 'Data.Semigroup.First' from "Data.Semigroup".++prop> resolve = const+-}+newtype ResolveAsFirst v = ResolveAsFirst {unResolveAsFirst:: v}+ deriving stock (Eq, Show)+ deriving newtype (SerialiseValue)++instance ResolveValue (ResolveAsFirst v) where+ resolve :: ResolveAsFirst v -> ResolveAsFirst v -> ResolveAsFirst v+ resolve = const+ resolveSerialised :: Proxy (ResolveAsFirst v) -> RawBytes -> RawBytes -> RawBytes+ resolveSerialised _p = const++deriving via ResolveViaSemigroup Void instance ResolveValue Void++deriving via (ResolveViaSemigroup (Sum v)) instance (Num v, SerialiseValue v) => ResolveValue (Sum v)
@@ -0,0 +1,51 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.UniqCounter (+ UniqCounter (..),+ newUniqCounter,+ incrUniqCounter,+ Unique,+ uniqueToInt,+ uniqueToRunNumber,+ uniqueToTableId,+ uniqueToCursorId,+) where++import Control.DeepSeq (NFData (..))+import Control.Monad.Primitive (PrimMonad, PrimState)+import Data.Primitive.PrimVar as P+import Database.LSMTree.Internal.RunNumber++-- | A unique value derived from a 'UniqCounter'.+newtype Unique = Unique Int++-- | Use specialised versions like 'uniqueToRunNumber' where possible.+uniqueToInt :: Unique -> Int+uniqueToInt (Unique n) = n++uniqueToRunNumber :: Unique -> RunNumber+uniqueToRunNumber (Unique n) = RunNumber n++uniqueToTableId :: Unique -> TableId+uniqueToTableId (Unique n) = TableId n++uniqueToCursorId :: Unique -> CursorId+uniqueToCursorId (Unique n) = CursorId n++-- | An atomic counter for producing 'Unique' values.+--+newtype UniqCounter m = UniqCounter (PrimVar (PrimState m) Int)++instance NFData (UniqCounter m) where+ rnf (UniqCounter (P.PrimVar mba)) = rnf mba++{-# INLINE newUniqCounter #-}+newUniqCounter :: PrimMonad m => Int -> m (UniqCounter m)+newUniqCounter = fmap UniqCounter . P.newPrimVar++{-# INLINE incrUniqCounter #-}+-- | Atomically, return the current state of the counter, and increment the+-- counter.+incrUniqCounter :: PrimMonad m => UniqCounter m -> m Unique+incrUniqCounter (UniqCounter uniqVar) =+ Unique <$> P.fetchAddInt uniqVar 1
@@ -0,0 +1,2229 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE DataKinds #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | This module brings together the internal parts to provide an API in terms+-- of untyped serialised keys, values and blobs.+--+-- Apart from defining the API, this module mainly deals with concurrency- and+-- exception-safe opening and closing of resources. Any other non-trivial logic+-- should live somewhere else.+--+module Database.LSMTree.Internal.Unsafe (+ -- * Exceptions+ SessionDirDoesNotExistError (..)+ , SessionDirLockedError (..)+ , SessionDirCorruptedError (..)+ , SessionClosedError (..)+ , TableClosedError (..)+ , TableCorruptedError (..)+ , TableTooLargeError (..)+ , TableUnionNotCompatibleError (..)+ , SnapshotExistsError (..)+ , SnapshotDoesNotExistError (..)+ , SnapshotCorruptedError (..)+ , SnapshotNotCompatibleError (..)+ , BlobRefInvalidError (..)+ , CursorClosedError (..)+ , FileFormat (..)+ , FileCorruptedError (..)+ , Paths.InvalidSnapshotNameError (..)+ -- * Tracing+ -- $traces+ , LSMTreeTrace (..)+ , SessionId (..)+ , SessionTrace (..)+ , TableTrace (..)+ , CursorTrace (..)+ -- * Session+ , Session (..)+ , SessionState (..)+ , SessionEnv (..)+ , withKeepSessionOpen+ -- ** Implementation of public API+ , withOpenSession+ , withNewSession+ , withRestoreSession+ , openSession+ , newSession+ , restoreSession+ , closeSession+ -- * Table+ , Table (..)+ , TableState (..)+ , TableEnv (..)+ , withKeepTableOpen+ -- ** Implementation of public API+ , ResolveSerialisedValue+ , withTable+ , new+ , close+ , lookups+ , rangeLookup+ , updates+ , retrieveBlobs+ -- ** Cursor API+ , Cursor (..)+ , CursorState (..)+ , CursorEnv (..)+ , OffsetKey (..)+ , withCursor+ , newCursor+ , closeCursor+ , readCursor+ , readCursorWhile+ -- * Snapshots+ , SnapshotLabel+ , saveSnapshot+ , openTableFromSnapshot+ , deleteSnapshot+ , doesSnapshotExist+ , listSnapshots+ -- * Multiple writable tables+ , duplicate+ -- * Table union+ , unions+ , UnionDebt (..)+ , remainingUnionDebt+ , UnionCredits (..)+ , supplyUnionCredits+ ) where++import qualified Codec.Serialise as S+import Control.ActionRegistry+import Control.Concurrent.Class.MonadMVar.Strict+import Control.Concurrent.Class.MonadSTM (MonadSTM (..))+import Control.Concurrent.Class.MonadSTM.RWVar (RWVar)+import qualified Control.Concurrent.Class.MonadSTM.RWVar as RW+import Control.DeepSeq+import Control.Monad (forM, unless, void, when, (<$!>))+import Control.Monad.Class.MonadAsync as Async+import Control.Monad.Class.MonadST (MonadST (..))+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive+import Control.RefCount+import Control.Tracer+import qualified Data.BloomFilter.Hash as Bloom+import Data.Either (fromRight)+import Data.Foldable+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (catMaybes, maybeToList)+import qualified Data.Set as Set+import Data.Text (Text)+import qualified Data.Text as Text+import Data.Typeable+import qualified Data.Vector as V+import Database.LSMTree.Internal.Arena (ArenaManager, newArenaManager)+import Database.LSMTree.Internal.BlobRef (WeakBlobRef (..))+import qualified Database.LSMTree.Internal.BlobRef as BlobRef+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.Config.Override (TableConfigOverride,+ overrideTableConfig)+import Database.LSMTree.Internal.CRC32C (FileCorruptedError (..),+ FileFormat (..))+import qualified Database.LSMTree.Internal.Cursor as Cursor+import Database.LSMTree.Internal.Entry (Entry)+import Database.LSMTree.Internal.IncomingRun (IncomingRun (..))+import Database.LSMTree.Internal.Lookup (TableCorruptedError (..),+ lookupsIO, lookupsIOWithWriteBuffer)+import Database.LSMTree.Internal.MergeSchedule+import Database.LSMTree.Internal.MergingRun (TableTooLargeError (..))+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.MergingTree+import qualified Database.LSMTree.Internal.MergingTree as MT+import qualified Database.LSMTree.Internal.MergingTree.Lookup as MT+import Database.LSMTree.Internal.Paths (SessionRoot (..),+ SnapshotMetaDataChecksumFile (..),+ SnapshotMetaDataFile (..), SnapshotName)+import qualified Database.LSMTree.Internal.Paths as Paths+import Database.LSMTree.Internal.Range (Range (..))+import Database.LSMTree.Internal.RawBytes (RawBytes)+import Database.LSMTree.Internal.Readers (OffsetKey (..))+import qualified Database.LSMTree.Internal.Readers as Readers+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue,+ SerialisedBlob (..), SerialisedKey (SerialisedKey),+ SerialisedValue)+import Database.LSMTree.Internal.Snapshot+import Database.LSMTree.Internal.Snapshot.Codec+import Database.LSMTree.Internal.UniqCounter+import qualified Database.LSMTree.Internal.Vector as V+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import qualified System.FS.API as FS+import System.FS.API (FsError, FsErrorPath (..), FsPath, HasFS)+import qualified System.FS.API.Lazy as FS+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++{-------------------------------------------------------------------------------+ Traces+-------------------------------------------------------------------------------}++{- $traces++ Trace messages are divided into three categories for each type of resource:+ sessions, tables, and cursors. The trace messages are structured so that:++ 1. Resources have (unique) identifiers, and these are included in each+ message.++ 2. Operations that modify, create, or close resources trace a start and end+ message. The reasoning here is that it's useful for troubleshooting+ purposes to know not only that an operation started but also that it ended.+ To an extent this would also be useful for read-only operations like+ lookups, but since read-only operations do not modify resources, we've left+ out end messages for those. They could be added if asked for by users.++ 3. When we are in the process of creating a new resource from existing+ resources, then the /child/ resource traces the identifier(s) of its+ /parent/ resource(s).++ These properties ensure that troubleshooting using tracers in a concurrent+ setting is possible. Messages can be interleaved, so it's important to find+ meaningful pairings of messages, like 'TraceCloseTable' and+ 'TraceClosedTable'.+-}++data LSMTreeTrace =+ -- | Trace messages related to sessions.+ TraceSession+ -- | Sessions are identified by the path to their root directory.+ SessionId+ SessionTrace++ -- | Trace messages related to tables.+ | TraceTable+ -- | Tables are identified by a unique number.+ TableId+ TableTrace++ -- | Trace messages related to cursors.+ | TraceCursor+ -- | Cursors are identified by a unique number.+ CursorId+ CursorTrace+ deriving stock (Show, Eq)++-- | Sessions are identified by the path to their root directory.+newtype SessionId = SessionId FsPath+ deriving stock (Eq, Show)++-- | Trace messages related to sessions.+data SessionTrace =+ -- | We are opening a session in the requested session directory. A+ -- 'TraceNewSession'\/'TraceRestoreSession' and 'TraceCreatedSession'+ -- message should follow if succesful.+ TraceOpenSession+ -- | We are creating a new, fresh session. A 'TraceCreatedSession' should+ -- follow if succesful.+ | TraceNewSession+ -- | We are restoring a session from the directory contents. A+ -- 'TraceCreatedSession' shold follow if succesful.+ | TraceRestoreSession+ -- | A session has been successfully created.+ | TraceCreatedSession++ -- | We are closing the session. A 'TraceClosedSession' message should+ -- follow if succesful.+ | TraceCloseSession+ -- | Closing the session was successful.+ | TraceClosedSession++ -- | We are deleting the snapshot with the given name. A+ -- 'TraceDeletedSnapshot' message should follow if succesful.+ | TraceDeleteSnapshot SnapshotName+ -- | We have successfully deleted the snapshot with the given name.+ | TraceDeletedSnapshot SnapshotName+ -- | We are listing snapshots.+ | TraceListSnapshots++ -- | We are retrieving blobs.+ | TraceRetrieveBlobs Int+ deriving stock (Show, Eq)++-- | Trace messages related to tables.+data TableTrace =+ -- | A table has been successfully created with the specified config.+ TraceCreatedTable+ -- | The parent session+ SessionId+ TableConfig+ -- | We are creating a new, fresh table with the specified config. A+ -- 'TraceCreatedTable' message should follow if successful.+ | TraceNewTable TableConfig++ -- | We are closing the table. A 'TraceClosedTable' message should follow if+ -- successful.+ | TraceCloseTable+ -- | Closing the table was succesful.+ | TraceClosedTable++ -- | We are performing a batch of lookups.+ | TraceLookups+ -- | The number of keys that are looked up.+ Int+ -- | We are performing a range lookup.+ | TraceRangeLookup (Range SerialisedKey)+ -- | We are performing a batch of updates.+ | TraceUpdates+ -- | The number of keys that will be updated.+ Int+ -- | We have successfully performed a batch of updates.+ | TraceUpdated+ -- | The number of keys that have been updated.+ Int++ -- | We are opening a table from a snapshot. A 'TraceCreatedTable' message+ -- should follow if successful.+ | TraceOpenTableFromSnapshot SnapshotName TableConfigOverride+ -- | We are saving a snapshot with the given name. A 'TraceSavedSnapshot'+ -- message should follow if successful.+ | TraceSaveSnapshot SnapshotName+ -- | A snapshot with the given name was saved successfully.+ | TraceSavedSnapshot SnapshotName++ -- | We are creating a duplicate of a table. A 'TraceCreatedTable' message+ -- should follow if successful.+ | TraceDuplicate+ -- | The parent table+ TableId++ -- | We are creating an incremental union of a list of tables. A+ -- 'TraceCreatedTable' message should follow if successful.+ | TraceIncrementalUnions+ -- | Unique identifiers for the tables that are used as inputs to the+ -- union.+ (NonEmpty TableId)+ -- | We are querying the remaining union debt.+ | TraceRemainingUnionDebt+ -- | We are supplying the given number of union credits.+ | TraceSupplyUnionCredits UnionCredits+ -- | We have supplied union credits.+ | TraceSuppliedUnionCredits+ -- | The number of input credits that have been supplied.+ UnionCredits+ -- | Leftover credits.+ UnionCredits++#ifdef DEBUG_TRACES+ -- | INTERNAL: debug traces for the merge schedule+ | TraceMerge (AtLevel MergeTrace)+#endif+ deriving stock (Show, Eq)++contramapTraceMerge :: forall m. Monad m => Tracer m TableTrace -> Tracer m (AtLevel MergeTrace)+#if MIN_VERSION_contra_tracer(0,2,0)+#ifdef DEBUG_TRACES+contramapTraceMerge t = TraceMerge `contramap` t+#else+contramapTraceMerge t = traceMaybe (const Nothing) t+#endif+#else+contramapTraceMerge _t = nullTracer+ where+ -- See #766+ _unused = pure @m ()+#endif++-- | Trace messages related to cursors.+data CursorTrace =+ -- | A cursor has been successfully created.+ TraceCreatedCursor+ -- | The parent session+ SessionId+ -- | We are creating a new, fresh cursor positioned at the given offset key.+ -- A 'TraceCreatedCursor' message should follow if successful.+ | TraceNewCursor+ -- | The parent table+ TableId+ -- | The optional serialised key that is used as the initial offset for+ -- the cursor.+ (Maybe RawBytes)++ -- | We are closing the cursor. A 'TraceClosedCursor' message should follow+ -- if successful.+ | TraceCloseCursor+ -- | Closing the cursor was succesful.+ | TraceClosedCursor++ -- | We are reading from the cursor. A 'TraceReadCursor' message should+ -- follow if successful.+ | TraceReadingCursor+ -- | Requested number of entries to read at most.+ Int+ -- | We have succesfully read from the cursor.+ | TraceReadCursor+ -- | Requested number of entries to read at most.+ Int+ -- | Actual number of entries read.+ Int+ deriving stock (Show, Eq)++{-------------------------------------------------------------------------------+ Session+-------------------------------------------------------------------------------}++-- | A session provides context that is shared across multiple tables.+--+-- For more information, see 'Database.LSMTree.Internal.Types.Session'.+data Session m h = Session {+ -- | The primary purpose of this 'RWVar' is to ensure consistent views of+ -- the open-/closedness of a session when multiple threads require access+ -- to the session's fields (see 'withKeepSessionOpen'). We use more+ -- fine-grained synchronisation for various mutable parts of an open+ -- session.+ --+ -- INVARIANT: once the session state is changed from 'SessionOpen' to+ -- 'SessionClosed', it is never changed back to 'SessionOpen' again.+ sessionState :: !(RWVar m (SessionState m h))+ , lsmTreeTracer :: !(Tracer m LSMTreeTrace)+ , sessionTracer :: !(Tracer m SessionTrace)++ -- | A session-wide shared, atomic counter that is used to produce unique+ -- names, for example: run names, table IDs.+ --+ -- This counter lives in 'Session' instead of 'SessionEnv' because it is an+ -- atomic counter, and so it does not need to be guarded by a lock.+ , sessionUniqCounter :: !(UniqCounter m)+ }++instance NFData (Session m h) where+ rnf (Session a b c d) = rnf a `seq` rwhnf b `seq` rwhnf c `seq` rnf d++data SessionState m h =+ SessionOpen !(SessionEnv m h)+ | SessionClosed++data SessionEnv m h = SessionEnv {+ -- | The path to the directory in which this session is live. This is a path+ -- relative to root of the 'HasFS' instance.+ --+ -- INVARIANT: the session root is never changed during the lifetime of a+ -- session.+ sessionRoot :: !SessionRoot+ -- | Session-wide salt for bloomfilter hashes+ --+ -- INVARIANT: all bloom filters in all tables in the session are created+ -- using the same salt, and all bloom filter are queried using that same+ -- salt.+ , sessionSalt :: !Bloom.Salt+ , sessionHasFS :: !(HasFS m h)+ , sessionHasBlockIO :: !(HasBlockIO m h)+ , sessionLockFile :: !(FS.LockFileHandle m)+ -- | Open tables are tracked here so they can be closed once the session is+ -- closed. Tables also become untracked when they are closed manually.+ --+ -- Tables are assigned unique identifiers using 'sessionUniqCounter' to+ -- ensure that modifications to the set of known tables are independent.+ -- Each identifier is added only once in 'new', 'openTableFromSnapshot', 'duplicate',+ -- 'union', or 'unions', and is deleted only once in 'close' or+ -- 'closeSession'.+ --+ -- * A new table may only insert its own identifier when it has acquired the+ -- 'sessionState' read-lock. This is to prevent races with 'closeSession'.+ --+ -- * A table 'close' may delete its own identifier from the set of open+ -- tables without restrictions, even concurrently with 'closeSession'.+ -- This is safe because 'close' is idempotent'.+ , sessionOpenTables :: !(StrictMVar m (Map TableId (Table m h)))+ -- | Similarly to tables, open cursors are tracked so they can be closed+ -- once the session is closed. See 'sessionOpenTables'.+ , sessionOpenCursors :: !(StrictMVar m (Map CursorId (Cursor m h)))+ }++{-# INLINE sessionId #-}+sessionId :: SessionEnv m h -> SessionId+sessionId = SessionId . getSessionRoot . sessionRoot++-- | The session is closed.+data SessionClosedError+ = ErrSessionClosed+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# INLINE withKeepSessionOpen #-}+{-# SPECIALISE withKeepSessionOpen ::+ Session IO h+ -> (SessionEnv IO h -> IO a)+ -> IO a #-}+-- | 'withKeepSessionOpen' ensures that the session stays open for the duration of+-- the provided continuation.+--+-- NOTE: any operation except 'sessionClose' can use this function.+withKeepSessionOpen ::+ (MonadSTM m, MonadThrow m)+ => Session m h+ -> (SessionEnv m h -> m a)+ -> m a+withKeepSessionOpen sesh action = RW.withReadAccess (sessionState sesh) $ \case+ SessionClosed -> throwIO ErrSessionClosed+ SessionOpen seshEnv -> action seshEnv++--+-- Implementation of public API+--++-- | The session directory does not exist.+data SessionDirDoesNotExistError+ = ErrSessionDirDoesNotExist !FsErrorPath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | The session directory is locked by another active session.+data SessionDirLockedError+ = ErrSessionDirLocked !FsErrorPath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | The session directory is corrupted, e.g., it misses required files or contains unexpected files.+data SessionDirCorruptedError+ = ErrSessionDirCorrupted !Text !FsErrorPath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# INLINE withOpenSession #-}+withOpenSession ::+ forall m h a.+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m, MonadEvaluate m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> FsPath -- ^ Path to the session directory+ -> (Session m h -> m a)+ -> m a+withOpenSession tr hfs hbio salt dir k = do+ bracket+ (openSession tr hfs hbio salt dir)+ closeSession+ k++{-# INLINE withNewSession #-}+withNewSession ::+ forall m h a.+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> FsPath -- ^ Path to the session directory+ -> (Session m h -> m a)+ -> m a+withNewSession tr hfs hbio salt dir k = do+ bracket+ (newSession tr hfs hbio salt dir)+ closeSession+ k++{-# INLINE withRestoreSession #-}+withRestoreSession ::+ forall m h a.+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m, MonadEvaluate m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> FsPath -- ^ Path to the session directory+ -> (Session m h -> m a)+ -> m a+withRestoreSession tr hfs hbio dir k = do+ bracket+ (restoreSession tr hfs hbio dir)+ closeSession+ k++{-# SPECIALISE openSession ::+ Tracer IO LSMTreeTrace+ -> HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> FsPath+ -> IO (Session IO h) #-}+-- | See 'Database.LSMTree.openSession'.+openSession ::+ forall m h.+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m, MonadEvaluate m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> FsPath -- ^ Path to the session directory+ -> m (Session m h)+openSession tr hfs hbio salt dir = do+ traceWith sessionTracer TraceOpenSession++ -- This is checked by 'newSession' and 'restoreSession' too, but it does not+ -- hurt to check it twice, and it's arguably simpler like this.+ dirExists <- FS.doesDirectoryExist hfs dir+ unless dirExists $+ throwIO (ErrSessionDirDoesNotExist (FS.mkFsErrorPath hfs dir))++ b <- isSessionDirEmpty hfs dir+ if b then+ newSession tr hfs hbio salt dir+ else+ restoreSession tr hfs hbio dir+ where+ sessionTracer = TraceSession (SessionId dir) `contramap` tr++{-# SPECIALISE newSession ::+ Tracer IO LSMTreeTrace+ -> HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> FsPath+ -> IO (Session IO h) #-}+-- | See 'Database.LSMTree.newSession'.+newSession ::+ forall m h.+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> Bloom.Salt+ -> FsPath -- ^ Path to the session directory+ -> m (Session m h)+newSession tr hfs hbio salt dir = do+ traceWith sessionTracer TraceNewSession+ -- We can not use modifyWithActionRegistry here, since there is no in-memory+ -- state to modify. We use withActionRegistry instead, which may have a tiny+ -- chance of leaking resources if openSession is not called in a masked+ -- state.+ withActionRegistry $ \reg -> do+ dirExists <- FS.doesDirectoryExist hfs dir+ unless dirExists $+ throwIO (ErrSessionDirDoesNotExist (FS.mkFsErrorPath hfs dir))++ -- Try to acquire the session file lock as soon as possible to reduce the+ -- risk of race conditions.+ --+ -- The lock is only released when an exception is raised, otherwise the+ -- lock is included in the returned Session.+ sessionFileLock <- acquireSessionLock hfs hbio reg lockFilePath++ -- If we're starting a new session, then the session directory should be+ -- non-empty.+ b <- isSessionDirEmpty hfs dir+ unless b $ do+ throwIO $ ErrSessionDirCorrupted+ (Text.pack "Session directory is non-empty")+ (FS.mkFsErrorPath hfs dir)++ withRollback_ reg+ (FS.withFile hfs metadataFilePath (FS.WriteMode FS.MustBeNew) $ \h ->+ void $ FS.hPutAll hfs h $ S.serialise salt)+ (FS.removeFile hfs metadataFilePath)+ withRollback_ reg+ (FS.createDirectory hfs activeDirPath)+ (FS.removeDirectoryRecursive hfs activeDirPath)+ withRollback_ reg+ (FS.createDirectory hfs snapshotsDirPath)+ (FS.removeDirectoryRecursive hfs snapshotsDirPath)++ mkSession tr hfs hbio reg root sessionFileLock salt+ where+ sessionTracer = TraceSession (SessionId dir) `contramap` tr++ root = Paths.SessionRoot dir+ lockFilePath = Paths.lockFile root+ metadataFilePath = Paths.metadataFile root+ activeDirPath = Paths.getActiveDir (Paths.activeDir root)+ snapshotsDirPath = Paths.snapshotsDir root++{-# SPECIALISE restoreSession ::+ Tracer IO LSMTreeTrace+ -> HasFS IO h+ -> HasBlockIO IO h+ -> FsPath+ -> IO (Session IO h) #-}+-- | See 'Database.LSMTree.restoreSession'.+restoreSession ::+ forall m h.+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m, MonadEvaluate m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> FsPath -- ^ Path to the session directory+ -> m (Session m h)+restoreSession tr hfs hbio dir = do+ traceWith sessionTracer TraceRestoreSession++ -- We can not use modifyWithActionRegistry here, since there is no in-memory+ -- state to modify. We use withActionRegistry instead, which may have a tiny+ -- chance of leaking resources if openSession is not called in a masked+ -- state.+ withActionRegistry $ \reg -> do+ dirExists <- FS.doesDirectoryExist hfs dir+ unless dirExists $+ throwIO (ErrSessionDirDoesNotExist (FS.mkFsErrorPath hfs dir))++ -- Try to acquire the session file lock as soon as possible to reduce the+ -- risk of race conditions.+ --+ -- The lock is only released when an exception is raised, otherwise the+ -- lock is included in the returned Session.+ sessionFileLock <- acquireSessionLock hfs hbio reg lockFilePath++ -- If we're restoring a session, then the session directory should be+ -- non-empty.+ b <- isSessionDirEmpty hfs dir+ when b $ do+ throwIO $ ErrSessionDirCorrupted+ (Text.pack "Session directory is empty")+ (FS.mkFsErrorPath hfs dir)++ -- If the layouts are wrong, we throw an exception+ checkTopLevelDirLayout++ salt <-+ FS.withFile hfs metadataFilePath FS.ReadMode $ \h -> do+ bs <- FS.hGetAll hfs h+ evaluate $ S.deserialise bs++ -- Clear the active directory by removing the directory and recreating+ -- it again.+ FS.removeDirectoryRecursive hfs activeDirPath+ `finally` FS.createDirectoryIfMissing hfs False activeDirPath++ checkActiveDirLayout+ checkSnapshotsDirLayout++ mkSession tr hfs hbio reg root sessionFileLock salt+ where+ sessionTracer = TraceSession (SessionId dir) `contramap` tr++ root = Paths.SessionRoot dir+ lockFilePath = Paths.lockFile root+ metadataFilePath = Paths.metadataFile root+ activeDirPath = Paths.getActiveDir (Paths.activeDir root)+ snapshotsDirPath = Paths.snapshotsDir root++ -- Check that the active directory and snapshots directory exist. We assume+ -- the lock file already exists at this point.+ --+ -- This checks only that the /expected/ files and directories exist.+ -- Unexpected files in the top-level directory are ignored for the layout+ -- check.+ checkTopLevelDirLayout = do+ FS.doesFileExist hfs metadataFilePath >>= \b ->+ unless b $ throwIO $+ ErrSessionDirCorrupted+ (Text.pack "Missing metadata file")+ (FS.mkFsErrorPath hfs metadataFilePath)+ FS.doesDirectoryExist hfs activeDirPath >>= \b ->+ unless b $ throwIO $+ ErrSessionDirCorrupted+ (Text.pack "Missing active directory")+ (FS.mkFsErrorPath hfs activeDirPath)+ FS.doesDirectoryExist hfs snapshotsDirPath >>= \b ->+ unless b $ throwIO $+ ErrSessionDirCorrupted+ (Text.pack "Missing snapshot directory")+ (FS.mkFsErrorPath hfs snapshotsDirPath)++ -- The active directory should be empty+ checkActiveDirLayout = do+ contents <- FS.listDirectory hfs activeDirPath+ unless (Set.null contents) $ throwIO $+ ErrSessionDirCorrupted+ (Text.pack "Active directory is non-empty")+ (FS.mkFsErrorPath hfs activeDirPath)++ -- Nothing to check: snapshots are verified when they are loaded, not when a+ -- session is restored.+ checkSnapshotsDirLayout = pure ()++{-# SPECIALISE closeSession :: Session IO h -> IO () #-}+-- | See 'Database.LSMTree.closeSession'.+--+-- A session's global resources will only be released once it is sure that no+-- tables or cursors are open anymore.+closeSession ::+ (MonadMask m, MonadSTM m, MonadMVar m, PrimMonad m)+ => Session m h+ -> m ()+closeSession Session{sessionState, sessionTracer} = do+ traceWith sessionTracer TraceCloseSession+ modifyWithActionRegistry_+ (RW.unsafeAcquireWriteAccess sessionState)+ (atomically . RW.unsafeReleaseWriteAccess sessionState)+ $ \reg -> \case+ SessionClosed -> pure SessionClosed+ SessionOpen seshEnv -> do+ -- Close tables and cursors first, so that we know none are open when we+ -- release session-wide resources.+ --+ -- If any tables or cursors have been closed already by a different+ -- thread, then the idempotent close functions will act like a no-op,+ -- and so we are not in trouble.+ --+ -- Since we have a write lock on the session state, we know that no+ -- tables or cursors will be added while we are closing the session+ -- (see sessionOpenTables), and that we are the only thread currently+ -- closing the session. .+ --+ -- We technically don't have to overwrite this with an empty Map, but+ -- why not.++ -- close cursors+ cursors <-+ withRollback reg+ (swapMVar (sessionOpenCursors seshEnv) Map.empty)+ (void . swapMVar (sessionOpenCursors seshEnv))+ mapM_ (delayedCommit reg . closeCursor) cursors++ -- close tables+ tables <-+ withRollback reg+ (swapMVar (sessionOpenTables seshEnv) Map.empty)+ (void . swapMVar (sessionOpenTables seshEnv))+ mapM_ (delayedCommit reg . close) tables++ delayedCommit reg $ FS.hUnlock (sessionLockFile seshEnv)++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $ traceWith sessionTracer TraceClosedSession++ pure SessionClosed++{-# SPECIALISE acquireSessionLock ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ActionRegistry IO+ -> FsPath+ -> IO (FS.LockFileHandle IO) #-}+acquireSessionLock ::+ forall m h. (MonadSTM m, PrimMonad m, MonadMask m)+ => HasFS m h+ -> HasBlockIO m h+ -> ActionRegistry m+ -> FsPath+ -> m (FS.LockFileHandle m)+acquireSessionLock hfs hbio reg lockFilePath = do+ elock <-+ withRollbackFun reg+ (fromRight Nothing)+ acquireLock+ releaseLock++ case elock of+ Left e+ | FS.FsResourceAlreadyInUse <- FS.fsErrorType e+ , fsep@(FsErrorPath _ fsp) <- FS.fsErrorPath e+ , fsp == lockFilePath+ -> throwIO (ErrSessionDirLocked fsep)+ Left e -> throwIO e -- rethrow unexpected errors+ Right Nothing -> throwIO (ErrSessionDirLocked (FS.mkFsErrorPath hfs lockFilePath))+ Right (Just sessionFileLock) -> pure sessionFileLock+ where+ acquireLock = try @m @FsError $ FS.tryLockFile hbio lockFilePath FS.ExclusiveLock++ releaseLock = FS.hUnlock++{-# SPECIALISE mkSession ::+ Tracer IO LSMTreeTrace+ -> HasFS IO h+ -> HasBlockIO IO h+ -> ActionRegistry IO+ -> SessionRoot+ -> FS.LockFileHandle IO+ -> Bloom.Salt+ -> IO (Session IO h) #-}+mkSession ::+ (PrimMonad m, MonadMVar m, MonadSTM m)+ => Tracer m LSMTreeTrace+ -> HasFS m h+ -> HasBlockIO m h+ -> ActionRegistry m+ -> SessionRoot+ -> FS.LockFileHandle m+ -> Bloom.Salt+ -> m (Session m h)+mkSession tr hfs hbio reg root@(SessionRoot dir) lockFile salt = do+ counterVar <- newUniqCounter 0+ openTablesVar <- newMVar Map.empty+ openCursorsVar <- newMVar Map.empty+ sessionVar <- RW.new $ SessionOpen $ SessionEnv {+ sessionRoot = root+ , sessionSalt = salt+ , sessionHasFS = hfs+ , sessionHasBlockIO = hbio+ , sessionLockFile = lockFile+ , sessionOpenTables = openTablesVar+ , sessionOpenCursors = openCursorsVar+ }++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith sessionTracer TraceCreatedSession++ pure $! Session {+ sessionState = sessionVar+ , lsmTreeTracer = tr+ , sessionTracer = sessionTracer+ , sessionUniqCounter = counterVar+ }+ where+ sessionTracer = TraceSession (SessionId dir) `contramap` tr++{-# INLINE isSessionDirEmpty #-}+isSessionDirEmpty :: Monad m => HasFS m h -> FsPath -> m Bool+isSessionDirEmpty hfs dir = do+ dirContents <- FS.listDirectory hfs dir+ pure $ Set.null dirContents || dirContents == Set.singleton Paths.lockFileName++{-------------------------------------------------------------------------------+ Table+-------------------------------------------------------------------------------}++-- | A handle to an on-disk key\/value table.+--+-- For more information, see 'Database.LSMTree.Table'.+data Table m h = Table {+ tableConfig :: !TableConfig+ -- | The primary purpose of this 'RWVar' is to ensure consistent views of+ -- the open-/closedness of a table when multiple threads require access to+ -- the table's fields (see 'withKeepTableOpen'). We use more fine-grained+ -- synchronisation for various mutable parts of an open table.+ , tableState :: !(RWVar m (TableState m h))+ , tableArenaManager :: !(ArenaManager (PrimState m))+ , tableTracer :: !(Tracer m TableTrace)+ -- | Session-unique identifier for this table.+ --+ -- INVARIANT: a table's identifier is never changed during the lifetime of+ -- the table.+ , tableId :: !TableId++ -- === Session-inherited++ -- | The session that this table belongs to.+ --+ -- INVARIANT: a table only ever belongs to one session, and can't be+ -- transferred to a different session.+ , tableSession :: !(Session m h)+ }++instance NFData (Table m h) where+ rnf (Table a b c d e f) =+ rnf a `seq` rnf b `seq` rnf c `seq` rwhnf d `seq` rnf e`seq` rwhnf f++-- | A table may assume that its corresponding session is still open as+-- long as the table is open. A session's global resources, and therefore+-- resources that are inherited by the table, will only be released once the+-- session is sure that no tables are open anymore.+data TableState m h =+ TableOpen !(TableEnv m h)+ | TableClosed++data TableEnv m h = TableEnv {+ -- === Session-inherited++ -- | Use this instead of 'tableSession' for easy access. An open table may+ -- assume that its session is open.+ tableSessionEnv :: !(SessionEnv m h)++ -- === Table-specific++ -- | All of the state being in a single 'StrictMVar' is a relatively simple+ -- solution, but there could be more concurrency. For example, while inserts+ -- are in progress, lookups could still look at the old state without+ -- waiting for the MVar.+ --+ -- TODO: switch to more fine-grained synchronisation approach+ , tableContent :: !(RWVar m (TableContent m h))+ }++{-# INLINE tableSessionRoot #-}+ -- | Inherited from session for ease of access.+tableSessionRoot :: TableEnv m h -> SessionRoot+tableSessionRoot = sessionRoot . tableSessionEnv++{-# INLINE tableSessionId #-}+ -- | Inherited from session for ease of access.+tableSessionId :: TableEnv m h -> SessionId+tableSessionId = sessionId . tableSessionEnv++{-# INLINE tableSessionSalt #-}+ -- | Inherited from session for ease of access.+tableSessionSalt :: TableEnv m h -> Bloom.Salt+tableSessionSalt = sessionSalt . tableSessionEnv++{-# INLINE tableHasFS #-}+-- | Inherited from session for ease of access.+tableHasFS :: TableEnv m h -> HasFS m h+tableHasFS = sessionHasFS . tableSessionEnv++{-# INLINE tableHasBlockIO #-}+-- | Inherited from session for ease of access.+tableHasBlockIO :: TableEnv m h -> HasBlockIO m h+tableHasBlockIO = sessionHasBlockIO . tableSessionEnv++{-# INLINE tableSessionUniqCounter #-}+-- | Inherited from session for ease of access.+tableSessionUniqCounter :: Table m h -> UniqCounter m+tableSessionUniqCounter = sessionUniqCounter . tableSession++{-# INLINE tableSessionUntrackTable #-}+{-# SPECIALISE tableSessionUntrackTable :: TableId -> TableEnv IO h -> IO () #-}+-- | Open tables are tracked in the corresponding session, so when a table is+-- closed it should become untracked (forgotten).+tableSessionUntrackTable :: MonadMVar m => TableId -> TableEnv m h -> m ()+tableSessionUntrackTable tableId tEnv =+ modifyMVar_ (sessionOpenTables (tableSessionEnv tEnv)) $ pure . Map.delete tableId++-- | The table is closed.+data TableClosedError+ = ErrTableClosed+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | 'withKeepTableOpen' ensures that the table stays open for the duration of the+-- provided continuation.+--+-- NOTE: any operation except 'close' can use this function.+{-# INLINE withKeepTableOpen #-}+{-# SPECIALISE withKeepTableOpen ::+ Table IO h+ -> (TableEnv IO h -> IO a)+ -> IO a #-}+withKeepTableOpen ::+ (MonadSTM m, MonadThrow m)+ => Table m h+ -> (TableEnv m h -> m a)+ -> m a+withKeepTableOpen t action = RW.withReadAccess (tableState t) $ \case+ TableClosed -> throwIO ErrTableClosed+ TableOpen tEnv -> action tEnv++--+-- Implementation of public API+--++{-# SPECIALISE withTable ::+ Session IO h+ -> TableConfig+ -> (Table IO h -> IO a)+ -> IO a #-}+-- | See 'Database.LSMTree.withTable'.+withTable ::+ (MonadMask m, MonadSTM m, MonadMVar m, PrimMonad m)+ => Session m h+ -> TableConfig+ -> (Table m h -> m a)+ -> m a+withTable sesh conf = bracket (new sesh conf) close++{-# SPECIALISE new ::+ Session IO h+ -> TableConfig+ -> IO (Table IO h) #-}+-- | See 'Database.LSMTree.new'.+new ::+ (MonadSTM m, MonadMVar m, PrimMonad m, MonadMask m)+ => Session m h+ -> TableConfig+ -> m (Table m h)+new sesh conf = do+ tableId <- uniqueToTableId <$> incrUniqCounter (sessionUniqCounter sesh)+ let tr = TraceTable tableId `contramap` lsmTreeTracer sesh+ traceWith tr $ TraceNewTable conf++ withKeepSessionOpen sesh $ \seshEnv ->+ withActionRegistry $ \reg -> do+ am <- newArenaManager+ tc <- newEmptyTableContent (sessionUniqCounter sesh) seshEnv reg+ newWith reg sesh seshEnv conf am tr tableId tc++{-# SPECIALISE newEmptyTableContent ::+ UniqCounter IO+ -> SessionEnv IO h+ -> ActionRegistry IO+ -> IO (TableContent IO h) #-}+newEmptyTableContent ::+ (PrimMonad m, MonadMask m, MonadMVar m)+ => UniqCounter m+ -> SessionEnv m h+ -> ActionRegistry m+ -> m (TableContent m h)+newEmptyTableContent uc seshEnv reg = do+ blobpath <- Paths.tableBlobPath (sessionRoot seshEnv) <$>+ incrUniqCounter uc+ let tableWriteBuffer = WB.empty+ tableWriteBufferBlobs+ <- withRollback reg+ (WBB.new (sessionHasFS seshEnv) blobpath)+ releaseRef+ let tableLevels = V.empty+ tableCache <- mkLevelsCache reg tableLevels+ pure TableContent {+ tableWriteBuffer+ , tableWriteBufferBlobs+ , tableLevels+ , tableCache+ , tableUnionLevel = NoUnion+ }+++{-# SPECIALISE newWith ::+ ActionRegistry IO+ -> Session IO h+ -> SessionEnv IO h+ -> TableConfig+ -> ArenaManager RealWorld+ -> Tracer IO TableTrace+ -> TableId+ -> TableContent IO h+ -> IO (Table IO h) #-}+newWith ::+ (MonadSTM m, MonadMVar m, PrimMonad m)+ => ActionRegistry m+ -> Session m h+ -> SessionEnv m h+ -> TableConfig+ -> ArenaManager (PrimState m)+ -> Tracer m TableTrace+ -> TableId+ -> TableContent m h+ -> m (Table m h)+newWith reg sesh seshEnv conf !am !tr !tableId !tc = do+ -- The session is kept open until we've updated the session's set of tracked+ -- tables. If 'closeSession' is called by another thread while this code+ -- block is being executed, that thread will block until it reads the+ -- /updated/ set of tracked tables.+ contentVar <- RW.new $ tc+ tableVar <- RW.new $ TableOpen $ TableEnv {+ tableSessionEnv = seshEnv+ , tableContent = contentVar+ }+ let !t = Table conf tableVar am tr tableId sesh+ -- Track the current table+ delayedCommit reg $+ modifyMVar_ (sessionOpenTables seshEnv) $+ pure . Map.insert tableId t++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith tr $ TraceCreatedTable (sessionId seshEnv) conf++ pure $! t++{-# SPECIALISE close :: Table IO h -> IO () #-}+-- | See 'Database.LSMTree.close'.+close ::+ (MonadMask m, MonadSTM m, MonadMVar m, PrimMonad m)+ => Table m h+ -> m ()+close t = do+ traceWith (tableTracer t) TraceCloseTable+ modifyWithActionRegistry_+ (RW.unsafeAcquireWriteAccess (tableState t))+ (atomically . RW.unsafeReleaseWriteAccess (tableState t)) $ \reg -> \case+ TableClosed -> pure TableClosed+ TableOpen tEnv -> do+ -- Since we have a write lock on the table state, we know that we are the+ -- only thread currently closing the table. We can safely make the session+ -- forget about this table.+ delayedCommit reg (tableSessionUntrackTable (tableId t) tEnv)+ RW.withWriteAccess_ (tableContent tEnv) $ \tc -> do+ releaseTableContent reg tc+ pure tc++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith (tableTracer t) TraceClosedTable++ pure TableClosed++{-# SPECIALISE lookups ::+ ResolveSerialisedValue+ -> V.Vector SerialisedKey+ -> Table IO h+ -> IO (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef IO h)))) #-}+-- | See 'Database.LSMTree.lookups'.+lookups ::+ (MonadAsync m, MonadMask m, MonadMVar m, MonadST m)+ => ResolveSerialisedValue+ -> V.Vector SerialisedKey+ -> Table m h+ -> m (V.Vector (Maybe (Entry SerialisedValue (WeakBlobRef m h))))+lookups resolve ks t = do+ traceWith (tableTracer t) $ TraceLookups (V.length ks)+ withKeepTableOpen t $ \tEnv ->+ RW.withReadAccess (tableContent tEnv) $ \tc -> do+ case tableUnionLevel tc of+ NoUnion -> lookupsRegular tEnv tc+ Union tree unionCache -> do+ isStructurallyEmpty tree >>= \case+ True -> lookupsRegular tEnv tc+ False -> if WB.null (tableWriteBuffer tc) && V.null (tableLevels tc)+ then lookupsUnion tEnv unionCache+ else lookupsRegularAndUnion tEnv tc unionCache+ where+ lookupsRegular tEnv tc = do+ let !cache = tableCache tc+ lookupsIOWithWriteBuffer+ (tableHasBlockIO tEnv)+ (tableArenaManager t)+ resolve+ (tableSessionSalt tEnv)+ (tableWriteBuffer tc)+ (tableWriteBufferBlobs tc)+ (cachedRuns cache)+ (cachedFilters cache)+ (cachedIndexes cache)+ (cachedKOpsFiles cache)+ ks++ lookupsUnion tEnv unionCache = do+ treeResults <- flip MT.mapMStrict (cachedTree unionCache) $ \runs ->+ Async.async $ lookupsUnionSingleBatch tEnv runs+ MT.foldLookupTree resolve treeResults++ lookupsRegularAndUnion tEnv tc unionCache = do+ -- asynchronously, so tree lookup batches can already be submitted+ -- without waiting for the regular level result.+ regularResult <- Async.async $ lookupsRegular tEnv tc+ treeResults <- flip MT.mapMStrict (cachedTree unionCache) $ \runs ->+ Async.async $ lookupsUnionSingleBatch tEnv runs+ MT.foldLookupTree resolve $+ MT.mkLookupNode MR.MergeLevel $ V.fromList+ [ MT.LookupBatch regularResult+ , treeResults+ ]++ lookupsUnionSingleBatch tEnv runs =+ lookupsIO+ (tableHasBlockIO tEnv)+ (tableArenaManager t)+ resolve+ (tableSessionSalt tEnv)+ runs+ (V.mapStrict (\(DeRef r) -> Run.runFilter r) runs)+ (V.mapStrict (\(DeRef r) -> Run.runIndex r) runs)+ (V.mapStrict (\(DeRef r) -> Run.runKOpsFile r) runs)+ ks++{-# SPECIALISE rangeLookup ::+ ResolveSerialisedValue+ -> Range SerialisedKey+ -> Table IO h+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef IO h) -> res)+ -> IO (V.Vector res) #-}+-- | See 'Database.LSMTree.rangeLookup'.+rangeLookup ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> Range SerialisedKey+ -> Table m h+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef m h) -> res)+ -- ^ How to map to a query result+ -> m (V.Vector res)+rangeLookup resolve range t fromEntry = do+ traceWith (tableTracer t) $ TraceRangeLookup range+ case range of+ FromToExcluding lb ub ->+ withCursor resolve (OffsetKey lb) t $ \cursor ->+ go cursor (< ub) []+ FromToIncluding lb ub ->+ withCursor resolve (OffsetKey lb) t $ \cursor ->+ go cursor (<= ub) []+ where+ -- TODO: tune!+ -- Also, such a high number means that many tests never cover the case+ -- of having multiple chunks. Expose through the public API as config?+ chunkSize = 500++ go cursor isInUpperBound !chunks = do+ chunk <- readCursorWhile resolve isInUpperBound chunkSize cursor fromEntry+ let !n = V.length chunk+ if n >= chunkSize+ then go cursor isInUpperBound (chunk : chunks)+ -- This requires an extra copy. If we had a size hint, we could+ -- directly write everything into the result vector.+ -- TODO(optimise): revisit+ else pure (V.concat (reverse (V.slice 0 n chunk : chunks)))++{-# SPECIALISE updates ::+ ResolveSerialisedValue+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> Table IO h+ -> IO () #-}+-- | See 'Database.LSMTree.updates'.+--+-- Does not enforce that upsert and BLOBs should not occur in the same table.+updates ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> Table m h+ -> m ()+updates resolve es t = do+ traceWith (tableTracer t) $ TraceUpdates (V.length es)+ let conf = tableConfig t+ withKeepTableOpen t $ \tEnv -> do+ let hfs = tableHasFS tEnv+ modifyWithActionRegistry_+ (RW.unsafeAcquireWriteAccess (tableContent tEnv))+ (atomically . RW.unsafeReleaseWriteAccess (tableContent tEnv)) $ \reg tc -> do+ tc' <-+ updatesWithInterleavedFlushes+ (contramapTraceMerge $ tableTracer t)+ conf+ resolve+ hfs+ (tableHasBlockIO tEnv)+ (tableSessionRoot tEnv)+ (tableSessionSalt tEnv)+ (tableSessionUniqCounter t)+ es+ reg+ tc++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith (tableTracer t) $ TraceUpdated (V.length es)++ pure tc'++{-------------------------------------------------------------------------------+ Blobs+-------------------------------------------------------------------------------}++{- | A 'BlobRef' used with 'retrieveBlobs' was invalid.++'BlobRef's are obtained from lookups in a 'Table', but they may be+invalidated by subsequent changes in that 'Table'. In general the+reliable way to retrieve blobs is not to change the 'Table' before+retrieving the blobs. To allow later retrievals, duplicate the table+before making modifications and keep the table open until all blob+retrievals are complete.+-}+data BlobRefInvalidError+ = -- | The 'Int' index indicates the first invalid 'BlobRef'.+ ErrBlobRefInvalid !Int+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE retrieveBlobs ::+ Session IO h+ -> V.Vector (WeakBlobRef IO h)+ -> IO (V.Vector SerialisedBlob) #-}+retrieveBlobs ::+ (MonadMask m, MonadST m, MonadSTM m)+ => Session m h+ -> V.Vector (WeakBlobRef m h)+ -> m (V.Vector SerialisedBlob)+retrieveBlobs sesh wrefs = do+ traceWith (sessionTracer sesh) $ TraceRetrieveBlobs (V.length wrefs)+ withKeepSessionOpen sesh $ \seshEnv ->+ let hbio = sessionHasBlockIO seshEnv in+ handle (\(BlobRef.WeakBlobRefInvalid i) ->+ throwIO (ErrBlobRefInvalid i)) $+ BlobRef.readWeakBlobRefs hbio wrefs++{-------------------------------------------------------------------------------+ Cursors+-------------------------------------------------------------------------------}++-- | A read-only view into the table state at the time of cursor creation.+--+-- For more information, see 'Database.LSMTree.Cursor'.+--+-- The representation of a cursor is similar to that of a 'Table', but+-- simpler, as it is read-only.+data Cursor m h = Cursor {+ -- | Mutual exclusion, only a single thread can read from a cursor at a+ -- given time.+ cursorState :: !(StrictMVar m (CursorState m h))+ , cursorTracer :: !(Tracer m CursorTrace)+ }++instance NFData (Cursor m h) where+ rnf (Cursor a b) = rwhnf a `seq` rwhnf b++data CursorState m h =+ CursorOpen !(CursorEnv m h)+ | CursorClosed -- ^ Calls to a closed cursor raise an exception.++data CursorEnv m h = CursorEnv {+ -- === Session-inherited++ -- | The session that this cursor belongs to.+ --+ -- NOTE: Consider using the 'cursorSessionEnv' field instead of acquiring+ -- the session lock.+ cursorSession :: !(Session m h)+ -- | Use this instead of 'cursorSession' for easy access. An open cursor may+ -- assume that its session is open. A session's global resources, and+ -- therefore resources that are inherited by the cursor, will only be+ -- released once the session is sure that no cursors are open anymore.+ , cursorSessionEnv :: !(SessionEnv m h)++ -- === Cursor-specific++ -- | Session-unique identifier for this cursor.+ , cursorId :: !CursorId+ -- | Readers are immediately discarded once they are 'Readers.Drained',+ -- so if there is a 'Just', we can assume that it has further entries.+ -- Note that, while the readers do retain references on the blob files+ -- while they are active, once they are drained they do not. This could+ -- invalidate any 'BlobRef's previously handed out. To avoid this, we+ -- explicitly retain references on the runs and write buffer blofs and+ -- only release them when the cursor is closed (see 'cursorWBB' and+ -- 'cursorRuns' below).+ , cursorReaders :: !(Maybe (Readers.Readers m h))++ -- TODO: the cursorRuns and cursorWBB could be replaced by just retaining+ -- the BlobFile from the runs and WBB, so that we retain less. Since we+ -- only retain these to keep BlobRefs valid until the cursor is closed.+ -- Alternatively: the Readers could be modified to keep the BlobFiles even+ -- once the readers are drained, and only release them when the Readers is+ -- itself closed.++ -- | The write buffer blobs, which like the runs, we have to keep open+ -- untile the cursor is closed.+ , cursorWBB :: !(Ref (WBB.WriteBufferBlobs m h))++ -- | The runs from regular levels that are held open by the cursor. We must+ -- release these references when the cursor gets closed.+ , cursorRuns :: !(V.Vector (Ref (Run m h)))++ -- | The runs from the union level that are held open by the cursor. We must+ -- release these references when the cursor gets closed.+ , cursorUnion :: !(Maybe (UnionCache m h))+ }++{-# SPECIALISE withCursor ::+ ResolveSerialisedValue+ -> OffsetKey+ -> Table IO h+ -> (Cursor IO h -> IO a)+ -> IO a #-}+-- | See 'Database.LSMTree.withCursor'.+withCursor ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> OffsetKey+ -> Table m h+ -> (Cursor m h -> m a)+ -> m a+withCursor resolve offsetKey t = bracket (newCursor resolve offsetKey t) closeCursor++{-# SPECIALISE newCursor ::+ ResolveSerialisedValue+ -> OffsetKey+ -> Table IO h+ -> IO (Cursor IO h) #-}+-- | See 'Database.LSMTree.newCursor'.+newCursor ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> OffsetKey+ -> Table m h+ -> m (Cursor m h)+newCursor !resolve !offsetKey t = withKeepTableOpen t $ \tEnv -> do+ let cursorSession = tableSession t+ let cursorSessionEnv = tableSessionEnv tEnv+ cursorId <- uniqueToCursorId <$>+ incrUniqCounter (tableSessionUniqCounter t)+ let cursorTracer = TraceCursor cursorId `contramap` lsmTreeTracer cursorSession+ traceWith cursorTracer $ TraceNewCursor (tableId t) $ case offsetKey of+ NoOffsetKey -> Nothing+ OffsetKey (SerialisedKey bytes) -> Just bytes++ -- We acquire a read-lock on the session open-state to prevent races, see+ -- 'sessionOpenTables'.+ withKeepSessionOpen cursorSession $ \_ -> do+ withActionRegistry $ \reg -> do+ (wb, wbblobs, cursorRuns, cursorUnion) <-+ dupTableContent reg (tableContent tEnv)+ let cursorSources =+ Readers.FromWriteBuffer wb wbblobs+ : fmap Readers.FromRun (V.toList cursorRuns)+ <> case cursorUnion of+ Nothing -> []+ Just (UnionCache treeCache) ->+ [lookupTreeToReaderSource treeCache]+ cursorReaders <-+ withRollbackMaybe reg+ (Readers.new resolve offsetKey cursorSources)+ Readers.close+ let cursorWBB = wbblobs+ cursorState <- newMVar (CursorOpen CursorEnv {..})+ let !cursor = Cursor {cursorState, cursorTracer}+ -- Track cursor, but careful: If now an exception is raised, all+ -- resources get freed by the registry, so if the session still+ -- tracks 'cursor' (which is 'CursorOpen'), it later double frees.+ -- Therefore, we only track the cursor if 'withActionRegistry' exits+ -- successfully, i.e. using 'delayedCommit'.+ delayedCommit reg $+ modifyMVar_ (sessionOpenCursors cursorSessionEnv) $+ pure . Map.insert cursorId cursor++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith cursorTracer $ TraceCreatedCursor (tableSessionId tEnv)++ pure $! cursor+ where+ -- The table contents escape the read access, but we just duplicate+ -- references to each run, so it is safe.+ dupTableContent reg contentVar = do+ RW.withReadAccess contentVar $ \content -> do+ let !wb = tableWriteBuffer content+ !wbblobs = tableWriteBufferBlobs content+ wbblobs' <- withRollback reg (dupRef wbblobs) releaseRef+ let runs = cachedRuns (tableCache content)+ runs' <- V.forM runs $ \r ->+ withRollback reg (dupRef r) releaseRef+ unionCache <- case tableUnionLevel content of+ NoUnion -> pure Nothing+ Union _ c -> Just <$!> duplicateUnionCache reg c+ pure (wb, wbblobs', runs', unionCache)++lookupTreeToReaderSource ::+ MT.LookupTree (V.Vector (Ref (Run m h))) -> Readers.ReaderSource m h+lookupTreeToReaderSource = \case+ MT.LookupBatch v ->+ case map Readers.FromRun (V.toList v) of+ [src] -> src+ srcs -> Readers.FromReaders Readers.MergeLevel srcs+ MT.LookupNode ty children ->+ Readers.FromReaders+ (convertMergeType ty)+ (map lookupTreeToReaderSource (V.toList children))+ where+ convertMergeType = \case+ MR.MergeUnion -> Readers.MergeUnion+ MR.MergeLevel -> Readers.MergeLevel++{-# SPECIALISE closeCursor :: Cursor IO h -> IO () #-}+-- | See 'Database.LSMTree.closeCursor'.+closeCursor ::+ (MonadMask m, MonadMVar m, MonadSTM m, PrimMonad m)+ => Cursor m h+ -> m ()+closeCursor Cursor {..} = do+ traceWith cursorTracer $ TraceCloseCursor+ modifyWithActionRegistry_ (takeMVar cursorState) (putMVar cursorState) $ \reg -> \case+ CursorClosed -> pure CursorClosed+ CursorOpen CursorEnv {..} -> do+ -- This should be safe-ish, but it's still not ideal, because it doesn't+ -- rule out sync exceptions in the cleanup operations.+ -- In that case, the cursor ends up closed, but resources might not have+ -- been freed. Probably better than the other way around, though.+ delayedCommit reg $+ modifyMVar_ (sessionOpenCursors cursorSessionEnv) $+ pure . Map.delete cursorId++ forM_ cursorReaders $ delayedCommit reg . Readers.close+ V.forM_ cursorRuns $ delayedCommit reg . releaseRef+ forM_ cursorUnion $ releaseUnionCache reg+ delayedCommit reg (releaseRef cursorWBB)++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith cursorTracer $ TraceClosedCursor++ pure CursorClosed++{-# SPECIALISE readCursor ::+ ResolveSerialisedValue+ -> Int+ -> Cursor IO h+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef IO h) -> res)+ -> IO (V.Vector res) #-}+-- | See 'Database.LSMTree.readCursor'.+readCursor ::+ forall m h res.+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> Int -- ^ Maximum number of entries to read+ -> Cursor m h+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef m h) -> res)+ -- ^ How to map to a query result+ -> m (V.Vector res)+readCursor resolve n cursor fromEntry =+ readCursorWhile resolve (const True) n cursor fromEntry++-- | The cursor is closed.+data CursorClosedError+ = ErrCursorClosed+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE readCursorWhile ::+ ResolveSerialisedValue+ -> (SerialisedKey -> Bool)+ -> Int+ -> Cursor IO h+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef IO h) -> res)+ -> IO (V.Vector res) #-}+-- | @readCursorWhile _ p n cursor _@ reads elements until either:+--+-- * @n@ elements have been read already+-- * @p@ returns @False@ for the key of an entry to be read+-- * the cursor is drained+--+-- Consequently, once a call returned fewer than @n@ elements, any subsequent+-- calls with the same predicate @p@ will return an empty vector.+readCursorWhile ::+ forall m h res.+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => ResolveSerialisedValue+ -> (SerialisedKey -> Bool) -- ^ Only read as long as this predicate holds+ -> Int -- ^ Maximum number of entries to read+ -> Cursor m h+ -> (SerialisedKey -> SerialisedValue -> Maybe (WeakBlobRef m h) -> res)+ -- ^ How to map to a query result+ -> m (V.Vector res)+readCursorWhile resolve keyIsWanted n Cursor {..} fromEntry = do+ traceWith cursorTracer $ TraceReadingCursor n+ res <- modifyMVar cursorState $ \case+ CursorClosed -> throwIO ErrCursorClosed+ state@(CursorOpen cursorEnv) -> do+ case cursorReaders cursorEnv of+ Nothing ->+ -- a drained cursor will just return an empty vector+ pure (state, V.empty)+ Just readers -> do+ (vec, hasMore) <- Cursor.readEntriesWhile resolve keyIsWanted fromEntry readers n+ -- if we drained the readers, remove them from the state+ let !state' = case hasMore of+ Readers.HasMore -> state+ Readers.Drained -> CursorOpen (cursorEnv {cursorReaders = Nothing})+ pure (state', vec)++ traceWith cursorTracer $ TraceReadCursor n (V.length res)++ pure res++{-------------------------------------------------------------------------------+ Snapshots+-------------------------------------------------------------------------------}++-- | The named snapshot already exists.+data SnapshotExistsError+ = ErrSnapshotExists !SnapshotName+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE saveSnapshot ::+ SnapshotName+ -> SnapshotLabel+ -> Table IO h+ -> IO () #-}+-- | See 'Database.LSMTree.saveSnapshot'.+saveSnapshot ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => SnapshotName+ -> SnapshotLabel+ -> Table m h+ -> m ()+saveSnapshot snap label t = do+ traceWith (tableTracer t) $ TraceSaveSnapshot snap+ withKeepTableOpen t $ \tEnv ->+ withActionRegistry $ \reg -> do -- TODO: use the action registry for all side effects+ let hfs = tableHasFS tEnv+ hbio = tableHasBlockIO tEnv+ activeUc = tableSessionUniqCounter t++ -- Guard that the snapshot does not exist already+ let snapDir = Paths.namedSnapshotDir (tableSessionRoot tEnv) snap+ snapshotExists <- doesSnapshotDirExist snap (tableSessionEnv tEnv)+ if snapshotExists then+ throwIO (ErrSnapshotExists snap)+ else+ -- we assume the snapshots directory already exists, so we just have+ -- to create the directory for this specific snapshot.+ withRollback_ reg+ (FS.createDirectory hfs (Paths.getNamedSnapshotDir snapDir))+ (FS.removeDirectoryRecursive hfs (Paths.getNamedSnapshotDir snapDir))++ -- Duplicate references to the table content, so that resources do not disappear+ -- from under our feet while taking a snapshot. These references are released+ -- again after the snapshot files/directories are written.+ content <- RW.withReadAccess (tableContent tEnv) (duplicateTableContent reg)++ -- Fresh UniqCounter so that we start numbering from 0 in the named+ -- snapshot directory+ snapUc <- newUniqCounter 0++ -- Snapshot the write buffer.+ let activeDir = Paths.activeDir (tableSessionRoot tEnv)+ let wb = tableWriteBuffer content+ let wbb = tableWriteBufferBlobs content+ snapWriteBufferNumber <- Paths.writeBufferNumber <$>+ snapshotWriteBuffer hfs hbio activeUc snapUc reg activeDir snapDir wb wbb++ -- Convert to snapshot format+ snapLevels <- toSnapLevels (tableLevels content)++ -- Hard link runs into the named snapshot directory+ snapLevels' <- traverse (snapshotRun hfs hbio snapUc reg snapDir) snapLevels++ -- If a merging tree exists, do the same hard-linking for the runs within+ mTreeOpt <- case tableUnionLevel content of+ NoUnion -> pure Nothing+ Union mTreeRef _cache -> do+ mTree <- toSnapMergingTree mTreeRef+ Just <$> traverse (snapshotRun hfs hbio snapUc reg snapDir) mTree++ releaseTableContent reg content++ let snapMetaData = SnapshotMetaData+ label+ (tableConfig t)+ snapWriteBufferNumber+ snapLevels'+ mTreeOpt+ SnapshotMetaDataFile contentPath = Paths.snapshotMetaDataFile snapDir+ SnapshotMetaDataChecksumFile checksumPath = Paths.snapshotMetaDataChecksumFile snapDir+ writeFileSnapshotMetaData hfs contentPath checksumPath snapMetaData++ -- Make the directory and its contents durable.+ FS.synchroniseDirectoryRecursive hfs hbio (Paths.getNamedSnapshotDir snapDir)++ -- Note: we're "abusing" the action registry to trace the success+ -- message as late as possible.+ delayedCommit reg $+ traceWith (tableTracer t) $ TraceSavedSnapshot snap++-- | The named snapshot does not exist.+data SnapshotDoesNotExistError+ = ErrSnapshotDoesNotExist !SnapshotName+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | The named snapshot is corrupted.+data SnapshotCorruptedError+ = ErrSnapshotCorrupted+ !SnapshotName+ !FileCorruptedError+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | The named snapshot is not compatible with the expected type.+data SnapshotNotCompatibleError+ = -- | The named snapshot is not compatible with the given label.+ ErrSnapshotWrongLabel+ !SnapshotName+ -- | Expected label+ !SnapshotLabel+ -- | Actual label+ !SnapshotLabel+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE openTableFromSnapshot ::+ TableConfigOverride+ -> Session IO h+ -> SnapshotName+ -> SnapshotLabel+ -> ResolveSerialisedValue+ -> IO (Table IO h) #-}+-- | See 'Database.LSMTree.openTableFromSnapshot'.+openTableFromSnapshot ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => TableConfigOverride+ -> Session m h+ -> SnapshotName+ -> SnapshotLabel -- ^ Expected label+ -> ResolveSerialisedValue+ -> m (Table m h)+openTableFromSnapshot policyOveride sesh snap label resolve = do+ tableId <- uniqueToTableId <$> incrUniqCounter (sessionUniqCounter sesh)+ let tr = TraceTable tableId `contramap` lsmTreeTracer sesh+ traceWith tr $ TraceOpenTableFromSnapshot snap policyOveride++ wrapFileCorruptedErrorAsSnapshotCorruptedError snap $ do+ withKeepSessionOpen sesh $ \seshEnv -> do+ withActionRegistry $ \reg -> do+ let hfs = sessionHasFS seshEnv+ hbio = sessionHasBlockIO seshEnv+ uc = sessionUniqCounter sesh++ -- Guard that the snapshot exists+ let snapDir = Paths.namedSnapshotDir (sessionRoot seshEnv) snap+ FS.doesDirectoryExist hfs (Paths.getNamedSnapshotDir snapDir) >>= \b ->+ unless b $ throwIO (ErrSnapshotDoesNotExist snap)++ let SnapshotMetaDataFile contentPath = Paths.snapshotMetaDataFile snapDir+ SnapshotMetaDataChecksumFile checksumPath = Paths.snapshotMetaDataChecksumFile snapDir++ snapMetaData <- readFileSnapshotMetaData hfs contentPath checksumPath++ let SnapshotMetaData label' conf snapWriteBuffer snapLevels mTreeOpt+ = overrideTableConfig policyOveride snapMetaData++ unless (label == label') $+ throwIO (ErrSnapshotWrongLabel snap label label')++ am <- newArenaManager++ let salt = sessionSalt seshEnv+ let activeDir = Paths.activeDir (sessionRoot seshEnv)++ -- Read write buffer+ let snapWriteBufferPaths = Paths.WriteBufferFsPaths (Paths.getNamedSnapshotDir snapDir) snapWriteBuffer+ (tableWriteBuffer, tableWriteBufferBlobs) <-+ openWriteBuffer reg resolve hfs hbio uc activeDir snapWriteBufferPaths++ -- Hard link runs into the active directory,+ snapLevels' <- traverse (openRun hfs hbio uc reg snapDir activeDir salt) snapLevels+ unionLevel <- case mTreeOpt of+ Nothing -> pure NoUnion+ Just mTree -> do+ snapTree <- traverse (openRun hfs hbio uc reg snapDir activeDir salt) mTree+ mt <- fromSnapMergingTree hfs hbio salt uc resolve activeDir reg snapTree+ isStructurallyEmpty mt >>= \case+ True ->+ pure NoUnion+ False -> do+ traverse_ (delayedCommit reg . releaseRef) snapTree+ cache <- mkUnionCache reg mt+ pure (Union mt cache)++ -- Convert from the snapshot format, restoring merge progress in the process+ tableLevels <- fromSnapLevels hfs hbio salt uc conf resolve reg activeDir snapLevels'+ traverse_ (delayedCommit reg . releaseRef) snapLevels'++ tableCache <- mkLevelsCache reg tableLevels+ newWith reg sesh seshEnv conf am tr tableId $! TableContent {+ tableWriteBuffer+ , tableWriteBufferBlobs+ , tableLevels+ , tableCache+ , tableUnionLevel = unionLevel+ }++{-# SPECIALISE wrapFileCorruptedErrorAsSnapshotCorruptedError ::+ SnapshotName+ -> IO a+ -> IO a+ #-}+wrapFileCorruptedErrorAsSnapshotCorruptedError ::+ forall m a.+ (MonadCatch m)+ => SnapshotName+ -> m a+ -> m a+wrapFileCorruptedErrorAsSnapshotCorruptedError snapshotName =+ mapExceptionWithActionRegistry (ErrSnapshotCorrupted snapshotName)++{-# SPECIALISE doesSnapshotExist ::+ Session IO h+ -> SnapshotName+ -> IO Bool #-}+-- | See 'Database.LSMTree.doesSnapshotExist'.+doesSnapshotExist ::+ (MonadMask m, MonadSTM m)+ => Session m h+ -> SnapshotName+ -> m Bool+doesSnapshotExist sesh snap = withKeepSessionOpen sesh (doesSnapshotDirExist snap)++-- | Internal helper: Variant of 'doesSnapshotExist' that does not take a session lock.+doesSnapshotDirExist :: SnapshotName -> SessionEnv m h -> m Bool+doesSnapshotDirExist snap seshEnv = do+ let snapDir = Paths.namedSnapshotDir (sessionRoot seshEnv) snap+ FS.doesDirectoryExist (sessionHasFS seshEnv) (Paths.getNamedSnapshotDir snapDir)++{-# SPECIALISE deleteSnapshot ::+ Session IO h+ -> SnapshotName+ -> IO () #-}+-- | See 'Database.LSMTree.deleteSnapshot'.+deleteSnapshot ::+ (MonadMask m, MonadSTM m)+ => Session m h+ -> SnapshotName+ -> m ()+deleteSnapshot sesh snap = do+ traceWith (sessionTracer sesh) $ TraceDeleteSnapshot snap+ withKeepSessionOpen sesh $ \seshEnv -> do+ let snapDir = Paths.namedSnapshotDir (sessionRoot seshEnv) snap+ snapshotExists <- doesSnapshotDirExist snap seshEnv+ unless snapshotExists $ throwIO (ErrSnapshotDoesNotExist snap)+ FS.removeDirectoryRecursive (sessionHasFS seshEnv) (Paths.getNamedSnapshotDir snapDir)+ traceWith (sessionTracer sesh) $ TraceDeletedSnapshot snap++{-# SPECIALISE listSnapshots :: Session IO h -> IO [SnapshotName] #-}+-- | See 'Database.LSMTree.listSnapshots'.+listSnapshots ::+ (MonadMask m, MonadSTM m)+ => Session m h+ -> m [SnapshotName]+listSnapshots sesh = do+ traceWith (sessionTracer sesh) TraceListSnapshots+ withKeepSessionOpen sesh $ \seshEnv -> do+ let hfs = sessionHasFS seshEnv+ root = sessionRoot seshEnv+ contents <- FS.listDirectory hfs (Paths.snapshotsDir (sessionRoot seshEnv))+ snaps <- mapM (checkSnapshot hfs root) $ Set.toList contents+ pure $ catMaybes snaps+ where+ checkSnapshot hfs root s = do+ -- TODO: rethrow 'ErrInvalidSnapshotName' as 'ErrSnapshotDirCorrupted'+ let snap = Paths.toSnapshotName s+ -- check that it is a directory+ b <- FS.doesDirectoryExist hfs+ (Paths.getNamedSnapshotDir $ Paths.namedSnapshotDir root snap)+ if b then pure $ Just snap+ else pure $ Nothing++{-------------------------------------------------------------------------------+ Multiple writable tables+-------------------------------------------------------------------------------}++{-# SPECIALISE duplicate :: Table IO h -> IO (Table IO h) #-}+-- | See 'Database.LSMTree.duplicate'.+duplicate ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => Table m h+ -> m (Table m h)+duplicate t@Table{..} = do+ childTableId <- uniqueToTableId <$> incrUniqCounter (tableSessionUniqCounter t)+ let childTableTracer = TraceTable childTableId `contramap` lsmTreeTracer tableSession+ parentTableId = tableId+ traceWith childTableTracer $ TraceDuplicate parentTableId++ withKeepTableOpen t $ \TableEnv{..} -> do+ -- We acquire a read-lock on the session open-state to prevent races, see+ -- 'sessionOpenTables'.+ withKeepSessionOpen tableSession $ \_ -> do+ withActionRegistry $ \reg -> do+ -- The table contents escape the read access, but we just added references+ -- to each run so it is safe.+ content <- RW.withReadAccess tableContent (duplicateTableContent reg)+ newWith+ reg+ tableSession+ tableSessionEnv+ tableConfig+ tableArenaManager+ childTableTracer+ childTableId+ content++{-------------------------------------------------------------------------------+ Table union+-------------------------------------------------------------------------------}++-- | A table union was constructed with two tables that are not compatible.+data TableUnionNotCompatibleError+ = ErrTableUnionHandleTypeMismatch+ -- | The index of the first table.+ !Int+ -- | The type of the filesystem handle of the first table.+ !TypeRep+ -- | The index of the second table.+ !Int+ -- | The type of the filesystem handle of the second table.+ !TypeRep+ | ErrTableUnionSessionMismatch+ -- | The index of the first table.+ !Int+ -- | The session directory of the first table.+ !FsErrorPath+ -- | The index of the second table.+ !Int+ -- | The session directory of the second table.+ !FsErrorPath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{-# SPECIALISE unions :: NonEmpty (Table IO h) -> IO (Table IO h) #-}+-- | See 'Database.LSMTree.unions'.+unions ::+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => NonEmpty (Table m h)+ -> m (Table m h)+unions ts = do+ sesh <- ensureSessionsMatch ts++ childTableId <- uniqueToTableId <$> incrUniqCounter (sessionUniqCounter sesh)+ let childTableTracer = TraceTable childTableId `contramap` lsmTreeTracer sesh+ traceWith childTableTracer $ TraceIncrementalUnions (NE.map tableId ts)++ -- The TableConfig for the new table is taken from the last table in the+ -- union. This corresponds to the "Data.Map.union updates baseMap" order,+ -- where we take the config from the base map, not the updates.+ --+ -- This could be modified to take the new config as an input. It works to+ -- pick any config because the new table is almost completely fresh. It+ -- will have an empty write buffer and no runs in the normal levels. All+ -- the existing runs get squashed down into a single run before rejoining+ -- as a last level.+ let conf = tableConfig (NE.last ts)++ -- We acquire a read-lock on the session open-state to prevent races, see+ -- 'sessionOpenTables'.+ modifyWithActionRegistry+ (atomically $ RW.unsafeAcquireReadAccess (sessionState sesh))+ (\_ -> atomically $ RW.unsafeReleaseReadAccess (sessionState sesh)) $+ \reg -> \case+ SessionClosed -> throwIO ErrSessionClosed+ seshState@(SessionOpen seshEnv) -> do+ t <- unionsInOpenSession reg sesh seshEnv conf childTableTracer childTableId ts+ pure (seshState, t)++{-# SPECIALISE unionsInOpenSession ::+ ActionRegistry IO+ -> Session IO h+ -> SessionEnv IO h+ -> TableConfig+ -> Tracer IO TableTrace+ -> TableId+ -> NonEmpty (Table IO h)+ -> IO (Table IO h) #-}+unionsInOpenSession ::+ (MonadSTM m, MonadMask m, MonadMVar m, MonadST m)+ => ActionRegistry m+ -> Session m h+ -> SessionEnv m h+ -> TableConfig+ -> Tracer m TableTrace+ -> TableId+ -> NonEmpty (Table m h)+ -> m (Table m h)+unionsInOpenSession reg sesh seshEnv conf tr !tableId ts = do+ mts <- forM (NE.toList ts) $ \t ->+ withKeepTableOpen t $ \tEnv ->+ RW.withReadAccess (tableContent tEnv) $ \tc ->+ -- tableContentToMergingTree duplicates all runs and merges+ -- so the ones from the tableContent here do not escape+ -- the read access.+ withRollback reg+ (tableContentToMergingTree (sessionUniqCounter sesh) seshEnv conf tc)+ releaseRef+ mt <- withRollback reg (newPendingUnionMerge mts) releaseRef++ -- The mts here is a temporary value, since newPendingUnionMerge+ -- will make its own references, so release mts at the end of+ -- the action registry bracket+ forM_ mts (delayedCommit reg . releaseRef)++ content <- MT.isStructurallyEmpty mt >>= \case+ True -> do+ -- no need to have an empty merging tree+ delayedCommit reg (releaseRef mt)+ newEmptyTableContent ((sessionUniqCounter sesh)) seshEnv reg+ False -> do+ empty <- newEmptyTableContent (sessionUniqCounter sesh) seshEnv reg+ cache <- mkUnionCache reg mt+ pure empty { tableUnionLevel = Union mt cache }++ -- Pick the arena manager to optimise the case of:+ -- someUpdates <> bigTableWithLotsOfLookups+ -- by reusing the arena manager from the last one.+ let am = tableArenaManager (NE.last ts)++ newWith reg sesh seshEnv conf am tr tableId content++{-# SPECIALISE tableContentToMergingTree ::+ UniqCounter IO+ -> SessionEnv IO h+ -> TableConfig+ -> TableContent IO h+ -> IO (Ref (MergingTree IO h)) #-}+tableContentToMergingTree ::+ forall m h.+ (MonadMask m, MonadMVar m, MonadST m, MonadSTM m)+ => UniqCounter m+ -> SessionEnv m h+ -> TableConfig+ -> TableContent m h+ -> m (Ref (MergingTree m h))+tableContentToMergingTree uc seshEnv conf+ tc@TableContent {+ tableLevels,+ tableUnionLevel+ } =+ bracket (writeBufferToNewRun uc seshEnv conf tc)+ (mapM_ releaseRef) $ \mwbr ->+ let runs :: [PreExistingRun m h]+ runs = maybeToList (fmap PreExistingRun mwbr)+ ++ concatMap levelToPreExistingRuns (V.toList tableLevels)+ -- any pre-existing union in the input table:+ unionmt = case tableUnionLevel of+ NoUnion -> Nothing+ Union mt _ -> Just mt -- we could reuse the cache, but it+ -- would complicate things+ in newPendingLevelMerge runs unionmt+ where+ levelToPreExistingRuns :: Level m h -> [PreExistingRun m h]+ levelToPreExistingRuns Level{incomingRun, residentRuns} =+ case incomingRun of+ Single r -> PreExistingRun r+ Merging _ _ _ mr -> PreExistingMergingRun mr+ : map PreExistingRun (V.toList residentRuns)++--TODO: can we share this or move it to MergeSchedule?+{-# SPECIALISE writeBufferToNewRun ::+ UniqCounter IO+ -> SessionEnv IO h+ -> TableConfig+ -> TableContent IO h+ -> IO (Maybe (Ref (Run IO h))) #-}+writeBufferToNewRun ::+ (MonadMask m, MonadST m, MonadSTM m)+ => UniqCounter m+ -> SessionEnv m h+ -> TableConfig+ -> TableContent m h+ -> m (Maybe (Ref (Run m h)))+writeBufferToNewRun uc+ SessionEnv {+ sessionRoot = root,+ sessionSalt = salt,+ sessionHasFS = hfs,+ sessionHasBlockIO = hbio+ }+ conf+ TableContent{+ tableWriteBuffer,+ tableWriteBufferBlobs+ }+ | WB.null tableWriteBuffer = pure Nothing+ | otherwise = Just <$> do+ !uniq <- incrUniqCounter uc+ let (!runParams, !runPaths) = mergingRunParamsForLevel+ (Paths.activeDir root) conf uniq (LevelNo 1)+ Run.fromWriteBuffer+ hfs hbio salt+ runParams runPaths+ tableWriteBuffer+ tableWriteBufferBlobs++{-# SPECIALISE ensureSessionsMatch ::+ NonEmpty (Table IO h)+ -> IO (Session IO h) #-}+-- | Check if all tables have the same session.+-- If so, return the session.+-- Otherwise, throw a 'TableUnionNotCompatibleError'.+ensureSessionsMatch ::+ (MonadSTM m, MonadThrow m)+ => NonEmpty (Table m h)+ -> m (Session m h)+ensureSessionsMatch (t :| ts) = do+ let sesh = tableSession t+ withKeepSessionOpen sesh $ \seshEnv -> do+ let root = FS.mkFsErrorPath (sessionHasFS seshEnv) (getSessionRoot (sessionRoot seshEnv))+ -- Check that the session roots for all tables are the same. There can only+ -- be one *open/active* session per directory because of cooperative file+ -- locks, so each unique *open* session has a unique session root. We check+ -- that all the table's sessions are open at the same time while comparing+ -- the session roots.+ for_ (zip [1..] ts) $ \(i, t') -> do+ let sesh' = tableSession t'+ withKeepSessionOpen sesh' $ \seshEnv' -> do+ let root' = FS.mkFsErrorPath (sessionHasFS seshEnv') (getSessionRoot (sessionRoot seshEnv'))+ -- TODO: compare LockFileHandle instead of SessionRoot (?).+ -- We can write an Eq instance for LockFileHandle based on pointer equality,+ -- just like base does for Handle.+ unless (root == root') $ throwIO $ ErrTableUnionSessionMismatch 0 root i root'+ pure sesh++{-------------------------------------------------------------------------------+ Table union: debt and credit+-------------------------------------------------------------------------------}++{- |+Union debt represents the amount of computation that must be performed before an incremental union is completed.+This includes the cost of completing incremental unions that were part of a union's input.++__Warning:__ The 'UnionDebt' returned by 'Database.LSMTree.remainingUnionDebt' is an /upper bound/ on the remaining union debt, not the exact union debt.+-}+newtype UnionDebt = UnionDebt Int+ deriving newtype (Show, Eq, Ord, Num)++{-# SPECIALISE remainingUnionDebt :: Table IO h -> IO UnionDebt #-}+-- | See 'Database.LSMTree.remainingUnionDebt'.+remainingUnionDebt ::+ (MonadSTM m, MonadMVar m, MonadThrow m, PrimMonad m)+ => Table m h -> m UnionDebt+remainingUnionDebt t = do+ traceWith (tableTracer t) TraceRemainingUnionDebt+ withKeepTableOpen t $ \tEnv -> do+ RW.withReadAccess (tableContent tEnv) $ \tableContent -> do+ case tableUnionLevel tableContent of+ NoUnion ->+ pure (UnionDebt 0)+ Union mt _ -> do+ (MergeDebt (MergeCredits c), _) <- MT.remainingMergeDebt mt+ pure (UnionDebt c)++{- |+Union credits are passed to 'Database.LSMTree.supplyUnionCredits' to perform some amount of computation to incrementally complete a union.+-}+newtype UnionCredits = UnionCredits Int+ deriving newtype (Show, Eq, Ord, Num)++{-# SPECIALISE supplyUnionCredits ::+ ResolveSerialisedValue -> Table IO h -> UnionCredits -> IO UnionCredits #-}+-- | See 'Database.LSMTree.supplyUnionCredits'.+supplyUnionCredits ::+ (MonadST m, MonadSTM m, MonadMVar m, MonadMask m)+ => ResolveSerialisedValue -> Table m h -> UnionCredits -> m UnionCredits+supplyUnionCredits resolve t credits = do+ traceWith (tableTracer t) $ TraceSupplyUnionCredits credits+ withKeepTableOpen t $ \tEnv -> do+ -- We also want to mutate the table content to re-build the union cache,+ -- but we don't need to hold a writer lock while we work on the tree+ -- itself.+ -- TODO: revisit the locking strategy here.+ leftovers <- RW.withReadAccess (tableContent tEnv) $ \tc -> do+ case tableUnionLevel tc of+ NoUnion ->+ pure (max 0 credits) -- all leftovers (but never negative)+ Union mt _ -> do+ let conf = tableConfig t+ let AllocNumEntries x = confWriteBufferAlloc conf+ -- We simply use the write buffer capacity as merge credit threshold, as+ -- the regular level merges also do.+ -- TODO: pick a more suitable threshold or make configurable?+ let thresh = MR.CreditThreshold (MR.UnspentCredits (MergeCredits x))+ MergeCredits leftovers <-+ MT.supplyCredits+ (tableHasFS tEnv)+ (tableHasBlockIO tEnv)+ resolve+ (tableSessionSalt tEnv)+ (runParamsForLevel conf UnionLevel)+ thresh+ (tableSessionRoot tEnv)+ (tableSessionUniqCounter t)+ mt+ (let UnionCredits c = credits in MergeCredits c)+ pure (UnionCredits leftovers)+ traceWith (tableTracer t) $ TraceSuppliedUnionCredits credits leftovers+ -- TODO: don't go into this section if we know there is NoUnion+ modifyWithActionRegistry_+ (RW.unsafeAcquireWriteAccess (tableContent tEnv))+ (atomically . RW.unsafeReleaseWriteAccess (tableContent tEnv))+ $ \reg tc ->+ case tableUnionLevel tc of+ NoUnion -> pure tc+ Union mt cache -> do+ unionLevel' <- MT.isStructurallyEmpty mt >>= \case+ True ->+ pure NoUnion+ False -> do+ cache' <- mkUnionCache reg mt+ releaseUnionCache reg cache+ pure (Union mt cache')+ pure tc { tableUnionLevel = unionLevel' }+ pure leftovers
@@ -0,0 +1,82 @@+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE RoleAnnotations #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Unsliced (+ -- * Unsliced raw bytes+ Unsliced+ -- * Unsliced keys+ , makeUnslicedKey+ , unsafeMakeUnslicedKey+ , fromUnslicedKey+ ) where++import Control.DeepSeq (NFData)+import Control.Exception (assert)+import Data.ByteString.Short (ShortByteString (SBS))+import Data.Primitive.ByteArray+import qualified Data.Vector.Primitive as VP+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise (SerialisedKey (..))+import Database.LSMTree.Internal.Vector (mkPrimVector,+ noRetainedExtraMemory)++-- | Unsliced string of bytes+type role Unsliced nominal+newtype Unsliced a = Unsliced ByteArray+ deriving newtype NFData++getByteArray :: RawBytes -> ByteArray+getByteArray (RawBytes (VP.Vector _ _ ba)) = ba++precondition :: RawBytes -> Bool+precondition (RawBytes pvec) = noRetainedExtraMemory pvec++makeUnsliced :: RawBytes -> Unsliced RawBytes+makeUnsliced bytes+ | precondition bytes = Unsliced (getByteArray bytes)+ | otherwise = Unsliced (getByteArray $ RB.copy bytes)++unsafeMakeUnsliced :: RawBytes -> Unsliced RawBytes+unsafeMakeUnsliced bytes = assert (precondition bytes) (Unsliced (getByteArray bytes))++fromUnsliced :: Unsliced RawBytes -> RawBytes+fromUnsliced (Unsliced ba) = RawBytes (mkPrimVector 0 (sizeofByteArray ba) ba)++{-------------------------------------------------------------------------------+ Unsliced keys+-------------------------------------------------------------------------------}++from :: Unsliced RawBytes -> Unsliced SerialisedKey+from (Unsliced ba) = Unsliced ba++to :: Unsliced SerialisedKey -> Unsliced RawBytes+to (Unsliced ba) = Unsliced ba++makeUnslicedKey :: SerialisedKey -> Unsliced SerialisedKey+makeUnslicedKey (SerialisedKey rb) = from (makeUnsliced rb)++unsafeMakeUnslicedKey :: SerialisedKey -> Unsliced SerialisedKey+unsafeMakeUnslicedKey (SerialisedKey rb) = from (unsafeMakeUnsliced rb)++fromUnslicedKey :: Unsliced SerialisedKey -> SerialisedKey+fromUnslicedKey x = SerialisedKey (fromUnsliced (to x))++instance Show (Unsliced SerialisedKey) where+ show x = show (fromUnslicedKey x)++instance Eq (Unsliced SerialisedKey) where+ Unsliced ba1 == Unsliced ba2 = SBS ba1' == SBS ba2'+ where+ !(ByteArray ba1') = ba1+ !(ByteArray ba2') = ba2++instance Ord (Unsliced SerialisedKey) where+ compare (Unsliced ba1) (Unsliced ba2) = compare (SBS ba1') (SBS ba2')+ where+ !(ByteArray ba1') = ba1+ !(ByteArray ba2') = ba2
@@ -0,0 +1,138 @@+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.Vector (+ mkPrimVector,+ byteVectorFromPrim,+ noRetainedExtraMemory,+ primArrayToPrimVector,+ mapStrict,+ mapMStrict,+ imapMStrict,+ forMStrict,+ zipWithStrict,+ binarySearchL,+ unsafeInsertWithMStrict,+ unfoldrNM',+) where++import Control.Monad+import Control.Monad.Primitive (PrimMonad, PrimState)+import qualified Data.Primitive as P+import Data.Primitive.ByteArray (ByteArray, newByteArray,+ runByteArray, sizeofByteArray, writeByteArray)+import Data.Primitive.Types (Prim (sizeOfType#), sizeOfType)+import Data.Proxy (Proxy (..))+import qualified Data.Vector as V+import qualified Data.Vector.Algorithms.Search as VA+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Primitive as VP+import Data.Word (Word8)+import Database.LSMTree.Internal.Assertions+import GHC.Exts (Int (..))+import GHC.ST (runST)++mkPrimVector :: forall a. Prim a => Int -> Int -> ByteArray -> VP.Vector a+mkPrimVector off len ba =+ assert (isValidSlice (off * sizeof) (len * sizeof) ba) $+ VP.Vector off len ba+ where+ sizeof = I# (sizeOfType# (Proxy @a))+{-# INLINE mkPrimVector #-}++byteVectorFromPrim :: forall a. Prim a => a -> VP.Vector Word8+byteVectorFromPrim prim = mkPrimVector 0 (sizeOfType @a) $+ runByteArray $ do+ rep <- newByteArray (sizeOfType @a)+ writeByteArray rep 0 prim+ pure rep+{-# INLINE byteVectorFromPrim #-}++noRetainedExtraMemory :: forall a. Prim a => VP.Vector a -> Bool+noRetainedExtraMemory (VP.Vector off len ba) =+ off == 0 && len * sizeof == sizeofByteArray ba+ where+ sizeof = I# (sizeOfType# (Proxy @a))++{-# INLINE primArrayToPrimVector #-}+primArrayToPrimVector :: Prim a => P.PrimArray a -> VP.Vector a+primArrayToPrimVector pa@(P.PrimArray ba) =+ VP.Vector 0 (P.sizeofPrimArray pa) (P.ByteArray ba)++{-# INLINE mapStrict #-}+-- | /( O(n) /) Like 'V.map', but strict in the produced elements of type @b@.+mapStrict :: forall a b. (a -> b) -> V.Vector a -> V.Vector b+mapStrict f v = runST (V.mapM (\x -> pure $! f x) v)++{-# INLINE mapMStrict #-}+-- | /( O(n) /) Like 'V.mapM', but strict in the produced elements of type @b@.+mapMStrict :: Monad m => (a -> m b) -> V.Vector a -> m (V.Vector b)+mapMStrict f v = V.mapM (f >=> (pure $!)) v++{-# INLINE imapMStrict #-}+-- | /( O(n) /) Like 'V.imapM', but strict in the produced elements of type @b@.+imapMStrict :: Monad m => (Int -> a -> m b) -> V.Vector a -> m (V.Vector b)+imapMStrict f v = V.imapM (\i -> f i >=> (pure $!)) v++{-# INLINE zipWithStrict #-}+-- | /( O(min(m,n)) /) Like 'V.zipWithM', but strict in the produced elements of+-- type @c@.+zipWithStrict :: forall a b c. (a -> b -> c) -> V.Vector a -> V.Vector b -> V.Vector c+zipWithStrict f xs ys = runST (V.zipWithM (\x y -> pure $! f x y) xs ys)++-- | /( O(n) /) Like 'V.forM', but strict in the produced elements of type @b@.+{-# INLINE forMStrict #-}+forMStrict :: Monad m => V.Vector a -> (a -> m b) -> m (V.Vector b)+forMStrict xs f = V.forM xs (f >=> (pure $!))++{-|+ Finds the lowest index in a given sorted vector at which the given element+ could be inserted while maintaining the sortedness.++ This is a variant of 'Data.Vector.Algorithms.Search.binarySearchL' for+ immutable vectors.+-}+binarySearchL :: Ord a => V.Vector a -> a -> Int+binarySearchL vec val = runST $ V.unsafeThaw vec >>= flip VA.binarySearchL val+{-# INLINE binarySearchL #-}++{-# INLINE unsafeInsertWithMStrict #-}+-- | Insert (in a broad sense) an entry in a mutable vector at a given index,+-- but if a @Just@ entry already exists at that index, combine the two entries+-- using @f@.+unsafeInsertWithMStrict ::+ PrimMonad m+ => VM.MVector (PrimState m) (Maybe a)+ -> (a -> a -> a) -- ^ function @f@, called as @f new old@+ -> Int+ -> a+ -> m ()+unsafeInsertWithMStrict mvec f i y = VM.unsafeModifyM mvec g i+ where+ g x = pure $! Just $! maybe y (`f` y) x++{-# INLINE unfoldrNM' #-}+-- | A version of 'V.unfoldrNM' that also returns the final state.+--+-- /O(n)/ Construct a vector by repeatedly applying the monadic generator+-- function to a seed. The generator function also yields 'Just' the next+-- element or 'Nothing' if there are no more elements.+--+-- The state as well as all elements of the result vector are forced to weak+-- head normal form.+unfoldrNM' :: PrimMonad m => Int -> (b -> m (Maybe a, b)) -> b -> m (V.Vector a, b)+unfoldrNM' len f = \b0 -> do+ vec <- VM.unsafeNew len+ go vec 0 b0+ where+ go !vec !n !b+ | n >= len = (, b) <$!> V.unsafeFreeze vec+ | otherwise =+ f b >>= \case+ (Nothing, !b') ->+ (, b') <$!> V.unsafeFreeze (VM.slice 0 n vec)+ (Just !a, !b') -> do+ VM.unsafeWrite vec n a+ go vec (n+1) b'
@@ -0,0 +1,131 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+{-# OPTIONS_HADDOCK not-home #-}+{- HLINT ignore "Avoid restricted alias" -}++-- | Vectors with support for appending elements.+module Database.LSMTree.Internal.Vector.Growing+(+ GrowingVector (GrowingVector),+ new,+ append,+ freeze,+ readMaybeLast+)+where++import Prelude hiding (init, last, length, read)++import Control.Monad (when)+import Control.Monad.ST.Strict (ST)+import Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,+ writePrimVar)+import Data.STRef.Strict (STRef, newSTRef, readSTRef, writeSTRef)+import Data.Vector (Vector)+import qualified Data.Vector as Mutable (freeze)+import Data.Vector.Mutable (MVector)+import qualified Data.Vector.Mutable as Mutable (grow, length, new, read, set,+ slice, take)++{-|+ A vector with support for appending elements.++ Internally, the elements of a growing vector are stored in a buffer. This+ buffer is automatically enlarged whenever this is needed for storing+ additional elements. On each such enlargement, the size of the buffer is+ multiplied by a power of 2, whose exponent is chosen just big enough to make+ the final buffer size at least as high as the new vector length.++ Note that, while buffer sizes and vector lengths are represented as 'Int'+ values internally, the above-described buffer enlargement scheme has the+ consequence that the largest possible buffer size and thus the largest+ possible vector length may not be the maximum 'Int' value. That said, they+ are always greater than half the maximum 'Int' value, which should be enough+ for all practical purposes.+-}+data GrowingVector s a = GrowingVector+ !(STRef s (MVector s a)) -- Reference to the buffer+ !(PrimVar s Int) -- Reference to the length++-- | Creates a new, initially empty, vector.+new :: Int -- ^ Initial buffer size+ -> ST s (GrowingVector s a) -- ^ Construction of the vector+new illegalInitialBufferSize | illegalInitialBufferSize <= 0+ = error "Initial buffer size not positive"+new initialBufferSize+ = do+ buffer <- Mutable.new initialBufferSize+ bufferRef <- newSTRef $! buffer+ lengthRef <- newPrimVar 0+ pure (GrowingVector bufferRef lengthRef)++{-|+ Appends a value a certain number of times to a vector. If a negative number+ is provided as the count, the vector is not changed.+-}+append :: forall s a . GrowingVector s a -> Int -> a -> ST s ()+append _ pseudoCount _ | pseudoCount <= 0+ = pure ()+append (GrowingVector bufferRef lengthRef) count val+ = do+ length <- readPrimVar lengthRef+ makeRoom+ buffer' <- readSTRef bufferRef+ Mutable.set (Mutable.slice length count buffer') $! val+ where++ makeRoom :: ST s ()+ makeRoom = do+ length <- readPrimVar lengthRef+ when (count > maxBound - length) (error "New length too large")+ buffer <- readSTRef bufferRef+ let++ bufferSize :: Int+ !bufferSize = Mutable.length buffer++ length' :: Int+ !length' = length + count++ when (bufferSize < length') $ do+ let++ higherBufferSizes :: [Int]+ higherBufferSizes = tail (init ++ [last]) where++ init :: [Int]+ last :: Int+ (init, last : _) = span (<= maxBound `div` 2) $+ iterate (* 2) bufferSize+ {-NOTE:+ In order to prevent overflow, we have to start with the+ current buffer size here, although we know that it is+ not sufficient.+ -}++ sufficientBufferSizes :: [Int]+ sufficientBufferSizes = dropWhile (< length') higherBufferSizes++ case sufficientBufferSizes of+ []+ -> error "No sufficient buffer size available"+ bufferSize' : _+ -> Mutable.grow buffer (bufferSize' - bufferSize) >>=+ (writeSTRef bufferRef $!)+ writePrimVar lengthRef length'++-- | Turns a growing vector into an ordinary vector.+freeze :: GrowingVector s a -> ST s (Vector a)+freeze (GrowingVector bufferRef lengthRef) = do+ buffer <- readSTRef bufferRef+ length <- readPrimVar lengthRef+ Mutable.freeze (Mutable.take length buffer)++-- | Reads the last element of a growing vector if it exists.+readMaybeLast :: GrowingVector s a -> ST s (Maybe a)+readMaybeLast (GrowingVector bufferRef lengthRef) = do+ length <- readPrimVar lengthRef+ if length == 0+ then pure Nothing+ else do+ buffer <- readSTRef bufferRef+ Just <$> Mutable.read buffer (pred length)
@@ -0,0 +1,132 @@+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | The in-memory LSM level 0.+--+module Database.LSMTree.Internal.WriteBuffer (+ WriteBuffer,+ empty,+ numEntries,+ fromMap,+ toMap,+ fromList,+ toList,+ addEntry,+ null,+ lookups,+ lookup,+ rangeLookups,+) where++import Control.DeepSeq (NFData (..))+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import qualified Data.Vector as V+import Database.LSMTree.Internal.BlobRef (BlobSpan)+import Database.LSMTree.Internal.Entry+import qualified Database.LSMTree.Internal.Map.Range as Map.R+import Database.LSMTree.Internal.Range (Range (..))+import Database.LSMTree.Internal.Serialise+import qualified Database.LSMTree.Internal.Vector as V+import Prelude hiding (lookup, null)++{-------------------------------------------------------------------------------+ Writebuffer type+-------------------------------------------------------------------------------}++newtype WriteBuffer =+ WB { unWB :: Map SerialisedKey (Entry SerialisedValue BlobSpan) }+ deriving stock (Eq, Show)+ deriving newtype NFData++empty :: WriteBuffer+empty = WB Map.empty++-- | \( O(1) \)+numEntries :: WriteBuffer -> NumEntries+numEntries (WB m) = NumEntries (Map.size m)++-- | \( O(1)) \)+fromMap ::+ Map SerialisedKey (Entry SerialisedValue BlobSpan)+ -> WriteBuffer+fromMap m = WB m++-- | \( O(1) \)+toMap :: WriteBuffer -> Map SerialisedKey (Entry SerialisedValue BlobSpan)+toMap = unWB++-- | \( O(n \log n) \)+fromList ::+ ResolveSerialisedValue -- ^ merge function+ -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]+ -> WriteBuffer+fromList f es = WB $ Map.fromListWith (combine f) es++-- | \( O(n) \)+toList :: WriteBuffer -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]+toList (WB m) = Map.assocs m++{-------------------------------------------------------------------------------+ Updates+-------------------------------------------------------------------------------}++addEntry ::+ ResolveSerialisedValue -- ^ merge function+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan+ -> WriteBuffer+ -> WriteBuffer+addEntry f k e (WB wb) =+ WB (Map.insertWith (combine f) k e wb)++{-------------------------------------------------------------------------------+ Querying+-------------------------------------------------------------------------------}++null :: WriteBuffer -> Bool+null (WB m) = Map.null m++-- We return an 'Entry' with serialised values, so it can be properly combined+-- with the lookups in other runs. Deserialisation only occurs afterwards.+--+-- Note: the entry may be 'Delete'.+--+lookups ::+ WriteBuffer+ -> V.Vector SerialisedKey+ -> V.Vector (Maybe (Entry SerialisedValue BlobSpan))+lookups (WB !m) !ks = V.mapStrict (`Map.lookup` m) ks++lookup ::+ WriteBuffer+ -> SerialisedKey+ -> Maybe (Entry SerialisedValue BlobSpan)+lookup (WB !m) !k = Map.lookup k m++{-------------------------------------------------------------------------------+ RangeQueries+-------------------------------------------------------------------------------}++-- | We return 'Entry' so we can properly combine lookup results.+--+-- Note: 'Delete's are not filtered out.+--+rangeLookups ::+ WriteBuffer+ -> Range SerialisedKey+ -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]+rangeLookups (WB m) r =+ [ (k, e)+ | let (lb, ub) = convertRange r+ , (k, e) <- Map.R.rangeLookup lb ub m+ ]++convertRange :: Range k -> (Map.R.Bound k, Map.R.Bound k)+convertRange (FromToExcluding lb ub) =+ ( Map.R.Bound lb Map.R.Inclusive+ , Map.R.Bound ub Map.R.Exclusive )+convertRange (FromToIncluding lb ub) =+ ( Map.R.Bound lb Map.R.Inclusive+ , Map.R.Bound ub Map.R.Inclusive )
@@ -0,0 +1,270 @@+{-# OPTIONS_GHC -Wno-partial-fields #-}+{-# OPTIONS_HADDOCK not-home #-}++-- | An on-disk store for blobs for the write buffer.+--+-- For table inserts with blobs, the blob get written out immediately to a+-- file, while the rest of the 'Entry' goes into the 'WriteBuffer'. The+-- 'WriteBufferBlobs' manages the storage of the blobs.+--+-- A single write buffer blob file can be shared between multiple tables. As a+-- consequence, the lifetime of the 'WriteBufferBlobs' must be managed using+-- 'new', and with the 'Ref' API: 'releaseRef' and 'dupRef'. When a table is+-- duplicated, the new table needs its own reference, so use 'dupRef' upon+-- duplication.+--+-- Blobs are copied from the write buffer blob file when the write buffer is+-- flushed to make a run. This is needed since the blob file is shared and so+-- not stable by the time one table wants to flush it.+--+-- Not all tables need a blob file so we defer opening the file until it+-- is needed.+--+module Database.LSMTree.Internal.WriteBufferBlobs (+ WriteBufferBlobs (..),+ new,+ open,+ addBlob,+ mkRawBlobRef,+ mkWeakBlobRef,+ -- * For tests+ FilePointer (..)+ ) where++import Control.DeepSeq (NFData (..))+import Control.Monad (void)+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive (PrimMonad, PrimState)+import Control.RefCount+import Data.Primitive.PrimVar as P+import Data.Word (Word64)+import Database.LSMTree.Internal.BlobFile+import qualified Database.LSMTree.Internal.BlobFile as BlobFile+import Database.LSMTree.Internal.BlobRef (RawBlobRef (..),+ WeakBlobRef (..))+import Database.LSMTree.Internal.Serialise+import qualified System.FS.API as FS+import System.FS.API (HasFS)++-- | A single 'WriteBufferBlobs' may be shared between multiple tables.+-- As a consequence of being shared, the management of the shared state has to+-- be quite careful.+--+-- In particular there is the blob file itself. We may have to write to this+-- blob file from multiple threads on behalf of independent tables.+-- The offset at which we write is thus shared mutable state. Our strategy for+-- the write offset is to explicitly track it (since we need to know the offset+-- to return correct 'BlobSpan's) and then to not use the file's own file+-- pointer. We do this by always writing at specific file offsets rather than+-- writing at the open file's normal file pointer. We use a 'PrimVar' with+-- atomic operations to manage the file offset.+--+-- A consequence of the blob file being shared between the write buffers of+-- many tables is that the blobs in the file will not all belong to one+-- table. The write buffer blob file is unsuitable to use as-is as the+-- blob file for a run when the write buffer is flushed. The run blob file must+-- be immutable and with a known CRC. Whereas because the write buffer blob+-- file is shared, it can still be appended to via inserts in one table+-- while another is trying to flush the write buffer. So there is no stable CRC+-- for the whole file (as required by the snapshot format). Further more we+-- cannot even incrementally calculate the blob file CRC without additional+-- expensive serialisation. To solve this we follow the design that the open+-- file handle for the blob file is only shared between multiple write buffers,+-- and is /not/ shared with the runs once flushed. This separates the lifetimes+-- of the files. Correspondingly, the reference counter is only for+-- tracking the lifetime of the read\/write mode file handle.+--+-- One concern with sharing blob files and the open blob file handle between+-- multiple write buffers is: can we guarantee that the blob file is eventually+-- closed?+--+-- A problematic example would be, starting from a root handle and then+-- repeatedly: duplicating; inserting (with blobs) into the duplicate; and then+-- closing it. This would use only a fixed number of tables at once, but+-- would keep inserting into the same the write buffer blob file. This could be+-- done indefinitely.+--+-- On the other hand, provided that there's a bound on the number of duplicates+-- that are created from any point, and each table is eventually closed,+-- then each write buffer blob file will eventually be closed.+--+-- The latter seems like the more realistic use case, and so the design here is+-- probably reasonable.+--+-- If not, an entirely different approach would be to manage blobs across all+-- runs (and the write buffer) differently: avoiding copying when blobs are+-- merged and using some kind of GC algorithm to recover space for blobs that+-- are not longer needed. There are LSM algorithms that do this for values+-- (i.e. copying keys only during merge and referring to values managed in a+-- separate disk heap), so the same could be applied to blobs.+--+data WriteBufferBlobs m h =+ WriteBufferBlobs {++ -- | The blob file+ --+ -- INVARIANT: the file may contain garbage bytes, but no blob reference+ -- ('RawBlobRef', 'WeakBlobRef', or 'StrongBlobRef) will reference these+ -- bytes.+ blobFile :: !(Ref (BlobFile m h))++ -- | The manually tracked file pointer.+ --+ -- INVARIANT: the file pointer points to a file offset at or beyond the+ -- file size.+ , blobFilePointer :: !(FilePointer m)++ -- The 'WriteBufferBlobs' is a shared reference-counted object type+ , writeBufRefCounter :: !(RefCounter m)+ }++instance NFData h => NFData (WriteBufferBlobs m h) where+ rnf (WriteBufferBlobs a b c) = rnf a `seq` rnf b `seq` rnf c++instance RefCounted m (WriteBufferBlobs m h) where+ getRefCounter = writeBufRefCounter++{-# SPECIALISE new :: HasFS IO h -> FS.FsPath -> IO (Ref (WriteBufferBlobs IO h)) #-}+-- | Create a new 'WriteBufferBlobs' with a new file.+--+-- REF: the resulting reference must be released once it is no longer used.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates/creates resources.+new ::+ (PrimMonad m, MonadMask m)+ => HasFS m h+ -> FS.FsPath+ -> m (Ref (WriteBufferBlobs m h))+new fs blobFileName = open fs blobFileName FS.MustBeNew++{-# SPECIALISE open :: HasFS IO h -> FS.FsPath -> FS.AllowExisting -> IO (Ref (WriteBufferBlobs IO h)) #-}+-- | Open a `WriteBufferBlobs` file and sets the file pointer to the end of the file.+--+-- REF: the resulting reference must be released once it is no longer used.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates/creates resources.+open ::+ (PrimMonad m, MonadMask m)+ => HasFS m h+ -> FS.FsPath+ -> FS.AllowExisting+ -> m (Ref (WriteBufferBlobs m h))+open fs blobFileName blobFileAllowExisting = do+ -- Must use read/write mode because we write blobs when adding, but+ -- we can also be asked to retrieve blobs at any time.+ bracketOnError+ (openBlobFile fs blobFileName (FS.ReadWriteMode blobFileAllowExisting))+ releaseRef+ (fromBlobFile fs)++{-# SPECIALISE fromBlobFile :: HasFS IO h -> Ref (BlobFile IO h) -> IO (Ref (WriteBufferBlobs IO h)) #-}+-- | Make a `WriteBufferBlobs` from a `BlobFile` and set the file pointer to the+-- end of the file.+--+-- REF: the resulting reference must be released once it is no longer used.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates/creates resources.+fromBlobFile ::+ (PrimMonad m, MonadMask m)+ => HasFS m h+ -> Ref (BlobFile m h)+ -> m (Ref (WriteBufferBlobs m h))+fromBlobFile fs blobFile = do+ blobFilePointer <- newFilePointer+ -- Set the blob file pointer to the end of the file+ blobFileSize <- withRef blobFile $ FS.hGetSize fs . blobFileHandle+ void . updateFilePointer blobFilePointer . fromIntegral $ blobFileSize+ newRef (releaseRef blobFile) $ \writeBufRefCounter ->+ WriteBufferBlobs {+ blobFile,+ blobFilePointer,+ writeBufRefCounter+ }++{-# SPECIALISE addBlob :: HasFS IO h -> Ref (WriteBufferBlobs IO h) -> SerialisedBlob -> IO BlobSpan #-}+-- | Append a blob.+--+-- If no exception is returned, then the file pointer will be set to exactly the+-- file size.+--+-- If an exception is returned, the file pointer points to a file+-- offset at or beyond the file size. The bytes between the old and new offset+-- might be garbage or missing.+addBlob :: (PrimMonad m, MonadThrow m)+ => HasFS m h+ -> Ref (WriteBufferBlobs m h)+ -> SerialisedBlob+ -> m BlobSpan+addBlob fs (DeRef WriteBufferBlobs {blobFile, blobFilePointer}) blob = do+ let blobsize = sizeofBlob blob+ -- If an exception happens after updating the file pointer, then no write+ -- takes place. The next 'addBlob' will start writing at the new file+ -- offset, so there are going to be some uninitialised bytes in the file.+ bloboffset <- updateFilePointer blobFilePointer blobsize+ -- If an exception happens while writing the blob, the bytes in the file+ -- might be corrupted.+ BlobFile.writeBlob fs blobFile blob bloboffset+ pure BlobSpan {+ blobSpanOffset = bloboffset,+ blobSpanSize = fromIntegral blobsize+ }++-- | Helper function to make a 'RawBlobRef' that points into a+-- 'WriteBufferBlobs'.+--+-- This function should only be used on the result of 'addBlob' on the same+-- 'WriteBufferBlobs'. For example:+--+-- @+-- 'addBlob' hfs wbb blob >>= \\span -> pure ('mkRawBlobRef' wbb span)+-- @+mkRawBlobRef :: Ref (WriteBufferBlobs m h)+ -> BlobSpan+ -> RawBlobRef m h+mkRawBlobRef (DeRef WriteBufferBlobs {blobFile = DeRef blobfile}) blobspan =+ RawBlobRef {+ rawBlobRefFile = blobfile,+ rawBlobRefSpan = blobspan+ }++-- | Helper function to make a 'WeakBlobRef' that points into a+-- 'WriteBufferBlobs'.+--+-- This function should only be used on the result of 'addBlob' on the same+-- 'WriteBufferBlobs'. For example:+--+-- @+-- 'addBlob' hfs wbb blob >>= \\span -> pure ('mkWeakBlobRef' wbb span)+-- @+mkWeakBlobRef :: Ref (WriteBufferBlobs m h)+ -> BlobSpan+ -> WeakBlobRef m h+mkWeakBlobRef (DeRef WriteBufferBlobs {blobFile}) blobspan =+ WeakBlobRef {+ weakBlobRefFile = mkWeakRef blobFile,+ weakBlobRefSpan = blobspan+ }+++-- | A mutable file offset, suitable to share between threads.+--+-- This pointer is limited to 31-bit file offsets on 32-bit systems. This should+-- be a sufficiently large limit that we never reach it in practice.+newtype FilePointer m = FilePointer (PrimVar (PrimState m) Int)++instance NFData (FilePointer m) where+ rnf (FilePointer var) = var `seq` ()++{-# SPECIALISE newFilePointer :: IO (FilePointer IO) #-}+newFilePointer :: PrimMonad m => m (FilePointer m)+newFilePointer = FilePointer <$> P.newPrimVar 0++{-# SPECIALISE updateFilePointer :: FilePointer IO -> Int -> IO Word64 #-}+-- | Update the file offset by a given amount and return the new offset. This+-- is safe to use concurrently.+--+updateFilePointer :: PrimMonad m => FilePointer m -> Int -> m Word64+updateFilePointer (FilePointer var) n = fromIntegral <$> P.fetchAddInt var n
@@ -0,0 +1,190 @@+{-# OPTIONS_HADDOCK not-home #-}++-- | A persisted write buffer that is being read incrementally from disk.+--+module Database.LSMTree.Internal.WriteBufferReader (+ readWriteBuffer+ ) where++import Control.Concurrent.Class.MonadMVar.Strict+import Control.Monad.Class.MonadST (MonadST (..))+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadMask, MonadThrow (..),+ bracketOnError)+import Control.Monad.Primitive (PrimMonad (..))+import Control.RefCount (Ref, dupRef, releaseRef)+import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar,+ writeMutVar)+import Data.Primitive.PrimVar+import Data.Word (Word16)+import Database.LSMTree.Internal.BlobFile (BlobFile)+import Database.LSMTree.Internal.BlobRef (RawBlobRef (..),+ mkRawBlobRef)+import qualified Database.LSMTree.Internal.Entry as E+import Database.LSMTree.Internal.Paths+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.RunReader (Entry (..), Result (..),+ mkEntryOverflow, readDiskPage, readOverflowPages,+ toFullEntry)+import Database.LSMTree.Internal.Serialise (ResolveSerialisedValue,+ SerialisedValue)+import Database.LSMTree.Internal.WriteBuffer (WriteBuffer)+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)++{-# SPECIALISE+ readWriteBuffer ::+ ResolveSerialisedValue+ -> HasFS IO h+ -> HasBlockIO IO h+ -> ForKOps FS.FsPath+ -> Ref (BlobFile IO h)+ -> IO WriteBuffer+ #-}+-- | Read a serialised `WriteBuffer` back into memory.+--+-- The argument blob file ('BlobFile') must be the file associated with the+-- argument key\/ops file ('ForKOps'). 'readWriteBuffer' does not check this.+readWriteBuffer ::+ (MonadMVar m, MonadMask m, MonadSTM m, MonadST m)+ => ResolveSerialisedValue+ -> HasFS m h+ -> HasBlockIO m h+ -> ForKOps FS.FsPath+ -> Ref (BlobFile m h)+ -> m WriteBuffer+readWriteBuffer resolve hfs hbio kOpsPath blobFile =+ bracket (new hfs hbio kOpsPath blobFile) close $ readEntries+ where+ readEntries reader = readEntriesAcc WB.empty+ where+ readEntriesAcc acc = next reader >>= \case+ Empty -> pure acc+ ReadEntry key entry -> readEntriesAcc $+ WB.addEntry resolve key (rawBlobRefSpan <$> toFullEntry entry) acc++-- | Allows reading the k\/ops of a serialised write buffer incrementally,+-- using its own read-only file handle and in-memory cache of the current disk page.+--+-- New pages are loaded when trying to read their first entry.+data WriteBufferReader m h = WriteBufferReader {+ -- | The disk page currently being read. If it is 'Nothing', the reader+ -- is considered closed.+ readerCurrentPage :: !(MutVar (PrimState m) (Maybe RawPage))+ -- | The index of the entry to be returned by the next call to 'next'.+ , readerCurrentEntryNo :: !(PrimVar (PrimState m) Word16)+ , readerKOpsHandle :: !(FS.Handle h)+ , readerBlobFile :: !(Ref (BlobFile m h))+ , readerHasFS :: !(HasFS m h)+ , readerHasBlockIO :: !(HasBlockIO m h)+ }++{-# SPECIALISE+ new ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ForKOps FS.FsPath+ -> Ref (BlobFile IO h)+ -> IO (WriteBufferReader IO h)+ #-}+-- | See 'Database.LSMTree.Internal.RunReader.new'.+--+-- REF: the resulting 'WriteBufferReader' must be closed once it is no longer+-- used.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- allocates/creates resources.+new :: forall m h.+ (MonadMVar m, MonadST m, MonadMask m)+ => HasFS m h+ -> HasBlockIO m h+ -> ForKOps FS.FsPath+ -> Ref (BlobFile m h)+ -> m (WriteBufferReader m h)+new readerHasFS readerHasBlockIO kOpsPath blobFile =+ bracketOnError openKOps (FS.hClose readerHasFS) $ \readerKOpsHandle -> do+ -- Double the file readahead window (only applies to this file descriptor)+ FS.hAdviseAll readerHasBlockIO readerKOpsHandle FS.AdviceSequential+ bracketOnError (dupRef blobFile) releaseRef $ \readerBlobFile -> do+ -- Load first page from disk, if it exists.+ readerCurrentEntryNo <- newPrimVar (0 :: Word16)+ firstPage <- readDiskPage readerHasFS readerKOpsHandle+ readerCurrentPage <- newMutVar firstPage+ pure $ WriteBufferReader{..}+ where+ openKOps = FS.hOpen readerHasFS (unForKOps kOpsPath) FS.ReadMode++{-# SPECIALISE+ next ::+ WriteBufferReader IO h+ -> IO (Result IO h)+ #-}+-- | See 'Database.LSMTree.Internal.RunReader.next'.+--+-- TODO: 'next' is currently only used in 'readWriteBuffer', where it is a safe+-- use of an unsafe function. If this function is ever exported and used+-- directly, the TODOs in the body of this function should be addressed first.+next :: forall m h.+ (MonadSTM m, MonadST m, MonadMask m)+ => WriteBufferReader m h+ -> m (Result m h)+next WriteBufferReader {..} = do+ readMutVar readerCurrentPage >>= \case+ Nothing ->+ pure Empty+ Just page -> do+ entryNo <- readPrimVar readerCurrentEntryNo+ go entryNo page+ where+ -- TODO: if 'readerCurrentEntryNo' is incremented but an exception is thrown+ -- before the 'Result' is used by the caller of 'next', then we'll lose that+ -- 'Result'. The following call to 'next' will not return the 'Result' we+ -- missed.+ go :: Word16 -> RawPage -> m (Result m h)+ go !entryNo !page =+ -- take entry from current page (resolve blob if necessary)+ case rawPageIndex page entryNo of+ IndexNotPresent -> do+ -- if it is past the last one, load a new page from disk, try again+ newPage <- readDiskPage readerHasFS readerKOpsHandle+ -- TODO: if the next disk page is read but an (async) exception is+ -- thrown just before updating the MutVar below, then we lose the+ -- disk page because 'readDiskPage' has already updated its file+ -- pointer.+ stToIO $ writeMutVar readerCurrentPage newPage+ case newPage of+ Nothing -> do+ pure Empty+ Just p -> do+ writePrimVar readerCurrentEntryNo 0+ go 0 p -- try again on the new page+ IndexEntry key entry -> do+ modifyPrimVar readerCurrentEntryNo (+1)+ let entry' :: E.Entry SerialisedValue (RawBlobRef m h)+ entry' = fmap (mkRawBlobRef readerBlobFile) entry+ let rawEntry = Entry entry'+ pure (ReadEntry key rawEntry)+ IndexEntryOverflow key entry lenSuffix -> do+ -- TODO: we know that we need the next page, could already load?+ modifyPrimVar readerCurrentEntryNo (+1)+ let entry' :: E.Entry SerialisedValue (RawBlobRef m h)+ entry' = fmap (mkRawBlobRef readerBlobFile) entry+ overflowPages <- readOverflowPages readerHasFS readerKOpsHandle lenSuffix+ let rawEntry = mkEntryOverflow entry' page lenSuffix overflowPages+ pure (ReadEntry key rawEntry)++{-# SPECIALISE close :: WriteBufferReader IO h -> IO () #-}+-- | Close the 'WriteBufferReader'.+--+-- ASYNC: this should be called with asynchronous exceptions masked because it+-- releases/removes resources.+close ::+ (MonadMask m, PrimMonad m)+ => WriteBufferReader m h+ -> m ()+close WriteBufferReader{..} = do+ FS.hClose readerHasFS readerKOpsHandle+ `finally` releaseRef readerBlobFile
@@ -0,0 +1,242 @@+{-# OPTIONS_HADDOCK not-home #-}++module Database.LSMTree.Internal.WriteBufferWriter+ ( writeWriteBuffer+ ) where++import Control.Exception (assert)+import Control.Monad (void, when)+import Control.Monad.Class.MonadST (MonadST (..))+import qualified Control.Monad.Class.MonadST as ST+import Control.Monad.Class.MonadSTM (MonadSTM (..))+import Control.Monad.Class.MonadThrow (MonadThrow)+import Control.Monad.Primitive (PrimMonad (..))+import Control.Monad.ST (ST)+import Control.RefCount (Ref)+import Data.Foldable (for_)+import Data.Maybe (maybeToList)+import Data.Primitive.PrimVar (PrimVar, newPrimVar)+import Data.Word (Word64)+import Database.LSMTree.Internal.BlobFile (BlobSpan)+import Database.LSMTree.Internal.BlobRef (RawBlobRef)+import Database.LSMTree.Internal.ChecksumHandle (ChecksumHandle,+ closeHandle, copyBlob, dropCache, makeHandle, readChecksum,+ writeRawOverflowPages, writeRawPage)+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.Entry (Entry)+import Database.LSMTree.Internal.PageAcc (PageAcc)+import qualified Database.LSMTree.Internal.PageAcc as PageAcc+import qualified Database.LSMTree.Internal.PageAcc1 as PageAcc+import Database.LSMTree.Internal.Paths (ForBlob (..), ForKOps (..),+ WriteBufferFsPaths, toChecksumsFileForWriteBufferFiles,+ writeBufferBlobPath, writeBufferChecksumsPath,+ writeBufferKOpsPath)+import Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)+import Database.LSMTree.Internal.RawPage (RawPage)+import Database.LSMTree.Internal.Serialise (SerialisedKey,+ SerialisedValue)+import Database.LSMTree.Internal.WriteBuffer (WriteBuffer)+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import Database.LSMTree.Internal.WriteBufferBlobs (WriteBufferBlobs)+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)+++{-# SPECIALISE+ writeWriteBuffer ::+ HasFS IO h+ -> HasBlockIO IO h+ -> WriteBufferFsPaths+ -> WriteBuffer+ -> Ref (WriteBufferBlobs IO h)+ -> IO ()+ #-}+-- | Write a 'WriteBuffer' to disk.+writeWriteBuffer ::+ (MonadSTM m, MonadST m, MonadThrow m)+ => HasFS m h+ -> HasBlockIO m h+ -> WriteBufferFsPaths+ -> WriteBuffer+ -> Ref (WriteBufferBlobs m h)+ -> m ()+writeWriteBuffer hfs hbio fsPaths buffer blobs = do+ writer <- new hfs hbio fsPaths+ for_ (WB.toList buffer) $ \(key, op) ->+ -- TODO: The fmap entry here reallocates even when there are no blobs.+ addKeyOp writer key (WBB.mkRawBlobRef blobs <$> op)+ void $ unsafeFinalise True writer++-- | The in-memory representation of an LSM 'WriteBuffer' that is in the process of being serialised to disk.+data WriteBufferWriter m h = WriteBufferWriter+ { -- | The file system paths for all the files used by the serialised write buffer.+ writerFsPaths :: !WriteBufferFsPaths,+ -- | The page accumulator.+ writerPageAcc :: !(PageAcc (PrimState m)),+ -- | The byte offset within the blob file for the next blob to be written.+ writerBlobOffset :: !(PrimVar (PrimState m) Word64),+ -- | The (write mode) file handles.+ writerKOpsHandle :: !(ForKOps (ChecksumHandle (PrimState m) h)),+ writerBlobHandle :: !(ForBlob (ChecksumHandle (PrimState m) h)),+ writerHasFS :: !(HasFS m h),+ writerHasBlockIO :: !(HasBlockIO m h)+ }++{-# SPECIALISE+ new ::+ HasFS IO h+ -> HasBlockIO IO h+ -> WriteBufferFsPaths+ -> IO (WriteBufferWriter IO h)+ #-}+-- | Create a 'WriteBufferWriter' to start serialising a 'WriteBuffer'.+--+-- See 'Database.LSMTree.Internal.RunBuilder.new'.+--+-- NOTE: 'new' assumes that the directory passed via 'WriteBufferFsPaths' exists.+new ::+ (MonadST m, MonadSTM m)+ => HasFS m h+ -> HasBlockIO m h+ -> WriteBufferFsPaths+ -> m (WriteBufferWriter m h)+new hfs hbio fsPaths = do+ writerPageAcc <- ST.stToIO PageAcc.newPageAcc+ writerBlobOffset <- newPrimVar 0+ writerKOpsHandle <- ForKOps <$> makeHandle hfs (writeBufferKOpsPath fsPaths)+ writerBlobHandle <- ForBlob <$> makeHandle hfs (writeBufferBlobPath fsPaths)+ pure WriteBufferWriter+ { writerFsPaths = fsPaths,+ writerHasFS = hfs,+ writerHasBlockIO = hbio,+ ..+ }++{-# SPECIALISE+ unsafeFinalise ::+ Bool+ -> WriteBufferWriter IO h+ -> IO (HasFS IO h, HasBlockIO IO h, WriteBufferFsPaths)+ #-}+-- | Finalise an incremental 'WriteBufferWriter'.+--+-- Do /not/ use a 'WriteBufferWriter' after finalising it.+--+-- See 'Database.LSMTree.Internal.RunBuilder.unsafeFinalise'.+--+-- TODO: Ensure proper cleanup even in presence of exceptions.+unsafeFinalise ::+ (MonadST m, MonadSTM m, MonadThrow m)+ => Bool -- ^ drop caches+ -> WriteBufferWriter m h+ -> m (HasFS m h, HasBlockIO m h, WriteBufferFsPaths)+unsafeFinalise dropCaches WriteBufferWriter {..} = do+ -- write final bits+ mPage <- ST.stToIO $ flushPageIfNonEmpty writerPageAcc+ for_ mPage $ writeRawPage writerHasFS writerKOpsHandle+ kOpsChecksum <- traverse readChecksum writerKOpsHandle+ blobChecksum <- traverse readChecksum writerBlobHandle+ let checksums = toChecksumsFileForWriteBufferFiles (kOpsChecksum, blobChecksum)+ FS.withFile writerHasFS (writeBufferChecksumsPath writerFsPaths) (FS.WriteMode FS.MustBeNew) $ \h -> do+ CRC.writeChecksumsFile' writerHasFS h checksums+ FS.hDropCacheAll writerHasBlockIO h+ -- drop the KOps and blobs files from the cache if asked for+ when dropCaches $ do+ dropCache writerHasBlockIO (unForKOps writerKOpsHandle)+ dropCache writerHasBlockIO (unForBlob writerBlobHandle)+ closeHandle writerHasFS (unForKOps writerKOpsHandle)+ closeHandle writerHasFS (unForBlob writerBlobHandle)+ pure (writerHasFS, writerHasBlockIO, writerFsPaths)+++{-# SPECIALISE+ addKeyOp ::+ WriteBufferWriter IO h+ -> SerialisedKey+ -> Entry SerialisedValue (RawBlobRef IO h)+ -> IO ()+ #-}+-- | See 'Database.LSMTree.Internal.RunBuilder.addKeyOp'.+addKeyOp ::+ (MonadST m, MonadSTM m, MonadThrow m)+ => WriteBufferWriter m h+ -> SerialisedKey+ -> Entry SerialisedValue (RawBlobRef m h)+ -> m ()+addKeyOp WriteBufferWriter{..} key op = do+ -- TODO: consider optimisation described in 'Database.LSMTree.Internal.RunBuilder.addKeyOp'.+ op' <- traverse (copyBlob writerHasFS writerBlobOffset writerBlobHandle) op+ if PageAcc.entryWouldFitInPage key op+ then do+ mPage <- ST.stToIO $ addSmallKeyOp writerPageAcc key op'+ for_ mPage $ writeRawPage writerHasFS writerKOpsHandle+ else do+ (pages, overflowPages) <- ST.stToIO $ addLargeKeyOp writerPageAcc key op'+ -- TODO: consider optimisation described in 'Database.LSMTree.Internal.RunBuilder.addKeyOp'.+ for_ pages $ writeRawPage writerHasFS writerKOpsHandle+ writeRawOverflowPages writerHasFS writerKOpsHandle overflowPages++-- | See 'Database.LSMTree.Internal.RunAcc.addSmallKeyOp'.+addSmallKeyOp ::+ PageAcc s+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan+ -> ST s (Maybe RawPage)+addSmallKeyOp pageAcc key op =+ assert (PageAcc.entryWouldFitInPage key op) $ do+ pageBoundaryNeeded <-+ -- Try adding the key/op to the page accumulator to see if it fits. If+ -- it does not fit, a page boundary is needed.+ not <$> PageAcc.pageAccAddElem pageAcc key op+ if pageBoundaryNeeded+ then do+ -- We need a page boundary. If the current page is empty then we have+ -- a boundary already, otherwise we need to flush the current page.+ mPage <- flushPageIfNonEmpty pageAcc+ -- The current page is now empty, either because it was already empty+ -- or because we just flushed it. Adding the new key/op to an empty+ -- page must now succeed, because we know it fits in a page.+ added <- PageAcc.pageAccAddElem pageAcc key op+ assert added $ pure mPage+ else do+ pure Nothing++-- | See 'Database.LSMTree.Internal.RunAcc.addLargeKeyOp'.+addLargeKeyOp ::+ PageAcc s+ -> SerialisedKey+ -> Entry SerialisedValue BlobSpan -- ^ the full value, not just a prefix+ -> ST s ([RawPage], [RawOverflowPage])+addLargeKeyOp pageAcc key op =+ assert (not (PageAcc.entryWouldFitInPage key op)) $ do+ -- If the existing page accumulator is non-empty, we flush it, since the+ -- new large key/op will need more than one page to itself.+ mPagePre <- flushPageIfNonEmpty pageAcc+ -- Make the new page and overflow pages. Add the span of pages to the index.+ let (page, overflowPages) = PageAcc.singletonPage key op+ -- Combine the results with anything we flushed before+ let !pages = selectPages mPagePre page+ pure (pages, overflowPages)++-- | Internal helper. See 'Database.LSMTree.Internal.RunAcc.flushPageIfNonEmpty'.+flushPageIfNonEmpty :: PageAcc s -> ST s (Maybe RawPage)+flushPageIfNonEmpty pageAcc = do+ keysCount <- PageAcc.keysCountPageAcc pageAcc+ if keysCount > 0+ then do+ -- Serialise the page and reset the accumulator+ page <- PageAcc.serialisePageAcc pageAcc+ PageAcc.resetPageAcc pageAcc+ pure $ Just page+ else pure Nothing++-- | Internal helper. See 'Database.LSMTree.Internal.RunAcc.selectPagesAndChunks'.+selectPages ::+ Maybe RawPage+ -> RawPage+ -> [RawPage]+selectPages mPagePre page =+ maybeToList mPagePre ++ [page]
@@ -0,0 +1,59 @@+module Database.LSMTree.Extras (+ showPowersOf10+ , showPowersOf+ , showRangesOf+ , groupsOfN+ , vgroupsOfN+ ) where++import Data.List (find)+import qualified Data.List as List+import Data.List.NonEmpty (NonEmpty)+import qualified Data.List.NonEmpty as NE+import Data.Maybe (fromJust)+import qualified Data.Vector as V+import Text.Printf++showPowersOf10 :: Int -> String+showPowersOf10 = showPowersOf 10++showPowersOf :: Int -> Int -> String+showPowersOf factor n+ | factor <= 1 = error "showPowersOf: factor must be larger than 1"+ | n < 0 = "n < 0"+ | n == 0 = "n == 0"+ | otherwise = printf "%d <= n < %d" lb ub+ where+ ub = fromJust (find (n <) (iterate (* factor) factor))+ lb = ub `div` factor++showRangesOf :: Int -> Int -> String+showRangesOf range n+ | range <= 0 = error "showRangesOf: range must be larger than 0"+ | n == 0 = "n == 0"+ | m == 0 = printf "%d < n < %d" lb ub+ | otherwise = printf "%d <= n < %d" lb ub+ where+ m = n `div` range+ lb = m * range+ ub = (m + 1) * range++-- | Make groups of @n@ elements from a list @xs@+groupsOfN :: Int -> [a] -> [NonEmpty a]+groupsOfN n+ | n <= 0 = error "groupsOfN: n <= 0"+ | otherwise = List.unfoldr f+ where f xs = let (ys, zs) = List.splitAt n xs+ in (,zs) <$> NE.nonEmpty ys++-- | Make groups of @n@ elements from a vector @xs@+vgroupsOfN :: Int -> V.Vector a -> V.Vector (V.Vector a)+vgroupsOfN n+ | n <= 0 = error "groupsOfN: n <= 0"+ | otherwise = V.unfoldr f+ where+ f xs+ | V.null xs+ = Nothing+ | otherwise+ = Just $ V.splitAt n xs
@@ -0,0 +1,573 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Database.LSMTree.Extras.Generators (+ -- * WithSerialised+ WithSerialised (..)+ -- * A (logical\/true) page+ -- ** A true page+ , TruePageSummary (..)+ , flattenLogicalPageSummary+ -- ** A logical page+ , LogicalPageSummary (..)+ , shrinkLogicalPageSummary+ , toAppend+ -- * Sequences of (logical\/true) pages+ , Pages (..)+ -- ** Sequences of true pages+ , TruePageSummaries+ , flattenLogicalPageSummaries+ -- ** Sequences of logical pages+ , LogicalPageSummaries+ , toAppends+ , labelPages+ , shrinkPages+ , genPages+ , mkPages+ , pagesInvariant+ -- * Chunking size+ , ChunkSize (..)+ , chunkSizeInvariant+ -- * Serialised keys\/values\/blobs+ , genRawBytes+ , genRawBytesN+ , genRawBytesSized+ , packRawBytesPinnedOrUnpinned+ , LargeRawBytes (..)+ , BiasedKey (..)+ -- * helpers+ , shrinkVec+ ) where++import Control.DeepSeq (NFData)+import Control.Exception (assert)+import Data.Coerce (coerce)+import Data.Containers.ListUtils (nubOrd)+import Data.Function ((&))+import Data.List (nub, sort)+import qualified Data.Primitive.ByteArray as BA+import qualified Data.Vector.Primitive as VP+import Data.Word+import qualified Database.LSMTree as Full+import Database.LSMTree.Extras+import Database.LSMTree.Extras.Index (Append (..))+import Database.LSMTree.Extras.Orphans ()+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry (Entry (..), NumEntries (..))+import qualified Database.LSMTree.Internal.Merge as Merge+import Database.LSMTree.Internal.Page (PageNo (..))+import Database.LSMTree.Internal.Range (Range (..))+import Database.LSMTree.Internal.RawBytes (RawBytes (RawBytes))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise+import qualified Database.LSMTree.Internal.Serialise.Class as S.Class+import Database.LSMTree.Internal.Unsliced (Unsliced, fromUnslicedKey,+ makeUnslicedKey)+import Database.LSMTree.Internal.Vector (mkPrimVector)+import GHC.Generics (Generic)+import qualified Test.QuickCheck as QC+import Test.QuickCheck (Arbitrary (..), Arbitrary1 (..),+ Arbitrary2 (..), Gen, Property, elements, frequency)+import Test.QuickCheck.Gen (genDouble)+import Test.QuickCheck.Instances ()++{-------------------------------------------------------------------------------+ Common LSMTree types+-------------------------------------------------------------------------------}++instance (Arbitrary v, Arbitrary b) => Arbitrary (Full.Update v b) where+ arbitrary = QC.arbitrary2+ shrink = QC.shrink2++instance Arbitrary2 Full.Update where+ liftArbitrary2 genVal genBlob = frequency+ [ (10, Full.Insert <$> genVal <*> liftArbitrary genBlob)+ , (5, Full.Upsert <$> genVal)+ , (1, pure Full.Delete)+ ]++ liftShrink2 shrinkVal shrinkBlob = \case+ Full.Insert v blob ->+ Full.Delete+ : map (uncurry Full.Insert)+ (liftShrink2 shrinkVal (liftShrink shrinkBlob) (v, blob))+ Full.Upsert v -> Full.Insert v Nothing : map Full.Upsert (shrinkVal v)+ Full.Delete -> []++instance (Arbitrary k, Ord k) => Arbitrary (Range k) where+ arbitrary = do+ key1 <- arbitrary+ key2 <- arbitrary `QC.suchThat` (/= key1)+ (lb, ub) <- frequency+ [ (1, pure (key1, key1)) -- lb == ub+ , (1, pure (max key1 key2, min key1 key2)) -- lb > ub+ , (8, pure (min key1 key2, max key1 key2)) -- lb < ub+ ]+ elements+ [ FromToExcluding lb ub+ , FromToIncluding lb ub+ ]++ shrink (FromToExcluding f t) =+ uncurry FromToExcluding <$> shrink (f, t)+ shrink (FromToIncluding f t) =+ uncurry FromToIncluding <$> shrink (f, t)++{-------------------------------------------------------------------------------+ Entry+-------------------------------------------------------------------------------}++instance (Arbitrary v, Arbitrary b) => Arbitrary (Entry v b) where+ arbitrary = QC.arbitrary2+ shrink = QC.shrink2++instance Arbitrary2 Entry where+ liftArbitrary2 genVal genBlob = frequency+ [ (5, Insert <$> genVal)+ , (1, InsertWithBlob <$> genVal <*> genBlob)+ , (1, Upsert <$> genVal)+ , (1, pure Delete)+ ]++ liftShrink2 shrinkVal shrinkBlob = \case+ Insert v -> Delete : (Insert <$> shrinkVal v)+ InsertWithBlob v b -> [Delete, Insert v]+ ++ fmap (uncurry InsertWithBlob)+ (liftShrink2 shrinkVal shrinkBlob (v, b))+ Upsert v -> Delete : Insert v : (Upsert <$> shrinkVal v)+ Delete -> []++{-------------------------------------------------------------------------------+ WithSerialised+-------------------------------------------------------------------------------}++-- | Cache serialised keys+--+-- Also useful for failing tests that have keys as inputs, because the printed+-- 'WithSerialised' values will show both keys and their serialised form.+data WithSerialised k = WithSerialised k SerialisedKey+ deriving stock Show++instance Eq k => Eq (WithSerialised k) where+ WithSerialised k1 _ == WithSerialised k2 _ = k1 == k2++instance Ord k => Ord (WithSerialised k) where+ WithSerialised k1 _ `compare` WithSerialised k2 _ = k1 `compare` k2++instance (Arbitrary k, SerialiseKey k) => Arbitrary (WithSerialised k) where+ arbitrary = do+ x <- arbitrary+ pure $ WithSerialised x (serialiseKey x)+ shrink (WithSerialised k _) = [WithSerialised k' (serialiseKey k') | k' <- shrink k]++instance SerialiseKey k => SerialiseKey (WithSerialised k) where+ serialiseKey (WithSerialised _ (SerialisedKey bytes)) = bytes+ deserialiseKey bytes = WithSerialised (S.Class.deserialiseKey bytes) (SerialisedKey bytes)++{-------------------------------------------------------------------------------+ Other number newtypes+-------------------------------------------------------------------------------}++instance Arbitrary PageNo where+ arbitrary = coerce (arbitrary @(QC.NonNegative Int))+ shrink = coerce (shrink @(QC.NonNegative Int))++instance Arbitrary NumEntries where+ arbitrary = coerce (arbitrary @(QC.NonNegative Int))+ shrink = coerce (shrink @(QC.NonNegative Int))++{-------------------------------------------------------------------------------+ True page+-------------------------------------------------------------------------------}++-- | A summary of min/max information for keys on a /true/ page.+--+-- A true page corresponds directly to a disk page. See 'LogicalPageSummary' for+-- contrast.+data TruePageSummary k = TruePageSummary { tpsMinKey :: k, tpsMaxKey :: k }++flattenLogicalPageSummary :: LogicalPageSummary k -> [TruePageSummary k]+flattenLogicalPageSummary = \case+ OnePageOneKey k -> [TruePageSummary k k]+ OnePageManyKeys k1 k2 -> [TruePageSummary k1 k2]+ MultiPageOneKey k n -> replicate (fromIntegral n+1) (TruePageSummary k k)++{-------------------------------------------------------------------------------+ Logical page+-------------------------------------------------------------------------------}++-- | A summary of min/max information for keys on a /logical/ page.+--+-- A key\/operation pair can fit onto a single page, or the operation is so+-- large that its bytes flow over into subsequent pages. A logical page makes+-- this overflow explicit. Making these cases explicit in the representation+-- makes generating and shrinking test cases easier.+data LogicalPageSummary k =+ OnePageOneKey k+ | OnePageManyKeys k k+ | MultiPageOneKey k Word32 -- ^ number of overflow pages+ deriving stock (Show, Generic, Functor)+ deriving anyclass NFData++toAppend :: LogicalPageSummary SerialisedKey -> Append+toAppend (OnePageOneKey k) = AppendSinglePage k k+toAppend (OnePageManyKeys k1 k2) = AppendSinglePage k1 k2+toAppend (MultiPageOneKey k n) = AppendMultiPage k n++shrinkLogicalPageSummary :: Arbitrary k => LogicalPageSummary k -> [LogicalPageSummary k]+shrinkLogicalPageSummary = \case+ OnePageOneKey k -> OnePageOneKey <$> shrink k+ OnePageManyKeys k1 k2 -> [+ OnePageManyKeys k1' k2'+ | (k1', k2') <- shrink (k1, k2)+ ]+ MultiPageOneKey k n -> [+ MultiPageOneKey k' n'+ | (k', n') <- shrink (k, n)+ ]++{-------------------------------------------------------------------------------+ Sequences of (logical\/true) pages+-------------------------------------------------------------------------------}++-- | Sequences of (logical\/true) pages+--+-- INVARIANT: The sequence consists of multiple pages in sorted order (keys are+-- sorted within a page and across pages).+newtype Pages fp k = Pages { getPages :: [fp k] }+ deriving stock (Show, Generic, Functor)+ deriving anyclass NFData++class TrueNumberOfPages fp where+ trueNumberOfPages :: Pages fp k -> Int++instance TrueNumberOfPages LogicalPageSummary where+ trueNumberOfPages :: LogicalPageSummaries k -> Int+ trueNumberOfPages = length . getPages . flattenLogicalPageSummaries++instance TrueNumberOfPages TruePageSummary where+ trueNumberOfPages :: TruePageSummaries k -> Int+ trueNumberOfPages = length . getPages++{-------------------------------------------------------------------------------+ Sequences of true pages+-------------------------------------------------------------------------------}++type TruePageSummaries k = Pages TruePageSummary k++flattenLogicalPageSummaries :: LogicalPageSummaries k -> TruePageSummaries k+flattenLogicalPageSummaries (Pages ps) = Pages (concatMap flattenLogicalPageSummary ps)++{-------------------------------------------------------------------------------+ Sequences of logical pages+-------------------------------------------------------------------------------}++type LogicalPageSummaries k = Pages LogicalPageSummary k++toAppends :: SerialiseKey k => LogicalPageSummaries k -> [Append]+toAppends (Pages ps) = fmap (toAppend . fmap serialiseKey) ps++--+-- Labelling+--++labelPages :: LogicalPageSummaries k -> (Property -> Property)+labelPages ps =+ QC.tabulate "# True pages" [showPowersOf10 nTruePages]+ . QC.tabulate "# Logical pages" [showPowersOf10 nLogicalPages]+ . QC.tabulate "# OnePageOneKey logical pages" [showPowersOf10 n1]+ . QC.tabulate "# OnePageManyKeys logical pages" [showPowersOf10 n2]+ . QC.tabulate "# MultiPageOneKey logical pages" [showPowersOf10 n3]+ where+ nLogicalPages = length $ getPages ps+ nTruePages = trueNumberOfPages ps++ (n1,n2,n3) = counts (getPages ps)++ counts :: [LogicalPageSummary k] -> (Int, Int, Int)+ counts [] = (0, 0, 0)+ counts (lp:lps) = let (x, y, z) = counts lps+ in case lp of+ OnePageOneKey{} -> (x+1, y, z)+ OnePageManyKeys{} -> (x, y+1, z)+ MultiPageOneKey{} -> (x, y, z+1)++--+-- Generation and shrinking+--++instance (Arbitrary k, Ord k)+ => Arbitrary (LogicalPageSummaries k) where+ arbitrary = genPages 0.03 (QC.choose (0, 16)) 0.01+ shrink = shrinkPages++shrinkPages ::+ (Arbitrary k, Ord k)+ => LogicalPageSummaries k+ -> [LogicalPageSummaries k]+shrinkPages (Pages ps) = [+ Pages ps'+ | ps' <- QC.shrinkList shrinkLogicalPageSummary ps, pagesInvariant (Pages ps')+ ]++genPages ::+ (Arbitrary k, Ord k)+ => Double -- ^ Probability of a value being larger-than-page+ -> Gen Word32 -- ^ Number of overflow pages for a larger-than-page value+ -> Double -- ^ Probability of generating a page with only one key and value,+ -- which does /not/ span multiple pages.+ -> Gen (LogicalPageSummaries k)+genPages p genN p' = do+ ks <- arbitrary+ mkPages p genN p' ks++mkPages ::+ forall k. Ord k+ => Double -- ^ Probability of a value being larger-than-page+ -> Gen Word32 -- ^ Number of overflow pages for a larger-than-page value+ -> Double -- ^ Probability of generating a page with only one key and value,+ -- which does /not/ span multiple pages.+ -> [k]+ -> Gen (LogicalPageSummaries k)+mkPages p genN p' =+ fmap Pages . go . nubOrd . sort+ where+ go :: [k] -> Gen [LogicalPageSummary k]+ go [] = pure []+ go [k] = do+ b <- largerThanPage+ if b then pure . MultiPageOneKey k <$> genN+ else pure [OnePageOneKey k]+ -- the min and max key are allowed to be the same+ go (k1:k2:ks) = do+ b <- largerThanPage+ b' <- onePageOneKey+ if b then (:) <$> (MultiPageOneKey k1 <$> genN) <*> go (k2 : ks)+ else if b' then (OnePageOneKey k1 :) <$> go (k2 : ks)+ else (OnePageManyKeys k1 k2 :) <$> go ks++ largerThanPage :: Gen Bool+ largerThanPage = genDouble >>= \x -> pure (x < p)++ onePageOneKey :: Gen Bool+ onePageOneKey = genDouble >>= \x -> pure (x < p')++pagesInvariant :: Ord k => LogicalPageSummaries k -> Bool+pagesInvariant (Pages ps0) =+ sort ks == ks+ && nubOrd ks == ks+ where+ ks = flatten ps0++ flatten :: Eq k => [LogicalPageSummary k] -> [k]+ flatten [] = []+ -- the min and max key are allowed to be the same+ flatten (p:ps) = case p of+ OnePageOneKey k -> k : flatten ps+ OnePageManyKeys k1 k2 -> k1 : k2 : flatten ps+ MultiPageOneKey k _ -> k : flatten ps++{-------------------------------------------------------------------------------+ Chunking size+-------------------------------------------------------------------------------}++newtype ChunkSize = ChunkSize Int+ deriving stock Show+ deriving newtype Num++instance Arbitrary ChunkSize where+ arbitrary = ChunkSize <$> QC.chooseInt (chunkSizeLB, chunkSizeUB)+ shrink (ChunkSize csize) = [+ ChunkSize csize'+ | csize' <- shrink csize+ , chunkSizeInvariant (ChunkSize csize')+ ]++chunkSizeLB, chunkSizeUB :: Int+chunkSizeLB = 1+chunkSizeUB = 20++chunkSizeInvariant :: ChunkSize -> Bool+chunkSizeInvariant (ChunkSize csize) = chunkSizeLB <= csize && csize <= chunkSizeUB++{-------------------------------------------------------------------------------+ Serialised keys/values/blobs+-------------------------------------------------------------------------------}++instance Arbitrary RawBytes where+ arbitrary = do+ QC.NonNegative (QC.Small prefixLength) <- arbitrary+ QC.NonNegative (QC.Small payloadLength) <- arbitrary+ QC.NonNegative (QC.Small suffixLength) <- arbitrary+ base <- genRawBytesN (prefixLength + payloadLength + suffixLength)+ pure (base & RB.drop prefixLength & RB.take payloadLength)+ shrink rb = shrinkSlice rb ++ shrinkRawBytes rb++genRawBytesN :: Int -> Gen RawBytes+genRawBytesN n =+ packRawBytesPinnedOrUnpinned <$> arbitrary <*> QC.vectorOf n arbitrary++genRawBytes :: Gen RawBytes+genRawBytes =+ packRawBytesPinnedOrUnpinned <$> arbitrary <*> QC.listOf arbitrary++genRawBytesSized :: Int -> Gen RawBytes+genRawBytesSized n = QC.resize n genRawBytes++packRawBytesPinnedOrUnpinned :: Bool -> [Word8] -> RawBytes+packRawBytesPinnedOrUnpinned False = RB.pack+packRawBytesPinnedOrUnpinned True = \ws ->+ let len = length ws in+ RB.RawBytes $ mkPrimVector 0 len $ BA.runByteArray $ do+ mba <- BA.newPinnedByteArray len+ sequence_ [ BA.writeByteArray mba i w | (i, w) <- zip [0..] ws ]+ pure mba++shrinkRawBytes :: RawBytes -> [RawBytes]+shrinkRawBytes (RawBytes pvec) =+ [ RawBytes pvec'+ | pvec' <- shrinkVec shrinkByte pvec+ ]+ where+ -- no need to try harder shrinking individual bytes+ shrinkByte b = nub (takeWhile (< b) [0, b `div` 2])++-- | Based on QuickCheck's 'shrinkList' (behaves identically, see tests).+shrinkVec :: VP.Prim a => (a -> [a]) -> VP.Vector a -> [VP.Vector a]+shrinkVec shr vec =+ concat [ removeBlockOf k | k <- takeWhile (> 0) (iterate (`div` 2) len) ]+ ++ shrinkOne+ where+ len = VP.length vec++ shrinkOne =+ [ vec VP.// [(i, x')]+ | i <- [0 .. len-1]+ , let x = vec VP.! i+ , x' <- shr x+ ]++ removeBlockOf k =+ [ VP.take i vec VP.++ VP.drop (i + k) vec+ | i <- [0, k .. len - k]+ ]++genSlice :: RawBytes -> Gen RawBytes+genSlice (RawBytes pvec) = do+ n <- QC.chooseInt (0, VP.length pvec)+ m <- QC.chooseInt (0, VP.length pvec - n)+ pure $ RawBytes (VP.slice m n pvec)++shrinkSlice :: RawBytes -> [RawBytes]+shrinkSlice (RawBytes pvec) =+ [ RawBytes (VP.take len' pvec)+ | len' <- QC.shrink len+ ] +++ [ RawBytes (VP.drop (len - len') pvec)+ | len' <- QC.shrink len+ ]+ where+ len = VP.length pvec++deriving newtype instance Arbitrary SerialisedKey++instance Arbitrary SerialisedValue where+ -- good mix of sizes, including larger than two pages, also some slices+ arbitrary = SerialisedValue <$> frequency+ [ (16, arbitrary)+ , ( 4, genRawBytesN =<< QC.chooseInt ( 100, 1000))+ , ( 2, genRawBytesN =<< QC.chooseInt (1000, 4000))+ , ( 1, genRawBytesN =<< QC.chooseInt (4000, 10000))+ , ( 1, genSlice =<< genRawBytesN =<< QC.chooseInt (0, 10000))+ ]+ shrink (SerialisedValue rb)+ | RB.size rb > 64 = coerce (shrink (LargeRawBytes rb))+ | otherwise = coerce (shrink rb)++deriving newtype instance Arbitrary SerialisedBlob++newtype LargeRawBytes = LargeRawBytes RawBytes+ deriving stock Show+ deriving newtype NFData++instance Arbitrary LargeRawBytes where+ arbitrary = genRawBytesSized (4096*3) >>= fmap LargeRawBytes . genSlice+ shrink (LargeRawBytes rb) =+ map LargeRawBytes (shrinkSlice rb)+ -- After shrinking length, don't shrink content using normal list shrink+ -- as that's too slow. We try zeroing out long suffixes of the bytes+ -- (since for large raw bytes in page format, the interesting information+ -- is at the start and the suffix is just the value.+ ++ [ LargeRawBytes (RawBytes pvec')+ | let (RawBytes pvec) = rb+ , n <- QC.shrink (VP.length pvec)+ , assert (n >= 0) True -- negative values would make pvec' longer+ , let pvec' = VP.take n pvec VP.++ VP.replicate (VP.length pvec - n) 0+ , assert (VP.length pvec' == VP.length pvec) $+ pvec' /= pvec+ ]++deriving newtype instance SerialiseValue LargeRawBytes++-- we try to make collisions and close keys more likely (very crudely)+arbitraryBiasedKey :: (RawBytes -> k) -> Gen RawBytes -> Gen k+arbitraryBiasedKey fromRB genUnbiased = fromRB <$> frequency+ [ (6, genUnbiased)+ , (1, do+ lastByte <- QC.sized $ skewedWithMax . fromIntegral+ pure (RB.pack ([1,3,3,7,0,1,7] <> [lastByte]))+ )+ ]+ where+ -- generates a value in range from 0 to ub, but skewed towards low end+ skewedWithMax ub0 = do+ ub1 <- QC.chooseBoundedIntegral (0, ub0)+ ub2 <- QC.chooseBoundedIntegral (0, ub1)+ QC.chooseBoundedIntegral (0, ub2)++newtype BiasedKey = BiasedKey { getBiasedKey :: RawBytes }+ deriving stock (Eq, Ord, Show)+ deriving newtype NFData++instance Arbitrary BiasedKey where+ arbitrary = arbitraryBiasedKey BiasedKey arbitrary++ shrink (BiasedKey rb) = [BiasedKey rb' | rb' <- shrink rb]++deriving newtype instance SerialiseKey BiasedKey++{-------------------------------------------------------------------------------+ Unsliced+-------------------------------------------------------------------------------}++instance Arbitrary (Unsliced SerialisedKey) where+ arbitrary = makeUnslicedKey <$> arbitrary+ shrink = fmap makeUnslicedKey . shrink . fromUnslicedKey++{-------------------------------------------------------------------------------+ BlobRef+-------------------------------------------------------------------------------}++instance Arbitrary BlobSpan where+ arbitrary = BlobSpan <$> arbitrary <*> arbitrary+ shrink (BlobSpan x y) = [ BlobSpan x' y' | (x', y') <- shrink (x, y) ]++{-------------------------------------------------------------------------------+ Merge+-------------------------------------------------------------------------------}++instance Arbitrary Merge.MergeType where+ arbitrary = QC.elements+ [Merge.MergeTypeMidLevel, Merge.MergeTypeLastLevel, Merge.MergeTypeUnion]+ shrink Merge.MergeTypeMidLevel = []+ shrink Merge.MergeTypeLastLevel = [Merge.MergeTypeMidLevel]+ shrink Merge.MergeTypeUnion = [Merge.MergeTypeLastLevel]++instance Arbitrary Merge.LevelMergeType where+ arbitrary = QC.elements [Merge.MergeMidLevel, Merge.MergeLastLevel]+ shrink Merge.MergeMidLevel = []+ shrink Merge.MergeLastLevel = [Merge.MergeMidLevel]++instance Arbitrary Merge.TreeMergeType where+ arbitrary = QC.elements [Merge.MergeLevel, Merge.MergeUnion]+ shrink Merge.MergeLevel = []+ shrink Merge.MergeUnion = [Merge.MergeLevel]
@@ -0,0 +1,101 @@+{-|+ Provides additional support for working with fence pointer indexes and their+ accumulators.+-}+module Database.LSMTree.Extras.Index+(+ Append (AppendSinglePage, AppendMultiPage),+ appendToCompact,+ appendToOrdinary,+ append+)+where++import Control.DeepSeq (NFData (rnf))+import Control.Monad.ST.Strict (ST)+import Data.Foldable (toList)+import Data.Word (Word32)+import Database.LSMTree.Internal.Chunk (Chunk)+import Database.LSMTree.Internal.Index (IndexAcc)+import qualified Database.LSMTree.Internal.Index as Index (appendMulti,+ appendSingle)+import Database.LSMTree.Internal.Index.CompactAcc (IndexCompactAcc)+import qualified Database.LSMTree.Internal.Index.CompactAcc as IndexCompact+ (appendMulti, appendSingle)+import Database.LSMTree.Internal.Index.OrdinaryAcc (IndexOrdinaryAcc)+import qualified Database.LSMTree.Internal.Index.OrdinaryAcc as IndexOrdinary+ (appendMulti, appendSingle)+import Database.LSMTree.Internal.Serialise (SerialisedKey)++-- | Instruction for appending pages, to be used in conjunction with indexes.+data Append+ = {-|+ Append a single page that fully comprises one or more key–value pairs.+ -}+ AppendSinglePage+ SerialisedKey -- ^ Minimum key+ SerialisedKey -- ^ Maximum key+ | {-|+ Append multiple pages that together comprise a single key–value pair.+ -}+ AppendMultiPage+ SerialisedKey -- ^ Sole key+ Word32 -- ^ Number of overflow pages++instance NFData Append where++ rnf (AppendSinglePage minKey maxKey)+ = rnf minKey `seq` rnf maxKey+ rnf (AppendMultiPage key overflowPageCount)+ = rnf key `seq` rnf overflowPageCount++{-|+ Adds information about appended pages to an index and outputs newly+ available chunks, using primitives specific to the type of the index.++ See the documentation of the 'IndexAcc' type for constraints to adhere to.+-}+appendWith :: ((SerialisedKey, SerialisedKey) -> j s -> ST s (Maybe Chunk))+ -> ((SerialisedKey, Word32) -> j s -> ST s [Chunk])+ -> Append+ -> j s+ -> ST s [Chunk]+appendWith appendSingle appendMulti instruction indexAcc = case instruction of+ AppendSinglePage minKey maxKey+ -> toList <$> appendSingle (minKey, maxKey) indexAcc+ AppendMultiPage key overflowPageCount+ -> appendMulti (key, overflowPageCount) indexAcc+{-# INLINABLE appendWith #-}++{-|+ Adds information about appended pages to a compact index and outputs newly+ available chunks.++ See the documentation of the 'IndexAcc' type for constraints to adhere to.+-}+appendToCompact :: Append -> IndexCompactAcc s -> ST s [Chunk]+appendToCompact = appendWith IndexCompact.appendSingle+ IndexCompact.appendMulti+{-# INLINE appendToCompact #-}++{-|+ Adds information about appended pages to an ordinary index and outputs newly+ available chunks.++ See the documentation of the 'IndexAcc' type for constraints to adhere to.+-}+appendToOrdinary :: Append -> IndexOrdinaryAcc s -> ST s [Chunk]+appendToOrdinary = appendWith IndexOrdinary.appendSingle+ IndexOrdinary.appendMulti+{-# INLINE appendToOrdinary #-}++{-|+ Adds information about appended pages to an index and outputs newly+ available chunks.++ See the documentation of the 'IndexAcc' type for constraints to adhere to.+-}+append :: Append -> IndexAcc s -> ST s [Chunk]+append = appendWith Index.appendSingle+ Index.appendMulti+{-# INLINE append #-}
@@ -0,0 +1,215 @@+-- | Utilities for generating 'MergingRun's. Tests and benchmarks should+-- preferably use these utilities instead of (re-)defining their own.+module Database.LSMTree.Extras.MergingRunData (+ -- * Create merging runs+ withMergingRun+ , unsafeCreateMergingRun+ -- * MergingRunData+ , MergingRunData (..)+ , mergingRunDataMergeType+ , mergingRunDataInvariant+ , mapMergingRunData+ , SerialisedMergingRunData+ , serialiseMergingRunData+ -- * QuickCheck+ , labelMergingRunData+ , genMergingRunData+ , shrinkMergingRunData+ ) where++import Control.Exception (bracket)+import Control.RefCount+import qualified Data.Vector as V+import Database.LSMTree.Extras (showPowersOf)+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Extras.RunData+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.MergingRun (MergingRun)+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.Paths+import qualified Database.LSMTree.Internal.Run as Run+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import System.FS.BlockIO.API (HasBlockIO)+import Test.QuickCheck as QC++{-------------------------------------------------------------------------------+ Create merging runs+-------------------------------------------------------------------------------}++-- | Create a temporary 'MergingRun' using 'unsafeCreateMergingRun'.+withMergingRun ::+ MR.IsMergeType t+ => HasFS IO h+ -> HasBlockIO IO h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> RunBuilder.RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> SerialisedMergingRunData t+ -> (Ref (MergingRun t IO h) -> IO a)+ -> IO a+withMergingRun hfs hbio resolve salt runParams path counter mrd = do+ bracket+ (unsafeCreateMergingRun hfs hbio resolve salt runParams path counter mrd)+ releaseRef++-- | Flush serialised merging run data to disk.+--+-- This might leak resources if not run with asynchronous exceptions masked.+-- Consider using 'withMergingRun' instead.+--+-- Use of this function should be paired with a 'releaseRef'.+unsafeCreateMergingRun ::+ MR.IsMergeType t+ => HasFS IO h+ -> HasBlockIO IO h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> RunBuilder.RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> SerialisedMergingRunData t+ -> IO (Ref (MergingRun t IO h))+unsafeCreateMergingRun hfs hbio resolve salt runParams path counter = \case+ CompletedMergeData _ rd -> do+ withRun hfs hbio salt runParams path counter rd $ \run -> do+ -- slightly hacky, generally it's larger+ let totalDebt = MR.numEntriesToMergeDebt (Run.size run)+ MR.newCompleted totalDebt run++ OngoingMergeData mergeType rds -> do+ withRuns hfs hbio salt runParams path counter (toRunData <$> rds)+ $ \runs -> do+ n <- incrUniqCounter counter+ let fsPaths = RunFsPaths path (RunNumber (uniqueToInt n))+ MR.new hfs hbio resolve salt runParams mergeType+ fsPaths (V.fromList runs)++{-------------------------------------------------------------------------------+ MergingRunData+-------------------------------------------------------------------------------}++-- | A data structure suitable for creating arbitrary 'MergingRun's.+--+-- Note: 'b ~ Void' should rule out blobs.+--+-- Currently, ongoing merges are always \"fresh\", i.e. there is no merge work+-- already performed.+--+-- TODO: Generate merge credits and supply them in 'unsafeCreateMergingRun',+-- similarly to how @ScheduledMergesTest@ does it.+data MergingRunData t k v b =+ CompletedMergeData t (RunData k v b)+ | OngoingMergeData t [NonEmptyRunData k v b] -- ^ at least 2 inputs+ deriving stock (Show, Eq)++mergingRunDataMergeType :: MergingRunData t k v b -> t+mergingRunDataMergeType = \case+ CompletedMergeData mt _ -> mt+ OngoingMergeData mt _ -> mt++-- | See @mergeInvariant@ in the prototype.+mergingRunDataInvariant :: MergingRunData t k v b -> Either String ()+mergingRunDataInvariant = \case+ CompletedMergeData _ _ -> Right ()+ OngoingMergeData _ rds -> do+ assertI "ongoing merges are non-trivial (at least two inputs)" $+ length rds >= 2+ where+ assertI msg False = Left msg+ assertI _ True = Right ()++mapMergingRunData ::+ Ord k'+ => (k -> k') -> (v -> v') -> (b -> b')+ -> MergingRunData t k v b -> MergingRunData t k' v' b'+mapMergingRunData f g h = \case+ CompletedMergeData t r ->+ CompletedMergeData t $ mapRunData f g h r+ OngoingMergeData t rs ->+ OngoingMergeData t $ map (mapNonEmptyRunData f g h) rs++type SerialisedMergingRunData t =+ MergingRunData t SerialisedKey SerialisedValue SerialisedBlob++serialiseMergingRunData ::+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => MergingRunData t k v b -> SerialisedMergingRunData t+serialiseMergingRunData =+ mapMergingRunData serialiseKey serialiseValue serialiseBlob++{-------------------------------------------------------------------------------+ QuickCheck+-------------------------------------------------------------------------------}++labelMergingRunData ::+ Show t => SerialisedMergingRunData t -> Property -> Property+labelMergingRunData (CompletedMergeData mt rd) =+ tabulate "merging run state" ["CompletedMerge"]+ . tabulate "merge type" [show mt]+ . labelRunData rd+labelMergingRunData (OngoingMergeData mt rds) =+ tabulate "merging run state" ["OngoingMerge"]+ . tabulate "merge type" [show mt]+ . tabulate "merging run inputs" [showPowersOf 2 (length rds)]+ . foldr ((.) . labelNonEmptyRunData) id rds++instance ( Arbitrary t, Ord k, Arbitrary k, Arbitrary v, Arbitrary b+ ) => Arbitrary (MergingRunData t k v b) where+ arbitrary = genMergingRunData arbitrary arbitrary arbitrary arbitrary+ shrink = shrinkMergingRunData shrink shrink shrink++genMergingRunData ::+ Ord k+ => Gen t+ -> Gen k+ -> Gen v+ -> Gen b+ -> Gen (MergingRunData t k v b)+genMergingRunData genMergeType genKey genVal genBlob =+ QC.oneof+ [ do+ mt <- genMergeType+ rd <- genRunData genKey genVal genBlob+ pure (CompletedMergeData mt rd)+ , do+ s <- QC.getSize+ mt <- genMergeType+ n <- QC.chooseInt (2, max 2 (s * 8 `div` 100)) -- 2 to 8+ rs <- QC.vectorOf n $+ -- Scaled, so overall number of entries is similar to a completed+ -- merge. However, the entries themselves should not be smaller.+ QC.scale (`div` n) $+ genNonEmptyRunData+ (resize s genKey)+ (resize s genVal)+ (resize s genBlob)+ pure (OngoingMergeData mt rs)+ ]++shrinkMergingRunData ::+ Ord k+ => (k -> [k])+ -> (v -> [v])+ -> (b -> [b])+ -> MergingRunData t k v b+ -> [MergingRunData t k v b]+shrinkMergingRunData shrinkKey shrinkVal shrinkBlob = \case+ CompletedMergeData mt rd ->+ [ CompletedMergeData mt rd'+ | rd' <- shrinkRunData shrinkKey shrinkVal shrinkBlob rd+ ]+ OngoingMergeData mt rds ->+ [ OngoingMergeData mt rds'+ | rds' <-+ liftShrink+ (shrinkNonEmptyRunData shrinkKey shrinkVal shrinkBlob)+ rds+ , length rds' >= 2+ ]
@@ -0,0 +1,431 @@+-- | Utilities for generating 'MergingTree's. Tests and benchmarks should+-- preferably use these utilities instead of (re-)defining their own.+module Database.LSMTree.Extras.MergingTreeData (+ -- * Create merging trees+ withMergingTree+ , unsafeCreateMergingTree+ -- * MergingTreeData+ , MergingTreeData (..)+ , PreExistingRunData (..)+ , mergingTreeDataInvariant+ , mapMergingTreeData+ , SerialisedMergingTreeData+ , serialiseMergingTreeData+ -- * QuickCheck+ , labelMergingTreeData+ , genMergingTreeData+ , shrinkMergingTreeData+ ) where++import Control.Exception (assert, bracket)+import Control.RefCount+import Data.Foldable (for_, toList)+import Database.LSMTree.Extras (showPowersOf)+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Extras.MergingRunData+import Database.LSMTree.Extras.RunData+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.MergingTree (MergingTree)+import qualified Database.LSMTree.Internal.MergingTree as MT+import Database.LSMTree.Internal.RunBuilder (RunParams)+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import System.FS.BlockIO.API (HasBlockIO)+import Test.QuickCheck as QC++{-------------------------------------------------------------------------------+ Create merging tree+-------------------------------------------------------------------------------}++-- | Create a temporary 'MergingTree' using 'unsafeCreateMergingTree'.+withMergingTree ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> SerialisedMergingTreeData+ -> (Ref (MergingTree IO h) -> IO a)+ -> IO a+withMergingTree hfs hbio resolve salt runParams path counter mrd = do+ bracket+ (unsafeCreateMergingTree hfs hbio resolve salt runParams path counter mrd)+ releaseRef++-- | Flush serialised merging tree data to disk.+--+-- This might leak resources if not run with asynchronous exceptions masked.+-- Consider using 'withMergingTree' instead.+--+-- Use of this function should be paired with a 'releaseRef'.+unsafeCreateMergingTree ::+ HasFS IO h+ -> HasBlockIO IO h+ -> ResolveSerialisedValue+ -> Bloom.Salt+ -> RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> SerialisedMergingTreeData+ -> IO (Ref (MergingTree IO h))+unsafeCreateMergingTree hfs hbio resolve salt runParams path counter = go+ where+ go = \case+ CompletedTreeMergeData rd ->+ withRun hfs hbio salt runParams path counter rd $ \run ->+ MT.newCompletedMerge run+ OngoingTreeMergeData mrd ->+ withMergingRun hfs hbio resolve salt runParams path counter mrd $ \mr ->+ MT.newOngoingMerge mr+ PendingLevelMergeData prds mtd ->+ withPreExistingRuns prds $ \prs ->+ withMaybeTree mtd $ \mt ->+ MT.newPendingLevelMerge prs mt+ PendingUnionMergeData mtds ->+ withTrees mtds $ \mts ->+ MT.newPendingUnionMerge mts++ withTrees [] act = act []+ withTrees (mtd:rest) act =+ bracket (go mtd) releaseRef $ \t ->+ withTrees rest $ \ts ->+ act (t:ts)++ withMaybeTree Nothing act = act Nothing+ withMaybeTree (Just mtd) act =+ bracket (go mtd) releaseRef $ \t ->+ act (Just t)++ withPreExistingRuns [] act = act []+ withPreExistingRuns (PreExistingRunData rd : rest) act =+ withRun hfs hbio salt runParams path counter rd $ \r ->+ withPreExistingRuns rest $ \prs ->+ act (MT.PreExistingRun r : prs)+ withPreExistingRuns (PreExistingMergingRunData mrd : rest) act =+ withMergingRun hfs hbio resolve salt runParams path counter mrd $ \mr ->+ withPreExistingRuns rest $ \prs ->+ act (MT.PreExistingMergingRun mr : prs)++{-------------------------------------------------------------------------------+ MergingTreeData+-------------------------------------------------------------------------------}++-- TODO: This module has quite a lot duplication with the prototype's+-- ScheduledMergesTest module. Maybe we can share some code?++-- | A data structure suitable for creating arbitrary 'MergingTree's.+--+-- Note: 'b ~ Void' should rule out blobs.+data MergingTreeData k v b =+ CompletedTreeMergeData (RunData k v b)+ | OngoingTreeMergeData (MergingRunData MR.TreeMergeType k v b)+ | PendingLevelMergeData+ [PreExistingRunData k v b]+ (Maybe (MergingTreeData k v b)) -- ^ not both empty!+ | PendingUnionMergeData [MergingTreeData k v b] -- ^ at least 2 children+ deriving stock (Show, Eq)++data PreExistingRunData k v b =+ PreExistingRunData (RunData k v b)+ | PreExistingMergingRunData (MergingRunData MR.LevelMergeType k v b)+ deriving stock (Show, Eq)++mergingTreeIsStructurallyEmpty :: MergingTreeData k v b -> Bool+mergingTreeIsStructurallyEmpty = \case+ CompletedTreeMergeData _ -> False -- could be, but we match MT+ OngoingTreeMergeData _ -> False+ PendingLevelMergeData ps t -> null ps && null t+ PendingUnionMergeData ts -> null ts++-- | See @treeInvariant@ in prototype.+mergingTreeDataInvariant :: MergingTreeData k v b -> Either String ()+mergingTreeDataInvariant mtd+ | mergingTreeIsStructurallyEmpty mtd = Right ()+ | otherwise = case mtd of+ CompletedTreeMergeData _rd ->+ Right ()+ OngoingTreeMergeData mr ->+ mergingRunDataInvariant mr+ PendingLevelMergeData prs t -> do+ assertI "pending level merges have at least one input" $+ length prs + length t > 0+ for_ prs $ \case+ PreExistingRunData _r -> Right ()+ PreExistingMergingRunData mr -> mergingRunDataInvariant mr+ for_ (drop 1 (reverse prs)) $ \case+ PreExistingRunData _r -> Right ()+ PreExistingMergingRunData mr ->+ assertI "only the last pre-existing run can be a last level merge" $+ mergingRunDataMergeType mr == MR.MergeMidLevel+ for_ t mergingTreeDataInvariant+ PendingUnionMergeData ts -> do+ assertI "pending union merges are non-trivial (at least two inputs)" $+ length ts >= 2+ for_ ts mergingTreeDataInvariant+ where+ assertI msg False = Left msg+ assertI _ True = Right ()++mapMergingTreeData ::+ Ord k'+ => (k -> k') -> (v -> v') -> (b -> b')+ -> MergingTreeData k v b -> MergingTreeData k' v' b'+mapMergingTreeData f g h = \case+ CompletedTreeMergeData r ->+ CompletedTreeMergeData $ mapRunData f g h r+ OngoingTreeMergeData mr ->+ OngoingTreeMergeData $ mapMergingRunData f g h mr+ PendingLevelMergeData prs t ->+ PendingLevelMergeData+ (map (mapPreExistingRunData f g h) prs)+ (fmap (mapMergingTreeData f g h) t)+ PendingUnionMergeData ts ->+ PendingUnionMergeData $ fmap (mapMergingTreeData f g h) ts++mapPreExistingRunData ::+ Ord k'+ => (k -> k') -> (v -> v') -> (b -> b')+ -> PreExistingRunData k v b -> PreExistingRunData k' v' b'+mapPreExistingRunData f g h = \case+ PreExistingRunData r ->+ PreExistingRunData (mapRunData f g h r)+ PreExistingMergingRunData mr ->+ PreExistingMergingRunData (mapMergingRunData f g h mr)++type SerialisedMergingTreeData =+ MergingTreeData SerialisedKey SerialisedValue SerialisedBlob++type SerialisedPreExistingRunData =+ PreExistingRunData SerialisedKey SerialisedValue SerialisedBlob++serialiseMergingTreeData ::+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => MergingTreeData k v b -> SerialisedMergingTreeData+serialiseMergingTreeData =+ mapMergingTreeData serialiseKey serialiseValue serialiseBlob++{-------------------------------------------------------------------------------+ QuickCheck+-------------------------------------------------------------------------------}++labelMergingTreeData :: SerialisedMergingTreeData -> Property -> Property+labelMergingTreeData = \rd ->+ tabulate "tree depth" [showPowersOf 2 (depthTree rd)] . go rd+ where+ go (CompletedTreeMergeData rd) =+ tabulate "merging tree state" ["CompletedTreeMerge"]+ . labelRunData rd+ go (OngoingTreeMergeData mrd) =+ tabulate "merging tree state" ["OngoingTreeMerge"]+ . labelMergingRunData mrd+ go (PendingLevelMergeData prds mtd) =+ tabulate "merging tree state" ["PendingLevelMerge"]+ . foldr ((.) . labelPreExistingRunData) id prds+ . maybe id go mtd+ go (PendingUnionMergeData mtds) =+ tabulate "merging tree state" ["PendingUnionMerge"]+ . foldr ((.) . go) id mtds++ -- the longest path from the root to a run+ depthTree = (+1) . \case -- maximum depth of children+ CompletedTreeMergeData _ -> 0+ OngoingTreeMergeData _ -> 0+ PendingLevelMergeData prds mtds ->+ maximum (0 : fmap (const 1) prds ++ map depthTree (toList mtds))+ PendingUnionMergeData mtds ->+ maximum (0 : map depthTree mtds)+++labelPreExistingRunData :: SerialisedPreExistingRunData -> Property -> Property+labelPreExistingRunData (PreExistingRunData rd) = labelRunData rd+labelPreExistingRunData (PreExistingMergingRunData mrd) = labelMergingRunData mrd++instance ( Ord k, Arbitrary k, Arbitrary v, Arbitrary b+ ) => Arbitrary (MergingTreeData k v b) where+ arbitrary = genMergingTreeData arbitrary arbitrary arbitrary+ shrink = shrinkMergingTreeData shrink shrink shrink++genMergingTreeData ::+ Ord k => Gen k -> Gen v -> Gen b -> Gen (MergingTreeData k v b)+genMergingTreeData genKey genVal genBlob =+ QC.frequency+ -- Only at the root, we can have pending merges with no children, see+ -- 'MR.newPendingLevelMerge' and 'MR.newPendingUnionMerge'.+ [ ( 1, pure $ PendingLevelMergeData [] Nothing)+ , ( 1, pure $ PendingUnionMergeData [])+ , (50, QC.sized $ \s -> do+ treeSize <- QC.chooseInt (1, 1 + (s `div` 4)) -- up to 26+ genMergingTreeDataOfSize genKey genVal genBlob treeSize)+ ]++-- | Minimal returned size will be 1. Doesn't generate structurally empty trees!+--+-- The size is measured by the number of MergingTreeData constructors.+genMergingTreeDataOfSize ::+ forall k v b. Ord k+ => Gen k -> Gen v -> Gen b -> Int -> Gen (MergingTreeData k v b)+genMergingTreeDataOfSize genKey genVal genBlob = \n0 -> do+ tree <- genMergingTree n0+ assert (mergingTreeDataSize tree == n0) $+ pure tree+ where+ genMergingTree n+ | n < 1+ = error ("arbitrary T: n == " <> show n)++ | n == 1+ = QC.oneof+ [ CompletedTreeMergeData <$> genRun+ , OngoingTreeMergeData <$> genMergingRun arbitrary+ , genPendingLevelMergeNoChild+ ]++ | n == 2+ = genPendingLevelMergeWithChild n++ | otherwise+ = QC.oneof [genPendingLevelMergeWithChild n, genPendingUnionMerge n]++ -- n == 1+ genPendingLevelMergeNoChild = do+ numPreExisting <- chooseIntSkewed (0, 6)+ initPreExisting <- QC.vectorOf numPreExisting $+ -- these can't be last level. we generate the last input below.+ genPreExistingRun (pure MR.MergeMidLevel)+ -- there must be at least one (last) input to the pending merge.+ lastPreExisting <- genPreExistingRun arbitrary+ let preExisting = initPreExisting ++ [lastPreExisting]+ pure (PendingLevelMergeData preExisting Nothing)++ -- n >= 2+ genPendingLevelMergeWithChild n = do+ numPreExisting <- chooseIntSkewed (0, 6)+ preExisting <- QC.vectorOf numPreExisting $+ -- there can't be a last level merge, child is last+ genPreExistingRun (pure MR.MergeMidLevel)+ tree <- genMergingTree (n - 1)+ pure (PendingLevelMergeData preExisting (Just tree))++ -- n >= 3, needs 1 constructor + 2 children+ genPendingUnionMerge n = do+ ns <- QC.shuffle =<< arbitraryPartition2 (n - 1)+ PendingUnionMergeData <$> traverse genMergingTree ns++ genRun = genScaled genRunData+ genMergingRun genType = genScaled (genMergingRunData genType)+ genPreExistingRun genType = genScaled (genPreExistingRunData genType)++ -- To avoid generating too large test cases, we reduce the number of+ -- entries for each run. The size of the individual entries is unaffected.+ genScaled :: forall r. (Gen k -> Gen v -> Gen b -> Gen r) -> Gen r+ genScaled gen =+ QC.sized $ \s ->+ QC.scale (`div` 2) $+ gen (QC.resize s genKey) (QC.resize s genVal) (QC.resize s genBlob)++ -- skewed towards smaller values+ chooseIntSkewed (lb, ub) = do+ ub' <- QC.chooseInt (lb, ub)+ QC.chooseInt (lb, ub')++mergingTreeDataSize :: MergingTreeData k v b -> Int+mergingTreeDataSize = \case+ CompletedTreeMergeData _ -> 1+ OngoingTreeMergeData _ -> 1+ PendingLevelMergeData _ tree -> 1 + maybe 0 mergingTreeDataSize tree+ PendingUnionMergeData trees -> 1 + sum (map mergingTreeDataSize trees)++-- Split into at least two smaller positive numbers. The input needs to be+-- greater than or equal to 2.+arbitraryPartition2 :: Int -> QC.Gen [Int]+arbitraryPartition2 n = assert (n >= 2) $ do+ first <- QC.chooseInt (1, n-1)+ (first :) <$> arbitraryPartition (n - first)++-- Split into smaller positive numbers.+arbitraryPartition :: Int -> QC.Gen [Int]+arbitraryPartition n+ | n < 1 = pure []+ | n == 1 = pure [1]+ | otherwise = do+ first <- QC.chooseInt (1, n)+ (first :) <$> arbitraryPartition (n - first)++-- TODO: Would it be useful to shrink by merging subtrees into a single run?+-- This would simplify the tree while preserving many errors that depend on the+-- specific content of the tree. See prototype tests.+shrinkMergingTreeData ::+ Ord k+ => (k -> [k])+ -> (v -> [v])+ -> (b -> [b])+ -> MergingTreeData k v b+ -> [MergingTreeData k v b]+shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob = \case+ CompletedTreeMergeData r ->+ [ CompletedTreeMergeData r'+ | r' <- shrinkRunData shrinkKey shrinkVal shrinkBlob r+ ]+ OngoingTreeMergeData mr ->+ [ OngoingTreeMergeData mr'+ | mr' <- shrinkMergingRunData shrinkKey shrinkVal shrinkBlob mr+ ]+ PendingLevelMergeData prs t ->+ -- just use the child tree, if present+ [ t' | Just t' <- [t] ]+ <>+ -- move completed child tree into regular levels+ [ PendingLevelMergeData (prs ++ [PreExistingRunData r]) Nothing+ | Just (CompletedTreeMergeData r) <- [t]+ ]+ <>+ [ PendingLevelMergeData prs' t'+ | (prs', t') <-+ liftShrink2+ (liftShrink (shrinkPreExistingRunData shrinkKey shrinkVal shrinkBlob))+ (liftShrink (shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob))+ (prs, t)+ , length prs' + length t' > 0+ ]+ PendingUnionMergeData ts ->+ ts+ <>+ [ PendingUnionMergeData ts'+ | ts' <- liftShrink (shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob) ts+ , length ts' >= 2+ ]++genPreExistingRunData ::+ Ord k+ => Gen MR.LevelMergeType+ -> Gen k+ -> Gen v+ -> Gen b+ -> Gen (PreExistingRunData k v b)+genPreExistingRunData genMergeType genKey genVal genBlob =+ QC.oneof+ [ PreExistingRunData <$> genRunData genKey genVal genBlob+ , PreExistingMergingRunData <$> genMergingRunData genMergeType genKey genVal genBlob+ ]++shrinkPreExistingRunData ::+ Ord k+ => (k -> [k])+ -> (v -> [v])+ -> (b -> [b])+ -> PreExistingRunData k v b+ -> [PreExistingRunData k v b]+shrinkPreExistingRunData shrinkKey shrinkVal shrinkBlob = \case+ PreExistingRunData r ->+ [ PreExistingRunData r'+ | r' <- shrinkRunData shrinkKey shrinkVal shrinkBlob r+ ]+ PreExistingMergingRunData mr ->+ [ PreExistingMergingRunData mr'+ | mr' <- shrinkMergingRunData shrinkKey shrinkVal shrinkBlob mr+ ]
@@ -0,0 +1,903 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE QuantifiedConstraints #-}+{-# LANGUAGE UndecidableInstances #-}++{-# OPTIONS_GHC -Wno-orphans #-}++-- | 'NoThunks' orphan instances+module Database.LSMTree.Extras.NoThunks (+ assertNoThunks+ , propUnsafeNoThunks+ , propNoThunks+ , NoThunksIOLike+ ) where++import Control.Concurrent.Class.MonadMVar.Strict+import Control.Concurrent.Class.MonadSTM.RWVar+import Control.Concurrent.Class.MonadSTM.Strict+import Control.Exception+import Control.Monad.Primitive+import Control.Monad.ST.Unsafe (unsafeIOToST, unsafeSTToIO)+import Control.RefCount+import Control.Tracer+import Data.Bit+import Data.Map.Strict+import Data.Primitive+import Data.Primitive.PrimVar+import Data.Proxy+import Data.STRef+import Data.Typeable+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word+import Database.LSMTree.Internal.Arena+import Database.LSMTree.Internal.BlobFile+import Database.LSMTree.Internal.BlobRef+import Database.LSMTree.Internal.BloomFilter (Bloom, MBloom)+import Database.LSMTree.Internal.ChecksumHandle+import Database.LSMTree.Internal.Chunk+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.CRC32C+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.IncomingRun+import Database.LSMTree.Internal.Index+import Database.LSMTree.Internal.Index.Compact+import Database.LSMTree.Internal.Index.CompactAcc+import Database.LSMTree.Internal.Index.Ordinary+import Database.LSMTree.Internal.Index.OrdinaryAcc+import Database.LSMTree.Internal.Merge+import qualified Database.LSMTree.Internal.Merge as Merge+import Database.LSMTree.Internal.MergeSchedule+import Database.LSMTree.Internal.MergingRun+import Database.LSMTree.Internal.MergingTree+import Database.LSMTree.Internal.MergingTree.Lookup (LookupTree (..))+import Database.LSMTree.Internal.Page+import Database.LSMTree.Internal.PageAcc+import Database.LSMTree.Internal.Paths+import Database.LSMTree.Internal.RawBytes+import Database.LSMTree.Internal.RawOverflowPage+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Readers+import Database.LSMTree.Internal.Run+import Database.LSMTree.Internal.RunAcc+import Database.LSMTree.Internal.RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.RunReader hiding (Entry)+import qualified Database.LSMTree.Internal.RunReader as Reader+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Types as Types+import Database.LSMTree.Internal.UniqCounter+import Database.LSMTree.Internal.Unsafe as Unsafe+import Database.LSMTree.Internal.Unsliced+import Database.LSMTree.Internal.Vector.Growing+import Database.LSMTree.Internal.WriteBuffer+import Database.LSMTree.Internal.WriteBufferBlobs+import GHC.Generics+import KMerge.Heap+import NoThunks.Class+import System.FS.API+import System.FS.BlockIO.API+import System.FS.IO+import System.FS.Sim.MockFS+import Test.QuickCheck (Property, Testable (..), counterexample)+import Unsafe.Coerce++assertNoThunks :: NoThunks a => a -> b -> b+assertNoThunks x = assert p+ where p = case unsafeNoThunks x of+ Nothing -> True+ Just thunkInfo -> error $ "Assertion failed: found thunk" <> show thunkInfo++propUnsafeNoThunks :: NoThunks a => a -> Property+propUnsafeNoThunks x =+ case unsafeNoThunks x of+ Nothing -> property True+ Just thunkInfo -> counterexample ("Found thunk " <> show thunkInfo) False++propNoThunks :: NoThunks a => a -> IO Property+propNoThunks x = do+ thunkInfoMay <- noThunks [] x+ pure $ case thunkInfoMay of+ Nothing -> property True+ Just thunkInfo -> counterexample ("Found thunk " <> show thunkInfo) False++{-------------------------------------------------------------------------------+ Public API+-------------------------------------------------------------------------------}++-- | Also checks 'NoThunks' for the tables that are known to be open in the session.+instance (NoThunksIOLike m, Typeable m, Typeable (PrimState m))+ => NoThunks (Types.Session m) where+ showTypeOf (_ :: Proxy (Types.Session m)) = "Database.LSMTree.Session"+ wNoThunks ctx (Types.Session s) = wNoThunks ctx s++-- | Does not check 'NoThunks' for the session that this table belongs to.+instance (NoThunksIOLike m, Typeable m, Typeable (PrimState m))+ => NoThunks (Types.Table m k v b) where+ showTypeOf (_ :: Proxy (Types.Table m k v b)) = "Database.LSMTree.Table"+ wNoThunks ctx (Types.Table t) = wNoThunks ctx t++{-------------------------------------------------------------------------------+ Unsafe+-------------------------------------------------------------------------------}++deriving stock instance Generic (Unsafe.Session m h)+-- | Also checks 'NoThunks' for the 'Unsafe.Table's that are known to be+-- open in the 'Unsafe.Session'.+deriving anyclass instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (Unsafe.Session m h)++deriving stock instance Generic (SessionState m h)+deriving anyclass instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (SessionState m h)++deriving stock instance Generic (SessionEnv m h)+deriving anyclass instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (SessionEnv m h)++deriving stock instance Generic (Unsafe.Table m h)+-- | Does not check 'NoThunks' for the 'Unsafe.Session' that this+-- 'Unsafe.Table' belongs to.+deriving via AllowThunksIn '["tableSession"] (Unsafe.Table m h)+ instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (Unsafe.Table m h)++deriving stock instance Generic (TableState m h)+deriving anyclass instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (TableState m h)++deriving stock instance Generic (TableEnv m h)+deriving via AllowThunksIn '["tableSessionEnv"] (TableEnv m h)+ instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (TableEnv m h)++-- | Does not check 'NoThunks' for the 'Unsafe.Session' that this+-- 'Unsafe.Cursor' belongs to.+deriving stock instance Generic (Unsafe.Cursor m h)+deriving anyclass instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (Unsafe.Cursor m h)++deriving stock instance Generic (CursorState m h)+deriving anyclass instance (NoThunksIOLike m, Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (CursorState m h)++deriving stock instance Generic (CursorEnv m h)+deriving via AllowThunksIn ["cursorSession", "cursorSessionEnv"] (CursorEnv m h)+ instance (Typeable m, Typeable h, Typeable (PrimState m))+ => NoThunks (CursorEnv m h)++deriving stock instance Generic TableId+deriving anyclass instance NoThunks TableId++deriving stock instance Generic CursorId+deriving anyclass instance NoThunks CursorId++{-------------------------------------------------------------------------------+ UniqCounter+-------------------------------------------------------------------------------}++deriving stock instance Generic (UniqCounter m)+deriving anyclass instance (NoThunks (PrimVar (PrimState m) Int))+ => NoThunks (UniqCounter m)++{-------------------------------------------------------------------------------+ Serialise+-------------------------------------------------------------------------------}++deriving stock instance Generic RawBytes+deriving anyclass instance NoThunks RawBytes++deriving stock instance Generic SerialisedKey+deriving anyclass instance NoThunks SerialisedKey++deriving stock instance Generic SerialisedValue+deriving anyclass instance NoThunks SerialisedValue++deriving stock instance Generic SerialisedBlob+deriving anyclass instance NoThunks SerialisedBlob++instance NoThunks (Unsliced a) where+ showTypeOf (_ :: Proxy (Unsliced a)) = "Unsliced"+ wNoThunks ctx (x :: Unsliced a) = noThunks ctx y+ where+ -- Unsliced is a newtype around a ByteArray, so we can unsafeCoerce+ -- safely. The bang pattern will only evaluate the coercion, because the+ -- byte array is already in WHNF.+ y :: ByteArray+ !y = unsafeCoerce x++{-------------------------------------------------------------------------------+ Run+-------------------------------------------------------------------------------}++deriving stock instance Generic (Run m h)+deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (Run m h)++deriving stock instance Generic RunParams+deriving anyclass instance NoThunks RunParams++deriving stock instance Generic RunBloomFilterAlloc+deriving anyclass instance NoThunks RunBloomFilterAlloc++deriving stock instance Generic RunDataCaching+deriving anyclass instance NoThunks RunDataCaching++deriving stock instance Generic IndexType+deriving anyclass instance NoThunks IndexType++{-------------------------------------------------------------------------------+ Paths+-------------------------------------------------------------------------------}++deriving stock instance Generic RunNumber+deriving anyclass instance NoThunks RunNumber++deriving stock instance Generic SessionRoot+deriving anyclass instance NoThunks SessionRoot++deriving stock instance Generic RunFsPaths+deriving anyclass instance NoThunks RunFsPaths++deriving stock instance Generic (ForKOps a)+deriving anyclass instance NoThunks a => NoThunks (ForKOps a)++deriving stock instance Generic (ForBlob a)+deriving anyclass instance NoThunks a => NoThunks (ForBlob a)++deriving stock instance Generic (ForFilter a)+deriving anyclass instance NoThunks a => NoThunks (ForFilter a)++deriving stock instance Generic (ForIndex a)+deriving anyclass instance NoThunks a => NoThunks (ForIndex a)++deriving stock instance Generic (ForRunFiles a)+deriving anyclass instance NoThunks a => NoThunks (ForRunFiles a)++{-------------------------------------------------------------------------------+ CRC32C+-------------------------------------------------------------------------------}++deriving stock instance Generic CRC32C+deriving anyclass instance NoThunks CRC32C++{-------------------------------------------------------------------------------+ WriteBuffer+-------------------------------------------------------------------------------}++instance NoThunks WriteBuffer where+ showTypeOf (_ :: Proxy WriteBuffer) = "WriteBuffer"+ wNoThunks ctx (x :: WriteBuffer) = noThunks ctx y+ where+ -- toMap simply unwraps the WriteBuffer newtype wrapper. The bang pattern+ -- will only evaluate the coercion, because the inner Map is already in+ -- WHNF.+ y :: Map SerialisedKey (Entry SerialisedValue BlobSpan)+ !y = toMap x++{-------------------------------------------------------------------------------+ BlobFile+-------------------------------------------------------------------------------}++deriving stock instance Generic (WriteBufferBlobs m h)+deriving anyclass instance (Typeable (PrimState m), Typeable m, Typeable h)+ => NoThunks (WriteBufferBlobs m h)++deriving stock instance Generic (FilePointer m)+deriving anyclass instance Typeable (PrimState m)+ => NoThunks (FilePointer m)++{-------------------------------------------------------------------------------+ Index+-------------------------------------------------------------------------------}++deriving stock instance Generic IndexCompact+deriving anyclass instance NoThunks IndexCompact++deriving stock instance Generic PageNo+deriving anyclass instance NoThunks PageNo++deriving stock instance Generic IndexOrdinary+deriving anyclass instance NoThunks IndexOrdinary++deriving stock instance Generic Index+deriving anyclass instance NoThunks Index++{-------------------------------------------------------------------------------+ MergeSchedule+-------------------------------------------------------------------------------}++deriving stock instance Generic (TableContent m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState LevelMergeType m h))+ , NoThunks (StrictMVar m (MergingTreeState m h))+ ) => NoThunks (TableContent m h)++deriving stock instance Generic (LevelsCache m h)+deriving anyclass instance+ (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (LevelsCache m h)++deriving stock instance Generic (Level m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState LevelMergeType m h))+ ) => NoThunks (Level m h)++deriving stock instance Generic (IncomingRun m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState LevelMergeType m h))+ ) => NoThunks (IncomingRun m h)++deriving stock instance Generic (UnionLevel m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingTreeState m h))+ ) => NoThunks (UnionLevel m h)++deriving stock instance Generic (UnionCache m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ ) => NoThunks (UnionCache m h)++deriving stock instance Generic MergePolicyForLevel+deriving anyclass instance NoThunks MergePolicyForLevel++deriving stock instance Generic NominalDebt+deriving anyclass instance NoThunks NominalDebt++deriving stock instance Generic NominalCredits+deriving anyclass instance NoThunks NominalCredits++{-------------------------------------------------------------------------------+ MergingRun+-------------------------------------------------------------------------------}++deriving stock instance Generic (MergingRun t m h)+deriving anyclass instance ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState t m h))+ ) => NoThunks (MergingRun t m h)++deriving stock instance Generic (MergingRunState t m h)+deriving anyclass instance ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks t+ ) => NoThunks (MergingRunState t m h)++deriving stock instance Generic MergeDebt+deriving anyclass instance NoThunks MergeDebt++deriving stock instance Generic MergeCredits+deriving anyclass instance NoThunks MergeCredits++deriving stock instance Generic (CreditsVar s)+deriving anyclass instance Typeable s => NoThunks (CreditsVar s)++deriving stock instance Generic MergeKnownCompleted+deriving anyclass instance NoThunks MergeKnownCompleted++{-------------------------------------------------------------------------------+ MergingTree+-------------------------------------------------------------------------------}++deriving stock instance Generic (MergingTree m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingTreeState m h))+ ) => NoThunks (MergingTree m h)++deriving stock instance Generic (MergingTreeState m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState LevelMergeType m h))+ , NoThunks (StrictMVar m (MergingRunState TreeMergeType m h))+ , NoThunks (StrictMVar m (MergingTreeState m h))+ ) => NoThunks (MergingTreeState m h)++deriving stock instance Generic (PendingMerge m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState LevelMergeType m h))+ , NoThunks (StrictMVar m (MergingTreeState m h))+ ) => NoThunks (PendingMerge m h)++deriving stock instance Generic (PreExistingRun m h)+deriving anyclass instance+ ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks (StrictMVar m (MergingRunState LevelMergeType m h))+ ) => NoThunks (PreExistingRun m h)++deriving stock instance Generic (LookupTree a)+deriving anyclass instance NoThunks a => NoThunks (LookupTree a)++{-------------------------------------------------------------------------------+ Entry+-------------------------------------------------------------------------------}++deriving stock instance Generic (Entry v b)+deriving anyclass instance (NoThunks v, NoThunks b)+ => NoThunks (Entry v b)++deriving stock instance Generic NumEntries+deriving anyclass instance NoThunks NumEntries++{-------------------------------------------------------------------------------+ RunBuilder+-------------------------------------------------------------------------------}++deriving stock instance Generic (RunBuilder m h)+deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (RunBuilder m h)++deriving stock instance Generic (ChecksumHandle s h)+deriving anyclass instance (Typeable s, Typeable h)+ => NoThunks (ChecksumHandle s h)++{-------------------------------------------------------------------------------+ RunAcc+-------------------------------------------------------------------------------}++deriving stock instance Generic (RunAcc s)+deriving anyclass instance Typeable s+ => NoThunks (RunAcc s)++{-------------------------------------------------------------------------------+ IndexAcc+-------------------------------------------------------------------------------}++deriving stock instance Generic (IndexCompactAcc s)+deriving anyclass instance Typeable s+ => NoThunks (IndexCompactAcc s)++deriving stock instance Generic (SMaybe a)+deriving anyclass instance NoThunks a => NoThunks (SMaybe a)++deriving stock instance Generic (IndexOrdinaryAcc s)+deriving anyclass instance Typeable s+ => NoThunks (IndexOrdinaryAcc s)++deriving stock instance Generic (IndexAcc s)+deriving anyclass instance Typeable s+ => NoThunks (IndexAcc s)++{-------------------------------------------------------------------------------+ GrowingVector+-------------------------------------------------------------------------------}++instance (NoThunks a, Typeable s, Typeable a) => NoThunks (GrowingVector s a) where+ showTypeOf (p :: Proxy (GrowingVector s a)) = show $ typeRep p+ wNoThunks ctx+ (GrowingVector (a :: STRef s (VM.MVector s a)) (b :: PrimVar s Int))+ = allNoThunks [+ noThunks ctx b+ -- Check that the STRef is in WHNF+ , noThunks ctx $ OnlyCheckWhnf a+ -- Check that the MVector is in WHNF+ , do+ mvec <- unsafeSTToIO $ readSTRef a+ noThunks ctx' $ OnlyCheckWhnf mvec+ -- Check that the vector elements contain no thunks. The vector+ -- contains undefined elements after the first @n@ elements+ , do+ n <- unsafeSTToIO $ readPrimVar b+ mvec <- unsafeSTToIO $ readSTRef a+ allNoThunks [+ unsafeSTToIO (VM.read mvec i) >>= \x -> noThunks ctx'' x+ | i <- [0..n-1]+ ]+ ]+ where+ ctx' = showTypeOf (Proxy @(STRef s (VM.MVector s a))) : ctx+ ctx'' = showTypeOf (Proxy @(VM.MVector s a)) : ctx'++{-------------------------------------------------------------------------------+ Baler+-------------------------------------------------------------------------------}++deriving stock instance Generic (Baler s)+deriving anyclass instance Typeable s+ => NoThunks (Baler s)++{-------------------------------------------------------------------------------+ PageAcc+-------------------------------------------------------------------------------}++deriving stock instance Generic (PageAcc s)+deriving anyclass instance Typeable s+ => NoThunks (PageAcc s)++{-------------------------------------------------------------------------------+ Merge+-------------------------------------------------------------------------------}++deriving stock instance Generic (Merge t m h)+deriving anyclass instance ( Typeable m, Typeable (PrimState m), Typeable h+ , NoThunks t+ ) => NoThunks (Merge t m h)++deriving stock instance Generic MergeType+deriving anyclass instance NoThunks MergeType++deriving stock instance Generic LevelMergeType+deriving anyclass instance NoThunks LevelMergeType++deriving stock instance Generic TreeMergeType+deriving anyclass instance NoThunks TreeMergeType++deriving stock instance Generic Merge.StepResult+deriving anyclass instance NoThunks Merge.StepResult++deriving stock instance Generic Merge.MergeState+deriving anyclass instance NoThunks Merge.MergeState++{-------------------------------------------------------------------------------+ Readers+-------------------------------------------------------------------------------}++deriving stock instance Generic (Readers m h)+deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (Readers m h)++deriving stock instance Generic (Reader m h)+instance (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (Reader m h) where+ showTypeOf (_ :: Proxy (Reader m h)) = "Reader"+ wNoThunks ctx = \case+ ReadRun r -> noThunks ctx r+ ReadBuffer var -> noThunks ctx (OnlyCheckWhnf var) -- contents intentionally lazy+ ReadReaders ty readers -> allNoThunks [+ noThunks ctx ty+ , noThunks ctx readers+ ]++deriving stock instance Generic ReadersMergeType+deriving anyclass instance NoThunks ReadersMergeType++deriving stock instance Generic HasMore+deriving anyclass instance NoThunks HasMore++deriving stock instance Generic ReaderNumber+deriving anyclass instance NoThunks ReaderNumber++deriving stock instance Generic (ReadCtx m h)+deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (ReadCtx m h)++{-------------------------------------------------------------------------------+ Reader+-------------------------------------------------------------------------------}++deriving stock instance Generic (RunReader m h)+deriving anyclass instance (Typeable m, Typeable (PrimState m), Typeable h)+ => NoThunks (RunReader m h)++-- | Allows thunks in the overflow pages+instance ( Typeable m, Typeable (PrimState m), Typeable h+ ) => NoThunks (Reader.Entry m h) where+ showTypeOf (p :: Proxy (Reader.Entry m h)) = show $ typeRep p+ wNoThunks ctx (Reader.Entry (e :: Entry SerialisedValue (RawBlobRef m h))) = noThunks ctx e+ wNoThunks ctx (EntryOverflow+ (entryPrefix :: Entry SerialisedValue (RawBlobRef m h))+ (page :: RawPage)+ (len :: Word32)+ (overflowPages :: [RawOverflowPage]) ) =+ allNoThunks [+ noThunks ctx entryPrefix+ , noThunks ctx page+ , noThunks ctx len+ , noThunks ctx (OnlyCheckWhnf overflowPages)+ ]++{-------------------------------------------------------------------------------+ RawPage+-------------------------------------------------------------------------------}++deriving stock instance Generic RawPage+deriving anyclass instance NoThunks RawPage++{-------------------------------------------------------------------------------+ RawPage+-------------------------------------------------------------------------------}++deriving stock instance Generic RawOverflowPage+deriving anyclass instance NoThunks RawOverflowPage++{-------------------------------------------------------------------------------+ BlobRef+-------------------------------------------------------------------------------}++deriving stock instance Generic BlobSpan+deriving anyclass instance NoThunks BlobSpan++deriving stock instance Generic (BlobFile m h)+deriving anyclass instance (Typeable h, Typeable (PrimState m))+ => NoThunks (BlobFile m h)++deriving stock instance Generic (RawBlobRef m h)+deriving anyclass instance (Typeable h, Typeable (PrimState m))+ => NoThunks (RawBlobRef m h)++deriving stock instance Generic (WeakBlobRef m h)+deriving anyclass instance (Typeable h, Typeable m, Typeable (PrimState m))+ => NoThunks (WeakBlobRef m h)++{-------------------------------------------------------------------------------+ Arena+-------------------------------------------------------------------------------}++-- TODO: proper instance+deriving via OnlyCheckWhnf (ArenaManager m)+ instance Typeable m => NoThunks (ArenaManager m)++{-------------------------------------------------------------------------------+ Config+-------------------------------------------------------------------------------}++deriving stock instance Generic TableConfig+deriving anyclass instance NoThunks TableConfig++deriving stock instance Generic MergePolicy+deriving anyclass instance NoThunks MergePolicy++deriving stock instance Generic SizeRatio+deriving anyclass instance NoThunks SizeRatio++deriving stock instance Generic WriteBufferAlloc+deriving anyclass instance NoThunks WriteBufferAlloc++deriving stock instance Generic BloomFilterAlloc+deriving anyclass instance NoThunks BloomFilterAlloc++deriving stock instance Generic FencePointerIndexType+deriving anyclass instance NoThunks FencePointerIndexType++deriving stock instance Generic DiskCachePolicy+deriving anyclass instance NoThunks DiskCachePolicy++deriving stock instance Generic MergeSchedule+deriving anyclass instance NoThunks MergeSchedule++deriving stock instance Generic MergeBatchSize+deriving anyclass instance NoThunks MergeBatchSize++{-------------------------------------------------------------------------------+ RWVar+-------------------------------------------------------------------------------}++deriving stock instance Generic (RWVar m a)+deriving anyclass instance NoThunks (StrictTVar m (RWState a)) => NoThunks (RWVar m a)++deriving stock instance Generic (RWState a)+deriving anyclass instance NoThunks a => NoThunks (RWState a)++{-------------------------------------------------------------------------------+ RefCounter+-------------------------------------------------------------------------------}++instance Typeable (PrimState m) => NoThunks (RefCounter m) where+ showTypeOf (_ :: Proxy (RefCounter m)) = "RefCounter"+ wNoThunks ctx+ (RefCounter (a :: PrimVar (PrimState m) Int) (b :: m ()))+ = allNoThunks [+ noThunks ctx a+ , noThunks ctx $ (OnlyCheckWhnfNamed b :: OnlyCheckWhnfNamed "finaliser" (m ()))+ ]++-- Ref constructor not exported, cannot derive Generic, use DeRef instead.+instance (NoThunks obj, Typeable obj) => NoThunks (Ref obj) where+ showTypeOf p@(_ :: Proxy (Ref obj)) = show $ typeRep p+ wNoThunks ctx (DeRef ref) = noThunks ctx ref++deriving stock instance Generic (WeakRef obj)+deriving anyclass instance (NoThunks obj, Typeable obj) => NoThunks (WeakRef obj)++{-------------------------------------------------------------------------------+ kmerge+-------------------------------------------------------------------------------}++instance (NoThunks a, Typeable s, Typeable a) => NoThunks (MutableHeap s a) where+ showTypeOf (p :: Proxy (MutableHeap s a)) = show $ typeRep p+ wNoThunks ctx+ (MH (a :: PrimVar s Int) (b :: SmallMutableArray s a))+ = allNoThunks [+ noThunks ctx a+ -- Check that the array is in WHNF+ , noThunks ctx (OnlyCheckWhnf b)+ -- Check that the array elements contain no thunks. The small array+ -- may contain undefined placeholder values after the first @n@+ -- elements in the array. The very first element of the array can also+ -- be undefined.+ , do+ n <- unsafeSTToIO (readPrimVar a)+ allNoThunks [+ unsafeSTToIO (readSmallArray b i) >>= \x -> noThunks ctx' x+ | i <- [1..n-1]+ ]+ ]+ where+ ctx' = showTypeOf (Proxy @(SmallMutableArray s a)) : ctx++{-------------------------------------------------------------------------------+ IOLike+-------------------------------------------------------------------------------}++-- | 'NoThunks' constraints for IO-like monads+--+-- Some constraints, like @NoThunks (MutVar s a)@ and @NoThunks (StrictTVar m+-- a)@, can not be satisfied for arbitrary @m@\/@s@, and must be instantiated+-- for a concrete @m@\/@s@, like @IO@\/@RealWorld@.+class ( forall a. (NoThunks a, Typeable a) => NoThunks (StrictTVar m a)+ , forall a. (NoThunks a, Typeable a) => NoThunks (StrictMVar m a)+ ) => NoThunksIOLike' m s++instance NoThunksIOLike' IO RealWorld++type NoThunksIOLike m = NoThunksIOLike' m (PrimState m)++instance (NoThunks a, Typeable a) => NoThunks (StrictTVar IO a) where+ showTypeOf (p :: Proxy (StrictTVar IO a)) = show $ typeRep p+ wNoThunks _ctx _var = do+ x <- readTVarIO _var+ noThunks _ctx x++-- TODO: in some cases, strict-mvar functions leave thunks behind, in particular+-- modifyMVarMasked and modifyMVarMasked_. So in some specific cases we evaluate+-- the contents of the MVar to WHNF, and keep checking nothunks from there. See+-- lsm-tree#444.+--+-- TODO: we tried using overlapping instances for @StrictMVar IO a@ and+-- @StrictMVar IO (MergingRunState IO h)@, but the quantified constraint in+-- NoThunksIOLike' will throw a compiler error telling us to mark the instances+-- for StrictMVar as incoherent. Marking them as incoherent makes the tests+-- fail... We are unsure if it can be overcome, but the current casting approach+-- works, so there is no priority to use rewrite this code to use overlapping+-- instances.+instance (NoThunks a, Typeable a) => NoThunks (StrictMVar IO a) where+ showTypeOf (p :: Proxy (StrictMVar IO a)) = show $ typeRep p+ wNoThunks ctx var+ -- TODO: Revisit which of these cases are still needed.+ | Just (Proxy :: Proxy (MergingRunState LevelMergeType IO HandleIO))+ <- gcast (Proxy @a)+ = workAroundCheck+ | Just (Proxy :: Proxy (MergingRunState TreeMergeType IO HandleIO))+ <- gcast (Proxy @a)+ = workAroundCheck+ | Just (Proxy :: Proxy (MergingRunState LevelMergeType IO HandleMock))+ <- gcast (Proxy @a)+ = workAroundCheck+ | Just (Proxy :: Proxy (MergingRunState TreeMergeType IO HandleMock))+ <- gcast (Proxy @a)+ = workAroundCheck+ | otherwise+ = properCheck+ where+ properCheck = do+ x <- readMVar var+ noThunks ctx x++ workAroundCheck = do+ !x <- readMVar var+ noThunks ctx x++{-------------------------------------------------------------------------------+ vector+-------------------------------------------------------------------------------}++-- TODO: upstream to @nothunks@+instance (NoThunks a, Typeable s, Typeable a) => NoThunks (VM.MVector s a) where+ showTypeOf (p :: Proxy (VM.MVector s a)) = show $ typeRep p+ wNoThunks ctx v =+ allNoThunks [+ unsafeSTToIO (VM.read v i >>= \ x -> unsafeIOToST (noThunks ctx x))+ | i <- [0.. VM.length v-1]+ ]++-- TODO: https://github.com/input-output-hk/nothunks/issues/57+deriving via OnlyCheckWhnf (VP.Vector a)+ instance Typeable a => NoThunks (VP.Vector a)++-- TODO: upstream to @nothunks@+deriving via OnlyCheckWhnf (VUM.MVector s Word64)+ instance Typeable s => NoThunks (VUM.MVector s Word64)++-- TODO: upstream to @nothunks@+deriving via OnlyCheckWhnf (VUM.MVector s Bit)+ instance Typeable s => NoThunks (VUM.MVector s Bit)++-- TODO: upstream to @nothunks@+deriving via OnlyCheckWhnf (VP.MVector s Word8)+ instance Typeable s => NoThunks (VP.MVector s Word8)++{-------------------------------------------------------------------------------+ ST+-------------------------------------------------------------------------------}++-- TODO: upstream to @nothunks@+instance NoThunks a => NoThunks (STRef s a) where+ showTypeOf (_ :: Proxy (STRef s a)) = "STRef"+ wNoThunks ctx ref = do+ x <- unsafeSTToIO $ readSTRef ref+ noThunks ctx x++{-------------------------------------------------------------------------------+ primitive+-------------------------------------------------------------------------------}++-- TODO: https://github.com/input-output-hk/nothunks/issues/56+instance NoThunks a => NoThunks (MutVar s a) where+ showTypeOf (_ :: Proxy (MutVar s a)) = "MutVar"+ wNoThunks ctx var = do+ x <- unsafeSTToIO $ readMutVar var+ noThunks ctx x++-- TODO: https://github.com/input-output-hk/nothunks/issues/56+deriving via OnlyCheckWhnf (PrimVar s a)+ instance (Typeable s, Typeable a) => NoThunks (PrimVar s a)++-- TODO: https://github.com/input-output-hk/nothunks/issues/56+instance NoThunks a => NoThunks (SmallMutableArray s a) where+ showTypeOf (_ :: Proxy (SmallMutableArray s a)) = "SmallMutableArray"+ wNoThunks ctx arr = do+ n <- unsafeSTToIO $ getSizeofSmallMutableArray arr+ allNoThunks [+ unsafeSTToIO (readSmallArray arr i) >>= \x -> noThunks ctx x+ | i <- [0..n-1]+ ]++-- TODO: https://github.com/input-output-hk/nothunks/issues/56+deriving via OnlyCheckWhnf (MutablePrimArray s a)+ instance (Typeable s, Typeable a) => NoThunks (MutablePrimArray s a)++-- TODO: https://github.com/input-output-hk/nothunks/issues/56+deriving via OnlyCheckWhnf ByteArray+ instance NoThunks ByteArray++{-------------------------------------------------------------------------------+ bloomfilter+-------------------------------------------------------------------------------}++-- TODO: check heap?+deriving via OnlyCheckWhnf (Bloom a)+ instance Typeable a => NoThunks (Bloom a)++-- TODO: check heap?+deriving via OnlyCheckWhnf (MBloom s a)+ instance (Typeable s, Typeable a) => NoThunks (MBloom s a)++{-------------------------------------------------------------------------------+ fs-api and fs-sim+-------------------------------------------------------------------------------}++-- TODO: check heap?+deriving via OnlyCheckWhnf (HasFS m h)+ instance (Typeable m, Typeable h) => NoThunks (HasFS m h)++-- TODO: check heap?+deriving via OnlyCheckWhnf (Handle h)+ instance Typeable h => NoThunks (Handle h)++-- TODO: check heap?+deriving via OnlyCheckWhnf FsPath+ instance NoThunks FsPath++{-------------------------------------------------------------------------------+ blockio+-------------------------------------------------------------------------------}++-- TODO: check heap?+deriving via OnlyCheckWhnf (HasBlockIO m h)+ instance (Typeable m, Typeable h) => NoThunks (HasBlockIO m h)++-- TODO: check heap?+deriving via OnlyCheckWhnf (LockFileHandle m)+ instance Typeable m => NoThunks (LockFileHandle m)++{-------------------------------------------------------------------------------+ contra-tracer+-------------------------------------------------------------------------------}++-- TODO: check heap?+deriving via OnlyCheckWhnf (Tracer m a)+ instance (Typeable m, Typeable a) => NoThunks (Tracer m a)
@@ -0,0 +1,176 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE TypeFamilies #-}++{-# OPTIONS_GHC -Wno-orphans #-}++module Database.LSMTree.Extras.Orphans (+ byteSwapWord256+ , indexWord8ArrayAsWord256+ , indexWord8ArrayAsWord128+ ) where++import Control.DeepSeq+import qualified Data.Primitive as P+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Generic.Mutable as VGM+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.WideWord.Word128 (Word128 (..), byteSwapWord128)+import Data.WideWord.Word256 (Word256 (..))+import Database.LSMTree.Internal.Primitive (indexWord8ArrayAsWord64)+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise (SerialisedBlob (..),+ SerialisedKey (..), SerialisedValue (..))+import Database.LSMTree.Internal.Serialise.Class+import Database.LSMTree.Internal.Vector+import GHC.Generics+import GHC.Word+import qualified System.FS.API as FS+import qualified System.FS.IO.Handle as FS+import System.Posix.Types (COff (..))+import System.Random (Uniform)+import Test.QuickCheck++{-------------------------------------------------------------------------------+ Word256+-------------------------------------------------------------------------------}++deriving anyclass instance Uniform Word256++instance SerialiseKey Word256 where+ serialiseKey w256 =+ RB.RawBytes $ mkPrimVector 0 32 $ P.runByteArray $ do+ ba <- P.newByteArray 32+ P.writeByteArray ba 0 $ byteSwapWord256 w256+ pure ba+ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word256" 32 len $+ byteSwapWord256 $ indexWord8ArrayAsWord256 ba off++instance SerialiseKeyOrderPreserving Word256++instance SerialiseValue Word256 where+ serialiseValue w256 =+ RB.RawBytes $ mkPrimVector 0 32 $ P.runByteArray $ do+ ba <- P.newByteArray 32+ P.writeByteArray ba 0 w256+ pure ba+ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word256" 32 len $+ indexWord8ArrayAsWord256 ba off++instance Arbitrary Word256 where+ arbitrary = Word256 <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary+ shrink w256 = [ w256'+ | let i256 = toInteger w256+ , i256' <- shrink i256+ , toInteger (minBound :: Word256) <= i256'+ , toInteger (maxBound :: Word256) >= i256'+ , let w256' = fromIntegral i256'+ ]++{-# INLINE byteSwapWord256 #-}+byteSwapWord256 :: Word256 -> Word256+byteSwapWord256 (Word256 a3 a2 a1 a0) =+ Word256 (byteSwap64 a0) (byteSwap64 a1) (byteSwap64 a2) (byteSwap64 a3)++{-# INLINE indexWord8ArrayAsWord256 #-}+indexWord8ArrayAsWord256 :: P.ByteArray -> Int -> Word256+indexWord8ArrayAsWord256 !ba !off =+ Word256 (indexWord8ArrayAsWord64 ba (off + 24))+ (indexWord8ArrayAsWord64 ba (off + 16))+ (indexWord8ArrayAsWord64 ba (off + 8))+ (indexWord8ArrayAsWord64 ba off)++newtype instance VUM.MVector s Word256 = MV_Word256 (VP.MVector s Word256)+newtype instance VU.Vector Word256 = V_Word256 (VP.Vector Word256)++deriving via VU.UnboxViaPrim Word256 instance VGM.MVector VU.MVector Word256+deriving via VU.UnboxViaPrim Word256 instance VG.Vector VU.Vector Word256++instance VUM.Unbox Word256++{-------------------------------------------------------------------------------+ Word128+-------------------------------------------------------------------------------}++deriving anyclass instance Uniform Word128++instance SerialiseKey Word128 where+ serialiseKey w128 =+ RB.RawBytes $ mkPrimVector 0 16 $ P.runByteArray $ do+ ba <- P.newByteArray 16+ P.writeByteArray ba 0 $ byteSwapWord128 w128+ pure ba+ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word128" 16 len $+ byteSwapWord128 $ indexWord8ArrayAsWord128 ba off++instance SerialiseKeyOrderPreserving Word128++instance SerialiseValue Word128 where+ serialiseValue w128 =+ RB.RawBytes $ mkPrimVector 0 16 $ P.runByteArray $ do+ ba <- P.newByteArray 16+ P.writeByteArray ba 0 w128+ pure ba+ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "Word128" 16 len $+ indexWord8ArrayAsWord128 ba off++instance Arbitrary Word128 where+ arbitrary = Word128 <$> arbitrary <*> arbitrary+ shrink w128 = [ w128'+ | let i128 = toInteger w128+ , i128' <- shrink i128+ , toInteger (minBound :: Word128) <= i128'+ , toInteger (maxBound :: Word128) >= i128'+ , let w128' = fromIntegral i128'+ ]++{-# INLINE indexWord8ArrayAsWord128 #-}+indexWord8ArrayAsWord128 :: P.ByteArray -> Int -> Word128+indexWord8ArrayAsWord128 !ba !off =+ Word128 (indexWord8ArrayAsWord64 ba (off + 8))+ (indexWord8ArrayAsWord64 ba off)++newtype instance VUM.MVector s Word128 = MV_Word128 (VP.MVector s Word128)+newtype instance VU.Vector Word128 = V_Word128 (VP.Vector Word128)++deriving via VU.UnboxViaPrim Word128 instance VGM.MVector VU.MVector Word128+deriving via VU.UnboxViaPrim Word128 instance VG.Vector VU.Vector Word128++instance VUM.Unbox Word128++{-------------------------------------------------------------------------------+ NFData+-------------------------------------------------------------------------------}++deriving stock instance Generic (FS.HandleOS h)+deriving anyclass instance NFData (FS.HandleOS h)+deriving newtype instance NFData FS.BufferOffset+deriving newtype instance NFData COff++{-------------------------------------------------------------------------------+ RawBytes+-------------------------------------------------------------------------------}++instance SerialiseKey RawBytes where+ serialiseKey = id+ deserialiseKey = id++instance SerialiseValue RawBytes where+ serialiseValue = id+ deserialiseValue = id++{-------------------------------------------------------------------------------+ SerialisedKey/Value/Blob+-------------------------------------------------------------------------------}++deriving newtype instance SerialiseKey SerialisedKey++deriving newtype instance SerialiseValue SerialisedValue++deriving newtype instance SerialiseValue SerialisedBlob
@@ -0,0 +1,120 @@+{-# LANGUAGE BangPatterns #-}++module Database.LSMTree.Extras.Random (+ -- * Sampling from uniform distributions+ uniformWithoutReplacement+ , uniformWithReplacement+ , sampleUniformWithoutReplacement+ , sampleUniformWithReplacement+ , withoutReplacement+ , withReplacement+ -- * Sampling from multiple distributions+ , frequency+ -- * Shuffling+ , shuffle+ -- * Generators for specific data types+ , randomByteStringR+ ) where++import qualified Data.ByteString as BS+import Data.List (sortBy, unfoldr)+import Data.Ord (comparing)+import qualified Data.Set as Set+import qualified System.Random as R+import System.Random (StdGen, Uniform, uniform, uniformR)+import Text.Printf (printf)++{-------------------------------------------------------------------------------+ Sampling from uniform distributions+-------------------------------------------------------------------------------}++uniformWithoutReplacement :: (Ord a, Uniform a) => StdGen -> Int -> [a]+uniformWithoutReplacement rng n = withoutReplacement rng n uniform++uniformWithReplacement :: Uniform a => StdGen -> Int -> [a]+uniformWithReplacement rng n = withReplacement rng n uniform++sampleUniformWithoutReplacement :: Ord a => StdGen -> Int -> [a] -> [a]+sampleUniformWithoutReplacement rng0 n (Set.fromList -> xs0)+ | n > Set.size xs0 =+ error $+ printf "sampleUniformWithoutReplacement: n > length xs0 for n=%d, \+ \ length xs0=%d"+ n+ (Set.size xs0)+ | otherwise =+ -- Could use 'withoutReplacement', but this is more efficient.+ take n $ go xs0 rng0+ where+ go !xs !rng = x : go xs' rng'+ where+ (i, rng') = uniformR (0, Set.size xs - 1) rng+ !x = Set.elemAt i xs+ !xs' = Set.deleteAt i xs++sampleUniformWithReplacement :: Ord a => StdGen -> Int -> [a] -> [a]+sampleUniformWithReplacement rng0 n (Set.fromList -> xs) =+ withReplacement rng0 n $ \rng ->+ let (i, rng') = uniformR (0, Set.size xs - 1) rng+ in (Set.elemAt i xs, rng')++withoutReplacement :: Ord a => StdGen -> Int -> (StdGen -> (a, StdGen)) -> [a]+withoutReplacement rng0 n0 sample = take n0 $+ go Set.empty rng0+ where+ go !seen !rng+ | Set.member x seen = go seen rng'+ | otherwise = x : go (Set.insert x seen) rng'+ where+ (!x, !rng') = sample rng++withReplacement :: StdGen -> Int -> (StdGen -> (a, StdGen)) -> [a]+withReplacement rng0 n0 sample =+ take n0 $ unfoldr (Just . sample) rng0++{-------------------------------------------------------------------------------+ Sampling from multiple distributions+-------------------------------------------------------------------------------}++-- | Chooses one of the given generators, with a weighted random distribution.+-- The input list must be non-empty, weights should be non-negative, and the sum+-- of weights should be non-zero (i.e., at least one weight should be positive).+--+-- Based on the implementation in @QuickCheck@.+frequency :: [(Int, StdGen -> (a, StdGen))] -> StdGen -> (a, StdGen)+frequency xs0 g+ | any ((< 0) . fst) xs0 = error "frequency: frequencies must be non-negative"+ | tot == 0 = error "frequency: at least one frequency should be non-zero"+ | otherwise = pick i xs0+ where+ (i, g') = uniformR (1, tot) g++ tot = sum (map fst xs0)++ pick n ((k,x):xs)+ | n <= k = x g'+ | otherwise = pick (n-k) xs+ pick _ _ = error "frequency: pick used with empty list"++{-------------------------------------------------------------------------------+ Shuffling+-------------------------------------------------------------------------------}++-- | Create a random permutation of a list.+--+-- Based on the implementation in @QuickCheck@.+shuffle :: [a] -> StdGen -> [a]+shuffle xs g =+ let ns = R.randoms @Int g+ in map snd (sortBy (comparing fst) (zip ns xs))++{-------------------------------------------------------------------------------+ Generators for specific data types+-------------------------------------------------------------------------------}++-- | Generates a random bytestring. Its length is uniformly distributed within+-- the provided range.+randomByteStringR :: (Int, Int) -> StdGen -> (BS.ByteString, StdGen)+randomByteStringR range g =+ let (!l, !g') = uniformR range g+ in R.uniformByteString l g'
@@ -0,0 +1,310 @@+module Database.LSMTree.Extras.ReferenceImpl (+ -- * Page types+ Key (..),+ Value (..),+ Operation (..),+ BlobRef (..),+ PageSerialised,+ PageIntermediate,++ -- * Page size+ PageSize (..),+ pageSizeEmpty,+ pageSizeAddElem,++ -- * Encoding and decoding+ DiskPageSize(..),+ encodePage,+ decodePage,+ serialisePage,+ deserialisePage,++ -- * Overflow pages+ pageOverflowPrefixSuffixLen,+ pageDiskPages,+ pageSerialisedChunks,++ -- * Conversions to real implementation types+ toRawPage,+ toRawPageMaybeOverfull,+ toEntry,+ toEntryPrefix,+ toSerialisedKey,+ toSerialisedValue,+ toBlobSpan,++ -- * Conversions from real implementation types+ fromRawPage,+ fromEntry,+ fromEntryPrefix,+ fromSerialisedKey,+ fromSerialisedValue,+ fromSerialisedValuePrefix,+ fromRawOverflowPages,+ fromBlobSpan,++ -- * Test case types and generators+ PageContentFits(..),+ pageContentFitsInvariant,+ PageContentOrdered(..),+ pageContentOrderedInvariant,+ PageContentMaybeOverfull(..),+ PageContentSingle(..),++ -- ** Generators and shrinkers+ genPageContentFits,+ genPageContentMaybeOverfull,+ genPageContentSingle,+ genPageContentNearFull,+ genPageContentMedium,+ MinKeySize(..),+ noMinKeySize,+ orderdKeyOps,+ shrinkKeyOps,+ shrinkOrderedKeyOps,+) where++import qualified Data.ByteString as BS+import qualified Data.ByteString.Short as SBS+import Data.Maybe (fromMaybe, isJust)+import Data.Primitive.ByteArray (ByteArray (..))++import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry (Entry)+import qualified Database.LSMTree.Internal.Entry as Entry+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Serialise+import FormatPage hiding (PageContentFits, PageContentMaybeOverfull,+ PageContentSingle)++import Test.QuickCheck++-- | A test case consisting of a key\/operation sequence that is guaranteed to+-- either fit into a single 4k disk page, or be a single entry that spans a+-- one primary 4k disk page and zero or more overflow disk pages.+--+-- The keys /are not/ ordered.+--+newtype PageContentFits = PageContentFits [(Key, Operation)]+ deriving stock (Eq, Show)++-- | A test case consisting of a key\/operation sequence that is guaranteed to+-- either fit into a single 4k disk page, or be a single entry that spans a+-- one primary 4k disk page and zero or more overflow disk pages.+--+-- The keys /are/ in strictly ascending order.+--+newtype PageContentOrdered = PageContentOrdered [(Key, Operation)]+ deriving stock (Eq, Show)++-- | A test case consisting of a key\/operation sequence that is /not/+-- guaranteed to fit into a a 4k disk page.+--+-- These test cases are useful for checking the boundary conditions for what+-- can fit into a disk page.+--+-- The keys /are not/ ordered.+--+newtype PageContentMaybeOverfull = PageContentMaybeOverfull [(Key, Operation)]+ deriving stock (Eq, Show)++-- | A tests case consisting of a single key\/operation pair.+--+data PageContentSingle = PageContentSingle Key Operation+ deriving stock (Eq, Show)++instance Arbitrary PageContentFits where+ arbitrary =+ PageContentFits <$>+ genPageContentFits DiskPage4k noMinKeySize++ shrink (PageContentFits kops) =+ map PageContentFits (shrinkKeyOps kops)++pageContentFitsInvariant :: PageContentFits -> Bool+pageContentFitsInvariant (PageContentFits kops) =+ isJust (calcPageSize DiskPage4k kops)++instance Arbitrary PageContentOrdered where+ arbitrary =+ PageContentOrdered . orderdKeyOps <$>+ genPageContentFits DiskPage4k noMinKeySize++ shrink (PageContentOrdered kops) =+ map PageContentOrdered (shrinkOrderedKeyOps kops)++-- | Strictly increasing, so no duplicates.+pageContentOrderedInvariant :: PageContentOrdered -> Bool+pageContentOrderedInvariant (PageContentOrdered kops) =+ pageContentFitsInvariant (PageContentFits kops)+ && let ks = map fst kops+ in and (zipWith (<) ks (drop 1 ks))++instance Arbitrary PageContentMaybeOverfull where+ arbitrary =+ PageContentMaybeOverfull <$>+ genPageContentMaybeOverfull DiskPage4k noMinKeySize++ shrink (PageContentMaybeOverfull kops) =+ map PageContentMaybeOverfull (shrinkKeyOps kops)++instance Arbitrary PageContentSingle where+ arbitrary =+ uncurry PageContentSingle <$>+ genPageContentSingle DiskPage4k noMinKeySize++ shrink (PageContentSingle k op) =+ map (uncurry PageContentSingle) (shrink (k, op))++-------------------------------------------------------------------------------+-- Conversions between reference and real implementation types+--++-- | Convert from 'PageContentFits' (from the reference implementation)+-- to 'RawPage' (from the real implementation).+--+toRawPage :: PageContentFits -> (RawPage, [RawOverflowPage])+toRawPage (PageContentFits kops) =+ fromMaybe overfull (toRawPageMaybeOverfull (PageContentMaybeOverfull kops))+ where+ overfull =+ error $ "toRawPage: encountered overfull page, but PageContentFits is "+ ++ "supposed to be guaranteed to fit, i.e. not to be overfull."++toRawPageMaybeOverfull :: PageContentMaybeOverfull+ -> Maybe (RawPage, [RawOverflowPage])+toRawPageMaybeOverfull (PageContentMaybeOverfull kops) = do+ serialised <- serialisePage <$> encodePage DiskPage4k kops+ (page : overflowPages) <- pure (pageSerialisedChunks DiskPage4k serialised)+ pure (makeRawPageBS page, map makeRawOverflowPageBS overflowPages)++makeRawPageBS :: BS.ByteString -> RawPage+makeRawPageBS bs =+ case SBS.toShort bs of+ SBS.SBS ba -> makeRawPage (ByteArray ba) 0++makeRawOverflowPageBS :: BS.ByteString -> RawOverflowPage+makeRawOverflowPageBS bs =+ case SBS.toShort bs of+ SBS.SBS ba -> makeRawOverflowPage (ByteArray ba) 0 (BS.length bs)++fromRawPage :: (RawPage, [RawOverflowPage]) -> PageContentFits+fromRawPage (page, overflowPages)+ | rawPageNumKeys page == 1+ = PageContentFits . (:[]) $+ case rawPageIndex page 0 of+ IndexEntry k e ->+ (fromSerialisedKey k, fromEntry e)++ IndexEntryOverflow k e suffixlen ->+ ( fromSerialisedKey k+ , fromEntryPrefix e (fromIntegral suffixlen) overflowPages+ )++ IndexNotPresent -> error "fromRawPage: 'rawPageIndex page 0' fails"++ | otherwise+ = PageContentFits+ [ case rawPageIndex page (fromIntegral i) of+ IndexEntry k e -> (fromSerialisedKey k, fromEntry e)+ _ -> error "fromRawPage: 'rawPageIndex page i' fails"+ | i <- [0 .. fromIntegral (rawPageNumKeys page) - 1 :: Int] ]++toEntry :: Operation -> Entry SerialisedValue BlobSpan+toEntry op =+ case op of+ Insert v Nothing ->+ Entry.Insert (toSerialisedValue v)++ Insert v (Just b) ->+ Entry.InsertWithBlob (toSerialisedValue v) (toBlobSpan b)++ Mupsert v ->+ Entry.Upsert (toSerialisedValue v)++ Delete ->+ Entry.Delete++toEntryPrefix :: Operation -> Int -> Entry SerialisedValue BlobSpan+toEntryPrefix op prefixlen =+ case op of+ Insert v Nothing ->+ Entry.Insert (toSerialisedValue (takeValue prefixlen v))++ Insert v (Just b) ->+ Entry.InsertWithBlob (toSerialisedValue (takeValue prefixlen v))+ (toBlobSpan b)++ Mupsert v ->+ Entry.Upsert (toSerialisedValue (takeValue prefixlen v))++ Delete ->+ Entry.Delete+ where+ takeValue n (Value v) = Value (BS.take n v)++fromEntry :: Entry SerialisedValue BlobSpan -> Operation+fromEntry e =+ case e of+ Entry.Insert v ->+ Insert (fromSerialisedValue v) Nothing++ Entry.InsertWithBlob v b ->+ Insert (fromSerialisedValue v) (Just (fromBlobSpan b))++ Entry.Upsert v ->+ Mupsert (fromSerialisedValue v)++ Entry.Delete ->+ Delete++fromEntryPrefix :: Entry SerialisedValue BlobSpan+ -> Int+ -> [RawOverflowPage]+ -> Operation+fromEntryPrefix e suffix overflow =+ case e of+ Entry.Insert v ->+ Insert (fromSerialisedValuePrefix v suffix overflow) Nothing++ Entry.InsertWithBlob v b ->+ Insert (fromSerialisedValuePrefix v suffix overflow)+ (Just (fromBlobSpan b))++ Entry.Upsert v ->+ Mupsert (fromSerialisedValuePrefix v suffix overflow)++ Entry.Delete ->+ Delete++toSerialisedKey :: Key -> SerialisedKey+toSerialisedKey (Key k) = SerialisedKey (RB.fromByteString k)++toSerialisedValue :: Value -> SerialisedValue+toSerialisedValue (Value v) = SerialisedValue (RB.fromByteString v)++fromSerialisedKey :: SerialisedKey -> Key+fromSerialisedKey (SerialisedKey k) = Key (RB.toByteString k)++fromSerialisedValue :: SerialisedValue -> Value+fromSerialisedValue (SerialisedValue v) = Value (RB.toByteString v)++fromSerialisedValuePrefix :: SerialisedValue -> Int -> [RawOverflowPage]+ -> Value+fromSerialisedValuePrefix (SerialisedValue v) suffixlen overflowPages =+ Value $ RB.toByteString v+ <> BS.take suffixlen (fromRawOverflowPages overflowPages)++fromRawOverflowPages :: [RawOverflowPage] -> BS.ByteString+fromRawOverflowPages =+ RB.toByteString+ . mconcat+ . map rawOverflowPageRawBytes++toBlobSpan :: BlobRef -> BlobSpan+toBlobSpan (BlobRef x y) = BlobSpan x y++fromBlobSpan :: BlobSpan -> BlobRef+fromBlobSpan (BlobSpan x y) = BlobRef x y
@@ -0,0 +1,350 @@+-- | Utilities for generating collections of key\/value pairs, and creating runs+-- from them. Tests and benchmarks should preferably use these utilities instead+-- of (re-)defining their own.+module Database.LSMTree.Extras.RunData (+ -- * Create runs+ withRun+ , withRunAt+ , withRuns+ , unsafeCreateRun+ , unsafeCreateRunAt+ , simplePath+ , simplePaths+ -- * Serialise write buffers+ , withRunDataAsWriteBuffer+ , withSerialisedWriteBuffer+ -- * RunData+ , RunData (..)+ , mapRunData+ , SerialisedRunData+ , serialiseRunData+ -- * NonEmptyRunData+ , NonEmptyRunData (..)+ , nonEmptyRunData+ , toRunData+ , mapNonEmptyRunData+ , SerialisedNonEmptyRunData+ -- * QuickCheck+ , labelRunData+ , labelNonEmptyRunData+ , genRunData+ , shrinkRunData+ , genNonEmptyRunData+ , shrinkNonEmptyRunData+ , liftArbitrary2Map+ , liftShrink2Map+ ) where++import Control.Exception (bracket, bracket_)+import Control.RefCount+import Data.Bifoldable (Bifoldable (bifoldMap))+import Data.Bifunctor+import Data.Foldable (for_)+import Data.Map.NonEmpty (NEMap)+import qualified Data.Map.NonEmpty as NEMap+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import qualified Data.Vector as V+import Database.LSMTree.Extras (showPowersOf10)+import Database.LSMTree.Extras.Generators ()+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.MergeSchedule (addWriteBufferEntries)+import Database.LSMTree.Internal.Paths+import qualified Database.LSMTree.Internal.Paths as Paths+import Database.LSMTree.Internal.Run (Run, RunParams (..))+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunAcc (entryWouldFitInPage)+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import Database.LSMTree.Internal.WriteBufferWriter (writeWriteBuffer)+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API (HasBlockIO)+import Test.QuickCheck+++{-------------------------------------------------------------------------------+ Create runs+-------------------------------------------------------------------------------}++-- | Create a temporary 'Run' using 'unsafeCreateRun'.+withRun ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> SerialisedRunData+ -> (Ref (Run IO h) -> IO a)+ -> IO a+withRun hfs hbio salt runParams path counter rd = do+ bracket+ (unsafeCreateRun hfs hbio salt runParams path counter rd)+ releaseRef++-- | Create a temporary 'Run' using 'unsafeCreateRunAt'.+withRunAt ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> SerialisedRunData+ -> (Ref (Run IO h) -> IO a)+ -> IO a+withRunAt hfs hbio salt runParams path rd = do+ bracket+ (unsafeCreateRunAt hfs hbio salt runParams path rd)+ releaseRef++{-# INLINABLE withRuns #-}+-- | Create temporary 'Run's using 'unsafeCreateRun'.+withRuns ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> [SerialisedRunData]+ -> ([Ref (Run IO h)] -> IO a)+ -> IO a+withRuns hfs hbio salt runParams path counter = go+ where+ go [] act = act []+ go (rd:rds) act =+ withRun hfs hbio salt runParams path counter rd $ \r ->+ go rds $ \rs ->+ act (r:rs)++-- | Like 'unsafeCreateRunAt', but uses a 'UniqCounter' to determine+-- the 'RunFsPaths', at a base file path.+unsafeCreateRun ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> FS.FsPath+ -> UniqCounter IO+ -> SerialisedRunData+ -> IO (Ref (Run IO h))+unsafeCreateRun fs hbio salt runParams path counter rd = do+ n <- incrUniqCounter counter+ let fsPaths = RunFsPaths path (uniqueToRunNumber n)+ unsafeCreateRunAt fs hbio salt runParams fsPaths rd++-- | Flush serialised run data to disk as if it were a write buffer.+--+-- This might leak resources if not run with asynchronous exceptions masked.+-- Use helper functions like 'withRun' or 'withRuns' instead.+--+-- Use of this function should be paired with a 'releaseRef'.+unsafeCreateRunAt ::+ HasFS IO h+ -> HasBlockIO IO h+ -> Bloom.Salt+ -> RunParams+ -> RunFsPaths+ -> SerialisedRunData+ -> IO (Ref (Run IO h))+unsafeCreateRunAt fs hbio salt runParams fsPaths (RunData m) = do+ -- the WBB file path doesn't have to be at a specific place relative to+ -- the run we want to create, but fsPaths should already point to a unique+ -- location, so we just append something to not conflict with that.+ let blobpath = FS.addExtension (runBlobPath fsPaths) ".wb"+ bracket (WBB.new fs blobpath) releaseRef $ \wbblobs -> do+ wb <- WB.fromMap <$> traverse (traverse (WBB.addBlob fs wbblobs)) m+ Run.fromWriteBuffer fs hbio salt runParams fsPaths wb wbblobs++-- | Create a 'RunFsPaths' using an empty 'FsPath'. The empty path corresponds+-- to the "root" or "mount point" of a 'HasFS' instance.+simplePath :: Int -> RunFsPaths+simplePath n = RunFsPaths (FS.mkFsPath []) (RunNumber n)++-- | Like 'simplePath', but for a list.+simplePaths :: [Int] -> [RunFsPaths]+simplePaths ns = fmap simplePath ns++{-------------------------------------------------------------------------------+ Serialise write buffers+-------------------------------------------------------------------------------}++-- | Use 'SerialisedRunData' to 'WriteBuffer' and 'WriteBufferBlobs'.+withRunDataAsWriteBuffer ::+ FS.HasFS IO h+ -> ResolveSerialisedValue+ -> WriteBufferFsPaths+ -> SerialisedRunData+ -> (WB.WriteBuffer -> Ref (WBB.WriteBufferBlobs IO h) -> IO a)+ -> IO a+withRunDataAsWriteBuffer hfs f fsPaths rd action = do+ let es = V.fromList . Map.toList $ unRunData rd+ let maxn = NumEntries $ V.length es+ let wbbPath = Paths.writeBufferBlobPath fsPaths+ bracket (WBB.new hfs wbbPath) releaseRef $ \wbb -> do+ (wb, _) <- addWriteBufferEntries hfs f wbb maxn WB.empty es+ action wb wbb++-- | Serialise a 'WriteBuffer' and 'WriteBufferBlobs' to disk and perform an+-- 'IO' action.+withSerialisedWriteBuffer ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> WriteBufferFsPaths+ -> WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs IO h)+ -> IO a+ -> IO a+withSerialisedWriteBuffer hfs hbio wbPaths wb wbb =+ bracket_ (writeWriteBuffer hfs hbio wbPaths wb wbb) $ do+ for_ [ Paths.writeBufferKOpsPath wbPaths+ , Paths.writeBufferBlobPath wbPaths+ , Paths.writeBufferChecksumsPath wbPaths+ ] $ FS.removeFile hfs++{-------------------------------------------------------------------------------+ RunData+-------------------------------------------------------------------------------}++-- | A collection of arbitrary key\/value pairs that are suitable for creating+-- 'Run's.+--+-- Note: 'b ~ Void' should rule out blobs.+newtype RunData k v b = RunData {+ unRunData :: Map k (Entry v b)+ }+ deriving stock (Eq, Show)++mapRunData ::+ Ord k'+ => (k -> k') -> (v -> v') -> (b -> b')+ -> RunData k v b -> RunData k' v' b'+mapRunData f g h = RunData . Map.mapKeys f . Map.map (bimap g h) . unRunData++type SerialisedRunData = RunData SerialisedKey SerialisedValue SerialisedBlob++serialiseRunData ::+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => RunData k v b -> SerialisedRunData+serialiseRunData = mapRunData serialiseKey serialiseValue serialiseBlob++{-------------------------------------------------------------------------------+ NonEmptyRunData+-------------------------------------------------------------------------------}++-- | A collection of arbitrary key\/value pairs that are suitable for creating+-- 'Run's.+--+-- Note: 'b ~ Void' should rule out blobs.+newtype NonEmptyRunData k v b =+ NonEmptyRunData { unNonEmptyRunData :: NEMap k (Entry v b) }+ deriving stock (Eq, Show)++nonEmptyRunData :: RunData k v b -> Maybe (NonEmptyRunData k v b)+nonEmptyRunData (RunData m) = NonEmptyRunData <$> NEMap.nonEmptyMap m++toRunData :: NonEmptyRunData k v b -> RunData k v b+toRunData (NonEmptyRunData m) = RunData (NEMap.toMap m)++mapNonEmptyRunData ::+ Ord k'+ => (k -> k') -> (v -> v') -> (b -> b')+ -> NonEmptyRunData k v b -> NonEmptyRunData k' v' b'+mapNonEmptyRunData f g h =+ NonEmptyRunData . NEMap.mapKeys f . NEMap.map (bimap g h) . unNonEmptyRunData++type SerialisedNonEmptyRunData =+ NonEmptyRunData SerialisedKey SerialisedValue SerialisedBlob++{-------------------------------------------------------------------------------+ QuickCheck+-------------------------------------------------------------------------------}++labelRunData :: SerialisedRunData -> Property -> Property+labelRunData (RunData m) =+ tabulate "value size" valSizes+ . tabulate "run length" [runLength]+ -- We use tabulate here, not label. Otherwise, if multiple RunDatas get+ -- labelled in a test case, each leads to a separate entry in the+ -- displayed statistics, which isn't that useful.+ . tabulate "run has large k/ops" [note]+ where+ kops = Map.toList m+ valSizes = map (showPowersOf10 . sizeofValue) vals+ runLength = showPowersOf10 (Map.size m)+ vals = concatMap (bifoldMap pure mempty . snd) kops+ note+ | any (not . uncurry entryWouldFitInPage) kops = "has large k/op"+ | otherwise = "no large k/op"++labelNonEmptyRunData :: SerialisedNonEmptyRunData -> Property -> Property+labelNonEmptyRunData = labelRunData . toRunData++instance ( Ord k, Arbitrary k, Arbitrary v, Arbitrary b+ ) => Arbitrary (RunData k v b) where+ arbitrary = genRunData arbitrary arbitrary arbitrary+ shrink = shrinkRunData shrink shrink shrink++instance ( Ord k, Arbitrary k, Arbitrary v, Arbitrary b+ ) => Arbitrary (NonEmptyRunData k v b) where+ arbitrary = genNonEmptyRunData arbitrary arbitrary arbitrary+ shrink = shrinkNonEmptyRunData shrink shrink shrink++genRunData :: Ord k => Gen k -> Gen v -> Gen b -> Gen (RunData k v b)+genRunData genKey genVal genBlob =+ RunData <$> liftArbitrary2Map genKey (liftArbitrary2 genVal genBlob)++shrinkRunData ::+ Ord k+ => (k -> [k])+ -> (v -> [v])+ -> (b -> [b])+ -> RunData k v b+ -> [RunData k v b]+shrinkRunData shrinkKey shrinkVal shrinkBlob =+ fmap RunData+ . liftShrink2Map shrinkKey (liftShrink2 shrinkVal shrinkBlob)+ . unRunData++-- | We cannot implement 'Arbitrary2' since we have constraints on @k@.+liftArbitrary2Map :: Ord k => Gen k -> Gen v -> Gen (Map k v)+liftArbitrary2Map genk genv = Map.fromList <$>+ liftArbitrary (liftArbitrary2 genk genv)++-- | We cannot implement 'Arbitrary2' since we have constraints @k@.+liftShrink2Map :: Ord k => (k -> [k]) -> (v -> [v]) -> Map k v -> [Map k v]+liftShrink2Map shrinkk shrinkv m = Map.fromList <$>+ liftShrink (liftShrink2 shrinkk shrinkv) (Map.toList m)++genNonEmptyRunData ::+ Ord k => Gen k -> Gen v -> Gen b -> Gen (NonEmptyRunData k v b)+genNonEmptyRunData genKey genVal genBlob = NonEmptyRunData <$>+ liftArbitrary2NEMap genKey (liftArbitrary2 genVal genBlob)++shrinkNonEmptyRunData ::+ Ord k+ => (k -> [k])+ -> (v -> [v])+ -> (b -> [b])+ -> NonEmptyRunData k v b+ -> [NonEmptyRunData k v b]+shrinkNonEmptyRunData shrinkKey shrinkVal shrinkBlob =+ fmap NonEmptyRunData+ . liftShrink2NEMap shrinkKey (liftShrink2 shrinkVal shrinkBlob)+ . unNonEmptyRunData++-- | We cannot implement 'Arbitrary2' since we have constraints on @k@.+liftArbitrary2NEMap :: Ord k => Gen k -> Gen v -> Gen (NEMap k v)+liftArbitrary2NEMap genk genv = NEMap.fromList <$>+ liftArbitrary (liftArbitrary2 genk genv)++-- | We cannot implement 'Arbitrary2' since we have constraints @k@.+liftShrink2NEMap :: Ord k => (k -> [k]) -> (v -> [v]) -> NEMap k v -> [NEMap k v]+liftShrink2NEMap shrinkk shrinkv m = NEMap.fromList <$>+ liftShrink (liftShrink2 shrinkk shrinkv) (NEMap.toList m)
@@ -0,0 +1,150 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE TypeFamilies #-}++module Database.LSMTree.Extras.UTxO (+ UTxOKey (..)+ , UTxOValue (..)+ , zeroUTxOValue+ , UTxOBlob (..)+ ) where++import Control.DeepSeq+import Data.Bits+import qualified Data.ByteString as BS+import qualified Data.Primitive as P+import qualified Data.Vector.Generic as VG+import qualified Data.Vector.Generic.Mutable as VGM+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.WideWord.Word128+import Data.WideWord.Word256+import Data.Word+import Database.LSMTree.Extras.Orphans+import Database.LSMTree.Internal.Primitive+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise ()+import Database.LSMTree.Internal.Serialise.Class as Class+import Database.LSMTree.Internal.Vector+import GHC.Generics+import System.Random+import Test.QuickCheck hiding ((.&.))++{-------------------------------------------------------------------------------+ UTxO keys+-------------------------------------------------------------------------------}++-- | A model of a UTxO key (34 bytes) after @TxIn@: a 256-bit hash, 16-bit identifier+data UTxOKey = UTxOKey {+ txId :: !Word256 -- no unpack, since the @TxId@ field doesn't have it+ , txIx :: {-# UNPACK #-} !Word16+ }+ deriving stock (Show, Eq, Ord, Generic)+ deriving anyclass (Uniform, NFData)++-- TODO: reduce number of allocations, optimise (by using unsafe functions)+instance SerialiseKey UTxOKey where+ serialiseKey UTxOKey{txId, txIx} =+ RB.RawBytes $ mkPrimVector 0 34 $ P.runByteArray $ do+ ba <- P.newByteArray 34+ let !cut = fromIntegral ((txId .&. (fromIntegral (0xffff :: Word16) `unsafeShiftL` 192)) `unsafeShiftR` 192)+ P.writeByteArray ba 0 $ byteSwapWord256 txId+ P.writeByteArray ba 3 $ byteSwap16 txIx+ P.writeByteArray ba 16 $ byteSwap16 cut+ pure ba+ deserialiseKey (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "UTxOKey" 34 len $+ let !cut = byteSwap16 $ indexWord8ArrayAsWord16 ba (off + 32)+ !txIx = byteSwap16 $ indexWord8ArrayAsWord16 ba (off + 6)+ !txId_ = byteSwapWord256 $ indexWord8ArrayAsWord256 ba off+ !txId = (txId_ .&. (complement (fromIntegral (0xffff :: Word16) `unsafeShiftL` 192)))+ .|. (fromIntegral cut `unsafeShiftL` 192)+ in UTxOKey{txId, txIx}++instance Arbitrary UTxOKey where+ arbitrary = UTxOKey <$> arbitrary <*> arbitrary+ shrink (UTxOKey a b) = [ UTxOKey a' b' | (a', b') <- shrink (a, b) ]++newtype instance VUM.MVector s UTxOKey = MV_UTxOKey (VU.MVector s (Word256, Word16))+newtype instance VU.Vector UTxOKey = V_UTxOKey (VU.Vector (Word256, Word16))++instance VU.IsoUnbox UTxOKey (Word256, Word16) where+ toURepr (UTxOKey a b) = (a, b)+ fromURepr (a, b) = UTxOKey a b+ {-# INLINE toURepr #-}+ {-# INLINE fromURepr #-}++deriving via VU.As UTxOKey (Word256, Word16) instance VGM.MVector VU.MVector UTxOKey+deriving via VU.As UTxOKey (Word256, Word16) instance VG.Vector VU.Vector UTxOKey++instance VUM.Unbox UTxOKey++{-------------------------------------------------------------------------------+ UTxO values+-------------------------------------------------------------------------------}++-- | A model of a UTxO value (60 bytes)+data UTxOValue = UTxOValue {+ utxoValue256 :: !Word256+ , utxoValue128 :: !Word128+ , utxoValue64 :: !Word64+ , utxoValue32 :: !Word32+ }+ deriving stock (Show, Eq, Ord, Generic)+ deriving anyclass (Uniform, NFData)++instance SerialiseValue UTxOValue where+ serialiseValue (UTxOValue {utxoValue256, utxoValue128, utxoValue64, utxoValue32}) =+ RB.RawBytes $ mkPrimVector 0 60 $ P.runByteArray $ do+ ba <- P.newByteArray 60+ P.writeByteArray ba 0 utxoValue256+ P.writeByteArray ba 2 utxoValue128+ P.writeByteArray ba 6 utxoValue64+ P.writeByteArray ba 14 utxoValue32+ pure ba+ deserialiseValue (RawBytes (VP.Vector off len ba)) =+ requireBytesExactly "UTxOValue" 60 len $+ UTxOValue (indexWord8ArrayAsWord256 ba off)+ (indexWord8ArrayAsWord128 ba (off + 32))+ (indexWord8ArrayAsWord64 ba (off + 48))+ (indexWord8ArrayAsWord32 ba (off + 56))++instance Arbitrary UTxOValue where+ arbitrary = UTxOValue <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary+ shrink (UTxOValue a b c d) = [ UTxOValue a' b' c' d'+ | (a', b', c', d') <- shrink (a, b, c, d) ]++newtype instance VUM.MVector s UTxOValue = MV_UTxOValue (VU.MVector s (Word256, Word128, Word64, Word32))+newtype instance VU.Vector UTxOValue = V_UTxOValue (VU.Vector (Word256, Word128, Word64, Word32))++instance VU.IsoUnbox UTxOValue (Word256, Word128, Word64, Word32) where+ toURepr (UTxOValue a b c d) = (a, b, c, d)+ fromURepr (a, b, c, d) = UTxOValue a b c d+ {-# INLINE toURepr #-}+ {-# INLINE fromURepr #-}++deriving via VU.As UTxOValue (Word256, Word128, Word64, Word32) instance VGM.MVector VU.MVector UTxOValue+deriving via VU.As UTxOValue (Word256, Word128, Word64, Word32) instance VG.Vector VU.Vector UTxOValue++instance VUM.Unbox UTxOValue++zeroUTxOValue :: UTxOValue+zeroUTxOValue = UTxOValue 0 0 0 0++{-------------------------------------------------------------------------------+ UTxO blobs+-------------------------------------------------------------------------------}++-- | A blob of arbitrary size+newtype UTxOBlob = UTxOBlob BS.ByteString+ deriving stock (Show, Eq, Ord, Generic)+ deriving anyclass NFData++instance SerialiseValue UTxOBlob where+ serialiseValue (UTxOBlob bs) = Class.serialiseValue bs+ deserialiseValue = error "deserialiseValue: UTxOBlob"
@@ -0,0 +1,180 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_GHC -fexpose-all-unfoldings #-}++-- | Mutable heap for k-merge algorithm.+--+-- This data-structure represents a min-heap with the root node *removed*.+-- (internally the filling of root value and sifting down is delayed).+--+-- Also there isn't *insert* operation, i.e. the heap can only shrink.+-- Other heap usual heap operations are *create-heap*, *extract-min* and *replace*.+-- However, as the 'MutableHeap' always represents a heap with its root (minimum value)+-- extracted, *extract-min* is "fused" to other operations.+module KMerge.Heap (+ MutableHeap (..),+ newMutableHeap,+ replaceRoot,+ extract,+) where++import Control.Monad (when)+import Control.Monad.Primitive (PrimMonad (PrimState), RealWorld)+import qualified Control.Monad.ST as Lazy+import qualified Control.Monad.ST as Strict+import Data.Bits (unsafeShiftL, unsafeShiftR)+import Data.Foldable.WithIndex (ifor_)+import Data.List.NonEmpty (NonEmpty (..))+import Data.Primitive (SmallMutableArray, newSmallArray,+ readSmallArray, writeSmallArray)+import Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,+ writePrimVar)+import Unsafe.Coerce (unsafeCoerce)++-- | Mutable heap for k-merge algorithm.+data MutableHeap s a = MH+ !(PrimVar s Int) -- ^ element count, size+ !(SmallMutableArray s a)++-- | Placeholder value used to fill the internal array.+placeholder :: a+placeholder = unsafeCoerce ()++-- | Create new heap, and immediately extract its minimum value.+newMutableHeap :: forall a m. (PrimMonad m, Ord a) => NonEmpty a -> m (MutableHeap (PrimState m) a, a)+newMutableHeap xs = do+ let !size = length xs++ arr <- newSmallArray size placeholder+ ifor_ xs $ \idx x -> do+ writeSmallArray arr idx x+ siftUp arr x idx++ sizeRef <- newPrimVar size++ -- This indexing is safe!+ -- Due to the required NonEmpty input type, there must be at least one element to read.+ x <- readSmallArray arr 0+ writeSmallArray arr 0 placeholder+ pure $! (MH sizeRef arr, x)++-- | Replace the minimum-value, and immediately extract the new minimum value.+replaceRoot :: forall a m. (PrimMonad m, Ord a) => MutableHeap (PrimState m) a -> a -> m a+replaceRoot (MH sizeRef arr) val = do+ size <- readPrimVar sizeRef+ if size <= 1+ then pure val+ else do+ writeSmallArray arr 0 val+ siftDown arr size val 0+ readSmallArray arr 0++{-# SPECIALISE replaceRoot :: forall a. Ord a => MutableHeap RealWorld a -> a -> IO a #-}+{-# SPECIALISE replaceRoot :: forall a s. Ord a => MutableHeap s a -> a -> Strict.ST s a #-}+{-# SPECIALISE replaceRoot :: forall a s. Ord a => MutableHeap s a -> a -> Lazy.ST s a #-}++-- | Extract the next minimum value.+extract :: forall a m. (PrimMonad m, Ord a) => MutableHeap (PrimState m) a -> m (Maybe a)+extract (MH sizeRef arr) = do+ size <- readPrimVar sizeRef+ if size <= 1+ then pure Nothing+ else do+ writePrimVar sizeRef $! size - 1+ val <- readSmallArray arr (size - 1)+ writeSmallArray arr 0 val+ siftDown arr size val 0+ x <- readSmallArray arr 0+ writeSmallArray arr (size - 1) placeholder+ pure $! Just x++{-# SPECIALISE extract :: forall a. Ord a => MutableHeap RealWorld a -> IO (Maybe a) #-}+{-# SPECIALISE extract :: forall a s. Ord a => MutableHeap s a -> Strict.ST s (Maybe a) #-}+{-# SPECIALISE extract :: forall a s. Ord a => MutableHeap s a -> Lazy.ST s (Maybe a) #-}++{-------------------------------------------------------------------------------+ Internal operations+-------------------------------------------------------------------------------}++siftUp :: forall a m. (PrimMonad m, Ord a) => SmallMutableArray (PrimState m) a -> a -> Int -> m ()+siftUp !arr !x = loop where+ loop !idx+ | idx <= 0+ = pure ()++ | otherwise+ = do+ let !parent = halfOf (idx - 1)+ p <- readSmallArray arr parent+ when (x < p) $ do+ writeSmallArray arr parent x+ writeSmallArray arr idx p+ loop parent++{-# SPECIALISE siftUp :: forall a. Ord a => SmallMutableArray RealWorld a -> a -> Int -> IO () #-}+{-# SPECIALISE siftUp :: forall a s. Ord a => SmallMutableArray s a -> a -> Int -> Strict.ST s () #-}+{-# SPECIALISE siftUp :: forall a s. Ord a => SmallMutableArray s a -> a -> Int -> Lazy.ST s () #-}++siftDown :: forall a m. (PrimMonad m, Ord a) => SmallMutableArray (PrimState m) a -> Int -> a -> Int -> m ()+siftDown !arr !size !x = loop where+ loop !idx+ | rgt < size+ = do+ l <- readSmallArray arr lft+ r <- readSmallArray arr rgt++ if x <= l+ then do+ if x <= r+ then pure ()+ else do+ -- r < x <= l; swap x and r+ writeSmallArray arr rgt x+ writeSmallArray arr idx r+ loop rgt+ else do+ if l <= r+ then do+ -- l < x, l <= r; swap x and l+ writeSmallArray arr idx l+ writeSmallArray arr lft x+ loop lft+ else do+ -- r < l <= x; swap x and r+ writeSmallArray arr rgt x+ writeSmallArray arr idx r+ loop rgt++ -- there's only left value+ | lft < size+ = do+ l <- readSmallArray arr lft+ if x <= l+ then pure ()+ else do+ writeSmallArray arr idx l+ writeSmallArray arr lft x+ -- there is no need to loop further, lft was the last value.++ | otherwise+ = pure ()+ where+ !lft = doubleOf idx + 1+ !rgt = doubleOf idx + 2++{-# SPECIALISE siftDown :: forall a. Ord a => SmallMutableArray RealWorld a -> Int -> a -> Int -> IO () #-}+{-# SPECIALISE siftDown :: forall a s. Ord a => SmallMutableArray s a -> Int -> a -> Int -> Strict.ST s () #-}+{-# SPECIALISE siftDown :: forall a s. Ord a => SmallMutableArray s a -> Int -> a -> Int -> Lazy.ST s () #-}++{-------------------------------------------------------------------------------+ Helpers+-------------------------------------------------------------------------------}++halfOf :: Int -> Int+halfOf i = unsafeShiftR i 1+{-# INLINE halfOf #-}++doubleOf :: Int -> Int+doubleOf i = unsafeShiftL i 1+{-# INLINE doubleOf #-}
@@ -0,0 +1,206 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_GHC -fexpose-all-unfoldings #-}++module KMerge.LoserTree (+ MutableLoserTree,+ newLoserTree,+ replace,+ remove,+) where++import Control.Monad.Primitive (PrimMonad (PrimState), RealWorld)+import qualified Control.Monad.ST as Lazy+import qualified Control.Monad.ST as Strict+import Data.Bits (unsafeShiftR)+import Data.List.NonEmpty (NonEmpty (..))+import Data.Primitive (MutablePrimArray, SmallMutableArray,+ newPrimArray, newSmallArray, readPrimArray, readSmallArray,+ setPrimArray, writePrimArray, writeSmallArray)+import Data.Primitive.PrimVar (PrimVar, newPrimVar, readPrimVar,+ writePrimVar)+import Unsafe.Coerce (unsafeCoerce)++-- | Mutable Loser Tree.+data MutableLoserTree s a = MLT+ !(PrimVar s Int) -- ^ element count, i.e. size.+ !(PrimVar s Int) -- ^ index of the hole (i.e. winner's initial index)+ !(MutablePrimArray s Int) -- ^ indices, we store the index of first match. -1 if there is no match.+ !(SmallMutableArray s a) -- ^ values++placeholder :: a+placeholder = unsafeCoerce ()++-- | Create new 'MutableLoserTree'.+--+-- The second half of a pair is the winner value (only losers are stored in the tree).+--+newLoserTree :: forall a m. (PrimMonad m, Ord a) => NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)+newLoserTree (x0 :| xs0) = do+ -- allocate array, we need one less than there are elements.+ -- one of the elements will be the winner.+ ids <- newPrimArray len+ arr <- newSmallArray len placeholder+ case xs0 of+ [] -> do+ sizeRef <- newPrimVar 0+ holeRef <- newPrimVar 0+ pure $! (MLT sizeRef holeRef ids arr, x0)+ _ -> do+ setPrimArray ids 0 len (-1)+ loop ids arr len $ x0 :| xs0+ where+ !len = length xs0++ loop :: MutablePrimArray (PrimState m) Int -> SmallMutableArray (PrimState m) a -> Int -> NonEmpty a -> m (MutableLoserTree (PrimState m) a, a)+ loop ids arr idx (x :| xs) = do+ sift ids arr (parentOf idx) (parentOf idx) x idx xs++ sift :: MutablePrimArray (PrimState m) Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> a -> Int -> [a] -> m (MutableLoserTree (PrimState m) a, a)+ sift !ids !arr !idxX !j !x !idx0 xs = do+ !idxY <- readPrimArray ids j+ y <- readSmallArray arr j+ -- NOTE: The length of xs is equal to number of uninitialised entries+ -- from this we can deduce that an entry at j is uninitialised implies+ -- that xs cannot be empty.+ -- We check this invariant here and throw an exception+ -- with a descriptive error message if it is violated.+ if idxY < 0+ then case xs of+ [] -> error $ unlines+ [ "Error in KMerge.LoserTree.newLoserTree"+ , unwords [ "Invariant violated at entry # j =", show j, "with xs = [] and idxY =", show idxY ]+ ]+ e:es -> do+ writePrimArray ids j idxX+ writeSmallArray arr j x+ loop ids arr (idx0 + 1) $ e :| es+ else+ if j <= 0+ then do+ if x <= y+ then do+ sizeRef <- newPrimVar len+ holeRef <- newPrimVar idxX+ pure (MLT sizeRef holeRef ids arr, x)+ else do+ writePrimArray ids j idxX+ writeSmallArray arr j x+ sizeRef <- newPrimVar len+ holeRef <- newPrimVar idxY+ pure (MLT sizeRef holeRef ids arr, y)+ else do+ if x < y+ then do+ sift ids arr idxX (parentOf j) x idx0 xs+ else do+ writePrimArray ids j idxX+ writeSmallArray arr j x+ sift ids arr idxY (parentOf j) y idx0 xs++{-# SPECIALISE newLoserTree :: forall a. Ord a => NonEmpty a -> IO (MutableLoserTree RealWorld a, a) #-}+{-# SPECIALISE newLoserTree :: forall a s. Ord a => NonEmpty a -> Strict.ST s (MutableLoserTree s a, a) #-}+{-# SPECIALISE newLoserTree :: forall a s. Ord a => NonEmpty a -> Lazy.ST s (MutableLoserTree s a, a) #-}++{-------------------------------------------------------------------------------+ Updates+-------------------------------------------------------------------------------}++{-# SPECIALISE replace :: forall a. Ord a => MutableLoserTree RealWorld a -> a -> IO a #-}+{-# SPECIALISE replace :: forall a s. Ord a => MutableLoserTree s a -> a -> Strict.ST s a #-}+{-# SPECIALISE replace :: forall a s. Ord a => MutableLoserTree s a -> a -> Lazy.ST s a #-}++{-# SPECIALISE remove :: forall a. Ord a => MutableLoserTree RealWorld a -> IO (Maybe a) #-}+{-# SPECIALISE remove :: forall a s. Ord a => MutableLoserTree s a -> Strict.ST s (Maybe a) #-}+{-# SPECIALISE remove :: forall a s. Ord a => MutableLoserTree s a -> Lazy.ST s (Maybe a) #-}++-- | Don't fill the winner "hole". Return a next winner of (smaller) tournament.+remove :: forall a m. (PrimMonad m, Ord a) => MutableLoserTree (PrimState m) a -> m (Maybe a)+remove (MLT sizeRef holeRef ids arr) = do+ size <- readPrimVar sizeRef+ if size <= 0+ then pure Nothing+ else do+ writePrimVar sizeRef (size - 1)+ hole <- readPrimVar holeRef+ siftEmpty hole+ where+ siftEmpty :: Int -> m (Maybe a)+ siftEmpty !j = do+ !idxY <- readPrimArray ids j+ y <- readSmallArray arr j+ if j <= 0+ then if idxY < 0+ then pure Nothing+ else do+ writePrimArray ids j (-1)+ writeSmallArray arr j placeholder+ writePrimVar holeRef idxY+ pure (Just y)+ else if idxY < 0+ then+ siftEmpty (parentOf j)+ else do+ writePrimArray ids j (-1)+ writeSmallArray arr j placeholder+ Just <$> siftUp ids arr holeRef (parentOf j) idxY y++-- | Fill the winner "hole" with a new element. Return a new tournament winner.+replace :: forall a m. (PrimMonad m, Ord a) => MutableLoserTree (PrimState m) a -> a -> m a+replace (MLT sizeRef holeRef ids arr) val = do+ size <- readPrimVar sizeRef+ if size <= 0+ then pure val+ else do+ hole <- readPrimVar holeRef+ siftUp ids arr holeRef hole hole val++{-# SPECIALISE siftUp :: forall a. Ord a => MutablePrimArray RealWorld Int -> SmallMutableArray RealWorld a -> PrimVar RealWorld Int -> Int -> Int -> a -> IO a #-}+{-# SPECIALISE siftUp :: forall a s. Ord a => MutablePrimArray s Int -> SmallMutableArray s a -> PrimVar s Int -> Int -> Int -> a -> Strict.ST s a #-}+{-# SPECIALISE siftUp :: forall a s. Ord a => MutablePrimArray s Int -> SmallMutableArray s a -> PrimVar s Int -> Int -> Int -> a -> Lazy.ST s a #-}++siftUp :: forall a m. (PrimMonad m, Ord a) => MutablePrimArray (PrimState m) Int -> SmallMutableArray (PrimState m) a -> PrimVar (PrimState m) Int -> Int -> Int -> a -> m a+siftUp ids arr holeRef = sift+ where+ sift :: Int -> Int -> a -> m a+ sift !j !idxX !x = do+ !idxY <- readPrimArray ids j+ y <- readSmallArray arr j+ if j <= 0+ then if idxY < 0+ then do+ writePrimVar holeRef idxX+ pure x+ else do+ if x <= y+ then do+ writePrimVar holeRef idxX+ pure x+ else do+ writePrimArray ids j idxX+ writeSmallArray arr j x+ writePrimVar holeRef idxY+ pure y+ else if idxY < 0+ then sift (parentOf j) idxX x+ else do+ if x <= y+ then do+ sift (parentOf j) idxX x+ else do+ writePrimArray ids j idxX+ writeSmallArray arr j x+ sift (parentOf j) idxY y++{-------------------------------------------------------------------------------+ Helpers+-------------------------------------------------------------------------------}++halfOf :: Int -> Int+halfOf i = unsafeShiftR i 1+{-# INLINE halfOf #-}++parentOf :: Int -> Int+parentOf i = halfOf (i - 1)+{-# INLINE parentOf #-}
@@ -0,0 +1,105 @@+module MCG (+ MCG,+ make,+ period,+ next,+ reject,+) where++import Data.Bits (countLeadingZeros, unsafeShiftR)+import Data.List (nub)+import Data.Numbers.Primes (isPrime, primeFactors)+import Data.Word (Word64)++-- $setup+-- >>> import Data.List (unfoldr, nub)++-- | https://en.wikipedia.org/wiki/Lehmer_random_number_generator+data MCG = MCG { m :: !Word64, a :: !Word64, x :: !Word64 }+ deriving stock Show++-- invariants: m is a prime+-- a is a primitive element of Z_m+-- x is in [1..m-1]++-- | Create a MCG+--+-- >>> make 20 04+-- MCG {m = 23, a = 11, x = 5}+--+-- >>> make 101_000_000 20240429+-- MCG {m = 101000023, a = 197265, x = 20240430}+--+make ::+ Word64 -- ^ a lower bound for the period+ -> Word64 -- ^ initial seed.+ -> MCG+make (max 4 -> period_) seed = MCG m a (mod (seed + 1) m)+ where+ -- start prime search from an odd number larger than asked period.+ m = findM (if odd period_ then period_ + 2 else period_ + 1)+ m' = m - 1+ qs = nub $ primeFactors m'++ a = findA (guessA m)++ findM p = if isPrime p then p else findM (p + 2)++ -- we find `a` using "brute-force" approach.+ -- luckily, many elements a prime factors, so we don't need to try too hard.+ -- and we only need to check prime factors of m - 1.+ findA x+ | all (\q -> mod (x ^ div m' q) m /= 1) qs+ = x++ | otherwise+ = findA (x + 1)++-- | Period of the MCG.+--+-- Period is usually a bit larger than asked for, we look for the next prime:+--+-- >>> let g = make 9 04+-- >>> period g+-- 10+--+-- >>> take 22 (unfoldr (Just . next) g)+-- [4,7,3,1,0,5,2,6,8,9,4,7,3,1,0,5,2,6,8,9,4,7]+--+period :: MCG -> Word64+period (MCG m _ _) = m - 1++-- | Generate next number.+next :: MCG -> (Word64, MCG)+next (MCG m a x) = (x - 1, MCG m a (mod (x * a) m))++-- | Generate next numbers until one less than given bound is generated.+--+-- Replacing 'next' with @'reject' n@ effectively cuts the period to @n@:+--+-- >>> let g = make 9 04+-- >>> period g+-- 10+--+-- >>> take 22 (unfoldr (Just . reject 9) g)+-- [4,7,3,1,0,5,2,6,8,4,7,3,1,0,5,2,6,8,4,7,3,1]+--+-- if @n@ is close enough to actual period of 'MCG', the rejection ratio+-- is very small.+--+reject :: Word64 -> MCG -> (Word64, MCG)+reject ub g = case next g of+ (x, g') -> if x < ub then (x, g') else reject ub g'++-------------------------------------------------------------------------------+-- guessing some initial a+-------------------------------------------------------------------------------++-- | calculate x -> log2 (x + 1) i.e. approximate how large the number is in bits.+word64Log2m1 :: Word64 -> Int+word64Log2m1 x = 64 - countLeadingZeros x++-- | we guess a such that a*a is larger than m:+-- we shift a number a little.+guessA :: Word64 -> Word64+guessA x = unsafeShiftR x (div (word64Log2m1 x) 3)
@@ -0,0 +1,884 @@+{-# LANGUAGE ParallelListComp #-}++-- | This accompanies the format-page.md documentation as a sanity check+-- and a precise reference. It is intended to demonstrate that the page+-- format works. It is also used as a reference implementation for tests of+-- the real implementation.+--+-- Logically, a page is a sequence of key,operation pairs (with optional+-- blobrefs), sorted by key, and its serialised form fits within a disk page.+--+-- This reference implementation covers serialisation and deserialisation+-- (not lookups) which do not rely on (or enforce) the keys being sorted.+--+module FormatPage (+ -- * Page types+ Key (..),+ Operation (..),+ opHasBlobRef,+ Value (..),+ BlobRef (..),+ PageSerialised,+ PageIntermediate (..),+ PageSizesOffsets (..),++ -- * Page size+ PageSize (..),+ pageSizeEmpty,+ pageSizeAddElem,+ calcPageSize,++ -- * Encoding and decoding+ DiskPageSize(..),+ diskPageSizeBytes,+ encodePage,+ decodePage,+ serialisePage,+ deserialisePage,+ fromBitmap,+ toBitmap,++ -- * Overflow pages+ pageOverflowPrefixSuffixLen,+ pageDiskPages,+ pageSerialisedChunks,++ -- * Generators and shrinkers+ PageContentFits (..),+ genPageContentFits,+ PageContentMaybeOverfull (..),+ genPageContentMaybeOverfull,+ PageContentSingle (..),+ genPageContentSingle,+ genPageContentNearFull,+ genPageContentMedium,+ MinKeySize(..),+ noMinKeySize,+ maxKeySize,+ orderdKeyOps,+ shrinkKeyOps,+ shrinkOrderedKeyOps,+) where++import Data.Bits+import Data.Function (on)+import qualified Data.List as List+import Data.Maybe (fromJust, fromMaybe)+import Data.Word++import qualified Data.Binary.Get as Bin+import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BB+import qualified Data.ByteString.Char8 as BSC+import qualified Data.ByteString.Lazy as BSL++import Control.Exception (assert)+import Control.Monad++import Test.QuickCheck hiding ((.&.))++-------------------------------------------------------------------------------+-- Page content types+--++newtype Key = Key { unKey :: ByteString } deriving stock (Eq, Ord, Show)+newtype Value = Value { unValue :: ByteString } deriving stock (Eq, Show)++data Operation = Insert Value (Maybe BlobRef)+ | Mupsert Value+ | Delete+ deriving stock (Eq, Show)++data BlobRef = BlobRef Word64 Word32+ deriving stock (Eq, Show)++opHasBlobRef :: Operation -> Bool+opHasBlobRef (Insert _ (Just _blobref)) = True+opHasBlobRef _ = False+++-------------------------------------------------------------------------------+-- Disk page size+--++-- | A serialised page fits within chunks of memory of 4k, 8k, 16k, 32k or 64k.+--+data DiskPageSize = DiskPage4k | DiskPage8k+ | DiskPage16k | DiskPage32k+ | DiskPage64k+ deriving stock (Eq, Show, Enum, Bounded)++diskPageSizeBytes :: DiskPageSize -> Int+diskPageSizeBytes DiskPage4k = 2^(12::Int)+diskPageSizeBytes DiskPage8k = 2^(13::Int)+diskPageSizeBytes DiskPage16k = 2^(14::Int)+diskPageSizeBytes DiskPage32k = 2^(15::Int)+diskPageSizeBytes DiskPage64k = 2^(16::Int)+++-------------------------------------------------------------------------------+-- Calculating the page size (incrementally)+--++data PageSize = PageSize {+ pageSizeElems :: !Int,+ pageSizeBlobs :: !Int,+ pageSizeBytes :: !Int,+ pageSizeDisk :: !DiskPageSize+ }+ deriving stock (Eq, Show)++pageSizeEmpty :: DiskPageSize -> PageSize+pageSizeEmpty = PageSize 0 0 10++pageSizeAddElem :: (Key, Operation) -> PageSize -> Maybe PageSize+pageSizeAddElem (Key key, op) (PageSize n b sz dpgsz)+ | sz' <= diskPageSizeBytes dpgsz || n' == 1+ = Just (PageSize n' b' sz' dpgsz)+ | otherwise = Nothing+ where+ n' = n+1+ b' | opHasBlobRef op = b+1+ | otherwise = b+ sz' = sz+ + (if n `mod` 64 == 0 then 8 else 0) -- blobrefs bitmap+ + (if n `mod` 32 == 0 then 8 else 0) -- operations bitmap+ + (if opHasBlobRef op then 12 else 0) -- blobref entry+ + 2 -- key offsets+ + (case n of { 0 -> 4; 1 -> 0; _ -> 2}) -- value offsets+ + BS.length key+ + (case op of+ Insert (Value v) _ -> BS.length v+ Mupsert (Value v) -> BS.length v+ Delete -> 0)++calcPageSize :: DiskPageSize -> [(Key, Operation)] -> Maybe PageSize+calcPageSize dpgsz kops =+ go (pageSizeEmpty dpgsz) kops+ where+ go !pgsz [] = Just pgsz+ go !pgsz ((key, op):kops') =+ case pageSizeAddElem (key, op) pgsz of+ Nothing -> Nothing+ Just pgsz' -> go pgsz' kops'+++-------------------------------------------------------------------------------+-- Page encoding and serialisation types+--++-- | A serialised page consists of either a single disk page or several+-- disk pages. The latter is a primary page followed by one or more overflow+-- pages. Each disk page (single or multi) uses the same 'DiskPageSize', which+-- should be known from context (e.g. configuration).+--+type PageSerialised = ByteString++data PageIntermediate =+ PageIntermediate {+ pageNumKeys :: !Word16,+ pageNumBlobs :: !Word16,+ pageSizesOffsets :: !PageSizesOffsets,+ pageBlobRefBitmap :: [Bool],+ pageOperations :: [OperationEnum],+ pageBlobRefs :: [BlobRef],+ pageKeyOffsets :: [Word16],+ pageValueOffsets :: Either [Word16] (Word16, Word32),+ pageKeys :: !ByteString,+ pageValues :: !ByteString,+ pagePadding :: !ByteString, -- ^ Padding to the 'DiskPageSize'+ pageDiskPageSize :: !DiskPageSize+ }+ deriving stock (Eq, Show)++data OperationEnum = OpInsert | OpMupsert | OpDelete+ deriving stock (Eq, Show)++data PageSizesOffsets =+ PageSizesOffsets {+ sizeDirectory :: !Word16,+ sizeBlobRefBitmap :: !Word16,+ sizeOperations :: !Word16,+ sizeBlobRefs :: !Word16,+ sizeKeyOffsets :: !Word16,+ sizeValueOffsets :: !Word16,+ sizeKeys :: !Word16,+ sizeValues :: !Word32,++ offBlobRefBitmap :: !Word16,+ offOperations :: !Word16,+ offBlobRefs :: !Word16,+ offKeyOffsets :: !Word16,+ offValueOffsets :: !Word16,+ offKeys :: !Word16,+ offValues :: !Word16,++ sizePageUsed :: !Word32, -- ^ The size in bytes actually used+ sizePagePadding :: !Word32, -- ^ The size in bytes of trailing padding+ sizePageDiskPage :: !Word32 -- ^ The size in bytes rounded up to a+ -- multiple of the disk page size.+ }+ deriving stock (Eq, Show)+++-------------------------------------------------------------------------------+-- Page encoding and serialisation+--++-- | Returns @Nothing@ if the size would be over-full for the given disk page+-- size.+--+calcPageSizeOffsets :: DiskPageSize -- ^ underlying page size: 4k, 8k ... 64k+ -> Int -- ^ number of keys\/entries+ -> Int -- ^ number of blobs+ -> Int -- ^ total size of the keys+ -> Int -- ^ total size of the values+ -> Maybe PageSizesOffsets+calcPageSizeOffsets dpgsz n b sizeKeys sizeValues+ | n < 0 || b < 0 || sizeKeys < 0 || sizeValues < 0+ = Nothing++ | n /= 1 --single entries can use multiple disk pages+ , sizePageUsed > diskPageSize+ = Nothing++ | otherwise+ = Just PageSizesOffsets {+ -- having checked for not over-full, we can now guarantee all+ -- these conversions into smaller types will not overflow:+ sizeDirectory = fromIntegralChecked sizeDirectory,+ sizeBlobRefBitmap = fromIntegralChecked sizeBlobRefBitmap,+ sizeOperations = fromIntegralChecked sizeOperations,+ sizeBlobRefs = fromIntegralChecked sizeBlobRefs,+ sizeKeyOffsets = fromIntegralChecked sizeKeyOffsets,+ sizeValueOffsets = fromIntegralChecked sizeValueOffsets,+ sizeKeys = fromIntegralChecked sizeKeys,+ sizeValues = fromIntegralChecked sizeValues,++ offBlobRefBitmap = fromIntegralChecked offBlobRefBitmap,+ offOperations = fromIntegralChecked offOperations,+ offBlobRefs = fromIntegralChecked offBlobRefs,+ offKeyOffsets = fromIntegralChecked offKeyOffsets,+ offValueOffsets = fromIntegralChecked offValueOffsets,+ offKeys = fromIntegralChecked offKeys,+ offValues = fromIntegralChecked offValues,++ sizePageUsed = fromIntegralChecked sizePageUsed,+ sizePagePadding = fromIntegralChecked sizePagePadding,+ sizePageDiskPage = fromIntegralChecked sizePageDiskPage+ }+ where+ sizeDirectory, sizeBlobRefBitmap,+ sizeOperations, sizeBlobRefs,+ sizeKeyOffsets, sizeValueOffsets :: Int+ sizeDirectory = 8+ sizeBlobRefBitmap = (n + 63) `shiftR` 3 .&. complement 0x7+ sizeOperations = (2 * n + 63) `shiftR` 3 .&. complement 0x7+ sizeBlobRefs = (4 + 8) * b+ sizeKeyOffsets = 2 * n+ sizeValueOffsets | n == 1 = 6+ | otherwise = 2 * (n+1)++ offBlobRefBitmap, offOperations, offBlobRefs,+ offKeyOffsets, offValueOffsets,+ offKeys, offValues :: Int+ offBlobRefBitmap = sizeDirectory+ offOperations = offBlobRefBitmap + sizeBlobRefBitmap+ offBlobRefs = offOperations + sizeOperations+ offKeyOffsets = offBlobRefs + sizeBlobRefs+ offValueOffsets = offKeyOffsets + sizeKeyOffsets+ offKeys = offValueOffsets + sizeValueOffsets+ offValues = offKeys + sizeKeys++ sizePageUsed, sizePagePadding,+ sizePageDiskPage, diskPageSize :: Int+ sizePageUsed = offValues + sizeValues+ sizePagePadding = case sizePageUsed `mod` diskPageSize of+ 0 -> 0+ p -> diskPageSize - p+ sizePageDiskPage = sizePageUsed + sizePagePadding+ diskPageSize = diskPageSizeBytes dpgsz++encodePage :: DiskPageSize -> [(Key, Operation)] -> Maybe PageIntermediate+encodePage dpgsz kops = do+ let pageNumKeys = length kops+ pageNumBlobs = length (filter (opHasBlobRef . snd) kops)+ keys = [ k | (k,_) <- kops ]+ values = [ v | (_,op) <- kops+ , let v = case op of+ Insert v' _ -> v'+ Mupsert v' -> v'+ Delete -> Value (BS.empty)+ ]++ pageSizesOffsets@PageSizesOffsets {+ offKeys, offValues, sizePagePadding+ } <- calcPageSizeOffsets+ dpgsz+ pageNumKeys pageNumBlobs+ (sum [ BS.length k | Key k <- keys ])+ (sum [ BS.length v | Value v <- values ])++ let pageBlobRefBitmap = [ opHasBlobRef op | (_,op) <- kops ]+ pageOperations = [ toOperationEnum op | (_,op) <- kops ]+ pageBlobRefs = [ blobref | (_,Insert _ (Just blobref)) <- kops ]++ pageKeyOffsets = init $ scanl (\o k -> o + keyLen16 k)+ offKeys keys+ pageValueOffsets = case values of+ [v] -> Right (offValues,+ fromIntegral offValues+ + valLen32 v)+ _ -> Left (scanl (\o v -> o + valLen16 v)+ offValues values)+ pageKeys = BS.concat [ k | Key k <- keys ]+ pageValues = BS.concat [ v | Value v <- values ]+ pagePadding = BS.replicate (fromIntegral sizePagePadding) 0+ pageDiskPageSize = dpgsz++ pure PageIntermediate {+ pageNumKeys = fromIntegralChecked pageNumKeys,+ pageNumBlobs = fromIntegralChecked pageNumBlobs,+ ..+ }+ where+ toOperationEnum Insert{} = OpInsert+ toOperationEnum Mupsert{} = OpMupsert+ toOperationEnum Delete{} = OpDelete++ keyLen16 :: Key -> Word16+ valLen16 :: Value -> Word16+ valLen32 :: Value -> Word32+ keyLen16 (Key k) = fromIntegral $ BS.length k+ valLen16 (Value v) = fromIntegral $ BS.length v+ valLen32 (Value v) = fromIntegral $ BS.length v++serialisePage :: PageIntermediate -> PageSerialised+serialisePage PageIntermediate{pageSizesOffsets = PageSizesOffsets{..}, ..} =+ BSL.toStrict . BB.toLazyByteString $++ -- the top level directory+ BB.word16LE pageNumKeys+ <> BB.word16LE pageNumBlobs+ <> BB.word16LE offKeyOffsets+ <> BB.word16LE 0 --spare+ <> mconcat [ BB.word64LE w | w <- toBitmap pageBlobRefBitmap ]+ <> mconcat [ BB.word64LE w | w <- toBitmap . concatMap opEnumToBits+ $ pageOperations ]+ <> mconcat [ BB.word64LE w64 | BlobRef w64 _w32 <- pageBlobRefs ]+ <> mconcat [ BB.word32LE w32 | BlobRef _w64 w32 <- pageBlobRefs ]+ <> mconcat [ BB.word16LE off | off <- pageKeyOffsets ]+ <> case pageValueOffsets of+ Left offsets -> mconcat [ BB.word16LE off | off <- offsets ]+ Right (offset1, offset2) -> BB.word16LE offset1+ <> BB.word32LE offset2+ <> BB.byteString pageKeys+ <> BB.byteString pageValues+ <> BB.byteString pagePadding+ where+ opEnumToBits OpInsert = [False, False]+ opEnumToBits OpMupsert = [True, False]+ opEnumToBits OpDelete = [False, True]++deserialisePage :: DiskPageSize -> PageSerialised -> PageIntermediate+deserialisePage dpgsz p =+ flip Bin.runGet (BSL.fromStrict p) $ do+ pageNumKeys <- fromIntegral <$> Bin.getWord16le+ pageNumBlobs <- fromIntegral <$> Bin.getWord16le+ offsetKeyOffsets <- Bin.getWord16le+ _ <- Bin.getWord16le++ let sizeWord64BlobRefBitmap :: Int+ sizeWord64BlobRefBitmap = (pageNumKeys + 63) `shiftR` 6++ sizeWord64Operations :: Int+ sizeWord64Operations = (2 * pageNumKeys + 63) `shiftR` 6++ pageBlobRefBitmap <- take pageNumKeys . fromBitmap <$>+ replicateM sizeWord64BlobRefBitmap Bin.getWord64le+ pageOperations <- take pageNumKeys . opBitsToEnum . fromBitmap <$>+ replicateM sizeWord64Operations Bin.getWord64le+ pageBlobRefsW64 <- replicateM pageNumBlobs Bin.getWord64le+ pageBlobRefsW32 <- replicateM pageNumBlobs Bin.getWord32le+ let pageBlobRefs = zipWith BlobRef pageBlobRefsW64 pageBlobRefsW32++ pageKeyOffsets <- replicateM pageNumKeys Bin.getWord16le+ pageValueOffsets <-+ if pageNumKeys == 1+ then Right <$> ((,) <$> Bin.getWord16le+ <*> Bin.getWord32le)+ else Left <$> replicateM (pageNumKeys + 1) Bin.getWord16le++ let sizeKeys :: Int+ sizeKeys+ | pageNumKeys > 0+ = fromIntegral (either head fst pageValueOffsets)+ - fromIntegral (head pageKeyOffsets)+ | otherwise = 0++ sizeValues :: Int+ sizeValues+ | pageNumKeys > 0+ = either (fromIntegral . last) (fromIntegral . snd) pageValueOffsets+ - fromIntegral (either head fst pageValueOffsets)+ | otherwise = 0+ pageKeys <- Bin.getByteString sizeKeys+ pageValues <- Bin.getByteString sizeValues+ pagePadding <- BSL.toStrict <$> Bin.getRemainingLazyByteString++ let pageSizesOffsets =+ fromMaybe (error "deserialisePage: disk page overflow") $+ calcPageSizeOffsets+ dpgsz+ pageNumKeys pageNumBlobs+ sizeKeys sizeValues++ assert (offsetKeyOffsets == offKeyOffsets pageSizesOffsets) $+ pure PageIntermediate {+ pageNumKeys = fromIntegral pageNumKeys,+ pageNumBlobs = fromIntegral pageNumBlobs,+ pageDiskPageSize = dpgsz,+ ..+ }+ where+ opBitsToEnum (False:False:bits) = OpInsert : opBitsToEnum bits+ opBitsToEnum (True: False:bits) = OpMupsert : opBitsToEnum bits+ opBitsToEnum (False:True :bits) = OpDelete : opBitsToEnum bits+ opBitsToEnum [] = []+ opBitsToEnum _ = error "opBitsToEnum"++decodePage :: PageIntermediate -> [(Key, Operation)]+decodePage PageIntermediate{pageSizesOffsets = PageSizesOffsets{..}, ..} =+ [ let op = case opEnum of+ OpInsert -> Insert (Value value) mblobref+ OpMupsert -> Mupsert (Value value)+ OpDelete -> Delete+ mblobref | hasBlobref = Just (pageBlobRefs !! idxBlobref)+ | otherwise = Nothing+ in (Key key, op)+ | opEnum <- pageOperations+ | hasBlobref <- pageBlobRefBitmap+ | idxBlobref <- scanl (\o b -> if b then o+1 else o) 0 pageBlobRefBitmap+ | keySpan <- spans $ pageKeyOffsets+ ++ either (take 1) ((:[]) . fst) pageValueOffsets+ , let key = BS.take (fromIntegral $ snd keySpan - fst keySpan)+ . BS.drop (fromIntegral $ fst keySpan - offKeys)+ $ pageKeys+ | valSpan <- spans $ case pageValueOffsets of+ Left offs -> map fromIntegral offs+ Right (off1, off2) -> [ fromIntegral off1, off2 ]+ , let value = BS.take (fromIntegral $ snd valSpan - fst valSpan)+ . BS.drop (fromIntegral $ fst valSpan - fromIntegral offValues)+ $ pageValues++ ]+ where+ spans xs = zip xs (drop 1 xs)++toBitmap :: [Bool] -> [Word64]+toBitmap =+ map toWord64 . group64+ where+ toWord64 :: [Bool] -> Word64+ toWord64 = List.foldl' (\w (n,b) -> if b then setBit w n else w) 0+ . zip [0 :: Int ..]+ group64 = List.unfoldr (\xs -> if null xs+ then Nothing+ else Just (splitAt 64 xs))++fromBitmap :: [Word64] -> [Bool]+fromBitmap =+ concatMap fromWord64+ where+ fromWord64 :: Word64 -> [Bool]+ fromWord64 w = [ testBit w i | i <- [0..63] ]++fromIntegralChecked :: (Integral a, Integral b) => a -> b+fromIntegralChecked x+ | let x' = fromIntegral x+ , fromIntegral x' == x+ = x'++ | otherwise+ = error "fromIntegralChecked: conversion failed"++-- | If a page uses overflow pages, return the:+--+-- 1. value prefix length (within the first page)+-- 2. value suffix length (within the overflow pages)+--+pageOverflowPrefixSuffixLen :: PageIntermediate -> Maybe (Int, Int)+pageOverflowPrefixSuffixLen p =+ case pageValueOffsets p of+ Right (offStart, offEnd)+ | let page1End = diskPageSizeBytes (pageDiskPageSize p)+ , fromIntegral offEnd > page1End+ , let prefixlen, suffixlen :: Int+ prefixlen = page1End - fromIntegral offStart+ suffixlen = fromIntegral offEnd - page1End+ -> Just (prefixlen, suffixlen)+ _ -> Nothing++-- | The total number of disk pages, including any overflow pages.+--+pageDiskPages :: PageIntermediate -> Int+pageDiskPages p =+ nbytes `div` diskPageSizeBytes (pageDiskPageSize p)+ where+ nbytes = fromIntegral (sizePageDiskPage (pageSizesOffsets p))++pageSerialisedChunks :: DiskPageSize -> PageSerialised -> [ByteString]+pageSerialisedChunks dpgsz =+ List.unfoldr (\p -> if BS.null p then Nothing+ else Just (BS.splitAt dpgszBytes p))+ where+ dpgszBytes = diskPageSizeBytes dpgsz++-------------------------------------------------------------------------------+-- Test types and generators+--++data PageContentFits = PageContentFits DiskPageSize [(Key, Operation)]+ deriving stock Show++data PageContentMaybeOverfull = PageContentMaybeOverfull DiskPageSize+ [(Key, Operation)]+ deriving stock Show++data PageContentSingle = PageContentSingle DiskPageSize Key Operation+ deriving stock Show++instance Arbitrary PageContentFits where+ arbitrary = do+ dpgsz <- arbitrary+ kops <- genPageContentFits dpgsz noMinKeySize+ pure (PageContentFits dpgsz kops)++instance Arbitrary PageContentMaybeOverfull where+ arbitrary = do+ dpgsz <- arbitrary+ kops <- genPageContentMaybeOverfull dpgsz noMinKeySize+ pure (PageContentMaybeOverfull dpgsz kops)++instance Arbitrary PageContentSingle where+ arbitrary = do+ dpgsz <- arbitrary+ (k,op) <- genPageContentSingle dpgsz noMinKeySize+ pure (PageContentSingle dpgsz k op)++-- | In some use cases it is necessary to generate 'Keys' that are at least of+-- some minimum length. Use 'noMinKeySize' if no such constraint is need.+newtype MinKeySize = MinKeySize Int+ deriving stock Show++-- | No minimum key size: @MinKeySize 0@.+noMinKeySize :: MinKeySize+noMinKeySize = MinKeySize 0++-- | Generate a test case consisting of a key\/operation sequence that is+-- guaranteed to fit into a given disk page size.+--+-- The distribution is designed to cover:+--+-- * small pages+-- * medium sized pages+-- * nearly full pages+-- * plus single key pages (possibly using one or more overflow pages)+-- * a corner case of a single large key\/operation pair followed by some small+-- key op pairs.+--+-- The keys are /not/ ordered: use 'orderdKeyOps' to sort and de-duplicate them+-- if that is needed (but note this will change the order of key sizes).+--+genPageContentFits :: DiskPageSize -> MinKeySize -> Gen [(Key, Operation)]+genPageContentFits dpgsz minkeysz =+ frequency+ [ (6, genPageContentMedium dpgsz genkey genval)+ , (2, (:[]) <$> genPageContentSingle dpgsz minkeysz)+ , (1, genPageContentLargeSmallFits dpgsz minkeysz)+ ]+ where+ genkey = genKeyMinSize dpgsz minkeysz+ genval = arbitrary++-- | Generate a test case consisting of a key\/operation sequence that is /not/+-- guaranteed to fit into a given disk page size.+--+-- These test cases are useful for checking the boundary conditions for what+-- can fit into a disk page. This covers a similar distribution to+-- 'genPageContentFits' but also includes about 20% of pages that are over full,+-- including the corner case of a large key ops pair followed by smaller key op+-- pairs (again possibly over full).+--+-- The keys are /not/ ordered: use 'orderdKeyOps' to sort and de-duplicate them+-- if that is needed.+--+genPageContentMaybeOverfull :: DiskPageSize+ -> MinKeySize -> Gen [(Key, Operation)]+genPageContentMaybeOverfull dpgsz minkeysz =+ frequency+ [ (6, genPageContentMedium dpgsz genkey genval)+ , (1, (:[]) <$> genPageContentSingle dpgsz minkeysz)+ , (1, genPageContentOverfull dpgsz genkey genval)+ , (1, genPageContentLargeSmallOverfull dpgsz minkeysz)+ ]+ where+ genkey = genKeyMinSize dpgsz minkeysz+ genval = arbitrary++-- | Generate a test case consisting of a single key\/operation pair.+--+genPageContentSingle :: DiskPageSize -> MinKeySize -> Gen (Key, Operation)+genPageContentSingle dpgsz minkeysz =+ oneof+ [ genPageContentSingleSmall genkey genval+ , genPageContentSingleNearFull dpgsz minkeysz+ , genPageContentSingleMultiPage dpgsz minkeysz+ ]+ where+ genkey = genKeyMinSize dpgsz minkeysz+ genval = arbitrary++-- | This generates a reasonable \"middle\" distribution of page sizes+-- (relative to the given disk page size). In particular it covers:+--+-- * small pages (~45% for 4k pages, ~15% for 64k pages)+-- * near-maximum pages (~20% for 4k pages, ~20% for 64k pages)+-- * some in between (~35% for 4k pages, ~60% for 64k pages)+--+-- The numbers above are when used with the normal 'arbitrary' 'Key' and+-- 'Value' generators. And with these generators, it tends to use lots of+-- small-to-medium size keys and values, rather than a few huge ones.+--+genPageContentMedium :: DiskPageSize+ -> Gen Key+ -> Gen Value+ -> Gen [(Key, Operation)]+genPageContentMedium dpgsz genkey genval =+ takePageContentFits dpgsz <$>+ scale scaleForDiskPageSize+ (listOf (genPageContentSingleSmall genkey genval))+ where+ scaleForDiskPageSize :: Int -> Int+ scaleForDiskPageSize sz =+ ceiling $+ fromIntegral sz ** (1.0 + fromIntegral (fromEnum dpgsz) / 10 :: Float)++takePageContentFits :: DiskPageSize -> [(Key, Operation)] -> [(Key, Operation)]+takePageContentFits dpgsz = go (pageSizeEmpty dpgsz)+ where+ go _sz [] = []+ go sz (kop:kops)+ | Just sz' <- pageSizeAddElem kop sz = kop : go sz' kops+ | otherwise = []++-- | Generate only pages that are nearly full. This isn't the maximum possible+-- size, but where adding one more randomly-chosen key\/op pair would not fit+-- (but perhaps a smaller pair would still fit).+--+-- Consider if you really need this: the 'genPageContentMedium' also includes+-- these cases naturally as part of its distribution. On the other hand, this+-- can be good for generating benchmark data.+--+genPageContentNearFull :: DiskPageSize+ -> Gen Key+ -> Gen Value+ -> Gen [(Key, Operation)]+genPageContentNearFull dpgsz genkey genval =+ --relies on first item being the one triggering over-full:+ drop 1 <$> genPageContentOverfull dpgsz genkey genval++-- | Generate pages that are just slightly over-full. This is where the last+-- key\/op pair takes it just over the disk page size (but this element is+-- first in the sequence).+--+genPageContentOverfull :: DiskPageSize+ -> Gen Key+ -> Gen Value+ -> Gen [(Key, Operation)]+genPageContentOverfull dpgsz genkey genval =+ go [] (pageSizeEmpty dpgsz)+ where+ go :: [(Key, Operation)] -> PageSize -> Gen [(Key, Operation)]+ go kops sz = do+ kop <- genPageContentSingleSmall genkey genval+ case pageSizeAddElem kop sz of+ -- include as the /first/ element, the one that will make it overfull:+ Nothing -> pure (kop:kops) -- not reversed!+ Just sz' -> go (kop:kops) sz'++genPageContentLargeSmallFits :: DiskPageSize+ -> MinKeySize+ -> Gen [(Key, Operation)]+genPageContentLargeSmallFits dpgsz minkeysz =+ takePageContentFits dpgsz <$>+ genPageContentLargeSmallOverfull dpgsz minkeysz++genPageContentLargeSmallOverfull :: DiskPageSize+ -> MinKeySize+ -> Gen [(Key, Operation)]+genPageContentLargeSmallOverfull dpgsz (MinKeySize minkeysz) =+ (\large small -> large : small)+ <$> genPageContentSingleOfSize genKeyValSizes+ <*> arbitrary+ where+ genKeyValSizes = do+ let size = maxKeySize dpgsz - 100+ split <- choose (minkeysz, size)+ pure (split, size - split)++genPageContentSingleOfSize :: Gen (Int, Int) -> Gen (Key, Operation)+genPageContentSingleOfSize genKeyValSizes = do+ (keySize, valSize) <- genKeyValSizes+ key <- Key . BS.pack <$> vectorOf keySize arbitrary+ val <- Value . BS.pack <$> vectorOf valSize arbitrary+ op <- oneof -- no delete+ [ Insert val <$> arbitrary+ , pure (Mupsert val) ]+ pure (key, op)++genPageContentSingleSmall :: Gen Key -> Gen Value -> Gen (Key, Operation)+genPageContentSingleSmall genkey genval =+ (,) <$> genkey <*> genOperation genval++-- | Generate pages around the disk page size, above and below.+--+-- The key is always within the min key size given and max key size for the+-- page size.+genPageContentSingleNearFull :: DiskPageSize+ -> MinKeySize+ -> Gen (Key, Operation)+genPageContentSingleNearFull dpgsz (MinKeySize minkeysize) =+ genPageContentSingleOfSize genKeyValSizes+ where+ genKeyValSizes = do+ let maxkeysize = maxKeySize dpgsz+ size <- choose (maxkeysize - 15, maxkeysize + 15)+ split <- choose (minkeysize, maxkeysize `min` size)+ pure (split, size - split)++genPageContentSingleMultiPage :: DiskPageSize+ -> MinKeySize+ -> Gen (Key, Operation)+genPageContentSingleMultiPage dpgsz (MinKeySize minkeysz) =+ genPageContentSingleOfSize genKeyValSizes+ where+ genKeyValSizes =+ (,) <$> choose (minkeysz, maxKeySize dpgsz)+ <*> choose (0, diskPageSizeBytes dpgsz * 3)++genKeyOfSize :: Gen Int -> Gen Key+genKeyOfSize genSize =+ genSize >>= \n -> Key . BS.pack <$!> vectorOf n arbitrary++genKeyMinSize :: DiskPageSize -> MinKeySize -> Gen Key+genKeyMinSize dpgsz (MinKeySize minsz) =+ genKeyOfSize+ (getSize >>= \sz -> chooseInt (minsz, sz `min` maxKeySize dpgsz))++instance Arbitrary Key where+ arbitrary =+ genKeyOfSize+ (getSize >>= \sz -> chooseInt (0, sz `min` maxKeySize DiskPage4k))++ shrink = shrinkMapBy Key unKey shrinkOpaqueByteString++genValueOfSize :: Gen Int -> Gen Value+genValueOfSize genSize =+ genSize >>= \n -> Value . BS.pack <$!> vectorOf n arbitrary++instance Arbitrary Value where+ arbitrary = genValueOfSize (getSize >>= \sz -> chooseInt (0, sz))++ shrink = shrinkMapBy Value unValue shrinkOpaqueByteString++genOperation :: Gen Value -> Gen Operation+genOperation genval =+ oneof+ [ Insert <$> genval <*> arbitrary+ , Mupsert <$> genval+ , pure Delete+ ]++instance Arbitrary Operation where+ arbitrary = genOperation arbitrary++ shrink :: Operation -> [Operation]+ shrink Delete = []+ shrink (Insert v mb) = Delete+ : [ Insert v' mb' | (v', mb') <- shrink (v, mb) ]+ shrink (Mupsert v) = Insert v Nothing+ : [ Mupsert v' | v' <- shrink v ]++instance Arbitrary BlobRef where+ arbitrary = BlobRef <$> arbitrary <*> arbitrary++ shrink (BlobRef 0 0) = []+ shrink (BlobRef _ _) = [BlobRef 0 0]++instance Arbitrary DiskPageSize where+ arbitrary = scale (`div` 5) $ growingElements [minBound..]+ shrink = shrinkBoundedEnum++-- | Sort and de-duplicate a key\/operation sequence to ensure the sequence is+-- strictly ascending by key.+--+-- If you need this in a QC generator, you will need 'shrinkOrderedKeyOps' in+-- the corresponding shrinker.+--+orderdKeyOps :: [(Key, Operation)] -> [(Key, Operation)]+orderdKeyOps =+ List.nubBy ((==) `on` fst)+ . List.sortBy (compare `on` fst)++-- | Shrink a key\/operation sequence (without regard to key order).+shrinkKeyOps :: [(Key, Operation)] -> [[(Key, Operation)]]+shrinkKeyOps = shrink+ -- It turns out that the generic list shrink is actually good enough,+ -- but only because we've got carefully chosen shrinkers for Key and Value.+ -- Without those special shrinkers, this one would blow up.++-- | Shrink a key\/operation sequence, preserving key order.+shrinkOrderedKeyOps :: [(Key, Operation)] -> [[(Key, Operation)]]+shrinkOrderedKeyOps = map orderdKeyOps . shrink++-- | Shrink 'ByteString's that are used as opaque blobs, where their value+-- is generally not used, except for ordering. We minimise the number of+-- alternative shrinks, to help minimise the number of shrinks when lots+-- of such values are used, e.g. in key\/value containers.+--+-- This tries only three alternatives:+--+-- * take the first half+-- * take everything but the final byte+-- * replace the last (non-space) character by a space+--+-- > > shrinkOpaqueByteString "hello world!"+-- > ["hello ","hello world", "hello world "]+--+-- Using space as the replacement character makes the resulting strings+-- printable and shorter than lots of @\NUL\NUL\NUL@, which makes for test+-- failure cases that are easier to read.+--+shrinkOpaqueByteString :: ByteString -> [ByteString]+shrinkOpaqueByteString bs =+ [ BS.take (BS.length bs `div` 2) bs | BS.length bs > 2 ]+ ++ [ BS.init bs | BS.length bs > 0 ]+ ++ case BSC.spanEnd (==' ') bs of+ (prefix, spaces)+ | BS.null prefix -> []+ | otherwise -> [ BS.init prefix <> BSC.cons ' ' spaces ]++-- | The maximum size of key that is guaranteed to always fit in an empty+-- 4k page. So this is a worst case maximum size: this size key will fit+-- irrespective of the corresponding operation, including the possibility+-- that the key\/op pair has a blob reference.+maxKeySize :: DiskPageSize -> Int+maxKeySize dpgsz = diskPageSizeBytes dpgsz - pageSizeOverhead++pageSizeOverhead :: Int+pageSizeOverhead =+ (pageSizeBytes . fromJust . calcPageSize DiskPage4k)+ [(Key BS.empty, Insert (Value BS.empty) (Just (BlobRef 0 0)))]+ -- the page size passed to calcPageSize here is irrelevant
@@ -0,0 +1,1981 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE UnboxedTuples #-}++{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+{-# OPTIONS_GHC -Wno-partial-fields #-}++-- | A prototype of an LSM with explicitly scheduled incremental merges.+--+-- The scheduled incremental merges is about ensuring that the merging+-- work (CPU and I\/O) can be spread out over time evenly. This also means+-- the LSM update operations have worst case complexity rather than amortised+-- complexity, because they do a fixed amount of merging work each.+--+-- Another thing this prototype demonstrates is a design for duplicating tables+-- and sharing ongoing incremental merges.+--+-- Finally, it demonstrates a design for table unions, including a+-- representation for in-progress merging trees.+--+-- The merging policy that this prototype uses is \"lazy levelling\".+-- Each level is T times bigger than the previous level.+-- Lazy levelling means we use tiering for every level except the last level+-- which uses levelling. Though note that the first level always uses tiering,+-- even if the first level is also the last level. This is to simplify flushing+-- the write buffer: if we used levelling on the first level we would need a+-- code path for merging the write buffer into the first level.+--+module ScheduledMerges (+ -- * Main API+ LSM,+ TableId (..),+ LSMConfig (..),+ Key (K), Value (V), resolveValue, Blob (B),+ new,+ newWith,+ LookupResult (..),+ lookup, lookups,+ Entry,+ Update (..),+ update, updates,+ insert, inserts,+ delete, deletes,+ mupsert, mupserts,+ supplyMergeCredits,+ duplicate,+ unions,+ Credit,+ Debt,+ remainingUnionDebt,+ supplyUnionCredits,++ -- * Test and trace+ MTree (..),+ logicalValue,+ Representation,+ dumpRepresentation,+ representationShape,+ Event,+ EventAt(..),+ EventDetail(..),+ MergingTree(..),+ MergingTreeState(..),+ PendingMerge(..),+ PreExistingRun(..),+ MergingRun(..),+ MergingRunState(..),+ MergePolicyForLevel(..),+ IsMergeType(..),+ TreeMergeType(..),+ LevelMergeType(..),+ MergeCredit(..),+ MergeDebt(..),+ NominalCredit(..),+ NominalDebt(..),+ Run,+ runSize,+ UnionCredits (..),+ supplyCreditsMergingTree,+ UnionDebt(..),+ remainingDebtMergingTree,+ mergek,+ mergeBatchSize,++ -- * Invariants+ Invariant,+ evalInvariant,+ treeInvariant,+ mergeDebtInvariant,++ -- * Run sizes+ levelNumberToMaxRunSize,+ runSizeToLevelNumber,+ maxWriteBufferSize,+ runSizeFitsInLevel,+ runSizeTooSmallForLevel,+ runSizeTooLargeForLevel,++ -- * Level capacity+ levelIsFull,+ ) where++import Prelude hiding (lookup)++import Data.Foldable (for_, toList, traverse_)+import Data.Functor.Contravariant+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (catMaybes)+import Data.Primitive.Types+import Data.STRef++import qualified Control.Exception as Exc (assert)+import Control.Monad (foldM, forM, when)+import Control.Monad.ST+import qualified Control.Monad.Trans.Except as E+import Control.Tracer+import GHC.Stack (HasCallStack, callStack)++import Text.Printf (printf)++import qualified Test.QuickCheck as QC++data LSM s = LSMHandle {+ tableId :: !TableId+ , _tableCounter :: !(STRef s Counter)+ , _tableConfig :: !LSMConfig+ , _tableContents :: !(STRef s (LSMContent s))+ }++-- | Identifiers for 'LSM' tables+newtype TableId = TableId Int+ deriving stock (Show, Eq, Ord)+ deriving newtype (Enum, Prim)++-- | Configuration options for individual LSM tables.+data LSMConfig = LSMConfig {+ configMaxWriteBufferSize :: !Int+ -- | Also known as the parameter @T@+ , configSizeRatio :: !Int+ }+ deriving stock (Show, Eq)++-- | A simple count of LSM operations to allow logging the operation+-- number in each event. This enables relating merge events to the+-- operation number (which is interesting for numerical representations+-- like this). We would not need this in the real implementation.+type Counter = Int++-- | The levels of the table, from most to least recently inserted.+data LSMContent s =+ LSMContent+ Buffer -- ^ write buffer is level 0 of the table, in-memory+ (Levels s) -- ^ \"regular\" levels 1+, on disk in real implementation+ (UnionLevel s) -- ^ a potential last level++type Levels s = [Level s]++-- | The number of the level. The write buffer lives at level 0, and all other+-- levels are numbered starting from 1.+type LevelNo = Int++-- | A level is a sequence of resident runs at this level, prefixed by an+-- incoming run, which is usually multiple runs that are being merged. Once+-- completed, the resulting run will become a resident run at this level.+data Level s = Level !(IncomingRun s) ![Run]++-- | We represent single runs specially, rather than putting them in as a+-- 'CompletedMerge'. This is for two reasons: to see statically that it's a+-- single run without having to read the 'STRef', and secondly to make it easier+-- to avoid supplying merge credits. It's not essential, but simplifies things+-- somewhat.+data IncomingRun s = Merging !MergePolicyForLevel+ !NominalDebt !(STRef s NominalCredit)+ !(MergingRun LevelMergeType s)+ | Single !Run++-- | The merge policy for a LSM level can be either tiering or levelling.+-- In this design we use levelling for the last level, and tiering for+-- all other levels. The first level always uses tiering however, even if+-- it's also the last level. So 'MergePolicyForLevel' and 'LevelMergeType' are+-- orthogonal, all combinations are possible.+--+data MergePolicyForLevel = LevelTiering | LevelLevelling+ deriving stock (Eq, Show)++-- | A \"merging run\" is a mutable representation of an incremental merge.+-- It is also a unit of sharing between duplicated tables.+--+data MergingRun t s = MergingRun !t !MergeDebt+ !(STRef s MergingRunState)++data MergingRunState = CompletedMerge !Run+ | OngoingMerge+ !MergeCredit+ ![Run] -- ^ inputs of the merge+ Run -- ^ output of the merge (lazily evaluated)++-- | Merges can exist in different parts of the LSM, each with different options+-- for the exact merge operation performed.+class Show t => IsMergeType t where+ isLastLevel :: t -> Bool+ isUnion :: t -> Bool++-- | Different types of merges created as part of a regular (non-union) level.+--+-- A last level merge behaves differently from a mid-level merge: last level+-- merges can actually remove delete entries, whereas mid-level merges must+-- preserve them. This is orthogonal to the 'MergePolicyForLevel'.+data LevelMergeType = MergeMidLevel | MergeLastLevel+ deriving stock (Eq, Show)++instance IsMergeType LevelMergeType where+ isLastLevel = \case+ MergeMidLevel -> False+ MergeLastLevel -> True+ isUnion = const False++-- | Different types of merges created as part of the merging tree.+--+-- Union merges follow the semantics of @Data.Map.unionWith (<>)@. Since+-- the input runs are semantically treated like @Data.Map@s, deletes are ignored+-- and inserts act like mupserts, so they need to be merged monoidally using+-- 'resolveValue'.+--+-- Trees can only exist on the union level, which is the last. Therefore, node+-- merges can always drop deletes.+data TreeMergeType = MergeLevel | MergeUnion+ deriving stock (Eq, Show)++instance IsMergeType TreeMergeType where+ isLastLevel = const True+ isUnion = \case+ MergeLevel -> False+ MergeUnion -> True++-- | An additional optional last level, created as a result of 'union'. It can+-- not only contain an ongoing merge of multiple runs, but a nested tree of+-- merges. See Note [Table Unions].+data UnionLevel s = NoUnion+ -- | We track the debt to make sure it never increases.+ | Union !(MergingTree s) !(STRef s Debt)++-- | A \"merging tree\" is a mutable representation of an incremental+-- tree-shaped nested merge. This allows to represent union merges of entire+-- tables, each of which itself first need to be merged to become a single run.+--+-- Trees have to support arbitrarily deep nesting, since each input to 'union'+-- might already contain an in-progress merging tree (which then becomes shared+-- between multiple tables).+--+-- See Note [Table Unions].+newtype MergingTree s = MergingTree (STRef s (MergingTreeState s))++data MergingTreeState s = CompletedTreeMerge !Run+ -- | Reuses MergingRun (with its STRef) to allow+ -- sharing existing merges.+ | OngoingTreeMerge !(MergingRun TreeMergeType s)+ | PendingTreeMerge !(PendingMerge s)++-- | A merge that is waiting for its inputs to complete.+--+-- The inputs can themselves be 'MergingTree's (with its STRef) to allow sharing+-- existing unions.+data PendingMerge s = -- | The inputs are entire content of a table, i.e. its+ -- (merging) runs and finally a union merge (if that table+ -- already contained a union).+ PendingLevelMerge ![PreExistingRun s] !(Maybe (MergingTree s))+ -- | Each input is a level merge of the entire content of+ -- a table.+ | PendingUnionMerge ![MergingTree s]++-- | This is much like an 'IncomingRun', and are created from them, but contain+-- only the essential information needed in a 'PendingLevelMerge'.+data PreExistingRun s = PreExistingRun !Run+ | PreExistingMergingRun !(MergingRun LevelMergeType s)++pendingContent :: PendingMerge s+ -> (TreeMergeType, [PreExistingRun s], [MergingTree s])+pendingContent = \case+ PendingLevelMerge prs t -> (MergeLevel, prs, toList t)+ PendingUnionMerge ts -> (MergeUnion, [], ts)++{-# COMPLETE PendingMerge #-}+pattern PendingMerge :: TreeMergeType+ -> [PreExistingRun s]+ -> [MergingTree s]+ -> PendingMerge s+pattern PendingMerge mt prs ts <- (pendingContent -> (mt, prs, ts))++type Run = Map Key Entry+type Buffer = Map Key Entry++bufferToRun :: Buffer -> Run+bufferToRun = id++runSize :: Run -> Int+runSize = Map.size++bufferSize :: Buffer -> Int+bufferSize = Map.size++type Entry = Update Value Blob++newtype Key = K Int+ deriving stock (Eq, Ord, Show)+ deriving newtype Enum++newtype Value = V Int+ deriving stock (Eq, Show)++resolveValue :: Value -> Value -> Value+resolveValue (V x) (V y) = V (x + y)++newtype Blob = B Int+ deriving stock (Eq, Show)++-- | We use levelling on the last level, unless that is also the first level.+mergePolicyForLevel :: Int -> [Level s] -> UnionLevel s -> MergePolicyForLevel+mergePolicyForLevel 1 _ _ = LevelTiering+mergePolicyForLevel _ [] NoUnion = LevelLevelling+mergePolicyForLevel _ _ _ = LevelTiering++-- | If there are no further levels provided, this level is the last one.+-- However, if a 'Union' is present, it acts as another (last) level.+mergeTypeForLevel :: [Level s] -> UnionLevel s -> LevelMergeType+mergeTypeForLevel [] NoUnion = MergeLastLevel+mergeTypeForLevel _ _ = MergeMidLevel++-- | Note that the invariants rely on the fact that levelling is only used on+-- the last level.+--+invariant :: forall s. LSMConfig -> LSMContent s -> ST s ()+invariant conf@LSMConfig{..} (LSMContent _ levels ul) = do+ levelsInvariant 1 levels+ case ul of+ NoUnion -> pure ()+ Union tree _ -> expectInvariant (treeInvariant tree)+ where+ levelsInvariant :: Int -> Levels s -> ST s ()+ levelsInvariant !_ [] = pure ()++ levelsInvariant !ln (Level ir rs : ls) = do+ mrs <- case ir of+ Single r ->+ pure (CompletedMerge r)+ Merging mp _ _ (MergingRun mt _ ref) -> do+ assertST $ ln > 1 -- no merges on level 1+ assertST $ mp == mergePolicyForLevel ln ls ul+ assertST $ mt == mergeTypeForLevel ls ul+ readSTRef ref++ assertST $ length rs <= configSizeRatio - 1+ expectedRunLengths ln rs ls+ expectedMergingRunLengths ln ir mrs ls++ levelsInvariant (ln+1) ls++ -- All runs within a level "proper" (as opposed to the incoming runs+ -- being merged) should be of the correct size for the level.+ expectedRunLengths :: Int -> [Run] -> [Level s] -> ST s ()+ expectedRunLengths ln rs ls =+ case mergePolicyForLevel ln ls ul of+ -- Levels using levelling have only one (incoming) run, which almost+ -- always consists of an ongoing merge. The exception is when a+ -- levelling run becomes too large and is promoted, in that case+ -- initially there's no merge, but it is still represented as an+ -- 'IncomingRun', using 'Single'. Thus there are no other resident runs.+ LevelLevelling -> assertST $ null rs && null ls+ -- Runs in tiering levels usually fit that size, but they can be one+ -- larger, if a run has been held back (creating a (T+1)-way merge).+ LevelTiering -> assertST $ all (\r -> runToLevelNumber LevelTiering conf r `elem` [ln, ln+1]) rs+ -- (This is actually still not really true, but will hold in practice.+ -- In the pathological case, all runs passed to the next level can be+ -- factor ((T+1)/T) too large, and there the same holding back can lead to+ -- factor ((T+2)/T) etc., until at level 12 a run is two levels too large.++ -- Incoming runs being merged also need to be of the right size, but the+ -- conditions are more complicated.+ expectedMergingRunLengths :: Int -> IncomingRun s -> MergingRunState+ -> [Level s] -> ST s ()+ expectedMergingRunLengths ln ir mrs ls =+ case mergePolicyForLevel ln ls ul of+ LevelLevelling -> do+ case (ir, mrs) of+ -- A single incoming run (which thus didn't need merging) must be+ -- of the expected size range already+ (Single r, m) -> do+ assertST $ case m of CompletedMerge{} -> True+ OngoingMerge{} -> False+ assertST $ runToLevelNumber LevelLevelling conf r == ln++ -- A completed merge for levelling can be of almost any size at all!+ -- It can be smaller, due to deletions in the last level. But it+ -- can't be bigger than would fit into the next level.+ (_, CompletedMerge r) ->+ assertST $ runToLevelNumber LevelLevelling conf r <= ln+1++ -- An ongoing merge for levelling should have T incoming runs of the+ -- right size for the level below (or slightly larger due to holding+ -- back underfull runs), and at most 1 run from this level. The run+ -- from this level can be of almost any size for the same reasons as+ -- above. Although if this is the first merge for a new level, it'll+ -- have only T runs.+ (_, OngoingMerge _ rs _) -> do+ assertST $ length rs `elem` [configSizeRatio, configSizeRatio + 1]+ assertST $ all (\r -> runSize r > 0) rs -- don't merge empty runs+ let incoming = take configSizeRatio rs+ let resident = drop configSizeRatio rs+ assertST $ all (\r -> runToLevelNumber LevelTiering conf r `elem` [ln-1, ln]) incoming+ assertST $ length resident `elem` [0, 1]+ assertST $ all (\r -> runToLevelNumber LevelLevelling conf r <= ln+1) resident++ LevelTiering ->+ case (ir, mrs, mergeTypeForLevel ls ul) of+ -- A single incoming run (which thus didn't need merging) must be+ -- of the expected size already+ (Single r, m, _) -> do+ assertST $ case m of CompletedMerge{} -> True+ OngoingMerge{} -> False+ assertST $ runToLevelNumber LevelTiering conf r == ln++ -- A completed last level run can be of almost any smaller size due+ -- to deletions, but it can't be bigger than the next level down.+ -- Note that tiering on the last level only occurs when there is+ -- a single level only.+ (_, CompletedMerge r, MergeLastLevel) -> do+ assertST $ ln == 1+ assertST $ runToLevelNumber LevelTiering conf r <= ln+1++ -- A completed mid level run is usually of the size for the+ -- level it is entering, but can also be one smaller (in which case+ -- it'll be held back and merged again) or one larger (because it+ -- includes a run that has been held back before).+ (_, CompletedMerge r, MergeMidLevel) ->+ assertST $ runToLevelNumber LevelTiering conf r `elem` [ln-1, ln, ln+1]++ -- An ongoing merge for tiering should have T incoming runs of the+ -- right size for the level below (or slightly larger due to holding+ -- back underfull runs), and at most 1 run held back due to being+ -- too small (which would thus also be of the size of the level+ -- below).+ (_, OngoingMerge _ rs _, _) -> do+ assertST $ length rs `elem` [configSizeRatio, configSizeRatio + 1]+ assertST $ all (\r -> runSize r > 0) rs -- don't merge empty runs+ let incoming = take configSizeRatio rs+ let heldBack = drop configSizeRatio rs+ assertST $ all (\r -> runToLevelNumber LevelTiering conf r `elem` [ln-1, ln]) incoming+ assertST $ length heldBack `elem` [0, 1]+ assertST $ all (\r -> runToLevelNumber LevelTiering conf r == ln-1) heldBack++-- We don't make many assumptions apart from what the types already enforce.+-- In particular, there are no invariants on the progress of the merges,+-- since union merge credits are independent from the tables' regular level+-- merges.+treeInvariant :: MergingTree s -> Invariant s ()+treeInvariant tree@(MergingTree treeState) = do+ liftI (readSTRef treeState) >>= \case+ CompletedTreeMerge _ ->+ -- We don't require the completed merges to be non-empty, since even+ -- a (last-level) merge of non-empty runs can end up being empty.+ -- In the prototype it would be possible to ensure that empty runs are+ -- immediately trimmed from the tree, but this kind of normalisation+ -- is complicated with sharing. For example, merging runs and+ -- trees are shared, so if one of them completes as an empty run,+ -- all tables referencing it suddenly contain an empty run and would+ -- need to be updated immediately.+ pure ()++ OngoingTreeMerge mr ->+ mergeInvariant mr++ PendingTreeMerge (PendingLevelMerge prs t) -> do+ -- Non-empty, but can be just one input (see 'newPendingLevelMerge').+ -- Note that children of a pending merge can be empty runs, as noted+ -- above for 'CompletedTreeMerge'.+ assertI "pending level merges have at least one input" $+ length prs + length t > 0+ for_ prs $ \case+ PreExistingRun _r -> pure ()+ PreExistingMergingRun mr -> mergeInvariant mr+ for_ t treeInvariant++ PendingTreeMerge (PendingUnionMerge ts) -> do+ assertI "pending union merges are non-trivial (at least two inputs)" $+ length ts > 1+ for_ ts treeInvariant++ (debt, _) <- liftI $ remainingDebtMergingTree tree+ when (debt <= 0) $ do+ _ <- isCompletedMergingTree tree+ pure ()++mergeInvariant :: MergingRun t s -> Invariant s ()+mergeInvariant (MergingRun _ mergeDebt ref) =+ liftI (readSTRef ref) >>= \case+ CompletedMerge _ -> pure ()+ OngoingMerge mergeCredit rs _ -> do+ assertI "merge debt & credit invariant" $+ mergeDebtInvariant mergeDebt mergeCredit+ assertI "inputs to ongoing merges aren't empty" $+ all (\r -> runSize r > 0) rs+ assertI "ongoing merges are non-trivial (at least two inputs)" $+ length rs > 1++isCompletedMergingRun :: MergingRun t s -> Invariant s Run+isCompletedMergingRun (MergingRun _ d ref) = do+ mrs <- liftI $ readSTRef ref+ case mrs of+ CompletedMerge r -> pure r+ OngoingMerge c _ _ -> failI $ "not completed: OngoingMerge with"+ ++ " remaining debt "+ ++ show (mergeDebtLeft d c)++isCompletedMergingTree :: MergingTree s -> Invariant s Run+isCompletedMergingTree (MergingTree ref) = do+ mts <- liftI $ readSTRef ref+ case mts of+ CompletedTreeMerge r -> pure r+ OngoingTreeMerge mr -> isCompletedMergingRun mr+ PendingTreeMerge _ -> failI $ "not completed: PendingTreeMerge"++type Invariant s = E.ExceptT String (ST s)++assertI :: String -> Bool -> Invariant s ()+assertI _ True = pure ()+assertI e False = failI e++failI :: String -> Invariant s a+failI = E.throwE++liftI :: ST s a -> Invariant s a+liftI = E.ExceptT . fmap Right++expectInvariant :: HasCallStack => Invariant s a -> ST s a+expectInvariant act = E.runExceptT act >>= either error pure++evalInvariant :: Invariant s a -> ST s (Either String a)+evalInvariant = E.runExceptT++-- 'callStack' just ensures that the 'HasCallStack' constraint is not redundant+-- when compiling with debug assertions disabled.+assert :: HasCallStack => Bool -> a -> a+assert p x = Exc.assert p (const x callStack)++assertST :: HasCallStack => Bool -> ST s ()+assertST p = assert p $ pure ()++assertWithMsg :: HasCallStack => Maybe String -> a -> a+assertWithMsg = assert . p+ where+ p Nothing = True+ p (Just msg) = error $ "Assertion failed: " <> msg++assertWithMsgM :: (HasCallStack, Monad m) => Maybe String -> m ()+assertWithMsgM mmsg = assertWithMsg mmsg $ pure ()++leq :: (Show a, Ord a) => a -> a -> Maybe String+leq x y = if x <= y then Nothing else Just $+ printf "Expected x <= y, but got %s > %s"+ (show x)+ (show y)++-------------------------------------------------------------------------------+-- Run sizes+--++-- | Compute the maximum size of a run for a given level.+--+-- The size of a tiering run at each level is allowed to be+-- @bufferSize*sizeRatio^(level-1) < size <= bufferSize*sizeRatio^level@.+--+-- >>> levelNumberToMaxRunSize LevelTiering (LSMConfig 2 4) <$> [0, 1, 2, 3, 4]+-- [0,2,8,32,128]+--+-- The @size@ of a levelling run at each level is allowed to be+-- @bufferSize*sizeRatio^level < size <= bufferSize*sizeRatio^(level+1)@. A+-- levelling run can take take up a whole level, so the maximum size of a run is+-- @sizeRatio@ tmes larger than the maximum size of a tiering run on the same+-- level.+--+-- >>> levelNumberToMaxRunSize LevelLevelling (LSMConfig 2 4) <$> [0, 1, 2, 3, 4]+-- [0,8,32,128,512]+levelNumberToMaxRunSize :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Int+levelNumberToMaxRunSize = \case+ LevelTiering -> levelNumberToMaxRunSizeTiering+ LevelLevelling -> levelNumberToMaxRunSizeLevelling++-- | See 'levelNumberToMaxRunSize'+levelNumberToMaxRunSizeTiering :: HasCallStack => LSMConfig -> LevelNo -> Int+levelNumberToMaxRunSizeTiering+ LSMConfig {configMaxWriteBufferSize = bufSize, configSizeRatio = sizeRatio}+ ln+ | ln < 0 = error "level number must be non-negative"+ | ln == 0 = 0+ | otherwise = fromIntegerChecked (toInteger bufSize * toInteger sizeRatio ^ pred (toInteger ln))+ -- Perform the computation with arbitrary precision using 'Integers', but+ -- throw an error if the result does not fit into an 'Int'.++-- | See 'levelNumberToMaxRunSize'+levelNumberToMaxRunSizeLevelling :: HasCallStack => LSMConfig -> LevelNo -> Int+levelNumberToMaxRunSizeLevelling conf ln+ | ln < 0 = error "level number must be non-negative"+ | ln == 0 = 0+ | otherwise = levelNumberToMaxRunSizeTiering conf (succ ln)++-- | See 'runSizeToLevelNumber'.+runToLevelNumber :: HasCallStack => MergePolicyForLevel -> LSMConfig -> Run -> LevelNo+runToLevelNumber mpl conf run = runSizeToLevelNumber mpl conf (runSize run)++-- | Compute the appropriate level for the size of the given run.+--+-- See 'levelNumberToMaxRunSize' for the bounds on (tiering or levelling) run+-- sizes at each level.+--+-- >>> runSizeToLevelNumber LevelTiering (LSMConfig 2 4) <$> [0,2,8,32,128]+-- [0,1,2,3,4]+--+-- >>> runSizeToLevelNumber LevelLevelling (LSMConfig 2 4) <$> [0,8,32,128,512]+-- [0,1,2,3,4]+runSizeToLevelNumber :: HasCallStack => MergePolicyForLevel -> LSMConfig -> Int -> LevelNo+runSizeToLevelNumber = \case+ LevelTiering -> runSizeToLevelNumberTiering+ LevelLevelling -> runSizeToLevelNumberLevelling++-- | See 'runSizeToLevelNumber'.+runSizeToLevelNumberTiering :: HasCallStack => LSMConfig -> Int -> LevelNo+runSizeToLevelNumberTiering conf n+ | n < 0 = error "run size must be positive"+ -- TODO: enumerating level numbers is potentially costly, but it does gives a+ -- precise answer, where we'd otherwise have to deal with Double rounding+ -- errors in computing @ln = logBase configSizeRatio (n / configMaxWriteBufferSize) + 1@+ | otherwise = head $ -- the list is guaranteed to be non-empty+ [ ln+ | ln <- [0..]+ , n <= levelNumberToMaxRunSizeTiering conf ln+ ]++-- | See 'runSizeToLevelNumber'.+runSizeToLevelNumberLevelling :: HasCallStack => LSMConfig -> Int -> LevelNo+runSizeToLevelNumberLevelling conf n+ | n < 0 = error "run size must be positive"+ -- TODO: enumerating level numbers is potentially costly, but it does gives a+ -- precise answer, where we'd otherwise have to deal with Double rounding+ -- errors in computing @ln = logBase configSizeRatio (n / configMaxWriteBufferSize)@+ | otherwise = head $ -- the list is guaranteed to be non-empty+ [ ln+ | ln <- [0..]+ , n <= levelNumberToMaxRunSizeLevelling conf ln+ ]++maxWriteBufferSize :: HasCallStack => LSMConfig -> Int+maxWriteBufferSize conf = levelNumberToMaxRunSizeTiering conf 1 -- equal to configMaxWriteBufferSize++{-# INLINABLE fromIntegerChecked #-}+-- | Like 'fromInteger', but throws an error when @(x :: Integer) /= toInteger+-- (fromInteger x :: b)@.+fromIntegerChecked :: (HasCallStack, Integral a) => Integer -> a+fromIntegerChecked x+ | x'' == x+ = x'+ | otherwise+ = error $ printf "fromIntegerChecked: conversion failed, %s /= %s" (show x) (show x'')+ where+ x' = fromInteger x+ x'' = toInteger x'++-- | See 'runSizeFitsInLevel'.+_runFitsInLevel :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Run -> Bool+_runFitsInLevel mpl conf ln r = runSizeFitsInLevel mpl conf ln (runSize r)++-- | Check wheter a run of the given size fits in the given level.+--+-- See 'levelNumberToMaxRunSize' for the bounds on (tiering or levelling) run+-- sizes at each level.+--+-- >>> runSizeFitsInLevel LevelTiering (LSMConfig 2 4) 3 <$> [8,9,16,32,33]+-- [False,True,True,True,False]+--+-- >>> runSizeFitsInLevel LevelLevelling (LSMConfig 2 4) 2 <$> [8,9,16,32,33]+-- [False,True,True,True,False]+runSizeFitsInLevel :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Int -> Bool+runSizeFitsInLevel mpl conf ln n+ | ln < 0 = error "level number must be non-negative"+ | ln == 0 = n == 0+ | otherwise =+ levelNumberToMaxRunSize mpl conf (pred ln) < n+ && n <= levelNumberToMaxRunSize mpl conf ln++-- | See 'runSizeTooSmallForLevel'.+runTooSmallForLevel :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Run -> Bool+runTooSmallForLevel mpl conf ln r = runSizeTooSmallForLevel mpl conf ln (runSize r)++-- | Check wheter a run of the given size is too small for the given level.+--+-- See 'levelNumberToMaxRunSize' for the bounds on (tiering or levelling) run+-- sizes at each level.+--+-- >>> runSizeTooSmallForLevel LevelTiering (LSMConfig 2 4) 3 <$> [8,9]+-- [True,False]+--+-- >>> runSizeTooSmallForLevel LevelLevelling (LSMConfig 2 4) 2 <$> [8,9]+-- [True,False]+runSizeTooSmallForLevel :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Int -> Bool+runSizeTooSmallForLevel mpl conf ln n+ | ln < 0 = error "level number must be non-negative"+ | ln == 0 = False+ | otherwise = case mpl of+ LevelTiering ->+ n <= levelNumberToMaxRunSize LevelTiering conf (pred ln)+ LevelLevelling ->+ n <= levelNumberToMaxRunSize LevelLevelling conf (pred ln)++-- | See 'runSizeTooLargeForLevel'.+runTooLargeForLevel :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Run -> Bool+runTooLargeForLevel mpl conf ln r = runSizeTooLargeForLevel mpl conf ln (runSize r)++-- | Check wheter a run of the given size is too large for the given level.+--+-- See 'levelNumberToMaxRunSize' for the bounds on (tiering or levelling) run+-- sizes at each level.+--+-- >>> runSizeTooLargeForLevel LevelTiering (LSMConfig 2 4) 2 <$> [8,9]+-- [False,True]+--+-- >>> runSizeTooLargeForLevel LevelLevelling (LSMConfig 2 4) 1 <$> [8,9]+-- [False,True]+runSizeTooLargeForLevel :: HasCallStack => MergePolicyForLevel -> LSMConfig -> LevelNo -> Int -> Bool+runSizeTooLargeForLevel mpl conf ln n+ | ln < 0 = error "level number must be non-negative"+ | ln == 0 = not (n == 0)+ | otherwise = case mpl of+ LevelTiering ->+ n > levelNumberToMaxRunSize LevelTiering conf ln+ LevelLevelling ->+ n > levelNumberToMaxRunSize LevelLevelling conf ln++-------------------------------------------------------------------------------+-- Level capacity+--++levelIsFull :: MergePolicyForLevel -> LSMConfig -> LevelNo -> [Run] -> [Run] -> Bool+levelIsFull mpl conf ln incoming resident = case mpl of+ LevelTiering -> levelIsFullTiering conf ln incoming resident+ LevelLevelling ->+ assert (length resident == 1) $+ levelIsFullLevelling conf ln incoming (head resident)++-- | Only based on run count, not their sizes.+levelIsFullTiering :: LSMConfig -> LevelNo -> [Run] -> [Run] -> Bool+levelIsFullTiering LSMConfig{..} _ln _incoming resident =+ length resident >= configSizeRatio++-- | The level is only considered full once the resident run is /too large/+-- for the level.+levelIsFullLevelling :: LSMConfig -> LevelNo -> [Run] -> Run -> Bool+levelIsFullLevelling conf ln _incoming resident =+ runTooLargeForLevel LevelLevelling conf ln resident++-------------------------------------------------------------------------------+-- Merging credits+--++-- | Credits for keeping track of merge progress. These credits correspond+-- directly to merge steps performed.+--+-- We also call these \"physical\" credits (since they correspond to steps+-- done), and as opposed to \"nominal\" credits in 'NominalCredit' and+-- 'NominalDebt'.+type Credit = Int++-- | Debt for keeping track of the total merge work to do.+type Debt = Int++data MergeCredit =+ MergeCredit {+ spentCredits :: !Credit, -- accumulating+ unspentCredits :: !Credit -- fluctuating+ }+ deriving stock Show++newtype MergeDebt =+ MergeDebt {+ totalDebt :: Debt -- fixed+ }+ deriving stock Show++zeroMergeCredit :: MergeCredit+zeroMergeCredit =+ MergeCredit {+ spentCredits = 0,+ unspentCredits = 0+ }++mergeDebtInvariant :: MergeDebt -> MergeCredit -> Bool+mergeDebtInvariant MergeDebt {totalDebt}+ MergeCredit {spentCredits, unspentCredits} =+ let suppliedCredits = spentCredits + unspentCredits+ in spentCredits >= 0+ -- unspentCredits could legitimately be negative, though that does not+ -- happen in this prototype+ && suppliedCredits >= 0+ && suppliedCredits <= totalDebt++mergeDebtLeft :: HasCallStack => MergeDebt -> MergeCredit -> Debt+mergeDebtLeft MergeDebt {totalDebt}+ MergeCredit {spentCredits, unspentCredits} =+ let suppliedCredits = spentCredits + unspentCredits+ in assert (suppliedCredits <= totalDebt)+ (totalDebt - suppliedCredits)++-- | As credits are paid, debt is reduced in batches when sufficient credits+-- have accumulated.+data MergeDebtPaydown =+ -- | This remaining merge debt is fully paid off, potentially with+ -- leftovers.+ MergeDebtDischarged !Debt !Credit++ -- | Credits were paid, but not enough for merge debt to be reduced by some+ -- batches of merging work.+ | MergeDebtPaydownCredited !MergeCredit++ -- | Enough credits were paid to reduce merge debt by performing some+ -- batches of merging work.+ | MergeDebtPaydownPerform !Debt !MergeCredit+ deriving stock Show++-- | Pay credits to merge debt, which might trigger performing some merge work+-- in batches. See 'MergeDebtPaydown'.+--+paydownMergeDebt :: MergeDebt -> MergeCredit -> Credit -> MergeDebtPaydown+paydownMergeDebt MergeDebt {totalDebt}+ MergeCredit {spentCredits, unspentCredits}+ c+ | suppliedCredits' >= totalDebt+ , let !leftover = suppliedCredits' - totalDebt+ !perform = c - leftover+ = assert (dischargePostcondition perform leftover) $+ MergeDebtDischarged perform leftover++ | unspentCredits' >= mergeBatchSize+ , let (!b, !r) = divMod unspentCredits' mergeBatchSize+ !perform = b * mergeBatchSize+ = assert (performPostcondition perform r) $+ MergeDebtPaydownPerform+ perform+ MergeCredit {+ spentCredits = spentCredits + perform,+ unspentCredits = unspentCredits' - perform+ }++ | otherwise+ = assert creditedPostcondition $+ MergeDebtPaydownCredited+ MergeCredit {+ spentCredits,+ unspentCredits = unspentCredits'+ }+ where+ suppliedCredits' = spentCredits + unspentCredits + c+ unspentCredits' = unspentCredits + c++ dischargePostcondition perform leftover =+ (c >= 0)+ && (perform >= 0 && leftover >= 0)+ && (c == perform + leftover)+ && (spentCredits + unspentCredits + perform == totalDebt)++ performPostcondition perform r =+ let spentCredits' = spentCredits + perform+ unspentCredits'' = unspentCredits' - perform+ in (c >= 0)+ && (unspentCredits'' == r)+ && (suppliedCredits' == spentCredits' + unspentCredits'')+ && (suppliedCredits' < totalDebt)++ creditedPostcondition =+ (c >= 0)+ && (suppliedCredits' < totalDebt)++mergeBatchSize :: Int+mergeBatchSize = 32+++-------------------------------------------------------------------------------+-- Merging run abstraction+--++newMergingRun :: IsMergeType t => t -> [Run] -> ST s (MergingRun t s)+newMergingRun mergeType runs = do+ assertST $ length runs > 1+ -- in some cases, no merging is required at all+ (debt, state) <- case filter (\r -> runSize r > 0) runs of+ [] -> let (r:_) = runs -- just reuse the empty input+ in pure (runSize r, CompletedMerge r)+ [r] -> pure (runSize r, CompletedMerge r)+ rs -> do+ -- The (physical) debt is always exactly the cost (merge steps),+ -- which is the sum of run lengths in elements.+ let !debt = sum (map runSize rs)+ let merged = mergek mergeType rs -- deliberately lazy+ pure (debt, OngoingMerge zeroMergeCredit rs merged)+ MergingRun mergeType (MergeDebt debt) <$> newSTRef state++mergek :: IsMergeType t => t -> [Run] -> Run+mergek t =+ (if isLastLevel t then Map.filter (/= Delete) else id)+ . Map.unionsWith (if isUnion t then combineUnion else combine)++-- | Combines two entries that have been performed after another. Therefore, the+-- newer one overwrites the old one (or modifies it for 'Mupsert'). Only take a+-- blob from the left entry.+combine :: Entry -> Entry -> Entry+combine new_ old = case new_ of+ Insert{} -> new_+ Delete{} -> new_+ Mupsert v -> case old of+ Insert v' _ -> Insert (resolveValue v v') Nothing+ Delete -> Insert v Nothing+ Mupsert v' -> Mupsert (resolveValue v v')++-- | Combines two entries of runs that have been 'union'ed together. If any one+-- has a value, the result should have a value (represented by 'Insert'). If+-- both have a value, these values get combined monoidally. Only take a blob+-- from the left entry.+--+-- See 'MergeUnion'.+combineUnion :: Entry -> Entry -> Entry+combineUnion Delete (Mupsert v) = Insert v Nothing+combineUnion Delete old = old+combineUnion (Mupsert u) Delete = Insert u Nothing+combineUnion new_ Delete = new_+combineUnion (Mupsert v') (Mupsert v ) = Insert (resolveValue v' v) Nothing+combineUnion (Mupsert v') (Insert v _) = Insert (resolveValue v' v) Nothing+combineUnion (Insert v' b') (Mupsert v) = Insert (resolveValue v' v) b'+combineUnion (Insert v' b') (Insert v _) = Insert (resolveValue v' v) b'++expectCompletedMergingRun :: HasCallStack => MergingRun t s -> ST s Run+expectCompletedMergingRun = expectInvariant . isCompletedMergingRun++supplyCreditsMergingRun :: Credit -> MergingRun t s -> ST s Credit+supplyCreditsMergingRun =+ checked remainingDebtMergingRun $ \credits (MergingRun _ mergeDebt ref) -> do+ mrs <- readSTRef ref+ case mrs of+ CompletedMerge{} -> pure credits+ OngoingMerge mergeCredit rs r ->+ case paydownMergeDebt mergeDebt mergeCredit credits of+ MergeDebtDischarged _ leftover -> do+ writeSTRef ref (CompletedMerge r)+ pure leftover++ MergeDebtPaydownCredited mergeCredit' -> do+ writeSTRef ref (OngoingMerge mergeCredit' rs r)+ pure 0++ MergeDebtPaydownPerform _mergeSteps mergeCredit' -> do+ -- we're not doing any actual merging+ -- just tracking what we would do+ writeSTRef ref (OngoingMerge mergeCredit' rs r)+ pure 0++suppliedCreditMergingRun :: MergingRun t s -> ST s Credit+suppliedCreditMergingRun (MergingRun _ d ref) =+ readSTRef ref >>= \case+ CompletedMerge{} ->+ let MergeDebt { totalDebt } = d in+ pure totalDebt+ OngoingMerge MergeCredit {spentCredits, unspentCredits} _ _ ->+ pure (spentCredits + unspentCredits)++-------------------------------------------------------------------------------+-- LSM handle+--++new :: Tracer (ST s) Event -> TableId -> ST s (LSM s)+new tr tid = newWith tr tid conf+ where+ -- 4 was the default for both the max write buffer size and size ratio+ -- before they were made configurable+ conf = LSMConfig {+ configMaxWriteBufferSize = 4+ , configSizeRatio = 4+ }++newWith :: Tracer (ST s) Event -> TableId -> LSMConfig -> ST s (LSM s)+newWith tr tid conf+ | configMaxWriteBufferSize conf <= 0 =+ error "newWith: configMaxWriteBufferSize should be positive"+ | configSizeRatio conf <= 1 =+ error "newWith: configSizeRatio should be larger than 1"+ | otherwise = do+ traceWith tr $ NewTableEvent tid conf+ c <- newSTRef 0+ lsm <- newSTRef (LSMContent Map.empty [] NoUnion)+ pure (LSMHandle tid c conf lsm)++inserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value, Maybe Blob)] -> ST s ()+inserts tr lsm kvbs = updates tr lsm [ (k, Insert v b) | (k, v, b) <- kvbs ]++insert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> Maybe Blob -> ST s ()+insert tr lsm k v b = update tr lsm k (Insert v b)++deletes :: Tracer (ST s) Event -> LSM s -> [Key] -> ST s ()+deletes tr lsm ks = updates tr lsm [ (k, Delete) | k <- ks ]++delete :: Tracer (ST s) Event -> LSM s -> Key -> ST s ()+delete tr lsm k = update tr lsm k Delete++mupserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value)] -> ST s ()+mupserts tr lsm kvbs = updates tr lsm [ (k, Mupsert v) | (k, v) <- kvbs ]++mupsert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> ST s ()+mupsert tr lsm k v = update tr lsm k (Mupsert v)++data Update v b =+ Insert !v !(Maybe b)+ | Mupsert !v+ | Delete+ deriving stock (Eq, Show)++updates :: Tracer (ST s) Event -> LSM s -> [(Key, Entry)] -> ST s ()+updates tr lsm = mapM_ (uncurry (update tr lsm))++update :: Tracer (ST s) Event -> LSM s -> Key -> Entry -> ST s ()+update tr (LSMHandle tid scr conf lsmr) k entry = do+ traceWith tr $ UpdateEvent tid k entry+ sc <- readSTRef scr+ content@(LSMContent wb ls unionLevel) <- readSTRef lsmr+ modifySTRef' scr (+1)+ supplyCreditsLevels (NominalCredit 1) ls+ invariant conf content+ let wb' = Map.insertWith combine k entry wb+ if bufferSize wb' >= maxWriteBufferSize conf+ then do+ ls' <- increment (LevelEvent tid >$< tr) sc conf (bufferToRun wb') ls unionLevel+ let content' = LSMContent Map.empty ls' unionLevel+ invariant conf content'+ writeSTRef lsmr content'+ else+ writeSTRef lsmr (LSMContent wb' ls unionLevel)++supplyMergeCredits :: LSM s -> NominalCredit -> ST s ()+supplyMergeCredits (LSMHandle _ scr conf lsmr) credits = do+ content@(LSMContent _ ls _) <- readSTRef lsmr+ modifySTRef' scr (+1)+ supplyCreditsLevels credits ls+ invariant conf content++data LookupResult v b =+ NotFound+ | Found !v !(Maybe b)+ deriving stock (Eq, Show)++lookups :: LSM s -> [Key] -> ST s [LookupResult Value Blob]+lookups (LSMHandle _ _ _conf lsmr) ks = do+ LSMContent wb ls ul <- readSTRef lsmr+ runs <- concat <$> flattenLevels ls+ traverse (doLookup wb runs ul) ks++lookup :: Tracer (ST s) Event -> LSM s -> Key -> ST s (LookupResult Value Blob)+lookup tr (LSMHandle tid _ _conf lsmr) k = do+ traceWith tr $ LookupEvent tid k+ LSMContent wb ls ul <- readSTRef lsmr+ runs <- concat <$> flattenLevels ls+ doLookup wb runs ul k++duplicate :: Tracer (ST s) Event -> TableId -> LSM s -> ST s (LSM s)+duplicate tr childTid (LSMHandle parentTid _scr conf lsmr) = do+ traceWith tr $ DuplicateEvent childTid parentTid+ scr' <- newSTRef 0+ lsmr' <- newSTRef =<< readSTRef lsmr+ pure (LSMHandle childTid scr' conf lsmr')+ -- it's that simple here, because we share all the pure value and all the+ -- STRefs and there's no ref counting to be done++-- | Similar to @Data.Map.unionWith@.+--+-- A call to 'union' itself is not expensive, as the input tables are not+-- immediately merged. Instead, it creates a representation of an in-progress+-- merge that can be performed incrementally (somewhat similar to a thunk).+--+-- The more merge work remains, the more expensive are lookups on the table.+unions :: Tracer (ST s) Event -> TableId -> [LSM s] -> ST s (LSM s)+unions tr childTid lsms = do+ traceWith tr $+ let parentTids = fmap tableId lsms+ in UnionsEvent childTid parentTids+ (confs, trees) <- fmap unzip $ forM lsms $ \(LSMHandle _ _ conf lsmr) ->+ (conf,) <$> (contentToMergingTree =<< readSTRef lsmr)+ -- Check that the configurations are equal+ conf <- case confs of+ [] -> error "unions: 0 tables"+ conf : _ -> assert (all (conf==) confs) $ pure conf+ -- TODO: if only one table is non-empty, we don't have to create a Union,+ -- we can just duplicate the table.+ unionLevel <- newPendingUnionMerge (catMaybes trees) >>= \case+ Nothing -> pure NoUnion+ Just tree -> do+ debt <- fst <$> remainingDebtMergingTree tree+ Union tree <$> newSTRef debt+ lsmr <- newSTRef (LSMContent Map.empty [] unionLevel)+ c <- newSTRef 0+ pure (LSMHandle childTid c conf lsmr)++-- | The /current/ upper bound on the number of 'UnionCredits' that have to be+-- supplied before a 'union' is completed.+--+-- The union debt is the number of merging steps that need to be performed /at+-- most/ until the delayed work of performing a 'union' is completed. This+-- includes the cost of completing merges that were part of the union's input+-- tables.+newtype UnionDebt = UnionDebt Debt+ deriving stock (Show, Eq, Ord)+ deriving newtype Num++-- | Return the current union debt. This debt can be reduced until it is paid+-- off using 'supplyUnionCredits'.+remainingUnionDebt :: LSM s -> ST s UnionDebt+remainingUnionDebt (LSMHandle _ _ _conf lsmr) = do+ LSMContent _ _ ul <- readSTRef lsmr+ UnionDebt <$> case ul of+ NoUnion -> pure 0+ Union tree d -> checkedUnionDebt tree d++-- | Credits are used to pay off 'UnionDebt', completing a 'union' in the+-- process.+--+-- A union credit corresponds to a single merging step being performed.+newtype UnionCredits = UnionCredits Credit+ deriving stock (Show, Eq, Ord)+ deriving newtype Num++-- | Supply union credits to reduce union debt.+--+-- Supplying union credits leads to union merging work being performed in+-- batches. This reduces the union debt returned by 'remainingUnionDebt'. Union+-- debt will be reduced by /at least/ the number of supplied union credits. It+-- is therefore advisable to query 'remainingUnionDebt' every once in a while to+-- see what the current debt is.+--+-- This function returns any surplus of union credits as /leftover/ credits when+-- a union has finished. In particular, if the returned number of credits is+-- non-negative, then the union is finished.+supplyUnionCredits :: LSM s -> UnionCredits -> ST s UnionCredits+supplyUnionCredits (LSMHandle _ scr conf lsmr) (UnionCredits credits)+ | credits <= 0 = pure (UnionCredits 0)+ | otherwise = do+ content@(LSMContent _ _ ul) <- readSTRef lsmr+ UnionCredits <$> case ul of+ NoUnion ->+ pure credits+ Union tree debtRef -> do+ modifySTRef' scr (+1)+ _debt <- checkedUnionDebt tree debtRef -- just to make sure it's checked+ c' <- supplyCreditsMergingTree credits tree+ debt' <- checkedUnionDebt tree debtRef+ when (debt' > 0) $+ assertST $ c' == 0 -- should have spent these credits+ invariant conf content+ pure c'++-- TODO: At some point the completed merging tree should to moved into the+-- regular levels, so it can be merged with other runs and last level merges can+-- happen again to drop deletes. Also, lookups then don't need to handle the+-- merging tree any more. There are two possible strategies:+--+-- 1. As soon as the merging tree completes, move the resulting run to the+-- regular levels. However, its size does generally not fit the last level,+-- which requires relaxing 'invariant' and adjusting 'increment'.+--+-- If the run is much larger than the resident and incoming runs of the last+-- level, it should also not be included into a merge yet, as that merge+-- would be expensive, but offer very little potential for compaction (the+-- run from the merging tree is already compacted after all). So it needs to+-- be bumped to the next level instead.+--+-- 2. Initially leave the completed run in the union level. Then every time a+-- new last level merge is created in 'increment', check if there is a+-- completed run in the union level with a size that fits the new merge. If+-- yes, move it over.++-- | Like 'remainingDebtMergingTree', but additionally asserts that the debt+-- never increases.+checkedUnionDebt :: MergingTree s -> STRef s Debt -> ST s Debt+checkedUnionDebt tree debtRef = do+ storedDebt <- readSTRef debtRef+ debt <- fst <$> remainingDebtMergingTree tree+ assertST $ debt <= storedDebt+ writeSTRef debtRef debt+ pure debt++-------------------------------------------------------------------------------+-- Lookups+--++type LookupAcc = Maybe Entry++updateAcc :: (Entry -> Entry -> Entry) -> LookupAcc -> Entry -> LookupAcc+updateAcc _ Nothing old = Just old+updateAcc f (Just new_) old = Just (f new_ old) -- acc has more recent Entry++mergeAcc :: TreeMergeType -> [LookupAcc] -> LookupAcc+mergeAcc mt = foldl (updateAcc com) Nothing . catMaybes+ where+ com = case mt of+ MergeLevel -> combine+ MergeUnion -> combineUnion++-- | We handle lookups by accumulating results by going through the runs from+-- most recent to least recent, starting with the write buffer.+--+-- In the real implementation, this is done not on an individual 'LookupAcc',+-- but one for each key, i.e. @Vector (Maybe Entry)@.+doLookup :: Buffer -> [Run] -> UnionLevel s -> Key -> ST s (LookupResult Value Blob)+doLookup wb runs ul k = do+ let acc0 = lookupBatch (Map.lookup k wb) k runs+ case ul of+ NoUnion ->+ pure (convertAcc acc0)+ Union tree _ -> do+ treeBatches <- buildLookupTree tree+ let treeResults = lookupBatch Nothing k <$> treeBatches+ pure $ convertAcc $ foldLookupTree $+ if null wb && null runs+ then treeResults+ else LookupNode MergeLevel [LookupBatch acc0, treeResults ]+ where+ convertAcc :: LookupAcc -> LookupResult Value Blob+ convertAcc = \case+ Nothing -> NotFound+ Just (Insert v b) -> Found v b+ Just (Mupsert v) -> Found v Nothing+ Just Delete -> NotFound++-- | Perform a batch of lookups, accumulating the result onto an initial+-- 'LookupAcc'.+--+-- In a real implementation, this would take all keys at once and be in IO.+lookupBatch :: LookupAcc -> Key -> [Run] -> LookupAcc+lookupBatch acc k rs =+ let entries = [entry | r <- rs, Just entry <- [Map.lookup k r]]+ in foldl (updateAcc combine) acc entries++data LookupTree a = LookupBatch a+ | LookupNode TreeMergeType [LookupTree a]+ deriving stock Functor++-- | Do lookups on runs at the leaves and recursively combine the resulting+-- 'LookupAcc's, either using 'mergeAcc' or 'unionAcc' depending on the merge+-- type.+--+-- Doing this naively would result in a call to 'lookupBatch' and creation of+-- a 'LookupAcc' for each run in the tree. However, when there are adjacent+-- 'Run's or 'MergingRuns' (with 'MergeLevel') as inputs to a level-merge, we+-- combine them into a single batch of runs.+--+-- For example, this means that if we union two tables (which themselves don't+-- have a union level) and then do lookups, two batches of lookups have to be+-- performed (plus a batch for the table's regular levels if it has been updated+-- after the union).+--+-- TODO: we can still improve the batching, for example combining the child of+-- PendingLevelMerge with the pre-existing runs when it is already completed.+buildLookupTree :: MergingTree s -> ST s (LookupTree [Run])+buildLookupTree = go+ where+ go :: MergingTree s -> ST s (LookupTree [Run])+ go (MergingTree treeState) = readSTRef treeState >>= \case+ CompletedTreeMerge r ->+ pure $ LookupBatch [r]+ OngoingTreeMerge (MergingRun mt _ mergeState) ->+ readSTRef mergeState >>= \case+ CompletedMerge r ->+ pure $ LookupBatch [r]+ OngoingMerge _ rs _ -> case mt of+ MergeLevel -> pure $ LookupBatch rs -- combine into batch+ MergeUnion -> pure $ LookupNode MergeUnion $ map (\r -> LookupBatch [r]) rs+ PendingTreeMerge (PendingLevelMerge prs tree) -> do+ preExisting <- LookupBatch . concat <$>+ traverse flattenPreExistingRun prs -- combine into batch+ case tree of+ Nothing -> pure preExisting+ Just t -> do+ lTree <- go t+ pure (LookupNode MergeLevel [preExisting, lTree])+ PendingTreeMerge (PendingUnionMerge trees) -> do+ LookupNode MergeUnion <$> traverse go trees++foldLookupTree :: LookupTree LookupAcc -> LookupAcc+foldLookupTree = \case+ LookupBatch acc -> acc+ LookupNode mt children -> mergeAcc mt (map foldLookupTree children)++-------------------------------------------------------------------------------+-- Nominal credits+--++-- | Nominal credit is the credit supplied to each level as we insert update+-- entries, one credit per update entry inserted.+--+-- Nominal credit must be supplied up to the 'NominalDebt' to ensure the merge+-- is complete.+--+-- Nominal credits are a similar order of magnitude to physical credits (see+-- 'Credit') but not the same, and we have to scale linearly to convert between+-- them. Physical credits are the actual number of inputs to the merge, which+-- may be somewhat more or somewhat less than the number of update entries+-- we will insert before we need the merge to be complete.+--+newtype NominalCredit = NominalCredit Credit+ deriving stock Show++-- | The nominal debt for a merging run is the worst case (minimum) number of+-- update entries we expect to insert before we expect the merge to be+-- complete.+--+-- We require that an equal amount of nominal credit is supplied before we can+-- expect a merge to be complete.+--+-- We scale linearly to convert nominal credits to physical credits, such that+-- the nominal debt and physical debt are both considered \"100%\", and so that+-- both debts are paid off at exactly the same time.+--+newtype NominalDebt = NominalDebt Credit+ deriving stock Show++-- TODO: If there is a UnionLevel, there is no (more expensive) last level merge+-- in the regular levels, so a little less merging work is required than if+-- there was no UnionLevel. It might be a good idea to spend this "saved" work+-- on the UnionLevel instead. This makes future lookups cheaper and ensures that+-- we can get rid of the UnionLevel at some point, even if a user just keeps+-- inserting without calling 'supplyUnionCredits'.+supplyCreditsLevels :: NominalCredit -> Levels s -> ST s ()+supplyCreditsLevels nominalDeposit =+ traverse_ $ \(Level ir _rs) -> do+ case ir of+ Single{} -> pure ()+ Merging _mp nominalDebt nominalCreditVar+ mr@(MergingRun _ physicalDebt _) -> do++ nominalCredit <- depositNominalCredit+ nominalDebt nominalCreditVar nominalDeposit+ physicalCredit <- suppliedCreditMergingRun mr+ let !physicalCredit' = scaleNominalToPhysicalCredit+ nominalDebt physicalDebt nominalCredit+ -- Our target physicalCredit' could actually be less than the+ -- actual current physicalCredit if other tables were contributing+ -- credits to the shared merge.+ !physicalDeposit = physicalCredit' - physicalCredit++ -- So we may have a zero or negative deposit, which we ignore.+ when (physicalDeposit > 0) $ do+ leftoverCredits <- supplyCreditsMergingRun physicalDeposit mr+ -- For merges at ordinary levels (not unions) we expect to hit the+ -- debt limit exactly and not exceed it. However if we had a race+ -- on supplying credit then we could go over (which is not a problem).+ -- We can detect such races if the credit afterwards is not the amount+ -- that we credited. This is all just for sanity checking.+ physicalCredit'' <- suppliedCreditMergingRun mr+ assert (leftoverCredits == 0 || physicalCredit' /= physicalCredit'')+ (pure ())++ -- There is a potential race here in between deciding how much physical+ -- credit to supply, and then supplying it. That's because we read the+ -- "current" (absolute) physical credits, decide how much extra+ -- (relative) credits to supply and then do the transaction to supply+ -- the extra (relative) credits. In between the reading and supplying+ -- the current (absolute) physical credits could have changed due to+ -- another thread doing a merge on a different table handle.+ --+ -- This race is relatively benign. When it happens, we will supply more+ -- credit to the merge than either thread intended, however, next time+ -- either thread comes round they'll find the merge has more physical+ -- credits and will thus supply less or none. The only minor problem is+ -- in asserting that we don't supply more physical credits than the+ -- debt limit.++ -- There is a trade-off, we could supply absolute physical credit to+ -- the merging run, and let it calculate the relative credit as part+ -- of the credit transaction. However, we would also need to support+ -- relative credit for the union merges, which do not have any notion+ -- of nominal credit and only work in terms of relative physical credit.+ -- So we can have a simple relative physical credit and rare benign+ -- races, or a more complex scheme for contributing physical credits+ -- either as absolute or relative values.++scaleNominalToPhysicalCredit ::+ NominalDebt+ -> MergeDebt+ -> NominalCredit+ -> Credit+scaleNominalToPhysicalCredit (NominalDebt nominalDebt)+ MergeDebt { totalDebt = physicalDebt }+ (NominalCredit nominalCredit) =+ floor $ toRational nominalCredit * toRational physicalDebt+ / toRational nominalDebt+ -- This specification using Rational as an intermediate representation can+ -- be implemented efficiently using only integer operations.++depositNominalCredit ::+ NominalDebt+ -> STRef s NominalCredit+ -> NominalCredit+ -> ST s NominalCredit+depositNominalCredit (NominalDebt nominalDebt)+ nominalCreditVar+ (NominalCredit deposit) = do+ NominalCredit before <- readSTRef nominalCreditVar+ -- Depositing _could_ leave the credit higher than the debt, because+ -- sometimes under-full runs mean we don't shuffle runs down the levels+ -- as quickly as the worst case. So here we do just drop excess nominal+ -- credits.+ let !after = NominalCredit (min (before + deposit) nominalDebt)+ writeSTRef nominalCreditVar after+ pure after++-------------------------------------------------------------------------------+-- Updates+--++increment :: forall s. Tracer (ST s) (EventAt EventDetail)+ -> Counter+ -> LSMConfig+ -> Run -> Levels s -> UnionLevel s -> ST s (Levels s)+increment tr sc conf run0 ls0 ul = do+ go 1 [run0] ls0+ where+ mergeTypeFor :: Levels s -> LevelMergeType+ mergeTypeFor ls = mergeTypeForLevel ls ul++ go :: Int -> [Run] -> Levels s -> ST s (Levels s)+ go !ln incoming [] = do+ traceWith tr' AddLevelEvent+ let mergePolicy = mergePolicyForLevel ln [] ul+ ir <- newLevelMerge tr' conf ln mergePolicy (mergeTypeFor []) incoming+ pure (Level ir [] : [])+ where+ tr' = contramap (EventAt sc ln) tr++ go !ln incoming (Level ir rs : ls) = do+ r <- case ir of+ Single r -> do+ traceWith tr' $ SingleRunCompletedEvent r+ pure r+ Merging mergePolicy _ _ mr -> do+ r <- expectCompletedMergingRun mr+ traceWith tr' LevelMergeCompletedEvent {+ mergePolicy,+ mergeType = let MergingRun mt _ _ = mr in mt,+ mergeSize = runSize r+ }+ pure r++ let resident = r:rs+ case mergePolicyForLevel ln ls ul of++ -- If r is still too small for this level then keep it and merge again+ -- with the incoming runs.+ LevelTiering | runTooSmallForLevel LevelTiering conf ln r -> do+ traceWith tr' $ RunTooSmallForLevelEvent LevelTiering r++ ir' <- newLevelMerge tr' conf ln LevelTiering (mergeTypeFor ls) (incoming ++ [r])+ pure (Level ir' rs : ls)++ -- This tiering level is now full. We take the completed merged run+ -- (the previous incoming runs), plus all the other runs on this level+ -- as a bundle and move them down to the level below. We start a merge+ -- for the new incoming runs. This level is otherwise empty.+ LevelTiering | levelIsFullTiering conf ln incoming resident -> do+ traceWith tr' $ LevelIsFullEvent LevelTiering++ ir' <- newLevelMerge tr' conf ln LevelTiering MergeMidLevel incoming+ ls' <- go (ln+1) resident ls+ pure (Level ir' [] : ls')++ -- This tiering level is not yet full. We move the completed merged run+ -- into the level proper, and start the new merge for the incoming runs.+ LevelTiering -> do+ traceWith tr' $ LevelIsNotFullEvent LevelTiering++ ir' <- newLevelMerge tr' conf ln LevelTiering (mergeTypeFor ls) incoming+ traceWith tr' (AddRunEvent resident)+ pure (Level ir' resident : ls)++ -- The final level is using levelling. If the existing completed merge+ -- run is too large for this level, we promote the run to the next+ -- level and start merging the incoming runs into this (otherwise+ -- empty) level .+ LevelLevelling | levelIsFullLevelling conf ln incoming r -> do+ traceWith tr' $ LevelIsFullEvent LevelLevelling++ assert (null rs && null ls) $ pure ()+ ir' <- newLevelMerge tr' conf ln LevelTiering MergeMidLevel incoming+ ls' <- go (ln+1) [r] []+ pure (Level ir' [] : ls')++ -- Otherwise we start merging the incoming runs into the run.+ LevelLevelling -> do+ traceWith tr' $ LevelIsNotFullEvent LevelLevelling++ assert (null rs && null ls) $ pure ()+ ir' <- newLevelMerge tr' conf ln LevelLevelling (mergeTypeFor ls)+ (incoming ++ [r])+ pure (Level ir' [] : [])++ where+ tr' = contramap (EventAt sc ln) tr++newLevelMerge :: Tracer (ST s) EventDetail+ -> LSMConfig+ -> Int -> MergePolicyForLevel -> LevelMergeType+ -> [Run] -> ST s (IncomingRun s)+newLevelMerge tr _ _ _ _ [r] = do+ traceWith tr $ NewSingleRunEvent r+ pure (Single r)+newLevelMerge tr conf@LSMConfig{..} level mergePolicy mergeType rs = do+ mergingRun@(MergingRun _ physicalDebt _) <- newMergingRun mergeType rs+ traceWith tr NewLevelMergeEvent {+ mergePolicy,+ mergeType,+ mergeDebt = totalDebt physicalDebt,+ mergeRuns = rs+ }+ assertST (length rs `elem` [configSizeRatio, configSizeRatio + 1])+ assertWithMsgM $ leq (totalDebt physicalDebt) maxPhysicalDebt+ nominalCreditVar <- newSTRef (NominalCredit 0)+ pure (Merging mergePolicy nominalDebt nominalCreditVar mergingRun)+ where+ -- The nominal debt equals the minimum of credits we will supply before we+ -- expect the merge to complete. This is the same as the number of updates+ -- in a run that gets moved to this level.+ nominalDebt = NominalDebt (levelNumberToMaxRunSize LevelTiering conf level)++ -- The physical debt is the number of actual merge steps we will need to+ -- perform before the merge is complete. This is always the sum of the+ -- lengths of the input runs.+ --+ -- As we supply nominal credit, we scale them and supply physical credits,+ -- such that we pay off the physical and nominal debts at the same time.+ --+ -- We can bound the worst case physical debt: this is the maximum amount of+ -- steps a merge at this level could need. See the+ -- 'expectedMergingRunLengths' where-clause of the 'invariant' function for+ -- the full reasoning.+ maxPhysicalDebt =+ case mergePolicy of+ LevelLevelling ->+ -- Incoming runs, which may be slightly overfull with respect to the+ -- previous level+ configSizeRatio * levelNumberToMaxRunSize LevelTiering conf level+ -- The single run that was already on this level+ + levelNumberToMaxRunSize LevelLevelling conf level+ LevelTiering ->+ -- Incoming runs, which may be slightly overfull with respect to the+ -- previous level+ configSizeRatio * levelNumberToMaxRunSize LevelTiering conf level+ -- Held back run that is underfull with respect to the current+ -- level+ + levelNumberToMaxRunSize LevelTiering conf (level - 1)++-------------------------------------------------------------------------------+-- MergingTree abstraction+--++-- Note [Table Unions]+-- ~~~~~~~~~~~~~~~~~~~+--+-- Semantically, tables are key-value stores like Haskell's @Map@. Table unions+-- then behave like @Map.unionWith (<>)@. If one of the input tables contains+-- a value at a particular key, the result will also contain it. If multiple+-- tables share that key, the values will be combined monoidally (using+-- 'resolveValue' in in this prototype).+--+-- Looking at the implementation, tables are not just key-value pairs, but+-- consist of runs. If each table was just a single run, unioning would involve+-- a run merge similar to the one used for compaction (when a level is full),+-- but with a different merge type 'MergeUnion' that differs semantically: Here,+-- runs don't represent updates (overwriting each other), but they each+-- represent the full state of a table. There is no distinction between no entry+-- and a 'Delete', between an 'Insert' and a 'Mupsert'.+--+-- To union two tables, we can therefore first merge down each table into a+-- single run (using regular level merges) and then union merge these.+--+-- However, we want to spread out the work required and perform these merges+-- incrementally. At first, we only create a new table that is empty except for+-- a data structure 'MergingTree', representing the merges that need to be done.+-- The usual operations can then be performed on the table while the merge is+-- in progress: Inserts go into the table as usual, not affecting its last level+-- ('UnionLevel'), lookups need to consider the tree (requiring some complexity+-- and runtime overhead), further unions incorporate the in-progress tree into+-- the resulting one, which also shares future merging work.+--+-- It seems necessary to represent the suspended merges using a tree. Other+-- approaches don't allow for full sharing of the incremental work (e.g. because+-- they effectively \"re-bracket\" nested unions). It also seems necessary to+-- first merge each input table into a single run, as there is no practical+-- distributive property between level and union merges.++-- | Ensures that the merge contains more than one input, avoiding creating a+-- pending merge where possible.+newPendingLevelMerge :: [IncomingRun s]+ -> Maybe (MergingTree s)+ -> ST s (Maybe (MergingTree s))+newPendingLevelMerge [] t = pure t+newPendingLevelMerge [Single r] Nothing =+ Just . MergingTree <$> newSTRef (CompletedTreeMerge r)+newPendingLevelMerge [Merging{}] Nothing =+ -- This case should never occur. If there is a single entry in the list,+ -- there can only be one level in the input table. At level 1 there are no+ -- merging runs, so it must be a PreExistingRun.+ error "newPendingLevelMerge: singleton Merging run"+newPendingLevelMerge irs tree = do+ let prs = map incomingToPreExistingRun irs+ st = PendingTreeMerge (PendingLevelMerge prs tree)+ Just . MergingTree <$> newSTRef st+ where+ incomingToPreExistingRun (Single r) = PreExistingRun r+ incomingToPreExistingRun (Merging _ _ _ mr) = PreExistingMergingRun mr++-- | Ensures that the merge contains more than one input.+newPendingUnionMerge :: [MergingTree s] -> ST s (Maybe (MergingTree s))+newPendingUnionMerge [] = pure Nothing+newPendingUnionMerge [t] = pure (Just t)+newPendingUnionMerge trees = do+ let st = PendingTreeMerge (PendingUnionMerge trees)+ Just . MergingTree <$> newSTRef st++contentToMergingTree :: LSMContent s -> ST s (Maybe (MergingTree s))+contentToMergingTree (LSMContent wb ls ul) =+ newPendingLevelMerge (buffers ++ levels) trees+ where+ -- flush the write buffer (but this should not modify the content)+ buffers+ | bufferSize wb == 0 = []+ | otherwise = [Single (bufferToRun wb)]++ levels = flip concatMap ls $ \(Level ir rs) -> ir : map Single rs++ trees = case ul of+ NoUnion -> Nothing+ Union t _ -> Just t++-- | When calculating (an upped bound of) the total debt of a recursive tree of+-- merges, we also need to return an upper bound on the size of the resulting+-- run. See 'remainingDebtPendingMerge'.+type Size = Int++remainingDebtMergingTree :: MergingTree s -> ST s (Debt, Size)+remainingDebtMergingTree (MergingTree ref) =+ readSTRef ref >>= \case+ CompletedTreeMerge r -> pure (0, runSize r)+ OngoingTreeMerge mr -> addDebtOne <$> remainingDebtMergingRun mr+ PendingTreeMerge pm -> addDebtOne <$> remainingDebtPendingMerge pm+ where+ -- An ongoing merge should never have 0 debt, even if the 'MergingRun' in it+ -- says it is completed. We still need to update it to 'CompletedTreeMerge'.+ -- Similarly, a pending merge needs some work to complete it, even if all+ -- its inputs are empty.+ --+ -- Note that we can't use @max 1@, as this would violate the property that+ -- supplying N credits reduces the remaining debt by at least N.+ addDebtOne (debt, size) = (debt + 1, size)++remainingDebtPendingMerge :: PendingMerge s -> ST s (Debt, Size)+remainingDebtPendingMerge (PendingMerge _ prs trees) = do+ (debts, sizes) <- unzip . concat <$> sequence+ [ traverse remainingDebtPreExistingRun prs+ , traverse remainingDebtMergingTree trees+ ]+ let totalSize = sum sizes+ let totalDebt = sum debts + totalSize+ pure (totalDebt, totalSize)+ where+ remainingDebtPreExistingRun = \case+ PreExistingRun r -> pure (0, runSize r)+ PreExistingMergingRun mr -> remainingDebtMergingRun mr++remainingDebtMergingRun :: MergingRun t s -> ST s (Debt, Size)+remainingDebtMergingRun (MergingRun _ d ref) =+ readSTRef ref >>= \case+ CompletedMerge r ->+ pure (0, runSize r)+ OngoingMerge c inputRuns _ ->+ pure (mergeDebtLeft d c, sum (map runSize inputRuns))++-- | For each of the @supplyCredits@ type functions, we want to check some+-- common properties.+checked :: HasCallStack+ => (a -> ST s (Debt, Size)) -- ^ how to calculate the current debt+ -> (Credit -> a -> ST s Credit) -- ^ how to supply the credits+ -> Credit -> a -> ST s Credit+checked query supply credits x = do+ assertST $ credits > 0 -- only call them when there are credits to spend+ debt <- fst <$> query x+ assertST $ debt >= 0 -- debt can't be negative+ c' <- supply credits x+ assertST $ c' <= credits -- can't have more leftovers than we started with+ assertST $ c' >= 0 -- leftovers can't be negative+ debt' <- fst <$> query x+ assertST $ debt' >= 0+ -- the debt was reduced sufficiently (amount of credits spent)+ assertST $ debt' <= debt - (credits - c')+ pure c'++supplyCreditsMergingTree :: Credit -> MergingTree s -> ST s Credit+supplyCreditsMergingTree = checked remainingDebtMergingTree $ \credits (MergingTree ref) -> do+ treeState <- readSTRef ref+ (!c', !treeState') <- supplyCreditsMergingTreeState credits treeState+ writeSTRef ref treeState'+ pure c'++supplyCreditsMergingTreeState :: Credit -> MergingTreeState s+ -> ST s (Credit, MergingTreeState s)+supplyCreditsMergingTreeState credits !state = do+ assertST (credits >= 0)+ case state of+ CompletedTreeMerge{} ->+ pure (credits, state)+ OngoingTreeMerge mr -> do+ c' <- supplyCreditsMergingRun credits mr+ if c' <= 0+ then pure (0, state)+ else do+ r <- expectCompletedMergingRun mr+ -- all work is done, we can't spend any more credits+ pure (c', CompletedTreeMerge r)+ PendingTreeMerge pm -> do+ c' <- supplyCreditsPendingMerge credits pm+ if c' <= 0+ then+ -- still remaining work in children, we can't do more for now+ pure (c', state)+ else do+ -- all children must be done, create new merge!+ (mergeType, rs) <- expectCompletedChildren pm+ case rs of+ [r] -> pure (c', CompletedTreeMerge r)+ _ -> do+ state' <- OngoingTreeMerge <$> newMergingRun mergeType rs+ -- use any remaining credits to progress the new merge+ supplyCreditsMergingTreeState c' state'++supplyCreditsPendingMerge :: Credit -> PendingMerge s -> ST s Credit+supplyCreditsPendingMerge = checked remainingDebtPendingMerge $ \credits -> \case+ PendingLevelMerge prs tree ->+ leftToRight supplyPreExistingRun prs credits+ >>= leftToRight supplyCreditsMergingTree (toList tree)+ PendingUnionMerge trees ->+ splitEqually supplyCreditsMergingTree trees credits+ where+ supplyPreExistingRun c = \case+ PreExistingRun _r -> pure c+ PreExistingMergingRun mr -> supplyCreditsMergingRun c mr++ -- supply credits left to right until they are used up+ leftToRight :: (Credit -> a -> ST s Credit) -> [a] -> Credit -> ST s Credit+ leftToRight _ _ 0 = pure 0+ leftToRight _ [] c = pure c+ leftToRight f (x:xs) c = f c x >>= leftToRight f xs++ -- approximately equal, being more precise would require more iterations+ splitEqually :: (Credit -> a -> ST s Credit) -> [a] -> Credit -> ST s Credit+ splitEqually f xs credits =+ -- first give each tree k = ceil(1/n) credits (last ones might get less).+ -- it's important we fold here to collect leftovers.+ -- any remainders go left to right.+ foldM supply credits xs >>= leftToRight f xs+ where+ !n = length xs+ !k = (credits + (n - 1)) `div` n++ supply 0 _ = pure 0+ supply c t = do+ let creditsToSpend = min k c+ leftovers <- f creditsToSpend t+ pure (c - creditsToSpend + leftovers)++expectCompletedChildren :: HasCallStack+ => PendingMerge s -> ST s (TreeMergeType, [Run])+expectCompletedChildren (PendingMerge mt prs trees) = do+ rs1 <- traverse expectCompletedPreExistingRun prs+ rs2 <- traverse expectCompletedMergingTree trees+ pure (mt, rs1 ++ rs2)+ where+ expectCompletedPreExistingRun = \case+ PreExistingRun r -> pure r+ PreExistingMergingRun mr -> expectCompletedMergingRun mr++expectCompletedMergingTree :: HasCallStack => MergingTree s -> ST s Run+expectCompletedMergingTree = expectInvariant . isCompletedMergingTree++-------------------------------------------------------------------------------+-- Measurements+--++data MTree r = MLeaf r+ | MNode TreeMergeType [MTree r]+ deriving stock (Eq, Foldable, Functor, Show)++allLevels :: LSM s -> ST s (Buffer, [[Run]], Maybe (MTree Run))+allLevels (LSMHandle _ _ _conf lsmr) = do+ LSMContent wb ls ul <- readSTRef lsmr+ rs <- flattenLevels ls+ tree <- case ul of+ NoUnion -> pure Nothing+ Union t _ -> Just <$> flattenTree t+ pure (wb, rs, tree)++flattenLevels :: Levels s -> ST s [[Run]]+flattenLevels = mapM flattenLevel++flattenLevel :: Level s -> ST s [Run]+flattenLevel (Level ir rs) = (++ rs) <$> flattenIncomingRun ir++flattenIncomingRun :: IncomingRun s -> ST s [Run]+flattenIncomingRun = \case+ Single r -> pure [r]+ Merging _ _ _ mr -> flattenMergingRun mr++flattenMergingRun :: MergingRun t s -> ST s [Run]+flattenMergingRun (MergingRun _ _ ref) = do+ mrs <- readSTRef ref+ case mrs of+ CompletedMerge r -> pure [r]+ OngoingMerge _ rs _ -> pure rs++flattenTree :: MergingTree s -> ST s (MTree Run)+flattenTree (MergingTree ref) = do+ mts <- readSTRef ref+ case mts of+ CompletedTreeMerge r ->+ pure (MLeaf r)+ OngoingTreeMerge (MergingRun mt _ mrs) ->+ readSTRef mrs >>= \case+ CompletedMerge r -> pure (MLeaf r)+ OngoingMerge _ rs _ -> pure (MNode mt (MLeaf <$> rs))+ PendingTreeMerge (PendingMerge mt irs trees) -> do+ irs' <- map MLeaf . concat <$> traverse flattenPreExistingRun irs+ trees' <- traverse flattenTree trees+ pure (MNode mt (irs' ++ trees'))++flattenPreExistingRun :: PreExistingRun s -> ST s [Run]+flattenPreExistingRun = \case+ PreExistingRun r -> pure [r]+ PreExistingMergingRun mr -> flattenMergingRun mr++logicalValue :: LSM s -> ST s (Map Key (Value, Maybe Blob))+logicalValue lsm = do+ (wb, levels, tree) <- allLevels lsm+ let r = mergek+ MergeLevel+ (wb : concat levels ++ toList (mergeTree <$> tree))+ pure (Map.mapMaybe justInsert r)+ where+ mergeTree :: MTree Run -> Run+ mergeTree (MLeaf r) = r+ mergeTree (MNode mt ts) = mergek mt (map mergeTree ts)++ justInsert (Insert v b) = Just (v, b)+ justInsert Delete = Nothing+ justInsert (Mupsert v) = Just (v, Nothing)++type Representation = (Run, [LevelRepresentation], Maybe (MTree Run))++type LevelRepresentation =+ (Maybe (MergePolicyForLevel, NominalDebt, NominalCredit,+ LevelMergeType, MergingRunState),+ [Run])++dumpRepresentation :: LSM s -> ST s Representation+dumpRepresentation (LSMHandle _ _ _conf lsmr) = do+ LSMContent wb ls ul <- readSTRef lsmr+ levels <- mapM dumpLevel ls+ tree <- case ul of+ NoUnion -> pure Nothing+ Union t _ -> Just <$> flattenTree t+ pure (wb, levels, tree)++dumpLevel :: Level s -> ST s LevelRepresentation+dumpLevel (Level (Single r) rs) =+ pure (Nothing, (r:rs))+dumpLevel (Level (Merging mp nd ncv (MergingRun mt _ ref)) rs) = do+ mrs <- readSTRef ref+ nc <- readSTRef ncv+ pure (Just (mp, nd, nc, mt, mrs), rs)++-- For each level:+-- 1. the runs involved in an ongoing merge+-- 2. the other runs (including completed merge)+representationShape :: Representation+ -> (Int, [([Int], [Int])], Maybe (MTree Int))+representationShape (wb, levels, tree) =+ (summaryRun wb, map summaryLevel levels, fmap (fmap summaryRun) tree)+ where+ summaryLevel (mmr, rs) =+ let (ongoing, complete) = summaryMR mmr+ in (ongoing, complete <> map summaryRun rs)++ summaryRun = runSize++ summaryMR = \case+ Nothing -> ([], [])+ Just (_, _, _, _, CompletedMerge r) -> ([], [summaryRun r])+ Just (_, _, _, _, OngoingMerge _ rs _) -> (map summaryRun rs, [])++-------------------------------------------------------------------------------+-- Tracing+--++-- TODO: these events are incomplete, in particular we should also trace what+-- happens in the union level.+data Event =+ NewTableEvent TableId LSMConfig+ | UpdateEvent TableId Key Entry+ | LookupEvent TableId Key+ | DuplicateEvent TableId TableId+ | UnionsEvent TableId [TableId]+ | LevelEvent TableId (EventAt EventDetail)+ deriving stock Show++data EventAt e = EventAt {+ eventAtStep :: Counter,+ eventAtLevel :: Int,+ eventDetail :: e+ }+ deriving stock Show++data EventDetail =+ AddLevelEvent+ | AddRunEvent {+ runsAtLevel :: [Run]+ }+ | NewLevelMergeEvent {+ mergePolicy :: MergePolicyForLevel,+ mergeType :: LevelMergeType,+ mergeDebt :: Debt,+ mergeRuns :: [Run]+ }+ | NewSingleRunEvent Run+ | LevelMergeCompletedEvent {+ mergePolicy :: MergePolicyForLevel,+ mergeType :: LevelMergeType,+ mergeSize :: Int+ }+ | SingleRunCompletedEvent Run++ | RunTooSmallForLevelEvent MergePolicyForLevel Run+ | LevelIsFullEvent MergePolicyForLevel+ | LevelIsNotFullEvent MergePolicyForLevel+ deriving stock Show++-------------------------------------------------------------------------------+-- Arbitrary+--++instance QC.Arbitrary Key where+ arbitrary = K <$> QC.arbitrarySizedNatural+ shrink (K v) = K <$> QC.shrink v++instance QC.Arbitrary Value where+ arbitrary = V <$> QC.arbitrarySizedNatural+ shrink (V v) = V <$> QC.shrink v++instance QC.Arbitrary Blob where+ arbitrary = B <$> QC.arbitrarySizedNatural+ shrink (B v) = B <$> QC.shrink v++instance (QC.Arbitrary v, QC.Arbitrary b) => QC.Arbitrary (Update v b) where+ arbitrary = QC.frequency+ [ (3, Insert <$> QC.arbitrary <*> QC.arbitrary)+ , (1, Mupsert <$> QC.arbitrary)+ , (1, pure Delete)+ ]++instance QC.Arbitrary LevelMergeType where+ arbitrary = QC.elements [MergeMidLevel, MergeLastLevel]++instance QC.Arbitrary TreeMergeType where+ arbitrary = QC.elements [MergeLevel, MergeUnion]
@@ -0,0 +1,313 @@+{-# LANGUAGE CPP #-}+module RocksDB (+ -- * Options+ Options,+ withOptions,+ optionsSetCreateIfMissing,+ optionsSetMaxOpenFiles,+ optionsSetCompression,+ optionsIncreaseParallelism,+ optionsOptimizeLevelStyleCompaction,+ -- * Read Options+ ReadOptions,+ withReadOptions,+ -- * Write options+ WriteOptions,+ withWriteOptions,+ writeOptionsDisableWAL,+ -- * DB operations+ RocksDB,+ withRocksDB,+ put,+ get,+ multiGet,+ delete,+ write,+ checkpoint,+ -- * Write batch+ WriteBatch,+ withWriteBatch,+ writeBatchPut,+ writeBatchDelete,+ -- * Block based table options+ BlockTableOptions,+ withBlockTableOptions,+ blockBasedOptionsSetFilterPolicy,+ optionsSetBlockBasedTableFactory,+ -- * Filter policy+ FilterPolicy,+ withFilterPolicyBloom,+) where++import Control.Exception (bracket)+import Control.Monad (forM)+import Data.ByteString (ByteString)+import Data.ByteString.Unsafe (unsafePackMallocCStringLen,+ unsafeUseAsCStringLen)+import Data.Foldable.WithIndex (ifor_)+import Data.Word (Word64)+import Foreign.C.String (peekCString, withCString)+import Foreign.C.Types (CInt, CSize)+import Foreign.Marshal.Alloc (alloca, free)+import Foreign.Marshal.Array (allocaArray)+import Foreign.Marshal.Utils (withMany)+import Foreign.Ptr (Ptr, nullPtr)+import Foreign.Storable (peek, peekElemOff, poke, pokeElemOff)++import RocksDB.FFI++#if MIN_VERSION_base(4,18,0)+import Foreign.C.ConstPtr (ConstPtr (..))+#endif++-------------------------------------------------------------------------------+-- constptr+-------------------------------------------------------------------------------++#if MIN_VERSION_base(4,18,0)+constPtr :: Ptr a -> ConstPtr a+constPtr = ConstPtr+#else+constPtr :: Ptr a -> Ptr a+constPtr = id+#endif++-------------------------------------------------------------------------------+-- error+-------------------------------------------------------------------------------++withErrPtr :: (ErrPtr -> IO r) -> IO r+withErrPtr kont = alloca $ \ptr -> do+ poke ptr nullPtr+ x <- kont ptr+ ptr' <- peek ptr+ if ptr' == nullPtr+ then pure ()+ else free ptr'+ pure x++assertErrPtr :: String -> ErrPtr -> IO ()+assertErrPtr fun ptr = do+ ptr' <- peek ptr+ if ptr' == nullPtr+ then pure ()+ else do+ msg <- peekCString ptr'+ fail $ fun ++ ": " ++ msg++-------------------------------------------------------------------------------+-- options+-------------------------------------------------------------------------------++newtype Options = Options (Ptr OPTIONS)++withOptions :: (Options -> IO r) -> IO r+withOptions kont = bracket+ rocksdb_options_create+ rocksdb_options_destroy+ (\ptr -> kont (Options ptr))++optionsSetCreateIfMissing :: Options -> Bool -> IO ()+optionsSetCreateIfMissing (Options ptr) v =+ rocksdb_options_set_create_if_missing ptr (if v then 1 else 0)++optionsIncreaseParallelism :: Options -> Int -> IO ()+optionsIncreaseParallelism (Options ptr) v =+ rocksdb_options_increase_parallelism ptr (intToCInt v)++optionsSetMaxOpenFiles :: Options -> Int -> IO ()+optionsSetMaxOpenFiles (Options ptr) v =+ rocksdb_options_set_max_open_files ptr (intToCInt v)++optionsOptimizeLevelStyleCompaction :: Options -> Word64 -> IO ()+optionsOptimizeLevelStyleCompaction (Options ptr) v =+ rocksdb_options_optimize_level_style_compaction ptr v++optionsSetCompression :: Options -> Int -> IO ()+optionsSetCompression (Options ptr) v =+ rocksdb_options_set_compression ptr (intToCInt v)++-------------------------------------------------------------------------------+-- read options+-------------------------------------------------------------------------------++newtype ReadOptions = ReadOptions (Ptr READOPTIONS)++withReadOptions :: (ReadOptions -> IO r) -> IO r+withReadOptions kont = bracket+ rocksdb_readoptions_create+ rocksdb_readoptions_destroy+ (\ptr -> kont (ReadOptions ptr))++-------------------------------------------------------------------------------+-- write options+-------------------------------------------------------------------------------++newtype WriteOptions = WriteOptions (Ptr WRITEOPTIONS)++withWriteOptions :: (WriteOptions -> IO r) -> IO r+withWriteOptions kont = bracket+ rocksdb_writeoptions_create+ rocksdb_writeoptions_destroy+ (\ptr -> kont (WriteOptions ptr))++writeOptionsDisableWAL :: WriteOptions -> Bool {- ^ Disable -} -> IO ()+writeOptionsDisableWAL (WriteOptions opts) disable =+ rocksdb_writeoptions_disable_WAL opts (if disable then 1 else 0)++-------------------------------------------------------------------------------+-- db operations+-------------------------------------------------------------------------------++data RocksDB = RocksDB (Ptr DB) ErrPtr++withRocksDB :: Options -> FilePath -> (RocksDB -> IO r) -> IO r+withRocksDB (Options opt) path kont =+ withCString path $ \path' ->+ withErrPtr $ \errptr ->+ bracket (rocksdb_open' path' errptr) rocksdb_close $ \ptr ->+ kont (RocksDB ptr errptr)+ where+ rocksdb_open' path' errptr = do+ ptr <- rocksdb_open opt path' errptr+ assertErrPtr "rocksdb_open" errptr+ pure ptr++put :: RocksDB -> WriteOptions -> ByteString -> ByteString -> IO ()+put (RocksDB dbPtr errPtr) (WriteOptions opts) key val =+ unsafeUseAsCStringLen key $ \(kp, kl) ->+ unsafeUseAsCStringLen val $ \(vp, vl) -> do+ rocksdb_put dbPtr opts kp (intToCSize kl) vp (intToCSize vl) errPtr+ assertErrPtr "rocksdb_put" errPtr++get :: RocksDB -> ReadOptions -> ByteString -> IO (Maybe ByteString)+get (RocksDB dbPtr errPtr) (ReadOptions opts) key =+ unsafeUseAsCStringLen key $ \(kp, kl) ->+ alloca $ \vlPtr -> do+ vp <- rocksdb_get dbPtr opts kp (intToCSize kl) vlPtr errPtr+ assertErrPtr "rocksdb_get" errPtr -- TODO: may leak+ vl <- peek vlPtr+ if vp == nullPtr+ then pure Nothing+ else Just <$> unsafePackMallocCStringLen (vp, csizeToInt vl)++multiGet :: RocksDB -> ReadOptions -> [ByteString] -> IO [Maybe ByteString]+multiGet (RocksDB db _errPtr) (ReadOptions opts) keys =+ allocaArray n $ \kps ->+ allocaArray n $ \kls ->+ allocaArray n $ \vps ->+ allocaArray n $ \vls ->+ allocaArray n $ \errs ->+ withMany unsafeUseAsCStringLen keys $ \keys' -> do+ -- populate keys+ ifor_ keys' $ \i (kp, kl) -> do+ pokeElemOff kps i (constPtr kp)+ pokeElemOff kls i (intToCSize kl)++ -- multi get+ rocksdb_multi_get db opts (intToCSize n) kps kls vps vls errs++ -- read keys+ forM [ 0 .. n - 1 ] $ \i -> do+ vp <- peekElemOff vps i+ vl <- peekElemOff vls i++ -- TODO: we don't check errs here. we should.++ if vp == nullPtr+ then pure Nothing+ else Just <$> unsafePackMallocCStringLen (vp, csizeToInt vl)++ where+ n = length keys++delete :: RocksDB -> WriteOptions -> ByteString -> IO ()+delete (RocksDB dbPtr errPtr) (WriteOptions opts) key =+ unsafeUseAsCStringLen key $ \(kp, kl) -> do+ rocksdb_delete dbPtr opts kp (intToCSize kl) errPtr+ assertErrPtr "rocksdb_delete" errPtr++write :: RocksDB -> WriteOptions -> WriteBatch -> IO ()+write (RocksDB dbPtr errPtr) (WriteOptions opts) (WriteBatch batch) = do+ rocksdb_write dbPtr opts batch errPtr+ assertErrPtr "rocksdb_write" errPtr++checkpoint :: RocksDB -> FilePath -> IO ()+checkpoint (RocksDB dbPtr errPtr) path =+ withCString path $ \path' ->+ bracket rocksdb_checkpoint_object_create' rocksdb_checkpoint_object_destroy $ \cp -> do+ rocksdb_checkpoint_create cp path' 0 errPtr+ assertErrPtr "rocksdb_checkpoint_create" errPtr+ where+ rocksdb_checkpoint_object_create' = do+ cp <- rocksdb_checkpoint_object_create dbPtr errPtr+ assertErrPtr "rocksdb_checkpoint_object_create" errPtr+ pure cp++-------------------------------------------------------------------------------+-- write batch+-------------------------------------------------------------------------------++newtype WriteBatch = WriteBatch (Ptr WRITEBATCH)++withWriteBatch :: (WriteBatch -> IO r) -> IO r+withWriteBatch kont = bracket+ rocksdb_writebatch_create+ rocksdb_writebatch_destroy+ (\ptr -> kont (WriteBatch ptr))++writeBatchPut :: WriteBatch -> ByteString -> ByteString -> IO ()+writeBatchPut (WriteBatch batch) key val =+ unsafeUseAsCStringLen key $ \(kp, kl) ->+ unsafeUseAsCStringLen val $ \(vp, vl) ->+ rocksdb_writebatch_put batch kp (intToCSize kl) vp (intToCSize vl)++writeBatchDelete :: WriteBatch -> ByteString -> IO ()+writeBatchDelete (WriteBatch batch) key =+ unsafeUseAsCStringLen key $ \(kp, kl) ->+ rocksdb_writebatch_delete batch kp (intToCSize kl)++-------------------------------------------------------------------------------+-- block based table options+-------------------------------------------------------------------------------++newtype BlockTableOptions = BlockTableOptions (Ptr BLOCKTABLEOPTIONS)++withBlockTableOptions :: (BlockTableOptions -> IO r) -> IO r+withBlockTableOptions kont = bracket+ rocksdb_block_based_options_create+ rocksdb_block_based_options_destroy+ (\ptr -> kont (BlockTableOptions ptr))++blockBasedOptionsSetFilterPolicy :: BlockTableOptions -> FilterPolicy -> IO ()+blockBasedOptionsSetFilterPolicy (BlockTableOptions opts) (FilterPolicy policy) =+ rocksdb_block_based_options_set_filter_policy opts policy++optionsSetBlockBasedTableFactory :: Options -> BlockTableOptions -> IO ()+optionsSetBlockBasedTableFactory (Options opts) (BlockTableOptions ptr) =+ rocksdb_options_set_block_based_table_factory opts ptr++-------------------------------------------------------------------------------+-- filter policy+-------------------------------------------------------------------------------++newtype FilterPolicy = FilterPolicy (Ptr FILTERPOLICY)++withFilterPolicyBloom :: Int -> (FilterPolicy -> IO r) -> IO r+withFilterPolicyBloom bits_per_key kont = bracket+ (rocksdb_filterpolicy_create_bloom (intToCInt bits_per_key))+ (\_ -> pure ()) -- rocksdb_filterpolicy_destroy, causes segfault if we free it.+ (\ptr -> kont (FilterPolicy ptr))++-------------------------------------------------------------------------------+-- utils+-------------------------------------------------------------------------------++intToCSize :: Int -> CSize+intToCSize = fromIntegral++csizeToInt :: CSize -> Int+csizeToInt = fromIntegral++intToCInt :: Int -> CInt+intToCInt = fromIntegral
@@ -0,0 +1,245 @@+{-# LANGUAGE CApiFFI #-}+{-# LANGUAGE CPP #-}+module RocksDB.FFI (+ -- * options+ OPTIONS,+ rocksdb_options_create,+ rocksdb_options_destroy,+ rocksdb_options_set_create_if_missing,+ rocksdb_options_set_max_open_files,+ rocksdb_options_increase_parallelism,+ rocksdb_options_optimize_level_style_compaction,+ rocksdb_options_set_compression,+ -- * read options+ READOPTIONS,+ rocksdb_readoptions_create,+ rocksdb_readoptions_destroy,+ -- * write options+ WRITEOPTIONS,+ rocksdb_writeoptions_create,+ rocksdb_writeoptions_destroy,+ rocksdb_writeoptions_disable_WAL,+ -- * db operations+ DB,+ rocksdb_open,+ rocksdb_close,+ rocksdb_get,+ rocksdb_multi_get,+ rocksdb_put,+ rocksdb_delete,+ rocksdb_write,+ -- ** checkpoints+ CHECKPOINT,+ rocksdb_checkpoint_object_create,+ rocksdb_checkpoint_object_destroy,+ rocksdb_checkpoint_create,+ -- * write batch+ WRITEBATCH,+ rocksdb_writebatch_create,+ rocksdb_writebatch_destroy,+ rocksdb_writebatch_put,+ rocksdb_writebatch_delete,+ -- * block table options+ BLOCKTABLEOPTIONS,+ rocksdb_block_based_options_create,+ rocksdb_block_based_options_destroy,+ rocksdb_block_based_options_set_filter_policy,+ rocksdb_options_set_block_based_table_factory,+ -- * filter policy+ FILTERPOLICY,+ rocksdb_filterpolicy_destroy,+ rocksdb_filterpolicy_create_bloom,+ -- * error+ ErrPtr,+) where++import Data.Word (Word64)+import Foreign.C.String (CString)+import Foreign.C.Types (CChar, CInt (..), CSize (..), CUChar (..))+import Foreign.Ptr (Ptr)++#if MIN_VERSION_base(4,18,0)+import Foreign.C.ConstPtr (ConstPtr (..))+#endif++-------------------------------------------------------------------------------+-- error+-------------------------------------------------------------------------------++type ErrPtr = Ptr CString++-------------------------------------------------------------------------------+-- options+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_options_t" #-} OPTIONS++foreign import capi "rocksdb/c.h rocksdb_options_create"+ rocksdb_options_create :: IO (Ptr OPTIONS)++foreign import capi "rocksdb/c.h rocksdb_options_destroy"+ rocksdb_options_destroy :: Ptr OPTIONS -> IO ()++foreign import capi "rocksdb/c.h rocksdb_options_increase_parallelism"+ rocksdb_options_increase_parallelism :: Ptr OPTIONS -> CInt -> IO ()++foreign import capi "rocksdb/c.h rocksdb_options_optimize_level_style_compaction"+ rocksdb_options_optimize_level_style_compaction :: Ptr OPTIONS -> Word64 -> IO ()++foreign import capi "rocksdb/c.h rocksdb_options_set_create_if_missing"+ rocksdb_options_set_create_if_missing :: Ptr OPTIONS -> CUChar -> IO ()++foreign import capi "rocksdb/c.h rocksdb_options_set_max_open_files"+ rocksdb_options_set_max_open_files :: Ptr OPTIONS -> CInt -> IO ()++foreign import capi "rocksdb/c.h rocksdb_options_set_compression"+ rocksdb_options_set_compression :: Ptr OPTIONS -> CInt -> IO ()++-------------------------------------------------------------------------------+-- read options+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_readoptions_t" #-} READOPTIONS++foreign import capi "rocksdb/c.h rocksdb_readoptions_create"+ rocksdb_readoptions_create :: IO (Ptr READOPTIONS)++foreign import capi "rocksdb/c.h rocksdb_readoptions_destroy"+ rocksdb_readoptions_destroy :: Ptr READOPTIONS -> IO ()++-------------------------------------------------------------------------------+-- write options+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_writeoptions_t" #-} WRITEOPTIONS++foreign import capi "rocksdb/c.h rocksdb_writeoptions_create"+ rocksdb_writeoptions_create :: IO (Ptr WRITEOPTIONS)++foreign import capi "rocksdb/c.h rocksdb_writeoptions_destroy"+ rocksdb_writeoptions_destroy :: Ptr WRITEOPTIONS -> IO ()++foreign import capi "rocksdb/c.h rocksdb_writeoptions_disable_WAL"+ rocksdb_writeoptions_disable_WAL :: Ptr WRITEOPTIONS -> CInt -> IO ()++-------------------------------------------------------------------------------+-- db operations+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_t" #-} DB++foreign import capi "rocksdb/c.h rocksdb_open"+ rocksdb_open :: Ptr OPTIONS -> CString -> ErrPtr -> IO (Ptr DB)++foreign import capi "rocksdb/c.h rocksdb_close"+ rocksdb_close :: Ptr DB -> IO ()++-- | Returns NULL if not found. A malloc()ed array otherwise.+-- Stores the length of the array in *vallen.+foreign import capi "rocksdb/c.h rocksdb_get"+ rocksdb_get :: Ptr DB -> Ptr READOPTIONS+ -> CString -> CSize+ -> Ptr CSize+ -> ErrPtr+ -> IO CString++#if MIN_VERSION_base(4,18,0)+foreign import capi "rocksdb/c.h rocksdb_multi_get"+ rocksdb_multi_get :: Ptr DB -> Ptr READOPTIONS+ -> CSize+ -> Ptr (ConstPtr CChar) -> Ptr CSize+ -> Ptr CString -> Ptr CSize+ -> Ptr CString+ -> IO ()+#else+foreign import ccall "rocksdb/c.h rocksdb_multi_get"+ rocksdb_multi_get :: Ptr DB -> Ptr READOPTIONS+ -> CSize+ -> Ptr (Ptr CChar) -> Ptr CSize+ -> Ptr CString -> Ptr CSize+ -> Ptr CString+ -> IO ()+#endif++foreign import capi "rocksdb/c.h rocksdb_put"+ rocksdb_put :: Ptr DB -> Ptr WRITEOPTIONS+ -> CString -> CSize+ -> CString -> CSize+ -> ErrPtr+ -> IO ()++foreign import capi "rocksdb/c.h rocksdb_delete"+ rocksdb_delete :: Ptr DB -> Ptr WRITEOPTIONS+ -> CString -> CSize+ -> ErrPtr+ -> IO ()++foreign import capi "rocksdb/c.h rocksdb_write"+ rocksdb_write :: Ptr DB -> Ptr WRITEOPTIONS+ -> Ptr WRITEBATCH+ -> ErrPtr+ -> IO ()++data {-# CTYPE "rocksdb/c.h" "rocksdb_checkpoint_t" #-} CHECKPOINT++foreign import capi "rocksdb/c.h rocksdb_checkpoint_object_create"+ rocksdb_checkpoint_object_create :: Ptr DB -> ErrPtr -> IO (Ptr CHECKPOINT)++foreign import capi "rocksdb/c.h rocksdb_checkpoint_object_destroy"+ rocksdb_checkpoint_object_destroy :: Ptr CHECKPOINT -> IO ()++foreign import capi "rocksdb/c.h rocksdb_checkpoint_create"+ rocksdb_checkpoint_create :: Ptr CHECKPOINT -> CString -> Word64 -> ErrPtr -> IO ()++-------------------------------------------------------------------------------+-- write batch+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_writebatch_t" #-} WRITEBATCH++foreign import capi "rocksdb/c.h rocksdb_writebatch_create"+ rocksdb_writebatch_create :: IO (Ptr WRITEBATCH)++foreign import capi "rocksdb/c.h rocksdb_writebatch_destroy"+ rocksdb_writebatch_destroy :: Ptr WRITEBATCH -> IO ()++foreign import capi "rocksdb/c.h rocksdb_writebatch_put"+ rocksdb_writebatch_put :: Ptr WRITEBATCH+ -> CString -> CSize+ -> CString -> CSize+ -> IO ()++foreign import capi "rocksdb/c.h rocksdb_writebatch_delete"+ rocksdb_writebatch_delete :: Ptr WRITEBATCH+ -> CString -> CSize+ -> IO ()++-------------------------------------------------------------------------------+-- block based table options+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_block_based_table_options_t" #-} BLOCKTABLEOPTIONS++foreign import capi "rocksdb/c.h rocksdb_block_based_options_create"+ rocksdb_block_based_options_create :: IO (Ptr BLOCKTABLEOPTIONS)++foreign import capi "rocksdb/c.h rocksdb_block_based_options_destroy"+ rocksdb_block_based_options_destroy :: Ptr BLOCKTABLEOPTIONS -> IO ()++foreign import capi "rocksdb/c.h rocksdb_block_based_options_set_filter_policy"+ rocksdb_block_based_options_set_filter_policy :: Ptr BLOCKTABLEOPTIONS -> Ptr FILTERPOLICY -> IO ()++foreign import capi "rocksdb/c.h rocksdb_options_set_block_based_table_factory"+ rocksdb_options_set_block_based_table_factory :: Ptr OPTIONS -> Ptr BLOCKTABLEOPTIONS -> IO ()++-------------------------------------------------------------------------------+-- filter policy+-------------------------------------------------------------------------------++data {-# CTYPE "rocksdb/c.h" "rocksdb_filterpolicy_t" #-} FILTERPOLICY++foreign import capi "rocksdb/c.h rocksdb_filterpolicy_destroy"+ rocksdb_filterpolicy_destroy :: Ptr FILTERPOLICY -> IO ()++foreign import capi "rocksdb/c.h rocksdb_filterpolicy_create_bloom"+ rocksdb_filterpolicy_create_bloom :: CInt -> IO (Ptr FILTERPOLICY)
@@ -0,0 +1,2819 @@+{- |+Module : Database.LSMTree+Copyright : (c) 2023-2025, Cardano Development Foundation+License : Apache-2.0+Stability : experimental+Portability : portable+-}+module Database.LSMTree (+ -- * Usage Notes+ -- $usage_notes++ -- ** Real and Simulated IO+ -- $real_and_simulated_io+ IOLike,++ -- * Examples+ -- $setup++ -- * Sessions+ Session,+ withOpenSession,+ withOpenSessionIO,+ withNewSession,+ withRestoreSession,+ openSession,+ newSession,+ restoreSession,+ closeSession,++ -- * Tables+ Table,+ withTable,+ withTableWith,+ newTable,+ newTableWith,+ closeTable,++ -- ** Table Lookups #table_lookups#+ member,+ members,+ LookupResult (..),+ getValue,+ getBlob,+ lookup,+ lookups,+ Entry (..),+ getEntryKey,+ getEntryValue,+ getEntryBlob,+ rangeLookup,++ -- ** Table Updates #table_updates#+ insert,+ inserts,+ upsert,+ upserts,+ delete,+ deletes,+ Update (..),+ update,+ updates,++ -- ** Table Duplication #table_duplication#+ withDuplicate,+ duplicate,++ -- ** Table Unions #table_unions#+ withUnion,+ withUnions,+ union,+ unions,+ withIncrementalUnion,+ withIncrementalUnions,+ incrementalUnion,+ incrementalUnions,+ remainingUnionDebt,+ supplyUnionCredits,++ -- * Blob References #blob_references#+ BlobRef,+ retrieveBlob,+ retrieveBlobs,++ -- * Cursors #cursor#+ Cursor,+ withCursor,+ withCursorAtOffset,+ newCursor,+ newCursorAtOffset,+ closeCursor,+ next,+ take,+ takeWhile,++ -- * Snapshots #snapshots#+ saveSnapshot,+ withTableFromSnapshot,+ withTableFromSnapshotWith,+ openTableFromSnapshot,+ openTableFromSnapshotWith,+ doesSnapshotExist,+ deleteSnapshot,+ listSnapshots,+ SnapshotName,+ isValidSnapshotName,+ toSnapshotName,+ SnapshotLabel (..),++ -- * Session Configuration #session_configuration#+ Salt,++ -- * Table Configuration #table_configuration#+ TableConfig (+ confMergePolicy,+ confSizeRatio,+ confWriteBufferAlloc,+ confBloomFilterAlloc,+ confFencePointerIndex,+ confDiskCachePolicy,+ confMergeSchedule,+ confMergeBatchSize+ ),+ defaultTableConfig,+ MergePolicy (LazyLevelling),+ MergeSchedule (..),+ SizeRatio (Four),+ WriteBufferAlloc (AllocNumEntries),+ BloomFilterAlloc (AllocFixed, AllocRequestFPR),+ FencePointerIndexType (OrdinaryIndex, CompactIndex),+ DiskCachePolicy (..),+ MergeBatchSize (..),++ -- ** Table Configuration Overrides #table_configuration_overrides#+ TableConfigOverride (..),+ noTableConfigOverride,++ -- * Ranges #ranges#+ Range (..),++ -- * Union Credit and Debt+ UnionCredits (..),+ UnionDebt (..),++ -- * Key\/Value Serialisation #key_value_serialisation#+ RawBytes (RawBytes),+ SerialiseKey (serialiseKey, deserialiseKey),+ SerialiseKeyOrderPreserving,+ SerialiseValue (serialiseValue, deserialiseValue),++ -- ** Key\/Value Serialisation Property Tests #key_value_serialisation_property_tests#+ serialiseKeyIdentity,+ serialiseKeyIdentityUpToSlicing,+ serialiseKeyPreservesOrdering,+ serialiseValueIdentity,+ serialiseValueIdentityUpToSlicing,+ packSlice,++ -- * Monoidal Value Resolution #monoidal_value_resolution#+ ResolveValue (..),+ ResolveViaSemigroup (..),+ ResolveAsFirst (..),++ -- ** Monoidal Value Resolution Property Tests #monoidal_value_resolution_property_tests#+ resolveCompatibility,+ resolveValidOutput,+ resolveAssociativity,++ -- * Errors #errors#+ SessionDirDoesNotExistError (..),+ SessionDirLockedError (..),+ SessionDirCorruptedError (..),+ SessionClosedError (..),+ TableClosedError (..),+ TableCorruptedError (..),+ TableTooLargeError (..),+ TableUnionNotCompatibleError (..),+ SnapshotExistsError (..),+ SnapshotDoesNotExistError (..),+ SnapshotCorruptedError (..),+ SnapshotNotCompatibleError (..),+ BlobRefInvalidError (..),+ CursorClosedError (..),+ InvalidSnapshotNameError (..),++ -- * Traces #traces#+ Tracer,+ LSMTreeTrace (..),+ SessionTrace (..),+ TableTrace (..),+ CursorTrace (..),+ SessionId (..),+ TableId (..),+ CursorId (..),+) where++import Control.Concurrent.Class.MonadMVar.Strict (MonadMVar)+import Control.Concurrent.Class.MonadSTM (MonadSTM (STM))+import Control.DeepSeq (NFData (..))+import Control.Exception.Base (assert)+import Control.Monad.Class.MonadAsync (MonadAsync)+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadThrow (MonadCatch (..), MonadEvaluate,+ MonadMask, MonadThrow (..))+import Control.Monad.Primitive (PrimMonad)+import Control.Tracer (Tracer)+import Data.Bifunctor (Bifunctor (..))+import Data.Coerce (coerce)+import Data.Kind (Constraint, Type)+import Data.List.NonEmpty (NonEmpty (..))+import Data.Maybe (fromMaybe, isJust)+import Data.Typeable (Proxy (..), Typeable, eqT, type (:~:) (Refl),+ typeRep)+import Data.Vector (Vector)+import qualified Data.Vector as V+import qualified Database.LSMTree.Internal.BlobRef as Internal+import Database.LSMTree.Internal.Config+ (BloomFilterAlloc (AllocFixed, AllocRequestFPR),+ DiskCachePolicy (..), FencePointerIndexType (..),+ MergeBatchSize (..), MergePolicy (..), MergeSchedule (..),+ SizeRatio (..), TableConfig (..), WriteBufferAlloc (..),+ defaultTableConfig)+import Database.LSMTree.Internal.Config.Override+ (TableConfigOverride (..), noTableConfigOverride)+import qualified Database.LSMTree.Internal.Entry as Entry+import Database.LSMTree.Internal.Paths (SnapshotName,+ isValidSnapshotName, toSnapshotName)+import Database.LSMTree.Internal.Range (Range (..))+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import Database.LSMTree.Internal.RunNumber (CursorId (..),+ TableId (..))+import qualified Database.LSMTree.Internal.Serialise as Internal+import Database.LSMTree.Internal.Serialise.Class (SerialiseKey (..),+ SerialiseKeyOrderPreserving, SerialiseValue (..),+ packSlice, serialiseKeyIdentity,+ serialiseKeyIdentityUpToSlicing,+ serialiseKeyPreservesOrdering, serialiseValueIdentity,+ serialiseValueIdentityUpToSlicing)+import Database.LSMTree.Internal.Snapshot (SnapshotLabel (..))+import Database.LSMTree.Internal.Types (BlobRef (..), Cursor (..),+ ResolveAsFirst (..), ResolveValue (..),+ ResolveViaSemigroup (..), Salt, Session (..), Table (..),+ resolveAssociativity, resolveCompatibility,+ resolveValidOutput)+import Database.LSMTree.Internal.Unsafe (BlobRefInvalidError (..),+ CursorClosedError (..), CursorTrace,+ InvalidSnapshotNameError (..), LSMTreeTrace (..),+ ResolveSerialisedValue, SessionClosedError (..),+ SessionDirCorruptedError (..),+ SessionDirDoesNotExistError (..),+ SessionDirLockedError (..), SessionId (..),+ SessionTrace (..), SnapshotCorruptedError (..),+ SnapshotDoesNotExistError (..), SnapshotExistsError (..),+ SnapshotNotCompatibleError (..), TableClosedError (..),+ TableCorruptedError (..), TableTooLargeError (..),+ TableTrace, TableUnionNotCompatibleError (..),+ UnionCredits (..), UnionDebt (..))+import qualified Database.LSMTree.Internal.Unsafe as Internal+import Prelude hiding (lookup, take, takeWhile)+import System.FS.API (FsPath, HasFS (..), MountPoint (..), mkFsPath)+import System.FS.BlockIO.API (HasBlockIO (..))+import System.FS.BlockIO.IO (defaultIOCtxParams, withIOHasBlockIO)+import System.FS.IO (HandleIO)+import System.Random (randomIO)++--------------------------------------------------------------------------------+-- Usage Notes+--------------------------------------------------------------------------------++{- $usage_notes+This section focuses on the differences between the full API as defined in this module and the simple API as defined in "Database.LSMTree.Simple".+It assumes that the reader is familiar with [Usage Notes for the simple API]("Database.LSMTree.Simple#g:usage_notes"), which discusses crucial topics such as [Resource Management]("Database.LSMTree.Simple#g:resource_management"), [Concurrency]("Database.LSMTree.Simple#g:concurrency"), [ACID properties]("Database.LSMTree.Simple#g:acid"), and [Sharing]("Database.LSMTree.Simple#g:sharing").+-}++{- $real_and_simulated_io+-}++type IOLike :: (Type -> Type) -> Constraint+type IOLike m =+ ( MonadAsync m+ , MonadMVar m+ , MonadThrow m+ , MonadThrow (STM m)+ , MonadCatch m+ , MonadMask m+ , PrimMonad m+ , MonadST m+ , MonadEvaluate m+ )++--------------------------------------------------------------------------------+-- Example+--------------------------------------------------------------------------------++{- $setup++The examples in this module use the preamble described in this section, which does three things:++1. It imports this module qualified, as intended, as well as any other relevant modules.+2. It defines types for keys, values, and BLOBs.+3. It defines a helper function that runs examples with access to an open session and fresh table.++=== Importing "Database.LSMTree"++This module is intended to be imported qualified, to avoid name clashes with Prelude functions.++>>> import Database.LSMTree (BlobRef, Cursor, RawBytes, ResolveValue (..), SerialiseKey (..), SerialiseValue (..), Session, Table)+>>> import qualified Database.LSMTree as LSMT++=== Defining key, value, and BLOB types++The examples in this module use the types @Key@, @Value@, and @Blob@ for keys, values and BLOBs.++>>> import Data.ByteString (ByteString)+>>> import Data.ByteString.Short (ShortByteString)+>>> import Data.Proxy (Proxy)+>>> import Data.String (IsString)+>>> import Data.Word (Word64)++The type @Key@ is a newtype wrapper around 'Data.Word.Word64'.+The required instance of 'SerialiseKey' is derived by @GeneralisedNewtypeDeriving@ from the preexisting instance for 'Data.Word.Word64'.++>>> :{+newtype Key = Key Word64+ deriving stock (Eq, Ord, Show)+ deriving newtype (Num, SerialiseKey)+:}++The type @Value@ is a newtype wrapper around 'Data.ByteString.Short.ShortByteString'.+The required instance of 'SerialiseValue' is derived by @GeneralisedNewtypeDeriving@ from the preexisting instance for 'Data.ByteString.Short.ShortByteString'.++>>> :{+newtype Value = Value ShortByteString+ deriving stock (Eq, Show)+ deriving newtype (IsString, SerialiseValue)+:}++The type @Value@ has an instance of @ResolveValue@ which appends the new value to the old value separated by a space.+It is sufficient to define either 'resolve' or 'resolveSerialised',+as each can be defined in terms of the other and 'serialiseValue'\/'deserialiseValue'.+For optimal performance, you should /always/ define 'resolveSerialised' manually.++__NOTE__:+The /first/ argument of 'resolve' and 'resolveSerialised' is the /new/ value and the /second/ argument is the /old/ value.++>>> :{+instance ResolveValue Value where+ resolve :: Value -> Value -> Value+ resolve (Value new) (Value old) = Value (new <> " " <> old)+ resolveSerialised :: Proxy Value -> RawBytes -> RawBytes -> RawBytes+ resolveSerialised _ new old = new <> " " <> old+:}++The type @Blob@ is a newtype wrapper around 'Data.ByteString.ByteString',+The required instance of 'SerialiseValue' is derived by @GeneralisedNewtypeDeriving@ from the preexisting instance for 'Data.ByteString.ByteString'.++>>> :{+newtype Blob = Blob ByteString+ deriving stock (Eq, Show)+ deriving newtype (IsString, SerialiseValue)+:}++=== Defining a helper function to run examples++The examples in this module are wrapped in a call to @runExample@,+which creates a temporary session directory and+runs the example with access to an open 'Session' and a fresh 'Table'.++>>> import Control.Exception (bracket, bracket_)+>>> import Data.Foldable (traverse_)+>>> import qualified System.Directory as Dir+>>> import System.FilePath ((</>))+>>> import System.Process (getCurrentPid)+>>> :{+runExample :: (Session IO -> Table IO Key Value Blob -> IO a) -> IO a+runExample action = do+ tmpDir <- Dir.getTemporaryDirectory+ pid <- getCurrentPid+ let sessionDir = tmpDir </> "doctest_Database_LSMTree" </> show pid+ let createSessionDir = Dir.createDirectoryIfMissing True sessionDir+ let removeSessionDir = Dir.removeDirectoryRecursive sessionDir+ bracket_ createSessionDir removeSessionDir $ do+ LSMT.withOpenSessionIO mempty sessionDir $ \session -> do+ LSMT.withTable session $ \table ->+ action session table+:}+-}++--------------------------------------------------------------------------------+-- Sessions+--------------------------------------------------------------------------------++-- NOTE: 'Session' is defined in 'Database.LSMTree.Internal.Types'++{- |+Run an action with access to a session opened from a session directory.++If the session directory is empty, a new session is created using the given salt.+Otherwise, the session directory is restored as an existing session ignoring the given salt.++If there are no open tables or cursors when the session terminates, then the disk I\/O complexity of this operation is \(O(1)\).+Otherwise, 'closeTable' is called for each open table and 'closeCursor' is called for each open cursor.+Consequently, the worst-case disk I\/O complexity of this operation depends on the merge policy of the open tables in the session.+The following assumes all tables in the session have the same merge policy:++['LazyLevelling']:+ \(O(o \: T \log_T \frac{n}{B})\).++The variable \(o\) refers to the number of open tables and cursors in the session.++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'openSession' and 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+{-# SPECIALISE+ withOpenSession ::+ Tracer IO LSMTreeTrace ->+ HasFS IO HandleIO ->+ HasBlockIO IO HandleIO ->+ Salt ->+ FsPath ->+ (Session IO -> IO a) ->+ IO a+ #-}+withOpenSession ::+ forall m h a.+ (IOLike m, Typeable h) =>+ Tracer m LSMTreeTrace ->+ HasFS m h ->+ HasBlockIO m h ->+ -- | The session salt.+ Salt ->+ -- | The session directory.+ FsPath ->+ (Session m -> m a) ->+ m a+withOpenSession tracer hasFS hasBlockIO sessionSalt sessionDir action = do+ Internal.withOpenSession tracer hasFS hasBlockIO sessionSalt sessionDir (action . Session)++-- | Variant of 'withOpenSession' that is specialised to 'IO' using the real filesystem.+withOpenSessionIO ::+ Tracer IO LSMTreeTrace ->+ FilePath ->+ (Session IO -> IO a) ->+ IO a+withOpenSessionIO tracer sessionDir action = do+ let mountPoint = MountPoint sessionDir+ let sessionDirFsPath = mkFsPath []+ sessionSalt <- randomIO+ withIOHasBlockIO mountPoint defaultIOCtxParams $ \hasFS hasBlockIO ->+ withOpenSession tracer hasFS hasBlockIO sessionSalt sessionDirFsPath action++{- |+Run an action with access to a new session.++The session directory must be empty.++If there are no open tables or cursors when the session terminates, then the disk I\/O complexity of this operation is \(O(1)\).+Otherwise, 'closeTable' is called for each open table and 'closeCursor' is called for each open cursor.+Consequently, the worst-case disk I\/O complexity of this operation depends on the merge policy of the open tables in the session.+The following assumes all tables in the session have the same merge policy:++['LazyLevelling']:+ \(O(o \: T \log_T \frac{n}{B})\).++The variable \(o\) refers to the number of open tables and cursors in the session.++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'newSession' and 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+{-# SPECIALISE+ withNewSession ::+ Tracer IO LSMTreeTrace ->+ HasFS IO HandleIO ->+ HasBlockIO IO HandleIO ->+ Salt ->+ FsPath ->+ (Session IO -> IO a) ->+ IO a+ #-}+withNewSession ::+ forall m h a.+ (IOLike m, Typeable h) =>+ Tracer m LSMTreeTrace ->+ HasFS m h ->+ HasBlockIO m h ->+ -- | The session salt.+ Salt ->+ -- | The session directory.+ FsPath ->+ (Session m -> m a) ->+ m a+withNewSession tracer hasFS hasBlockIO sessionSalt sessionDir action = do+ Internal.withNewSession tracer hasFS hasBlockIO sessionSalt sessionDir (action . Session)++{- |+Run an action with access to a restored session.++The session directory must be non-empty: a session must have previously been+opened and closed in this directory.++If there are no open tables or cursors when the session terminates, then the disk I\/O complexity of this operation is \(O(1)\).+Otherwise, 'closeTable' is called for each open table and 'closeCursor' is called for each open cursor.+Consequently, the worst-case disk I\/O complexity of this operation depends on the merge policy of the open tables in the session.+The following assumes all tables in the session have the same merge policy:++['LazyLevelling']:+ \(O(o \: T \log_T \frac{n}{B})\).++The variable \(o\) refers to the number of open tables and cursors in the session.++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'restoreSession' and 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+{-# SPECIALISE+ withRestoreSession ::+ Tracer IO LSMTreeTrace ->+ HasFS IO HandleIO ->+ HasBlockIO IO HandleIO ->+ FsPath ->+ (Session IO -> IO a) ->+ IO a+ #-}+withRestoreSession ::+ forall m h a.+ (IOLike m, Typeable h) =>+ Tracer m LSMTreeTrace ->+ HasFS m h ->+ HasBlockIO m h ->+ -- | The session directory.+ FsPath ->+ (Session m -> m a) ->+ m a+withRestoreSession tracer hasFS hasBlockIO sessionDir action = do+ Internal.withRestoreSession tracer hasFS hasBlockIO sessionDir (action . Session)++{- |+Open a session from a session directory.++If the session directory is empty, a new session is created using the given salt.+Otherwise, the session directory is restored as an existing session ignoring the given salt.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ Sessions hold open resources and must be closed using 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+{-# SPECIALISE+ openSession ::+ Tracer IO LSMTreeTrace ->+ HasFS IO HandleIO ->+ HasBlockIO IO HandleIO ->+ Salt ->+ FsPath ->+ IO (Session IO)+ #-}+openSession ::+ forall m h.+ (IOLike m, Typeable h) =>+ Tracer m LSMTreeTrace ->+ HasFS m h ->+ HasBlockIO m h ->+ -- | The session salt.+ Salt ->+ -- | The session directory.+ FsPath ->+ m (Session m)+openSession tracer hasFS hasBlockIO sessionSalt sessionDir =+ Session <$> Internal.openSession tracer hasFS hasBlockIO sessionSalt sessionDir++{- |+Create a new session.++The session directory must be empty.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ Sessions hold open resources and must be closed using 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+{-# SPECIALISE+ newSession ::+ Tracer IO LSMTreeTrace ->+ HasFS IO HandleIO ->+ HasBlockIO IO HandleIO ->+ Salt ->+ FsPath ->+ IO (Session IO)+ #-}+newSession ::+ forall m h.+ (IOLike m, Typeable h) =>+ Tracer m LSMTreeTrace ->+ HasFS m h ->+ HasBlockIO m h ->+ -- | The session salt.+ Salt ->+ -- | The session directory.+ FsPath ->+ m (Session m)+newSession tracer hasFS hasBlockIO sessionSalt sessionDir =+ Session <$> Internal.newSession tracer hasFS hasBlockIO sessionSalt sessionDir++{- |+Restore a session from a session directory.++The session directory must be non-empty: a session must have previously been+opened (and closed) in this directory.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ Sessions hold open resources and must be closed using 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+{-# SPECIALISE+ restoreSession ::+ Tracer IO LSMTreeTrace ->+ HasFS IO HandleIO ->+ HasBlockIO IO HandleIO ->+ FsPath ->+ IO (Session IO)+ #-}+restoreSession ::+ forall m h.+ (IOLike m, Typeable h) =>+ Tracer m LSMTreeTrace ->+ HasFS m h ->+ HasBlockIO m h ->+ -- | The session directory.+ FsPath ->+ m (Session m)+restoreSession tracer hasFS hasBlockIO sessionDir =+ Session <$> Internal.restoreSession tracer hasFS hasBlockIO sessionDir++{- |+Close a session.++If there are no open tables or cursors in the session, then the disk I\/O complexity of this operation is \(O(1)\).+Otherwise, 'closeTable' is called for each open table and 'closeCursor' is called for each open cursor.+Consequently, the worst-case disk I\/O complexity of this operation depends on the merge policy of the tables in the session.+The following assumes all tables in the session have the same merge policy:++['LazyLevelling']:+ \(O(o \: T \log_T \frac{n}{B})\).++The variable \(o\) refers to the number of open tables and cursors in the session.++Closing is idempotent, i.e., closing a closed session does nothing.+All other operations on a closed session will throw an exception.+-}+{-# SPECIALISE+ closeSession ::+ Session IO ->+ IO ()+ #-}+closeSession ::+ forall m.+ (IOLike m) =>+ Session m ->+ m ()+closeSession (Session session) =+ Internal.closeSession session++--------------------------------------------------------------------------------+-- Tables+--------------------------------------------------------------------------------++-- NOTE: 'Table' is defined in 'Database.LSMTree.Internal.Types'++{- |+Run an action with access to an empty table.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'newTable' and 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+{-# SPECIALISE+ withTable ::+ Session IO ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withTable ::+ forall m k v b a.+ (IOLike m) =>+ Session m ->+ (Table m k v b -> m a) ->+ m a+withTable session =+ withTableWith defaultTableConfig session++-- | Variant of 'withTable' that accepts [table configuration](#g:table_configuration).+{-# SPECIALISE+ withTableWith ::+ TableConfig ->+ Session IO ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withTableWith ::+ forall m k v b a.+ (IOLike m) =>+ TableConfig ->+ Session m ->+ (Table m k v b -> m a) ->+ m a+withTableWith tableConfig (Session session) action =+ Internal.withTable session tableConfig (action . Table)++{- |+Create an empty table.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ Tables hold open resources and must be closed using 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+{-# SPECIALISE+ newTable ::+ Session IO ->+ IO (Table IO k v b)+ #-}+newTable ::+ forall m k v b.+ (IOLike m) =>+ Session m ->+ m (Table m k v b)+newTable session =+ newTableWith defaultTableConfig session++{- |+Variant of 'newTable' that accepts [table configuration](#g:table_configuration).+-}+{-# SPECIALISE+ newTableWith ::+ TableConfig ->+ Session IO ->+ IO (Table IO k v b)+ #-}+newTableWith ::+ forall m k v b.+ (IOLike m) =>+ TableConfig ->+ Session m ->+ m (Table m k v b)+newTableWith tableConfig (Session session) =+ Table <$> Internal.new session tableConfig++{- |+Close a table.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Closing is idempotent, i.e., closing a closed table does nothing.+All other operations on a closed table will throw an exception.++__Warning:__ Tables are ephemeral. Once you close a table, its data is lost forever. To persist tables, use [snapshots](#g:snapshots).+-}+{-# SPECIALISE+ closeTable ::+ Table IO k v b ->+ IO ()+ #-}+closeTable ::+ forall m k v b.+ (IOLike m) =>+ Table m k v b ->+ m ()+closeTable (Table table) =+ Internal.close table++--------------------------------------------------------------------------------+-- Lookups+--------------------------------------------------------------------------------++{- |+Check if the key is a member of the table.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ print =<< LSMT.member table 0+:}+True++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Membership tests can be performed concurrently from multiple Haskell threads.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableCorruptedError']:+ If the table data is corrupted.+-}+{-# SPECIALISE+ member ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table IO k v b ->+ k ->+ IO Bool+ #-}+member ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table m k v b ->+ k ->+ m Bool+member =+ -- Technically, this does not need the 'SerialiseValue' constraint.+ (fmap (isJust . getValue) .) . lookup++{- |+Variant of 'member' for batch membership tests.+The batch of keys corresponds in-order to the batch of results.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(b \: T \log_T \frac{n}{B})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> members table keys = traverse (member table) keys+-}+{-# SPECIALISE+ members ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table IO k v b ->+ Vector k ->+ IO (Vector Bool)+ #-}+members ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table m k v b ->+ Vector k ->+ m (Vector Bool)+members =+ -- Technically, this does not need the 'SerialiseValue' constraint.+ (fmap (fmap (isJust . getValue)) .) . lookups++data LookupResult v b+ = NotFound+ | Found !v+ | FoundWithBlob !v !b+ deriving stock (Eq, Show, Functor, Foldable, Traversable)++{- |+Get the field of type @v@ from a @'LookupResult' v b@, if any.+-}+getValue :: LookupResult v b -> Maybe v+getValue = \case+ NotFound -> Nothing+ Found !v -> Just v+ FoundWithBlob !v !_b -> Just v++{- |+Get the field of type @b@ from a @'LookupResult' v b@, if any.++The following property holds:++prop> isJust (getBlob result) <= isJust (getValue result)+-}+getBlob :: LookupResult v b -> Maybe b+getBlob = \case+ NotFound -> Nothing+ Found !_v -> Nothing+ FoundWithBlob !_v !b -> Just b++instance (NFData v, NFData b) => NFData (LookupResult v b) where+ rnf :: LookupResult v b -> ()+ rnf = \case+ NotFound -> ()+ Found v -> rnf v+ FoundWithBlob v b -> rnf v `seq` rnf b++instance Bifunctor LookupResult where+ bimap :: (v -> v') -> (b -> b') -> LookupResult v b -> LookupResult v' b'+ bimap f g = \case+ NotFound -> NotFound+ Found v -> Found (f v)+ FoundWithBlob v b -> FoundWithBlob (f v) (g b)++{- |+Look up the value associated with a key.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ print =<< LSMT.lookup table 0+:}+Found (Value "Hello")++If the key is not associated with any value, 'lookup' returns 'NotFound'.++>>> :{+runExample $ \session table -> do+ LSMT.lookup table 0+:}+NotFound++If the key has an associated BLOB, the result contains a 'BlobRef'.+The full BLOB can be retrieved by passing that 'BlobRef' to 'retrieveBlob'.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" (Just "World")+ print+ =<< traverse (LSMT.retrieveBlob session)+ =<< LSMT.lookup table 0+:}+FoundWithBlob (Value "Hello") (Blob "World")++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Lookups can be performed concurrently from multiple Haskell threads.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableCorruptedError']:+ If the table data is corrupted.+-}+{-# SPECIALISE+ lookup ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table IO k v b ->+ k ->+ IO (LookupResult v (BlobRef IO b))+ #-}+lookup ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table m k v b ->+ k ->+ m (LookupResult v (BlobRef m b))+lookup table k = do+ mvs <- lookups table (V.singleton k)+ let mmv = fst <$> V.uncons mvs+ pure $ fromMaybe NotFound mmv++{- |+Variant of 'lookup' for batch lookups.+The batch of keys corresponds in-order to the batch of results.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(b \: T \log_T \frac{n}{B})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> lookups table keys = traverse (lookup table) keys+-}+{-# SPECIALISE+ lookups ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table IO k v b ->+ Vector k ->+ IO (Vector (LookupResult v (BlobRef IO b)))+ #-}+lookups ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table m k v b ->+ Vector k ->+ m (Vector (LookupResult v (BlobRef m b)))+lookups (Table table :: Table m k v b) keys = do+ maybeEntries <- Internal.lookups (_getResolveSerialisedValue (Proxy @v)) (fmap Internal.serialiseKey keys) table+ pure $ maybe NotFound entryToLookupResult <$> maybeEntries+ where+ entryToLookupResult = \case+ Entry.Insert !v -> Found (Internal.deserialiseValue v)+ Entry.InsertWithBlob !v !b -> FoundWithBlob (Internal.deserialiseValue v) (BlobRef b)+ Entry.Upsert !v -> Found (Internal.deserialiseValue v)+ Entry.Delete -> NotFound++data Entry k v b+ = Entry !k !v+ | EntryWithBlob !k !v !b+ deriving stock (Eq, Show, Functor, Foldable, Traversable)++{- |+Get the field of type @k@ from an @'Entry' k v b@.+-}+getEntryKey :: Entry k v b -> k+getEntryKey (Entry !k !_v) = k+getEntryKey (EntryWithBlob !k !_v !_b) = k++{- |+Get the field of type @v@ from an @'Entry' k v b@.+-}+getEntryValue :: Entry k v b -> v+getEntryValue (Entry !_k !v) = v+getEntryValue (EntryWithBlob !_k !v !_b) = v++{- |+Get the field of type @b@ from an @'Entry' k v b@, if any.+-}+getEntryBlob :: Entry k v b -> Maybe b+getEntryBlob (Entry !_k !_v) = Nothing+getEntryBlob (EntryWithBlob !_k !_v !b) = Just b++instance (NFData k, NFData v, NFData b) => NFData (Entry k v b) where+ rnf :: Entry k v b -> ()+ rnf = \case+ Entry k v -> rnf k `seq` rnf v+ EntryWithBlob k v b -> rnf k `seq` rnf v `seq` rnf b++instance Bifunctor (Entry k) where+ bimap :: (v -> v') -> (b -> b') -> Entry k v b -> Entry k v' b'+ bimap f g = \case+ Entry k v -> Entry k (f v)+ EntryWithBlob k v b -> EntryWithBlob k (f v) (g b)++{- |+Look up a batch of values associated with keys in the given range.++The worst-case disk I\/O complexity of this operation is \(O(T \log_T \frac{n}{B} + \frac{b}{P})\),+where the variable \(b\) refers to the length of the /output/ vector.++Range lookups can be performed concurrently from multiple Haskell threads.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableCorruptedError']:+ If the table data is corrupted.+-}+{-# SPECIALISE+ rangeLookup ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table IO k v b ->+ Range k ->+ IO (Vector (Entry k v (BlobRef IO b)))+ #-}+rangeLookup ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Table m k v b ->+ Range k ->+ m (Vector (Entry k v (BlobRef m b)))+rangeLookup (Table table :: Table m k v b) range =+ Internal.rangeLookup (_getResolveSerialisedValue (Proxy @v)) (Internal.serialiseKey <$> range) table $ \ !k !v -> \case+ Just !b -> EntryWithBlob (Internal.deserialiseKey k) (Internal.deserialiseValue v) (BlobRef b)+ Nothing -> Entry (Internal.deserialiseKey k) (Internal.deserialiseValue v)++--------------------------------------------------------------------------------+-- Updates+--------------------------------------------------------------------------------++{- |+Insert associates the given value and BLOB with the given key in the table.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ print =<< LSMT.lookup table 0+:}+Found (Value "Hello")++Insert may optionally associate a BLOB value with the given key.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" (Just "World")+ print+ =<< traverse (retrieveBlob session)+ =<< LSMT.lookup table 0+:}+FoundWithBlob (Value "Hello") (Blob "World")++Insert overwrites any value and BLOB previously associated with the given key,+even if the given BLOB is 'Nothing'.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" (Just "World")+ LSMT.insert table 0 "Goodbye" Nothing+ print+ =<< traverse (retrieveBlob session)+ =<< LSMT.lookup table 0+:}+Found (Value "Goodbye")++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ insert ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ k ->+ v ->+ Maybe b ->+ IO ()+ #-}+insert ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ k ->+ v ->+ Maybe b ->+ m ()+insert table k v b =+ inserts table (V.singleton (k, v, b))++{- |+Variant of 'insert' for batch insertions.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> inserts table entries = traverse_ (uncurry $ insert table) entries+-}+{-# SPECIALISE+ inserts ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ Vector (k, v, Maybe b) ->+ IO ()+ #-}+inserts ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ Vector (k, v, Maybe b) ->+ m ()+inserts table entries =+ updates table (fmap (\(k, v, mb) -> (k, Insert v mb)) entries)++{- |+If the given key is not a member of the table, 'upsert' associates the given value with the given key in the table.+Otherwise, 'upsert' updates the value associated with the given key by combining it with the given value using 'resolve'.++>>> :{+runExample $ \session table -> do+ LSMT.upsert table 0 "Hello"+ LSMT.upsert table 0 "Goodbye"+ print =<< LSMT.lookup table 0+:}+Found (Value "Goodbye Hello")++__Warning:__+Upsert deletes any BLOB previously associated with the given key.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" (Just "World")+ LSMT.upsert table 0 "Goodbye"+ print+ =<< traverse (LSMT.retrieveBlob session)+ =<< LSMT.lookup table 0+:}+Found (Value "Goodbye Hello")++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.++The following property holds in the absence of races:++@+upsert table k v = do+ r <- lookup table k+ let v' = maybe v (resolve v) (getValue r)+ insert table k v' Nothing+@+-}+{-# SPECIALISE+ upsert ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ k ->+ v ->+ IO ()+ #-}+upsert ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ k ->+ v ->+ m ()+upsert table k v =+ upserts table (V.singleton (k, v))++{- |+Variant of 'upsert' for batch insertions.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> upserts table entries = traverse_ (uncurry $ upsert table) entries+-}+{-# SPECIALISE+ upserts ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ Vector (k, v) ->+ IO ()+ #-}+upserts ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ Vector (k, v) ->+ m ()+upserts table entries =+ updates table (second Upsert <$> entries)++{- |+Delete a key from the table.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.delete table 0+ print =<< LSMT.lookup table 0+:}+NotFound++If the key is not a member of the table, the table is left unchanged.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.delete table 1+ print =<< LSMT.lookup table 0+:}+Found (Value "Hello")++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ delete ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ k ->+ IO ()+ #-}+delete ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ k ->+ m ()+delete table k =+ deletes table (V.singleton k)++{- |+Variant of 'delete' for batch deletions.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> deletes table keys = traverse_ (delete table) keys+-}+{-# SPECIALISE+ deletes ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ Vector k ->+ IO ()+ #-}+deletes ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ Vector k ->+ m ()+deletes table entries =+ updates table (fmap (,Delete) entries)++type Update :: Type -> Type -> Type+data Update v b+ = Insert !v !(Maybe b)+ | Delete+ | Upsert !v+ deriving stock (Show, Eq)++instance (NFData v, NFData b) => NFData (Update v b) where+ rnf :: Update v b -> ()+ rnf = \case+ Insert v mb -> rnf v `seq` rnf mb+ Delete -> ()+ Upsert v -> rnf v++{- |+Update generalises 'insert', 'delete', and 'upsert'.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++The following properties hold:++prop> update table k (Insert v mb) = insert table k v mb+prop> update table k Delete = delete table k+prop> update table k (Upsert v) = upsert table k v++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ update ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ k ->+ Update v b ->+ IO ()+ #-}+update ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ k ->+ Update v b ->+ m ()+update table k mv =+ updates table (V.singleton (k, mv))++{- |+Variant of 'update' for batch updates.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> updates table entries = traverse_ (uncurry $ update table) entries+-}+{-# SPECIALISE+ updates ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table IO k v b ->+ Vector (k, Update v b) ->+ IO ()+ #-}+updates ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v, SerialiseValue b) =>+ Table m k v b ->+ Vector (k, Update v b) ->+ m ()+updates (Table table :: Table m k v b) entries =+ Internal.updates (_getResolveSerialisedValue (Proxy @v)) (serialiseEntry <$> entries) table+ where+ serialiseEntry (k, u) = (Internal.serialiseKey k, serialiseUpdate u)+ serialiseUpdate = \case+ Insert v (Just b) -> Entry.InsertWithBlob (Internal.serialiseValue v) (Internal.serialiseBlob b)+ Insert v Nothing -> Entry.Insert (Internal.serialiseValue v)+ Delete -> Entry.Delete+ Upsert v -> Entry.Upsert (Internal.serialiseValue v)++--------------------------------------------------------------------------------+-- Duplication+--------------------------------------------------------------------------------++{- |+Run an action with access to the duplicate of a table.++The duplicate is an independent copy of the given table.+Subsequent updates to the original table do not affect the duplicate, and vice versa.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.withDuplicate table $ \table' -> do+ print =<< LSMT.lookup table' 0+ LSMT.insert table' 0 "Goodbye" Nothing+ print =<< LSMT.lookup table' 0+ LSMT.lookup table 0+ print =<< LSMT.lookup table 0+:}+Found (Value "Hello")+Found (Value "Goodbye")+Found (Value "Hello")++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'duplicate' and 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ withDuplicate ::+ Table IO k v b ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withDuplicate ::+ forall m k v b a.+ (IOLike m) =>+ Table m k v b ->+ (Table m k v b -> m a) ->+ m a+withDuplicate table =+ bracket (duplicate table) closeTable++{- |+Duplicate a table.++The duplicate is an independent copy of the given table.+Subsequent updates to the original table do not affect the duplicate, and vice versa.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ bracket (LSMT.duplicate table) LSMT.closeTable $ \table' -> do+ print =<< LSMT.lookup table' 0+ LSMT.insert table' 0 "Goodbye" Nothing+ print =<< LSMT.lookup table' 0+ LSMT.lookup table 0+ print =<< LSMT.lookup table 0+:}+Found (Value "Hello")+Found (Value "Goodbye")+Found (Value "Hello")++The worst-case disk I\/O complexity of this operation is \(O(0)\).++__Warning:__ The duplicate must be independently closed using 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ duplicate ::+ Table IO k v b ->+ IO (Table IO k v b)+ #-}+duplicate ::+ forall m k v b.+ (IOLike m) =>+ Table m k v b ->+ m (Table m k v b)+duplicate (Table table) =+ Table <$> Internal.duplicate table++--------------------------------------------------------------------------------+-- Union+--------------------------------------------------------------------------------++{- |+Run an action with access to a table that contains the union of the entries of the given tables.++>>> :{+runExample $ \session table1 -> do+ LSMT.insert table1 0 "Hello" Nothing+ LSMT.withTable session $ \table2 -> do+ LSMT.insert table2 0 "World" Nothing+ LSMT.insert table2 1 "Goodbye" Nothing+ LSMT.withUnion table1 table2 $ \table3 -> do+ print =<< LSMT.lookup table3 0+ print =<< LSMT.lookup table3 1+ print =<< LSMT.lookup table1 0+ print =<< LSMT.lookup table2 0+:}+Found (Value "Hello World")+Found (Value "Goodbye")+Found (Value "Hello")+Found (Value "World")++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'union' and 'closeTable'.++__Warning:__ Both input tables must be from the same 'Session'.++__Warning:__ This is a relatively expensive operation that may take some time to complete.+See 'withIncrementalUnion' for an incremental alternative.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+{-# SPECIALISE+ withUnion ::+ (ResolveValue v) =>+ Table IO k v b ->+ Table IO k v b ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withUnion ::+ forall m k v b a.+ (IOLike m) =>+ (ResolveValue v) =>+ Table m k v b ->+ Table m k v b ->+ (Table m k v b -> m a) ->+ m a+withUnion table1 table2 =+ bracket (table1 `union` table2) closeTable++{- |+Variant of 'withUnions' that takes any number of tables.+-}+{-# SPECIALISE+ withUnions ::+ (ResolveValue v) =>+ NonEmpty (Table IO k v b) ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withUnions ::+ forall m k v b a.+ (IOLike m) =>+ (ResolveValue v) =>+ NonEmpty (Table m k v b) ->+ (Table m k v b -> m a) ->+ m a+withUnions tables =+ bracket (unions tables) closeTable++{- |+Create a table that contains the union of the entries of the given tables.++If the given key is a member of a single input table, then the same key and value occur in the output table.+Otherwise, the values for duplicate keys are combined using 'resolve' from left to right.+If the 'resolve' function behaves like 'const', then this computes a left-biased union.++>>> :{+runExample $ \session table1 -> do+ LSMT.insert table1 0 "Hello" Nothing+ LSMT.withTable session $ \table2 -> do+ LSMT.insert table2 0 "World" Nothing+ LSMT.insert table2 1 "Goodbye" Nothing+ bracket (LSMT.union table1 table2) LSMT.closeTable $ \table3 -> do+ print =<< LSMT.lookup table3 0+ print =<< LSMT.lookup table3 1+ print =<< LSMT.lookup table1 0+ print =<< LSMT.lookup table2 0+:}+Found (Value "Hello World")+Found (Value "Goodbye")+Found (Value "Hello")+Found (Value "World")++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++__Warning:__ The new table must be independently closed using 'closeTable'.++__Warning:__ Both input tables must be from the same 'Session'.++__Warning:__ This is a relatively expensive operation that may take some time to complete.+See 'incrementalUnion' for an incremental alternative.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+{-# SPECIALISE+ union ::+ (ResolveValue v) =>+ Table IO k v b ->+ Table IO k v b ->+ IO (Table IO k v b)+ #-}+union ::+ forall m k v b.+ (IOLike m) =>+ (ResolveValue v) =>+ Table m k v b ->+ Table m k v b ->+ m (Table m k v b)+union table1 table2 =+ unions (table1 :| table2 : [])++{- |+Variant of 'union' that takes any number of tables.+-}+{-# SPECIALISE+ unions ::+ (ResolveValue v) =>+ NonEmpty (Table IO k v b) ->+ IO (Table IO k v b)+ #-}+unions ::+ forall m k v b.+ (IOLike m) =>+ (ResolveValue v) =>+ NonEmpty (Table m k v b) ->+ m (Table m k v b)+unions tables = do+ bracketOnError (incrementalUnions tables) closeTable $ \table -> do+ UnionDebt debt <- remainingUnionDebt table+ UnionCredits leftovers <- supplyUnionCredits table (UnionCredits debt)+ assert (leftovers >= 0) $ pure ()+ pure table++{- |+Run an action with access to a table that incrementally computes the union of the given tables.++>>> :{+runExample $ \session table1 -> do+ LSMT.insert table1 0 "Hello" Nothing+ LSMT.withTable session $ \table2 -> do+ LSMT.insert table2 0 "World" Nothing+ LSMT.insert table2 1 "Goodbye" Nothing+ LSMT.withIncrementalUnion table1 table2 $ \table3 -> do+ print =<< LSMT.lookup table3 0+ print =<< LSMT.lookup table3 1+ print =<< LSMT.lookup table1 0+ print =<< LSMT.lookup table2 0+:}+Found (Value "Hello World")+Found (Value "Goodbye")+Found (Value "Hello")+Found (Value "World")++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'incrementalUnion' and 'closeTable'.++The created table has a /union debt/ which represents the amount of computation that remains. See 'remainingUnionDebt'.+The union debt can be paid off by supplying /union credit/ which performs an amount of computation proportional to the amount of union credit. See 'supplyUnionCredits'.+While a table has unresolved union debt, operations may become more expensive by a factor that scales with the number of unresolved unions.++__Warning:__ Both input tables must be from the same 'Session'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+{-# SPECIALISE+ withIncrementalUnion ::+ Table IO k v b ->+ Table IO k v b ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withIncrementalUnion ::+ forall m k v b a.+ (IOLike m) =>+ Table m k v b ->+ Table m k v b ->+ (Table m k v b -> m a) ->+ m a+withIncrementalUnion table1 table2 =+ bracket (incrementalUnion table1 table2) closeTable++{- |+Variant of 'withIncrementalUnion' that takes any number of tables.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B} + b)\).++The variable \(b\) refers to the number of input tables.+-}+{-# SPECIALISE+ withIncrementalUnions ::+ NonEmpty (Table IO k v b) ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withIncrementalUnions ::+ forall m k v b a.+ (IOLike m) =>+ NonEmpty (Table m k v b) ->+ (Table m k v b -> m a) ->+ m a+withIncrementalUnions tables =+ bracket (incrementalUnions tables) closeTable++{- |+Create a table that incrementally computes the union of the given tables.++>>> :{+runExample $ \session table1 -> do+ LSMT.insert table1 0 "Hello" Nothing+ LSMT.withTable session $ \table2 -> do+ LSMT.insert table2 0 "World" Nothing+ LSMT.insert table2 1 "Goodbye" Nothing+ bracket (LSMT.incrementalUnion table1 table2) LSMT.closeTable $ \table3 -> do+ print =<< LSMT.lookup table3 0+ print =<< LSMT.lookup table3 1+ print =<< LSMT.lookup table1 0+ print =<< LSMT.lookup table2 0+:}+Found (Value "Hello World")+Found (Value "Goodbye")+Found (Value "Hello")+Found (Value "World")++The worst-case disk I\/O complexity of this operation is \(O(1)\).++The created table has a /union debt/ which represents the amount of computation that remains. See 'remainingUnionDebt'.+The union debt can be paid off by supplying /union credit/ which performs an amount of computation proportional to the amount of union credit. See 'supplyUnionCredits'.+While a table has unresolved union debt, operations may become more expensive by a factor that scales with the number of unresolved unions.++__Warning:__ The new table must be independently closed using 'closeTable'.++__Warning:__ Both input tables must be from the same 'Session'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+{-# SPECIALISE+ incrementalUnion ::+ Table IO k v b ->+ Table IO k v b ->+ IO (Table IO k v b)+ #-}+incrementalUnion ::+ forall m k v b.+ (IOLike m) =>+ Table m k v b ->+ Table m k v b ->+ m (Table m k v b)+incrementalUnion table1 table2 = do+ incrementalUnions (table1 :| table2 : [])++{- |+Variant of 'incrementalUnion' for any number of tables.++The worst-case disk I\/O complexity of this operation is \(O(b)\),+where the variable \(b\) refers to the number of input tables.+-}+{-# SPECIALISE+ incrementalUnions ::+ NonEmpty (Table IO k v b) ->+ IO (Table IO k v b)+ #-}+incrementalUnions ::+ forall m k v b.+ (IOLike m) =>+ NonEmpty (Table m k v b) ->+ m (Table m k v b)+incrementalUnions tables@(Table _ :| _) =+ _withInternalTables tables (fmap Table . Internal.unions)++-- | Internal helper. Run an action with access to the underlying tables.+{-# SPECIALISE+ _withInternalTables ::+ NonEmpty (Table IO k v b) ->+ (forall h. (Typeable h) => NonEmpty (Internal.Table IO h) -> IO a) ->+ IO a+ #-}+_withInternalTables ::+ forall m k v b a.+ (IOLike m) =>+ NonEmpty (Table m k v b) ->+ (forall h. (Typeable h) => NonEmpty (Internal.Table m h) -> m a) ->+ m a+_withInternalTables (Table (table :: Internal.Table m h) :| tables) action =+ action . (table :|) =<< traverse assertTableType (zip [1 ..] tables)+ where+ assertTableType :: (Int, Table m k v b) -> m (Internal.Table m h)+ assertTableType (i, Table (table' :: Internal.Table m h'))+ | Just Refl <- eqT @h @h' = pure table'+ | otherwise = throwIO $ ErrTableUnionHandleTypeMismatch 0 (typeRep $ Proxy @h) i (typeRep $ Proxy @h')++{- |+Get an /upper bound/ for the amount of remaining union debt.+This includes the union debt of any table that was part of the union's input.++>>> :{+runExample $ \session table1 -> do+ LSMT.insert table1 0 "Hello" Nothing+ LSMT.withTable session $ \table2 -> do+ LSMT.insert table2 0 "World" Nothing+ LSMT.insert table2 1 "Goodbye" Nothing+ bracket (LSMT.incrementalUnion table1 table2) LSMT.closeTable $ \table3 -> do+ putStrLn . ("UnionDebt: "<>) . show =<< LSMT.remainingUnionDebt table3+:}+UnionDebt: 4++The worst-case disk I\/O complexity of this operation is \(O(0)\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ remainingUnionDebt ::+ Table IO k v b ->+ IO UnionDebt+ #-}+remainingUnionDebt ::+ forall m k v b.+ (IOLike m) =>+ Table m k v b ->+ m UnionDebt+remainingUnionDebt (Table table) =+ Internal.remainingUnionDebt table++{- |+Supply the given amount of union credits.++This reduces the union debt by /at least/ the number of supplied union credits.+It is therefore advisable to query 'remainingUnionDebt' every once in a while to get an upper bound on the current debt.++This function returns any surplus of union credits as /leftover/ credits when a union has finished.+In particular, if the returned number of credits is positive, then the union is finished.++>>> :{+runExample $ \session table1 -> do+ LSMT.insert table1 0 "Hello" Nothing+ LSMT.withTable session $ \table2 -> do+ LSMT.insert table2 0 "World" Nothing+ LSMT.insert table2 1 "Goodbye" Nothing+ bracket (LSMT.incrementalUnion table1 table2) LSMT.closeTable $ \table3 -> do+ putStrLn . ("UnionDebt: "<>) . show =<< LSMT.remainingUnionDebt table3+ putStrLn . ("Leftovers: "<>) . show =<< LSMT.supplyUnionCredits table3 2+ putStrLn . ("UnionDebt: "<>) . show =<< LSMT.remainingUnionDebt table3+ putStrLn . ("Leftovers: "<>) . show =<< LSMT.supplyUnionCredits table3 4+:}+UnionDebt: 4+Leftovers: 0+UnionDebt: 2+Leftovers: 3++__NOTE:__+The 'remainingUnionDebt' functions gets an /upper bound/ for the amount of remaning union debt.+In the example above, the second call to 'remainingUnionDebt' reports @2@, but the union debt is @1@.+Therefore, the second call to 'supplyUnionCredits' returns more leftovers than expected.++The worst-case disk I\/O complexity of this operation is \(O(\frac{b}{P})\),+where the variable \(b\) refers to the amount of credits supplied.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ supplyUnionCredits ::+ (ResolveValue v) =>+ Table IO k v b ->+ UnionCredits ->+ IO UnionCredits+ #-}+supplyUnionCredits ::+ forall m k v b.+ (IOLike m) =>+ (ResolveValue v) =>+ Table m k v b ->+ UnionCredits ->+ m UnionCredits+supplyUnionCredits (Table table :: Table m k v b) credits =+ Internal.supplyUnionCredits (_getResolveSerialisedValue (Proxy @v)) table credits++--------------------------------------------------------------------------------+-- Blob References+--------------------------------------------------------------------------------++-- NOTE: 'BlobRef' is defined in 'Database.LSMTree.Internal.Types'++{- |+Retrieve the blob value from a blob reference.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" (Just "World")+ print+ =<< traverse (LSMT.retrieveBlob session)+ =<< LSMT.lookup table 0+:}+FoundWithBlob (Value "Hello") (Blob "World")++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ A blob reference is /not stable/. Any operation that modifies the table,+cursor, or session that corresponds to a blob reference may cause it to be invalidated.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['BlobRefInvalidError']:+ If the blob reference has been invalidated.+-}+{-# SPECIALISE+ retrieveBlob ::+ (SerialiseValue b) =>+ Session IO ->+ BlobRef IO b ->+ IO b+ #-}+retrieveBlob ::+ forall m b.+ (IOLike m, SerialiseValue b) =>+ Session m ->+ BlobRef m b ->+ m b+retrieveBlob session blobRef = do+ blobs <- retrieveBlobs session (V.singleton blobRef)+ pure $ V.head blobs++{- |+Variant of 'retrieveBlob' for batch retrieval.+The batch of blob references corresponds in-order to the batch of results.++The worst-case disk I\/O complexity of this operation is \(O(b)\),+where the variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> retrieveBlobs session blobRefs = traverse (retrieveBlob session) blobRefs+-}+{-# SPECIALISE+ retrieveBlobs ::+ (SerialiseValue b) =>+ Session IO ->+ Vector (BlobRef IO b) ->+ IO (Vector b)+ #-}+retrieveBlobs ::+ forall m b.+ (IOLike m, SerialiseValue b) =>+ Session m ->+ Vector (BlobRef m b) ->+ m (Vector b)+retrieveBlobs (Session (session :: Internal.Session m h)) blobRefs = do+ let numBlobRefs = V.length blobRefs+ let blobRefNums = V.enumFromTo 0 (numBlobRefs - 1)+ weakBlobRefs <- traverse assertBlobRefHandleType (V.zip blobRefNums blobRefs)+ serialisedBlobs <- Internal.retrieveBlobs session weakBlobRefs+ pure $ Internal.deserialiseBlob <$> serialisedBlobs+ where+ assertBlobRefHandleType :: (Int, BlobRef m b) -> m (Internal.WeakBlobRef m h)+ assertBlobRefHandleType (i, BlobRef (weakBlobRef :: Internal.WeakBlobRef m h'))+ | Just Refl <- eqT @h @h' = pure weakBlobRef+ | otherwise = throwIO $ ErrBlobRefInvalid i++--------------------------------------------------------------------------------+-- Cursors+--------------------------------------------------------------------------------++-- NOTE: 'Cursor' is defined in 'Database.LSMTree.Internal.Types'++{- |+Run an action with access to a cursor.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.withCursor table $ \cursor -> do+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 0) (Value "Hello")+Entry (Key 1) (Value "World")++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'newCursor' and 'closeCursor'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ withCursor ::+ (ResolveValue v) =>+ Table IO k v b ->+ (Cursor IO k v b -> IO a) ->+ IO a+ #-}+withCursor ::+ forall m k v b a.+ (IOLike m) =>+ (ResolveValue v) =>+ Table m k v b ->+ (Cursor m k v b -> m a) ->+ m a+withCursor (Table table) action =+ Internal.withCursor (_getResolveSerialisedValue (Proxy @v)) Internal.NoOffsetKey table (action . Cursor)++{- |+Variant of 'withCursor' that starts at a given key.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.withCursorAtOffset table 1 $ \cursor -> do+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 1) (Value "World")+-}+{-# SPECIALISE+ withCursorAtOffset ::+ (SerialiseKey k, ResolveValue v) =>+ Table IO k v b ->+ k ->+ (Cursor IO k v b -> IO a) ->+ IO a+ #-}+withCursorAtOffset ::+ forall m k v b a.+ (IOLike m) =>+ (SerialiseKey k, ResolveValue v) =>+ Table m k v b ->+ k ->+ (Cursor m k v b -> m a) ->+ m a+withCursorAtOffset (Table table) offsetKey action =+ Internal.withCursor (_getResolveSerialisedValue (Proxy @v)) (Internal.OffsetKey $ Internal.serialiseKey offsetKey) table (action . Cursor)++{- |+Create a cursor for the given table.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ bracket (LSMT.newCursor table) LSMT.closeCursor $ \cursor -> do+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 0) (Value "Hello")+Entry (Key 1) (Value "World")++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++__Warning:__ Cursors hold open resources and must be closed using 'closeCursor'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+{-# SPECIALISE+ newCursor ::+ (ResolveValue v) =>+ Table IO k v b ->+ IO (Cursor IO k v b)+ #-}+newCursor ::+ forall m k v b.+ (IOLike m) =>+ (ResolveValue v) =>+ Table m k v b ->+ m (Cursor m k v b)+newCursor (Table table) =+ Cursor <$> Internal.newCursor (_getResolveSerialisedValue (Proxy @v)) Internal.NoOffsetKey table++{- |+Variant of 'newCursor' that starts at a given key.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ bracket (LSMT.newCursorAtOffset table 1) LSMT.closeCursor $ \cursor -> do+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 1) (Value "World")+-}+{-# SPECIALISE+ newCursorAtOffset ::+ (SerialiseKey k, ResolveValue v) =>+ Table IO k v b ->+ k ->+ IO (Cursor IO k v b)+ #-}+newCursorAtOffset ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, ResolveValue v) =>+ Table m k v b ->+ k ->+ m (Cursor m k v b)+newCursorAtOffset (Table table) offsetKey =+ Cursor <$> Internal.newCursor (_getResolveSerialisedValue (Proxy @v)) (Internal.OffsetKey $ Internal.serialiseKey offsetKey) table++{- |+Close a cursor.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Closing is idempotent, i.e., closing a closed cursor does nothing.+All other operations on a closed cursor will throw an exception.+-}+{-# SPECIALISE+ closeCursor ::+ Cursor IO k v b ->+ IO ()+ #-}+closeCursor ::+ forall m k v b.+ (IOLike m) =>+ Cursor m k v b ->+ m ()+closeCursor (Cursor cursor) =+ Internal.closeCursor cursor++{- |+Read the next table entry from the cursor.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.withCursor table $ \cursor -> do+ print =<< LSMT.next cursor+ print =<< LSMT.next cursor+ print =<< LSMT.next cursor+:}+Just (Entry (Key 0) (Value "Hello"))+Just (Entry (Key 1) (Value "World"))+Nothing++The worst-case disk I\/O complexity of this operation is \(O(\frac{1}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['CursorClosedError']:+ If the cursor is closed.+-}+{-# SPECIALISE+ next ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Cursor IO k v b ->+ IO (Maybe (Entry k v (BlobRef IO b)))+ #-}+next ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Cursor m k v b ->+ m (Maybe (Entry k v (BlobRef m b)))+next iterator = do+ -- TODO: implement this function in terms of 'readEntry'+ entries <- take 1 iterator+ pure $ fst <$> V.uncons entries++{- |+Read the next batch of table entries from the cursor.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.withCursor table $ \cursor -> do+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 0) (Value "Hello")+Entry (Key 1) (Value "World")++The worst-case disk I\/O complexity of this operation is \(O(\frac{b}{P})\),+where the variable \(b\) refers to the length of the /output/ vector,+which is /at most/ equal to the given number.+In practice, the length of the output vector is only less than the given number+once the cursor reaches the end of the table.++The following property holds:++prop> take n cursor = catMaybes <$> replicateM n (next cursor)++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['CursorClosedError']:+ If the cursor is closed.+-}+{-# SPECIALISE+ take ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Int ->+ Cursor IO k v b ->+ IO (Vector (Entry k v (BlobRef IO b)))+ #-}+take ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Int ->+ Cursor m k v b ->+ m (Vector (Entry k v (BlobRef m b)))+take n (Cursor cursor :: Cursor m k v b) =+ Internal.readCursor (_getResolveSerialisedValue (Proxy @v)) n cursor $ \ !k !v -> \case+ Just !b -> EntryWithBlob (Internal.deserialiseKey k) (Internal.deserialiseValue v) (BlobRef b)+ Nothing -> Entry (Internal.deserialiseKey k) (Internal.deserialiseValue v)++{- |+Variant of 'take' that accepts an additional predicate to determine whether or not to continue reading.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.withCursor table $ \cursor -> do+ traverse_ print+ =<< LSMT.takeWhile 32 (<1) cursor+:}+Entry (Key 0) (Value "Hello")++The worst-case disk I\/O complexity of this operation is \(O(\frac{b}{P})\),+where the variable \(b\) refers to the length of the /output/ vector,+which is /at most/ equal to the given number.+In practice, the length of the output vector is only less than the given number+when the predicate returns false or the cursor reaches the end of the table.++The following properties hold:++prop> takeWhile n (const True) cursor = take n cursor+prop> takeWhile n (const False) cursor = pure empty++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['CursorClosedError']:+ If the cursor is closed.+-}+{-# SPECIALISE+ takeWhile ::+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Int ->+ (k -> Bool) ->+ Cursor IO k v b ->+ IO (Vector (Entry k v (BlobRef IO b)))+ #-}+takeWhile ::+ forall m k v b.+ (IOLike m) =>+ (SerialiseKey k, SerialiseValue v, ResolveValue v) =>+ Int ->+ (k -> Bool) ->+ Cursor m k v b ->+ m (Vector (Entry k v (BlobRef m b)))+takeWhile n p (Cursor cursor :: Cursor m k v b) =+ -- TODO: implement this function using a variant of 'readCursorWhile' that does not take the maximum batch size+ Internal.readCursorWhile (_getResolveSerialisedValue (Proxy @v)) (p . Internal.deserialiseKey) n cursor $ \ !k !v -> \case+ Just !b -> EntryWithBlob (Internal.deserialiseKey k) (Internal.deserialiseValue v) (BlobRef b)+ Nothing -> Entry (Internal.deserialiseKey k) (Internal.deserialiseValue v)++--------------------------------------------------------------------------------+-- Snapshots+--------------------------------------------------------------------------------++{- |+Save the current state of the table to disk as a snapshot under the given+snapshot name. This is the /only/ mechanism that persists a table. Each snapshot+must have a unique name, which may be used to restore the table from that snapshot+using 'openTableFromSnapshot'.+Saving a snapshot /does not/ close the table.++Saving a snapshot is /relatively/ cheap when compared to opening a snapshot.+However, it is not so cheap that one should use it after every operation.++>>> :{+runExample $ \session table -> do+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.saveSnapshot "example" "Key Value Blob" table+:}++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['SnapshotExistsError']:+ If a snapshot with the same name already exists.+-}+{-# SPECIALISE+ saveSnapshot ::+ SnapshotName ->+ SnapshotLabel ->+ Table IO k v b ->+ IO ()+ #-}+saveSnapshot ::+ forall m k v b.+ (IOLike m) =>+ SnapshotName ->+ SnapshotLabel ->+ Table m k v b ->+ m ()+saveSnapshot snapName snapLabel (Table table) =+ Internal.saveSnapshot snapName snapLabel table++{- |+Run an action with access to a table from a snapshot.++>>> :{+runExample $ \session table -> do+ -- Save snapshot+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.saveSnapshot "example" "Key Value Blob" table+ -- Open snapshot+ LSMT.withTableFromSnapshot @_ @Key @Value @Blob session "example" "Key Value Blob" $ \table' -> do+ LSMT.withCursor table' $ \cursor ->+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 0) (Value "Hello")+Entry (Key 1) (Value "World")++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'openTableFromSnapshot' and 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['SnapshotDoesNotExistError']+ If no snapshot with the given name exists.+['SnapshotCorruptedError']:+ If the snapshot data is corrupted.+['SnapshotNotCompatibleError']:+ If the snapshot has a different label or is a different table type.+-}+{-# SPECIALISE+ withTableFromSnapshot ::+ (ResolveValue v) =>+ Session IO ->+ SnapshotName ->+ SnapshotLabel ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withTableFromSnapshot ::+ forall m k v b a.+ (IOLike m) =>+ (ResolveValue v) =>+ Session m ->+ SnapshotName ->+ SnapshotLabel ->+ (Table m k v b -> m a) ->+ m a+withTableFromSnapshot session snapName snapLabel =+ bracket (openTableFromSnapshot session snapName snapLabel) closeTable++{- |+Variant of 'withTableFromSnapshot' that accepts [table configuration overrides](#g:table_configuration_overrides).+-}+{-# SPECIALISE+ withTableFromSnapshotWith ::+ forall k v b a.+ (ResolveValue v) =>+ TableConfigOverride ->+ Session IO ->+ SnapshotName ->+ SnapshotLabel ->+ (Table IO k v b -> IO a) ->+ IO a+ #-}+withTableFromSnapshotWith ::+ forall m k v b a.+ (IOLike m) =>+ (ResolveValue v) =>+ TableConfigOverride ->+ Session m ->+ SnapshotName ->+ SnapshotLabel ->+ (Table m k v b -> m a) ->+ m a+withTableFromSnapshotWith tableConfigOverride session snapName snapLabel =+ bracket (openTableFromSnapshotWith tableConfigOverride session snapName snapLabel) closeTable++{- |+Open a table from a named snapshot.++>>> :{+runExample $ \session table -> do+ -- Save snapshot+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.saveSnapshot "example" "Key Value Blob" table+ -- Open snapshot+ bracket+ (LSMT.openTableFromSnapshot @_ @Key @Value @Blob session "example" "Key Value Blob")+ LSMT.closeTable $ \table' -> do+ LSMT.withCursor table' $ \cursor ->+ traverse_ print+ =<< LSMT.take 32 cursor+:}+Entry (Key 0) (Value "Hello")+Entry (Key 1) (Value "World")++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++__Warning:__ The new table must be independently closed using 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['SnapshotDoesNotExistError']+ If no snapshot with the given name exists.+['SnapshotCorruptedError']:+ If the snapshot data is corrupted.+['SnapshotNotCompatibleError']:+ If the snapshot has a different label or is a different table type.+-}+{-# SPECIALISE+ openTableFromSnapshot ::+ forall k v b.+ (ResolveValue v) =>+ Session IO ->+ SnapshotName ->+ SnapshotLabel ->+ IO (Table IO k v b)+ #-}+openTableFromSnapshot ::+ forall m k v b.+ (IOLike m) =>+ (ResolveValue v) =>+ Session m ->+ SnapshotName ->+ SnapshotLabel ->+ m (Table m k v b)+openTableFromSnapshot session snapName snapLabel =+ openTableFromSnapshotWith noTableConfigOverride session snapName snapLabel++{- |+Variant of 'openTableFromSnapshot' that accepts [table configuration overrides](#g:table_configuration_overrides).+-}+{-# SPECIALISE+ openTableFromSnapshotWith ::+ forall k v b.+ (ResolveValue v) =>+ TableConfigOverride ->+ Session IO ->+ SnapshotName ->+ SnapshotLabel ->+ IO (Table IO k v b)+ #-}+openTableFromSnapshotWith ::+ forall m k v b.+ (IOLike m) =>+ (ResolveValue v) =>+ TableConfigOverride ->+ Session m ->+ SnapshotName ->+ SnapshotLabel ->+ m (Table m k v b)+openTableFromSnapshotWith tableConfigOverride (Session session) snapName snapLabel =+ Table <$> Internal.openTableFromSnapshot tableConfigOverride session snapName snapLabel (_getResolveSerialisedValue (Proxy @v))++{- |+Delete the named snapshot.++>>> :{+runExample $ \session table -> do+ -- Save snapshot+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.saveSnapshot "example" "Key Value Blob" table+ -- Delete snapshot+ LSMT.deleteSnapshot session "example"+:}++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['SnapshotDoesNotExistError']:+ If no snapshot with the given name exists.+-}+{-# SPECIALISE+ deleteSnapshot ::+ Session IO ->+ SnapshotName ->+ IO ()+ #-}+deleteSnapshot ::+ forall m.+ (IOLike m) =>+ Session m ->+ SnapshotName ->+ m ()+deleteSnapshot (Session session) =+ Internal.deleteSnapshot session++{- |+Check if the named snapshot exists.++>>> :{+runExample $ \session table -> do+ -- Save snapshot+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.saveSnapshot "example" "Key Value Blob" table+ -- Check snapshots+ print =<< doesSnapshotExist session "example"+ print =<< doesSnapshotExist session "this_snapshot_does_not_exist"+:}+True+False++The worst-case disk I\/O complexity of this operation is \(O(1)\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+{-# SPECIALISE+ doesSnapshotExist ::+ Session IO ->+ SnapshotName ->+ IO Bool+ #-}+doesSnapshotExist ::+ forall m.+ (IOLike m) =>+ Session m ->+ SnapshotName ->+ m Bool+doesSnapshotExist (Session session) =+ Internal.doesSnapshotExist session++{- |+List the names of all snapshots.++>>> :{+runExample $ \session table -> do+ -- Save snapshot+ LSMT.insert table 0 "Hello" Nothing+ LSMT.insert table 1 "World" Nothing+ LSMT.saveSnapshot "example" "Key Value Blob" table+ -- List snapshots+ traverse_ print+ =<< listSnapshots session+:}+"example"++The worst-case disk I\/O complexity of this operation is \(O(s)\),+where \(s\) refers to the number of snapshots in the session.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+{-# SPECIALISE+ listSnapshots ::+ Session IO ->+ IO [SnapshotName]+ #-}+listSnapshots ::+ forall m.+ (IOLike m) =>+ Session m ->+ m [SnapshotName]+listSnapshots (Session session) =+ Internal.listSnapshots session++-- | Internal helper. Get 'resolveSerialised' at type 'ResolveSerialisedValue'.+_getResolveSerialisedValue ::+ forall v.+ (ResolveValue v) =>+ Proxy v ->+ ResolveSerialisedValue+_getResolveSerialisedValue = coerce . resolveSerialised
@@ -0,0 +1,1634 @@+{- |+Module : Database.LSMTree.Simple+Copyright : (c) 2023-2025, Cardano Development Foundation+License : Apache-2.0+Stability : experimental+Portability : portable+-}+module Database.LSMTree.Simple (+ -- * Example+ -- $example++ -- * Usage Notes #usage_notes#++ -- ** Resource Management #resource_management#+ -- $resource_management++ -- ** Concurrency #concurrency#+ -- $concurrency++ -- ** ACID properties #acid#+ -- $acid++ -- ** Sharing #sharing#+ -- $sharing++ -- * Sessions #sessions#+ Session,+ withOpenSession,+ openSession,+ closeSession,++ -- * Tables #tables#+ Table,+ withTable,+ withTableWith,+ newTable,+ newTableWith,+ closeTable,++ -- ** Table Lookups #table_lookups#+ member,+ members,+ lookup,+ lookups,+ rangeLookup,++ -- ** Table Updates #table_updates#+ insert,+ inserts,+ delete,+ deletes,+ update,+ updates,++ -- ** Table Duplication #table_duplication#+ withDuplicate,+ duplicate,++ -- ** Table Unions #table_unions#+ withUnion,+ withUnions,+ union,+ unions,+ withIncrementalUnion,+ withIncrementalUnions,+ incrementalUnion,+ incrementalUnions,+ remainingUnionDebt,+ supplyUnionCredits,++ -- * Cursors #cursor#+ Cursor,+ withCursor,+ withCursorAtOffset,+ newCursor,+ newCursorAtOffset,+ closeCursor,+ next,+ take,+ takeWhile,++ -- * Snapshots #snapshots#+ saveSnapshot,+ withTableFromSnapshot,+ withTableFromSnapshotWith,+ openTableFromSnapshot,+ openTableFromSnapshotWith,+ doesSnapshotExist,+ deleteSnapshot,+ listSnapshots,+ SnapshotName,+ isValidSnapshotName,+ toSnapshotName,+ SnapshotLabel (..),++ -- * Table Configuration #table_configuration#+ TableConfig (+ confMergePolicy,+ confSizeRatio,+ confWriteBufferAlloc,+ confBloomFilterAlloc,+ confFencePointerIndex,+ confDiskCachePolicy,+ confMergeSchedule+ ),+ MergePolicy (LazyLevelling),+ SizeRatio (Four),+ WriteBufferAlloc (AllocNumEntries),+ BloomFilterAlloc (AllocFixed, AllocRequestFPR),+ FencePointerIndexType (OrdinaryIndex, CompactIndex),+ DiskCachePolicy (..),+ MergeSchedule (..),+ MergeBatchSize (..),++ -- ** Table Configuration Overrides #table_configuration_overrides#+ TableConfigOverride (..),+ noTableConfigOverride,++ -- * Ranges #ranges#+ Range (..),++ -- * Union Credit and Debt+ UnionCredits (..),+ UnionDebt (..),++ -- * Key\/Value Serialisation #key_value_serialisation#+ RawBytes (RawBytes),+ SerialiseKey (serialiseKey, deserialiseKey),+ SerialiseKeyOrderPreserving,+ SerialiseValue (serialiseValue, deserialiseValue),++ -- ** Key\/Value Serialisation Property Tests #key_value_serialisation_property_tests#+ serialiseKeyIdentity,+ serialiseKeyIdentityUpToSlicing,+ serialiseKeyPreservesOrdering,+ serialiseValueIdentity,+ serialiseValueIdentityUpToSlicing,+ packSlice,++ -- * Errors #errors#+ SessionDirDoesNotExistError (..),+ SessionDirLockedError (..),+ SessionDirCorruptedError (..),+ SessionClosedError (..),+ TableClosedError (..),+ TableCorruptedError (..),+ TableTooLargeError (..),+ TableUnionNotCompatibleError (..),+ SnapshotExistsError (..),+ SnapshotDoesNotExistError (..),+ SnapshotCorruptedError (..),+ SnapshotNotCompatibleError (..),+ CursorClosedError (..),+ InvalidSnapshotNameError (..),+) where++import Control.ActionRegistry (mapExceptionWithActionRegistry)+import Control.Exception (Exception, SomeException (..),+ bracketOnError, finally)+import Data.Bifunctor (Bifunctor (..))+import Data.Coerce (coerce)+import Data.Kind (Type)+import Data.List.NonEmpty (NonEmpty (..))+import Data.Text (Text)+import Data.Typeable (TypeRep)+import Data.Vector (Vector)+import Data.Void (Void)+import Database.LSMTree (BloomFilterAlloc, CursorClosedError (..),+ DiskCachePolicy, FencePointerIndexType,+ InvalidSnapshotNameError (..), MergeBatchSize, MergePolicy,+ MergeSchedule, Range (..), RawBytes, ResolveAsFirst (..),+ SerialiseKey (..), SerialiseKeyOrderPreserving,+ SerialiseValue (..), SessionClosedError (..), SizeRatio,+ SnapshotCorruptedError (..),+ SnapshotDoesNotExistError (..), SnapshotExistsError (..),+ SnapshotLabel (..), SnapshotName,+ SnapshotNotCompatibleError (..), TableClosedError (..),+ TableConfig (..), TableConfigOverride (..),+ TableCorruptedError (..), TableTooLargeError (..),+ UnionCredits (..), UnionDebt (..), WriteBufferAlloc,+ isValidSnapshotName, noTableConfigOverride, packSlice,+ serialiseKeyIdentity, serialiseKeyIdentityUpToSlicing,+ serialiseKeyPreservesOrdering, serialiseValueIdentity,+ serialiseValueIdentityUpToSlicing, toSnapshotName)+import qualified Database.LSMTree as LSMT+import qualified Database.LSMTree.Internal.Types as LSMT+import qualified Database.LSMTree.Internal.Unsafe as Internal+import Prelude hiding (lookup, take, takeWhile)+import System.FS.API (MountPoint (..), mkFsPath)+import System.FS.BlockIO.API (HasBlockIO (..))+import System.FS.BlockIO.IO (defaultIOCtxParams, ioHasBlockIO)+import System.Random (randomIO)++--------------------------------------------------------------------------------+-- Example+--------------------------------------------------------------------------------++{- $setup+>>> import Prelude hiding (lookup)+>>> import Data.ByteString.Short (ShortByteString)+>>> import Data.String (IsString)+>>> import Data.Word (Word64)+>>> import System.Directory (createDirectoryIfMissing, getTemporaryDirectory)+>>> import System.FilePath ((</>))++>>> :{+newtype Key = Key Word64+ deriving stock (Eq, Show)+ deriving newtype (Num, SerialiseKey)+:}++>>> :{+newtype Value = Value ShortByteString+ deriving stock (Eq, Show)+ deriving newtype (IsString, SerialiseValue)+:}++>>> :{+runExample :: (Session -> Table Key Value -> IO a) -> IO a+runExample action = do+ tmpDir <- getTemporaryDirectory+ let sessionDir = tmpDir </> "doctest_Database_LSMTree_Simple"+ createDirectoryIfMissing True sessionDir+ withOpenSession sessionDir $ \session ->+ withTable session $ \table ->+ action session table+:}+-}++{- $example++>>> :{+runExample $ \session table -> do+ insert table 0 "Hello"+ insert table 1 "World"+ lookup table 0+:}+Just (Value "Hello")+-}++--------------------------------------------------------------------------------+-- Resource Management+--------------------------------------------------------------------------------++{- $resource_management+This package uses explicit resource management. The 'Session', 'Table', and 'Cursor'+handles hold open resources, such as file handles, which must be explicitly released.+Every operation that allocates a resource is paired with another operation to releases+that resource. For each pair of allocate and release operations there is a bracketed+function that combines the two.+++------------+--------------------------+-------------------------+-------------------++| Resource | Bracketed #bracketed# | Allocate #allocate# | Release #release# |++============+==========================+=========================+===================++| 'Session' | 'withOpenSession' | 'openSession' | 'closeSession' |++------------+--------------------------+-------------------------+-------------------++| 'Table' | 'withTable' | 'newTable' | 'closeTable' |++ +--------------------------+-------------------------+ ++| | 'withDuplicate' | 'duplicate' | |++ +--------------------------+-------------------------+ ++| | 'withUnion' | 'union' | |++ +--------------------------+-------------------------+ ++| | 'withIncrementalUnion' | 'incrementalUnion' | |++ +--------------------------+-------------------------+ ++| | 'withTableFromSnapshot' | 'openTableFromSnapshot' | |++------------+--------------------------+-------------------------+-------------------++| 'Cursor' | 'withCursor' | 'newCursor' | 'closeCursor' |++------------+--------------------------+-------------------------+-------------------+++To prevent resource and memory leaks due to asynchronous exceptions,+it is recommended to use the [bracketed](#bracketed) functions whenever+possible, and otherwise:++* Run functions that allocate and release a resource with asynchronous+ exceptions masked.+* Ensure that every use allocate operation is followed by the corresponding release+ operation even in the presence of asynchronous exceptions, e.g., using 'bracket'.+-}++--------------------------------------------------------------------------------+-- Concurrency+--------------------------------------------------------------------------------++{- $concurrency+Table handles may be used concurrently from multiple Haskell threads,+and doing read operations concurrently may result in improved throughput,+as it can take advantage of CPU and I\/O parallelism. However, concurrent+use of write operations may introduces races. Specifically:++* It is a race to read and write the same table concurrently.+* It is a race to write and write the same table concurrently.+* It is /not/ a race to read and read the same table concurrently.+* It is /not/ a race to read or write /separate/ tables concurrently.++For the purposes of the above rules:++* The read operations are 'lookup', 'rangeLookup', 'duplicate', `union`, 'saveSnapshot', 'newCursor', and their variants.+* The write operations are 'insert', 'delete', 'update', 'closeTable', and their variants.++It is possible to read from a stable view of a table while concurrently writing to+the table by using 'duplicate' and performing the read operations on the duplicate.+However, this requires that the 'duplicate' operation /happens before/ the subsequent+writes, as it is a race to duplicate concurrently with any writes.+As this package does not provide any construct for synchronisation or atomic+operations, this ordering of operations must be accomplished by the user through+other means.++A 'Cursor' creates a stable view of a table and can safely be read while+modifying the original table. However, reading the 'next' key\/value pair from+a cursor locks the view, so concurrent reads on the same cursor block.+This is because 'next' updates the cursor's current position.++Session handles may be used concurrently from multiple Haskell threads,+but concurrent use of read and write operations may introduce races.+Specifically, it is a race to use `listSnapshots` and `deleteSnapshots`+with the same session handle concurrently.+-}++--------------------------------------------------------------------------------+-- ACID properties+--------------------------------------------------------------------------------++{- $acid+This text copies liberally from https://en.wikipedia.org/wiki/ACID and related wiki pages.++Atomicity, consistency, isolation, and durability (ACID) are important+properties of database transactions.+They guarantee data validity despite errors, power failures, and other mishaps.+A /transaction/ is a sequence of database operations that satisfy the ACID properties.++@lsm-tree@ does not support transactions in the typical sense that many relational databases do,+where transactions can be built from smaller components/actions,+e.g., reads and writes of individual cells.+Instead, the public API only exposes functions that individually form a transaction;+there are no smaller building blocks.+An example of such a transaction is 'updates'.++An @lsm-tree@ transaction still perform multiple database actions /internally/,+but transactions themselves are not composable into larger transactions,+so it should be expected that table contents can change between transactions in a concurrent setting.+A consistent view of a table can be created,+so that independent transactions have access to their own version of the database state (see [concurrency](#g:concurreny)).++All @lsm-tree@ transactions are designed for atomicity, consistency, and isolation (ACI),+assuming that users of the library perform proper [resource management](#g:resource-management).+Durability is only guaranteed when saving a [snapshot](#g:snapshots),+which is the only method of stopping and restarting tables.++We currently cannot guarantee consistency in the presence of synchronous and asynchronous exceptions,+eventhough major strides were made to make it so.+The safest course of action when an internal exception is encountered is to stop and restart:+close the session along with all its tables and cursors, reopen the session,+and load a previous saved table snapshot.+-}++--------------------------------------------------------------------------------+-- Sharing+--------------------------------------------------------------------------------++{- $sharing+Tables created via 'duplicate' or 'union' will initially share as much of their+in-memory and on-disk data as possible with the tables they were created from.+Over time as these related tables are modified, the contents of the tables will+diverge, which means that the tables will share less and less.++Sharing of in-memory data is not preserved by snapshots, but sharing of on-disk+data is partially preserved.+Existing files for runs are shared, but files for ongoing merges are not.+Opening a table from a snapshot (using 'openTableFromSnapshot' or+'withTableFromSnapshot') is expensive, but creating a snapshot (using+'saveSnapshot') is relatively cheap.+-}++--------------------------------------------------------------------------------+-- Sessions+--------------------------------------------------------------------------------++{- |+A session stores context that is shared by multiple tables.++Each session is associated with one session directory where the files+containing table data are stored. Each session locks its session directory.+There can only be one active session for each session directory at a time.+If a database is must be accessed from multiple parts of a program,+one session should be opened and shared between those parts of the program.+Session directories cannot be shared between OS processes.+-}+type Session :: Type+newtype Session = Session (LSMT.Session IO)++{- |+Run an action with access to a session opened from a session directory.++If the session directory is empty, a new session is created.+Otherwise, the session directory is restored as an existing session.++If there are no open tables or cursors when the session terminates, then the disk I\/O complexity of this operation is \(O(1)\).+Otherwise, 'closeTable' is called for each open table and 'closeCursor' is called for each open cursor.+Consequently, the worst-case disk I\/O complexity of this operation depends on the merge policy of the open tables in the session.+The following assumes all tables in the session have the same merge policy:++['LazyLevelling']:+ \(O(o \: T \log_T \frac{n}{B})\).++The variable \(o\) refers to the number of open tables and cursors in the session.++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'openSession' and 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+withOpenSession ::+ forall a.+ -- | The session directory.+ FilePath ->+ (Session -> IO a) ->+ IO a+withOpenSession dir action = do+ let tracer = mempty+ _convertSessionDirErrors dir $+ LSMT.withOpenSessionIO tracer dir (action . Session)++{- |+Open a session from a session directory.++If the session directory is empty, a new session is created.+Otherwise, the session directory is restored as an existing session.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ Sessions hold open resources and must be closed using 'closeSession'.++Throws the following exceptions:++['SessionDirDoesNotExistError']:+ If the session directory does not exist.+['SessionDirLockedError']:+ If the session directory is locked by another process.+['SessionDirCorruptedError']:+ If the session directory is malformed.+-}+openSession ::+ -- | The session directory.+ FilePath ->+ IO Session+openSession dir = do+ let tracer = mempty+ _convertSessionDirErrors dir $ do+ let mountPoint = MountPoint dir+ let sessionDirFsPath = mkFsPath []+ sessionSalt <- randomIO+ let acquireHasBlockIO = ioHasBlockIO mountPoint defaultIOCtxParams+ let releaseHasBlockIO (_, HasBlockIO{close}) = close+ bracketOnError acquireHasBlockIO releaseHasBlockIO $ \(hasFS, hasBlockIO) ->+ Session <$> LSMT.openSession tracer hasFS hasBlockIO sessionSalt sessionDirFsPath++{- |+Close a session.++If there are no open tables or cursors in the session, then the disk I\/O complexity of this operation is \(O(1)\).+Otherwise, 'closeTable' is called for each open table and 'closeCursor' is called for each open cursor.+Consequently, the worst-case disk I\/O complexity of this operation depends on the merge policy of the tables in the session.+The following assumes all tables in the session have the same merge policy:++['LazyLevelling']:+ \(O(o \: T \log_T \frac{n}{B})\).++The variable \(o\) refers to the number of open tables and cursors in the session.++Closing is idempotent, i.e., closing a closed session does nothing.+All other operations on a closed session will throw an exception.+-}+closeSession ::+ Session ->+ IO ()+closeSession (Session session@(LSMT.Session session')) = do+ HasBlockIO{close} <- Internal.withKeepSessionOpen session' $+ pure . Internal.sessionHasBlockIO+ LSMT.closeSession session `finally` close++--------------------------------------------------------------------------------+-- Tables+--------------------------------------------------------------------------------++{- |+A table is a handle to an individual LSM-tree key\/value store with both in-memory and on-disk parts.++__Warning:__ Tables are ephemeral. Once you close a table, its data is lost forever. To persist tables, use [snapshots](#g:snapshots).+-}+type role Table nominal nominal++type Table :: Type -> Type -> Type+newtype Table k v = Table (LSMT.Table IO k (LSMT.ResolveAsFirst v) Void)++{- |+Run an action with access to an empty table.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'newTable' and 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+withTable ::+ forall k v a.+ Session ->+ (Table k v -> IO a) ->+ IO a+withTable (Session session) action =+ LSMT.withTable session (action . Table)++-- | Variant of 'withTable' that accepts [table configuration](#g:table_configuration).+withTableWith ::+ forall k v a.+ TableConfig ->+ Session ->+ (Table k v -> IO a) ->+ IO a+withTableWith tableConfig (Session session) action =+ LSMT.withTableWith tableConfig session (action . Table)++{- |+Create an empty table.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++__Warning:__ Tables hold open resources and must be closed using 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+newTable ::+ forall k v.+ Session ->+ IO (Table k v)+newTable (Session session) =+ Table <$> LSMT.newTable session++{- |+Variant of 'newTable' that accepts [table configuration](#g:table_configuration).+-}+newTableWith ::+ forall k v.+ TableConfig ->+ Session ->+ IO (Table k v)+newTableWith tableConfig (Session session) =+ Table <$> LSMT.newTableWith tableConfig session++{- |+Close a table.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Closing is idempotent, i.e., closing a closed table does nothing.+All other operations on a closed table will throw an exception.++__Warning:__ Tables are ephemeral. Once you close a table, its data is lost forever. To persist tables, use [snapshots](#g:snapshots).+-}+closeTable ::+ forall k v.+ Table k v ->+ IO ()+closeTable (Table table) =+ LSMT.closeTable table++--------------------------------------------------------------------------------+-- Lookups+--------------------------------------------------------------------------------++{- |+Check if the key is a member of the table.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Membership tests can be performed concurrently from multiple Haskell threads.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableCorruptedError']:+ If the table data is corrupted.+-}+member ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ k ->+ IO Bool+member (Table table) =+ LSMT.member table++{- |+Variant of 'member' for batch membership tests.+The batch of keys corresponds in-order to the batch of results.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(b \: T \log_T \frac{n}{B})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> members table keys = traverse (member table) keys+-}+members ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ Vector k ->+ IO (Vector Bool)+members (Table table) =+ LSMT.members table++-- | Internal helper. Get the value from a 'LSMT.LookupResult' from the full API.+getValue :: LSMT.LookupResult (ResolveAsFirst v) (LSMT.BlobRef IO Void) -> Maybe v+getValue =+ fmap LSMT.unResolveAsFirst . LSMT.getValue++{- |+Look up the value associated with a key.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Lookups can be performed concurrently from multiple Haskell threads.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableCorruptedError']:+ If the table data is corrupted.+-}+lookup ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ k ->+ IO (Maybe v)+lookup (Table table) =+ fmap getValue . LSMT.lookup table++{- |+Variant of 'lookup' for batch lookups.+The batch of keys corresponds in-order to the batch of results.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(b \: T \log_T \frac{n}{B})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> lookups table keys = traverse (lookup table) keys+-}+lookups ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ Vector k ->+ IO (Vector (Maybe v))+lookups (Table table) =+ fmap (fmap getValue) . LSMT.lookups table++-- | Internal helper. Get a key\/value pair from an 'LSMT.Entry' from the full API.+getKeyValue :: LSMT.Entry k (ResolveAsFirst v) (LSMT.BlobRef IO Void) -> (k, v)+getKeyValue (LSMT.Entry k v) = (k, LSMT.unResolveAsFirst v)+getKeyValue (LSMT.EntryWithBlob k v !_b) = (k, LSMT.unResolveAsFirst v)++{- |+Look up a batch of values associated with keys in the given range.++The worst-case disk I\/O complexity of this operation is \(O(T \log_T \frac{n}{B} + \frac{b}{P})\),+where the variable \(b\) refers to the length of the /output/ vector.++Range lookups can be performed concurrently from multiple Haskell threads.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableCorruptedError']:+ If the table data is corrupted.+-}+rangeLookup ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ Range k ->+ IO (Vector (k, v))+rangeLookup (Table table) =+ fmap (fmap getKeyValue) . LSMT.rangeLookup table++--------------------------------------------------------------------------------+-- Updates+--------------------------------------------------------------------------------++{- |+Insert a new key and value in the table.+If the key is already present in the table, the associated value is replaced with the given value.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+insert ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ k ->+ v ->+ IO ()+insert (Table table) k v =+ LSMT.insert table k (LSMT.ResolveAsFirst v) Nothing++{- |+Variant of 'insert' for batch insertions.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> inserts table entries = traverse_ (uncurry $ insert table) entries+-}+inserts ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ Vector (k, v) ->+ IO ()+inserts (Table table) entries =+ LSMT.inserts table (fmap (\(k, v) -> (k, LSMT.ResolveAsFirst v, Nothing)) entries)++{- |+Delete a key and its value from the table.+If the key is not present in the table, the table is left unchanged.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+delete ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ k ->+ IO ()+delete (Table table) =+ LSMT.delete table++{- |+Variant of 'delete' for batch deletions.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> deletes table keys = traverse_ (delete table) keys+-}+deletes ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ Vector k ->+ IO ()+deletes (Table table) =+ LSMT.deletes table++-- | Internal helper. Convert from @'Maybe' v@ to an 'LSMT.Update'.+maybeValueToUpdate :: Maybe v -> LSMT.Update (ResolveAsFirst v) Void+maybeValueToUpdate =+ maybe LSMT.Delete (\v -> LSMT.Insert (LSMT.ResolveAsFirst v) Nothing)++{- |+Update the value at a specific key:++* If the given value is 'Just', this operation acts as 'insert'.+* If the given value is 'Nothing', this operation acts as 'delete'.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(\frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{n}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+update ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ k ->+ Maybe v ->+ IO ()+update (Table table) k =+ LSMT.update table k . maybeValueToUpdate++{- |+Variant of 'update' for batch updates.++The worst-case disk I\/O complexity of this operation depends on the merge policy and the merge schedule of the table:++['LazyLevelling'\/'Incremental']:+ \(O(b \: \frac{1}{P} \: \log_T \frac{n}{B})\).+['LazyLevelling'\/'OneShot']:+ \(O(\frac{b}{P} \log_T \frac{b}{B} + \frac{n}{P})\).++The variable \(b\) refers to the length of the input vector.++The following property holds in the absence of races:++prop> updates table entries = traverse_ (uncurry $ update table) entries+-}+updates ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Table k v ->+ Vector (k, Maybe v) ->+ IO ()+updates (Table table) =+ LSMT.updates table . fmap (second maybeValueToUpdate)++--------------------------------------------------------------------------------+-- Duplication+--------------------------------------------------------------------------------++{- |+Run an action with access to the duplicate of a table.++The duplicate is an independent copy of the given table.+The duplicate is unaffected by subsequent updates to the given table and vice versa.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'duplicate' and 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+withDuplicate ::+ forall k v a.+ Table k v ->+ (Table k v -> IO a) ->+ IO a+withDuplicate (Table table) action =+ LSMT.withDuplicate table (action . Table)++{- |+Duplicate a table.++The duplicate is an independent copy of the given table.+The duplicate is unaffected by subsequent updates to the given table and vice versa.++The worst-case disk I\/O complexity of this operation is \(O(0)\).++__Warning:__ The duplicate must be independently closed using 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+duplicate ::+ forall k v.+ Table k v ->+ IO (Table k v)+duplicate (Table table) =+ Table <$> LSMT.duplicate table++--------------------------------------------------------------------------------+-- Union+--------------------------------------------------------------------------------++{- |+Run an action with access to a table that contains the union of the entries of the given tables.++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'union' and 'closeTable'.++__Warning:__ Both input tables must be from the same 'Session'.++__Warning:__ This is a relatively expensive operation that may take some time to complete.+See 'withIncrementalUnion' for an incremental alternative.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+withUnion ::+ forall k v a.+ Table k v ->+ Table k v ->+ (Table k v -> IO a) ->+ IO a+withUnion (Table table1) (Table table2) action =+ LSMT.withUnion table1 table2 (action . Table)++{- |+Variant of 'withUnions' that takes any number of tables.+-}+withUnions ::+ forall k v a.+ NonEmpty (Table k v) ->+ (Table k v -> IO a) ->+ IO a+withUnions tables action =+ LSMT.withUnions (coerce tables) (action . Table)++{- |+Create a table that contains the left-biased union of the entries of the given tables.++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++__Warning:__ The new table must be independently closed using 'closeTable'.++__Warning:__ Both input tables must be from the same 'Session'.++__Warning:__ This is a relatively expensive operation that may take some time to complete.+See 'incrementalUnion' for an incremental alternative.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+union ::+ forall k v.+ Table k v ->+ Table k v ->+ IO (Table k v)+union (Table table1) (Table table2) =+ Table <$> LSMT.union table1 table2++{- |+Variant of 'union' that takes any number of tables.+-}+unions ::+ forall k v.+ NonEmpty (Table k v) ->+ IO (Table k v)+unions tables =+ Table <$> LSMT.unions (coerce tables)++{- |+Run an action with access to a table that incrementally computes the union of the given tables.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'incrementalUnion' and 'closeTable'.++The created table has a /union debt/ which represents the amount of computation that remains. See 'remainingUnionDebt'.+The union debt can be paid off by supplying /union credit/ which performs an amount of computation proportional to the amount of union credit. See 'supplyUnionCredits'.+While a table has unresolved union debt, operations may become more expensive by a factor that scales with the number of unresolved unions.++__Warning:__ Both input tables must be from the same 'Session'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+withIncrementalUnion ::+ forall k v a.+ Table k v ->+ Table k v ->+ (Table k v -> IO a) ->+ IO a+withIncrementalUnion (Table table1) (Table table2) action =+ LSMT.withIncrementalUnion table1 table2 (action . Table)++{- |+Variant of 'withIncrementalUnion' that takes any number of tables.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B} + b)\).++The variable \(b\) refers to the number of input tables.+-}+withIncrementalUnions ::+ forall k v a.+ NonEmpty (Table k v) ->+ (Table k v -> IO a) ->+ IO a+withIncrementalUnions tables action =+ LSMT.withIncrementalUnions (coerce tables) (action . Table)++{- |+Create a table that incrementally computes the union of the given tables.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++The created table has a /union debt/ which represents the amount of computation that remains. See 'remainingUnionDebt'.+The union debt can be paid off by supplying /union credit/ which performs an amount of computation proportional to the amount of union credit. See 'supplyUnionCredits'.+While a table has unresolved union debt, operations may become more expensive by a factor that scales with the number of unresolved unions.++__Warning:__ The new table must be independently closed using 'closeTable'.++__Warning:__ Both input tables must be from the same 'Session'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['TableUnionNotCompatibleError']:+ If both tables are not from the same 'Session'.+-}+incrementalUnion ::+ forall k v.+ Table k v ->+ Table k v ->+ IO (Table k v)+incrementalUnion (Table table1) (Table table2) = do+ Table <$> LSMT.incrementalUnion table1 table2++{- |+Variant of 'incrementalUnion' for any number of tables.++The worst-case disk I\/O complexity of this operation is \(O(b)\),+where the variable \(b\) refers to the number of input tables.+-}+incrementalUnions ::+ forall k v.+ NonEmpty (Table k v) ->+ IO (Table k v)+incrementalUnions tables = do+ Table <$> LSMT.unions (coerce tables)++{- |+Get the amount of remaining union debt.+This includes the union debt of any table that was part of the union's input.++The worst-case disk I\/O complexity of this operation is \(O(0)\).+-}+remainingUnionDebt ::+ forall k v.+ Table k v ->+ IO UnionDebt+remainingUnionDebt (Table table) =+ LSMT.remainingUnionDebt table++{- |+Supply the given amount of union credits.++This reduces the union debt by /at least/ the number of supplied union credits.+It is therefore advisable to query 'remainingUnionDebt' every once in a while to see what the current debt is.++This function returns any surplus of union credits as /leftover/ credits when a union has finished.+In particular, if the returned number of credits is positive, then the union is finished.++The worst-case disk I\/O complexity of this operation is \(O(\frac{b}{P})\),+where the variable \(b\) refers to the amount of credits supplied.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+supplyUnionCredits ::+ forall k v.+ Table k v ->+ UnionCredits ->+ IO UnionCredits+supplyUnionCredits (Table table) credits =+ LSMT.supplyUnionCredits table credits++--------------------------------------------------------------------------------+-- Cursors+--------------------------------------------------------------------------------++{- |+A cursor is a stable read-only iterator for a table.++A cursor iterates over the entries in a table following the order of the+serialised keys. After the cursor is created, updates to the referenced table+do not affect the cursor.++The name of this type references [database cursors](https://en.wikipedia.org/wiki/Cursor_(databases\)), not, e.g., text editor cursors.+-}+type role Cursor nominal nominal++type Cursor :: Type -> Type -> Type+newtype Cursor k v = Cursor (LSMT.Cursor IO k (ResolveAsFirst v) Void)++{- |+Run an action with access to a cursor.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'newCursor' and 'closeCursor'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+withCursor ::+ forall k v a.+ Table k v ->+ (Cursor k v -> IO a) ->+ IO a+withCursor (Table table) action =+ LSMT.withCursor table (action . Cursor)++{- |+Variant of 'withCursor' that starts at a given key.+-}+withCursorAtOffset ::+ forall k v a.+ (SerialiseKey k) =>+ Table k v ->+ k ->+ (Cursor k v -> IO a) ->+ IO a+withCursorAtOffset (Table table) offsetKey action =+ LSMT.withCursorAtOffset table offsetKey (action . Cursor)++{- |+Create a cursor for the given table.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++__Warning:__ Cursors hold open resources and must be closed using 'closeCursor'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+-}+newCursor ::+ forall k v.+ Table k v ->+ IO (Cursor k v)+newCursor (Table table) =+ Cursor <$> LSMT.newCursor table++{- |+Variant of 'newCursor' that starts at a given key.+-}+newCursorAtOffset ::+ forall k v.+ (SerialiseKey k) =>+ Table k v ->+ k ->+ IO (Cursor k v)+newCursorAtOffset (Table table) offsetKey =+ Cursor <$> LSMT.newCursorAtOffset table offsetKey++{- |+Close a cursor.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Closing is idempotent, i.e., closing a closed cursor does nothing.+All other operations on a closed cursor will throw an exception.+-}+closeCursor ::+ forall k v.+ Cursor k v ->+ IO ()+closeCursor (Cursor cursor) =+ LSMT.closeCursor cursor++{- |+Read the next table entry from the cursor.++The worst-case disk I\/O complexity of this operation is \(O(\frac{1}{P})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['CursorClosedError']:+ If the cursor is closed.+-}+next ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Cursor k v ->+ IO (Maybe (k, v))+next (Cursor cursor) =+ fmap getKeyValue <$> LSMT.next cursor++{- |+Read the next batch of table entries from the cursor.++The worst-case disk I\/O complexity of this operation is \(O(\frac{b}{P})\),+where the variable \(b\) refers to the length of the /output/ vector,+which is /at most/ equal to the given number.+In practice, the length of the output vector is only less than the given number+once the cursor reaches the end of the table.++The following property holds:++prop> take n cursor = catMaybes <$> replicateM n (next cursor)++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['CursorClosedError']:+ If the cursor is closed.+-}+take ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Int ->+ Cursor k v ->+ IO (Vector (k, v))+take n (Cursor cursor) =+ fmap getKeyValue <$> LSMT.take n cursor++{- |+Variant of 'take' that accepts an additional predicate to determine whether or not to continue reading.++The worst-case disk I\/O complexity of this operation is \(O(\frac{b}{P})\),+where the variable \(b\) refers to the length of the /output/ vector,+which is /at most/ equal to the given number.+In practice, the length of the output vector is only less than the given number+when the predicate returns false or the cursor reaches the end of the table.++The following properties hold:++prop> takeWhile n (const True) cursor = take n cursor+prop> takeWhile n (const False) cursor = pure empty++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['CursorClosedError']:+ If the cursor is closed.+-}+takeWhile ::+ forall k v.+ (SerialiseKey k, SerialiseValue v) =>+ Int ->+ (k -> Bool) ->+ Cursor k v ->+ IO (Vector (k, v))+takeWhile n p (Cursor cursor) =+ fmap getKeyValue <$> LSMT.takeWhile n p cursor++--------------------------------------------------------------------------------+-- Snapshots+--------------------------------------------------------------------------------++{- |+Save the current state of the table to disk as a snapshot under the given+snapshot name. This is the /only/ mechanism that persists a table. Each snapshot+must have a unique name, which may be used to restore the table from that snapshot+using 'openTableFromSnapshot'.+Saving a snapshot /does not/ close the table.++Saving a snapshot is /relatively/ cheap when compared to opening a snapshot.+However, it is not so cheap that one should use it after every operation.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['SnapshotExistsError']:+ If a snapshot with the same name already exists.+-}+saveSnapshot ::+ forall k v.+ SnapshotName ->+ SnapshotLabel ->+ Table k v ->+ IO ()+saveSnapshot snapName snapLabel (Table table) =+ LSMT.saveSnapshot snapName snapLabel table++{- |+Run an action with access to a table from a snapshot.++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++This function is exception-safe for both synchronous and asynchronous exceptions.++It is recommended to use this function instead of 'openTableFromSnapshot' and 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['SnapshotDoesNotExistError']+ If no snapshot with the given name exists.+['SnapshotCorruptedError']:+ If the snapshot data is corrupted.+['SnapshotNotCompatibleError']:+ If the snapshot has a different label or is a different table type.+-}+withTableFromSnapshot ::+ forall k v a.+ Session ->+ SnapshotName ->+ SnapshotLabel ->+ (Table k v -> IO a) ->+ IO a+withTableFromSnapshot (Session session) snapName snapLabel action =+ LSMT.withTableFromSnapshot session snapName snapLabel (action . Table)++{- |+Variant of 'withTableFromSnapshot' that accepts [table configuration overrides](#g:table_configuration_overrides).+-}+withTableFromSnapshotWith ::+ forall k v a.+ TableConfigOverride ->+ Session ->+ SnapshotName ->+ SnapshotLabel ->+ (Table k v -> IO a) ->+ IO a+withTableFromSnapshotWith tableConfigOverride (Session session) snapName snapLabel action =+ LSMT.withTableFromSnapshotWith tableConfigOverride session snapName snapLabel (action . Table)++{- |+Open a table from a named snapshot.++The worst-case disk I\/O complexity of this operation is \(O(\frac{n}{P})\).++__Warning:__ The new table must be independently closed using 'closeTable'.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['TableClosedError']:+ If the table is closed.+['SnapshotDoesNotExistError']+ If no snapshot with the given name exists.+['SnapshotCorruptedError']:+ If the snapshot data is corrupted.+['SnapshotNotCompatibleError']:+ If the snapshot has a different label or is a different table type.+-}+openTableFromSnapshot ::+ forall k v.+ Session ->+ SnapshotName ->+ SnapshotLabel ->+ IO (Table k v)+openTableFromSnapshot (Session session) snapName snapLabel =+ Table <$> LSMT.openTableFromSnapshot @IO @k @(ResolveAsFirst v) session snapName snapLabel++{- |+Variant of 'openTableFromSnapshot' that accepts [table configuration overrides](#g:table_configuration_overrides).+-}+openTableFromSnapshotWith ::+ forall k v.+ TableConfigOverride ->+ Session ->+ SnapshotName ->+ SnapshotLabel ->+ IO (Table k v)+openTableFromSnapshotWith tableConfigOverride (Session session) snapName snapLabel =+ Table <$> LSMT.openTableFromSnapshotWith tableConfigOverride session snapName snapLabel++{- |+Delete the named snapshot.++The worst-case disk I\/O complexity of this operation depends on the merge policy of the table:++['LazyLevelling']:+ \(O(T \log_T \frac{n}{B})\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+['SnapshotDoesNotExistError']:+ If no snapshot with the given name exists.+-}+deleteSnapshot ::+ Session ->+ SnapshotName ->+ IO ()+deleteSnapshot (Session session) =+ LSMT.deleteSnapshot session++{- |+Check if the named snapshot exists.++The worst-case disk I\/O complexity of this operation is \(O(1)\).++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+doesSnapshotExist ::+ Session ->+ SnapshotName ->+ IO Bool+doesSnapshotExist (Session session) =+ LSMT.doesSnapshotExist session++{- |+List the names of all snapshots.++The worst-case disk I\/O complexity of this operation is \(O(s)\),+where the variable \(s\) refers to the number of snapshots in the session.++Throws the following exceptions:++['SessionClosedError']:+ If the session is closed.+-}+listSnapshots ::+ Session ->+ IO [SnapshotName]+listSnapshots (Session session) =+ LSMT.listSnapshots session++--------------------------------------------------------------------------------+-- Errors+--------------------------------------------------------------------------------++-- | The session directory does not exist.+data SessionDirDoesNotExistError+ = ErrSessionDirDoesNotExist !FilePath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | The session directory is locked by another active session.+data SessionDirLockedError+ = ErrSessionDirLocked !FilePath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++-- | The session directory is corrupted, e.g., it misses required files or contains unexpected files.+data SessionDirCorruptedError+ = ErrSessionDirCorrupted !Text !FilePath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{- | Internal helper. Convert:++* t'LSMT.SessionDirDoesNotExistError' to t'SessionDirDoesNotExistError';+* t'LSMT.SessionDirLockedError' to t'SessionDirLockedError'; and+* t'LSMT.SessionDirCorruptedError' to t'SessionDirCorruptedError'.+-}+_convertSessionDirErrors ::+ forall a.+ FilePath ->+ IO a ->+ IO a+_convertSessionDirErrors sessionDir =+ mapExceptionWithActionRegistry (\(LSMT.ErrSessionDirDoesNotExist _fsErrorPath) -> SomeException $ ErrSessionDirDoesNotExist sessionDir)+ . mapExceptionWithActionRegistry (\(LSMT.ErrSessionDirLocked _fsErrorPath) -> SomeException $ ErrSessionDirLocked sessionDir)+ . mapExceptionWithActionRegistry (\(LSMT.ErrSessionDirCorrupted reason _fsErrorPath) -> SomeException $ ErrSessionDirCorrupted reason sessionDir)++{-------------------------------------------------------------------------------+ Table union+-------------------------------------------------------------------------------}++-- | A table union was constructed with two tables that are not compatible.+data TableUnionNotCompatibleError+ = ErrTableUnionHandleTypeMismatch+ -- | The index of the first table.+ !Int+ -- | The type of the filesystem handle of the first table.+ !TypeRep+ -- | The index of the second table.+ !Int+ -- | The type of the filesystem handle of the second table.+ !TypeRep+ | ErrTableUnionSessionMismatch+ -- | The index of the first table.+ !Int+ -- | The session directory of the first table.+ !FilePath+ -- | The index of the second table.+ !Int+ -- | The session directory of the second table.+ !FilePath+ deriving stock (Show, Eq)+ deriving anyclass (Exception)++{- | Internal helper. Convert:++* t'LSMT.TableUnionNotCompatibleError' to t'TableUnionNotCompatibleError';+-}+_convertTableUnionNotCompatibleError ::+ forall a.+ (Int -> FilePath) ->+ IO a ->+ IO a+_convertTableUnionNotCompatibleError sessionDirFor =+ mapExceptionWithActionRegistry $ \case+ LSMT.ErrTableUnionHandleTypeMismatch i1 typeRep1 i2 typeRep2 ->+ ErrTableUnionHandleTypeMismatch i1 typeRep1 i2 typeRep2+ LSMT.ErrTableUnionSessionMismatch i1 _fsErrorPath1 i2 _fsErrorPath2 ->+ ErrTableUnionSessionMismatch i1 (sessionDirFor i1) i2 (sessionDirFor i2)
@@ -0,0 +1,13 @@+module Main (main) where++import qualified Test.Control.ActionRegistry+import qualified Test.Control.Concurrent.Class.MonadSTM.RWVar+import qualified Test.Control.RefCount+import Test.Tasty++main :: IO ()+main = defaultMain $ testGroup "control"+ [ Test.Control.ActionRegistry.tests+ , Test.Control.Concurrent.Class.MonadSTM.RWVar.tests+ , Test.Control.RefCount.tests+ ]
@@ -0,0 +1,45 @@+module Test.Control.ActionRegistry (tests) where++import Control.ActionRegistry+import Control.Monad.Class.MonadThrow+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Control.ActionRegistry" [+ testProperty "prop_commitActionRegistryError" prop_commitActionRegistryError+ , testProperty "prop_abortActionRegistryError" prop_abortActionRegistryError+ ]++-- | An example where an exception happens while an action registry is being+-- committed.+prop_commitActionRegistryError :: Property+prop_commitActionRegistryError = once $ ioProperty $ do+ eith <-+ try @_ @CommitActionRegistryError $+ withActionRegistry $ \reg -> do+ delayedCommit reg+ (throwIO (userError "delayed action failed"))+ pure $ case eith of+ Left e ->+ tabulate "displayException" [displayExceptionNewline e] $ property True+ Right () -> property False++-- | An example where an exception happens while an action registry is being+-- aborted.+prop_abortActionRegistryError :: Property+prop_abortActionRegistryError = once $ ioProperty $ do+ eith <-+ try @_ @AbortActionRegistryError $+ withActionRegistry $ \reg -> do+ withRollback reg+ (pure ())+ (\_ -> throwIO (userError "rollback action failed"))+ throwIO (userError "error in withActionRegistry scope")+ pure $ case eith of+ Left e ->+ tabulate "displayException" [displayExceptionNewline e] $ property True+ Right () -> property False++displayExceptionNewline :: Exception e => e -> String+displayExceptionNewline e = '\n':displayException e
@@ -0,0 +1,84 @@+module Test.Control.Concurrent.Class.MonadSTM.RWVar (tests) where++import qualified Control.Concurrent.Class.MonadSTM.RWVar as RW+import Control.Monad.Class.MonadAsync+import Control.Monad.Class.MonadSay (MonadSay (say))+import Control.Monad.Class.MonadTest (MonadTest (exploreRaces))+import Control.Monad.Class.MonadThrow+import Control.Monad.IOSim+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Control.Concurrent.Class.MonadSTM.RWVar" [+ testProperty "prop_noRace" prop_noRace+ ]++data Action a = Read a | Incr a+ deriving stock (Show, Eq, Read)++instance Arbitrary a => Arbitrary (Action a) where+ arbitrary = oneof [+ Read <$> arbitrary+ , Incr <$> arbitrary+ ]+ shrink (Read x) = [Read x' | x' <- shrink x]+ shrink (Incr x) = [Incr x' | x' <- shrink x]++-- | Performing reads and increments on an @'RWVar' Int@ in parallel should not+-- lead to races. We let IOSimPOR check that there are no deadlocks. We also+-- look at the trace to see if the value inside the 'RWVar' increases+-- monotonically.+prop_noRace ::+ Action ()+ -> Action ()+ -> Action ()+ -> Property+prop_noRace a1 a2 a3 = exploreSimTrace modifyOpts prop $ \_ tr ->+ case traceResult False tr of+ Left e -> counterexample (show e) (property False)+ _ -> propSayMonotonic tr+ where+ modifyOpts = id++ prop :: IOSim s ()+ prop = do+ exploreRaces+ var <- RW.new (0 :: Int)+ let g = \case+ Read () ->+ RW.withReadAccess var $ \x -> do+ say (show (Read x))+ Incr () ->+ RW.withWriteAccess_ var $ \x -> do+ let x' = x + 1+ say (show (Incr x'))+ pure x'+ t1 <- async (g a1)+ t2 <- async (g a2)+ t3 <- async (g a3)+ async (cancel t1) >>= wait+ (_ :: Either AsyncCancelled ()) <- try (wait t1)+ (_ :: Either AsyncCancelled ()) <- try (wait t2)+ (_ :: Either AsyncCancelled ()) <- try (wait t3)+ pure ()++propSayMonotonic :: SimTrace () -> Property+propSayMonotonic simTrace = propResultsMonotonic (actionResults simTrace)++propResultsMonotonic :: [Action Int] -> Property+propResultsMonotonic as =+ counterexample+ ("Action results are not monotonic: " ++ show as)+ (resultsMonotonic as)++resultsMonotonic :: [Action Int] -> Bool+resultsMonotonic = go 0+ where+ go _ [] = True+ go prev (Read x : as) = prev == x && go x as+ go prev (Incr x : as) = prev + 1 == x && go x as++actionResults :: SimTrace () -> [Action Int]+actionResults = map read . selectTraceEventsSay
@@ -0,0 +1,252 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE TypeFamilies #-}++module Test.Control.RefCount (tests) where++import Control.Concurrent.Class.MonadMVar+import Control.Exception (AssertionFailed (..))+import Control.Monad+import Control.Monad.Class.MonadThrow+import Control.Monad.IOSim (IOSim, runSimOrThrow)+import Control.Monad.Primitive+import Control.RefCount+import Data.Primitive.PrimVar+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck++#ifdef NO_IGNORE_ASSERTS+import Data.Primitive+#endif++tests :: TestTree+tests = testGroup "Test.Control.RefCount" [+ testProperty "prop_RefCounter @IO" $+ ioPropertyOnce prop_RefCounter+#ifndef NO_IGNORE_ASSERTS+ -- prop_RefCounter throws and catches AssertionFailed exceptions, but we+ -- can only catch these exceptions in IO. In IOSim, these uncaught+ -- exceptions will lead to property failures. So, we only run the property+ -- in IOSim if assertions are turned off.+ , testProperty "prop_RefCounter @IOSim" $+ ioSimPropertyOnce prop_RefCounter+#endif+#ifdef NO_IGNORE_ASSERTS+ -- All of these tests below are checking that the debug implementation of+ -- Ref does indeed detect all the violations (double-free, use-after-free,+ -- never-free). But this obviously depends on the debug implementation+ -- being in use. Hence only tested when NO_IGNORE_ASSERTS.+ , testGroup "IO" [+ testProperty "prop_ref_double_free" $+ ioPropertyOnce prop_ref_double_free+ , testProperty "prop_ref_use_after_free" $+ ioPropertyOnce $ prop_ref_use_after_free True+ , testProperty "prop_ref_never_released0" $+ ioPropertyOnce prop_ref_never_released0+ , testProperty "prop_ref_never_released1" $+ ioPropertyOnce prop_ref_never_released1+ , testProperty "prop_ref_never_released2" $+ ioPropertyOnce prop_ref_never_released2+ , testProperty "prop_release_ref_exception" $+ ioPropertyOnce prop_release_ref_exception+ ]+ , testGroup "IOSim" [+ testProperty "prop_ref_double_free" $+ ioSimPropertyOnce prop_ref_double_free+ , testProperty "prop_ref_use_after_free" $+ -- Exceptions thrown by the DeRef pattern can only be caught from+ -- IO, so we do not test the DeRef pattern in IOSim.+ ioSimPropertyOnce $ prop_ref_use_after_free False+ , testProperty "prop_ref_never_released0" $+ ioSimPropertyOnce prop_ref_never_released0+ , testProperty "prop_ref_never_released1" $+ ioSimPropertyOnce prop_ref_never_released1+ , testProperty "prop_ref_never_released2" $+ ioSimPropertyOnce prop_ref_never_released2+ , testProperty "prop_release_ref_exception" $+ ioSimPropertyOnce prop_release_ref_exception+ ]+#endif+ ]++ioPropertyOnce :: Testable prop => IO prop -> Property+ioPropertyOnce p = once $ ioProperty p++ioSimPropertyOnce :: Testable prop => (forall s. IOSim s prop) -> Property+ioSimPropertyOnce p = once $ property $ runSimOrThrow p++-- | Test for the low level RefCounter API+prop_RefCounter ::+ (MonadMVar m, PrimMonad m, MonadMask m)+ => m Property+prop_RefCounter = do+ obj <- newMVar False+ ref <- newRefCounter (void $ modifyMVar_ obj (\x -> pure (not x)) )++ incrementRefCounter ref+ n1 <- readRefCount ref -- 2+ b1 <- readMVar obj -- False++ decrementRefCounter ref+ n2 <- readRefCount ref -- 1+ b2 <- readMVar obj -- False++ e2w <- tryIncrementRefCounter ref -- ok, True+ n2w <- readRefCount ref -- 2+ b2w <- readMVar obj -- False+ decrementRefCounter ref++ decrementRefCounter ref+ n3 <- readRefCount ref -- 0+ b3 <- readMVar obj -- True, finaliser ran++ e4 <- try (decrementRefCounter ref) -- error+ n4 <- readRefCount ref -- -1+ b4 <- readMVar obj -- True, finaliser did not run again++ e5 <- try (incrementRefCounter ref) -- error+ n5 <- readRefCount ref -- 0+ b5 <- readMVar obj -- True, finaliser did not run again++ e6 <- tryIncrementRefCounter ref+ n6 <- readRefCount ref -- 0+ b6 <- readMVar obj -- True, finaliser did not run again++ pure $+ counterexample "n1" (n1 == 2) .&&.+ counterexample "b1" (not b1) .&&.++ counterexample "n2" (n2 == 1) .&&.+ counterexample "b2" (not b2) .&&.++ counterexample "e2w" (e2w == True) .&&.+ counterexample "n2w" (n2w == 2) .&&.+ counterexample "b2w" (not b2w) .&&.++ counterexample "n3" (n3 == 0) .&&.+ counterexample "b3" b3 .&&.++ counterexample "e4" (check e4) .&&.+ counterexample "n4" (n4 == (-1)) .&&.+ counterexample "b4" b4 .&&.++ counterexample "e5" (check e5) .&&.+ counterexample "n5" (n5 == 0) .&&.+ counterexample "b5" b5 .&&.++ counterexample "e6" (e6 == False) .&&.+ counterexample "n6" (n6 == 0) .&&.+ counterexample "b6" b6+ where+#ifdef NO_IGNORE_ASSERTS+ check = \case Left (AssertionFailed _) -> True; Right () -> False+#else+ check = \case Left (AssertionFailed _) -> False; Right () -> True+#endif++-- | Warning: reading the current reference count is inherently racy as there+-- is no way to reliably act on the information. It is only useful for tests or+-- debugging.+--+readRefCount :: PrimMonad m => RefCounter m -> m Int+readRefCount (RefCounter countVar _) = readPrimVar countVar++#ifdef NO_IGNORE_ASSERTS+data TestObject m = TestObject !(RefCounter m)++instance RefCounted m (TestObject m) where+ getRefCounter (TestObject rc) = rc++data TestObject2 m = TestObject2 (Ref (TestObject m))++instance RefCounted m (TestObject2 m) where+ getRefCounter (TestObject2 (DeRef to1)) = getRefCounter to1++prop_ref_double_free ::+ (PrimMonad m, MonadMask m, MonadFail m)+ => m Property+prop_ref_double_free = do+ finalised <- newMutVar False+ ref <- newRef (writeMutVar finalised True) TestObject+ releaseRef ref+ True <- readMutVar finalised+ Left e@RefDoubleRelease{} <- try $ releaseRef ref+ checkForgottenRefs+ -- Print the displayed exception as an example+ pure $ tabulate "displayException" [displayException e] ()++prop_ref_use_after_free ::+ (PrimMonad m, MonadMask m, MonadFail m)+ => Bool -- ^ Test the DeRef pattern+ -> m Property+prop_ref_use_after_free testDeRef = do+ finalised <- newMutVar False+ ref <- newRef (writeMutVar finalised True) TestObject+ releaseRef ref+ True <- readMutVar finalised+ Left e@RefUseAfterRelease{} <- try $ withRef ref return+ when testDeRef $ do+ Left RefUseAfterRelease{} <- try $ case ref of DeRef _ -> pure ()+ pure ()+ Left RefUseAfterRelease{} <- try $ dupRef ref+ checkForgottenRefs+ -- Print the displayed exception as an example+ pure $ tabulate "displayException" [displayException e] ()++prop_ref_never_released0 ::+ (PrimMonad m, MonadMask m)+ => m ()+prop_ref_never_released0 = do+ finalised <- newMutVar False+ ref <- newRef (writeMutVar finalised True) TestObject+ _ <- case ref of DeRef _ -> pure ()+ checkForgottenRefs+ -- ref is still being used, so check should not fail+ _ <- case ref of DeRef _ -> pure ()+ releaseRef ref++prop_ref_never_released1 ::+ (PrimMonad m, MonadMask m)+ => m Property+prop_ref_never_released1 =+ handle expectRefNeverReleased $ do+ finalised <- newMutVar False+ ref <- newRef (writeMutVar finalised True) TestObject+ _ <- withRef ref return+ _ <- case ref of DeRef _ -> pure ()+ -- ref is never released, so should fail+ checkForgottenRefs+ pure (counterexample "no forgotten refs detected" $ property False)++prop_ref_never_released2 ::+ (PrimMonad m, MonadMask m)+ => m Property+prop_ref_never_released2 =+ handle expectRefNeverReleased $ do+ finalised <- newMutVar False+ ref <- newRef (writeMutVar finalised True) TestObject+ ref2 <- dupRef ref+ releaseRef ref+ _ <- withRef ref2 return+ _ <- case ref2 of DeRef _ -> pure ()+ -- ref2 is never released, so should fail+ checkForgottenRefs+ pure (counterexample "no forgotten refs detected" $ property False)++expectRefNeverReleased :: Monad m => RefException -> m Property+expectRefNeverReleased e@RefNeverReleased{} =+ -- Print the displayed exception as an example+ pure (tabulate "displayException" [displayException e] (property True))+expectRefNeverReleased e =+ pure (counterexample (displayException e) $ property False)++-- | If a finaliser throws an exception, then the 'RefTracker' is still released+prop_release_ref_exception ::+ (PrimMonad m, MonadMask m)+ => m ()+prop_release_ref_exception = do+ finalised <- newMutVar False+ ref <- newRef (writeMutVar finalised True >> throwIO (userError "oops")) TestObject+ _ <- try @_ @SomeException (releaseRef ref)+ checkForgottenRefs+#endif+
@@ -0,0 +1,16 @@+module Main (main) where++import Test.Tasty++import qualified Test.FormatPage+import qualified Test.ScheduledMerges+import qualified Test.ScheduledMerges.RunSizes+import qualified Test.ScheduledMergesQLS++main :: IO ()+main = defaultMain $ testGroup "prototypes" [+ Test.FormatPage.tests+ , Test.ScheduledMerges.tests+ , Test.ScheduledMerges.RunSizes.tests+ , Test.ScheduledMergesQLS.tests+ ]
@@ -0,0 +1,229 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}++module Test.FormatPage (tests) where++import qualified Data.ByteString as BS++import FormatPage++import Test.QuickCheck hiding ((.&.))+import Test.Tasty+import Test.Tasty.QuickCheck (testProperty)++-------------------------------------------------------------------------------+-- Tests+--++tests :: TestTree+tests = testGroup "FormatPage"+ [ testProperty "to/from bitmap" prop_toFromBitmap+ , testProperty "maxKeySize" prop_maxKeySize++ , let dpgsz = DiskPage4k in+ testGroup "size distribution"+ [ testProperty "genPageContentFits" $+ checkCoverage $+ coverTable "page size in bytes"+ [("0 <= n < 512",10)+ ,("3k < n <= 4k", 5)] $+ coverTable "page size in disk pages"+ [("1 page", 50)+ ,("2 pages", 0.5)+ ,("3+ pages", 0.5)] $+ forAll (genPageContentFits dpgsz noMinKeySize) $+ prop_size_distribution dpgsz (property False)++ , testProperty "genPageContentMaybeOverfull" $+ checkCoverage $+ forAll (genPageContentMaybeOverfull dpgsz noMinKeySize) $+ prop_size_distribution dpgsz+ (cover 10 True "over-full" (property True))++ , testProperty "genPageContentSingle" $+ checkCoverage $+ coverTable "page size in disk pages"+ [("1 page", 10)+ ,("2 pages", 2)+ ,("3+ pages", 2)] $+ forAll ((:[]) <$> genPageContentSingle dpgsz noMinKeySize) $+ prop_size_distribution dpgsz (property False)+ ]+ , testProperty "size 0" prop_size0+ , testProperty "size 1" prop_size1+ , testProperty "size 2" prop_size2+ , testProperty "size 3" prop_size3+ , testProperty "encode/decode" prop_encodeDecode+ , testProperty "serialise/deserialise" prop_serialiseDeserialise+ , testProperty "encode/serialise/deserialise/decode"+ prop_encodeSerialiseDeserialiseDecode+ , testProperty "overflow pages" prop_overflowPages+ ]++prop_toFromBitmap :: [Bool] -> Bool+prop_toFromBitmap bits =+ bits == take (length bits) (roundTrip bits)+ where+ roundTrip = fromBitmap . toBitmap++prop_size_distribution :: DiskPageSize+ -> Property -- ^ over-full sub-property+ -> [(Key, Operation)]+ -> Property+prop_size_distribution dpgsz propOverfull p =+ case calcPageSize dpgsz p of+ Nothing -> propOverfull+ Just PageSize{pageSizeElems, pageSizeBlobs, pageSizeBytes} ->+ tabulate "page size in elements"+ [ showNumElems pageSizeElems ] $+ tabulate "page number of blobs"+ [ showNumElems pageSizeBlobs ] $+ tabulate "page size in bytes"+ [ showPageSizeBytes pageSizeBytes ] $+ tabulate "page size in disk pages"+ [ showPageSizeDiskPages pageSizeBytes ] $+ tabulate "key size in bytes"+ [ showKeyValueSizeBytes (BS.length k) | (Key k, _) <- p ] $+ tabulate "value size in bytes"+ [ showKeyValueSizeBytes (BS.length v)+ | (_, op) <- p+ , Value v <- case op of+ Insert v _ -> [v]+ Mupsert v -> [v]+ Delete -> []+ ] $+ property $ (if pageSizeElems > 1+ then pageSizeBytes <= dpgszBytes+ else True)+ && (pageSizeElems == length p)+ && (pageSizeBlobs == length (filter (opHasBlobRef . snd) p))+ where+ dpgszBytes = diskPageSizeBytes dpgsz++ showNumElems :: Int -> String+ showNumElems n+ | n <= 1 = show n+ | n < 10 = "1 < n < 10"+ | otherwise = nearest 10 n++ showPageSizeBytes :: Int -> String+ showPageSizeBytes n+ | n > 4096 = nearest4k n+ | n > 1024 = nearest1k n+ | otherwise = nearest 512 n++ showPageSizeDiskPages :: Int -> String+ showPageSizeDiskPages n+ | npgs == 1 = "1 page"+ | npgs == 2 = "2 pages"+ | otherwise = "3+ pages"+ where+ npgs = (n + dpgszBytes - 1) `div` dpgszBytes++ showKeyValueSizeBytes :: Int -> String+ showKeyValueSizeBytes n+ | n < 20 = nearest 5 n+ | n < 100 = "20 <= n < 100"+ | n < 1024 = nearest 100 n+ | otherwise = nearest1k n++ nearest :: Int -> Int -> String+ nearest m n = show ((n `div` m) * m) ++ " <= n < "+ ++ show ((n `div` m) * m + m)++ nearest1k, nearest4k :: Int -> String+ nearest1k n = show ((n-1) `div` 1024) ++ "k < n <= "+ ++ show ((n-1) `div` 1024 + 1) ++ "k"+ nearest4k n = show (((n-1) `div` 4096) * 4) ++ "k < n <= "+ ++ show (((n-1) `div` 4096) * 4 + 4) ++ "k"++prop_maxKeySize :: Bool+prop_maxKeySize = maxKeySize DiskPage4k == 4052++-- | The 'calcPageSize' and 'calcPageSizeOffsets' (used by 'encodePage') had+-- better agree with each other!+--+-- The 'calcPageSize' uses the incremental 'PageSize' API, to work out the page+-- size, element by element, while 'calcPageSizeOffsets' is a bulk operation+-- used by 'encodePage'. It's critical that they agree on how many elements can+-- fit into a page.+--+prop_size0 :: PageContentMaybeOverfull -> Bool+prop_size0 (PageContentMaybeOverfull dpgsz p) =+ case (calcPageSize dpgsz p, encodePage dpgsz p) of+ (Nothing, Nothing) -> True+ (Nothing, Just{}) -> False -- they disagree!+ (Just{}, Nothing) -> False -- they disagree!+ (Just PageSize{..},+ Just PageIntermediate{pageSizesOffsets = PageSizesOffsets{..}, ..}) ->+ pageSizeElems == fromIntegral pageNumKeys+ && pageSizeBlobs == fromIntegral pageNumBlobs+ && pageSizeBytes == fromIntegral sizePageUsed+ && pageSizeDisk == pageDiskPageSize+ && pageSizeDisk == dpgsz++prop_size1 :: PageContentFits -> Bool+prop_size1 (PageContentFits dpgsz p) =+ sizePageUsed > 0+ && sizePageUsed + sizePagePadding == sizePageDiskPage+ && if pageNumKeys p' == 1+ then fromIntegral sizePageDiskPage `mod` diskPageSizeBytes dpgsz == 0+ else fromIntegral sizePageDiskPage == diskPageSizeBytes dpgsz+ where+ Just p' = encodePage dpgsz p+ PageSizesOffsets{..} = pageSizesOffsets p'++prop_size2 :: PageContentFits -> Bool+prop_size2 (PageContentFits dpgsz p) =+ BS.length (serialisePage p')+ == fromIntegral (sizePageDiskPage (pageSizesOffsets p'))+ where+ Just p' = encodePage dpgsz p++prop_size3 :: PageContentFits -> Bool+prop_size3 (PageContentFits dpgsz p) =+ case (calcPageSize dpgsz p, encodePage dpgsz p) of+ (Just PageSize{pageSizeBytes}, Just p') ->+ pageSizeBytes == (fromIntegral . sizePageUsed . pageSizesOffsets) p'+ _ -> False++prop_encodeDecode :: PageContentFits -> Property+prop_encodeDecode (PageContentFits dpgsz p) =+ p === decodePage p'+ where+ Just p' = encodePage dpgsz p++prop_serialiseDeserialise :: PageContentFits -> Bool+prop_serialiseDeserialise (PageContentFits dpgsz p) =+ p' == roundTrip p'+ where+ Just p' = encodePage dpgsz p+ roundTrip = deserialisePage dpgsz . serialisePage++prop_encodeSerialiseDeserialiseDecode :: PageContentFits -> Bool+prop_encodeSerialiseDeserialiseDecode (PageContentFits dpgsz p) =+ p == roundTrip p'+ where+ Just p' = encodePage dpgsz p+ roundTrip = decodePage . deserialisePage dpgsz . serialisePage++prop_overflowPages :: PageContentSingle -> Property+prop_overflowPages (PageContentSingle dpgsz k op) =+ label ("pages " ++ show (length ps)) $+ all ((== diskPageSizeBytes dpgsz) . BS.length) ps+ .&&. pageDiskPages p === length ps+ .&&. case pageOverflowPrefixSuffixLen p of+ Nothing -> length ps === 1+ Just (prefixlen, suffixlen) ->+ prefixlen + suffixlen === BS.length (unValue v)+ .&&. Value (BS.drop (dpgszBytes - prefixlen) (head ps)+ <> BS.take suffixlen (BS.concat (drop 1 ps)))+ === v+ where+ Just p = encodePage dpgsz [(k, op)]+ ps = pageSerialisedChunks dpgsz (serialisePage p)+ v = case op of+ Insert v' _ -> v'+ Mupsert v' -> v'+ Delete -> error "unexpected"+ dpgszBytes = diskPageSizeBytes dpgsz+
@@ -0,0 +1,576 @@+module Test.ScheduledMerges (tests) where++import Control.Exception+import Control.Monad (replicateM_, when)+import Control.Monad.ST+import Control.Tracer (Tracer (Tracer))+import qualified Control.Tracer as Tracer+import Data.Foldable (find, traverse_)+import Data.Maybe (fromJust)+import Data.STRef+import Text.Printf (printf)++import ScheduledMerges as LSM++import qualified Test.QuickCheck as QC+import Test.QuickCheck (Arbitrary (arbitrary, shrink), Property)+import Test.QuickCheck.Exception (isDiscard)+import Test.Tasty+import Test.Tasty.HUnit (HasCallStack, testCase)+import Test.Tasty.QuickCheck (QuickCheckMaxSize (..),+ QuickCheckTests (..), testProperty, (=/=), (===))++tests :: TestTree+tests = testGroup "Test.ScheduledMerges"+ [ testCase "test_regression_empty_run" test_regression_empty_run+ , testCase "test_merge_again_with_incoming" test_merge_again_with_incoming+ , testProperty "prop_union" prop_union+ , testGroup "T"+ [ localOption (QuickCheckTests 1000) $ -- super quick, run more+ testProperty "Arbitrary satisfies invariant" prop_arbitrarySatisfiesInvariant+ , localOption (QuickCheckMaxSize 60) $ -- many shrinks for huge trees+ testProperty "Shrinking satisfies invariant" prop_shrinkSatisfiesInvariant+ ]+ , testProperty "prop_MergingTree" prop_MergingTree+ ]++-- | Results in an empty run on level 2.+test_regression_empty_run :: IO ()+test_regression_empty_run =+ runWithTracer $ \tracer -> do+ stToIO $ do+ lsm <- LSM.new tracer (LSM.TableId 0)+ let ins k = LSM.insert tracer lsm (K k) (V 0) Nothing+ let del k = LSM.delete tracer lsm (K k)+ -- run 1+ ins 0+ ins 1+ ins 2+ ins 3+ -- run 2+ ins 0+ ins 1+ ins 2+ ins 3+ -- run 3+ ins 0+ ins 1+ ins 2+ ins 3+ -- run 4, deletes all previous elements+ del 0+ del 1+ del 2+ del 3++ expectShape lsm+ 0+ [ ([], [4,4,4,4])+ ]++ -- run 5, results in last level merge of run 1-4+ ins 0+ ins 1+ ins 2+ ins 3++ expectShape lsm+ 0+ [ ([], [4])+ , ([4,4,4,4], [])+ ]++ -- finish merge+ LSM.supplyMergeCredits lsm (NominalCredit 16)++ expectShape lsm+ 0+ [ ([], [4])+ , ([], [0])+ ]++ -- insert more data, so the empty run becomes input to a merge+ traverse_ ins [101..112]++ expectShape lsm+ 0+ [ ([], [4,4,4,4]) -- about to trigger a new last level merge+ , ([], [0])+ ]++ traverse_ ins [113..116]++ expectShape lsm+ 0+ [ ([], [4])+ , ([4,4,4,4], []) -- merge started, empty run has been dropped+ ]++-- | Covers the case where a run ends up too small for a level, so it gets+-- merged again with the next incoming runs.+-- That 5-way merge gets completed by supplying credits That merge gets+-- completed by supplying credits and then becomes part of another merge.+test_merge_again_with_incoming :: IO ()+test_merge_again_with_incoming =+ runWithTracer $ \tracer -> do+ stToIO $ do+ lsm <- LSM.new tracer (LSM.TableId 0)+ let ins k = LSM.insert tracer lsm (K k) (V 0) Nothing+ -- get something to 3rd level (so 2nd level is not levelling)+ -- (needs 5 runs to go to level 2 so the resulting run becomes too big)+ traverse_ ins [101..100+(5*16)]++ -- also get a very small run (4 elements) to 2nd level by re-using keys+ replicateM_ 4 $+ traverse_ ins [201..200+4]++ expectShape lsm+ 0+ [ ([], [4,4,4,4]) -- these runs share keys, will compact down to 4+ , ([4,4,4,4,64], []) -- this run will end up in level 3+ ]++ -- get another run to 2nd level, which the small run can be merged with+ traverse_ ins [301..300+16]++ expectShape lsm+ 0+ [ ([], [4,4,4,4])+ , ([4,4,4,4], [])+ , ([], [80])+ ]++ -- add just one more run so the 5-way merge on 2nd level gets created+ traverse_ ins [401..400+4]++ expectShape lsm+ 0+ [ ([], [4])+ , ([4,4,4,4,4], [])+ , ([], [80])+ ]++ -- complete the merge (20 entries, but credits get scaled up by 1.25)+ LSM.supplyMergeCredits lsm (NominalCredit 16)++ expectShape lsm+ 0+ [ ([], [4])+ , ([], [20])+ , ([], [80])+ ]++ -- get 3 more runs to 2nd level, so the 5-way merge completes+ -- and becomes part of a new merge.+ -- (actually 4, as runs only move once a fifth run arrives...)+ traverse_ ins [501..500+(4*16)]++ expectShape lsm+ 0+ [ ([], [4])+ , ([4,4,4,4], [])+ , ([16,16,16,20,80], [])+ ]++-------------------------------------------------------------------------------+-- properties+--++-- | Supplying enough credits for the remaining debt completes the union merge.+prop_union :: [[(LSM.Key, LSM.Entry)]] -> Property+prop_union kess = length (filter (not . null) kess) > 1 QC.==>+ QC.ioProperty $ runWithTracer $ \tr ->+ stToIO $ do+ ts <- traverse (uncurry $ mkTable tr) (zip [LSM.TableId 0..] kess)+ t <- LSM.unions tr (LSM.TableId (length kess)) ts++ debt@(UnionDebt x) <- LSM.remainingUnionDebt t+ _ <- LSM.supplyUnionCredits t (UnionCredits x)+ debt' <- LSM.remainingUnionDebt t++ rep <- dumpRepresentation t+ pure $ QC.counterexample (show (debt, debt')) $ QC.conjoin+ [ debt =/= UnionDebt 0+ , debt' === UnionDebt 0+ , hasUnionWith isCompleted rep+ ]+ where+ isCompleted = \case+ MLeaf{} -> True+ MNode{} -> False++mkTable :: Tracer (ST s) Event -> LSM.TableId -> [(LSM.Key, LSM.Entry)] -> ST s (LSM s)+mkTable tr tid ks = do+ t <- LSM.new tr tid+ LSM.updates tr t ks+ pure t++-------------------------------------------------------------------------------+-- tests for MergingTree+--++prop_MergingTree :: T -> QC.InfiniteList SmallCredit -> Property+prop_MergingTree TCompleted{} _ = QC.discard+prop_MergingTree (TOngoing MCompleted{}) _ = QC.discard+prop_MergingTree t credits =+ QC.ioProperty $ runWithTracer $ \_tr ->+ stToIO $ do+ tree <- fromT t+ res <- go tree (QC.getInfiniteList credits)+ pure $+ res === Right ()+ where+ -- keep supplying until there is an error or the tree merge is completed+ go :: MergingTree s -> [SmallCredit] -> ST s (Either String ())+ go tree (SmallCredit c : cs) = do+ c' <- LSM.supplyCreditsMergingTree c tree+ evalInvariant (treeInvariant tree) >>= \case+ Left e -> pure (Left e)+ Right () -> if c' > 0 then pure (Right ())+ else go tree cs+ go _ [] = error "infinite list is finite"++newtype SmallCredit = SmallCredit Credit+ deriving stock Show++instance Arbitrary SmallCredit where+ arbitrary = SmallCredit <$> QC.chooseInt (1, 10)+ shrink (SmallCredit c) = [SmallCredit c' | c' <- shrink c, c' > 0]++-- simplified non-ST version of MergingTree+data T = TCompleted Run+ | TOngoing (M TreeMergeType)+ | TPendingLevel [P] (Maybe T) -- not both empty!+ | TPendingUnion [T] -- at least 2 children+ deriving stock Show++-- simplified non-ST version of PreExistingRun+data P = PRun Run+ | PMergingRun (M LevelMergeType)+ deriving stock Show++-- simplified non-ST version of MergingRun+data M t = MCompleted t MergeDebt Run+ | MOngoing+ t+ MergeDebt -- debt bounded by input sizes+ MergeCredit+ [NonEmptyRun] -- at least 2 inputs+ deriving stock Show++newtype NonEmptyRun = NonEmptyRun { getNonEmptyRun :: Run }+ deriving stock Show++invariantT :: T -> Either String ()+invariantT t = runST $ do+ tree <- fromT t+ evalInvariant (treeInvariant tree)++-- | Size is the number of T and P constructors.+sizeT :: T -> Int+sizeT (TCompleted _) = 1+sizeT (TOngoing _) = 1+sizeT (TPendingLevel ps mt) = sum (fmap sizeP ps) + maybe 0 sizeT mt+sizeT (TPendingUnion ts) = sum (fmap sizeT ts)++sizeP :: P -> Int+sizeP (PRun _) = 1+sizeP (PMergingRun _) = 1++-- | Depth is the longest path through the tree from the root to a leaf using T+-- and P constructors.+depthT :: T -> Int+depthT (TCompleted _) = 0+depthT (TOngoing _) = 0+depthT (TPendingLevel ps mt) =+ let depthPs = case ps of+ [] -> 0+ _ -> maximum (fmap depthP ps)+ depthMt = maybe 0 depthT mt+ in 1 + max depthPs depthMt+depthT (TPendingUnion ts) = case ts of+ [] -> 0+ _ -> 1 + maximum (fmap depthT ts)++depthP :: P -> Int+depthP (PRun _) = 0+depthP (PMergingRun _) = 0++fromT :: T -> ST s (MergingTree s)+fromT t = do+ state <- case t of+ TCompleted r -> pure (CompletedTreeMerge r)+ TOngoing mr -> OngoingTreeMerge <$> fromM mr+ TPendingLevel ps mt ->+ fmap PendingTreeMerge $+ PendingLevelMerge <$> traverse fromP ps <*> traverse fromT mt+ TPendingUnion ts -> do+ fmap PendingTreeMerge $ PendingUnionMerge <$> traverse fromT ts+ MergingTree <$> newSTRef state++fromP :: P -> ST s (PreExistingRun s)+fromP (PRun r) = pure (PreExistingRun r)+fromP (PMergingRun m) = PreExistingMergingRun <$> fromM m++fromM :: IsMergeType t => M t -> ST s (MergingRun t s)+fromM m = do+ let (mergeType, mergeDebt, state) = case m of+ MCompleted mt md r -> (mt, md, CompletedMerge r)+ MOngoing mt md mc rs -> (mt, md, OngoingMerge mc rs' (mergek mt rs'))+ where rs' = map getNonEmptyRun rs+ MergingRun mergeType mergeDebt <$> newSTRef state++completeT :: T -> Run+completeT (TCompleted r) = r+completeT (TOngoing m) = completeM m+completeT (TPendingLevel is t) =+ mergek MergeLevel (map completeP is <> maybe [] (pure . completeT) t)+completeT (TPendingUnion ts) =+ mergek MergeUnion (map completeT ts)++completeP :: P -> Run+completeP (PRun r) = r+completeP (PMergingRun m) = completeM m++completeM :: IsMergeType t => M t -> Run+completeM (MCompleted _ _ r) = r+completeM (MOngoing mt _ _ rs) = mergek mt (map getNonEmptyRun rs)++-------------------------------------------------------------------------------+-- Generators+--++instance Arbitrary T where+ arbitrary = QC.sized $ \s -> do+ n <- QC.chooseInt (1, max 1 s)+ go n+ where+ -- n is the number of constructors of T and P+ go n | n < 1 = error ("arbitrary T: n == " <> show n)+ go n | n == 1 =+ QC.frequency+ [ (1, TCompleted <$> arbitrary)+ , (1, TOngoing <$> arbitrary)+ ]+ go n =+ QC.frequency+ [ (1, do+ -- pending level merge without child+ preExisting <- QC.vector (n - 1) -- 1 for constructor itself+ pure (TPendingLevel preExisting Nothing))+ , (1, do+ -- pending level merge with child+ numPreExisting <- QC.chooseInt (0, min 20 (n - 2))+ preExisting <- QC.vector numPreExisting+ tree <- go (n - numPreExisting - 1)+ pure (TPendingLevel preExisting (Just tree)))+ , (2, do+ -- pending union merge+ ns <- QC.shuffle =<< arbitraryPartition2 n+ TPendingUnion <$> traverse go ns)+ ]++ -- Split into at least two smaller positive numbers. The input needs to be+ -- greater than or equal to 2.+ arbitraryPartition2 :: Int -> QC.Gen [Int]+ arbitraryPartition2 n = assert (n >= 2) $ do+ first <- QC.chooseInt (1, n-1)+ (first :) <$> arbitraryPartition (n - first)++ -- Split into smaller positive numbers.+ arbitraryPartition :: Int -> QC.Gen [Int]+ arbitraryPartition n+ | n < 1 = pure []+ | n == 1 = pure [1]+ | otherwise = do+ first <- QC.chooseInt (1, n)+ (first :) <$> arbitraryPartition (n - first)++ shrink (TCompleted r) =+ [ TCompleted r'+ | r' <- shrink r+ ]+ shrink tree@(TOngoing m) =+ [ TCompleted (completeT tree) ]+ <> [ TOngoing m'+ | m' <- shrink m+ ]+ shrink tree@(TPendingLevel ps t) =+ [ TCompleted (completeT tree) ]+ <> [ t' | Just t' <- [t] ]+ <> [ TPendingLevel (ps ++ [PRun r]) Nothing -- move into regular levels+ | Just (TCompleted r) <- [t]+ ]+ <> [ TPendingLevel ps' t'+ | (ps', t') <- shrink (ps, t)+ , length ps' + length t' > 0+ ]+ shrink tree@(TPendingUnion ts) =+ [ TCompleted (completeT tree) ]+ <> ts+ <> [ TPendingUnion ts'+ | ts' <- shrink ts+ , length ts' > 1+ ]++instance Arbitrary P where+ arbitrary = QC.oneof [PRun <$> arbitrary, PMergingRun <$> arbitrary]+ shrink (PRun r) = [PRun r' | r' <- shrink r]+ shrink (PMergingRun m) = [PRun (completeM m)]+ <> [PMergingRun m' | m' <- shrink m]++instance (Arbitrary t, IsMergeType t) => Arbitrary (M t) where+ arbitrary = QC.oneof+ [ do (mt, r) <- arbitrary+ let md = MergeDebt (runSize r)+ pure (MCompleted mt md r)+ , do mt <- arbitrary+ n <- QC.chooseInt (2, 8)+ rs <- QC.vectorOf n (QC.scale (`div` n) arbitrary)+ (md, mc) <- genMergeCreditForRuns rs+ pure (MOngoing mt md mc rs)+ ]++ shrink (MCompleted mt md r) =+ [ MCompleted mt md r' | r' <- shrink r ]+ shrink m@(MOngoing mt md mc rs) =+ [ MCompleted mt md (completeM m) ]+ <> [ MOngoing mt md' mc' rs'+ | rs' <- shrink rs+ , length rs' > 1+ , (md', mc') <- shrinkMergeCreditForRuns rs' mc+ ]++-- | The 'MergeDebt' and 'MergeCredit' must maintain a couple invariants:+--+-- * the total debt must be the same as the sum of the input run sizes;+-- * the supplied credit is less than the total merge debt.+--+genMergeCreditForRuns :: [NonEmptyRun] -> QC.Gen (MergeDebt, MergeCredit)+genMergeCreditForRuns rs = do+ let totalDebt = sum (map (length . getNonEmptyRun) rs)+ suppliedCredits <- QC.chooseInt (0, totalDebt-1)+ unspentCredits <- QC.chooseInt (0, min (mergeBatchSize-1) suppliedCredits)+ let spentCredits = suppliedCredits - unspentCredits+ md = MergeDebt {+ totalDebt+ }+ mc = MergeCredit {+ unspentCredits,+ spentCredits+ }+ assert (mergeDebtInvariant md mc) $+ pure (md, mc)++-- | Shrink the 'MergeDebt' and 'MergeCredit' given the old 'MergeCredit' and+-- the already-shrunk runs.+--+-- Thus must maintain invariants, see 'genMergeCreditForDebt'.+--+shrinkMergeCreditForRuns :: [NonEmptyRun]+ -> MergeCredit -> [(MergeDebt, MergeCredit)]+shrinkMergeCreditForRuns rs' MergeCredit {spentCredits, unspentCredits} =+ [ assert (mergeDebtInvariant md' mc')+ (md', mc')+ | let totalDebt' = sum (map (length . getNonEmptyRun) rs')+ , suppliedCredits' <- shrink (min (spentCredits+unspentCredits)+ (totalDebt'-1))+ , unspentCredits' <- shrink (min unspentCredits suppliedCredits')+ , let spentCredits' = suppliedCredits' - unspentCredits'+ md' = MergeDebt {+ totalDebt = totalDebt'+ }+ mc' = MergeCredit {+ spentCredits = spentCredits',+ unspentCredits = unspentCredits'+ }+ ]++instance Arbitrary NonEmptyRun where+ arbitrary = NonEmptyRun <$> (arbitrary `QC.suchThat` (not . null))+ shrink (NonEmptyRun r) = [NonEmptyRun r' | r' <- shrink r, not (null r')]++prop_arbitrarySatisfiesInvariant :: T -> Property+prop_arbitrarySatisfiesInvariant t =+ QC.tabulate "Tree size" [showPowersOf 2 $ sizeT t] $+ QC.tabulate "Tree depth" [showPowersOf 2 $ depthT t] $+ Right () === invariantT t++prop_shrinkSatisfiesInvariant :: T -> Property+prop_shrinkSatisfiesInvariant t =+ QC.forAll (genShrinkTrace 4 t) $ \trace ->+ QC.tabulate "Trace length" [showPowersOf 2 $ length trace] $+ QC.conjoin $ flip map trace $ \(numAlternatives, t') ->+ QC.tabulate "Shrink alternatives" [showPowersOf 2 numAlternatives] $+ Right () === invariantT t'++-- | Iterative shrinks, and how many alternatives were possible at each point.+genShrinkTrace :: Arbitrary a => Int -> a -> QC.Gen [(Int, a)]+genShrinkTrace !n x+ | n <= 0 = pure []+ | otherwise =+ case shrink x of+ [] -> pure []+ xs -> do+ -- like QC.elements, but we want access to the length+ let len = length xs+ x' <- (xs !!) <$> QC.chooseInt (0, len - 1)+ ((len, x') :) <$> genShrinkTrace (n - 1) x'++-------------------------------------------------------------------------------+-- tracing and expectations on LSM shape+--++-- | Provides a tracer and will add the log of traced events to the reported+-- failure.+runWithTracer :: (Tracer (ST RealWorld) Event -> IO a) -> IO a+runWithTracer action = do+ events <- stToIO $ newSTRef []+ let tracer = Tracer $ Tracer.emit $ \e -> modifySTRef events (e :)+ action tracer `catch` \e -> do+ if isDiscard e -- don't intercept these+ then throwIO e+ else do+ ev <- reverse <$> stToIO (readSTRef events)+ throwIO (Traced e ev)++data TracedException = Traced SomeException [Event]+ deriving stock (Show)++instance Exception TracedException where+ displayException (Traced e ev) =+ displayException e <> "\ntrace:\n" <> unlines (map show ev)++expectShape :: HasCallStack => LSM s -> Int -> [([Int], [Int])] -> ST s ()+expectShape lsm expectedWb expectedLevels = do+ let expected = (expectedWb, expectedLevels, Nothing)+ shape <- representationShape <$> dumpRepresentation lsm+ when (shape /= expected) $+ error $ unlines+ [ "expected shape: " <> show expected+ , "actual shape: " <> show shape+ ]++hasUnionWith :: (MTree Int -> Bool) -> Representation -> Property+hasUnionWith p rep = do+ let (_, _, shape) = representationShape rep+ QC.counterexample "expected suitable Union" $+ QC.counterexample (show shape) $+ case shape of+ Nothing -> False+ Just t -> p t++-------------------------------------------------------------------------------+-- Printing utils+--++-- | Copied from @lsm-tree:extras.Database.LSMTree.Extras@+showPowersOf :: Int -> Int -> String+showPowersOf factor n+ | factor <= 1 = error "showPowersOf: factor must be larger than 1"+ | n < 0 = "n < 0"+ | n == 0 = "n == 0"+ | otherwise = printf "%d <= n < %d" lb ub+ where+ ub = fromJust (find (n <) (iterate (* factor) factor))+ lb = ub `div` factor
@@ -0,0 +1,118 @@+module Test.ScheduledMerges.RunSizes (tests) where++import qualified ScheduledMerges as Proto+import ScheduledMerges hiding (MergePolicyForLevel)+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.ScheduledMerges.RunSizes" [+ testProperty "prop_roundTrip_levelNumberMaxRunSize" $+ \mpl conf ln -> levelNumberInvariant mpl conf ln ==>+ prop_roundTrip_levelNumberMaxRunSize mpl conf ln+ , testProperty "prop_roundTrip_runSizeLevelNumber" prop_roundTrip_runSizeLevelNumber+ , testProperty "prop_maxWriteBufferSize" prop_maxWriteBufferSize+ , testProperty "prop_runSizeFitsInLevel" $+ \mpl conf ln -> levelNumberInvariant mpl conf ln ==>+ prop_runSizeFitsInLevel mpl conf ln+ ]++-- | Test 'levelNumberToMaxRunSize' roundtrips with 'runSizeToLevelNumber'.+prop_roundTrip_levelNumberMaxRunSize :: MergePolicyForLevel -> Config -> LevelNo -> Property+prop_roundTrip_levelNumberMaxRunSize (MergePolicyForLevel mpl) (Config conf) (LevelNo ln) =+ let n = levelNumberToMaxRunSize mpl conf ln+ ln' = runSizeToLevelNumber mpl conf n+ in ln === ln'++-- | Test that 'runSizeToLevelNumber'roundtrips with 'levelNumberToMaxRunSize'.+prop_roundTrip_runSizeLevelNumber :: MergePolicyForLevel -> Config -> RunSize -> Property+prop_roundTrip_runSizeLevelNumber (MergePolicyForLevel mpl) (Config conf) (RunSize n) =+ let ln = runSizeToLevelNumber mpl conf n+ in if ln == 0+ then 0 === n+ else let n1 = levelNumberToMaxRunSize mpl conf (ln - 1)+ n2 = levelNumberToMaxRunSize mpl conf ln+ in property (n1 < n && n <= n2)++-- | Test that 'maxWriteBufferSize' equals the configured 'configMaxWriteBufferSize'.+prop_maxWriteBufferSize :: Config -> Property+prop_maxWriteBufferSize (Config conf) =+ configMaxWriteBufferSize conf === maxWriteBufferSize conf++-- | If a run is neither too small or too large for a level, then it fits in+-- that level.+prop_runSizeFitsInLevel :: MergePolicyForLevel -> Config -> LevelNo -> RunSize -> Property+prop_runSizeFitsInLevel (MergePolicyForLevel mpl) (Config conf) (LevelNo ln) (RunSize n) =+ classify (runSizeFitsInLevel mpl conf ln n) "Run size fits in level" $+ (not tooSmall && not tooLarge) === fits+ where+ tooSmall = runSizeTooSmallForLevel mpl conf ln n+ tooLarge = runSizeTooLargeForLevel mpl conf ln n+ fits = runSizeFitsInLevel mpl conf ln n++{-------------------------------------------------------------------------------+ Generators and shrinkers+-------------------------------------------------------------------------------}++newtype MergePolicyForLevel = MergePolicyForLevel Proto.MergePolicyForLevel+ deriving stock (Show, Eq)++instance Arbitrary MergePolicyForLevel where+ arbitrary = MergePolicyForLevel <$> elements [Proto.LevelTiering, Proto.LevelLevelling]+ shrink (MergePolicyForLevel x) = MergePolicyForLevel <$> case x of+ Proto.LevelTiering -> []+ Proto.LevelLevelling -> [Proto.LevelTiering]++newtype Config = Config LSMConfig+ deriving stock (Show, Eq)++instance Arbitrary Config where+ arbitrary = Config <$> do+ bufSize <- (getSmall <$> arbitrary) `suchThat` (>0)+ sizeRatio <- (getSmall <$> arbitrary) `suchThat` (>1)+ pure $ LSMConfig {+ configMaxWriteBufferSize = bufSize+ , configSizeRatio = sizeRatio+ }+ shrink (Config conf@LSMConfig{..}) =+ [ Config conf{configMaxWriteBufferSize = bufSize'}+ | bufSize' <- shrink configMaxWriteBufferSize+ , bufSize' > 0+ ]+ ++ [ Config conf{configSizeRatio = sizeRatio'}+ | sizeRatio' <- shrink configSizeRatio+ , sizeRatio' > 1+ ]+++newtype LevelNo = LevelNo Int+ deriving stock (Show, Eq, Ord)+ deriving Arbitrary via NonNegative Int++-- | The maximum size of a run on a level scales exponentially with the level+-- number, and linearly with the maximum write buffer size. There is a+-- possibility for 'Int' overflow primarily because of the exponential factor,+-- which would make the @ScheduledMerges@ prototype throw an error if using+-- @fromIntegerChecked@. The 'noOverflow' predicate makes sure that we do not+-- generate too large level numbers that could cause such overflows. But this+-- overflow condition also depends on the maximum write buffer size and merge+-- policy, so both are arguments to this function as well.+levelNumberInvariant :: MergePolicyForLevel -> Config -> LevelNo -> Bool+levelNumberInvariant+ (MergePolicyForLevel mpl)+ (Config LSMConfig{configMaxWriteBufferSize, configSizeRatio})+ (LevelNo ln)+ | ln < 0 = False+ | ln == 0 = True+ | otherwise = case mpl of+ Proto.LevelTiering ->+ toInteger configMaxWriteBufferSize * (toInteger configSizeRatio ^ toInteger (pred ln))+ <= toInteger (maxBound :: Int)+ Proto.LevelLevelling ->+ toInteger configMaxWriteBufferSize * (toInteger configSizeRatio ^ toInteger ln)+ <= toInteger (maxBound :: Int)++newtype RunSize = RunSize Int+ deriving stock (Show, Eq, Ord)+ deriving Arbitrary via NonNegative Int
@@ -0,0 +1,463 @@+{-# LANGUAGE TypeFamilies #-}++{-# OPTIONS_GHC -Wno-orphans #-}+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}++module Test.ScheduledMergesQLS (tests) where++import Control.Monad.Reader (ReaderT (..))+import Control.Monad.ST+import Control.Tracer (Tracer, nullTracer)+import Data.Constraint (Dict (..))+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromJust)+import Data.Primitive.PrimVar+import Data.Proxy+import Data.Semigroup (First (..))+import Prelude hiding (lookup)++import ScheduledMerges as LSM++import Test.QuickCheck hiding (Some)+import Test.QuickCheck.StateModel hiding (lookUpVar)+import Test.QuickCheck.StateModel.Lockstep hiding (ModelOp)+import qualified Test.QuickCheck.StateModel.Lockstep.Defaults as Lockstep+import qualified Test.QuickCheck.StateModel.Lockstep.Run as Lockstep+import Test.Tasty+import Test.Tasty.QuickCheck (testProperty)++tests :: TestTree+tests = testGroup "Test.ScheduledMergesQLS" [+ testProperty "ScheduledMerges vs model" $ mapSize (*10) prop_LSM -- still <10s+ ]++-- TODO: add tagging, e.g. how often ASupplyUnion makes progress or completes a+-- union merge.+prop_LSM :: Actions (Lockstep Model) -> Property+prop_LSM =+ Lockstep.runActionsBracket+ (Proxy :: Proxy Model)+ (newPrimVar (TableId 0))+ (\_ -> pure ())+ runReaderT++-------------------------------------------------------------------------------+-- QLS infrastructure+--++type ModelLSM = Int++newtype Model = Model { mlsms :: Map ModelLSM Table }+ deriving stock (Show)++data Table = Table {+ tableContent :: !(Map Key (Value, Maybe Blob))+ , tableHasUnion :: !Bool+ }+ deriving stock Show++emptyTable :: Table+emptyTable = Table Map.empty False++onContent ::+ (Map Key (Value, Maybe Blob) -> Map Key (Value, Maybe Blob))+ -> ModelLSM+ -> Model+ -> Model+onContent f k (Model m) = Model (Map.adjust f' k m)+ where+ f' t = t { tableContent = f (tableContent t) }++type ModelOp r = Model -> (r, Model)++initModel :: Model+initModel = Model { mlsms = Map.empty }++resolveValueAndBlob :: (Value, Maybe Blob) -> (Value, Maybe Blob) -> (Value, Maybe Blob)+resolveValueAndBlob (v, b) (v', b') = (resolveValue v v', getFirst (First b <> First b'))++modelNew :: ModelOp ModelLSM+modelInsert :: ModelLSM -> Key -> Value -> Maybe Blob -> ModelOp ()+modelDelete :: ModelLSM -> Key -> ModelOp ()+modelMupsert :: ModelLSM -> Key -> Value -> ModelOp ()+modelLookup :: ModelLSM -> Key -> ModelOp (LookupResult Value Blob)+modelDuplicate :: ModelLSM -> ModelOp ModelLSM+modelUnions :: [ModelLSM] -> ModelOp ModelLSM+modelSupplyUnion :: ModelLSM -> NonNegative UnionCredits -> ModelOp ()+modelDump :: ModelLSM -> ModelOp (Map Key (Value, Maybe Blob))++modelNew Model {mlsms} =+ (mlsm, Model { mlsms = Map.insert mlsm emptyTable mlsms })+ where+ mlsm = Map.size mlsms++modelInsert mlsm k v b model =+ ((), onContent (Map.insert k (v, b)) mlsm model)++modelDelete mlsm k model =+ ((), onContent (Map.delete k) mlsm model)++modelMupsert mlsm k v model =+ ((), onContent (Map.insertWith resolveValueAndBlob k (v, Nothing)) mlsm model)++modelLookup mlsm k model@Model {mlsms} =+ (result, model)+ where+ Just mval = Map.lookup mlsm mlsms+ result = case Map.lookup k (tableContent mval) of+ Nothing -> NotFound+ Just (v, mb) -> Found v mb++modelDuplicate mlsm Model {mlsms} =+ (mlsm', Model { mlsms = Map.insert mlsm' mval mlsms })+ where+ Just mval = Map.lookup mlsm mlsms+ mlsm' = Map.size mlsms++modelUnions inputs Model {mlsms} =+ (mlsm', Model { mlsms = Map.insert mlsm' mval' mlsms })+ where+ contents = map (\i -> tableContent (fromJust (Map.lookup i mlsms))) inputs+ hasUnion = True+ mval' = Table (Map.unionsWith resolveValueAndBlob contents) hasUnion+ mlsm' = Map.size mlsms++modelSupplyUnion _mlsm _credit model =+ ((), model)++modelDump mlsm model@Model {mlsms} =+ (mval, model)+ where+ Just (Table mval _) = Map.lookup mlsm mlsms++instance StateModel (Lockstep Model) where+ data Action (Lockstep Model) a where+ ANew :: LSMConfig -> Action (Lockstep Model) (LSM RealWorld)++ AInsert :: ModelVar Model (LSM RealWorld)+ -> Either (ModelVar Model Key) Key -- to refer to a prior key+ -> Value+ -> Maybe Blob+ -> Action (Lockstep Model) (Key)++ ADelete :: ModelVar Model (LSM RealWorld)+ -> Either (ModelVar Model Key) Key+ -> Action (Lockstep Model) ()++ AMupsert :: ModelVar Model (LSM RealWorld)+ -> Either (ModelVar Model Key) Key+ -> Value+ -> Action (Lockstep Model) (Key)++ ALookup :: ModelVar Model (LSM RealWorld)+ -> Either (ModelVar Model Key) Key+ -> Action (Lockstep Model) (LookupResult Value Blob)++ ADuplicate :: ModelVar Model (LSM RealWorld)+ -> Action (Lockstep Model) (LSM RealWorld)++ AUnions :: [ModelVar Model (LSM RealWorld)]+ -> Action (Lockstep Model) (LSM RealWorld)++ ASupplyUnion :: ModelVar Model (LSM RealWorld)+ -> NonNegative UnionCredits+ -> Action (Lockstep Model) ()++ ADump :: ModelVar Model (LSM RealWorld)+ -> Action (Lockstep Model) (Map Key (Value, Maybe Blob))++ initialState = Lockstep.initialState initModel+ nextState = Lockstep.nextState+ precondition = Lockstep.precondition+ arbitraryAction = Lockstep.arbitraryAction+ shrinkAction = Lockstep.shrinkAction++instance RunModel (Lockstep Model) (ReaderT (PrimVar RealWorld TableId) IO) where+ perform = \_state -> runActionIO+ postcondition = Lockstep.postcondition+ monitoring = Lockstep.monitoring (Proxy :: Proxy (ReaderT (PrimVar RealWorld TableId) IO))++instance InLockstep Model where+ data ModelValue Model a where+ MLSM :: ModelLSM+ -> ModelValue Model (LSM RealWorld)+ MUnit :: ()+ -> ModelValue Model ()+ MInsert :: Key+ -> ModelValue Model Key+ MLookup :: LookupResult Value Blob+ -> ModelValue Model (LookupResult Value Blob)+ MDump :: Map Key (Value, Maybe Blob)+ -> ModelValue Model (Map Key (Value, Maybe Blob))++ data Observable Model a where+ ORef :: Observable Model (LSM RealWorld)+ OId :: (Show a, Eq a) => a -> Observable Model a++ observeModel (MLSM _) = ORef+ observeModel (MUnit x) = OId x+ observeModel (MInsert x) = OId x+ observeModel (MLookup x) = OId x+ observeModel (MDump x) = OId x++ usedVars ANew{} = []+ usedVars (AInsert v evk _ _) = SomeGVar v+ : case evk of Left vk -> [SomeGVar vk]; _ -> []+ usedVars (ADelete v evk) = SomeGVar v+ : case evk of Left vk -> [SomeGVar vk]; _ -> []+ usedVars (AMupsert v evk _) = SomeGVar v+ : case evk of Left vk -> [SomeGVar vk]; _ -> []+ usedVars (ALookup v evk) = SomeGVar v+ : case evk of Left vk -> [SomeGVar vk]; _ -> []+ usedVars (ADuplicate v) = [SomeGVar v]+ usedVars (AUnions vs) = [SomeGVar v | v <- vs]+ usedVars (ASupplyUnion v _) = [SomeGVar v]+ usedVars (ADump v) = [SomeGVar v]++ modelNextState = runModel++ arbitraryWithVars ctx model =+ case findVars ctx (Proxy :: Proxy (LSM RealWorld)) of+ [] ->+ -- Generate a write buffer size and size ratio in the range [3,5] most+ -- of the time, sometimes in the range [1,10] to hit edge cases. 4 was+ -- the hard-coded default for both before it was made configurable.+ fmap Some $ ANew <$> (+ LSMConfig <$> frequency [(10, choose (1,10)), (90, choose (3,5))]+ <*> frequency [(10, choose (2,10)), (90, choose (3,5))]+ )+ vars ->+ let kvars = findVars ctx (Proxy :: Proxy Key)+ existingKey = Left <$> elements kvars+ freshKey = Right <$> arbitrary @Key+ in frequency $+ -- inserts of potentially fresh keys+ [ (3, fmap Some $+ AInsert <$> elements vars+ <*> freshKey+ <*> arbitrary @Value+ <*> arbitrary @(Maybe Blob))+ ]+ -- inserts of the same keys as used earlier+ ++ [ (1, fmap Some $+ AInsert <$> elements vars+ <*> existingKey+ <*> arbitrary @Value+ <*> arbitrary @(Maybe Blob))+ | not (null kvars)+ ]+ -- deletes of arbitrary keys:+ ++ [ (1, fmap Some $+ ADelete <$> elements vars+ <*> freshKey)+ ]+ -- deletes of the same key as inserted earlier:+ ++ [ (1, fmap Some $+ ADelete <$> elements vars+ <*> existingKey)+ | not (null kvars)+ ]+ -- mupserts of potentially fresh keys+ ++ [ (1, fmap Some $+ AMupsert <$> elements vars+ <*> freshKey+ <*> arbitrary @Value)+ ]+ -- mupserts of the same keys as used earlier+ ++ [ (1, fmap Some $+ AMupsert <$> elements vars+ <*> existingKey+ <*> arbitrary @Value)+ | not (null kvars)+ ]+ -- lookup of arbitrary keys:+ ++ [ (1, fmap Some $+ ALookup <$> elements vars+ <*> freshKey)+ ]+ -- lookup of the same key as inserted earlier:+ ++ [ (3, fmap Some $+ ALookup <$> elements vars+ <*> existingKey)+ | not (null kvars)+ ]+ ++ [ (1, fmap Some $+ ADuplicate <$> elements vars)+ ]+ ++ [ (1, fmap Some $ do+ -- bias towards binary, only go high when many tables exist+ len <- elements ([2, 2] ++ take (length vars) [1..5])+ AUnions <$> vectorOf len (elements vars))+ ]+ -- only supply to tables with unions+ ++ [ (2, fmap Some $+ ASupplyUnion <$> elements varsWithUnion+ <*> arbitrary)+ | let hasUnion v = let MLSM m = lookupVar ctx v in+ case Map.lookup m (mlsms model) of+ Nothing -> False+ Just t -> tableHasUnion t+ , let varsWithUnion = filter hasUnion vars+ , not (null varsWithUnion)+ ]+ ++ [ (1, fmap Some $+ ADump <$> elements vars)+ ]++ shrinkWithVars _ctx _model (ANew conf) =+ [ Some $ ANew conf { configMaxWriteBufferSize = mwbs' }+ | mwbs' <- shrink mwbs+ , mwbs' >= 1, mwbs' <= 10+ ]+ ++ [ Some $ ANew conf { configSizeRatio = sr' }+ | sr' <- shrink sr+ , sr' >= 2, sr' <= 10+ ]+ where+ LSMConfig mwbs sr = conf++ shrinkWithVars _ctx _model (AInsert var (Right k) v b) =+ [ Some $ AInsert var (Right k') v' b' | (k', v', b') <- shrink (k, v, b) ]++ shrinkWithVars _ctx _model (AInsert var (Left _kv) v b) =+ [ Some $ AInsert var (Right k') v' b' | (k', v', b') <- shrink (K 100, v, b) ]++ shrinkWithVars _ctx _model (ADelete var (Right k)) =+ [ Some $ ADelete var (Right k') | k' <- shrink k ]++ shrinkWithVars _ctx _model (ADelete var (Left _kv)) =+ [ Some $ ADelete var (Right k) | k <- shrink (K 100) ]++ shrinkWithVars _ctx _model (AMupsert var (Right k) v) =+ [ Some $ AInsert var (Right k) v Nothing ] +++ [ Some $ AMupsert var (Right k') v' | (k', v') <- shrink (k, v) ]++ shrinkWithVars _ctx _model (AMupsert var (Left kv) v) =+ [ Some $ AInsert var (Left kv) v Nothing ] +++ [ Some $ AMupsert var (Right k') v' | (k', v') <- shrink (K 100, v) ]++ shrinkWithVars _ctx _model (AUnions [var]) =+ [ Some $ ADuplicate var ]++ shrinkWithVars _ctx _model (AUnions vars) =+ [ Some $ AUnions vs | vs <- shrinkList (const []) vars, not (null vs) ]++ shrinkWithVars _ctx _model (ASupplyUnion var c) =+ [ Some $ ASupplyUnion var c' | c' <- shrink c ]++ shrinkWithVars _ctx _model _action = []++deriving newtype instance Arbitrary UnionCredits++instance RunLockstep Model (ReaderT (PrimVar RealWorld TableId) IO) where+ observeReal _ action result =+ case (action, result) of+ (ANew{}, _) -> ORef+ (AInsert{}, x) -> OId x+ (ADelete{}, x) -> OId x+ (AMupsert{}, x) -> OId x+ (ALookup{}, x) -> OId x+ (ADuplicate{}, _) -> ORef+ (AUnions{}, _) -> ORef+ (ASupplyUnion{}, x) -> OId x+ (ADump{}, x) -> OId x++ showRealResponse _ ANew{} = Nothing+ showRealResponse _ AInsert{} = Just Dict+ showRealResponse _ ADelete{} = Just Dict+ showRealResponse _ AMupsert{} = Just Dict+ showRealResponse _ ALookup{} = Just Dict+ showRealResponse _ ADuplicate{} = Nothing+ showRealResponse _ AUnions{} = Nothing+ showRealResponse _ ASupplyUnion{} = Nothing+ showRealResponse _ ADump{} = Just Dict++deriving stock instance Show (Action (Lockstep Model) a)+deriving stock instance Show (Observable Model a)+deriving stock instance Show (ModelValue Model a)++deriving stock instance Eq (Action (Lockstep Model) a)+deriving stock instance Eq (Observable Model a)+deriving stock instance Eq (ModelValue Model a)+++runActionIO :: Action (Lockstep Model) a+ -> LookUp+ -> ReaderT (PrimVar RealWorld TableId) IO a+runActionIO action lookUp = ReaderT $ \tidVar -> do+ case action of+ ANew conf -> do+ tid <- incrTidVar tidVar+ stToIO $ newWith tr tid conf+ AInsert var evk v b -> stToIO $ insert tr (lookUpVar var) k v b >> pure k+ where k = either lookUpVar id evk+ ADelete var evk -> stToIO$ delete tr (lookUpVar var) k >> pure ()+ where k = either lookUpVar id evk+ AMupsert var evk v -> stToIO $ mupsert tr (lookUpVar var) k v >> pure k+ where k = either lookUpVar id evk+ ALookup var evk -> stToIO $ lookup tr (lookUpVar var) k+ where k = either lookUpVar id evk+ ADuplicate var -> do+ tid <- incrTidVar tidVar+ stToIO $ duplicate tr tid (lookUpVar var)+ AUnions vars -> do+ tid <- incrTidVar tidVar+ stToIO $ unions tr tid (map lookUpVar vars)+ ASupplyUnion var c -> stToIO $ supplyUnionCredits (lookUpVar var) (getNonNegative c) >> pure ()+ ADump var -> stToIO $ logicalValue (lookUpVar var)+ where+ lookUpVar :: ModelVar Model a -> a+ lookUpVar = realLookupVar lookUp++ tr :: Tracer (ST RealWorld) Event+ tr = nullTracer++ incrTidVar :: PrimVar RealWorld TableId -> IO TableId+ incrTidVar tidVar = do+ tid@(TableId x) <- readPrimVar tidVar+ writePrimVar tidVar (TableId (x + 1))+ pure tid++runModel :: Action (Lockstep Model) a+ -> ModelVarContext Model+ -> Model+ -> (ModelValue Model a, Model)+runModel action ctx m =+ case action of+ ANew{} -> (MLSM mlsm, m')+ where (mlsm, m') = modelNew m++ AInsert var evk v b -> (MInsert k, m')+ where ((), m') = modelInsert (lookUpLsMVar var) k v b m+ k = either lookUpKeyVar id evk++ ADelete var evk -> (MUnit (), m')+ where ((), m') = modelDelete (lookUpLsMVar var) k m+ k = either lookUpKeyVar id evk++ AMupsert var evk v -> (MInsert k, m')+ where ((), m') = modelMupsert (lookUpLsMVar var) k v m+ k = either lookUpKeyVar id evk++ ALookup var evk -> (MLookup mv, m')+ where (mv, m') = modelLookup (lookUpLsMVar var) k m+ k = either lookUpKeyVar id evk++ ADuplicate var -> (MLSM mlsm', m')+ where (mlsm', m') = modelDuplicate (lookUpLsMVar var) m++ AUnions vars -> (MLSM mlsm', m')+ where (mlsm', m') = modelUnions (map lookUpLsMVar vars) m++ ASupplyUnion var c -> (MUnit (), m')+ where ((), m') = modelSupplyUnion (lookUpLsMVar var) c m++ ADump var -> (MDump mapping, m)+ where (mapping, _) = modelDump (lookUpLsMVar var) m+ where+ lookUpLsMVar :: ModelVar Model (LSM RealWorld) -> ModelLSM+ lookUpLsMVar var = case lookupVar ctx var of MLSM r -> r++ lookUpKeyVar :: ModelVar Model Key -> Key+ lookUpKeyVar var = case lookupVar ctx var of MInsert k -> k
@@ -0,0 +1,299 @@+{-# LANGUAGE TypeFamilies #-}++-- | An abstraction of the LSM API, instantiated by both the real+-- implementation and a model (see "Database.LSMTree.Model.IO").+module Database.LSMTree.Class (+ IsTable (..)+ , withTableNew+ , withTableFromSnapshot+ , withTableDuplicate+ , withTableUnion+ , withTableUnions+ , withCursor+ , module Common+ , module Types+ ) where++import Control.Monad (void)+import Control.Monad.Class.MonadThrow (MonadThrow (..))+import Data.Kind (Constraint, Type)+import Data.List.NonEmpty (NonEmpty)+import Data.Typeable (Proxy (..))+import qualified Data.Vector as V+import Database.LSMTree as Types (Entry (..), LookupResult (..),+ ResolveAsFirst (..), ResolveValue (..), Update (..))+import qualified Database.LSMTree as R+import Database.LSMTree.Class.Common as Common+import qualified Database.LSMTree.Internal.Paths as RIP+import qualified Database.LSMTree.Internal.Types as RT (Table (..))+import qualified Database.LSMTree.Internal.Unsafe as RU (SessionEnv (..),+ Table (..), withKeepSessionOpen)+import Test.Util.FS (flipRandomBitInRandomFileHardlinkSafe)+import Test.Util.QC (Choice)++-- | Class abstracting over table operations.+--+type IsTable :: ((Type -> Type) -> Type -> Type -> Type -> Type) -> Constraint+class (IsSession (Session h)) => IsTable h where+ type Session h :: (Type -> Type) -> Type+ type TableConfig h :: Type+ type BlobRef h :: (Type -> Type) -> Type -> Type+ type Cursor h :: (Type -> Type) -> Type -> Type -> Type -> Type++ newTableWith ::+ ( IOLike m+ , C k v b+ )+ => TableConfig h+ -> Session h m+ -> m (h m k v b)++ closeTable ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> m ()++ lookups ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> V.Vector k+ -> m (V.Vector (LookupResult v (BlobRef h m b)))++ rangeLookup ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> Range k+ -> m (V.Vector (Entry k v (BlobRef h m b)))++ newCursor ::+ ( IOLike m+ , C k v b+ )+ => Maybe k+ -> h m k v b+ -> m (Cursor h m k v b)++ closeCursor ::+ ( IOLike m+ , C k v b+ )+ => proxy h+ -> Cursor h m k v b+ -> m ()++ readCursor ::+ ( IOLike m+ , C k v b+ )+ => proxy h+ -> Int+ -> Cursor h m k v b+ -> m (V.Vector (Entry k v (BlobRef h m b)))++ retrieveBlobs ::+ ( IOLike m+ , CB b+ )+ => proxy h+ -> Session h m+ -> V.Vector (BlobRef h m b)+ -> m (V.Vector b)++ updates ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> V.Vector (k, Update v b)+ -> m ()++ inserts ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> V.Vector (k, v, Maybe b)+ -> m ()++ deletes ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> V.Vector k+ -> m ()++ upserts ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> V.Vector (k, v)+ -> m ()++ saveSnapshot ::+ ( IOLike m+ , C k v b+ )+ => SnapshotName+ -> SnapshotLabel+ -> h m k v b+ -> m ()++ corruptSnapshot ::+ ( IOLike m+ )+ => Choice+ -> SnapshotName+ -> h m k v b+ -> m ()++ openTableFromSnapshot ::+ ( IOLike m+ , C k v b+ )+ => Session h m+ -> SnapshotName+ -> SnapshotLabel+ -> m (h m k v b)++ duplicate ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> m (h m k v b)++ union ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> h m k v b+ -> m (h m k v b)++ unions ::+ ( IOLike m+ , C k v b+ )+ => NonEmpty (h m k v b)+ -> m (h m k v b)++ remainingUnionDebt ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> m UnionDebt++ supplyUnionCredits ::+ ( IOLike m+ , C k v b+ )+ => h m k v b+ -> UnionCredits+ -> m UnionCredits++withTableNew :: forall h m k v b a.+ (IOLike m, IsTable h, C k v b)+ => Session h m+ -> TableConfig h+ -> (h m k v b -> m a)+ -> m a+withTableNew sesh conf = bracket (newTableWith conf sesh) closeTable++withTableFromSnapshot :: forall h m k v b a.+ (IOLike m, IsTable h, C k v b)+ => Session h m+ -> SnapshotLabel+ -> SnapshotName+ -> (h m k v b -> m a)+ -> m a+withTableFromSnapshot sesh label snap = bracket (openTableFromSnapshot sesh snap label) closeTable++withTableDuplicate :: forall h m k v b a.+ (IOLike m, IsTable h, C k v b)+ => h m k v b+ -> (h m k v b -> m a)+ -> m a+withTableDuplicate table = bracket (duplicate table) closeTable++withTableUnion :: forall h m k v b a.+ (IOLike m, IsTable h, C k v b)+ => h m k v b+ -> h m k v b+ -> (h m k v b -> m a)+ -> m a+withTableUnion table1 table2 = bracket (table1 `union` table2) closeTable++withTableUnions :: forall h m k v b a.+ (IOLike m, IsTable h, C k v b)+ => NonEmpty (h m k v b)+ -> (h m k v b -> m a)+ -> m a+withTableUnions tables = bracket (unions tables) closeTable++withCursor :: forall h m k v b a.+ (IOLike m, IsTable h, C k v b)+ => Maybe k+ -> h m k v b+ -> (Cursor h m k v b -> m a)+ -> m a+withCursor offset tbl = bracket (newCursor offset tbl) (closeCursor (Proxy @h))++{-------------------------------------------------------------------------------+ Real instance+-------------------------------------------------------------------------------}++-- | Snapshot corruption for the real instance.+-- Implemented here, instead of as part of the public API.+rCorruptSnapshot ::+ IOLike m+ => Choice+ -> SnapshotName+ -> R.Table m k v b+ -> m ()+rCorruptSnapshot choice name (RT.Table t) =+ RU.withKeepSessionOpen (RU.tableSession t) $ \seshEnv ->+ let hfs = RU.sessionHasFS seshEnv+ root = RU.sessionRoot seshEnv+ namedSnapDir = RIP.getNamedSnapshotDir (RIP.namedSnapshotDir root name)+ in void $ flipRandomBitInRandomFileHardlinkSafe hfs choice namedSnapDir++instance IsTable R.Table where+ type Session R.Table = R.Session+ type TableConfig R.Table = R.TableConfig+ type BlobRef R.Table = R.BlobRef+ type Cursor R.Table = R.Cursor++ newTableWith = R.newTableWith+ closeTable = R.closeTable+ lookups = R.lookups+ updates = R.updates+ inserts = R.inserts+ deletes = R.deletes+ upserts = R.upserts++ rangeLookup = R.rangeLookup+ retrieveBlobs _ = R.retrieveBlobs++ newCursor = maybe R.newCursor (flip R.newCursorAtOffset)+ closeCursor _ = R.closeCursor+ readCursor _ = R.take++ saveSnapshot = R.saveSnapshot+ corruptSnapshot = rCorruptSnapshot+ openTableFromSnapshot = R.openTableFromSnapshot++ duplicate = R.duplicate++ union = R.incrementalUnion+ unions = R.incrementalUnions+ remainingUnionDebt = R.remainingUnionDebt+ supplyUnionCredits = R.supplyUnionCredits
@@ -0,0 +1,94 @@+{-# LANGUAGE TypeFamilies #-}++module Database.LSMTree.Class.Common (+ C, CK, CV, CB, C_+ , IsSession (..)+ , SessionArgs (..)+ , withSession+ , module Types+ ) where++import Control.Monad.Class.MonadThrow (MonadThrow (..))+import Control.Tracer (nullTracer)+import Data.Kind (Constraint, Type)+import Data.Typeable (Typeable)+import Database.LSMTree (ResolveValue)+import Database.LSMTree as Types (IOLike, Range (..), SerialiseKey,+ SerialiseValue, SnapshotLabel (..), SnapshotName,+ UnionCredits (..), UnionDebt (..))+import qualified Database.LSMTree as R+import System.FS.API (FsPath, HasFS)+import System.FS.BlockIO.API (HasBlockIO)++{-------------------------------------------------------------------------------+ Constraints+-------------------------------------------------------------------------------}++-- | Constraints for keys, values, and blobs+type C k v b = (CK k, CV v, CB b)++-- | Constraints for keys+type CK k = (C_ k, SerialiseKey k)++-- | Constraints for values+type CV v = (C_ v, SerialiseValue v, ResolveValue v)++-- | Constraints for blobs+type CB b = (C_ b, SerialiseValue b)++-- | Model-specific constraints for keys, values, and blobs+type C_ a = (Show a, Eq a, Typeable a)++{-------------------------------------------------------------------------------+ Session+-------------------------------------------------------------------------------}++-- | Class abstracting over session operations.+--+type IsSession :: ((Type -> Type) -> Type) -> Constraint+class IsSession s where+ data SessionArgs s :: (Type -> Type) -> Type++ openSession ::+ IOLike m+ => SessionArgs s m+ -> m (s m)++ closeSession ::+ IOLike m+ => s m+ -> m ()++ deleteSnapshot ::+ IOLike m+ => s m+ -> SnapshotName+ -> m ()++ listSnapshots ::+ IOLike m+ => s m+ -> m [SnapshotName]++withSession :: (IOLike m, IsSession s) => SessionArgs s m -> (s m -> m a) -> m a+withSession seshArgs = bracket (openSession seshArgs) closeSession++{-------------------------------------------------------------------------------+ Real instance+-------------------------------------------------------------------------------}++testSalt :: R.Salt+testSalt = 4++instance IsSession R.Session where+ data SessionArgs R.Session m where+ SessionArgs ::+ forall m h. Typeable h+ => HasFS m h -> HasBlockIO m h -> FsPath+ -> SessionArgs R.Session m++ openSession (SessionArgs hfs hbio dir) = do+ R.openSession nullTracer hfs hbio testSalt dir+ closeSession = R.closeSession+ deleteSnapshot = R.deleteSnapshot+ listSnapshots = R.listSnapshots
@@ -0,0 +1,21 @@+-- | The are three kinds of model, each depending on the previous one:+--+-- * [@Model.Table@]:+-- Pure model of a single table.+-- @+-- updates :: [_] -> Table k v b -> Table k v b+-- @+--+-- * [@Model.Session@]:+-- Pure model of a session (containing multiple tables).+-- @+-- updates :: MonadState Model m => [_] -> Table k v b -> m ()+-- @+--+-- * [@Model.IO@]:+-- STM-based model allowing multiple (potentially closed) sessions.+-- @+-- updates :: MonadSTM m => MSession m -> [_] -> Table k v b -> m ()+-- @+--+module Database.LSMTree.Model () where
@@ -0,0 +1,100 @@+{-# LANGUAGE TypeFamilies #-}++-- | An instance of `Class.IsTable`, modelling potentially closed sessions in+-- @IO@ by lifting the pure session model from "Database.LSMTree.Model.Session".+module Database.LSMTree.Model.IO (+ Err (..)+ , Session (..)+ , Class.SessionArgs (NoSessionArgs)+ , Table (..)+ , TableConfig (..)+ , BlobRef (..)+ , Cursor (..)+ -- * helpers+ , runInOpenSession+ ) where++import Control.Concurrent.Class.MonadSTM.Strict+import Control.Exception (Exception)+import Control.Monad.Class.MonadThrow (MonadThrow (..))+import qualified Data.List.NonEmpty as NE+import qualified Database.LSMTree.Class as Class+import Database.LSMTree.Model.Session (TableConfig (..))+import qualified Database.LSMTree.Model.Session as Model++newtype Session m = Session (StrictTVar m (Maybe Model.Model))++data Table m k v b = Table {+ thSession :: !(Session m)+ , thTable :: !(Model.Table k v b)+ }++data BlobRef m b = BlobRef {+ brSession :: !(Session m)+ , brBlobRef :: !(Model.BlobRef b)+ }++data Cursor m k v b = Cursor {+ cSession :: !(Session m)+ , cCursor :: !(Model.Cursor k v b)+ }++newtype Err = Err (Model.Err)+ deriving stock Show+ deriving anyclass Exception++runInOpenSession :: (MonadSTM m, MonadThrow (STM m)) => Session m -> Model.ModelM a -> m a+runInOpenSession (Session var) action = atomically $ do+ readTVar var >>= \case+ Nothing -> error "session closed"+ Just m -> do+ let (r, m') = Model.runModelM action m+ case r of+ Left e -> throwSTM (Err e)+ Right x -> writeTVar var (Just m') >> pure x++instance Class.IsSession Session where+ data SessionArgs Session m = NoSessionArgs+ openSession NoSessionArgs = Session <$> newTVarIO (Just $! Model.initModel)+ closeSession (Session var) = atomically $ writeTVar var Nothing+ deleteSnapshot s x = runInOpenSession s $ Model.deleteSnapshot x+ listSnapshots s = runInOpenSession s $ Model.listSnapshots++instance Class.IsTable Table where+ type Session Table = Session+ type TableConfig Table = Model.TableConfig+ type BlobRef Table = BlobRef+ type Cursor Table = Cursor++ newTableWith x s = Table s <$> runInOpenSession s (Model.new x)+ closeTable (Table s t) = runInOpenSession s (Model.close t)+ lookups (Table s t) x1 = fmap (fmap (BlobRef s)) <$>+ runInOpenSession s (Model.lookups x1 t)+ updates (Table s t) x1 = runInOpenSession s (Model.updates Model.getResolve x1 t)+ inserts (Table s t) x1 = runInOpenSession s (Model.inserts Model.getResolve x1 t)+ deletes (Table s t) x1 = runInOpenSession s (Model.deletes Model.getResolve x1 t)+ upserts (Table s t) x1 = runInOpenSession s (Model.upserts Model.getResolve x1 t)++ rangeLookup (Table s t) x1 = fmap (fmap (BlobRef s)) <$>+ runInOpenSession s (Model.rangeLookup x1 t)+ retrieveBlobs _ s x1 = runInOpenSession s (Model.retrieveBlobs (fmap brBlobRef x1))++ newCursor k (Table s t) = Cursor s <$> runInOpenSession s (Model.newCursor k t)+ closeCursor _ (Cursor s c) = runInOpenSession s (Model.closeCursor c)+ readCursor _ x1 (Cursor s c) = fmap (fmap (BlobRef s)) <$>+ runInOpenSession s (Model.readCursor x1 c)++ saveSnapshot x1 x2 (Table s t) = runInOpenSession s (Model.saveSnapshot x1 x2 t)+ corruptSnapshot _ x (Table s _t) = runInOpenSession s (Model.corruptSnapshot x)+ openTableFromSnapshot s x1 x2 = Table s <$> runInOpenSession s (Model.openTableFromSnapshot x1 x2)++ duplicate (Table s t) = Table s <$> runInOpenSession s (Model.duplicate t)++ union (Table s1 t1) (Table _s2 t2) =+ Table s1 <$> runInOpenSession s1 (Model.union Model.getResolve t1 t2)+ unions ts =+ Table s <$> runInOpenSession s (Model.unions Model.getResolve (fmap thTable ts))+ where+ Table s _ = NE.head ts+ remainingUnionDebt (Table s t) = runInOpenSession s (Model.remainingUnionDebt t)+ supplyUnionCredits (Table s t) credits = runInOpenSession s (Model.supplyUnionCredits t credits)
@@ -0,0 +1,875 @@+{-# LANGUAGE PatternSynonyms #-}+-- | A pure model of a single session containing multiple tables.+--+-- The session model builds on top of "Database.LSMTree.Model.Table", adding+-- table and snapshot administration.+module Database.LSMTree.Model.Session (+ -- * Model+ Model (..)+ , initModel+ , UpdateCounter (..)+ -- ** SomeTable, for testing+ , SomeTable (..)+ , toSomeTable+ , fromSomeTable+ , withSomeTable+ , TableID+ , tableID+ , isUnionDescendant+ , Model.size+ -- ** Constraints+ , C+ , C_+ , Model.SerialiseKey (..)+ , Model.SerialiseValue (..)+ -- ** ModelT and ModelM+ , ModelT (..)+ , runModelT+ , ModelM+ , runModelM+ , runModelMWithInjectedErrors+ -- ** Errors+ , Err (..)+ , isDiskFault+ , isSnapshotCorrupted+ , isOther+ -- * Tables+ , Table+ , TableConfig (..)+ , new+ , close+ -- * Monoidal value resolution+ , ResolveSerialisedValue (..)+ , getResolve+ , noResolve+ -- * Table querying and updates+ -- ** Queries+ , Range (..)+ , LookupResult (..)+ , lookups+ , Entry (..)+ , rangeLookup+ -- ** Cursor+ , Cursor+ , CursorID+ , cursorID+ , newCursor+ , closeCursor+ , readCursor+ -- ** Updates+ , Update (..)+ , updates+ , inserts+ , deletes+ , upserts+ -- ** Blobs+ , BlobRef+ , retrieveBlobs+ , invalidateBlobRefs+ -- * Snapshots+ , SnapshotName+ , saveSnapshot+ , openTableFromSnapshot+ , corruptSnapshot+ , deleteSnapshot+ , listSnapshots+ -- * Multiple writable tables+ , duplicate+ -- * Table union+ , IsUnionDescendant (..)+ , union+ , unions+ , UnionDebt (..)+ , remainingUnionDebt+ , UnionCredits (..)+ , supplyUnionCredits+ , supplyPortionOfDebt+ ) where++import Control.Monad (forM, when)+import Control.Monad.Except (ExceptT (..), MonadError (..),+ runExceptT)+import Control.Monad.Identity (Identity (runIdentity))+import Control.Monad.State.Strict (MonadState (..), StateT (..), gets,+ modify)+import Data.Data+import Data.Dynamic+import qualified Data.List as L+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromJust)+import qualified Data.Vector as V+import Data.Word+import Database.LSMTree (SerialiseKey (..), SerialiseValue (..),+ SnapshotLabel (..), SnapshotName, UnionCredits (..),+ UnionDebt (..))+import Database.LSMTree.Model.Table (Entry (..), LookupResult (..),+ Range (..), ResolveSerialisedValue (..), Update (..),+ getResolve, noResolve)+import qualified Database.LSMTree.Model.Table as Model++{-------------------------------------------------------------------------------+ Model+-------------------------------------------------------------------------------}++data Model = Model {+ tables :: Map TableID (UpdateCounter, SomeTable)+ , cursors :: Map CursorID SomeCursor+ , nextID :: Int+ , snapshots :: Map SnapshotName Snapshot+ }+ deriving stock Show++initModel :: Model+initModel = Model {+ tables = Map.empty+ , cursors = Map.empty+ , nextID = 0+ , snapshots = Map.empty+ }++-- | We conservatively model blob reference invalidation: each update after+-- acquiring a blob reference will invalidate it. We use 'UpdateCounter' to+-- track updates.+--+-- Supplying union credits is also considered an update, though this can only+-- invalidate a blob reference that is associated with a (descendant of a) union+-- table.+newtype UpdateCounter = UpdateCounter Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (Num)++data SomeTable where+ SomeTable :: (Typeable k, Typeable v, Typeable b)+ => Model.Table k v b -> SomeTable++instance Show SomeTable where+ show (SomeTable table) = show table++toSomeTable ::+ (Typeable k, Typeable v, Typeable b)+ => Model.Table k v b+ -> SomeTable+toSomeTable = SomeTable++fromSomeTable ::+ (Typeable k, Typeable v, Typeable b)+ => SomeTable+ -> Maybe (Model.Table k v b)+fromSomeTable (SomeTable tbl) = cast tbl++withSomeTable ::+ (forall k v b. (Typeable k, Typeable v, Typeable b)+ => Model.Table k v b -> a)+ -> SomeTable+ -> a+withSomeTable f (SomeTable tbl) = f tbl++newtype SomeCursor = SomeCursor Dynamic++instance Show SomeCursor where+ show (SomeCursor c) = show c++toSomeCursor ::+ (Typeable k, Typeable v, Typeable b)+ => Model.Cursor k v b+ -> SomeCursor+toSomeCursor = SomeCursor . toDyn++fromSomeCursor ::+ (Typeable k, Typeable v, Typeable b)+ => SomeCursor+ -> Maybe (Model.Cursor k v b)+fromSomeCursor (SomeCursor c) = fromDynamic c++--+-- Constraints+--++-- | Common constraints for keys, values and blobs+type C_ a = (Show a, Eq a, Typeable a)+type C k v b = (C_ k, C_ v, C_ b)++--+-- ModelT and ModelM+--++newtype ModelT m a = ModelT { _runModelT :: ExceptT Err (StateT Model m) a }+ deriving stock Functor+ deriving newtype ( Applicative+ , Monad+ , MonadState Model+ , MonadError Err+ )++runModelT :: ModelT m a -> Model -> m (Either Err a, Model)+runModelT = runStateT . runExceptT . _runModelT++type ModelM = ModelT Identity++runModelM :: ModelM a -> Model -> (Either Err a, Model)+runModelM m = runIdentity . runModelT m++-- | @'runModelMWithInjectedErrors' merrs onNoErrors onErrors@ runs the model+-- with injected disk faults @merrs@.+--+-- The model may run different actions based on whether there are injections or+-- not. 'onNoErrors' runs in case @merrs == Nothing@, and 'onErrors@ runs in+-- case @merrs == Just _@.+--+-- Typically, @onNoErrors@ will be a model function like 'lookups', and+-- @onErrors@ should be the identity state transition. This models the approach+-- of the @lsm-tree@ library: the *logical* state of the database should remain+-- unchanged if there were any disk faults.+--+-- The model functions in the remainder of the module only describe the happy+-- path with respect to disk faults: they assume there aren't any such faults.+-- The intent of 'runModelMWithInjectedErrors' is to augment the model with+-- responses to disk fault injection.+--+-- The real system's is not guaranteed to fail with an error if there are+-- injected disk faults. However, the approach taken in+-- 'runModelMWithInjectedErrors' is to *always* throw a modelled disk error in+-- case there are injections, *even if* the real system happened to not fail+-- completely. This makes 'runModelMWithInjectedErrors' an approximation of the+-- real system's behaviour in case of injections. For the model to more+-- accurately model the system's behaviour, the model would have to know for+-- each API function precisely which errors are injected in which order, and+-- whether they are handled internally by the library or not. This is simply+-- infeasible: the model would become very complex.+runModelMWithInjectedErrors ::+ Maybe e -- ^ The errors that are injected+ -> ModelM a -- ^ Action to run on 'Nothing'+ -> ModelM () -- ^ Action to run on 'Just'+ -> Model -> (Either Err a, Model)+runModelMWithInjectedErrors Nothing onNoErrors _ st =+ runModelM onNoErrors st+runModelMWithInjectedErrors (Just _) _ onErrors st =+ runModelM (onErrors >> throwError ModelErrDiskFault) st++-- | The default 'ErrDiskFault' that model operations will throw.+pattern ModelErrDiskFault :: Err+pattern ModelErrDiskFault = ErrDiskFault "model does not produce an error message"++-----+-- Errors+--++data Err+ = ErrSessionDirDoesNotExist+ | ErrSessionDirLocked+ | ErrSessionDirCorrupted+ | ErrSessionClosed+ | ErrTableClosed+ | ErrTableCorrupted+ | ErrTableUnionHandleTypeMismatch+ | ErrTableUnionSessionMismatch+ | ErrSnapshotExists !SnapshotName+ | ErrSnapshotDoesNotExist !SnapshotName+ | ErrSnapshotCorrupted !SnapshotName+ | ErrSnapshotWrongLabel !SnapshotName+ | ErrBlobRefInvalid+ | ErrCursorClosed+ | ErrDiskFault !String+ | ErrFsError !String+ | ErrCommitActionRegistry !(NonEmpty Err)+ | ErrAbortActionRegistry !(Maybe Err) !(NonEmpty Err)+ | ErrOther !String+ deriving stock (Eq, Show)++isSnapshotCorrupted :: Err -> Bool+isSnapshotCorrupted (ErrSnapshotCorrupted _) = True+isSnapshotCorrupted (ErrFsError _) = True+isSnapshotCorrupted (ErrCommitActionRegistry es) = firstNotOtherSatisfies isSnapshotCorrupted es+isSnapshotCorrupted (ErrAbortActionRegistry e es) = firstNotOtherSatisfies isSnapshotCorrupted (maybe id NE.cons e es)+isSnapshotCorrupted _ = False++isDiskFault :: Err -> Bool+isDiskFault (ErrDiskFault _) = True+isDiskFault (ErrFsError _) = True+isDiskFault (ErrCommitActionRegistry es) = firstNotOtherSatisfies isDiskFault es+isDiskFault (ErrAbortActionRegistry e es) = firstNotOtherSatisfies isDiskFault (maybe id NE.cons e es)+isDiskFault _ = False++isOther :: Err -> Bool+isOther (ErrOther _) = True+isOther (ErrCommitActionRegistry es) = all isOther es+isOther (ErrAbortActionRegistry e es) = all isOther (maybe id NE.cons e es)+isOther _ = False++firstNotOtherSatisfies :: (Err -> Bool) -> NonEmpty Err -> Bool+firstNotOtherSatisfies p =+ maybe False (p . fst) . L.uncons . NE.dropWhile isOther++{-------------------------------------------------------------------------------+ Tables+-------------------------------------------------------------------------------}++type TableID = Int++--+-- API+--++data Table k v b = Table {+ tableID :: TableID+ , config :: TableConfig+ , isUnionDescendant :: IsUnionDescendant+ }+ deriving stock Show++-- | The model does not distinguish multiple values of the config type. This is+-- the right thing because the table config is never observed directly, only+-- via the influence it might have on other operation results. And we /do not+-- want/ table configs to have observable value-level effects. The table config+-- should only have performance effects. Thus having only one config value is+-- the right abstraction for the model.+--+data TableConfig = TableConfig+ deriving stock (Show, Eq)++new ::+ forall k v b m. (MonadState Model m, C k v b)+ => TableConfig+ -> m (Table k v b)+new config = newTableWith config IsNotUnionDescendant Model.empty++-- |+--+-- This is idempotent.+close :: MonadState Model m => Table k v b -> m ()+close Table{..} = state $ \Model{..} ->+ let tables' = Map.delete tableID tables+ model' = Model {+ tables = tables'+ , ..+ }+ in ((), model')++--+-- Utility+--++guardTableIsOpen ::+ forall k v b m. (+ MonadState Model m, MonadError Err m+ , Typeable k, Typeable v, Typeable b+ )+ => Table k v b+ -> m (UpdateCounter, Model.Table k v b)+guardTableIsOpen Table{..} =+ gets (Map.lookup tableID . tables) >>= \case+ Nothing ->+ throwError ErrTableClosed+ Just (updc, tbl) ->+ pure (updc, fromJust $ fromSomeTable tbl)++newTableWith ::+ (MonadState Model m, C k v b)+ => TableConfig+ -> IsUnionDescendant+ -> Model.Table k v b+ -> m (Table k v b)+newTableWith config isUnionDescendant tbl = state $ \Model{..} ->+ let table = Table {+ tableID = nextID+ , config+ , isUnionDescendant+ }+ someTable = toSomeTable tbl+ tables' = Map.insert nextID (0, someTable) tables+ nextID' = nextID + 1+ model' = Model {+ tables = tables'+ , nextID = nextID'+ , ..+ }+ in (table, model')++{-------------------------------------------------------------------------------+ Lookups+-------------------------------------------------------------------------------}++lookups ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , C k v b+ )+ => V.Vector k+ -> Table k v b+ -> m (V.Vector (Model.LookupResult v (BlobRef b)))+lookups ks t = do+ (updc, table) <- guardTableIsOpen t+ pure $ liftBlobRefs (SomeTableID updc (tableID t)) $ Model.lookups ks table++rangeLookup ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , C k v b+ )+ => Range k+ -> Table k v b+ -> m (V.Vector (Model.Entry k v (BlobRef b)))+rangeLookup r t = do+ (updc, table) <- guardTableIsOpen t+ pure $ liftBlobRefs (SomeTableID updc (tableID t)) $ Model.rangeLookup r table++{-------------------------------------------------------------------------------+ Updates+-------------------------------------------------------------------------------}++updates ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , SerialiseValue b+ , C k v b+ )+ => ResolveSerialisedValue v+ -> V.Vector (k, Model.Update v b)+ -> Table k v b+ -> m ()+updates r ups t@Table{..} = do+ (updc, table) <- guardTableIsOpen t+ let table' = Model.updates r ups table+ modify (\m -> m {+ tables = Map.insert tableID (updc + 1, toSomeTable table') (tables m)+ })++inserts ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , SerialiseValue b+ , C k v b+ )+ => ResolveSerialisedValue v+ -> V.Vector (k, v, Maybe b)+ -> Table k v b+ -> m ()+inserts r = updates r . fmap (\(k, v, blob) -> (k, Model.Insert v blob))++deletes ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , SerialiseValue b+ , C k v b+ )+ => ResolveSerialisedValue v+ -> V.Vector k+ -> Table k v b+ -> m ()+deletes r = updates r . fmap (,Model.Delete)++upserts ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , SerialiseValue b+ , C k v b+ )+ => ResolveSerialisedValue v+ -> V.Vector (k, v)+ -> Table k v b+ -> m ()+upserts r = updates r . fmap (fmap Model.Upsert)++{-------------------------------------------------------------------------------+ Blobs+-------------------------------------------------------------------------------}++-- | For more details: 'Database.LSMTree.Internal.BlobRef' describes the+-- intended semantics of blob references.+data BlobRef b = BlobRef {+ handleRef :: !(SomeHandleID b)+ , innerBlob :: !(Model.BlobRef b)+ }++deriving stock instance Show b => Show (BlobRef b)++retrieveBlobs ::+ forall m b. ( MonadState Model m+ , MonadError Err m+ , SerialiseValue b+ )+ => V.Vector (BlobRef b)+ -> m (V.Vector b)+retrieveBlobs refs = Model.retrieveBlobs <$> V.mapM guard refs+ where+ guard BlobRef{..} = do+ m <- get+ -- In the real implementation, a blob reference /could/ be invalidated+ -- every time you modify a table or cursor. This model takes the most+ -- conservative approach: a blob reference is immediately invalidated+ -- every time a table/cursor is modified.+ case handleRef of+ SomeTableID createdAt tableID ->+ case Map.lookup tableID (tables m) of+ -- If the table is now closed, it means the table has been modified.+ Nothing -> errInvalid+ -- If the table is open, we check timestamps (i.e., UpdateCounter)+ -- to see if any modifications were made.+ Just (updc, _) -> do+ when (updc /= createdAt) $ errInvalid+ pure innerBlob+ SomeCursorID cursorID ->+ -- The only modification to a cursor is that it can be closed.+ case Map.lookup cursorID (cursors m) of+ Nothing -> errInvalid+ Just _ -> pure innerBlob++ errInvalid :: m a+ errInvalid = throwError ErrBlobRefInvalid++data SomeHandleID b where+ SomeTableID :: !UpdateCounter -> !TableID -> SomeHandleID b+ SomeCursorID :: !CursorID -> SomeHandleID b+ deriving stock Show++liftBlobRefs ::+ (Functor f, Functor g)+ => SomeHandleID b+ -> g (f (Model.BlobRef b))+ -> g (f (BlobRef b))+liftBlobRefs hid = fmap (fmap (BlobRef hid))++-- | Invalidate blob references that were created from the given table. This+-- function assumes that it is called on an open table.+--+-- This is useful in tests where blob references should be invalidated for other+-- reasons than normal operation of the model.+invalidateBlobRefs ::+ MonadState Model m+ => Table k v b+ -> m ()+invalidateBlobRefs Table{..} = do+ gets (Map.lookup tableID . tables) >>= \case+ Nothing -> error "invalidateBlobRefs: table is closed!"+ Just (updc, tbl) -> do+ modify (\m -> m {+ tables = Map.insert tableID (updc + 1, tbl) (tables m)+ })++{-------------------------------------------------------------------------------+ Snapshots+-------------------------------------------------------------------------------}++data Snapshot = Snapshot+ { snapshotConfig :: TableConfig+ , snapshotLabel :: SnapshotLabel+ , snapshotTable :: SomeTable+ , snapshotIsUnionDescendant :: IsUnionDescendant+ , snapshotCorrupted :: Bool+ }+ deriving stock Show++saveSnapshot ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => SnapshotName+ -> SnapshotLabel+ -> Table k v b+ -> m ()+saveSnapshot name label t@Table{..} = do+ (_updc, table) <- guardTableIsOpen t+ snaps <- gets snapshots+ when (Map.member name snaps) $+ throwError $ ErrSnapshotExists name+ let snap =+ Snapshot+ config label+ (toSomeTable $ Model.snapshot table)+ isUnionDescendant False+ modify (\m -> m {+ snapshots = Map.insert name snap (snapshots m)+ })++openTableFromSnapshot ::+ forall k v b m.(+ MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => SnapshotName+ -> SnapshotLabel+ -> m (Table k v b)+openTableFromSnapshot name label = do+ snaps <- gets snapshots+ case Map.lookup name snaps of+ Nothing ->+ throwError $ ErrSnapshotDoesNotExist name+ Just (Snapshot conf label' tbl snapshotIsUnion corrupted) -> do+ when corrupted $+ throwError $ ErrSnapshotCorrupted name+ when (label /= label') $+ throwError $ ErrSnapshotWrongLabel name+ case fromSomeTable tbl of+ Nothing ->+ -- The label should contain enough information to type snapshots,+ -- but it's up to the user to pick labels. If the labels are picked+ -- badly, then the SUT would fail with any number of and type of+ -- errors. The model simply does not allow this to occur: if we fail+ -- to cast modelled tables, then we consider it to be a bug in the+ -- test setup, and so we use @error@ instead of @throwError@.+ error "openTableFromSnapshot: snapshot opened at wrong type"+ Just table' ->+ newTableWith conf snapshotIsUnion table'++-- To match the implementation of the real table, this should not corrupt the+-- snapshot if there are _no non-empty files_; however, since there are no such+-- snapshots, this is probably fine.+corruptSnapshot ::+ (MonadState Model m, MonadError Err m)+ => SnapshotName+ -> m ()+corruptSnapshot name = do+ snapshots <- gets snapshots+ if Map.notMember name snapshots+ then throwError $ ErrSnapshotDoesNotExist name+ else modify $ \m -> m {snapshots = Map.adjust corruptSnapshotEntry name snapshots}+ where+ corruptSnapshotEntry (Snapshot c l t u _) = Snapshot c l t u True++deleteSnapshot ::+ (MonadState Model m, MonadError Err m)+ => SnapshotName+ -> m ()+deleteSnapshot name = do+ snaps <- gets snapshots+ case Map.lookup name snaps of+ Nothing ->+ throwError $ ErrSnapshotDoesNotExist name+ Just _ ->+ modify (\m -> m {+ snapshots = Map.delete name snaps+ })++listSnapshots ::+ MonadState Model m+ => m [SnapshotName]+listSnapshots = gets (Map.keys . snapshots)++{-------------------------------------------------------------------------------+ Multiple writable tables+-------------------------------------------------------------------------------}++duplicate ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => Table k v b+ -> m (Table k v b)+duplicate t@Table{..} = do+ table <- snd <$> guardTableIsOpen t+ newTableWith config isUnionDescendant $ Model.duplicate table++{-------------------------------------------------------------------------------+ Cursor+-------------------------------------------------------------------------------}++type CursorID = Int++data Cursor k v b = Cursor {+ cursorID :: !CursorID+ }+ deriving stock Show++newCursor ::+ forall k v b m. (+ MonadState Model m, MonadError Err m+ , SerialiseKey k+ , C k v b+ )+ => Maybe k+ -> Table k v b+ -> m (Cursor k v b)+newCursor offset t = do+ table <- snd <$> guardTableIsOpen t+ state $ \Model{..} ->+ let cursor = Cursor { cursorID = nextID }+ someCursor = toSomeCursor $ Model.newCursor offset table+ cursors' = Map.insert nextID someCursor cursors+ nextID' = nextID + 1+ model' = Model {+ cursors = cursors'+ , nextID = nextID'+ , ..+ }+ in (cursor, model')++closeCursor :: MonadState Model m => Cursor k v b -> m ()+closeCursor Cursor {..} = state $ \Model{..} ->+ let cursors' = Map.delete cursorID cursors+ model' = Model {+ cursors = cursors'+ , ..+ }+ in ((), model')++readCursor ::+ ( MonadState Model m+ , MonadError Err m+ , SerialiseKey k+ , SerialiseValue v+ , C k v b+ )+ => Int+ -> Cursor k v b+ -> m (V.Vector (Model.Entry k v (BlobRef b)))+readCursor n c = do+ cursor <- guardCursorIsOpen c+ let (qrs, cursor') = Model.readCursor n cursor+ modify (\m -> m {+ cursors = Map.insert (cursorID c) (toSomeCursor cursor') (cursors m)+ })+ pure $ liftBlobRefs (SomeCursorID (cursorID c)) $ qrs++guardCursorIsOpen ::+ forall k v b m. (+ MonadState Model m, MonadError Err m+ , Typeable k, Typeable v, Typeable b+ )+ => Cursor k v b+ -> m (Model.Cursor k v b)+guardCursorIsOpen Cursor{..} =+ gets (Map.lookup cursorID . cursors) >>= \case+ Nothing ->+ throwError ErrCursorClosed+ Just c ->+ pure (fromJust $ fromSomeCursor c)++{-------------------------------------------------------------------------------+ Table union+-------------------------------------------------------------------------------}++-- Is this a (descendant of a) union table?+--+-- This is important for invalidating blob references: if a table is a+-- (descendant of a) union table, then 'supplyUnionCredits' can invalidate blob+-- references.+data IsUnionDescendant = IsUnionDescendant | IsNotUnionDescendant+ deriving stock (Show, Eq)++union ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => ResolveSerialisedValue v+ -> Table k v b+ -> Table k v b+ -> m (Table k v b)+union r th1 th2 = do+ (_, t1) <- guardTableIsOpen th1+ (_, t2) <- guardTableIsOpen th2+ newTableWith TableConfig IsUnionDescendant $ Model.union r t1 t2++unions ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => ResolveSerialisedValue v+ -> NonEmpty (Table k v b)+ -> m (Table k v b)+unions r tables = do+ tables' <- forM tables $ \table -> do+ (_, table') <- guardTableIsOpen table+ pure table'+ newTableWith TableConfig IsUnionDescendant $ Model.unions r tables'++-- | The model can not accurately predict union debt without considerable+-- knowledge about the implementation of /real/ tables. Therefore the model+-- considers unions to be finished right away, and the resulting debt will+-- always be 0.+remainingUnionDebt ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => Table k v b+ -> m UnionDebt+remainingUnionDebt t = do+ (_updc, _table) <- guardTableIsOpen t+ pure (UnionDebt 0)++-- | The union debt is always 0, so supplying union credits has no effect on the+-- tables, except for invalidating its blob references in some cases.+--+-- In the /real/ implementation, blob references can be associated with a run in+-- a regular level, or in a union level. In the former case, only updates can+-- invalidate the blob reference. In the latter case, only supplying union+-- credits can invalidate the blob reference.+--+-- Without considerable knowledge about the /real/ implementation, the model can+-- not /always/ accurately predict which of two cases a blob reference belongs+-- to. However, there is one case where a table is guaranteed not to contain a+-- union level: if the table is /not/ a (descendant of a) union table.+--+-- There is another caveat: without considerable knowledge about the real+-- implementation, the model can not accurately predict after how many supplied+-- union credits /real/ blob references are invalidated. Therefore, we model+-- invalidation conservatively in a similar way to 'updates': any supply of+-- @>=1@ union credits is enough to invalidate union blob references.+--+-- To summarise, 'supplyUnionCredits' will invalidate blob references associated+-- with the input table if:+--+-- * The table is a (descendant of a) union table+--+-- * The number of supplied union credits is at least 1.+supplyUnionCredits ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => Table k v b+ -> UnionCredits+ -> m UnionCredits+supplyUnionCredits t@Table{..} c@(UnionCredits credits)+ | credits <= 0 = do+ _ <- guardTableIsOpen t+ pure (UnionCredits 0) -- always 0, not negative+ | otherwise = do+ (updc, table) <- guardTableIsOpen t+ when (isUnionDescendant == IsUnionDescendant) $+ modify (\m -> m {+ tables = Map.insert tableID (updc + 1, toSomeTable table) (tables m)+ })+ pure c++-- | A version of 'supplyUnionCredits' that supplies a portion of the current debt.+--+-- The debt in the model is always 0, so any portion of that debt is also 0. The+-- real system might have non-zero debt, but the model does not know how much,+-- so it has to assume that *any* portion (assuming it's a positive portion)+-- leads to invalidated blob references.+supplyPortionOfDebt ::+ ( MonadState Model m+ , MonadError Err m+ , C k v b+ )+ => Table k v b+ -> portion+ -> m UnionCredits+supplyPortionOfDebt t@Table{..} _ = do+ (updc, table) <- guardTableIsOpen t+ when (isUnionDescendant == IsUnionDescendant) $+ modify (\m -> m {+ tables = Map.insert tableID (updc + 1, toSomeTable table) (tables m)+ })+ pure (UnionCredits 0)
@@ -0,0 +1,323 @@+{-# LANGUAGE TypeFamilies #-}++-- | A pure model of a single table, supporting both blobs and mupserts.+module Database.LSMTree.Model.Table (+ -- * Serialisation+ SerialiseKey (..)+ , SerialiseValue (..)+ -- * Tables+ , Table (..)+ , empty+ -- * Monoidal value resolution+ , ResolveSerialisedValue (..)+ , getResolve+ , noResolve+ -- * Table querying and updates+ -- ** Queries+ , Range (..)+ , LookupResult (..)+ , lookups+ , Entry (..)+ , rangeLookup+ -- ** Cursor+ , Cursor+ , newCursor+ , readCursor+ -- ** Updates+ , Update (..)+ , updates+ , inserts+ , deletes+ , mupserts+ -- ** Blobs+ , BlobRef+ , retrieveBlobs+ -- * Snapshots+ , snapshot+ -- * Multiple writable tables+ , duplicate+ -- * Table union+ , union+ , unions+ -- * Testing+ , size+ ) where++import qualified Crypto.Hash.SHA256 as SHA256+import Data.Bifunctor+import qualified Data.ByteString as BS+import Data.Kind (Type)+import Data.List.NonEmpty (NonEmpty (..))+import Data.Map (Map)+import qualified Data.Map.Strict as Map+import Data.Proxy (Proxy (Proxy))+import Data.Semigroup (First (..))+import qualified Data.Vector as V+import Database.LSMTree (Entry (..), LookupResult (..), Range (..),+ ResolveValue (..), SerialiseKey (..), SerialiseValue (..),+ Update (..))+import qualified Database.LSMTree.Internal.Map.Range as Map.R+import Database.LSMTree.Internal.RawBytes (RawBytes)+import GHC.Exts (IsList (..))++newtype ResolveSerialisedValue v =+ Resolve { resolveSerialisedValue :: RawBytes -> RawBytes -> RawBytes }++getResolve :: forall v. ResolveValue v => ResolveSerialisedValue v+getResolve = Resolve (resolveSerialised (Proxy @v))++noResolve :: ResolveSerialisedValue v+noResolve = Resolve const++resolveValueAndBlob ::+ ResolveSerialisedValue v+ -> (RawBytes, Maybe b)+ -> (RawBytes, Maybe b)+ -> (RawBytes, Maybe b)+resolveValueAndBlob r (v1, bMay1) (v2, bMay2) =+ (resolveSerialisedValue r v1 v2, getFirst (First bMay1 <> First bMay2))++{-------------------------------------------------------------------------------+ Tables+-------------------------------------------------------------------------------}++type Table :: Type -> Type -> Type -> Type+data Table k v b = Table+ { values :: Map RawBytes (RawBytes, Maybe (BlobRef b))+ }++type role Table nominal nominal nominal++-- | An empty table.+empty :: Table k v b+empty = Table Map.empty++size :: Table k v b -> Int+size (Table m) = Map.size m++-- | This instance is for testing and debugging only.+instance+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => IsList (Table k v b)+ where+ type Item (Table k v b) = (k, v, Maybe b)+ fromList xs = Table $ Map.fromList+ [ (serialiseKey k, (serialiseValue v, mkBlobRef <$> mblob))+ | (k, v, mblob) <- xs+ ]++ toList (Table m) =+ [ (deserialiseKey k, deserialiseValue v, getBlobFromRef <$> mbref)+ | (k, (v, mbref)) <- Map.toList m+ ]++-- | This instance is for testing and debugging only.+instance Show (Table k v b) where+ showsPrec d (Table tbl) = showParen (d > 10)+ $ showString "fromList "+ . showsPrec 11 (toList (Table @BS.ByteString @BS.ByteString @BS.ByteString tbl'))+ where+ tbl' :: Map RawBytes (RawBytes, Maybe (BlobRef BS.ByteString))+ tbl' = fmap (fmap (fmap coerceBlobRef)) tbl++-- | This instance is for testing and debugging only.+deriving stock instance Eq (Table k v b)++{-------------------------------------------------------------------------------+ Lookups+-------------------------------------------------------------------------------}++lookups ::+ (SerialiseKey k, SerialiseValue v)+ => V.Vector k+ -> Table k v b+ -> V.Vector (LookupResult v (BlobRef b))+lookups ks tbl = flip V.map ks $ \k ->+ case Map.lookup (serialiseKey k) (values tbl) of+ Nothing -> NotFound+ Just (v, Nothing) -> Found (deserialiseValue v)+ Just (v, Just br) -> FoundWithBlob (deserialiseValue v) br++rangeLookup :: forall k v b.+ (SerialiseKey k, SerialiseValue v)+ => Range k+ -> Table k v b+ -> V.Vector (Entry k v (BlobRef b))+rangeLookup r tbl = V.fromList+ [ case v of+ (v', Nothing) -> Entry (deserialiseKey k) (deserialiseValue v')+ (v', Just br) -> EntryWithBlob (deserialiseKey k) (deserialiseValue v') br+ | let (lb, ub) = convertRange r+ , (k, v) <- Map.R.rangeLookup lb ub (values tbl)+ ]+ where+ convertRange :: Range k -> (Map.R.Bound RawBytes, Map.R.Bound RawBytes)+ convertRange (FromToExcluding lb ub) =+ ( Map.R.Bound (serialiseKey lb) Map.R.Inclusive+ , Map.R.Bound (serialiseKey ub) Map.R.Exclusive )+ convertRange (FromToIncluding lb ub) =+ ( Map.R.Bound (serialiseKey lb) Map.R.Inclusive+ , Map.R.Bound (serialiseKey ub) Map.R.Inclusive )++{-------------------------------------------------------------------------------+ Updates+-------------------------------------------------------------------------------}++updates :: forall k v b.+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => ResolveSerialisedValue v+ -> V.Vector (k, Update v b)+ -> Table k v b+ -> Table k v b+updates r ups tbl0 = V.foldl' update tbl0 ups where+ update :: Table k v b -> (k, Update v b) -> Table k v b+ update tbl (k, Delete) = tbl+ { values = Map.delete (serialiseKey k) (values tbl) }+ update tbl (k, Insert v Nothing) = tbl+ { values = Map.insert (serialiseKey k) (serialiseValue v, Nothing) (values tbl) }+ update tbl (k, Insert v (Just blob)) = tbl+ { values = Map.insert (serialiseKey k) (serialiseValue v, Just (mkBlobRef blob)) (values tbl)+ }+ update tbl (k, Upsert v) = tbl+ { values = mapUpsert (serialiseKey k) e f (values tbl) }+ where+ e = (serialiseValue v, Nothing)+ f = resolveValueAndBlob r e++mapUpsert :: Ord k => k -> v -> (v -> v) -> Map k v -> Map k v+mapUpsert k v f = Map.alter (Just . g) k where+ g Nothing = v+ g (Just v') = f v'++inserts ::+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => ResolveSerialisedValue v+ -> V.Vector (k, v, Maybe b)+ -> Table k v b+ -> Table k v b+inserts r = updates r . fmap (\(k, v, blob) -> (k, Insert v blob))++deletes ::+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => ResolveSerialisedValue v+ -> V.Vector k+ -> Table k v b+ -> Table k v b+deletes r = updates r . fmap (,Delete)++mupserts ::+ (SerialiseKey k, SerialiseValue v, SerialiseValue b)+ => ResolveSerialisedValue v+ -> V.Vector (k, v)+ -> Table k v b+ -> Table k v b+mupserts r = updates r . fmap (second Upsert)++{-------------------------------------------------------------------------------+ Blobs+-------------------------------------------------------------------------------}++retrieveBlobs ::+ SerialiseValue b+ => V.Vector (BlobRef b)+ -> V.Vector b+retrieveBlobs refs = V.map getBlobFromRef refs++data BlobRef b = BlobRef+ !BS.ByteString -- ^ digest+ !RawBytes -- ^ actual contents+ deriving stock (Show)++type role BlobRef nominal++mkBlobRef :: SerialiseValue b => b -> BlobRef b+mkBlobRef blob = BlobRef (SHA256.hash bs) rb+ where+ !rb = serialiseValue blob+ !bs = deserialiseValue rb :: BS.ByteString++coerceBlobRef :: BlobRef b -> BlobRef b'+coerceBlobRef (BlobRef d b) = BlobRef d b++getBlobFromRef :: SerialiseValue b => BlobRef b -> b+getBlobFromRef (BlobRef _ rb) = deserialiseValue rb++instance Eq (BlobRef b) where+ BlobRef x _ == BlobRef y _ = x == y++{-------------------------------------------------------------------------------+ Snapshots+-------------------------------------------------------------------------------}++snapshot ::+ Table k v b+ -> Table k v b+snapshot = id++{-------------------------------------------------------------------------------+ Multiple writable tables+-------------------------------------------------------------------------------}++duplicate ::+ Table k v b+ -> Table k v b+duplicate = id++{-------------------------------------------------------------------------------+ Cursors+-------------------------------------------------------------------------------}++type Cursor :: Type -> Type -> Type -> Type+data Cursor k v b = Cursor+ { -- | these entries are already resolved, they do not contain duplicate keys.+ _cursorValues :: [(RawBytes, (RawBytes, Maybe (BlobRef b)))]+ }++type role Cursor nominal nominal nominal++newCursor ::+ SerialiseKey k+ => Maybe k+ -> Table k v b+ -> Cursor k v b+newCursor offset tbl = Cursor (skip $ Map.toList $ values tbl)+ where+ skip = case offset of+ Nothing -> id+ Just k -> dropWhile ((< serialiseKey k) . fst)++readCursor ::+ (SerialiseKey k, SerialiseValue v)+ => Int+ -> Cursor k v b+ -> (V.Vector (Entry k v (BlobRef b)), Cursor k v b)+readCursor n c =+ ( V.fromList+ [ case v of+ (v', Nothing) -> Entry (deserialiseKey k) (deserialiseValue v')+ (v', Just br) -> EntryWithBlob (deserialiseKey k) (deserialiseValue v') br+ | (k, v) <- take n (_cursorValues c)+ ]+ , Cursor $ drop n (_cursorValues c)+ )++{-------------------------------------------------------------------------------+ Table union+-------------------------------------------------------------------------------}++-- | Union two full tables, creating a new table.+union ::+ ResolveSerialisedValue v+ -> Table k v b+ -> Table k v b+ -> Table k v b+union r (Table xs) (Table ys) =+ Table (Map.unionWith (resolveValueAndBlob r) xs ys)++-- | Like 'union', but for @n@ tables.+unions ::+ ResolveSerialisedValue v+ -> NonEmpty (Table k v b)+ -> Table k v b+unions r tables =+ Table (Map.unionsWith (resolveValueAndBlob r) (fmap values tables))
@@ -0,0 +1,113 @@+-- | Tests for the @lsm-tree@ library.+--+module Main (main) where++import qualified Control.RefCount++import qualified Test.Database.LSMTree+import qualified Test.Database.LSMTree.Class+import qualified Test.Database.LSMTree.Generators+import qualified Test.Database.LSMTree.Internal+import qualified Test.Database.LSMTree.Internal.Arena+import qualified Test.Database.LSMTree.Internal.BlobFile.FS+import qualified Test.Database.LSMTree.Internal.BloomFilter+import qualified Test.Database.LSMTree.Internal.Chunk+import qualified Test.Database.LSMTree.Internal.CRC32C+import qualified Test.Database.LSMTree.Internal.Entry+import qualified Test.Database.LSMTree.Internal.Index.Compact+import qualified Test.Database.LSMTree.Internal.Index.Ordinary+import qualified Test.Database.LSMTree.Internal.Lookup+import qualified Test.Database.LSMTree.Internal.Merge+import qualified Test.Database.LSMTree.Internal.MergingRun+import qualified Test.Database.LSMTree.Internal.MergingTree+import qualified Test.Database.LSMTree.Internal.PageAcc+import qualified Test.Database.LSMTree.Internal.PageAcc1+import qualified Test.Database.LSMTree.Internal.RawBytes+import qualified Test.Database.LSMTree.Internal.RawOverflowPage+import qualified Test.Database.LSMTree.Internal.RawPage+import qualified Test.Database.LSMTree.Internal.Readers+import qualified Test.Database.LSMTree.Internal.Run+import qualified Test.Database.LSMTree.Internal.RunAcc+import qualified Test.Database.LSMTree.Internal.RunBloomFilterAlloc+import qualified Test.Database.LSMTree.Internal.RunBuilder+import qualified Test.Database.LSMTree.Internal.RunReader+import qualified Test.Database.LSMTree.Internal.Serialise+import qualified Test.Database.LSMTree.Internal.Serialise.Class+import qualified Test.Database.LSMTree.Internal.Snapshot.Codec+import qualified Test.Database.LSMTree.Internal.Snapshot.Codec.Golden+import qualified Test.Database.LSMTree.Internal.Snapshot.FS+import qualified Test.Database.LSMTree.Internal.Unsliced+import qualified Test.Database.LSMTree.Internal.Vector+import qualified Test.Database.LSMTree.Internal.Vector.Growing+import qualified Test.Database.LSMTree.Internal.WriteBufferBlobs.FS+import qualified Test.Database.LSMTree.Internal.WriteBufferReader.FS+import qualified Test.Database.LSMTree.Model.Table+import qualified Test.Database.LSMTree.Resolve+import qualified Test.Database.LSMTree.StateMachine+import qualified Test.Database.LSMTree.StateMachine.DL+import qualified Test.Database.LSMTree.Tracer.Golden+import qualified Test.Database.LSMTree.UnitTests+import qualified Test.FS+import Test.Tasty++main :: IO ()+main = do+ defaultMain $ testGroup "lsm-tree"+ [ Test.Database.LSMTree.tests+ , Test.Database.LSMTree.Internal.Arena.tests+ , Test.Database.LSMTree.Class.tests+ , Test.Database.LSMTree.Generators.tests+ , Test.Database.LSMTree.Internal.tests+ , Test.Database.LSMTree.Internal.BlobFile.FS.tests+ , Test.Database.LSMTree.Internal.BloomFilter.tests+ , Test.Database.LSMTree.Internal.Chunk.tests+ , Test.Database.LSMTree.Internal.CRC32C.tests+ , Test.Database.LSMTree.Internal.Entry.tests+ , Test.Database.LSMTree.Internal.Index.Compact.tests+ , Test.Database.LSMTree.Internal.Index.Ordinary.tests+ , Test.Database.LSMTree.Internal.Lookup.tests+ , Test.Database.LSMTree.Internal.Merge.tests+ , Test.Database.LSMTree.Internal.MergingRun.tests+ , Test.Database.LSMTree.Internal.MergingTree.tests+ , Test.Database.LSMTree.Internal.PageAcc.tests+ , Test.Database.LSMTree.Internal.PageAcc1.tests+ , Test.Database.LSMTree.Internal.RawBytes.tests+ , Test.Database.LSMTree.Internal.RawOverflowPage.tests+ , Test.Database.LSMTree.Internal.RawPage.tests+ , Test.Database.LSMTree.Internal.Readers.tests+ , Test.Database.LSMTree.Internal.Run.tests+ , Test.Database.LSMTree.Internal.RunAcc.tests+ , Test.Database.LSMTree.Internal.RunBloomFilterAlloc.tests+ , Test.Database.LSMTree.Internal.RunBuilder.tests+ , Test.Database.LSMTree.Internal.RunReader.tests+ , Test.Database.LSMTree.Internal.Serialise.tests+ , Test.Database.LSMTree.Internal.Serialise.Class.tests+ , Test.Database.LSMTree.Internal.Snapshot.Codec.tests+ , Test.Database.LSMTree.Internal.Snapshot.Codec.Golden.tests+ , Test.Database.LSMTree.Internal.Snapshot.FS.tests+ , Test.Database.LSMTree.Internal.Unsliced.tests+ , Test.Database.LSMTree.Internal.Vector.tests+ , Test.Database.LSMTree.Internal.Vector.Growing.tests+ , Test.Database.LSMTree.Internal.WriteBufferBlobs.FS.tests+ , Test.Database.LSMTree.Internal.WriteBufferReader.FS.tests+ , Test.Database.LSMTree.Model.Table.tests+ , Test.Database.LSMTree.Resolve.tests+ , Test.Database.LSMTree.StateMachine.tests+ , Test.Database.LSMTree.StateMachine.DL.tests+ , Test.Database.LSMTree.Tracer.Golden.tests+ , Test.Database.LSMTree.UnitTests.tests+ , Test.FS.tests+ ]+ Control.RefCount.checkForgottenRefs+ -- This use of checkForgottenRefs is a last resort. Refs that are forgotten+ -- before being released are detected by the first Ref operation after a+ -- major GC. So they may be thrown during the run of individual tests (though+ -- depending on GC timing this may be during a subsequent test to the one+ -- that triggered the bug). As a last resort, checkForgottenRefs does a last+ -- major GC and will trigger any forgotten refs. So this will reliably catch+ -- the errors, but will not identify where they come from, not even which+ -- test!+ --+ -- If this exception occurs, it may be necessary to put proper use of+ -- checkForgottenRefs into the tests suspected of being the culprit to+ -- identiyfy which one is really failing.
@@ -0,0 +1,314 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RecordWildCards #-}++module Test.Database.LSMTree (tests) where++import Control.Exception+import Control.Tracer+import Data.Function (on)+import Data.IORef+import Data.Monoid+import Data.Typeable (Typeable)+import qualified Data.Vector as V+import qualified Data.Vector.Algorithms as VA+import Data.Void+import Data.Word+import Database.LSMTree+import Database.LSMTree.Extras (showRangesOf)+import Database.LSMTree.Extras.Generators ()+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.FS++tests :: TestTree+tests = testGroup "Test.Database.LSMTree" [+ testGroup "Session" [+ -- openSession+ testProperty "prop_openSession_newSession" prop_openSession_newSession+ , testProperty "prop_openSession_restoreSession" prop_openSession_restoreSession+ -- happy path+ , testProperty "prop_newSession_restoreSession_happyPath" prop_newSession_restoreSession_happyPath+ -- missing session directory+ , testProperty "prop_sessionDirDoesNotExist" prop_sessionDirDoesNotExist+ -- session directory already locked+ , testProperty "prop_sessionDirLocked" prop_sessionDirLocked+ -- malformed session directory+ , testProperty "prop_sessionDirCorrupted" prop_sessionDirCorrupted+ -- salt+ , testProperty "prop_goodAndBadSessionSalt" prop_goodAndBadSessionSalt+ ]+ ]++{-------------------------------------------------------------------------------+ Test types and utilities+-------------------------------------------------------------------------------}++newtype Key = Key Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (Arbitrary, SerialiseKey)++newtype Value = Value Word64+ deriving stock (Show, Eq)+ deriving newtype (Arbitrary, SerialiseValue)+ deriving ResolveValue via Sum Word64++newtype Blob = Blob Word64+ deriving stock (Show, Eq)+ deriving newtype (Arbitrary, SerialiseValue)++data NewOrRestore = New | Restore+ deriving stock (Show, Eq, Bounded, Enum)++instance Arbitrary NewOrRestore where+ arbitrary = arbitraryBoundedEnum+ shrink = shrinkBoundedEnum++-- | If 'New', use 'newSession', otherwise if 'Restore', use 'restoreSession'.+--+-- This allows us to run properties on both 'newSession' and 'restoreSession',+-- without having to write almost identical code twice.+--+-- In a sense, this is somewhat similar to 'openSession', but whereas+-- 'openSession' would defer to 'newSession' or 'restoreSession' based on the+-- directory contents, here the user gets to pick whether to use 'newSession' or+-- 'restoreSession'.+withNewSessionOrRestoreSession ::+ (IOLike m, Typeable h)+ => NewOrRestore+ -> Tracer m LSMTreeTrace+ -> FS.HasFS m h+ -> FS.HasBlockIO m h+ -> Salt+ -> FS.FsPath+ -> (Session m -> m a)+ -> m a+withNewSessionOrRestoreSession newOrRestore tr hfs hbio salt path =+ case newOrRestore of+ New -> withNewSession tr hfs hbio salt path+ Restore -> withRestoreSession tr hfs hbio path++{-------------------------------------------------------------------------------+ Session: openSession+-------------------------------------------------------------------------------}++-- | When the session directory is empty, 'openSession' will call 'newSession'+prop_openSession_newSession :: Property+prop_openSession_newSession =+ ioProperty $+ withTempIOHasBlockIO "prop_openSession_newSession" $ \hfs hbio -> do+ -- Use resultsVar to record which session functions were called+ resultsVar <- newIORef []+ withOpenSession+ (mkSessionOpenModeTracer resultsVar) hfs hbio+ testSalt (FS.mkFsPath [])+ $ \_session -> pure ()+ results <- readIORef resultsVar+ -- Check that we first called openSession, then newSession+ pure $ results === ["Created", "New", "Open"]+ where+ testSalt = 6++-- | When the session directory is non-empty, 'openSession' will call 'restoreSession'+prop_openSession_restoreSession :: Property+prop_openSession_restoreSession =+ ioProperty $+ withTempIOHasBlockIO "prop_openSession_restoreSession" $ \hfs hbio -> do+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath [])+ $ \_session1 -> pure ()+ -- Use resultsVar to record which session functions were called+ resultsVar <- newIORef []+ withOpenSession+ (mkSessionOpenModeTracer resultsVar) hfs hbio+ testSalt (FS.mkFsPath [])+ $ \_session2 -> pure ()+ results <- readIORef resultsVar+ -- Check that we first called openSession, then restoreSession+ pure $ results === ["Created", "Restore", "Open"]+ where+ testSalt = 6++-- | A tracer that records session open, session new, and session restore+-- messages in a mutable variable.+mkSessionOpenModeTracer :: IORef [String] -> Tracer IO LSMTreeTrace+mkSessionOpenModeTracer var = Tracer $ emit $ \case+ TraceSession _ TraceOpenSession{} -> modifyIORef var ("Open" :)+ TraceSession _ TraceNewSession{} -> modifyIORef var ("New" :)+ TraceSession _ TraceRestoreSession{} -> modifyIORef var ("Restore" :)+ TraceSession _ TraceCreatedSession{} -> modifyIORef var ("Created" :)+ _ -> pure ()++{-------------------------------------------------------------------------------+ Session: happy path+-------------------------------------------------------------------------------}++prop_newSession_restoreSession_happyPath ::+ Positive (Small Int)+ -> V.Vector (Key, Value)+ -> Property+prop_newSession_restoreSession_happyPath (Positive (Small bufferSize)) ins =+ ioProperty $+ withTempIOHasBlockIO "prop_newSession_restoreSession_happyPath" $ \hfs hbio -> do+ withNewSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \session1 ->+ withTableWith conf session1 $ \(table :: Table IO Key Value Blob) -> do+ inserts table $ V.map (\(k, v) -> (k, v, Nothing)) ins+ saveSnapshot "snap" "KeyValueBlob" table+ withRestoreSession nullTracer hfs hbio (FS.mkFsPath []) $ \session2 ->+ withTableFromSnapshot session2 "snap" "KeyValueBlob"+ $ \(_ :: Table IO Key Value Blob) -> pure ()+ where+ testSalt = 6+ conf = defaultTableConfig {+ confWriteBufferAlloc = AllocNumEntries bufferSize+ }++{-------------------------------------------------------------------------------+ Session: missing session directory+-------------------------------------------------------------------------------}++prop_sessionDirDoesNotExist :: NewOrRestore -> Property+prop_sessionDirDoesNotExist newOrRestore =+ ioProperty $+ withTempIOHasBlockIO "prop_sessionDirDoesNotExist" $ \hfs hbio -> do+ result <- try @SessionDirDoesNotExistError $+ withNewSessionOrRestoreSession+ newOrRestore+ nullTracer hfs hbio testSalt (FS.mkFsPath ["missing-dir"])+ $ \_session -> pure ()+ pure+ $ counterexample+ ("Expecting an ErrSessionDirDoesNotExist error, but got: " ++ show result)+ $ case result of+ Left ErrSessionDirDoesNotExist{} -> True+ _ -> False+ where+ testSalt = 6++{-------------------------------------------------------------------------------+ Session: session directory already locked+-------------------------------------------------------------------------------}++prop_sessionDirLocked :: NewOrRestore -> Property+prop_sessionDirLocked newOrRestore =+ ioProperty $+ withTempIOHasBlockIO "prop_sessionDirLocked" $ \hfs hbio -> do+ result <-+ withNewSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \_session1 -> do+ try @SessionDirLockedError $+ withNewSessionOrRestoreSession+ newOrRestore+ nullTracer hfs hbio testSalt (FS.mkFsPath [])+ $ \_session2 -> pure ()+ pure+ $ counterexample+ ("Expecting an ErrSessionDirLocked error, but got: " ++ show result)+ $ case result of+ Left ErrSessionDirLocked{} -> True+ _ -> False+ where+ testSalt = 6++{-------------------------------------------------------------------------------+ Session: malformed session directory+-------------------------------------------------------------------------------}++prop_sessionDirCorrupted :: NewOrRestore -> Property+prop_sessionDirCorrupted newOrRestore =+ ioProperty $+ withTempIOHasBlockIO "sessionDirCorrupted" $ \hfs hbio -> do+ FS.createDirectory hfs (FS.mkFsPath ["unexpected-directory"])+ result <- try @SessionDirCorruptedError $+ withNewSessionOrRestoreSession+ newOrRestore+ nullTracer hfs hbio testSalt (FS.mkFsPath [])+ $ \_session -> pure ()+ pure+ $ counterexample+ ("Expecting an ErrSessionDirCorrupted error, but got: " ++ show result)+ $ case result of+ Left ErrSessionDirCorrupted{} -> True+ _ -> False+ where+ testSalt = 6++{-------------------------------------------------------------------------------+ Session: salt+-------------------------------------------------------------------------------}++-- | When we call 'openSession' on an existing session directory, then the salt+-- value we pass in is ignored and the actual salt is restored from a metatada+-- file instead. This property verifies that we indeed ignore the salt value by+-- checking that lookups return the right results, which wouldn't happen if the+-- wrong salt was used.+--+-- NOTE: this only tests with /positive/ lookups, i.e., lookups for keys that+-- are known to exist in the tables.+prop_goodAndBadSessionSalt ::+ Positive (Small Int)+ -> V.Vector (Key, Value)+ -> Property+prop_goodAndBadSessionSalt (Positive (Small bufferSize)) ins =+ checkCoverage $+ ioProperty $+ withTempIOHasBlockIO "prop_sessionSalt" $ \hfs hbio -> do+ -- Open a session and create a snapshot for some arbitrary table contents+ withOpenSession nullTracer hfs hbio goodSalt sessionDir $ \session ->+ withTableWith conf session $ \(table :: Table IO Key Value Void) -> do+ inserts table $ V.map (\(k, v) -> (k, v, Nothing)) insWithoutDupKeys+ saveSnapshot "snap" "KeyValueBlob" table++ -- Determine the expected results of key lookups+ let+ expectedValues :: V.Vector (Maybe Value)+ expectedValues = V.map (Just . snd) insWithoutDupKeys++ -- Open the session using the good salt, open the snapshot, perform lookups+ goodSaltLookups <-+ withOpenSession nullTracer hfs hbio goodSalt sessionDir $ \session ->+ withTableFromSnapshot session "snap" "KeyValueBlob" $ \(table :: Table IO Key Value Void) -> do+ lookups table $ V.map fst insWithoutDupKeys++ -- Determine the result of key lookups using the good salt+ let+ goodSaltValues :: V.Vector (Maybe Value)+ goodSaltValues = V.map getValue goodSaltLookups++ -- Open the session using a bad salt, open the snapshot, perform lookups+ badSaltLookups <-+ withOpenSession nullTracer hfs hbio badSalt sessionDir $ \session ->+ withTableFromSnapshot session "snap" "KeyValueBlob" $ \(table :: Table IO Key Value Void) -> do+ lookups table $ V.map fst insWithoutDupKeys++ -- Determine the result of key lookups using a bad salt+ let+ badSaltValues :: V.Vector (Maybe Value)+ badSaltValues = V.map getValue badSaltLookups++ pure $+ tabulate "number of keys" [ showRangesOf 10 (V.length insWithoutDupKeys) ] $+ -- Regardless of whether the salt we passed to 'openSession' was a good+ -- or bad salt, the lookup results are correct.+ expectedValues === badSaltValues .&&.+ expectedValues === goodSaltValues+ where+ -- Duplicate keys in inserts make the property more complicated, because+ -- keys that are inserted /earlier/ (towards the head of the vector) are+ -- overridden by keys that are inserted /later/ (towards the tail of the+ -- vector). So, we remove duplicate keys instead+ insWithoutDupKeys :: V.Vector (Key, Value)+ insWithoutDupKeys = VA.nubBy (compare `on` fst) ins++ goodSalt :: Salt+ goodSalt = 17++ badSalt :: Salt+ badSalt = 19++ sessionDir = FS.mkFsPath []++ conf = defaultTableConfig {+ confWriteBufferAlloc = AllocNumEntries bufferSize+ }
@@ -0,0 +1,737 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Database.LSMTree.Class (+ tests+ ) where+import Control.Exception (SomeException, assert, try)+import Control.Monad (forM, when)+import Control.Monad.ST.Strict (runST)+import Control.Monad.Trans.State+import qualified Data.ByteString as BS+import Data.Foldable (toList)+import Data.Functor.Compose (Compose (..))+import qualified Data.List as List+import qualified Data.List.NonEmpty as NE+import qualified Data.Proxy as Proxy+import qualified Data.Vector as V+import qualified Data.Vector.Algorithms.Merge as VA+import Data.Word (Word64)+import qualified Database.LSMTree as R+import Database.LSMTree.Class+import Database.LSMTree.Extras.Generators ()+import qualified Database.LSMTree.Model.IO as ModelIO+import qualified System.FS.API as FS+import Test.Database.LSMTree.StateMachine ()+import Test.QuickCheck.Monadic (monadicIO, monitor, run)+import Test.Tasty (TestName, TestTree, testGroup)+import qualified Test.Tasty.QuickCheck as QC+import Test.Tasty.QuickCheck hiding (label)+import qualified Test.Util.FS as FS++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Class"+ [ testGroup "Model" $ zipWith ($) (props tbl1) expectFailures1+ , testGroup "Real" $ zipWith ($) (props tbl2) expectFailures2+ ]+ where+ tbl1 :: RunSetup ModelIO.Table IO+ tbl1 = RunSetup $ \conf -> Setup {+ testTableConfig = conf+ , testWithSessionArgs = \action -> action ModelIO.NoSessionArgs+ }++ expectFailures1 = repeat False++ tbl2 :: RunSetup R.Table IO+ tbl2 = RunSetup $ \conf -> Setup {+ testTableConfig = conf+ , testWithSessionArgs = \action ->+ FS.withTempIOHasBlockIO "R" $ \hfs hbio ->+ action (SessionArgs hfs hbio (FS.mkFsPath []))+ }++ expectFailures2 = repeat False++ props RunSetup {..} =+ [ testProperty' "lookup-insert" $ prop_lookupInsert . runSetup+ , testProperty' "lookup-insert-else" $ prop_lookupInsertElse . runSetup+ , testProperty' "lookup-insert-blob" $ prop_lookupInsertBlob . runSetup+ , testProperty' "lookup-delete" $ prop_lookupDelete . runSetup+ , testProperty' "lookup-delete-else" $ prop_lookupDeleteElse . runSetup+ , testProperty' "insert-insert" $ prop_insertInsert . runSetup+ , testProperty' "insert-insert-blob" $ prop_insertInsertBlob . runSetup+ , testProperty' "insert-commutes" $ prop_insertCommutes . runSetup+ , testProperty' "insert-commutes-blob" $ prop_insertCommutesBlob . runSetup+ , testProperty' "invalidated-blob-references" $ prop_updatesMayInvalidateBlobRefs . runSetup+ , testProperty' "dup-insert-insert" $ prop_dupInsertInsert . runSetup+ , testProperty' "dup-insert-comm" $ prop_dupInsertCommutes . runSetup+ , testProperty' "dup-nochanges" $ prop_dupNoChanges . runSetup+ , testProperty' "lookupRange-like-lookups" $ prop_lookupRangeLikeLookups . runSetup+ , testProperty' "lookupRange-insert" $ prop_insertLookupRange . runSetup+ , testProperty' "readCursor-sorted" $ prop_readCursorSorted . runSetup+ , testProperty' "readCursor-num-results" $ prop_readCursorNumResults . runSetup+ , testProperty' "readCursor-insert" $ prop_readCursorInsert . runSetup+ , testProperty' "readCursor-delete" $ prop_readCursorDelete . runSetup+ , testProperty' "readCursor-delete-else" $ prop_readCursorDeleteElse . runSetup+ , testProperty' "readCursor-stable-view" $ prop_readCursorStableView . runSetup+ , testProperty' "readCursor-offset" $ prop_readCursorOffset . runSetup+ , testProperty' "snapshot-nochanges" $ prop_snapshotNoChanges . runSetup+ , testProperty' "snapshot-nochanges2" $ prop_snapshotNoChanges2 . runSetup+ , testProperty' "lookup-mupsert" $ prop_lookupUpdate . runSetup+ , testProperty' "union" $ prop_union . runSetup+ ]++testProperty' :: forall a. Testable a => TestName -> a -> Bool -> TestTree+testProperty' name prop = \b ->+ testProperty name ((if b then expectFailure else property) prop)++-------------------------------------------------------------------------------+-- test setup and helpers+-------------------------------------------------------------------------------++type Key = Word64+type Blob = BS.ByteString++newtype Value = Value BS.ByteString+ deriving stock (Eq, Show)+ deriving newtype (Arbitrary, Semigroup, SerialiseValue)+ deriving ResolveValue via (R.ResolveViaSemigroup Value)++label :: SnapshotLabel+label = SnapshotLabel "Word64 ByteString ByteString"++type Proxy h = Setup h IO++newtype RunSetup h m = RunSetup {+ runSetup :: TableConfig h -> Setup h m+ }++data Setup h m = Setup {+ testTableConfig :: TableConfig h+ , testWithSessionArgs :: forall a. (SessionArgs (Session h) m -> m a) -> m a+ }++-- | create session, table, and populate it with some data.+withSessionAndTableNew :: forall h m a.+ ( IsTable h+ , IOLike m+ )+ => Setup h m+ -> [(Key, Update Value Blob)]+ -> (Session h m -> h m Key Value Blob -> m a)+ -> m a+withSessionAndTableNew Setup{..} ups action =+ testWithSessionArgs $ \args ->+ withSession args $ \sesh ->+ withTableNew sesh testTableConfig $ \table -> do+ updates table (V.fromList ups)+ action sesh table++-- | Like 'retrieveBlobs' but works for any 'Traversable'.+--+-- Like 'partsOf' in @lens@ this uses state monad.+retrieveBlobsTrav ::+ ( IsTable h+ , IOLike m+ , SerialiseValue b+ , Traversable t+ , C_ b+ )+ => proxy h+ -> Session h m+ -> t (BlobRef h m b)+ -> m (t b)+retrieveBlobsTrav tbl ses brefs = do+ blobs <- retrieveBlobs tbl ses (V.fromList $ toList brefs)+ evalStateT (traverse (\_ -> state un) brefs) (V.toList blobs)+ where+ un [] = error "invalid traversal"+ un (x:xs) = (x, xs)++lookupsWithBlobs :: forall h m k v b.+ ( IsTable h+ , IOLike m+ , C k v b+ )+ => h m k v b+ -> Session h m+ -> V.Vector k+ -> m (V.Vector (LookupResult v b))+lookupsWithBlobs tbl ses ks = do+ res <- lookups tbl ks+ getCompose <$> retrieveBlobsTrav (Proxy.Proxy @h) ses (Compose res)++rangeLookupWithBlobs :: forall h m k v b.+ ( IsTable h+ , IOLike m+ , C k v b+ )+ => h m k v b+ -> Session h m+ -> Range k+ -> m (V.Vector (Entry k v b))+rangeLookupWithBlobs tbl ses r = do+ res <- rangeLookup tbl r+ getCompose <$> retrieveBlobsTrav (Proxy.Proxy @h) ses (Compose res)++readCursorWithBlobs :: forall h m k v b proxy.+ ( IsTable h+ , IOLike m+ , C k v b+ )+ => proxy h+ -> Session h m+ -> Cursor h m k v b+ -> Int+ -> m (V.Vector (Entry k v b))+readCursorWithBlobs tbl ses cursor n = do+ res <- readCursor tbl n cursor+ getCompose <$> retrieveBlobsTrav tbl ses (Compose res)++readCursorAllWithBlobs :: forall h m k v b proxy.+ ( IsTable h+ , IOLike m+ , C k v b+ )+ => proxy h+ -> Session h m+ -> Cursor h m k v b+ -> CursorReadSchedule+ -> m [V.Vector (Entry k v b)]+readCursorAllWithBlobs tbl ses cursor = go . getCursorReadSchedule+ where+ go [] = error "readCursorAllWithBlobs: finite infinite list"+ go (n : ns) = do+ res <- readCursorWithBlobs tbl ses cursor n+ if V.null res+ then pure [res]+ else (res :) <$> go ns++type CursorReadSchedule = InfiniteList (Positive Int)++getCursorReadSchedule :: CursorReadSchedule -> [Int]+getCursorReadSchedule = map getPositive . getInfiniteList++-------------------------------------------------------------------------------+-- implement classic QC tests for basic k/v properties+-------------------------------------------------------------------------------++-- | You can lookup what you inserted.+prop_lookupInsert ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Property+prop_lookupInsert h ups k v = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do++ -- the main dish+ inserts tbl (V.singleton (k, v, Nothing))+ res <- lookupsWithBlobs tbl ses (V.singleton k)++ pure $ res === V.singleton (Found v)++-- | Insert doesn't change the lookup results of other keys.+prop_lookupInsertElse ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> [Key] -> Property+prop_lookupInsertElse h ups k v testKeys = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do++ let testKeys' = V.fromList $ filter (/= k) testKeys+ res1 <- lookupsWithBlobs tbl ses testKeys'+ inserts tbl (V.singleton (k, v, Nothing))+ res2 <- lookupsWithBlobs tbl ses testKeys'++ pure $ res1 === res2++-- | You cannot lookup what you have just deleted+prop_lookupDelete ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Property+prop_lookupDelete h ups k = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ deletes tbl (V.singleton k)+ res <- lookupsWithBlobs tbl ses (V.singleton k)+ pure $ res === V.singleton NotFound++-- | Delete doesn't change the lookup results of other keys+prop_lookupDeleteElse ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> [Key] -> Property+prop_lookupDeleteElse h ups k testKeys = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do++ let testKeys' = V.fromList $ filter (/= k) testKeys+ res1 <- lookupsWithBlobs tbl ses testKeys'+ deletes tbl (V.singleton k)+ res2 <- lookupsWithBlobs tbl ses testKeys'++ pure $ res1 === res2++-- | Last insert wins.+prop_insertInsert ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Value -> Property+prop_insertInsert h ups k v1 v2 = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ inserts tbl (V.fromList [(k, v1, Nothing), (k, v2, Nothing)])+ res <- lookupsWithBlobs tbl ses (V.singleton k)+ pure $ res === V.singleton (Found v2)++-- | Inserts with different keys don't interfere.+prop_insertCommutes ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Key -> Value -> Property+prop_insertCommutes h ups k1 v1 k2 v2 = k1 /= k2 ==> ioProperty do+ withSessionAndTableNew h ups $ \ses tbl -> do+ inserts tbl (V.fromList [(k1, v1, Nothing), (k2, v2, Nothing)])++ res <- lookupsWithBlobs tbl ses (V.fromList [k1,k2])+ pure $ res === V.fromList [Found v1, Found v2]++-------------------------------------------------------------------------------+-- implement classic QC tests for cursors+-------------------------------------------------------------------------------++-- | Cursor read results are sorted by key.+prop_readCursorSorted ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Maybe Key+ -> CursorReadSchedule+ -> Property+prop_readCursorSorted h ups offset ns = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ res <- withCursor offset tbl $ \cursor -> do+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ let keys = map queryResultKey (V.toList res)+ pure $ keys === List.sort keys++-- | Cursor reads return the requested number of results, until the end.+prop_readCursorNumResults ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Maybe Key+ -> CursorReadSchedule+ -> Property+prop_readCursorNumResults h ups offset ns = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ res <- withCursor offset tbl $ \cursor -> do+ readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ let elemsRead = map V.length res+ let numFullReads = length res - 2+ pure $ last elemsRead === 0+ .&&. take numFullReads elemsRead+ === take numFullReads (getCursorReadSchedule ns)++-- | You can read what you inserted.+prop_readCursorInsert ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> CursorReadSchedule+ -> Key -> Value -> Property+prop_readCursorInsert h ups ns k v = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ inserts tbl (V.singleton (k, v, Nothing))+ res <- withCursor Nothing tbl $ \cursor ->+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ pure $ V.find (\r -> queryResultKey r == k) res+ === Just (Entry k v)++-- | You can't read what you deleted.+prop_readCursorDelete ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> CursorReadSchedule+ -> Key -> Property+prop_readCursorDelete h ups ns k = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ deletes tbl (V.singleton k)+ res <- withCursor Nothing tbl $ \cursor -> do+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ pure $ V.find (\r -> queryResultKey r == k) res === Nothing++-- | Updates don't change the cursor read results of other keys.+prop_readCursorDeleteElse ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Maybe Key+ -> CursorReadSchedule+ -> [(Key, Update Value Blob)] -> Property+prop_readCursorDeleteElse h ups offset ns ups2 = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ res1 <- withCursor offset tbl $ \cursor -> do+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ updates tbl (V.fromList ups2)+ res2 <- withCursor offset tbl $ \cursor -> do+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ let updatedKeys = map fst ups2+ pure $ V.filter (\r -> queryResultKey r `notElem` updatedKeys) res1+ === V.filter (\r -> queryResultKey r `notElem` updatedKeys) res2++-- | Updates don't affect previously created cursors.+prop_readCursorStableView ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Maybe Key+ -> CursorReadSchedule+ -> [(Key, Update Value Blob)] -> Property+prop_readCursorStableView h ups offset ns ups2 = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ res1 <- withCursor offset tbl $ \cursor -> do+ readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ res2 <- withCursor offset tbl $ \cursor -> do+ updates tbl (V.fromList ups2)+ readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ pure $ res1 === res2++-- | Creating a cursor at an offset simply skips a prefix.+prop_readCursorOffset ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key+ -> CursorReadSchedule+ -> Property+prop_readCursorOffset h ups offset ns = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ res1 <- withCursor (Just offset) tbl $ \cursor -> do+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ res2 <- withCursor Nothing tbl $ \cursor -> do+ V.concat <$> readCursorAllWithBlobs (Proxy.Proxy @h) ses cursor ns+ pure $ res1 === V.dropWhile ((< offset) . queryResultKey) res2++-------------------------------------------------------------------------------+-- implement classic QC tests for range lookups+-------------------------------------------------------------------------------++evalRange :: Ord k => Range k -> k -> Bool+evalRange (FromToExcluding lo hi) x = lo <= x && x < hi+evalRange (FromToIncluding lo hi) x = lo <= x && x <= hi++queryResultKey :: Entry k v b -> k+queryResultKey (Entry k _) = k+queryResultKey (EntryWithBlob k _ _ ) = k++queryResultFromLookup :: k -> LookupResult v b -> Maybe (Entry k v b)+queryResultFromLookup k = \case+ NotFound -> Nothing+ Found v -> Just (Entry k v)+ FoundWithBlob v b -> Just (EntryWithBlob k v b)++-- | A range lookup behaves like many point lookups.+prop_lookupRangeLikeLookups ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Range Key+ -> Property+prop_lookupRangeLikeLookups h ups r = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ res1 <- rangeLookupWithBlobs tbl ses r++ let testKeys = V.fromList $ nubSort $ filter (evalRange r) $ map fst ups+ res2 <- fmap (V.catMaybes . V.zipWith queryResultFromLookup testKeys) $+ lookupsWithBlobs tbl ses testKeys++ pure $ res1 === res2++ where+ nubSort = map NE.head . NE.group . List.sort++-- | Last insert wins.+prop_insertLookupRange ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Range Key -> Property+prop_insertLookupRange h ups k v r = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do++ res <- rangeLookupWithBlobs tbl ses r++ inserts tbl (V.singleton (k, v, Nothing))++ res' <- rangeLookupWithBlobs tbl ses r++ let p :: Entry Key Value b -> Bool+ p rlr = queryResultKey rlr /= k++ if evalRange r k+ then pure $ vsortOn queryResultKey (V.cons (Entry k v) (V.filter p res)) === res'+ else pure $ res === res'++ where+ vsortOn f vec = runST $ do+ mvec <- V.thaw vec+ VA.sortBy (\e1 e2 -> f e1 `compare` f e2) mvec+ V.unsafeFreeze mvec++-------------------------------------------------------------------------------+-- implement classic QC tests for split-value BLOB retrieval+-------------------------------------------------------------------------------++-- | You can lookup what you inserted.+prop_lookupInsertBlob ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Blob -> Property+prop_lookupInsertBlob h ups k v blob = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do++ -- the main dish+ inserts tbl (V.singleton (k, v, Just blob))+ res <- lookupsWithBlobs tbl ses (V.singleton k)++ pure $ res === V.singleton (FoundWithBlob v blob)++-- | Last insert wins.+prop_insertInsertBlob ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Value -> Maybe Blob -> Maybe Blob -> Property+prop_insertInsertBlob h ups k v1 v2 mblob1 mblob2 = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ inserts tbl (V.fromList [(k, v1, mblob1), (k, v2, mblob2)])+ res <- lookupsWithBlobs tbl ses (V.singleton k)+ pure $ res === case mblob2 of+ Nothing -> V.singleton (Found v2)+ Just blob2 -> V.singleton (FoundWithBlob v2 blob2)++-- | Inserts with different keys don't interfere.+prop_insertCommutesBlob ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Maybe Blob+ -> Key -> Value -> Maybe Blob -> Property+prop_insertCommutesBlob h ups k1 v1 mblob1 k2 v2 mblob2 = k1 /= k2 ==> ioProperty do+ withSessionAndTableNew h ups $ \ses tbl -> do+ inserts tbl (V.fromList [(k1, v1, mblob1), (k2, v2, mblob2)])++ res <- lookupsWithBlobs tbl ses $ V.fromList [k1,k2]+ pure $ res === case (mblob1, mblob2) of+ (Nothing, Nothing) -> V.fromList [Found v1, Found v2]+ (Just blob1, Nothing) -> V.fromList [FoundWithBlob v1 blob1, Found v2]+ (Nothing, Just blob2) -> V.fromList [Found v1, FoundWithBlob v2 blob2]+ (Just blob1, Just blob2) -> V.fromList [FoundWithBlob v1 blob1, FoundWithBlob v2 blob2]++-- | A blob reference may be invalidated by an update.+prop_updatesMayInvalidateBlobRefs ::+ forall h. IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Blob+ -> [(Key, Update Value Blob)]+ -> Property+prop_updatesMayInvalidateBlobRefs h ups k1 v1 blob1 ups' = monadicIO $ do+ (res, blobs, res') <- run $ do+ withSessionAndTableNew h ups $ \ses tbl -> do+ inserts tbl (V.singleton (k1, v1, Just blob1))+ res <- lookups tbl (V.singleton k1)+ blobs <- getCompose <$> retrieveBlobsTrav (Proxy.Proxy @h) ses (Compose res)+ updates tbl (V.fromList ups')+ res' <- try @SomeException (getCompose <$> retrieveBlobsTrav (Proxy.Proxy @h) ses (Compose res))+ pure (res, blobs, res')++ case (V.toList res, V.toList blobs) of+ ([FoundWithBlob{}], [FoundWithBlob _ x])+ | Left _ <- res' ->+ monitor (QC.label "blob reference invalidated") >> pure True+ | Right (V.toList -> [FoundWithBlob _ y]) <- res' ->+ monitor (QC.label "blob reference valid") >> pure (x == y)+ _ -> monitor (counterexample "insert before lookup failed, somehow...") >> pure False++-------------------------------------------------------------------------------+-- implement classic QC tests for monoidal updates+-------------------------------------------------------------------------------++-- | You can lookup what you inserted.+prop_lookupUpdate ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Maybe Blob -> Value -> Property+prop_lookupUpdate h ups k v1 mb1 v2 = ioProperty $ do+ withSessionAndTableNew h ups $ \s tbl -> do++ -- the main dish+ inserts tbl (V.singleton (k, v1, mb1))+ upserts tbl (V.singleton (k, v2))+ res <- lookupsWithBlobs tbl s (V.singleton k)++ -- notice the order.+ pure $ res === V.singleton (Found (resolve v2 v1))++-------------------------------------------------------------------------------+-- implement classic QC tests for monoidal table unions+-------------------------------------------------------------------------------++prop_union :: forall h.+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)] -> [(Key, Update Value Blob)]+ -> [Key]+ -> Positive (Small Int) -- ^ Number of batches for supplying union credits+ -> Property+prop_union h ups1 ups2+ (V.fromList -> testKeys)+ (Positive (Small nSupplyBatches))+ = ioProperty $ do+ withSessionAndTableNew h ups1 $ \s tbl1 -> do+ withTableNew s (testTableConfig h) $ \tbl2 -> do+ updates tbl2 $ V.fromList ups2++ -- union them.+ withTableUnion tbl1 tbl2 $ \tbl3 -> do+ propBegin <- compareLookups s tbl1 tbl2 tbl3++ -- Supply union credits in @nSupplyBatches@ batches+ props <- forM (reverse [1 .. nSupplyBatches]) $ \i -> do+ -- Try to keep the batch sizes roughly the same size+ UnionDebt debt <- remainingUnionDebt tbl3+ _ <- supplyUnionCredits tbl3 (UnionCredits (debt `div` i))++ -- In case @i == 0@, then @debt `div` i == debt@, so the last supply+ -- should have finished the union.+ when (i == 1) $ do+ finalDebt <- remainingUnionDebt tbl3+ assert (finalDebt == UnionDebt 0) $ pure ()++ -- Check that the lookup results are still the same after each batch+ -- of union credits.+ compareLookups s tbl1 tbl2 tbl3++ pure (propBegin .&&. conjoin props)+ where+ compareLookups s tbl1 tbl2 tbl3 = do+ -- results in parts and the union table+ res1 <- lookupsWithBlobs tbl1 s testKeys+ res2 <- lookupsWithBlobs tbl2 s testKeys+ res3 <- lookupsWithBlobs tbl3 s testKeys++ let unionResult ::+ LookupResult Value Blob+ -> LookupResult Value Blob+ -> LookupResult Value Blob++ unionResult r@NotFound NotFound = r+ unionResult NotFound r@(Found _) = r+ unionResult NotFound r@(FoundWithBlob _ _) = r++ unionResult r@(Found _) NotFound+ = r+ unionResult (Found v1) (Found v2)+ = Found (resolve v1 v2)+ unionResult (Found v1) (FoundWithBlob v2 _)+ = Found (resolve v1 v2)++ unionResult r@(FoundWithBlob _ _) NotFound+ = r+ unionResult (FoundWithBlob v1 b1) (Found v2)+ = FoundWithBlob (resolve v1 v2) b1+ unionResult (FoundWithBlob v1 b1) (FoundWithBlob v2 _b2)+ = FoundWithBlob (resolve v1 v2) b1++ pure $ V.zipWith unionResult res1 res2 === res3++-------------------------------------------------------------------------------+-- implement classic QC tests for snapshots+-------------------------------------------------------------------------------++-- changes to handle would not change the snapshot+prop_snapshotNoChanges :: forall h.+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> [(Key, Update Value Blob)] -> [Key] -> Property+prop_snapshotNoChanges h ups ups' testKeys = ioProperty $ do+ withSessionAndTableNew h ups $ \ses tbl1 -> do++ res <- lookupsWithBlobs tbl1 ses $ V.fromList testKeys++ let name = R.toSnapshotName "foo"++ saveSnapshot name label tbl1+ updates tbl1 (V.fromList ups')++ withTableFromSnapshot @h ses label name$ \tbl2 -> do++ res' <- lookupsWithBlobs tbl2 ses $ V.fromList testKeys++ pure $ res == res'++-- same snapshot may be opened multiple times,+-- and the handles are separate.+prop_snapshotNoChanges2 :: forall h.+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> [(Key, Update Value Blob)] -> [Key] -> Property+prop_snapshotNoChanges2 h ups ups' testKeys = ioProperty $ do+ withSessionAndTableNew h ups $ \sess tbl0 -> do+ let name = R.toSnapshotName "foo"+ saveSnapshot name label tbl0++ withTableFromSnapshot @h sess label name $ \tbl1 ->+ withTableFromSnapshot @h sess label name $ \tbl2 -> do++ res <- lookupsWithBlobs tbl1 sess $ V.fromList testKeys+ updates tbl1 (V.fromList ups')+ res' <- lookupsWithBlobs tbl2 sess $ V.fromList testKeys++ pure $ res == res'++-------------------------------------------------------------------------------+-- implement classic QC tests for multiple writable tables+-- - results of insert/delete, monoidal update, lookups, and range lookups should+-- be equal if applied to two duplicated tables+-- - changes to one table should not cause any visible changes in any others+-------------------------------------------------------------------------------++-- | Last insert wins.+prop_dupInsertInsert ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Value -> [Key] -> Property+prop_dupInsertInsert h ups k v1 v2 testKeys = ioProperty $ do+ withSessionAndTableNew h ups $ \sess tbl1 -> do+ withTableDuplicate tbl1 $ \tbl2 -> do++ inserts tbl1 (V.fromList [(k, v1, Nothing), (k, v2, Nothing)])+ inserts tbl2 (V.fromList [(k, v2, Nothing)])++ res1 <- lookupsWithBlobs tbl1 sess $ V.fromList testKeys+ res2 <- lookupsWithBlobs tbl2 sess $ V.fromList testKeys+ pure $ res1 === res2++-- | Different key inserts commute.+prop_dupInsertCommutes ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> Key -> Value -> Key -> Value -> [Key] -> Property+prop_dupInsertCommutes h ups k1 v1 k2 v2 testKeys = k1 /= k2 ==> ioProperty do+ withSessionAndTableNew h ups $ \sess tbl1 -> do+ withTableDuplicate tbl1 $ \tbl2 -> do++ inserts tbl1 (V.fromList [(k1, v1, Nothing), (k2, v2, Nothing)])+ inserts tbl2 (V.fromList [(k2, v2, Nothing), (k1, v1, Nothing)])++ res1 <- lookupsWithBlobs tbl1 sess $ V.fromList testKeys+ res2 <- lookupsWithBlobs tbl2 sess $ V.fromList testKeys+ pure $ res1 === res2++-- changes to one handle should not cause any visible changes in any others+prop_dupNoChanges ::+ IsTable h+ => Proxy h -> [(Key, Update Value Blob)]+ -> [(Key, Update Value Blob)] -> [Key] -> Property+prop_dupNoChanges h ups ups' testKeys = ioProperty $ do+ withSessionAndTableNew h ups $ \sess tbl1 -> do++ res <- lookupsWithBlobs tbl1 sess $ V.fromList testKeys++ withTableDuplicate tbl1 $ \tbl2 -> do+ updates tbl2 (V.fromList ups')++ -- lookup tbl1 again.+ res' <- lookupsWithBlobs tbl1 sess $ V.fromList testKeys++ pure $ res == res'
@@ -0,0 +1,200 @@+module Test.Database.LSMTree.Generators (+ tests+ ) where++import Data.Bifoldable (bifoldMap)+import Data.Coerce (Coercible, coerce)+import qualified Data.Map.Strict as Map+import qualified Data.Vector.Primitive as VP+import Data.Word (Word64, Word8)+import Database.LSMTree.Extras (showPowersOf)+import Database.LSMTree.Extras.Generators+import Database.LSMTree.Extras.MergingRunData+import Database.LSMTree.Extras.MergingTreeData+import Database.LSMTree.Extras.ReferenceImpl+import Database.LSMTree.Extras.RunData+import Database.LSMTree.Internal.BlobRef (BlobSpan)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry+import qualified Database.LSMTree.Internal.Index as Index+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.PageAcc (entryWouldFitInPage,+ sizeofEntry)+import Database.LSMTree.Internal.Paths (RunFsPaths (..))+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber (RunNumber (..))+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.Sim.MockFS as MockFS+import qualified Test.QuickCheck as QC+import Test.QuickCheck (Property)+import Test.Tasty (TestTree, localOption, testGroup)+import Test.Tasty.QuickCheck (QuickCheckMaxSize (..), testProperty,+ (===))+import Test.Util.Arbitrary+import Test.Util.FS (propNoOpenHandles, withSimHasBlockIO)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Generators" [+ testGroup "PageContentFits" $+ prop_arbitraryAndShrinkPreserveInvariant noTags+ pageContentFitsInvariant+ , testGroup "PageContentOrdered" $+ prop_arbitraryAndShrinkPreserveInvariant noTags+ pageContentOrderedInvariant+ , localOption (QuickCheckMaxSize 20) $ -- takes too long!+ testGroup "LogicalPageSummaries" $+ prop_arbitraryAndShrinkPreserveInvariant noTags $+ pagesInvariant @Word64+ , testGroup "Chunk size" $+ prop_arbitraryAndShrinkPreserveInvariant noTags+ chunkSizeInvariant+ , testGroup "RawBytes" $+ [ testProperty "packRawBytesPinnedOrUnpinned"+ prop_packRawBytesPinnedOrUnpinned+ ]+ ++ prop_arbitraryAndShrinkPreserveInvariant labelRawBytes+ (deepseqInvariant @RawBytes)+ , testGroup "LargeRawBytes" $+ prop_arbitraryAndShrinkPreserveInvariant+ (\(LargeRawBytes rb) -> labelRawBytes rb)+ (deepseqInvariant @LargeRawBytes)+ , testGroup "BiasedKey" $+ prop_arbitraryAndShrinkPreserveInvariant+ (labelTestKOps @BiasedKey)+ deepseqInvariant+ , testGroup "helpers"+ [ testProperty "prop_shrinkVec" $ \vec ->+ shrinkVec (QC.shrink @Int) vec === map VP.fromList (QC.shrink (VP.toList vec))+ ]+ , testGroup "RunData" $+ prop_arbitraryAndShrinkPreserveInvariant+ labelRunData+ noInvariant+ ++ [ testProperty "withRun doesn't leak resources" $ \rd ->+ QC.ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_withRunDoesntLeak hfs hbio rd+ ]+ , testGroup "NonEmptyRunData" $+ prop_arbitraryAndShrinkPreserveInvariant+ labelNonEmptyRunData+ noInvariant+ , testGroup "MergingRunData" $+ prop_arbitraryAndShrinkPreserveInvariant+ @(SerialisedMergingRunData MR.LevelMergeType)+ labelMergingRunData+ ((=== Right ()) . mergingRunDataInvariant)+ ++ [ testProperty "withMergingRun doesn't leak resources" $ \mrd ->+ QC.ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_withMergingRunDoesntLeak hfs hbio mrd+ ]+ , testGroup "MergingTreeData" $+ prop_arbitraryAndShrinkPreserveInvariant+ labelMergingTreeData+ ((=== Right ()) . mergingTreeDataInvariant)+ ++ [ testProperty "withMergingTree doesn't leak resources" $ \mtd ->+ QC.ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_withMergingTreeDoesntLeak hfs hbio mtd+ ]+ ]++testSalt :: Bloom.Salt+testSalt = 4++runParams :: Index.IndexType -> RunBuilder.RunParams+runParams indexType =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = indexType+ }++prop_packRawBytesPinnedOrUnpinned :: Bool -> [Word8] -> Bool+prop_packRawBytesPinnedOrUnpinned pinned ws =+ packRawBytesPinnedOrUnpinned pinned ws == RawBytes (VP.fromList ws)++labelRawBytes :: RawBytes -> Property -> Property+labelRawBytes rb =+ QC.tabulate "size" [showPowersOf 2 (RB.size rb)]++type TestEntry = Entry SerialisedValue BlobSpan+type TestKOp k = (k, TestEntry)++labelTestKOps ::+ Coercible k SerialisedKey+ => [TestKOp k]+ -> Property+ -> Property+labelTestKOps kops' =+ QC.tabulate "key occurrences (>1 is collision)" (map (show . snd) (Map.assocs keyCounts))+ . QC.tabulate "key sizes" (map (showPowersOf 4 . sizeofKey) keys)+ . QC.tabulate "value sizes" (map (showPowersOf 4 . sizeofValue) values)+ . QC.tabulate "k/op sizes" (map (showPowersOf 4 . uncurry sizeofEntry) kops)+ . QC.tabulate "k/op is large" (map (show . isLarge) kops)+ . QC.checkCoverage+ . QC.cover 50 (any isLarge kops) "any k/op is large"+ . QC.cover 1 (ratioUniqueKeys < (0.9 :: Double)) ">10% of keys collide"+ . QC.cover 5 (any (> 2) keyCounts) "has key with >2 collisions"+ where+ kops = coerce kops' :: [(SerialisedKey, Entry SerialisedValue BlobSpan)]+ keys = map fst kops+ keyCounts = Map.fromListWith (+) [(k, (1 :: Int)) | k <- keys]+ uniqueKeys = Map.keys keyCounts+ ratioUniqueKeys = fromIntegral (length uniqueKeys) / fromIntegral (length keys)+ values = foldMap (bifoldMap pure mempty . snd) kops++ isLarge = not . uncurry entryWouldFitInPage++prop_withRunDoesntLeak ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> SerialisedRunData+ -> IO Property+prop_withRunDoesntLeak hfs hbio rd = do+ let indexType = Index.Ordinary+ let path = FS.mkFsPath ["something-1"]+ let fsPaths = RunFsPaths path (RunNumber 0)+ FS.createDirectory hfs path+ withRunAt hfs hbio testSalt (runParams indexType) fsPaths rd $ \_run -> do+ pure (QC.property True)++prop_withMergingRunDoesntLeak ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> SerialisedMergingRunData MR.LevelMergeType+ -> IO Property+prop_withMergingRunDoesntLeak hfs hbio mrd = do+ let indexType = Index.Ordinary+ let path = FS.mkFsPath ["something-2"]+ FS.createDirectory hfs path+ counter <- newUniqCounter 0+ withMergingRun hfs hbio resolveVal testSalt (runParams indexType) path counter mrd $+ \_mr -> do+ pure (QC.property True)++-- TODO: This only tests the happy path. For everything else, we'd need to+-- inject errors, e.g. with @simErrorHasBlockIO@.+prop_withMergingTreeDoesntLeak ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> SerialisedMergingTreeData+ -> IO Property+prop_withMergingTreeDoesntLeak hfs hbio mrd = do+ let indexType = Index.Ordinary+ let path = FS.mkFsPath ["something-3"]+ FS.createDirectory hfs path+ counter <- newUniqCounter 0+ withMergingTree hfs hbio resolveVal testSalt (runParams indexType) path counter mrd $+ \_tree -> do+ pure (QC.property True)++resolveVal :: SerialisedValue -> SerialisedValue -> SerialisedValue+resolveVal (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)
@@ -0,0 +1,122 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RecordWildCards #-}++module Test.Database.LSMTree.Internal (tests) where++import Control.Tracer+import Data.Coerce (coerce)+import qualified Data.Map.Strict as Map+import Data.Maybe (isJust, mapMaybe)+import qualified Data.Vector as V+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.BlobRef+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Unsafe+import qualified System.FS.API as FS+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.FS++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal" [+ testGroup "Cursor" [+ testProperty "prop_roundtripCursor" $ withMaxSuccess 500 $+ prop_roundtripCursor+ ]+ ]+++testSalt :: Bloom.Salt+testSalt = 4++testTableConfig :: TableConfig+testTableConfig = defaultTableConfig {+ -- Write buffer size is small on purpose, so that the test actually+ -- flushes and merges.+ confWriteBufferAlloc = AllocNumEntries 3+ }++-- | Check that reading from a cursor returns exactly the entries that have+-- been inserted into the table. Roughly:+--+-- @+-- readCursor . newCursor . inserts == id+-- @+--+-- If lower and upper bound are provided:+--+-- @+-- readCursorWhile (<= ub) . newCursorAt lb . inserts+-- == takeWhile ((<= ub) . key) . dropWhile ((< lb) . key)+-- @+prop_roundtripCursor ::+ Maybe SerialisedKey -- ^ Inclusive lower bound+ -> Maybe SerialisedKey -- ^ Inclusive upper bound+ -> V.Vector (SerialisedKey, Entry SerialisedValue SerialisedBlob)+ -> Property+prop_roundtripCursor lb ub kops = ioProperty $+ withTempIOHasBlockIO "prop_roundtripCursor" $ \hfs hbio -> do+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sesh -> do+ withTable sesh conf $ \t -> do+ updates resolve (coerce kops) t+ fromCursor <- withCursor resolve (toOffsetKey lb) t $ \c ->+ fetchBlobs hfs =<< readCursorUntil ub c+ pure $+ tabulate "duplicates" (show <$> Map.elems duplicates) $+ tabulate "any blobs" [show (any (isJust . snd . snd) fromCursor)] $+ expected === fromCursor+ where+ conf = testTableConfig++ fetchBlobs :: FS.HasFS IO h+ -> V.Vector (k, (v, Maybe (WeakBlobRef IO h)))+ -> IO (V.Vector (k, (v, Maybe SerialisedBlob)))+ fetchBlobs hfs = traverse (traverse (traverse (traverse (readWeakBlobRef hfs))))++ toOffsetKey = maybe NoOffsetKey (OffsetKey . coerce)++ expected =+ V.fromList . mapMaybe (traverse entryToValue) $+ maybe id (\k -> takeWhile ((<= k) . fst)) ub $+ maybe id (\k -> dropWhile ((< k) . fst)) lb $+ Map.assocs . Map.fromListWith (combine resolve) $+ V.toList kops++ entryToValue :: Entry v b -> Maybe (v, Maybe b)+ entryToValue = \case+ Insert v -> Just (v, Nothing)+ InsertWithBlob v b -> Just (v, Just b)+ Upsert v -> Just (v, Nothing)+ Delete -> Nothing++ duplicates :: Map.Map SerialisedKey Int+ duplicates =+ Map.filter (> 1) $+ Map.fromListWith (+) . map (\(k, _) -> (k, 1)) $+ V.toList kops++readCursorUntil ::+ Maybe SerialisedKey -- Inclusive upper bound+ -> Cursor IO h+ -> IO (V.Vector (SerialisedKey,+ (SerialisedValue,+ Maybe (WeakBlobRef IO h))))+readCursorUntil ub cursor = go V.empty+ where+ chunkSize = 50+ toResult k v b = (coerce k, (v, b))++ go !acc = do+ res <- case ub of+ Nothing -> readCursor resolve chunkSize cursor toResult+ Just k -> readCursorWhile resolve (<= coerce k) chunkSize cursor toResult+ if V.length res < chunkSize then pure (acc <> res)+ else go (acc <> res)++resolve :: ResolveSerialisedValue+resolve (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)
@@ -0,0 +1,36 @@+{-# LANGUAGE CPP #-}+module Test.Database.LSMTree.Internal.Arena (+ tests,+) where++import Control.Monad.ST (runST)+import Data.Primitive.ByteArray+import Data.Word (Word8)+import Database.LSMTree.Internal.Arena+import GHC.Exts (toList)+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (testCaseSteps, (@?=))++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Arena"+ [ testCaseSteps "safe" $ \_info -> do+ let !ba = runST $ withUnmanagedArena $ \arena -> do+ (off', mba) <- allocateFromArena arena 32 8+ setByteArray mba off' 32 (1 :: Word8)+ freezeByteArray mba off' 32++ toList ba @?= replicate 32 (1 :: Word8)++ , testCaseSteps "unsafe" $ \_info -> do+ let !(off, ba) = runST $ withUnmanagedArena $ \arena -> do+ (off', mba) <- allocateFromArena arena 32 8+ setByteArray mba off' 32 (1 :: Word8)+ ba' <- unsafeFreezeByteArray mba+ pure (off', ba')++#if NO_IGNORE_ASSERTS+ take 32 (drop off (toList ba)) @?= replicate 32 (0x77 :: Word8)+#else+ take 32 (drop off (toList ba)) @?= replicate 32 (1 :: Word8)+#endif+ ]
@@ -0,0 +1,67 @@+module Test.Database.LSMTree.Internal.BlobFile.FS (tests) where++import Control.Concurrent.Class.MonadSTM.Strict+import Control.Monad+import Control.Monad.Class.MonadThrow+import Control.RefCount+import Database.LSMTree.Internal.BlobFile+import System.FS.API+import System.FS.Sim.Error hiding (genErrors)+import qualified System.FS.Sim.MockFS as MockFS+import Test.Tasty+import Test.Tasty.QuickCheck as QC+import Test.Util.FS++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.BlobFile.FS" [+ testProperty "prop_fault_openRelease" prop_fault_openRelease+ ]++-- Test that opening and releasing a blob file properly cleans handles and files+-- in the presence of disk faults.+prop_fault_openRelease ::+ Bool -- ^ create the file or not+ -> OpenMode+ -> NoCleanupErrors+ -> NoCleanupErrors+ -> Property+prop_fault_openRelease doCreateFile om+ (NoCleanupErrors openErrors)+ (NoCleanupErrors releaseErrors) =+ ioProperty $+ withSimErrorHasFS propPost MockFS.empty emptyErrors $ \hfs fsVar errsVar -> do+ when doCreateFile $+ withFile hfs path (WriteMode MustBeNew) $ \_ -> pure ()+ eith <- try @_ @FsError $+ bracket (acquire hfs errsVar) (release errsVar) $ \_blobFile -> do+ fs' <- atomically $ readTMVar fsVar+ pure $ propNumOpenHandles 1 fs' .&&. propNumDirEntries (mkFsPath []) 1 fs'+ pure $ case eith of+ Left{} ->+ label "FsError" $ property True+ Right prop ->+ label "Success" $ prop+ where+ root = mkFsPath []+ path = mkFsPath ["blobfile"]++ acquire hfs errsVar =+ withErrors errsVar openErrors $ openBlobFile hfs path om++ release errsVar blobFile =+ withErrors errsVar releaseErrors $ releaseRef blobFile++ propPost fs = propNoOpenHandles fs .&&.+ if doCreateFile then+ case allowExisting om of+ AllowExisting ->+ -- TODO: fix, see the TODO on openBlobFile+ propNoDirEntries root fs .||. propNumDirEntries root 1 fs+ MustBeNew ->+ propNumDirEntries root 1 fs+ MustExist ->+ -- TODO: fix, see the TODO on openBlobFile+ propNoDirEntries root fs .||. propNumDirEntries root 1 fs+ else+ propNoDirEntries root fs+
@@ -0,0 +1,170 @@+module Test.Database.LSMTree.Internal.BloomFilter (tests) where++import Control.DeepSeq (deepseq)+import Control.Exception (Exception (..), displayException)+import Control.Monad (void)+import qualified Control.Monad.IOSim as IOSim+import Data.Bits ((.&.))+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BS.Builder+import qualified Data.ByteString.Builder.Extra as BS.Builder+import qualified Data.ByteString.Lazy as LBS+import qualified Data.List as List+import qualified Data.Set as Set+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import Data.Word (Word32, Word64, byteSwap32)+import qualified System.FS.API as FS+import qualified System.FS.API.Strict as FS+import qualified System.FS.Sim.MockFS as MockFS+import qualified System.FS.Sim.STM as FSSim++import Test.QuickCheck.Gen (genDouble)+import Test.QuickCheck.Instances ()+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck hiding ((.&.))++import qualified Data.BloomFilter.Blocked as Bloom+import Database.LSMTree.Internal.BloomFilter+import Database.LSMTree.Internal.CRC32C (FileCorruptedError (..),+ FileFormat (..))+import Database.LSMTree.Internal.Serialise (SerialisedKey,+ serialiseKey)++--TODO: add a golden test for the BloomFilter format vs the 'formatVersion'+-- to ensure we don't change the format without conciously bumping the version.+tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.BloomFilter"+ [ testProperty "roundtrip" roundtrip_prop+ -- a specific case: 300 bits is just under 5x 64 bit words+ , testProperty "roundtrip-3-300" $ roundtrip_prop (Positive (Small 3)) (Positive 300)+ , testProperty "total-deserialisation" $ withMaxSuccess 10000 $+ prop_total_deserialisation+ , testProperty "total-deserialisation-whitebox" $ withMaxSuccess 10000 $+ prop_total_deserialisation_whitebox+ , testProperty "bloomQueries (bulk)" $+ prop_bloomQueries+ ]++testSalt :: Bloom.Salt+testSalt = 4++roundtrip_prop :: Positive (Small Int) -> Positive Int -> [Word64] -> Property+roundtrip_prop (Positive (Small hfN)) (Positive bits) ws =+ counterexample (show bs) $+ case bloomFilterFromBS bs of+ Left err -> counterexample (displayException err) $ property False+ Right rhs -> lhs === rhs+ where+ sz = Bloom.BloomSize { sizeBits = limitBits bits,+ sizeHashes = hfN }+ lhs = Bloom.create sz testSalt (\b -> mapM_ (Bloom.insert b) ws)+ bs = LBS.toStrict (bloomFilterToLBS lhs)++limitBits :: Int -> Int+limitBits b = b .&. 0xffffff++prop_total_deserialisation :: BS.ByteString -> Property+prop_total_deserialisation bs =+ case bloomFilterFromBS bs of+ Left err ->+ label (mkLabel err) $ property True+ Right bf -> label "parsed successfully" $ property $+ bf `deepseq` True+ where+ mkLabel err = case err of+ IOSim.FailureException e+ | Just (ErrFileFormatInvalid fsep FormatBloomFilterFile msg) <- fromException e+ , let msg' = "Expected salt does not match actual salt"+ , msg' `List.isPrefixOf` msg+ -> displayException $ ErrFileFormatInvalid fsep FormatBloomFilterFile msg'+ _ -> displayException err++-- | Write the bytestring to a file in the mock file system and then use+-- 'bloomFilterFromFile'.+bloomFilterFromBS :: BS.ByteString -> Either IOSim.Failure (Bloom a)+bloomFilterFromBS bs =+ IOSim.runSim $ do+ hfs <- FSSim.simHasFS' MockFS.empty+ let file = FS.mkFsPath ["filter"]+ -- write the bytestring+ FS.withFile hfs file (FS.WriteMode FS.MustBeNew) $ \h -> do+ void $ FS.hPutAllStrict hfs h bs+ -- deserialise from file+ FS.withFile hfs file FS.ReadMode $ \h ->+ bloomFilterFromFile hfs testSalt h++-- Length is in bits. A large length would require significant amount of+-- memory, so we make it 'Small'.+prop_total_deserialisation_whitebox :: Word32 -> Small Word64 -> Property+prop_total_deserialisation_whitebox hsn (Small nbits) =+ forAll genBytes $ \bytes ->+ forAll genVersion $ \version ->+ prop_total_deserialisation (prefix version <> BS.pack bytes)+ where+ genBytes = do n <- choose (-1,1)+ vector (nbytes' + n)+ nbytes = (fromIntegral nbits+7) `div` 8+ nbytes' = ((nbytes + 63) `div` 64) * 64 -- rounded to nearest 64 bytes+ genVersion = frequency [+ (6, pure bloomFilterVersion),+ (1, pure (byteSwap32 bloomFilterVersion)),+ (1, arbitrary)+ ]+ prefix version =+ LBS.toStrict $ BS.Builder.toLazyByteString $+ BS.Builder.word32Host version+ <> BS.Builder.word32Host hsn+ <> BS.Builder.word64Host nbits++newtype FPR = FPR Double deriving stock Show++instance Arbitrary FPR where+ arbitrary =+ FPR <$> frequency+ [ (1, pure 0.999)+ , (9, (fmap (/2) genDouble) `suchThat` \fpr -> fpr > 0) ]++prop_bloomQueries :: FPR+ -> [[Small Word64]]+ -> [Small Word64]+ -> Property+prop_bloomQueries (FPR fpr) filters keys =+ let filters' :: [Bloom SerialisedKey]+ filters' = map (Bloom.fromList (Bloom.policyForFPR fpr) testSalt+ . map (\(Small k) -> serialiseKey k))+ filters++ keys' :: [SerialisedKey]+ keys' = map (\(Small k) -> serialiseKey k) keys++ referenceResults :: [(Int, Int)]+ referenceResults =+ [ (f_i, k_i)+ | (k, k_i) <- zip keys' [0..]+ , (f, f_i) <- zip filters' [0..]+ , Bloom.elem k f+ ]++ filterSets = map (Set.fromList . map (\(Small k) -> serialiseKey k)) filters+ referenceCmp =+ [ (Bloom.elem k f, k `Set.member` f')+ | (f, f') <- zip filters' filterSets+ , k <- keys'+ ]+ truePositives = [ "true positives" | (True, True) <- referenceCmp ]+ falsePositives = [ "false positives" | (True, False) <- referenceCmp ]+ trueNegatives = [ "true negatives" | (False, False) <- referenceCmp ]+ falseNegatives = [ "false negatives" | (False, True) <- referenceCmp ]+ distribution = truePositives ++ falsePositives+ ++ trueNegatives ++ falseNegatives++ -- To get coverage of bloomQueries array resizing we want some+ -- cases with high FPRs.+ in tabulate "FPR" [show (round (fpr * 10) * 10 :: Int) ++ "%"] $+ coverTable "FPR" [("100%", 5)] $+ tabulate "distribution of true/false positives/negatives" distribution $+ referenceResults+ ===+ map (\(RunIxKeyIx rix kix) -> (rix, kix))+ (VP.toList (bloomQueries testSalt (V.fromList filters') (V.fromList keys')))
@@ -0,0 +1,172 @@+{-# LANGUAGE NamedFieldPuns #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.CRC32C (+ -- * Main test tree+ tests+ ) where++import Control.Monad+import qualified Control.Monad.IOSim as IOSim+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as BSL+import qualified Data.ByteString.Short as SBS+import qualified Data.Digest.CRC32C+import qualified Data.Foldable as Fold+import Database.LSMTree.Extras (showPowersOf)+import Database.LSMTree.Internal.CRC32C+import System.FS.API+import System.FS.API.Lazy+import System.FS.API.Strict+import qualified System.FS.Sim.Error as FsSim+import qualified System.FS.Sim.MockFS as FsSim+import qualified System.FS.Sim.Stream as FsSim.Stream+import System.Posix.Types+import Test.QuickCheck+import Test.QuickCheck.Instances ()+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)+import Test.Util.FS (withTempIOHasFS)+++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.CRC32C" [+ testProperty "incremental" prop_incremental+ , testProperty "all splits" prop_splits+ , testProperty "put/get" prop_putGet+ , testProperty "write/read" prop_writeRead+ , testProperty "prop_hGetExactlyCRC32C_SBS" prop_hGetExactlyCRC32C_SBS+ , testProperty "prop_hGetAllCRC32C'" prop_hGetAllCRC32C'+ ]++{-------------------------------------------------------------------------------+ Properties+-------------------------------------------------------------------------------}++prop_incremental :: [BS.ByteString] -> Bool+prop_incremental bss =+ Fold.foldl' (flip updateCRC32C) initialCRC32C bss+ == CRC32C (Data.Digest.CRC32C.crc32c (BS.concat bss))++prop_splits :: BS.ByteString -> Bool+prop_splits bs =+ and [ updateCRC32C b2 (updateCRC32C b1 initialCRC32C) == expected+ | let expected = CRC32C (Data.Digest.CRC32C.crc32c bs)+ , n <- [0 .. BS.length bs]+ , let (b1, b2) = BS.splitAt n bs ]++prop_putGet :: PartialIOErrors -> [BS.ByteString] -> Bool+prop_putGet (PartialIOErrors errs) bss =+ fst $ IOSim.runSimOrThrow $+ FsSim.runSimErrorFS FsSim.empty errs $ \_ fs -> do+ let path = mkFsPath ["foo"]+ crc1 <- withFile fs path (WriteMode MustBeNew) $ \h ->+ foldM (\crc bs -> snd <$> hPutAllCRC32C fs h bs crc)+ initialCRC32C bss+ crc2 <- withFile fs path ReadMode $ \h ->+ let go crc = do+ (bs,crc') <- hGetSomeCRC32C fs h 42 crc+ if BS.null bs then pure crc' else go crc'+ in go initialCRC32C+ crc3 <- readFileCRC32C fs path+ let expected = CRC32C (Data.Digest.CRC32C.crc32c (BS.concat bss))+ pure (all (expected==) [crc1, crc2, crc3])++-- TODO: test like prop_putGet but with put corruption, to detect that the+-- crc doesn't match, to detect the corruption.++prop_writeRead :: PartialIOErrors -> ChecksumsFile -> Bool+prop_writeRead (PartialIOErrors errs) chks =+ fst $ IOSim.runSimOrThrow $+ FsSim.runSimErrorFS FsSim.empty errs $ \_ fs -> do+ let path = mkFsPath ["foo.checksums"]+ writeChecksumsFile fs path chks+ chks' <- readChecksumsFile fs path+ pure (chks == chks')++instance Arbitrary CRC32C where+ arbitrary = CRC32C <$> arbitraryBoundedIntegral++instance Arbitrary ChecksumsFileName where+ arbitrary = ChecksumsFileName <$> name+ where+ name = fmap BS.pack (listOf nameChar)+ nameChar = arbitrary `suchThat` \w8 ->+ let c = toEnum (fromIntegral w8)+ in c /= '(' && c /= ')' && c /= '\n'++-- | File system simulated errors for partial reads and writes. No other+-- silent corruption or noisy exceptions. Used to test that handling of+-- partial read\/writes is robust.+--+newtype PartialIOErrors = PartialIOErrors FsSim.Errors+ deriving stock Show++instance Arbitrary PartialIOErrors where+ arbitrary = do+ hGetSomeE <- genPartialErrorStream+ hPutSomeE <- genPartialErrorStream+ pure $ PartialIOErrors FsSim.emptyErrors {+ FsSim.hGetSomeE,+ FsSim.hPutSomeE+ }++ shrink (PartialIOErrors errs) = map PartialIOErrors (shrink errs)++genPartialErrorStream :: Gen (FsSim.Stream.Stream (Either a FsSim.Partial))+genPartialErrorStream =+ FsSim.Stream.genInfinite+ . FsSim.Stream.genMaybe 2 1 -- 2:1 normal:partial+ . fmap Right+ $ arbitrary++{-------------------------------------------------------------------------------+ Properties for functions with user-supplied buffers+-------------------------------------------------------------------------------}++newtype BS = BS BS.ByteString+ deriving stock (Show, Eq)++instance Arbitrary BS where+ arbitrary = BS . BS.pack <$> scale (2*) arbitrary+ shrink = fmap (BS . BS.pack) . shrink . BS.unpack . unBS+ where unBS (BS x) = x++prop_hGetExactlyCRC32C_SBS :: BS -> Property+prop_hGetExactlyCRC32C_SBS (BS bs) =+ ioProperty $ withTempIOHasFS "tempDir" $ \hfs -> do+ withFile hfs (mkFsPath ["temp.file"]) (WriteMode MustBeNew) $ \h ->+ void $ hPutAllStrict hfs h bs+ x <- withFile hfs (mkFsPath ["temp.file"]) ReadMode $ \h ->+ simpl hfs h initialCRC32C+ y <- withFile hfs (mkFsPath ["temp.file"]) ReadMode $ \h ->+ hGetExactlyCRC32C_SBS hfs h (fromIntegral $ BS.length bs) initialCRC32C++ pure $ x === y+ where+ simpl fs h crc = do+ lbs <- hGetAll fs h+ let !crc' = BSL.foldlChunks (flip updateCRC32C) crc lbs+ pure (SBS.toShort (BS.toStrict lbs), crc')+++prop_hGetAllCRC32C':: Positive (Small ByteCount) -> BS -> Property+prop_hGetAllCRC32C' (Positive (Small chunkSize)) (BS bs) =+ ioProperty $ withTempIOHasFS "tempDir" $ \hfs -> do+ withFile hfs (mkFsPath ["temp.file"]) (WriteMode MustBeNew) $ \h ->+ void $ hPutAllStrict hfs h bs+ x <- withFile hfs (mkFsPath ["temp.file"]) ReadMode $ \h ->+ simpl hfs h initialCRC32C+ y <- withFile hfs (mkFsPath ["temp.file"]) ReadMode $ \h ->+ hGetAllCRC32C' hfs h (ChunkSize chunkSize) initialCRC32C++ pure+ $ tabulate "number of chunks"+ [ showPowersOf 2 $ ceiling @Double $+ fromIntegral (BS.length bs) / (fromIntegral chunkSize) ]+ $ x === y+ where+ simpl fs h crc = do+ lbs <- hGetAll fs h+ let !crc' = BSL.foldlChunks (flip updateCRC32C) crc lbs+ pure crc'
@@ -0,0 +1,168 @@+module Test.Database.LSMTree.Internal.Chunk (tests) where++import Prelude hiding (concat, drop, length)++import Control.Category ((>>>))+import Control.Monad.ST.Strict (runST)+import qualified Data.List as List (concat, drop, length)+import Data.Maybe (catMaybes, fromJust, isJust, isNothing)+import Data.Vector.Primitive (Vector, concat, length)+import Data.Word (Word8)+import Database.LSMTree.Extras.Generators ()+ -- for @Arbitrary@ instantiation of @Vector@+import Database.LSMTree.Internal.Chunk (Chunk, createBaler, feedBaler,+ toByteVector, unsafeEndBaler)+import Test.QuickCheck (Arbitrary (arbitrary, shrink),+ NonEmptyList (NonEmpty), Positive (Positive, getPositive),+ Property, Small (Small, getSmall), Testable, scale,+ shrinkMap, tabulate, (===), (==>))+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)++-- * Tests++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Chunk" $+ [+ testProperty "Content is preserved"+ prop_contentIsPreserved,+ testProperty "No remnant after output"+ prop_noRemnantAfterOutput,+ testProperty "Common chunks are large"+ prop_commonChunksAreLarge,+ testProperty "Remnant chunk is non-empty"+ prop_remnantChunkIsNonEmpty,+ testProperty "Remnant chunk is small"+ prop_remnantChunkIsSmall+ ]++-- * Properties to test++{-|+ Feeds a freshly created baler a sequence of data portions and ends it+ afterwards, yielding all output.+-}+balingOutput :: Int -- Minimum chunk size+ -> [[Vector Word8]] -- Data portions to be fed+ -> ([Maybe Chunk], Maybe Chunk) -- Feeding output and remnant+balingOutput minChunkSize food = runST $ do+ baler <- createBaler minChunkSize+ commonChunks <- mapM (flip feedBaler baler) food+ remnant <- unsafeEndBaler baler+ pure (commonChunks, remnant)++{-|+ Supplies the output of a complete baler run for constructing a property.++ The resulting property additionally provides statistics about the lengths of+ buildup phases, where a buildup phase is a sequence of feedings that does+ not result in chunks and is followed by an ultimate chunk production, which+ happens either due to another feeding or due to the baler run ending and+ producing a remnant chunk.+-}+withBalingOutput ::+ Testable prop+ => Int -- Minimum chunk size+ -> [[Vector Word8]] -- Data portions to be fed+ -> ([Maybe Chunk] -> Maybe Chunk -> prop) -- Property from baler output+ -> Property -- Resulting property+withBalingOutput minChunkSize food consumer+ = tabulate "Lengths of buildup phases"+ (map show (buildupPhasesLengths commonChunks))+ (consumer commonChunks remnant)+ where++ commonChunks :: [Maybe Chunk]+ remnant :: Maybe Chunk+ (commonChunks, remnant) = balingOutput minChunkSize food++ buildupPhasesLengths :: [Maybe Chunk] -> [Int]+ buildupPhasesLengths [] = []+ buildupPhasesLengths chunks = List.length buildupOutput :+ buildupPhasesLengths (List.drop 1 followUp)+ where++ buildupOutput, followUp :: [Maybe Chunk]+ (buildupOutput, followUp) = span isNothing chunks++prop_contentIsPreserved :: MinChunkSize -> [[Vector Word8]] -> Property+prop_contentIsPreserved (MinChunkSize minChunkSize) food+ = withBalingOutput minChunkSize food $ \ commonChunks remnant ->+ let++ input :: Vector Word8+ input = concat (List.concat food)++ output :: Vector Word8+ output = concat $+ toByteVector <$> catMaybes (commonChunks ++ [remnant])++ in input === output++prop_noRemnantAfterOutput :: MinChunkSize+ -> NonEmptyList [Vector Word8]+ -> Property+prop_noRemnantAfterOutput (MinChunkSize minChunkSize) (NonEmpty food)+ = withBalingOutput minChunkSize food $ \ commonChunks remnant ->+ isJust (last commonChunks) ==> isNothing remnant++prop_commonChunksAreLarge :: MinChunkSize -> [[Vector Word8]] -> Property+prop_commonChunksAreLarge (MinChunkSize minChunkSize) food+ = withBalingOutput minChunkSize food $ \ commonChunks _ ->+ all (toByteVector >>> length >>> (>= minChunkSize)) $+ catMaybes commonChunks++remnantChunkSizeIs :: (Int -> Bool) -> Int -> [[Vector Word8]] -> Property+remnantChunkSizeIs constraint minChunkSize food+ = withBalingOutput minChunkSize food $ \ _ remnant ->+ isJust remnant ==> constraint (length (toByteVector (fromJust remnant)))++prop_remnantChunkIsNonEmpty :: MinChunkSize -> [[Vector Word8]] -> Property+prop_remnantChunkIsNonEmpty (MinChunkSize minChunkSize)+ = remnantChunkSizeIs (> 0) minChunkSize++prop_remnantChunkIsSmall :: MinChunkSize -> [[Vector Word8]] -> Property+prop_remnantChunkIsSmall (MinChunkSize minChunkSize)+ = remnantChunkSizeIs (< minChunkSize) minChunkSize++-- * Test case generation and shrinking++{-|+ The type of minimum chunk sizes.++ This type is isomorphic to 'Int' but has a different way of generating test+ cases. Only small, positive integers are generated, and they are generated+ using \(2 \cdot s^{2}\) as the size parameter, where \(s\) refers to the+ original size parameter.++ The reasons for the modification of the size parameter in the+ above-mentioned way are somewhat subtle.++ First, we want the ratio between the average minimum chunk size and the+ average size of a data portion that we feed to a baler to be independent of+ the size parameter. Each data portion is a list of primitive vectors of+ bytes, and arbitrarily generated lists and byte vectors have lengths that+ are small, positive integers. Such integers are \(s/2\) on average. As a+ result, the average size of data fed to a baler is \(s^{2}/4\). By+ generating minimum chunk sizes with \(a \cdot s^{2}\) as the size parameter+ for some constant \(a\), the average minimum chunk size is \(a/2 \cdot+ s^{2}\) and therefore $2a$ times the average size of a data portion fed,+ independently of \(s\).++ Second, we want prompt chunk generation as well as chunk generation after+ only two or more feedings to occur reasonably often. To achieve this to some+ degree, we can tune the parameter \(a\). It appears that \(a\) being \(2\)+ leads to reasonable results.+-}+newtype MinChunkSize = MinChunkSize Int deriving stock Show++fromMinChunkSize :: MinChunkSize -> Int+fromMinChunkSize (MinChunkSize minChunkSize) = minChunkSize++instance Arbitrary MinChunkSize where++ arbitrary = scale (\ size -> 2 * size ^ (2 :: Int)) $+ MinChunkSize <$> getSmall <$> getPositive <$> arbitrary++ shrink = shrinkMap (MinChunkSize . getSmall . getPositive)+ (Positive . Small . fromMinChunkSize)
@@ -0,0 +1,221 @@+{-# LANGUAGE DerivingVia #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module Test.Database.LSMTree.Internal.Entry (tests) where++import Data.List.NonEmpty (NonEmpty)+import Data.Semigroup hiding (First)+import qualified Data.Semigroup as S+import qualified Database.LSMTree as Full+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.Entry+import Test.QuickCheck+import Test.QuickCheck.Classes (semigroupLaws)+import Test.Tasty+import Test.Tasty.QuickCheck (QuickCheckMaxSize (QuickCheckMaxSize),+ QuickCheckTests (QuickCheckTests), testProperty)+import Test.Util.QC++tests :: TestTree+tests = adjustOption (\_ -> QuickCheckTests 10000) $+ adjustOption (\_ -> QuickCheckMaxSize 1000) $+ testGroup "Test.Database.LSMTree.Internal.Entry" [+ -- * Class laws++ testClassLaws "Regular Entry" $+ semigroupLaws (Proxy @(Regular (Entry (Sum Int) String)))+ , testClassLaws "Regular Update" $+ semigroupLaws (Proxy @(Regular (Full.Update (Sum Int) String)))++ , testClassLaws "Union Entry" $+ semigroupLaws (Proxy @(Union (Entry (Sum Int) String)))+ , testClassLaws "Union Update" $+ semigroupLaws (Proxy @(Union (Full.Update (Sum Int) String)))++ -- * Semantics++ , testProperty "prop_regularSemantics" $+ prop_regularSemantics @(Sum Int) @String++ , testProperty "prop_unionSemantics" $+ prop_unionSemantics @(Sum Int) @String+ ]++-- TODO: it would be nice to write down how the semantic tests below relate to+-- the semantics of operations on the public API.++-- | @sconcat == fromEntry . sconcat . toEntry@ with regular semantics.+prop_regularSemantics ::+ (Show v, Show b, Eq v, Eq b, Semigroup v)+ => NonEmpty (Full.Update v b)+ -> Property+prop_regularSemantics es = expected === real+ where+ expected = from . sconcat . fmap to $ es+ where+ to :: Full.Update v b -> Regular (Full.Update v b)+ to = Regular++ from :: Regular (Full.Update v b) -> Full.Update v b+ from = unRegular++ real = from . sconcat . fmap to $ es+ where+ to :: Full.Update v b -> Regular (Entry v b)+ to = Regular . updateToEntry++ from :: Regular (Entry v b) -> Full.Update v b+ from = entryToUpdate . unRegular++-- | @sconcat == fromEntry . sconcat . toEntry@ with union semantics.+prop_unionSemantics ::+ (Show v, Show b, Eq v, Eq b, Semigroup v)+ => NonEmpty (Full.Update v b)+ -> Property+prop_unionSemantics es = expected === real+ where+ expected = from . sconcat . fmap to $ es+ where+ to :: Full.Update v b -> Union (Full.Update v b)+ to = Union++ from :: Union (Full.Update v b) -> Full.Update v b+ from = unUnion++ real = from . sconcat . fmap to $ es+ where+ to :: Full.Update v b -> Union (Entry v b)+ to = Union . updateToEntry++ from :: Union (Entry v b) -> Full.Update v b+ from = entryToUpdate . unUnion++{-------------------------------------------------------------------------------+ Regular semantics+-------------------------------------------------------------------------------}++newtype Regular a = Regular { unRegular :: a }+ deriving stock (Show, Eq)++--+-- Update+--++deriving newtype instance (Arbitrary v, Arbitrary b)+ => Arbitrary (Regular (Full.Update v b))++instance Semigroup v => Semigroup (Regular (Full.Update v b)) where+ Regular up1 <> Regular up2 = Regular $ case (up1, up2) of+ (Full.Delete , _ ) -> up1+ (Full.Insert{} , _ ) -> up1+ (Full.Upsert v1 , Full.Delete ) -> Full.Insert v1 Nothing+ (Full.Upsert v1 , Full.Insert v2 _) -> Full.Insert (v1 <> v2) Nothing+ (Full.Upsert v1 , Full.Upsert v2 ) -> Full.Upsert (v1 <> v2)++--+-- Entry+--++deriving via Uniform (Entry v b)+ instance (Arbitrary v, Arbitrary b) => Arbitrary (Regular (Entry v b))++-- | Semigroup instance using 'combine'.+instance Semigroup v => Semigroup (Regular (Entry v b)) where+ Regular e1 <> Regular e2 = Regular $ combine (<>) e1 e2++{-------------------------------------------------------------------------------+ Union semantics+-------------------------------------------------------------------------------}++newtype Union a = Union { unUnion :: a }+ deriving stock (Show, Eq)++--+-- Update+--++deriving newtype instance (Arbitrary v, Arbitrary b)+ => Arbitrary (Union (Full.Update v b))++instance Semigroup v => Semigroup (Union (Full.Update v b)) where+ Union up1 <> Union up2 = Union $ fromModel $ toModel up1 <> toModel up2+ where+ toModel :: Full.Update v b -> Maybe (v, S.First (Maybe b))+ toModel Full.Delete = Nothing+ toModel (Full.Insert v mb) = Just (v, S.First mb)+ toModel (Full.Upsert v) = Just (v, S.First Nothing)++ fromModel :: Maybe (v, S.First (Maybe b)) -> Full.Update v b+ fromModel Nothing = Full.Delete+ fromModel (Just (v, S.First mb)) = Full.Insert v mb++--+-- Entry+--++deriving via Uniform (Entry v b)+ instance (Arbitrary v, Arbitrary b) => Arbitrary (Union (Entry v b))++-- | Semigroup instance using 'combineUnion'.+instance Semigroup v => Semigroup (Union (Entry v b)) where+ Union e1 <> Union e2 = Union $ combineUnion (<>) e1 e2++{-------------------------------------------------------------------------------+ Utility+-------------------------------------------------------------------------------}++-- | A wrapper type with a 'Semigroup' instance that always throws an error.+newtype Unlawful a = Unlawful a+ deriving stock (Show, Eq)+ deriving newtype Arbitrary++-- | A 'Semigroup' instance that always throws an error.+instance Semigroup (Unlawful a) where+ _ <> _ = error "unlawful"++newtype Uniform a = Uniform a++-- | We do not use the @'Arbitrary' 'Entry'@ instance here, because we want to+-- generate each constructor with equal probability.+instance (Arbitrary v, Arbitrary b) => Arbitrary (Uniform (Entry v b)) where+ arbitrary = Uniform <$> liftArbitrary2Entry arbitrary arbitrary+ shrink (Uniform x) = Uniform <$> liftShrink2Entry shrink shrink x++liftArbitrary2Entry :: Gen v -> Gen b -> Gen (Entry v b)+liftArbitrary2Entry genVal genBlob = frequency+ [ (1, Insert <$> genVal)+ , (1, InsertWithBlob <$> genVal <*> genBlob)+ , (1, Upsert <$> genVal)+ , (1, pure Delete)+ ]++liftShrink2Entry :: (v -> [v]) -> (b -> [b]) -> Entry v b -> [Entry v b]+liftShrink2Entry shrinkVal shrinkBlob = \case+ Insert v -> Delete : (Insert <$> shrinkVal v)+ InsertWithBlob v b -> [Delete, Insert v]+ ++ [ InsertWithBlob v' b'+ | (v', b') <- liftShrink2 shrinkVal shrinkBlob (v, b)+ ]+ Upsert v -> Delete : Insert v : (Upsert <$> shrinkVal v)+ Delete -> []++{-------------------------------------------------------------------------------+ Injections/projections+-------------------------------------------------------------------------------}++updateToEntry :: Full.Update v b -> Entry v b+updateToEntry = \case+ Full.Insert v Nothing -> Insert v+ Full.Insert v (Just b) -> InsertWithBlob v b+ Full.Upsert v -> Upsert v+ Full.Delete -> Delete++entryToUpdate :: Entry v b -> Full.Update v b+entryToUpdate = \case+ Insert v -> Full.Insert v Nothing+ InsertWithBlob v b -> Full.Insert v (Just b)+ Upsert v -> Full.Upsert v+ Delete -> Full.Delete
@@ -0,0 +1,492 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE RecordWildCards #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.Index.Compact (tests) where++import Control.DeepSeq (deepseq)+import Control.Monad (foldM)+import Control.Monad.ST (runST)+import Control.Monad.State.Strict (MonadState (..), State, evalState,+ get, put)+import Data.Bit (Bit (..))+import qualified Data.Bit as BV+import qualified Data.ByteString.Lazy as LBS+import qualified Data.ByteString.Short as SBS+import Data.Coerce (coerce)+import Data.Foldable (Foldable (..))+import Data.List.Split (chunksOf)+import qualified Data.Map.Merge.Strict as Merge+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Primitive.ByteArray (ByteArray (..), byteArrayFromList,+ sizeofByteArray)+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Base as VU+import Data.Word+import Database.LSMTree.Extras+import Database.LSMTree.Extras.Generators (ChunkSize (..),+ LogicalPageSummaries, LogicalPageSummary (..), Pages (..),+ genRawBytes, labelPages, toAppends)+import Database.LSMTree.Extras.Index (Append (..), appendToCompact)+import Database.LSMTree.Internal.BitMath+import Database.LSMTree.Internal.Chunk as Chunk (toByteString)+import Database.LSMTree.Internal.Entry (NumEntries (..))+import Database.LSMTree.Internal.Index.Compact+import Database.LSMTree.Internal.Index.CompactAcc+import Database.LSMTree.Internal.Page (PageNo (PageNo), PageSpan,+ multiPage, singlePage)+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise+import Numeric (showHex)+import Prelude hiding (max, min, pi)+import qualified Test.QuickCheck as QC+import Test.QuickCheck+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (assertEqual, testCase)+import Test.Tasty.QuickCheck (testProperty)+import Test.Util.Arbitrary (noTags,+ prop_arbitraryAndShrinkPreserveInvariant)+import Test.Util.Orphans ()+import Text.Printf (printf)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Index.Compact" [+ testProperty "prop_distribution @TestKey" $+ prop_distribution @TestKey+ , testProperty "prop_searchMinMaxKeysAfterConstruction" $+ prop_searchMinMaxKeysAfterConstruction @TestKey 100+ , testProperty "prop_differentChunkSizesSameResults" $+ prop_differentChunkSizesSameResults @TestKey+ , testProperty "prop_singlesEquivMulti" $+ prop_singlesEquivMulti @TestKey+ , testGroup "(De)serialisation" [+ testGroup "Chunks generator" $+ prop_arbitraryAndShrinkPreserveInvariant noTags chunksInvariant++ , testCase "index-2-clash" $ do+ let k1 = SerialisedKey' (VP.replicate 16 0x00)+ let k2 = SerialisedKey' (VP.replicate 16 0x11)+ let k3 = SerialisedKey' (VP.replicate 15 0x11 <> VP.replicate 1 0x12)+ let (chunks, index) = runST $ do+ ica <- new 16+ ch1 <- flip appendToCompact ica $ AppendSinglePage k1 k2+ ch2 <- flip appendToCompact ica $ AppendSinglePage k3 k3+ (mCh3, idx) <- unsafeEnd ica+ pure (ch1 <> ch2 <> toList mCh3, idx)++ let expectedVersion :: [Word8]+ expectedVersion = word32toBytesLE 0x0000_0001 <> word32toBytesLE 0x0000_0000+ let expectedPrimary :: [Word8]+ expectedPrimary = foldMap word64toBytesLE+ -- 1. primary array: two pages (64 bits LE)+ [ 0x0000_0000_0000_0000, 0x1111_1111_1111_1111+ ]+ let expectedRest :: [Word8]+ expectedRest = foldMap word32toBytesLE+ [ -- 3. clash indicator: two pages, second one has bit+ 0x0000_0002, 0+ -- 4. larger-than-page: two pages, no bits+ , 0, 0+ -- 5. clash map: maps k3 to page 1+ , 1, 0 {- size = 1 (64 bit LE) -}+ , 1, 16 {- page 1, key size 16 byte -}+ , 0x1111_1111, 0x1111_1111 {- k3 -}+ , 0x1111_1111, 0x1211_1111+ -- 6.1 number of pages in the primary array (64 bits LE)+ , 2, 0+ -- 6.2 number of keys (64bit LE)+ , 7, 0+ ]++ let header = LBS.unpack headerLBS+ let primary = LBS.unpack $+ LBS.fromChunks (map Chunk.toByteString chunks)+ let rest = LBS.unpack (finalLBS (NumEntries 7) index)++ let comparison msg xs ys = unlines $+ (msg <> ":")+ : zipWith (\x y -> x <> " | " <> y)+ (showBytes xs <> repeat (replicate 17 '.'))+ (showBytes ys)++ assertEqual (comparison "header" expectedVersion header)+ expectedVersion header+ assertEqual (comparison "primary" expectedPrimary primary)+ expectedPrimary primary+ assertEqual (comparison "rest" expectedRest rest)+ expectedRest rest++ , testProperty "prop_roundtrip_chunks" $+ prop_roundtrip_chunks+ , testProperty "prop_roundtrip" $+ prop_roundtrip @TestKey+ , testProperty "prop_total_deserialisation" $ withMaxSuccess 10000+ prop_total_deserialisation+ , testProperty "prop_total_deserialisation_whitebox" $ withMaxSuccess 10000+ prop_total_deserialisation_whitebox+ ]+ ]++{-------------------------------------------------------------------------------+ Test key+-------------------------------------------------------------------------------}++-- | Key type for compact index tests+--+-- Tests outside this module don't have to worry about generating clashing keys.+-- We can assume that the compact index handles clashes correctly, because we+-- test this extensively in this module already.+newtype TestKey = TestKey RawBytes+ deriving stock (Show, Eq, Ord)+ deriving newtype SerialiseKey++-- | Generate keys with a non-neglible probability of clashes. This generates+-- sliced keys too.+--+-- Note: recall that keys /clash/ only if their primary bits (first 8 bytes)+-- match. It does not matter whether the other bytes do not match.+instance Arbitrary TestKey where+ arbitrary = do+ -- Generate primary bits from a relatively small distribution. This+ -- ensures that we get clashes between keys with a non-negligible+ -- probability.+ primBits <- do+ lastPrefixByte <- QC.getSmall <$> arbitrary+ pure $ RB.pack ([0,0,0,0,0,0,0] <> [lastPrefixByte])+ -- The rest of the bits after the primary bits can be anything+ restBits <- genRawBytes+ -- The compact index should store keys without retaining unused memory.+ -- Therefore, we generate slices of keys too.+ prefix <- elements [RB.pack [], RB.pack [0]]+ suffix <- elements [RB.pack [], RB.pack [0]]+ -- Combine the bytes and make sure to take out only the slice we need.+ let bytes = prefix <> primBits <> restBits <> suffix+ n = RB.size primBits + RB.size restBits+ bytes' = RB.take n $ RB.drop (RB.size prefix) bytes+ pure $ TestKey bytes'++ -- Shrink keys extensively: most failures will occur in small counterexamples,+ -- so we don't have to limit the number of shrinks as much.+ shrink (TestKey bytes) = [+ testkey'+ | let RawBytes vec = bytes+ , vec' <- VP.fromList <$> shrink (VP.toList vec)+ , let testkey' = TestKey $ RawBytes vec'+ ]+++{-------------------------------------------------------------------------------+ Properties+-------------------------------------------------------------------------------}++--+-- Search+--++type CounterM a = State Int a++evalCounterM :: CounterM a -> Int -> a+evalCounterM = evalState++incrCounter :: CounterM Int+incrCounter = get >>= \c -> put (c+1) >> pure c++plusCounter :: Int -> CounterM Int+plusCounter n = get >>= \c -> put (c+n) >> pure c++-- | After construction, searching for the minimum/maximum key of every page+-- @pageNo@ returns the @pageNo@.+prop_searchMinMaxKeysAfterConstruction ::+ forall k. (SerialiseKey k, Show k, Ord k)+ => ChunkSize+ -> LogicalPageSummaries k+ -> Property+prop_searchMinMaxKeysAfterConstruction csize ps = eqMapProp real model+ where+ model = evalCounterM (foldM modelSearch Map.empty $ getPages ps) 0++ modelSearch :: Map k PageSpan -> LogicalPageSummary k -> CounterM (Map k PageSpan)+ modelSearch m = \case+ OnePageOneKey k -> do+ c <- incrCounter+ pure $ Map.insert k (singlePage (PageNo c)) m+ OnePageManyKeys k1 k2 -> do+ c <- incrCounter+ pure $ Map.insert k1 (singlePage (PageNo c)) $ Map.insert k2 (singlePage (PageNo c)) m+ MultiPageOneKey k n -> do+ let incr = 1 + fromIntegral n+ c <- plusCounter incr+ pure $ if incr == 1 then Map.insert k (singlePage (PageNo c)) m+ else Map.insert k (multiPage (PageNo c) (PageNo $ c + fromIntegral n)) m++ real = foldMap' realSearch (getPages ps)++ ic = fromPageSummaries (coerce csize) ps++ realSearch :: LogicalPageSummary k -> Map k PageSpan+ realSearch = \case+ OnePageOneKey k -> Map.singleton k (search (serialiseKey k) ic)+ OnePageManyKeys k1 k2 -> Map.fromList [ (k1, search (serialiseKey k1) ic)+ , (k2, search (serialiseKey k2) ic)]+ MultiPageOneKey k _ -> Map.singleton k (search (serialiseKey k) ic)++--+-- Construction+--++prop_differentChunkSizesSameResults ::+ SerialiseKey k+ => ChunkSize+ -> ChunkSize+ -> LogicalPageSummaries k+ -> Property+prop_differentChunkSizesSameResults csize1 csize2 ps =+ fromPageSummaries csize1 ps === fromPageSummaries csize2 ps++-- | Constructing an index using only 'appendSingle' is equivalent to using a+-- mix of 'appendSingle' and 'appendMulti'.+prop_singlesEquivMulti ::+ SerialiseKey k+ => ChunkSize+ -> ChunkSize+ -> LogicalPageSummaries k+ -> Property+prop_singlesEquivMulti csize1 csize2 ps = ic1 === ic2+ where+ apps1 = toAppends ps+ apps2 = concatMap toSingles apps1+ ic1 = fromList (coerce csize1) apps1+ ic2 = fromListSingles (coerce csize2) apps2++ toSingles :: Append -> [(SerialisedKey, SerialisedKey)]+ toSingles (AppendSinglePage k1 k2) = [(k1, k2)]+ toSingles (AppendMultiPage k n) = replicate (fromIntegral n + 1) (k, k)++-- | Distribution of generated test data+prop_distribution :: SerialiseKey k => LogicalPageSummaries k -> Property+prop_distribution ps =+ labelPages ps $ labelIndex (fromPageSummaries 100 ps) $ property True++-- Generate and serialise chunks directly.+-- This gives more direct shrinking of individual parts and has fewer invariants+-- (covering a larger space).+prop_roundtrip_chunks :: Chunks -> NumEntries -> Property+prop_roundtrip_chunks (Chunks chunks index) numEntries =+ counterexample (show (SBS.length sbs) <> " bytes") $+ counterexample ("header:\n" <> showBS bsVersion) $+ counterexample ("primary:\n" <> showBS bsPrimary) $+ counterexample ("rest:\n" <> showBS bsRest) $+ Right (numEntries, index) === fromSBS sbs+ where+ bsVersion = headerLBS+ bsPrimary = LBS.fromChunks $+ map (Chunk.toByteString . word64VectorToChunk) chunks+ bsRest = finalLBS numEntries index+ sbs = SBS.toShort (LBS.toStrict (bsVersion <> bsPrimary <> bsRest))++ showBS = unlines . showBytes . LBS.unpack++-- Generate the compact index from logical pages.+prop_roundtrip :: SerialiseKey k => ChunkSize -> LogicalPageSummaries k -> NumEntries -> Property+prop_roundtrip csize ps numEntries =+ counterexample (show (SBS.length sbs) <> " bytes") $+ counterexample ("header:\n" <> showBS bsVersion) $+ counterexample ("primary:\n" <> showBS bsPrimary) $+ counterexample ("rest:\n" <> showBS bsRest) $+ Right (numEntries, index) === fromSBS sbs+ where+ index = fromPageSummaries csize ps++ (bsVersion, bsPrimary, bsRest) = writeIndexCompact numEntries csize ps+ sbs = SBS.toShort (LBS.toStrict (bsVersion <> bsPrimary <> bsRest))++ showBS = unlines . showBytes . LBS.unpack++prop_total_deserialisation :: [Word32] -> Property+prop_total_deserialisation word32s =+ let !(ByteArray ba) = byteArrayFromList word32s+ in case fromSBS (SBS.SBS ba) of+ Left err -> label err $ property True+ Right (numEntries, ic) -> label "parsed successfully" $ property $+ -- Just forcing the index is not enough. The underlying vectors might+ -- point to outside of the byte array, so we check they are valid.+ (numEntries, ic) `deepseq`+ vec64IsValid (icPrimary ic)+ && bitVecIsValid (icClashes ic)+ && bitVecIsValid (icLargerThanPage ic)+ where+ vec64IsValid (VU.V_Word64 (VP.Vector off len ba)) =+ off >= 0 && len >= 0 && mul8 (off + len) <= sizeofByteArray ba+ bitVecIsValid (BV.BitVec off len ba) =+ off >= 0 && len >= 0 && ceilDiv8 (off + len) <= sizeofByteArray ba++prop_total_deserialisation_whitebox :: Small Word16 -> Small Word16 -> [Word32] -> Property+prop_total_deserialisation_whitebox numEntries numPages word32s =+ prop_total_deserialisation $+ [1] -- version+ <> word32s -- primary array, clash bits, LTP bits, clash map+ <> [0 | even (length word32s)] -- padding+ <> [ fromIntegral numPages, 0+ , fromIntegral numEntries, 0+ ]++{-------------------------------------------------------------------------------+ Util+-------------------------------------------------------------------------------}++writeIndexCompact :: SerialiseKey k => NumEntries -> ChunkSize -> LogicalPageSummaries k -> (LBS.ByteString, LBS.ByteString, LBS.ByteString)+writeIndexCompact numEntries (ChunkSize csize) ps = runST $ do+ ica <- new csize+ cs <- mapM (`appendToCompact` ica) (toAppends ps)+ (c, index) <- unsafeEnd ica+ pure+ ( headerLBS+ , LBS.fromChunks $+ foldMap (map Chunk.toByteString) $ cs <> pure (toList c)+ , finalLBS numEntries index+ )++fromPageSummaries :: SerialiseKey k => ChunkSize -> LogicalPageSummaries k -> IndexCompact+fromPageSummaries (ChunkSize csize) ps =+ fromList csize (toAppends ps)++fromList :: Int -> [Append] -> IndexCompact+fromList maxcsize apps = runST $ do+ ica <- new maxcsize+ mapM_ (`appendToCompact` ica) apps+ (_, index) <- unsafeEnd ica+ pure index++-- | One-shot construction using only 'appendSingle'.+fromListSingles :: Int -> [(SerialisedKey, SerialisedKey)] -> IndexCompact+fromListSingles maxcsize apps = runST $ do+ ica <- new maxcsize+ mapM_ (`appendSingle` ica) apps+ (_, index) <- unsafeEnd ica+ pure index++labelIndex :: IndexCompact -> (Property -> Property)+labelIndex ic =+ checkCoverage+ . QC.tabulate "# Clashes" [showPowersOf 2 nclashes]+ . QC.cover 60 (nclashes > 0) "Has clashes"+ . QC.tabulate "# Contiguous clash runs" [showPowersOf 2 (length nscontig)]+ . QC.cover 30 (not (null nscontig)) "Has contiguous clash runs"+ . QC.tabulate "Length of contiguous clash runs" (fmap (showPowersOf 2 . snd) nscontig)+ . QC.tabulate "Contiguous clashes contain multi-page values" (fmap (show . fst) nscontig)+ . QC.cover 3 (any fst nscontig) "Has contiguous clashes that contain multi-page values"+ . QC.cover 0.1 (multiPageValuesClash ic) "Has clashing multi-page values"+ where nclashes = countClashes ic+ nscontig = countContiguousClashes ic++multiPageValuesClash :: IndexCompact -> Bool+multiPageValuesClash ic+ | VU.length (icClashes ic) < 3 = False+ | otherwise = VU.any p $ VU.zip4 v1 v2 v3 v4+ where+ -- note: @i = j - 1@ and @k = j + 1@. This gives us a local view of a+ -- triplet of contiguous LTP bits, and a C bit corresponding to the middle+ -- of the triplet.+ p (cj, ltpi, ltpj, ltpk) =+ -- two multi-page values border each other+ unBit ltpi && not (unBit ltpj) && unBit ltpk+ -- and they clash+ && unBit cj+ v1 = VU.tail (icClashes ic)+ v2 = icLargerThanPage ic+ v3 = VU.tail v2+ v4 = VU.tail v3++-- Returns the number of entries and whether any multi-page values were in the+-- contiguous clashes+countContiguousClashes :: IndexCompact -> [(Bool, Int)]+countContiguousClashes ic = actualContigClashes+ where+ -- filtered is a list of maximal sub-vectors that have only only contiguous+ -- clashes+ zipped = VU.zip (icClashes ic) (icLargerThanPage ic)+ grouped = VU.groupBy (\x y -> fst x == fst y) zipped+ filtered = filter (VU.all (\(c, _ltp) -> c == Bit True)) grouped+ -- clashes that are part of a multi-page value shouldn't be counted towards+ -- the total number of /actual/ clashes. We only care about /actual/ clashes+ -- if they total more than 1 (otherwise it's just a single clash)+ actualContigClashes = filter (\(_, n) -> n > 1) $+ fmap (\v -> (countLTP v > 0, countC v - countLTP v)) filtered+ countC = VU.length+ countLTP = BV.countBits . VU.map snd++-- | Point-wise equality test for two maps, returning counterexamples for each+-- mismatch.+eqMapProp :: (Ord k, Eq v, Show k, Show v) => Map k v -> Map k v -> Property+eqMapProp m1 m2 = conjoin . Map.elems $+ Merge.merge+ (Merge.mapMissing onlyLeft)+ (Merge.mapMissing onlyRight)+ (Merge.zipWithMatched both)+ m1+ m2+ where+ onlyLeft k x = flip counterexample (property False) $+ printf "Key-value pair only on the left, (k, x) = (%s, %s)" (show k) (show x)+ onlyRight k y = flip counterexample (property False) $+ printf "Key-value pair only on the right, (k, y) = (%s, %s)" (show k) (show y)+ both k x y = flip counterexample (property (x == y)) $+ printf "Mismatch between x and y, (k, x, y) = (%s, %s, %s)" (show k) (show x) (show y)++showBytes :: [Word8] -> [String]+showBytes = map (unwords . map (foldMap showByte) . chunksOf 4) . chunksOf 8++showByte :: Word8 -> String+showByte b = let str = showHex b "" in replicate (2 - length str) '0' <> str++word32toBytesLE :: Word32 -> [Word8]+word32toBytesLE = take 4 . map fromIntegral . iterate (`div` 256)++word64toBytesLE :: Word64 -> [Word8]+word64toBytesLE = take 8 . map fromIntegral . iterate (`div` 256)++data Chunks = Chunks [VU.Vector Word64] IndexCompact+ deriving stock (Show)++-- | The concatenated chunks must correspond to the primary array of the index.+-- Apart from that, the only invariant we make sure to uphold is that the length+-- of the vectors match each other, as this is required for correct+-- deserialisation. These invariants do not guarantee that the 'IndexCompact' is+-- valid in other ways (e.g. can successfully be queried).+chunksInvariant :: Chunks -> Bool+chunksInvariant (Chunks chunks IndexCompact {..}) =+ VU.length icPrimary == sum (map VU.length chunks)+ && VU.length icClashes == VU.length icPrimary+ && VU.length icLargerThanPage == VU.length icPrimary++instance Arbitrary Chunks where+ arbitrary = do+ chunks <- map VU.fromList <$> arbitrary+ let icPrimary = mconcat chunks+ let numPages = VU.length icPrimary++ icClashes <- VU.fromList . map Bit <$> vector numPages+ icLargerThanPage <- VU.fromList . map Bit <$> vector numPages+ icTieBreaker <- arbitrary+ pure (Chunks chunks IndexCompact {..})++ shrink (Chunks chunks index) =+ -- shrink number of pages+ [ Chunks chunks' index+ { icPrimary = primary'+ , icClashes = VU.slice 0 numPages' (icClashes index)+ , icLargerThanPage = VU.slice 0 numPages' (icLargerThanPage index)+ }+ | chunks' <- shrink chunks+ , let primary' = mconcat chunks'+ , let numPages' = VU.length primary'+ ] +++ -- shrink tie breaker+ [ Chunks chunks index+ { icTieBreaker = tieBreaker'+ }+ | tieBreaker' <- shrink (icTieBreaker index)+ ]
@@ -0,0 +1,569 @@+{-# LANGUAGE MagicHash #-}++{-# OPTIONS_GHC -Wno-orphans #-}++{- HLINT ignore "Avoid restricted alias" -}++module Test.Database.LSMTree.Internal.Index.Ordinary (tests) where++import Prelude hiding (all, head, last, length, notElem, splitAt,+ tail, takeWhile)++import Control.Arrow (first, (>>>))+import Control.Monad.ST.Strict (runST)+import qualified Data.ByteString.Lazy as LazyByteString (unpack)+import Data.ByteString.Short (ShortByteString (SBS))+import qualified Data.ByteString.Short as ShortByteString (length, pack)+import Data.Either (isLeft)+import Data.List (genericReplicate)+import qualified Data.List as List (tail)+import Data.Maybe (maybeToList)+import Data.Primitive.ByteArray (ByteArray (ByteArray), ByteArray#)+import Data.Vector (Vector, all, fromList, head, last, length,+ notElem, splitAt, tail, takeWhile, toList, (!))+import qualified Data.Vector.Primitive as Primitive (Vector (Vector), concat,+ force, length, singleton, toList)+import Data.Word (Word16, Word32, Word64, Word8)+import Database.LSMTree.Extras.Generators (LogicalPageSummaries,+ toAppends)+import Database.LSMTree.Extras.Index+ (Append (AppendMultiPage, AppendSinglePage),+ appendToOrdinary)+import qualified Database.LSMTree.Internal.Chunk as Chunk (toByteVector)+import Database.LSMTree.Internal.Entry (NumEntries (NumEntries))+import Database.LSMTree.Internal.Index.Ordinary+ (IndexOrdinary (IndexOrdinary), finalLBS, fromSBS,+ headerLBS, search, toUnslicedLastKeys)+import Database.LSMTree.Internal.Index.OrdinaryAcc (new, unsafeEnd)+import Database.LSMTree.Internal.Page (PageNo (PageNo),+ PageSpan (PageSpan))+import Database.LSMTree.Internal.Serialise+ (SerialisedKey (SerialisedKey'))+import Database.LSMTree.Internal.Unsliced (Unsliced, fromUnslicedKey,+ makeUnslicedKey)+import Database.LSMTree.Internal.Vector (byteVectorFromPrim)+import Test.Database.LSMTree.Internal.Chunk ()+import Test.QuickCheck (Arbitrary (arbitrary, shrink), Gen,+ NonNegative (NonNegative), Property,+ Small (Small, getSmall), Testable, chooseInt,+ counterexample, frequency, getPositive, getSorted,+ shrinkMap, suchThat, vector, (.&&.), (===), (==>))+import Test.QuickCheck.Instances.ByteString ()+ -- for @Arbitrary ShortByteString@+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)++-- * Tests++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Index.Ordinary" $+ [+ testGroup "Search" $+ [+ testProperty+ "Search for mentioned key works"+ prop_searchForMentionedKeyWorks,+ testProperty+ "Search for unmentioned key in range works"+ prop_searchForUnmentionedKeyInRangeWorks,+ testProperty+ "Search for unmentioned key beyond range works"+ prop_searchForUnmentionedKeyBeyondRangeWorks+ ],+ testGroup "Header and footer construction" $+ [+ testProperty+ "Header construction works"+ prop_headerConstructionWorks,+ testProperty+ "Footer construction works"+ prop_footerConstructionWorks+ ],+ testGroup "Deserialisation" $+ [+ testProperty+ "Number of entries from serialised index works"+ prop_numberOfEntriesFromSerialisedIndexWorks,+ testProperty+ "Index from serialised index works"+ prop_indexFromSerialisedIndexWorks,+ testProperty+ "Too short input makes deserialisation fail"+ prop_tooShortInputMakesDeserialisationFail,+ testProperty+ "Type-and-version error makes deserialisation fail"+ prop_typeAndVersionErrorMakesDeserialisationFail,+ testProperty+ "Partial key size block makes deserialisation fail"+ prop_partialKeySizeBlockMakesDeserialisationFail,+ testProperty+ "Partial key block makes deserialisation fail"+ prop_partialKeyBlockMakesDeserialisationFail+ ],+ testGroup "Incremental construction" $+ [+ testProperty+ "Incremental index construction works"+ prop_incrementalIndexConstructionWorks,+ testProperty+ "Incremental serialised key list construction works"+ prop_incrementalSerialisedKeyListConstructionWorks+ ]+ ]++-- * Tested type and version++-- | The type and version of the index to which these tests refer.+testedTypeAndVersion :: Word32+testedTypeAndVersion = 0x0101++-- * Utilities++-- ** Simple index construction++{-|+ Constructs an index from a list of unsliced keys. The keys have the same+ meaning as in the vector given to the 'IndexOrdinary' data constructor.+-}+indexFromUnslicedLastKeyList :: [Unsliced SerialisedKey] -> IndexOrdinary+indexFromUnslicedLastKeyList = IndexOrdinary . fromList++-- ** Partitioning according to search++{-|+ Invokes the 'search' function to obtain a page span and partitions the key+ list of the index into a prefix, selection, and suffix part, which+ correspond to the pages before, within, and after the page span,+ respectively.+-}+searchPartitioning ::+ Unsliced SerialisedKey+ -> IndexOrdinary+ -> (,,) (Vector (Unsliced SerialisedKey))+ (Vector (Unsliced SerialisedKey))+ (Vector (Unsliced SerialisedKey))+searchPartitioning unslicedKey index@(IndexOrdinary unslicedLastKeys)+ = (prefix, selection, suffix)+ where++ start, end :: Int+ PageSpan (PageNo start) (PageNo end) = search (fromUnslicedKey unslicedKey)+ index++ prefix, selection, suffix :: Vector (Unsliced SerialisedKey)+ ((prefix, selection), suffix) = first (splitAt start) $+ splitAt (succ end) $+ unslicedLastKeys++-- | Adds a search partitioning to a counterexample.+searchPartitioningCounterexample ::+ Testable prop+ => Vector (Unsliced SerialisedKey)+ -> Vector (Unsliced SerialisedKey)+ -> Vector (Unsliced SerialisedKey)+ -> prop+ -> Property+searchPartitioningCounterexample prefix selection suffix+ = counterexample $ "Prefix: " ++ show prefix ++ "; " +++ "Selection: " ++ show selection ++ "; " +++ "Suffix: " ++ show suffix++-- ** Construction of serialised indexes, possibly with errors++{-|+ Constructs a potential serialised index from a block for representing the+ type and the version of the index, a list of blocks for representing the+ keys list of the index, and a block for representing the number of entries+ of the run. Since all blocks can be arbitrary sequences of bytes, it is+ possible to construct an incorrect serialised index.+-}+potentialSerialisedIndex :: Primitive.Vector Word8+ -> [Primitive.Vector Word8]+ -> Primitive.Vector Word8+ -> ShortByteString+potentialSerialisedIndex typeAndVersionBlock+ potentialLastKeysBlocks+ potentialEntryCountBlock+ = SBS unliftedByteArray+ where++ primitiveVector :: Primitive.Vector Word8+ primitiveVector = Primitive.concat $ [typeAndVersionBlock] +++ potentialLastKeysBlocks +++ [potentialEntryCountBlock]++ unliftedByteArray :: ByteArray#+ !(Primitive.Vector _ _ (ByteArray unliftedByteArray))+ = Primitive.force primitiveVector++{-|+ The serialisation of the type and version of the index to which these tests+ refer.+-}+testedTypeAndVersionBlock :: Primitive.Vector Word8+testedTypeAndVersionBlock = byteVectorFromPrim testedTypeAndVersion++{-|+ Constructs blocks that constitute the serialisation of the key list of an+ index.+-}+lastKeysBlocks :: [Unsliced SerialisedKey] -> [Primitive.Vector Word8]+lastKeysBlocks unslicedLastKeys = concatMap lastKeyBlocks unslicedLastKeys where++ lastKeyBlocks :: Unsliced SerialisedKey -> [Primitive.Vector Word8]+ lastKeyBlocks unslicedLastKey = [lastKeySizeBlock, lastKeyBlock] where++ SerialisedKey' lastKeyBlock = fromUnslicedKey unslicedLastKey++ lastKeySizeBlock :: Primitive.Vector Word8+ lastKeySizeBlock = byteVectorFromPrim lastKeySizeAsWord16 where++ lastKeySizeAsWord16 :: Word16+ lastKeySizeAsWord16 = fromIntegral $+ Primitive.length lastKeyBlock++-- | Constructs the serialisation of the number of entries of a run.+entryCountBlock :: NumEntries -> Primitive.Vector Word8+entryCountBlock (NumEntries entryCount)+ = byteVectorFromPrim entryCountAsWord64+ where++ entryCountAsWord64 :: Word64+ entryCountAsWord64 = fromIntegral entryCount++-- | Constructs a correct serialised index.+serialisedIndex :: NumEntries -> [Unsliced SerialisedKey] -> ShortByteString+serialisedIndex entryCount unslicedLastKeys+ = potentialSerialisedIndex testedTypeAndVersionBlock+ (lastKeysBlocks unslicedLastKeys)+ (entryCountBlock entryCount)++-- ** Construction via appending++-- | Yields the keys that an append operation adds to an index.+appendedKeys :: Append -> [Unsliced SerialisedKey]+appendedKeys (AppendSinglePage _ lastKey)+ = [makeUnslicedKey lastKey]+appendedKeys (AppendMultiPage key overflowPageCount)+ = unslicedKey : genericReplicate overflowPageCount unslicedKey+ where++ unslicedKey :: Unsliced SerialisedKey+ unslicedKey = makeUnslicedKey key++{-|+ Yields the index that results from performing a sequence of append+ operations, starting with no keys.+-}+indexFromAppends :: [Append] -> IndexOrdinary+indexFromAppends appends = indexFromUnslicedLastKeyList $+ concatMap appendedKeys appends++{-|+ Yields the serialized key list that results from performing a sequence of+ append operations, starting with no keys.+-}+lastKeysBlockFromAppends :: [Append] -> Primitive.Vector Word8+lastKeysBlockFromAppends appends = lastKeysBlock where++ unslicedLastKeys :: [Unsliced SerialisedKey]+ unslicedLastKeys = concatMap appendedKeys appends++ lastKeysBlock :: Primitive.Vector Word8+ lastKeysBlock = Primitive.concat (lastKeysBlocks unslicedLastKeys)++{-|+ Incrementally constructs an index, using the functions 'new', 'append', and+ 'unsafeEnd'.+-}+incrementalConstruction :: [Append] -> (IndexOrdinary, Primitive.Vector Word8)+incrementalConstruction appends = runST $ do+ acc <- new initialKeyBufferSize minChunkSize+ commonChunks <- concat <$> mapM (flip appendToOrdinary acc) appends+ (remnant, unserialised) <- unsafeEnd acc+ let++ serialised :: Primitive.Vector Word8+ serialised = Primitive.concat $+ map Chunk.toByteVector $+ commonChunks ++ maybeToList remnant++ pure (unserialised, serialised)+ where++ {-+ We do not need to vary the initial key buffer size, since we are not+ testing growing vectors here.+ -}+ initialKeyBufferSize :: Int+ initialKeyBufferSize = 0x100++ {-+ We do not need to vary the minimum chunk size, since we are not testing+ chunk building here.+ -}+ minChunkSize :: Int+ minChunkSize = 0x1000++-- * Properties to test++-- ** Search++prop_searchForMentionedKeyWorks :: NonNegative (Small Int)+ -> SearchableIndex+ -> Property+prop_searchForMentionedKeyWorks (NonNegative (Small pageNo))+ (SearchableIndex index)+ = searchPartitioningCounterexample prefix selection suffix $+ pageNo < length unslicedLastKeys ==> all (< unslicedKey) prefix .&&.+ all (== unslicedKey) selection .&&.+ all (> unslicedKey) suffix+ where++ unslicedLastKeys :: Vector (Unsliced SerialisedKey)+ unslicedLastKeys = toUnslicedLastKeys index++ unslicedKey :: Unsliced SerialisedKey+ unslicedKey = unslicedLastKeys ! pageNo++ prefix, selection, suffix :: Vector (Unsliced SerialisedKey)+ (prefix, selection, suffix) = searchPartitioning unslicedKey index++prop_searchForUnmentionedKeyInRangeWorks :: Unsliced SerialisedKey+ -> SearchableIndex+ -> Property+prop_searchForUnmentionedKeyInRangeWorks unslicedKey (SearchableIndex index)+ = searchPartitioningCounterexample prefix selection suffix $+ unslicedKey `notElem` unslicedLastKeys &&+ unslicedKey <= last unslicedLastKeys ==>+ unslicedKey < selectionHead .&&.+ all (< unslicedKey) prefix .&&.+ all (== selectionHead) (tail selection) .&&.+ all (> selectionHead) suffix+ where++ unslicedLastKeys :: Vector (Unsliced SerialisedKey)+ unslicedLastKeys = toUnslicedLastKeys index++ prefix, selection, suffix :: Vector (Unsliced SerialisedKey)+ (prefix, selection, suffix) = searchPartitioning unslicedKey index++ selectionHead :: Unsliced SerialisedKey+ selectionHead = head selection++prop_searchForUnmentionedKeyBeyondRangeWorks :: Unsliced SerialisedKey+ -> SearchableIndex+ -> Property+prop_searchForUnmentionedKeyBeyondRangeWorks unslicedKey (SearchableIndex index)+ = searchPartitioningCounterexample prefix selection suffix $+ not (null lesserUnslicedLastKeys) ==>+ all (< selectionHead) prefix .&&.+ all (== selectionHead) (tail selection) .&&.+ all (> selectionHead) suffix+ where++ unslicedLastKeys :: Vector (Unsliced SerialisedKey)+ unslicedLastKeys = toUnslicedLastKeys index++ lesserUnslicedLastKeys :: Vector (Unsliced SerialisedKey)+ lesserUnslicedLastKeys = takeWhile (< unslicedKey) unslicedLastKeys++ prefix, selection, suffix :: Vector (Unsliced SerialisedKey)+ (prefix, selection, suffix) = searchPartitioning unslicedKey $+ IndexOrdinary lesserUnslicedLastKeys++ selectionHead :: Unsliced SerialisedKey+ selectionHead = head selection++-- ** Header and footer construction++prop_headerConstructionWorks :: Property+prop_headerConstructionWorks+ = LazyByteString.unpack headerLBS+ ===+ Primitive.toList testedTypeAndVersionBlock++prop_footerConstructionWorks :: NumEntries -> IndexOrdinary -> Property+prop_footerConstructionWorks entryCount index+ = LazyByteString.unpack (finalLBS entryCount index)+ ===+ Primitive.toList (entryCountBlock entryCount)++-- ** Deserialisation++prop_numberOfEntriesFromSerialisedIndexWorks :: NumEntries+ -> [Unsliced SerialisedKey]+ -> Property+prop_numberOfEntriesFromSerialisedIndexWorks entryCount unslicedLastKeys+ = errorMsgOrEntryCount === noErrorMsgButCorrectEntryCount+ where++ errorMsgOrEntryCount :: Either String NumEntries+ errorMsgOrEntryCount = fst <$>+ fromSBS (serialisedIndex entryCount unslicedLastKeys)++ noErrorMsgButCorrectEntryCount :: Either String NumEntries+ noErrorMsgButCorrectEntryCount = Right entryCount++prop_indexFromSerialisedIndexWorks :: NumEntries+ -> [Unsliced SerialisedKey]+ -> Property+prop_indexFromSerialisedIndexWorks entryCount unslicedLastKeys+ = errorMsgOrIndex === noErrorMsgButCorrectIndex+ where++ errorMsgOrIndex :: Either String IndexOrdinary+ errorMsgOrIndex = snd <$>+ fromSBS (serialisedIndex entryCount unslicedLastKeys)++ noErrorMsgButCorrectIndex :: Either String IndexOrdinary+ noErrorMsgButCorrectIndex = Right $+ indexFromUnslicedLastKeyList unslicedLastKeys++prop_tooShortInputMakesDeserialisationFail :: TooShortByteString -> Bool+prop_tooShortInputMakesDeserialisationFail+ = isLeft . fromSBS . fromTooShortByteString++prop_typeAndVersionErrorMakesDeserialisationFail :: Word32+ -> [Unsliced SerialisedKey]+ -> NumEntries+ -> Property+prop_typeAndVersionErrorMakesDeserialisationFail typeAndVersion+ unslicedLastKeys+ entryCount+ = typeAndVersion /= testedTypeAndVersion ==> isLeft errorMsgOrResult+ where++ errorMsgOrResult :: Either String (NumEntries, IndexOrdinary)+ errorMsgOrResult+ = fromSBS $+ potentialSerialisedIndex (byteVectorFromPrim typeAndVersion)+ (lastKeysBlocks unslicedLastKeys)+ (entryCountBlock entryCount)++prop_partialKeySizeBlockMakesDeserialisationFail :: [Unsliced SerialisedKey]+ -> Word8+ -> NumEntries+ -> Bool+prop_partialKeySizeBlockMakesDeserialisationFail unslicedLastKeys+ partialKeySizeByte+ entryCount+ = isLeft $+ fromSBS $+ potentialSerialisedIndex+ testedTypeAndVersionBlock+ (correctBlocks ++ [Primitive.singleton partialKeySizeByte])+ (entryCountBlock entryCount)+ where++ correctBlocks :: [Primitive.Vector Word8]+ correctBlocks = lastKeysBlocks unslicedLastKeys++prop_partialKeyBlockMakesDeserialisationFail :: [Unsliced SerialisedKey]+ -> Small Word16+ -> Primitive.Vector Word8+ -> NumEntries+ -> Property+prop_partialKeyBlockMakesDeserialisationFail unslicedLastKeys+ (Small statedSize)+ partialKeyBlock+ entryCount+ = fromIntegral statedSize > Primitive.length partialKeyBlock ==>+ isLeft (fromSBS input)+ where++ correctBlocks :: [Primitive.Vector Word8]+ correctBlocks = lastKeysBlocks unslicedLastKeys++ statedSizeBlock :: Primitive.Vector Word8+ statedSizeBlock = byteVectorFromPrim statedSize++ input :: ShortByteString+ input = potentialSerialisedIndex+ testedTypeAndVersionBlock+ (correctBlocks ++ [statedSizeBlock, partialKeyBlock])+ (entryCountBlock entryCount)++-- ** Incremental construction++prop_incrementalIndexConstructionWorks ::+ LogicalPageSummaries SerialisedKey -> Property+prop_incrementalIndexConstructionWorks logicalPageSummaries+ = fst (incrementalConstruction appends) === indexFromAppends appends+ where++ appends :: [Append]+ appends = toAppends logicalPageSummaries++prop_incrementalSerialisedKeyListConstructionWorks ::+ LogicalPageSummaries SerialisedKey -> Property+prop_incrementalSerialisedKeyListConstructionWorks logicalPageSummaries+ = snd (incrementalConstruction appends) === lastKeysBlockFromAppends appends+ where++ appends :: [Append]+ appends = toAppends logicalPageSummaries++-- * Test case generation and shrinking++{-+ For 'NumEntries' and 'Unsliced SerialisedKey', we use the 'Arbitrary'+ instantiations from "Database.LSMTree.Extras.Generators". For the+ correctness of the above tests, we need to assume the following properties+ of said instantiations, which are not strictly guaranteed:++ * Generated 'NumEntry' values are smaller than @2 ^ 64@.+ * The lengths of generated 'Unsliced SerialisedKey' values are smaller+ than @2 ^ 16@.+-}++instance Arbitrary IndexOrdinary where++ arbitrary = IndexOrdinary <$> fromList <$> arbitrary++ shrink = shrinkMap (IndexOrdinary . fromList) (toList . toUnslicedLastKeys)++newtype SearchableIndex = SearchableIndex IndexOrdinary deriving stock Show++instance Arbitrary SearchableIndex where++ arbitrary = do+ availableUnslicedLastKeys+ <- (getSorted <$> arbitrary) `suchThat` (not . null)+ unslicedLastKeys+ <- concat <$>+ mapM ((<$> genKeyCount) . flip replicate)+ availableUnslicedLastKeys+ pure $ SearchableIndex (indexFromUnslicedLastKeyList unslicedLastKeys)+ where++ genKeyCount :: Gen Int+ genKeyCount = frequency $+ [(4, pure 1), (1, getSmall <$> getPositive <$> arbitrary)]++ shrink (SearchableIndex (IndexOrdinary unslicedLastKeysVec))+ = [+ SearchableIndex (indexFromUnslicedLastKeyList unslicedLastKeys') |+ unslicedLastKeys' <- shrink (toList unslicedLastKeysVec),+ not (null unslicedLastKeys'),+ and $+ zipWith (<=) unslicedLastKeys' (List.tail unslicedLastKeys')+ ]++-- | A byte string that is too short to be a serialised index.+newtype TooShortByteString = TooShortByteString ShortByteString+ deriving stock Show++fromTooShortByteString :: TooShortByteString -> ShortByteString+fromTooShortByteString (TooShortByteString shortByteString) = shortByteString++instance Arbitrary TooShortByteString where++ arbitrary = do+ len <- chooseInt (0, 11)+ TooShortByteString <$> ShortByteString.pack <$> vector len++ shrink = fromTooShortByteString >>>+ shrink >>>+ filter (ShortByteString.length >>> (<= 11)) >>>+ map TooShortByteString
@@ -0,0 +1,606 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE NamedFieldPuns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TupleSections #-}+{-# LANGUAGE TypeApplications #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.Lookup (+ tests+ -- * internals+ , InMemLookupData (..)+ , SmallList (..)+ ) where++import Control.DeepSeq+import Control.Exception+import Control.Monad.ST.Strict+import Control.RefCount+import Data.Bifunctor+import Data.BloomFilter.Blocked (Bloom)+import qualified Data.BloomFilter.Blocked as Bloom+import Data.Coerce (coerce)+import Data.Either (rights)+import qualified Data.Foldable as F+import qualified Data.List as List+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe+import Data.Monoid (Endo (..))+import Data.Set (Set)+import qualified Data.Set as Set+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import qualified Data.Vector.Unboxed as VU+import Data.Word+import Database.LSMTree.Extras+import Database.LSMTree.Extras.Generators+import Database.LSMTree.Extras.RunData (RunData (..),+ liftArbitrary2Map, liftShrink2Map, withRuns)+import Database.LSMTree.Internal.Arena (newArenaManager,+ withUnmanagedArena)+import Database.LSMTree.Internal.BlobRef+import Database.LSMTree.Internal.Entry as Entry+import Database.LSMTree.Internal.Index (Index, IndexType)+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Ordinary),+ search)+import Database.LSMTree.Internal.Lookup+import Database.LSMTree.Internal.Page (PageNo (PageNo), PageSpan (..))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage+import Database.LSMTree.Internal.RawPage+import qualified Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunAcc as Run+import Database.LSMTree.Internal.RunBuilder+ (RunDataCaching (CacheRunData), RunParams (RunParams))+import Database.LSMTree.Internal.Serialise as Serialise+import Database.LSMTree.Internal.Serialise.Class+import Database.LSMTree.Internal.UniqCounter+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import GHC.Generics+import qualified System.FS.API as FS+import System.FS.API (Handle (..), mkFsPath)+import qualified System.FS.BlockIO.API as FS+import System.FS.BlockIO.API+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.Arbitrary (deepseqInvariant, noTags,+ prop_arbitraryAndShrinkPreserveInvariant,+ prop_forAllArbitraryAndShrinkPreserveInvariant)+import Test.Util.FS (withTempIOHasBlockIO)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Lookup" [+ testGroup "models" [+ testProperty "prop_bloomQueriesModel" $+ prop_bloomQueriesModel+ , testProperty "prop_indexSearchesModel" $+ prop_indexSearchesModel+ , testProperty "prop_prepLookupsModel" $+ prop_prepLookupsModel+ , testProperty "input distribution" $ \dats ->+ tabulateInMemLookupDataN (getSmallList dats) True+ ]+ , testGroup "With multi-page values" [+ testGroup "InMemLookupData" $+ prop_arbitraryAndShrinkPreserveInvariant noTags $+ deepseqInvariant @(InMemLookupData SerialisedKey SerialisedValue BlobSpan)+ , localOption (QuickCheckMaxSize 1000) $+ testProperty "prop_inMemRunLookupAndConstruction" prop_inMemRunLookupAndConstruction+ ]+ , testGroup "Without multi-page values" [+ testGroup "InMemLookupData" $+ prop_forAllArbitraryAndShrinkPreserveInvariant noTags+ genNoMultiPage+ shrinkNoMultiPage+ (deepseqInvariant @(InMemLookupData SerialisedKey SerialisedValue BlobSpan))+ , localOption (QuickCheckMaxSize 1000) $+ testProperty "prop_inMemRunLookupAndConstruction" $+ forAllShrink genNoMultiPage shrinkNoMultiPage prop_inMemRunLookupAndConstruction+ ]++ , testProperty "prop_roundtripFromWriteBufferLookupIO" $+ prop_roundtripFromWriteBufferLookupIO+ ]+ where+ genNoMultiPage = liftArbitrary2 arbitrary arbitrary+ shrinkNoMultiPage = liftShrink2 shrink shrink++runParams :: Index.IndexType -> RunParams+runParams indexType =+ RunParams {+ runParamCaching = CacheRunData,+ runParamAlloc = RunAllocFixed 10,+ runParamIndex = indexType+ }++testSalt :: Bloom.Salt+testSalt = 4++{-------------------------------------------------------------------------------+ Models+-------------------------------------------------------------------------------}++prop_bloomQueriesModel ::+ SmallList (InMemLookupData SerialisedKey SerialisedValue BlobSpan)+ -> Property+prop_bloomQueriesModel dats =+ -- The model never returns false positives, but the real bloom filter does,+ -- so the model result should be a subsequence of the real result.+ counterexample (show model ++ " is not a subsequence of " ++ show real) $+ property (model `List.isSubsequenceOf` real)+ where+ runDatas = getSmallList $ fmap runData dats+ runs = fmap mkTestRun runDatas+ blooms = fmap snd3 runs+ lookupss = concatMap lookups $ getSmallList dats+ real = map (\(RunIxKeyIx rix kix) -> (rix,kix)) $ VP.toList $+ bloomQueries testSalt (V.fromList blooms) (V.fromList lookupss)+ model = bloomQueriesModel (fmap Map.keysSet runDatas) lookupss++-- | A bloom filter is a probablistic set that can return false positives, but+-- not false negatives. The simplest model of a bloom filter is therefore a+-- non-probablistic set: a set that only returns true positives or negatives.+bloomQueriesModel :: [Set SerialisedKey] -> [SerialisedKey] -> [(RunIx, KeyIx)]+bloomQueriesModel blooms ks = [+ (rix, kix)+ | (kix, k) <- ks'+ , (rix, b) <- rs'+ , Set.member k b+ ]+ where+ rs' = zip [0..] blooms+ ks' = zip [0..] ks++prop_indexSearchesModel ::+ SmallList (InMemLookupData SerialisedKey SerialisedValue BlobSpan)+ -> Property+prop_indexSearchesModel dats =+ forAllShrink (rkixsGen rkixsAll) shrink $ \rkixs ->+ real (VU.fromList rkixs) === model rkixs+ .&&. real (VU.fromList rkixs) === model rkixs+ .&&. length rkixs === length (model rkixs)+ where+ rkixsAll = [(rix, kix) | (rix,_) <- zip [0..] runs, (kix,_) <- zip [0..] lookupss]+ rkixsGen [] = pure []+ rkixsGen xs = listOf (elements xs)++ runs = getSmallList $ fmap (mkTestRun . runData) dats+ lookupss = concatMap lookups $ getSmallList dats+ real rkixs = runST $ withUnmanagedArena $ \arena -> do+ let rs = V.fromList (fmap runWithHandle runs)+ ks = V.fromList lookupss+ res <- indexSearches arena (V.map thrd3 rs) (V.map fst3 rs) ks+ ((V.convert . V.map (uncurry RunIxKeyIx) . V.convert) rkixs)+ pure $ V.map ioopPageSpan res+ model rkixs = V.fromList $ indexSearchesModel (fmap thrd3 runs) lookupss $ rkixs++indexSearchesModel ::+ [Index]+ -> [SerialisedKey]+ -> [(RunIx, KeyIx)]+ -> [PageSpan]+indexSearchesModel cs ks rkixs =+ flip fmap rkixs $ \(rix, kix) ->+ let c = cs List.!! rix+ k = ks List.!! kix+ in Index.search k c++prop_prepLookupsModel ::+ SmallList (InMemLookupData SerialisedKey SerialisedValue BlobSpan)+ -> Property+prop_prepLookupsModel dats = real === model+ where+ runs = getSmallList $ fmap (mkTestRun . runData) dats+ lookupss = concatMap lookups $ getSmallList dats+ real = runST $ withUnmanagedArena $ \arena -> do+ let rs = V.fromList (fmap runWithHandle runs)+ ks = V.fromList lookupss+ (kixs, ioops) <- prepLookups+ arena+ testSalt+ (V.map snd3 rs)+ (V.map thrd3 rs)+ (V.map fst3 rs) ks+ pure ( map (\(RunIxKeyIx r k) -> (r,k)) (VP.toList kixs)+ , map ioopPageSpan (V.toList ioops)+ )+ model = prepLookupsModel (fmap (\x -> (snd3 x, thrd3 x)) runs) lookupss++prepLookupsModel ::+ [(Bloom SerialisedKey, Index)]+ -> [SerialisedKey]+ -> ([(RunIx, KeyIx)], [PageSpan])+prepLookupsModel rs ks = unzip+ [ ((rix, kix), pspan)+ | (kix, k) <- zip [0..] ks+ , (rix, (b, c)) <- zip [0..] rs+ , Bloom.elem k b+ , let pspan = Index.search k c+ ]++{-------------------------------------------------------------------------------+ Round-trip+-------------------------------------------------------------------------------}++-- | Construct a run incrementally, then test a number of positive and negative lookups.+prop_inMemRunLookupAndConstruction ::+ InMemLookupData SerialisedKey SerialisedValue BlobSpan+ -> Property+prop_inMemRunLookupAndConstruction dat =+ tabulateInMemLookupData dat run+ $ conjoinF (fmap checkMaybeInRun keysMaybeInRun) .&&. conjoinF (fmap checkNotInRun keysNotInRun)+ where+ InMemLookupData{runData, lookups} = dat++ run = mkTestRun runData+ keys = V.fromList lookups+ -- prepLookups says that a key /could be/ in the given page+ keysMaybeInRun = runST $ withUnmanagedArena $ \arena -> do+ (kixs, ioops) <- let r = V.singleton (runWithHandle run)+ in prepLookups+ arena+ testSalt+ (V.map snd3 r)+ (V.map thrd3 r)+ (V.map fst3 r)+ keys+ let ks = V.map (V.fromList lookups V.!)+ (V.convert (VP.map (\(RunIxKeyIx _r k) -> k) kixs))+ pss = V.map (handleRaw . ioopHandle) ioops+ pspans = V.map (ioopPageSpan) ioops+ pure $ zip3 (V.toList ks) (V.toList pss) (V.toList pspans)+ -- prepLookups says that a key /is definitely not/ in the given page+ keysNotInRun = Set.toList (Set.fromList lookups Set.\\ Set.fromList (fmap (\(x,_,_) -> x) keysMaybeInRun))++ -- Check that a key /is definitely not/ in the given page.+ checkNotInRun :: SerialisedKey -> Property+ checkNotInRun k =+ tabulate1Pre (classifyBin (isJust truth) False)+ $ counterexample ("checkNotInRun: " <> show k)+ $ truth === test+ where truth = Map.lookup k runData+ test = Nothing++ -- | Check that a key /could be/ in the given page+ checkMaybeInRun :: ( SerialisedKey+ , Map Int (Either RawPage RawOverflowPage)+ , PageSpan )+ -> Property+ checkMaybeInRun (k, ps, PageSpan (PageNo i) (PageNo j))+ | i <= j = tabulate "PageSpan size" [showPowersOf10 $ j - i + 1]+ $ tabulate1Pre (classifyBin (isJust truth) True)+ $ counterexample ("checkMaybeInRun: " <> show k)+ $ truth === test+ | otherwise = error "impossible: end of a page span can not come before its start"+ where+ truth = Map.lookup k runData+ test =+ case ps Map.! i of+ Left rawPage ->+ case rawPageLookup rawPage (coerce k) of+ LookupEntryNotPresent -> Nothing+ LookupEntry entry -> Just entry+ LookupEntryOverflow entry n -> Just (first (concatOverflow n) entry)+ Right _rawOverflowPage ->+ error "looked up overflow page!"++ -- read remaining bytes for a multi-page value, and append it to the+ -- prefix we already have+ concatOverflow :: Word32 -> SerialisedValue -> SerialisedValue+ concatOverflow = coerce $ \(n :: Word32) (v :: RawBytes) ->+ v <> RB.take (fromIntegral n)+ (mconcat $ map rawOverflowPageRawBytes overflowPages)+ where+ start = i + 1+ size = j - i+ overflowPages :: [RawOverflowPage]+ overflowPages = rights+ . Map.elems+ . Map.take size+ . Map.drop start $ ps++ tabulate1Pre :: BinaryClassification -> Property -> Property+ tabulate1Pre x = tabulate "Lookup classification: pre intra-page lookup" [show x]+++{-------------------------------------------------------------------------------+ Round-trip lookups in IO+-------------------------------------------------------------------------------}++prop_roundtripFromWriteBufferLookupIO ::+ SmallList (InMemLookupData SerialisedKey SerialisedValue SerialisedBlob)+ -> Property+prop_roundtripFromWriteBufferLookupIO (SmallList dats) =+ ioProperty $+ withTempIOHasBlockIO "prop_roundtripFromWriteBufferLookupIO" $ \hfs hbio ->+ withWbAndRuns hfs hbio Index.Ordinary dats $ \wb wbblobs runs -> do+ let model :: Map SerialisedKey (Entry SerialisedValue SerialisedBlob)+ model = Map.unionsWith (Entry.combine resolveV) (map runData dats)+ keys = V.fromList [ k | InMemLookupData{lookups} <- dats+ , k <- lookups ]+ modelres :: V.Vector (Maybe (Entry SerialisedValue SerialisedBlob))+ modelres = V.map (\k -> Map.lookup k model) keys+ arenaManager <- newArenaManager+ realres <-+ fetchBlobs hfs =<< -- retrieve blobs to match type of model result+ lookupsIOWithWriteBuffer+ hbio+ arenaManager+ resolveV+ testSalt+ wb wbblobs+ runs+ (V.map (\(DeRef r) -> Run.runFilter r) runs)+ (V.map (\(DeRef r) -> Run.runIndex r) runs)+ (V.map (\(DeRef r) -> Run.runKOpsFile r) runs)+ keys+ pure $ modelres === realres+ where+ resolveV = \(SerialisedValue v1) (SerialisedValue v2) -> SerialisedValue (v1 <> v2)++ fetchBlobs :: FS.HasFS IO h+ -> (V.Vector (Maybe (Entry v (WeakBlobRef IO h))))+ -> IO (V.Vector (Maybe (Entry v SerialisedBlob)))+ fetchBlobs hfs = traverse (traverse (traverse (readWeakBlobRef hfs)))++-- | Given a bunch of 'InMemLookupData', prepare the data into the form needed+-- for 'lookupsIOWithWriteBuffer': a write buffer (and blobs) and a vector of+-- on-disk runs. Also passes the model and the keys to look up to the inner+-- action.+--+withWbAndRuns :: FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> IndexType+ -> [InMemLookupData SerialisedKey SerialisedValue SerialisedBlob]+ -> ( WB.WriteBuffer+ -> Ref (WBB.WriteBufferBlobs IO h)+ -> V.Vector (Ref (Run.Run IO h))+ -> IO a)+ -> IO a+withWbAndRuns hfs _ _ [] action =+ bracket+ (WBB.new hfs (FS.mkFsPath ["wbblobs"]))+ releaseRef+ (\wbblobs -> action WB.empty wbblobs V.empty)++withWbAndRuns hfs hbio indexType (wbdat:rundats) action =+ bracket (WBB.new hfs (FS.mkFsPath ["wbblobs"])) releaseRef $ \wbblobs -> do+ wbkops <- traverse (traverse (WBB.addBlob hfs wbblobs))+ (runData wbdat)+ let wb = WB.fromMap wbkops+ let rds = map (RunData . runData) rundats+ counter <- newUniqCounter 1+ withRuns hfs hbio testSalt (runParams indexType) (FS.mkFsPath []) counter rds $+ \runs ->+ action wb wbblobs (V.fromList runs)++{-------------------------------------------------------------------------------+ Utils+-------------------------------------------------------------------------------}++newtype SmallList a = SmallList { getSmallList :: [a] }+ deriving stock (Show, Eq)+ deriving newtype (Functor, Foldable)++instance Arbitrary a => Arbitrary (SmallList a) where+ arbitrary = do+ n <- chooseInt (0, 5)+ SmallList <$> vectorOf n arbitrary+ shrink = fmap SmallList . shrink . getSmallList++conjoinF :: (Testable prop, Foldable f) => f prop -> Property+conjoinF = conjoin . F.toList++ioopPageSpan :: IOOp s h -> PageSpan+ioopPageSpan ioop =+ assert (fst x `mod` 4096 == 0) $+ assert (snd x `mod` 4096 == 0) $+ assert (start >= 0) $+ assert (end >= start)+ PageSpan {+ pageSpanStart = PageNo start+ , pageSpanEnd = PageNo end+ }+ where+ start = fst x `div` 4096+ size = (snd x `div` 4096) - 1+ end = start + size++ x = case ioop of+ IOOpRead _ foff _ _ c -> (fromIntegral foff, fromIntegral c)+ IOOpWrite _ foff _ _ c -> (fromIntegral foff, fromIntegral c)++fst3 :: (a, b, c) -> a+fst3 (h, _, _) = h++snd3 :: (a, b, c) -> b+snd3 (_, b, _) = b++thrd3 :: (a, b, c) -> c+thrd3 (_, _, c) = c++{-------------------------------------------------------------------------------+ Test run+-------------------------------------------------------------------------------}++runWithHandle ::+ TestRun+ -> ( Handle (Map Int (Either RawPage RawOverflowPage))+ , Bloom SerialisedKey, Index+ )+runWithHandle (rawPages, b, ic) = (Handle rawPages (mkFsPath ["do not use"]), b, ic)++type TestRun = (Map Int (Either RawPage RawOverflowPage), Bloom SerialisedKey, Index)++mkTestRun :: Map SerialisedKey (Entry SerialisedValue BlobSpan) -> TestRun+mkTestRun dat = (rawPages, b, ic)+ where+ nentries = NumEntries (Map.size dat)++ -- one-shot run construction+ (pages, b, ic) = runST $ do+ racc <- Run.new nentries (RunAllocFixed 10) testSalt Index.Ordinary+ let kops = Map.toList dat+ psopss <- traverse (uncurry (Run.addKeyOp racc)) kops+ (mp, _ , b', ic', _) <- Run.unsafeFinalise racc+ let pages' = [ p | (ps, ops, _) <- psopss+ , p <- map Left ps ++ map Right ops ]+ ++ [ Left p | p <- maybeToList mp ]+ pure (pages', b', ic')++ -- create a mapping of page numbers to raw pages, which we can use to do+ -- intra-page lookups on after first probing the bloom filter and index+ rawPages :: Map Int (Either RawPage RawOverflowPage)+ rawPages = Map.fromList (zip [0..] pages)++{-------------------------------------------------------------------------------+ Labelling+-------------------------------------------------------------------------------}++-- | Binary classification of truth vs. a test result+classifyBin :: Bool -> Bool -> BinaryClassification+classifyBin truth test+ | truth && test = TruePositive+ | not truth && test = FalsePositive+ | truth && not test = FalseNegative+ -- not truth && not test =+ | otherwise = TrueNegative++data BinaryClassification =+ TruePositive | FalsePositive+ | FalseNegative | TrueNegative+ deriving stock Show++tabulateInMemLookupDataN ::+ forall prop. Testable prop+ => [InMemLookupData SerialisedKey SerialisedValue BlobSpan]+ -> (prop -> Property)+tabulateInMemLookupDataN dats = appEndo (foldMap Endo [+ let run = mkTestRun (runData dat)+ in tabulateInMemLookupData dat run+ | dat <- dats+ ])+ . tabulate "Number of runs" [show $ length dats]++tabulateInMemLookupData ::+ forall prop. Testable prop+ => InMemLookupData SerialisedKey SerialisedValue BlobSpan+ -> TestRun+ -> (prop -> Property)+tabulateInMemLookupData dat run =+ tabulateKeySizes+ . tabulateValueSizes+ . tabulateNumKeyEntryPairs+ . tabulateNumPages+ . tabulateNumLookups+ . tabulateEntryType+ where+ InMemLookupData{runData, lookups} = dat+ tabulateKeySizes = tabulate "Size of key in run" [showPowersOf10 $ sizeofKey k | k <- Map.keys runData ]+ tabulateValueSizes = tabulate "Size of value in run" [showPowersOf10 $ onValue 0 sizeofValue e | e <- Map.elems runData]+ tabulateNumKeyEntryPairs = tabulate "Number of key-entry pairs" [showPowersOf10 (Map.size runData) ]+ tabulateNumPages = tabulate "Number of pages" [showPowersOf10 (Map.size ps) | let (ps,_,_) = run]+ tabulateNumLookups = tabulate "Number of lookups" [showPowersOf10 (length lookups)]+ tabulateEntryType = tabulate "Entry type" (map (takeWhile (/= ' ') . show) (Map.elems runData))++{-------------------------------------------------------------------------------+ Arbitrary+-------------------------------------------------------------------------------}++data InMemLookupData k v b = InMemLookupData {+ -- | Data for constructing a run+ runData :: Map k (Entry v b)+ -- | Keys to look up. Expected lookup results are obtained by querying+ -- runData.+ , lookups :: [k]+ }+ deriving stock (Show, Generic)+ deriving anyclass NFData++instance Arbitrary (InMemLookupData SerialisedKey SerialisedValue BlobSpan) where+ arbitrary = liftArbitrary3InMemLookupData genSerialisedKey genSerialisedValue genBlobSpan+ shrink = liftShrink3InMemLookupData shrinkSerialisedKey shrinkSerialisedValue shrinkBlobSpan++instance Arbitrary1 (InMemLookupData SerialisedKey SerialisedValue) where+ liftArbitrary = liftArbitrary3InMemLookupData genSerialisedKey genSerialisedValue++instance Arbitrary2 (InMemLookupData SerialisedKey) where+ liftArbitrary2 = liftArbitrary3InMemLookupData genSerialisedKey++liftArbitrary3InMemLookupData ::+ Ord k+ => Gen k+ -> Gen v+ -> Gen b+ -> Gen (InMemLookupData k v b)+liftArbitrary3InMemLookupData genKey genValue genBlob = do+ kops <- liftArbitrary2Map genKey (liftArbitrary genEntry)+ `suchThat` (\x -> Map.size (Map.filter isJust x) > 0)+ let runData = Map.mapMaybe id kops+ lookups <- (sublistOf (Map.keys kops) >>= shuffle)+ pure InMemLookupData{ runData, lookups }+ where+ genEntry = liftArbitrary2 genValue genBlob++liftShrink3InMemLookupData ::+ Ord k+ => (k -> [k])+ -> (v -> [v])+ -> (b -> [b])+ -> InMemLookupData k v b+ -> [InMemLookupData k v b]+liftShrink3InMemLookupData shrinkKey shrinkValue shrinkBlob InMemLookupData{ runData, lookups } =+ [ InMemLookupData runData' lookups+ | runData' <- liftShrink2Map shrinkKey shrinkEntry runData+ , Map.size runData' > 0 ]+ ++ [ InMemLookupData runData lookups'+ | lookups' <- liftShrink shrinkKey lookups ]+ where+ shrinkEntry = liftShrink2 shrinkValue shrinkBlob++genSerialisedKey :: Gen SerialisedKey+genSerialisedKey = Serialise.serialiseKey <$> arbitraryBoundedIntegral @Word64++shrinkSerialisedKey :: SerialisedKey -> [SerialisedKey]+shrinkSerialisedKey k = Serialise.serialiseKey <$> shrink (Serialise.deserialiseKey k :: Word64)++genSerialisedValue :: Gen SerialisedValue+genSerialisedValue = frequency [ (50, arbitrary), (1, genLongValue) ]+ where genLongValue = coerce ((<>) <$> genRawBytesSized 65536 <*> arbitrary)++-- Shrinking as lists can normally be quite slow, so if the value is larger than+-- a threshold, use less exhaustive shrinking.+shrinkSerialisedValue :: SerialisedValue -> [SerialisedValue]+shrinkSerialisedValue v+ | sizeofValue v > 64 = -- shrink towards fewer bytes+ [ coerce RB.take n' v | n' <- shrinkIntegral n ]+ -- shrink towards a value of all 0-bytes+ ++ [ v' | let v' = coerce (VP.fromList $ replicate n 0), v' /= v ]+ | otherwise = shrink v -- expensive, but thorough+ where n = sizeofValue v++genBlobSpan :: Gen BlobSpan+genBlobSpan = arbitrary++shrinkBlobSpan :: BlobSpan -> [BlobSpan]+shrinkBlobSpan = shrink++instance Arbitrary (InMemLookupData SerialisedKey SerialisedValue SerialisedBlob) where+ arbitrary = liftArbitrary3InMemLookupData genSerialisedKey genSerialisedValue genSerialisedBlob+ shrink = liftShrink3InMemLookupData shrinkSerialisedKey shrinkSerialisedValue shrinkSerialisedBlob++genSerialisedBlob :: Gen SerialisedBlob+genSerialisedBlob = arbitrary++shrinkSerialisedBlob :: SerialisedBlob -> [SerialisedBlob]+shrinkSerialisedBlob = shrink
@@ -0,0 +1,238 @@+module Test.Database.LSMTree.Internal.Merge (tests) where++import Control.Exception (evaluate)+import Control.RefCount+import Data.Bifoldable (bifoldMap)+import qualified Data.BloomFilter.Blocked as Bloom+import Data.Foldable (traverse_)+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (isJust)+import qualified Data.Vector as V+import Database.LSMTree.Extras+import Database.LSMTree.Extras.RunData+import qualified Database.LSMTree.Internal.BlobFile as BlobFile+import qualified Database.LSMTree.Internal.Entry as Entry+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Ordinary))+import Database.LSMTree.Internal.Merge (MergeType (..))+import qualified Database.LSMTree.Internal.Merge as Merge+import Database.LSMTree.Internal.PageAcc (entryWouldFitInPage)+import Database.LSMTree.Internal.Paths (RunFsPaths (..),+ pathsForRunFiles)+import qualified Database.LSMTree.Internal.Run as Run+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified System.FS.API as FS+import qualified System.FS.API.Lazy as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.Sim as FsSim+import qualified System.FS.Sim.Error as FsSim+import qualified System.FS.Sim.MockFS as FsSim+import Test.Database.LSMTree.Internal.RunReader (readKOps)+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Merge"+ [ testProperty "prop_MergeDistributes" $ \mergeType stepSize rds ->+ ioPropertyWithMockFS $ \fs hbio ->+ prop_MergeDistributes fs hbio mergeType stepSize rds+ , testProperty "prop_AbortMerge" $ \level stepSize rds ->+ ioPropertyWithMockFS $ \fs hbio ->+ prop_AbortMerge fs hbio level stepSize rds+ ]+ where+ ioPropertyWithMockFS ::+ Testable p+ => (FS.HasFS IO FsSim.HandleMock -> FS.HasBlockIO IO FsSim.HandleMock -> IO p)+ -> Property+ ioPropertyWithMockFS prop = ioProperty $ do+ (res, mockFS, _) <- FsSim.runSimErrorHasBlockIO FsSim.empty FsSim.emptyErrors prop+ pure $ res+ .&&. counterexample "open handles"+ (FsSim.numOpenHandles mockFS === 0)++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Ordinary+ }++testSalt :: Bloom.Salt+testSalt = 4++-- | Creating multiple runs from write buffers and merging them leads to the+-- same run as merging the write buffers and creating a run.+--+-- @mergeRuns . map flush === flush . mergeWriteBuffers@+prop_MergeDistributes ::+ FS.HasFS IO h ->+ FS.HasBlockIO IO h ->+ MergeType ->+ StepSize ->+ SmallList (RunData SerialisedKey SerialisedValue SerialisedBlob) ->+ IO Property+prop_MergeDistributes fs hbio mergeType stepSize (SmallList rds) = do+ let path = FS.mkFsPath []+ counter <- newUniqCounter 0+ withRuns fs hbio testSalt runParams path counter rds' $ \runs -> do+ let stepsNeeded = sum (map (Map.size . unRunData) rds)++ fsPathLhs <- RunFsPaths path . uniqueToRunNumber <$> incrUniqCounter counter+ (stepsDone, lhs) <- mergeRuns fs hbio mergeType stepSize fsPathLhs runs+ let runData = RunData $ mergeWriteBuffers mergeType $ fmap unRunData rds'+ withRun fs hbio testSalt runParams path counter runData $ \rhs -> do++ (lhsSize, lhsFilter, lhsIndex, lhsKOps,+ lhsKOpsFileContent, lhsBlobFileContent) <- getRunContent lhs++ (rhsSize, rhsFilter, rhsIndex, rhsKOps,+ rhsKOpsFileContent, rhsBlobFileContent) <- getRunContent rhs++ -- cleanup+ releaseRef lhs++ pure $ stats $+ counterexample "numEntries"+ (lhsSize === rhsSize)+ .&&. -- we can't just test bloom filter equality, their sizes may differ.+ counterexample "runFilter"+ ( Bloom.sizeBits (Bloom.size lhsFilter)+ >= Bloom.sizeBits (Bloom.size rhsFilter))+ .&&. -- the index is equal, but only because range finder precision is+ -- always 0 for the numbers of entries we are dealing with.+ counterexample "runIndex"+ (lhsIndex === rhsIndex)+ .&&. counterexample "kops"+ (lhsKOps === rhsKOps)+ .&&. counterexample "kopsFile"+ (lhsKOpsFileContent === rhsKOpsFileContent)+ .&&. counterexample "blobFile"+ (lhsBlobFileContent === rhsBlobFileContent)+ .&&. counterexample ("step counting")+ (stepsDone === stepsNeeded)+ where+ stats = tabulate "value size" (map (showPowersOf10 . sizeofValue) vals)+ . tabulate "entry type" (map (takeWhile (/= ' ') . show . snd) kops)+ . label (if any isLarge kops then "has large k/op" else "no large k/op")+ . label ("number of runs: " <> showPowersOf 2 (length rds'))+ rds' = fmap serialiseRunData rds+ kops = foldMap (Map.toList . unRunData) rds'+ vals = concatMap (bifoldMap pure mempty . snd) kops+ isLarge = not . uncurry entryWouldFitInPage++ getRunContent run@(DeRef Run.Run {+ Run.runFilter,+ Run.runIndex,+ Run.runKOpsFile,+ Run.runBlobFile+ }) = do+ runSize <- evaluate (Run.size run)+ runKOps <- readKOps Nothing run+ kopsFileContent <- FS.hGetAll fs runKOpsFile+ blobFileContent <- withRef runBlobFile $+ FS.hGetAll fs . BlobFile.blobFileHandle+ pure ( runSize+ , runFilter+ , runIndex+ , runKOps+ , kopsFileContent+ , blobFileContent+ )++-- | After merging for a few steps, we can prematurely abort the merge, which+-- should clean up properly.+prop_AbortMerge ::+ FS.HasFS IO h ->+ FS.HasBlockIO IO h ->+ MergeType ->+ StepSize ->+ SmallList (RunData SerialisedKey SerialisedValue SerialisedBlob) ->+ IO Property+prop_AbortMerge fs hbio mergeType (Positive stepSize) (SmallList wbs) = do+ let path = FS.mkFsPath []+ let pathOut = RunFsPaths path (RunNumber 0)+ counter <- newUniqCounter 1+ withRuns fs hbio testSalt runParams path counter wbs' $ \runs -> do+ mergeToClose <- makeInProgressMerge pathOut runs+ traverse_ Merge.abort mergeToClose++ filesExist <- traverse (FS.doesFileExist fs) (pathsForRunFiles pathOut)++ pure $+ counterexample ("run files exist: " <> show filesExist) $+ isJust mergeToClose ==> all not filesExist+ where+ wbs' = fmap serialiseRunData wbs++ makeInProgressMerge path runs =+ Merge.new fs hbio testSalt runParams mergeType resolveVal+ path (V.fromList runs) >>= \case+ Nothing -> pure Nothing -- not in progress+ Just merge -> do+ -- just do a few steps once, ideally not completing the merge+ Merge.steps merge stepSize >>= \case+ (_, Merge.MergeDone) -> do+ Merge.abort merge -- run not needed+ pure Nothing -- not in progress+ (_, Merge.MergeInProgress) ->+ pure (Just merge)++{-------------------------------------------------------------------------------+ Utilities+-------------------------------------------------------------------------------}++type StepSize = Positive Int++mergeRuns ::+ FS.HasFS IO h ->+ FS.HasBlockIO IO h ->+ MergeType ->+ StepSize ->+ RunFsPaths ->+ [Ref (Run.Run IO h)] ->+ IO (Int, Ref (Run.Run IO h))+mergeRuns fs hbio mergeType (Positive stepSize) fsPath runs = do+ Merge.new fs hbio testSalt runParams mergeType resolveVal+ fsPath (V.fromList runs)+ >>= \case+ Just m -> Merge.stepsToCompletionCounted m stepSize+ Nothing -> (,) 0 <$> unsafeCreateRunAt fs hbio testSalt runParams fsPath+ (RunData Map.empty)++type SerialisedEntry = Entry.Entry SerialisedValue SerialisedBlob++mergeWriteBuffers :: MergeType+ -> [Map SerialisedKey SerialisedEntry]+ -> Map SerialisedKey SerialisedEntry+mergeWriteBuffers = \case+ MergeTypeMidLevel -> Map.unionsWith (Entry.combine resolveVal)+ MergeTypeLastLevel -> Map.filter (not . isDelete)+ . Map.unionsWith (Entry.combine resolveVal)+ MergeTypeUnion -> Map.filter (not . isDelete)+ . Map.unionsWith (Entry.combineUnion resolveVal)+ where+ isDelete Entry.Delete = True+ isDelete _ = False++resolveVal :: ResolveSerialisedValue+resolveVal (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)++newtype SmallList a = SmallList { getSmallList :: [a] }+ deriving stock (Show, Eq)+ deriving newtype (Functor, Foldable)++-- | Skewed towards short lists, but still generates longer ones.+instance Arbitrary a => Arbitrary (SmallList a) where+ arbitrary = do+ ub <- sized $ \s -> chooseInt (5, s `div` 3)+ n <- chooseInt (1, ub)+ SmallList <$> vectorOf n arbitrary++ shrink = fmap SmallList . shrink . getSmallList
@@ -0,0 +1,53 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.MergingRun (tests) where++import Database.LSMTree.Internal.MergingRun+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.MergingRun"+ [ testProperty "prop_CreditsPair" prop_CreditsPair+ ]++-- | The representation of CreditsPair should round trip properly. This is+-- non-trivial because it uses a packed bitfield representation.+--+prop_CreditsPair :: SpentCredits -> UnspentCredits -> Property+prop_CreditsPair spentCredits unspentCredits =+ tabulate "bounds" [spentCreditsBound, unspentCreditsBound] $+ let cp :: Int+ !cp = CreditsPair spentCredits unspentCredits+ in case cp of+ CreditsPair spentCredits' unspentCredits' ->+ (spentCredits, unspentCredits) === (spentCredits', unspentCredits')+ where+ spentCreditsBound+ | spentCredits == minBound = "spentCredits == minBound"+ | spentCredits == maxBound = "spentCredits == maxBound"+ | otherwise = "spentCredits == other"++ unspentCreditsBound+ | unspentCredits == minBound = "unspentCredits == minBound"+ | unspentCredits == maxBound = "unspentCredits == maxBound"+ | otherwise = "unspentCredits == other"++deriving newtype instance Enum SpentCredits+deriving newtype instance Enum UnspentCredits++instance Arbitrary SpentCredits where+ arbitrary =+ frequency [ (1, pure minBound)+ , (1, pure maxBound)+ , (10, arbitraryBoundedEnum)+ ]++instance Arbitrary UnspentCredits where+ arbitrary =+ frequency [ (1, pure minBound)+ , (1, pure maxBound)+ , (10, arbitraryBoundedEnum)+ ]+
@@ -0,0 +1,278 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.MergingTree (tests) where++import Control.ActionRegistry+import Control.Exception (bracket)+import Control.Monad.Class.MonadAsync as Async+import Control.RefCount+import Data.Coerce (coerce)+import Data.Foldable (toList, traverse_)+import Data.List.NonEmpty (NonEmpty)+import Data.Map (Map)+import qualified Data.Map as Map+import Data.Traversable (for)+import qualified Data.Vector as V+import Database.LSMTree.Extras (showPowersOf10)+import Database.LSMTree.Extras.MergingRunData+import Database.LSMTree.Extras.MergingTreeData+import Database.LSMTree.Extras.RunData+import Database.LSMTree.Internal.Arena (newArenaManager)+import Database.LSMTree.Internal.BlobRef+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry (Entry)+import qualified Database.LSMTree.Internal.Entry as Entry+import qualified Database.LSMTree.Internal.Index as Index+import qualified Database.LSMTree.Internal.Lookup as Lookup+import qualified Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.MergingTree+import Database.LSMTree.Internal.MergingTree.Lookup+import qualified Database.LSMTree.Internal.Paths as Paths+import qualified Database.LSMTree.Internal.Run as Run+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.Sim.MockFS as MockFS+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.FS (propNoOpenHandles, withSimHasBlockIO)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.MergingTree"+ [ testProperty "prop_isStructurallyEmpty" prop_isStructurallyEmpty+ , testProperty "prop_lookupTree" $ \keys mtd ->+ ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_lookupTree hfs hbio keys mtd+ , testProperty "prop_supplyCredits" $ \threshold credits mtd ->+ ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_supplyCredits hfs hbio threshold credits mtd+ ]++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Ordinary+ }++testSalt :: Bloom.Salt+testSalt = 4++-- | Check that the merging tree constructor functions preserve the property+-- that if the inputs are obviously empty, the output is also obviously empty.+--+prop_isStructurallyEmpty :: EmptyMergingTree -> Property+prop_isStructurallyEmpty emt =+ ioProperty $+ bracket (mkEmptyMergingTree emt)+ releaseRef+ isStructurallyEmpty++-- | An expression to specify the shape of an empty 'MergingTree'+--+data EmptyMergingTree = ObviouslyEmptyLevelMerge+ | ObviouslyEmptyUnionMerge+ | NonObviouslyEmptyLevelMerge EmptyMergingTree+ | NonObviouslyEmptyUnionMerge [EmptyMergingTree]+ deriving stock (Eq, Show)++instance Arbitrary EmptyMergingTree where+ arbitrary =+ sized $ \sz ->+ frequency $+ take (1 + sz)+ [ (1, pure ObviouslyEmptyLevelMerge)+ , (1, pure ObviouslyEmptyUnionMerge)+ , (2, NonObviouslyEmptyLevelMerge <$> resize (sz `div` 2) arbitrary)+ , (2, NonObviouslyEmptyUnionMerge <$> resize (sz `div` 2) arbitrary)+ ]+ shrink ObviouslyEmptyLevelMerge = []+ shrink ObviouslyEmptyUnionMerge = [ObviouslyEmptyLevelMerge]+ shrink (NonObviouslyEmptyLevelMerge mt) = ObviouslyEmptyLevelMerge+ : [ NonObviouslyEmptyLevelMerge mt'+ | mt' <- shrink mt ]+ shrink (NonObviouslyEmptyUnionMerge mt) = ObviouslyEmptyUnionMerge+ : [ NonObviouslyEmptyUnionMerge mt'+ | mt' <- shrink mt ]++mkEmptyMergingTree :: EmptyMergingTree -> IO (Ref (MergingTree IO h))+mkEmptyMergingTree ObviouslyEmptyLevelMerge = newPendingLevelMerge [] Nothing+mkEmptyMergingTree ObviouslyEmptyUnionMerge = newPendingUnionMerge []+mkEmptyMergingTree (NonObviouslyEmptyLevelMerge emt) = do+ mt <- mkEmptyMergingTree emt+ mt' <- newPendingLevelMerge [] (Just mt)+ releaseRef mt+ pure mt'+mkEmptyMergingTree (NonObviouslyEmptyUnionMerge emts) = do+ mts <- mapM mkEmptyMergingTree emts+ mt' <- newPendingUnionMerge mts+ mapM_ releaseRef mts+ pure mt'++{-------------------------------------------------------------------------------+ Lookup+-------------------------------------------------------------------------------}++prop_lookupTree ::+ forall h.+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> V.Vector SerialisedKey+ -> MergingTreeData SerialisedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_lookupTree hfs hbio keys mtd = do+ let path = FS.mkFsPath []+ counter <- newUniqCounter 0+ withMergingTree hfs hbio resolveVal testSalt runParams path counter mtd $ \tree -> do+ arenaManager <- newArenaManager+ withActionRegistry $ \reg -> do+ res <- fetchBlobs =<< lookupsIO reg arenaManager tree+ pure $+ normalise res+ === normalise (modelLookup (modelFoldMergingTree mtd) keys)+ where+ fetchBlobs ::+ V.Vector (Maybe (Entry v (WeakBlobRef IO h)))+ -> IO (V.Vector (Maybe (Entry v SerialisedBlob)))+ fetchBlobs = traverse (traverse (traverse (readWeakBlobRef hfs)))++ -- the lookup accs might be different between implementation and model+ -- (Nothing vs. Just Delete, Insert vs. Mupsert), but this doesn't matter+ -- for the final result of the lookup+ normalise = V.map toLookupResult++ toLookupResult Nothing = Nothing+ toLookupResult (Just e) = case e of+ Entry.Insert v -> Just (v, Nothing)+ Entry.InsertWithBlob v b -> Just (v, Just b)+ Entry.Upsert v -> Just (v, Nothing)+ Entry.Delete -> Nothing++ lookupsIO reg mgr tree =+ isStructurallyEmpty tree >>= \case+ True ->+ -- if the tree was empty, then the model should also have no results+ pure $ V.map (const Nothing) keys+ False -> do+ batches <- buildLookupTree reg tree+ results <- mapMStrict (performLookups mgr) batches+ acc <- foldLookupTree resolveVal results+ traverse_ (traverse_ (delayedCommit reg . releaseRef)) batches+ pure acc++ performLookups mgr runs =+ Async.async $+ Lookup.lookupsIO+ hbio+ mgr+ resolveVal+ testSalt+ runs+ (fmap (\(DeRef r) -> Run.runFilter r) runs)+ (fmap (\(DeRef r) -> Run.runIndex r) runs)+ (fmap (\(DeRef r) -> Run.runKOpsFile r) runs)+ keys++type SerialisedEntry = Entry SerialisedValue SerialisedBlob+type LookupAcc' = V.Vector (Maybe (Entry SerialisedValue SerialisedBlob))++modelLookup :: Map SerialisedKey SerialisedEntry -> V.Vector SerialisedKey -> LookupAcc'+modelLookup m = V.map (\k -> Map.lookup k m)++modelFoldMergingTree :: SerialisedMergingTreeData -> Map SerialisedKey SerialisedEntry+modelFoldMergingTree = goMergingTree+ where+ goMergingTree :: SerialisedMergingTreeData -> Map SerialisedKey SerialisedEntry+ goMergingTree = \case+ CompletedTreeMergeData r ->+ unRunData r+ OngoingTreeMergeData mr ->+ goMergingRun mr+ PendingLevelMergeData prs t ->+ modelMerge MR.MergeLevel (map goPreExistingRun prs <> map goMergingTree (toList t))+ PendingUnionMergeData ts ->+ modelMerge MR.MergeUnion (map goMergingTree ts)++ goPreExistingRun = \case+ PreExistingRunData r -> unRunData r+ PreExistingMergingRunData mr -> goMergingRun mr++ goMergingRun :: MR.IsMergeType t => SerialisedMergingRunData t -> Map SerialisedKey SerialisedEntry+ goMergingRun = \case+ CompletedMergeData _ r -> unRunData r+ OngoingMergeData mt rs -> modelMerge mt (map (unRunData . toRunData) rs)++modelMerge :: (Ord k, MR.IsMergeType t) => t -> [Map k SerialisedEntry] -> Map k SerialisedEntry+modelMerge mt = handleDeletes . Map.unionsWith (combine resolveVal)+ where+ handleDeletes = if MR.isLastLevel mt then Map.filter (/= Entry.Delete) else id+ combine = if MR.isUnion mt then Entry.combineUnion else Entry.combine++resolveVal :: SerialisedValue -> SerialisedValue -> SerialisedValue+resolveVal (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)++{-------------------------------------------------------------------------------+ Supplying Credits+-------------------------------------------------------------------------------}++prop_supplyCredits ::+ forall h.+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> MR.CreditThreshold+ -> NonEmpty MR.MergeCredits+ -> MergingTreeData SerialisedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_supplyCredits hfs hbio threshold credits mtd = do+ FS.createDirectory hfs setupPath+ FS.createDirectory hfs (FS.mkFsPath ["active"])+ counter <- newUniqCounter 0+ withMergingTree hfs hbio resolveVal testSalt runParams setupPath counter mtd $ \tree -> do+ (MR.MergeDebt initialDebt, _) <- remainingMergeDebt tree+ props <- for credits $ \c -> do+ (MR.MergeDebt debt, _) <- remainingMergeDebt tree+ if debt <= 0+ then+ pure $ property True+ else do+ leftovers <-+ supplyCredits hfs hbio resolveVal testSalt runParams threshold root counter tree c+ (MR.MergeDebt debt', _) <- remainingMergeDebt tree+ pure $+ -- semi-useful, but mainly tells us in how many steps we supplied+ tabulate "supplied credits" [showPowersOf10 (fromIntegral c)] $+ counterexample (show (debt, leftovers, debt')) $ conjoin [+ counterexample "negative values" $+ debt >= 0 && leftovers >= 0 && debt' >= 0+ , counterexample "did not reduce debt sufficiently" $+ debt' <= debt - (c - leftovers)+ ]+ (MR.MergeDebt finalDebt, _) <- remainingMergeDebt tree+ pure $+ labelDebt initialDebt finalDebt $+ conjoin (toList props)+ where+ root = Paths.SessionRoot (FS.mkFsPath [])+ setupPath = FS.mkFsPath ["setup"] -- separate dir, so file paths in errors+ -- are identifiable as created in setup+ --+ labelDebt initial final+ | initial == 0 = label "trivial"+ | final == 0 = label "completed"+ | otherwise = label "incomplete"++instance Arbitrary MR.MergeCredits where+ arbitrary = MR.MergeCredits . getPositive <$> arbitrary+ shrink (MR.MergeCredits c) = [MR.MergeCredits c' | c' <- shrink c, c' > 0]++instance Arbitrary MR.CreditThreshold where+ arbitrary = coerce (arbitrary @MR.MergeCredits)+ shrink = coerce (shrink @MR.MergeCredits)+ -- TODO: does this make sense? in a way a larger threshold is "simpler".
@@ -0,0 +1,206 @@+{-# LANGUAGE OverloadedStrings #-}+module Test.Database.LSMTree.Internal.PageAcc (tests) where++import Control.Exception+import Control.Monad (forM_)+import Control.Monad.ST.Strict (ST, runST, stToIO)+import qualified Data.ByteString as BS+import Data.Maybe (isJust)+import Data.Word++import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry (Entry (..))+import Database.LSMTree.Internal.PageAcc+import Database.LSMTree.Internal.PageAcc1+import Database.LSMTree.Internal.RawPage (RawPage)+import Database.LSMTree.Internal.Serialise++import Database.LSMTree.Extras.NoThunks (propNoThunks)+import qualified Database.LSMTree.Extras.ReferenceImpl as Ref+import Test.Util.RawPage (propEqualRawPages)++import Test.QuickCheck+import Test.QuickCheck.Instances ()+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)++tests :: TestTree+tests =+ testGroup "Database.LSMTree.Internal.PageAcc" $+ [ testProperty "vs reference impl" prop_vsReferenceImpl+ , testProperty "maxPageKeys == 759" (maxPageKeys === 759)+ ]++ ++ [ testProperty+ ("example-" ++ show (n :: Int) ++ [ a | length exs > 1 ])+ (prop_vsReferenceImpl (Ref.PageContentMaybeOverfull kops))+ | (n, exs) <- zip [0..] examples+ , (a, kops) <- zip ['a'..] exs+ ]++ ++ [ testProperty "+PageAcc1" prop_vsRefWithPageAcc1+ ]++ ++ [ testProperty "prop_noThunks_newPageAcc" prop_noThunks_newPageAcc+ , testProperty "prop_noThunks_pageAccAddElem" prop_noThunks_pageAccAddElem+ , testProperty "prop_noThunks_resetPageAcc" prop_noThunks_resetPageAcc+ ]++ where+ examples = example0123 ++ [example4s, example5s, example6s, example7s]++ example0123 =+ map (:[])+ [ []+ , [(Ref.Key "foobar", Ref.Delete)]+ , [(Ref.Key "foobar", Ref.Insert (Ref.Value "value")+ (Just (Ref.BlobRef 111 333)))]+ , [(Ref.Key "\NUL", Ref.Delete), (Ref.Key "\SOH", Ref.Delete)]+ ]++ example4s = [ [(Ref.Key "", Ref.Insert (Ref.Value (BS.replicate sz 120))+ Nothing)]+ | sz <- [4063..4065] ]++ example5s = [ [ (Ref.Key "",Ref.Delete)+ , (Ref.Key "k", Ref.Insert (Ref.Value (BS.replicate sz 120))+ Nothing) ]+ | sz <- [4060..4062] ]++ example6s = [ [(Ref.Key "", Ref.Insert (Ref.Value (BS.replicate sz 120))+ (Just (Ref.BlobRef 111 333))) ]+ | sz <- [4051..4053] ]++ example7s = [ (replicate maxPageKeys (Ref.Key " ",Ref.Delete))+ , (replicate (maxPageKeys+1) (Ref.Key " ",Ref.Delete))+ , (replicate (maxPageKeys+1) (Ref.Key "", Ref.Delete))+ ]++maxPageKeys :: Int+maxPageKeys =+ go 0 (Ref.pageSizeEmpty Ref.DiskPage4k)+ where+ go s ps =+ case Ref.pageSizeAddElem (Ref.Key " ", Ref.Delete) ps of+ Nothing -> s+ Just ps' -> go (s + 1) ps'++prop_vsReferenceImpl :: Ref.PageContentMaybeOverfull -> Property+prop_vsReferenceImpl (Ref.PageContentMaybeOverfull kops) =+ case (refImpl, realImpl) of+ (Just (lhs, _), Just rhs) -> propEqualRawPages lhs rhs+ (Nothing, Nothing) -> label "overflow" $+ property True++ -- Special cases where we allow a test pass.+ (Just _, Nothing)+ -- The PageAcc does not support single-key/op pairs that overflow onto+ -- multiple pages. That case is handled by PageAcc1.+ | [_] <- kops+ , Just page <- Ref.encodePage Ref.DiskPage4k kops+ , Ref.pageDiskPages page > 1+ -> label "PageAcc1 special case" $+ property True++ -- PageAcc (quite reasonably) assumes that keys are not all empty+ -- (since in practice they'll be distinct) and thus it can impose an+ -- upper bound on the number of keys in a page. It's possible to+ -- construct test cases with empty keys that exceed the buffer size.+ | length kops >= maxPageKeys+ -> label "max number of keys reached" $+ property True++ _ -> property False+ where+ refImpl = Ref.toRawPageMaybeOverfull (Ref.PageContentMaybeOverfull kops)+ realImpl = toRawPageViaPageAcc [ (Ref.toSerialisedKey k, Ref.toEntry op)+ | (k,op) <- kops ]+++-- | This is like 'prop_vsReferenceImpl' bus used _both_ @PageAcc@ and+-- @PageAcc1@ together to fill in the special cases.+--+prop_vsRefWithPageAcc1 :: Ref.PageContentMaybeOverfull -> Property+prop_vsRefWithPageAcc1 (Ref.PageContentMaybeOverfull kops) =+ case (refImpl, realImpl) of+ (Just (lhs, loverflow),+ Just (rhs, roverflow)) ->+ label (show (length loverflow) ++ " overflow pages") $+ (if isJust (pageAcc1SpecialCase kops) then label "PageAcc1" else id)+ (propEqualRawPages lhs rhs)+ .&&. counterexample "overflow pages do not match"+ (loverflow === roverflow)++ (Nothing, Nothing) ->+ label "overfull" $ property True++ -- Special cases are a subset of those above in 'prop_vsReferenceImpl'.+ (Just _, Nothing) | length kops >= maxPageKeys ->+ label "max number of keys reached" $ property True++ _ -> property False+ where+ refImpl = Ref.toRawPageMaybeOverfull (Ref.PageContentMaybeOverfull kops)++ -- Use whichever implementation is appropriate:+ realImpl+ | Just (k,op) <- pageAcc1SpecialCase kops+ = Just (singletonPage (Ref.toSerialisedKey k) (Ref.toEntry op))++ | otherwise+ = (\rp -> (rp, [])) <$>+ toRawPageViaPageAcc [ (Ref.toSerialisedKey k, Ref.toEntry op)+ | (k,op) <- kops ]++ pageAcc1SpecialCase [(k, op)] | op /= Ref.Delete = Just (k, op)+ pageAcc1SpecialCase _ = Nothing+++-- | Use a 'PageAcc' to try to make a 'RawPage' from key\/op pairs. It will+-- return @Nothing@ if the key\/op pairs would not all fit in a page.+--+toRawPageViaPageAcc :: [(SerialisedKey, Entry SerialisedValue BlobSpan)]+ -> Maybe RawPage+toRawPageViaPageAcc kops0 =+ runST $ do+ acc <- newPageAcc+ go acc kops0+ where+ go acc [] = Just <$> serialisePageAcc acc+ go acc ((k,op):kops) = do+ added <- pageAccAddElem acc k op+ if added+ then go acc kops+ else pure Nothing++{-------------------------------------------------------------------------------+ NoThunks+-------------------------------------------------------------------------------}++prop_noThunks_newPageAcc :: Property+prop_noThunks_newPageAcc = once $ ioProperty $ do+ pa <- stToIO newPageAcc+ propNoThunks pa++prop_noThunks_pageAccAddElem :: Property+prop_noThunks_pageAccAddElem = once $ ioProperty $ do+ pa <- stToIO $ do+ pa <- newPageAcc+ pageAccAddElemN pa 10+ pure pa+ propNoThunks pa++prop_noThunks_resetPageAcc :: Property+prop_noThunks_resetPageAcc = once $ ioProperty $ do+ pa <- stToIO $ do+ pa <- newPageAcc+ pageAccAddElemN pa 10+ resetPageAcc pa+ pure pa+ propNoThunks pa++pageAccAddElemN :: PageAcc s -> Word64 -> ST s ()+pageAccAddElemN pa n = do+ forM_ [1..n] $ \(x :: Word64) -> do+ b <- pageAccAddElem pa (serialiseKey x) (Insert (serialiseValue x))+ assert b $ pure ()
@@ -0,0 +1,49 @@+{-# LANGUAGE OverloadedStrings #-}+module Test.Database.LSMTree.Internal.PageAcc1 (tests) where++import qualified Data.ByteString as BS++import Database.LSMTree.Internal.PageAcc1++import qualified Database.LSMTree.Extras.ReferenceImpl as Ref++import Test.QuickCheck+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)+import Test.Util.RawPage++tests :: TestTree+tests =+ testGroup "Database.LSMTree.Internal.PageAcc1" $+ [ testProperty "vs reference impl" prop_vsReferenceImpl ]++ ++ [ testProperty+ ("example-" ++ show (n :: Int) ++ [a])+ (prop_vsReferenceImpl (Ref.PageContentSingle (Ref.Key "") op))+ | (n,exs) <- zip [1..] examples+ , (a, op) <- zip ['a'..] exs+ ]+ where+ examples :: [[Ref.Operation]]+ examples = [example1s, example2s, example3s]+ example1s = [ Ref.Insert (Ref.Value (BS.replicate sz 120)) Nothing+ | sz <- [4064..4066] ]++ example2s = [ Ref.Insert (Ref.Value (BS.replicate sz 120))+ (Just (Ref.BlobRef 3 5))+ | sz <- [4050..4052] ]++ example3s = [ Ref.Mupsert (Ref.Value (BS.replicate sz 120))+ | sz <- [4064..4066] ]++prop_vsReferenceImpl :: Ref.PageContentSingle -> Property+prop_vsReferenceImpl (Ref.PageContentSingle k op) =+ op /= Ref.Delete ==>+ label (show (length loverflow) ++ " overflow pages") $+ propEqualRawPages lhs rhs+ .&&. counterexample "overflow pages do not match"+ (loverflow === roverflow)+ where+ (lhs, loverflow) = Ref.toRawPage $ Ref.PageContentFits [(k, op)]+ (rhs, roverflow) = singletonPage (Ref.toSerialisedKey k) (Ref.toEntry op)+
@@ -0,0 +1,120 @@+{-# LANGUAGE OverloadedLists #-}++module Test.Database.LSMTree.Internal.RawBytes (tests) where++import Data.Bits (Bits (shiftL))+import qualified Data.List as List+import qualified Data.Vector.Primitive as VP+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.RawBytes (RawBytes (RawBytes))+import qualified Database.LSMTree.Internal.RawBytes as RB+import Test.QuickCheck (Property, classify, collect, mapSize,+ withDiscardRatio, withMaxSuccess, (.||.), (===), (==>))+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)++-- * Tests++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.RawBytes" $+ [+ testGroup "Eq laws" $+ [+ testProperty "Reflexivity" prop_eqReflexivity,+ testProperty "Symmetry" prop_eqSymmetry,+ testProperty "Transitivity" prop_eqTransitivity,+ testProperty "Negation" prop_eqNegation+ ],+ testGroup "Ord laws" $+ [+ testProperty "Comparability" prop_ordComparability,+ testProperty "Transitivity" prop_ordTransitivity,+ testProperty "Reflexivity" prop_ordReflexivity,+ testProperty "Antisymmetry" prop_ordAntisymmetry+ ],+ testProperty "prop_topBits64" prop_topBits64,+ testProperty "prop_topBits64_default0s" prop_topBits64_default0s+ ]++-- * Utilities++twoBlocksProp :: String -> RawBytes -> RawBytes -> Property -> Property+twoBlocksProp msgAddition block1 block2+ = withMaxSuccess 10000 .+ classify (block1 == block2) ("equal blocks" ++ msgAddition)++withFirstBlockSizeInfo :: RawBytes -> Property -> Property+withFirstBlockSizeInfo firstBlock+ = collect ("Size of first block is " ++ show (RB.size firstBlock))++-- * Properties to test++-- ** 'Eq' laws++prop_eqReflexivity :: RawBytes -> Property+prop_eqReflexivity block = block === block++prop_eqSymmetry :: RawBytes -> RawBytes -> Property+prop_eqSymmetry block1 block2 = twoBlocksProp "" block1 block2 $+ (block1 == block2) === (block2 == block1)++prop_eqTransitivity :: Property+prop_eqTransitivity = mapSize (const 3) $+ withDiscardRatio 1000 $+ untunedProp+ where++ untunedProp :: RawBytes -> RawBytes -> RawBytes -> Property+ untunedProp block1 block2 block3+ = withFirstBlockSizeInfo block1 $+ block1 == block2 && block2 == block3 ==> block1 === block3++prop_eqNegation :: RawBytes -> RawBytes -> Property+prop_eqNegation block1 block2 = twoBlocksProp "" block1 block2 $+ (block1 /= block2) === not (block1 == block2)++-- ** 'Ord' laws++prop_ordComparability :: RawBytes -> RawBytes -> Property+prop_ordComparability block1 block2 = twoBlocksProp "" block1 block2 $+ block1 <= block2 .||. block2 <= block1++prop_ordTransitivity :: RawBytes -> RawBytes -> RawBytes -> Property+prop_ordTransitivity block1 block2 block3+ = twoBlocksProp " front-side" block1 block2 $+ twoBlocksProp " rear-side" block2 block3 $+ twoBlocksProp " at the edges" block1 block3 $+ block1 <= block2 && block2 <= block3 ==> block1 <= block3++prop_ordReflexivity :: RawBytes -> Bool+prop_ordReflexivity block = block <= block++prop_ordAntisymmetry :: Property+prop_ordAntisymmetry = mapSize (const 4) $+ withDiscardRatio 100 $+ untunedProp+ where++ untunedProp :: RawBytes -> RawBytes -> Property+ untunedProp block1 block2+ = withFirstBlockSizeInfo block1 $+ block1 <= block2 && block2 <= block1 ==> block1 === block2++{-------------------------------------------------------------------------------+ Accessors+-------------------------------------------------------------------------------}++-- | Compare 'topBits64' against a model+prop_topBits64 :: RawBytes -> Property+prop_topBits64 x@(RawBytes v) =+ expected === RB.topBits64 x+ where+ expected =+ let ws = take 8 (VP.toList v ++ repeat 0)+ in List.foldl' (\acc w -> acc `shiftL` 8 + fromIntegral w) 0 ws++-- | If @x@ has fewer than 8 bytes, then all missing bits in the result default+-- to 0s.+prop_topBits64_default0s :: RawBytes -> Property+prop_topBits64_default0s x =+ RB.topBits64 x === RB.topBits64 (x <> mconcat (replicate 8 [0]))
@@ -0,0 +1,57 @@+module Test.Database.LSMTree.Internal.RawOverflowPage (+ -- * Main test tree+ tests,+) where++import qualified Data.Primitive.ByteArray as BA+import qualified Data.Vector.Primitive as VP++import Test.QuickCheck+import Test.QuickCheck.Instances ()+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck++import Database.LSMTree.Extras.Generators (LargeRawBytes (..))+import Database.LSMTree.Internal.BitMath+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage++tests :: TestTree+tests =+ testGroup "Database.LSMTree.Internal.RawOverflowPage"+ [ testProperty "RawBytes prefix to RawOverflowPage"+ prop_rawBytesToRawOverflowPage+ , testProperty "RawBytes to [RawOverflowPage]"+ prop_rawBytesToRawOverflowPages+ ]++-- | Converting up to the first 4096 bytes of a 'RawBytes' to an+-- 'RawOverflowPage' and back gives us the original, padded with zeros to a+-- multiple of the page size.+prop_rawBytesToRawOverflowPage :: LargeRawBytes -> Property+prop_rawBytesToRawOverflowPage+ (LargeRawBytes bytes@(RB.RawBytes (VP.Vector off len ba))) =+ label (if RB.size bytes >= 4096 then "large" else "small") $+ label (if BA.isByteArrayPinned ba then "pinned" else "unpinned") $+ label (if off == 0 then "offset 0" else "offset non-0") $++ rawOverflowPageRawBytes (makeRawOverflowPage ba off (min len 4096))+ === RB.take 4096 bytes <> padding+ where+ padding = RB.fromVector (VP.replicate paddinglen 0)+ paddinglen = 4096 - (min len 4096)+++-- | Converting the bytes to @[RawOverflowPage]@ and back gives us the original+-- bytes, padded with zeros to a multiple of the page size.+--+prop_rawBytesToRawOverflowPages :: LargeRawBytes -> Property+prop_rawBytesToRawOverflowPages (LargeRawBytes bytes) =+ length pages === roundUpToPageSize (RB.size bytes) `div` 4096+ .&&. mconcat (map rawOverflowPageRawBytes pages) === bytes <> padding+ where+ pages = rawBytesToOverflowPages bytes+ padding = RB.fromVector (VP.replicate paddinglen 0)+ paddinglen = let trailing = RB.size bytes `mod` 4096 in+ if trailing == 0 then 0 else 4096 - trailing+
@@ -0,0 +1,537 @@+{-# LANGUAGE BlockArguments #-}+{-# LANGUAGE OverloadedStrings #-}+module Test.Database.LSMTree.Internal.RawPage (+ -- * Main test tree+ tests,+) where++import Control.DeepSeq (deepseq)+import Control.Monad (when)+import Control.Monad.State (MonadState (..), StateT, evalStateT, lift)+import qualified Data.ByteString as BS+import Data.Functor ((<&>))+import Data.Maybe (fromMaybe)+import Data.Primitive.ByteArray (byteArrayFromList)+import qualified Data.Vector as V+import qualified Data.Vector.Primitive as VP+import Data.Word (Word16, Word64, Word8)+import GHC.Word (byteSwap16)+import Test.QuickCheck.Instances ()+import Test.Tasty (TestName, TestTree, localOption, testGroup)+import Test.Tasty.HUnit (testCase, (@=?))+import Test.Tasty.QuickCheck+import Test.Util.RawPage++import qualified Database.LSMTree.Extras.ReferenceImpl as Ref+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import qualified Database.LSMTree.Internal.Entry as Entry+import Database.LSMTree.Internal.RawBytes (RawBytes (..))+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Serialise++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.RawPage"+ [ testCase "empty" $ do+ -- directory:+ -- * 0 keys,+ -- * 0 blobs+ -- * key's offset offset at 8 bytes (as bitmaps are empty)+ -- * nil spare+ --+ -- and+ -- offset past last value.+ let bytes :: [Word16]+ bytes = [0, 0, 8, 0, 10]++ let page = makeRawPage (byteArrayFromList bytes) 0++ (page, []) @=? Ref.toRawPage (Ref.PageContentFits [])+ rawPageNumKeys page @=? 0+ rawPageNumBlobs page @=? 0+ rawPageKeyOffsets page @=? VP.fromList [10]+ rawPageValueOffsets page @=? VP.fromList [10]++ , testCase "single-insert" $ do+ let bytes :: [Word16]+ bytes =+ [ 1, 0, 24, 0 -- directory+ , 0, 0, 0, 0, 0, 0, 0, 0 -- ...+ , 32 -- key offsets+ , 34, 36 -- value offsets+ , 0x00 -- single key case (top bits of 32bit offset)+ , byteSwap16 0x4243 -- key+ , byteSwap16 0x8899 -- value+ ]++ let page = makeRawPage (byteArrayFromList bytes) 0++ assertEqualRawPages+ page+ ((fst . Ref.toRawPage . Ref.PageContentFits)+ [(Ref.Key "\x42\x43", Ref.Insert (Ref.Value "\x88\x99") Nothing)])+ rawPageNumKeys page @=? 1+ rawPageNumBlobs page @=? 0+ rawPageKeyOffsets page @=? VP.fromList [32, 34]+ rawPageValueOffsets1 page @=? (34, 36)+ rawPageHasBlobSpanAt page 0 @=? 0+ rawPageOpAt page 0 @=? 0+ rawPageKeys page @=? V.singleton (SerialisedKey "\x42\x43")++ rawPageLookup page (SerialisedKey "\x42\x43")+ @=? LookupEntry (Entry.Insert (SerialisedValue "\x88\x99"))++ , testCase "single-insert-blobspan" $ do+ let bytes :: [Word16]+ bytes =+ [ 1, 1, 36, 0 -- directory+ , 1, 0, 0, 0 -- has blobspan+ , 0, 0, 0, 0 -- ops+ , 0xff, 0, 0, 0+ , 0xfe, 0+ , 44 -- key offsets+ , 46, 48 -- value offsets+ , 0x00 -- single key case (top bits of 32bit offset)+ , byteSwap16 0x4243 -- key+ , byteSwap16 0x8899 -- value+ ]++ let page = makeRawPage (byteArrayFromList bytes) 0++ (page, [])+ @=? (Ref.toRawPage . Ref.PageContentFits)+ [(Ref.Key "\x42\x43",+ Ref.Insert (Ref.Value "\x88\x99")+ (Just (Ref.BlobRef 0xff 0xfe)))]+ rawPageNumKeys page @=? 1+ rawPageNumBlobs page @=? 1+ rawPageKeyOffsets page @=? VP.fromList [44, 46]+ rawPageValueOffsets1 page @=? (46, 48)+ rawPageHasBlobSpanAt page 0 @=? 1+ rawPageBlobSpanIndex page 0 @=? BlobSpan 0xff 0xfe+ rawPageOpAt page 0 @=? 0+ rawPageKeys page @=? V.singleton (SerialisedKey "\x42\x43")++ rawPageLookup page (SerialisedKey "\x42\x43")+ @=? LookupEntry (Entry.InsertWithBlob (SerialisedValue "\x88\x99")+ (BlobSpan 0xff 0xfe))++ , testCase "single-delete" $ do+ let bytes :: [Word16]+ bytes =+ [ 1, 0, 24, 0 -- directory+ , 0, 0, 0, 0, 2, 0, 0, 0 -- ...+ , 32 -- key offsets+ , 34, 34 -- value offsets+ , 0x00 -- single key case (top bits of 32bit offset)+ , byteSwap16 0x4243 -- key+ ]++ let page = makeRawPage (byteArrayFromList bytes) 0++ (page, [])+ @=? (Ref.toRawPage . Ref.PageContentFits)+ [(Ref.Key "\x42\x43", Ref.Delete)]+ rawPageNumKeys page @=? 1+ rawPageNumBlobs page @=? 0+ rawPageKeyOffsets page @=? VP.fromList [32, 34]+ rawPageValueOffsets1 page @=? (34, 34)+ rawPageHasBlobSpanAt page 0 @=? 0+ rawPageOpAt page 0 @=? 2+ rawPageKeys page @=? V.singleton (SerialisedKey "\x42\x43")++ rawPageLookup page (SerialisedKey "\x42\x43")+ @=? LookupEntry Entry.Delete++ , testCase "double-mupsert" $ do+ let bytes :: [Word16]+ bytes =+ [ 2, 0, 24, 0 -- directory+ , 0, 0, 0, 0, 5, 0, 0, 0 -- ...+ , 34, 36 -- key offsets+ , 38, 40, 42 -- value offsets+ , byteSwap16 0x4243 -- key 1+ , byteSwap16 0x5253 -- key 2+ , byteSwap16 0x4445 -- value 1+ , byteSwap16 0x5455 -- value 2+ ]++ let page = makeRawPage (byteArrayFromList bytes) 0++ (page, [])+ @=? (Ref.toRawPage . Ref.PageContentFits)+ [(Ref.Key "\x42\x43", Ref.Mupsert (Ref.Value "\x44\x45")),+ (Ref.Key "\x52\x53", Ref.Mupsert (Ref.Value "\x54\x55"))]+ rawPageNumKeys page @=? 2+ rawPageNumBlobs page @=? 0+ rawPageKeyOffsets page @=? VP.fromList [34, 36, 38]+ rawPageValueOffsets page @=? VP.fromList [38, 40, 42]+ rawPageHasBlobSpanAt page 0 @=? 0+ rawPageHasBlobSpanAt page 1 @=? 0+ rawPageOpAt page 0 @=? 1+ rawPageOpAt page 1 @=? 1+ rawPageKeys page @=? V.fromList [SerialisedKey "\x42\x43",+ SerialisedKey "\x52\x53"]+ rawPageValues page @=? V.fromList [SerialisedValue "\x44\x45",+ SerialisedValue "\x54\x55"]++ rawPageLookup page (SerialisedKey "\x52\x53")+ @=? LookupEntry (Entry.Upsert (SerialisedValue "\x54\x55"))+ rawPageLookup page (SerialisedKey "\x99\x99")+ @=? LookupEntryNotPresent++ , testProperty "toRawPage" prop_toRawPage+ , testProperty "keys" prop_keys+ , testProperty "values" prop_values+ , testProperty "hasblobspans" prop_hasblobspans+ , testProperty "blobspans" prop_blobspans+ , testProperty "ops" prop_ops+ , testProperty "entries" prop_entries_exists+ , testProperty "missing" prop_entries_all+ , testProperty "big-insert" prop_big_insert+ , testProperty "entry" prop_single_entry+ , testProperty "rawPageOverflowPages" prop_rawPageOverflowPages+ , testProperty "from/to reference impl" prop_fromToReferenceImpl+ , testPropertyExtraCoverage "rawPageFindKey >" prop_findKey_index_law_GT+ , testPropertyExtraCoverage "rawPageFindKey <=" prop_findKey_index_law_LT+ , testPropertyExtraCoverage "rawPageFindKey missing" prop_findKey_index_law_missing+ ]++testPropertyExtraCoverage :: Testable prop => TestName -> prop -> TestTree+testPropertyExtraCoverage tName = localOption (QuickCheckTests 1024) . testProperty tName++prop_toRawPage :: Ref.PageContentFits -> Property+prop_toRawPage p =+ let (_rawpage, overflowPages) = Ref.toRawPage p+ in tabulate "toRawPage number of overflowPages"+ [ show (length overflowPages) ] $+ property (deepseq overflowPages True)++prop_keys :: Ref.PageContentFits -> Property+prop_keys (Ref.PageContentFits kops) =+ fromIntegral (rawPageNumKeys rawpage) === length kops+ .&&.+ rawPageKeys rawpage+ === V.fromList [ Ref.toSerialisedKey k | (k, _) <- kops ]+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++prop_values :: Ref.PageContentFits -> Property+prop_values (Ref.PageContentFits kops) =+ length kops /= 1 ==>+ rawPageValues rawpage+ === V.fromList [ Ref.toSerialisedValue (extractValue op) | (_, op) <- kops ]+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++ extractValue (Ref.Insert v _) = v+ extractValue (Ref.Mupsert v) = v+ extractValue Ref.Delete = Ref.Value BS.empty++prop_blobspans :: Ref.PageContentFits -> Property+prop_blobspans (Ref.PageContentFits kops) =+ [ rawPageBlobSpanIndex rawpage i+ | i <- [0 .. fromIntegral (rawPageNumBlobs rawpage) - 1] ]+ === [ Ref.toBlobSpan b | (_, Ref.Insert _ (Just b)) <- kops ]+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++prop_hasblobspans :: Ref.PageContentFits -> Property+prop_hasblobspans (Ref.PageContentFits kops) =+ [ rawPageHasBlobSpanAt rawpage i /= 0 | i <- [0 .. length kops - 1] ]+ === [ opHasBlobSpan op | (_, op) <- kops ]+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++ opHasBlobSpan (Ref.Insert _ (Just _)) = True+ opHasBlobSpan _ = False++prop_ops :: Ref.PageContentFits -> Property+prop_ops (Ref.PageContentFits kops) =+ [ rawPageOpAt rawpage i | i <- [0 .. length kops - 1] ]+ === [ fromOp op | (_, op) <- kops ]+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++ fromOp :: Ref.Operation -> Word64+ fromOp Ref.Insert {} = 0+ fromOp Ref.Delete {} = 2+ fromOp Ref.Mupsert {} = 1++prop_rawPageOverflowPages :: Ref.PageContentFits -> Property+prop_rawPageOverflowPages (Ref.PageContentFits kops) =+ rawPageOverflowPages page === length overflowPages+ where+ (page, overflowPages) = Ref.toRawPage (Ref.PageContentFits kops)++prop_fromToReferenceImpl :: Ref.PageContentFits -> Property+prop_fromToReferenceImpl (Ref.PageContentFits kops) =+ -- serialise using the reference impl+ -- deserialise using the real impl and convert back to the reference types+ Ref.fromRawPage (Ref.toRawPage (Ref.PageContentFits kops))+ === Ref.PageContentFits kops++prop_entries_exists :: Ref.PageContentOrdered -> Property+prop_entries_exists (Ref.PageContentOrdered kops) =+ length kops > 1 ==>+ foldr1 (.&&.)+ [ rawPageLookup rawpage (Ref.toSerialisedKey k)+ === LookupEntry (Ref.toEntry op)+ | (k,op) <- kops ]+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++prop_entries_all :: Ref.PageContentOrdered -> Ref.Key -> Property+prop_entries_all (Ref.PageContentOrdered kops) k =+ length kops /= 1 ==>+ rawPageLookup rawpage (Ref.toSerialisedKey k)+ === maybe LookupEntryNotPresent+ (LookupEntry . Ref.toEntry)+ (lookup k kops)+ where+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++prop_big_insert :: Ref.Key -> Maybe Ref.BlobRef -> Property+prop_big_insert k mblobref =+ rawPageLookup rawpage (Ref.toSerialisedKey k)+ ===+ let (prefixlen, suffixlen) =+ fromMaybe+ (error "expected overflow pages")+ (Ref.pageOverflowPrefixSuffixLen+ =<< Ref.encodePage Ref.DiskPage4k kops)+ in LookupEntryOverflow (Ref.toEntryPrefix op prefixlen)+ (fromIntegral suffixlen)+ where+ v = Ref.Value (BS.replicate 5000 42)+ op = Ref.Insert v mblobref+ kops = [(k, op)]+ rawpage = fst $ Ref.toRawPage (Ref.PageContentFits kops)++prop_single_entry :: Ref.PageContentSingle -> Property+prop_single_entry (Ref.PageContentSingle k op) =+ label ("pages " ++ show (length overflowPages)) $++ rawPageLookup rawpage (Ref.toSerialisedKey k)+ ===+ case Ref.pageOverflowPrefixSuffixLen+ =<< Ref.encodePage Ref.DiskPage4k [(k, op)] of+ Nothing ->+ LookupEntry (Ref.toEntry op)++ Just (prefixlen, suffixlen) ->+ LookupEntryOverflow (Ref.toEntryPrefix op prefixlen)+ (fromIntegral suffixlen)+ where+ (rawpage, overflowPages) = Ref.toRawPage (Ref.PageContentFits [(k, op)])++prop_findKey_index_law_GT :: RawPageAndOffset -> Property+prop_findKey_index_law_GT (RawPageAndOffset rawpage key locInfo) =+ recordOffsetLocation locInfo $+ findKeyIndexCounterExample+ key+ rawpage+ "∀ key page. maybe True (key > ) (getRawPageIndexKey . rawPageIndex page . pred =<< rawPageFindKey page key)" $+ maybe+ (discard)+ (property . (key <=))+ (getRawPageIndexKey . rawPageIndex rawpage =<< rawPageFindKey rawpage key)++prop_findKey_index_law_LT :: RawPageAndOffset -> Property+prop_findKey_index_law_LT (RawPageAndOffset rawpage key locInfo) =+ recordOffsetLocation locInfo $+ findKeyIndexCounterExample+ key+ rawpage+ "∀ key page. maybe True (key <=) (getRawPageIndexKey . rawPageIndex page . id =<< rawPageFindKey page key)" $+ maybe+ (discard)+ (property . (key >))+ (getRawPageIndexKey . rawPageIndex rawpage =<< pred' =<< rawPageFindKey rawpage key)+ where+ pred' n+ | n <= 0 = Nothing+ | otherwise = Just $ n - 1++prop_findKey_index_law_missing :: RawPageAndOffset -> Property+prop_findKey_index_law_missing (RawPageAndOffset rawpage key locInfo) =+ recordOffsetLocation locInfo $+ findKeyIndexCounterExample+ key+ rawpage+ "∀ key page. maybe (maximum (rawPageKeys page) < key) (rawPageFindKey page key)" $+ maybe+ (property $ rawPageKeys rawpage `maximumIsLessThan` key)+ (const discard)+ (rawPageFindKey rawpage key)+ where+ maximumIsLessThan iVec obj+ | null iVec = True -- When there are no elements, the test passes+ | otherwise = maximum iVec < obj++findKeyIndexCounterExample :: Testable prop => SerialisedKey -> RawPage -> String -> prop -> Property+findKeyIndexCounterExample key rawpage lawStr = counterexample msg+ where+ pKeys = rawPageKeys rawpage+ msg = case rawPageFindKey rawpage key of+ Nothing -> unlines+ [ "No key found"+ , show $ pKeys+ , if null pKeys+ then "<NONE>"+ else show $ maximum pKeys+ , show key+ , if null pKeys+ then "<NONE>"+ else show $ maximum pKeys < key+ ]+ Just loc -> concat+ [ "Relational law violated:\n "+ , lawStr+ , "\n"+ , unwords+ [ "At entry№"+ , show loc+ , "found next {"+ , show key+ , "} > {"+ , show $ rawPageIndex rawpage loc+ , "} of page index "+ , show loc+ ]+ , show $ pKeys+ , if null pKeys+ then "<NONE>"+ else show $ maximum pKeys+ , show key+ , if null pKeys+ then "<NONE>"+ else show $ maximum pKeys < key+ ]++recordOffsetLocation :: Testable prop => OffsetRelativeToRawPage -> prop -> Property+recordOffsetLocation = label <$> \case+ BeforePage -> "Key was before page"+ BehindPage -> "Key was behind page"+ PresentInPage -> "Key was within page"+ MissingInPage -> "Key not within page"+++-- |+-- Helper data-type for generating arbitrary pages and offset keys+-- with a useful distribution of hit and misses for testing purposes.+data RawPageAndOffset = RawPageAndOffset RawPage SerialisedKey OffsetRelativeToRawPage+ deriving stock (Eq, Show)++-- | Exists for test-case reporting purposes.+data OffsetRelativeToRawPage+ = BeforePage+ | BehindPage+ | MissingInPage+ | PresentInPage+ deriving stock (Eq, Show)++instance Arbitrary RawPageAndOffset where++ arbitrary = do+ Ref.PageContentOrdered kops <- arbitrary+ let rawpage = fst . Ref.toRawPage $ Ref.PageContentFits kops+ (offsetKey, locInfo) <- case kops of+ -- If there are no k/ops, then generate any offset key.+ [] -> arbitrary <&> \x -> (Ref.toSerialisedKey x, BehindPage)+ -- Otherwise, if at least one k/op exists then generate the offset key according to+ -- the following probability distribution which is conditional on kOps:+ -- * (1/2): Offset Key exists within the page+ -- * (1/6): Offset Key goes before the page+ -- * (1/6): Offset Key goes after the page+ -- * (1/6): Offset Key missing within the page+ -- (Assuming that there are two or more keys)+ xs ->+ let keys = fst <$> xs+ keySmall = Ref.toSerialisedKey $ minimum keys+ keyLarge = Ref.toSerialisedKey $ maximum keys+ arbitrarySerialisedKeyExisting :: Gen (SerialisedKey, OffsetRelativeToRawPage)+ arbitrarySerialisedKeyExisting =+ elements keys <&> \x -> (Ref.toSerialisedKey x, PresentInPage)+ missingWithinRange :: [(Int, Gen (SerialisedKey, OffsetRelativeToRawPage))]+ missingWithinRange+ | keySmall == keyLarge = []+ | otherwise =+ let missingKey = do+ afterIndex <- chooseInt (0, length keys - 2)+ let loKey = Ref.toSerialisedKey $ keys !! afterIndex+ hiKey = Ref.toSerialisedKey $ keys !! (afterIndex + 1)+ arbitrarySerialisedKeyBetween loKey hiKey+ in [ (1, missingKey) ]+ in frequency $+ [ (3, arbitrarySerialisedKeyExisting)+ , (1, arbitrarySerialisedKeyLessThan keySmall)+ , (1, arbitrarySerialisedKeyMoreThan keyLarge)+ ] <> missingWithinRange++ pure $ RawPageAndOffset rawpage offsetKey locInfo++-- |+-- Generate an arbitrary 'SerialisedKey' between the given two keys.+--+-- The first key is assumed to be the "smaller" of the pair.+--+-- NOTE: Under pathological input conditions, there is an infinitesimally small chance+-- that this will generate /the same/ key as the larger value. This will not effect the+-- validity of 'Arbitrary' instances which depend on this function.+arbitrarySerialisedKeyBetween :: SerialisedKey -> SerialisedKey -> Gen (SerialisedKey, OffsetRelativeToRawPage)+arbitrarySerialisedKeyBetween (SerialisedKey (RawBytes loBytes)) (SerialisedKey (RawBytes hiBytes)) =+ finalizer <$> VP.zipWithM withinBytes loBytes hiBytes `evalStateT` (True, True)+ where+ finalizer x = (SerialisedKey $ RawBytes x, MissingInPage)+ withinBytes :: Word8 -> Word8 -> StateT (Bool, Bool) Gen Word8+ withinBytes !lhs !rhs = do+ (!loMatch, !hiMatch) <- get+ let loBound+ | loMatch = min lhs rhs+ | otherwise = minBound+ hiBound+ | hiMatch = max lhs rhs+ | otherwise = maxBound+ resultByte <- lift $ chooseEnum (loBound, hiBound)+ put (loMatch && resultByte == lhs, hiMatch && resultByte == rhs)+ pure resultByte++-- |+-- Generate an arbitrary 'SerialisedKey' less than the given key.+--+-- NOTE: Under pathological input conditions, there is an infinitesimally small chance+-- that this will generate /the same/ key as the input. This will not effect the+-- validity of 'Arbitrary' instances which depend on this function.+arbitrarySerialisedKeyLessThan :: SerialisedKey -> Gen (SerialisedKey, OffsetRelativeToRawPage)+arbitrarySerialisedKeyLessThan (SerialisedKey (RawBytes bytes)) =+ finalizer <$> VP.mapM lesserByte bytes `evalStateT` True+ where+ finalizer x = (SerialisedKey $ RawBytes x, BeforePage)+ lesserByte byte = do+ samePrefix <- get+ let loBound+ | samePrefix = byte+ | otherwise = minBound+ resultByte <- lift $ chooseEnum (loBound, maxBound)+ when samePrefix . put $ resultByte == byte+ pure resultByte++-- |+-- Generate an arbitrary 'SerialisedKey' greater than the given key.+--+-- NOTE: Under pathological input conditions, there is an infinitesimally small chance+-- that this will generate /the same/ key as the input. This will not effect the+-- validity of 'Arbitrary' instances which depend on this function.+arbitrarySerialisedKeyMoreThan :: SerialisedKey -> Gen (SerialisedKey, OffsetRelativeToRawPage)+arbitrarySerialisedKeyMoreThan (SerialisedKey (RawBytes bytes)) =+ finalizer <$> VP.mapM greaterByte bytes `evalStateT` True+ where+ finalizer x = (SerialisedKey $ RawBytes x, BehindPage)+ greaterByte byte = do+ samePrefix <- get+ let hiBound+ | samePrefix = byte+ | otherwise = maxBound+ resultByte <- lift $ chooseEnum (maxBound, hiBound)+ when samePrefix . put $ resultByte == byte+ pure resultByte
@@ -0,0 +1,487 @@+{-# LANGUAGE TypeFamilies #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.Readers (tests) where++import Control.Exception (assert)+import Control.Monad (forM)+import Control.Monad.IO.Class (liftIO)+import Control.Monad.Trans.Reader (ReaderT (..))+import Control.Monad.Trans.State (StateT (..), get, put)+import Control.RefCount+import Data.Bifunctor (bimap, first)+import Data.Coerce (coerce)+import Data.Foldable (traverse_)+import qualified Data.Map.Strict as Map+import Data.Proxy (Proxy (..))+import Data.Tree (Tree)+import qualified Data.Tree as Tree+import Database.LSMTree.Extras (showPowersOf)+import Database.LSMTree.Extras.Generators (BiasedKey (..))+import Database.LSMTree.Extras.RunData+import Database.LSMTree.Internal.BlobRef+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Ordinary))+import Database.LSMTree.Internal.Readers (HasMore (Drained, HasMore),+ Readers)+import qualified Database.LSMTree.Internal.Readers as Readers+import qualified Database.LSMTree.Internal.Run as Run+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import qualified Database.LSMTree.Internal.RunReader as Reader+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.UniqCounter+import qualified Database.LSMTree.Internal.WriteBuffer as WB+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.Sim as FsSim+import qualified System.FS.Sim.MockFS as MockFS+import qualified Test.QuickCheck as QC+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck hiding (Some)+import Test.Util.Orphans ()++import Test.QuickCheck.StateModel+import Test.QuickCheck.StateModel.Lockstep+import qualified Test.QuickCheck.StateModel.Lockstep.Defaults as Lockstep+import qualified Test.QuickCheck.StateModel.Lockstep.Run as Lockstep++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.Readers"+ [ testProperty "prop_lockstep" $+ Lockstep.runActionsBracket (Proxy @ReadersState)+ mempty mempty $ \act () -> do+ (prop, mockFS) <- FsSim.runSimHasBlockIO MockFS.empty $ \hfs hbio -> do+ (prop, RealState _ mCtx) <- runRealMonad hfs hbio+ (RealState 0 Nothing) act+ traverse_ closeReadersCtx mCtx -- close current readers+ pure prop++ -- ensure that all handles have been closed+ pure $ prop+ .&&. counterexample "file handles" (MockFS.numOpenHandles mockFS === 0)+ ]++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Ordinary+ }++testSalt :: Bloom.Salt+testSalt = 4++--------------------------------------------------------------------------------++type SerialisedEntry = Entry SerialisedValue SerialisedBlob++type Handle = MockFS.HandleMock++data ReaderSourceData =+ FromWriteBufferData+ !(RunData BiasedKey SerialisedValue SerialisedBlob)+ | FromRunData+ !(RunData BiasedKey SerialisedValue SerialisedBlob)+ | FromReadersData+ !Readers.ReadersMergeType ![ReaderSourceData]+ deriving stock (Eq, Show)++depth :: ReaderSourceData -> Int+depth = \case+ FromWriteBufferData _ -> 1+ FromRunData _ -> 1+ FromReadersData _ rds -> 1 + maximum (0 : map depth rds)++instance Arbitrary ReaderSourceData where+ arbitrary = QC.oneof [genLeaf, genTree]+ where+ genLeaf = QC.oneof [+ FromWriteBufferData <$> arbitrary+ , FromRunData <$> arbitrary+ ]++ genTree =+ -- first generate a tree shape of up to 10 nodes, then the actual data+ QC.scale (`div` 10) (arbitrary @(Tree ())) >>= enrich++ enrich :: Tree () -> Gen ReaderSourceData+ enrich (Tree.Node () []) =+ genLeaf+ enrich (Tree.Node () children) =+ FromReadersData <$> arbitrary <*> traverse enrich children++ shrink (FromWriteBufferData rd) =+ [ FromWriteBufferData rd' | rd' <- shrink rd+ ]+ shrink (FromRunData rd) =+ [ FromWriteBufferData rd+ ] <>+ [ FromRunData rd' | rd' <- shrink rd+ ]+ shrink s@(FromReadersData ty rds) =+ rds+ <>+ [ FromRunData (RunData (sourceKOps s))+ ] <>+ [ FromReadersData ty' rds' | (ty', rds') <- shrink (ty, rds)+ ]++instance Arbitrary Readers.ReadersMergeType where+ arbitrary = QC.elements [Readers.MergeLevel, Readers.MergeUnion]+ where+ _coveredAllCases :: Readers.ReadersMergeType -> ()+ _coveredAllCases = \case+ Readers.MergeLevel -> ()+ Readers.MergeUnion -> ()++sourceKOps :: ReaderSourceData -> Map.Map BiasedKey SerialisedEntry+sourceKOps (FromWriteBufferData rd) = unRunData rd+sourceKOps (FromRunData rd) = unRunData rd+sourceKOps (FromReadersData ty rds) = Map.unionsWith f (map sourceKOps rds)+ where+ f = case ty of+ Readers.MergeLevel -> combine resolve+ Readers.MergeUnion -> combineUnion resolve++resolve :: ResolveSerialisedValue+resolve (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)++--------------------------------------------------------------------------------+-- Mock++newtype MockReaders = MockReaders+ { mockEntries :: [((SerialisedKey, RunNumber), SerialisedEntry)] }+ deriving stock Show++isEmpty :: MockReaders -> Bool+isEmpty (MockReaders xs) = null xs++size :: MockReaders -> Int+size (MockReaders xs) = length xs++newMock :: Maybe SerialisedKey+ -> [ReaderSourceData]+ -> MockReaders+newMock offset =+ MockReaders . Map.toList . Map.unions+ . zipWith (\i -> Map.mapKeysMonotonic (\k -> (k, RunNumber i))) [0..]+ . map (skip . Map.mapKeysMonotonic serialiseKey . sourceKOps)+ where+ skip = maybe id (\k -> Map.dropWhileAntitone (< k)) offset++peekKeyMock :: MockReaders -> Either () SerialisedKey+peekKeyMock (MockReaders xs) = case xs of+ [] -> Left ()+ (((k, _), _) : _) -> Right k++-- | Drops the first @n@ entries, returning the last of them.+popMock :: Int -> MockReaders -> (Either () (SerialisedKey, SerialisedEntry, HasMore), MockReaders)+popMock n (MockReaders xs) = assert (n >= 1) $+ case drop (n - 1) xs of+ [] ->+ (Left (), MockReaders []) --popping too many still modifies state+ (((k, _), e) : rest) ->+ (Right (k, e, toHasMore rest), MockReaders rest)++dropWhileKeyMock :: SerialisedKey -> MockReaders -> (Either () (Int, HasMore), MockReaders)+dropWhileKeyMock k m@(MockReaders xs)+ | null xs = (Left (), m)+ | otherwise =+ let (dropped, xs') = span ((<= k) . fst . fst) xs+ in (Right (length dropped, toHasMore xs'), MockReaders xs')++toHasMore :: [a] -> HasMore+toHasMore xs = if null xs then Drained else HasMore+++data ReadersState = ReadersState MockReaders+ deriving stock Show++initReadersState :: ReadersState+initReadersState = ReadersState (newMock Nothing [])++--------------------------------------------------------------------------------++type ReadersAct a = Action (Lockstep ReadersState) (Either () a)++deriving stock instance Show (Action (Lockstep ReadersState) a)+deriving stock instance Eq (Action (Lockstep ReadersState) a)++instance StateModel (Lockstep ReadersState) where+ data Action (Lockstep ReadersState) a where+ New :: Maybe BiasedKey -- ^ optional offset+ -> [ReaderSourceData]+ -> ReadersAct ()+ PeekKey :: ReadersAct SerialisedKey+ Pop :: Int -- allow popping many at once to drain faster+ -> ReadersAct (SerialisedKey, SerialisedEntry, HasMore)+ DropWhileKey :: SerialisedKey+ -> ReadersAct (Int, HasMore)++ initialState = Lockstep.initialState initReadersState+ nextState = Lockstep.nextState+ precondition = Lockstep.precondition+ arbitraryAction = Lockstep.arbitraryAction+ shrinkAction = Lockstep.shrinkAction++type ReadersVal a = ModelValue ReadersState a+type ReadersObs a = Observable ReadersState a++deriving stock instance Show (ReadersVal a)+deriving stock instance Show (ReadersObs a)+deriving stock instance Eq (ReadersObs a)++instance InLockstep ReadersState where+ data ModelValue ReadersState a where+ MEntry :: SerialisedEntry -> ReadersVal SerialisedEntry+ MKey :: SerialisedKey -> ReadersVal SerialisedKey+ MHasMore:: HasMore -> ReadersVal HasMore+ MInt :: Int -> ReadersVal Int+ MUnit :: () -> ReadersVal ()+ MEither :: Either (ReadersVal a) (ReadersVal b) -> ReadersVal (Either a b)+ MTuple2 :: (ReadersVal a, ReadersVal b) -> ReadersVal (a, b)+ MTuple3 :: (ReadersVal a, ReadersVal b, ReadersVal c) -> ReadersVal (a, b, c)++ data Observable ReadersState a where+ OId :: (Eq a, Show a) => a -> ReadersObs a+ OEither :: Either (ReadersObs a) (ReadersObs b) -> ReadersObs (Either a b)+ OTuple2 :: (ReadersObs a, ReadersObs b) -> ReadersObs (a, b)+ OTuple3 :: (ReadersObs a, ReadersObs b, ReadersObs c) -> ReadersObs (a, b, c)++ observeModel :: ReadersVal a -> ReadersObs a+ observeModel = \case+ MEntry e -> OId e+ MKey k -> OId k+ MHasMore h -> OId h+ MInt n -> OId n+ MUnit () -> OId ()+ MEither x -> OEither $ bimap observeModel observeModel x+ MTuple2 x -> OTuple2 $ bimap observeModel observeModel x+ MTuple3 x -> OTuple3 $ trimap observeModel observeModel observeModel x++ modelNextState :: forall a.+ LockstepAction ReadersState a+ -> ModelVarContext ReadersState+ -> ReadersState+ -> (ReadersVal a, ReadersState)+ modelNextState action _ctx (ReadersState mock) =+ ReadersState <$> runMock action mock++ usedVars :: LockstepAction ReadersState a -> [AnyGVar (ModelOp ReadersState)]+ usedVars = const []++ arbitraryWithVars ::+ ModelVarContext ReadersState+ -> ReadersState+ -> Gen (Any (LockstepAction ReadersState))+ arbitraryWithVars _ctx (ReadersState mock)+ | isEmpty mock = do+ -- It's not allowed to keep using a drained RunReaders,+ -- we can only create a new one.+ sources <- vector =<< chooseInt (1, 10)+ let keys = concatMap (Map.keys . sourceKOps) sources+ offset <-+ if null keys+ then pure Nothing+ else oneof+ [ liftArbitrary (elements (coerce keys)) -- existing key+ , arbitrary -- any key+ ]+ pure $ Some $ New offset sources+ | otherwise =+ QC.frequency $+ [ (5, pure (Some PeekKey))+ , (8, pure (Some (Pop 1)))+ , (1, Some . Pop <$> chooseInt (1, size mock)) -- drain a significant amount+ , (1, pure (Some (Pop (max 1 (size mock - 3))))) -- drain almost everything+ , (1, pure (Some (Pop (size mock + 1)))) -- drain /more/ than available+ , (1, Some . DropWhileKey <$> arbitrary) -- might drop a lot+ ] <>+ [ (4, pure (Some (DropWhileKey k))) -- drops at least one key+ | Right k <- [peekKeyMock mock]+ ]++ shrinkWithVars ::+ ModelVarContext ReadersState+ -> ReadersState+ -> LockstepAction ReadersState a+ -> [Any (LockstepAction ReadersState)]+ shrinkWithVars _ctx _st = \case+ New k sources -> [ Some (New k' sources')+ | (k', sources') <- shrink (k, sources)+ ]+ -- arbitraryWithVars does /not/ have an invariant that n is less than+ -- the number of elements available. The only invariant to preserve when+ -- shrinking is to keep n greater than 0.+ Pop n -> Some . Pop <$> filter (> 0) (shrink n)+ _ -> []++ tagStep ::+ (ReadersState, ReadersState)+ -> LockstepAction ReadersState a+ -> ReadersVal a+ -> [String]+ tagStep (ReadersState before, ReadersState after) action _result = concat+ -- Directly using strings, since there is only a small number of tags.+ [ [ "NewEntries " <> showPowersOf 10 numEntries+ | New _ sources <- [action]+ , let numEntries = sum (map (Map.size . sourceKOps) sources)+ ]+ , [ "NewEntriesKeyDuplicates " <> showPowersOf 2 keyCount+ | New _ sources <- [action]+ , let keyCounts = Map.unionsWith (+) (map (Map.map (const 1) . sourceKOps) sources)+ , keyCount <- Map.elems keyCounts+ , keyCount > 1+ ]+ , [ "NewDepth " <> showPowersOf 2 (maximum (0 : map depth sources))+ | New _ sources <- [action]+ ]+ , [ "ReadersFullyDrained"+ | not (isEmpty before), isEmpty after+ ]+ , [ "DropWhileKeyDropped " <> showPowersOf 2 (length dropped)+ | DropWhileKey key <- [action]+ , let dropped = takeWhile ((== key) . fst . fst) (mockEntries before)+ ]+ ]++runMock ::+ Action (Lockstep ReadersState) a+ -> MockReaders+ -> (ReadersVal a, MockReaders)+runMock = \case+ New k sources -> const $ wrap MUnit (Right (), newMock (coerce k) sources)+ PeekKey -> \m -> wrap MKey (peekKeyMock m, m)+ Pop n -> wrap wrapPop . popMock n+ DropWhileKey k -> wrap wrapDrop . dropWhileKeyMock k+ where+ wrap :: (a -> ReadersVal b) -> (Either () a, MockReaders) -> (ReadersVal (Either () b), MockReaders)+ wrap f = first (MEither . bimap MUnit f)++ wrapPop = MTuple3 . trimap MKey MEntry MHasMore++ wrapDrop = MTuple2 . bimap MInt MHasMore++trimap :: (a -> a') -> (b -> b') -> (c -> c') -> (a, b, c) -> (a', b', c')+trimap f g h (a, b, c) = (f a, g b, h c)++type RealMonad = ReaderT (FS.HasFS IO MockFS.HandleMock,+ FS.HasBlockIO IO MockFS.HandleMock)+ (StateT RealState IO)++runRealMonad :: FS.HasFS IO MockFS.HandleMock+ -> FS.HasBlockIO IO MockFS.HandleMock+ -> RealState+ -> RealMonad a+ -> IO (a, RealState)+runRealMonad hfs hbio st = (`runStateT` st) . (`runReaderT` (hfs, hbio))++data RealState =+ RealState+ !Int -- ^ number of runs created so far (to generate fresh run numbers)+ !(Maybe ReadersCtx)++-- | Readers, together with their sources, so they can be cleaned up at the end+type ReadersCtx = ( [Readers.ReaderSource IO MockFS.HandleMock]+ , Readers IO Handle+ )++closeReadersCtx :: ReadersCtx -> IO ()+closeReadersCtx (sources, readers) = do+ Readers.close readers+ traverse_ closeReaderSource sources++closeReaderSource :: Readers.ReaderSource IO h -> IO ()+closeReaderSource = \case+ Readers.FromWriteBuffer _ wbblobs -> releaseRef wbblobs+ Readers.FromRun run -> releaseRef run+ Readers.FromReaders _ sources -> traverse_ closeReaderSource sources++instance RunModel (Lockstep ReadersState) RealMonad where+ perform = \_st -> runIO+ postcondition = Lockstep.postcondition+ monitoring = Lockstep.monitoring (Proxy @RealMonad)++instance RunLockstep ReadersState RealMonad where+ observeReal _proxy = \case+ New {} -> OEither . bimap OId OId+ PeekKey {} -> OEither . bimap OId OId+ Pop {} -> OEither . bimap OId (OTuple3 . trimap OId OId OId)+ DropWhileKey {} -> OEither . bimap OId (OTuple2 . bimap OId OId)++runIO :: LockstepAction ReadersState a -> LookUp -> RealMonad a+runIO act lu = case act of+ New offset srcDatas -> ReaderT $ \(hfs, hbio) -> do+ RealState numRuns mCtx <- get+ -- if runs are still being read, they need to be cleaned up+ traverse_ (liftIO . closeReadersCtx) mCtx+ counter <- liftIO $ newUniqCounter numRuns+ sources <- liftIO $ forM srcDatas (fromSourceData hfs hbio counter)+ newReaders <- liftIO $ do+ let offsetKey = maybe Readers.NoOffsetKey (Readers.OffsetKey . coerce) offset+ mreaders <- Readers.new resolve offsetKey sources+ -- TODO: tidy up cleanup code?+ case mreaders of+ Nothing -> do+ traverse_ closeReaderSource sources+ pure Nothing+ Just readers ->+ pure $ Just (sources, readers)+ -- TODO: unnecessarily convoluted, should we just drop the State monad?+ numRuns' <- liftIO $ uniqueToInt <$> incrUniqCounter counter+ put (RealState numRuns' newReaders)+ pure (Right ())+ PeekKey -> expectReaders $ \_ r -> do+ (,) HasMore <$> Readers.peekKey r+ Pop n | n <= 1 -> pop+ Pop n -> pop >>= \case+ Left () -> pure (Left ())+ Right (_, _, Drained) -> pure (Left ()) -- n > 1, so n too big+ Right (_, _, HasMore) -> runIO (Pop (n-1)) lu+ DropWhileKey k -> expectReaders $ \_ r -> do+ (n, hasMore) <- Readers.dropWhileKey resolve r k+ pure (hasMore, (n, hasMore))+ where+ fromSourceData hfs hbio counter = \case+ FromWriteBufferData rd -> do+ n <- incrUniqCounter counter+ wbblobs <- WBB.new hfs (FS.mkFsPath [show (uniqueToInt n) <> ".wb.blobs"])+ let kops = unRunData (serialiseRunData rd)+ wb <- WB.fromMap <$> traverse (traverse (WBB.addBlob hfs wbblobs)) kops+ pure $ Readers.FromWriteBuffer wb wbblobs+ FromRunData rd -> do+ r <- unsafeCreateRun hfs hbio testSalt runParams (FS.mkFsPath []) counter $ serialiseRunData rd+ pure $ Readers.FromRun r+ FromReadersData ty rds -> do+ Readers.FromReaders ty <$> traverse (fromSourceData hfs hbio counter) rds++ pop = expectReaders $ \hfs r -> do+ (key, e, hasMore) <- Readers.pop resolve r+ fullEntry <- toMockEntry hfs e+ pure (hasMore, (key, fullEntry, hasMore))++ expectReaders ::+ (FS.HasFS IO MockFS.HandleMock -> Readers IO MockFS.HandleMock -> IO (HasMore, a))+ -> RealMonad (Either () a)+ expectReaders f =+ ReaderT $ \(hfs, _hbio) -> do+ get >>= \case+ RealState _ Nothing -> pure (Left ())+ RealState n (Just (sources, readers)) -> do+ (hasMore, x) <- liftIO $ f hfs readers+ case hasMore of+ HasMore ->+ pure (Right x)+ Drained -> do+ -- Readers is drained, clean up the sources+ liftIO $ traverse_ closeReaderSource sources+ put (RealState n Nothing)+ pure (Right x)++ toMockEntry :: FS.HasFS IO MockFS.HandleMock -> Reader.Entry IO MockFS.HandleMock -> IO SerialisedEntry+ toMockEntry hfs = traverse (readRawBlobRef hfs) . Reader.toFullEntry
@@ -0,0 +1,296 @@+{-# LANGUAGE OverloadedStrings #-}++module Test.Database.LSMTree.Internal.Run (+ -- * Main test tree+ tests,+) where++import Control.ActionRegistry (withActionRegistry)+import Control.RefCount+import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Short as SBS+import Data.Coerce (coerce)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromJust)+import qualified Data.Primitive.ByteArray as BA+import Database.LSMTree.Extras.RunData+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import qualified Database.LSMTree.Internal.CRC32C as CRC+import Database.LSMTree.Internal.Entry+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Ordinary))+import Database.LSMTree.Internal.Paths (RunFsPaths (..),+ WriteBufferFsPaths (..))+import qualified Database.LSMTree.Internal.Paths as Paths+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawPage+import Database.LSMTree.Internal.Run as Run+import Database.LSMTree.Internal.RunAcc as RunAcc+import Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Snapshot+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import Database.LSMTree.Internal.WriteBufferReader (readWriteBuffer)+import qualified FormatPage as Proto+import System.FilePath+import qualified System.FS.API as FS+import qualified System.FS.API.Lazy as FSL+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.BlockIO.IO as FS+import qualified System.FS.Sim.MockFS as MockFS+import qualified System.IO.Temp as Temp+import Test.Database.LSMTree.Internal.RunReader (readKOps)+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (assertEqual, testCase, (@=?), (@?))+import Test.Tasty.QuickCheck+import Test.Util.FS (propNoOpenHandles, withSimHasBlockIO)+++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.Run"+ [ testGroup "Write buffer to disk"+ [ testCase "Single insert (small)" $ do+ withSessionDir $ \sessionRoot ->+ testSingleInsert sessionRoot+ (mkKey "test-key")+ (mkVal "test-value")+ Nothing+ , testCase "Single insert (blob)" $ do+ withSessionDir $ \sessionRoot -> do+ testSingleInsert sessionRoot+ (mkKey "test-key")+ (mkVal "test-value")+ (Just (mkBlob "test-blob"))+ , testCase "Single insert (larger-than-page)" $ do+ withSessionDir $ \sessionRoot -> do+ testSingleInsert sessionRoot+ (mkKey "test-key")+ (mkVal ("test-value-" <> BS.concat (replicate 500 "0123456789")))+ Nothing+ , testProperty "prop_WriteAndOpen" $ \wb ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_WriteAndOpen hfs hbio wb+ , testProperty "prop_WriteNumEntries" $ \wb ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_WriteNumEntries hfs hbio wb+ , testProperty "prop_WriteAndOpenWriteBuffer" $ \wb ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_WriteAndOpenWriteBuffer hfs hbio wb+ , testProperty "prop_WriteRunEqWriteWriteBuffer" $ \wb ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ ->+ prop_WriteRunEqWriteWriteBuffer hfs hbio wb+ ]+ ]+ where+ withSessionDir = Temp.withSystemTempDirectory "session-run"++ mkKey = SerialisedKey . RB.fromByteString+ mkVal = SerialisedValue . RB.fromByteString+ mkBlob = SerialisedBlob . RB.fromByteString++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Ordinary+ }++testSalt :: Bloom.Salt+testSalt = 4++-- | Runs in IO, with a real file system.+testSingleInsert :: FilePath -> SerialisedKey -> SerialisedValue -> Maybe SerialisedBlob -> IO ()+testSingleInsert sessionRoot key val mblob =+ FS.withIOHasBlockIO (FS.MountPoint sessionRoot) FS.defaultIOCtxParams $ \fs hbio -> do+ -- flush write buffer+ let e = case mblob of Nothing -> Insert val; Just blob -> InsertWithBlob val blob+ wb = Map.singleton key e+ withRunAt fs hbio testSalt runParams (simplePath 42) (RunData wb) $ \_ -> do+ -- check all files have been written+ let activeDir = sessionRoot+ bsKOps <- BS.readFile (activeDir </> "42.keyops")+ bsBlobs <- BS.readFile (activeDir </> "42.blobs")+ bsFilter <- BS.readFile (activeDir </> "42.filter")+ bsIndex <- BS.readFile (activeDir </> "42.index")+ not (BS.null bsKOps) @? "k/ops file is empty"+ null mblob @=? BS.null bsBlobs -- blob file might be empty+ not (BS.null bsFilter) @? "filter file is empty"+ not (BS.null bsIndex) @? "index file is empty"+ -- checksums+ checksums <- CRC.readChecksumsFile fs (FS.mkFsPath ["42.checksums"])+ Map.lookup (CRC.ChecksumsFileName "keyops") checksums+ @=? Just (CRC.updateCRC32C bsKOps CRC.initialCRC32C)+ Map.lookup (CRC.ChecksumsFileName "blobs") checksums+ @=? Just (CRC.updateCRC32C bsBlobs CRC.initialCRC32C)+ Map.lookup (CRC.ChecksumsFileName "filter") checksums+ @=? Just (CRC.updateCRC32C bsFilter CRC.initialCRC32C)+ Map.lookup (CRC.ChecksumsFileName "index") checksums+ @=? Just (CRC.updateCRC32C bsIndex CRC.initialCRC32C)+ -- check page+ let page = rawPageFromByteString bsKOps 0+ 1 @=? rawPageNumKeys page++ let pagesize :: Int+ pagesize = fromJust $+ Proto.pageSizeBytes <$> Proto.calcPageSize Proto.DiskPage4k+ [ ( Proto.Key (coerce RB.toByteString key)+ , Proto.Insert (Proto.Value (coerce RB.toByteString val))+ Nothing+ ) ]+ suffix, prefix :: Int+ suffix = max 0 (pagesize - 4096)+ prefix = coerce RB.size val - suffix+ let expectedEntry = case mblob of+ Nothing -> Insert (coerce RB.take prefix val)+ Just b -> InsertWithBlob (coerce RB.take prefix val) b+ let expectedResult+ | suffix > 0 = LookupEntryOverflow expectedEntry (fromIntegral suffix)+ | otherwise = LookupEntry expectedEntry++ let actualEntry = fmap (readBlobFromBS bsBlobs) <$> rawPageLookup page key++ -- the lookup result is as expected, possibly with a prefix of the value+ expectedResult @=? actualEntry++ -- the value is as expected, including any overflow suffix+ let valPrefix = coerce RB.take prefix val+ valSuffix = (RB.fromByteString . BS.take suffix . BS.drop 4096) bsKOps+ val @=? SerialisedValue (valPrefix <> valSuffix)++ -- blob sanity checks+ length mblob @=? fromIntegral (rawPageNumBlobs page)++rawPageFromByteString :: ByteString -> Int -> RawPage+rawPageFromByteString bs off =+ makeRawPage (toBA bs) off+ where+ -- ensure that the resulting RawPage has no trailing data that could+ -- accidentally be read.+ toBA = (\(SBS.SBS ba) -> BA.ByteArray ba) . SBS.toShort . BS.take (off+4096)++readBlobFromBS :: ByteString -> BlobSpan -> SerialisedBlob+readBlobFromBS bs (BlobSpan off sz) =+ serialiseBlob $ BS.take (fromIntegral sz) (BS.drop (fromIntegral off) bs)++{-------------------------------------------------------------------------------+ Properties+-------------------------------------------------------------------------------}++-- | Flushing a write buffer results in a run with the desired elements.+--+-- To assert that the run really contains the right entries, there are+-- additional tests in "Test.Database.LSMTree.Internal.RunReader".+--+prop_WriteNumEntries ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData SerialisedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_WriteNumEntries fs hbio wb@(RunData m) =+ withRunAt fs hbio testSalt runParams (simplePath 42) wb' $ \run -> do+ let !runSize = Run.size run++ pure . labelRunData wb' $+ NumEntries (Map.size m) === runSize+ where+ wb' = serialiseRunData wb++-- | Loading a run (written out from a write buffer) from disk gives the same+-- in-memory representation as the original run.+--+-- @openFromDisk . flush === flush@+prop_WriteAndOpen ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData SerialisedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_WriteAndOpen fs hbio wb =+ withRunAt fs hbio testSalt runParams (simplePath 1337) (serialiseRunData wb) $ \written ->+ withActionRegistry $ \reg -> do+ let paths = Run.runFsPaths written+ paths' = paths { runNumber = RunNumber 17}+ hardLinkRunFiles fs hbio reg paths paths'+ loaded <- openFromDisk fs hbio (runParamCaching runParams)+ (runParamIndex runParams) testSalt+ (simplePath 17)++ Run.size written @=? Run.size loaded+ withRef written $ \written' ->+ withRef loaded $ \loaded' -> do+ runFilter written' @=? runFilter loaded'+ runIndex written' @=? runIndex loaded'++ writtenKOps <- readKOps Nothing written+ loadedKOps <- readKOps Nothing loaded++ assertEqual "k/ops" writtenKOps loadedKOps++ -- make sure runs get closed again+ releaseRef loaded++ -- TODO: return a proper Property instead of using assertEqual etc.+ pure (property True)++-- | Writing and loading a 'WriteBuffer' gives the same in-memory+-- representation as the original write buffer.+prop_WriteAndOpenWriteBuffer ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData SerialisedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_WriteAndOpenWriteBuffer hfs hbio rd = do+ -- Serialise run data as write buffer:+ let srd = serialiseRunData rd+ let inPaths = WrapRunFsPaths $ simplePath 1111+ let resolve (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)+ withRunDataAsWriteBuffer hfs resolve inPaths srd $ \wb wbb -> do+ -- Write write buffer to disk:+ let wbPaths = WrapRunFsPaths $ simplePath 1312+ withSerialisedWriteBuffer hfs hbio wbPaths wb wbb $ do+ -- Read write buffer from disk:+ let kOpsPath = Paths.ForKOps (Paths.writeBufferKOpsPath wbPaths)+ withRef wbb $ \wbb' -> do+ wb' <- readWriteBuffer resolve hfs hbio kOpsPath (WBB.blobFile wbb')+ assertEqual "k/ops" wb wb'+ -- TODO: return a proper Property instead of using assertEqual etc.+ pure (property True)++-- | Writing run data to the disk via 'writeWriteBuffer' gives the same key/ops+-- and blob files as when written out as a run.+prop_WriteRunEqWriteWriteBuffer ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData SerialisedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_WriteRunEqWriteWriteBuffer hfs hbio rd = do+ -- Serialise run data as run:+ let srd = serialiseRunData rd+ let rdPaths = simplePath 1337+ let rdKOpsFile = Paths.runKOpsPath rdPaths+ let rdBlobFile = Paths.runBlobPath rdPaths+ withRunAt hfs hbio testSalt runParams rdPaths srd $ \_run -> do+ -- Serialise run data as write buffer:+ let f (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)+ let inPaths = WrapRunFsPaths $ simplePath 1111+ withRunDataAsWriteBuffer hfs f inPaths srd $ \wb wbb -> do+ let wbPaths = WrapRunFsPaths $ simplePath 1312+ let wbKOpsPath = Paths.writeBufferKOpsPath wbPaths+ let wbBlobPath = Paths.writeBufferBlobPath wbPaths+ withSerialisedWriteBuffer hfs hbio wbPaths wb wbb $ do+ -- Compare KOps:+ FS.withFile hfs rdKOpsFile FS.ReadMode $ \rdKOpsHandle ->+ FS.withFile hfs wbKOpsPath FS.ReadMode $ \wbKOpsHandle -> do+ rdKOps <- FSL.hGetAll hfs rdKOpsHandle+ wbKOps <- FSL.hGetAll hfs wbKOpsHandle+ assertEqual "writtenRunKOps/writtenWriteBufferKOps" rdKOps wbKOps+ -- Compare Blob:+ FS.withFile hfs rdBlobFile FS.ReadMode $ \rdBlobHandle ->+ FS.withFile hfs wbBlobPath FS.ReadMode $ \wbBlobHandle -> do+ rdBlob <- FSL.hGetAll hfs rdBlobHandle+ wbBlob <- FSL.hGetAll hfs wbBlobHandle+ assertEqual "writtenRunBlob/writtenWriteBufferBlob" rdBlob wbBlob+ -- TODO: return a proper Property instead of using assertEqual etc.+ pure (property True)
@@ -0,0 +1,158 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE TupleSections #-}++{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}++module Test.Database.LSMTree.Internal.RunAcc (tests) where++import Control.Exception (assert)+import Control.Monad.ST+import Data.Bifunctor (Bifunctor (..))+import qualified Data.BloomFilter.Blocked as Bloom+import qualified Data.ByteString as BS+import qualified Data.ByteString.Short as SBS+import Data.Maybe+import qualified Data.Vector.Primitive as VP+import Database.LSMTree.Internal.BlobRef (BlobSpan (..))+import Database.LSMTree.Internal.Entry+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Ordinary),+ search)+import Database.LSMTree.Internal.Page (PageNo (PageNo), singlePage)+import qualified Database.LSMTree.Internal.PageAcc as PageAcc+import qualified Database.LSMTree.Internal.PageAcc1 as PageAcc+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawOverflowPage (RawOverflowPage)+import qualified Database.LSMTree.Internal.RawOverflowPage as RawOverflowPage+import Database.LSMTree.Internal.RawPage (RawPage)+import qualified Database.LSMTree.Internal.RawPage as RawPage+import Database.LSMTree.Internal.RunAcc+import Database.LSMTree.Internal.Serialise+import qualified FormatPage as Proto+import Test.Tasty+import Test.Tasty.HUnit hiding (assert)+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.RunAcc" [+ testGroup "RunAcc" [+ testCase "test_singleKeyRun" $ test_singleKeyRun+ ]+ , testGroup "PageAcc" [+ largerTestCases $+ testProperty "prop_paddedToDiskPageSize" $+ prop_paddedToDiskPageSize+ , largerTestCases $+ testProperty "prop_runAccMatchesPrototype" prop_runAccMatchesPrototype+ ]+ ]+ where largerTestCases = localOption (QuickCheckMaxSize 500) . localOption (QuickCheckTests 10000)++testSalt :: Bloom.Salt+testSalt = 4++{-------------------------------------------------------------------------------+ RunAcc+-------------------------------------------------------------------------------}++test_singleKeyRun :: Assertion+test_singleKeyRun = do+ let !k = SerialisedKey' (VP.fromList [37, 37, 37, 37, 37, 37, 37, 37])+ !e = InsertWithBlob (SerialisedValue' (VP.fromList [48, 19])) (BlobSpan 55 77)++ (addRes, (mp, mc, b, ic, _numEntries)) <- stToIO $ do+ racc <- new (NumEntries 1) (RunAllocFixed 10) testSalt Index.Ordinary+ addRes <- addKeyOp racc k e+ (addRes,) <$> unsafeFinalise racc++ ([], [], []) @=? addRes+ Just (fst (PageAcc.singletonPage k e)) @=? mp+ isJust mc @? "expected a chunk"+ True @=? Bloom.elem k b+ singlePage (PageNo 0) @=? Index.search k ic++{-------------------------------------------------------------------------------+ PageAcc+-------------------------------------------------------------------------------}++prop_paddedToDiskPageSize :: PageLogical' -> Property+prop_paddedToDiskPageSize page =+ counterexample "expected number of output bytes to be of disk page size" $+ tabulate "page size in bytes" [show $ BS.length bytes] $+ BS.length bytes `rem` 4096 === 0+ where+ bytes = uncurry pagesToByteString $ fromListPageAcc (getRealKOps page)++prop_runAccMatchesPrototype :: PageLogical' -> Property+prop_runAccMatchesPrototype page =+ counterexample "real /= model" $+ real === model+ where+ Just model = Proto.serialisePage <$>+ Proto.encodePage Proto.DiskPage4k (getPrototypeKOps page)+ real = trunc $ uncurry pagesToByteString $ fromListPageAcc (getRealKOps page)++ -- truncate padding on the real page+ trunc = BS.take (BS.length model)++{-------------------------------------------------------------------------------+ Util+-------------------------------------------------------------------------------}++fromListPageAcc :: [(SerialisedKey, Entry SerialisedValue BlobSpan)]+ -> (RawPage, [RawOverflowPage])+fromListPageAcc ((k,e):kops)+ | not (PageAcc.entryWouldFitInPage k e) =+ assert (null kops) $+ PageAcc.singletonPage k e++fromListPageAcc kops =+ runST (do+ pacc <- PageAcc.newPageAcc+ sequence_+ [ do added <- PageAcc.pageAccAddElem pacc k e+ -- we expect the kops to all fit in one page+ assert added $ pure ()+ | (k,e) <- kops ]+ page <- PageAcc.serialisePageAcc pacc+ pure (page, []))++pagesToByteString :: RawPage -> [RawOverflowPage] -> BS.ByteString+pagesToByteString rp rops =+ RB.toByteString+ . mconcat+ $ RawPage.rawPageRawBytes rp+ : map RawOverflowPage.rawOverflowPageRawBytes rops++fromProtoKOp ::+ (Proto.Key, Proto.Operation)+ -> (SerialisedKey, Entry SerialisedValue BlobSpan)+fromProtoKOp (k, op) =+ (fromProtoKey k, bimap fromProtoValue fromProtoBlobRef e)+ where e = case op of+ Proto.Insert v Nothing -> Insert v+ Proto.Insert v (Just br) -> InsertWithBlob v br+ Proto.Mupsert v -> Upsert v+ Proto.Delete -> Delete++fromProtoKey :: Proto.Key -> SerialisedKey+fromProtoKey (Proto.Key bs) = SerialisedKey . RB.fromShortByteString $ SBS.toShort bs++fromProtoValue :: Proto.Value -> SerialisedValue+fromProtoValue (Proto.Value bs) = SerialisedValue . RB.fromShortByteString $ SBS.toShort bs++fromProtoBlobRef :: Proto.BlobRef -> BlobSpan+fromProtoBlobRef (Proto.BlobRef x y) = BlobSpan x y++-- | Wrapper around 'PageLogical' that generates nearly-full pages.+newtype PageLogical' = PageLogical' { getPrototypeKOps :: [(Proto.Key, Proto.Operation)] }+ deriving stock Show++getRealKOps :: PageLogical' -> [(SerialisedKey, Entry SerialisedValue BlobSpan)]+getRealKOps = fmap fromProtoKOp . getPrototypeKOps++instance Arbitrary PageLogical' where+ arbitrary = PageLogical' <$>+ Proto.genPageContentFits Proto.DiskPage4k Proto.noMinKeySize+ shrink (PageLogical' page) =+ [ PageLogical' page' | page' <- shrink page ]+
@@ -0,0 +1,304 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE NumericUnderscores #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.RunBloomFilterAlloc (+ -- * Main test tree+ tests+ -- * Bloom filter construction+ --+ -- A common interface to bloom filter construction, based on expected false+ -- positive rates.+ , BloomMaker+ , mkBloomFromAlloc+ -- * Verifying FPRs+ , measureApproximateFPR+ , measureExactFPR+ ) where++import Control.Exception (assert)+import Control.Monad.ST+import Data.BloomFilter.Blocked (Bloom)+import qualified Data.BloomFilter.Blocked as Bloom+import Data.BloomFilter.Hash (Hashable)+import Data.Foldable (Foldable (..))+import Data.Proxy (Proxy (..))+import Data.Set (Set)+import qualified Data.Set as Set+import Data.Word (Word64)+import Database.LSMTree.Extras.Random+import qualified Database.LSMTree.Internal.Entry as LSMT+import Database.LSMTree.Internal.RunAcc (RunBloomFilterAlloc (..),+ newMBloom)+import System.Random hiding (Seed)+import Test.QuickCheck+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck+import Test.Util.Arbitrary (noTags,+ prop_arbitraryAndShrinkPreserveInvariant)+import Text.Printf (printf)++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.RunBloomFilterAlloc" [+ testProperty "prop_noFalseNegatives" $ prop_noFalseNegatives (Proxy @Word64)+ , testProperty "prop_verifyFPR" $ prop_verifyFPR (Proxy @Word64)+ , testGroup "RunBloomFilterAlloc" $+ prop_arbitraryAndShrinkPreserveInvariant noTags allocInvariant+ , testGroup "NumEntries" $+ prop_arbitraryAndShrinkPreserveInvariant noTags numEntriesInvariant+ ]++testSalt :: Bloom.Salt+testSalt = 4++{-------------------------------------------------------------------------------+ Properties+-------------------------------------------------------------------------------}++prop_noFalseNegatives :: forall a proxy. Hashable a+ => proxy a+ -> RunBloomFilterAlloc+ -> UniformWithoutReplacement a+ -> Property+prop_noFalseNegatives _ alloc (UniformWithoutReplacement xs) =+ let xsBloom = mkBloomFromAlloc alloc xs+ in property $ all (`Bloom.elem` xsBloom) xs++prop_verifyFPR ::+ (Ord a, Uniform a, Hashable a)+ => proxy a+ -> RunBloomFilterAlloc+ -> NumEntries -- ^ @numEntries@+ -> Seed -- ^ 'StdGen' seed+ -> Property+prop_verifyFPR p alloc (NumEntries numEntries) (Seed seed) =+ let stdgen = mkStdGen seed+ measuredFPR = measureApproximateFPR p (mkBloomFromAlloc alloc) numEntries stdgen+ expectedFPR = case alloc of+ RunAllocFixed bits -> Bloom.policyFPR (Bloom.policyForBits bits)+ RunAllocRequestFPR requestedFPR -> requestedFPR+ -- error margins+ lb = expectedFPR - 0.1+ ub = expectedFPR + 0.03+ in counterexample (printf "expected %f <= %f <= %f" lb measuredFPR ub) $+ lb <= measuredFPR .&&. measuredFPR <= ub++{-------------------------------------------------------------------------------+ Modifiers+-------------------------------------------------------------------------------}++--+-- Alloc+--++instance Arbitrary RunBloomFilterAlloc where+ arbitrary = oneof [+ RunAllocFixed <$> genFixed+ , RunAllocRequestFPR <$> genFPR+ ]+ shrink (RunAllocFixed x) = RunAllocFixed <$> shrinkFixed x+ shrink (RunAllocRequestFPR x) = RunAllocRequestFPR <$> shrinkFPR x++allocInvariant :: RunBloomFilterAlloc -> Bool+allocInvariant (RunAllocFixed x) = fixedInvariant x+allocInvariant (RunAllocRequestFPR x) = fprInvariant x++genFixed :: Gen Double+genFixed = choose (fixedLB, fixedUB)++shrinkFixed :: Double -> [Double]+shrinkFixed x = [ x' | x' <- shrink x, fixedInvariant x']++fixedInvariant :: Double -> Bool+fixedInvariant x = fixedLB <= x && x <= fixedUB++fixedLB :: Double+fixedLB = 3 -- bits per entry++fixedUB :: Double+fixedUB = 24 -- bits per entry++genFPR :: Gen Double+genFPR = do m <- choose (1, 9.99) -- not less than 1 or it's a different exponent+ e <- choose (fpr_exponentLB, fpr_exponentUB)+ pure (m * 10 ^^ e)+ `suchThat` fprInvariant++fpr_exponentLB :: Int+fpr_exponentLB = -5 -- 1 in 10,000++fpr_exponentUB :: Int+fpr_exponentUB = -1 -- 1 in 10++shrinkFPR :: Double -> [Double]+shrinkFPR x = [ x' | x' <- shrink x, fprInvariant x']++-- | The FPR calculations are only accurate over the range 0.25 down to 0.00006+-- which corresponds to bits in the range 3 .. 24.+fprInvariant :: Double -> Bool+fprInvariant x = 6e-5 < x && x < 2.5e-1++--+-- NumEntries+--++newtype NumEntries = NumEntries { getNumEntries :: Int }+ deriving stock Show++instance Arbitrary NumEntries where+ arbitrary = NumEntries <$> chooseInt (numEntriesLB, numEntriesUB)+ shrink (NumEntries x) = [+ x''+ | x' <- shrink x+ , let x'' = NumEntries x'+ , numEntriesInvariant x''+ ]++numEntriesLB :: Int+numEntriesLB = 50_000++numEntriesUB :: Int+numEntriesUB = 100_000++numEntriesInvariant :: NumEntries -> Bool+numEntriesInvariant (NumEntries x) = x >= numEntriesLB && x <= numEntriesUB++--+-- Seed+--++newtype Seed = Seed { getSeed :: Int }+ deriving stock Show++instance Arbitrary Seed where+ arbitrary = Seed <$> arbitraryBoundedIntegral+ shrink (Seed x) = Seed <$> shrink x++--+-- UniformWithoutReplacement+--++newtype UniformWithoutReplacement a = UniformWithoutReplacement [a]++instance Show (UniformWithoutReplacement a) where+ show (UniformWithoutReplacement xs) = "UniformWithoutReplacement " <> show (length xs)++instance (Ord a, Uniform a) => Arbitrary (UniformWithoutReplacement a) where+ arbitrary = do+ stdgen <- mkStdGen . getSeed <$> arbitrary+ numEntries <- getNumEntries <$> arbitrary+ pure $ UniformWithoutReplacement $ uniformWithoutReplacement stdgen numEntries++{-------------------------------------------------------------------------------+ Verifying FPRs+-------------------------------------------------------------------------------}++-- | Measure the /approximate/ FPR for a bloom filter.+--+-- Ensure that @a@ is large enough to draw @2 * numEntries@ uniformly random+-- values, or the computation will get stuck.+--+-- REF: based on https://stackoverflow.com/questions/74999807/how-to-measure-the-rate-of-false-positives-in-a-bloom-filter+--+-- REF: https://en.wikipedia.org/wiki/False_positive_rate+measureApproximateFPR ::+ forall a proxy. (Ord a, Uniform a, Hashable a)+ => proxy a -- ^ The types of values to generate.+ -> BloomMaker a -- ^ How to construct the bloom filter.+ -> Int -- ^ @numEntries@: number of entries to put into the bloom filter.+ -> StdGen+ -> Double+measureApproximateFPR _ mkBloom numEntries stdgen =+ let !xs = uniformWithoutReplacement @a stdgen (2 * numEntries)+ (!ys, !zs) = splitAt numEntries xs+ !ysBloom = mkBloom ys+ !ysSet = Set.fromList ys+ oneIfElem z = assert (not $ Set.member z ysSet)+ $ if Bloom.elem z ysBloom then 1 else 0+ !fp = foldl' (\acc x -> acc + oneIfElem x) (0 :: Int) zs+ !fp' = fromIntegral fp :: Double+ in fp' / fromIntegral numEntries -- FPR = FP / FP + TN++-- | Measure the /exact/ FPR for a bloom filter.+--+-- Ensure that @a@ is small enough that we can enumare it within reasonable+-- time. For example, a 'Word16' would be fine, but a 'Word32' would take much+-- too long.+measureExactFPR ::+ forall a proxy. (Ord a, Enum a, Bounded a, Uniform a, Hashable a)+ => proxy a -- ^ The types of values to generate.+ -> BloomMaker a -- ^ How to construct the bloom filter.+ -> Int -- ^ @numEntries@: number of entries to put into the bloom filter.+ -> StdGen+ -> Double+measureExactFPR _ mkBloom numEntries stdgen =+ let !xs = uniformWithoutReplacement @a stdgen numEntries+ !xsBloom = mkBloom xs+ !xsSet = Set.fromList xs+ !aEnumerated = [minBound .. maxBound]+ Counts _ !fp !tn _ = foldMap' (fromTest . analyse xsBloom xsSet) aEnumerated+ fp' = fromIntegral fp :: Double+ tn' = fromIntegral tn :: Double+ in fp' / (fp' + tn') -- FPR = FP / FP + TN++data Test =+ TruePositive+ | FalsePositive+ | TrueNegative+ | FalseNegative++analyse :: (Ord a, Hashable a) => Bloom a -> Set a -> a -> Test+analyse xsBloom xsSet y+ | isBloomMember && isTrueMember = TruePositive+ | isBloomMember && not isTrueMember = FalsePositive+ | not isBloomMember && not isTrueMember = TrueNegative+ | otherwise = FalseNegative+ where+ isBloomMember = Bloom.elem y xsBloom+ isTrueMember = Set.member y xsSet++fromTest :: Test -> Counts+fromTest = \case+ TruePositive -> Counts 1 0 0 0+ FalsePositive -> Counts 0 1 0 0+ TrueNegative -> Counts 0 0 1 0+ FalseNegative -> Counts 0 0 0 1+++data Counts = Counts {+ _cTruePositives :: !Int+ , _cFalsePositives :: !Int+ , _cTrueNegatives :: !Int+ , _cFalseNegatives :: !Int+ }++instance Semigroup Counts where+ (<>) :: Counts -> Counts -> Counts+ (Counts tp1 fp1 tn1 fn1) <> (Counts tp2 fp2 tn2 fn2) =+ Counts (tp1 + tp2) (fp1 + fp2) (tn1 + tn2) (fn1 + fn2)++instance Monoid Counts where+ mempty :: Counts+ mempty = Counts 0 0 0 0++{-------------------------------------------------------------------------------+ Bloom filter construction+-------------------------------------------------------------------------------}++type BloomMaker a = [a] -> Bloom a++-- | Create a bloom filter through the 'newMBloom' interface. Tunes the bloom+-- filter according to 'RunBloomFilterAlloc'.+mkBloomFromAlloc :: Hashable a => RunBloomFilterAlloc -> BloomMaker a+mkBloomFromAlloc alloc xs = runST $ do+ mb <- newMBloom n alloc testSalt+ mapM_ (Bloom.insert mb) xs+ Bloom.unsafeFreeze mb+ where+ n = LSMT.NumEntries $ length xs
@@ -0,0 +1,95 @@+{-# LANGUAGE LambdaCase #-}++module Test.Database.LSMTree.Internal.RunBuilder (tests) where++import Control.Monad.Class.MonadThrow+import Data.Foldable (traverse_)+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry (NumEntries (..))+import qualified Database.LSMTree.Internal.Index as Index+import Database.LSMTree.Internal.Paths (RunFsPaths (..))+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import Database.LSMTree.Internal.RunNumber+import qualified System.FS.API as FS+import System.FS.API (HasFS)+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.Sim.MockFS as MockFS+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.FS (propNoOpenHandles, withSimHasBlockIO,+ withTempIOHasBlockIO)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.RunBuilder" [+ testGroup "ioHasFS" [+ testProperty "prop_newInExistingDir" $ ioProperty $+ withTempIOHasBlockIO "prop_newInExistingDir" prop_newInExistingDir+ , testProperty "prop_newInNonExistingDir" $ ioProperty $+ withTempIOHasBlockIO "prop_newInNonExistingDir" prop_newInNonExistingDir+ , testProperty "prop_newTwice" $ ioProperty $+ withTempIOHasBlockIO "prop_newTwice" prop_newTwice+ ]+ , testGroup "simHasFS" [+ testProperty "prop_newInExistingDir" $ ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $+ \hfs hbio _ -> prop_newInExistingDir hfs hbio+ , testProperty "prop_newInNonExistingDir" $ ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $+ \hfs hbio _ -> prop_newInNonExistingDir hfs hbio+ , testProperty "prop_newTwice" $ ioProperty $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $+ \hfs hbio _ -> prop_newTwice hfs hbio+ ]+ ]++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Ordinary+ }++testSalt :: Bloom.Salt+testSalt = 4++-- | 'new' in an existing directory should be successful.+prop_newInExistingDir :: HasFS IO h -> FS.HasBlockIO IO h -> IO Property+prop_newInExistingDir hfs hbio = do+ let runDir = FS.mkFsPath ["a", "b", "c"]+ FS.createDirectoryIfMissing hfs True runDir+ bracket+ (try (RunBuilder.new hfs hbio testSalt runParams (RunFsPaths runDir (RunNumber 17)) (NumEntries 0)))+ (traverse_ RunBuilder.close) $ pure . \case+ Left e@FS.FsError{} ->+ counterexample ("expected a success, but got: " <> show e) $ property False+ Right _ -> property True++-- | 'new' in a non-existing directory should throw an error.+prop_newInNonExistingDir :: HasFS IO h -> FS.HasBlockIO IO h -> IO Property+prop_newInNonExistingDir hfs hbio = do+ let runDir = FS.mkFsPath ["a", "b", "c"]+ bracket+ (try (RunBuilder.new hfs hbio testSalt runParams (RunFsPaths runDir (RunNumber 17)) (NumEntries 0)))+ (traverse_ RunBuilder.close) $ pure . \case+ Left FS.FsError{} -> property True+ Right _ ->+ counterexample ("expected an FsError, but got a RunBuilder") $ property False++-- | Calling 'new' twice with the same arguments should throw an error.+--+-- TODO: maybe in this case a custom error should be thrown? Does the thrown+-- 'FsError' cause file resources to leak?+prop_newTwice :: HasFS IO h -> FS.HasBlockIO IO h -> IO Property+prop_newTwice hfs hbio = do+ let runDir = FS.mkFsPath []+ bracket+ (RunBuilder.new hfs hbio testSalt runParams (RunFsPaths runDir (RunNumber 17)) (NumEntries 0))+ RunBuilder.close $ \_ ->+ bracket+ (try (RunBuilder.new hfs hbio testSalt runParams (RunFsPaths runDir (RunNumber 17)) (NumEntries 0)))+ (traverse_ RunBuilder.close) $ pure . \case+ Left FS.FsError{} -> property True+ Right _ ->+ counterexample ("expected an FsError, but got a RunBuilder") $ property False
@@ -0,0 +1,201 @@+module Test.Database.LSMTree.Internal.RunReader (+ -- * Main test tree+ tests,+ -- * Utilities+ readKOps,+) where++import Control.RefCount+import Data.Coerce (coerce)+import qualified Data.Map as Map+import Database.LSMTree.Extras.Generators (BiasedKey (..))+import Database.LSMTree.Extras.RunData+import Database.LSMTree.Internal.BlobRef+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Entry (Entry)+import qualified Database.LSMTree.Internal.Index as Index (IndexType (Ordinary))+import Database.LSMTree.Internal.Run (Run)+import qualified Database.LSMTree.Internal.RunAcc as RunAcc+import qualified Database.LSMTree.Internal.RunBuilder as RunBuilder+import qualified Database.LSMTree.Internal.RunReader as Reader+import Database.LSMTree.Internal.Serialise+import qualified System.FS.API as FS+import qualified System.FS.BlockIO.API as FS+import qualified System.FS.Sim.MockFS as MockFS+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck+import Test.Util.FS (propNoOpenHandles, withSimHasBlockIO,+ withTempIOHasBlockIO)+import Test.Util.Orphans ()++tests :: TestTree+tests = testGroup "Database.LSMTree.Internal.RunReader"+ [ testGroup "MockFS"+ [ testProperty "prop_read" $ \wb ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ -> do+ prop_readAtOffset hfs hbio wb Nothing+ , testProperty "prop_readAtOffset" $ \wb offset ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ -> do+ prop_readAtOffset hfs hbio wb (Just offset)+ , testProperty "prop_readAtOffsetExisting" $ \wb i ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ -> do+ prop_readAtOffsetExisting hfs hbio wb i+ , testProperty "prop_readAtOffsetIdempotence" $ \wb i ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ -> do+ prop_readAtOffsetIdempotence hfs hbio wb i+ , testProperty "prop_readAtOffsetReadHead" $ \wb ->+ ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ -> do+ prop_readAtOffsetReadHead hfs hbio wb+ ]+ , testGroup "RealFS"+ [ testProperty "prop_read" $ \wb ->+ ioProperty $ withTempIOHasBlockIO "tmp_RunReader" $ \hfs hbio -> do+ prop_readAtOffset hfs hbio wb Nothing+ , testProperty "prop_readAtOffset" $ \wb offset ->+ ioProperty $ withTempIOHasBlockIO "tmp_RunReader" $ \hfs hbio -> do+ prop_readAtOffset hfs hbio wb (Just offset)+ , testProperty "prop_readAtOffsetExisting" $ \wb i ->+ ioProperty $ withTempIOHasBlockIO "tmp_RunReader" $ \hfs hbio -> do+ prop_readAtOffsetExisting hfs hbio wb i+ , testProperty "prop_readAtOffsetIdempotence" $ \wb i ->+ ioProperty $ withTempIOHasBlockIO "tmp_RunReader" $ \hfs hbio -> do+ prop_readAtOffsetIdempotence hfs hbio wb i+ , testProperty "prop_readAtOffsetReadHead" $ \wb ->+ ioProperty $ withTempIOHasBlockIO "tmp_RunReader" $ \hfs hbio -> do+ prop_readAtOffsetReadHead hfs hbio wb+ ]+ ]++runParams :: RunBuilder.RunParams+runParams =+ RunBuilder.RunParams {+ runParamCaching = RunBuilder.CacheRunData,+ runParamAlloc = RunAcc.RunAllocFixed 10,+ runParamIndex = Index.Ordinary+ }++testSalt :: Bloom.Salt+testSalt = 4++-- | Creating a run from a write buffer and reading from the run yields the+-- original elements.+--+-- @id === read . new . flush@+--+-- If there is an offset:+--+-- @dropWhile ((< offset) . key) === read . newAtOffset offset . flush@+--+prop_readAtOffset ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData BiasedKey SerialisedValue SerialisedBlob+ -> Maybe BiasedKey+ -> IO Property+prop_readAtOffset fs hbio rd offsetKey =+ withRunAt fs hbio testSalt runParams (simplePath 42) rd' $ \run -> do+ rhs <- readKOps (coerce offsetKey) run++ pure . labelRunData rd' $+ counterexample ("entries expected: " <> show (length lhs)) $+ counterexample ("entries found: " <> show (length rhs)) $+ lhs === rhs+ where+ rd' = serialiseRunData rd+ kops = Map.toList (unRunData rd')+ lhs = case offsetKey of+ Nothing -> kops+ Just k -> dropWhile ((< coerce k) . fst) kops++-- | A version of 'prop_readAtOffset' where the offset key is always one+-- of the keys that exist in the run.+prop_readAtOffsetExisting ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData BiasedKey SerialisedValue SerialisedBlob+ -> NonNegative Int+ -> IO Property+prop_readAtOffsetExisting fs hbio rd (NonNegative index)+ | null kops = pure discard+ | otherwise =+ prop_readAtOffset fs hbio rd (Just (keys !! (index `mod` length keys)))+ where+ keys :: [BiasedKey]+ keys = coerce (fst <$> kops)+ kops = Map.toList (unRunData rd)++-- | Idempotence of 'readAtOffset'.+-- Reading at an offset should not perform any stateful effects which alter+-- the result of a subsequent read at the same offset.+--+-- @readAtOffset offset *> readAtOffset offset === readAtOffset offset@+--+prop_readAtOffsetIdempotence ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData BiasedKey SerialisedValue SerialisedBlob+ -> Maybe BiasedKey+ -> IO Property+prop_readAtOffsetIdempotence fs hbio rd offsetKey =+ withRunAt fs hbio testSalt runParams (simplePath 42) rd' $ \run -> do+ lhs <- readKOps (coerce offsetKey) run+ rhs <- readKOps (coerce offsetKey) run++ pure . labelRunData rd' $+ counterexample ("entries expected: " <> show (length lhs)) $+ counterexample ("entries found: " <> show (length rhs)) $+ lhs === rhs+ where+ rd' = serialiseRunData rd++-- | Head of 'read' equals 'readAtOffset' of head of 'read'.+-- Reading the first key from the run initialized without an offset+-- should be the same as reading from a run initialized at the first key.+-- the result of a subsequent read at the same offset.+--+-- @read === readAtOffset (head read)@+--+prop_readAtOffsetReadHead ::+ FS.HasFS IO h+ -> FS.HasBlockIO IO h+ -> RunData BiasedKey SerialisedValue SerialisedBlob+ -> IO Property+prop_readAtOffsetReadHead fs hbio rd =+ withRunAt fs hbio testSalt runParams (simplePath 42) rd' $ \run -> do+ lhs <- readKOps Nothing run+ rhs <- case lhs of+ [] -> pure []+ (key,_):_ -> readKOps (Just key) run++ pure . labelRunData rd' $+ counterexample ("entries expected: " <> show (length lhs)) $+ counterexample ("entries found: " <> show (length rhs)) $+ lhs === rhs+ where+ rd' = serialiseRunData rd++{-------------------------------------------------------------------------------+ Utilities+-------------------------------------------------------------------------------}++type SerialisedEntry = Entry SerialisedValue SerialisedBlob+type SerialisedKOp = (SerialisedKey, SerialisedEntry)++readKOps ::+ Maybe (SerialisedKey) -- ^ offset+ -> Ref (Run IO h)+ -> IO [SerialisedKOp]+readKOps offset run = do+ reader <- Reader.new offsetKey run+ go reader+ where+ offsetKey = maybe Reader.NoOffsetKey (Reader.OffsetKey . coerce) offset++ go reader = do+ Reader.next reader >>= \case+ Reader.Empty -> pure []+ Reader.ReadEntry key e -> do+ let fs = Reader.readerHasFS reader+ e' <- traverse (readRawBlobRef fs) $ Reader.toFullEntry e+ ((key, e') :) <$> go reader+
@@ -0,0 +1,37 @@+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.Serialise (tests) where++import Data.Bits+import qualified Data.ByteString as BS+import qualified Data.ByteString.Builder as BB+import qualified Data.ByteString.Short as SBS+import qualified Data.Vector.Primitive as VP+import Data.Word+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.Serialise+import Test.Tasty+import Test.Tasty.HUnit++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Serialise" [+ testCase "example keyTopBits64" $ do+ let k = SerialisedKey' (VP.fromList [0, 0, 0, 0, 37, 42, 204, 130])+ expected :: Word64+ expected = 37 `shiftL` 24 + 42 `shiftL` 16 + 204 `shiftL` 8 + 130+ expected @=? keyTopBits64 k+ , testCase "example keyTopBits64 on sliced byte array" $ do+ let pvec = VP.fromList [0, 0, 0, 0, 0, 37, 42, 204, 130]+ k = SerialisedKey' (VP.slice 1 (VP.length pvec - 1) pvec)+ expected :: Word64+ expected = 37 `shiftL` 24 + 42 `shiftL` 16 + 204 `shiftL` 8 + 130+ expected @=? keyTopBits64 k+ , testCase "example unsafeFromByteString and fromShortByteString" $ do+ let bb = mconcat [BB.word64LE x | x <- [0..100]]+ bs = BS.toStrict . BB.toLazyByteString $ bb+ k1 = RB.unsafeFromByteString bs+ k2 = RB.fromShortByteString (SBS.toShort bs)+ k1 @=? k2+ ]
@@ -0,0 +1,130 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}++module Test.Database.LSMTree.Internal.Serialise.Class (tests) where++import Data.ByteString (ByteString)+import Data.ByteString.Lazy (LazyByteString)+import Data.ByteString.Short (ShortByteString)+import Data.Constraint (Dict (..))+import Data.Int+import Data.Monoid (Sum)+import Data.Primitive (ByteArray)+import Data.WideWord (Word128, Word256)+import Data.Word+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Extras.UTxO (UTxOKey, UTxOValue)+import Database.LSMTree.Internal.Serialise.Class+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Serialise.Class" [+ testGroup "String" (allProperties @String (Just Dict))++ -- Int+ , testGroup "Int8" (allProperties @Int8 Nothing)+ , testGroup "Int16" (allProperties @Int16 Nothing)+ , testGroup "Int32" (allProperties @Int32 Nothing)+ , testGroup "Int64" (allProperties @Int64 Nothing)+ , testGroup "Int" (allProperties @Int Nothing)++ -- Word+ , testGroup "Word8" (allProperties @Word8 (Just Dict))+ , testGroup "Word16" (allProperties @Word16 (Just Dict))+ , testGroup "Word32" (allProperties @Word32 (Just Dict))+ , testGroup "Word64" (allProperties @Word64 (Just Dict))+ , testGroup "Word" (allProperties @Word (Just Dict))++ -- ByteString+ , testGroup "ByteString" (allProperties @ByteString (Just Dict))+ , testGroup "LazyByteString" (allProperties @LazyByteString (Just Dict))+ , testGroup "ShortByteString" (allProperties @ShortByteString (Just Dict))+ , testGroup "ByteArray" (allProperties @ByteArray Nothing )++ -- UTXO+ , testGroup "Word128" (allProperties @Word128 (Just Dict))+ , testGroup "Word256" (allProperties @Word256 (Just Dict))+ , testGroup "UTxOKey" (keyProperties @UTxOKey Nothing )+ , testGroup "UTxOValue" (valueProperties @UTxOValue )++ -- Modifiers+ , testGroup "Sum Word" (valueProperties @(Sum Word))+ ]++allProperties ::+ forall a. (Ord a, Show a, Arbitrary a, SerialiseKey a, SerialiseValue a)+ => Maybe (Dict (SerialiseKeyOrderPreserving a))+ -> [TestTree]+allProperties orderPreserving = keyProperties @a orderPreserving <> valueProperties @a++keyProperties ::+ forall a. (Ord a, Show a, Arbitrary a, SerialiseKey a)+ => Maybe (Dict (SerialiseKeyOrderPreserving a))+ -> [TestTree]+keyProperties orderPreserving =+ [ testProperty "prop_roundtripSerialiseKey" $+ prop_roundtripSerialiseKey @a+ , testProperty "prop_roundtripSerialiseKeyUpToSlicing" $+ prop_roundtripSerialiseKeyUpToSlicing @a+ , testProperty "prop_orderPreservationSerialiseKey" $ \k1 k2 ->+ let modifier = case orderPreserving of+ Just Dict -> id+ Nothing -> expectFailure+ in modifier $ prop_orderPreservationSerialiseKey @a k1 k2++ ]++valueProperties :: forall a. (Ord a, Show a, Arbitrary a, SerialiseValue a) => [TestTree]+valueProperties =+ [ testProperty "prop_roundtripSerialiseValue" $+ prop_roundtripSerialiseValue @a+ , testProperty "prop_roundtripSerialiseValueUpToSlicing" $+ prop_roundtripSerialiseValueUpToSlicing @a+ ]++prop_roundtripSerialiseKey :: forall k. (Eq k, Show k, SerialiseKey k) => k -> Property+prop_roundtripSerialiseKey k =+ counterexample ("serialised: " <> show (serialiseKey k)) $+ counterexample ("deserialised: " <> show @k (deserialiseKey (serialiseKey k))) $+ serialiseKeyIdentity k++prop_roundtripSerialiseKeyUpToSlicing ::+ forall k. (Eq k, Show k, SerialiseKey k)+ => RawBytes -> k -> RawBytes -> Property+prop_roundtripSerialiseKeyUpToSlicing prefix x suffix =+ counterexample ("serialised: " <> show @RawBytes k) $+ counterexample ("serialised and sliced: " <> show @RawBytes k') $+ counterexample ("deserialised: " <> show @k x') $+ serialiseKeyIdentityUpToSlicing prefix x suffix+ where+ k = serialiseKey x+ k' = packSlice prefix k suffix+ x' = deserialiseKey k'++prop_orderPreservationSerialiseKey :: forall k. (Ord k, SerialiseKey k) => k -> k -> Property+prop_orderPreservationSerialiseKey x y =+ counterexample ("serialised: " <> show (serialiseKey x, serialiseKey y)) $+ counterexample ("compare: " <> show (compare x y)) $+ counterexample ("compare serialised: " <> show (compare (serialiseKey x) (serialiseKey y))) $+ serialiseKeyPreservesOrdering x y++prop_roundtripSerialiseValue :: forall v. (Eq v, Show v, SerialiseValue v) => v -> Property+prop_roundtripSerialiseValue v =+ counterexample ("serialised: " <> show (serialiseValue v)) $+ counterexample ("deserialised: " <> show @v (deserialiseValue (serialiseValue v))) $+ serialiseValueIdentity v++prop_roundtripSerialiseValueUpToSlicing ::+ forall v. (Eq v, Show v, SerialiseValue v)+ => RawBytes -> v -> RawBytes -> Property+prop_roundtripSerialiseValueUpToSlicing prefix x suffix =+ counterexample ("serialised: " <> show v) $+ counterexample ("serialised and sliced: " <> show @RawBytes v') $+ counterexample ("deserialised: " <> show @v x') $+ serialiseValueIdentityUpToSlicing prefix x suffix+ where+ v = serialiseValue x+ v' = packSlice prefix v suffix+ x' = deserialiseValue v'
@@ -0,0 +1,473 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.Snapshot.Codec (tests) where++import Codec.CBOR.Decoding+import Codec.CBOR.Encoding+import Codec.CBOR.FlatTerm+import Codec.CBOR.Read+import Codec.CBOR.Write+import Control.DeepSeq (NFData)+import qualified Data.ByteString.Lazy as BSL+import Data.Proxy+import qualified Data.Text as Text+import Data.Typeable+import qualified Data.Vector as V+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.MergeSchedule+import Database.LSMTree.Internal.MergingRun+import Database.LSMTree.Internal.RunBuilder (IndexType (..),+ RunBloomFilterAlloc (..), RunDataCaching (..))+import Database.LSMTree.Internal.RunNumber+import Database.LSMTree.Internal.Snapshot+import Database.LSMTree.Internal.Snapshot.Codec+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.Arbitrary++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Snapshot.Codec" [+ testGroup "SnapshotVersion" [+ testProperty "roundtripCBOR" $ roundtripCBOR (Proxy @SnapshotVersion)+ , testProperty "roundtripFlatTerm" $ roundtripFlatTerm (Proxy @SnapshotVersion)+ ]+ , testGroup "Versioned SnapshotMetaData" [+ testProperty "roundtripCBOR" $ roundtripCBOR (Proxy @(Versioned SnapshotMetaData))+ , testProperty "roundtripFlatTerm" $ roundtripFlatTerm (Proxy @(Versioned SnapshotMetaData))+ ]+ , testGroup "roundtripCBOR'" $+ propAll roundtripCBOR'+ , testGroup "roundtripFlatTerm'" $+ propAll roundtripFlatTerm'+ -- Test generators and shrinkers+ , testGroup "Generators and shrinkers are finite" $+ testAll $ \(p :: Proxy a) ->+ testGroup (show $ typeRep p) $+ prop_arbitraryAndShrinkPreserveInvariant @a noTags deepseqInvariant+ ]++{-------------------------------------------------------------------------------+ Properties+-------------------------------------------------------------------------------}++-- | @decode . encode = id@+explicitRoundtripCBOR ::+ (Eq a, Show a)+ => (a -> Encoding)+ -> (forall s. Decoder s a)+ -> a+ -> Property+explicitRoundtripCBOR enc dec x = case back (there x) of+ Left e ->+ counterexample+ ("Decoding failed: " <> show e)+ (property False)+ Right (bs, y) ->+ counterexample+ ("Expected value and roundtripped value do not match")+ (x === y)+ .&&. counterexample+ ("Found trailing bytes")+ (property (BSL.null bs))+ where+ there = toLazyByteString . enc+ back = deserialiseFromBytes dec++-- | See 'explicitRoundtripCBOR'.+roundtripCBOR :: (Encode a, Decode a, Eq a, Show a) => Proxy a -> a -> Property+roundtripCBOR _ = explicitRoundtripCBOR encode decode++-- | See 'explicitRoundtripCBOR'.+roundtripCBOR' :: (Encode a, DecodeVersioned a, Eq a, Show a) => Proxy a -> a -> Property+roundtripCBOR' _ = explicitRoundtripCBOR encode (decodeVersioned currentSnapshotVersion)++-- | @fromFlatTerm . toFlatTerm = id@+--+-- This will also check @validFlatTerm@ on the result of @toFlatTerm@.+explicitRoundtripFlatTerm ::+ (Eq a, Show a)+ => (a -> Encoding)+ -> (forall s. Decoder s a)+ -> a+ -> Property+explicitRoundtripFlatTerm enc dec x = case back flatTerm of+ Left e ->+ counterexample+ ("Decoding failed: " <> show e)+ (property False)+ Right y ->+ counterexample+ ("Expected value and roundtripped value do not match")+ (x === y)+ .&&. counterexample+ ("Invalid flat term")+ (property (validFlatTerm flatTerm))+ where+ flatTerm = there x++ there = toFlatTerm . enc+ back = fromFlatTerm dec++-- | See 'explicitRoundtripFlatTerm'.+roundtripFlatTerm ::+ (Encode a, Decode a, Eq a, Show a)+ => Proxy a+ -> a+ -> Property+roundtripFlatTerm _ = explicitRoundtripFlatTerm encode decode++-- | See 'explicitRoundtripFlatTerm'.+roundtripFlatTerm' ::+ (Encode a, DecodeVersioned a, Eq a, Show a)+ => Proxy a+ -> a+ -> Property+roundtripFlatTerm' _ = explicitRoundtripFlatTerm encode (decodeVersioned currentSnapshotVersion)++{-------------------------------------------------------------------------------+ Test and property runners+-------------------------------------------------------------------------------}++type Constraints a = (+ Eq a, Show a, Typeable a, Arbitrary a+ , Encode a, DecodeVersioned a, NFData a+ )++-- | Run a property on all types in the snapshot metadata hierarchy.+propAll ::+ (forall a. Constraints a => Proxy a -> a -> Property)+ -> [TestTree]+propAll prop = testAll mkTest+ where+ mkTest :: forall a. Constraints a => Proxy a -> TestTree+ mkTest pa = testProperty (show $ typeRep pa) (prop pa)++-- | Run a test on all types in the snapshot metadata hierarchy.+testAll ::+ (forall a. Constraints a => Proxy a -> TestTree)+ -> [TestTree]+testAll test = [+ -- SnapshotMetaData+ test (Proxy @SnapshotMetaData)+ , test (Proxy @SnapshotLabel)+ , test (Proxy @SnapshotRun)+ -- TableConfig+ , test (Proxy @TableConfig)+ , test (Proxy @MergePolicy)+ , test (Proxy @SizeRatio)+ , test (Proxy @WriteBufferAlloc)+ , test (Proxy @BloomFilterAlloc)+ , test (Proxy @FencePointerIndexType)+ , test (Proxy @DiskCachePolicy)+ , test (Proxy @MergeSchedule)+ -- SnapLevels+ , test (Proxy @(SnapLevels SnapshotRun))+ , test (Proxy @(SnapLevel SnapshotRun))+ , test (Proxy @(V.Vector SnapshotRun))+ , test (Proxy @RunNumber)+ , test (Proxy @(SnapIncomingRun SnapshotRun))+ , test (Proxy @MergePolicyForLevel)+ , test (Proxy @RunDataCaching)+ , test (Proxy @RunBloomFilterAlloc)+ , test (Proxy @IndexType)+ , test (Proxy @RunParams)+ , test (Proxy @(SnapMergingRun LevelMergeType SnapshotRun))+ , test (Proxy @MergeDebt)+ , test (Proxy @MergeCredits)+ , test (Proxy @NominalDebt)+ , test (Proxy @NominalCredits)+ , test (Proxy @LevelMergeType)+ , test (Proxy @TreeMergeType)+ , test (Proxy @(SnapMergingTree SnapshotRun))+ , test (Proxy @(SnapMergingTreeState SnapshotRun))+ , test (Proxy @(SnapMergingRun TreeMergeType SnapshotRun))+ , test (Proxy @(SnapPendingMerge SnapshotRun))+ , test (Proxy @(SnapPreExistingRun SnapshotRun))+ ]++{-------------------------------------------------------------------------------+ Arbitrary: versioning+-------------------------------------------------------------------------------}++instance Arbitrary SnapshotVersion where+ arbitrary = elements [V0, V1]+ shrink V0 = []+ shrink V1 = [V0]++deriving newtype instance Arbitrary a => Arbitrary (Versioned a)++{-------------------------------------------------------------------------------+ Arbitrary: SnapshotMetaData+-------------------------------------------------------------------------------}++instance Arbitrary SnapshotMetaData where+ arbitrary = SnapshotMetaData <$>+ arbitrary <*> arbitrary <*> arbitrary <*>+ arbitrary <*> arbitrary+ shrink (SnapshotMetaData a b c d e) =+ [ SnapshotMetaData a' b' c' d' e'+ | (a', b', c', d', e') <- shrink (a, b, c, d, e)]++instance Arbitrary SnapshotLabel where+ -- Ensure that the labeling string is not excessively long.+ -- If too long, negatively effects the number of shrinks required to reach the+ -- minimum example value.+ arbitrary = do+ prefix <- arbitraryPrintableChar+ suffix <- vectorOfUpTo 3 arbitraryPrintableChar+ pure . SnapshotLabel . Text.pack $ prefix : suffix+ shrink (SnapshotLabel txt) = SnapshotLabel <$> shrink txt++instance Arbitrary SnapshotRun where+ arbitrary = SnapshotRun <$> arbitrary <*> arbitrary <*> arbitrary+ shrink (SnapshotRun a b c) =+ [ SnapshotRun a' b' c'+ | (a', b', c') <- shrink (a, b, c)]++{-------------------------------------------------------------------------------+ Arbitrary: TableConfig+-------------------------------------------------------------------------------}++instance Arbitrary TableConfig where+ arbitrary =+ TableConfig <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary+ <*> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary+ shrink (TableConfig a b c d e f g h) =+ [ TableConfig a' b' c' d' e' f' g' h'+ | (a', b', c', d', e', f', g', h') <- shrink (a, b, c, d, e, f, g, h) ]++instance Arbitrary MergePolicy where+ arbitrary = pure LazyLevelling+ shrink LazyLevelling = []++instance Arbitrary SizeRatio where+ arbitrary = pure Four+ shrink Four = []++instance Arbitrary WriteBufferAlloc where+ arbitrary = AllocNumEntries <$> arbitrary+ shrink (AllocNumEntries x) = AllocNumEntries <$> shrink x++instance Arbitrary BloomFilterAlloc where+ arbitrary = oneof [+ AllocFixed <$> arbitrary+ , AllocRequestFPR <$> arbitrary+ ]+ shrink (AllocFixed x) = AllocFixed <$> shrink x+ shrink (AllocRequestFPR x) = AllocRequestFPR <$> shrink x++instance Arbitrary FencePointerIndexType where+ arbitrary = elements [CompactIndex, OrdinaryIndex]+ shrink _ = []++instance Arbitrary DiskCachePolicy where+ arbitrary = oneof [+ pure DiskCacheAll+ , DiskCacheLevelOneTo <$> arbitrary+ , pure DiskCacheNone+ ]+ shrink (DiskCacheLevelOneTo x) = DiskCacheLevelOneTo <$> shrink x+ shrink _ = []++instance Arbitrary MergeSchedule where+ arbitrary = elements [OneShot, Incremental]+ shrink _ = []++instance Arbitrary MergeBatchSize where+ arbitrary = MergeBatchSize <$> arbitrary+ shrink (MergeBatchSize n) = map MergeBatchSize (shrink n)++{-------------------------------------------------------------------------------+ Arbitrary: SnapLevels+-------------------------------------------------------------------------------}++instance Arbitrary r => Arbitrary (SnapLevels r) where+ arbitrary = SnapLevels <$> arbitraryShortVector+ shrink (SnapLevels x) = SnapLevels . V.fromList <$> shrink (V.toList x)++instance Arbitrary r => Arbitrary (SnapLevel r) where+ arbitrary = SnapLevel <$> arbitrary <*> arbitraryShortVector+ shrink (SnapLevel a b) = [SnapLevel a' b' | (a', b') <- shrink (a, b)]++arbitraryShortVector :: Arbitrary a => Gen (V.Vector a)+arbitraryShortVector = V.fromList <$> vectorOfUpTo 5 arbitrary++vectorOfUpTo :: Int -> Gen a -> Gen [a]+vectorOfUpTo maxlen gen = do+ len <- chooseInt (0, maxlen)+ vectorOf len gen++instance Arbitrary RunNumber where+ arbitrary = RunNumber <$> arbitrarySizedNatural+ shrink (RunNumber n) =+ -- fewer shrinks+ [RunNumber 0 | n > 0]+ ++ [RunNumber (n `div` 2) | n >= 2]++instance Arbitrary r => Arbitrary (SnapIncomingRun r) where+ arbitrary = oneof [+ SnapIncomingMergingRun <$> arbitrary <*> arbitrary+ <*> arbitrary <*> arbitrary+ , SnapIncomingSingleRun <$> arbitrary+ ]+ shrink (SnapIncomingMergingRun a b c d) =+ [ SnapIncomingMergingRun a' b' c' d'+ | (a', b', c', d') <- shrink (a, b, c, d) ]+ shrink (SnapIncomingSingleRun a) = SnapIncomingSingleRun <$> shrink a++instance Arbitrary MergePolicyForLevel where+ arbitrary = elements [LevelTiering, LevelLevelling]+ shrink _ = []++instance (Arbitrary t, Arbitrary r) => Arbitrary (SnapMergingRun t r) where+ arbitrary = oneof [+ SnapCompletedMerge <$> arbitrary <*> arbitrary+ , SnapOngoingMerge <$> arbitrary <*> arbitrary+ <*> arbitraryShortVector <*> arbitrary+ ]+ shrink (SnapCompletedMerge a b) =+ [ SnapCompletedMerge a' b'+ | (a', b') <- shrink (a, b) ]+ shrink (SnapOngoingMerge a b c d) =+ [ SnapOngoingMerge a' b' c' d'+ | (a', b', c', d') <- shrink (a, b, c, d) ]++deriving newtype instance Arbitrary MergeDebt+deriving newtype instance Arbitrary MergeCredits+deriving newtype instance Arbitrary NominalDebt+deriving newtype instance Arbitrary NominalCredits++{-------------------------------------------------------------------------------+ RunParams+-------------------------------------------------------------------------------}++instance Arbitrary RunParams where+ arbitrary = RunParams <$> arbitrary <*> arbitrary <*> arbitrary+ shrink (RunParams a b c) =+ [ RunParams a' b' c'+ | (a', b', c') <- shrink (a, b, c) ]++instance Arbitrary RunDataCaching where+ arbitrary = elements [CacheRunData, NoCacheRunData]+ shrink _ = []++instance Arbitrary IndexType where+ arbitrary = elements [Ordinary, Compact]+ shrink _ = []++instance Arbitrary RunBloomFilterAlloc where+ arbitrary = oneof [+ RunAllocFixed <$> arbitrary+ , RunAllocRequestFPR <$> arbitrary+ ]+ shrink (RunAllocFixed x) = RunAllocFixed <$> shrink x+ shrink (RunAllocRequestFPR x) = RunAllocRequestFPR <$> shrink x++{-------------------------------------------------------------------------------+ Show+-------------------------------------------------------------------------------}++deriving stock instance Show SnapshotMetaData+deriving stock instance Show SnapshotRun+deriving stock instance Show r => Show (SnapLevels r)+deriving stock instance Show r => Show (SnapLevel r)+deriving stock instance Show r => Show (SnapIncomingRun r)+deriving stock instance (Show t, Show r) => Show (SnapMergingRun t r)++deriving stock instance Show r => Show (SnapMergingTree r)+deriving stock instance Show r => Show (SnapMergingTreeState r)+deriving stock instance Show r => Show (SnapPendingMerge r)+deriving stock instance Show r => Show (SnapPreExistingRun r)++deriving stock instance Show MergeDebt+deriving stock instance Show NominalDebt+deriving stock instance Show NominalCredits++{-------------------------------------------------------------------------------+ Arbitrary: SnapshotMetaData+-------------------------------------------------------------------------------}++deriving newtype instance Arbitrary r => Arbitrary (SnapMergingTree r)++instance Arbitrary r => Arbitrary (SnapMergingTreeState r) where+ arbitrary = genMergingTreeState mergingTreeDepthLimit+ shrink (SnapCompletedTreeMerge a) = SnapCompletedTreeMerge <$> shrink a+ shrink (SnapPendingTreeMerge a) = SnapPendingTreeMerge <$> shrink a+ shrink (SnapOngoingTreeMerge a) = SnapOngoingTreeMerge <$> shrink a++instance Arbitrary r => Arbitrary (SnapPendingMerge r) where+ arbitrary = genPendingTreeMerge mergingTreeDepthLimit+ shrink (SnapPendingUnionMerge a) = SnapPendingUnionMerge <$> shrinkList shrink a+ shrink (SnapPendingLevelMerge a b) =+ [ SnapPendingLevelMerge a' b' | (a', b') <- shrink (a, b)]++instance Arbitrary r => Arbitrary (SnapPreExistingRun r) where+ arbitrary = oneof [+ SnapPreExistingRun <$> arbitrary+ , SnapPreExistingMergingRun <$> arbitrary+ ]+ shrink (SnapPreExistingRun a) = SnapPreExistingRun <$> shrink a+ shrink (SnapPreExistingMergingRun a) = SnapPreExistingMergingRun <$> shrink a++-- | The 'SnapMergingTree' is an inductive data-type and therefore we must limit+-- the recursive depth at which new 'Arbitrary' sub-trees are generated. Hence+-- the need for this limit. This limit is the "gas" for the inductive functions.+-- At reach recursive call, the "gas" value decremented until it reaches zero.+-- Each inductive function ensures it never create a forest of sub-trees greater+-- than the /monotonically decreasing/ gas parameter it received.+mergingTreeDepthLimit :: Int+mergingTreeDepthLimit = 4++-- | Do not generate a number of direct child sub-trees greater than the this+-- branching limit. This simplifies the topology of trees generated+branchingLimit :: Int+branchingLimit = 3++-- | Generate an 'Arbitrary', "gas-limited" 'SnapMergingTreeState~'.+genMergingTreeState :: Arbitrary a => Int -> Gen (SnapMergingTreeState a)+genMergingTreeState gas =+ let perpetualCase = [+ SnapCompletedTreeMerge <$> arbitrary+ , SnapOngoingTreeMerge <$> arbitrary+ ]+ recursiveCase+ | gas == 0 = []+ | otherwise = [ SnapPendingTreeMerge <$> genPendingTreeMerge gas ]+ in oneof $ perpetualCase <> recursiveCase++-- | Generate an 'Arbitrary', "gas-limited" 'SnapPendingMerge'.+genPendingTreeMerge :: Arbitrary a => Int -> Gen (SnapPendingMerge a)+genPendingTreeMerge gas =+ oneof [+ SnapPendingLevelMerge <$> genPreExistings <*> genMaybeSubTree+ , SnapPendingUnionMerge <$> genListSubtrees+ ]+ where+ -- Decrement the gas for the recursive calls+ nextGas = max 0 $ gas - 1+ subGen = SnapMergingTree <$> genMergingTreeState nextGas++ -- No recursive subtrees within here, so not constrained by gas.+ genPreExistings = vectorOfUpTo branchingLimit arbitrary++ -- Define custom generators to ensure that the sub-trees are less than+ -- or equal to the lesser of the "gas" parameter and the branching limit.+ genMaybeSubTree+ | gas == 0 = pure Nothing+ | otherwise = oneof [ pure Nothing, Just <$> subGen ]++ genListSubtrees = case gas of+ 0 -> vectorOf 0 subGen+ _ ->+ -- This frequency distribution will uniformly at random select an+ -- n-ary tree topology with a specified branching factor.+ let recursiveOptions branching = \case+ 0 -> 1+ depth ->+ let sub = recursiveOptions branching $ depth - 1+ in sum $ (sub ^) <$> [ 0 .. branching ]+ probability e =+ let basis = recursiveOptions branchingLimit nextGas+ in (basis ^ e, vectorOf e subGen)+ in frequency $ probability <$> [ 0 .. branchingLimit ]
@@ -0,0 +1,528 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE UndecidableInstances #-}++module Test.Database.LSMTree.Internal.Snapshot.Codec.Golden (+ tests+ , EnumGolden (..)+ , Annotation+ ) where++import Codec.CBOR.Write (toLazyByteString)+import Control.Monad (when)+import qualified Data.ByteString.Lazy as BSL (writeFile)+import qualified Data.Set as Set+import Data.Typeable+import qualified Data.Vector as V+import Database.LSMTree.Internal.Config (BloomFilterAlloc (..),+ DiskCachePolicy (..), FencePointerIndexType (..),+ MergeBatchSize (..), MergePolicy (..), MergeSchedule (..),+ SizeRatio (..), TableConfig (..), WriteBufferAlloc (..))+import Database.LSMTree.Internal.MergeSchedule+ (MergePolicyForLevel (..), NominalCredits (..),+ NominalDebt (..))+import Database.LSMTree.Internal.MergingRun as MR+import Database.LSMTree.Internal.RunBuilder (IndexType (..),+ RunBloomFilterAlloc (..), RunDataCaching (..))+import Database.LSMTree.Internal.RunNumber (RunNumber (..))+import Database.LSMTree.Internal.Snapshot+import Database.LSMTree.Internal.Snapshot.Codec+import Paths_lsm_tree+import qualified System.Directory as Dir+import System.FilePath+import qualified System.FS.API as FS+import System.FS.API.Types (MountPoint (..))+import System.FS.IO (ioHasFS)+import System.IO.Unsafe+import Test.QuickCheck (Property, counterexample, ioProperty, once,+ (.&&.))+import qualified Test.Tasty as Tasty+import Test.Tasty (TestTree, testGroup)+import qualified Test.Tasty.Golden as Au+import Test.Tasty.QuickCheck (testProperty)++tests :: TestTree+tests =+ handleOutputFiles $+ testGroup "Test.Database.LSMTree.Internal.Snapshot.Codec.Golden" $+ concat (forallSnapshotTypes snapshotCodecGoldenTest)+ ++ [testProperty "prop_noUnexpectedOrMissingGoldenFiles" prop_noUnexpectedOrMissingGoldenFiles]++{-------------------------------------------------------------------------------+ Configuration+-------------------------------------------------------------------------------}++-- | The location of the golden file data directory relative to the project root.+goldenDataFilePath :: FilePath+goldenDataFilePath = unsafePerformIO getDataDir </> "snapshot-codec"++goldenDataMountPoint :: MountPoint+goldenDataMountPoint = MountPoint goldenDataFilePath++-- | Delete output files on test-case success.+-- Change the option here if this is undesirable.+handleOutputFiles :: TestTree -> TestTree+handleOutputFiles = Tasty.localOption Au.OnPass++{-------------------------------------------------------------------------------+ Golden tests+-------------------------------------------------------------------------------}++-- | Compare the serialization of snapshot metadata with a known reference file.+snapshotCodecGoldenTest ::+ forall a. (Typeable a, EnumGolden a, Encode a)+ => Proxy a+ -> [TestTree]+snapshotCodecGoldenTest proxy = [+ go annotation datum+ | (annotation, datum) <- enumGoldenAnnotated' proxy+ ]+ where+ go ann datum =+ let v = currentSnapshotVersion+ outputFilePath = goldenDataFilePath </> filePathOutput proxy ann v+ goldenFilePath = goldenDataFilePath </> filePathGolden proxy ann v++ -- IO actions+ removeIfExists :: FilePath -> IO ()+ removeIfExists fp =+ Dir.doesFileExist fp >>= (`when` (Dir.removeFile fp))+ outputAction :: IO ()+ outputAction = do+ -- Ensure that if the output file already exists, we remove it and+ -- re-write out the serialized data. This ensures that there are no+ -- false-positives, false-negatives, or irrelevant I/O exceptions.+ removeIfExists outputFilePath+ BSL.writeFile outputFilePath . toLazyByteString $ encode datum++ in Au.goldenVsFile (nameGolden proxy ann v) goldenFilePath outputFilePath outputAction++-- | Check that are no missing or unexpected files in the output directory+prop_noUnexpectedOrMissingGoldenFiles :: Property+prop_noUnexpectedOrMissingGoldenFiles = once $ ioProperty $ do+ let expectedFiles = Set.fromList $ concat $ forallSnapshotTypes $ \p -> concat [+ filePathsGolden p v+ | v <- supportedVersions p+ ]++ let hfs = ioHasFS goldenDataMountPoint+ actualDirectoryEntries <- FS.listDirectory hfs (FS.mkFsPath [])++ let missingFiles = expectedFiles Set.\\ actualDirectoryEntries+ propMissing =+ counterexample ("Missing expected files: " ++ show missingFiles)+ $ counterexample ("Run the golden tests to regenerate the missing files")+ $ (Set.null missingFiles)++ let unexpectedFiles = actualDirectoryEntries Set.\\ expectedFiles+ propUnexpected =+ counterexample ("Found unexpected files: " ++ show unexpectedFiles)+ $ counterexample ("Delete the unexpected files manually from " ++ goldenDataFilePath)+ (Set.null unexpectedFiles)++ pure $ propMissing .&&. propUnexpected++{-------------------------------------------------------------------------------+ Mapping+-------------------------------------------------------------------------------}++type Constraints a = (Typeable a, Encode a, EnumGolden a)++-- | Do something for all snapshot types and collect the results+forallSnapshotTypes ::+ (forall a. Constraints a => Proxy a -> b)+ -> [b]+forallSnapshotTypes f = [+ -- SnapshotMetaData+ f (Proxy @SnapshotMetaData)+ , f (Proxy @SnapshotLabel)+ , f (Proxy @SnapshotRun)+ -- TableConfig+ , f (Proxy @TableConfig)+ , f (Proxy @MergePolicy)+ , f (Proxy @SizeRatio)+ , f (Proxy @WriteBufferAlloc)+ , f (Proxy @BloomFilterAlloc)+ , f (Proxy @FencePointerIndexType)+ , f (Proxy @DiskCachePolicy)+ , f (Proxy @MergeSchedule)+ , f (Proxy @MergeBatchSize)+ -- SnapLevels+ , f (Proxy @(SnapLevels SnapshotRun))+ , f (Proxy @(SnapLevel SnapshotRun))+ , f (Proxy @(V.Vector SnapshotRun))+ , f (Proxy @RunNumber)+ , f (Proxy @(SnapIncomingRun SnapshotRun))+ , f (Proxy @MergePolicyForLevel)+ , f (Proxy @RunDataCaching)+ , f (Proxy @RunBloomFilterAlloc)+ , f (Proxy @IndexType)+ , f (Proxy @RunParams)+ , f (Proxy @(SnapMergingRun LevelMergeType SnapshotRun))+ , f (Proxy @MergeDebt)+ , f (Proxy @MergeCredits)+ , f (Proxy @NominalDebt)+ , f (Proxy @NominalCredits)+ , f (Proxy @LevelMergeType)+ , f (Proxy @TreeMergeType)+ , f (Proxy @(SnapMergingTree SnapshotRun))+ , f (Proxy @(SnapMergingTreeState SnapshotRun))+ , f (Proxy @(SnapMergingRun TreeMergeType SnapshotRun))+ , f (Proxy @(SnapPendingMerge SnapshotRun))+ , f (Proxy @(SnapPreExistingRun SnapshotRun))+ ]++{-------------------------------------------------------------------------------+ Enumeration class+-------------------------------------------------------------------------------}++-- | Enumerate values of type @a@ for golden tests+--+-- To prevent combinatorial explosion, the enumeration should generally be+-- /shallow/: the different constructors for type @a@ should be enumerated+-- without recursively enumerating the constructors' fields. For example,+-- enumerating @Maybe Int@ should give us something like:+--+-- > enumGolden @(Maybe Int) = [ Just 17, Nothing ]+--+-- This is generally a suitable approach, since the snapshot encodings do not+-- encode types differently depending on values in the constructor fields.+--+-- Example (recursive) instances that prevent combinatorial explosion:+--+-- @+-- instance EnumGolden a => EnumGolden (Maybe a) where+-- enumGolden = [ Just (singGolden @a), Nothing ]+-- instance EnumGolden Int where+-- enumGolden = [17, -72] -- singGolden = 17+-- @+--+-- If there are encoders that do require more elaborate (recursive)+-- enumerations, then it is okay to deviate from shallow enumerations, but take+-- care not to explode the combinatorics ;)+class EnumGolden a where+ {-# MINIMAL enumGolden | enumGoldenAnnotated | singGolden #-}++ -- | Enumerated values. The enumeration should be /shallow/.+ --+ -- The default implementation is to return a singleton list containing+ -- 'singGolden'.+ enumGolden :: [a]+ enumGolden = [ singGolden ]++ -- | Enumerated values with an annotation for naming purposes. The enumeration+ -- should be /shallow/, and the annotations should be unique.+ --+ -- The default implementation is to annotate 'enumGolden' with capital letters+ -- starting with @\'A\'@.+ enumGoldenAnnotated :: [(Annotation, a)]+ enumGoldenAnnotated = zip [[c] | c <- ['A' .. 'Z']] enumGolden++ -- | A singleton enumerated value. This is mainly useful for superclass+ -- instances.+ --+ -- The default implementation is to take the 'head' of 'enumGoldenAnnotated'.+ singGolden :: a+ singGolden = snd $ head enumGoldenAnnotated++ supportedVersions :: Proxy a -> [SnapshotVersion]+ supportedVersions _ = allCompatibleSnapshotVersions++type Annotation = String++enumGoldenAnnotated' :: EnumGolden a => Proxy a -> [(Annotation, a)]+enumGoldenAnnotated' _ = enumGoldenAnnotated++{-------------------------------------------------------------------------------+ Enumeration class: names and file paths+-------------------------------------------------------------------------------}++nameGolden :: Typeable a => Proxy a -> Annotation -> SnapshotVersion -> String+nameGolden p ann v = show v ++ "." ++ map spaceToUnderscore (show $ typeRep p) ++ "." ++ ann++spaceToUnderscore :: Char -> Char+spaceToUnderscore ' ' = '_'+spaceToUnderscore c = c++filePathsGolden ::+ (EnumGolden a, Typeable a)+ => Proxy a+ -> SnapshotVersion+ -> [String]+filePathsGolden p v = [+ filePathGolden p annotation v+ | (annotation, _) <- enumGoldenAnnotated' p+ ]++filePathOutput :: Typeable a => Proxy a -> String -> SnapshotVersion -> String+filePathOutput p ann v = nameGolden p ann v ++ ".snapshot"++filePathGolden :: Typeable a => Proxy a -> String -> SnapshotVersion -> String+filePathGolden p ann v = nameGolden p ann v ++ ".snapshot.golden"++{-------------------------------------------------------------------------------+ Enumeration class: instances+-------------------------------------------------------------------------------}++instance EnumGolden SnapshotMetaData where+ singGolden = SnapshotMetaData singGolden singGolden singGolden singGolden singGolden+ where+ _coveredAllCases = \case+ SnapshotMetaData{} -> ()++instance EnumGolden SnapshotLabel where+ enumGolden = [+ SnapshotLabel "UserProvidedLabel"+ , SnapshotLabel ""+ ]+ where+ _coveredAllCases = \case+ SnapshotLabel{} -> ()++instance EnumGolden TableConfig where+ singGolden = TableConfig singGolden singGolden singGolden singGolden+ singGolden singGolden singGolden singGolden+ where+ _coveredAllCases = \case+ TableConfig{} -> ()++instance EnumGolden MergePolicy where+ singGolden = LazyLevelling+ where+ _coveredAllCases = \case+ LazyLevelling{} -> ()+++instance EnumGolden SizeRatio where+ singGolden = Four+ where+ _coveredAllCases = \case+ Four{} -> ()++instance EnumGolden WriteBufferAlloc where+ singGolden = AllocNumEntries magicNumber2+ where+ _coveredAllCases = \case+ AllocNumEntries{} -> ()++instance EnumGolden BloomFilterAlloc where+ enumGolden = [ AllocFixed magicNumber3, AllocRequestFPR pi ]+ where+ _coveredAllCases = \case+ AllocFixed{} -> ()+ AllocRequestFPR{} -> ()++instance EnumGolden FencePointerIndexType where+ enumGolden = [ CompactIndex, OrdinaryIndex ]+ where+ _coveredAllCases = \case+ CompactIndex{} -> ()+ OrdinaryIndex{} -> ()++instance EnumGolden DiskCachePolicy where+ enumGolden = [ DiskCacheAll, DiskCacheLevelOneTo magicNumber3, DiskCacheNone ]+ where+ _coveredAllCases = \case+ DiskCacheAll{} -> ()+ DiskCacheLevelOneTo{} -> ()+ DiskCacheNone{} -> ()++instance EnumGolden MergeSchedule where+ enumGolden = [ OneShot, Incremental ]+ where+ _coveredAllCases = \case+ OneShot{} -> ()+ Incremental{} -> ()++instance EnumGolden MergeBatchSize where+ enumGolden = map MergeBatchSize [ 1, 1000 ]+ supportedVersions _ = [V1]++instance EnumGolden (SnapLevels SnapshotRun) where+ singGolden = SnapLevels singGolden+ where+ _coveredAllCases = \case+ SnapLevels{} -> ()++instance EnumGolden (SnapLevel SnapshotRun) where+ singGolden = SnapLevel singGolden singGolden+ where+ _coveredAllCases = \case+ SnapLevel{} -> ()++instance EnumGolden (SnapIncomingRun SnapshotRun) where+ enumGolden = [+ SnapIncomingMergingRun singGolden singGolden singGolden singGolden+ , SnapIncomingSingleRun singGolden+ ]+ where+ _coveredAllCases = \case+ SnapIncomingMergingRun{} -> ()+ SnapIncomingSingleRun{} -> ()++instance EnumGolden MergePolicyForLevel where+ enumGolden = [ LevelTiering, LevelLevelling ]+ where+ _coveredAllCases = \case+ LevelTiering -> ()+ LevelLevelling -> ()++instance EnumGolden LevelMergeType where+ enumGolden = [ MergeMidLevel, MergeLastLevel ]+ where+ _coveredAllCases = \case+ MergeMidLevel{} -> ()+ MergeLastLevel{} -> ()++instance EnumGolden (SnapMergingTree SnapshotRun) where+ singGolden = SnapMergingTree singGolden+ where+ _coveredAllCases = \case+ SnapMergingTree{} -> ()++instance EnumGolden (SnapMergingTreeState SnapshotRun) where+ enumGolden = [+ SnapCompletedTreeMerge singGolden+ , SnapPendingTreeMerge singGolden+ , SnapOngoingTreeMerge singGolden+ ]+ where+ _coveredAllCases = \case+ SnapCompletedTreeMerge{} -> ()+ SnapPendingTreeMerge{} -> ()+ SnapOngoingTreeMerge{} -> ()++instance EnumGolden (SnapPendingMerge SnapshotRun) where+ enumGolden = [+ SnapPendingLevelMerge singGolden singGolden+ , SnapPendingUnionMerge singGolden+ ]+ where+ _coveredAllCases = \case+ SnapPendingLevelMerge{} -> ()+ SnapPendingUnionMerge{} -> ()++instance EnumGolden (SnapPreExistingRun SnapshotRun) where+ enumGolden = [+ SnapPreExistingRun singGolden+ , SnapPreExistingMergingRun singGolden+ ]+ where+ _coveredAllCases = \case+ SnapPreExistingRun{} -> ()+ SnapPreExistingMergingRun{} -> ()++instance EnumGolden TreeMergeType where+ enumGolden = [ MergeLevel, MergeUnion ]+ where+ _coveredAllCases = \case+ MergeLevel{} -> ()+ MergeUnion{} -> ()++instance EnumGolden a => EnumGolden (Maybe a) where+ enumGolden = [ Just singGolden, Nothing ]+ where+ _coveredAllCases = \case+ Just{} -> ()+ Nothing{} -> ()++instance EnumGolden a => EnumGolden (V.Vector a) where+ enumGolden = [+ V.fromList [ singGolden, singGolden ]+ , mempty+ , V.fromList [ singGolden ]+ ]++instance EnumGolden a => EnumGolden [a] where+ enumGolden = [+ [singGolden, singGolden]+ , []+ , [singGolden]+ ]++instance EnumGolden RunParams where+ singGolden = RunParams singGolden singGolden singGolden+ where+ _coveredAllCases = \case+ RunParams{} -> ()++instance EnumGolden t => EnumGolden (SnapMergingRun t SnapshotRun) where+ enumGolden = [+ SnapCompletedMerge singGolden singGolden+ , SnapOngoingMerge singGolden singGolden singGolden singGolden+ ]+ where+ _coveredAllCases = \case+ SnapCompletedMerge{} -> ()+ SnapOngoingMerge{} -> ()++instance EnumGolden RunBloomFilterAlloc where+ enumGolden = [+ RunAllocFixed magicNumber3+ , RunAllocRequestFPR pi+ ]+ where+ _coveredAllCases = \case+ RunAllocFixed{} -> ()+ RunAllocRequestFPR{} -> ()++instance EnumGolden RunNumber where+ singGolden = RunNumber magicNumber3+ where+ _coveredAllCases = \case+ RunNumber{} -> ()++instance EnumGolden IndexType where+ enumGolden = [+ Compact+ , Ordinary+ ]+ where+ _coveredAllCases = \case+ Compact{} -> ()+ Ordinary{} -> ()++instance EnumGolden RunDataCaching where+ enumGolden = [+ CacheRunData+ , NoCacheRunData+ ]+ where+ _coveredAllCases = \case+ CacheRunData{} -> ()+ NoCacheRunData{} -> ()++instance EnumGolden SnapshotRun where+ singGolden = SnapshotRun singGolden singGolden singGolden+ where+ _coveredAllCases = \case+ SnapshotRun{} -> ()++instance EnumGolden MergeDebt where+ singGolden = MergeDebt magicNumber2+ where+ _coveredAllCases = \case+ MergeDebt{} -> ()++instance EnumGolden MergeCredits where+ singGolden = MergeCredits magicNumber2+ where+ _coveredAllCases = \case+ MergeCredits{} -> ()++instance EnumGolden NominalDebt where+ singGolden = NominalDebt magicNumber2+ where+ _coveredAllCases = \case+ NominalDebt{} -> ()++instance EnumGolden NominalCredits where+ singGolden = NominalCredits magicNumber1+ where+ _coveredAllCases = \case+ NominalCredits{} -> ()++ -- Randomly chosen numbers+magicNumber1, magicNumber2, magicNumber3 :: Enum e => e+magicNumber1 = toEnum 42+magicNumber2 = toEnum 88+magicNumber3 = toEnum 1024
@@ -0,0 +1,230 @@+{-# LANGUAGE OverloadedStrings #-}++-- | Tests for snapshots and their interaction with the file system+module Test.Database.LSMTree.Internal.Snapshot.FS (tests) where++import Control.Monad.Class.MonadThrow+import Control.Monad.IOSim (runSimOrThrow)+import Control.Tracer+import Data.Bifunctor (Bifunctor (..))+import qualified Data.Vector as V+import Data.Word+import Database.LSMTree.Extras (showPowersOf10)+import Database.LSMTree.Extras.Generators ()+import qualified Database.LSMTree.Internal.BloomFilter as Bloom+import Database.LSMTree.Internal.Config+import Database.LSMTree.Internal.Config.Override+ (noTableConfigOverride)+import Database.LSMTree.Internal.Entry+import Database.LSMTree.Internal.Paths+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Snapshot+import Database.LSMTree.Internal.Snapshot.Codec+import Database.LSMTree.Internal.Unsafe+import qualified System.FS.API as FS+import System.FS.API+import System.FS.Sim.Error hiding (genErrors)+import qualified System.FS.Sim.MockFS as MockFS+import Test.Database.LSMTree.Internal.Snapshot.Codec ()+import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck+import Test.Util.FS+import Test.Util.QC (Choice)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Snapshot.FS" [+ testProperty "prop_fsRoundtripSnapshotMetaData" $+ prop_fsRoundtripSnapshotMetaData+ , testProperty "prop_fault_fsRoundtripSnapshotMetaData"+ prop_fault_fsRoundtripSnapshotMetaData+ , testProperty "prop_flipSnapshotBit" prop_flipSnapshotBit+ ]++-- | @readFileSnapshotMetaData . writeFileSnapshotMetaData = id@+--+-- NOTE: prop_fault_fsRoundtripSnapshotMetaData with empty errors is equivalent+-- to prop_fsRoundtripSnapshotMetaData. I (Joris) chose to keep the properties+-- separate, so that there are fewer cases to account for (like @allNull@+-- errors) in prop_fault_fsRoundtripSnapshotMetaData+prop_fsRoundtripSnapshotMetaData :: SnapshotMetaData -> Property+prop_fsRoundtripSnapshotMetaData metadata =+ ioProperty $+ withTempIOHasFS "temp" $ \hfs -> do+ writeFileSnapshotMetaData hfs contentPath checksumPath metadata+ snapshotMetaData' <-+ try @_ @FileCorruptedError (readFileSnapshotMetaData hfs contentPath checksumPath)+ pure $ case snapshotMetaData' of+ Left e -> counterexample (show e) False+ Right metadata' -> metadata === metadata'+ where+ contentPath = mkFsPath ["content"]+ checksumPath = mkFsPath ["checksum"]++-- | @readFileSnapshotMetaData . writeFileSnapshotMetaData = id@, even if+-- exceptions happened.+--+-- NOTE: we can be more precise about the success scenarios for this property,+-- but it complicates the test a lot, so I (Joris) decided not to include it for+-- now. For example, if the read part fails with a deserialise failure, then we+-- *could* check that file corruption took place during the write part.+prop_fault_fsRoundtripSnapshotMetaData ::+ TestErrors+ -> SnapshotMetaData+ -> Property+prop_fault_fsRoundtripSnapshotMetaData testErrs metadata =+ ioProperty $+ withSimErrorHasFS propNoOpenHandles MockFS.empty emptyErrors $ \hfs _fsVar errsVar -> do+ writeResult <-+ try @_ @FsError $+ withErrors errsVar (writeErrors testErrs) $+ writeFileSnapshotMetaData hfs metadataPath checksumPath metadata++ readResult <-+ try @_ @SomeException $+ withErrors errsVar (readErrors testErrs) $+ readFileSnapshotMetaData hfs metadataPath checksumPath++ let+ -- Regardless of whether the write part failed with an exception, if+ -- metadata was returned from read+deserialise it should be exactly+ -- equal to the metadata that was written to disk.+ prop =+ case readResult of+ Right metadata' -> metadata === metadata'+ _ -> property True++ pure $+ -- TODO: there are some scenarios that we never hit, like deserialise+ -- failures. We could tweak the error(stream) generator distribution to+ -- hit these cases more often. One neat idea would be to "prime" the+ -- generator for errors as follows:+ --+ -- 1. run the property without errors, but count how often each+ -- primitive is used+ -- 2. generate errors based on the counts/distribution we obtained in step 1+ -- 3. run the property with these errors+ tabulate "Write result vs. read result" [mkLabel writeResult readResult] $+ prop+ where+ metadataPath = mkFsPath ["metadata"]+ checksumPath = mkFsPath ["checksum"]++ -- This label is mainly there to print the success/failure of the write+ -- part, the read part, and the deserialisation part. The concrete error+ -- contents are not printed.+ mkLabel ::+ Either FsError ()+ -> Either SomeException SnapshotMetaData+ -> String+ mkLabel writeResult readResult =+ "(" <> mkLabelWriteResult writeResult <>+ ", " <> mkLabelReadResult readResult <>+ ")"++ mkLabelWriteResult :: Either FsError () -> String+ mkLabelWriteResult = \case+ Left FsError{} -> "Left FSError"+ Right () -> "Right ()"++ mkLabelReadResult ::+ Either SomeException SnapshotMetaData+ -> String+ mkLabelReadResult = \case+ Left e+ | Just FsError{} <- fromException e ->+ "Left FSError"+ | Just ErrFileFormatInvalid{} <- fromException e ->+ "Left ErrFileFormatInvalid"+ | Just ErrFileChecksumMismatch{} <- fromException e ->+ "Left ErrFileChecksumMismatch"+ | otherwise ->+ error ("impossible: " <> show e)+ Right (_ :: SnapshotMetaData) ->+ "Right SnapshotMetaData"++data TestErrors = TestErrors {+ writeErrors :: Errors+ , readErrors :: Errors+ }+ deriving stock Show++instance Arbitrary TestErrors where+ arbitrary = TestErrors <$> arbitrary <*> arbitrary+ shrink TestErrors{writeErrors, readErrors} =+ [ TestErrors writeErrors' readErrors'+ | (writeErrors', readErrors') <- shrink (writeErrors, readErrors)+ ]++{-------------------------------------------------------------------------------+ Snapshot corruption+-------------------------------------------------------------------------------}++testSalt :: Bloom.Salt+testSalt = 4++-- TODO: an alternative to generating a Choice a priori is to run the monadic+-- code in @PropertyM (IOSim s)@, and then we can do quantification inside the+-- monadic property using @pick@. This complicates matters, however, because+-- functions like @withSimHasBlockIO@ and @withTable@ would have to run in+-- @PropertyM (IOSim s)@ as well. It's not clear whether the refactoring is+-- worth it.+prop_flipSnapshotBit ::+ Positive (Small Int)+ -> V.Vector (Word64, Entry Word64 Word64)+ -> Choice -- ^ Used to pick which file/bit to corrupt.+ -> Property+prop_flipSnapshotBit (Positive (Small bufferSize)) es pickFileBit =+ runSimOrThrow $+ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _fsVar ->+ withOpenSession nullTracer hfs hbio testSalt root $ \s ->+ withTable s conf $ \t -> do+ -- Create a table, populate it, and create a snapshot+ updates resolve es' t+ createSnap t++ -- Corrupt the snapshot+ flipRandomBitInRandomFile hfs pickFileBit (getNamedSnapshotDir namedSnapDir) >>= \case+ Nothing -> pure $ property False+ Just (path, j) -> do+ -- Some info for the test output+ let tabCorruptedFile = tabulate "Corrupted file" [show path]+ counterCorruptedFile = counterexample ("Corrupted file: " ++ show path)+ tabFlippedBit = tabulate "Flipped bit" [showPowersOf10 j]+ counterFlippedBit = counterexample ("Flipped bit: " ++ show j)++ t' <- try @_ @SomeException $ bracket (openSnap s) close $ \_ -> pure ()+ pure $+ tabCorruptedFile $ counterCorruptedFile $ tabFlippedBit $ counterFlippedBit $+ case t' of+ -- If we find an error, we detected corruption. Success!+ Left e ->+ tabulate+ "Result"+ ["Corruption detected: " <> getConstructorName e]+ True+ -- The corruption was not detected. Failure!+ Right _ -> property False+ where+ root = FS.mkFsPath []+ namedSnapDir = namedSnapshotDir (SessionRoot root) snapName++ conf = defaultTableConfig {+ confWriteBufferAlloc = AllocNumEntries bufferSize+ , confFencePointerIndex = CompactIndex+ }+ es' = fmap (bimap serialiseKey (bimap serialiseValue serialiseBlob)) es++ resolve (SerialisedValue x) (SerialisedValue y) =+ SerialisedValue (x <> y)++ snapName = toSnapshotName "snap"+ snapLabel = SnapshotLabel "label"++ createSnap t =+ saveSnapshot snapName snapLabel t++ openSnap s =+ openTableFromSnapshot noTableConfigOverride s snapName snapLabel resolve++ getConstructorName e = takeWhile (/= ' ') (show e)
@@ -0,0 +1,51 @@+module Test.Database.LSMTree.Internal.Unsliced (tests) where++import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.Serialise+import Database.LSMTree.Internal.Unsliced+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Unsliced" [+ testProperty "prop_makeUnslicedKeyPreservesEq" prop_makeUnslicedKeyPreservesEq+ , testProperty "prop_fromUnslicedKeyPreservesEq" prop_fromUnslicedKeyPreservesEq+ , testProperty "prop_makeUnslicedKeyPreservesOrd" prop_makeUnslicedKeyPreservesOrd+ , testProperty "prop_fromUnslicedKeyPreservesOrd" prop_fromUnslicedKeyPreservesOrd+ ]++-- 'Eq' on serialised keys is preserved when converting to /unsliced/ serialised+-- keys.+prop_makeUnslicedKeyPreservesEq :: SerialisedKey -> SerialisedKey -> Property+prop_makeUnslicedKeyPreservesEq k1 k2 = checkCoverage $+ cover 1 lhs "k1 == k2" $ lhs === rhs+ where+ lhs = k1 == k2+ rhs = makeUnslicedKey k1 == makeUnslicedKey k2++-- 'Eq' on /unsliced/ serialised keys is preserved when converting to serialised+-- keys.+prop_fromUnslicedKeyPreservesEq :: Unsliced SerialisedKey -> Unsliced SerialisedKey -> Property+prop_fromUnslicedKeyPreservesEq k1 k2 = checkCoverage $+ cover 1 lhs "k1 == k2" $ lhs === rhs+ where+ lhs = k1 == k2+ rhs = fromUnslicedKey k1 == fromUnslicedKey k2++-- 'Ord' on serialised keys is preserved when converting to /unsliced/+-- serialised keys.+prop_makeUnslicedKeyPreservesOrd :: SerialisedKey -> SerialisedKey -> Property+prop_makeUnslicedKeyPreservesOrd k1 k2 = checkCoverage $+ cover 50 lhs "k1 <= k2" $ lhs === rhs+ where+ lhs = k1 <= k2+ rhs = makeUnslicedKey k1 <= makeUnslicedKey k2++-- 'Ord' on /unsliced/ serialised keys is preserved when converting to serialised+-- keys.+prop_fromUnslicedKeyPreservesOrd :: Unsliced SerialisedKey -> Unsliced SerialisedKey -> Property+prop_fromUnslicedKeyPreservesOrd k1 k2 = checkCoverage $+ cover 50 lhs "k1 <= k2" $ lhs === rhs+ where+ lhs = k1 <= k2+ rhs = fromUnslicedKey k1 <= fromUnslicedKey k2
@@ -0,0 +1,120 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.Internal.Vector (tests) where++import Control.Monad (forM_)+import Control.Monad.ST+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Unboxed.Mutable as VUM+import Data.Word+import Database.LSMTree.Extras+import Database.LSMTree.Internal.Index.CompactAcc+import Database.LSMTree.Internal.Map.Range+import Test.QuickCheck+import Test.QuickCheck.Instances ()+import Test.QuickCheck.Monadic (PropertyM, monadicST, run)+import Test.Tasty (TestTree, localOption, testGroup)+import Test.Tasty.QuickCheck (QuickCheckTests (QuickCheckTests),+ testProperty)+import Test.Util.Orphans ()+import Text.Printf (printf)++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Vector" [+ localOption (QuickCheckTests 400) $+ testProperty "propWriteRange" $ \v lb ub (x :: Word8) -> monadicST $ do+ mv <- run $ VU.thaw v+ propWriteRange mv lb ub x+ , localOption (QuickCheckTests 400) $+ testProperty "propUnsafeWriteRange" $ \v lb ub (x :: Word8) -> monadicST $ do+ mv <- run $ VU.thaw v+ propUnsafeWriteRange mv lb ub x+ ]++instance Arbitrary (Bound Int) where+ arbitrary = oneof [+ pure NoBound+ , BoundInclusive <$> arbitrary+ , BoundExclusive <$> arbitrary+ ]+ shrink = \case+ NoBound -> []+ BoundInclusive x -> NoBound : (BoundInclusive <$> shrink x)+ BoundExclusive x -> NoBound : (BoundInclusive <$> shrink x)+ ++ (BoundExclusive <$> shrink x)++intToInclusiveLowerBound :: Bound Int -> Int+intToInclusiveLowerBound = \case+ NoBound -> 0+ BoundInclusive i -> i+ BoundExclusive i -> i + 1++intToInclusiveUpperBound :: VUM.Unbox a => VU.Vector a -> Bound Int -> Int+intToInclusiveUpperBound xs = \case+ NoBound -> VU.length xs - 1+ BoundInclusive i -> i+ BoundExclusive i -> i - 1++-- | Safe version of 'unsafeWriteRange', used to test the unsafe version+-- against.+writeRange :: VU.Unbox a => VU.MVector s a -> Bound Int -> Bound Int -> a -> ST s Bool+writeRange !v !lb !ub !x+ | 0 <= lb' && lb' < VUM.length v+ , 0 <= ub' && ub' < VUM.length v+ , lb' <= ub'+ = forM_ [lb' .. ub'] (\j -> VUM.write v j x) >> pure True+ | otherwise = pure False+ where+ !lb' = vectorLowerBound lb+ !ub' = mvectorUpperBound v ub++propWriteRange :: forall s a. (VUM.Unbox a, Eq a, Show a)+ => VU.MVector s a+ -> Bound Int+ -> Bound Int+ -> a+ -> PropertyM (ST s) Property+propWriteRange mv1 lb ub x = run $ do+ v1 <- VU.unsafeFreeze mv1+ v2 <- VU.freeze mv1+ b <- writeRange mv1 lb ub x++ let xs1 = zip [0 :: Int ..] $ VU.toList v1+ xs2 = zip [0..] $ VU.toList v2+ lb' = intToInclusiveLowerBound lb+ ub' = intToInclusiveUpperBound v1 ub++ pure $ tabulate "range size" [showPowersOf10 (ub' - lb' + 1)] $+ tabulate "vector size" [showPowersOf10 (VU.length v1)] $+ if not b then+ label "no suitable range" $ xs1 === xs2+ else+ counterexample (printf "lb=%d" lb') $+ counterexample (printf "ub=%d" ub') $+ conjoin [+ counterexample "mismatch in prefix" $+ take (lb' - 1) xs1 === take (lb' - 1) xs2+ , counterexample "mismatch in suffix" $+ drop (ub' + 1) xs1 === drop (ub' + 1) xs2+ , counterexample "mimsatch in infix" $+ fmap snd (drop lb' (take (ub' + 1) xs1)) ===+ replicate (ub' - lb' + 1) x+ ]++propUnsafeWriteRange ::+ forall s a. (VUM.Unbox a, Eq a, Show a)+ => VU.MVector s a+ -> Bound Int+ -> Bound Int+ -> a+ -> PropertyM (ST s) Property+propUnsafeWriteRange mv1 lb ub x = run $ do+ v1 <- VU.unsafeFreeze mv1+ v2 <- VU.freeze mv1+ mv2 <- VU.unsafeThaw v2+ b <- writeRange mv1 lb ub x+ if not b then+ pure $ label "no suitable range" True+ else do+ unsafeWriteRange mv2 lb ub x+ pure $ v1 === v2
@@ -0,0 +1,222 @@+module Test.Database.LSMTree.Internal.Vector.Growing (tests) where++import Prelude hiding (head, length, tail)++import Control.Category ((>>>))+import Control.Monad.ST.Strict (ST, runST)+import qualified Data.List as List (length)+import Data.Vector (Vector, toList)+import Database.LSMTree.Internal.Vector.Growing (GrowingVector,+ append, freeze, new)+import Test.QuickCheck (Arbitrary (arbitrary, shrink), Gen,+ NonNegative (NonNegative, getNonNegative),+ Positive (Positive), Property, Small (Small), Testable,+ chooseInt, coverTable, shrinkList, shrinkMap, shrinkMapBy,+ sized, tabulate, (===))+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)++-- * Tests++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.Vector.Growing" $+ [+ testProperty "Final vector is correct"+ prop_finalVectorIsCorrect+ ]++-- * Utilities++-- ** Segments++data Segment a = Replicate Int a deriving stock Show++segmentLength :: Segment a -> Int+segmentLength (Replicate count _) = max 0 count++segmentToList :: Segment a -> [a]+segmentToList (Replicate count val) = replicate count val++appendSegment :: GrowingVector s a -> Segment a -> ST s ()+appendSegment vec (Replicate count val) = append vec count val++-- ** Vector construction++{-|+ Constructs an ordinary vector by creating a growing vector, appending+ segments to it, and finally freezing it.+-}+finalVector :: Int -- Initial buffer size+ -> [Segment a] -- Segments+ -> Vector a -- Final vector+finalVector initialBufferSize segments = runST $ do+ vec <- new initialBufferSize+ mapM_ (appendSegment vec) segments+ freeze vec++{-|+ Supplies the final contents of a growing vector for constructing a property.++ The resulting property additionally provides information about the+ occurrence of buffer size scaling exponents, where a buffer size scaling+ exponent is the binary logarithm of the ratio between the buffer size after+ and before appending a segment. If buffer enlargement for an append would+ happen in cycles where in each cycle the buffer size is doubled, then the+ buffer size scaling exponent would be equal to the number of such cycles.+ Therefore, a buffer size scaling exponent tells if the buffer is enlarged+ and, if yes, how much.+-}+withFinalVector ::+ Testable prop+ => Int -- Initial buffer size+ -> [Segment a] -- Segments+ -> (Vector a -> prop) -- Property from final vector+ -> Property -- Resulting property+withFinalVector initialBufferSize segments consumer+ = coverTable+ "Buffer size scaling exponents"+ [(show scalingExp, 10) | scalingExp <- [0 .. 3] :: [Int]]+ $+ tabulate+ "Buffer size scaling exponents"+ (show <$> bufferSizeScalingExponents availableBufferSizes 0 segments)+ $+ consumer (finalVector initialBufferSize segments)+ where++ availableBufferSizes :: [Int]+ availableBufferSizes = iterate (* 2) initialBufferSize+ {-+ We do not need to account for overflow here, because the lengths of the+ vectors we construct are small compared to @maxBound :: Int@.+ -}++ bufferSizeScalingExponents :: [Int] -> Int -> [Segment a] -> [Int]+ bufferSizeScalingExponents _ _ []+ = []+ bufferSizeScalingExponents availableSizes length (head : tail)+ = List.length insufficient :+ bufferSizeScalingExponents sufficient length' tail+ where++ length' :: Int+ !length' = length + segmentLength head++ insufficient, sufficient :: [Int]+ (insufficient, sufficient) = span (< length') availableSizes++-- * Properties to test++prop_finalVectorIsCorrect :: Positive (Small Int)+ -> Segments Integer+ -> Property+prop_finalVectorIsCorrect (Positive (Small initialBufferSize))+ (Segments segments)+ = withFinalVector initialBufferSize segments $ \ vec ->+ toList vec === concatMap segmentToList segments++-- * Test case generation and shrinking++generateSegment :: Arbitrary a => Int -> Gen (Segment a)+generateSegment maxCount = Replicate <$> chooseInt (0, maxCount) <*> arbitrary++shrinkSegment :: Arbitrary a => Segment a -> [Segment a]+shrinkSegment (Replicate count val)+ = [Replicate count' val | count' <- shrinkNat count] +++ [Replicate count val' | val' <- shrink val]+ where++ shrinkNat :: Int -> [Int]+ shrinkNat = shrinkMap getNonNegative NonNegative++{-|+ A list of segments to be appended to an initially empty vector. 'Segments a'+ is isomorphic to '[Segment a]' but has a special way of generating test+ cases, which avoids skewing with respect to buffer size scaling exponents.++ The key issue of segment list generation is how to generate the length of an+ individual segment. It seems worthwhile to randomly pick a natural number+ smaller than or equal to some maximum length, using a uniform distribution.+ The question to be answered is how to choose the maximum lengths for the+ different segments.++ A naïve approach is to use a common maximum length for all segments in a+ particular list, possibly computed from the size parameter. However, this+ results in most appends not enlarging the buffer, meaning that the buffer+ size scaling exponent is zero in most cases. This is because, as the buffer+ size increases, the number of elements needed to cause the next buffer+ enlargement increases too.++ Note that it does not help much to choose the maximum segment length large+ compared to the initial buffer size. If this is done, then the buffer is+ enlarged very quickly, likely during the first append, and the distributions+ of buffer size scaling exponents up to this first enlargement may actually+ be very good. However, after the first enlargement, the buffer size is+ typically of the same order of magnitude as the maximum segment length, so+ that the buffer size scaling exponents of the following appends are likely+ to be zero or close to it, with further buffer enlargements making the+ situation only worse.++ A better idea is to choose the maximum length of each segment but the first+ proportional to the length of the vector just before appending it (for the+ first segment, such a choice is not appropriate, because it would lead to a+ maximum length of zero). With this approach, the distribution of buffer size+ scaling exponents stays roughly the same as soon as the buffer has been+ enlarged for the first time, because then the ratio between the current+ buffer size and the current vector length is always in the interval \([1,+ 2)\) and thus varies only slightly. Reasonable distributions of buffer size+ scaling exponents can be achieved by choosing the ratio between maximum+ segment lengths and corresponding current-vector lengths appropriately.++ A downside of this approach is that it requires tracking the current vector+ length. This tracking can be avoided by considering only the /expected/+ current vector length and choosing the maximum segment length proportional+ to it. The expected length of a vector is the sum of the expected lengths of+ the segments constituting it, and the expected length of a segment is+ proportional to the maximum length used for generating it. Therefore, this+ modified solution boils down to choosing the maximum length of each segment+ but the first proportional to the sum of the maximum lengths of the segments+ preceding it.++ Note that, with the approach just set out, the maximum lengths of the+ segments in a segment list almost form an exponential sequence, whose base+ is determined by the chosen ratio between maximum segment lengths and+ corresponding sums of maximum predecessor segment lengths. Therefore, a+ similarly good, but simpler, solution is to have the maximum segment lengths+ precisely form an exponential sequence. This is the approach that we use in+ our segment list generation algorithm.++ The concrete exponential sequences that we employ start with the size+ parameter and use 3 as the base of the exponentiation. Since initial buffer+ sizes are generated using a uniform distribution with the size parameter as+ the maximum, the choice of the size parameter as the first segment’s maximum+ length leads to non-trivial distributions of buffer size scaling exponents+ for the appends up to the one that causes the first buffer enlargement. The+ choice of 3 as the exponentiation’s base leads to a good overall+ distribution of buffer size scaling exponents, as experiments have shown.++ Because of the exponential growth of maximum segment lengths, vectors+ constructed during testing get very large for longer segment lists.+ Therefore, we do not pick the lengths of segment lists in the usual way but+ instead make all segment lists have a common, small length. Concretely, we+ use 9 as this common length, because this leads to vectors that are not+ trivial but still consume reasonable amounts of memory.+-}+newtype Segments a = Segments [Segment a] deriving stock Show++instance Arbitrary a => Arbitrary (Segments a) where++ arbitrary = sized $ iterate (* scalingFactor) >>>+ take count >>>+ mapM generateSegment >>>+ fmap Segments+ where++ scalingFactor :: Int+ scalingFactor = 3++ count :: Int+ count = 9++ shrink = shrinkMapBy Segments (\ (Segments segments) -> segments) $+ shrinkList shrinkSegment
@@ -0,0 +1,87 @@+module Test.Database.LSMTree.Internal.WriteBufferBlobs.FS (tests) where++import Control.Concurrent.Class.MonadSTM.Strict+import Control.Monad+import Control.Monad.Class.MonadThrow+import Control.RefCount+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.BlobRef (readRawBlobRef)+import Database.LSMTree.Internal.Serialise (SerialisedBlob)+import Database.LSMTree.Internal.WriteBufferBlobs+import System.FS.API+import System.FS.Sim.Error hiding (genErrors)+import qualified System.FS.Sim.MockFS as MockFS+import Test.Tasty+import Test.Tasty.QuickCheck as QC+import Test.Util.FS++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.WriteBufferBlobs.FS" [+ testProperty "prop_fault_WriteBufferBlobs" prop_fault_WriteBufferBlobs+ ]++-- Test that opening and releasing a 'WriteBufferBlobs' properly cleans handles+-- and files in the presence of disk faults. Also test that we can write then+-- read blobs correctly in the presence of disk faults.+--+-- By testing 'open', we also test 'new'.+prop_fault_WriteBufferBlobs ::+ Bool -- ^ create the file or not+ -> AllowExisting+ -> NoCleanupErrors+ -> Errors+ -> NoCleanupErrors+ -> SerialisedBlob+ -> SerialisedBlob+ -> Property+prop_fault_WriteBufferBlobs doCreateFile ae+ (NoCleanupErrors openErrors)+ errs+ (NoCleanupErrors releaseErrors)+ b1 b2 =+ ioProperty $+ withSimErrorHasFS propPost MockFS.empty emptyErrors $ \hfs fsVar errsVar -> do+ when doCreateFile $+ withFile hfs path (WriteMode MustBeNew) $ \_ -> pure ()+ eith <- try @_ @FsError $+ bracket (acquire hfs errsVar) (release errsVar) $ \wbb -> do+ fs' <- atomically $ readTMVar fsVar+ let prop = propNumOpenHandles 1 fs' .&&. propNumDirEntries root 1 fs'+ props <- blobRoundtrips hfs errsVar wbb+ pure (prop .&&. props)+ pure $ case eith of+ Left{} -> do+ label "FsError" $ property True+ Right prop ->+ label "Success" $ prop+ where+ root = mkFsPath []+ path = mkFsPath ["wbb"]++ acquire hfs errsVar = withErrors errsVar openErrors $ open hfs path ae++ -- Test that we can roundtrip blobs+ blobRoundtrips hfs errsVar wbb = withErrors errsVar errs $ do+ props <-+ forM [b1, b2] $ \b -> do+ bspan <- addBlob hfs wbb b+ let bref = mkRawBlobRef wbb bspan+ b' <- readRawBlobRef hfs bref+ pure (b === b')+ pure $ conjoin props++ release errsVar wbb = withErrors errsVar releaseErrors $ releaseRef wbb++ propPost fs = propNoOpenHandles fs .&&.+ if doCreateFile then+ case ae of+ AllowExisting ->+ -- TODO: fix, see the TODO on openBlobFile+ propNoDirEntries root fs .||. propNumDirEntries root 1 fs+ MustBeNew ->+ propNumDirEntries root 1 fs+ MustExist ->+ -- TODO: fix, see the TODO on openBlobFile+ propNoDirEntries root fs .||. propNumDirEntries root 1 fs+ else+ propNoDirEntries root fs
@@ -0,0 +1,76 @@+module Test.Database.LSMTree.Internal.WriteBufferReader.FS (tests) where+++import Control.Concurrent.Class.MonadSTM.Strict (MonadSTM (..))+import Control.Concurrent.Class.MonadSTM.Strict.TMVar+import Control.Monad.Class.MonadThrow+import Control.RefCount+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Extras.RunData (RunData,+ withRunDataAsWriteBuffer, withSerialisedWriteBuffer)+import Database.LSMTree.Internal.Paths (ForKOps (ForKOps),+ WriteBufferFsPaths (WriteBufferFsPaths),+ writeBufferKOpsPath)+import Database.LSMTree.Internal.RunNumber (RunNumber (RunNumber))+import Database.LSMTree.Internal.Serialise (SerialisedBlob,+ SerialisedKey, SerialisedValue (..))+import qualified Database.LSMTree.Internal.WriteBufferBlobs as WBB+import Database.LSMTree.Internal.WriteBufferReader+import System.FS.API+import System.FS.Sim.Error hiding (genErrors)+import qualified System.FS.Sim.MockFS as MockFS+import qualified System.FS.Sim.Stream as Stream+import Test.Tasty+import Test.Tasty.QuickCheck as QC+import Test.Util.FS as FS++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Internal.WriteBufferReader.FS" [+ testProperty "prop_fault_WriteBufferReader" prop_fault_WriteBufferReader+ ]++-- | Test that 'writeWriteBuffer' roundtrips with 'readWriteBuffer', and test+-- that the presence of disk faults for the latter does not leak file handles+-- and files.+prop_fault_WriteBufferReader ::+ NoCleanupErrors+ -> RunData SerialisedKey SerialisedValue SerialisedBlob+ -> Property+prop_fault_WriteBufferReader (NoCleanupErrors readErrors) rdata =+ ioProperty $+ withSimErrorHasBlockIO propPost MockFS.empty emptyErrors $ \hfs hbio fsVar errsVar ->+ withRunDataAsWriteBuffer hfs resolve inPath rdata $ \wb wbb ->+ withSerialisedWriteBuffer hfs hbio outPath wb wbb $ do+ fsBefore <- atomically $ readTMVar fsVar+ eith <-+ try @_ @FsError $+ withErrors errsVar readErrors' $+ withRef wbb $ \wbb' -> do+ wb' <- readWriteBuffer resolve hfs hbio outKOpsPath (WBB.blobFile wbb')+ pure (wb === wb')++ fsAfter <- atomically $ readTMVar fsVar+ pure $+ case eith of+ Left{} -> do+ label "FsError" $ property True+ Right prop ->+ label "Success" $ prop .&&. propEqNumDirEntries root fsBefore fsAfter+ where+ root = mkFsPath []+ -- The run number for the original write buffer. Primarily used to name the+ -- 'WriteBufferBlobs' corresponding to the write buffer.+ inPath = WriteBufferFsPaths root (RunNumber 0)+ -- The run number for the serialised write buffer. Used to name all files+ -- that are the result of serialising the write buffer.+ outPath = WriteBufferFsPaths root (RunNumber 1)+ outKOpsPath = ForKOps (writeBufferKOpsPath outPath)+ resolve (SerialisedValue x) (SerialisedValue y) = SerialisedValue (x <> y)+ propPost fs = propNoOpenHandles fs .&&. propNoDirEntries root fs++ -- TODO: fix, see the TODO on readDiskPage+ readErrors' = readErrors {+ hGetBufSomeE = Stream.filter (not . isFsReachedEOFError) (hGetBufSomeE readErrors)+ }++ isFsReachedEOFError = maybe False (either isFsReachedEOF (const False))
@@ -0,0 +1,91 @@+{-# OPTIONS_GHC -Wno-orphans #-}+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}++module Test.Database.LSMTree.Model.Table (tests) where++import qualified Data.ByteString as BS+import qualified Data.Vector as V+import Database.LSMTree (ResolveValue (..), ResolveViaSemigroup (..),+ SerialiseKey (..), SerialiseValue (..))+import Database.LSMTree.Model.Table (LookupResult (..), Table,+ Update (..), lookups)+import qualified Database.LSMTree.Model.Table as Model+import GHC.Exts (IsList (..))+import Test.QuickCheck.Instances ()+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Database.LSMTree.Model.Table"+ [ testProperty "lookup-insert" prop_lookupInsert+ , testProperty "lookup-delete" prop_lookupDelete+ , testProperty "insert-insert" prop_insertInsert+ , testProperty "upsert-insert" prop_upsertInsert+ , testProperty "upsert=lookup+insert" prop_upsertDef+ , testProperty "insert-commutes" prop_insertCommutes+ ]++type Key = BS.ByteString++newtype Value = Value BS.ByteString+ deriving stock (Eq, Show)+ deriving newtype (Arbitrary, Semigroup, SerialiseValue)+ deriving ResolveValue via (ResolveViaSemigroup Value)++type Blob = BS.ByteString++type Tbl = Table Key Value Blob++updates :: V.Vector (Key, Update Value Blob) -> Table Key Value Blob -> Table Key Value Blob+updates = Model.updates Model.getResolve++inserts :: V.Vector (Key, Value, Maybe Blob) -> Table Key Value Blob -> Table Key Value Blob+inserts = Model.inserts Model.getResolve++deletes :: V.Vector Key -> Table Key Value Blob -> Table Key Value Blob+deletes = Model.deletes Model.getResolve++mupserts :: V.Vector (Key, Value) -> Table Key Value Blob -> Table Key Value Blob+mupserts = Model.mupserts Model.getResolve++-- | You can lookup what you inserted.+prop_lookupInsert :: Key -> Value -> Tbl -> Property+prop_lookupInsert k v tbl =+ lookups (V.singleton k) (inserts (V.singleton (k, v, Nothing)) tbl) === V.singleton (Found v)++-- | You cannot lookup what you have deleted+prop_lookupDelete :: Key -> Tbl -> Property+prop_lookupDelete k tbl =+ lookups (V.singleton k) (deletes (V.singleton k) tbl) === V.singleton NotFound++-- | Last insert wins.+prop_insertInsert :: Key -> Value -> Value -> Tbl -> Property+prop_insertInsert k v1 v2 tbl =+ inserts (V.fromList [(k, v1, Nothing), (k, v2, Nothing)]) tbl === inserts (V.singleton (k, v2, Nothing)) tbl++-- | Updating after insert is the same as inserting merged value.+--+-- Note: the order of merge.+prop_upsertInsert :: Key -> Value -> Value -> Tbl -> Property+prop_upsertInsert k v1 v2 tbl =+ updates (V.fromList [(k, Insert v1 Nothing), (k, Upsert v2)]) tbl+ === inserts (V.singleton (k, resolve v2 v1, Nothing)) tbl++-- | Upsert is the same as lookup followed by an insert.+prop_upsertDef :: Key -> Value -> Tbl -> Property+prop_upsertDef k v tbl =+ tbl' === mupserts (V.singleton (k, v)) tbl+ where+ tbl' = case toList (lookups (V.singleton k) tbl) of+ [Found v'] -> inserts (V.singleton (k, resolve v v', Nothing)) tbl+ [FoundWithBlob v' _] -> inserts (V.singleton (k, resolve v v', Nothing)) tbl+ _ -> inserts (V.singleton (k, v, Nothing)) tbl++-- | Different key inserts commute.+prop_insertCommutes :: Key -> Value -> Key -> Value -> Tbl -> Property+prop_insertCommutes k1 v1 k2 v2 tbl = k1 /= k2 ==>+ inserts (V.fromList [(k1, v1, Nothing), (k2, v2, Nothing)]) tbl === inserts (V.fromList [(k2, v2, Nothing), (k1, v1, Nothing)]) tbl++instance (SerialiseKey k, SerialiseValue v, SerialiseValue b, Arbitrary k, Arbitrary v, Arbitrary b) => Arbitrary (Table k v b) where+ arbitrary = fromList <$> arbitrary+ shrink t = fromList <$> shrink (toList t)
@@ -0,0 +1,40 @@+{-# LANGUAGE AllowAmbiguousTypes #-}++module Test.Database.LSMTree.Resolve (tests) where++import Control.DeepSeq (NFData)+import Data.Monoid (Sum (..))+import Data.Word+import Database.LSMTree+import Database.LSMTree.Extras.Generators ()+import Test.Tasty+import Test.Tasty.QuickCheck++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.Resolve"+ [ testGroup "Sum Word64" (allProperties @(Sum Word64))+ ]++allProperties ::+ forall v. (Show v, Arbitrary v, NFData v, SerialiseValue v, ResolveValue v)+ => [TestTree]+allProperties =+ [ testProperty "prop_resolveValidOutput" $ withMaxSuccess 1000 $+ prop_resolveValidOutput @v+ , testProperty "prop_resolveAssociativity" $ withMaxSuccess 1000 $+ prop_resolveAssociativity @v+ ]++prop_resolveValidOutput ::+ forall v. (Show v, NFData v, SerialiseValue v, ResolveValue v)+ => v -> v -> Property+prop_resolveValidOutput x y =+ counterexample ("inputs: " <> show (x, y)) $+ resolveValidOutput x y++prop_resolveAssociativity ::+ forall v. (Show v, SerialiseValue v, ResolveValue v)+ => v -> v -> v -> Property+prop_resolveAssociativity x y z =+ counterexample ("inputs: " <> show (x, y)) $+ resolveAssociativity x y z
@@ -0,0 +1,2937 @@+{-# LANGUAGE CPP #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE QuantifiedConstraints #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++#if MIN_VERSION_GLASGOW_HASKELL(9,8,1,0)+{-# LANGUAGE TypeAbstractions #-}+#endif++{-# OPTIONS_GHC -Wno-orphans #-}++{-+ TODO: improve generation and shrinking of dependencies. See+ https://github.com/IntersectMBO/lsm-tree/pull/4#discussion_r1334295154.++ TODO: The 'RetrieveBlobs' action will currently only retrieve blob references+ that come from a single batch lookup or range lookup. It would be interesting+ to also retrieve blob references from a mix of different batch lookups and/or+ range lookups. This would require some non-trivial changes, such as changes to+ 'Op' to also include expressions for manipulating lists, such that we can map+ @'Var' ['R.BlobRef' b]@ to @'Var' ('R.BlobRef' b)@. 'RetrieveBlobs'+ would then hold a list of variables (e.g., @['Var' ('R.BlobRef b')]@)+ instead of a variable of a list (@'Var' ['R.BlobRef' b]@).++ TODO: it is currently not correctly modelled what happens if blob references+ are retrieved from an incorrect table.+-}+module Test.Database.LSMTree.StateMachine (+ tests+ , labelledExamples+ -- * Properties+ , propLockstep_ModelIOImpl+ , propLockstep_RealImpl_RealFS_IO+ , propLockstep_RealImpl_MockFS_IO+ , propLockstep_RealImpl_MockFS_IOSim+ , CheckCleanup (..)+ , CheckFS (..)+ , CheckRefs (..)+ -- * Lockstep+ , ModelState (..)+ , Key (..)+ , Value (..)+ , Blob (..)+ , StateModel (..)+ , Action (..)+ , Action' (..)+ ) where++import Control.ActionRegistry (AbortActionRegistryError (..),+ CommitActionRegistryError (..), getActionError,+ getReasonExitCaseException)+import Control.Concurrent.Class.MonadMVar.Strict+import Control.Concurrent.Class.MonadSTM.Strict+import Control.Exception (assert)+import Control.Monad (forM_, (<=<))+import Control.Monad.Class.MonadST (MonadST)+import Control.Monad.Class.MonadThrow (Exception (..), Handler (..),+ MonadCatch (..), MonadMask, MonadThrow (..), SomeException,+ catches, displayException, fromException)+import Control.Monad.IOSim+import Control.Monad.Primitive+import Control.Monad.Reader (ReaderT (..))+import Control.RefCount (RefException, checkForgottenRefs,+ ignoreForgottenRefs)+import Control.Tracer (Tracer, nullTracer)+import Data.Bifunctor (Bifunctor (..))+import Data.Constraint (Dict (..))+import Data.Either (partitionEithers)+import Data.Kind (Type)+import qualified Data.List as List+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (catMaybes, fromMaybe)+import Data.Monoid (First (..))+import Data.Primitive.MutVar+import Data.Set (Set)+import qualified Data.Set as Set+import Data.Typeable (Proxy (..), Typeable, cast)+import qualified Data.Vector as V+import Database.LSMTree (BlobRefInvalidError (..),+ CursorClosedError (..), SessionClosedError (..),+ SessionDirCorruptedError (..),+ SessionDirDoesNotExistError (..),+ SessionDirLockedError (..), SnapshotCorruptedError (..),+ SnapshotDoesNotExistError (..), SnapshotExistsError (..),+ SnapshotNotCompatibleError (..), TableClosedError (..),+ TableCorruptedError (..),+ TableUnionNotCompatibleError (..))+import qualified Database.LSMTree as R+import Database.LSMTree.Class (Entry (..), LookupResult (..))+import qualified Database.LSMTree.Class as Class+import Database.LSMTree.Extras (showPowersOf)+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Extras.NoThunks (propNoThunks)+import qualified Database.LSMTree.Internal.Config as R (TableConfig (..))+import Database.LSMTree.Internal.Serialise (SerialisedBlob,+ SerialisedKey, SerialisedValue)+import qualified Database.LSMTree.Internal.Types as R.Types+import qualified Database.LSMTree.Internal.Unsafe as R.Unsafe+import qualified Database.LSMTree.Model.IO as ModelIO+import qualified Database.LSMTree.Model.Session as Model+import NoThunks.Class+import Prelude hiding (init)+import System.Directory (removeDirectoryRecursive)+import System.FS.API (FsError (..), HasFS, MountPoint (..), mkFsPath)+import System.FS.BlockIO.API (HasBlockIO, close)+import System.FS.BlockIO.IO (defaultIOCtxParams, ioHasBlockIO)+import System.FS.IO (HandleIO)+import qualified System.FS.Sim.Error as FSSim+import System.FS.Sim.Error (Errors)+import qualified System.FS.Sim.MockFS as MockFS+import System.FS.Sim.MockFS (MockFS)+import System.FS.Sim.Stream (Stream)+import System.IO.Temp (createTempDirectory,+ getCanonicalTemporaryDirectory)+import Test.Database.LSMTree.StateMachine.Op (HasBlobRef (getBlobRef),+ Op (..))+import qualified Test.QuickCheck as QC+import Test.QuickCheck (Arbitrary, Gen, Property)+import Test.QuickCheck.StateModel hiding (Var)+import Test.QuickCheck.StateModel.Lockstep+import qualified Test.QuickCheck.StateModel.Lockstep.Defaults as Lockstep.Defaults+import qualified Test.QuickCheck.StateModel.Lockstep.Run as Lockstep.Run+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (testProperty)+import Test.Util.FS+import Test.Util.FS.Error+import Test.Util.PrettyProxy+import Test.Util.QC (Choice)+import qualified Test.Util.QLS as QLS+import Test.Util.TypeFamilyWrappers (WrapBlob (..), WrapBlobRef (..),+ WrapCursor (..), WrapTable (..))++{-------------------------------------------------------------------------------+ Test tree+-------------------------------------------------------------------------------}++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.StateMachine" [+ testProperty "propLockstep_ModelIOImpl"+ propLockstep_ModelIOImpl++ , testProperty "propLockstep_RealImpl_RealFS_IO" $+ propLockstep_RealImpl_RealFS_IO nullTracer++ , testProperty "propLockstep_RealImpl_MockFS_IO" $+ propLockstep_RealImpl_MockFS_IO nullTracer CheckCleanup CheckFS CheckRefs++ , testProperty "propLockstep_RealImpl_MockFS_IOSim" $+ propLockstep_RealImpl_MockFS_IOSim nullTracer CheckCleanup CheckFS CheckRefs+ ]++labelledExamples :: IO ()+labelledExamples = QC.labelledExamples $ Lockstep.Run.tagActions (Proxy @(ModelState IO R.Table))++{-------------------------------------------------------------------------------+ propLockstep: reference implementation+-------------------------------------------------------------------------------}++instance Arbitrary Model.TableConfig where+ arbitrary :: Gen Model.TableConfig+ arbitrary = pure Model.TableConfig++deriving via AllowThunk (ModelIO.Session IO)+ instance NoThunks (ModelIO.Session IO)++propLockstep_ModelIOImpl ::+ Actions (Lockstep (ModelState IO ModelIO.Table))+ -> QC.Property+propLockstep_ModelIOImpl =+ runActionsBracket+ (Proxy @(ModelState IO ModelIO.Table))+ CheckCleanup+ NoCheckRefs -- there are no references to check for in the ModelIO implementation+ acquire+ release+ (\r (session, errsVar, logVar) -> do+ faultsVar <- newMutVar []+ let+ env :: RealEnv ModelIO.Table IO+ env = RealEnv {+ envSession = session+ , envHandlers = [handler]+ , envErrors = errsVar+ , envErrorsLog = logVar+ , envInjectFaultResults = faultsVar+ }+ prop <- runReaderT r env+ faults <- readMutVar faultsVar+ pure $ QC.tabulate "Fault results" (fmap show faults) prop+ )+ tagFinalState'+ where+ acquire :: IO (Class.Session ModelIO.Table IO, StrictTVar IO Errors, StrictTVar IO ErrorsLog)+ acquire = do+ session <- Class.openSession ModelIO.NoSessionArgs+ errsVar <- newTVarIO FSSim.emptyErrors+ logVar <- newTVarIO emptyLog+ pure (session, errsVar, logVar)++ release :: (Class.Session ModelIO.Table IO, StrictTVar IO Errors, StrictTVar IO ErrorsLog) -> IO ()+ release (session, _, _) = Class.closeSession session++ handler :: Handler IO (Maybe Model.Err)+ handler = Handler $ pure . handler'+ where+ handler' :: ModelIO.Err -> Maybe Model.Err+ handler' (ModelIO.Err err) = Just err++{-------------------------------------------------------------------------------+ propLockstep: real implementation+-------------------------------------------------------------------------------}++instance Arbitrary R.TableConfig where+ arbitrary = do+ confMergeSchedule <- QC.frequency [+ (1, pure R.OneShot)+ , (4, pure R.Incremental)+ ]+ confWriteBufferAlloc <- QC.arbitrary+ confFencePointerIndex <- QC.arbitrary+ confMergeBatchSize <- QC.sized $ \sz ->+ R.MergeBatchSize <$> QC.chooseInt (1, sz)+ pure $ R.TableConfig {+ R.confMergePolicy = R.LazyLevelling+ , R.confSizeRatio = R.Four+ , confWriteBufferAlloc+ , R.confBloomFilterAlloc = R.AllocFixed 10+ , confFencePointerIndex+ , R.confDiskCachePolicy = R.DiskCacheNone+ , confMergeSchedule+ , confMergeBatchSize+ }++ shrink R.TableConfig{..} =+ [ R.TableConfig {+ confWriteBufferAlloc = confWriteBufferAlloc'+ , confFencePointerIndex = confFencePointerIndex'+ , ..+ }+ | ( confWriteBufferAlloc', confFencePointerIndex')+ <- QC.shrink (confWriteBufferAlloc, confFencePointerIndex)+ ]++-- TODO: the current generator is suboptimal, and should be improved. There are+-- some aspects to consider, and since they are at tension with each other, we+-- should try find a good balance.+--+-- Say the write buffer allocation is @n@.+--+-- * If @n@ is too large, then the tests degrade to only testing the write+-- buffer, because there are no flushes/merges.+--+-- * If @n@ is too small, then the resulting runs are also small. They might+-- consist of only a few disk pages, or they might consist of only 1 underfull+-- disk page. This means there are some parts of the code that we would rarely+-- cover in the state machine tests.+--+-- * If @n@ is too small, then we flush/merge very frequently, which may exhaust+-- the number of open file handles when using the real file system. Specially+-- if the test fails and the shrinker kicks in, then @n@ can currently become+-- very small. This is good for finding a minimal counterexample, but it might+-- also make the real filesystem run out of file descriptors. Arguably, you+-- could overcome this specific issue by only generating or shrinking to small+-- @n@ when we use the mocked file system. This would require some boilerplate+-- to add type level tags to distinguish between the two cases.+instance Arbitrary R.WriteBufferAlloc where+ arbitrary = QC.scale (max 30) $ do+ QC.Positive x <- QC.arbitrary+ pure (R.AllocNumEntries x)++ shrink (R.AllocNumEntries x) =+ [ R.AllocNumEntries x'+ | QC.Positive x' <- QC.shrink (QC.Positive x)+ ]++deriving stock instance Enum R.FencePointerIndexType+deriving stock instance Bounded R.FencePointerIndexType+instance Arbitrary R.FencePointerIndexType where+ arbitrary = QC.arbitraryBoundedEnum+ shrink R.OrdinaryIndex = []+ shrink R.CompactIndex = [R.OrdinaryIndex]++propLockstep_RealImpl_RealFS_IO ::+ Tracer IO R.LSMTreeTrace+ -> QC.Fixed R.Salt+ -> Actions (Lockstep (ModelState IO R.Table))+ -> QC.Property+propLockstep_RealImpl_RealFS_IO tr (QC.Fixed salt) =+ runActionsBracket+ (Proxy @(ModelState IO R.Table))+ CheckCleanup+ CheckRefs+ acquire+ release+ (\r (_, session, _, errsVar, logVar) -> do+ faultsVar <- newMutVar []+ let+ env :: RealEnv R.Table IO+ env = RealEnv {+ envSession = session+ , envHandlers = realErrorHandlers @IO+ , envErrors = errsVar+ , envErrorsLog = logVar+ , envInjectFaultResults = faultsVar+ }+ prop <- runReaderT r env+ faults <- readMutVar faultsVar+ pure $ QC.tabulate "Fault results" (fmap show faults) prop+ )+ tagFinalState'+ where+ acquire :: IO (FilePath, Class.Session R.Table IO, HasBlockIO IO HandleIO, StrictTVar IO Errors, StrictTVar IO ErrorsLog)+ acquire = do+ (tmpDir, hasFS, hasBlockIO) <- createSystemTempDirectory "prop_lockstepIO_RealImpl_RealFS"+ session <- R.openSession tr hasFS hasBlockIO salt (mkFsPath [])+ errsVar <- newTVarIO FSSim.emptyErrors+ logVar <- newTVarIO emptyLog+ pure (tmpDir, session, hasBlockIO, errsVar, logVar)++ release :: (FilePath, Class.Session R.Table IO, HasBlockIO IO HandleIO, StrictTVar IO Errors, StrictTVar IO ErrorsLog) -> IO Property+ release (tmpDir, !session, hasBlockIO, _, _) = do+ !prop <- propNoThunks session+ R.closeSession session+ close hasBlockIO+ removeDirectoryRecursive tmpDir+ pure prop++propLockstep_RealImpl_MockFS_IO ::+ Tracer IO R.LSMTreeTrace+ -> CheckCleanup+ -> CheckFS+ -> CheckRefs+ -> QC.Fixed R.Salt+ -> Actions (Lockstep (ModelState IO R.Table))+ -> QC.Property+propLockstep_RealImpl_MockFS_IO tr cleanupFlag fsFlag refsFlag (QC.Fixed salt) =+ runActionsBracket+ (Proxy @(ModelState IO R.Table))+ cleanupFlag+ refsFlag+ (acquire_RealImpl_MockFS tr salt)+ (release_RealImpl_MockFS fsFlag)+ (\r (_, session, errsVar, logVar) -> do+ faultsVar <- newMutVar []+ let+ env :: RealEnv R.Table IO+ env = RealEnv {+ envSession = session+ , envHandlers = realErrorHandlers @IO+ , envErrors = errsVar+ , envErrorsLog = logVar+ , envInjectFaultResults = faultsVar+ }+ prop <- runReaderT r env+ faults <- readMutVar faultsVar+ pure $ QC.tabulate "Fault results" (fmap show faults) prop+ )+ tagFinalState'++propLockstep_RealImpl_MockFS_IOSim ::+ (forall s. Tracer (IOSim s) R.LSMTreeTrace)+ -> CheckCleanup+ -> CheckFS+ -> CheckRefs+ -> QC.Fixed R.Salt+ -> Actions (Lockstep (ModelState (IOSim RealWorld) R.Table))+ -> QC.Property+propLockstep_RealImpl_MockFS_IOSim tr cleanupFlag fsFlag refsFlag (QC.Fixed salt) =+ runActionsBracket+ (Proxy @(ModelState (IOSim RealWorld) R.Table))+ cleanupFlag+ refsFlag+ (acquire_RealImpl_MockFS tr salt)+ (release_RealImpl_MockFS fsFlag)+ (\r (_, session, errsVar, logVar) -> do+ faultsVar <- newMutVar []+ let+ env :: RealEnv R.Table (IOSim RealWorld)+ env = RealEnv {+ envSession = session+ , envHandlers = realErrorHandlers @(IOSim RealWorld)+ , envErrors = errsVar+ , envErrorsLog = logVar+ , envInjectFaultResults = faultsVar+ }+ prop <- runReaderT r env+ faults <- readMutVar faultsVar+ pure $ QC.tabulate "Fault results" (fmap show faults) prop+ )+ tagFinalState'++acquire_RealImpl_MockFS ::+ R.IOLike m+ => Tracer m R.LSMTreeTrace+ -> R.Salt+ -> m (StrictTMVar m MockFS, Class.Session R.Table m, StrictTVar m Errors, StrictTVar m ErrorsLog)+acquire_RealImpl_MockFS tr salt = do+ fsVar <- newTMVarIO MockFS.empty+ errsVar <- newTVarIO FSSim.emptyErrors+ logVar <- newTVarIO emptyLog+ (hfs, hbio) <- simErrorHasBlockIOLogged fsVar errsVar logVar+ session <- R.openSession tr hfs hbio salt (mkFsPath [])+ pure (fsVar, session, errsVar, logVar)++-- | Flag that turns on\/off file system checks.+data CheckFS = CheckFS | NoCheckFS++release_RealImpl_MockFS ::+ (R.IOLike m)+ => CheckFS+ -> (StrictTMVar m MockFS, Class.Session R.Table m, StrictTVar m Errors, StrictTVar m ErrorsLog)+ -> m Property+release_RealImpl_MockFS fsFlag (fsVar, session, _, _) = do+ sts <- getAllSessionTables session+ forM_ sts $ \(SomeTable t) -> R.closeTable t+ scs <- getAllSessionCursors session+ forM_ scs $ \(SomeCursor c) -> R.closeCursor c+ mockfs1 <- atomically $ readTMVar fsVar+ R.closeSession session+ mockfs2 <- atomically $ readTMVar fsVar+ pure $ case fsFlag of+ CheckFS -> propNumOpenHandles 1 mockfs1 QC..&&. propNoOpenHandles mockfs2+ NoCheckFS -> QC.property ()++data SomeTable m+ = SomeTable (forall k v b. R.Table m k v b)++data SomeCursor m+ = SomeCursor (forall k v b. R.Cursor m k v b)++getAllSessionTables ::+ (MonadSTM m, MonadThrow m, MonadMVar m)+ => R.Session m+ -> m [SomeTable m]+getAllSessionTables (R.Types.Session s) = do+ R.Unsafe.withKeepSessionOpen s $ \seshEnv -> do+ ts <- readMVar (R.Unsafe.sessionOpenTables seshEnv)+ pure ((\x -> SomeTable (R.Types.Table x)) <$> Map.elems ts)++getAllSessionCursors ::+ (MonadSTM m, MonadThrow m, MonadMVar m)+ => R.Session m+ -> m [SomeCursor m]+getAllSessionCursors (R.Types.Session s) =+ R.Unsafe.withKeepSessionOpen s $ \seshEnv -> do+ cs <- readMVar (R.Unsafe.sessionOpenCursors seshEnv)+ pure ((\x -> SomeCursor (R.Types.Cursor x)) <$> Map.elems cs)++createSystemTempDirectory :: [Char] -> IO (FilePath, HasFS IO HandleIO, HasBlockIO IO HandleIO)+createSystemTempDirectory prefix = do+ systemTempDir <- getCanonicalTemporaryDirectory+ tempDir <- createTempDirectory systemTempDir prefix+ (hasFS, hasBlockIO) <- ioHasBlockIO (MountPoint tempDir) defaultIOCtxParams+ pure (tempDir, hasFS, hasBlockIO)++{-------------------------------------------------------------------------------+ Error handlers+-------------------------------------------------------------------------------}++realErrorHandlers :: Applicative f => [Handler f (Maybe Model.Err)]+realErrorHandlers = [Handler $ pure . Just . handleSomeException]++handleSomeException :: SomeException -> Model.Err+handleSomeException e =+ fromMaybe (Model.ErrOther $ displayException e) . getFirst . mconcat . fmap First $+ [ handleCommitActionRegistryError <$> fromException e+ , handleAbortActionRegistryError <$> fromException e+ , handleSessionDirDoesNotExistError <$> fromException e+ , handleSessionDirLockedError <$> fromException e+ , handleSessionDirCorruptedError <$> fromException e+ , handleSessionClosedError <$> fromException e+ , handleTableClosedError <$> fromException e+ , handleTableCorruptedError <$> fromException e+ , handleTableUnionNotCompatibleError <$> fromException e+ , handleSnapshotExistsError <$> fromException e+ , handleSnapshotDoesNotExistError <$> fromException e+ , handleSnapshotCorruptedError <$> fromException e+ , handleSnapshotNotCompatibleError <$> fromException e+ , handleBlobRefInvalidError <$> fromException e+ , handleCursorClosedError <$> fromException e+ , handleFsError <$> fromException e+ ]++handleCommitActionRegistryError :: CommitActionRegistryError -> Model.Err+handleCommitActionRegistryError = \case+ CommitActionRegistryError actionErrors ->+ Model.ErrCommitActionRegistry $+ handleSomeException <$> (getActionError <$> actionErrors)++handleAbortActionRegistryError :: AbortActionRegistryError -> Model.Err+handleAbortActionRegistryError = \case+ AbortActionRegistryError abortReason actionErrors ->+ Model.ErrAbortActionRegistry+ (handleSomeException <$> getReasonExitCaseException abortReason)+ (handleSomeException . getActionError <$> actionErrors)++handleSessionDirDoesNotExistError :: SessionDirDoesNotExistError -> Model.Err+handleSessionDirDoesNotExistError = \case+ ErrSessionDirDoesNotExist _dir -> Model.ErrSessionDirDoesNotExist++handleSessionDirLockedError :: SessionDirLockedError -> Model.Err+handleSessionDirLockedError = \case+ ErrSessionDirLocked _dir -> Model.ErrSessionDirLocked++handleSessionDirCorruptedError :: SessionDirCorruptedError -> Model.Err+handleSessionDirCorruptedError = \case+ ErrSessionDirCorrupted _reason _dir -> Model.ErrSessionDirCorrupted++handleSessionClosedError :: SessionClosedError -> Model.Err+handleSessionClosedError = \case+ ErrSessionClosed -> Model.ErrSessionClosed++handleTableClosedError :: TableClosedError -> Model.Err+handleTableClosedError = \case+ ErrTableClosed -> Model.ErrTableClosed++handleTableCorruptedError :: TableCorruptedError -> Model.Err+handleTableCorruptedError = \case+ ErrLookupByteCountDiscrepancy{} -> Model.ErrTableCorrupted++handleTableUnionNotCompatibleError :: TableUnionNotCompatibleError -> Model.Err+handleTableUnionNotCompatibleError = \case+ ErrTableUnionHandleTypeMismatch{} -> Model.ErrTableUnionHandleTypeMismatch+ ErrTableUnionSessionMismatch{} -> Model.ErrTableUnionSessionMismatch++handleSnapshotExistsError :: SnapshotExistsError -> Model.Err+handleSnapshotExistsError = \case+ ErrSnapshotExists name -> Model.ErrSnapshotExists name++handleSnapshotDoesNotExistError :: SnapshotDoesNotExistError -> Model.Err+handleSnapshotDoesNotExistError = \case+ ErrSnapshotDoesNotExist name -> Model.ErrSnapshotDoesNotExist name++handleSnapshotCorruptedError :: SnapshotCorruptedError -> Model.Err+handleSnapshotCorruptedError = \case+ ErrSnapshotCorrupted name _ -> Model.ErrSnapshotCorrupted name++handleSnapshotNotCompatibleError :: SnapshotNotCompatibleError -> Model.Err+handleSnapshotNotCompatibleError = \case+ ErrSnapshotWrongLabel name _ _ -> Model.ErrSnapshotWrongLabel name++handleBlobRefInvalidError :: BlobRefInvalidError -> Model.Err+handleBlobRefInvalidError = \case+ ErrBlobRefInvalid _ -> Model.ErrBlobRefInvalid++handleCursorClosedError :: CursorClosedError -> Model.Err+handleCursorClosedError = \case+ ErrCursorClosed -> Model.ErrCursorClosed++handleFsError :: FsError -> Model.Err+handleFsError = Model.ErrFsError . displayException++{-------------------------------------------------------------------------------+ Key and value types+-------------------------------------------------------------------------------}++newtype Key = Key SerialisedKey+ deriving stock (Show, Eq, Ord)+ deriving newtype (Arbitrary, R.SerialiseKey)++newtype Value = Value SerialisedValue+ deriving stock (Show, Eq)+ deriving newtype (Arbitrary, R.SerialiseValue)++newtype Blob = Blob SerialisedBlob+ deriving stock (Show, Eq)+ deriving newtype (Arbitrary, R.SerialiseValue)++keyValueBlobLabel :: R.SnapshotLabel+keyValueBlobLabel = R.SnapshotLabel "Key Value Blob"++instance R.ResolveValue Value where+ resolveSerialised _ = (<>)++{-------------------------------------------------------------------------------+ Model state+-------------------------------------------------------------------------------}++type ModelState :: (Type -> Type) -> ((Type -> Type) -> Type -> Type -> Type -> Type) -> Type+data ModelState m h = ModelState Model.Model Stats+ deriving stock Show++initModelState :: ModelState m h+initModelState = ModelState Model.initModel initStats++{-------------------------------------------------------------------------------+ Type synonyms+-------------------------------------------------------------------------------}++type Act m h a = Action (Lockstep (ModelState m h)) a+type Act' m h a = Action' m h (Either Model.Err a)+type Var m h a = ModelVar (ModelState m h) a+type Val m h a = ModelValue (ModelState m h) a+type Obs m h a = Observable (ModelState m h) a++type K a = (+ Class.C_ a+ , R.SerialiseKey a+ , Arbitrary a+ )++type V a = (+ Class.C_ a+ , R.SerialiseValue a+ , R.ResolveValue a+ , Arbitrary a+ )++type B a = (+ Class.C_ a+ , R.SerialiseValue a+ , Arbitrary a+ )++-- | Common constraints for keys, values and blobs+type C k v b = (K k, V v, B b)++{-------------------------------------------------------------------------------+ StateModel+-------------------------------------------------------------------------------}++newtype SilentCorruption = SilentCorruption {bitChoice :: Choice}+ deriving stock (Eq, Show)+ deriving newtype (Arbitrary)++instance ( Show (Class.TableConfig h)+ , Eq (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable m+ ) => StateModel (Lockstep (ModelState m h)) where+ data instance Action (Lockstep (ModelState m h)) a where+ Action :: Maybe Errors -> Action' m h a -> Act m h a++ initialState = Lockstep.Defaults.initialState initModelState+ nextState = Lockstep.Defaults.nextState+ precondition = Lockstep.Defaults.precondition+ arbitraryAction = Lockstep.Defaults.arbitraryAction+ shrinkAction = Lockstep.Defaults.shrinkAction++ -- Does not use the default implementation, because that would print "Action".+ -- We print the name of the inner 'Action'' instead.+ actionName (Action _ action') = head . words . show $ action'++deriving stock instance Show (Class.TableConfig h)+ => Show (LockstepAction (ModelState m h) a)++instance ( Eq (Class.TableConfig h)+ , Typeable h+ , Typeable m+ ) => Eq (LockstepAction (ModelState m h) a) where+ (==) :: LockstepAction (ModelState m h) a -> LockstepAction (ModelState m h) a -> Bool+ Action merrs1 x == Action merrs2 y = merrs1 == merrs2 && x == y+ where+ _coveredAllCases :: Action (Lockstep (ModelState m h)) a -> ()+ _coveredAllCases = \case+ Action{} -> ()++data Action' m h a where+ -- Tables+ NewTableWith ::+ C k v b+ => {-# UNPACK #-} !(PrettyProxy (k, v, b))+ -> Class.TableConfig h+ -> Act' m h (WrapTable h m k v b)+ CloseTable ::+ C k v b+ => Var m h (WrapTable h m k v b)+ -> Act' m h ()+ -- Queries+ Lookups ::+ C k v b+ => V.Vector k -> Var m h (WrapTable h m k v b)+ -> Act' m h (V.Vector (LookupResult v (WrapBlobRef h m b)))+ RangeLookup ::+ (C k v b, Ord k)+ => R.Range k -> Var m h (WrapTable h m k v b)+ -> Act' m h (V.Vector (Entry k v (WrapBlobRef h m b)))+ -- Cursor+ NewCursor ::+ C k v b+ => Maybe k+ -> Var m h (WrapTable h m k v b)+ -> Act' m h (WrapCursor h m k v b)+ CloseCursor ::+ C k v b+ => Var m h (WrapCursor h m k v b)+ -> Act' m h ()+ ReadCursor ::+ C k v b+ => Int+ -> Var m h (WrapCursor h m k v b)+ -> Act' m h (V.Vector (Entry k v (WrapBlobRef h m b)))+ -- Updates+ Updates ::+ C k v b+ => V.Vector (k, R.Update v b) -> Var m h (WrapTable h m k v b)+ -> Act' m h ()+ Inserts ::+ C k v b+ => V.Vector (k, v, Maybe b) -> Var m h (WrapTable h m k v b)+ -> Act' m h ()+ Deletes ::+ C k v b+ => V.Vector k -> Var m h (WrapTable h m k v b)+ -> Act' m h ()+ Upserts ::+ C k v b+ => V.Vector (k, v) -> Var m h (WrapTable h m k v b)+ -> Act' m h ()+ -- Blobs+ RetrieveBlobs ::+ B b+ => Var m h (V.Vector (WrapBlobRef h m b))+ -> Act' m h (V.Vector (WrapBlob b))+ -- Snapshots+ SaveSnapshot ::+ C k v b+ => Maybe SilentCorruption+ -> R.SnapshotName+ -> R.SnapshotLabel+ -> Var m h (WrapTable h m k v b)+ -> Act' m h ()+ OpenTableFromSnapshot ::+ C k v b+ => {-# UNPACK #-} !(PrettyProxy (k, v, b))+ -> R.SnapshotName+ -> R.SnapshotLabel+ -> Act' m h (WrapTable h m k v b)+ DeleteSnapshot :: R.SnapshotName -> Act' m h ()+ ListSnapshots :: Act' m h [R.SnapshotName]+ -- Duplicate tables+ Duplicate ::+ C k v b+ => Var m h (WrapTable h m k v b)+ -> Act' m h (WrapTable h m k v b)+ -- Table union+ Union ::+ C k v b+ => Var m h (WrapTable h m k v b)+ -> Var m h (WrapTable h m k v b)+ -> Act' m h (WrapTable h m k v b)+ Unions ::+ C k v b+ => NonEmpty (Var m h (WrapTable h m k v b))+ -> Act' m h (WrapTable h m k v b)+ RemainingUnionDebt ::+ C k v b+ => Var m h (WrapTable h m k v b)+ -> Act' m h R.UnionDebt+ SupplyUnionCredits ::+ C k v b+ => Var m h (WrapTable h m k v b)+ -> R.UnionCredits+ -> Act' m h R.UnionCredits+ -- | Alternative version of 'SupplyUnionCredits' that supplies a portion of+ -- the table's current union debt as union credits.+ --+ -- 'SupplyUnionCredits' gets no information about union debt, so the union+ -- credits we generate are arbitrary, and it would require precarious, manual+ -- tuning to make sure the debt is ever paid off by an action sequence.+ -- 'SupplyUnionCredits' supplies a portion (if not all) of the current debt,+ -- so that unions are more likely to finish during a sequence of actions.+ SupplyPortionOfDebt ::+ C k v b+ => Var m h (WrapTable h m k v b)+ -> Portion+ -> Act' m h R.UnionCredits++portionOf :: Portion -> R.UnionDebt -> R.UnionCredits+portionOf (Portion denominator) (R.UnionDebt debt)+ | denominator <= 0 = error "portion: denominator should be positive"+ | debt < 0 = error "portion: debt should be non-negative"+ | otherwise = R.UnionCredits (debt `div` denominator)++newtype Portion = Portion Int -- ^ Denominator: should be non-negative+ deriving stock (Show, Eq)++deriving stock instance Show (Class.TableConfig h)+ => Show (Action' m h a)++instance ( Eq (Class.TableConfig h)+ , Typeable h+ , Typeable m+ ) => Eq (Action' m h a) where+ x == y = go x y+ where+ go :: Action' m h a -> Action' m h a -> Bool+ go+ (NewTableWith (PrettyProxy :: PrettyProxy kvb) conf1)+ (NewTableWith (PrettyProxy :: PrettyProxy kvb) conf2) =+ conf1 == conf2+ go (CloseTable var1) (CloseTable var2) =+ Just var1 == cast var2+ go (Lookups ks1 var1) (Lookups ks2 var2) =+ Just ks1 == cast ks2 && Just var1 == cast var2+ go (RangeLookup range1 var1) (RangeLookup range2 var2) =+ range1 == range2 && var1 == var2+ go (NewCursor k1 var1) (NewCursor k2 var2) =+ k1 == k2 && var1 == var2+ go (CloseCursor var1) (CloseCursor var2) =+ Just var1 == cast var2+ go (ReadCursor n1 var1) (ReadCursor n2 var2) =+ n1 == n2 && var1 == var2+ go (Updates ups1 var1) (Updates ups2 var2) =+ Just ups1 == cast ups2 && Just var1 == cast var2+ go (Inserts inss1 var1) (Inserts inss2 var2) =+ Just inss1 == cast inss2 && Just var1 == cast var2+ go (Deletes ks1 var1) (Deletes ks2 var2) =+ Just ks1 == cast ks2 && Just var1 == cast var2+ go (Upserts mups1 var1) (Upserts mups2 var2) =+ Just mups1 == cast mups2 && Just var1 == cast var2+ go (RetrieveBlobs vars1) (RetrieveBlobs vars2) =+ Just vars1 == cast vars2+ go (SaveSnapshot mcorr1 name1 label1 var1) (SaveSnapshot mcorr2 name2 label2 var2) =+ mcorr1 == mcorr2 && name1 == name2 && label1 == label2 && Just var1 == cast var2+ go+ (OpenTableFromSnapshot (PrettyProxy :: PrettyProxy kvb) name1 label1)+ (OpenTableFromSnapshot (PrettyProxy :: PrettyProxy kvb) name2 label2) =+ label1 == label2 && name1 == name2+ go (DeleteSnapshot name1) (DeleteSnapshot name2) =+ name1 == name2+ go ListSnapshots ListSnapshots =+ True+ go (Duplicate var1) (Duplicate var2) =+ Just var1 == cast var2+ go (Union var1_1 var1_2) (Union var2_1 var2_2) =+ Just var1_1 == cast var2_1 && Just var1_2 == cast var2_2+ go (Unions vars1) (Unions vars2) =+ Just vars1 == cast vars2+ go (RemainingUnionDebt var1) (RemainingUnionDebt var2) =+ Just var1 == cast var2+ go (SupplyUnionCredits var1 credits1) (SupplyUnionCredits var2 credits2) =+ Just var1 == cast var2 && credits1 == credits2+ go (SupplyPortionOfDebt var1 portion1) (SupplyPortionOfDebt var2 portion2) =+ Just var1 == cast var2 && portion1 == portion2+ go _ _ = False++ _coveredAllCases :: Action' m h a -> ()+ _coveredAllCases = \case+ NewTableWith{} -> ()+ CloseTable{} -> ()+ Lookups{} -> ()+ RangeLookup{} -> ()+ NewCursor{} -> ()+ CloseCursor{} -> ()+ ReadCursor{} -> ()+ Updates{} -> ()+ Inserts{} -> ()+ Deletes{} -> ()+ Upserts{} -> ()+ RetrieveBlobs{} -> ()+ SaveSnapshot{} -> ()+ OpenTableFromSnapshot{} -> ()+ DeleteSnapshot{} -> ()+ ListSnapshots{} -> ()+ Duplicate{} -> ()+ Union{} -> ()+ Unions{} -> ()+ RemainingUnionDebt{} -> ()+ SupplyUnionCredits{} -> ()+ SupplyPortionOfDebt{} -> ()++-- | This is not a fully lawful instance, because it uses 'approximateEqStream'.+instance Eq a => Eq (Stream a) where+ (==) = approximateEqStream++-- | This is not a fully lawful instance, because it uses 'approximateEqStream'.+deriving stock instance Eq Errors+deriving stock instance Eq FSSim.Partial+deriving stock instance Eq FSSim.PutCorruption+deriving stock instance Eq FSSim.Blob++{-------------------------------------------------------------------------------+ InLockstep+-------------------------------------------------------------------------------}++instance ( Eq (Class.TableConfig h)+ , Show (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable m+ ) => InLockstep (ModelState m h) where+ type instance ModelOp (ModelState m h) = Op++ data instance ModelValue (ModelState m h) a where+ -- handle-like+ MTable ::+ Model.Table k v b ->+ Val m h (WrapTable h m k v b)+ MCursor :: Model.Cursor k v b -> Val m h (WrapCursor h m k v b)+ MBlobRef ::+ (Class.C_ b) =>+ Model.BlobRef b ->+ Val m h (WrapBlobRef h m b)+ -- values+ MLookupResult ::+ (Class.C_ v, Class.C_ b) =>+ LookupResult v (Val m h (WrapBlobRef h m b)) ->+ Val m h (LookupResult v (WrapBlobRef h m b))+ MEntry ::+ (Class.C k v b) =>+ Entry k v (Val m h (WrapBlobRef h m b)) ->+ Val m h (Entry k v (WrapBlobRef h m b))+ MBlob ::+ (Show b, Typeable b, Eq b) =>+ WrapBlob b ->+ Val m h (WrapBlob b)+ MSnapshotName :: R.SnapshotName -> Val m h R.SnapshotName+ MUnionDebt :: R.UnionDebt -> Val m h R.UnionDebt+ MUnionCredits :: R.UnionCredits -> Val m h R.UnionCredits+ MUnionCreditsPortion :: R.UnionCredits -> Val m h R.UnionCredits+ MErr :: Model.Err -> Val m h Model.Err+ -- combinators+ MUnit :: () -> Val m h ()+ MPair :: (Val m h a, Val m h b) -> Val m h (a, b)+ MEither :: Either (Val m h a) (Val m h b) -> Val m h (Either a b)+ MList :: [Val m h a] -> Val m h [a]+ MVector :: V.Vector (Val m h a) -> Val m h (V.Vector a)++ data instance Observable (ModelState m h) a where+ -- handle-like (opaque)+ OTable :: Obs m h (WrapTable h m k v b)+ OCursor :: Obs m h (WrapCursor h m k v b)+ OBlobRef :: Obs m h (WrapBlobRef h m b)+ -- values+ OLookupResult ::+ (Class.C_ v, Class.C_ b) =>+ LookupResult v (Obs m h (WrapBlobRef h m b)) ->+ Obs m h (LookupResult v (WrapBlobRef h m b))+ OEntry ::+ (Class.C k v b) =>+ Entry k v (Obs m h (WrapBlobRef h m b)) ->+ Obs m h (Entry k v (WrapBlobRef h m b))+ OBlob ::+ (Show b, Typeable b, Eq b) =>+ WrapBlob b ->+ Obs m h (WrapBlob b)+ OUnionCredits :: R.UnionCredits -> Obs m h R.UnionCredits+ OUnionCreditsPortion :: R.UnionCredits -> Obs m h R.UnionCredits+ OId :: (Show a, Typeable a, Eq a) => a -> Obs m h a+ -- combinators+ OPair :: (Obs m h a, Obs m h b) -> Obs m h (a, b)+ OEither :: Either (Obs m h a) (Obs m h b) -> Obs m h (Either a b)+ OList :: [Obs m h a] -> Obs m h [a]+ OVector :: V.Vector (Obs m h a) -> Obs m h (V.Vector a)++ observeModel :: Val m h a -> Obs m h a+ observeModel = \case+ MTable _ -> OTable+ MCursor _ -> OCursor+ MBlobRef _ -> OBlobRef+ MLookupResult x -> OLookupResult $ fmap observeModel x+ MEntry x -> OEntry $ fmap observeModel x+ MSnapshotName x -> OId x+ MBlob x -> OBlob x+ MUnionDebt x -> OId x+ MUnionCredits x -> OUnionCredits x+ MUnionCreditsPortion x -> OUnionCreditsPortion x+ MErr x -> OId x+ MUnit x -> OId x+ MPair x -> OPair $ bimap observeModel observeModel x+ MEither x -> OEither $ bimap observeModel observeModel x+ MList x -> OList $ map observeModel x+ MVector x -> OVector $ V.map observeModel x++ modelNextState :: forall a.+ LockstepAction (ModelState m h) a+ -> ModelVarContext (ModelState m h)+ -> ModelState m h+ -> (ModelValue (ModelState m h) a, ModelState m h)+ modelNextState action ctx (ModelState state stats) =+ auxStats $ runModel (lookupVar ctx) action state+ where+ auxStats :: (Val m h a, Model.Model) -> (Val m h a, ModelState m h)+ auxStats (result, state') = (result, ModelState state' stats')+ where+ stats' = updateStats action (lookupVar ctx) state state' result stats++ usedVars :: LockstepAction (ModelState m h) a -> [AnyGVar (ModelOp (ModelState m h))]+ usedVars (Action _ action') = case action' of+ NewTableWith _ _ -> []+ CloseTable tableVar -> [SomeGVar tableVar]+ Lookups _ tableVar -> [SomeGVar tableVar]+ RangeLookup _ tableVar -> [SomeGVar tableVar]+ NewCursor _ tableVar -> [SomeGVar tableVar]+ CloseCursor cursorVar -> [SomeGVar cursorVar]+ ReadCursor _ cursorVar -> [SomeGVar cursorVar]+ Updates _ tableVar -> [SomeGVar tableVar]+ Inserts _ tableVar -> [SomeGVar tableVar]+ Deletes _ tableVar -> [SomeGVar tableVar]+ Upserts _ tableVar -> [SomeGVar tableVar]+ RetrieveBlobs blobsVar -> [SomeGVar blobsVar]+ SaveSnapshot _ _ _ tableVar -> [SomeGVar tableVar]+ OpenTableFromSnapshot{} -> []+ DeleteSnapshot _ -> []+ ListSnapshots -> []+ Duplicate tableVar -> [SomeGVar tableVar]+ Union table1Var table2Var -> [SomeGVar table1Var, SomeGVar table2Var]+ Unions tableVars -> [SomeGVar tableVar | tableVar <- NE.toList tableVars]+ RemainingUnionDebt tableVar -> [SomeGVar tableVar]+ SupplyUnionCredits tableVar _ -> [SomeGVar tableVar]+ SupplyPortionOfDebt tableVar _ -> [SomeGVar tableVar]++ arbitraryWithVars ::+ ModelVarContext (ModelState m h)+ -> ModelState m h+ -> Gen (Any (LockstepAction (ModelState m h)))+ arbitraryWithVars ctx st =+ QC.scale (max 100) $+ arbitraryActionWithVars (Proxy @(Key, Value, Blob)) keyValueBlobLabel ctx st++ shrinkWithVars ::+ ModelVarContext (ModelState m h)+ -> ModelState m h+ -> LockstepAction (ModelState m h) a+ -> [Any (LockstepAction (ModelState m h))]+ shrinkWithVars = shrinkActionWithVars++ tagStep ::+ (ModelState m h, ModelState m h)+ -> LockstepAction (ModelState m h) a+ -> Val m h a+ -> [String]+ tagStep states action = map show . tagStep' states action++deriving stock instance Show (Class.TableConfig h) => Show (Val m h a)+deriving stock instance Show (Obs m h a)++instance Eq (Obs m h a) where+ obsReal == obsModel = case (obsReal, obsModel) of+ -- The model is conservative about blob retrieval: the model invalidates a+ -- blob reference immediately after an update to the table, and if the SUT+ -- returns a blob, then that's okay. If both return a blob or both return+ -- an error, then those must match exactly.+ (OEither (Right (OVector vec)), OEither (Left (OId y)))+ | Just Model.ErrBlobRefInvalid <- cast y ->+ flip all vec $ \case+ OBlob (WrapBlob _) -> True+ _ -> False++ -- When disk faults are injected, the model only knows /something/ went+ -- wrong, but the SUT can throw much more specific errors. We allow this.+ --+ -- See also 'Model.runModelMWithInjectedErrors' and+ -- 'runRealWithInjectedErrors'.+ (OEither (Left (OId lhs)), OEither (Left (OId rhs)))+ | Just (e :: Model.Err) <- cast lhs+ , not $ Model.isOther e+ , Just (Model.ErrDiskFault _) <- cast rhs+ -> True++ -- When snapshots are corrupted, the model only knows that the snapshot+ -- was corrupted, but not how, but the SUT can throw much more specific+ -- errors. We allow this.+ (OEither (Left (OId lhs)), OEither (Left (OId rhs)))+ | Just (e :: Model.Err) <- cast lhs+ , not $ Model.isOther e+ , Just (Model.ErrSnapshotCorrupted _) <- cast rhs+ -> True++ -- RemainingUnionDebt+ --+ -- Debt in the model is always 0, while debt in the real implementation+ -- may be larger than 0.+ (OEither (Right (OId lhs)), OEither (Right (OId rhs)))+ | Just (lhsDebt :: R.UnionDebt) <- cast lhs+ , Just (rhsDebt :: R.UnionDebt) <- cast rhs+ -> lhsDebt >= R.UnionDebt 0 && rhsDebt == R.UnionDebt 0++ -- SupplyUnionCredits+ --+ -- In the model, all supplied union credits are returned as leftovers,+ -- whereas the real implementation may use up some union credits.+ (OEither (Right (OUnionCredits lhsLeftovers)), OEither (Right (OUnionCredits rhsLeftovers)))+ -> lhsLeftovers <= rhsLeftovers++ -- SupplyPortionOfDebt+ --+ -- In the model, a portion of the debt is always 0, whereas in the real+ -- implementation the portion of the debt is non-negative.+ (OEither (Right (OUnionCreditsPortion lhsLeftovers)), OEither (Right (OUnionCreditsPortion rhsLeftovers)))+ -> lhsLeftovers >= rhsLeftovers++ -- default equalities+ (OTable, OTable) -> True+ (OCursor, OCursor) -> True+ (OBlobRef, OBlobRef) -> True+ (OLookupResult x, OLookupResult y) -> x == y+ (OEntry x, OEntry y) -> x == y+ (OBlob x, OBlob y) -> x == y+ (OUnionCredits x, OUnionCredits y) -> x == y+ (OUnionCreditsPortion x, OUnionCreditsPortion y) -> x == y+ (OId x, OId y) -> x == y+ (OPair x, OPair y) -> x == y+ (OEither x, OEither y) -> x == y+ (OList x, OList y) -> x == y+ (OVector x, OVector y) -> x == y+ (_, _) -> False+ where+ _coveredAllCases :: Obs m h a -> ()+ _coveredAllCases = \case+ OTable{} -> ()+ OCursor{} -> ()+ OBlobRef{} -> ()+ OLookupResult{} -> ()+ OEntry{} -> ()+ OBlob{} -> ()+ OUnionCredits{} -> ()+ OUnionCreditsPortion{} -> ()+ OId{} -> ()+ OPair{} -> ()+ OEither{} -> ()+ OList{} -> ()+ OVector{} -> ()++{-------------------------------------------------------------------------------+ Real monad+-------------------------------------------------------------------------------}++type RealMonad h m = ReaderT (RealEnv h m) m++-- | An environment for implementations @h@ of the public API to run actions in+-- (see 'perform', 'runIO', 'runIOSim').+data RealEnv h m = RealEnv {+ -- | The session to run actions in.+ envSession :: !(Class.Session h m)+ -- | Error handlers to convert thrown exceptions into pure error values.+ --+ -- Uncaught exceptions make the tests fail, so some handlers should be+ -- provided that catch the exceptions and turn them into pure error values+ -- if possible. The state machine infrastructure can then also compare the+ -- error values obtained from running @h@ to the error values obtained by+ -- running the model.+ , envHandlers :: [Handler m (Maybe Model.Err)]+ -- | A variable holding simulated disk faults,+ --+ -- This variable shared with the simulated file system (if in use). This+ -- variable can be used to enable/disable errors locally, for example on a+ -- per-action basis.+ , envErrors :: !(StrictTVar m Errors)+ -- | A variable holding a log of simulated disk faults.+ --+ -- Errors that are injected into the simulated file system using 'envErrors'+ -- are logged here.+ , envErrorsLog :: !(StrictTVar m ErrorsLog)+ -- | The results of fault injection+ , envInjectFaultResults :: !(MutVar (PrimState m) [InjectFaultResult])+ }++data InjectFaultResult =+ -- | No faults were injected.+ InjectFaultNone+ String -- ^ Action name+ -- | Faults were injected, but the action accidentally succeeded, so the+ -- action had to be rolled back+ | InjectFaultAccidentalSuccess+ String -- ^ Action name+ -- | Faults were injected, which made the action fail with an error.+ | InjectFaultInducedError+ String -- ^ Action name+ deriving stock Show++{-------------------------------------------------------------------------------+ RunLockstep+-------------------------------------------------------------------------------}++instance ( Eq (Class.TableConfig h)+ , Class.IsTable h+ , Show (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ ) => RunLockstep (ModelState IO h) (RealMonad h IO) where+ observeReal ::+ Proxy (RealMonad h IO)+ -> LockstepAction (ModelState IO h) a+ -> a+ -> Obs IO h a+ observeReal _proxy (Action _ action') result = case action' of+ NewTableWith{} -> OEither $ bimap OId (const OTable) result+ CloseTable{} -> OEither $ bimap OId OId result+ Lookups{} -> OEither $+ bimap OId (OVector . fmap (OLookupResult . fmap (const OBlobRef))) result+ RangeLookup{} -> OEither $+ bimap OId (OVector . fmap (OEntry . fmap (const OBlobRef))) result+ NewCursor{} -> OEither $ bimap OId (const OCursor) result+ CloseCursor{} -> OEither $ bimap OId OId result+ ReadCursor{} -> OEither $+ bimap OId (OVector . fmap (OEntry . fmap (const OBlobRef))) result+ Updates{} -> OEither $ bimap OId OId result+ Inserts{} -> OEither $ bimap OId OId result+ Deletes{} -> OEither $ bimap OId OId result+ Upserts{} -> OEither $ bimap OId OId result+ RetrieveBlobs{} -> OEither $ bimap OId (OVector . fmap OBlob) result+ SaveSnapshot{} -> OEither $ bimap OId OId result+ OpenTableFromSnapshot{} -> OEither $ bimap OId (const OTable) result+ DeleteSnapshot{} -> OEither $ bimap OId OId result+ ListSnapshots{} -> OEither $ bimap OId (OList . fmap OId) result+ Duplicate{} -> OEither $ bimap OId (const OTable) result+ Union{} -> OEither $ bimap OId (const OTable) result+ Unions{} -> OEither $ bimap OId (const OTable) result+ RemainingUnionDebt{} -> OEither $ bimap OId OId result+ SupplyUnionCredits{} -> OEither $ bimap OId OUnionCredits result+ SupplyPortionOfDebt{} -> OEither $ bimap OId OUnionCreditsPortion result++ showRealResponse ::+ Proxy (RealMonad h IO)+ -> LockstepAction (ModelState IO h) a+ -> Maybe (Dict (Show a))+ showRealResponse _ (Action _ action') = case action' of+ NewTableWith{} -> Nothing+ CloseTable{} -> Just Dict+ Lookups{} -> Nothing+ RangeLookup{} -> Nothing+ NewCursor{} -> Nothing+ CloseCursor{} -> Just Dict+ ReadCursor{} -> Nothing+ Updates{} -> Just Dict+ Inserts{} -> Just Dict+ Deletes{} -> Just Dict+ Upserts{} -> Just Dict+ RetrieveBlobs{} -> Just Dict+ SaveSnapshot{} -> Just Dict+ OpenTableFromSnapshot{} -> Nothing+ DeleteSnapshot{} -> Just Dict+ ListSnapshots{} -> Just Dict+ Duplicate{} -> Nothing+ Union{} -> Nothing+ Unions{} -> Nothing+ RemainingUnionDebt{} -> Just Dict+ SupplyUnionCredits{} -> Just Dict+ SupplyPortionOfDebt{} -> Just Dict++instance ( Eq (Class.TableConfig h)+ , Class.IsTable h+ , Show (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable s+ ) => RunLockstep (ModelState (IOSim s) h) (RealMonad h (IOSim s)) where+ observeReal ::+ Proxy (RealMonad h (IOSim s))+ -> LockstepAction (ModelState (IOSim s) h) a+ -> a+ -> Obs (IOSim s) h a+ observeReal _proxy (Action _ action') result = case action' of+ NewTableWith{} -> OEither $ bimap OId (const OTable) result+ CloseTable{} -> OEither $ bimap OId OId result+ Lookups{} -> OEither $+ bimap OId (OVector . fmap (OLookupResult . fmap (const OBlobRef))) result+ RangeLookup{} -> OEither $+ bimap OId (OVector . fmap (OEntry . fmap (const OBlobRef))) result+ NewCursor{} -> OEither $ bimap OId (const OCursor) result+ CloseCursor{} -> OEither $ bimap OId OId result+ ReadCursor{} -> OEither $+ bimap OId (OVector . fmap (OEntry . fmap (const OBlobRef))) result+ Updates{} -> OEither $ bimap OId OId result+ Inserts{} -> OEither $ bimap OId OId result+ Deletes{} -> OEither $ bimap OId OId result+ Upserts{} -> OEither $ bimap OId OId result+ RetrieveBlobs{} -> OEither $ bimap OId (OVector . fmap OBlob) result+ SaveSnapshot{} -> OEither $ bimap OId OId result+ OpenTableFromSnapshot{} -> OEither $ bimap OId (const OTable) result+ DeleteSnapshot{} -> OEither $ bimap OId OId result+ ListSnapshots{} -> OEither $ bimap OId (OList . fmap OId) result+ Duplicate{} -> OEither $ bimap OId (const OTable) result+ Union{} -> OEither $ bimap OId (const OTable) result+ Unions{} -> OEither $ bimap OId (const OTable) result+ RemainingUnionDebt{} -> OEither $ bimap OId OId result+ SupplyUnionCredits{} -> OEither $ bimap OId OUnionCredits result+ SupplyPortionOfDebt{} -> OEither $ bimap OId OUnionCreditsPortion result++ showRealResponse ::+ Proxy (RealMonad h (IOSim s))+ -> LockstepAction (ModelState (IOSim s) h) a+ -> Maybe (Dict (Show a))+ showRealResponse _ (Action _ action') = case action' of+ NewTableWith{} -> Nothing+ CloseTable{} -> Just Dict+ Lookups{} -> Nothing+ RangeLookup{} -> Nothing+ NewCursor{} -> Nothing+ CloseCursor{} -> Just Dict+ ReadCursor{} -> Nothing+ Updates{} -> Just Dict+ Inserts{} -> Just Dict+ Deletes{} -> Just Dict+ Upserts{} -> Just Dict+ RetrieveBlobs{} -> Just Dict+ SaveSnapshot{} -> Just Dict+ OpenTableFromSnapshot{} -> Nothing+ DeleteSnapshot{} -> Just Dict+ ListSnapshots{} -> Just Dict+ Duplicate{} -> Nothing+ Union{} -> Nothing+ Unions{} -> Nothing+ RemainingUnionDebt{} -> Just Dict+ SupplyUnionCredits{} -> Just Dict+ SupplyPortionOfDebt{} -> Just Dict++{-------------------------------------------------------------------------------+ RunModel+-------------------------------------------------------------------------------}++instance ( Eq (Class.TableConfig h)+ , Class.IsTable h+ , Show (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ ) => RunModel (Lockstep (ModelState IO h)) (RealMonad h IO) where+ perform _ = runIO+ postcondition = Lockstep.Defaults.postcondition+ monitoring = Lockstep.Defaults.monitoring (Proxy @(RealMonad h IO))++instance ( Eq (Class.TableConfig h)+ , Class.IsTable h+ , Show (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable s+ ) => RunModel (Lockstep (ModelState (IOSim s) h)) (RealMonad h (IOSim s)) where+ perform _ = runIOSim+ postcondition = Lockstep.Defaults.postcondition+ monitoring = Lockstep.Defaults.monitoring (Proxy @(RealMonad h (IOSim s)))++{-------------------------------------------------------------------------------+ Interpreter for the model+-------------------------------------------------------------------------------}++-- TODO: there are a bunch of TODO(err) in 'runModel' on the last argument to+-- 'Model.runModelMWithInjectedErrors'. This last argument defines how the model+-- should respond to injected errors. Since we don't generate injected errors+-- for most of these actions yet, they are left open. We will fill these in as+-- we start generating injected errors for these actions and testing with them.++runModel ::+ ModelLookUp (ModelState m h)+ -> LockstepAction (ModelState m h) a+ -> Model.Model -> (Val m h a, Model.Model)+runModel lookUp (Action merrs action') = case action' of+ NewTableWith _ _cfg ->+ wrap MTable+ . Model.runModelMWithInjectedErrors merrs+ (Model.new Model.TableConfig)+ (pure ()) -- TODO(err)+ CloseTable tableVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.close (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ Lookups ks tableVar ->+ wrap (MVector . fmap (MLookupResult . fmap MBlobRef))+ . Model.runModelMWithInjectedErrors merrs+ (Model.lookups ks (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ RangeLookup range tableVar ->+ wrap (MVector . fmap (MEntry . fmap MBlobRef))+ . Model.runModelMWithInjectedErrors merrs+ (Model.rangeLookup range (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ NewCursor offset tableVar ->+ wrap MCursor+ . Model.runModelMWithInjectedErrors merrs+ (Model.newCursor offset (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ CloseCursor cursorVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.closeCursor (getCursor $ lookUp cursorVar))+ (pure ()) -- TODO(err)+ ReadCursor n cursorVar ->+ wrap (MVector . fmap (MEntry . fmap MBlobRef))+ . Model.runModelMWithInjectedErrors merrs+ (Model.readCursor n (getCursor $ lookUp cursorVar))+ (pure ()) -- TODO(err)+ Updates kups tableVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.updates Model.getResolve kups (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ Inserts kins tableVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.inserts Model.getResolve kins (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ Deletes kdels tableVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.deletes Model.getResolve kdels (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ Upserts kmups tableVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.upserts Model.getResolve kmups (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ RetrieveBlobs blobsVar ->+ wrap (MVector . fmap (MBlob . WrapBlob))+ . Model.runModelMWithInjectedErrors merrs+ (Model.retrieveBlobs (getBlobRefs . lookUp $ blobsVar))+ (pure ()) -- TODO(err)+ SaveSnapshot mcorr name label tableVar ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (do Model.saveSnapshot name label (getTable $ lookUp tableVar)+ forM_ mcorr $ \_ -> Model.corruptSnapshot name)+ (pure ()) -- TODO(err)+ OpenTableFromSnapshot _ name label ->+ wrap MTable+ . Model.runModelMWithInjectedErrors merrs+ (Model.openTableFromSnapshot name label)+ (pure ())+ DeleteSnapshot name ->+ wrap MUnit+ . Model.runModelMWithInjectedErrors merrs+ (Model.deleteSnapshot name)+ (pure ()) -- TODO(err)+ ListSnapshots ->+ wrap (MList . fmap MSnapshotName)+ . Model.runModelMWithInjectedErrors merrs+ Model.listSnapshots+ (pure ()) -- TODO(err)+ Duplicate tableVar ->+ wrap MTable+ . Model.runModelMWithInjectedErrors merrs+ (Model.duplicate (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ Union table1Var table2Var ->+ wrap MTable+ . Model.runModelMWithInjectedErrors merrs+ (Model.union Model.getResolve (getTable $ lookUp table1Var) (getTable $ lookUp table2Var))+ (pure ()) -- TODO(err)+ Unions tableVars ->+ wrap MTable+ . Model.runModelMWithInjectedErrors merrs+ (Model.unions Model.getResolve (fmap (getTable . lookUp) tableVars))+ (pure ()) -- TODO(err)+ RemainingUnionDebt tableVar ->+ wrap MUnionDebt+ . Model.runModelMWithInjectedErrors merrs+ (Model.remainingUnionDebt (getTable $ lookUp tableVar))+ (pure ()) -- TODO(err)+ SupplyUnionCredits tableVar credits ->+ wrap MUnionCredits+ . Model.runModelMWithInjectedErrors merrs+ (Model.supplyUnionCredits (getTable $ lookUp tableVar) credits)+ (pure ()) -- TODO(err)+ SupplyPortionOfDebt tableVar portion ->+ wrap MUnionCreditsPortion+ . Model.runModelMWithInjectedErrors merrs+ (do let table = getTable $ lookUp tableVar+ Model.supplyPortionOfDebt table portion+ )+ (pure ()) -- TODO(err)+ where+ getTable ::+ ModelValue (ModelState m h) (WrapTable h m k v b)+ -> Model.Table k v b+ getTable (MTable t) = t++ getCursor ::+ ModelValue (ModelState m h) (WrapCursor h m k v b)+ -> Model.Cursor k v b+ getCursor (MCursor t) = t++ getBlobRefs :: ModelValue (ModelState m h) (V.Vector (WrapBlobRef h m b)) -> V.Vector (Model.BlobRef b)+ getBlobRefs (MVector brs) = fmap (\(MBlobRef br) -> br) brs++wrap ::+ (a -> Val m h b)+ -> (Either Model.Err a, Model.Model)+ -> (Val m h (Either Model.Err b), Model.Model)+wrap f = first (MEither . bimap MErr f)++{-------------------------------------------------------------------------------+ Interpreters for @'IOLike' m@+-------------------------------------------------------------------------------}++-- TODO: there are a bunch of TODO(err) in 'runIO' and 'runIOSim' below on the+-- last argument to 'Model.runModelMWithInjectedErrors'. This last argument+-- defines how the SUT should recover from actions that accidentally succeeded+-- in the presence of disk faults. Since we don't generate injected errors for+-- most of these actions yet, they are left open. We will fill these in as we+-- start generating injected errors for these actions and testing with them.++runIO ::+ forall a h. Class.IsTable h+ => LockstepAction (ModelState IO h) a+ -> LookUp+ -> RealMonad h IO a+runIO action lookUp = ReaderT $ \ !env -> do+ aux env action+ where+ aux ::+ RealEnv h IO+ -> LockstepAction (ModelState IO h) a+ -> IO a+ aux env (Action merrs action') = case action' of+ NewTableWith _ cfg ->+ runRealWithInjectedErrors "NewTableWith" env merrs+ (WrapTable <$> Class.newTableWith cfg session)+ (\_ -> pure ()) -- TODO(err)+ CloseTable tableVar ->+ runRealWithInjectedErrors "CloseTable" env merrs+ (Class.closeTable (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ Lookups ks tableVar ->+ runRealWithInjectedErrors "Lookups" env merrs+ (fmap (fmap WrapBlobRef) <$> Class.lookups (unwrapTable $ lookUp' tableVar) ks)+ (\_ -> pure ()) -- TODO(err)+ RangeLookup range tableVar ->+ runRealWithInjectedErrors "RangeLookup" env merrs+ (fmap (fmap WrapBlobRef) <$> Class.rangeLookup (unwrapTable $ lookUp' tableVar) range)+ (\_ -> pure ()) -- TODO(err)+ NewCursor offset tableVar ->+ runRealWithInjectedErrors "NewCursor" env merrs+ (WrapCursor <$> Class.newCursor offset (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ CloseCursor cursorVar ->+ runRealWithInjectedErrors "CloseCursor" env merrs+ (Class.closeCursor (Proxy @h) (unwrapCursor $ lookUp' cursorVar))+ (\_ -> pure ()) -- TODO(err)+ ReadCursor n cursorVar ->+ runRealWithInjectedErrors "ReadCursor" env merrs+ (fmap (fmap WrapBlobRef) <$> Class.readCursor (Proxy @h) n (unwrapCursor $ lookUp' cursorVar))+ (\_ -> pure ()) -- TODO(err)+ Updates kups tableVar ->+ runRealWithInjectedErrors "Updates" env merrs+ (Class.updates (unwrapTable $ lookUp' tableVar) kups)+ (\_ -> pure ()) -- TODO(err)+ Inserts kins tableVar ->+ runRealWithInjectedErrors "Inserts" env merrs+ (Class.inserts (unwrapTable $ lookUp' tableVar) kins)+ (\_ -> pure ()) -- TODO(err)+ Deletes kdels tableVar ->+ runRealWithInjectedErrors "Deletes" env merrs+ (Class.deletes (unwrapTable $ lookUp' tableVar) kdels)+ (\_ -> pure ()) -- TODO(err)+ Upserts kmups tableVar ->+ runRealWithInjectedErrors "Upserts" env merrs+ (Class.upserts (unwrapTable $ lookUp' tableVar) kmups)+ (\_ -> pure ()) -- TODO(err)+ RetrieveBlobs blobRefsVar ->+ runRealWithInjectedErrors "RetrieveBlobs" env merrs+ (fmap WrapBlob <$> Class.retrieveBlobs (Proxy @h) session (unwrapBlobRef <$> lookUp' blobRefsVar))+ (\_ -> pure ()) -- TODO(err)+ SaveSnapshot mcorr name label tableVar ->+ let table = unwrapTable $ lookUp' tableVar in+ runRealWithInjectedErrors "SaveSnapshot" env merrs+ (do Class.saveSnapshot name label table+ forM_ mcorr $ \corr -> Class.corruptSnapshot (bitChoice corr) name table)+ (\() -> Class.deleteSnapshot session name) -- TODO(err)+ OpenTableFromSnapshot _ name label ->+ runRealWithInjectedErrors "OpenTableFromSnapshot" env merrs+ (WrapTable <$> Class.openTableFromSnapshot session name label)+ (\(WrapTable t) -> Class.closeTable t)+ DeleteSnapshot name ->+ runRealWithInjectedErrors "DeleteSnapshot" env merrs+ (Class.deleteSnapshot session name)+ (\_ -> pure ()) -- TODO(err)+ ListSnapshots ->+ runRealWithInjectedErrors "ListSnapshots" env merrs+ (Class.listSnapshots session)+ (\_ -> pure ()) -- TODO(err)+ Duplicate tableVar ->+ runRealWithInjectedErrors "Duplicate" env merrs+ (WrapTable <$> Class.duplicate (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ Union table1Var table2Var ->+ runRealWithInjectedErrors "Union" env merrs+ (WrapTable <$> Class.union (unwrapTable $ lookUp' table1Var) (unwrapTable $ lookUp' table2Var))+ (\_ -> pure ()) -- TODO(err)+ Unions tableVars ->+ runRealWithInjectedErrors "Unions" env merrs+ (WrapTable <$> Class.unions (fmap (unwrapTable . lookUp') tableVars))+ (\_ -> pure ()) -- TODO(err)+ RemainingUnionDebt tableVar ->+ runRealWithInjectedErrors "RemainingUnionDebt" env merrs+ (Class.remainingUnionDebt (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ SupplyUnionCredits tableVar credits ->+ runRealWithInjectedErrors "SupplyUnionCredits" env merrs+ (Class.supplyUnionCredits (unwrapTable $ lookUp' tableVar) credits)+ (\_ -> pure ()) -- TODO(err)+ SupplyPortionOfDebt tableVar portion ->+ runRealWithInjectedErrors "SupplyPortionOfDebt" env merrs+ (do let table = unwrapTable $ lookUp' tableVar+ debt <- Class.remainingUnionDebt table+ Class.supplyUnionCredits table (portion `portionOf` debt))+ (\_ -> pure ()) -- TODO(err)+ where+ session = envSession env++ lookUp' :: Var m h x -> x+ lookUp' = realLookupVar lookUp++runIOSim ::+ forall s a h. Class.IsTable h+ => LockstepAction (ModelState (IOSim s) h) a+ -> LookUp+ -> RealMonad h (IOSim s) a+runIOSim action lookUp = ReaderT $ \ !env -> do+ aux env action+ where+ aux ::+ RealEnv h (IOSim s)+ -> LockstepAction (ModelState (IOSim s) h) a+ -> IOSim s a+ aux env (Action merrs action') = case action' of+ NewTableWith _ cfg ->+ runRealWithInjectedErrors "NewTableWith" env merrs+ (WrapTable <$> Class.newTableWith cfg session)+ (\_ -> pure ()) -- TODO(err)+ CloseTable tableVar ->+ runRealWithInjectedErrors "CloseTable" env merrs+ (Class.closeTable (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ Lookups ks tableVar ->+ runRealWithInjectedErrors "Lookups" env merrs+ (fmap (fmap WrapBlobRef) <$> Class.lookups (unwrapTable $ lookUp' tableVar) ks)+ (\_ -> pure ()) -- TODO(err)+ RangeLookup range tableVar ->+ runRealWithInjectedErrors "RangeLookup" env merrs+ (fmap (fmap WrapBlobRef) <$> Class.rangeLookup (unwrapTable $ lookUp' tableVar) range)+ (\_ -> pure ()) -- TODO(err)+ NewCursor offset tableVar ->+ runRealWithInjectedErrors "NewCursor" env merrs+ (WrapCursor <$> Class.newCursor offset (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ CloseCursor cursorVar ->+ runRealWithInjectedErrors "CloseCursor" env merrs+ (Class.closeCursor (Proxy @h) (unwrapCursor $ lookUp' cursorVar))+ (\_ -> pure ()) -- TODO(err)+ ReadCursor n cursorVar ->+ runRealWithInjectedErrors "ReadCursor" env merrs+ (fmap (fmap WrapBlobRef) <$> Class.readCursor (Proxy @h) n (unwrapCursor $ lookUp' cursorVar))+ (\_ -> pure ()) -- TODO(err)+ Updates kups tableVar ->+ runRealWithInjectedErrors "Updates" env merrs+ (Class.updates (unwrapTable $ lookUp' tableVar) kups)+ (\_ -> pure ()) -- TODO(err)+ Inserts kins tableVar ->+ runRealWithInjectedErrors "Inserts" env merrs+ (Class.inserts (unwrapTable $ lookUp' tableVar) kins)+ (\_ -> pure ()) -- TODO(err)+ Deletes kdels tableVar ->+ runRealWithInjectedErrors "Deletes" env merrs+ (Class.deletes (unwrapTable $ lookUp' tableVar) kdels)+ (\_ -> pure ()) -- TODO(err)+ Upserts kmups tableVar ->+ runRealWithInjectedErrors "Upserts" env merrs+ (Class.upserts (unwrapTable $ lookUp' tableVar) kmups)+ (\_ -> pure ()) -- TODO(err)+ RetrieveBlobs blobRefsVar ->+ runRealWithInjectedErrors "RetrieveBlobs" env merrs+ (fmap WrapBlob <$> Class.retrieveBlobs (Proxy @h) session (unwrapBlobRef <$> lookUp' blobRefsVar))+ (\_ -> pure ()) -- TODO(err)+ SaveSnapshot mcorr name label tableVar ->+ let table = unwrapTable $ lookUp' tableVar in+ runRealWithInjectedErrors "SaveSnapshot" env merrs+ (do Class.saveSnapshot name label table+ forM_ mcorr $ \corr -> Class.corruptSnapshot (bitChoice corr) name table)+ (\() -> Class.deleteSnapshot session name) -- TODO(err)+ OpenTableFromSnapshot _ name label ->+ runRealWithInjectedErrors "OpenTableFromSnapshot" env merrs+ (WrapTable <$> Class.openTableFromSnapshot session name label)+ (\(WrapTable t) -> Class.closeTable t)+ DeleteSnapshot name ->+ runRealWithInjectedErrors "DeleteSnapshot" env merrs+ (Class.deleteSnapshot session name)+ (\_ -> pure ()) -- TODO(err)+ ListSnapshots ->+ runRealWithInjectedErrors "ListSnapshots" env merrs+ (Class.listSnapshots session)+ (\_ -> pure ()) -- TODO(err)+ Duplicate tableVar ->+ runRealWithInjectedErrors "Duplicate" env merrs+ (WrapTable <$> Class.duplicate (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ Union table1Var table2Var ->+ runRealWithInjectedErrors "Union" env merrs+ (WrapTable <$> Class.union (unwrapTable $ lookUp' table1Var) (unwrapTable $ lookUp' table2Var))+ (\_ -> pure ()) -- TODO(err)+ Unions tableVars ->+ runRealWithInjectedErrors "Unions" env merrs+ (WrapTable <$> Class.unions (fmap (unwrapTable . lookUp') tableVars))+ (\_ -> pure ()) -- TODO(err)+ RemainingUnionDebt tableVar ->+ runRealWithInjectedErrors "RemainingUnionDebt" env merrs+ (Class.remainingUnionDebt (unwrapTable $ lookUp' tableVar))+ (\_ -> pure ()) -- TODO(err)+ SupplyUnionCredits tableVar credits ->+ runRealWithInjectedErrors "SupplyUnionCredits" env merrs+ (Class.supplyUnionCredits (unwrapTable $ lookUp' tableVar) credits)+ (\_ -> pure ()) -- TODO(err)+ SupplyPortionOfDebt tableVar portion ->+ runRealWithInjectedErrors "SupplyPortionOfDebt" env merrs+ (do let table = unwrapTable $ lookUp' tableVar+ debt <- Class.remainingUnionDebt table+ Class.supplyUnionCredits table (portion `portionOf` debt))+ (\_ -> pure ()) -- TODO(err)+ where+ session = envSession env++ lookUp' :: Var m h x -> x+ lookUp' = realLookupVar lookUp++-- | @'runRealWithInjectedErrors' _ errsVar merrs action rollback@ runs @action@+-- with injected errors if available in @merrs@.+--+-- See 'Model.runModelMWithInjectedErrors': the real system is not guaranteed to+-- fail with an error if there are injected disk faults, but the model *always*+-- fails. In case the real system accidentally succeeded in running @action@+-- when there were errors to inject, then we roll back the success using+-- @rollback@. This ensures that the model and system stay in sync. For example:+-- if creating a snapshot accidentally succeeded, then the rollback action is to+-- delete that snapshot.+runRealWithInjectedErrors ::+ (MonadCatch m, MonadSTM m, PrimMonad m)+ => String -- ^ Name of the action+ -> RealEnv h m+ -> Maybe Errors+ -> m t-- ^ Action to run+ -> (t -> m ()) -- ^ Rollback if the action *accidentally* succeeded+ -> m (Either Model.Err t)+runRealWithInjectedErrors s env merrs k rollback =+ case merrs of+ Nothing -> do+ modifyMutVar faultsVar (InjectFaultNone s :)+ catchErr handlers k+ Just errs -> do+ atomically $ writeTVar logVar emptyLog+ eith <- catchErr handlers $ FSSim.withErrors errsVar errs k+ errsLog <- readTVarIO logVar+ case eith of+ Left e | Model.isDiskFault e -> do+ modifyMutVar faultsVar (InjectFaultInducedError s :)+ if countNoisyErrors errsLog == 0 then+ pure $ Left $ Model.ErrOther $+ -- If we injected 0 disk faults, but we still found an+ -- ErrDiskFault, then there is a bug in our code. ErrDiskFaults+ -- should not occur on the happy path.+ "Found a disk fault error, but no disk faults were injected: " <> show e+ else+ pure eith+ Left e -> do+ if countNoisyErrors errsLog > 0 then+ pure $ Left $ Model.ErrOther $+ -- If we injected 1 or more disk faults, but we did not find an+ -- ErrDiskFault, then there is a bug in our code. An injected disk+ -- fault should always lead to an ErrDiskFault.+ "Found an error that isn't a disk fault error, but disk faults were injected: " <> show e+ else+ pure eith+ Right x -> do+ modifyMutVar faultsVar (InjectFaultAccidentalSuccess s :)+ rollback x+ if (countNoisyErrors errsLog > 0) then+ pure $ Left $ Model.ErrOther $+ -- If we injected 1 or more disk faults, but the action+ -- accidentally succeeded, then 1 or more errors were swallowed+ -- that should have been found as ErrDiskFault.+ "Action succeeded, but disk faults were injected. Errors were swallowed!"+ else+ pure $ Left $ Model.ErrDiskFault ("dummy: " <> s)+ where+ errsVar = envErrors env+ logVar = envErrorsLog env+ faultsVar = envInjectFaultResults env+ handlers = envHandlers env++catchErr ::+ forall m a e. MonadCatch m+ => [Handler m (Maybe e)] -> m a -> m (Either e a)+catchErr hs action = catches (Right <$> action) (fmap f hs)+ where+ f (Handler h) = Handler $ \e -> maybe (throwIO e) (pure . Left) =<< h e++{-------------------------------------------------------------------------------+ Generator and shrinking+-------------------------------------------------------------------------------}++arbitraryActionWithVars ::+ forall m h k v b. (+ C k v b+ , Ord k+ , Eq (Class.TableConfig h)+ , Show (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable m+ )+ => Proxy (k, v, b)+ -> R.SnapshotLabel+ -> ModelVarContext (ModelState m h)+ -> ModelState m h+ -> Gen (Any (LockstepAction (ModelState m h)))+arbitraryActionWithVars _ label ctx (ModelState st _stats) =+ QC.frequency $+ concat+ [ genActionsSession+ , genActionsTables+ , genActionsUnion+ , genActionsCursor+ , genActionsBlobRef+ ]+ where+ _coveredAllCases :: LockstepAction (ModelState m h) a -> ()+ _coveredAllCases (Action _ action') = _coveredAllCases' action'++ _coveredAllCases' :: Action' m h a -> ()+ _coveredAllCases' = \case+ NewTableWith{} -> ()+ CloseTable{} -> ()+ Lookups{} -> ()+ RangeLookup{} -> ()+ NewCursor{} -> ()+ CloseCursor{} -> ()+ ReadCursor{} -> ()+ Updates{} -> ()+ Inserts{} -> ()+ Deletes{} -> ()+ Upserts{} -> ()+ RetrieveBlobs{} -> ()+ SaveSnapshot{} -> ()+ DeleteSnapshot{} -> ()+ ListSnapshots{} -> ()+ OpenTableFromSnapshot{} -> ()+ Duplicate{} -> ()+ Union{} -> ()+ Unions{} -> ()+ RemainingUnionDebt{} -> ()+ SupplyUnionCredits{} -> ()+ SupplyPortionOfDebt{} -> ()++ genTableVar = QC.elements tableVars++ tableVars :: [Var m h (WrapTable h m k v b)]+ tableVars =+ [ fromRight v+ | v <- findVars ctx Proxy+ , case lookupVar ctx v of+ MEither (Left _) -> False+ MEither (Right (MTable t)) ->+ Map.member (Model.tableID t) (Model.tables st)+ ]++ genUnionDescendantTableVar = QC.elements unionDescendantTableVars++ unionDescendantTableVars :: [Var m h (WrapTable h m k v b)]+ unionDescendantTableVars =+ [ v+ | v <- tableVars+ , let t = case lookupVar ctx v of+ MTable t' -> t'+ , Model.isUnionDescendant t == Model.IsUnionDescendant+ ]++ genCursorVar = QC.elements cursorVars++ cursorVars :: [Var m h (WrapCursor h m k v b)]+ cursorVars =+ [ fromRight v+ | v <- findVars ctx Proxy+ , case lookupVar ctx v of+ MEither (Left _) -> False+ MEither (Right (MCursor c)) ->+ Map.member (Model.cursorID c) (Model.cursors st)+ ]++ genBlobRefsVar = QC.elements blobRefsVars++ blobRefsVars :: [Var m h (V.Vector (WrapBlobRef h m b))]+ blobRefsVars = fmap (mapGVar (OpComp OpLookupResults)) lookupResultVars+ ++ fmap (mapGVar (OpComp OpEntrys)) queryResultVars+ where+ lookupResultVars :: [Var m h (V.Vector (LookupResult v (WrapBlobRef h m b)))]+ queryResultVars :: [Var m h (V.Vector (Entry k v (WrapBlobRef h m b)))]++ lookupResultVars = fromRight <$> findVars ctx Proxy+ queryResultVars = fromRight <$> findVars ctx Proxy++ genUsedSnapshotName = QC.elements usedSnapshotNames+ genUnusedSnapshotName = QC.elements unusedSnapshotNames++ usedSnapshotNames, unusedSnapshotNames :: [R.SnapshotName]+ (usedSnapshotNames, unusedSnapshotNames) =+ partitionEithers+ [ if Map.member snapshotname (Model.snapshots st)+ then Left snapshotname -- used+ else Right snapshotname -- unused+ | name <- ["snap1", "snap2", "snap3" ]+ , let snapshotname = R.toSnapshotName name+ ]++ genActionsSession :: [(Int, Gen (Any (LockstepAction (ModelState m h))))]+ genActionsSession =+ [ (1, fmap Some $ (Action <$> genErrors <*>) $+ NewTableWith @k @v @b <$> pure PrettyProxy <*> QC.arbitrary)+ | length tableVars <= 5 -- no more than 5 tables at once+ , let genErrors = pure Nothing -- TODO: generate errors+ ]++ ++ [ (2, fmap Some $ (Action <$> genErrors <*>) $+ OpenTableFromSnapshot @k @v @b PrettyProxy <$> genUsedSnapshotName <*> pure label)+ | length tableVars <= 5 -- no more than 5 tables at once+ , not (null usedSnapshotNames)+ , -- TODO: bloomFilterFromFile swallows an FsReachedEOF error+ let isFsReachedEOFError = maybe False (either isFsReachedEOF (const False))+ , let genErrors = QC.frequency [+ (3, pure Nothing)+ , (1, Just . filterHGetBufSomeE (not . isFsReachedEOFError)+ . noRemoveDirectoryRecursiveE+ <$> QC.arbitrary)+ ]+ ]++ ++ [ (1, fmap Some $ (Action <$> genErrors <*>) $+ DeleteSnapshot <$> genUsedSnapshotName)+ | not (null usedSnapshotNames)+ , let genErrors = pure Nothing -- TODO: generate errors+ ]++ ++ [ (1, fmap Some $ (Action <$> genErrors <*>) $+ pure ListSnapshots)+ | not (null tableVars) -- otherwise boring!+ , let genErrors = pure Nothing -- TODO: generate errors+ ]++ genActionsTables :: [(Int, Gen (Any (LockstepAction (ModelState m h))))]+ genActionsTables+ | null tableVars = []+ | otherwise =+ [ (1, fmap Some $ (Action <$> genErrors <*>) $+ CloseTable <$> genTableVar)+ | let genErrors = pure Nothing+ ]+ ++ [ (10, fmap Some $ (Action <$> genErrors <*>) $+ Lookups <$> genLookupKeys <*> genTableVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (5, fmap Some $ (Action <$> genErrors <*>) $+ RangeLookup <$> genRange <*> genTableVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (10, fmap Some $ (Action <$> genErrors <*>) $+ Updates <$> genUpdates <*> genTableVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (10, fmap Some $ (Action <$> genErrors <*>) $+ Inserts <$> genInserts <*> genTableVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (10, fmap Some $ (Action <$> genErrors <*>) $+ Deletes <$> genDeletes <*> genTableVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (10, fmap Some $ (Action <$> genErrors <*>) $+ Upserts <$> genMupserts <*> genTableVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (3, fmap Some $ (Action <$> genErrors <*>) $+ NewCursor <$> QC.arbitrary <*> genTableVar)+ | length cursorVars <= 5 -- no more than 5 cursors at once+ , let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (2, fmap Some $ (Action <$> genErrors <*>) $+ SaveSnapshot <$> genCorruption <*> genUnusedSnapshotName <*> pure label <*> genTableVar)+ | not (null unusedSnapshotNames)+ -- TODO: should errors and corruption be generated at the same time,+ -- or should they be mutually exclusive?+ , let genErrors = pure Nothing -- TODO: generate errors+ , let genCorruption = QC.frequency [+ (3, pure Nothing)+ , (1, Just <$> QC.arbitrary)+ ]+ ]+ ++ [ (5, fmap Some $ (Action <$> genErrors <*>) $+ Duplicate <$> genTableVar)+ | length tableVars <= 5 -- no more than 5 tables at once+ , let genErrors = pure Nothing -- TODO: generate errors+ ]++ -- | Generate table actions that have to do with unions.+ genActionsUnion :: [(Int, Gen (Any (LockstepAction (ModelState m h))))]+ genActionsUnion+ | null tableVars = []+ | otherwise =+ [ (2, fmap Some $ (Action <$> genErrors <*>) $+ Union <$> genTableVar <*> genTableVar)+ | length tableVars <= 5 -- no more than 5 tables at once+ , let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (2, fmap Some $ (Action <$> genErrors <*>) $ do+ -- Generate at least a 2-way union, and at most a 3-way union.+ --+ -- Tests for 1-way unions are included in the UnitTests module.+ -- n-way unions for n>3 lead to large unions, which are less+ -- likely to be finished before the end of an action sequence.+ n <- QC.chooseInt (2, 3)+ Unions . NE.fromList <$> QC.vectorOf n genTableVar)+ | length tableVars <= 5 -- no more than 5 tables at once+ , let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (2, fmap Some $ (Action <$> genErrors <*>) $+ RemainingUnionDebt <$> genUnionDescendantTableVar)+ -- Tables not derived from unions are covered in UnitTests.+ | not (null unionDescendantTableVars)+ , let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (8, fmap Some $ (Action <$> genErrors <*>) $+ SupplyUnionCredits <$> genUnionDescendantTableVar <*> genUnionCredits)+ -- Tables not derived from unions are covered in UnitTests.+ | not (null unionDescendantTableVars)+ , let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (2, fmap Some $ (Action <$> genErrors <*>) $+ SupplyPortionOfDebt <$> genUnionDescendantTableVar <*> genPortion)+ -- Tables not derived from unions are covered in UnitTests.+ | not (null unionDescendantTableVars)+ , let genErrors = pure Nothing -- TODO: generate errors+ ]+ where+ -- The typical, interesting case is to supply a positive number of+ -- union credits. Supplying 0 or less credits is a no-op. We cover+ -- it in UnitTests so we don't have to cover it here.+ genUnionCredits = R.UnionCredits . QC.getPositive <$> QC.arbitrary++ -- TODO: tweak distribution once table unions are implemented+ genPortion = Portion <$> QC.elements [1, 2, 3]++ genActionsCursor :: [(Int, Gen (Any (LockstepAction (ModelState m h))))]+ genActionsCursor+ | null cursorVars = []+ | otherwise =+ [ (2, fmap Some $ (Action <$> genErrors <*>) $+ CloseCursor <$> genCursorVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]+ ++ [ (10, fmap Some $ (Action <$> genErrors <*>) $+ ReadCursor <$> (QC.getNonNegative <$> QC.arbitrary) <*> genCursorVar)+ | let genErrors = pure Nothing -- TODO: generate errors+ ]++ genActionsBlobRef :: [(Int, Gen (Any (LockstepAction (ModelState m h))))]+ genActionsBlobRef =+ [ (5, fmap Some $ (Action <$> genErrors <*>) $+ RetrieveBlobs <$> genBlobRefsVar)+ | not (null blobRefsVars)+ , let genErrors = pure Nothing -- TODO: generate errors+ ]++ fromRight ::+ Var m h (Either Model.Err a)+ -> Var m h a+ fromRight = mapGVar (\op -> OpFromRight `OpComp` op)++ genLookupKeys :: Gen (V.Vector k)+ genLookupKeys = QC.arbitrary++ genRange :: Gen (R.Range k)+ genRange = QC.arbitrary++ genUpdates :: Gen (V.Vector (k, R.Update v b))+ genUpdates = QC.liftArbitrary ((,) <$> QC.arbitrary <*> QC.oneof [+ R.Insert <$> QC.arbitrary <*> genBlob+ , R.Upsert <$> QC.arbitrary+ , pure R.Delete+ ])+ where+ _coveredAllCases :: R.Update v b -> ()+ _coveredAllCases = \case+ R.Insert{} -> ()+ R.Upsert{} -> ()+ R.Delete{} -> ()++ genInserts :: Gen (V.Vector (k, v, Maybe b))+ genInserts = QC.liftArbitrary ((,,) <$> QC.arbitrary <*> QC.arbitrary <*> genBlob)++ genDeletes :: Gen (V.Vector k)+ genDeletes = QC.arbitrary++ genMupserts :: Gen (V.Vector (k, v))+ genMupserts = QC.liftArbitrary ((,) <$> QC.arbitrary <*> QC.arbitrary)++ genBlob :: Gen (Maybe b)+ genBlob = QC.arbitrary++shrinkActionWithVars ::+ forall m h a. (+ Eq (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable m+ )+ => ModelVarContext (ModelState m h)+ -> ModelState m h+ -> LockstepAction (ModelState m h) a+ -> [Any (LockstepAction (ModelState m h))]+shrinkActionWithVars _ctx _st (Action merrs action') =+ [ -- TODO: it's somewhat unfortunate, but we have to dynamically+ -- construct evidence that @a@ is typeable, which is a requirement+ -- coming from the @Some@ existential. Could we find a different way+ -- to solve this using just regular constraints?+ case dictIsTypeable action' of+ Dict -> Some $ Action merrs' action'+ | merrs' <- QC.shrink merrs ]+ ++ [ Some $ Action merrs action''+ | Some action'' <- shrinkAction'WithVars _ctx _st action'+ ]++-- | Dynamically construct evidence that the result type @a@ of an action is+-- typeable.+dictIsTypeable :: (Typeable m, Typeable h) => Action' m h a -> Dict (Typeable a)+dictIsTypeable = \case+ NewTableWith{} -> Dict+ CloseTable{} -> Dict+ Lookups{} -> Dict+ RangeLookup{} -> Dict+ NewCursor{} -> Dict+ CloseCursor{} -> Dict+ ReadCursor{} -> Dict+ Updates{} -> Dict+ Inserts{} -> Dict+ Deletes{} -> Dict+ Upserts{} -> Dict+ RetrieveBlobs{} -> Dict+ SaveSnapshot{} -> Dict+ OpenTableFromSnapshot{} -> Dict+ DeleteSnapshot{} -> Dict+ ListSnapshots{} -> Dict+ Duplicate{} -> Dict+ Union{} -> Dict+ Unions{} -> Dict+ RemainingUnionDebt{} -> Dict+ SupplyUnionCredits{} -> Dict+ SupplyPortionOfDebt{} -> Dict++shrinkAction'WithVars ::+ forall m h a. (+ Eq (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable m+ )+ => ModelVarContext (ModelState m h)+ -> ModelState m h+ -> Action' m h a+ -> [Any (Action' m h)]+shrinkAction'WithVars _ctx _st a = case a of+ NewTableWith p conf -> [+ Some $ NewTableWith p conf'+ | conf' <- QC.shrink conf+ ]++ -- Shrink inserts and deletes towards updates.+ Updates upds tableVar -> [+ Some $ Updates upds' tableVar+ | upds' <- QC.shrink upds+ ]+ Inserts kvbs tableVar -> [+ Some $ Inserts kvbs' tableVar+ | kvbs' <- QC.shrink kvbs+ ] <> [+ Some $ Updates (V.map f kvbs) tableVar+ | let f (k, v, mb) = (k, R.Insert v mb)+ ]+ Deletes ks tableVar -> [+ Some $ Deletes ks' tableVar+ | ks' <- QC.shrink ks+ ] <> [+ Some $ Updates (V.map f ks) tableVar+ | let f k = (k, R.Delete)+ ]++ Upserts kvs tableVar -> [+ Some $ Upserts kvs' tableVar+ | kvs' <- QC.shrink kvs+ ] <> [+ Some $ Updates (V.map f kvs) tableVar+ | let f (k, v) = (k, R.Upsert v)+ ]++ Lookups ks tableVar -> [+ Some $ Lookups ks' tableVar+ | ks' <- QC.shrink ks+ ]++ -- * Unions++ Unions tableVars -> [+ Some $ Unions tableVars'+ | tableVars' <- QC.liftShrink (const []) tableVars+ ]++ SupplyUnionCredits tableVar (R.UnionCredits x) -> [+ Some $ SupplyUnionCredits tableVar (R.UnionCredits x')+ | x' <- QC.shrink x+ ]++ -- Shrink portions to absolute union credits. The credits are all /positive/+ -- powers of 2 that fit into a 64-bit Int. This choice is arbitrary, but the+ -- model has no accurate debt information, so the best we can do is provide+ -- a sensible distribution of union credits. This particular distribution is+ -- skewed to smaller numbers, because the union debt is also likely to be on+ -- the smaller side. Still, all larger powers of 2 are also included in case+ -- the union debt is larger.+ SupplyPortionOfDebt tableVar (Portion x) -> [+ Some $ SupplyUnionCredits tableVar (R.UnionCredits x')+ | x' <- [2 ^ i - 1 | (i :: Int) <- [0..20]]+ , assert (x' >= 0) True+ ] ++ [+ Some $ SupplyPortionOfDebt tableVar (Portion x')+ | x' <- QC.shrink x+ , x' > 0+ ]++ _ -> []++{-------------------------------------------------------------------------------+ Interpret 'Op' against 'ModelValue'+-------------------------------------------------------------------------------}++instance InterpretOp Op (ModelValue (ModelState m h)) where+ intOp = \case+ OpId -> Just+ OpFst -> \case MPair x -> Just (fst x)+ OpSnd -> \case MPair x -> Just (snd x)+ OpLeft -> Just . MEither . Left+ OpRight -> Just . MEither . Right+ OpFromLeft -> \case MEither x -> either Just (const Nothing) x+ OpFromRight -> \case MEither x -> either (const Nothing) Just x+ OpComp g f -> intOp g <=< intOp f+ OpLookupResults -> Just . MVector+ . V.mapMaybe (\case MLookupResult x -> getBlobRef x)+ . \case MVector x -> x+ OpEntrys -> Just . MVector+ . V.mapMaybe (\case MEntry x -> getBlobRef x)+ . \case MVector x -> x++{-------------------------------------------------------------------------------+ Statistics, labelling/tagging+-------------------------------------------------------------------------------}++data Stats = Stats {+ -- === Tags+ -- | Names for which snapshots exist+ snapshotted :: !(Set R.SnapshotName)+ -- === Final tags (per action sequence, across all tables)+ -- | Number of successful lookups and their results+ , numLookupsResults :: {-# UNPACK #-} !(Int, Int, Int)+ -- (NotFound, Found, FoundWithBlob)+ -- | Number of successful updates+ , numUpdates :: {-# UNPACK #-} !(Int, Int, Int, Int)+ -- (Insert, InsertWithBlob, Delete, Mupsert)+ -- | Actions that succeeded+ , successActions :: [String]+ -- | Actions that failed with an error+ , failActions :: [(String, Model.Err)]+ -- === Final tags (per action sequence, per table)+ -- | Number of actions per table (successful or failing)+ , numActionsPerTable :: !(Map Model.TableID Int)+ -- | The state of model tables at the point they were closed. This is used+ -- to augment the tables from the final model state (which of course has+ -- only tables still open in the final state).+ , closedTables :: !(Map Model.TableID Model.SomeTable)+ -- | The ultimate parents for each table. These are the 'TableId's of tables+ -- created using 'new' or 'open'.+ , parentTable :: !(Map Model.TableID [Model.TableID])+ -- | Track the interleavings of operations via different but related tables.+ -- This is a map from each ultimate parent table to a summary log of which+ -- tables (derived from that parent table via duplicate or union) have had+ -- \"interesting\" actions performed on them. We record only the+ -- interleavings of different tables not multiple actions on the same table.+ , dupTableActionLog :: !(Map Model.TableID [Model.TableID])+ -- | The subset of tables (open or closed) that were created as a result+ -- of a union operation. This can be used for example to select subsets of+ -- the other per-table tracking maps above, or the state from the model.+ -- The map value is the size of the union table at the point it was created,+ -- so we can distinguish trivial empty unions from non-trivial.+ , unionTables :: !(Map Model.TableID Int)+ }+ deriving stock Show++initStats :: Stats+initStats = Stats {+ -- === Tags+ snapshotted = Set.empty+ -- === Final tags+ , numLookupsResults = (0, 0, 0)+ , numUpdates = (0, 0, 0, 0)+ , successActions = []+ , failActions = []+ , numActionsPerTable = Map.empty+ , closedTables = Map.empty+ , parentTable = Map.empty+ , dupTableActionLog = Map.empty+ , unionTables = Map.empty+ }++updateStats ::+ forall m h a. ( Show (Class.TableConfig h)+ , Eq (Class.TableConfig h)+ , Arbitrary (Class.TableConfig h)+ , Typeable h+ , Typeable m+ )+ => LockstepAction (ModelState m h) a+ -> ModelLookUp (ModelState m h)+ -> Model.Model+ -> Model.Model+ -> Val m h a+ -> Stats+ -> Stats+updateStats action@(Action _merrs action') lookUp modelBefore modelAfter result =+ -- === Tags+ updSnapshotted+ -- === Final tags+ . updNumLookupsResults+ . updNumUpdates+ . updSuccessActions+ . updFailActions+ . updNumActionsPerTable+ . updClosedTables+ . updDupTableActionLog+ . updParentTable+ . updUnionTables+ where+ -- === Tags++ updSnapshotted stats = case (action', result) of+ (SaveSnapshot _ name _ _, MEither (Right (MUnit ()))) -> stats {+ snapshotted = Set.insert name (snapshotted stats)+ }+ (DeleteSnapshot name, MEither (Right (MUnit ()))) -> stats {+ snapshotted = Set.delete name (snapshotted stats)+ }+ _ -> stats++ -- === Final tags++ updNumLookupsResults stats = case (action', result) of+ (Lookups _ _, MEither (Right (MVector lrs))) -> stats {+ numLookupsResults =+ let count :: (Int, Int, Int)+ -> Val m h (LookupResult v (WrapBlobRef h m blob))+ -> (Int, Int, Int)+ count (nf, f, fwb) (MLookupResult x) = case x of+ NotFound -> (nf+1, f , fwb )+ Found{} -> (nf , f+1, fwb )+ FoundWithBlob{} -> (nf , f , fwb+1)+ in V.foldl' count (numLookupsResults stats) lrs+ }+ _ -> stats++ updNumUpdates stats = case (action', result) of+ (Updates upds _, MEither (Right (MUnit ()))) -> stats {+ numUpdates = countAll upds+ }+ (Inserts ins _, MEither (Right (MUnit ()))) -> stats {+ numUpdates = countAll $ V.map (\(k, v, b) -> (k, R.Insert v b)) ins+ }+ (Deletes ks _, MEither (Right (MUnit ()))) -> stats {+ numUpdates = countAll $ V.map (\k -> (k, R.Delete)) ks+ }+ (Upserts mups _, MEither (Right (MUnit ()))) -> stats {+ numUpdates = countAll $ V.map (second R.Upsert) mups+ }+ _ -> stats+ where+ countAll :: forall k v b. V.Vector (k, R.Update v b) -> (Int, Int, Int, Int)+ countAll upds =+ let count :: (Int, Int, Int, Int)+ -> (k, R.Update v blob)+ -> (Int, Int, Int, Int)+ count (i, iwb, d, m) (_, upd) = case upd of+ R.Insert _ Nothing -> (i+1, iwb , d , m )+ R.Insert _ Just{} -> (i , iwb+1, d , m )+ R.Delete{} -> (i , iwb , d+1, m )+ R.Upsert{} -> (i , iwb , d , m+1)+ in V.foldl' count (numUpdates stats) upds++ updSuccessActions stats = case result of+ MEither (Right _) -> stats {+ successActions = actionName action : successActions stats+ }+ _ -> stats++ updFailActions stats = case result of+ MEither (Left (MErr e)) -> stats {+ failActions = (actionName action, e) : failActions stats+ }+ _ -> stats++ updNumActionsPerTable :: Stats -> Stats+ updNumActionsPerTable stats = case action' of+ NewTableWith{}+ | MEither (Right (MTable table)) <- result -> initCount table+ | otherwise -> stats+ OpenTableFromSnapshot{}+ | MEither (Right (MTable table)) <- result -> initCount table+ | otherwise -> stats+ Duplicate{}+ | MEither (Right (MTable table)) <- result -> initCount table+ | otherwise -> stats+ Union{}+ | MEither (Right (MTable table)) <- result -> initCount table+ | otherwise -> stats+ Unions{}+ | MEither (Right (MTable table)) <- result -> initCount table+ | otherwise -> stats++ -- Note that for the other actions we don't count success vs failure.+ -- We don't need that level of detail. We just want to see the+ -- distribution. Success / failure is detailed elsewhere.+ Lookups _ tableVar -> updateCount tableVar+ RangeLookup _ tableVar -> updateCount tableVar+ NewCursor _ tableVar -> updateCount tableVar+ Updates _ tableVar -> updateCount tableVar+ Inserts _ tableVar -> updateCount tableVar+ Deletes _ tableVar -> updateCount tableVar+ Upserts _ tableVar -> updateCount tableVar+ RemainingUnionDebt tableVar -> updateCount tableVar+ SupplyUnionCredits tableVar _ -> updateCount tableVar+ SupplyPortionOfDebt tableVar _ -> updateCount tableVar+ -- Note that we don't remove tracking map entries for tables that get+ -- closed. We want to know actions per table of all tables used, not+ -- just those that were still open at the end of the sequence of+ -- actions. We do also count CloseTable itself as an action.+ CloseTable tableVar -> updateCount tableVar++ -- The others are not counted as table actions. We list them here+ -- explicitly so we don't miss any new ones we might add later.+ CloseCursor{} -> stats+ ReadCursor{} -> stats+ RetrieveBlobs{} -> stats+ SaveSnapshot{} -> stats+ DeleteSnapshot{} -> stats+ ListSnapshots{} -> stats+ where+ -- Init to 0 so we get an accurate count of tables with no actions.+ initCount :: forall k v b. Model.Table k v b -> Stats+ initCount table =+ let tid = Model.tableID table+ in stats {+ numActionsPerTable = Map.insert tid 0 (numActionsPerTable stats)+ }++ -- Note that batches (of inserts lookups etc) count as one action.+ updateCount :: forall k v b.+ Var m h (WrapTable h m k v b)+ -> Stats+ updateCount tableVar =+ let tid = getTableId (lookUp tableVar)+ in stats {+ numActionsPerTable = Map.insertWith (+) tid 1+ (numActionsPerTable stats)+ }++ updClosedTables stats = case (action', result) of+ (CloseTable tableVar, MEither (Right (MUnit ())))+ | MTable t <- lookUp tableVar+ , let tid = Model.tableID t+ -- This lookup can fail if the table was already closed:+ , Just (_, table) <- Map.lookup tid (Model.tables modelBefore)+ -> stats {+ closedTables = Map.insert tid table (closedTables stats)+ }+ _ -> stats++ updParentTable stats = case (action', result) of+ (NewTableWith{}, MEither (Right (MTable tbl))) ->+ insertParentTableNew tbl stats+ (OpenTableFromSnapshot{}, MEither (Right (MTable tbl))) ->+ insertParentTableNew tbl stats+ (Duplicate ptblVar, MEither (Right (MTable tbl))) ->+ insertParentTableDerived [ptblVar] tbl stats+ (Union ptblVar1 ptblVar2, MEither (Right (MTable tbl))) ->+ insertParentTableDerived [ptblVar1, ptblVar2] tbl stats+ (Unions ptblVars, MEither (Right (MTable tbl))) ->+ insertParentTableDerived (NE.toList ptblVars) tbl stats+ _ -> stats++ -- insert an entry into the parentTable for a completely new table+ insertParentTableNew :: forall k v b. Model.Table k v b -> Stats -> Stats+ insertParentTableNew tbl stats =+ stats {+ parentTable = Map.insert (Model.tableID tbl)+ [Model.tableID tbl]+ (parentTable stats)+ }++ -- insert an entry into the parentTable for a table derived from a parent+ insertParentTableDerived :: forall k v b.+ [GVar Op (WrapTable h m k v b)]+ -> Model.Table k v b -> Stats -> Stats+ insertParentTableDerived ptblVars tbl stats =+ let uptblIds :: [Model.TableID] -- the set of ultimate parent table ids+ uptblIds = List.nub [+ uptblId+ | ptblVar <- ptblVars+ -- immediate and ultimate parent table id:+ , let iptblId = getTableId (lookUp ptblVar)+ , uptblId <- parentTable stats Map.! iptblId+ ]+ in stats {+ parentTable = Map.insert (Model.tableID tbl)+ uptblIds+ (parentTable stats)+ }++ updDupTableActionLog stats | MEither (Right _) <- result =+ case action' of+ Lookups ks tableVar+ | not (null ks) -> updateLastActionLog tableVar+ | otherwise -> stats+ RangeLookup r tableVar+ | not (emptyRange r) -> updateLastActionLog tableVar+ | otherwise -> stats+ NewCursor _ tableVar -> updateLastActionLog tableVar+ Updates upds tableVar+ | not (null upds) -> updateLastActionLog tableVar+ | otherwise -> stats+ Inserts ins tableVar+ | not (null ins) -> updateLastActionLog tableVar+ | otherwise -> stats+ Deletes ks tableVar+ | not (null ks) -> updateLastActionLog tableVar+ | otherwise -> stats+ Upserts mups tableVar+ | not (null mups) -> updateLastActionLog tableVar+ | otherwise -> stats+ CloseTable tableVar -> updateLastActionLog tableVar+ -- Uninteresting actions+ NewTableWith{} -> stats+ CloseCursor{} -> stats+ ReadCursor{} -> stats+ RetrieveBlobs{} -> stats+ SaveSnapshot{} -> stats+ OpenTableFromSnapshot{} -> stats+ DeleteSnapshot{} -> stats+ ListSnapshots{} -> stats+ Duplicate{} -> stats+ Union{} -> stats+ Unions{} -> stats+ RemainingUnionDebt{} -> stats+ SupplyUnionCredits{} -> stats+ SupplyPortionOfDebt{} -> stats+ where+ -- add the current table to the front of the list of tables, if it's+ -- not the latest one already+ updateLastActionLog :: GVar Op (WrapTable h m k v b) -> Stats+ updateLastActionLog tableVar =+ stats {+ dupTableActionLog = List.foldl'+ (flip (Map.alter extendLog))+ (dupTableActionLog stats)+ (parentTable stats Map.! thid)+ }+ where+ thid = getTableId (lookUp tableVar)++ extendLog :: Maybe [Model.TableID] -> Maybe [Model.TableID]+ extendLog (Just alog@(thid' : _)) | thid == thid'+ = Just alog -- action via same table+ extendLog (Just alog) = Just (thid : alog) -- action via different table+ extendLog Nothing = Just (thid : []) -- first action on table++ emptyRange (R.FromToExcluding l u) = l >= u+ emptyRange (R.FromToIncluding l u) = l > u++ updDupTableActionLog stats = stats++ updUnionTables stats = case action' of+ Union{} -> insertUnionTable+ Unions{} -> insertUnionTable+ _ -> stats+ where+ insertUnionTable+ | MEither (Right (MTable t)) <- result+ , let tid = Model.tableID t+ , Just (_,tbl) <- Map.lookup tid (Model.tables modelAfter)+ , let sz = Model.withSomeTable Model.size tbl+ = stats {+ unionTables = Map.insert tid sz (unionTables stats)+ }+ | otherwise+ = stats++ getTableId :: ModelValue (ModelState m h) (WrapTable h m k v b)+ -> Model.TableID+ getTableId (MTable t) = Model.tableID t++-- | Tags for every step+data Tag =+ -- | Snapshot with a name that already exists+ SnapshotTwice+ -- | Open an existing snapshot+ | OpenExistingSnapshot+ -- | Open a missing snapshot+ | OpenMissingSnapshot+ -- | Delete an existing snapshot+ | DeleteExistingSnapshot+ -- | Delete a missing snapshot+ | DeleteMissingSnapshot+ -- | A corrupted snapshot was created successfully+ | SaveSnapshotCorrupted R.SnapshotName+ -- | An /un/corrupted snapshot was created successfully+ | SaveSnapshotUncorrupted R.SnapshotName+ -- | A snapshot failed to open because we detected that the snapshot was+ -- corrupt+ | OpenTableFromSnapshotDetectsCorruption R.SnapshotName+ -- | Opened a snapshot (successfully) for a table involving a union+ | OpenTableFromSnapshotUnion+ deriving stock (Show, Eq, Ord)++-- | This is run for after every action+tagStep' ::+ (ModelState m h, ModelState m h)+ -> LockstepAction (ModelState m h) a+ -> Val m h a+ -> [Tag]+tagStep' (ModelState _stateBefore statsBefore,+ ModelState _stateAfter _statsAfter)+ (Action _ action') result =+ catMaybes [+ tagSnapshotTwice+ , tagOpenExistingSnapshot+ , tagOpenExistingSnapshot+ , tagOpenMissingSnapshot+ , tagDeleteExistingSnapshot+ , tagDeleteMissingSnapshot+ , tagSaveSnapshotCorruptedOrUncorrupted+ , tagOpenTableFromSnapshotDetectsCorruption+ , tagOpenTableFromSnapshotUnion+ ]+ where+ tagSnapshotTwice+ | SaveSnapshot _ name _ _ <- action'+ , name `Set.member` snapshotted statsBefore+ = Just SnapshotTwice+ | otherwise+ = Nothing++ tagOpenExistingSnapshot+ | OpenTableFromSnapshot _ name _ <- action'+ , name `Set.member` snapshotted statsBefore+ = Just OpenExistingSnapshot+ | otherwise+ = Nothing++ tagOpenMissingSnapshot+ | OpenTableFromSnapshot _ name _ <- action'+ , not (name `Set.member` snapshotted statsBefore)+ = Just OpenMissingSnapshot+ | otherwise+ = Nothing++ tagDeleteExistingSnapshot+ | DeleteSnapshot name <- action'+ , name `Set.member` snapshotted statsBefore+ = Just DeleteExistingSnapshot+ | otherwise+ = Nothing++ tagDeleteMissingSnapshot+ | DeleteSnapshot name <- action'+ , not (name `Set.member` snapshotted statsBefore)+ = Just DeleteMissingSnapshot+ | otherwise+ = Nothing++ tagSaveSnapshotCorruptedOrUncorrupted+ | SaveSnapshot mcorr name _ _ <- action'+ , MEither (Right (MUnit ())) <- result+ = Just $ case mcorr of+ Just (_ :: SilentCorruption) -> SaveSnapshotCorrupted name+ _ -> SaveSnapshotUncorrupted name+ | otherwise+ = Nothing++ tagOpenTableFromSnapshotDetectsCorruption+ | OpenTableFromSnapshot _ name _ <- action'+ , MEither (Left (MErr (Model.ErrSnapshotCorrupted _))) <- result+ = Just (OpenTableFromSnapshotDetectsCorruption name)+ | otherwise+ = Nothing++ tagOpenTableFromSnapshotUnion+ | OpenTableFromSnapshot{} <- action'+ , MEither (Right (MTable t)) <- result+ , Model.isUnionDescendant t == Model.IsUnionDescendant+ = Just OpenTableFromSnapshotUnion+ | otherwise+ = Nothing++-- | Tags for the final state+data FinalTag =+ -- | Total number of lookup results that were 'SUT.NotFound'+ NumLookupsNotFound String+ -- | Total number of lookup results that were 'SUT.Found'+ | NumLookupsFound String+ -- | Total number of lookup results that were 'SUT.FoundWithBlob'+ | NumLookupsFoundWithBlob String+ -- | Number of 'Class.Insert's successfully submitted to a table+ -- (this includes submissions through both 'Class.updates' and+ -- 'Class.inserts')+ | NumInserts String+ -- | Number of 'Class.InsertWithBlob's successfully submitted to a table+ -- (this includes submissions through both 'Class.updates' and+ -- 'Class.inserts')+ | NumInsertsWithBlobs String+ -- | Number of 'Class.Delete's successfully submitted to a table+ -- (this includes submissions through both 'Class.updates' and+ -- 'Class.deletes')+ | NumDeletes String+ -- | Number of 'Class.Mupsert's successfully submitted to a table+ -- (this includes submissions through both 'Class.updates' and+ -- 'Class.upserts')+ | NumMupserts String+ -- | Total number of actions (failing, succeeding, either)+ | NumActions String+ -- | Which actions succeeded+ | ActionSuccess String+ -- | Which actions failed+ | ActionFail String Model.Err+ -- | Number of tables created (new, open or duplicate)+ | NumTables String+ -- | Number of actions on each table+ | NumTableActions String+ -- | Total /logical/ size of a table+ | TableSize String+ -- | Number of interleaved actions on duplicate tables+ | DupTableActionLog String+ deriving stock Show++-- | This is run only after completing every action+tagFinalState' :: Lockstep (ModelState m h) -> [(String, [FinalTag])]+tagFinalState' (getModel -> ModelState finalState finalStats) = concat [+ tagNumLookupsResults+ , tagNumUpdates+ , tagNumActions+ , tagSuccessActions+ , tagFailActions+ , tagNumTables+ , tagNumTableActions+ , tagTableSizes+ , tagDupTableActionLog+ , tagNumUnionTables+ , tagNumUnionTableActions+ ]+ where+ tagNumLookupsResults = [+ ("Lookups not found" , [NumLookupsNotFound $ showPowersOf 10 nf])+ , ("Lookups found" , [NumLookupsFound $ showPowersOf 10 f])+ , ("Lookups found with blob", [NumLookupsFoundWithBlob $ showPowersOf 10 fwb])+ ]+ where (nf, f, fwb) = numLookupsResults finalStats++ tagNumUpdates = [+ ("Inserts" , [NumInserts $ showPowersOf 10 i])+ , ("Inserts with blobs" , [NumInsertsWithBlobs $ showPowersOf 10 iwb])+ , ("Deletes" , [NumDeletes $ showPowersOf 10 d])+ , ("Upserts" , [NumMupserts $ showPowersOf 10 m])+ ]+ where (i, iwb, d, m) = numUpdates finalStats++ tagNumActions =+ [ let n = length (successActions finalStats) in+ ("Actions that succeeded total", [NumActions (showPowersOf 10 n)])+ , let n = length (failActions finalStats) in+ ("Actions that failed total", [NumActions (showPowersOf 10 n)])+ , let n = length (successActions finalStats)+ + length (failActions finalStats) in+ ("Actions total", [NumActions (showPowersOf 10 n)])+ ]++ tagSuccessActions =+ [ ("Actions that succeeded", [ActionSuccess c])+ | c <- successActions finalStats ]++ tagFailActions =+ [ ("Actions that failed", [ActionFail c e])+ | (c, e) <- failActions finalStats ]++ tagNumTables =+ [ ("Number of tables", [NumTables (showPowersOf 2 n)])+ | let n = Map.size (numActionsPerTable finalStats)+ ]++ tagNumTableActions =+ [ ("Number of actions per table", [ NumTableActions (showPowersOf 2 n) ])+ | n <- Map.elems (numActionsPerTable finalStats)+ ]++ tagTableSizes =+ [ ("Table sizes", [ TableSize (showPowersOf 2 size) ])+ | let openSizes, closedSizes :: Map Model.TableID Int+ openSizes = Model.withSomeTable Model.size . snd <$>+ Model.tables finalState+ closedSizes = Model.withSomeTable Model.size <$>+ closedTables finalStats+ , size <- Map.elems (openSizes `Map.union` closedSizes)+ ]++ tagDupTableActionLog =+ [ ("Interleaved actions on table duplicates or unions",+ [DupTableActionLog (showPowersOf 2 n)])+ | (_, alog) <- Map.toList (dupTableActionLog finalStats)+ , let n = length alog+ ]++ tagNumUnionTables =+ [ ("Number of union tables (empty)",+ [NumTables (showPowersOf 2 (Map.size trivial))])+ , ("Number of union tables (non-empty)",+ [NumTables (showPowersOf 2 (Map.size nonTrivial))])+ ]+ where+ (nonTrivial, trivial) = Map.partition (> 0) (unionTables finalStats)++ tagNumUnionTableActions =+ [ ("Number of actions per table with non-empty unions",+ [ NumTableActions (showPowersOf 2 n) ])+ | n <- Map.elems $ numActionsPerTable finalStats+ `Map.intersection`+ Map.filter (> 0) (unionTables finalStats)+ ]++{-------------------------------------------------------------------------------+ Utils+-------------------------------------------------------------------------------}++-- | Version of 'QLS.runActionsBracket' with tagging of the final state.+--+-- The 'tagStep' feature tags each step (i.e., 'Action'), but there are cases+-- where one wants to tag a /list of/ 'Action's. For example, if one wants to+-- count how often something happens over the course of running these actions,+-- then we would want to only tag the final state, not intermediate steps.+runActionsBracket ::+ forall state st m n e prop. (+ RunLockstep state m+ , e ~ Error (Lockstep state) m+ , forall a. IsPerformResult e a+ , QC.Testable prop+ , MonadMask n+ , MonadST n+ , QLS.IOProperty n+ )+ => Proxy state+ -> CheckCleanup+ -> CheckRefs+ -> n st+ -> (st -> n prop)+ -> (m QC.Property -> st -> n QC.Property)+ -> (Lockstep state -> [(String, [FinalTag])])+ -> Actions (Lockstep state) -> QC.Property+runActionsBracket p cleanupFlag refsFlag init cleanup runner tagger actions =+ tagFinalState actions tagger+ $ QLS.runActionsBracket p init cleanup' runner actions+ where+ cleanup' st = do+ -- We want to run forgotten reference checks after cleanup, since cleanup+ -- itself may lead to forgotten refs. The reference checks have the+ -- crucial side effect of resetting the forgotten refs state. If we don't+ -- do this then the next test run (e.g. during shrinking) will encounter a+ -- false/stale forgotten refs exception. But we also have to make sure+ -- that if cleanup itself fails, that we still run the reference checks.+ -- 'propCleanup' will make sure to catch any exceptions that are thrown by+ -- the 'cleanup' action. 'propRefs' will then definitely run afterwards so+ -- that the forgotten reference checks are definitely performed.+ x <- propCleanup cleanupFlag $ cleanup st+ y <- propRefs refsFlag+ pure (x QC..&&. y)++tagFinalState ::+ forall state. StateModel (Lockstep state)+ => Actions (Lockstep state)+ -> (Lockstep state -> [(String, [FinalTag])])+ -> Property+ -> Property+tagFinalState actions tagger =+ flip (foldr (\(key, values) -> QC.tabulate key (fmap show values))) finalTags+ where+ finalAnnState :: Annotated (Lockstep state)+ finalAnnState = stateAfter @(Lockstep state) actions++ finalTags = tagger $ underlyingState finalAnnState++propException :: (Exception e, QC.Testable prop) => String -> Either e prop -> Property+propException s (Left e) = QC.counterexample (s <> displayException e) False+propException _ (Right prop) = QC.property prop++{-------------------------------------------------------------------------------+ Cleanup exceptions+-------------------------------------------------------------------------------}++-- | Flag that turns on\/off cleanup checks.+--+-- If injected errors left the database in an inconsistent state, then property+-- cleanup might throw exceptions. If 'CheckCleanup' is used, this will lead to+-- failing properties, otherwise the exceptions are ignored.+data CheckCleanup = CheckCleanup | NoCheckCleanup++propCleanup :: (MonadCatch m, QC.Testable prop) => CheckCleanup -> m prop -> m Property+propCleanup flag cleanupAction =+ propException "Cleanup exception: " <$> checkCleanupM flag cleanupAction++checkCleanupM :: (MonadCatch m, QC.Testable prop) => CheckCleanup -> m prop -> m (Either SomeException Property)+checkCleanupM flag cleanupAction = do+ eith <- try @_ @SomeException cleanupAction+ case flag of+ CheckCleanup -> pure $ QC.property <$> eith+ NoCheckCleanup -> pure (Right $ QC.property ())++{-------------------------------------------------------------------------------+ Reference checks+-------------------------------------------------------------------------------}++-- | Flag that turns on\/off reference checks.+--+-- If injected errors left the database in an inconsistent state, then some+-- references might be forgotten, which leads to reference exceptions. If+-- 'CheckRefs' is used, this will lead to failing properties, otherwise the+-- exceptions are ignored.+data CheckRefs = CheckRefs | NoCheckRefs++propRefs :: (PrimMonad m, MonadCatch m) => CheckRefs -> m Property+propRefs flag = propException "Reference exception: " <$> checkRefsM flag++checkRefsM :: (PrimMonad m, MonadCatch m) => CheckRefs -> m (Either RefException ())+checkRefsM flag = case flag of+ CheckRefs -> try checkForgottenRefs+ NoCheckRefs -> Right <$> ignoreForgottenRefs
@@ -0,0 +1,295 @@+{-# LANGUAGE TemplateHaskell #-}++{-# OPTIONS_GHC -Wno-unused-do-bind #-}+{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Database.LSMTree.StateMachine.DL (+ tests+ ) where++import Control.Monad (void)+import Control.RefCount+import Control.Tracer+import qualified Data.Map.Strict as Map+import Data.Typeable (Typeable)+import qualified Data.Vector as V+import Database.LSMTree as R+import qualified Database.LSMTree.Internal.Config as R (TableConfig (..))+import qualified Database.LSMTree.Model.Session as Model (fromSomeTable, tables)+import qualified Database.LSMTree.Model.Table as Model (values)+import Prelude+import SafeWildCards+import System.FS.API.Types+import System.FS.Sim.Error hiding (Blob)+import qualified System.FS.Sim.Stream as Stream+import System.FS.Sim.Stream (Stream)+import Test.Database.LSMTree.StateMachine hiding (tests)+import Test.Database.LSMTree.StateMachine.Op+import Test.QuickCheck as QC hiding (Some, label)+import Test.QuickCheck.DynamicLogic+import qualified Test.QuickCheck.Gen as QC+import qualified Test.QuickCheck.Random as QC+import Test.QuickCheck.StateModel.Lockstep+import qualified Test.QuickCheck.StateModel.Lockstep.Defaults as QLS+import Test.QuickCheck.StateModel.Variables+import Test.Tasty (TestTree, testGroup, withResource)+import qualified Test.Tasty.QuickCheck as QC+import Test.Util.FS+import Test.Util.PrettyProxy++tests :: TestTree+tests = testGroup "Test.Database.LSMTree.StateMachine.DL" [+ QC.testProperty "prop_example" prop_example+ , test_noSwallowedExceptions+ ]++instance DynLogicModel (Lockstep (ModelState IO R.Table))++-- | An example of how dynamic logic formulas can be run.+--+-- 'dl_example' is a manually created formula, but the same method of running a+-- formula also applies to counterexamples produced by a state machine property+-- test, such as 'propLockstep_RealImpl_RealFS_IO'. Such counterexamples are+-- (often almost) valid 'DL' expression. They can be copy-pasted into a Haskell+-- module with minor tweaks to make the compiler accept the copied code, and+-- then they can be run as any other 'DL' expression.+prop_example :: Property+prop_example =+ -- Run the example ...+ forAllDL dl_example $+ -- ... with the given lockstep property+ propLockstep_RealImpl_MockFS_IO tr CheckCleanup CheckFS CheckRefs (QC.Fixed 17)+ where+ -- To enable tracing, use something like @show `contramap` stdoutTracer@+ -- instead+ tr = nullTracer++-- | Create an initial "large" table+dl_example :: Typeable m => DL (Lockstep (ModelState m R.Table)) ()+dl_example = do+ -- Create an initial table and fill it with some inserts+ var3 <- action $ Action Nothing $ NewTableWith (PrettyProxy @((Key, Value, Blob))) (R.TableConfig {+ confMergePolicy = LazyLevelling+ , confSizeRatio = Four+ , confWriteBufferAlloc = AllocNumEntries 4+ , confBloomFilterAlloc = AllocFixed 10+ , confFencePointerIndex = OrdinaryIndex+ , confDiskCachePolicy = DiskCacheNone+ , confMergeSchedule = OneShot+ , confMergeBatchSize = MergeBatchSize 4+ })+ let kvs :: Map.Map Key Value+ kvs = Map.fromList $+ QC.unGen (QC.vectorOf 37 $ (,) <$> QC.arbitrary <*> QC.arbitrary)+ (QC.mkQCGen 42) 30+ ups :: V.Vector (Key, Update Value Blob)+ ups = V.fromList+ . map (\(k,v) -> (k, Insert v Nothing))+ . Map.toList $ kvs+ action $ Action Nothing $ Updates ups (unsafeMkGVar var3 (OpFromRight `OpComp` OpId))+ -- This is a rather ugly assertion, and could be improved using some helper+ -- function(s). However, it does serve its purpose as checking that the+ -- insertions we just did were successful.+ assertModel "table size" $ \s ->+ let (ModelState s' _) = getModel s in+ case Map.elems (Model.tables s') of+ [(_, smTbl)]+ | Just tbl <- (Model.fromSomeTable @Key @Value @Blob smTbl)+ -> Map.size (Model.values tbl) == Map.size kvs+ _ -> False++{-------------------------------------------------------------------------------+ Swallowed exceptions+-------------------------------------------------------------------------------}++-- | See 'prop_noSwallowedExceptions'.+--+-- Forgotten reference checks are disabled completely, because we allow bugs+-- (like forgotten references) in exception unsafe code where we inject disk+-- faults.+test_noSwallowedExceptions :: TestTree+test_noSwallowedExceptions =+ withResource+ (checkForgottenRefs >> disableForgottenRefChecks)+ (\_ -> enableForgottenRefChecks) $ \ !_ ->+ QC.testProperty "prop_noSwallowedExceptions" prop_noSwallowedExceptions++-- | Test that the @lsm-tree@ library does not swallow exceptions.+--+-- A functional requirement for the @lsm-tree@ library is that all exceptions+-- are properly communicated to the user. An alternative way of stating this+-- requirement is that no exceptions should be /swallowed/ by the library. We+-- test this requirement by running the state machine test with injected disk+-- errors using @fs-sim@, and asserting that no exceptions are swallowed.+--+-- The state machine test compares the SUT against a model by checking that+-- their responses to @lsm-tree@ actions are the same. As of writing this+-- property, not all of these actions on the SUT are guaranteed to be fully+-- exception safe. As a result, an exception might leave the database (i.e.,+-- session, tables, cursors) in an inconsistent state. The results of subsequent+-- operations on the inconsistent database should be considered undefined. As+-- such, it is highly likely that the SUT and model will thereafter disagree,+-- leading to a failing property.+--+-- Still, we want to run the swallowed error assertion on /all/ operations,+-- regardless of whether they are exception safe. We overcome this problem by+-- /definitely/ injecting errors (and running a swallowed error assertion) for+-- the last action in a sequence of actions. This may leave the final database+-- state inconsistent, but that is okay. However, we'll also have to disable+-- sanity checks like 'NoCheckCleanup', 'NoCheckFS', and 'NoCheckRefs', because+-- they are likely to fail if the database is an inconsistent state.+--+-- TODO: running only one swallowed exception assertion per action sequence is+-- restrictive, but this automatically improves as we make more actions+-- exceptions safe. When we generate injected errors for these errors by default+-- (in @arbitraryWithVars@), the swallowed exception assertion automatically+-- runs for those actions as well.+prop_noSwallowedExceptions :: QC.Fixed Salt -> Property+prop_noSwallowedExceptions salt = forAllDL dl_noSwallowExceptions runner+ where+ -- disable all file system and reference checks+ runner = propLockstep_RealImpl_MockFS_IO tr NoCheckCleanup NoCheckFS NoCheckRefs salt+ tr = nullTracer++-- | Run any number of actions using the default actions generator, and finally+-- run a single action with errors *definitely* enabled.+dl_noSwallowExceptions :: Typeable m => DL (Lockstep (ModelState m R.Table)) ()+dl_noSwallowExceptions = do+ -- Run any number of actions as normal+ anyActions_++ -- Generate a single action as normal+ varCtx <- getVarContextDL+ st <- getModelStateDL+ let+ gen = QLS.arbitraryAction varCtx st+ predicate (Some a) = QLS.precondition st a+ shr (Some a) = QLS.shrinkAction varCtx st a+ Some a <- forAllQ $ withGenQ gen predicate shr++ -- Overwrite the maybe errors of the generated action with *definitely* just+ -- errors.+ case a of+ Action _merrs a' -> do+ HasNoVariables errs <-+ forAllQ $ hasNoVariablesQ $ withGenQ arbitraryErrors (\_ -> True) shrinkErrors+ -- Run the modified action+ void $ action $ Action (Just errs) a'++-- | Generate an 'Errors' with arbitrary probabilities of exceptions.+--+-- The default 'genErrors' from @fs-sim@ generates streams of 'Maybe' exceptions+-- with a fixed probability for a 'Just' or 'Nothing'. The version here+-- generates an arbitrary probability for each stream, which should generate a+-- larger variety of 'Errors' structures.+--+-- TODO: upstream to @fs-sim@ to replace the default 'genErrors'?+arbitraryErrors :: Gen Errors+arbitraryErrors = do+ dumpStateE <- genStream arbitrary+ hCloseE <- genStream arbitrary+ hTruncateE <- genStream arbitrary+ doesDirectoryExistE <- genStream arbitrary+ doesFileExistE <- genStream arbitrary+ hOpenE <- genStream arbitrary+ hSeekE <- genStream arbitrary+ hGetSomeE <- genErrorStreamGetSome+ hGetSomeAtE <- genErrorStreamGetSome+ hPutSomeE <- genErrorStreamPutSome+ hGetSizeE <- genStream arbitrary+ createDirectoryE <- genStream arbitrary+ createDirectoryIfMissingE <- genStream arbitrary+ listDirectoryE <- genStream arbitrary+ removeDirectoryRecursiveE <- genStream arbitrary+ removeFileE <- genStream arbitrary+ renameFileE <- genStream arbitrary+ hGetBufSomeE <- genErrorStreamGetSome+ hGetBufSomeAtE <- genErrorStreamGetSome+ hPutBufSomeE <- genErrorStreamPutSome+ hPutBufSomeAtE <- genErrorStreamPutSome+ pure $ filterErrors Errors {..}+ where+ -- Generate a stream using 'genLikelihoods' for its 'Maybe' elements.+ genStream :: forall a. Gen a -> Gen (Stream a)+ genStream genA = do+ (pNothing, pJust) <- genLikelihoods+ Stream.genInfinite $ Stream.genMaybe pNothing pJust genA++ -- Generate two integer likelihoods for 'Nothing' and 'Just' constructors.+ genLikelihoods :: Gen (Int, Int)+ genLikelihoods = do+ NonNegative pNothing <- arbitrary+ NonNegative pJust <- arbitrary+ if pNothing == 0 then+ pure (0, 1)+ else if pJust == 0 then+ pure (1, 0)+ else+ pure (pNothing, pJust)++ genErrorStreamGetSome :: Gen ErrorStreamGetSome+ genErrorStreamGetSome = genStream $ liftArbitrary2 arbitrary arbitrary++ genErrorStreamPutSome :: Gen ErrorStreamPutSome+ genErrorStreamPutSome = genStream $ flip liftArbitrary2 arbitrary $ do+ errorType <- arbitrary+ maybePutCorruption <- liftArbitrary genPutCorruption+ pure (errorType, maybePutCorruption)++ genPutCorruption :: Gen PutCorruption+ genPutCorruption = oneof [+ PartialWrite <$> arbitrary+ , SubstituteWithJunk <$> arbitrary+ ]+ where+ _coveredAllCases x = case x of+ PartialWrite{} -> pure ()+ SubstituteWithJunk{} -> pure ()++ -- TODO: there is one case where an 'FsReachEOF' error is swallowed. Is that+ -- valid behaviour, or should we change it?+ filterErrors errs = errs {+ hGetBufSomeE = Stream.filter (not . isFsReachedEOFError) (hGetBufSomeE errs)+ }++ isFsReachedEOFError = maybe False (either isFsReachedEOF (const False))++-- | Shrink each error stream and all error stream elements.+--+-- The default 'shrink' from @fs-sim@ shrinks only the stream structure, but not+-- the elements contained in those streams.+--+-- TODO: upstream to @fs-sim@ to replace the default 'shrink'?+shrinkErrors :: Errors -> [Errors]+shrinkErrors err@($(fields 'Errors))+ | allNull err = []+ | otherwise = emptyErrors : concatMap (filter (not . allNull))+ [ (\s' -> err { dumpStateE = s' }) <$> Stream.liftShrinkStream shrink dumpStateE+ , (\s' -> err { hOpenE = s' }) <$> Stream.liftShrinkStream shrink hOpenE+ , (\s' -> err { hCloseE = s' }) <$> Stream.liftShrinkStream shrink hCloseE+ , (\s' -> err { hSeekE = s' }) <$> Stream.liftShrinkStream shrink hSeekE+ , (\s' -> err { hGetSomeE = s' }) <$> Stream.liftShrinkStream shrink hGetSomeE+ , (\s' -> err { hGetSomeAtE = s' }) <$> Stream.liftShrinkStream shrink hGetSomeAtE+ , (\s' -> err { hPutSomeE = s' }) <$> Stream.liftShrinkStream shrink hPutSomeE+ , (\s' -> err { hTruncateE = s' }) <$> Stream.liftShrinkStream shrink hTruncateE+ , (\s' -> err { hGetSizeE = s' }) <$> Stream.liftShrinkStream shrink hGetSizeE+ , (\s' -> err { createDirectoryE = s' }) <$> Stream.liftShrinkStream shrink createDirectoryE+ , (\s' -> err { createDirectoryIfMissingE = s' }) <$> Stream.liftShrinkStream shrink createDirectoryIfMissingE+ , (\s' -> err { listDirectoryE = s' }) <$> Stream.liftShrinkStream shrink listDirectoryE+ , (\s' -> err { doesDirectoryExistE = s' }) <$> Stream.liftShrinkStream shrink doesDirectoryExistE+ , (\s' -> err { doesFileExistE = s' }) <$> Stream.liftShrinkStream shrink doesFileExistE+ , (\s' -> err { removeDirectoryRecursiveE = s' }) <$> Stream.liftShrinkStream shrink removeDirectoryRecursiveE+ , (\s' -> err { removeFileE = s' }) <$> Stream.liftShrinkStream shrink removeFileE+ , (\s' -> err { renameFileE = s' }) <$> Stream.liftShrinkStream shrink renameFileE+ , (\s' -> err { hGetBufSomeE = s' }) <$> Stream.liftShrinkStream shrink hGetBufSomeE+ , (\s' -> err { hGetBufSomeAtE = s' }) <$> Stream.liftShrinkStream shrink hGetBufSomeAtE+ , (\s' -> err { hPutBufSomeE = s' }) <$> Stream.liftShrinkStream shrink hPutBufSomeE+ , (\s' -> err { hPutBufSomeAtE = s' }) <$> Stream.liftShrinkStream shrink hPutBufSomeAtE+ ]++deriving stock instance Enum FsErrorType+deriving stock instance Bounded FsErrorType++instance Arbitrary FsErrorType where+ arbitrary = arbitraryBoundedEnum+ shrink = shrinkBoundedEnum
@@ -0,0 +1,134 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiParamTypeClasses #-}++-- | SumProd Op extended with BlobRef extraction+module Test.Database.LSMTree.StateMachine.Op (+ -- * 'Op'+ Op (..)+ , intOpId+ -- * 'HasBlobRef' class+ , HasBlobRef (..)+ ) where++import Control.Monad ((<=<))+import Control.Monad.Identity (Identity (..))+import qualified Data.Vector as V+import qualified Database.LSMTree.Class as Class+import GHC.Show (appPrec)+import Test.QuickCheck.StateModel.Lockstep (InterpretOp, Operation)+import qualified Test.QuickCheck.StateModel.Lockstep.Op as Op+import Test.Util.Orphans ()+import Test.Util.TypeFamilyWrappers (WrapBlobRef)++{-------------------------------------------------------------------------------+ 'Op'+-------------------------------------------------------------------------------}++data Op a b where+ OpId :: Op a a+ OpFst :: Op (a, b) a+ OpSnd :: Op (b, a) a+ OpLeft :: Op a (Either a b)+ OpRight :: Op b (Either a b)+ OpFromLeft :: Op (Either a b) a+ OpFromRight :: Op (Either a b) b+ OpComp :: Op b c -> Op a b -> Op a c+ OpLookupResults :: Op (V.Vector (Class.LookupResult v (WrapBlobRef h m b))) (V.Vector (WrapBlobRef h m b))+ OpEntrys :: Op (V.Vector (Class.Entry k v (WrapBlobRef h m b))) (V.Vector (WrapBlobRef h m b))++intOpId :: Op a b -> a -> Maybe b+intOpId OpId = Just+intOpId OpFst = Just . fst+intOpId OpSnd = Just . snd+intOpId OpLeft = Just . Left+intOpId OpRight = Just . Right+intOpId OpFromLeft = either Just (const Nothing)+intOpId OpFromRight = either (const Nothing) Just+intOpId (OpComp g f) = intOpId g <=< intOpId f+intOpId OpLookupResults = Just . V.mapMaybe getBlobRef+intOpId OpEntrys = Just . V.mapMaybe getBlobRef++{-------------------------------------------------------------------------------+ 'InterpretOp' instances+-------------------------------------------------------------------------------}++instance Operation Op where+ opIdentity = OpId++instance InterpretOp Op Identity where+ intOp = Op.intOpIdentity intOpId++{-------------------------------------------------------------------------------+ 'Show' and 'Eq' instances+-------------------------------------------------------------------------------}++sameOp :: Op a b -> Op c d -> Bool+sameOp = go+ where+ go :: Op a b -> Op c d -> Bool+ go OpId OpId = True+ go OpFst OpFst = True+ go OpSnd OpSnd = True+ go OpLeft OpLeft = True+ go OpRight OpRight = True+ go OpFromLeft OpFromLeft = True+ go OpFromRight OpFromRight = True+ go (OpComp g f) (OpComp g' f') = go g g' && go f f'+ go OpLookupResults OpLookupResults = True+ go OpEntrys OpEntrys = True+ go _ _ = False++ _coveredAllCases :: Op a b -> ()+ _coveredAllCases = \case+ OpId -> ()+ OpFst -> ()+ OpSnd -> ()+ OpLeft -> ()+ OpRight -> ()+ OpFromLeft -> ()+ OpFromRight -> ()+ OpComp{} -> ()+ OpLookupResults{} -> ()+ OpEntrys{} -> ()+++instance Eq (Op a b) where+ (==) = sameOp++instance Show (Op a b) where+ showsPrec :: Int -> Op a b -> ShowS+ showsPrec p = \op -> case op of+ OpComp{} -> showParen (p > appPrec) (go op)+ _ -> go op+ where+ go :: Op x y -> String -> String+ go OpId = showString "OpId"+ go OpFst = showString "OpFst"+ go OpSnd = showString "OpSnd"+ go OpLeft = showString "OpLeft"+ go OpRight = showString "OpRight"+ go OpFromLeft = showString "OpFromLeft"+ go OpFromRight = showString "OpFromRight"+ go (OpComp g f) = go g . showString " `OpComp` " . go f+ go OpLookupResults = showString "OpLookupResults"+ go OpEntrys = showString "OpEntrys"++{-------------------------------------------------------------------------------+ 'HasBlobRef' class+-------------------------------------------------------------------------------}++class HasBlobRef f where+ getBlobRef :: f b -> Maybe b++instance HasBlobRef (Class.LookupResult v) where+ getBlobRef Class.NotFound{} = Nothing+ getBlobRef Class.Found{} = Nothing+ getBlobRef (Class.FoundWithBlob _ blobref) = Just blobref++instance HasBlobRef (Class.Entry k v) where+ getBlobRef Class.Entry{} = Nothing+ getBlobRef (Class.EntryWithBlob _ _ blobref) = Just blobref
@@ -0,0 +1,159 @@+{-# LANGUAGE OverloadedLists #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Database.LSMTree.Tracer.Golden (tests) where++import Control.Exception (mask_)+import Control.Monad (void)+import Control.Tracer+import qualified Data.List.NonEmpty as NE+import Data.Monoid (Sum (..))+import Data.Word+import Database.LSMTree as R+import Paths_lsm_tree+import Prelude hiding (lookup)+import System.FilePath+import qualified System.FS.API as FS+import qualified System.IO as IO+import System.IO.Unsafe+import Test.Tasty (TestTree, localOption, testGroup)+import Test.Tasty.Golden+import Test.Util.FS (withTempIOHasBlockIO)+++tests :: TestTree+tests =+ localOption OnPass $+ testGroup "Test.Database.LSMTree.Tracer.Golden" [+ goldenVsFile+ "golden_traceMessages"+ goldenDataFile+ dataFile+ (golden_traceMessages dataFile)+ ]++goldenDataDirectory :: FilePath+goldenDataDirectory = unsafePerformIO getDataDir </> "tracer"++-- | Path to the /golden/ file for 'golden_traceMessages'+goldenDataFile :: FilePath+goldenDataFile = dataFile <.> "golden"++-- | Path to the output file for 'golden_traceMessages'+dataFile :: FilePath+dataFile = goldenDataDirectory </> "golden_traceMessages"++-- | A monadic action that traces messages emitted by the public API of+-- @lsm-tree@ to a file.+--+-- The idea is to cover as much of the public API as possible to trace most (if+-- not all) unique trace messages at least once. The trace output is useful in+-- two ways:+--+-- * By using golden tests based on the trace output we can help to ensure that+-- trace datatypes and their printing do not change unexpectedly.+--+-- * The trace output serves as an example that can be manually inspected to see+-- whether the trace messages are printed correctly and whether they are+-- useful.+golden_traceMessages :: FilePath -> IO ()+golden_traceMessages file =+ withTempIOHasBlockIO "golden_traceMessages" $ \hfs hbio ->+ withBinaryFileTracer file $ \(contramap show -> tr) -> do++ let sessionPath = FS.mkFsPath ["session-dir"]+ FS.createDirectory hfs sessionPath++ -- Open and close a session+ withOpenSession tr hfs hbio 4 sessionPath $ \s -> do+ -- Open and close a table+ withTable s $ \(t1 :: Table IO K V B) -> do+ -- Try out all the update variants+ update t1 (K 0) (Insert (V 0) Nothing)+ insert t1 (K 1) (V 1) (Just (B 1))+ delete t1 (K 2)+ upsert t1 (K 3) (V 3)++ -- Perform a lookup and its member variant+ FoundWithBlob (V 1) bref <- lookup t1 (K 1)+ void $ member t1 (K 2)++ void $ rangeLookup t1 (FromToExcluding (K 0) (K 3))++ -- Open and close a cursor+ withCursorAtOffset t1 (K 1) $ \c -> do+ -- Read from a cursor+ [EntryWithBlob (K 1) (V 1) _] <- R.takeWhile 2 (< K 3) c+ void $ R.next c+ -- Retrieve blobs (we didn't try this before this part)+ B 1 <- retrieveBlob s bref+ pure ()++ let snapA = "snapshot_a"+ snapB = "snapshot_b"++ -- Try out a few snapshot operations, and keep one saved snapshot for+ -- later+ saveSnapshot snapA snapLabel t1+ saveSnapshot snapB snapLabel t1+ void $ listSnapshots s+ deleteSnapshot s snapB+ void $ doesSnapshotExist s snapA++ -- Open a table from a snapshot and close it again+ withTableFromSnapshot s snapA snapLabel $ \ t2 -> do++ -- Create a table duplicate and close it again+ withDuplicate t1 $ \t3 -> do++ let unionInputs = NE.fromList [t1, t2, t3]++ -- One-shot union+ withUnions unionInputs $ \_ -> pure ()++ -- Incremental union+ withIncrementalUnions unionInputs $ \t4 -> do+ UnionDebt d <- remainingUnionDebt t4+ void $ supplyUnionCredits t4 (UnionCredits (d - 1))+ UnionDebt d' <- remainingUnionDebt t4+ void $ supplyUnionCredits t4 (UnionCredits (d' + 1))+ void $ remainingUnionDebt t4++ -- Open a table and cursor, but close them automatically as part of+ -- closing the session instead of closing them manually.+ (t :: Table IO K V B) <- mask_ $ newTable s+ _c <- mask_ $ newCursor t++ pure ()++-- | A tracer that emits trace messages to a file handle+fileHandleTracer :: IO.Handle -> Tracer IO String+fileHandleTracer h = Tracer $ emit (IO.hPutStrLn h)+++-- | A tracer that emits trace messages to a binary file+withBinaryFileTracer :: FilePath -> (Tracer IO String -> IO a) -> IO a+withBinaryFileTracer fp k =+ IO.withBinaryFile fp IO.WriteMode $ \h ->+ k (fileHandleTracer h)++{-------------------------------------------------------------------------------+ Key and value types+-------------------------------------------------------------------------------}++newtype K = K Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseKey)++newtype V = V Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseValue)++deriving via Sum Word64 instance ResolveValue V++newtype B = B Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseValue)++snapLabel :: SnapshotLabel+snapLabel = SnapshotLabel "KVBs"
@@ -0,0 +1,292 @@+{-# LANGUAGE OverloadedLists #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Database.LSMTree.UnitTests (tests) where++import Control.Tracer (nullTracer)+import Data.ByteString (ByteString)+import qualified Data.ByteString.Char8 as BS+import Data.Foldable (for_)+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.Map as Map+import qualified Data.Vector as V+import Data.Word+import qualified System.FS.API as FS++import Database.LSMTree as R++import Control.Exception (Exception, bracket, try)+import Database.LSMTree.Extras.Generators ()+import Database.LSMTree.Internal.Serialise (SerialisedKey)+import qualified Test.QuickCheck.Arbitrary as QC+import qualified Test.QuickCheck.Gen as QC+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit+import Test.Util.FS (withTempIOHasBlockIO)+++tests :: TestTree+tests =+ testGroup "Test.Database.LSMTree.UnitTests"+ [ testCaseSteps "unit_blobs" unit_blobs+ , testCase "unit_closed_table" unit_closed_table+ , testCase "unit_closed_cursor" unit_closed_cursor+ , testCase "unit_twoTableTypes" unit_twoTableTypes+ , testCase "unit_snapshots" unit_snapshots+ , testCase "unit_unions_1" unit_unions_1+ , testCase "unit_union_credits" unit_union_credits+ , testCase "unit_union_credit_0" unit_union_credit_0+ , testCase "unit_union_blobref_invalidation" unit_union_blobref_invalidation+ ]++testSalt :: R.Salt+testSalt = 4++unit_blobs :: (String -> IO ()) -> Assertion+unit_blobs info =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess -> do+ table <- newTable @_ @ByteString @(ResolveAsFirst ByteString) @ByteString sess+ inserts table [("key1", ResolveAsFirst "value1", Just "blob1")]++ res <- lookups table ["key1"]+ info (show res)++ case res of+ [FoundWithBlob val bref] -> do+ val @?= ResolveAsFirst "value1"+ blob <- retrieveBlobs sess [bref]+ info (show blob)+ blob @?= ["blob1"]+ _ -> assertFailure "expected FoundWithBlob"++unit_closed_table :: Assertion+unit_closed_table =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess -> do+ table <- newTable @_ @Key1 @Value1 @Blob1 sess+ inserts table [(Key1 42, Value1 42, Nothing)]+ r1 <- lookups table [Key1 42]+ V.map ignoreBlobRef r1 @?= [Found (Value1 42)]+ closeTable table+ -- Expect ErrTableClosed for operations after closeTable+ assertException ErrTableClosed $+ lookups table [Key1 42] >> pure ()+ -- But closing again is idempotent+ closeTable table++unit_closed_cursor :: Assertion+unit_closed_cursor =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess -> do+ table <- newTable @_ @Key1 @Value1 @Blob1 sess+ inserts table [(Key1 42, Value1 42, Nothing), (Key1 43, Value1 43, Nothing)]+ cur <- newCursor table+ -- closing the table should not affect the cursor+ closeTable table+ r1 <- next cur+ fmap ignoreBlobRef r1 @?= Just (Entry (Key1 42) (Value1 42))+ closeCursor cur+ -- Expect ErrCursorClosed for operations after closeCursor+ assertException ErrCursorClosed $+ next cur >> pure ()+ -- But closing again is idempotent+ closeCursor cur++unit_twoTableTypes :: Assertion+unit_twoTableTypes =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess -> do+ let tableConfig =+ defaultTableConfig {+ confWriteBufferAlloc = AllocNumEntries 10+ }+ table1 <- newTableWith @_ @Key1 @Value1 @Blob1 tableConfig sess+ table2 <- newTableWith @_ @Key2 @Value2 @Blob2 tableConfig sess++ kvs1 <- QC.generate (Map.fromList <$> QC.vectorOf 100 QC.arbitrary)+ kvs2 <- QC.generate (Map.fromList <$> QC.vectorOf 100 QC.arbitrary)+ let ins1 = V.fromList [ (Key1 k, Value1 v, Nothing)+ | (k,v) <- Map.toList kvs1 ]+ ins2 = V.fromList [ (Key2 k, Value2 v, Nothing)+ | (k,v) <- Map.toList kvs2 ]+ inserts table1 ins1+ inserts table2 ins2++ saveSnapshot snap1 label1 table1+ saveSnapshot snap2 label2 table2+ table1' <- openTableFromSnapshot @_ @Key1 @Value1 @Blob1 sess snap1 label1+ table2' <- openTableFromSnapshot @_ @Key2 @Value2 @Blob2 sess snap2 label2++ vs1 <- lookups table1' ((\(k,_,_)->k) <$> ins1)+ vs2 <- lookups table2' ((\(k,_,_)->k) <$> ins2)++ V.map ignoreBlobRef vs1 @?= ((\(_,v,_) -> Found v) <$> ins1)+ V.map ignoreBlobRef vs2 @?= ((\(_,v,_) -> Found v) <$> ins2)+ where+ snap1, snap2 :: SnapshotName+ snap1 = "table1"+ snap2 = "table2"++unit_snapshots :: Assertion+unit_snapshots =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess -> do+ table <- newTable @_ @Key1 @Value1 @Blob1 sess++ assertException (ErrSnapshotDoesNotExist snap2) $+ deleteSnapshot sess snap2++ saveSnapshot snap1 label1 table+ assertException (ErrSnapshotExists snap1) $+ saveSnapshot snap1 label1 table++ assertException (ErrSnapshotWrongLabel snap1+ (SnapshotLabel "Key2 Value2 Blob2")+ (SnapshotLabel "Key1 Value1 Blob1")) $ do+ _ <- openTableFromSnapshot @_ @Key2 @Value2 @Blob2 sess snap1 label2+ pure ()++ assertException (ErrSnapshotDoesNotExist snap2) $ do+ _ <- openTableFromSnapshot @_ @Key1 @Value1 @Blob1 sess snap2 label2+ pure ()+ where+ snap1, snap2 :: SnapshotName+ snap1 = "table1"+ snap2 = "table2"++-- | Unions of 1 table are equivalent to duplicate+unit_unions_1 :: Assertion+unit_unions_1 =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess ->+ withTable @_ @Key1 @Value1 @Blob1 sess $ \table -> do+ inserts table [(Key1 17, Value1 42, Nothing)]++ bracket (unions $ table :| []) closeTable $ \table' ->+ bracket (duplicate table) closeTable $ \table'' -> do+ inserts table' [(Key1 17, Value1 43, Nothing)]+ inserts table'' [(Key1 17, Value1 44, Nothing)]++ -- The original table is unmodified+ r <- lookups table [Key1 17]+ V.map ignoreBlobRef r @?= [Found (Value1 42)]++ -- The unioned table sees an updated value+ r' <- lookups table' [Key1 17]+ V.map ignoreBlobRef r' @?= [Found (Value1 43)]++ -- The duplicated table sees a different updated value+ r'' <- lookups table'' [Key1 17]+ V.map ignoreBlobRef r'' @?= [Found (Value1 44)]++-- | Querying or supplying union credits to non-union tables is trivial.+unit_union_credits :: Assertion+unit_union_credits =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess ->+ withTable @_ @Key1 @Value1 @Blob1 sess $ \table -> do+ inserts table [(Key1 17, Value1 42, Nothing)]++ -- The table is not the result of a union, so the debt is always 0,+ UnionDebt debt <- remainingUnionDebt table+ debt @?= 0++ -- and supplying credits returns them all as leftovers.+ UnionCredits leftover <- supplyUnionCredits table (UnionCredits 42)+ leftover @?= 42++-- | Supplying zero or negative credits to union tables works, but does nothing.+unit_union_credit_0 :: Assertion+unit_union_credit_0 =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess ->+ withTable @_ @Key1 @Value1 @Blob1 sess $ \table -> do+ inserts table [(Key1 17, Value1 42, Nothing)]++ bracket (table `union` table) closeTable $ \table' -> do+ -- Supplying 0 credits works and returns 0 leftovers.+ UnionCredits leftover <- supplyUnionCredits table' (UnionCredits 0)+ leftover @?= 0++ -- Supplying negative credits also works and returns 0 leftovers.+ UnionCredits leftover' <- supplyUnionCredits table' (UnionCredits (-42))+ leftover' @?= 0++ -- And the table is still vaguely cromulent+ r <- lookups table' [Key1 17]+ V.map ignoreBlobRef r @?= [Found (Value1 42)]++-- | Blob refs into a union don't get invalidated when updating the union's+-- input tables.+unit_union_blobref_invalidation :: Assertion+unit_union_blobref_invalidation =+ withTempIOHasBlockIO "test" $ \hfs hbio ->+ withOpenSession nullTracer hfs hbio testSalt (FS.mkFsPath []) $ \sess -> do+ t1 <- newTableWith @_ @Key1 @Value1 @Blob1 config sess+ for_ ([0..99] :: [Word64]) $ \i ->+ inserts t1 [(Key1 i, Value1 i, Just (Blob1 i))]+ t2 <- t1 `union` t1++ -- do lookups on the union table (the result contains blob refs)+ res <- lookups t2 (Key1 <$> [0..99])++ -- progress original table (supplying merge credits would be most direct),+ -- so merges complete+ inserts t1 (fmap (\i -> (Key1 i, Value1 i, Nothing)) [1000..2000])+ -- closeTable it, so it doesn't hold open extra references+ closeTable t1++ -- try to resolve the blob refs we obtained earlier+ _blobs <- retrieveBlobs sess (V.mapMaybe R.getBlob res)+ pure ()+ where+ config = defaultTableConfig {+ confWriteBufferAlloc = AllocNumEntries 4+ }++ignoreBlobRef :: Functor f => f (BlobRef m b) -> f ()+ignoreBlobRef = fmap (const ())++assertException :: (Exception e, Eq e) => e -> Assertion -> Assertion+assertException e assertion = do+ r <- try assertion+ r @?= Left e++{-------------------------------------------------------------------------------+ Key and value types+-------------------------------------------------------------------------------}++newtype Key1 = Key1 Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseKey)++newtype Value1 = Value1 Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseValue)++deriving via ResolveAsFirst Word64 instance ResolveValue Value1++newtype Blob1 = Blob1 Word64+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseValue)++label1 :: SnapshotLabel+label1 = SnapshotLabel "Key1 Value1 Blob1"++newtype Key2 = Key2 SerialisedKey+ deriving stock (Show, Eq, Ord)+ deriving newtype (QC.Arbitrary, SerialiseKey)++newtype Value2 = Value2 BS.ByteString+ deriving stock (Show, Eq, Ord)+ deriving newtype (QC.Arbitrary, SerialiseValue)++deriving via ResolveAsFirst BS.ByteString instance ResolveValue Value2++newtype Blob2 = Blob2 BS.ByteString+ deriving stock (Show, Eq, Ord)+ deriving newtype (SerialiseValue)++label2 :: SnapshotLabel+label2 = SnapshotLabel "Key2 Value2 Blob2"
@@ -0,0 +1,210 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- TODO: upstream to fs-sim+module Test.FS (tests) where++import Control.Concurrent.Class.MonadSTM (MonadSTM (atomically))+import Control.Concurrent.Class.MonadSTM.Strict.TMVar+import Control.Monad+import Control.Monad.Class.MonadThrow+import Control.Monad.IOSim (runSimOrThrow)+import Control.Monad.ST (runST)+import Data.Bit (cloneFromByteString, cloneToByteString, flipBit)+import Data.ByteString (ByteString)+import qualified Data.ByteString as BS+import Data.Set (Set)+import qualified Data.Set as Set+import Data.Vector.Unboxed (thaw, unsafeFreeze)+import GHC.Generics (Generic)+import System.FS.API+import System.FS.API.Lazy+import System.FS.API.Strict+import System.FS.Sim.Error+import qualified System.FS.Sim.MockFS as MockFS+import qualified System.FS.Sim.Stream as S+import System.FS.Sim.Stream (InternalInfo (..), Stream (..))+import Test.QuickCheck+import Test.QuickCheck.Classes (eqLaws)+import Test.QuickCheck.Instances ()+import Test.Tasty+import Test.Tasty.QuickCheck (testProperty)+import Test.Util.FS+import Test.Util.QC++tests :: TestTree+tests = testGroup "Test.FS" [+ -- * Simulated file system properties+ testProperty "prop_numOpenHandles" prop_numOpenHandles+ , testProperty "prop_numDirEntries" prop_numDirEntries+ -- * Corruption+ , testProperty "prop_flipFileBit" prop_flipFileBit+ -- * Equality+ , testClassLaws "Stream" $+ eqLaws (Proxy @(Stream Int))+ , testClassLaws "Errors" $+ eqLaws (Proxy @Errors)+ ]++{-------------------------------------------------------------------------------+ Simulated file system properties+-------------------------------------------------------------------------------}++-- | Sanity check for 'propNoOpenHandles' and 'propNumOpenHandles'+prop_numOpenHandles :: Set FsPathComponent -> Property+prop_numOpenHandles (Set.toList -> paths) = runSimOrThrow $+ withSimHasFS propTrivial MockFS.empty $ \hfs fsVar -> do+ -- No open handles initially+ fs <- atomically $ readTMVar fsVar+ let prop = propNoOpenHandles fs++ -- Open n handles+ hs <- forM paths $ \(fsPathComponentFsPath -> p) -> hOpen hfs p (WriteMode MustBeNew)++ -- Now there should be precisely n open handles+ fs' <- atomically $ readTMVar fsVar+ let prop' = propNumOpenHandles n fs'++ -- Close all previously opened handles+ forM_ hs $ hClose hfs++ -- No open handles again+ fs'' <- atomically $ readTMVar fsVar+ let prop'' = propNoOpenHandles fs''++ pure (prop .&&. prop' .&&. prop'')+ where+ n = length paths++-- | Sanity check for 'propNoDirEntries' and 'propNumDirEntries'+prop_numDirEntries ::+ FsPathComponent+ -> InfiniteList Bool+ -> Set FsPathComponent+ -> Property+prop_numDirEntries (fsPathComponentFsPath -> dir) isFiles (Set.toList -> paths) = runSimOrThrow $+ withSimHasFS propTrivial MockFS.empty $ \hfs fsVar -> do+ createDirectoryIfMissing hfs False dir++ -- No entries initially+ fs <- atomically $ readTMVar fsVar+ let prop = propNoDirEntries dir fs++ -- Create n entries+ forM_ xs $ \(isFile, fsPathComponentFsPath -> p) ->+ if isFile+ then createFile hfs (dir </> p)+ else createDirectory hfs (dir </> p)++ -- Now there should be precisely n entries+ fs' <- atomically $ readTMVar fsVar+ let prop' = propNumDirEntries dir n fs'++ -- Remove n entries+ forM_ xs $ \(isFile, fsPathComponentFsPath -> p) ->+ if isFile+ then removeFile hfs (dir </> p)+ else removeDirectoryRecursive hfs (dir </> p)++ -- No entries again+ fs'' <- atomically $ readTMVar fsVar+ let prop'' = propNoDirEntries dir fs''++ pure (prop .&&. prop' .&&. prop'')+ where+ n = length paths+ xs = zip (getInfiniteList isFiles) paths++createFile :: MonadThrow m => HasFS m h -> FsPath -> m ()+createFile hfs p = withFile hfs p (WriteMode MustBeNew) $ \_ -> pure ()++{-------------------------------------------------------------------------------+ Corruption+-------------------------------------------------------------------------------}++data WithBitOffset a = WithBitOffset Int a+ deriving stock Show++bitLength :: BS.ByteString -> Int+bitLength bs = BS.length bs * 8++instance Arbitrary (WithBitOffset ByteString) where+ arbitrary = do+ bs <- arbitrary `suchThat` (\bs -> BS.length bs > 0)+ bitOffset <- chooseInt (0, bitLength bs - 1)+ pure $ WithBitOffset bitOffset bs+ shrink (WithBitOffset bitOffset bs) =+ [ WithBitOffset bitOffset' bs'+ | bs' <- shrink bs+ , BS.length bs' > 0+ , let bitOffset' = max 0 $ min (bitLength bs' - 1) bitOffset+ ] ++ [+ WithBitOffset bitOffset' bs+ | bitOffset' <- max 0 <$> shrink bitOffset+ , bitOffset' >= 0+ ]++prop_flipFileBit :: WithBitOffset ByteString -> Property+prop_flipFileBit (WithBitOffset bitOffset bs) =+ ioProperty $+ withSimHasFS propTrivial MockFS.empty $ \hfs _fsVar -> do+ void $ withFile hfs path (WriteMode MustBeNew) $ \h -> hPutAllStrict hfs h bs+ flipFileBit hfs path bitOffset+ bs' <- withFile hfs path ReadMode $ \h -> BS.toStrict <$> hGetAll hfs h+ pure (spec_flipFileBit bs bitOffset === bs')+ where+ path = mkFsPath ["file"]++spec_flipFileBit :: ByteString -> Int -> ByteString+spec_flipFileBit bs bitOffset = runST $ do+ mv <- thaw $ cloneFromByteString bs+ flipBit mv bitOffset+ v <- unsafeFreeze mv+ pure $ cloneToByteString v++{-------------------------------------------------------------------------------+ Equality+-------------------------------------------------------------------------------}++-- | This is not a fully lawful instance, because it uses 'approximateEqStream'.+instance Eq a => Eq (Stream a) where+ (==) = approximateEqStream++instance Arbitrary a => Arbitrary (Stream a) where+ arbitrary = oneof [+ S.genFinite arbitrary+ , S.genInfinite arbitrary+ ]+ shrink s = S.liftShrinkStream shrink s++deriving stock instance Generic (Stream a)+deriving anyclass instance CoArbitrary a => CoArbitrary (Stream a)+deriving anyclass instance Function a => Function (Stream a)++deriving stock instance Generic InternalInfo+deriving anyclass instance Function InternalInfo+deriving anyclass instance CoArbitrary InternalInfo++-- | This is not a fully lawful instance, because it uses 'approximateEqStream'.+deriving stock instance Eq Errors+deriving stock instance Generic Errors+deriving anyclass instance Function Errors+deriving anyclass instance CoArbitrary Errors++deriving stock instance Generic FsErrorType+deriving anyclass instance Function FsErrorType+deriving anyclass instance CoArbitrary FsErrorType++deriving stock instance Eq Partial+deriving stock instance Generic Partial+deriving anyclass instance Function Partial+deriving anyclass instance CoArbitrary Partial++deriving stock instance Eq PutCorruption+deriving stock instance Generic PutCorruption+deriving anyclass instance Function PutCorruption+deriving anyclass instance CoArbitrary PutCorruption++deriving stock instance Eq Blob+deriving stock instance Generic Blob+deriving anyclass instance Function Blob+deriving anyclass instance CoArbitrary Blob
@@ -0,0 +1,46 @@+module Test.Util.Arbitrary (+ prop_arbitraryAndShrinkPreserveInvariant+ , prop_forAllArbitraryAndShrinkPreserveInvariant+ , deepseqInvariant+ , noInvariant+ , noTags+ ) where++import Control.DeepSeq (NFData, deepseq)+import Database.LSMTree.Extras (showPowersOf10)+import Test.QuickCheck+import Test.Tasty (TestTree)+import Test.Tasty.QuickCheck (testProperty)++prop_arbitraryAndShrinkPreserveInvariant ::+ forall a prop. (Arbitrary a, Show a, Testable prop)+ => (a -> Property -> Property) -> (a -> prop) -> [TestTree]+prop_arbitraryAndShrinkPreserveInvariant tag =+ prop_forAllArbitraryAndShrinkPreserveInvariant tag arbitrary shrink++prop_forAllArbitraryAndShrinkPreserveInvariant ::+ forall a prop. (Show a, Testable prop)+ => (a -> Property -> Property) -> Gen a -> (a -> [a]) -> (a -> prop) -> [TestTree]+prop_forAllArbitraryAndShrinkPreserveInvariant tag gen shr inv =+ [ testProperty "Arbitrary satisfies invariant" $+ forAllShrink gen shr $ \x ->+ tag x $ property $ inv x+ , testProperty "Shrinking satisfies invariant" $+ -- We don't use forallShrink here. If this property fails, it means that+ -- the shrinker is broken, so we don't want to rely on it.+ forAll gen $ \x ->+ case shr x of+ [] -> label "no shrinks" $ property True+ xs -> tabulate "number of shrinks" [showPowersOf10 (length xs)] $+ forAll (elements xs) inv -- TODO: check more than one?+ ]++-- | Trivial invariant, but checks that the value is finite+deepseqInvariant :: NFData a => a -> Bool+deepseqInvariant x = x `deepseq` True++noInvariant :: a -> Bool+noInvariant _ = True++noTags :: a -> Property -> Property+noTags _ = id
@@ -0,0 +1,634 @@+{-# OPTIONS_GHC -Wno-orphans #-}+++module Test.Util.FS (+ -- * Real file system+ withTempIOHasFS+ , withTempIOHasBlockIO+ -- * Simulated file system+ , withSimHasFS+ , withSimHasBlockIO+ -- * Simulated file system with errors+ , withSimErrorHasFS+ , withSimErrorHasBlockIO+ -- * Simulated file system properties+ , propTrivial+ , propNumOpenHandles+ , propNoOpenHandles+ , numDirEntries+ , propNumDirEntries+ , propNoDirEntries+ , propEqNumDirEntries+ , assertNoOpenHandles+ , assertNumOpenHandles+ -- * Equality+ , approximateEqStream+ -- * List directory+ , DirEntry (..)+ , listDirectoryFiles+ , listDirectoryRecursive+ , listDirectoryRecursiveFiles+ -- * Corruption+ , flipRandomBitInRandomFileHardlinkSafe+ , flipRandomBitInRandomFile+ , flipFileBitHardlinkSafe+ , flipFileBit+ , hFlipBit+ -- * Errors+ , noHCloseE+ , noRemoveFileE+ , noRemoveDirectoryRecursiveE+ , filterHGetBufSomeE+ , isFsReachedEOF+ -- * Arbitrary+ , FsPathComponent (..)+ , fsPathComponentFsPath+ , fsPathComponentString+ -- ** Modifiers+ , NoCleanupErrors (..)+ -- ** Orphans+ , isPathChar+ , pathChars+ ) where++import Control.Concurrent.Class.MonadMVar+import Control.Concurrent.Class.MonadSTM.Strict+import Control.Exception (assert)+import Control.Monad (void)+import Control.Monad.Class.MonadThrow (MonadCatch, MonadThrow)+import Control.Monad.IOSim (runSimOrThrow)+import Control.Monad.Primitive (PrimMonad)+import Data.Bit (MVector (..), flipBit)+import Data.Char (isAscii, isDigit, isLetter)+import Data.Foldable (Foldable (..), foldlM, for_)+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.List.NonEmpty as NE+import Data.Primitive.ByteArray (newPinnedByteArray, setByteArray)+import Data.Primitive.Types (sizeOf)+import Data.Set (Set)+import qualified Data.Set as Set+import qualified Data.Text as T+import Data.Traversable (for)+import Database.LSMTree.Internal.CRC32C (CRC32C (..), readFileCRC32C)+import GHC.Stack+import System.FS.API as FS+import qualified System.FS.API.Lazy as FSL+import System.FS.BlockIO.API+import System.FS.BlockIO.IO hiding (unsafeFromHasFS)+import System.FS.BlockIO.Sim (unsafeFromHasFS)+import System.FS.IO+import System.FS.Sim.Error+import System.FS.Sim.MockFS (HandleMock, MockFS, numOpenHandles,+ openHandles, pretty)+import System.FS.Sim.STM+import qualified System.FS.Sim.Stream as Stream+import System.FS.Sim.Stream (InternalInfo (..), Stream (..))+import System.IO.Temp+import Test.QuickCheck+import Test.QuickCheck.Instances ()+import Test.Util.QC (Choice, getChoice)+import Text.Printf++{-------------------------------------------------------------------------------+ Real file system+-------------------------------------------------------------------------------}++withTempIOHasFS :: FilePath -> (HasFS IO HandleIO -> IO a) -> IO a+withTempIOHasFS path action = withSystemTempDirectory path $ \dir -> do+ let hfs = ioHasFS (MountPoint dir)+ action hfs++withTempIOHasBlockIO :: FilePath -> (HasFS IO HandleIO -> HasBlockIO IO HandleIO -> IO a) -> IO a+withTempIOHasBlockIO path action = withSystemTempDirectory path $ \dir -> do+ withIOHasBlockIO (MountPoint dir) defaultIOCtxParams action++{-------------------------------------------------------------------------------+ Simulated file system+-------------------------------------------------------------------------------}++{-# INLINABLE withSimHasFS #-}+withSimHasFS ::+ (MonadSTM m, MonadThrow m, PrimMonad m, Testable prop1, Testable prop2)+ => (MockFS -> prop1)+ -> MockFS+ -> ( HasFS m HandleMock+ -> StrictTMVar m MockFS+ -> m prop2+ )+ -> m Property+withSimHasFS post fs k = do+ var <- newTMVarIO fs+ let hfs = simHasFS var+ x <- k hfs var+ fs' <- atomically $ readTMVar var+ pure (x .&&. post fs')++{-# INLINABLE withSimHasBlockIO #-}+withSimHasBlockIO ::+ (MonadMVar m, MonadSTM m, MonadCatch m, PrimMonad m, Testable prop1, Testable prop2)+ => (MockFS -> prop1)+ -> MockFS+ -> ( HasFS m HandleMock+ -> HasBlockIO m HandleMock+ -> StrictTMVar m MockFS+ -> m prop2+ )+ -> m Property+withSimHasBlockIO post fs k = do+ withSimHasFS post fs $ \hfs fsVar -> do+ hbio <- unsafeFromHasFS hfs+ k hfs hbio fsVar++{-------------------------------------------------------------------------------+ Simulated file system with errors+-------------------------------------------------------------------------------}++{-# INLINABLE withSimErrorHasFS #-}+withSimErrorHasFS ::+ (MonadSTM m, MonadThrow m, PrimMonad m, Testable prop1, Testable prop2)+ => (MockFS -> prop1)+ -> MockFS+ -> Errors+ -> ( HasFS m HandleMock+ -> StrictTMVar m MockFS+ -> StrictTVar m Errors+ -> m prop2+ )+ -> m Property+withSimErrorHasFS post fs errs k = do+ fsVar <- newTMVarIO fs+ errVar <- newTVarIO errs+ let hfs = simErrorHasFS fsVar errVar+ x <- k hfs fsVar errVar+ fs' <- atomically $ readTMVar fsVar+ pure (x .&&. post fs')++{-# INLINABLE withSimErrorHasBlockIO #-}+withSimErrorHasBlockIO ::+ ( MonadSTM m, MonadCatch m, MonadMVar m, PrimMonad m+ , Testable prop1, Testable prop2+ )+ => (MockFS -> prop1)+ -> MockFS+ -> Errors+ -> ( HasFS m HandleMock+ -> HasBlockIO m HandleMock+ -> StrictTMVar m MockFS+ -> StrictTVar m Errors+ -> m prop2+ )+ -> m Property+withSimErrorHasBlockIO post fs errs k =+ withSimErrorHasFS post fs errs $ \hfs fsVar errsVar -> do+ hbio <- unsafeFromHasFS hfs+ k hfs hbio fsVar errsVar++{-------------------------------------------------------------------------------+ Simulated file system properties+-------------------------------------------------------------------------------}++propTrivial :: MockFS -> Property+propTrivial _ = property True++{-# INLINABLE propNumOpenHandles #-}+propNumOpenHandles :: Int -> MockFS -> Property+propNumOpenHandles expected fs =+ counterexample (printf "Expected %d open handles, but found %d" expected actual) $+ counterexample ("Open handles: " <> show (openHandles fs)) $+ printMockFSOnFailure fs $+ expected == actual+ where actual = numOpenHandles fs++{-# INLINABLE propNoOpenHandles #-}+propNoOpenHandles :: MockFS -> Property+propNoOpenHandles fs = propNumOpenHandles 0 fs++numDirEntries :: FsPath -> MockFS -> Int+numDirEntries path fs = Set.size contents+ where+ (contents, _) =+ runSimOrThrow $ runSimFS fs $ \hfs -> FS.listDirectory hfs path++{-# INLINABLE propNumDirEntries #-}+propNumDirEntries :: FsPath -> Int -> MockFS -> Property+propNumDirEntries path expected fs =+ counterexample+ (printf "Expected %d entries in the directory at %s, but found %d"+ expected+ (show path) actual) $+ printMockFSOnFailure fs $+ expected === actual+ where actual = numDirEntries path fs++{-# INLINABLE propNoDirEntries #-}+propNoDirEntries :: FsPath -> MockFS -> Property+propNoDirEntries path fs = propNumDirEntries path 0 fs++{-# INLINABLE propEqNumDirEntries #-}+propEqNumDirEntries :: FsPath -> MockFS -> MockFS -> Property+propEqNumDirEntries path lhsFs rhsFs =+ counterexample+ (printf "The LHS has %d entries in the directory at %s, but the RHS has %d"+ lhs (show path) rhs) $+ printMockFSOnFailureWith "Mocked file system (LHS)" lhsFs $+ printMockFSOnFailureWith "Mocked file system (RHS)" rhsFs $+ lhs === rhs+ where+ lhs = numDirEntries path lhsFs+ rhs = numDirEntries path rhsFs++printMockFSOnFailure :: Testable prop => MockFS -> prop -> Property+printMockFSOnFailure = printMockFSOnFailureWith "Mocked file system"++printMockFSOnFailureWith :: Testable prop => String -> MockFS -> prop -> Property+printMockFSOnFailureWith s fs = counterexample (s <> ": " <> pretty fs)++assertNoOpenHandles :: HasCallStack => MockFS -> a -> a+assertNoOpenHandles fs = assertNumOpenHandles fs 0++assertNumOpenHandles :: HasCallStack => MockFS -> Int -> a -> a+assertNumOpenHandles fs m =+ assert $+ if n /= m then+ error (printf "Expected %d open handles, but found %d" m n)+ else+ True+ where n = numOpenHandles fs++{-------------------------------------------------------------------------------+ Equality+-------------------------------------------------------------------------------}++-- | Approximate equality for streams.+--+-- Equality is checked as follows:+-- * Infinite streams are equal: any infinity is as good as another infinity+-- * Finite streams are checked for pointwise equality on their elements.+-- * Other streams are trivially unequal: they do not have matching finiteness+--+-- This approximate equality satisfies the __Reflexivity__, __Symmetry__,+-- __Transitivity__ and __Negation__ laws for the 'Eq' class, but not+-- __Substitutivity.+--+-- TODO: upstream to fs-sim+approximateEqStream :: Eq a => Stream a -> Stream a -> Bool+approximateEqStream (UnsafeStream infoXs xs) (UnsafeStream infoYs ys) =+ case (infoXs, infoYs) of+ (Infinite, Infinite) -> True+ (Finite, Finite) -> xs == ys+ (_, _) -> False++{-------------------------------------------------------------------------------+ List directory+-------------------------------------------------------------------------------}++-- TODO: test the list directory functions++data DirEntry a = Directory a | File a+ deriving stock (Show, Eq, Ord, Functor)++-- | List all files and directories in the given directory and recursively in all+-- sub-directories.+listDirectoryRecursive ::+ Monad m+ => HasFS m h+ -> FsPath+ -> m (Set (DirEntry FsPath))+listDirectoryRecursive hfs = go Set.empty (mkFsPath [])+ where+ go !acc relPath absPath = do+ pcs <- listDirectory hfs absPath+ foldlM (\acc' -> go' acc' relPath absPath) acc pcs++ go' !acc relPath absPath pc = do+ let+ p = mkFsPath [pc]+ relPath' = relPath </> p+ absPath' = absPath </> p+ isFile <- doesFileExist hfs absPath'+ if isFile then+ pure (File relPath' `Set.insert` acc)+ else do+ isDirectory <- doesDirectoryExist hfs absPath'+ if isDirectory then+ go (Directory relPath' `Set.insert` acc) relPath' absPath'+ else+ error $ printf+ "listDirectoryRecursive: %s is not a file or directory"+ (show relPath')++-- | List files in the given directory and recursively in all sub-directories.+listDirectoryRecursiveFiles ::+ Monad m+ => HasFS m h+ -> FsPath+ -> m (Set FsPath)+listDirectoryRecursiveFiles hfs dir = do+ dirEntries <- listDirectoryRecursive hfs dir+ foldlM f Set.empty dirEntries+ where+ f !acc (File p) = pure $ Set.insert p acc+ f !acc _ = pure acc++-- | List files in the given directory+listDirectoryFiles ::+ Monad m+ => HasFS m h+ -> FsPath+ -> m (Set FsPath)+listDirectoryFiles hfs = go Set.empty+ where+ go !acc absPath = do+ pcs <- listDirectory hfs absPath+ foldlM go' acc pcs++ go' !acc pc = do+ let path = mkFsPath [pc]+ isFile <- doesFileExist hfs path+ if isFile then+ pure (path `Set.insert` acc)+ else+ pure acc+{-------------------------------------------------------------------------------+ Corruption+-------------------------------------------------------------------------------}++-- | Flip a random bit in a random file in a given directory.+flipRandomBitInRandomFile ::+ (PrimMonad m, MonadThrow m)+ => HasFS m h+ -> Choice+ -> FsPath+ -> m (Maybe (FsPath, Int))+flipRandomBitInRandomFile hfs bitChoice dir = do+ maybeFileBit <- pickRandomBitInRandomFile hfs bitChoice dir+ for_ maybeFileBit $ \(file, bit) -> flipFileBit hfs file bit+ pure maybeFileBit++-- | Flip a random bit in a random file in a given directory.+flipRandomBitInRandomFileHardlinkSafe ::+ (PrimMonad m, MonadThrow m)+ => HasFS m h+ -> Choice+ -> FsPath+ -> m (Maybe (FsPath, Int))+flipRandomBitInRandomFileHardlinkSafe hfs bitChoice dir = do+ maybeFileBit <- pickRandomBitInRandomFile hfs bitChoice dir+ for_ maybeFileBit $ \(file, bit) -> flipFileBitHardlinkSafe hfs file bit+ pure maybeFileBit++-- | Pick a random bit in a random file in a given directory.+pickRandomBitInRandomFile ::+ (PrimMonad m, MonadThrow m)+ => HasFS m h+ -> Choice+ -> FsPath+ -> m (Maybe (FsPath, Int))+pickRandomBitInRandomFile hfs bitChoice dir = do+ -- List all files+ files <- fmap (dir </>) . toList <$> listDirectoryRecursiveFiles hfs dir+ -- Handle the situation where there are no files+ if null files then pure Nothing else do+ filesAndFileSizeBits <-+ for files $ \file -> do+ fileSizeBytes <- withFile hfs file ReadMode (hGetSize hfs)+ pure (file, fileSizeBytes * 8)+ let totalFileSizeBits = sum (snd <$> filesAndFileSizeBits)+ -- Handle the situation where there are no non-empty files+ if totalFileSizeBits == 0 then pure Nothing else do+ assert (totalFileSizeBits > 0) $ pure ()+ -- Internal helper: find the file/bit that a choice points to.+ let pickFileBitAt bitIndex [] =+ error $ printf "flipFileBitAt: bit index out of bounds (%d)" bitIndex+ pickFileBitAt bitIndex ((file, fileSize) : filesAndSizes)+ | bitIndex < fileSize = pure (file, fromIntegral $ bitIndex `min` fromIntegral (maxBound @Int))+ | otherwise = pickFileBitAt (bitIndex - fileSize) filesAndSizes+ -- Interpret `index` to point to a bit between `0` and `totalFileSize - 1`+ let bitIndex = getChoice bitChoice (0, totalFileSizeBits - 1)+ Just <$> pickFileBitAt bitIndex filesAndFileSizeBits++-- | Flip a single bit in the given file, ensuring that it is not hardlinked.+flipFileBitHardlinkSafe ::+ (PrimMonad m, MonadThrow m)+ => HasFS m h+ -> FsPath+ -> Int -- ^ Bit offset.+ -> m ()+flipFileBitHardlinkSafe hfs fileOrig bitOffset = do+ -- Compute the CRC of fileOrig:+ CRC32C crc <- readFileCRC32C hfs fileOrig+ -- Copy fileOrig to fileCorr:+ let copyFile fileFrom fileTo =+ withFile hfs fileFrom ReadMode $ \hFrom ->+ withFile hfs fileTo (WriteMode MustBeNew) $ \hTo -> do+ bs <- FSL.hGetAll hfs hFrom+ void $ FSL.hPutAll hfs hTo bs+ let fileCorr = (fileOrig <.> show crc) <.> "corrupted"+ copyFile fileOrig fileCorr+ -- Corrupt fileCorr:+ flipFileBit hfs fileCorr bitOffset+ -- Hardlink fileCorr over fileOrig:+ removeFile hfs fileOrig+ renameFile hfs fileCorr fileOrig++-- | Flip a single bit in the given file.+flipFileBit :: (MonadThrow m, PrimMonad m) => HasFS m h -> FsPath -> Int -> m ()+flipFileBit hfs p bitOffset =+ withFile hfs p (ReadWriteMode AllowExisting) $ \h -> hFlipBit hfs h bitOffset++-- | Flip a single bit in the file pointed to by the given handle.+hFlipBit ::+ (MonadThrow m, PrimMonad m)+ => HasFS m h+ -> Handle h+ -> Int -- ^ Bit offset+ -> m ()+hFlipBit hfs h bitOffset = do+ -- Check that the bit offset is within the file+ fileSize <- hGetSize hfs h+ let fileSizeBits = 8 * fileSize+ assert (bitOffset >= 0) $ pure ()+ assert (bitOffset < fromIntegral fileSizeBits) $ pure ()+ -- Create an empty buffer initialised to all 0 bits. The buffer must have at+ -- least the size of a machine word.+ let n = sizeOf (0 :: Word)+ buf <- newPinnedByteArray n+ setByteArray buf 0 1 (0 :: Word)+ -- Read the bit at the given offset+ let (byteOffset, i) = bitOffset `quotRem` 8+ bufOff = BufferOffset 0+ count = 1+ off = AbsOffset (fromIntegral byteOffset)+ -- Check that the byte offset is within the file+ assert (byteOffset >= 0) $ pure ()+ assert (byteOffset < fromIntegral fileSize) $ pure ()+ assert (i >= 0 && i < 8) $ pure ()+ void $ hGetBufExactlyAt hfs h buf bufOff count off+ -- Flip the bit in memory, and then write it back+ let bvec = BitMVec 0 8 buf+ flipBit bvec i+ void $ hPutBufExactlyAt hfs h buf bufOff count off++{-------------------------------------------------------------------------------+ Errors+-------------------------------------------------------------------------------}++noHCloseE :: Errors -> Errors+noHCloseE errs = errs { hCloseE = Stream.empty }++noRemoveFileE :: Errors -> Errors+noRemoveFileE errs = errs { removeFileE = Stream.empty }++noRemoveDirectoryRecursiveE :: Errors -> Errors+noRemoveDirectoryRecursiveE errs = errs { removeDirectoryRecursiveE = Stream.empty }++filterHGetBufSomeE :: (Maybe (Either FsErrorType Partial) -> Bool) -> Errors -> Errors+filterHGetBufSomeE p e = e {+ hGetBufSomeE = Stream.filter p (hGetBufSomeE e)+ }++isFsReachedEOF :: FsErrorType -> Bool+isFsReachedEOF FsReachedEOF = True+isFsReachedEOF _ = False++{-------------------------------------------------------------------------------+ Arbitrary+-------------------------------------------------------------------------------}++--+-- FsPathComponent+--++-- | A single component in an 'FsPath'.+--+-- If we have a path @a/b/c/d@, then @a@, @b@ and @c@ are components, but for+-- example @a/b@ is not.+newtype FsPathComponent = FsPathComponent (NonEmpty Char)+ deriving stock (Eq, Ord)++instance Show FsPathComponent where+ show = show . fsPathComponentFsPath++fsPathComponentFsPath :: FsPathComponent -> FsPath+fsPathComponentFsPath (FsPathComponent s) = FS.mkFsPath [NE.toList s]++fsPathComponentString :: FsPathComponent -> String+fsPathComponentString (FsPathComponent s) = NE.toList s++instance Arbitrary FsPathComponent where+ arbitrary = resize 5 $ -- path components don't have to be very long+ FsPathComponent <$> liftArbitrary genPathChar+ shrink :: FsPathComponent -> [FsPathComponent]+ shrink (FsPathComponent s) = FsPathComponent <$> liftShrink shrinkPathChar s++{-------------------------------------------------------------------------------+ Arbitrary: modifiers+-------------------------------------------------------------------------------}++--+-- NoCleanupErrors+--++-- | No errors on closing file handles and removing files+newtype NoCleanupErrors = NoCleanupErrors Errors+ deriving stock Show++mkNoCleanupErrors :: Errors -> NoCleanupErrors+mkNoCleanupErrors errs = NoCleanupErrors $+ noHCloseE+ $ noRemoveFileE+ $ noRemoveDirectoryRecursiveE+ $ errs++instance Arbitrary NoCleanupErrors where+ arbitrary = do+ errs <- arbitrary+ pure $ mkNoCleanupErrors errs++ -- The shrinker for 'Errors' does not re-introduce 'hCloseE' and 'removeFile'.+ shrink (NoCleanupErrors errs) = NoCleanupErrors <$> shrink errs++{-------------------------------------------------------------------------------+ Arbitrary: orphans+-------------------------------------------------------------------------------}++instance Arbitrary FsPath where+ arbitrary = scale (`div` 10) $ -- paths don't have to be very long+ FS.mkFsPath <$> listOf (fsPathComponentString <$> arbitrary)+ shrink p =+ let ss = T.unpack <$> fsPathToList p+ in FS.mkFsPath <$> shrinkList shrinkAsComponent ss+ where+ shrinkAsComponent s = fsPathComponentString <$>+ shrink (FsPathComponent $ NE.fromList s)++-- >>> all isPathChar pathChars+-- True+isPathChar :: Char -> Bool+isPathChar c = isAscii c && (isLetter c || isDigit c)++-- >>> pathChars+-- "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"+pathChars :: [Char]+pathChars = concat [['a'..'z'], ['A'..'Z'], ['0'..'9']]++genPathChar :: Gen Char+genPathChar = elements pathChars++shrinkPathChar :: Char -> [Char]+shrinkPathChar c = [ c' | c' <- shrink c, isPathChar c']++instance Arbitrary OpenMode where+ arbitrary = genOpenMode+ shrink = shrinkOpenMode++genOpenMode :: Gen OpenMode+genOpenMode = oneof [+ pure ReadMode+ , WriteMode <$> genAllowExisting+ , ReadWriteMode <$> genAllowExisting+ , AppendMode <$> genAllowExisting+ ]+ where+ _coveredAllCases x = case x of+ ReadMode{} -> ()+ WriteMode{} -> ()+ ReadWriteMode{} -> ()+ AppendMode{} -> ()++shrinkOpenMode :: OpenMode -> [OpenMode]+shrinkOpenMode = \case+ ReadMode -> []+ WriteMode ae ->+ ReadMode+ : (WriteMode <$> shrinkAllowExisting ae)+ ReadWriteMode ae ->+ ReadMode+ : WriteMode ae+ : (ReadWriteMode <$> shrinkAllowExisting ae)+ AppendMode ae ->+ ReadMode+ : WriteMode ae+ : ReadWriteMode ae+ : (AppendMode <$> shrinkAllowExisting ae)++instance Arbitrary AllowExisting where+ arbitrary = genAllowExisting+ shrink = shrinkAllowExisting++genAllowExisting :: Gen AllowExisting+genAllowExisting = elements [+ AllowExisting+ , MustBeNew+ , MustExist+ ]+ where+ _coveredAllCases x = case x of+ AllowExisting -> ()+ MustBeNew -> ()+ MustExist -> ()++shrinkAllowExisting :: AllowExisting -> [AllowExisting]+shrinkAllowExisting AllowExisting = []+shrinkAllowExisting MustBeNew = [AllowExisting]+shrinkAllowExisting MustExist = [AllowExisting, MustBeNew]
@@ -0,0 +1,226 @@+module Test.Util.FS.Error (+ -- * Errors log+ ErrorsLog+ , emptyLog+ , countNoisyErrors+ -- * Logged HasFS and HasBlockIO+ , simErrorHasBlockIOLogged+ , simErrorHasFS+ ) where++import Control.Concurrent.Class.MonadMVar (MonadMVar)+import Control.Concurrent.Class.MonadSTM.Strict+import Control.Monad.Class.MonadThrow+import Control.Monad.Primitive+import Data.Functor.Barbie+import Data.Monoid+import GHC.Generics+import System.FS.API+import System.FS.BlockIO.API+import System.FS.BlockIO.Sim+import System.FS.Sim.Error+import System.FS.Sim.MockFS (HandleMock, MockFS)+import System.FS.Sim.Stream++{-------------------------------------------------------------------------------+ Higher-kinded datatype+-------------------------------------------------------------------------------}++-- | A higher-kinded datatype (HKD) version of the 'Errors' type. 'Errors' is+-- equivalent to @'HKD' 'Stream'@.+--+-- 'HKD' has instances for @barbies@ classes, which lets us manipulate 'HKD'+-- with a small set of combinators like 'bpure' and 'bmap'. This makes writing+-- functions operating on 'HKD' less arduous than writing the functions out+-- fully.+--+-- TODO: admittedly, we are not getting the full benefis from higher-kinded+-- datatypes because 'Errors' and 'HasFS' are not HKDs, meaning we still have to+-- write some stuff out manually, like in 'simErrorHasFSLogged'. We could take+-- the HKD approach even further, but it is likely going to require quite a bit+-- of boilerplate to set up before we can start programming with HKDs. At that+-- point one might ask whether the costs of defining the HKD boilerplate+-- outweighs the benefits of programming with HKDs.+data HKD f = HKD {+ dumpStateHKD :: f FsErrorType+ , hOpenHKD :: f FsErrorType+ , hCloseHKD :: f FsErrorType+ , hSeekHKD :: f FsErrorType+ , hGetSomeHKD :: f (Either FsErrorType Partial)+ , hGetSomeAtHKD :: f (Either FsErrorType Partial)+ , hPutSomeHKD :: f (Either (FsErrorType, Maybe PutCorruption) Partial)+ , hTruncateHKD :: f FsErrorType+ , hGetSizeHKD :: f FsErrorType+ , createDirectoryHKD :: f FsErrorType+ , createDirectoryIfMissingHKD:: f FsErrorType+ , listDirectoryHKD :: f FsErrorType+ , doesDirectoryExistHKD :: f FsErrorType+ , doesFileExistHKD :: f FsErrorType+ , removeDirectoryRecursiveHKD:: f FsErrorType+ , removeFileHKD :: f FsErrorType+ , renameFileHKD :: f FsErrorType+ , hGetBufSomeHKD :: f (Either FsErrorType Partial)+ , hGetBufSomeAtHKD :: f (Either FsErrorType Partial)+ , hPutBufSomeHKD :: f (Either (FsErrorType, Maybe PutCorruption) Partial)+ , hPutBufSomeAtHKD :: f (Either (FsErrorType, Maybe PutCorruption) Partial)+ }+ deriving stock Generic++deriving anyclass instance FunctorB HKD+deriving anyclass instance ApplicativeB HKD+deriving anyclass instance TraversableB HKD+deriving anyclass instance ConstraintsB HKD++{-------------------------------------------------------------------------------+ Errors log+-------------------------------------------------------------------------------}++type ErrorsLog = HKD []++emptyLog :: ErrorsLog+emptyLog = bpure []++{-------------------------------------------------------------------------------+ Count noisy errors in the log+-------------------------------------------------------------------------------}++class IsNoisy a where+ isNoisy :: a -> Bool++instance IsNoisy FsErrorType where+ isNoisy _ = True++instance IsNoisy (Either FsErrorType b) where+ isNoisy (Left x) = isNoisy x+ isNoisy (Right _) = False++instance IsNoisy (Either (FsErrorType, b) c) where+ isNoisy (Left (x, _)) = isNoisy x+ isNoisy (Right _) = False++countNoisyErrors :: ErrorsLog -> Int+countNoisyErrors = getSum . bfoldMapC @IsNoisy (Sum . length . Prelude.filter isNoisy)++{-------------------------------------------------------------------------------+ Logged HasFS and HasBlockIO+-------------------------------------------------------------------------------}++-- | Like 'simErrorHasFSLogged', but also produces a simulated 'HasBlockIO'.+simErrorHasBlockIOLogged ::+ forall m. (MonadCatch m, MonadMVar m, PrimMonad m, MonadSTM m)+ => StrictTMVar m MockFS+ -> StrictTVar m Errors+ -> StrictTVar m (HKD [])+ -> m (HasFS m HandleMock, HasBlockIO m HandleMock)+simErrorHasBlockIOLogged fsVar errorsVar logVar = do+ let hfs = simErrorHasFSLogged fsVar errorsVar logVar+ hbio <- unsafeFromHasFS hfs+ pure (hfs, hbio)++-- | Produce a simulated file system with injected errors and a logger for those+-- errors.+--+-- Every time a 'HasFS' primitive is used and an error from 'Errors' is used, it+-- will be logged in 'ErrorsLog'.+simErrorHasFSLogged ::+ forall m. (MonadSTM m, MonadThrow m, PrimMonad m)+ => StrictTMVar m MockFS+ -> StrictTVar m Errors+ -> StrictTVar m ErrorsLog+ -> HasFS m HandleMock+simErrorHasFSLogged fsVar errorsVar logVar =+ HasFS {+ dumpState = do+ addToLog dumpStateE dumpStateHKD (\l x -> l { dumpStateHKD = x} )+ dumpState hfs+ , hOpen = \a b -> do+ addToLog hOpenE hOpenHKD (\l x -> l { hOpenHKD = x} )+ hOpen hfs a b+ , hClose = \a -> do+ addToLog hCloseE hCloseHKD (\l x -> l { hCloseHKD = x} )+ hClose hfs a+ , hIsOpen = \a ->+ hIsOpen hfs a+ , hSeek = \a b c -> do+ addToLog hSeekE hSeekHKD (\l x -> l { hSeekHKD = x} )+ hSeek hfs a b c+ , hGetSome = \a b -> do+ addToLog hGetSomeE hGetSomeHKD (\l x -> l { hGetSomeHKD = x} )+ hGetSome hfs a b+ , hGetSomeAt = \a b c -> do+ addToLog hGetSomeAtE hGetSomeAtHKD (\l x -> l { hGetSomeAtHKD = x} )+ hGetSomeAt hfs a b c+ , hPutSome = \a b -> do+ addToLog hPutSomeE hPutSomeHKD (\l x -> l { hPutSomeHKD = x} )+ hPutSome hfs a b+ , hTruncate = \a b -> do+ addToLog hTruncateE hTruncateHKD (\l x -> l { hTruncateHKD = x} )+ hTruncate hfs a b+ , hGetSize = \a -> do+ addToLog hGetSizeE hGetSizeHKD (\l x -> l { hGetSizeHKD = x} )+ hGetSize hfs a+ , createDirectory = \a -> do+ addToLog createDirectoryE createDirectoryHKD (\l x -> l { createDirectoryHKD = x} )+ createDirectory hfs a+ , createDirectoryIfMissing = \a b -> do+ addToLog createDirectoryIfMissingE createDirectoryIfMissingHKD (\l x -> l { createDirectoryIfMissingHKD = x} )+ createDirectoryIfMissing hfs a b+ , listDirectory = \a -> do+ addToLog listDirectoryE listDirectoryHKD (\l x -> l { listDirectoryHKD = x} )+ listDirectory hfs a+ , doesDirectoryExist = \a -> do+ addToLog doesDirectoryExistE doesDirectoryExistHKD (\l x -> l { doesDirectoryExistHKD = x} )+ doesDirectoryExist hfs a+ , doesFileExist = \a -> do+ addToLog doesFileExistE doesFileExistHKD (\l x -> l { doesFileExistHKD = x} )+ doesFileExist hfs a+ , removeDirectoryRecursive = \a -> do+ addToLog removeDirectoryRecursiveE removeDirectoryRecursiveHKD (\l x -> l { removeDirectoryRecursiveHKD = x} )+ removeDirectoryRecursive hfs a+ , removeFile = \a -> do+ addToLog removeFileE removeFileHKD (\l x -> l { removeFileHKD = x} )+ removeFile hfs a+ , renameFile = \a b -> do+ addToLog renameFileE renameFileHKD (\l x -> l { renameFileHKD = x} )+ renameFile hfs a b+ , mkFsErrorPath = mkFsErrorPath hfs+ , unsafeToFilePath = unsafeToFilePath hfs+ , hGetBufSome = \a b c d -> do+ addToLog hGetBufSomeE hGetBufSomeHKD (\l x -> l { hGetBufSomeHKD = x} )+ hGetBufSome hfs a b c d+ , hGetBufSomeAt = \a b c d e -> do+ addToLog hGetBufSomeAtE hGetBufSomeAtHKD (\l x -> l { hGetBufSomeAtHKD = x} )+ hGetBufSomeAt hfs a b c d e+ , hPutBufSome = \a b c d -> do+ addToLog hPutBufSomeE hPutBufSomeHKD (\l x -> l { hPutBufSomeHKD = x} )+ hPutBufSome hfs a b c d+ , hPutBufSomeAt = \a b c d e -> do+ addToLog hPutBufSomeAtE hPutBufSomeAtHKD (\l x -> l { hPutBufSomeAtHKD = x} )+ hPutBufSomeAt hfs a b c d e+ }+ where+ hfs = simErrorHasFS fsVar errorsVar++ -- Peek at the first element from an error stream and add it to the errors+ -- log if it is 'Just' an error.+ addToLog ::+ (Errors -> Stream e)+ -> (ErrorsLog -> [e])+ -> (ErrorsLog -> [e] -> ErrorsLog)+ -> m ()+ addToLog getE getL setL = do+ errors <- readTVarIO errorsVar+ logs <- readTVarIO logVar+ let s = getE errors+ let x = peek s+ case x of+ Nothing -> pure ()+ Just y ->+ atomically $ writeTVar logVar $+ setL logs $ (++ [y]) $ getL logs++ -- Peek at the first element from an error stream.+ peek :: Stream a -> Maybe a+ peek (UnsafeStream _ xs) = case xs of+ [] -> Nothing+ (x:_) -> x
@@ -0,0 +1,23 @@+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE StandaloneDeriving #-}+{-# LANGUAGE StandaloneKindSignatures #-}+{-# LANGUAGE TypeFamilies #-}++{-# OPTIONS_GHC -Wno-orphans #-}++module Test.Util.Orphans () where++import Database.LSMTree (SerialiseKey, SerialiseValue)+import Test.QuickCheck.Modifiers (Small (..))++{-------------------------------------------------------------------------------+ QuickCheck+-------------------------------------------------------------------------------}++deriving newtype instance SerialiseKey a => SerialiseKey (Small a)+deriving newtype instance SerialiseValue a => SerialiseValue (Small a)
@@ -0,0 +1,19 @@+{-# LANGUAGE PolyKinds #-}++module Test.Util.PrettyProxy (+ PrettyProxy (..)+ ) where++import Data.Typeable++-- | A version of 'Proxy' that also shows its type parameter.+data PrettyProxy a = PrettyProxy++-- | Shows the type parameter @a@ as an explicit type application.+instance Typeable a => Show (PrettyProxy a) where+ showsPrec d p =+ showParen (d > app_prec)+ $ showString "PrettyProxy @("+ . showString (show (typeRep p))+ . showString ")"+ where app_prec = 10
@@ -0,0 +1,48 @@+module Test.Util.QC (+ testClassLaws+ , testClassLawsWith+ , Proxy (..)+ , Choice+ , getChoice+ ) where++import Data.Proxy (Proxy (..))+import Data.Word (Word64)+import Test.QuickCheck.Classes (Laws (..))+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.QuickCheck (Arbitrary (..), Property,+ arbitraryBoundedIntegral, shrinkIntegral, testProperty)++testClassLaws :: String -> Laws -> TestTree+testClassLaws typename laws = testClassLawsWith typename laws testProperty++testClassLawsWith ::+ String -> Laws+ -> (String -> Property -> TestTree)+ -> TestTree+testClassLawsWith typename Laws {lawsTypeclass, lawsProperties} k =+ testGroup ("class laws" ++ lawsTypeclass ++ " " ++ typename)+ [ k name prop+ | (name, prop) <- lawsProperties ]+++-- | A 'Choice' of a uniform random number in a range where shrinking picks smaller numbers.+newtype Choice = Choice Word64+ deriving stock (Show, Eq)++instance Arbitrary Choice where+ arbitrary = Choice <$> arbitraryBoundedIntegral+ shrink (Choice x) = Choice <$> shrinkIntegral x++-- | Use a 'Choice' to get a concrete 'Integral' in range @(a, a)@ inclusive.+--+-- The choice of integral is uniform as long as the range is smaller than or+-- equal to the maximum bound of `Word64`, i.e., 18446744073709551615.+getChoice :: (Integral a) => Choice -> (a, a) -> a+getChoice (Choice n) (l, u) = fromIntegral (((ni * (ui - li)) `div` mi) + li)+ where+ ni = toInteger n+ li = toInteger l+ ui = toInteger u+ mi = toInteger (maxBound :: Word64)+
@@ -0,0 +1,83 @@+{-# LANGUAGE QuantifiedConstraints #-}+{-# LANGUAGE TypeFamilies #-}++-- | Utilities for the @quickcheck-lockstep@ package+--+-- TODO: it might be nice to upstream these utilities to @quickcheck-lockstep@+-- at some point.+module Test.Util.QLS (+ IOProperty (..)+ , runActionsBracket+ ) where++import Prelude hiding (init)++import Control.Monad (void)+import Control.Monad.Class.MonadST+import Control.Monad.Class.MonadThrow+import Control.Monad.IOSim (IOSim, runSimTraceST, traceResult)+import Control.Monad.Primitive (RealWorld)+import qualified Control.Monad.ST.Lazy as ST+import Data.Typeable+import qualified Test.QuickCheck as QC+import Test.QuickCheck (Property, Testable)+import Test.QuickCheck.Monadic+import qualified Test.QuickCheck.StateModel as StateModel+import Test.QuickCheck.StateModel hiding (runActions)+import Test.QuickCheck.StateModel.Lockstep++class IOProperty m where+ ioProperty :: m Property -> Property++instance IOProperty IO where+ ioProperty = QC.ioProperty++instance IOProperty (IOSim RealWorld) where+ ioProperty x = QC.ioProperty $ do+ trac <- ST.stToIO $ runSimTraceST x+ case traceResult False trac of+ Left e -> pure $ QC.counterexample (show e) False+ Right p -> pure p++runActionsBracket ::+ ( RunLockstep state m+ , e ~ Error (Lockstep state) m+ , forall a. IsPerformResult e a+ , Testable prop+ , MonadMask n+ , MonadST n+ , IOProperty n+ )+ => Proxy state+ -> n st -- ^ Initialisation+ -> (st -> n prop) -- ^ Cleanup+ -> (m Property -> st -> n Property) -- ^ Runner+ -> Actions (Lockstep state)+ -> Property+runActionsBracket _ init cleanup runner actions =+ monadicBracketIO init cleanup runner $+ void $ StateModel.runActions actions++ioPropertyBracket ::+ (Testable a, Testable b, MonadMask n, MonadST n, IOProperty n)+ => n st+ -> (st -> n b)+ -> (m a -> st -> n a)+ -> m a+ -> Property+ioPropertyBracket init cleanup runner prop =+ ioProperty $ mask $ \restore -> do+ st <- init+ a <- restore (runner prop st) `onException` cleanup st+ b <- cleanup st+ pure $ a QC..&&. b++monadicBracketIO :: forall st a b m n.+ (Monad m, Testable a, Testable b, MonadMask n, MonadST n, IOProperty n)+ => n st+ -> (st -> n b)+ -> (m Property -> st -> n Property)+ -> PropertyM m a+ -> Property+monadicBracketIO init cleanup runner =+ monadic (ioPropertyBracket init cleanup runner)
@@ -0,0 +1,103 @@+module Test.Util.RawPage (+ assertEqualRawPages,+ propEqualRawPages,+) where++import Control.Monad (unless)+import Data.Align (align)+import Data.List.Split (chunksOf)+import Data.These (These (..))+import Data.Word (Word8)+import qualified System.Console.ANSI as ANSI++import Database.LSMTree.Internal.BitMath (div16, mod16)+import qualified Database.LSMTree.Internal.RawBytes as RB+import Database.LSMTree.Internal.RawPage (RawPage, rawPageRawBytes)++import Test.Tasty.HUnit (Assertion, assertFailure)+import Test.Tasty.QuickCheck (Property, counterexample)++assertEqualRawPages :: RawPage -> RawPage -> Assertion+assertEqualRawPages a b = unless (a == b) $ do+ assertFailure $ "unequal pages:\n" ++ ANSI.setSGRCode [ANSI.Reset] ++ compareBytes (RB.unpack (rawPageRawBytes a)) (RB.unpack (rawPageRawBytes b))++propEqualRawPages :: RawPage -> RawPage -> Property+propEqualRawPages a b = counterexample+ (ANSI.setSGRCode [ANSI.Reset] ++ compareBytes (RB.unpack (rawPageRawBytes a)) (RB.unpack (rawPageRawBytes b)))+ (a == b)++-- Print two bytestreams next to each other highlighting the differences.+compareBytes :: [Word8] -> [Word8] -> String+compareBytes xs ys = unlines $ go (grouping chunks)+ where+ go :: [Either [()] [([Word8], [Word8])]] -> [String]+ go [] = []+ go (Left _ : zs) = "..." : go zs+ go (Right diff : zs) = map (uncurry showDiff) diff ++ go zs++ showDiff :: [Word8] -> [Word8] -> String+ showDiff = aux id id where+ aux :: ShowS -> ShowS -> [Word8] -> [Word8] -> String+ aux accl accr [] [] = accl . showString " " . accr $ ""+ aux accl accr [] rs = accl . showString " " . accr . green (concatMapS showsWord8 rs) $ ""+ aux accl accr ls [] = accl . red (concatMapS showsWord8 ls) . showString " " . accr $ ""+ aux accl accr (l:ls) (r:rs)+ | l == r = aux (accl . showsWord8 l) (accr . showsWord8 r) ls rs+ | otherwise = aux (accl . red (showsWord8 l)) (accr . green (showsWord8 r)) ls rs++ -- chunks are either equal, or not+ chunks :: [Either () ([Word8], [Word8])]+ chunks =+ [ case b of+ These x y+ | x == y -> Left ()+ | otherwise -> Right (x, y)+ This x -> Right (x, [])+ That y -> Right ([], y)+ | b <- align (chunksOf 16 xs) (chunksOf 16 ys)+ ]++sgr :: [ANSI.SGR] -> ShowS+sgr cs = (ANSI.setSGRCode cs ++)++red :: ShowS -> ShowS+red s = sgr [ANSI.SetColor ANSI.Foreground ANSI.Vivid ANSI.Red] . s . sgr [ANSI.Reset]++green :: ShowS -> ShowS+green s = sgr [ANSI.SetColor ANSI.Foreground ANSI.Vivid ANSI.Green] . s . sgr [ANSI.Reset]++concatMapS :: (a -> ShowS) -> [a] -> ShowS+concatMapS f xs = \acc -> foldr (\a acc' -> f a acc') acc xs++showsWord8 :: Word8 -> ShowS+showsWord8 w = \acc -> hexdigit (div16 w) : hexdigit (mod16 w) : acc++grouping :: [Either a b] -> [Either [a] [b]]+grouping = foldr add []+ where+ add (Left x) [] = [Left [x]]+ add (Right y) [] = [Right [y]]+ add (Left x) (Left xs : rest) = Left (x : xs) : rest+ add (Right y) rest@(Left _ : _) = Right [y] : rest+ add (Left x) rest@(Right _ : _ ) = Left [x] : rest+ add (Right y) (Right ys : rest) = Right (y:ys) : rest+++hexdigit :: (Num a, Eq a) => a -> Char+hexdigit 0x0 = '0'+hexdigit 0x1 = '1'+hexdigit 0x2 = '2'+hexdigit 0x3 = '3'+hexdigit 0x4 = '4'+hexdigit 0x5 = '5'+hexdigit 0x6 = '6'+hexdigit 0x7 = '7'+hexdigit 0x8 = '8'+hexdigit 0x9 = '9'+hexdigit 0xA = 'a'+hexdigit 0xB = 'b'+hexdigit 0xC = 'c'+hexdigit 0xD = 'd'+hexdigit 0xE = 'e'+hexdigit 0xF = 'f'+hexdigit _ = '?'
@@ -0,0 +1,62 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE StandaloneDeriving #-}+{-# LANGUAGE StandaloneKindSignatures #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Type family wrappers are useful for a variety of reasons:+--+-- * Type families can not be partially applied, but type wrappers can.+--+-- * Type family synonyms can not appear in a class head, but type wrappers can.+--+-- * Wrappers can be used to direct type family reduction. For an example, see+-- the uses of 'WrapTable' and co in the definition of 'RealizeIOSim',+-- which can be found in "Test.Util.Orphans".+module Test.Util.TypeFamilyWrappers (+ WrapSession (..)+ , WrapTable (..)+ , WrapCursor (..)+ , WrapBlobRef (..)+ , WrapBlob (..)+ ) where++import Data.Kind (Type)+import qualified Database.LSMTree.Class as SUT.Class++type WrapSession ::+ ((Type -> Type) -> Type -> Type -> Type -> Type)+ -> (Type -> Type) -> Type+newtype WrapSession h m = WrapSession {+ unwrapSession :: SUT.Class.Session h m+ }++type WrapTable ::+ ((Type -> Type) -> Type -> Type -> Type -> Type)+ -> (Type -> Type) -> Type -> Type -> Type -> Type+newtype WrapTable h m k v b = WrapTable {+ unwrapTable :: h m k v b+ }+ deriving stock (Show, Eq)++type WrapCursor ::+ ((Type -> Type) -> Type -> Type -> Type -> Type)+ -> (Type -> Type) -> Type -> Type -> Type -> Type+newtype WrapCursor h m k v b = WrapCursor {+ unwrapCursor :: SUT.Class.Cursor h m k v b+ }++type WrapBlobRef ::+ ((Type -> Type) -> Type -> Type -> Type -> Type)+ -> (Type -> Type) -> Type -> Type+newtype WrapBlobRef h m b = WrapBlobRef {+ unwrapBlobRef :: SUT.Class.BlobRef h m b+ }++deriving stock instance Show (SUT.Class.BlobRef h m b) => Show (WrapBlobRef h m b)+deriving stock instance Eq (SUT.Class.BlobRef h m b) => Eq (WrapBlobRef h m b)++type WrapBlob :: Type -> Type+newtype WrapBlob b = WrapBlob {+ unwrapBlob :: b+ }+ deriving stock (Show, Eq)
@@ -0,0 +1,468 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE DerivingStrategies #-}+{-# LANGUAGE GeneralisedNewtypeDeriving #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# OPTIONS_GHC -fspecialize-aggressively #-}+module Main (main) where++import Control.DeepSeq (NFData (..), force)+import Control.Exception (evaluate)+import Control.Monad.ST.Strict (ST, runST)+import Data.Bits (unsafeShiftR)+import qualified Data.Heap as Heap+import Data.IORef+import qualified Data.List as L+import Data.List.NonEmpty (NonEmpty (..))+import Data.Primitive.ByteArray (compareByteArrays)+import qualified Data.Vector.Primitive as VP+import Data.Word (Word64, Word8)+import System.IO.Unsafe (unsafePerformIO)+import qualified System.Random.SplitMix as SM+import Test.Tasty (TestName, TestTree, defaultMainWithIngredients,+ testGroup)+import qualified Test.Tasty.Bench as B+import Test.Tasty.HUnit (testCase, (@?=))+import Test.Tasty.QuickCheck (testProperty, (===))++import qualified KMerge.Heap as K.Heap+import qualified KMerge.LoserTree as K.Tree+++-- tests and benchmarks for various k-way merge implementations.+-- in short: loser tree is optimal in comparison counts performed,+-- but mutable heap implementation has lower constant factors.+--+-- Noteworthy, maybe not obvious observations:+-- - mutable heap does a similar amount of comparisons as persistent heap+-- (from @heaps@ package) on full trees with evenly sized input lists,+-- but performs more comparisons when these constraints get lifted.+-- - tree-shaped iterative two-way merge performs optimal amount of comparisons+-- loser tree is an explicit state variant of that.+-- - on skewed input sizes, the heap does benefit a little, as two consecutive+-- outputs often come from the same input, which then is already at the root,+-- only requiring two comparisons. sadly, this benefit does not translate as+-- quite as nicely to the mutable implementation.+-- - the loser tree with its balanced tree structure is not optimal for skewed+-- merges, but it can be if the tree structure is managed explicitly.+-- for a hacky proof of concept, see 'loserTreeMerge\'', where we make sure+-- that one side of the tree only contains the large input plus dummy inputs,+-- allowing a path to the root using a single comparison.+-- - 'listMerge' performs very well for skewed inputs since it merges the first+-- (i.e. long) input only once. If the last input is largest, it gets very bad.+--+main :: IO ()+main = do+ _ <- evaluate $ force input8+ _ <- evaluate $ force input7+ _ <- evaluate $ force input5+++ defaultMainWithIngredients B.benchIngredients $ testGroup "kmerge"+ [ testGroup "tests"+ [ testGroup "merge"+ [ mergeProperty "listMerge" listMerge+ , mergeProperty "treeMerge" treeMerge+ , mergeProperty "heapMerge" heapMerge+ , mergeProperty "loserTreeMerge" loserTreeMerge+ , mergeProperty "mutHeapMerge" mutHeapMerge+ ]+ , testGroup "count"+ [ testGroup "eight"+ -- loserTree comparison upper bounds for 8 inputs is 3 x element count.+ -- for 8 100-element lists, i.e. 800 elements the total comparison count is 2400+ -- loserTree (and tree merge) implementations hit exactly that number.+ --+ -- (because the input values are unformly random,+ -- there shouldn't be a lot of "cheap" leftovers elements,+ -- i.e. when other inputs are exhausted, but there are few)+ [ testCount "sortConcat" comparisons8 (L.sort . concat) input8+ , testCount "listMerge" 3479 listMerge input8+ , testCount "treeMerge" 2391 treeMerge input8+ , testCount "heapMerge" 3168 heapMerge input8+ , testCount "loserTreeMerge" 2391 loserTreeMerge input8+ , testCount "mutHeapMerge" 3169 mutHeapMerge input8+ ]+ -- seven inputs: we have 6x100 elements with 3 comparisons+ -- and 1x100 elements with just 2.+ -- i.e. target is 2000 total comparisons.+ --+ -- The difference here and in five-input case between+ -- treeMerge and loserTreeMerge is caused by+ -- different "tournament bracket" assignments done by the+ -- algorithms.+ --+ -- In particular in five case, the treeMerge bracket looks like+ --+ -- *+ -- / \+ -- * 5+ -- / \+ -- * *+ -- / \ / \+ -- 1 2 3 4+ --+ -- But the LoserTree is balanced:+ --+ -- *+ -- / \+ -- * *+ -- / \ / \+ -- * 3 4 5+ -- / \+ -- 1 2+ --+ -- (maybe treeMerge can be better balanced too,+ -- but I'm too lazy to think how to do that)+ --+ , testGroup "seven"+ [ testCount "sortConcat" comparisons7 (L.sort . concat) input7+ , testCount "listMerge" 2682 listMerge input7+ , testCount "treeMerge" 1992 treeMerge input7+ , testCount "heapMerge" 2645 heapMerge input7+ , testCount "loserTreeMerge" 1989 loserTreeMerge input7+ , testCount "mutHeapMerge" 2570 mutHeapMerge input7+ ]+ -- five inputs: we have 3x100 elements with 2 comparisons+ -- and 2x100 with 3 comparisons.+ -- i.e. target is 1200 total comparisons.+ , testGroup "five"+ [ testCount "sortConcat" comparisons5 (L.sort . concat) input5+ , testCount "listMerge" 1389 listMerge input5+ , testCount "treeMerge" 1291 treeMerge input5+ , testCount "heapMerge" 1485 heapMerge input5+ , testCount "loserTreeMerge" 1191 loserTreeMerge input5+ , testCount "mutHeapMerge" 1592 mutHeapMerge input5+ ]+ -- minimal skew for a levelling merge of 1000 elements.+ -- with a tree that gives the long input a short path:+ -- 1x500 elements with 1 comparison+ -- 4x125 elements with 3 comparisons+ -- i.e. target is 2000 total comparisons.+ , testGroup "levelling-min"+ [ testCount "sortConcat" comparisonsMin (L.sort . concat) inputLevellingMin+ , testCount "listMerge" 2112 listMerge inputLevellingMin+ , testCount "treeMerge" 2730 treeMerge inputLevellingMin+ , testCount "heapMerge" 2655 heapMerge inputLevellingMin+ , testCount "loserTreeMerge" 2235 loserTreeMerge inputLevellingMin+ , testCount "loserTreeMerge'" 1999 loserTreeMerge inputLevellingMin'+ , testCount "mutHeapMerge" 3021 mutHeapMerge inputLevellingMin+ ]+ -- maximal skew for a levelling merge of 1000 elements.+ -- with a tree that gives the long input a short path:+ -- 1x800 elements with 1 comparison+ -- 4x 50 elements with 3 comparisons+ -- i.e. target is 1400 total comparisons.+ , testGroup "levelling-max"+ [ testCount "sortConcat" comparisonsMax (L.sort . concat) inputLevellingMax+ , testCount "listMerge" 1440 listMerge inputLevellingMax+ , testCount "treeMerge" 2873 treeMerge inputLevellingMax+ , testCount "heapMerge" 1784 heapMerge inputLevellingMax+ , testCount "loserTreeMerge" 2081 loserTreeMerge inputLevellingMax+ , testCount "loserTreeMerge'" 1400 loserTreeMerge inputLevellingMax'+ , testCount "mutHeapMerge" 2493 mutHeapMerge inputLevellingMax+ ]+ ]+ ]+#ifdef KMERGE_BENCHMARKS+ , testGroup "bench"+ [ testGroup "eight"+ [ B.bench "sortConcat" $ B.nf (L.sort . concat) input8+ , B.bench "listMerge" $ B.nf listMerge input8+ , B.bench "treeMerge" $ B.nf treeMerge input8+ , B.bench "heapMerge" $ B.nf heapMerge input8+ , B.bench "loserTreeMerge" $ B.nf loserTreeMerge input8+ , B.bench "mutHeapMerge" $ B.nf mutHeapMerge input8+ ]+ , testGroup "seven"+ [ B.bench "sortConcat" $ B.nf (L.sort . concat) input7+ , B.bench "listMerge" $ B.nf listMerge input7+ , B.bench "treeMerge" $ B.nf treeMerge input7+ , B.bench "heapMerge" $ B.nf heapMerge input7+ , B.bench "loserTreeMerge" $ B.nf loserTreeMerge input7+ , B.bench "mutHeapMerge" $ B.nf mutHeapMerge input7+ ]+ , testGroup "five"+ [ B.bench "sortConcat" $ B.nf (L.sort . concat) input5+ , B.bench "listMerge" $ B.nf listMerge input5+ , B.bench "treeMerge" $ B.nf treeMerge input5+ , B.bench "heapMerge" $ B.nf heapMerge input5+ , B.bench "loserTreeMerge" $ B.nf loserTreeMerge input5+ , B.bench "mutHeapMerge" $ B.nf mutHeapMerge input5+ ]+ , testGroup "levelling-min"+ [ B.bench "sortConcat" $ B.nf (L.sort . concat) inputLevellingMin+ , B.bench "listMerge" $ B.nf listMerge inputLevellingMin+ , B.bench "treeMerge" $ B.nf treeMerge inputLevellingMin+ , B.bench "heapMerge" $ B.nf heapMerge inputLevellingMin+ , B.bench "loserTreeMerge" $ B.nf loserTreeMerge inputLevellingMin+ , B.bench "loserTreeMerge'"$ B.nf loserTreeMerge inputLevellingMin'+ , B.bench "mutHeapMerge" $ B.nf mutHeapMerge inputLevellingMin+ ]+ , testGroup "levelling-max"+ [ B.bench "sortConcat" $ B.nf (L.sort . concat) inputLevellingMax+ , B.bench "listMerge" $ B.nf listMerge inputLevellingMax+ , B.bench "treeMerge" $ B.nf treeMerge inputLevellingMax+ , B.bench "heapMerge" $ B.nf heapMerge inputLevellingMax+ , B.bench "loserTreeMerge" $ B.nf loserTreeMerge inputLevellingMax+ , B.bench "loserTreeMerge'"$ B.nf loserTreeMerge inputLevellingMax'+ , B.bench "mutHeapMerge" $ B.nf mutHeapMerge inputLevellingMax+ ]+ ]+#endif+ ]++{-------------------------------------------------------------------------------+ Test utils+-------------------------------------------------------------------------------}++counter :: IORef Int+counter = unsafePerformIO $ newIORef 0+{-# NOINLINE counter #-}++newtype Wrapped a = Wrap a -- { unwrap :: Word256 }++instance Eq a => Eq (Wrapped a) where+ Wrap x == Wrap y = unsafePerformIO $ do+ atomicModifyIORef' counter $ \n -> (1 + n, ())+ pure $! x == y+ {-# NOINLINE (==) #-}++instance Ord a => Ord (Wrapped a) where+ compare (Wrap x) (Wrap y) = unsafePerformIO $ do+ atomicModifyIORef' counter $ \n -> (1 + n, ())+ pure $! compare x y+ Wrap x < Wrap y = unsafePerformIO $ do+ atomicModifyIORef' counter $ \n -> (1 + n, ())+ pure $! x < y+ Wrap x <= Wrap y = unsafePerformIO $ do+ atomicModifyIORef' counter $ \n -> (1 + n, ())+ pure $! x <= y++ {-# NOINLINE compare #-}+ {-# NOINLINE (<) #-}+ {-# NOINLINE (<=) #-}++instance NFData a => NFData (Wrapped a) where+ rnf (Wrap x) = rnf x++testCount :: (NFData b, Ord b) => TestName -> Int -> (forall a. Ord a => [[a]] -> [a]) -> [[b]] -> TestTree+testCount name expected f input = testCase name $ do+ n <- readIORef counter+ _ <- evaluate $ force $ f $ map (map Wrap) input+ m <- readIORef counter+ m - n @?= expected+{-# NOINLINE testCount #-}++mergeProperty :: TestName -> (forall a. Ord a => [[a]] -> [a]) -> TestTree+mergeProperty name f = testProperty name $ \xss ->+ let lhs = L.sort (concat xss)+ rhs = f $ map L.sort (xss :: [[Word64]])+ in lhs === rhs++{-------------------------------------------------------------------------------+ Element type+-------------------------------------------------------------------------------}++-- This type corresponds to the @SerialisedKey@ type we are using (or rather the+-- @RawBytes@ it wraps), so the cost of comparisons should be similar.+-- We expect key lengths of 32 bytes.+newtype Element = Element (VP.Vector Word8)+ deriving newtype (Show, NFData)++instance Eq Element where+ bs1 == bs2 = compareBytes bs1 bs2 == EQ++-- | Lexicographical 'Ord' instance.+instance Ord Element where+ compare = compareBytes++-- | Based on @Ord 'ShortByteString'@.+compareBytes :: Element -> Element -> Ordering+compareBytes rb1@(Element vec1) rb2@(Element vec2) =+ let !len1 = sizeofElement rb1+ !len2 = sizeofElement rb2+ !len = min len1 len2+ in case compareByteArrays ba1 off1 ba2 off2 len of+ EQ | len1 < len2 -> LT+ | len1 > len2 -> GT+ o -> o+ where+ VP.Vector off1 _size1 ba1 = vec1+ VP.Vector off2 _size2 ba2 = vec2++sizeofElement :: Element -> Int+sizeofElement (Element pvec) = VP.length pvec++genElement :: SM.SMGen -> (Element, SM.SMGen)+genElement g0 = (Element (VP.fromListN 32 bytes), g4)+ where+ -- we expect a shared 16 bit prefix+ bytes = 0 : 0 : concatMap toBytes [w1, w2, w3, w4]+ (!w1, g1) = SM.nextWord64 g0+ (!w2, g2) = SM.nextWord64 g1+ (!w3, g3) = SM.nextWord64 g2+ (!w4, g4) = SM.nextWord64 g3+ toBytes = reverse . take 8 . map fromIntegral . iterate (`unsafeShiftR` 8)++minElement :: Element+minElement = Element (VP.fromListN 32 (L.repeat 0))++{-------------------------------------------------------------------------------+ Inputs+-------------------------------------------------------------------------------}++input8 :: [[Element]]+input8 = take 8 $ inputs 100++-- Seven inputs is not optimal case for "binary tree" patterns.+input7 :: [[Element]]+input7 = take 7 $ inputs 100++-- Five inputs is bad case for "binary tree" patterns.+input5 :: [[Element]]+input5 = take 5 $ inputs 100++-- This input corresponds to a levelling merge with minimal or maximal skew.+-- For each, there are 500 elements total, just as 'input5'.+inputLevellingMin, inputLevellingMax :: [[Element]]+inputLevellingMin =+ head (inputs (4*n))+ : take 4 (tail (inputs n))+ where+ n = 1000 `div` (4+4)+inputLevellingMax =+ head (inputs (16*n))+ : take 4 (tail (inputs n))+ where+ n = 1000 `div` (16+4)++inputLevellingMin', inputLevellingMax' :: [[Element]]+inputLevellingMin' = arrangeInputForLoserTree inputLevellingMin+inputLevellingMax' = arrangeInputForLoserTree inputLevellingMax++-- A hacky way to create a degenerate loser tree where one side of the whole+-- tournament tree effectively only consists of a single (large) input,+-- so it can immediately "play in the final" and get chosen with just one+-- comparison.+arrangeInputForLoserTree :: [[Element]] -> [[Element]]+arrangeInputForLoserTree input =+ head input+ : replicate 3 [minElement] -- non-empty to be considered during tree building+ ++ tail input++inputs :: Int -> [[Element]]+inputs n =+ [ L.sort $ take n $ L.unfoldr (Just . genElement) $ SM.mkSMGen seed+ | seed <- iterate (3 +) 42+ ]++{-------------------------------------------------------------------------------+ Recursive 2-way merge+-------------------------------------------------------------------------------}++listMerge :: Ord a => [[a]] -> [a]+listMerge [] = []+listMerge [xs] = xs+listMerge (xs:xss) = merge xs (listMerge xss)++merge :: Ord a => [a] -> [a] -> [a]+merge [] [] = []+merge [] ys = ys+merge xs [] = xs+merge xs@(x:xs') ys@(y:ys')+ | x <= y = x : merge xs' ys+ | otherwise = y : merge xs ys'++{-------------------------------------------------------------------------------+ Recursive 2-way merge, tree shape+-------------------------------------------------------------------------------}++-- | Like 'listMerge', but merges in binary-tree pattern.+--+-- Given inputs of about the same length, there will be less work in merges.+treeMerge :: Ord a => [[a]] -> [a]+treeMerge [] = []+treeMerge [xs] = xs+treeMerge (xs:ys:xss) = treeMerge (merge xs ys : go xss) where+ go [] = []+ go [vs] = [vs]+ go (vs:ws:vss) = merge vs ws : go vss++{-------------------------------------------------------------------------------+ Direct k-way merge using heaps Data.Heap.Heap+-------------------------------------------------------------------------------}++heapMerge :: forall a. Ord a => [[a]] -> [a]+heapMerge xss = go $ Heap.fromList+ [ Heap.Entry x xs+ | x:xs <- xss+ ]+ where+ go :: Heap.Heap (Heap.Entry a [a]) -> [a]+ go heap = case Heap.viewMin heap of+ Nothing -> []+ Just (Heap.Entry x xs, heap') -> x : case xs of+ [] -> go heap'+ x':xs' -> go (Heap.insert (Heap.Entry x' xs') heap')++{-------------------------------------------------------------------------------+ Direct k-way merge using LoserTree+-------------------------------------------------------------------------------}++loserTreeMerge :: forall a. Ord a => [[a]] -> [a]+loserTreeMerge xss = case [ Heap.Entry x xs | x : xs <- xss ] of+ [] -> []+ e:es -> runST $ do+ -- we reuse Heap.Entry structure here.+ (tree, element) <- K.Tree.newLoserTree $ e :| es+ go tree $ Just element+ where+ go :: K.Tree.MutableLoserTree s (Heap.Entry a [a]) -> Maybe (Heap.Entry a [a]) -> ST s [a]+ go !_ Nothing = pure []+ go !tree (Just (Heap.Entry x xs)) = fmap (x :) $ case xs of+ [] -> K.Tree.remove tree >>= go tree+ x':xs' -> K.Tree.replace tree (Heap.Entry x' xs') >>= go tree . Just++{-------------------------------------------------------------------------------+ Direct k-way merge using MutableHeap+-------------------------------------------------------------------------------}++mutHeapMerge :: forall a. Ord a => [[a]] -> [a]+mutHeapMerge xss = case [ Heap.Entry x xs | x : xs <- xss ] of+ [] -> []+ e:es -> runST $ do+ -- we reuse Heap.Entry structure here.+ (heap, element) <- K.Heap.newMutableHeap $ e :| es+ go heap $ Just element+ where+ go :: K.Heap.MutableHeap s (Heap.Entry a [a]) -> Maybe (Heap.Entry a [a]) -> ST s [a]+ go !_ Nothing = pure []+ go !heap (Just (Heap.Entry x xs)) = fmap (x :) $ case xs of+ [] -> K.Heap.extract heap >>= go heap+ x':xs' -> K.Heap.replaceRoot heap (Heap.Entry x' xs') >>= go heap . Just++{-------------------------------------------------------------------------------+ Account for differing sort comparisons across base versions+-------------------------------------------------------------------------------}++-- | The 'sort' and 'sortBy' implementations changed as of @base-4.21@.+-- The new implementation performs fewer comparisons on longer lists.+--+-- Because of this, we fall back to the old sort method when the version of+-- @base@ is @4.21@ or greater.+comparisons5, comparisons7, comparisons8, comparisonsMin, comparisonsMax :: Int+#if MIN_VERSION_base(4,21,0)+comparisons5 = 1692+comparisons7 = 2691+comparisons8 = 3389+comparisonsMin = 3606+comparisonsMax = 3820+#else+comparisons5 = 1790+comparisons7 = 2691+comparisons8 = 3190+comparisonsMin = 3729+comparisonsMax = 3872+#endif
@@ -0,0 +1,77 @@+{-# LANGUAGE OverloadedStrings #-}+{-# OPTIONS_GHC -Wno-orphans #-}+module Main (main) where++import Data.ByteString (ByteString)+import Data.Map (Map)+import qualified Data.Map as Map+import Database.LSMTree.Internal.Map.Range (Bound (..), Clusive (..),+ rangeLookup)+import Test.QuickCheck (Arbitrary (..), Property, elements, frequency,+ (===))+import Test.Tasty (defaultMain, testGroup)+import Test.Tasty.HUnit (testCase, (@?=))+import Test.Tasty.QuickCheck (testProperty)++main :: IO ()+main = defaultMain $ testGroup "map-range-test"+ [ testProperty "model" prop+ , testCase "example1" $ do+ let m = Map.fromList [(0 :: Int, 'x'), (2, 'y')]++ rangeLookup (Bound 0 Inclusive) (Bound 1 Inclusive) m @?= [(0, 'x')]++ , testCase "example2" $ do+ let m = Map.fromList [("\NUL\NUL\NUL\NUL\NUL\NUL\NUL\SOH" :: ByteString,'x')]++ let lb = Bound "\NUL\NUL\NUL\NUL\NUL\NUL\NUL\SOH" Inclusive+ ub = Bound "\NUL\NUL\NUL\NUL\NUL\NUL\NUL\STX" Inclusive++ rangeLookup lb ub m @?= [("\NUL\NUL\NUL\NUL\NUL\NUL\NUL\SOH",'x')]++ , testCase "unordered-bounds" $ do+ let m = Map.fromList [('x', "ex"), ('y', "why" :: String)]++ -- if lower bound is greater than upper bound empty list is returned.+ rangeLookup (Bound 'z' Inclusive) (Bound 'a' Inclusive) m @?= []+ naiveRangeLookup (Bound 'z' Inclusive) (Bound 'a' Inclusive) m @?= []+ ]++prop :: Bound Int -> Bound Int -> Map Int Int -> Property+prop lb ub m =+ rangeLookup lb ub m === naiveRangeLookup lb ub m++naiveRangeLookup ::+ Ord k+ => Bound k -- ^ lower bound+ -> Bound k -- ^ upper bound+ -> Map k v+ -> [(k, v)]+naiveRangeLookup lb ub m =+ [ p+ | p@(k, _) <- Map.toList m+ , evalLowerBound lb k+ , evalUpperBound ub k+ ]++evalLowerBound :: Ord k => Bound k -> k -> Bool+evalLowerBound NoBound _ = True+evalLowerBound (Bound b Exclusive) k = b < k+evalLowerBound (Bound b Inclusive) k = b <= k++evalUpperBound :: Ord k => Bound k -> k -> Bool+evalUpperBound NoBound _ = True+evalUpperBound (Bound b Exclusive) k = k < b+evalUpperBound (Bound b Inclusive) k = k <= b++instance Arbitrary k => Arbitrary (Bound k) where+ arbitrary = frequency+ [ (1, pure NoBound)+ , (20, Bound <$> arbitrary <*> arbitrary)+ ]++ shrink NoBound = []+ shrink (Bound k b) = NoBound : map (`Bound` b) (shrink k)++instance Arbitrary Clusive where+ arbitrary = elements [Exclusive, Inclusive]