diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for `lsm-tree` + +## 1.0.0.0 -- 2025-08-06 + +* First released version. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/NOTICE b/NOTICE new file mode 100644 --- /dev/null +++ b/NOTICE @@ -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. diff --git a/app/Database/LSMTree/Demo.hs b/app/Database/LSMTree/Demo.hs new file mode 100644 --- /dev/null +++ b/app/Database/LSMTree/Demo.hs @@ -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 diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -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 diff --git a/bench-unions/Bench/Unions.hs b/bench-unions/Bench/Unions.hs new file mode 100644 --- /dev/null +++ b/bench-unions/Bench/Unions.hs @@ -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 diff --git a/bench-unions/Main.hs b/bench-unions/Main.hs new file mode 100644 --- /dev/null +++ b/bench-unions/Main.hs @@ -0,0 +1,6 @@ +module Main (main) where + +import qualified Bench.Unions + +main :: IO () +main = Bench.Unions.main diff --git a/bench/macro/lsm-tree-bench-bloomfilter.hs b/bench/macro/lsm-tree-bench-bloomfilter.hs new file mode 100644 --- /dev/null +++ b/bench/macro/lsm-tree-bench-bloomfilter.hs @@ -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 diff --git a/bench/macro/lsm-tree-bench-lookups.hs b/bench/macro/lsm-tree-bench-lookups.hs new file mode 100644 --- /dev/null +++ b/bench/macro/lsm-tree-bench-lookups.hs @@ -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 "" + -- 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 "" + + 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 "" + let !x = classifyLookups blooms keyRng0 benchmarkNumLookups + putStrLn "" + 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' diff --git a/bench/macro/utxo-bench.hs b/bench/macro/utxo-bench.hs new file mode 100644 --- /dev/null +++ b/bench/macro/utxo-bench.hs @@ -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 diff --git a/bench/macro/utxo-rocksdb-bench.hs b/bench/macro/utxo-rocksdb-bench.hs new file mode 100644 --- /dev/null +++ b/bench/macro/utxo-rocksdb-bench.hs @@ -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 diff --git a/bench/micro/Bench/Database/LSMTree.hs b/bench/micro/Bench/Database/LSMTree.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree.hs @@ -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 diff --git a/bench/micro/Bench/Database/LSMTree/Internal/BloomFilter.hs b/bench/micro/Bench/Database/LSMTree/Internal/BloomFilter.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/BloomFilter.hs @@ -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 diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Index.hs b/bench/micro/Bench/Database/LSMTree/Internal/Index.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/Index.hs @@ -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 diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Index/Compact.hs b/bench/micro/Bench/Database/LSMTree/Internal/Index/Compact.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/Index/Compact.hs @@ -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) diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs b/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/Lookup.hs @@ -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 diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Merge.hs b/bench/micro/Bench/Database/LSMTree/Internal/Merge.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/Merge.hs @@ -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) diff --git a/bench/micro/Bench/Database/LSMTree/Internal/RawPage.hs b/bench/micro/Bench/Database/LSMTree/Internal/RawPage.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/RawPage.hs @@ -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 + diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Serialise.hs b/bench/micro/Bench/Database/LSMTree/Internal/Serialise.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/Serialise.hs @@ -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) diff --git a/bench/micro/Bench/Database/LSMTree/Internal/WriteBuffer.hs b/bench/micro/Bench/Database/LSMTree/Internal/WriteBuffer.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Bench/Database/LSMTree/Internal/WriteBuffer.hs @@ -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'') diff --git a/bench/micro/Main.hs b/bench/micro/Main.hs new file mode 100644 --- /dev/null +++ b/bench/micro/Main.hs @@ -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 + ] diff --git a/lsm-tree.cabal b/lsm-tree.cabal new file mode 100644 --- /dev/null +++ b/lsm-tree.cabal @@ -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 diff --git a/src-control/Control/ActionRegistry.hs b/src-control/Control/ActionRegistry.hs new file mode 100644 --- /dev/null +++ b/src-control/Control/ActionRegistry.hs @@ -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) diff --git a/src-control/Control/Concurrent/Class/MonadSTM/RWVar.hs b/src-control/Control/Concurrent/Class/MonadSTM/RWVar.hs new file mode 100644 --- /dev/null +++ b/src-control/Control/Concurrent/Class/MonadSTM/RWVar.hs @@ -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) diff --git a/src-control/Control/RefCount.hs b/src-control/Control/RefCount.hs new file mode 100644 --- /dev/null +++ b/src-control/Control/RefCount.hs @@ -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 _ = "" + +-- | 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 diff --git a/src-core/Database/LSMTree/Internal/Arena.hs b/src-core/Database/LSMTree/Internal/Arena.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Arena.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Assertions.hs b/src-core/Database/LSMTree/Internal/Assertions.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Assertions.hs @@ -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' diff --git a/src-core/Database/LSMTree/Internal/BitMath.hs b/src-core/Database/LSMTree/Internal/BitMath.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/BitMath.hs @@ -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 #-} diff --git a/src-core/Database/LSMTree/Internal/BlobFile.hs b/src-core/Database/LSMTree/Internal/BlobFile.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/BlobFile.hs @@ -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 () diff --git a/src-core/Database/LSMTree/Internal/BlobRef.hs b/src-core/Database/LSMTree/Internal/BlobRef.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/BlobRef.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/BloomFilter.hs b/src-core/Database/LSMTree/Internal/BloomFilter.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/BloomFilter.hs @@ -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 + diff --git a/src-core/Database/LSMTree/Internal/ByteString.hs b/src-core/Database/LSMTree/Internal/ByteString.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/ByteString.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/CRC32C.hs b/src-core/Database/LSMTree/Internal/CRC32C.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/CRC32C.hs @@ -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: +{-# 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 + diff --git a/src-core/Database/LSMTree/Internal/ChecksumHandle.hs b/src-core/Database/LSMTree/Internal/ChecksumHandle.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/ChecksumHandle.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Chunk.hs b/src-core/Database/LSMTree/Internal/Chunk.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Chunk.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/Config.hs b/src-core/Database/LSMTree/Internal/Config.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Config.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Config/Override.hs b/src-core/Database/LSMTree/Internal/Config/Override.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Config/Override.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/Cursor.hs b/src-core/Database/LSMTree/Internal/Cursor.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Cursor.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Entry.hs b/src-core/Database/LSMTree/Internal/Entry.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Entry.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/IncomingRun.hs b/src-core/Database/LSMTree/Internal/IncomingRun.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/IncomingRun.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Index.hs b/src-core/Database/LSMTree/Internal/Index.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Index.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Index/Compact.hs b/src-core/Database/LSMTree/Internal/Index/Compact.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Index/Compact.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Index/CompactAcc.hs b/src-core/Database/LSMTree/Internal/Index/CompactAcc.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Index/CompactAcc.hs @@ -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 . +-- +-- 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 diff --git a/src-core/Database/LSMTree/Internal/Index/Ordinary.hs b/src-core/Database/LSMTree/Internal/Index/Ordinary.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Index/Ordinary.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/Index/OrdinaryAcc.hs b/src-core/Database/LSMTree/Internal/Index/OrdinaryAcc.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Index/OrdinaryAcc.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/Lookup.hs b/src-core/Database/LSMTree/Internal/Lookup.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Lookup.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/Map/Range.hs b/src-core/Database/LSMTree/Internal/Map/Range.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Map/Range.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Merge.hs b/src-core/Database/LSMTree/Internal/Merge.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Merge.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/MergeSchedule.hs b/src-core/Database/LSMTree/Internal/MergeSchedule.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/MergeSchedule.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/MergingRun.hs b/src-core/Database/LSMTree/Internal/MergingRun.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/MergingRun.hs @@ -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" diff --git a/src-core/Database/LSMTree/Internal/MergingTree.hs b/src-core/Database/LSMTree/Internal/MergingTree.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/MergingTree.hs @@ -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" diff --git a/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs b/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/Page.hs b/src-core/Database/LSMTree/Internal/Page.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Page.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/PageAcc.hs b/src-core/Database/LSMTree/Internal/PageAcc.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/PageAcc.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/PageAcc1.hs b/src-core/Database/LSMTree/Internal/PageAcc1.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/PageAcc1.hs @@ -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" diff --git a/src-core/Database/LSMTree/Internal/Paths.hs b/src-core/Database/LSMTree/Internal/Paths.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Paths.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Primitive.hs b/src-core/Database/LSMTree/Internal/Primitive.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Primitive.hs @@ -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#) diff --git a/src-core/Database/LSMTree/Internal/Range.hs b/src-core/Database/LSMTree/Internal/Range.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Range.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/RawBytes.hs b/src-core/Database/LSMTree/Internal/RawBytes.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RawBytes.hs @@ -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 +#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#) diff --git a/src-core/Database/LSMTree/Internal/RawOverflowPage.hs b/src-core/Database/LSMTree/Internal/RawOverflowPage.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RawOverflowPage.hs @@ -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 +-- +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 + diff --git a/src-core/Database/LSMTree/Internal/RawPage.hs b/src-core/Database/LSMTree/Internal/RawPage.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RawPage.hs @@ -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)) diff --git a/src-core/Database/LSMTree/Internal/Readers.hs b/src-core/Database/LSMTree/Internal/Readers.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Readers.hs @@ -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' + , .. + } diff --git a/src-core/Database/LSMTree/Internal/Run.hs b/src-core/Database/LSMTree/Internal/Run.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Run.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/RunAcc.hs b/src-core/Database/LSMTree/Internal/RunAcc.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RunAcc.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/RunBuilder.hs b/src-core/Database/LSMTree/Internal/RunBuilder.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RunBuilder.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/RunNumber.hs b/src-core/Database/LSMTree/Internal/RunNumber.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RunNumber.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/RunReader.hs b/src-core/Database/LSMTree/Internal/RunReader.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/RunReader.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Serialise.hs b/src-core/Database/LSMTree/Internal/Serialise.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Serialise.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Serialise/Class.hs b/src-core/Database/LSMTree/Internal/Serialise/Class.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Serialise/Class.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Snapshot.hs b/src-core/Database/LSMTree/Internal/Snapshot.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Snapshot.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Snapshot/Codec.hs b/src-core/Database/LSMTree/Internal/Snapshot/Codec.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Snapshot/Codec.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/StrictArray.hs b/src-core/Database/LSMTree/Internal/StrictArray.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/StrictArray.hs @@ -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#) diff --git a/src-core/Database/LSMTree/Internal/Types.hs b/src-core/Database/LSMTree/Internal/Types.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Types.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/UniqCounter.hs b/src-core/Database/LSMTree/Internal/UniqCounter.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/UniqCounter.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Unsafe.hs b/src-core/Database/LSMTree/Internal/Unsafe.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Unsafe.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Unsliced.hs b/src-core/Database/LSMTree/Internal/Unsliced.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Unsliced.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/Vector.hs b/src-core/Database/LSMTree/Internal/Vector.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Vector.hs @@ -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' diff --git a/src-core/Database/LSMTree/Internal/Vector/Growing.hs b/src-core/Database/LSMTree/Internal/Vector/Growing.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/Vector/Growing.hs @@ -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) diff --git a/src-core/Database/LSMTree/Internal/WriteBuffer.hs b/src-core/Database/LSMTree/Internal/WriteBuffer.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/WriteBuffer.hs @@ -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 ) diff --git a/src-core/Database/LSMTree/Internal/WriteBufferBlobs.hs b/src-core/Database/LSMTree/Internal/WriteBufferBlobs.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/WriteBufferBlobs.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/WriteBufferReader.hs b/src-core/Database/LSMTree/Internal/WriteBufferReader.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/WriteBufferReader.hs @@ -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 diff --git a/src-core/Database/LSMTree/Internal/WriteBufferWriter.hs b/src-core/Database/LSMTree/Internal/WriteBufferWriter.hs new file mode 100644 --- /dev/null +++ b/src-core/Database/LSMTree/Internal/WriteBufferWriter.hs @@ -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] diff --git a/src-extras/Database/LSMTree/Extras.hs b/src-extras/Database/LSMTree/Extras.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras.hs @@ -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 diff --git a/src-extras/Database/LSMTree/Extras/Generators.hs b/src-extras/Database/LSMTree/Extras/Generators.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/Generators.hs @@ -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] diff --git a/src-extras/Database/LSMTree/Extras/Index.hs b/src-extras/Database/LSMTree/Extras/Index.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/Index.hs @@ -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 #-} diff --git a/src-extras/Database/LSMTree/Extras/MergingRunData.hs b/src-extras/Database/LSMTree/Extras/MergingRunData.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/MergingRunData.hs @@ -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 + ] diff --git a/src-extras/Database/LSMTree/Extras/MergingTreeData.hs b/src-extras/Database/LSMTree/Extras/MergingTreeData.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/MergingTreeData.hs @@ -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 + ] diff --git a/src-extras/Database/LSMTree/Extras/NoThunks.hs b/src-extras/Database/LSMTree/Extras/NoThunks.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/NoThunks.hs @@ -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) diff --git a/src-extras/Database/LSMTree/Extras/Orphans.hs b/src-extras/Database/LSMTree/Extras/Orphans.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/Orphans.hs @@ -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 diff --git a/src-extras/Database/LSMTree/Extras/Random.hs b/src-extras/Database/LSMTree/Extras/Random.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/Random.hs @@ -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' diff --git a/src-extras/Database/LSMTree/Extras/ReferenceImpl.hs b/src-extras/Database/LSMTree/Extras/ReferenceImpl.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/ReferenceImpl.hs @@ -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 diff --git a/src-extras/Database/LSMTree/Extras/RunData.hs b/src-extras/Database/LSMTree/Extras/RunData.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/RunData.hs @@ -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) diff --git a/src-extras/Database/LSMTree/Extras/UTxO.hs b/src-extras/Database/LSMTree/Extras/UTxO.hs new file mode 100644 --- /dev/null +++ b/src-extras/Database/LSMTree/Extras/UTxO.hs @@ -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" diff --git a/src-kmerge/KMerge/Heap.hs b/src-kmerge/KMerge/Heap.hs new file mode 100644 --- /dev/null +++ b/src-kmerge/KMerge/Heap.hs @@ -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 #-} diff --git a/src-kmerge/KMerge/LoserTree.hs b/src-kmerge/KMerge/LoserTree.hs new file mode 100644 --- /dev/null +++ b/src-kmerge/KMerge/LoserTree.hs @@ -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 #-} diff --git a/src-mcg/MCG.hs b/src-mcg/MCG.hs new file mode 100644 --- /dev/null +++ b/src-mcg/MCG.hs @@ -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) diff --git a/src-prototypes/FormatPage.hs b/src-prototypes/FormatPage.hs new file mode 100644 --- /dev/null +++ b/src-prototypes/FormatPage.hs @@ -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 diff --git a/src-prototypes/ScheduledMerges.hs b/src-prototypes/ScheduledMerges.hs new file mode 100644 --- /dev/null +++ b/src-prototypes/ScheduledMerges.hs @@ -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] diff --git a/src-rocksdb/RocksDB.hs b/src-rocksdb/RocksDB.hs new file mode 100644 --- /dev/null +++ b/src-rocksdb/RocksDB.hs @@ -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 diff --git a/src-rocksdb/RocksDB/FFI.hs b/src-rocksdb/RocksDB/FFI.hs new file mode 100644 --- /dev/null +++ b/src-rocksdb/RocksDB/FFI.hs @@ -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) diff --git a/src/Database/LSMTree.hs b/src/Database/LSMTree.hs new file mode 100644 --- /dev/null +++ b/src/Database/LSMTree.hs @@ -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 diff --git a/src/Database/LSMTree/Simple.hs b/src/Database/LSMTree/Simple.hs new file mode 100644 --- /dev/null +++ b/src/Database/LSMTree/Simple.hs @@ -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) diff --git a/test-control/Main.hs b/test-control/Main.hs new file mode 100644 --- /dev/null +++ b/test-control/Main.hs @@ -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 + ] diff --git a/test-control/Test/Control/ActionRegistry.hs b/test-control/Test/Control/ActionRegistry.hs new file mode 100644 --- /dev/null +++ b/test-control/Test/Control/ActionRegistry.hs @@ -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 diff --git a/test-control/Test/Control/Concurrent/Class/MonadSTM/RWVar.hs b/test-control/Test/Control/Concurrent/Class/MonadSTM/RWVar.hs new file mode 100644 --- /dev/null +++ b/test-control/Test/Control/Concurrent/Class/MonadSTM/RWVar.hs @@ -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 diff --git a/test-control/Test/Control/RefCount.hs b/test-control/Test/Control/RefCount.hs new file mode 100644 --- /dev/null +++ b/test-control/Test/Control/RefCount.hs @@ -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 + diff --git a/test-prototypes/Main.hs b/test-prototypes/Main.hs new file mode 100644 --- /dev/null +++ b/test-prototypes/Main.hs @@ -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 + ] diff --git a/test-prototypes/Test/FormatPage.hs b/test-prototypes/Test/FormatPage.hs new file mode 100644 --- /dev/null +++ b/test-prototypes/Test/FormatPage.hs @@ -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 + diff --git a/test-prototypes/Test/ScheduledMerges.hs b/test-prototypes/Test/ScheduledMerges.hs new file mode 100644 --- /dev/null +++ b/test-prototypes/Test/ScheduledMerges.hs @@ -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 diff --git a/test-prototypes/Test/ScheduledMerges/RunSizes.hs b/test-prototypes/Test/ScheduledMerges/RunSizes.hs new file mode 100644 --- /dev/null +++ b/test-prototypes/Test/ScheduledMerges/RunSizes.hs @@ -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 diff --git a/test-prototypes/Test/ScheduledMergesQLS.hs b/test-prototypes/Test/ScheduledMergesQLS.hs new file mode 100644 --- /dev/null +++ b/test-prototypes/Test/ScheduledMergesQLS.hs @@ -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 diff --git a/test/Database/LSMTree/Class.hs b/test/Database/LSMTree/Class.hs new file mode 100644 --- /dev/null +++ b/test/Database/LSMTree/Class.hs @@ -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 diff --git a/test/Database/LSMTree/Class/Common.hs b/test/Database/LSMTree/Class/Common.hs new file mode 100644 --- /dev/null +++ b/test/Database/LSMTree/Class/Common.hs @@ -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 diff --git a/test/Database/LSMTree/Model.hs b/test/Database/LSMTree/Model.hs new file mode 100644 --- /dev/null +++ b/test/Database/LSMTree/Model.hs @@ -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 diff --git a/test/Database/LSMTree/Model/IO.hs b/test/Database/LSMTree/Model/IO.hs new file mode 100644 --- /dev/null +++ b/test/Database/LSMTree/Model/IO.hs @@ -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) diff --git a/test/Database/LSMTree/Model/Session.hs b/test/Database/LSMTree/Model/Session.hs new file mode 100644 --- /dev/null +++ b/test/Database/LSMTree/Model/Session.hs @@ -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) diff --git a/test/Database/LSMTree/Model/Table.hs b/test/Database/LSMTree/Model/Table.hs new file mode 100644 --- /dev/null +++ b/test/Database/LSMTree/Model/Table.hs @@ -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)) diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -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. diff --git a/test/Test/Database/LSMTree.hs b/test/Test/Database/LSMTree.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree.hs @@ -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 + } diff --git a/test/Test/Database/LSMTree/Class.hs b/test/Test/Database/LSMTree/Class.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Class.hs @@ -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' diff --git a/test/Test/Database/LSMTree/Generators.hs b/test/Test/Database/LSMTree/Generators.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Generators.hs @@ -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) diff --git a/test/Test/Database/LSMTree/Internal.hs b/test/Test/Database/LSMTree/Internal.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal.hs @@ -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) diff --git a/test/Test/Database/LSMTree/Internal/Arena.hs b/test/Test/Database/LSMTree/Internal/Arena.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Arena.hs @@ -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 + ] diff --git a/test/Test/Database/LSMTree/Internal/BlobFile/FS.hs b/test/Test/Database/LSMTree/Internal/BlobFile/FS.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/BlobFile/FS.hs @@ -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 + diff --git a/test/Test/Database/LSMTree/Internal/BloomFilter.hs b/test/Test/Database/LSMTree/Internal/BloomFilter.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/BloomFilter.hs @@ -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'))) diff --git a/test/Test/Database/LSMTree/Internal/CRC32C.hs b/test/Test/Database/LSMTree/Internal/CRC32C.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/CRC32C.hs @@ -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' diff --git a/test/Test/Database/LSMTree/Internal/Chunk.hs b/test/Test/Database/LSMTree/Internal/Chunk.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Chunk.hs @@ -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) diff --git a/test/Test/Database/LSMTree/Internal/Entry.hs b/test/Test/Database/LSMTree/Internal/Entry.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Entry.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Index/Compact.hs b/test/Test/Database/LSMTree/Internal/Index/Compact.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Index/Compact.hs @@ -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) + ] diff --git a/test/Test/Database/LSMTree/Internal/Index/Ordinary.hs b/test/Test/Database/LSMTree/Internal/Index/Ordinary.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Index/Ordinary.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Lookup.hs b/test/Test/Database/LSMTree/Internal/Lookup.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Lookup.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Merge.hs b/test/Test/Database/LSMTree/Internal/Merge.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Merge.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/MergingRun.hs b/test/Test/Database/LSMTree/Internal/MergingRun.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/MergingRun.hs @@ -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) + ] + diff --git a/test/Test/Database/LSMTree/Internal/MergingTree.hs b/test/Test/Database/LSMTree/Internal/MergingTree.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/MergingTree.hs @@ -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". diff --git a/test/Test/Database/LSMTree/Internal/PageAcc.hs b/test/Test/Database/LSMTree/Internal/PageAcc.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/PageAcc.hs @@ -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 () diff --git a/test/Test/Database/LSMTree/Internal/PageAcc1.hs b/test/Test/Database/LSMTree/Internal/PageAcc1.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/PageAcc1.hs @@ -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) + diff --git a/test/Test/Database/LSMTree/Internal/RawBytes.hs b/test/Test/Database/LSMTree/Internal/RawBytes.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RawBytes.hs @@ -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])) diff --git a/test/Test/Database/LSMTree/Internal/RawOverflowPage.hs b/test/Test/Database/LSMTree/Internal/RawOverflowPage.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RawOverflowPage.hs @@ -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 + diff --git a/test/Test/Database/LSMTree/Internal/RawPage.hs b/test/Test/Database/LSMTree/Internal/RawPage.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RawPage.hs @@ -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 "" + else show $ maximum pKeys + , show key + , if null pKeys + then "" + 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 "" + else show $ maximum pKeys + , show key + , if null pKeys + then "" + 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 diff --git a/test/Test/Database/LSMTree/Internal/Readers.hs b/test/Test/Database/LSMTree/Internal/Readers.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Readers.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Run.hs b/test/Test/Database/LSMTree/Internal/Run.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Run.hs @@ -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) diff --git a/test/Test/Database/LSMTree/Internal/RunAcc.hs b/test/Test/Database/LSMTree/Internal/RunAcc.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RunAcc.hs @@ -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 ] + diff --git a/test/Test/Database/LSMTree/Internal/RunBloomFilterAlloc.hs b/test/Test/Database/LSMTree/Internal/RunBloomFilterAlloc.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RunBloomFilterAlloc.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/RunBuilder.hs b/test/Test/Database/LSMTree/Internal/RunBuilder.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RunBuilder.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/RunReader.hs b/test/Test/Database/LSMTree/Internal/RunReader.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/RunReader.hs @@ -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 + diff --git a/test/Test/Database/LSMTree/Internal/Serialise.hs b/test/Test/Database/LSMTree/Internal/Serialise.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Serialise.hs @@ -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 + ] diff --git a/test/Test/Database/LSMTree/Internal/Serialise/Class.hs b/test/Test/Database/LSMTree/Internal/Serialise/Class.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Serialise/Class.hs @@ -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' diff --git a/test/Test/Database/LSMTree/Internal/Snapshot/Codec.hs b/test/Test/Database/LSMTree/Internal/Snapshot/Codec.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Snapshot/Codec.hs @@ -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 ] diff --git a/test/Test/Database/LSMTree/Internal/Snapshot/Codec/Golden.hs b/test/Test/Database/LSMTree/Internal/Snapshot/Codec/Golden.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Snapshot/Codec/Golden.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Snapshot/FS.hs b/test/Test/Database/LSMTree/Internal/Snapshot/FS.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Snapshot/FS.hs @@ -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) diff --git a/test/Test/Database/LSMTree/Internal/Unsliced.hs b/test/Test/Database/LSMTree/Internal/Unsliced.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Unsliced.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Vector.hs b/test/Test/Database/LSMTree/Internal/Vector.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Vector.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/Vector/Growing.hs b/test/Test/Database/LSMTree/Internal/Vector/Growing.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/Vector/Growing.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/WriteBufferBlobs/FS.hs b/test/Test/Database/LSMTree/Internal/WriteBufferBlobs/FS.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/WriteBufferBlobs/FS.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Internal/WriteBufferReader/FS.hs b/test/Test/Database/LSMTree/Internal/WriteBufferReader/FS.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Internal/WriteBufferReader/FS.hs @@ -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)) diff --git a/test/Test/Database/LSMTree/Model/Table.hs b/test/Test/Database/LSMTree/Model/Table.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Model/Table.hs @@ -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) diff --git a/test/Test/Database/LSMTree/Resolve.hs b/test/Test/Database/LSMTree/Resolve.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Resolve.hs @@ -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 diff --git a/test/Test/Database/LSMTree/StateMachine.hs b/test/Test/Database/LSMTree/StateMachine.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/StateMachine.hs @@ -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 diff --git a/test/Test/Database/LSMTree/StateMachine/DL.hs b/test/Test/Database/LSMTree/StateMachine/DL.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/StateMachine/DL.hs @@ -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 diff --git a/test/Test/Database/LSMTree/StateMachine/Op.hs b/test/Test/Database/LSMTree/StateMachine/Op.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/StateMachine/Op.hs @@ -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 diff --git a/test/Test/Database/LSMTree/Tracer/Golden.hs b/test/Test/Database/LSMTree/Tracer/Golden.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/Tracer/Golden.hs @@ -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" diff --git a/test/Test/Database/LSMTree/UnitTests.hs b/test/Test/Database/LSMTree/UnitTests.hs new file mode 100644 --- /dev/null +++ b/test/Test/Database/LSMTree/UnitTests.hs @@ -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" diff --git a/test/Test/FS.hs b/test/Test/FS.hs new file mode 100644 --- /dev/null +++ b/test/Test/FS.hs @@ -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 diff --git a/test/Test/Util/Arbitrary.hs b/test/Test/Util/Arbitrary.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/Arbitrary.hs @@ -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 diff --git a/test/Test/Util/FS.hs b/test/Test/Util/FS.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/FS.hs @@ -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] diff --git a/test/Test/Util/FS/Error.hs b/test/Test/Util/FS/Error.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/FS/Error.hs @@ -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 diff --git a/test/Test/Util/Orphans.hs b/test/Test/Util/Orphans.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/Orphans.hs @@ -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) diff --git a/test/Test/Util/PrettyProxy.hs b/test/Test/Util/PrettyProxy.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/PrettyProxy.hs @@ -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 diff --git a/test/Test/Util/QC.hs b/test/Test/Util/QC.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/QC.hs @@ -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) + diff --git a/test/Test/Util/QLS.hs b/test/Test/Util/QLS.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/QLS.hs @@ -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) diff --git a/test/Test/Util/RawPage.hs b/test/Test/Util/RawPage.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/RawPage.hs @@ -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 _ = '?' diff --git a/test/Test/Util/TypeFamilyWrappers.hs b/test/Test/Util/TypeFamilyWrappers.hs new file mode 100644 --- /dev/null +++ b/test/Test/Util/TypeFamilyWrappers.hs @@ -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) diff --git a/test/kmerge-test.hs b/test/kmerge-test.hs new file mode 100644 --- /dev/null +++ b/test/kmerge-test.hs @@ -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 diff --git a/test/map-range-test.hs b/test/map-range-test.hs new file mode 100644 --- /dev/null +++ b/test/map-range-test.hs @@ -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]