@convex-dev/sharded-counter 0.1.3 → 0.1.4

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

@@ -1,5 +1,3 @@
-/* prettier-ignore-start */
-
 /* eslint-disable */
 /**
  * Generated `api` utility.
@@ -37,6 +35,18 @@ export type Mounts = {
       null
     >;
     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
+    >;
   };
 };
 // For now fullApiWithMounts is only fullApi which provides
@@ -54,5 +64,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

@@ -26,7 +26,7 @@ describe("counter", () => {
 
 fcTest.prop({
   updates: fc.array(fc.record({
-    v: fc.integer({ min: 0, max: 10000 }).map((i) => i / 100),
+    v: fc.integer({ min: -10000, max: 10000 }).map((i) => i / 100),
     key: fc.string(),
     shards: fc.option(fc.integer({ min: 1, max: 100 })),
   })),
@@ -41,5 +41,15 @@ fcTest.prop({
       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

@@ -41,3 +41,65 @@ 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 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;
+}