Skip to content

Commit 1b3d3d2

Browse files
mheinzelmheinzel
committed
prototype: support mupsert
1 parent bba3791 commit 1b3d3d2

File tree

3 files changed

+98
-26
lines changed

3 files changed

+98
-26
lines changed

‎lsm-tree.cabal‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -699,7 +699,6 @@ library prototypes
699699
, constraints
700700
, containers
701701
, contra-tracer
702-
, lsm-tree
703702
, QuickCheck
704703
, quickcheck-dynamic
705704
, quickcheck-lockstep

‎prototypes/ScheduledMerges.hs‎

Lines changed: 54 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ module ScheduledMerges (
2727
update, updates,
2828
insert, inserts,
2929
delete, deletes,
30+
mupsert, mupserts,
3031
supply,
3132
duplicate,
3233

@@ -52,8 +53,6 @@ import Control.Monad.ST
5253
import Control.Tracer (Tracer, contramap, traceWith)
5354
import GHC.Stack (HasCallStack, callStack)
5455

55-
import Database.LSMTree.Normal (LookupResult (..), Update (..))
56-
5756

5857
data LSM s = LSMHandle !(STRef s Counter)
5958
!(STRef s (LSMContent s))
@@ -329,13 +328,23 @@ newMerge tr level mergepolicy mergelast rs = do
329328
MergeLastLevel -> lastLevelMerge (mergek rs)
330329

331330
mergek :: [Run] -> Run
332-
mergek = Map.unions
331+
mergek = Map.unionsWith combine
332+
333+
combine :: Op -> Op -> Op
334+
combine x y = case x of
335+
Insert{} -> x
336+
Delete{} -> x
337+
Mupsert v -> case y of
338+
Insert v' mb -> Insert (resolveValue v v') mb
339+
Delete -> Insert v Nothing
340+
Mupsert v' -> Mupsert (resolveValue v v')
333341

334342
lastLevelMerge :: Run -> Run
335-
lastLevelMerge = Map.filter isInsert
343+
lastLevelMerge = Map.filter (not . isDelete)
336344
where
337-
isInsert Insert{} = True
338-
isInsert Delete = False
345+
isDelete Delete = True
346+
isDelete Insert{} = False
347+
isDelete Mupsert{} = False
339348

340349
expectCompletedMerge :: HasCallStack
341350
=> Tracer (ST s) EventDetail
@@ -429,18 +438,29 @@ new = do
429438
lsm <- newSTRef (LSMContent Map.empty [])
430439
return (LSMHandle c lsm)
431440

432-
433441
inserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value, Maybe Blob)] -> ST s ()
434442
inserts tr lsm kvbs = updates tr lsm [ (k, Insert v b) | (k, v, b) <- kvbs ]
435443

436444
insert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> Maybe Blob -> ST s ()
437445
insert tr lsm k v b = update tr lsm k (Insert v b)
438446

447+
deletes :: Tracer (ST s) Event -> LSM s -> [Key] -> ST s ()
448+
deletes tr lsm ks = updates tr lsm [ (k, Delete) | k <- ks ]
449+
439450
delete :: Tracer (ST s) Event -> LSM s -> Key -> ST s ()
440451
delete tr lsm k = update tr lsm k Delete
441452

442-
deletes :: Tracer (ST s) Event -> LSM s -> [Key] -> ST s ()
443-
deletes tr lsm ks = updates tr lsm [ (k, Delete) | k <- ks ]
453+
mupserts :: Tracer (ST s) Event -> LSM s -> [(Key, Value)] -> ST s ()
454+
mupserts tr lsm kvbs = updates tr lsm [ (k, Mupsert v) | (k, v) <- kvbs ]
455+
456+
mupsert :: Tracer (ST s) Event -> LSM s -> Key -> Value -> ST s ()
457+
mupsert tr lsm k v = update tr lsm k (Mupsert v)
458+
459+
data Update v b =
460+
Insert !v !(Maybe b)
461+
| Mupsert !v
462+
| Delete
463+
deriving stock (Eq, Show)
444464

445465
updates :: Tracer (ST s) Event -> LSM s -> [(Key, Op)] -> ST s ()
446466
updates tr lsm = mapM_ (uncurry (update tr lsm))
@@ -452,7 +472,7 @@ update tr (LSMHandle scr lsmr) k op = do
452472
modifySTRef' scr (+1)
453473
supplyCredits 1 ls
454474
invariant ls
455-
let wb' = Map.insert k op wb
475+
let wb' = Map.insertWith combine k op wb
456476
if bufferSize wb' >= maxBufferSize
457477
then do
458478
ls' <- increment tr sc (bufferToRun wb') ls
@@ -468,21 +488,32 @@ supply (LSMHandle scr lsmr) credits = do
468488
supplyCredits credits ls
469489
invariant ls
470490

