From c4ddad99dbdcbbb2659bc88b7271bf3c8525c7e2 Mon Sep 17 00:00:00 2001 From: Charlie Lye Date: Sat, 30 Mar 2024 10:54:11 +0000 Subject: [PATCH 1/3] feat: in memory kv store. --- yarn-project/kv-store/package.json | 1 + yarn-project/kv-store/src/lmdb/array.test.ts | 81 +-------- yarn-project/kv-store/src/lmdb/map.test.ts | 89 +--------- yarn-project/kv-store/src/mem/array.test.ts | 7 + yarn-project/kv-store/src/mem/array.ts | 70 ++++++++ yarn-project/kv-store/src/mem/counter.test.ts | 7 + yarn-project/kv-store/src/mem/counter.ts | 50 ++++++ yarn-project/kv-store/src/mem/index.ts | 1 + yarn-project/kv-store/src/mem/map.test.ts | 7 + yarn-project/kv-store/src/mem/map.ts | 163 ++++++++++++++++++ yarn-project/kv-store/src/mem/mem_db.ts | 53 ++++++ .../kv-store/src/mem/singleton.test.ts | 7 + yarn-project/kv-store/src/mem/singleton.ts | 27 +++ yarn-project/kv-store/src/mem/store.test.ts | 66 +++++++ yarn-project/kv-store/src/mem/store.ts | 54 ++++++ .../kv-store/src/tests/aztec_array_tests.ts | 81 +++++++++ .../kv-store/src/tests/aztec_counter_tests.ts | 120 +++++++++++++ .../kv-store/src/tests/aztec_map_tests.ts | 108 ++++++++++++ .../src/tests/aztec_singleton_tests.ts | 28 +++ yarn-project/kv-store/src/utils.ts | 5 +- 20 files changed, 857 insertions(+), 168 deletions(-) create mode 100644 yarn-project/kv-store/src/mem/array.test.ts create mode 100644 yarn-project/kv-store/src/mem/array.ts create mode 100644 yarn-project/kv-store/src/mem/counter.test.ts create mode 100644 yarn-project/kv-store/src/mem/counter.ts create mode 100644 yarn-project/kv-store/src/mem/index.ts create mode 100644 yarn-project/kv-store/src/mem/map.test.ts create mode 100644 yarn-project/kv-store/src/mem/map.ts create mode 100644 yarn-project/kv-store/src/mem/mem_db.ts create mode 100644 yarn-project/kv-store/src/mem/singleton.test.ts create mode 100644 yarn-project/kv-store/src/mem/singleton.ts create mode 100644 yarn-project/kv-store/src/mem/store.test.ts create mode 100644 yarn-project/kv-store/src/mem/store.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_array_tests.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_counter_tests.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_map_tests.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_singleton_tests.ts diff --git a/yarn-project/kv-store/package.json b/yarn-project/kv-store/package.json index d5bfb66aec49..c5ee9d3c536f 100644 --- a/yarn-project/kv-store/package.json +++ b/yarn-project/kv-store/package.json @@ -5,6 +5,7 @@ "exports": { ".": "./dest/interfaces/index.js", "./lmdb": "./dest/lmdb/index.js", + "./mem": "./dest/mem/index.js", "./utils": "./dest/utils.js" }, "scripts": { diff --git a/yarn-project/kv-store/src/lmdb/array.test.ts b/yarn-project/kv-store/src/lmdb/array.test.ts index 140d620a87d5..cbd5bb845ebd 100644 --- a/yarn-project/kv-store/src/lmdb/array.test.ts +++ b/yarn-project/kv-store/src/lmdb/array.test.ts @@ -1,91 +1,14 @@ import { type Database, open } from 'lmdb'; +import { addArrayTests } from '../tests/aztec_array_tests.js'; import { LmdbAztecArray } from './array.js'; describe('LmdbAztecArray', () => { let db: Database; - let arr: LmdbAztecArray; beforeEach(() => { db = open({} as any); - arr = new LmdbAztecArray(db, 'test'); }); - it('should be able to push and pop values', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect(arr.length).toEqual(3); - expect(await arr.pop()).toEqual(3); - expect(await arr.pop()).toEqual(2); - expect(await arr.pop()).toEqual(1); - expect(await arr.pop()).toEqual(undefined); - }); - - it('should be able to get values by index', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect(arr.at(0)).toEqual(1); - expect(arr.at(1)).toEqual(2); - expect(arr.at(2)).toEqual(3); - expect(arr.at(3)).toEqual(undefined); - expect(arr.at(-1)).toEqual(3); - expect(arr.at(-2)).toEqual(2); - expect(arr.at(-3)).toEqual(1); - expect(arr.at(-4)).toEqual(undefined); - }); - - it('should be able to set values by index', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect(await arr.setAt(0, 4)).toEqual(true); - expect(await arr.setAt(1, 5)).toEqual(true); - expect(await arr.setAt(2, 6)).toEqual(true); - - expect(await arr.setAt(3, 7)).toEqual(false); - - expect(arr.at(0)).toEqual(4); - expect(arr.at(1)).toEqual(5); - expect(arr.at(2)).toEqual(6); - expect(arr.at(3)).toEqual(undefined); - - expect(await arr.setAt(-1, 8)).toEqual(true); - expect(await arr.setAt(-2, 9)).toEqual(true); - expect(await arr.setAt(-3, 10)).toEqual(true); - - expect(await arr.setAt(-4, 11)).toEqual(false); - - expect(arr.at(-1)).toEqual(8); - expect(arr.at(-2)).toEqual(9); - expect(arr.at(-3)).toEqual(10); - expect(arr.at(-4)).toEqual(undefined); - }); - - it('should be able to iterate over values', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - expect([...arr.values()]).toEqual([1, 2, 3]); - expect([...arr.entries()]).toEqual([ - [0, 1], - [1, 2], - [2, 3], - ]); - }); - - it('should be able to restore state', async () => { - await arr.push(1); - await arr.push(2); - await arr.push(3); - - const arr2 = new LmdbAztecArray(db, 'test'); - expect(arr2.length).toEqual(3); - expect([...arr2.values()]).toEqual([...arr.values()]); - }); + addArrayTests(() => new LmdbAztecArray(db, 'test')); }); diff --git a/yarn-project/kv-store/src/lmdb/map.test.ts b/yarn-project/kv-store/src/lmdb/map.test.ts index a13aa16e1266..5025dfa47544 100644 --- a/yarn-project/kv-store/src/lmdb/map.test.ts +++ b/yarn-project/kv-store/src/lmdb/map.test.ts @@ -1,99 +1,14 @@ import { type Database, open } from 'lmdb'; +import { addMapTests } from '../tests/aztec_map_tests.js'; import { LmdbAztecMap } from './map.js'; describe('LmdbAztecMap', () => { let db: Database; - let map: LmdbAztecMap; beforeEach(() => { db = open({ dupSort: true } as any); - map = new LmdbAztecMap(db, 'test'); }); - it('should be able to set and get values', async () => { - await map.set('foo', 'bar'); - await map.set('baz', 'qux'); - - expect(map.get('foo')).toEqual('bar'); - expect(map.get('baz')).toEqual('qux'); - expect(map.get('quux')).toEqual(undefined); - }); - - it('should be able to set values if they do not exist', async () => { - expect(await map.setIfNotExists('foo', 'bar')).toEqual(true); - expect(await map.setIfNotExists('foo', 'baz')).toEqual(false); - - expect(map.get('foo')).toEqual('bar'); - }); - - it('should be able to delete values', async () => { - await map.set('foo', 'bar'); - await map.set('baz', 'qux'); - - await map.delete('foo'); - - expect(map.get('foo')).toEqual(undefined); - expect(map.get('baz')).toEqual('qux'); - }); - - it('should be able to iterate over entries', async () => { - await map.set('foo', 'bar'); - await map.set('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', 'quux'); - - 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(['baz', 'foo']); - }); - - it('should be able to get multiple values for a single key', async () => { - await map.set('foo', 'bar'); - await map.set('foo', 'baz'); - - 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']); - }); + addMapTests(() => new LmdbAztecMap(db, 'test')); }); diff --git a/yarn-project/kv-store/src/mem/array.test.ts b/yarn-project/kv-store/src/mem/array.test.ts new file mode 100644 index 000000000000..e01dd4c57119 --- /dev/null +++ b/yarn-project/kv-store/src/mem/array.test.ts @@ -0,0 +1,7 @@ +import { addArrayTests } from '../tests/aztec_array_tests.js'; +import { MemAztecArray } from './array.js'; +import { MemDb } from './mem_db.js'; + +describe('MemAztecArray', () => { + addArrayTests(() => new MemAztecArray('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/array.ts b/yarn-project/kv-store/src/mem/array.ts new file mode 100644 index 000000000000..3350fda7f51e --- /dev/null +++ b/yarn-project/kv-store/src/mem/array.ts @@ -0,0 +1,70 @@ +import { type AztecArray } from '../interfaces/array.js'; +import type { MemDb } from './mem_db.js'; + +/** + * An persistent array backed by mem. + */ +export class MemAztecArray implements AztecArray { + private slot: string; + + constructor(private name: string, private db: MemDb) { + this.slot = JSON.stringify(['array', this.name]); + } + + get length(): number { + return this.db.get(this.slot)?.length || 0; + } + + push(...vals: T[]): Promise { + const arr = this.db.get(this.slot); + if (arr) { + this.db.set(this.slot, [...arr, ...vals]); + } else { + this.db.set(this.slot, [...vals]); + } + return Promise.resolve(this.length); + } + + pop(): Promise { + const arr = [...this.db.get(this.slot)]; + const result = arr.pop(); + this.db.set(this.slot, arr); + return Promise.resolve(result); + } + + at(index: number): T | undefined { + const arr = this.db.get(this.slot) || []; + if (index < 0) { + return arr[arr.length + index]; + } else { + return arr[index]; + } + } + + setAt(index: number, val: T): Promise { + if (index < 0) { + index = this.length + index; + } + + if (index < 0 || index >= this.length) { + return Promise.resolve(false); + } + + const arr = [...this.db.get(this.slot)]; + arr[index] = val; + this.db.set(this.slot, arr); + return Promise.resolve(true); + } + + entries(): IterableIterator<[number, T]> { + return this.db.get(this.slot)?.entries() || []; + } + + values(): IterableIterator { + return this.db.get(this.slot)?.values() || []; + } + + [Symbol.iterator](): IterableIterator { + return this.values(); + } +} diff --git a/yarn-project/kv-store/src/mem/counter.test.ts b/yarn-project/kv-store/src/mem/counter.test.ts new file mode 100644 index 000000000000..17df911b542a --- /dev/null +++ b/yarn-project/kv-store/src/mem/counter.test.ts @@ -0,0 +1,7 @@ +import { addCounterTests } from '../tests/aztec_counter_tests.js'; +import { MemAztecCounter } from './counter.js'; +import { MemDb } from './mem_db.js'; + +describe('MemAztecCounter', () => { + addCounterTests(() => new MemAztecCounter('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/counter.ts b/yarn-project/kv-store/src/mem/counter.ts new file mode 100644 index 000000000000..13a71b65e9cc --- /dev/null +++ b/yarn-project/kv-store/src/mem/counter.ts @@ -0,0 +1,50 @@ +import type { Key, Range } from '../interfaces/common.js'; +import type { AztecCounter } from '../interfaces/counter.js'; +import { MemAztecMap } from './map.js'; +import type { MemDb } from './mem_db.js'; + +export class MemAztecCounter implements AztecCounter { + private map: MemAztecMap; + + constructor(name: string, db: MemDb) { + this.map = new MemAztecMap(name, db); + } + + async set(key: Key, value: number): Promise { + if (value) { + return this.map.set(key, value); + } else { + await this.map.delete(key); + return true; + } + } + + async update(key: Key, delta = 1): Promise { + const current = this.map.get(key) ?? 0; + const next = current + delta; + + if (next < 0) { + throw new Error(`Cannot update ${key} in counter below zero`); + } + + await this.map.delete(key); + + if (next > 0) { + await this.map.set(key, next); + } + + return true; + } + + get(key: Key): number { + return this.map.get(key) ?? 0; + } + + entries(range: Range = {}): IterableIterator<[Key, number]> { + return this.map.entries(range); + } + + keys(range: Range = {}): IterableIterator { + return this.map.keys(range); + } +} diff --git a/yarn-project/kv-store/src/mem/index.ts b/yarn-project/kv-store/src/mem/index.ts new file mode 100644 index 000000000000..f5aee46d6d71 --- /dev/null +++ b/yarn-project/kv-store/src/mem/index.ts @@ -0,0 +1 @@ +export { AztecMemStore } from './store.js'; diff --git a/yarn-project/kv-store/src/mem/map.test.ts b/yarn-project/kv-store/src/mem/map.test.ts new file mode 100644 index 000000000000..c9c437db8d48 --- /dev/null +++ b/yarn-project/kv-store/src/mem/map.test.ts @@ -0,0 +1,7 @@ +import { addMapTests } from '../tests/aztec_map_tests.js'; +import { MemAztecMap } from './map.js'; +import { MemDb } from './mem_db.js'; + +describe('MemAztecMap', () => { + addMapTests(() => new MemAztecMap('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/map.ts b/yarn-project/kv-store/src/mem/map.ts new file mode 100644 index 000000000000..c4f7fa916757 --- /dev/null +++ b/yarn-project/kv-store/src/mem/map.ts @@ -0,0 +1,163 @@ +import type { Key, Range } from '../interfaces/common.js'; +import type { AztecMultiMap } from '../interfaces/map.js'; +import type { MemDb } from './mem_db.js'; + +// Comparator function for keys already parsed from JSON +function compareKeys(a: Key, b: Key) { + // Handle array (tuple) comparison + if (Array.isArray(a) && Array.isArray(b)) { + for (let i = 0; i < Math.min(a.length, b.length); i++) { + if (a[i] < b[i]) { + return -1; + } + if (a[i] > b[i]) { + return 1; + } + } + return a.length - b.length; + } + + // Fallback to normal comparison for non-array keys + if (a < b) { + return -1; + } + if (a > b) { + return 1; + } + return 0; +} + +/** + * A map backed by mem. + */ +export class MemAztecMap implements AztecMultiMap { + constructor(private name: string, private db: MemDb) {} + + close(): Promise { + return Promise.resolve(); + } + + get(key: Key): V | undefined { + const r = this.db.get(this.slot(key)); + return r ? r[r.length - 1] : undefined; + } + + getValues(key: Key): IterableIterator { + const r = this.db.get(this.slot(key)); + return r ? r.values() : new Array().values(); + } + + has(key: Key): boolean { + const r = this.db.get(this.slot(key)); + return r ? r.length > 0 : false; + } + + set(key: Key, val: V): Promise { + const r = this.db.get(this.slot(key)); + if (r) { + this.db.set(this.slot(key), [...r, val]); + } else { + this.db.set(this.slot(key), [val]); + } + return Promise.resolve(true); + } + + swap(key: Key, fn: (val: V | undefined) => V): Promise { + const entry = this.get(key); + const newValue = fn(entry); + this.db.set(this.slot(key), [newValue]); + return Promise.resolve(true); + } + + async setIfNotExists(key: Key, val: V): Promise { + const r = this.get(key); + if (!r) { + await this.set(key, val); + return true; + } + return false; + } + + delete(key: Key): Promise { + const r = this.db.get(this.slot(key)); + if (r?.length) { + this.db.set(this.slot(key), []); + return Promise.resolve(true); + } + return Promise.resolve(false); + } + + deleteValue(key: Key, val: V): Promise { + const r = this.db.get(this.slot(key)); + if (r) { + const i = r.indexOf(val); + if (i != -1) { + this.db.set(this.slot(key), [...r.slice(0, i), ...r.slice(i + 1)]); + } + } + return Promise.resolve(); + } + + *entries(range: Range = {}): IterableIterator<[Key, V]> { + let { limit } = range; + const { start, end, reverse = false } = range; + + // TODO: Horrifically inefficient as backing db is not an ordered map. + // Make it so. + const keys = this.db + .keys() + .map(key => JSON.parse(key)) + .filter(key => key[0] == 'map' && key[1] == this.name) + .map(key => key[2]) + .sort(compareKeys); + + if (reverse) { + keys.reverse(); + } + + for (const key of keys) { + if (reverse) { + if (end !== undefined && compareKeys(key, end) > 0) { + continue; + } + if (start !== undefined && compareKeys(key, start) <= 0) { + break; + } + } else { + if (start !== undefined && compareKeys(key, start) < 0) { + continue; + } + if (end !== undefined && compareKeys(key, end) >= 0) { + break; + } + } + + const values = this.db.get(this.slot(key)); + if (!values) { + return; + } + for (const value of values) { + yield [key, value]; + if (limit && --limit <= 0) { + return; + } + } + } + } + + *values(range: Range = {}): IterableIterator { + for (const [_, value] of this.entries(range)) { + yield value; + } + } + + *keys(range: Range = {}): IterableIterator { + for (const [key, _] of this.entries(range)) { + yield key; + } + } + + private slot(key: Key): string { + return JSON.stringify(['map', this.name, key]); + } +} diff --git a/yarn-project/kv-store/src/mem/mem_db.ts b/yarn-project/kv-store/src/mem/mem_db.ts new file mode 100644 index 000000000000..aca1567fd21b --- /dev/null +++ b/yarn-project/kv-store/src/mem/mem_db.ts @@ -0,0 +1,53 @@ +export class MemDb { + private data: { [key: string]: any } = {}; + private tx: { [key: string]: any } | undefined; + + set(key: string, value: any) { + if (this.tx) { + this.tx[key] = value; + } else { + this.data[key] = value; + } + } + + get(key: string) { + if (this.tx && this.tx[key] !== undefined) { + return this.tx[key]; + } + return this.data[key]; + } + + del(key: string) { + if (this.tx) { + this.tx[key] = undefined; + return; + } + delete this.data[key]; + } + + keys() { + return Array.from(new Set([...Object.keys(this.data), ...Object.keys(this.tx || {})])); + } + + commit() { + if (!this.tx) { + throw new Error('tx not in progress.'); + } + Object.assign(this.data, this.tx); + this.tx = undefined; + } + + rollback() { + if (!this.tx) { + throw new Error('tx not in progress.'); + } + this.tx = undefined; + } + + startTx() { + if (this.tx) { + throw new Error('MemDb can only handle 1 tx at a time.'); + } + this.tx = {}; + } +} diff --git a/yarn-project/kv-store/src/mem/singleton.test.ts b/yarn-project/kv-store/src/mem/singleton.test.ts new file mode 100644 index 000000000000..9ea8533bde70 --- /dev/null +++ b/yarn-project/kv-store/src/mem/singleton.test.ts @@ -0,0 +1,7 @@ +import { addSingletonTests } from '../tests/aztec_singleton_tests.js'; +import { MemDb } from './mem_db.js'; +import { MemAztecSingleton } from './singleton.js'; + +describe('MemAztecSingleton', () => { + addSingletonTests(() => new MemAztecSingleton('test', new MemDb())); +}); diff --git a/yarn-project/kv-store/src/mem/singleton.ts b/yarn-project/kv-store/src/mem/singleton.ts new file mode 100644 index 000000000000..147fc94ac4db --- /dev/null +++ b/yarn-project/kv-store/src/mem/singleton.ts @@ -0,0 +1,27 @@ +import type { AztecSingleton } from '../interfaces/singleton.js'; +import type { MemDb } from './mem_db.js'; + +/** + * Stores a single value in mem. + */ +export class MemAztecSingleton implements AztecSingleton { + private slot: string; + + constructor(private name: string, private db: MemDb) { + this.slot = JSON.stringify(['array', this.name]); + } + + get(): T | undefined { + return this.db.get(this.slot); + } + + set(val: T): Promise { + this.db.set(this.slot, val); + return Promise.resolve(true); + } + + delete(): Promise { + this.db.del(this.slot); + return Promise.resolve(true); + } +} diff --git a/yarn-project/kv-store/src/mem/store.test.ts b/yarn-project/kv-store/src/mem/store.test.ts new file mode 100644 index 000000000000..7ffd3b87a8b1 --- /dev/null +++ b/yarn-project/kv-store/src/mem/store.test.ts @@ -0,0 +1,66 @@ +import type { AztecArray, AztecCounter, AztecMap, AztecSingleton } from '../interfaces/index.js'; +import { AztecMemStore } from './store.js'; + +describe('AztecMemStore', () => { + let store: AztecMemStore; + let array: AztecArray; + let multimap: AztecMap; + let counter: AztecCounter; + let singleton: AztecSingleton; + + beforeEach(async () => { + store = new AztecMemStore(); + + array = store.openArray('test-array'); + multimap = store.openMultiMap('test-multimap'); + counter = store.openCounter('test-counter'); + singleton = store.openSingleton('test-singleton'); + + await array.push(1, 2, 3); + await multimap.set('key-1', 1); + await multimap.set('key-2', 2); + await counter.set('counter-1', 3); + await singleton.set(4); + }); + + it('check initial state', () => { + expect(array.at(2)).toBe(3); + expect(multimap.get('key-2')).toBe(2); + expect(counter.get('counter-1')).toBe(3); + expect(singleton.get()).toBe(4); + }); + + it('state should update with successful tx', async () => { + await store.transaction(() => { + void array.setAt(2, 10); + void multimap.set('key-2', 20); + void counter.set('counter-1', 30); + void singleton.set(40); + }); + void multimap.set('key-2', 20); + + expect(array.at(2)).toBe(10); + expect(multimap.get('key-2')).toBe(20); + expect(counter.get('counter-1')).toBe(30); + expect(singleton.get()).toBe(40); + }); + + it('state should rollback with unsuccessful tx', async () => { + try { + await store.transaction(() => { + void array.setAt(2, 10); + void multimap.set('key-2', 20); + void counter.set('counter-1', 30); + void singleton.set(40); + throw new Error(); + }); + } catch (err) { + // swallow + } + + expect(array.at(2)).toBe(3); + expect(multimap.get('key-2')).toBe(2); + expect(counter.get('counter-1')).toBe(3); + expect(singleton.get()).toBe(4); + }); +}); diff --git a/yarn-project/kv-store/src/mem/store.ts b/yarn-project/kv-store/src/mem/store.ts new file mode 100644 index 000000000000..c35f47a5718d --- /dev/null +++ b/yarn-project/kv-store/src/mem/store.ts @@ -0,0 +1,54 @@ +import type { AztecArray } from '../interfaces/array.js'; +import type { AztecCounter } from '../interfaces/counter.js'; +import type { AztecMap, AztecMultiMap } from '../interfaces/map.js'; +import type { AztecSingleton } from '../interfaces/singleton.js'; +import type { AztecKVStore } from '../interfaces/store.js'; +import { MemAztecArray } from './array.js'; +import { MemAztecCounter } from './counter.js'; +import { MemAztecMap } from './map.js'; +import { MemDb } from './mem_db.js'; +import { MemAztecSingleton } from './singleton.js'; + +/** + * A key-value store backed by mem. + */ +export class AztecMemStore implements AztecKVStore { + private data = new MemDb(); + + openMap(name: string): AztecMap { + return new MemAztecMap(name, this.data) as any; + } + + openMultiMap(name: string): AztecMultiMap { + return new MemAztecMap(name, this.data) as any; + } + + openCounter>(name: string): AztecCounter { + return new MemAztecCounter(name, this.data) as any; + } + + openArray(name: string): AztecArray { + return new MemAztecArray(name, this.data) as any; + } + + openSingleton(name: string): AztecSingleton { + return new MemAztecSingleton(name, this.data) as any; + } + + transaction(callback: () => T): Promise { + this.data.startTx(); + try { + const result = callback(); + this.data.commit(); + return Promise.resolve(result); + } catch (err) { + this.data.rollback(); + throw err; + } + } + + clear() { + this.data = new MemDb(); + return Promise.resolve(); + } +} diff --git a/yarn-project/kv-store/src/tests/aztec_array_tests.ts b/yarn-project/kv-store/src/tests/aztec_array_tests.ts new file mode 100644 index 000000000000..19fa70b1c3d0 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_array_tests.ts @@ -0,0 +1,81 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type AztecArray } from '../interfaces/index.js'; + +export function addArrayTests(getArray: () => AztecArray) { + let arr: AztecArray; + + describe('AztecArray', () => { + beforeEach(() => { + arr = getArray(); + }); + + it('should be able to push and pop values', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect(arr.length).toEqual(3); + expect(await arr.pop()).toEqual(3); + expect(await arr.pop()).toEqual(2); + expect(await arr.pop()).toEqual(1); + expect(await arr.pop()).toEqual(undefined); + }); + + it('should be able to get values by index', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect(arr.at(0)).toEqual(1); + expect(arr.at(1)).toEqual(2); + expect(arr.at(2)).toEqual(3); + expect(arr.at(3)).toEqual(undefined); + expect(arr.at(-1)).toEqual(3); + expect(arr.at(-2)).toEqual(2); + expect(arr.at(-3)).toEqual(1); + expect(arr.at(-4)).toEqual(undefined); + }); + + it('should be able to set values by index', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect(await arr.setAt(0, 4)).toEqual(true); + expect(await arr.setAt(1, 5)).toEqual(true); + expect(await arr.setAt(2, 6)).toEqual(true); + + expect(await arr.setAt(3, 7)).toEqual(false); + + expect(arr.at(0)).toEqual(4); + expect(arr.at(1)).toEqual(5); + expect(arr.at(2)).toEqual(6); + expect(arr.at(3)).toEqual(undefined); + + expect(await arr.setAt(-1, 8)).toEqual(true); + expect(await arr.setAt(-2, 9)).toEqual(true); + expect(await arr.setAt(-3, 10)).toEqual(true); + + expect(await arr.setAt(-4, 11)).toEqual(false); + + expect(arr.at(-1)).toEqual(8); + expect(arr.at(-2)).toEqual(9); + expect(arr.at(-3)).toEqual(10); + expect(arr.at(-4)).toEqual(undefined); + }); + + it('should be able to iterate over values', async () => { + await arr.push(1); + await arr.push(2); + await arr.push(3); + + expect([...arr.values()]).toEqual([1, 2, 3]); + expect([...arr.entries()]).toEqual([ + [0, 1], + [1, 2], + [2, 3], + ]); + }); + }); +} diff --git a/yarn-project/kv-store/src/tests/aztec_counter_tests.ts b/yarn-project/kv-store/src/tests/aztec_counter_tests.ts new file mode 100644 index 000000000000..0a5a1b3aa2ad --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_counter_tests.ts @@ -0,0 +1,120 @@ +import { randomBytes } from '@aztec/foundation/crypto'; + +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type AztecCounter } from '../interfaces/index.js'; + +export function addCounterTests(get: () => AztecCounter) { + describe('AztecCounter', () => { + describe.each([ + ['floating point number', () => Math.random()], + ['integers', () => (Math.random() * 1000) | 0], + ['strings', () => randomBytes(8).toString('hex')], + ['strings', () => [Math.random(), randomBytes(8).toString('hex')]], + ])('counts occurrences of %s values', (_, genKey) => { + let counter: AztecCounter; + + beforeEach(() => { + counter = get(); + }); + + it('returns 0 for unknown keys', () => { + expect(counter.get(genKey())).toEqual(0); + }); + + it('increments values', async () => { + const key = genKey(); + await counter.update(key, 1); + + expect(counter.get(key)).toEqual(1); + }); + + it('decrements values', async () => { + const key = genKey(); + await counter.update(key, 1); + await counter.update(key, -1); + + 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); + await counter.update(key, 2); + + expect(counter.get(key)).toEqual(3); + }); + + it('resets the counter', async () => { + const key = genKey(); + await counter.update(key, 1); + await counter.update(key, 2); + await counter.set(key, 0); + + expect(counter.get(key)).toEqual(0); + }); + + it('iterates over entries', async () => { + const key = genKey(); + await counter.update(key, 1); + await counter.update(key, 2); + + expect([...counter.entries({})]).toEqual([[key, 3]]); + }); + }); + + it.each([ + [ + [ + ['c', 2342], + ['a', 8], + ['b', 1], + ], + [ + ['a', 8], + ['b', 1], + ['c', 2342], + ], + ], + [ + [ + [10, 2], + [18, 1], + [1, 2], + ], + [ + [1, 2], + [10, 2], + [18, 1], + ], + ], + [ + [ + [[10, 'a'], 1], + [[10, 'c'], 2], + [[11, 'b'], 1], + [[9, 'f'], 1], + [[10, 'b'], 1], + ], + [ + [[9, 'f'], 1], + [[10, 'a'], 1], + [[10, 'b'], 1], + [[10, 'c'], 2], + [[11, 'b'], 1], + ], + ], + ])('iterates in key order', async (insertOrder, expectedOrder) => { + const counter = get(); + await Promise.all(insertOrder.map(([key, value]) => counter.update(key, value as number))); + expect([...counter.entries({})]).toEqual(expectedOrder); + }); + }); +} diff --git a/yarn-project/kv-store/src/tests/aztec_map_tests.ts b/yarn-project/kv-store/src/tests/aztec_map_tests.ts new file mode 100644 index 000000000000..1c5efe071999 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_map_tests.ts @@ -0,0 +1,108 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type Key } from '../interfaces/common.js'; +import { type AztecMultiMap } from '../interfaces/index.js'; + +export function addMapTests(get: () => AztecMultiMap) { + describe('AztecMap', () => { + let map: AztecMultiMap; + + beforeEach(() => { + map = get(); + }); + + it('should be able to set and get values', async () => { + await map.set('foo', 'bar'); + await map.set('baz', 'qux'); + + expect(map.get('foo')).toEqual('bar'); + expect(map.get('baz')).toEqual('qux'); + expect(map.get('quux')).toEqual(undefined); + }); + + it('should be able to update values', async () => { + await map.set('foo', 'bar'); + expect(map.get('foo')).toEqual('bar'); + + await map.set('foo', 'qux'); + expect(map.get('foo')).toEqual('qux'); + }); + + it('should be able to set values if they do not exist', async () => { + expect(await map.setIfNotExists('foo', 'bar')).toEqual(true); + expect(await map.setIfNotExists('foo', 'baz')).toEqual(false); + + expect(map.get('foo')).toEqual('bar'); + }); + + it('should be able to delete values', async () => { + await map.set('foo', 'bar'); + await map.set('baz', 'qux'); + + await map.delete('foo'); + + expect(map.get('foo')).toEqual(undefined); + expect(map.get('baz')).toEqual('qux'); + }); + + it('should be able to iterate over entries', async () => { + await map.set('foo', 'bar'); + await map.set('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', 'quux'); + + 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(['baz', 'foo']); + }); + + it('should be able to get multiple values for a single key', async () => { + await map.set('foo', 'bar'); + await map.set('foo', 'baz'); + + expect([...map.getValues('foo')]).toEqual(['bar', 'baz']); + }); + + it('supports tuple keys', async () => { + const map = get(); + + 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/tests/aztec_singleton_tests.ts b/yarn-project/kv-store/src/tests/aztec_singleton_tests.ts new file mode 100644 index 000000000000..5247f1ad3d49 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_singleton_tests.ts @@ -0,0 +1,28 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import { type Key } from '../interfaces/common.js'; +import { type AztecSingleton } from '../interfaces/singleton.js'; + +export function addSingletonTests(get: () => AztecSingleton) { + describe('AztecSingleton', () => { + let singleton: AztecSingleton; + beforeEach(() => { + singleton = get(); + }); + + it('returns undefined if the value is not set', () => { + expect(singleton.get()).toEqual(undefined); + }); + + it('should be able to set and get values', async () => { + expect(await singleton.set('foo')).toEqual(true); + expect(singleton.get()).toEqual('foo'); + }); + + it('overwrites the value if it is set again', async () => { + expect(await singleton.set('foo')).toEqual(true); + expect(await singleton.set('bar')).toEqual(true); + expect(singleton.get()).toEqual('bar'); + }); + }); +} diff --git a/yarn-project/kv-store/src/utils.ts b/yarn-project/kv-store/src/utils.ts index d527e43a8b5c..535445ae8c90 100644 --- a/yarn-project/kv-store/src/utils.ts +++ b/yarn-project/kv-store/src/utils.ts @@ -2,7 +2,7 @@ import { type EthAddress } from '@aztec/foundation/eth-address'; import { type Logger } from '@aztec/foundation/log'; import { type AztecKVStore } from './interfaces/store.js'; -import { AztecLmdbStore } from './lmdb/store.js'; +import { AztecMemStore } from './mem/store.js'; /** * Clears the store if the rollup address does not match the one stored in the database. @@ -38,5 +38,6 @@ export async function initStoreForRollup( * @returns A new store */ export function openTmpStore(ephemeral: boolean = false): AztecKVStore { - return AztecLmdbStore.open(undefined, ephemeral); + return new AztecMemStore(); + // return AztecLmdbStore.open(undefined, ephemeral); } From f220e91fbd14a8c5496eb76962d5f132e7a7d561 Mon Sep 17 00:00:00 2001 From: Charlie Lye Date: Sat, 30 Mar 2024 12:30:37 +0000 Subject: [PATCH 2/3] Initial mem store. --- .../aztec-node/src/aztec-node/server.ts | 11 +-- .../aztec/src/cli/cmds/start_archiver.ts | 12 ++-- yarn-project/kv-store/src/lmdb/store.test.ts | 6 ++ yarn-project/kv-store/src/mem/array.ts | 6 +- yarn-project/kv-store/src/mem/counter.ts | 2 +- yarn-project/kv-store/src/mem/map.ts | 6 +- yarn-project/kv-store/src/mem/store.test.ts | 64 +---------------- yarn-project/kv-store/src/mem/store.ts | 4 +- .../kv-store/src/tests/aztec_array_tests.ts | 1 + .../kv-store/src/tests/aztec_map_tests.ts | 10 +-- .../kv-store/src/tests/aztec_store_tests.ts | 71 +++++++++++++++++++ yarn-project/kv-store/src/utils.ts | 1 - .../pxe/src/pxe_service/create_pxe_service.ts | 11 +-- 13 files changed, 110 insertions(+), 95 deletions(-) create mode 100644 yarn-project/kv-store/src/lmdb/store.test.ts create mode 100644 yarn-project/kv-store/src/tests/aztec_store_tests.ts diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts index e00c0b93f145..1db336f124b0 100644 --- a/yarn-project/aztec-node/src/aztec-node/server.ts +++ b/yarn-project/aztec-node/src/aztec-node/server.ts @@ -47,6 +47,7 @@ import { padArrayEnd } from '@aztec/foundation/collection'; import { createDebugLogger } from '@aztec/foundation/log'; import { type AztecKVStore } from '@aztec/kv-store'; import { AztecLmdbStore } from '@aztec/kv-store/lmdb'; +import { AztecMemStore } from '@aztec/kv-store/mem'; import { initStoreForRollup, openTmpStore } from '@aztec/kv-store/utils'; import { SHA256Trunc, StandardTree } from '@aztec/merkle-tree'; import { AztecKVTxPool, type P2P, createP2PClient } from '@aztec/p2p'; @@ -121,11 +122,11 @@ export class AztecNodeService implements AztecNode { const log = createDebugLogger('aztec:node'); const storeLog = createDebugLogger('aztec:node:lmdb'); - const store = await initStoreForRollup( - AztecLmdbStore.open(config.dataDirectory, false, storeLog), - config.l1Contracts.rollupAddress, - storeLog, - ); + const storeDb = config.dataDirectory + ? AztecLmdbStore.open(config.dataDirectory, false, storeLog) + : new AztecMemStore(); + + const store = await initStoreForRollup(storeDb, config.l1Contracts.rollupAddress, storeLog); let archiver: ArchiveSource; if (!config.archiverUrl) { diff --git a/yarn-project/aztec/src/cli/cmds/start_archiver.ts b/yarn-project/aztec/src/cli/cmds/start_archiver.ts index 567f2b160b04..c9c980dd2136 100644 --- a/yarn-project/aztec/src/cli/cmds/start_archiver.ts +++ b/yarn-project/aztec/src/cli/cmds/start_archiver.ts @@ -5,9 +5,10 @@ import { createArchiverRpcServer, getConfigEnvVars as getArchiverConfigEnvVars, } from '@aztec/archiver'; -import { createDebugLogger } from '@aztec/aztec.js'; import { type ServerList } from '@aztec/foundation/json-rpc/server'; +import { createDebugLogger } from '@aztec/foundation/log'; import { AztecLmdbStore } from '@aztec/kv-store/lmdb'; +import { AztecMemStore } from '@aztec/kv-store/mem'; import { initStoreForRollup } from '@aztec/kv-store/utils'; import { mergeEnvVarsAndCliOptions, parseModuleOptions } from '../util.js'; @@ -23,11 +24,10 @@ export const startArchiver = async (options: any, signalHandlers: (() => Promise const archiverConfig = mergeEnvVarsAndCliOptions(archiverConfigEnvVars, archiverCliOptions, true); const storeLog = createDebugLogger('aztec:archiver:lmdb'); - const store = await initStoreForRollup( - AztecLmdbStore.open(archiverConfig.dataDirectory, false, storeLog), - archiverConfig.l1Contracts.rollupAddress, - storeLog, - ); + const storeDb = archiverConfig.dataDirectory + ? AztecLmdbStore.open(archiverConfig.dataDirectory, false, storeLog) + : new AztecMemStore(); + const store = await initStoreForRollup(storeDb, archiverConfig.l1Contracts.rollupAddress, storeLog); const archiverStore = new KVArchiverDataStore(store, archiverConfig.maxLogs); const archiver = await Archiver.createAndSync(archiverConfig, archiverStore, true); diff --git a/yarn-project/kv-store/src/lmdb/store.test.ts b/yarn-project/kv-store/src/lmdb/store.test.ts new file mode 100644 index 000000000000..3409aa02aa8b --- /dev/null +++ b/yarn-project/kv-store/src/lmdb/store.test.ts @@ -0,0 +1,6 @@ +import { addStoreTests } from '../tests/aztec_store_tests.js'; +import { AztecLmdbStore } from './store.js'; + +describe('AztecLmdbStore', () => { + addStoreTests(() => AztecLmdbStore.open()); +}); diff --git a/yarn-project/kv-store/src/mem/array.ts b/yarn-project/kv-store/src/mem/array.ts index 3350fda7f51e..b21eae6cbe3f 100644 --- a/yarn-project/kv-store/src/mem/array.ts +++ b/yarn-project/kv-store/src/mem/array.ts @@ -57,11 +57,13 @@ export class MemAztecArray implements AztecArray { } entries(): IterableIterator<[number, T]> { - return this.db.get(this.slot)?.entries() || []; + const arr = this.db.get(this.slot) || []; + return arr.entries(); } values(): IterableIterator { - return this.db.get(this.slot)?.values() || []; + const arr = this.db.get(this.slot) || []; + return arr.values(); } [Symbol.iterator](): IterableIterator { diff --git a/yarn-project/kv-store/src/mem/counter.ts b/yarn-project/kv-store/src/mem/counter.ts index 13a71b65e9cc..6be581ca3871 100644 --- a/yarn-project/kv-store/src/mem/counter.ts +++ b/yarn-project/kv-store/src/mem/counter.ts @@ -7,7 +7,7 @@ export class MemAztecCounter implements AztecCounter { private map: MemAztecMap; constructor(name: string, db: MemDb) { - this.map = new MemAztecMap(name, db); + this.map = new MemAztecMap(name, db, false); } async set(key: Key, value: number): Promise { diff --git a/yarn-project/kv-store/src/mem/map.ts b/yarn-project/kv-store/src/mem/map.ts index c4f7fa916757..5d7a15698a26 100644 --- a/yarn-project/kv-store/src/mem/map.ts +++ b/yarn-project/kv-store/src/mem/map.ts @@ -31,7 +31,7 @@ function compareKeys(a: Key, b: Key) { * A map backed by mem. */ export class MemAztecMap implements AztecMultiMap { - constructor(private name: string, private db: MemDb) {} + constructor(private name: string, private db: MemDb, private allowDups = true) {} close(): Promise { return Promise.resolve(); @@ -39,7 +39,7 @@ export class MemAztecMap implements AztecMultiMap { get(key: Key): V | undefined { const r = this.db.get(this.slot(key)); - return r ? r[r.length - 1] : undefined; + return r ? r[0] : undefined; } getValues(key: Key): IterableIterator { @@ -54,7 +54,7 @@ export class MemAztecMap implements AztecMultiMap { set(key: Key, val: V): Promise { const r = this.db.get(this.slot(key)); - if (r) { + if (r && this.allowDups) { this.db.set(this.slot(key), [...r, val]); } else { this.db.set(this.slot(key), [val]); diff --git a/yarn-project/kv-store/src/mem/store.test.ts b/yarn-project/kv-store/src/mem/store.test.ts index 7ffd3b87a8b1..81d27da6212e 100644 --- a/yarn-project/kv-store/src/mem/store.test.ts +++ b/yarn-project/kv-store/src/mem/store.test.ts @@ -1,66 +1,6 @@ -import type { AztecArray, AztecCounter, AztecMap, AztecSingleton } from '../interfaces/index.js'; +import { addStoreTests } from '../tests/aztec_store_tests.js'; import { AztecMemStore } from './store.js'; describe('AztecMemStore', () => { - let store: AztecMemStore; - let array: AztecArray; - let multimap: AztecMap; - let counter: AztecCounter; - let singleton: AztecSingleton; - - beforeEach(async () => { - store = new AztecMemStore(); - - array = store.openArray('test-array'); - multimap = store.openMultiMap('test-multimap'); - counter = store.openCounter('test-counter'); - singleton = store.openSingleton('test-singleton'); - - await array.push(1, 2, 3); - await multimap.set('key-1', 1); - await multimap.set('key-2', 2); - await counter.set('counter-1', 3); - await singleton.set(4); - }); - - it('check initial state', () => { - expect(array.at(2)).toBe(3); - expect(multimap.get('key-2')).toBe(2); - expect(counter.get('counter-1')).toBe(3); - expect(singleton.get()).toBe(4); - }); - - it('state should update with successful tx', async () => { - await store.transaction(() => { - void array.setAt(2, 10); - void multimap.set('key-2', 20); - void counter.set('counter-1', 30); - void singleton.set(40); - }); - void multimap.set('key-2', 20); - - expect(array.at(2)).toBe(10); - expect(multimap.get('key-2')).toBe(20); - expect(counter.get('counter-1')).toBe(30); - expect(singleton.get()).toBe(40); - }); - - it('state should rollback with unsuccessful tx', async () => { - try { - await store.transaction(() => { - void array.setAt(2, 10); - void multimap.set('key-2', 20); - void counter.set('counter-1', 30); - void singleton.set(40); - throw new Error(); - }); - } catch (err) { - // swallow - } - - expect(array.at(2)).toBe(3); - expect(multimap.get('key-2')).toBe(2); - expect(counter.get('counter-1')).toBe(3); - expect(singleton.get()).toBe(4); - }); + addStoreTests(() => new AztecMemStore()); }); diff --git a/yarn-project/kv-store/src/mem/store.ts b/yarn-project/kv-store/src/mem/store.ts index c35f47a5718d..249ca89581dd 100644 --- a/yarn-project/kv-store/src/mem/store.ts +++ b/yarn-project/kv-store/src/mem/store.ts @@ -16,11 +16,11 @@ export class AztecMemStore implements AztecKVStore { private data = new MemDb(); openMap(name: string): AztecMap { - return new MemAztecMap(name, this.data) as any; + return new MemAztecMap(name, this.data, false) as any; } openMultiMap(name: string): AztecMultiMap { - return new MemAztecMap(name, this.data) as any; + return new MemAztecMap(name, this.data, true) as any; } openCounter>(name: string): AztecCounter { diff --git a/yarn-project/kv-store/src/tests/aztec_array_tests.ts b/yarn-project/kv-store/src/tests/aztec_array_tests.ts index 19fa70b1c3d0..a6b376eb68fd 100644 --- a/yarn-project/kv-store/src/tests/aztec_array_tests.ts +++ b/yarn-project/kv-store/src/tests/aztec_array_tests.ts @@ -20,6 +20,7 @@ export function addArrayTests(getArray: () => AztecArray) { expect(await arr.pop()).toEqual(2); expect(await arr.pop()).toEqual(1); expect(await arr.pop()).toEqual(undefined); + expect(Array.from(arr)).toEqual([1, 2, 3]); }); it('should be able to get values by index', async () => { diff --git a/yarn-project/kv-store/src/tests/aztec_map_tests.ts b/yarn-project/kv-store/src/tests/aztec_map_tests.ts index 1c5efe071999..9ce9ac2ba964 100644 --- a/yarn-project/kv-store/src/tests/aztec_map_tests.ts +++ b/yarn-project/kv-store/src/tests/aztec_map_tests.ts @@ -20,14 +20,6 @@ export function addMapTests(get: () => AztecMultiMap) { expect(map.get('quux')).toEqual(undefined); }); - it('should be able to update values', async () => { - await map.set('foo', 'bar'); - expect(map.get('foo')).toEqual('bar'); - - await map.set('foo', 'qux'); - expect(map.get('foo')).toEqual('qux'); - }); - it('should be able to set values if they do not exist', async () => { expect(await map.setIfNotExists('foo', 'bar')).toEqual(true); expect(await map.setIfNotExists('foo', 'baz')).toEqual(false); @@ -73,6 +65,8 @@ export function addMapTests(get: () => AztecMultiMap) { await map.set('foo', 'bar'); await map.set('foo', 'baz'); + expect(map.get('foo')).toEqual('bar'); + expect([...map.getValues('foo')]).toEqual(['bar', 'baz']); }); diff --git a/yarn-project/kv-store/src/tests/aztec_store_tests.ts b/yarn-project/kv-store/src/tests/aztec_store_tests.ts new file mode 100644 index 000000000000..859a76528376 --- /dev/null +++ b/yarn-project/kv-store/src/tests/aztec_store_tests.ts @@ -0,0 +1,71 @@ +import { beforeEach, describe, expect, it } from '@jest/globals'; + +import type { AztecArray, AztecCounter, AztecKVStore, AztecMultiMap, AztecSingleton } from '../interfaces/index.js'; + +export function addStoreTests(get: () => AztecKVStore) { + describe('AztecStore', () => { + let store: AztecKVStore; + let array: AztecArray; + let multimap: AztecMultiMap; + let counter: AztecCounter; + let singleton: AztecSingleton; + + beforeEach(async () => { + store = get(); + + array = store.openArray('test-array'); + multimap = store.openMultiMap('test-multimap'); + counter = store.openCounter('test-counter'); + singleton = store.openSingleton('test-singleton'); + + await array.push(1, 2, 3); + await multimap.set('key-1', 1); + await multimap.set('key-2', 2); + await counter.set('counter-1', 3); + await singleton.set(4); + }); + + it('check initial state', () => { + expect(array.at(2)).toBe(3); + expect(multimap.get('key-2')).toBe(2); + expect([...multimap.getValues('key-2')]).toEqual([2]); + expect(counter.get('counter-1')).toBe(3); + expect(singleton.get()).toBe(4); + }); + + it('state should update with successful tx', async () => { + await store.transaction(() => { + void array.setAt(2, 10); + void multimap.set('key-2', 20); + void counter.set('counter-1', 30); + void singleton.set(40); + }); + + expect(array.at(2)).toBe(10); + expect(multimap.get('key-2')).toBe(2); + expect([...multimap.getValues('key-2')]).toEqual([2, 20]); + expect(counter.get('counter-1')).toBe(30); + expect(singleton.get()).toBe(40); + }); + + it.skip('state should rollback with unsuccessful tx', async () => { + try { + await store.transaction(() => { + void array.setAt(2, 10); + void multimap.set('key-2', 20); + void counter.set('counter-1', 30); + void singleton.set(40); + throw new Error(); + }); + } catch (err) { + // swallow + } + + expect(array.at(2)).toBe(3); + expect(multimap.get('key-2')).toBe(2); + expect([...multimap.getValues('key-2')]).toEqual([2]); + expect(counter.get('counter-1')).toBe(3); + expect(singleton.get()).toBe(4); + }); + }); +} diff --git a/yarn-project/kv-store/src/utils.ts b/yarn-project/kv-store/src/utils.ts index 535445ae8c90..5c7e1691b0f5 100644 --- a/yarn-project/kv-store/src/utils.ts +++ b/yarn-project/kv-store/src/utils.ts @@ -39,5 +39,4 @@ export async function initStoreForRollup( */ export function openTmpStore(ephemeral: boolean = false): AztecKVStore { return new AztecMemStore(); - // return AztecLmdbStore.open(undefined, ephemeral); } diff --git a/yarn-project/pxe/src/pxe_service/create_pxe_service.ts b/yarn-project/pxe/src/pxe_service/create_pxe_service.ts index 7daead96c667..ca42986bc2c2 100644 --- a/yarn-project/pxe/src/pxe_service/create_pxe_service.ts +++ b/yarn-project/pxe/src/pxe_service/create_pxe_service.ts @@ -3,6 +3,7 @@ import { Grumpkin } from '@aztec/circuits.js/barretenberg'; import { randomBytes } from '@aztec/foundation/crypto'; import { TestKeyStore } from '@aztec/key-store'; import { AztecLmdbStore } from '@aztec/kv-store/lmdb'; +import { AztecMemStore } from '@aztec/kv-store/mem'; import { initStoreForRollup } from '@aztec/kv-store/utils'; import { getCanonicalClassRegisterer } from '@aztec/protocol-contracts/class-registerer'; import { getCanonicalGasToken } from '@aztec/protocol-contracts/gas-token'; @@ -37,11 +38,11 @@ export async function createPXEService( const keyStorePath = config.dataDirectory ? join(config.dataDirectory, 'pxe_key_store') : undefined; const l1Contracts = await aztecNode.getL1ContractAddresses(); - const keyStore = new TestKeyStore( - new Grumpkin(), - await initStoreForRollup(AztecLmdbStore.open(keyStorePath), l1Contracts.rollupAddress), - ); - const db = new KVPxeDatabase(await initStoreForRollup(AztecLmdbStore.open(pxeDbPath), l1Contracts.rollupAddress)); + const keyStoreDb = keyStorePath ? AztecLmdbStore.open(keyStorePath) : new AztecMemStore(); + const pxeDb = pxeDbPath ? AztecLmdbStore.open(pxeDbPath) : new AztecMemStore(); + + const keyStore = new TestKeyStore(new Grumpkin(), await initStoreForRollup(keyStoreDb, l1Contracts.rollupAddress)); + const db = new KVPxeDatabase(await initStoreForRollup(pxeDb, l1Contracts.rollupAddress)); const server = new PXEService(keyStore, aztecNode, db, config, logSuffix); for (const contract of [ From 957bda714e9d8a18961c31934b7b44a4ef1243c5 Mon Sep 17 00:00:00 2001 From: Charlie Lye Date: Fri, 5 Apr 2024 13:06:27 +0000 Subject: [PATCH 3/3] fix --- yarn-project/kv-store/src/mem/counter.ts | 7 ++----- yarn-project/kv-store/src/mem/map.ts | 13 ++++++------- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/yarn-project/kv-store/src/mem/counter.ts b/yarn-project/kv-store/src/mem/counter.ts index 6be581ca3871..9a95831ad21b 100644 --- a/yarn-project/kv-store/src/mem/counter.ts +++ b/yarn-project/kv-store/src/mem/counter.ts @@ -10,16 +10,15 @@ export class MemAztecCounter implements AztecCounter { this.map = new MemAztecMap(name, db, false); } - async set(key: Key, value: number): Promise { + async set(key: Key, value: number): Promise { if (value) { return this.map.set(key, value); } else { await this.map.delete(key); - return true; } } - async update(key: Key, delta = 1): Promise { + async update(key: Key, delta = 1): Promise { const current = this.map.get(key) ?? 0; const next = current + delta; @@ -32,8 +31,6 @@ export class MemAztecCounter implements AztecCounter { if (next > 0) { await this.map.set(key, next); } - - return true; } get(key: Key): number { diff --git a/yarn-project/kv-store/src/mem/map.ts b/yarn-project/kv-store/src/mem/map.ts index 5d7a15698a26..d74f50b5bc5b 100644 --- a/yarn-project/kv-store/src/mem/map.ts +++ b/yarn-project/kv-store/src/mem/map.ts @@ -52,21 +52,21 @@ export class MemAztecMap implements AztecMultiMap { return r ? r.length > 0 : false; } - set(key: Key, val: V): Promise { + set(key: Key, val: V): Promise { const r = this.db.get(this.slot(key)); if (r && this.allowDups) { this.db.set(this.slot(key), [...r, val]); } else { this.db.set(this.slot(key), [val]); } - return Promise.resolve(true); + return Promise.resolve(); } - swap(key: Key, fn: (val: V | undefined) => V): Promise { + swap(key: Key, fn: (val: V | undefined) => V): Promise { const entry = this.get(key); const newValue = fn(entry); this.db.set(this.slot(key), [newValue]); - return Promise.resolve(true); + return Promise.resolve(); } async setIfNotExists(key: Key, val: V): Promise { @@ -78,13 +78,12 @@ export class MemAztecMap implements AztecMultiMap { return false; } - delete(key: Key): Promise { + delete(key: Key): Promise { const r = this.db.get(this.slot(key)); if (r?.length) { this.db.set(this.slot(key), []); - return Promise.resolve(true); } - return Promise.resolve(false); + return Promise.resolve(); } deleteValue(key: Key, val: V): Promise {