opa-react: tweak batch indexing/hashing (#139)

This way, our payloads sent to the Batch API become a lot smaller, and less
redundant. The WeakMap should ensure that our memory usage doesn't grow too
much from the extra accounting.

The keys are random, and I think the probability of a clash within a batch is
indistinguishable from zero.

Also:
* useMemo on useAuthz query pieces
* adjust the tests so that they don't depend on the hashes -- and implementation
   detail they shouldn't depend on.

packages/opa-react/src/authz-provider.tsx

@@ -22,8 +22,21 @@ type EvalQuery = {
   fromResult: ((_?: Result) => boolean) | undefined;
 };
 
-function key({ path, input }: EvalQuery): string {
-  return stringify({ path, input }); // Note: omit fromResult
+const table = new WeakMap();
+
+// key is used to index the individual Batch API batches when splitting
+// an incoming batch of EvalQuery. The same `x` is later passed to the
+// batch item resolver in the results object. We use the WeakMap to
+// ensure that the same `x` gets the same number.
+function key(x: EvalQuery): string {
+  const r = table.get(x);
+  if (r) return r;
+
+  const num = new Uint32Array(1);
+  crypto.getRandomValues(num);
+  const rand = num.toString();
+  table.set(x, rand);
+  return rand;
 }
 
 const evals = (sdk: OPAClient) =>
@@ -133,7 +146,7 @@ export default function AuthzProvider({
   defaultPath,
   defaultInput,
   defaultFromResult,
-  retry = 0, // Debugging
+  retry = 0,
   batch = false,
 }: AuthzProviderProps) {
   const batcher = useMemo(
@@ -201,53 +214,3 @@ export default function AuthzProvider({
     <AuthzContext.Provider value={context}>{children}</AuthzContext.Provider>
   );
 }
-
-// Taken from fast-json-stable-hash, MIT-licensed:
-// https://github.com/zkldi/fast-json-stable-hash/blob/31b3081e942c1ce491f9698fd0bf527847093036/index.js
-// That module was tricky to import because it's using `crypto` for hashing.
-// We only need a stable string.
-function stringify(obj: any) {
-  const type = typeof obj;
-  if (obj === undefined) return "_";
-
-  if (type === "string") {
-    return JSON.stringify(obj);
-  } else if (Array.isArray(obj)) {
-    let str = "[";
-
-    let al = obj.length - 1;
-
-    for (let i = 0; i < obj.length; i++) {
-      str += stringify(obj[i]);
-
-      if (i !== al) {
-        str += ",";
-      }
-    }
-
-    return `${str}]`;
-  } else if (type === "object" && obj !== null) {
-    let str = "{";
-    let keys = Object.keys(obj).sort();
-
-    let kl = keys.length - 1;
-
-    for (let i = 0; i < keys.length; i++) {
-      let key = keys[i];
-      str += `${JSON.stringify(key)}:${stringify(obj[key])}`;
-
-      if (i !== kl) {
-        str += ",";
-      }
-    }
-
-    return `${str}}`;
-  } else if (type === "number" || type === "boolean" || obj === null) {
-    // bool, num, null have correct auto-coercions
-    return `${obj}`;
-  } else {
-    throw new TypeError(
-      `Invalid JSON type of ${type}, value ${obj}. Can only hash JSON objects.`,
-    );
-  }
-}

packages/opa-react/src/use-authz.ts

@@ -1,4 +1,4 @@
-import { useContext } from "react";
+import { useContext, useMemo } from "react";
 import { AuthzContext } from "./authz-provider.js";
 import { type Input, type Result } from "@styra/opa";
 import merge from "lodash.merge";
@@ -33,9 +33,15 @@ export default function useAuthz(
     queryClient,
     opaClient,
   } = context;
-  const p = path ?? defaultPath;
-  const i = mergeInput(input, defaultInput);
-  const fromR = fromResult ?? defaultFromResult;
+
+  const queryKey = useMemo(
+    () => [path ?? defaultPath, mergeInput(input, defaultInput)],
+    [path, defaultPath, input, defaultInput],
+  );
+  const meta = useMemo(
+    () => ({ fromResult: fromResult ?? defaultFromResult }),
+    [fromResult, defaultFromResult],
+  );
 
   const {
     // NOTE(sr): we're ignoring 'status'
@@ -44,8 +50,8 @@ export default function useAuthz(
     isFetching: isLoading,
   } = useQuery<Result>(
     {
-      queryKey: [p, i],
-      meta: { fromResult: fromR },
+      queryKey,
+      meta,
       enabled: !!opaClient,
     },
     queryClient,

packages/opa-react/tests/use-authz.test.tsx

@@ -313,10 +313,13 @@ describe("useAuthz Hook", () => {
     const batch = true;
 
     it("works without input, without fromResult", async () => {
-      const hash = '{"input":_,"path":"path/allow"}';
+      let hash: string;
       const evaluateSpy = vi
         .spyOn(opa, "evaluateBatch")
-        .mockResolvedValue({ [hash]: false });
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          [hash] = Object.keys(inputs);
+          return Promise.resolve({ [hash]: false });
+        });
 
       const { result } = renderHook(
         () => useAuthz("path/allow"),
@@ -337,13 +340,16 @@ describe("useAuthz Hook", () => {
     });
 
     it("works without input, with fromResult", async () => {
-      const hash = '{"input":_,"path":"path/allow"}';
+      let hash: string;
       const evaluateSpy = vi
         .spyOn(opa, "evaluateBatch")
-        .mockResolvedValue({ [hash]: { foo: false } });
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          [hash] = Object.keys(inputs);
+          return Promise.resolve({ [hash]: { foo: false } });
+        });
 
       const { result } = renderHook(
-        () => useAuthz("path/allow", undefined, (x) => (x as any).foo),
+        () => useAuthz("path/allow", undefined, (x?: Result) => (x as any).foo),
         wrapper({ batch }),
       );
       await waitFor(() =>
@@ -361,7 +367,6 @@ describe("useAuthz Hook", () => {
     });
 
     it("rejects evals without path", async () => {
-      const hash = '{"input":_,"path":"path/allow"}';
       const evaluateSpy = vi.spyOn(opa, "evaluateBatch");
 
       const { result } = renderHook(
@@ -379,13 +384,16 @@ describe("useAuthz Hook", () => {
     });
 
     it("works with input, with fromResult", async () => {
-      const hash = '{"input":"foo","path":"path/allow"}';
+      let hash: string;
       const evaluateSpy = vi
         .spyOn(opa, "evaluateBatch")
-        .mockResolvedValue({ [hash]: { foo: false } });
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          [hash] = Object.keys(inputs);
+          return Promise.resolve({ [hash]: { foo: false } });
+        });
 
       const { result } = renderHook(
-        () => useAuthz("path/allow", "foo", (x) => (x as any).foo),
+        () => useAuthz("path/allow", "foo", (x?: Result) => (x as any).foo),
         wrapper({ batch }),
       );
       await waitFor(() =>
@@ -403,11 +411,23 @@ describe("useAuthz Hook", () => {
     });
 
     it("batches multiple requests with different inputs", async () => {
-      const hash1 = '{"input":"foo","path":"path/allow"}';
-      const hash2 = '{"input":"bar","path":"path/allow"}';
+      let hash1: string;
+      let hash2: string;
       const evaluateSpy = vi
         .spyOn(opa, "evaluateBatch")
-        .mockResolvedValue({ [hash1]: false, [hash2]: true });
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          const res = Object.fromEntries(
+            Object.entries(inputs).map(([k, inp]) => {
+              if (inp === "foo") {
+                hash1 = k;
+              } else {
+                hash2 = k;
+              }
+              return [k, inp != "foo"];
+            }),
+          );
+          return Promise.resolve({ [hash1]: false, [hash2]: true });
+        });
 
       const { result } = renderHook(() => {
         return {
@@ -439,17 +459,23 @@ describe("useAuthz Hook", () => {
     });
 
     it("batches multiple requests with different paths+fromResults, too", async () => {
-      const hash1 = '{"input":"foo","path":"path/foo"}';
-      const hash2 = '{"input":"bar","path":"path/bar"}';
+      let hash1: string;
+      let hash2: string;
       const evaluateSpy = vi
         .spyOn(opa, "evaluateBatch")
-        .mockResolvedValueOnce({ [hash1]: { a: false } })
-        .mockResolvedValueOnce({ [hash2]: { b: true } });
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          [hash1] = Object.keys(inputs);
+          return Promise.resolve({ [hash1]: { a: false } });
+        })
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          [hash2] = Object.keys(inputs);
+          return Promise.resolve({ [hash2]: { b: true } });
+        });
 
       const { result } = renderHook(() => {
         return {
-          first: useAuthz("path/foo", "foo", (x) => (x as any).a),
-          second: useAuthz("path/bar", "bar", (x) => (x as any).b),
+          first: useAuthz("path/foo", "foo", (x?: Result) => (x as any).a),
+          second: useAuthz("path/bar", "bar", (x?: Result) => (x as any).b),
         };
       }, wrapper({ batch }));
 
@@ -481,13 +507,16 @@ describe("useAuthz Hook", () => {
     });
 
     it("coalesces multiple requests with the same path input (disregarding fromResult)", async () => {
-      const hash = '{"input":"foo","path":"path/allow"}';
+      let hash: string;
       // NOTE(sr): Unlike the non-batching case, we're handling the application of `fromResult`
       // in the code of opa-react. So the functions that are passed are evaluated on the mocked
       // result returned here.
       const evaluateSpy = vi
         .spyOn(opa, "evaluateBatch")
-        .mockResolvedValue({ [hash]: { foo: false } });
+        .mockImplementationOnce((_path, inputs, _opts) => {
+          [hash] = Object.keys(inputs);
+          return Promise.resolve({ [hash]: { foo: false } });
+        });
 
       const { result } = renderHook(() => {
         return {