packages/core: add transaction flag to ExecutionContext methods

canvasxyz · Feb 6, 2025 · 934b96b · 934b96b
1 parent 11bcf49
commit 934b96b
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 49 deletions.
diff --git a/packages/core/src/ExecutionContext.ts b/packages/core/src/ExecutionContext.ts
@@ -48,7 +48,11 @@ export class ExecutionContext {
 		return await this.messageLog.isAncestor(this.root, ancestor)
 	}
 
-	public async getModelValue<T extends ModelValue = ModelValue>(model: string, key: string): Promise<null | T> {
+	public async getModelValue<T extends ModelValue = ModelValue>(
+		model: string,
+		key: string,
+		transaction: boolean,
+	): Promise<null | T> {
 		if (this.modelEntries[model] === undefined) {
 			const { name } = this.message.payload
 			throw new Error(`could not access model db.${model} inside runtime action ${name}`)
@@ -94,7 +98,7 @@ export class ExecutionContext {
 		}
 	}
 
-	public setModelValue(model: string, value: ModelValue): void {
+	public setModelValue(model: string, value: ModelValue, transaction: boolean): void {
 		assert(this.db.models[model] !== undefined, "model not found")
 		validateModelValue(this.db.models[model], value)
 		const {
@@ -105,30 +109,38 @@ export class ExecutionContext {
 		this.modelEntries[model][key] = value
 	}
 
-	public deleteModelValue(model: string, key: string): void {
+	public deleteModelValue(model: string, key: string, transaction: boolean): void {
 		assert(this.db.models[model] !== undefined, "model not found")
 		this.modelEntries[model][key] = null
 	}
 
-	public async updateModelValue(model: string, value: Record<string, PropertyValue | undefined>): Promise<void> {
+	public async updateModelValue(
+		model: string,
+		value: Record<string, PropertyValue | undefined>,
+		transaction: boolean,
+	): Promise<void> {
 		assert(this.db.models[model] !== undefined, "model not found")
 		const {
 			primaryKey: [primaryKey],
 		} = this.db.models[model]
 		const key = value[primaryKey] as string
-		const previousValue = await this.getModelValue(model, key)
+		const previousValue = await this.getModelValue(model, key, transaction)
 		const result = updateModelValues(value, previousValue)
 		validateModelValue(this.db.models[model], result)
 		this.modelEntries[model][key] = result
 	}
 
-	public async mergeModelValue(model: string, value: Record<string, PropertyValue | undefined>): Promise<void> {
+	public async mergeModelValue(
+		model: string,
+		value: Record<string, PropertyValue | undefined>,
+		transaction: boolean,
+	): Promise<void> {
 		assert(this.db.models[model] !== undefined, "model not found")
 		const {
 			primaryKey: [primaryKey],
 		} = this.db.models[model]
 		const key = value[primaryKey] as string
-		const previousValue = await this.getModelValue(model, key)
+		const previousValue = await this.getModelValue(model, key, transaction)
 		const result = mergeModelValues(value, previousValue)
 		validateModelValue(this.db.models[model], result)
 		this.modelEntries[model][key] = result

diff --git a/packages/core/src/runtime/ContractRuntime.ts b/packages/core/src/runtime/ContractRuntime.ts
@@ -71,6 +71,7 @@ export class ContractRuntime extends AbstractRuntime {
 
 	#context: ExecutionContext | null = null
 	#thisHandle: QuickJSHandle | null = null
+	#transaction: boolean = false
 
 	constructor(
 		public readonly topic: string,
@@ -85,55 +86,42 @@ export class ContractRuntime extends AbstractRuntime {
 				get: vm.wrapFunction((model, key) => {
 					assert(typeof model === "string", 'expected typeof model === "string"')
 					assert(typeof key === "string", 'expected typeof key === "string"')
-					return this.context.getModelValue(model, key)
+					return this.context.getModelValue(model, key, this.#transaction)
 				}),
-				set: vm.context.newFunction("set", (modelHandle, valueHandle) => {
-					const model = vm.context.getString(modelHandle)
-					const value = this.vm.unwrapValue(valueHandle) as ModelValue
-					this.context.setModelValue(model, value)
+
+				set: vm.wrapFunction((model, value) => {
+					assert(typeof model === "string", 'expected typeof model === "string"')
+					this.context.setModelValue(model, value as ModelValue, this.#transaction)
 				}),
-				update: vm.context.newFunction("update", (modelHandle, valueHandle) => {
-					const model = vm.context.getString(modelHandle)
-					const value = this.vm.unwrapValue(valueHandle) as ModelValue
 
-					const promise = vm.context.newPromise()
+				update: vm.wrapFunction(async (model, value) => {
+					assert(typeof model === "string", 'expected typeof model === "string"')
+					await this.context.updateModelValue(model, value as ModelValue, this.#transaction)
+				}),
 
-					// TODO: Ensure concurrent merges into the same value don't create a race condition
-					// if the user doesn't call db.update() with await.
-					this.context
-						.updateModelValue(model, value)
-						.then(() => promise.resolve())
-						.catch((err) => promise.reject())
+				merge: vm.wrapFunction(async (model, value) => {
+					assert(typeof model === "string", 'expected typeof model === "string"')
+					await this.context.mergeModelValue(model, value as ModelValue, this.#transaction)
+				}),
 
-					promise.settled.then(vm.runtime.executePendingJobs)
-					return promise.handle
+				delete: vm.wrapFunction((model, key) => {
+					assert(typeof model === "string", 'expected typeof model === "string"')
+					assert(typeof key === "string", 'expected typeof key === "string"')
+					this.context.deleteModelValue(model, key, this.#transaction)
 				}),
-				merge: vm.context.newFunction("merge", (modelHandle, valueHandle) => {
-					const model = vm.context.getString(modelHandle)
-					const value = this.vm.unwrapValue(valueHandle) as ModelValue
 
+				transaction: vm.context.newFunction("transaction", (callbackHandle) => {
 					const promise = vm.context.newPromise()
 
-					// TODO: Ensure concurrent merges into the same value don't create a race condition
-					// if the user doesn't call db.update() with await.
-					this.context
-						.mergeModelValue(model, value)
-						.then(() => promise.resolve())
-						.catch((err) => promise.reject())
+					this.#transaction = true
+					this.vm
+						.callAsync(callbackHandle, this.thisHandle, [])
+						.then(promise.resolve, promise.reject)
+						.finally(() => void (this.#transaction = false))
 
 					promise.settled.then(vm.runtime.executePendingJobs)
 					return promise.handle
 				}),
-
-				delete: vm.context.newFunction("delete", (modelHandle, keyHandle) => {
-					const model = vm.context.getString(modelHandle)
-					const key = vm.context.getString(keyHandle)
-					this.context.deleteModelValue(model, key)
-				}),
-
-				transaction: vm.context.newFunction("transaction", (callbackHandle) => {
-					this.vm.call(callbackHandle, this.thisHandle, [])
-				}),
 			})
 			.consume(vm.cache)
 	}
@@ -180,6 +168,8 @@ export class ContractRuntime extends AbstractRuntime {
 			timestamp,
 		})
 
+		this.vm.context.setProp(thisHandle, "db", this.#databaseAPI)
+
 		const argHandles = Array.isArray(args) ? args.map(this.vm.wrapValue) : [this.vm.wrapValue(args)]
 
 		try {

diff --git a/packages/core/src/runtime/FunctionRuntime.ts b/packages/core/src/runtime/FunctionRuntime.ts
@@ -26,6 +26,7 @@ export class FunctionRuntime<ModelsT extends ModelSchema> extends AbstractRuntim
 	}
 
 	#context: ExecutionContext | null = null
+	#transaction: boolean = false
 	#thisValue: ActionContext<DeriveModelTypes<ModelsT>> | null = null
 	#queue = new PQueue({ concurrency: 1 })
 	#db: ModelAPI<DeriveModelTypes<ModelsT>>
@@ -40,17 +41,24 @@ export class FunctionRuntime<ModelsT extends ModelSchema> extends AbstractRuntim
 
 		this.#db = {
 			get: async <T extends keyof DeriveModelTypes<ModelsT> & string>(model: T, key: string) => {
-				const result = await this.#queue.add(() => this.context.getModelValue<DeriveModelTypes<ModelsT>[T]>(model, key))
+				const result = await this.#queue.add(() =>
+					this.context.getModelValue<DeriveModelTypes<ModelsT>[T]>(model, key, this.#transaction),
+				)
 				return result ?? null
 			},
 
-			set: (model, value) => this.#queue.add(() => this.context.setModelValue(model, value)),
-			update: (model, value) => this.#queue.add(() => this.context.updateModelValue(model, value)),
-			merge: (model, value) => this.#queue.add(() => this.context.mergeModelValue(model, value)),
-			delete: (model, key) => this.#queue.add(() => this.context.deleteModelValue(model, key)),
+			set: (model, value) => this.#queue.add(() => this.context.setModelValue(model, value, this.#transaction)),
+			update: (model, value) => this.#queue.add(() => this.context.updateModelValue(model, value, this.#transaction)),
+			merge: (model, value) => this.#queue.add(() => this.context.mergeModelValue(model, value, this.#transaction)),
+			delete: (model, key) => this.#queue.add(() => this.context.deleteModelValue(model, key, this.#transaction)),
 
 			transaction: async (callback) => {
-				await callback.apply(this.thisValue, [])
+				try {
+					this.#transaction = true
+					await callback.apply(this.thisValue, [])
+				} finally {
+					this.#transaction = false
+				}
 			},
 		}
 	}