491+
data LookupResult v b =
492+
NotFound
493+
| Found !v !(Maybe b)
494+
deriving stock (Eq, Show)
495+
471496
lookups :: LSM s -> [Key] -> ST s [(Key, LookupResult Value Blob)]
472-
lookups lsm = mapM (\k -> (k,) <$> lookup lsm k)
497+
lookups lsm ks = do
498+
runs <- concat <$> allLayers lsm
499+
return $ map (\k -> (k, doLookup k runs)) ks
473500

474501
lookup :: LSM s -> Key -> ST s (LookupResult Value Blob)
475502
lookup lsm k = do
476-
rss <- allLayers lsm
477-
return $!
478-
foldr (\lookures continue ->
479-
case lookures of
480-
Nothing -> continue
481-
Just (Insert v Nothing) -> Found v
482-
Just (Insert v (Just b)) -> FoundWithBlob v b
483-
Just Delete -> NotFound)
484-
NotFound
485-
[ Map.lookup k r | rs <- rss, r <- rs ]
503+
runs <- concat <$> allLayers lsm
504+
return $ doLookup k runs
505+
506+
doLookup :: Key -> [Run] -> LookupResult Value Blob
507+
doLookup k =
508+
foldr (\run continue ->
509+
case Map.lookup k run of
510+
Nothing -> continue
511+
Just (Insert v mb) -> Found v mb
512+
Just Delete -> NotFound
513+
Just (Mupsert v) -> case continue of
514+
NotFound -> Found v Nothing
515+
Found v' mb -> Found (resolveValue v v') mb)
516+
NotFound
486517

487518
bufferToRun :: Buffer -> Run
488519
bufferToRun = id
@@ -648,11 +679,12 @@ flattenIncomingRun (Merging (MergingRun _ _ mr)) = do
648679
OngoingMerge _ rs _ -> return rs
649680

650681
logicalValue :: LSM s -> ST s (Map Key (Value, Maybe Blob))
651-
logicalValue = fmap (Map.mapMaybe justInsert . Map.unions . concat)
682+
logicalValue = fmap (Map.mapMaybe justInsert . Map.unionsWith combine . concat)
652683
. allLayers
653684
where
654685
justInsert (Insert v b) = Just (v, b)
655686
justInsert Delete = Nothing
687+
justInsert (Mupsert v) = Just (v, Nothing)
656688

657689
dumpRepresentation :: LSM s
658690
-> ST s [(Maybe (MergePolicy, MergeLastLevel, MergingRunState), [Run])]

‎prototypes/ScheduledMergesTestQLS.hs‎

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ type ModelOp r = Model -> (r, Model)
4343
modelNew :: ModelOp ModelLSM
4444
modelInsert :: ModelLSM -> Key -> Value -> Maybe Blob -> ModelOp ()
4545
modelDelete :: ModelLSM -> Key -> ModelOp ()
46+
modelMupsert :: ModelLSM -> Key -> Value -> ModelOp ()
4647
modelLookup :: ModelLSM -> Key -> ModelOp (LookupResult Value Blob)
4748
modelDuplicate :: ModelLSM -> ModelOp ModelLSM
4849
modelDump :: ModelLSM -> ModelOp (Map Key (Value, Maybe Blob))
@@ -61,14 +62,18 @@ modelInsert mlsm k v b Model {mlsms} =
6162
modelDelete mlsm k Model {mlsms} =
6263
((), Model { mlsms = Map.adjust (Map.delete k) mlsm mlsms })
6364

65+
modelMupsert mlsm k v Model {mlsms} =
66+
((), Model { mlsms = Map.adjust (Map.insertWith f k (v, Nothing)) mlsm mlsms })
67+
where
68+
f _ (vOld, b) = (resolveValue v vOld, b)
69+
6470
modelLookup mlsm k model@Model {mlsms} =
6571
(result, model)
6672
where
6773
Just mval = Map.lookup mlsm mlsms
6874
result = case Map.lookup k mval of
69-
Nothing -> NotFound
70-
Just (v, Nothing) -> Found v
71-
Just (v, Just b) -> FoundWithBlob v b
75+
Nothing -> NotFound
76+
Just (v, mb) -> Found v mb
7277

7378
modelDuplicate mlsm Model {mlsms} =
7479
(mlsm', Model { mlsms = Map.insert mlsm' mval mlsms })
@@ -95,6 +100,11 @@ instance StateModel (Lockstep Model) where
95100
-> Either (ModelVar Model Key) Key
96101
-> Action (Lockstep Model) ()
97102

