1
0
Fork 0
mirror of https://github.com/denoland/deno.git synced 2024-12-25 00:29:09 -05:00
denoland-deno/ext/kv/01_db.ts
2023-07-01 09:24:15 +02:00

660 lines
16 KiB
TypeScript

// Copyright 2018-2023 the Deno authors. All rights reserved. MIT license.
// @ts-ignore internal api
const {
AsyncGeneratorPrototype,
BigIntPrototypeToString,
ObjectFreeze,
ObjectGetPrototypeOf,
ObjectPrototypeIsPrototypeOf,
StringPrototypeReplace,
SymbolFor,
SymbolToStringTag,
Uint8ArrayPrototype,
} = globalThis.__bootstrap.primordials;
const core = Deno.core;
const ops = core.ops;
const encodeCursor: (
selector: [Deno.KvKey | null, Deno.KvKey | null, Deno.KvKey | null],
boundaryKey: Deno.KvKey,
) => string = (selector, boundaryKey) =>
ops.op_kv_encode_cursor(selector, boundaryKey);
async function openKv(path: string) {
const rid = await core.opAsync("op_kv_database_open", path);
return new Kv(rid, kvSymbol);
}
const millisecondsInOneWeek = 7 * 24 * 60 * 60 * 1000;
function validateQueueDelay(delay: number) {
if (delay < 0) {
throw new TypeError("delay cannot be negative");
}
if (delay > millisecondsInOneWeek) {
throw new TypeError("delay cannot be greater than one week");
}
if (isNaN(delay)) {
throw new TypeError("delay cannot be NaN");
}
}
interface RawKvEntry {
key: Deno.KvKey;
value: RawValue;
versionstamp: string;
}
type RawValue = {
kind: "v8";
value: Uint8Array;
} | {
kind: "bytes";
value: Uint8Array;
} | {
kind: "u64";
value: bigint;
};
const kvSymbol = Symbol("KvRid");
class Kv {
#rid: number;
#closed: boolean;
constructor(rid: number = undefined, symbol: symbol = undefined) {
if (kvSymbol !== symbol) {
throw new TypeError(
"Deno.Kv can not be constructed, use Deno.openKv instead.",
);
}
this.#rid = rid;
this.#closed = false;
}
atomic() {
return new AtomicOperation(this.#rid);
}
async get(key: Deno.KvKey, opts?: { consistency?: Deno.KvConsistencyLevel }) {
const [entries]: [RawKvEntry[]] = await core.opAsync(
"op_kv_snapshot_read",
this.#rid,
[[
null,
key,
null,
1,
false,
null,
]],
opts?.consistency ?? "strong",
);
if (!entries.length) {
return {
key,
value: null,
versionstamp: null,
};
}
return deserializeValue(entries[0]);
}
async getMany(
keys: Deno.KvKey[],
opts?: { consistency?: Deno.KvConsistencyLevel },
): Promise<Deno.KvEntry<unknown>[]> {
const ranges: RawKvEntry[][] = await core.opAsync(
"op_kv_snapshot_read",
this.#rid,
keys.map((key) => [
null,
key,
null,
1,
false,
null,
]),
opts?.consistency ?? "strong",
);
return ranges.map((entries, i) => {
if (!entries.length) {
return {
key: keys[i],
value: null,
versionstamp: null,
};
}
return deserializeValue(entries[0]);
});
}
async set(key: Deno.KvKey, value: unknown) {
value = serializeValue(value);
const checks: Deno.AtomicCheck[] = [];
const mutations = [
[key, "set", value],
];
const versionstamp = await core.opAsync(
"op_kv_atomic_write",
this.#rid,
checks,
mutations,
[],
);
if (versionstamp === null) throw new TypeError("Failed to set value");
return { ok: true, versionstamp };
}
async delete(key: Deno.KvKey) {
const checks: Deno.AtomicCheck[] = [];
const mutations = [
[key, "delete", null],
];
const result = await core.opAsync(
"op_kv_atomic_write",
this.#rid,
checks,
mutations,
[],
);
if (!result) throw new TypeError("Failed to set value");
}
list(
selector: Deno.KvListSelector,
options: {
limit?: number;
batchSize?: number;
cursor?: string;
reverse?: boolean;
consistency?: Deno.KvConsistencyLevel;
} = {},
): KvListIterator {
if (options.limit !== undefined && options.limit <= 0) {
throw new Error("limit must be positive");
}
let batchSize = options.batchSize ?? (options.limit ?? 100);
if (batchSize <= 0) throw new Error("batchSize must be positive");
if (options.batchSize === undefined && batchSize > 500) batchSize = 500;
return new KvListIterator({
limit: options.limit,
selector,
cursor: options.cursor,
reverse: options.reverse ?? false,
consistency: options.consistency ?? "strong",
batchSize,
pullBatch: this.#pullBatch(batchSize),
});
}
#pullBatch(batchSize: number): (
selector: Deno.KvListSelector,
cursor: string | undefined,
reverse: boolean,
consistency: Deno.KvConsistencyLevel,
) => Promise<Deno.KvEntry<unknown>[]> {
return async (selector, cursor, reverse, consistency) => {
const [entries]: [RawKvEntry[]] = await core.opAsync(
"op_kv_snapshot_read",
this.#rid,
[[
"prefix" in selector ? selector.prefix : null,
"start" in selector ? selector.start : null,
"end" in selector ? selector.end : null,
batchSize,
reverse,
cursor,
]],
consistency,
);
return entries.map(deserializeValue);
};
}
async enqueue(
message: unknown,
opts?: { delay?: number; keysIfUndelivered?: Deno.KvKey[] },
) {
if (opts?.delay !== undefined) {
validateQueueDelay(opts?.delay);
}
const enqueues = [
[
core.serialize(message, { forStorage: true }),
opts?.delay ?? 0,
opts?.keysIfUndelivered ?? [],
null,
],
];
const versionstamp = await core.opAsync(
"op_kv_atomic_write",
this.#rid,
[],
[],
enqueues,
);
if (versionstamp === null) throw new TypeError("Failed to enqueue value");
return { ok: true, versionstamp };
}
async listenQueue(
handler: (message: unknown) => Promise<void> | void,
): Promise<void> {
const finishMessageOps = new Map<number, Promise<void>>();
while (!this.#closed) {
// Wait for the next message.
let next: { 0: Uint8Array; 1: number };
try {
next = await core.opAsync(
"op_kv_dequeue_next_message",
this.#rid,
);
} catch (error) {
if (this.#closed) {
break;
} else {
throw error;
}
}
// Deserialize the payload.
const { 0: payload, 1: handleId } = next;
const deserializedPayload = core.deserialize(payload, {
forStorage: true,
});
// Dispatch the payload.
(async () => {
let success = false;
try {
const result = handler(deserializedPayload);
const _res = result instanceof Promise ? (await result) : result;
success = true;
} catch (error) {
console.error("Exception in queue handler", error);
} finally {
if (this.#closed) {
core.close(handleId);
} else {
const promise: Promise<void> = core.opAsync(
"op_kv_finish_dequeued_message",
handleId,
success,
);
finishMessageOps.set(handleId, promise);
try {
await promise;
} finally {
finishMessageOps.delete(handleId);
}
}
}
})();
}
for (const promise of finishMessageOps.values()) {
await promise;
}
finishMessageOps.clear();
}
close() {
core.close(this.#rid);
this.#closed = true;
}
}
class AtomicOperation {
#rid: number;
#checks: [Deno.KvKey, string | null][] = [];
#mutations: [Deno.KvKey, string, RawValue | null][] = [];
#enqueues: [Uint8Array, number, Deno.KvKey[], number[] | null][] = [];
constructor(rid: number) {
this.#rid = rid;
}
check(...checks: Deno.AtomicCheck[]): this {
for (const check of checks) {
this.#checks.push([check.key, check.versionstamp]);
}
return this;
}
mutate(...mutations: Deno.KvMutation[]): this {
for (const mutation of mutations) {
const key = mutation.key;
let type: string;
let value: RawValue | null;
switch (mutation.type) {
case "delete":
type = "delete";
if (mutation.value) {
throw new TypeError("invalid mutation 'delete' with value");
}
break;
case "set":
case "sum":
case "min":
case "max":
type = mutation.type;
if (!("value" in mutation)) {
throw new TypeError(`invalid mutation '${type}' without value`);
}
value = serializeValue(mutation.value);
break;
default:
throw new TypeError("Invalid mutation type");
}
this.#mutations.push([key, type, value]);
}
return this;
}
sum(key: Deno.KvKey, n: bigint): this {
this.#mutations.push([key, "sum", serializeValue(new KvU64(n))]);
return this;
}
min(key: Deno.KvKey, n: bigint): this {
this.#mutations.push([key, "min", serializeValue(new KvU64(n))]);
return this;
}
max(key: Deno.KvKey, n: bigint): this {
this.#mutations.push([key, "max", serializeValue(new KvU64(n))]);
return this;
}
set(key: Deno.KvKey, value: unknown): this {
this.#mutations.push([key, "set", serializeValue(value)]);
return this;
}
delete(key: Deno.KvKey): this {
this.#mutations.push([key, "delete", null]);
return this;
}
enqueue(
message: unknown,
opts?: { delay?: number; keysIfUndelivered?: Deno.KvKey[] },
): this {
if (opts?.delay !== undefined) {
validateQueueDelay(opts?.delay);
}
this.#enqueues.push([
core.serialize(message, { forStorage: true }),
opts?.delay ?? 0,
opts?.keysIfUndelivered ?? [],
null,
]);
return this;
}
async commit(): Promise<Deno.KvCommitResult | Deno.KvCommitError> {
const versionstamp = await core.opAsync(
"op_kv_atomic_write",
this.#rid,
this.#checks,
this.#mutations,
this.#enqueues,
);
if (versionstamp === null) return { ok: false };
return { ok: true, versionstamp };
}
then() {
throw new TypeError(
"`Deno.AtomicOperation` is not a promise. Did you forget to call `commit()`?",
);
}
}
const MIN_U64 = BigInt("0");
const MAX_U64 = BigInt("0xffffffffffffffff");
class KvU64 {
value: bigint;
constructor(value: bigint) {
if (typeof value !== "bigint") {
throw new TypeError("value must be a bigint");
}
if (value < MIN_U64) {
throw new RangeError("value must be a positive bigint");
}
if (value > MAX_U64) {
throw new RangeError("value must fit in a 64-bit unsigned integer");
}
this.value = value;
Object.freeze(this);
}
valueOf() {
return this.value;
}
toString() {
return BigIntPrototypeToString(this.value);
}
get [SymbolToStringTag]() {
return "Deno.KvU64";
}
[SymbolFor("Deno.privateCustomInspect")](inspect, inspectOptions) {
return StringPrototypeReplace(
inspect(Object(this.value), inspectOptions),
"BigInt",
"Deno.KvU64",
);
}
}
function deserializeValue(entry: RawKvEntry): Deno.KvEntry<unknown> {
const { kind, value } = entry.value;
switch (kind) {
case "v8":
return {
...entry,
value: core.deserialize(value, { forStorage: true }),
};
case "bytes":
return {
...entry,
value,
};
case "u64":
return {
...entry,
value: new KvU64(value),
};
default:
throw new TypeError("Invalid value type");
}
}
function serializeValue(value: unknown): RawValue {
if (ObjectPrototypeIsPrototypeOf(Uint8ArrayPrototype, value)) {
return {
kind: "bytes",
value,
};
} else if (ObjectPrototypeIsPrototypeOf(KvU64.prototype, value)) {
return {
kind: "u64",
value: value.valueOf(),
};
} else {
return {
kind: "v8",
value: core.serialize(value, { forStorage: true }),
};
}
}
// This gets the %AsyncIteratorPrototype% object (which exists but is not a
// global). We extend the KvListIterator iterator from, so that we immediately
// support async iterator helpers once they land. The %AsyncIterator% does not
// yet actually exist however, so right now the AsyncIterator binding refers to
// %Object%. I know.
// Once AsyncIterator is a global, we can just use it (from primordials), rather
// than doing this here.
const AsyncIteratorPrototype = ObjectGetPrototypeOf(AsyncGeneratorPrototype);
const AsyncIterator = AsyncIteratorPrototype.constructor;
class KvListIterator extends AsyncIterator
implements AsyncIterator<Deno.KvEntry<unknown>> {
#selector: Deno.KvListSelector;
#entries: Deno.KvEntry<unknown>[] | null = null;
#cursorGen: (() => string) | null = null;
#done = false;
#lastBatch = false;
#pullBatch: (
selector: Deno.KvListSelector,
cursor: string | undefined,
reverse: boolean,
consistency: Deno.KvConsistencyLevel,
) => Promise<Deno.KvEntry<unknown>[]>;
#limit: number | undefined;
#count = 0;
#reverse: boolean;
#batchSize: number;
#consistency: Deno.KvConsistencyLevel;
constructor(
{ limit, selector, cursor, reverse, consistency, batchSize, pullBatch }: {
limit?: number;
selector: Deno.KvListSelector;
cursor?: string;
reverse: boolean;
batchSize: number;
consistency: Deno.KvConsistencyLevel;
pullBatch: (
selector: Deno.KvListSelector,
cursor: string | undefined,
reverse: boolean,
consistency: Deno.KvConsistencyLevel,
) => Promise<Deno.KvEntry<unknown>[]>;
},
) {
super();
let prefix: Deno.KvKey | undefined;
let start: Deno.KvKey | undefined;
let end: Deno.KvKey | undefined;
if ("prefix" in selector && selector.prefix !== undefined) {
prefix = ObjectFreeze([...selector.prefix]);
}
if ("start" in selector && selector.start !== undefined) {
start = ObjectFreeze([...selector.start]);
}
if ("end" in selector && selector.end !== undefined) {
end = ObjectFreeze([...selector.end]);
}
if (prefix) {
if (start && end) {
throw new TypeError(
"Selector can not specify both 'start' and 'end' key when specifying 'prefix'.",
);
}
if (start) {
this.#selector = { prefix, start };
} else if (end) {
this.#selector = { prefix, end };
} else {
this.#selector = { prefix };
}
} else {
if (start && end) {
this.#selector = { start, end };
} else {
throw new TypeError(
"Selector must specify either 'prefix' or both 'start' and 'end' key.",
);
}
}
Object.freeze(this.#selector);
this.#pullBatch = pullBatch;
this.#limit = limit;
this.#reverse = reverse;
this.#consistency = consistency;
this.#batchSize = batchSize;
this.#cursorGen = cursor ? () => cursor : null;
}
get cursor(): string {
if (this.#cursorGen === null) {
throw new Error("Cannot get cursor before first iteration");
}
return this.#cursorGen();
}
async next(): Promise<IteratorResult<Deno.KvEntry<unknown>>> {
// Fused or limit exceeded
if (
this.#done ||
(this.#limit !== undefined && this.#count >= this.#limit)
) {
return { done: true, value: undefined };
}
// Attempt to fill the buffer
if (!this.#entries?.length && !this.#lastBatch) {
const batch = await this.#pullBatch(
this.#selector,
this.#cursorGen ? this.#cursorGen() : undefined,
this.#reverse,
this.#consistency,
);
// Reverse the batch so we can pop from the end
batch.reverse();
this.#entries = batch;
// Last batch, do not attempt to pull more
if (batch.length < this.#batchSize) {
this.#lastBatch = true;
}
}
const entry = this.#entries?.pop();
if (!entry) {
this.#done = true;
this.#cursorGen = () => "";
return { done: true, value: undefined };
}
this.#cursorGen = () => {
const selector = this.#selector;
return encodeCursor([
"prefix" in selector ? selector.prefix : null,
"start" in selector ? selector.start : null,
"end" in selector ? selector.end : null,
], entry.key);
};
this.#count++;
return {
done: false,
value: entry,
};
}
[Symbol.asyncIterator](): AsyncIterator<Deno.KvEntry<unknown>> {
return this;
}
}
export { AtomicOperation, Kv, KvListIterator, KvU64, openKv };