@convex-dev/sharded-counter 0.1.3 → 0.1.5

package/package.json

@@ -7,7 +7,7 @@
     "email": "support@convex.dev",
     "url": "https://github.com/get-convex/sharded-counter/issues"
   },
-  "version": "0.1.3",
+  "version": "0.1.5",
   "license": "Apache-2.0",
   "keywords": [
     "convex",
@@ -22,7 +22,7 @@
     "dev": "cd example; npm run dev",
     "typecheck": "tsc --noEmit",
     "prepare": "npm run build",
-    "test": "vitest run",
+    "test": "vitest",
     "test:debug": "vitest --inspect-brk --no-file-parallelism",
     "test:coverage": "vitest run --coverage --coverage.reporter=text"
   },

package/src/client/index.ts

@@ -1,24 +1,30 @@
 import {
+  DocumentByName,
   Expand,
   FunctionReference,
   GenericDataModel,
   GenericMutationCtx,
   GenericQueryCtx,
+  TableNamesInDataModel,
 } from "convex/server";
 import { GenericId } from "convex/values";
 import { api } from "../component/_generated/api";
 
-export class ShardedCounter<Shards extends Record<string, number>> {
+/**
+ * A sharded counter is a map from string -> counter, where each counter can
+ * be incremented or decremented atomically.
+ */
+export class ShardedCounter<ShardsKey extends string> {
   /**
    * A sharded counter is a map from string -> counter, where each counter can
    * be incremented or decremented.
-   * 
+   *
    * The counter is sharded into multiple documents to allow for higher
    * throughput of updates. The default number of shards is 16.
-   * 
+   *
    * - More shards => higher throughput of updates.
    * - Fewer shards => lower latency when querying the counter.
-   * 
+   *
    * @param options.shards The number of shards for each counter, for fixed
    *   keys.
    * @param options.defaultShards The number of shards for each counter, for
@@ -26,38 +32,130 @@ export class ShardedCounter<Shards extends Record<string, number>> {
    */
   constructor(
     private component: UseApi<typeof api>,
-    private options?: { shards?: Shards; defaultShards?: number }
-  ) {}
+    options?: {
+      shards?: Partial<Record<ShardsKey, number>>;
+      defaultShards?: number;
+    },
+  ) {
+    this.stickyShard = {};
+    const defaultShards = options?.defaultShards;
+    this.shardsForKey = (name: ShardsKey) => {
+      const explicitShards = options?.shards?.[name];
+      return explicitShards ?? defaultShards;
+    };
+  }
+
+  private shardsForKey: (name: ShardsKey) => number | undefined;
+
+  // Keep track of the shard for each key, so multiple mutations on the same key
+  // will use the same shard.
+  private stickyShard: Record<string, number>;
+
   /**
    * Increase the counter for key `name` by `count`.
    * If `count` is negative, the counter will decrease.
-   * 
+   *
    * @param name The key to update the counter for.
    * @param count The amount to increment the counter by. Defaults to 1.
    */
-  async add<Name extends string = keyof Shards & string>(
+  async add<Name extends ShardsKey>(
     ctx: RunMutationCtx,
     name: Name,
-    count: number = 1
+    count: number = 1,
   ) {
-    const shards = this.options?.shards?.[name] ?? this.options?.defaultShards;
-    return ctx.runMutation(this.component.public.add, {
+    const shard = await ctx.runMutation(this.component.public.add, {
       name,
       count,
-      shards,
+      shard: this.stickyShard?.[name],
+      shards: this.shardsForKey(name),
     });
+    this.stickyShard[name] = shard;
   }
+
+  /**
+   * Decrease the counter for key `name` by `count`.
+   */
+  async subtract<Name extends ShardsKey>(
+    ctx: RunMutationCtx,
+    name: Name,
+    count: number = 1,
+  ) {
+    return this.add(ctx, name, -count);
+  }
+
+  /**
+   * Increment the counter for key `name` by 1.
+   */
+  async inc<Name extends ShardsKey>(ctx: RunMutationCtx, name: Name) {
+    return this.add(ctx, name, 1);
+  }
+
+  /**
+   * Decrement the counter for key `name` by 1.
+   */
+  async dec<Name extends ShardsKey>(ctx: RunMutationCtx, name: Name) {
+    return this.add(ctx, name, -1);
+  }
+
   /**
    * Gets the counter for key `name`.
    *
    * NOTE: this reads from all shards. If used in a mutation, it will contend
    * with all mutations that update the counter for this key.
    */
-  async count<Name extends string = keyof Shards & string>(
+  async count<Name extends ShardsKey>(ctx: RunQueryCtx, name: Name) {
+    return ctx.runQuery(this.component.public.count, { name });
+  }
+
+  /**
+   * Redistribute counts evenly across the counter's shards.
+   *
+   * If there were more shards for this counter at some point, those shards
+   * will be removed.
+   *
+   * If there were fewer shards for this counter, or if the random distribution
+   * of counts is uneven, the counts will be redistributed evenly.
+   *
+   * This operation reads and writes all shards, so it can cause contention if
+   * called too often.
+   */
+  async rebalance<Name extends ShardsKey>(ctx: RunMutationCtx, name: Name) {
+    await ctx.runMutation(this.component.public.rebalance, {
+      name,
+      shards: this.shardsForKey(name),
+    });
+  }
+
+  /**
+   * Clear the counter for key `name`.
+   *
+   * @param name The key to clear the counter for.
+   */
+  async reset<Name extends ShardsKey>(ctx: RunMutationCtx, name: Name) {
+    await ctx.runMutation(this.component.public.reset, { name });
+  }
+
+  /**
+   * Estimate the count of a counter by only reading from a subset of shards,
+   * and extrapolating the total count.
+   *
+   * After a `rebalance`, or if there were a lot of data points to yield a
+   * random distribution across shards, this should be a good approximation of
+   * the total count. If there are few data points, which are not evenly
+   * distributed across shards, this will be a poor approximation.
+   *
+   * Use this to reduce contention when reading the counter.
+   */
+  async estimateCount<Name extends ShardsKey>(
     ctx: RunQueryCtx,
-    name: Name
+    name: Name,
+    readFromShards: number = 1,
   ) {
-    return ctx.runQuery(this.component.public.count, { name });
+    return await ctx.runQuery(this.component.public.estimateCount, {
+      name,
+      shards: this.shardsForKey(name),
+      readFromShards,
+    });
   }
   /**
    * Returns an object with methods to update and query the counter for key
@@ -74,7 +172,7 @@ export class ShardedCounter<Shards extends Record<string, number>> {
    * });
    * ```
    */
-  for<Name extends string = keyof Shards & string>(name: Name) {
+  for<Name extends ShardsKey>(name: Name) {
     return {
       /**
        * Add `count` to the counter.
@@ -96,17 +194,76 @@ export class ShardedCounter<Shards extends Record<string, number>> {
       dec: async (ctx: RunMutationCtx) => this.add(ctx, name, -1),
       /**
        * Get the current value of the counter.
-       * 
+       *
        * NOTE: this reads from all shards. If used in a mutation, it will
        * contend with all mutations that update the counter for this key.
        */
       count: async (ctx: RunQueryCtx) => this.count(ctx, name),
+      /**
+       * Reset the counter for this key.
+       */
+      reset: async (ctx: RunMutationCtx) => this.reset(ctx, name),
+      /**
+       * Redistribute counts evenly across the counter's shards.
+       *
+       * This operation reads and writes all shards, so it can cause contention
+       * if called too often.
+       */
+      rebalance: async (ctx: RunMutationCtx) => this.rebalance(ctx, name),
+      /**
+       * Estimate the counter by only reading from a subset of shards,
+       * and extrapolating the total count.
+       *
+       * Use this to reduce contention when reading the counter.
+       */
+      estimateCount: async (ctx: RunQueryCtx, readFromShards: number = 1) =>
+        this.estimateCount(ctx, name, readFromShards),
+    };
+  }
+  trigger<Ctx extends RunMutationCtx, Name extends ShardsKey>(
+    name: Name,
+  ): Trigger<Ctx, GenericDataModel, TableNamesInDataModel<GenericDataModel>> {
+    return async (ctx, change) => {
+      if (change.operation === "insert") {
+        await this.inc(ctx, name);
+      } else if (change.operation === "delete") {
+        await this.dec(ctx, name);
+      }
     };
   }
 }
 
 /* Type utils follow */
 
+export type Trigger<
+  Ctx,
+  DataModel extends GenericDataModel,
+  TableName extends TableNamesInDataModel<DataModel>,
+> = (ctx: Ctx, change: Change<DataModel, TableName>) => Promise<void>;
+
+export type Change<
+  DataModel extends GenericDataModel,
+  TableName extends TableNamesInDataModel<DataModel>,
+> = {
+  id: GenericId<TableName>;
+} & (
+  | {
+      operation: "insert";
+      oldDoc: null;
+      newDoc: DocumentByName<DataModel, TableName>;
+    }
+  | {
+      operation: "update";
+      oldDoc: DocumentByName<DataModel, TableName>;
+      newDoc: DocumentByName<DataModel, TableName>;
+    }
+  | {
+      operation: "delete";
+      oldDoc: DocumentByName<DataModel, TableName>;
+      newDoc: null;
+    }
+);
+
 type RunQueryCtx = {
   runQuery: GenericQueryCtx<GenericDataModel>["runQuery"];
 };

package/src/component/_generated/api.d.ts

@@ -1,5 +1,3 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 /**
  * Generated `api` utility.
@@ -33,10 +31,23 @@ export type Mounts = {
     add: FunctionReference<
       "mutation",
       "public",
-      { count: number; name: string; shards?: number },
-      null
+      { count: number; name: string; shard?: number; shards?: number },
+      number
     >;
     count: FunctionReference<"query", "public", { name: string }, number>;
+    estimateCount: FunctionReference<
+      "query",
+      "public",
+      { name: string; readFromShards?: number; shards?: number },
+      any
+    >;
+    rebalance: FunctionReference<
+      "mutation",
+      "public",
+      { name: string; shards?: number },
+      any
+    >;
+    reset: FunctionReference<"mutation", "public", { name: string }, any>;
   };
 };
 // For now fullApiWithMounts is only fullApi which provides
@@ -54,5 +65,3 @@ export declare const internal: FilterApi<
 >;
 
 export declare const components: {};
-
-/* prettier-ignore-end */

package/src/component/_generated/api.js

@@ -1,5 +1,3 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 /**
  * Generated `api` utility.
@@ -23,5 +21,3 @@ import { anyApi, componentsGeneric } from "convex/server";
 export const api = anyApi;
 export const internal = anyApi;
 export const components = componentsGeneric();
-
-/* prettier-ignore-end */

package/src/component/_generated/dataModel.d.ts

@@ -1,5 +1,3 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 /**
  * Generated data model types.
@@ -60,5 +58,3 @@ export type Id<TableName extends TableNames | SystemTableNames> =
  * `mutationGeneric` to make them type-safe.
  */
 export type DataModel = DataModelFromSchemaDefinition<typeof schema>;
-
-/* prettier-ignore-end */

package/src/component/_generated/server.d.ts

@@ -1,5 +1,3 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 /**
  * Generated utilities for implementing server-side Convex query and mutation functions.
@@ -149,5 +147,3 @@ export type DatabaseReader = GenericDatabaseReader<DataModel>;
  * for the guarantees Convex provides your functions.
  */
 export type DatabaseWriter = GenericDatabaseWriter<DataModel>;
-
-/* prettier-ignore-end */

package/src/component/_generated/server.js

@@ -1,5 +1,3 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 /**
  * Generated utilities for implementing server-side Convex query and mutation functions.
@@ -90,5 +88,3 @@ export const internalAction = internalActionGeneric;
  * @returns The wrapped endpoint function. Route a URL path to this function in `convex/http.js`.
  */
 export const httpAction = httpActionGeneric;
-
-/* prettier-ignore-end */

package/src/component/counter.test.ts

@@ -22,14 +22,39 @@ describe("counter", () => {
     expect(await t.query(api.public.count, { name: "beans" })).toEqual(10);
     expect(await t.query(api.public.count, { name: "friends" })).toEqual(11);
   });
+  test("respects shard argument", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.public.add, { name: "beans", count: 10, shard: 1 });
+    await t.mutation(api.public.add, { name: "beans", count: 5, shard: 2 });
+    const values = await t.run(async (ctx) => {
+      const shard1 = await ctx.db
+        .query("counters")
+        .withIndex("name", (q) => q.eq("name", "beans").eq("shard", 1))
+        .unique();
+      const shard2 = await ctx.db
+        .query("counters")
+        .withIndex("name", (q) => q.eq("name", "beans").eq("shard", 2))
+        .unique();
+      return [shard1?.value, shard2?.value];
+    });
+    expect(values).toEqual([10, 5]);
+  });
+  test("reset", async () => {
+    const t = convexTest(schema, modules);
+    await t.mutation(api.public.add, { name: "beans", count: 10 });
+    await t.mutation(api.public.reset, { name: "beans" });
+    expect(await t.query(api.public.count, { name: "beans" })).toEqual(0);
+  });
 });
 
 fcTest.prop({
-  updates: fc.array(fc.record({
-    v: fc.integer({ min: 0, max: 10000 }).map((i) => i / 100),
-    key: fc.string(),
-    shards: fc.option(fc.integer({ min: 1, max: 100 })),
-  })),
+  updates: fc.array(
+    fc.record({
+      v: fc.integer({ min: -10000, max: 10000 }).map((i) => i / 100),
+      key: fc.string(),
+      shards: fc.option(fc.integer({ min: 1, max: 100 })),
+    }),
+  ),
 })(
   "updates to counter should match in-memory counter which ignores sharding",
   async ({ updates }) => {
@@ -37,9 +62,26 @@ fcTest.prop({
     const counter = new Map<string, number>();
     for (const { v, key, shards } of updates) {
       counter.set(key, (counter.get(key) ?? 0) + v);
-      await t.mutation(api.public.add, { name: key, count: v, shards: shards ?? undefined });
+      await t.mutation(api.public.add, {
+        name: key,
+        count: v,
+        shards: shards ?? undefined,
+      });
       const count = await t.query(api.public.count, { name: key });
       expect(count).toBeCloseTo(counter.get(key)!);
     }
+    for (const [key, value] of counter.entries()) {
+      // Rebalancing keeps count the same and makes estimateCount accurate.
+      await t.mutation(api.public.rebalance, { name: key });
+      const count = await t.query(api.public.count, { name: key });
+      expect(count).toBeCloseTo(value);
+      for (let i = 1; i <= 16; i++) {
+        const estimate = await t.query(api.public.estimateCount, {
+          name: key,
+          readFromShards: i,
+        });
+        expect(estimate).toBeCloseTo(value);
+      }
+    }
   },
 );

package/src/component/public.ts

@@ -7,11 +7,14 @@ export const add = mutation({
   args: {
     name: v.string(),
     count: v.number(),
+    shard: v.optional(v.number()),
     shards: v.optional(v.number()),
   },
-  returns: v.null(),
+  returns: v.number(),
   handler: async (ctx, args) => {
-    const shard = Math.floor(Math.random() * (args.shards ?? DEFAULT_SHARD_COUNT));
+    const shard =
+      args.shard ??
+      Math.floor(Math.random() * (args.shards ?? DEFAULT_SHARD_COUNT));
     const counter = await ctx.db
       .query("counters")
       .withIndex("name", (q) => q.eq("name", args.name).eq("shard", shard))
@@ -27,6 +30,7 @@ export const add = mutation({
         shard,
       });
     }
+    return shard;
   },
 });
 
@@ -41,3 +45,83 @@ export const count = query({
     return counters.reduce((sum, counter) => sum + counter.value, 0);
   },
 });
+
+export const rebalance = mutation({
+  args: { name: v.string(), shards: v.optional(v.number()) },
+  handler: async (ctx, args) => {
+    const counters = await ctx.db
+      .query("counters")
+      .withIndex("name", (q) => q.eq("name", args.name))
+      .collect();
+    const count = counters.reduce((sum, counter) => sum + counter.value, 0);
+    const shardCount = args.shards ?? DEFAULT_SHARD_COUNT;
+    const value = count / shardCount;
+    for (let i = 0; i < shardCount; i++) {
+      const shard = counters.find((c) => c.shard === i);
+      if (shard) {
+        await ctx.db.patch(shard._id, { value });
+      } else {
+        await ctx.db.insert("counters", {
+          name: args.name,
+          value,
+          shard: i,
+        });
+      }
+    }
+    const toDelete = counters.filter((c) => c.shard >= shardCount);
+    for (const counter of toDelete) {
+      await ctx.db.delete(counter._id);
+    }
+  },
+});
+
+export const reset = mutation({
+  args: { name: v.string() },
+  handler: async (ctx, args) => {
+    await ctx.db
+      .query("counters")
+      .withIndex("name", (q) => q.eq("name", args.name))
+      .collect()
+      .then((counters) =>
+        Promise.all(counters.map((c) => ctx.db.delete(c._id))),
+      );
+  },
+});
+
+export const estimateCount = query({
+  args: {
+    name: v.string(),
+    readFromShards: v.optional(v.number()),
+    shards: v.optional(v.number()),
+  },
+  handler: async (ctx, args) => {
+    const shardCount = args.shards ?? DEFAULT_SHARD_COUNT;
+    const readFromShards = Math.min(
+      Math.max(1, args.readFromShards ?? 1),
+      shardCount,
+    );
+    const shards = shuffle(
+      Array.from({ length: shardCount }, (_, i) => i),
+    ).slice(0, readFromShards);
+    let readCount = 0;
+    for (const shard of shards) {
+      const counter = await ctx.db
+        .query("counters")
+        .withIndex("name", (q) => q.eq("name", args.name).eq("shard", shard))
+        .unique();
+      if (counter) {
+        readCount += counter.value;
+      }
+    }
+    return (readCount * shardCount) / readFromShards;
+  },
+});
+
+// Fisher-Yates shuffle
+function shuffle<T>(array: T[]): T[] {
+  for (let i = array.length - 1; i > 0; i--) {
+    const j = Math.floor(Math.random() * (i + 1));
+    [array[i], array[j]] = [array[j], array[i]];
+  }
+  return array;
+}