103+
AMupsert :: ModelVar Model (LSM RealWorld)
104+
-> Either (ModelVar Model Key) Key
105+
-> Value
106+
-> Action (Lockstep Model) (Key)
107+
98108
ALookup :: ModelVar Model (LSM RealWorld)
99109
-> Either (ModelVar Model Key) Key
100110
-> Action (Lockstep Model) (LookupResult Value Blob)
@@ -144,6 +154,8 @@ instance InLockstep Model where
144154
: case evk of Left vk -> [SomeGVar vk]; _ -> []
145155
usedVars (ADelete v evk) = SomeGVar v
146156
: case evk of Left vk -> [SomeGVar vk]; _ -> []
157+
usedVars (AMupsert v evk _) = SomeGVar v
158+
: case evk of Left vk -> [SomeGVar vk]; _ -> []
147159
usedVars (ALookup v evk) = SomeGVar v
148160
: case evk of Left vk -> [SomeGVar vk]; _ -> []
149161
usedVars (ADuplicate v) = [SomeGVar v]
@@ -185,6 +197,19 @@ instance InLockstep Model where
185197
<*> existingKey)
186198
| not (null kvars)
187199
]
200+
-- mupserts of potentially fresh keys
201+
++ [ (1, fmap Some $
202+
AMupsert <$> elements vars
203+
<*> freshKey
204+
<*> arbitrary @Value)
205+
]
206+
-- mupserts of the same keys as used earlier
207+
++ [ (1, fmap Some $
208+
AMupsert <$> elements vars
209+
<*> existingKey
210+
<*> arbitrary @Value)
211+
| not (null kvars)
212+
]
188213
-- lookup of arbitrary keys:
189214
++ [ (1, fmap Some $
190215
ALookup <$> elements vars
@@ -215,6 +240,14 @@ instance InLockstep Model where
215240
shrinkWithVars _ctx _model (ADelete var (Left _kv)) =
216241
[ Some $ ADelete var (Right k) | k <- shrink (K 100) ]
217242

243+
shrinkWithVars _ctx _model (AMupsert var (Right k) v) =
244+
[ Some $ AInsert var (Right k) v Nothing ] ++
245+
[ Some $ AMupsert var (Right k') v' | (k', v') <- shrink (k, v) ]
246+
247+
shrinkWithVars _ctx _model (AMupsert var (Left kv) v) =
248+
[ Some $ AInsert var (Left kv) v Nothing ] ++
249+
[ Some $ AMupsert var (Right k') v' | (k', v') <- shrink (K 100, v) ]
250+
218251
shrinkWithVars _ctx _model _action = []
219252

220253

@@ -224,13 +257,15 @@ instance RunLockstep Model IO where
224257
(ANew, _) -> ORef
225258
(AInsert{}, x) -> OId x
226259
(ADelete{}, x) -> OId x
260+
(AMupsert{}, x) -> OId x
227261
(ALookup{}, x) -> OId x
228262
(ADump{}, x) -> OId x
229263
(ADuplicate{}, _) -> ORef
230264

231265
showRealResponse _ ANew = Nothing
232266
showRealResponse _ AInsert{} = Just Dict
233267
showRealResponse _ ADelete{} = Just Dict
268+
showRealResponse _ AMupsert{} = Just Dict
234269
showRealResponse _ ALookup{} = Just Dict
235270
showRealResponse _ ADump{} = Just Dict
236271
showRealResponse _ ADuplicate{} = Nothing
@@ -255,6 +290,8 @@ runActionIO action lookUp =
255290
where k = either lookUpVar id evk
256291
ADelete var evk -> delete tr (lookUpVar var) k >> return ()
257292
where k = either lookUpVar id evk
293+
AMupsert var evk v -> mupsert tr (lookUpVar var) k v >> return k
294+
where k = either lookUpVar id evk
258295
ALookup var evk -> lookup (lookUpVar var) k
259296
where k = either lookUpVar id evk
260297
ADuplicate var -> duplicate (lookUpVar var)
@@ -283,6 +320,10 @@ runModel action ctx m =
283320
where ((), m') = modelDelete (lookUpLsMVar var) k m
284321
k = either lookUpKeyVar id evk
285322

323+
AMupsert var evk v -> (MInsert k, m')
324+
where ((), m') = modelMupsert (lookUpLsMVar var) k v m
325+
k = either lookUpKeyVar id evk
326+
286327
ALookup var evk -> (MLookup mv, m')
287328
where (mv, m') = modelLookup (lookUpLsMVar var) k m
288329
k = either lookUpKeyVar id evk

0 commit comments

Comments
 (0)