From 95c30f8b54a4b7a185c7eea906c6d4f10bc8a8e4 Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Thu, 18 Jan 2024 09:45:22 +0000 Subject: [PATCH] refactor: counter uses map --- yarn-project/kv-store/src/interfaces/map.ts | 8 +-- .../kv-store/src/lmdb/counter.test.ts | 7 ++ yarn-project/kv-store/src/lmdb/counter.ts | 53 ++++---------- yarn-project/kv-store/src/lmdb/map.test.ts | 45 +++++++++--- yarn-project/kv-store/src/lmdb/map.ts | 71 +++++++++++-------- 5 files changed, 104 insertions(+), 80 deletions(-) diff --git a/yarn-project/kv-store/src/interfaces/map.ts b/yarn-project/kv-store/src/interfaces/map.ts index 5e1b52d2a86..0916146a4ab 100644 --- a/yarn-project/kv-store/src/interfaces/map.ts +++ b/yarn-project/kv-store/src/interfaces/map.ts @@ -45,19 +45,19 @@ export interface AztecMap { delete(key: K): Promise; /** - * Iterates over the map's key-value entries + * Iterates over the map's key-value entries in the key's natural order * @param range - The range of keys to iterate over */ entries(range?: Range): IterableIterator<[K, V]>; /** - * Iterates over the map's values + * Iterates over the map's values in the key's natural order * @param range - The range of keys to iterate over */ values(range?: Range): IterableIterator; /** - * Iterates over the map's keys + * Iterates over the map's keys in the key's natural order * @param range - The range of keys to iterate over */ keys(range?: Range): IterableIterator; @@ -66,7 +66,7 @@ export interface AztecMap { /** * A map backed by a persistent store that can have multiple values for a single key. */ -export interface AztecMultiMap extends AztecMap { +export interface AztecMultiMap extends AztecMap { /** * Gets all the values at the given key. * @param key - The key to get the values from diff --git a/yarn-project/kv-store/src/lmdb/counter.test.ts b/yarn-project/kv-store/src/lmdb/counter.test.ts index 53ba60cdd02..fdd1204b15c 100644 --- a/yarn-project/kv-store/src/lmdb/counter.test.ts +++ b/yarn-project/kv-store/src/lmdb/counter.test.ts @@ -40,6 +40,13 @@ describe('LmdbAztecCounter', () => { expect(counter.get(key)).toEqual(0); }); + it('throws when decrementing below zero', async () => { + const key = genKey(); + await counter.update(key, 1); + + await expect(counter.update(key, -2)).rejects.toThrow(); + }); + it('increments values by a delta', async () => { const key = genKey(); await counter.update(key, 1); diff --git a/yarn-project/kv-store/src/lmdb/counter.ts b/yarn-project/kv-store/src/lmdb/counter.ts index 4604eb74458..74886e89dbf 100644 --- a/yarn-project/kv-store/src/lmdb/counter.ts +++ b/yarn-project/kv-store/src/lmdb/counter.ts @@ -2,36 +2,29 @@ import { Key as BaseKey, Database } from 'lmdb'; import { Key, Range } from '../interfaces/common.js'; import { AztecCounter } from '../interfaces/counter.js'; - -/** The slot where a key-value entry would be stored */ -type CountMapKey = ['count_map', string, 'slot', K]; +import { LmdbAztecMap } from './map.js'; /** * A counter implementation backed by LMDB */ export class LmdbAztecCounter implements AztecCounter { - #db: Database<[K, number], CountMapKey>; + #db: Database; #name: string; - - #startSentinel: CountMapKey; - #endSentinel: CountMapKey; + #map: LmdbAztecMap; constructor(db: Database, name: string) { + this.#db = db; this.#name = name; - this.#db = db as Database<[K, number], CountMapKey>; - - this.#startSentinel = ['count_map', this.#name, 'slot', Buffer.from([])]; - this.#endSentinel = ['count_map', this.#name, 'slot', Buffer.from([255])]; + this.#map = new LmdbAztecMap(db, name); } set(key: K, value: number): Promise { - return this.#db.put(this.#slot(key), [key, value]); + return this.#map.set(key, value); } update(key: K, delta = 1): Promise { return this.#db.childTransaction(() => { - const slot = this.#slot(key); - const [_, current] = this.#db.get(slot) ?? [key, 0]; + const current = this.#map.get(key) ?? 0; const next = current + delta; if (next < 0) { @@ -39,11 +32,11 @@ export class LmdbAztecCounter implements AztecCounter { } if (next === 0) { - void this.#db.remove(slot); + void this.#map.delete(key); } else { // store the key inside the entry because LMDB might return an internal representation // of the key when iterating over the database - void this.#db.put(slot, [key, next]); + void this.#map.set(key, next); } return true; @@ -51,32 +44,14 @@ export class LmdbAztecCounter implements AztecCounter { } get(key: K): number { - return (this.#db.get(this.#slot(key)) ?? [key, 0])[1]; - } - - *entries(range: Range = {}): IterableIterator<[K, number]> { - const { start, end, reverse, limit } = range; - const cursor = this.#db.getRange({ - start: start ? this.#slot(start) : reverse ? this.#endSentinel : this.#startSentinel, - end: end ? this.#slot(end) : reverse ? this.#startSentinel : this.#endSentinel, - reverse, - limit, - }); - - for (const { - value: [key, value], - } of cursor) { - yield [key, value]; - } + return this.#map.get(key) ?? 0; } - *keys(range: Range = {}): IterableIterator { - for (const [key] of this.entries(range)) { - yield key; - } + entries(range: Range = {}): IterableIterator<[K, number]> { + return this.#map.entries(range); } - #slot(key: K): CountMapKey { - return ['count_map', this.#name, 'slot', key]; + keys(range: Range = {}): IterableIterator { + return this.#map.keys(range); } } diff --git a/yarn-project/kv-store/src/lmdb/map.test.ts b/yarn-project/kv-store/src/lmdb/map.test.ts index 5319e0a26c3..007b4c4eb8f 100644 --- a/yarn-project/kv-store/src/lmdb/map.test.ts +++ b/yarn-project/kv-store/src/lmdb/map.test.ts @@ -41,26 +41,24 @@ describe('LmdbAztecMap', () => { await map.set('foo', 'bar'); await map.set('baz', 'qux'); - expect([...map.entries()]).toEqual( - expect.arrayContaining([ - ['foo', 'bar'], - ['baz', 'qux'], - ]), - ); + expect([...map.entries()]).toEqual([ + ['baz', 'qux'], + ['foo', 'bar'], + ]); }); it('should be able to iterate over values', async () => { await map.set('foo', 'bar'); - await map.set('baz', 'qux'); + await map.set('baz', 'quux'); - expect([...map.values()]).toEqual(expect.arrayContaining(['bar', 'qux'])); + expect([...map.values()]).toEqual(['quux', 'bar']); }); it('should be able to iterate over keys', async () => { await map.set('foo', 'bar'); await map.set('baz', 'qux'); - expect([...map.keys()]).toEqual(expect.arrayContaining(['foo', 'baz'])); + expect([...map.keys()]).toEqual(['baz', 'foo']); }); it('should be able to get multiple values for a single key', async () => { @@ -69,4 +67,33 @@ describe('LmdbAztecMap', () => { expect([...map.getValues('foo')]).toEqual(['bar', 'baz']); }); + + it('supports tuple keys', async () => { + const map = new LmdbAztecMap<[number, string], string>(db, 'test'); + + await map.set([5, 'bar'], 'val'); + await map.set([0, 'foo'], 'val'); + + expect([...map.keys()]).toEqual([ + [0, 'foo'], + [5, 'bar'], + ]); + + expect(map.get([5, 'bar'])).toEqual('val'); + }); + + it('supports range queries', async () => { + await map.set('a', 'a'); + await map.set('b', 'b'); + await map.set('c', 'c'); + await map.set('d', 'd'); + + expect([...map.keys({ start: 'b', end: 'c' })]).toEqual(['b']); + expect([...map.keys({ start: 'b' })]).toEqual(['b', 'c', 'd']); + expect([...map.keys({ end: 'c' })]).toEqual(['a', 'b']); + expect([...map.keys({ start: 'b', end: 'c', reverse: true })]).toEqual(['c']); + expect([...map.keys({ start: 'b', limit: 1 })]).toEqual(['b']); + expect([...map.keys({ start: 'b', reverse: true })]).toEqual(['d', 'c']); + expect([...map.keys({ end: 'b', reverse: true })]).toEqual(['b', 'a']); + }); }); diff --git a/yarn-project/kv-store/src/lmdb/map.ts b/yarn-project/kv-store/src/lmdb/map.ts index cc6df3e2c94..6e5fa67ef1e 100644 --- a/yarn-project/kv-store/src/lmdb/map.ts +++ b/yarn-project/kv-store/src/lmdb/map.ts @@ -1,24 +1,24 @@ -import { Database, Key } from 'lmdb'; +import { Database, RangeOptions } from 'lmdb'; -import { Range } from '../interfaces/common.js'; +import { Key, Range } from '../interfaces/common.js'; import { AztecMultiMap } from '../interfaces/map.js'; /** The slot where a key-value entry would be stored */ -type MapKeyValueSlot = ['map', string, 'slot', K]; +type MapValueSlot = ['map', string, 'slot', K]; /** * A map backed by LMDB. */ -export class LmdbAztecMap implements AztecMultiMap { - protected db: Database>; +export class LmdbAztecMap implements AztecMultiMap { + protected db: Database<[K, V], MapValueSlot>; protected name: string; - #startSentinel: MapKeyValueSlot; - #endSentinel: MapKeyValueSlot; + #startSentinel: MapValueSlot; + #endSentinel: MapValueSlot; - constructor(rootDb: Database, mapName: string) { + constructor(rootDb: Database, mapName: string) { this.name = mapName; - this.db = rootDb as Database>; + this.db = rootDb as Database<[K, V], MapValueSlot>; // sentinels are used to define the start and end of the map // with LMDB's key encoding, no _primitive value_ can be "less than" an empty buffer or greater than Byte 255 @@ -32,13 +32,13 @@ export class LmdbAztecMap implements AztecMultiMap } get(key: K): V | undefined { - return this.db.get(this.#slot(key)) as V | undefined; + return this.db.get(this.#slot(key))?.[1]; } *getValues(key: K): IterableIterator { const values = this.db.getValues(this.#slot(key)); for (const value of values) { - yield value; + yield value?.[1]; } } @@ -47,14 +47,14 @@ export class LmdbAztecMap implements AztecMultiMap } set(key: K, val: V): Promise { - return this.db.put(this.#slot(key), val); + return this.db.put(this.#slot(key), [key, val]); } swap(key: K, fn: (val: V | undefined) => V): Promise { return this.db.childTransaction(() => { const slot = this.#slot(key); - const val = this.db.get(slot); - void this.db.put(slot, fn(val)); + const entry = this.db.get(slot); + void this.db.put(slot, [key, fn(entry?.[1])]); return true; }); @@ -63,7 +63,7 @@ export class LmdbAztecMap implements AztecMultiMap setIfNotExists(key: K, val: V): Promise { const slot = this.#slot(key); return this.db.ifNoExists(slot, () => { - void this.db.put(slot, val); + void this.db.put(slot, [key, val]); }); } @@ -72,27 +72,42 @@ export class LmdbAztecMap implements AztecMultiMap } async deleteValue(key: K, val: V): Promise { - await this.db.remove(this.#slot(key), val); + await this.db.remove(this.#slot(key), [key, val]); } *entries(range: Range = {}): IterableIterator<[K, V]> { - const { start, end, reverse = false, limit } = range; + const { reverse = false, limit } = range; // LMDB has a quirk where it expects start > end when reverse=true // in that case, we need to swap the start and end sentinels - const iterator = this.db.getRange({ - start: start ? this.#slot(start) : reverse ? this.#endSentinel : this.#startSentinel, - end: end ? this.#slot(end) : reverse ? this.#startSentinel : this.#endSentinel, + const start = reverse + ? range.end + ? this.#slot(range.end) + : this.#endSentinel + : range.start + ? this.#slot(range.start) + : this.#startSentinel; + + const end = reverse + ? range.start + ? this.#slot(range.start) + : this.#startSentinel + : range.end + ? this.#slot(range.end) + : this.#endSentinel; + + const lmdbRange: RangeOptions = { + start, + end, reverse, limit, - }); + }; - for (const { key, value } of iterator) { - if (key[0] !== 'map' || key[1] !== this.name) { - break; - } + const iterator = this.db.getRange(lmdbRange); - const originalKey = key[3]; - yield [originalKey, value]; + for (const { + value: [key, value], + } of iterator) { + yield [key, value]; } } @@ -108,7 +123,7 @@ export class LmdbAztecMap implements AztecMultiMap } } - #slot(key: K): MapKeyValueSlot { + #slot(key: K): MapValueSlot { return ['map', this.name, 'slot', key]; } }