From b6a8b4d6301a0b56454e95a7c9082f4f89d2d503 Mon Sep 17 00:00:00 2001 From: Matthias Heinzel Date: Tue, 15 Sep 2026 12:54:39 +0200 Subject: [PATCH 1/5] make MergingTree non-structurally-empty --- .../Database/LSMTree/Internal/MergingTree.hs | 95 ++++++++++--------- .../LSMTree/Internal/MergingTree/Lookup.hs | 6 +- .../Database/LSMTree/Internal/Snapshot.hs | 21 +++- .../Database/LSMTree/Internal/Unsafe.hs | 17 ++-- .../LSMTree/Extras/MergingTreeData.hs | 57 +++++------ .../Database/LSMTree/Internal/MergingTree.hs | 57 +---------- 6 files changed, 103 insertions(+), 150 deletions(-) diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs index 3fbcea847..caff05a16 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs @@ -35,6 +35,8 @@ import Data.Foldable (toList, traverse_) import Data.List (foldl') -- foldl' is included in the Prelude from base 4.20 onwards #endif +import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.Maybe (isJust) import Data.Vector (Vector) import qualified Data.Vector as V import qualified Database.LSMTree.Internal.BloomFilter as Bloom @@ -117,14 +119,27 @@ 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 + PendingLevelMerge_ !(Vector (PreExistingRun m h)) !(Maybe (Ref (MergingTree m h))) -- | Each input is the entire content of a table (as a merging tree). - | PendingUnionMerge + | PendingUnionMerge_ !(Vector (Ref (MergingTree m h))) +-- TODO: we could avoid returing Maybe if we had a NonEmpty-like type where the +-- last element can have a different type. +mkPendingLevelMerge :: + Vector (PreExistingRun m h) + -> Maybe (Ref (MergingTree m h)) + -> Maybe (PendingMerge m h) +mkPendingLevelMerge prs mmt + | not (V.null prs) || isJust mmt = Just (PendingLevelMerge_ prs mmt) + | otherwise = Nothing + +mkPendingUnionMerge :: NonEmpty (Ref (MergingTree m h)) -> PendingMerge m h +mkPendingUnionMerge mts = PendingUnionMerge_ (V.fromList (toList mts)) + pendingContent :: PendingMerge m h -> ( MR.TreeMergeType @@ -132,8 +147,8 @@ pendingContent :: , 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) + PendingLevelMerge_ prs t -> (MR.MergeLevel, prs, maybe V.empty V.singleton t) + PendingUnionMerge_ ts -> (MR.MergeUnion, V.empty, ts) {-# COMPLETE PendingMerge #-} pattern PendingMerge :: @@ -176,15 +191,14 @@ newOngoingMerge refCtx mr = mkMergingTree refCtx . OngoingTreeMerge =<< dupRef m RefCtx -> [PreExistingRun IO h] -> Maybe (Ref (MergingTree IO h)) - -> IO (Ref (MergingTree IO h)) #-} + -> IO (Maybe (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'. +-- If the list of runs is empty and the optional input tree is not present, the +-- resulting tree would be empty, so we return @Nothing@. -- -- Resource tracking: -- * This allocates a new 'Ref' which the caller is responsible for releasing @@ -200,8 +214,8 @@ newPendingLevelMerge :: => RefCtx -> [PreExistingRun m h] -> Maybe (Ref (MergingTree m h)) - -> m (Ref (MergingTree m h)) -newPendingLevelMerge _ [] (Just t) = dupRef t + -> m (Maybe (Ref (MergingTree m h))) +newPendingLevelMerge _ [] (Just t) = Just <$> dupRef t newPendingLevelMerge refCtx [PreExistingRun r] Nothing = do -- No need to create a pending merge here. -- @@ -214,41 +228,26 @@ newPendingLevelMerge refCtx [PreExistingRun r] Nothing = do 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 refCtx (CompletedTreeMerge r') + Just <$> mkMergingTree refCtx (CompletedTreeMerge r') newPendingLevelMerge refCtx 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 + mmt' <- traverse dupRef mmt prs' <- traverse dupPreExistingRun (V.fromList prs) - mkMergingTree refCtx (PendingTreeMerge (PendingLevelMerge prs' mmt')) + traverse (mkMergingTree refCtx . PendingTreeMerge) (mkPendingLevelMerge 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 :: RefCtx - -> [Ref (MergingTree IO h)] + -> NonEmpty (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. @@ -260,17 +259,14 @@ newPendingLevelMerge refCtx prs mmt = do newPendingUnionMerge :: (MonadMVar m, MonadMask m, PrimMonad m) => RefCtx - -> [Ref (MergingTree m h)] + -> NonEmpty (Ref (MergingTree m h)) -> m (Ref (MergingTree m h)) +newPendingUnionMerge _ (mt :| []) = do + -- No need to create a new node, directly use the single input tree. + dupRef mt newPendingUnionMerge refCtx 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 refCtx (PendingTreeMerge (PendingUnionMerge mts'')) + state <- mkPendingUnionMerge <$> mapM dupRef mts + mkMergingTree refCtx (PendingTreeMerge state) {-# SPECIALISE getCompleted :: Ref (MergingTree IO h) @@ -303,17 +299,22 @@ getCompleted (DeRef MergingTree {mergeState}) = -- This is not the same as being empty due to a pending or ongoing merge -- happening to produce an empty run. -- +-- TODO: remove, this should now always be False due to invariants isStructurallyEmpty :: MonadMVar m => Ref (MergingTree m h) -> m Bool -isStructurallyEmpty (DeRef MergingTree {mergeState}) = - isStructurallyEmptyState <$> readMVar mergeState +isStructurallyEmpty (DeRef MergingTree {mergeState}) = do + b <- isStructurallyEmptyState <$> readMVar mergeState + assert (not b) $ pure False 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. + CompletedTreeMerge _ -> False + OngoingTreeMerge _ -> False + PendingTreeMerge (PendingLevelMerge_ prs mmt) -> V.null prs && null mmt + PendingTreeMerge (PendingUnionMerge_ mts) -> V.null mts + +-- TODO: add invariant (not structurally empty plus maybe invariants of fields?) {-# SPECIALISE mkMergingTree :: RefCtx @@ -347,8 +348,8 @@ finalise mergeState = releaseMTS =<< readMVar mergeState releaseMTS (OngoingTreeMerge mr) = releaseRef mr releaseMTS (PendingTreeMerge ptm) = case ptm of - PendingUnionMerge mts -> traverse_ releaseRef mts - PendingLevelMerge prs mmt -> traverse_ releasePER prs + PendingUnionMerge_ mts -> traverse_ releaseRef mts + PendingLevelMerge_ prs mmt -> traverse_ releasePER prs >> traverse_ releaseRef mmt releasePER (PreExistingRun r) = releaseRef r @@ -500,10 +501,10 @@ supplyCredits hfs hbio refCtx resolve salt runParams threshold root uc = \mt0 c0 supplyPending = MR.supplyChecked remainingMergeDebtPendingMerge $ \pm credits -> do case pm of - PendingLevelMerge prs mt -> + PendingLevelMerge_ prs mt -> leftToRight supplyPreExisting (V.toList prs) credits >>= leftToRight (flip supplyTree) (toList mt) - PendingUnionMerge mts -> + PendingUnionMerge_ mts -> splitEqually (flip supplyTree) (V.toList mts) credits supplyPreExisting c = \case diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs index 1d92a2ab4..7bfb76854 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs @@ -120,10 +120,10 @@ buildLookupTree reg (DeRef mt) = !rs' <- mapM V.singletonMStrict rs pure $ mkLookupNode MR.MergeUnion -- separate (LookupBatch <$!> rs') - MT.PendingTreeMerge (MT.PendingLevelMerge prs Nothing) -> do + 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 + MT.PendingTreeMerge (MT.PendingLevelMerge_ prs (Just tree)) -> do !child <- buildLookupTree reg tree if V.null prs then pure child @@ -132,7 +132,7 @@ buildLookupTree reg (DeRef mt) = LookupBatch . V.concatMap id <$!> -- combine runs V.mapMStrict duplicatePreExistingRun prs pure $ mkLookupNode MR.MergeLevel $ V.fromList [preExisting, child] - MT.PendingTreeMerge (MT.PendingUnionMerge trees) -> + MT.PendingTreeMerge (MT.PendingUnionMerge_ trees) -> mkLookupNode MR.MergeUnion <$!> V.mapMStrict (buildLookupTree reg) trees where dupRun r = withRollback reg (dupRef r) releaseRef diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs index 812a9f798..8f1eec521 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs @@ -44,6 +44,7 @@ import Control.Monad.Class.MonadThrow (MonadMask, bracket, import Control.Monad.Primitive (PrimMonad) import Control.RefCount import Data.Foldable (sequenceA_, traverse_) +import Data.List.NonEmpty (nonEmpty) import Data.String (IsString) import Data.Text (Text) import qualified Data.Vector as V @@ -276,16 +277,26 @@ fromSnapMergingTree hfs hbio refCtx salt uc resolve dir = (SnapPendingLevelMerge prs mmt))) = do prs' <- traverse (fromSnapPreExistingRun reg) prs mmt' <- traverse (go reg) mmt - mt <- withRollback reg + mt <- withRollbackMaybe reg (MT.newPendingLevelMerge refCtx prs' mmt') releaseRef traverse_ (delayedCommit reg . releasePER) prs' traverse_ (delayedCommit reg . releaseRef) mmt' - pure mt + case mt of + Nothing -> + -- TODO: fail properly, or enforce invariant in SnapPendingLevelMerge + error "structurally empty pending level merge in snapshot" + Just mt' -> + pure mt' go reg (SnapMergingTree (SnapPendingTreeMerge (SnapPendingUnionMerge mts))) = do - mts' <- traverse (go reg) mts + mts' <- case nonEmpty mts of + Nothing -> + -- TODO: fail properly, or enforce invariant in SnapPendingUnionMerge + error "structurally empty pending union merge in snapshot" + Just nonEmptyTrees -> + traverse (go reg) nonEmptyTrees mt <- withRollback reg (MT.newPendingUnionMerge refCtx mts') releaseRef @@ -345,9 +356,9 @@ toSnapPendingMerge :: (PrimMonad m, MonadMVar m) => MT.PendingMerge m h -> m (SnapPendingMerge (Ref (Run m h))) -toSnapPendingMerge (MT.PendingUnionMerge mts) = +toSnapPendingMerge (MT.PendingUnionMerge_ mts) = SnapPendingUnionMerge <$> traverse toSnapMergingTree (V.toList mts) -toSnapPendingMerge (MT.PendingLevelMerge pes mmt) = do +toSnapPendingMerge (MT.PendingLevelMerge_ pes mmt) = do pes' <- traverse toSnapPreExistingRun pes mmt' <- traverse toSnapMergingTree mmt pure $ SnapPendingLevelMerge (V.toList pes') mmt' diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs index 9311f4b75..db83ebb09 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs @@ -2188,22 +2188,21 @@ unionsInOpenSession reg sesh seshEnv conf tr !tableId ts = do -- tableContentToMergingTree duplicates all runs and merges -- so the ones from the tableContent here do not escape -- the read access. - withRollback reg + withRollbackMaybe reg (tableContentToMergingTree (sessionUniqCounter sesh) seshEnv conf tc) releaseRef - mt <- withRollback reg (newPendingUnionMerge (sessionRefCtx seshEnv) 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) + forM_ mts (traverse_ (delayedCommit reg . releaseRef)) - content <- MT.isStructurallyEmpty mt >>= \case - True -> do + content <- case NE.nonEmpty (catMaybes mts) of + Nothing -> do -- no need to have an empty merging tree - delayedCommit reg (releaseRef mt) newEmptyTableContent ((sessionUniqCounter sesh)) seshEnv reg - False -> do + Just mts' -> do + mt <- withRollback reg (newPendingUnionMerge (sessionRefCtx seshEnv) mts') releaseRef empty <- newEmptyTableContent (sessionUniqCounter sesh) seshEnv reg cache <- mkUnionCache reg mt pure empty { tableUnionLevel = Union mt cache } @@ -2220,7 +2219,7 @@ unionsInOpenSession reg sesh seshEnv conf tr !tableId ts = do -> SessionEnv IO h -> TableConfig -> TableContent IO h - -> IO (Ref (MergingTree IO h)) #-} + -> IO (Maybe (Ref (MergingTree IO h))) #-} tableContentToMergingTree :: forall m h. (MonadMask m, MonadMVar m, MonadST m, MonadSTM m) @@ -2228,7 +2227,7 @@ tableContentToMergingTree :: -> SessionEnv m h -> TableConfig -> TableContent m h - -> m (Ref (MergingTree m h)) + -> m (Maybe (Ref (MergingTree m h))) tableContentToMergingTree uc seshEnv conf tc@TableContent { tableLevels, diff --git a/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs b/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs index 850d73dfd..568dd8590 100644 --- a/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs +++ b/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs @@ -20,6 +20,9 @@ module Database.LSMTree.Extras.MergingTreeData ( import Control.Exception (assert, bracket) import Control.RefCount import Data.Foldable (for_, toList) +import Data.List.NonEmpty (NonEmpty ((:|))) +import qualified Data.List.NonEmpty as NE +import Data.Maybe (fromMaybe) import Database.LSMTree.Extras (showPowersOf) import Database.LSMTree.Extras.Generators () import Database.LSMTree.Extras.MergingRunData @@ -87,16 +90,20 @@ unsafeCreateMergingTree hfs hbio refCtx resolve salt runParams path counter = go PendingLevelMergeData prds mtd -> withPreExistingRuns prds $ \prs -> withMaybeTree mtd $ \mt -> - MT.newPendingLevelMerge refCtx prs mt + fromMaybe (error "PendingLevelMergeData: invalid") <$> + MT.newPendingLevelMerge refCtx prs mt PendingUnionMergeData mtds -> withTrees mtds $ \mts -> MT.newPendingUnionMerge refCtx mts - withTrees [] act = act [] - withTrees (mtd:rest) act = + withTrees (mtd :| rest) act = bracket (go mtd) releaseRef $ \t -> - withTrees rest $ \ts -> - act (t:ts) + case NE.nonEmpty rest of + Nothing -> + act (t :| []) + Just rest' -> + withTrees rest' $ \ts -> + act (NE.cons t ts) withMaybeTree Nothing act = act Nothing withMaybeTree (Just mtd) act = @@ -129,7 +136,7 @@ data MergingTreeData 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 + | PendingUnionMergeData (NonEmpty (MergingTreeData k v b)) -- ^ at least 2 children deriving stock (Show, Eq) data PreExistingRunData k v b = @@ -137,18 +144,10 @@ data PreExistingRunData 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 +mergingTreeDataInvariant mtd = + case mtd of CompletedTreeMergeData _rd -> Right () OngoingTreeMergeData mr -> @@ -240,7 +239,7 @@ labelMergingTreeData = \rd -> PendingLevelMergeData prds mtds -> maximum (0 : fmap (const 1) prds ++ map depthTree (toList mtds)) PendingUnionMergeData mtds -> - maximum (0 : map depthTree mtds) + maximum (0 : map depthTree (NE.toList mtds)) labelPreExistingRunData :: SerialisedPreExistingRunData -> Property -> Property @@ -255,15 +254,9 @@ instance ( Ord k, Arbitrary k, Arbitrary v, Arbitrary b 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) - ] + 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! -- @@ -315,8 +308,10 @@ genMergingTreeDataOfSize genKey genVal genBlob = \n0 -> do -- n >= 3, needs 1 constructor + 2 children genPendingUnionMerge n = do - ns <- QC.shuffle =<< arbitraryPartition2 (n - 1) + ns <- shuffleNE =<< arbitraryPartition2 (n - 1) PendingUnionMergeData <$> traverse genMergingTree ns + where + shuffleNE = fmap NE.fromList . QC.shuffle . NE.toList genRun = genScaled genRunData genMergingRun genType = genScaled (genMergingRunData genType) @@ -340,14 +335,14 @@ mergingTreeDataSize = \case CompletedTreeMergeData _ -> 1 OngoingTreeMergeData _ -> 1 PendingLevelMergeData _ tree -> 1 + maybe 0 mergingTreeDataSize tree - PendingUnionMergeData trees -> 1 + sum (map mergingTreeDataSize trees) + PendingUnionMergeData trees -> 1 + sum (fmap 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 :: Int -> QC.Gen (NonEmpty Int) arbitraryPartition2 n = assert (n >= 2) $ do first <- QC.chooseInt (1, n-1) - (first :) <$> arbitraryPartition (n - first) + (first :|) <$> arbitraryPartition (n - first) -- Split into smaller positive numbers. arbitraryPartition :: Int -> QC.Gen [Int] @@ -395,7 +390,7 @@ shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob = \case , length prs' + length t' > 0 ] PendingUnionMergeData ts -> - ts + NE.toList ts <> [ PendingUnionMergeData ts' | ts' <- liftShrink (shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob) ts diff --git a/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs b/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs index f8b258993..8ad8c5ad5 100644 --- a/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs +++ b/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs @@ -3,7 +3,6 @@ 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) @@ -43,8 +42,7 @@ 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 -> + [ testProperty "prop_lookupTree" $ \keys mtd -> ioProperty $ withSimHasBlockIO propNoOpenHandles MockFS.empty $ \hfs hbio _ -> prop_lookupTree hfs hbio keys mtd @@ -65,57 +63,6 @@ runParams = 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 $ withRefCtx $ \refCtx -> - bracket (mkEmptyMergingTree refCtx 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 :: RefCtx -> EmptyMergingTree -> IO (Ref (MergingTree IO h)) -mkEmptyMergingTree refCtx ObviouslyEmptyLevelMerge = newPendingLevelMerge refCtx [] Nothing -mkEmptyMergingTree refCtx ObviouslyEmptyUnionMerge = newPendingUnionMerge refCtx [] -mkEmptyMergingTree refCtx (NonObviouslyEmptyLevelMerge emt) = do - mt <- mkEmptyMergingTree refCtx emt - mt' <- newPendingLevelMerge refCtx [] (Just mt) - releaseRef mt - pure mt' -mkEmptyMergingTree refCtx (NonObviouslyEmptyUnionMerge emts) = do - mts <- mapM (mkEmptyMergingTree refCtx) emts - mt' <- newPendingUnionMerge refCtx mts - mapM_ releaseRef mts - pure mt' - {------------------------------------------------------------------------------- Lookup -------------------------------------------------------------------------------} @@ -199,7 +146,7 @@ modelFoldMergingTree = goMergingTree PendingLevelMergeData prs t -> modelMerge MR.MergeLevel (map goPreExistingRun prs <> map goMergingTree (toList t)) PendingUnionMergeData ts -> - modelMerge MR.MergeUnion (map goMergingTree ts) + modelMerge MR.MergeUnion (map goMergingTree (toList ts)) goPreExistingRun = \case PreExistingRunData r -> unRunData r From 9a9f5bb345f7962033691e22ec8fe32f4a97c0cf Mon Sep 17 00:00:00 2001 From: Matthias Heinzel Date: Wed, 16 Sep 2026 14:39:45 +0200 Subject: [PATCH 2/5] remove isStructurallyEmpty --- .../Database/LSMTree/Internal/MergingTree.hs | 39 ------------------- .../LSMTree/Internal/MergingTree/Lookup.hs | 2 - .../Database/LSMTree/Internal/Unsafe.hs | 31 +++++---------- .../Database/LSMTree/Internal/MergingTree.hs | 17 +++----- 4 files changed, 16 insertions(+), 73 deletions(-) diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs index caff05a16..3bf15740a 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs @@ -13,7 +13,6 @@ module Database.LSMTree.Internal.MergingTree ( , newPendingLevelMerge , newPendingUnionMerge , getCompleted - , isStructurallyEmpty , remainingMergeDebt , supplyCredits -- * Internal state @@ -294,28 +293,6 @@ getCompleted (DeRef MergingTree {mergeState}) = OngoingTreeMerge{} -> pure Nothing PendingTreeMerge{} -> pure Nothing -{-# 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. --- --- TODO: remove, this should now always be False due to invariants -isStructurallyEmpty :: MonadMVar m => Ref (MergingTree m h) -> m Bool -isStructurallyEmpty (DeRef MergingTree {mergeState}) = do - b <- isStructurallyEmptyState <$> readMVar mergeState - assert (not b) $ pure False - -isStructurallyEmptyState :: MergingTreeState m h -> Bool -isStructurallyEmptyState = \case - -- It may also turn out to be useful to consider CompletedTreeMerge with - -- a zero length runs as empty. - CompletedTreeMerge _ -> False - OngoingTreeMerge _ -> False - PendingTreeMerge (PendingLevelMerge_ prs mmt) -> V.null prs && null mmt - PendingTreeMerge (PendingUnionMerge_ mts) -> V.null mts - --- TODO: add invariant (not structurally empty plus maybe invariants of fields?) - {-# SPECIALISE mkMergingTree :: RefCtx -> MergingTreeState IO h @@ -468,22 +445,6 @@ supplyCredits hfs hbio refCtx resolve salt runParams threshold root uc = \mt0 c0 -- 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 refCtx salt runParams runPaths) - releaseRef - pure (CompletedTreeMerge run, credits) - PendingTreeMerge pm -> do leftovers <- supplyPending pm credits if leftovers <= 0 diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs index 7bfb76854..cf745ce32 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs @@ -91,8 +91,6 @@ mergeLookupAcc resolve mt accs = -- | 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 :: diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs index db83ebb09..0b71b8a3f 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs @@ -1192,12 +1192,10 @@ lookups resolve ks t = do 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 + Union _tree unionCache -> do + 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 @@ -1820,13 +1818,9 @@ openTableFromSnapshot policyOveride sesh snap label resolve = do Just mTree -> do snapTree <- traverse (openRun hfs hbio (sessionRefCtx seshEnv) uc reg snapDir activeDir salt) mTree mt <- fromSnapMergingTree hfs hbio (sessionRefCtx seshEnv) 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) + 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 (sessionRefCtx seshEnv) salt uc conf resolve reg activeDir snapLevels' @@ -2408,12 +2402,7 @@ supplyUnionCredits resolve t credits = do 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' } + cache' <- mkUnionCache reg mt + releaseUnionCache reg cache + pure tc' { tableUnionLevel = Union mt cache' } pure leftovers diff --git a/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs b/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs index 8ad8c5ad5..a24e75513 100644 --- a/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs +++ b/lsm-tree/test/Test/Database/LSMTree/Internal/MergingTree.hs @@ -103,17 +103,12 @@ prop_lookupTree hfs hbio keys mtd = withRefCtx $ \refCtx -> do 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 + lookupsIO reg mgr tree = 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 $ From a446502b7a919e666e94d57d3e785a8a71815013 Mon Sep 17 00:00:00 2001 From: Matthias Heinzel Date: Thu, 17 Sep 2026 13:41:39 +0200 Subject: [PATCH 3/5] pattern synonyms --- .../Database/LSMTree/Internal/MergingTree.hs | 76 +++++++++++++------ .../LSMTree/Internal/MergingTree/Lookup.hs | 6 +- .../Database/LSMTree/Internal/Snapshot.hs | 4 +- .../LSMTree/Extras/MergingTreeData.hs | 30 ++++---- 4 files changed, 71 insertions(+), 45 deletions(-) diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs index 3bf15740a..76df01403 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs @@ -18,6 +18,8 @@ module Database.LSMTree.Internal.MergingTree ( -- * Internal state , MergingTreeState (..) , PendingMerge (..) + , pattern PendingLevelMerge + , pattern PendingUnionMerge ) where import Control.ActionRegistry @@ -35,7 +37,6 @@ import Data.List (foldl') -- foldl' is included in the Prelude from base 4.20 onwards #endif import Data.List.NonEmpty (NonEmpty ((:|))) -import Data.Maybe (isJust) import Data.Vector (Vector) import qualified Data.Vector as V import qualified Database.LSMTree.Internal.BloomFilter as Bloom @@ -114,30 +115,55 @@ data MergingTreeState m h = !(PendingMerge m h) -- | A merge that is waiting for its inputs to complete. +-- +-- The constructors are unsafe and only exported to allow deriving instances in +-- other modules. Use the patterns 'PendingLevelMerge' and 'PendingUnionMerge' +-- instead. 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_ + -- + -- INVARIANT: At least one 'PreExistingRun' must be present. This is + -- enforced by in 'mkPendingLevelMerge'. Don't use the constructor directly. + UnsafePendingLevelMerge !(Vector (PreExistingRun m h)) !(Maybe (Ref (MergingTree m h))) -- | Each input is the entire content of a table (as a merging tree). - | PendingUnionMerge_ + -- + -- INVARIANT: At least two inputs ('MergingTree') must be present. This is + -- enforced in 'newPendingUnionMerge'. Don't use the constructor directly. + | UnsafePendingUnionMerge !(Vector (Ref (MergingTree m h))) --- TODO: we could avoid returing Maybe if we had a NonEmpty-like type where the --- last element can have a different type. -mkPendingLevelMerge :: +{-# COMPLETE PendingLevelMerge, PendingUnionMerge #-} +pattern PendingLevelMerge :: Vector (PreExistingRun m h) -> Maybe (Ref (MergingTree m h)) - -> Maybe (PendingMerge m h) -mkPendingLevelMerge prs mmt - | not (V.null prs) || isJust mmt = Just (PendingLevelMerge_ prs mmt) - | otherwise = Nothing + -> PendingMerge m h +pattern PendingLevelMerge prs mmt <- UnsafePendingLevelMerge prs mmt + +pattern PendingUnionMerge :: + Vector (Ref (MergingTree m h)) + -> PendingMerge m h +pattern PendingUnionMerge mts <- UnsafePendingUnionMerge mts + +mkPendingLevelMerge :: + PreExistingRun m h + -> [PreExistingRun m h] + -> Maybe (Ref (MergingTree m h)) + -> PendingMerge m h +mkPendingLevelMerge pr prs mmt = + UnsafePendingLevelMerge (V.fromList (pr : prs)) mmt -mkPendingUnionMerge :: NonEmpty (Ref (MergingTree m h)) -> PendingMerge m h -mkPendingUnionMerge mts = PendingUnionMerge_ (V.fromList (toList mts)) +mkPendingUnionMerge :: + Ref (MergingTree m h) + -> Ref (MergingTree m h) + -> [Ref (MergingTree m h)] + -> PendingMerge m h +mkPendingUnionMerge mt1 mt2 mts = + UnsafePendingUnionMerge (V.fromList (mt1 : mt2 : mts)) pendingContent :: PendingMerge m h @@ -146,8 +172,8 @@ pendingContent :: , 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) + PendingLevelMerge prs t -> (MR.MergeLevel, prs, maybe V.empty V.singleton t) + PendingUnionMerge ts -> (MR.MergeUnion, V.empty, ts) {-# COMPLETE PendingMerge #-} pattern PendingMerge :: @@ -214,7 +240,7 @@ newPendingLevelMerge :: -> [PreExistingRun m h] -> Maybe (Ref (MergingTree m h)) -> m (Maybe (Ref (MergingTree m h))) -newPendingLevelMerge _ [] (Just t) = Just <$> dupRef t +newPendingLevelMerge _ [] mt = traverse dupRef mt newPendingLevelMerge refCtx [PreExistingRun r] Nothing = do -- No need to create a pending merge here. -- @@ -229,10 +255,12 @@ newPendingLevelMerge refCtx [PreExistingRun r] Nothing = do -- exceptions are masked then there can be no async exceptions here at all. Just <$> mkMergingTree refCtx (CompletedTreeMerge r') -newPendingLevelMerge refCtx prs mmt = do +newPendingLevelMerge refCtx (pr : prs) mmt = do + pr' <- dupPreExistingRun pr + prs' <- traverse dupPreExistingRun prs mmt' <- traverse dupRef mmt - prs' <- traverse dupPreExistingRun (V.fromList prs) - traverse (mkMergingTree refCtx . PendingTreeMerge) (mkPendingLevelMerge prs' mmt') + let pending = mkPendingLevelMerge pr' prs' mmt' + Just <$> mkMergingTree refCtx (PendingTreeMerge pending) where dupPreExistingRun (PreExistingRun r) = PreExistingRun <$!> dupRef r @@ -263,8 +291,8 @@ newPendingUnionMerge :: newPendingUnionMerge _ (mt :| []) = do -- No need to create a new node, directly use the single input tree. dupRef mt -newPendingUnionMerge refCtx mts = do - state <- mkPendingUnionMerge <$> mapM dupRef mts +newPendingUnionMerge refCtx (mt1 :| (mt2 : mts)) = do + state <- mkPendingUnionMerge <$> dupRef mt1 <*> dupRef mt2 <*> mapM dupRef mts mkMergingTree refCtx (PendingTreeMerge state) {-# SPECIALISE getCompleted :: @@ -325,8 +353,8 @@ finalise mergeState = releaseMTS =<< readMVar mergeState releaseMTS (OngoingTreeMerge mr) = releaseRef mr releaseMTS (PendingTreeMerge ptm) = case ptm of - PendingUnionMerge_ mts -> traverse_ releaseRef mts - PendingLevelMerge_ prs mmt -> traverse_ releasePER prs + PendingUnionMerge mts -> traverse_ releaseRef mts + PendingLevelMerge prs mmt -> traverse_ releasePER prs >> traverse_ releaseRef mmt releasePER (PreExistingRun r) = releaseRef r @@ -462,10 +490,10 @@ supplyCredits hfs hbio refCtx resolve salt runParams threshold root uc = \mt0 c0 supplyPending = MR.supplyChecked remainingMergeDebtPendingMerge $ \pm credits -> do case pm of - PendingLevelMerge_ prs mt -> + PendingLevelMerge prs mt -> leftToRight supplyPreExisting (V.toList prs) credits >>= leftToRight (flip supplyTree) (toList mt) - PendingUnionMerge_ mts -> + PendingUnionMerge mts -> splitEqually (flip supplyTree) (V.toList mts) credits supplyPreExisting c = \case diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs index cf745ce32..58d4d6b84 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree/Lookup.hs @@ -118,10 +118,10 @@ buildLookupTree reg (DeRef mt) = !rs' <- mapM V.singletonMStrict rs pure $ mkLookupNode MR.MergeUnion -- separate (LookupBatch <$!> rs') - MT.PendingTreeMerge (MT.PendingLevelMerge_ prs Nothing) -> do + 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 + MT.PendingTreeMerge (MT.PendingLevelMerge prs (Just tree)) -> do !child <- buildLookupTree reg tree if V.null prs then pure child @@ -130,7 +130,7 @@ buildLookupTree reg (DeRef mt) = LookupBatch . V.concatMap id <$!> -- combine runs V.mapMStrict duplicatePreExistingRun prs pure $ mkLookupNode MR.MergeLevel $ V.fromList [preExisting, child] - MT.PendingTreeMerge (MT.PendingUnionMerge_ trees) -> + MT.PendingTreeMerge (MT.PendingUnionMerge trees) -> mkLookupNode MR.MergeUnion <$!> V.mapMStrict (buildLookupTree reg) trees where dupRun r = withRollback reg (dupRef r) releaseRef diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs index 8f1eec521..7d0abcb78 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs @@ -356,9 +356,9 @@ toSnapPendingMerge :: (PrimMonad m, MonadMVar m) => MT.PendingMerge m h -> m (SnapPendingMerge (Ref (Run m h))) -toSnapPendingMerge (MT.PendingUnionMerge_ mts) = +toSnapPendingMerge (MT.PendingUnionMerge mts) = SnapPendingUnionMerge <$> traverse toSnapMergingTree (V.toList mts) -toSnapPendingMerge (MT.PendingLevelMerge_ pes mmt) = do +toSnapPendingMerge (MT.PendingLevelMerge pes mmt) = do pes' <- traverse toSnapPreExistingRun pes mmt' <- traverse toSnapMergingTree mmt pure $ SnapPendingLevelMerge (V.toList pes') mmt' diff --git a/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs b/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs index 568dd8590..ee620360e 100644 --- a/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs +++ b/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs @@ -128,6 +128,10 @@ unsafeCreateMergingTree hfs hbio refCtx resolve salt runParams path counter = go -- ScheduledMergesTest module. Maybe we can share some code? -- | A data structure suitable for creating arbitrary 'MergingTree's. +-- Its invariants are slightly less strict as for 'MergingTree' itself, as it +-- allows for pending nodes containing only a single sub-tree. +-- 'MT.newPendingLevelMerge' and 'MT.newPendingUnionMerge' accept these inputs, +-- but will return the input tree itself without creating a single-child node. -- -- Note: 'b ~ Void' should rule out blobs. data MergingTreeData k v b = @@ -136,7 +140,7 @@ data MergingTreeData k v b = | PendingLevelMergeData [PreExistingRunData k v b] (Maybe (MergingTreeData k v b)) -- ^ not both empty! - | PendingUnionMergeData (NonEmpty (MergingTreeData k v b)) -- ^ at least 2 children + | PendingUnionMergeData (NonEmpty (MergingTreeData k v b)) deriving stock (Show, Eq) data PreExistingRunData k v b = @@ -165,8 +169,6 @@ mergingTreeDataInvariant mtd = 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 @@ -280,9 +282,6 @@ genMergingTreeDataOfSize genKey genVal genBlob = \n0 -> do , genPendingLevelMergeNoChild ] - | n == 2 - = genPendingLevelMergeWithChild n - | otherwise = QC.oneof [genPendingLevelMergeWithChild n, genPendingUnionMerge n] @@ -297,7 +296,7 @@ genMergingTreeDataOfSize genKey genVal genBlob = \n0 -> do let preExisting = initPreExisting ++ [lastPreExisting] pure (PendingLevelMergeData preExisting Nothing) - -- n >= 2 + -- n >= 2, needs 1 constructor + 1 child genPendingLevelMergeWithChild n = do numPreExisting <- chooseIntSkewed (0, 6) preExisting <- QC.vectorOf numPreExisting $ @@ -306,9 +305,9 @@ genMergingTreeDataOfSize genKey genVal genBlob = \n0 -> do tree <- genMergingTree (n - 1) pure (PendingLevelMergeData preExisting (Just tree)) - -- n >= 3, needs 1 constructor + 2 children + -- n >= 2, needs 1 constructor + 1 child genPendingUnionMerge n = do - ns <- shuffleNE =<< arbitraryPartition2 (n - 1) + ns <- shuffleNE =<< arbitraryPartition1 (n - 1) PendingUnionMergeData <$> traverse genMergingTree ns where shuffleNE = fmap NE.fromList . QC.shuffle . NE.toList @@ -337,12 +336,11 @@ mergingTreeDataSize = \case PendingLevelMergeData _ tree -> 1 + maybe 0 mergingTreeDataSize tree PendingUnionMergeData trees -> 1 + sum (fmap 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 (NonEmpty Int) -arbitraryPartition2 n = assert (n >= 2) $ do - first <- QC.chooseInt (1, n-1) - (first :|) <$> arbitraryPartition (n - first) +-- Split into at least one positive numbers. The input needs to be greater than +-- or equal to 1. +arbitraryPartition1 :: Int -> QC.Gen (NonEmpty Int) +arbitraryPartition1 n = assert (n >= 1) $ + NE.fromList <$> arbitraryPartition n -- Split into smaller positive numbers. arbitraryPartition :: Int -> QC.Gen [Int] @@ -394,7 +392,7 @@ shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob = \case <> [ PendingUnionMergeData ts' | ts' <- liftShrink (shrinkMergingTreeData shrinkKey shrinkVal shrinkBlob) ts - , length ts' >= 2 + , not (null ts') ] genPreExistingRunData :: From 449811b70394b231e9dcb95925dac65d45d9fd59 Mon Sep 17 00:00:00 2001 From: Matthias Heinzel Date: Thu, 17 Sep 2026 15:49:12 +0200 Subject: [PATCH 4/5] make newPendingUnionLevel return Maybe for consistency --- .../Database/LSMTree/Internal/MergingTree.hs | 22 +++++++----- .../Database/LSMTree/Internal/Snapshot.hs | 21 ++++++----- .../Database/LSMTree/Internal/Unsafe.hs | 36 +++++++++++-------- .../LSMTree/Extras/MergingTreeData.hs | 18 +++++----- 4 files changed, 54 insertions(+), 43 deletions(-) diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs index 76df01403..49b910c7f 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/MergingTree.hs @@ -36,7 +36,6 @@ import Data.Foldable (toList, traverse_) import Data.List (foldl') -- foldl' is included in the Prelude from base 4.20 onwards #endif -import Data.List.NonEmpty (NonEmpty ((:|))) import Data.Vector (Vector) import qualified Data.Vector as V import qualified Database.LSMTree.Internal.BloomFilter as Bloom @@ -269,12 +268,15 @@ newPendingLevelMerge refCtx (pr : prs) mmt = do {-# SPECIALISE newPendingUnionMerge :: RefCtx - -> NonEmpty (Ref (MergingTree IO h)) - -> IO (Ref (MergingTree IO h)) #-} + -> [Ref (MergingTree IO h)] + -> IO (Maybe (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). -- +-- If the list of inputs is empty, the resulting tree would be empty, so we +-- return @Nothing@. +-- -- Resource tracking: -- * This allocates a new 'Ref' which the caller is responsible for releasing -- eventually. @@ -286,14 +288,16 @@ newPendingLevelMerge refCtx (pr : prs) mmt = do newPendingUnionMerge :: (MonadMVar m, MonadMask m, PrimMonad m) => RefCtx - -> NonEmpty (Ref (MergingTree m h)) - -> m (Ref (MergingTree m h)) -newPendingUnionMerge _ (mt :| []) = do + -> [Ref (MergingTree m h)] + -> m (Maybe (Ref (MergingTree m h))) +newPendingUnionMerge _ [] = + pure Nothing +newPendingUnionMerge _ [mt] = -- No need to create a new node, directly use the single input tree. - dupRef mt -newPendingUnionMerge refCtx (mt1 :| (mt2 : mts)) = do + Just <$> dupRef mt +newPendingUnionMerge refCtx (mt1 : mt2 : mts) = do state <- mkPendingUnionMerge <$> dupRef mt1 <*> dupRef mt2 <*> mapM dupRef mts - mkMergingTree refCtx (PendingTreeMerge state) + Just <$> mkMergingTree refCtx (PendingTreeMerge state) {-# SPECIALISE getCompleted :: Ref (MergingTree IO h) diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs index 7d0abcb78..1a7e41ea7 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs @@ -44,7 +44,6 @@ import Control.Monad.Class.MonadThrow (MonadMask, bracket, import Control.Monad.Primitive (PrimMonad) import Control.RefCount import Data.Foldable (sequenceA_, traverse_) -import Data.List.NonEmpty (nonEmpty) import Data.String (IsString) import Data.Text (Text) import qualified Data.Vector as V @@ -204,9 +203,13 @@ instance NFData r => NFData (SnapMergingTreeState r) where rnf (SnapOngoingTreeMerge a) = rnf a data SnapPendingMerge r = + -- | INVARIANT: There is always at least one 'SnapPreExistingRun'. See + -- 'MT.PendingMerge'. SnapPendingLevelMerge ![SnapPreExistingRun r] !(Maybe (SnapMergingTree r)) + -- | INVARIANT: There is always at least one 'SnapMergingRun'. See + -- 'MT.PendingMerge'. | SnapPendingUnionMerge ![SnapMergingTree r] deriving stock (Eq, Functor, Foldable, Traversable) @@ -291,17 +294,17 @@ fromSnapMergingTree hfs hbio refCtx salt uc resolve dir = go reg (SnapMergingTree (SnapPendingTreeMerge (SnapPendingUnionMerge mts))) = do - mts' <- case nonEmpty mts of - Nothing -> - -- TODO: fail properly, or enforce invariant in SnapPendingUnionMerge - error "structurally empty pending union merge in snapshot" - Just nonEmptyTrees -> - traverse (go reg) nonEmptyTrees - mt <- withRollback reg + mts' <- traverse (go reg) mts + mt <- withRollbackMaybe reg (MT.newPendingUnionMerge refCtx mts') releaseRef traverse_ (delayedCommit reg . releaseRef) mts' - pure mt + case mt of + Nothing -> + -- TODO: fail properly, or enforce invariant in SnapPendingUnionMerge + error "structurally empty pending union merge in snapshot" + Just mt' -> + pure mt' go reg (SnapMergingTree (SnapOngoingTreeMerge smrs)) = do mr <- withRollback reg diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs index 0b71b8a3f..0eac7e693 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs @@ -2176,30 +2176,36 @@ unionsInOpenSession :: -> 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. - withRollbackMaybe reg - (tableContentToMergingTree (sessionUniqCounter sesh) seshEnv conf tc) - releaseRef + mts <- + fmap catMaybes $ + 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. + withRollbackMaybe reg + (tableContentToMergingTree (sessionUniqCounter sesh) seshEnv conf tc) + 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 (traverse_ (delayedCommit reg . releaseRef)) + forM_ mts (delayedCommit reg . releaseRef) + + mt <- + withRollbackMaybe reg + (newPendingUnionMerge (sessionRefCtx seshEnv) mts) + releaseRef - content <- case NE.nonEmpty (catMaybes mts) of + content <- case mt of Nothing -> do -- no need to have an empty merging tree newEmptyTableContent ((sessionUniqCounter sesh)) seshEnv reg - Just mts' -> do - mt <- withRollback reg (newPendingUnionMerge (sessionRefCtx seshEnv) mts') releaseRef + Just mt' -> do empty <- newEmptyTableContent (sessionUniqCounter sesh) seshEnv reg - cache <- mkUnionCache reg mt - pure empty { tableUnionLevel = Union mt cache } + cache <- mkUnionCache reg mt' + pure empty { tableUnionLevel = Union mt' cache } -- Pick the arena manager to optimise the case of: -- someUpdates <> bigTableWithLotsOfLookups diff --git a/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs b/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs index ee620360e..8d7633aca 100644 --- a/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs +++ b/lsm-tree/src-extras/Database/LSMTree/Extras/MergingTreeData.hs @@ -20,7 +20,7 @@ module Database.LSMTree.Extras.MergingTreeData ( import Control.Exception (assert, bracket) import Control.RefCount import Data.Foldable (for_, toList) -import Data.List.NonEmpty (NonEmpty ((:|))) +import Data.List.NonEmpty (NonEmpty) import qualified Data.List.NonEmpty as NE import Data.Maybe (fromMaybe) import Database.LSMTree.Extras (showPowersOf) @@ -93,17 +93,15 @@ unsafeCreateMergingTree hfs hbio refCtx resolve salt runParams path counter = go fromMaybe (error "PendingLevelMergeData: invalid") <$> MT.newPendingLevelMerge refCtx prs mt PendingUnionMergeData mtds -> - withTrees mtds $ \mts -> - MT.newPendingUnionMerge refCtx mts + withTrees (toList mtds) $ \mts -> + fromMaybe (error "PendingUnionMergeData: invalid") <$> + MT.newPendingUnionMerge refCtx mts - withTrees (mtd :| rest) act = + withTrees [] act = act [] + withTrees (mtd : rest) act = bracket (go mtd) releaseRef $ \t -> - case NE.nonEmpty rest of - Nothing -> - act (t :| []) - Just rest' -> - withTrees rest' $ \ts -> - act (NE.cons t ts) + withTrees rest $ \ts -> + act (t : ts) withMaybeTree Nothing act = act Nothing withMaybeTree (Just mtd) act = From 376edc3e165561dae67b6832d63048d04532625c Mon Sep 17 00:00:00 2001 From: Matthias Heinzel Date: Thu, 17 Sep 2026 17:47:49 +0200 Subject: [PATCH 5/5] throw proper exception --- .../Database/LSMTree/Internal/Snapshot.hs | 31 ++++++++++--------- .../Database/LSMTree/Internal/Unsafe.hs | 2 +- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs index 1a7e41ea7..2045930bc 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Snapshot.hs @@ -49,7 +49,7 @@ 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 Database.LSMTree.Internal.CRC32C (checkCRC, expectValidFile) import qualified Database.LSMTree.Internal.CRC32C as CRC import qualified Database.LSMTree.Internal.FS as FS import Database.LSMTree.Internal.IncomingRun @@ -238,6 +238,7 @@ instance NFData r => NFData (SnapPreExistingRun r) where -> Bloom.Salt -> UniqCounter IO -> ResolveSerialisedValue + -> FS.FsPath -> ActiveDir -> ActionRegistry IO -> SnapMergingTree (Ref (Run IO h)) @@ -254,11 +255,12 @@ fromSnapMergingTree :: -> Bloom.Salt -> UniqCounter m -> ResolveSerialisedValue + -> FS.FsPath -- ^ Path to source file for error reporting -> ActiveDir -> ActionRegistry m -> SnapMergingTree (Ref (Run m h)) -> m (Ref (MT.MergingTree m h)) -fromSnapMergingTree hfs hbio refCtx salt uc resolve dir = +fromSnapMergingTree hfs hbio refCtx salt uc resolve contentPath dir = go where -- Reference strategy: @@ -285,12 +287,12 @@ fromSnapMergingTree hfs hbio refCtx salt uc resolve dir = releaseRef traverse_ (delayedCommit reg . releasePER) prs' traverse_ (delayedCommit reg . releaseRef) mmt' - case mt of - Nothing -> - -- TODO: fail properly, or enforce invariant in SnapPendingLevelMerge - error "structurally empty pending level merge in snapshot" - Just mt' -> - pure mt' + expectValidFile hfs contentPath CRC.FormatSnapshotMetaData $ + case mt of + Nothing -> + Left "structurally empty pending level merge in snapshot" + Just mt' -> + Right mt' go reg (SnapMergingTree (SnapPendingTreeMerge (SnapPendingUnionMerge mts))) = do @@ -299,12 +301,13 @@ fromSnapMergingTree hfs hbio refCtx salt uc resolve dir = (MT.newPendingUnionMerge refCtx mts') releaseRef traverse_ (delayedCommit reg . releaseRef) mts' - case mt of - Nothing -> - -- TODO: fail properly, or enforce invariant in SnapPendingUnionMerge - error "structurally empty pending union merge in snapshot" - Just mt' -> - pure mt' + expectValidFile hfs contentPath CRC.FormatSnapshotMetaData $ + case mt of + Nothing -> + -- TODO: fail properly, or enforce invariant in SnapPendingUnionMerge + Left "structurally empty pending union merge in snapshot" + Just mt' -> + Right mt' go reg (SnapMergingTree (SnapOngoingTreeMerge smrs)) = do mr <- withRollback reg diff --git a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs index 0eac7e693..3c8e5aced 100644 --- a/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs +++ b/lsm-tree/src-core/Database/LSMTree/Internal/Unsafe.hs @@ -1817,7 +1817,7 @@ openTableFromSnapshot policyOveride sesh snap label resolve = do Nothing -> pure NoUnion Just mTree -> do snapTree <- traverse (openRun hfs hbio (sessionRefCtx seshEnv) uc reg snapDir activeDir salt) mTree - mt <- fromSnapMergingTree hfs hbio (sessionRefCtx seshEnv) salt uc resolve activeDir reg snapTree + mt <- fromSnapMergingTree hfs hbio (sessionRefCtx seshEnv) salt uc resolve contentPath activeDir reg snapTree traverse_ (delayedCommit reg . releaseRef) snapTree cache <- mkUnionCache reg mt pure (Union mt cache)