allow non-string keys

README.md

@@ -7,7 +7,7 @@
 <!-- START: Include on https://convex.dev/components -->
 
 This component adds counters to Convex. It acts as a key-value store from
-string to number, with sharding to increase throughput when updating values.
+key to number, with sharding to increase throughput when updating values.
 
 Since it's built on Convex, everything is automatically consistent, reactive,
 and cached. Since it's built with [Components](https://convex.dev/components),
@@ -143,20 +143,27 @@ const counter = new ShardedCounter(components.shardedCounter, {
 Or by setting a default that applies to all keys not specified in `shards`:
 
 ```ts
-const counter = new ShardedCounter(components.shardedCounter, {
+const counter = new ShardedCounter<string>(components.shardedCounter, {
   shards: { checkboxes: 100 },
   defaultShards: 20,
 });
 ```
 
 The default number of shards if none is specified is 16.
 
-Note your keys can be a subtype of string. e.g. if you want to store a count of
+## Keys can be anything
+
+The keys for your counters can be string literals,
+a subtype of string, or any arbitrary
+[Convex value](https://docs.convex.dev/database/types). You can enforce correct
+key usage with the generic Typescript type parameter.
+
+For example, if you want to store a count of
 friends for each user, and you don't care about throughput for a single user,
 you would declare ShardedCounter like so:
 
 ```ts
-const friendCounts = new ShardedCounter<Record<Id<"users">, number>>(
+const friendCounts = new ShardedCounter<Id<"users">>(
   components.shardedCounter,
   { defaultShards: 1 },
 );
@@ -165,6 +172,31 @@ const friendCounts = new ShardedCounter<Record<Id<"users">, number>>(
 await friendsCount.dec(ctx, userId);
 ```
 
+If you want to store two counts for each user: a count of followers and a count
+of followees, you can use a tuple key:
+
+```ts
+const followCounts = new ShardedCounter<[Id<"users">, "followers" | "follows"]>(
+  components.shardedCounter,
+  { defaultShards: 3 },
+);
+
+// Create a follower relationship. Note Convex mutations are atomic so don't
+// worry about the counts getting out of sync.
+await followCounts.inc(ctx, [follower, "follows"]);
+await followCounts.inc(ctx, [followee, "followers"]);
+```
+
+You can also install the component multiple times to count different things:
+
+```ts
+app.use(shardedCounter, { name: "friendCounter" });
+app.use(shardedCounter, { name: "followCounter" });
+```
+
+See [example/convex/nested.ts](example/convex/nested.ts) for this example in
+practice.
+
 ## Reduce contention on reads
 
 Reading the count with `counter.count(ctx, "checkboxes")` reads from all shards

example/convex/_generated/api.d.ts

@@ -9,6 +9,7 @@
  */
 
 import type * as example from "../example.js";
+import type * as nested from "../nested.js";
 
 import type {
   ApiFromModules,
@@ -25,6 +26,7 @@ import type {
  */
 declare const fullApi: ApiFromModules<{
   example: typeof example;
+  nested: typeof nested;
 }>;
 declare const fullApiWithMounts: typeof fullApi;
 
@@ -43,20 +45,43 @@ export declare const components: {
       add: FunctionReference<
         "mutation",
         "internal",
-        { count: number; name: string; shards?: number },
+        { count: number; name: any; shards?: number },
         null
       >;
-      count: FunctionReference<"query", "internal", { name: string }, number>;
+      count: FunctionReference<"query", "internal", { name: any }, number>;
       estimateCount: FunctionReference<
         "query",
         "internal",
-        { name: string; readFromShards?: number; shards?: number },
+        { name: any; readFromShards?: number; shards?: number },
         any
       >;
       rebalance: FunctionReference<
         "mutation",
         "internal",
-        { name: string; shards?: number },
+        { name: any; shards?: number },
+        any
+      >;
+    };
+  };
+  nestedCounter: {
+    public: {
+      add: FunctionReference<
+        "mutation",
+        "internal",
+        { count: number; name: any; shards?: number },
+        null
+      >;
+      count: FunctionReference<"query", "internal", { name: any }, number>;
+      estimateCount: FunctionReference<
+        "query",
+        "internal",
+        { name: any; readFromShards?: number; shards?: number },
+        any
+      >;
+      rebalance: FunctionReference<
+        "mutation",
+        "internal",
+        { name: any; shards?: number },
         any
       >;
     };

example/convex/convex.config.ts

@@ -4,6 +4,7 @@ import migrations from "@convex-dev/migrations/convex.config";
 
 const app = defineApp();
 app.use(shardedCounter);
+app.use(shardedCounter, { name: "nestedCounter" });
 app.use(migrations);
 
 export default app;

example/convex/example.ts

@@ -9,7 +9,7 @@ import { customCtx, customMutation } from "convex-helpers/server/customFunctions
 
 /// Example of ShardedCounter initialization.
 
-const counter = new ShardedCounter(components.shardedCounter, {
+const counter = new ShardedCounter<"beans" | "users" | "accomplishments">(components.shardedCounter, {
   shards: { beans: 10, users: 3 },
 });
 const numUsers = counter.for("users");

example/convex/nested.test.ts

@@ -0,0 +1,32 @@
+/// <reference types="vite/client" />
+
+import { convexTest } from "convex-test";
+import { describe, expect, test } from "vitest"
+import schema from "./schema";
+import componentSchema from "../../src/component/schema";
+import { api } from "./_generated/api";
+
+const modules = import.meta.glob("./**/*.ts");
+const componentModules = import.meta.glob("../../src/component/**/*.ts");
+
+describe("nested sharded counter", () => {
+  async function setupTest() {
+    const t = convexTest(schema, modules);
+    t.registerComponent("nestedCounter", componentSchema, componentModules);
+    return t;
+  }
+
+  test("follower and follows count", async () => {
+    const t = await setupTest();
+    const users = await t.run((ctx) => {
+      return Promise.all([
+        ctx.db.insert("users", { name: "1" }),
+        ctx.db.insert("users", { name: "2" }),
+      ]);
+    });
+    await t.mutation(api.nested.addFollower, { follower: users[0], followee: users[1] });
+    expect(await t.query(api.nested.countFollows, { user: users[0] })).toStrictEqual(1);
+    expect(await t.query(api.nested.countFollowers, { user: users[0] })).toStrictEqual(0);
+    expect(await t.query(api.nested.countFollowers, { user: users[1] })).toStrictEqual(1);
+  })
+});

example/convex/nested.ts

@@ -0,0 +1,39 @@
+/// Example of a hierarchical counter.
+/// This is the same as a regular sharded counter (see example.ts), but the keys
+/// are tuples instead of strings.
+/// You cannot accumulate across a prefix of the tuple; if you want to do that,
+/// use an additional counter for each prefix, or use the Aggregate component:
+/// https://www.npmjs.com/package/@convex-dev/aggregate
+
+import { ShardedCounter } from "@convex-dev/sharded-counter";
+import { components } from "./_generated/api";
+import { Id } from "./_generated/dataModel";
+import { v } from "convex/values";
+import { mutation, query } from "./_generated/server";
+
+const nestedCounter = new ShardedCounter<[Id<"users">, "follows" | "followers"]>(
+  components.nestedCounter,
+  { defaultShards: 3 },
+);
+
+export const addFollower = mutation({
+  args: { follower: v.id("users"), followee: v.id("users") },
+  handler: async (ctx, { follower, followee }) => {
+    await nestedCounter.inc(ctx, [followee, "followers"]);
+    await nestedCounter.inc(ctx, [follower, "follows"]);
+  },
+});
+
+export const countFollows = query({
+  args: { user: v.id("users") },
+  handler: async (ctx, { user }) => {
+    return await nestedCounter.count(ctx, [user, "follows"]);
+  },
+});
+
+export const countFollowers = query({
+  args: { user: v.id("users") },
+  handler: async (ctx, { user }) => {
+    return await nestedCounter.count(ctx, [user, "followers"]);
+  },
+});

example/package.json

@@ -12,7 +12,7 @@
   "dependencies": {
     "@convex-dev/migrations": "^0.1.6",
     "@convex-dev/sharded-counter": "file:..",
-    "convex": "^1.17.0",
+    "convex": "file:../node_modules/convex",
     "convex-helpers": "^0.1.61",
     "react": "^18.3.1",
     "react-dom": "^18.3.1"
@@ -23,6 +23,7 @@
     "@types/react": "^18.3.3",
     "@types/react-dom": "^18.3.0",
     "@vitejs/plugin-react": "^4.3.1",
+    "convex-test": "^0.0.33",
     "eslint": "^9.9.0",
     "eslint-plugin-react-hooks": "^5.1.0-rc.0",
     "eslint-plugin-react-refresh": "^0.4.9",

package.json

@@ -57,15 +57,15 @@
   },
   "devDependencies": {
     "@eslint/js": "^9.9.1",
+    "@fast-check/vitest": "^0.1.3",
     "@types/node": "^18.17.0",
+    "convex-test": "^0.0.33",
     "eslint": "^9.9.1",
-    "convex-test": "^0.0.31",
-    "@fast-check/vitest": "^0.1.3",
     "globals": "^15.9.0",
-    "vitest": "^2.1.1",
     "prettier": "3.2.5",
     "typescript": "~5.0.3",
-    "typescript-eslint": "^8.4.0"
+    "typescript-eslint": "^8.4.0",
+    "vitest": "^2.1.1"
   },
   "main": "./dist/commonjs/client/index.js",
   "types": "./dist/commonjs/client/index.d.ts",

src/client/index.ts

@@ -7,14 +7,14 @@ import {
   GenericQueryCtx,
   TableNamesInDataModel,
 } from "convex/server";
-import { GenericId } from "convex/values";
+import { GenericId, Value as ConvexValue } from "convex/values";
 import { api } from "../component/_generated/api";
 
 /**
  * A sharded counter is a map from string -> counter, where each counter can
  * be incremented or decremented atomically.
  */
-export class ShardedCounter<ShardsKey extends string> {
+export class ShardedCounter<ShardsKey extends ConvexValue> {
   /**
    * A sharded counter is a map from string -> counter, where each counter can
    * be incremented or decremented.
@@ -32,7 +32,7 @@ export class ShardedCounter<ShardsKey extends string> {
    */
   constructor(
     private component: UseApi<typeof api>,
-    private options?: { shards?: Record<ShardsKey, number>; defaultShards?: number }
+    private options?: { shards?: Partial<Record<ShardsKey & string, number>>; defaultShards?: number }
   ) {}
   /**
    * Increase the counter for key `name` by `count`.
@@ -41,7 +41,7 @@ export class ShardedCounter<ShardsKey extends string> {
    * @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 = ShardsKey>(
+  async add<Name extends ShardsKey>(
     ctx: RunMutationCtx,
     name: Name,
     count: number = 1
@@ -55,7 +55,7 @@ export class ShardedCounter<ShardsKey extends string> {
   /**
    * Decrease the counter for key `name` by `count`.
    */
-  async subtract<Name extends string = ShardsKey>(
+  async subtract<Name extends ShardsKey>(
     ctx: RunMutationCtx,
     name: Name,
     count: number = 1
@@ -65,13 +65,13 @@ export class ShardedCounter<ShardsKey extends string> {
   /**
    * Increment the counter for key `name` by 1.
    */
-  async inc<Name extends string = ShardsKey>(ctx: RunMutationCtx, name: Name) {
+  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 string = ShardsKey>(ctx: RunMutationCtx, name: Name) {
+  async dec<Name extends ShardsKey>(ctx: RunMutationCtx, name: Name) {
     return this.add(ctx, name, -1);
   }
   /**
@@ -80,7 +80,7 @@ export class ShardedCounter<ShardsKey extends string> {
    * 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 = ShardsKey>(
+  async count<Name extends ShardsKey>(
     ctx: RunQueryCtx,
     name: Name
   ) {
@@ -98,7 +98,7 @@ export class ShardedCounter<ShardsKey extends string> {
    * This operation reads and writes all shards, so it can cause contention if
    * called too often.
    */
-  async rebalance<Name extends string = ShardsKey>(
+  async rebalance<Name extends ShardsKey>(
     ctx: RunMutationCtx,
     name: Name,
   ) {
@@ -118,7 +118,7 @@ export class ShardedCounter<ShardsKey extends string> {
    * 
    * Use this to reduce contention when reading the counter.
    */
-  async estimateCount<Name extends string = ShardsKey>(
+  async estimateCount<Name extends ShardsKey>(
     ctx: RunQueryCtx,
     name: Name,
     readFromShards: number = 1,
@@ -144,7 +144,7 @@ export class ShardedCounter<ShardsKey extends string> {
    * });
    * ```
    */
-  for<Name extends string = ShardsKey>(name: Name) {
+  for<Name extends ShardsKey>(name: Name) {
     return {
       /**
        * Add `count` to the counter.
@@ -190,7 +190,7 @@ export class ShardedCounter<ShardsKey extends string> {
   }
   trigger<
     Ctx extends RunMutationCtx,
-    Name extends string = ShardsKey,
+    Name extends ShardsKey,
   >(
     name: Name,
   ): Trigger<Ctx, GenericDataModel, TableNamesInDataModel<GenericDataModel>> {
@@ -202,8 +202,8 @@ export class ShardedCounter<ShardsKey extends string> {
       }
     };
   }
-  private shardsForKey<Name extends string = ShardsKey>(name: Name) {
-    const explicitShards = this.options?.shards?.[name as string as ShardsKey];
+  private shardsForKey<Name extends ShardsKey>(name: Name) {
+    const explicitShards = this.options?.shards?.[name as unknown as string & ShardsKey];
     return explicitShards ?? this.options?.defaultShards;
   }
 }