implement indexes in store to keep whole history

bot.ts

@@ -1,8 +1,10 @@
-import { autoQuote, bold, Bot, Context, hydrateReply, ParseModeFlavor } from "./deps.ts";
+import { autoQuote, bold, Bot, Context, hydrateReply, log, ParseModeFlavor } from "./deps.ts";
 import { fmt } from "./intl.ts";
 import { getAllJobs, pushJob } from "./queue.ts";
 import { mySession, MySessionFlavor } from "./session.ts";
 
+const logger = () => log.getLogger();
+
 export type MyContext = ParseModeFlavor<Context> & MySessionFlavor;
 export const bot = new Bot<MyContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
 bot.use(autoQuote);
@@ -39,7 +41,8 @@ bot.use(async (ctx, next) => {
 
 bot.api.setMyShortDescription("I can generate furry images from text");
 bot.api.setMyDescription(
-  "I can generate furry images from text. Send /txt2img to generate an image.",
+  "I can generate furry images from text. " +
+    "Send /txt2img to generate an image.",
 );
 bot.api.setMyCommands([
   { command: "txt2img", description: "Generate an image" },
@@ -72,18 +75,16 @@ bot.command("txt2img", async (ctx) => {
   if (!ctx.match) {
     return ctx.reply("Please describe what you want to see after the command");
   }
-  pushJob({
+  const statusMessage = await ctx.reply("Accepted. You are now in queue.");
+  await pushJob({
     params: { prompt: ctx.match },
     user: ctx.from,
     chat: ctx.chat,
     requestMessage: ctx.message,
+    statusMessage,
     status: { type: "idle" },
   });
-  console.log(
-    `Enqueued job ${jobs.length + 1} for ${ctx.from.first_name} in ${ctx.chat.type} chat:`,
-    ctx.match.replace(/\s+/g, " "),
-    "\n",
-  );
+  logger().info("Job enqueued", ctx.from.first_name, ctx.chat.type, ctx.match.replace(/\s+/g, " "));
 });
 
 bot.command("queue", async (ctx) => {
@@ -268,5 +269,5 @@ bot.catch((err) => {
     msg += ` in ${chat.title}`;
     if (chat.type === "supergroup" && chat.username) msg += ` (@${chat.username})`;
   }
-  console.error(msg, err.error);
+  logger().error("handling update failed", from?.first_name, chat?.type, err);
 });

deps.ts

@@ -4,3 +4,4 @@ export * from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
 export * from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
 export * as types from "https://deno.land/x/grammy_types@v3.2.0/mod.ts";
 export * from "https://deno.land/x/ulid@v0.3.0/mod.ts";
+export * as log from "https://deno.land/std@0.201.0/log/mod.ts";

main.ts

@@ -1,6 +1,21 @@
 import "https://deno.land/std@0.201.0/dotenv/load.ts";
 import { bot } from "./bot.ts";
 import { processQueue, returnHangedJobs } from "./queue.ts";
+import { log } from "./deps.ts";
+
+log.setup({
+  handlers: {
+    console: new log.handlers.ConsoleHandler("INFO", {
+      formatter: (record) =>
+        `[${record.levelName}] ${record.msg} ${
+          record.args.map((arg) => JSON.stringify(arg)).join(" ")
+        } (${record.datetime.toISOString()})`,
+    }),
+  },
+  loggers: {
+    default: { level: "INFO", handlers: ["console"] },
+  },
+});
 
 await Promise.all([
   bot.start(),

queue.ts

@@ -1,30 +1,39 @@
-import { InputFile, InputMediaBuilder, types } from "./deps.ts";
+import { InputFile, InputMediaBuilder, log, types } from "./deps.ts";
 import { bot } from "./bot.ts";
 import { getGlobalSession } from "./session.ts";
 import { formatOrdinal } from "./intl.ts";
-import { SdRequest, txt2img } from "./sd.ts";
+import { SdTxt2ImgRequest, SdTxt2ImgResponse, txt2img } from "./sd.ts";
 import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";
-import { Model, Store } from "./store.ts";
+import { Model, Schema, Store } from "./store.ts";
+
+const logger = () => log.getLogger();
 
 interface Job {
-  params: Partial<SdRequest>;
+  params: Partial<SdTxt2ImgRequest>;
   user: types.User;
   chat: types.Chat.PrivateChat | types.Chat.GroupChat | types.Chat.SupergroupChat;
   requestMessage: types.Message & types.Message.TextMessage;
   statusMessage?: types.Message & types.Message.TextMessage;
-  status: { type: "idle" } | { type: "processing"; progress: number; updatedDate: Date };
+  status:
+    | { type: "idle" }
+    | { type: "processing"; progress: number; updatedDate: Date };
 }
 
 const db = await Deno.openKv("./app.db");
 
-const jobStore = new Store<Job>(db, "job");
+const jobStore = new Store(db, "job", {
+  schema: new Schema<Job>(),
+  indices: ["status.type", "user.id", "chat.id"],
+});
+
+jobStore.getBy("user.id", 123).then(() => {});
 
 export async function pushJob(job: Job) {
   await jobStore.create(job);
 }
 
 async function takeJob(): Promise<Model<Job> | null> {
-  const jobs = await jobStore.list();
+  const jobs = await jobStore.getAll();
   const job = jobs.find((job) => job.value.status.type === "idle");
   if (!job) return null;
   await job.update({ status: { type: "processing", progress: 0, updatedDate: new Date() } });
@@ -32,18 +41,20 @@ async function takeJob(): Promise<Model<Job> | null> {
 }
 
 export async function getAllJobs(): Promise<Array<Job>> {
-  return await jobStore.list().then((jobs) => jobs.map((job) => job.value));
+  return await jobStore.getAll().then((jobs) => jobs.map((job) => job.value));
 }
 
 export async function processQueue() {
   while (true) {
-    const job = await takeJob();
+    const job = await takeJob().catch((err) =>
+      void logger().warning("failed getting job", err.message)
+    );
     if (!job) {
       await new Promise((resolve) => setTimeout(resolve, 1000));
       continue;
     }
     let place = 0;
-    for (const job of await jobStore.list()) {
+    for (const job of await jobStore.getAll().catch(() => [])) {
       if (job.value.status.type === "idle") place += 1;
       if (place === 0) continue;
       const statusMessageText = `You are ${formatOrdinal(place)} in queue.`;
@@ -51,7 +62,7 @@ export async function processQueue() {
         await bot.api.sendMessage(job.value.chat.id, statusMessageText, {
           reply_to_message_id: job.value.requestMessage.message_id,
         }).catch(() => undefined)
-          .then((message) => job.update({ statusMessage: message }));
+          .then((message) => job.update({ statusMessage: message })).catch(() => undefined);
       } else {
         await bot.api.editMessageText(
           job.value.chat.id,
@@ -96,9 +107,10 @@ export async function processQueue() {
           }
         },
       );
-      console.log(
-        `Finished job for ${job.value.user.first_name} in ${job.value.chat.type} chat`,
-      );
+      const jobCount = (await jobStore.getAll()).filter((job) =>
+        job.value.status.type !== "processing"
+      ).length;
+      logger().info("Job finished", job.value.user.first_name, job.value.chat.type, { jobCount });
       if (job.value.statusMessage) {
         await bot.api.editMessageText(
           job.value.chat.id,
@@ -126,9 +138,7 @@ export async function processQueue() {
       });
       await job.delete();
     } catch (err) {
-      console.error(
-        `Failed to generate an image for ${job.value.user.first_name} in ${job.value.chat.type} chat: ${err}`,
-      );
+      logger().error("Job failed", job.value.user.first_name, job.value.chat.type, err);
       const errorMessage = await bot.api
         .sendMessage(job.value.chat.id, err.toString(), {
           reply_to_message_id: job.value.requestMessage.message_id,
@@ -138,12 +148,16 @@ export async function processQueue() {
         if (job.value.statusMessage) {
           await bot.api
             .deleteMessage(job.value.chat.id, job.value.statusMessage.message_id)
-            .catch(() => undefined)
-            .then(() => job.update({ statusMessage: undefined }));
+            .then(() => job.update({ statusMessage: undefined }))
+            .catch(() => void logger().warning("failed deleting status message", err.message));
         }
-        job.update({ status: { type: "idle" } });
+        await job.update({ status: { type: "idle" } }).catch((err) =>
+          void logger().warning("failed returning job", err.message)
+        );
       } else {
-        await job.delete();
+        await job.delete().catch((err) =>
+          void logger().warning("failed deleting job", err.message)
+        );
       }
     }
   }
@@ -152,15 +166,15 @@ export async function processQueue() {
 export async function returnHangedJobs() {
   while (true) {
     await new Promise((resolve) => setTimeout(resolve, 5000));
-    const jobs = await jobStore.list();
+    const jobs = await jobStore.getAll().catch(() => []);
     for (const job of jobs) {
-      if (job.value.status.type === "idle") continue;
+      if (job.value.status.type !== "processing") continue;
       // if job wasn't updated for 1 minute, return it to the queue
       if (job.value.status.updatedDate.getTime() < Date.now() - 60 * 1000) {
-        console.log(
-          `Returning hanged job for ${job.value.user.first_name} in ${job.value.chat.type} chat`,
+        logger().warning("Hanged job returned", job.value.user.first_name, job.value.chat.type);
+        await job.update({ status: { type: "idle" } }).catch((err) =>
+          void logger().warning("failed returning job", err.message)
         );
-        await job.update({ status: { type: "idle" } });
       }
     }
   }

sd.ts

@@ -1,9 +1,9 @@
 export async function txt2img(
   apiUrl: string,
-  params: Partial<SdRequest>,
+  params: Partial<SdTxt2ImgRequest>,
   onProgress?: (progress: SdProgressResponse) => void,
   signal?: AbortSignal,
-): Promise<SdResponse> {
+): Promise<SdTxt2ImgResponse> {
   let response: Response | undefined;
   let error: unknown;
 
@@ -26,7 +26,7 @@ export async function txt2img(
       }
       if (response != null) {
         if (response.ok) {
-          const result = (await response.json()) as SdResponse;
+          const result = (await response.json()) as SdTxt2ImgResponse;
           return result;
         } else {
           throw new Error(`Request failed: ${response.status} ${response.statusText}`);
@@ -44,7 +44,7 @@ export async function txt2img(
   }
 }
 
-export interface SdRequest {
+export interface SdTxt2ImgRequest {
   denoising_strength: number;
   prompt: string;
   seed: number;
@@ -60,9 +60,9 @@ export interface SdRequest {
   save_images: boolean;
 }
 
-export interface SdResponse {
+export interface SdTxt2ImgResponse {
   images: string[];
-  parameters: SdRequest;
+  parameters: SdTxt2ImgRequest;
   /** Contains serialized JSON */
   info: string;
 }

session.ts

@@ -1,5 +1,5 @@
 import { Context, DenoKVAdapter, session, SessionFlavor } from "./deps.ts";
-import { SdRequest } from "./sd.ts";
+import { SdTxt2ImgRequest } from "./sd.ts";
 
 export type MySessionFlavor = SessionFlavor<SessionData>;
 
@@ -15,7 +15,7 @@ export interface GlobalData {
   sdApiUrl: string;
   maxUserJobs: number;
   maxJobs: number;
-  defaultParams?: Partial<SdRequest>;
+  defaultParams?: Partial<SdTxt2ImgRequest>;
 }
 
 export interface ChatData {

store.test.ts

@@ -0,0 +1,96 @@
+import { assert } from "https://deno.land/std@0.198.0/assert/assert.ts";
+import { Schema, Store } from "./store.ts";
+import { log } from "./deps.ts";
+
+const db = await Deno.openKv();
+
+log.setup({
+  handlers: {
+    console: new log.handlers.ConsoleHandler("DEBUG", {}),
+  },
+  loggers: {
+    kvStore: { level: "DEBUG", handlers: ["console"] },
+  },
+});
+
+interface PointSchema {
+  x: number;
+  y: number;
+}
+
+interface JobSchema {
+  name: string;
+  params: {
+    a: number;
+    b: number | null;
+  };
+  status: { type: "idle" } | { type: "processing"; progress: number } | { type: "done" };
+  lastUpdateDate: Date;
+}
+
+const pointStore = new Store(db, "points", {
+  schema: new Schema<PointSchema>(),
+  indices: ["x", "y"],
+});
+const jobStore = new Store(db, "jobs", {
+  schema: new Schema<JobSchema>(),
+  indices: ["name", "status.type"],
+});
+
+Deno.test("create and delete", async () => {
+  await pointStore.deleteAll();
+  const point1 = await pointStore.create({ x: 1, y: 2 });
+  const point2 = await pointStore.create({ x: 3, y: 4 });
+  assert((await pointStore.getAll()).length === 2);
+  const point3 = await pointStore.create({ x: 5, y: 6 });
+  assert((await pointStore.getAll()).length === 3);
+  assert((await pointStore.get(point2.id))?.value.y === 4);
+  await point1.delete();
+  assert((await pointStore.getAll()).length === 2);
+  await point2.delete();
+  await point3.delete();
+  assert((await pointStore.getAll()).length === 0);
+});
+
+Deno.test("list by index", async () => {
+  await jobStore.deleteAll();
+
+  const test = await jobStore.create({
+    name: "test",
+    params: { a: 1, b: null },
+    status: { type: "idle" },
+    lastUpdateDate: new Date(),
+  });
+  assert((await jobStore.getBy("name", "test"))[0].value.params.a === 1);
+  assert((await jobStore.getBy("status.type", "idle"))[0].value.params.a === 1);
+
+  await test.update({ status: { type: "processing", progress: 33 } });
+  assert((await jobStore.getBy("status.type", "processing"))[0].value.params.a === 1);
+
+  await test.update({ status: { type: "done" } });
+  assert((await jobStore.getBy("status.type", "done"))[0].value.params.a === 1);
+  assert((await jobStore.getBy("status.type", "processing")).length === 0);
+
+  await test.delete();
+  assert((await jobStore.getBy("status.type", "done")).length === 0);
+  assert((await jobStore.getBy("name", "test")).length === 0);
+});
+
+Deno.test("fail on concurrent update", async () => {
+  await jobStore.deleteAll();
+
+  const test = await jobStore.create({
+    name: "test",
+    params: { a: 1, b: null },
+    status: { type: "idle" },
+    lastUpdateDate: new Date(),
+  });
+
+  const result = await Promise.all([
+    test.update({ status: { type: "processing", progress: 33 } }),
+    test.update({ status: { type: "done" } }),
+  ]).catch(() => true);
+  assert(result === true);
+
+  await test.delete();
+});

store.ts

@@ -1,76 +1,241 @@
-import { ulid } from "./deps.ts";
+import { log, ulid } from "./deps.ts";
 
-export class Store<T extends object> {
-  constructor(
-    private readonly db: Deno.Kv,
-    private readonly storeKey: Deno.KvKeyPart,
-  ) {
+const logger = () => log.getLogger("kvStore");
+
+export type validIndexKey<T> = {
+  [K in keyof T]: K extends string ? (T[K] extends Deno.KvKeyPart ? K
+      : T[K] extends readonly unknown[] ? never
+      : T[K] extends object ? `${K}.${validIndexKey<T[K]>}`
+      : never)
+    : never;
+}[keyof T];
+
+export type indexValue<T, I extends validIndexKey<T>> = I extends `${infer K}.${infer Rest}`
+  ? K extends keyof T ? Rest extends validIndexKey<T[K]> ? indexValue<T[K], Rest>
+    : never
+  : never
+  : I extends keyof T ? T[I]
+  : never;
+
+export class Schema<T> {}
+
+interface StoreOptions<T, I> {
+  readonly schema: Schema<T>;
+  readonly indices: readonly I[];
+}
+
+export class Store<T, I extends validIndexKey<T>> {
+  readonly #db: Deno.Kv;
+  readonly #key: Deno.KvKeyPart;
+  readonly #indices: readonly I[];
+
+  constructor(db: Deno.Kv, key: Deno.KvKeyPart, options: StoreOptions<T, I>) {
+    this.#db = db;
+    this.#key = key;
+    this.#indices = options.indices;
   }
 
   async create(value: T): Promise<Model<T>> {
     const id = ulid();
-    await this.db.set([this.storeKey, id], value);
-    return new Model(this.db, this.storeKey, id, value);
+    await this.#db.set([this.#key, "id", id], value);
+    logger().debug(["created", this.#key, "id", id].join(" "));
+    for (const index of this.#indices) {
+      const indexValue: Deno.KvKeyPart = index
+        .split(".")
+        .reduce((value, key) => value[key], value as any);
+      await this.#db.set([this.#key, index, indexValue, id], value);
+      logger().debug(["created", this.#key, index, indexValue, id].join(" "));
+    }
+    return new Model(this.#db, this.#key, this.#indices, id, value);
   }
 
   async get(id: Deno.KvKeyPart): Promise<Model<T> | null> {
-    const entry = await this.db.get<T>([this.storeKey, id]);
+    const entry = await this.#db.get<T>([this.#key, "id", id]);
     if (entry.versionstamp == null) return null;
-    return new Model(this.db, this.storeKey, id, entry.value);
+    return new Model(this.#db, this.#key, this.#indices, id, entry.value);
   }
 
-  async list(): Promise<Array<Model<T>>> {
-    const models: Array<Model<T>> = [];
-    for await (const entry of this.db.list<T>({ prefix: [this.storeKey] })) {
-      models.push(new Model(this.db, this.storeKey, entry.key[1], entry.value));
+  async getBy<J extends I>(
+    index: J,
+    value: indexValue<T, J>,
+    options?: Deno.KvListOptions,
+  ): Promise<Array<Model<T>>> {
+    const models: Model<T>[] = [];
+    for await (
+      const entry of this.#db.list<T>(
+        { prefix: [this.#key, index, value as Deno.KvKeyPart] },
+        options,
+      )
+    ) {
+      models.push(new Model(this.#db, this.#key, this.#indices, entry.key[3], entry.value));
     }
     return models;
   }
+
+  async getAll(
+    opts?: { limit?: number; reverse?: boolean },
+  ): Promise<Array<Model<T>>> {
+    const { limit, reverse } = opts ?? {};
+    const models: Array<Model<T>> = [];
+    for await (
+      const entry of this.#db.list<T>({
+        prefix: [this.#key, "id"],
+      }, { limit, reverse })
+    ) {
+      models.push(new Model(this.#db, this.#key, this.#indices, entry.key[2], entry.value));
+    }
+    return models;
+  }
+
+  async deleteAll(): Promise<void> {
+    for await (const entry of this.#db.list({ prefix: [this.#key] })) {
+      await this.#db.delete(entry.key);
+      logger().debug(["deleted", ...entry.key].join(" "));
+    }
+  }
 }
 
-export class Model<T extends object> {
-  #value: T;
+export class Model<T> {
+  readonly #db: Deno.Kv;
+  readonly #key: Deno.KvKeyPart;
+  readonly #indices: readonly string[];
+  readonly #id: Deno.KvKeyPart;
+  value: T;
 
   constructor(
-    private readonly db: Deno.Kv,
-    private readonly storeKey: Deno.KvKeyPart,
-    private readonly entryKey: Deno.KvKeyPart,
+    db: Deno.Kv,
+    key: Deno.KvKeyPart,
+    indices: readonly string[],
+    id: Deno.KvKeyPart,
     value: T,
   ) {
-    this.#value = value;
+    this.#db = db;
+    this.#key = key;
+    this.#indices = indices;
+    this.#id = id;
+    this.value = value;
   }
 
-  get value(): T {
-    return this.#value;
+  get id(): Deno.KvKeyPart {
+    return this.#id;
   }
 
-  async get(): Promise<T | null> {
-    const entry = await this.db.get<T>([this.storeKey, this.entryKey]);
-    if (entry.versionstamp == null) return null;
-    this.#value = entry.value;
-    return entry.value;
-  }
+  async update(updater: Partial<T> | ((value: T) => T)): Promise<T | null> {
+    // get current main entry
+    const oldEntry = await this.#db.get<T>([this.#key, "id", this.#id]);
 
-  async set(value: T): Promise<T> {
-    await this.db.set([this.storeKey, this.entryKey], value);
-    this.#value = value;
-    return value;
-  }
-
-  async update(value: Partial<T> | ((value: T) => T)): Promise<T | null> {
-    const entry = await this.db.get<T>([this.storeKey, this.entryKey]);
-    if (entry.versionstamp == null) return null;
-    if (typeof value === "function") {
-      entry.value = value(entry.value);
-    } else {
-      entry.value = { ...entry.value, ...value };
+    // get all current index entries
+    const oldIndexEntries: Record<string, Deno.KvEntryMaybe<T>> = {};
+    for (const index of this.#indices) {
+      const indexKey: Deno.KvKeyPart = index
+        .split(".")
+        .reduce((value, key) => value[key], oldEntry.value as any);
+      oldIndexEntries[index] = await this.#db.get<T>([this.#key, index, indexKey, this.#id]);
     }
-    await this.db.set([this.storeKey, this.entryKey], entry.value);
-    this.#value = entry.value;
-    return entry.value;
+
+    // compute new value
+    if (typeof updater === "function") {
+      this.value = updater(this.value);
+    } else {
+      this.value = { ...this.value, ...updater };
+    }
+
+    // begin transaction
+    const transaction = this.#db.atomic();
+
+    // set the main entry
+    transaction
+      .check(oldEntry)
+      .set([this.#key, "id", this.#id], this.value);
+    logger().debug(["updated", this.#key, "id", this.#id].join(" "));
+
+    // delete and create all changed index entries
+    for (const index of this.#indices) {
+      const oldIndexKey: Deno.KvKeyPart = index
+        .split(".")
+        .reduce((value, key) => value[key], oldIndexEntries[index].value as any);
+      const newIndexKey: Deno.KvKeyPart = index
+        .split(".")
+        .reduce((value, key) => value[key], this.value as any);
+      if (newIndexKey !== oldIndexKey) {
+        transaction
+          .check(oldIndexEntries[index])
+          .delete([this.#key, index, oldIndexKey, this.#id])
+          .set([this.#key, index, newIndexKey, this.#id], this.value);
+        logger().debug(["deleted", this.#key, index, oldIndexKey, this.#id].join(" "));
+        logger().debug(["created", this.#key, index, newIndexKey, this.#id].join(" "));
+      }
+    }
+
+    // commit
+    const result = await transaction.commit();
+    if (!result.ok) throw new Error(`Failed to update ${this.#key} ${this.#id}`);
+    return this.value;
   }
 
   async delete(): Promise<void> {
-    await this.db.delete([this.storeKey, this.entryKey]);
+    // get current main entry
+    const entry = await this.#db.get<T>([this.#key, "id", this.#id]);
+
+    // begin transaction
+    const transaction = this.#db.atomic();
+
+    // delete main entry
+    transaction
+      .check(entry)
+      .delete([this.#key, "id", this.#id]);
+    logger().debug(["deleted", this.#key, "id", this.#id].join(" "));
+
+    // delete all index entries
+    for (const index of this.#indices) {
+      const indexKey: Deno.KvKeyPart = index
+        .split(".")
+        .reduce((value, key) => value[key], entry.value as any);
+      transaction
+        .delete([this.#key, index, indexKey, this.#id]);
+      logger().debug(["deleted", this.#key, index, indexKey, this.#id].join(" "));
+    }
+
+    // commit
+    const result = await transaction.commit();
+    if (!result.ok) throw new Error(`Failed to delete ${this.#key} ${this.#id}`);
   }
 }
+
+export async function retry<T>(
+  fn: () => Promise<T>,
+  options: { maxAttempts?: number; delayMs?: number } = {},
+): Promise<T> {
+  const { maxAttempts = 3, delayMs = 1000 } = options;
+  let error: unknown;
+  for (let attempt = 0; attempt < maxAttempts; attempt++) {
+    try {
+      return await fn();
+    } catch (err) {
+      error = err;
+      await new Promise((resolve) => setTimeout(resolve, delayMs));
+    }
+  }
+  throw error;
+}
+
+export async function collectIterator<T>(
+  iterator: AsyncIterableIterator<T>,
+  options: { maxItems?: number; timeoutMs?: number } = {},
+): Promise<T[]> {
+  const { maxItems = 1000, timeoutMs = 2000 } = options;
+  const result: T[] = [];
+  const timeout = setTimeout(() => iterator.return?.(), timeoutMs);
+  try {
+    for await (const item of iterator) {
+      result.push(item);
+      if (result.length >= maxItems) {
+        iterator.return?.();
+        break;
+      }
+    }
+  } finally {
+    clearTimeout(timeout);
+  }
+  return result;
+}