Align BatchWriteItem affinity query plan

src/affinity.ts

@@ -119,12 +119,12 @@ export class KeyRouteAffinityPlanner {
       });
     }
 
-    const preferred = selectPreferredNode(votes);
-    if (!preferred) {
+    const preferredNodes = selectPreferredNodes(votes);
+    if (preferredNodes.length === 0) {
       return undefined;
     }
 
-    return new AlternatorQueryPlan(nodes, [], preferred, true);
+    return new AlternatorQueryPlan(nodes, [], preferredNodes, true);
   }
 
   private triggerPartitionKeyDiscovery(tableName: string): void {
@@ -286,10 +286,15 @@ function batchWriteRoutingCandidates(
       if (!isRecord(write)) {
         continue;
       }
-      if (isRecord(write.PutRequest) && isRecord(write.PutRequest.Item)) {
+      const hasPutRequest = "PutRequest" in write;
+      const hasDeleteRequest = "DeleteRequest" in write;
+      if (hasPutRequest === hasDeleteRequest) {
+        continue;
+      }
+      if (hasPutRequest && isRecord(write.PutRequest) && isRecord(write.PutRequest.Item)) {
         candidates.push({ tableName, values: write.PutRequest.Item });
       }
-      if (isRecord(write.DeleteRequest) && isRecord(write.DeleteRequest.Key)) {
+      if (hasDeleteRequest && isRecord(write.DeleteRequest) && isRecord(write.DeleteRequest.Key)) {
         candidates.push({ tableName, values: write.DeleteRequest.Key });
       }
     }
@@ -298,24 +303,14 @@ function batchWriteRoutingCandidates(
   return candidates;
 }
 
-function selectPreferredNode(votes: Map<string, { node: AlternatorNode; votes: number }>): AlternatorNode | undefined {
-  let preferred: AlternatorNode | undefined;
-  let preferredVotes = 0;
-  let tied = false;
-
-  for (const vote of votes.values()) {
-    if (vote.votes > preferredVotes) {
-      preferred = vote.node;
-      preferredVotes = vote.votes;
-      tied = false;
-      continue;
-    }
-    if (vote.votes === preferredVotes) {
-      tied = true;
-    }
-  }
-
-  return preferredVotes > 0 && !tied ? preferred : undefined;
+function selectPreferredNodes(votes: Map<string, { node: AlternatorNode; votes: number }>): AlternatorNode[] {
+  return [...votes.values()]
+    .filter((vote) => vote.votes > 0)
+    .sort((left, right) => {
+      const voteOrder = right.votes - left.votes;
+      return voteOrder === 0 ? left.node.url.localeCompare(right.node.url) : voteOrder;
+    })
+    .map((vote) => vote.node);
 }
 
 function isRecord(value: unknown): value is Record<string, unknown> {

src/query-plan.ts

@@ -4,17 +4,19 @@ import { SeededRandom } from "./seeded-random.js";
 export class AlternatorQueryPlan {
   private activeNodes: AlternatorNode[];
   private quarantinedNodes: AlternatorNode[];
+  private readonly preferredNodes: AlternatorNode[];
   private readonly random: SeededRandom | undefined;
 
   constructor(
     activeNodes: readonly AlternatorNode[],
     quarantinedNodes: readonly AlternatorNode[] = [],
-    private readonly preferredNode?: AlternatorNode,
+    preferredNodes?: AlternatorNode | readonly AlternatorNode[],
     private readonly deterministicOrder = false,
     random?: SeededRandom,
     sortBeforeSelection = deterministicOrder || random !== undefined,
   ) {
     this.random = random;
+    this.preferredNodes = normalizePreferredNodes(preferredNodes);
     this.activeNodes = sortBeforeSelection
       ? sortNodes(activeNodes)
       : [...activeNodes];
@@ -32,8 +34,12 @@ export class AlternatorQueryPlan {
   }
 
   next(): AlternatorNode | undefined {
-    if (this.preferredNode) {
-      const preferred = popNode(this.activeNodes, this.preferredNode);
+    while (this.preferredNodes.length > 0) {
+      const preferredNode = this.preferredNodes.shift();
+      if (!preferredNode) {
+        continue;
+      }
+      const preferred = popNode(this.activeNodes, preferredNode);
       if (preferred) {
         return preferred;
       }
@@ -90,3 +96,21 @@ function popNode(nodes: AlternatorNode[], preferredNode: AlternatorNode): Altern
   const [node] = nodes.splice(index, 1);
   return node;
 }
+
+function normalizePreferredNodes(
+  preferredNodes: AlternatorNode | readonly AlternatorNode[] | undefined,
+): AlternatorNode[] {
+  if (!preferredNodes) {
+    return [];
+  }
+  if (isNodeList(preferredNodes)) {
+    return [...preferredNodes];
+  }
+  return [preferredNodes];
+}
+
+function isNodeList(
+  preferredNodes: AlternatorNode | readonly AlternatorNode[],
+): preferredNodes is readonly AlternatorNode[] {
+  return Array.isArray(preferredNodes);
+}

test/affinity.test.ts

@@ -1,7 +1,9 @@
 import { BatchWriteItemCommand, PutItemCommand } from "@aws-sdk/client-dynamodb";
 import { describe, expect, it } from "vitest";
-import { hashAttributeValue } from "../src/affinity.js";
+import { hashAttributeValue, KeyRouteAffinityPlanner } from "../src/affinity.js";
 import { AlternatorDynamoDBClient } from "../src/index.js";
+import { firstNodeWithSeed } from "../src/query-plan.js";
+import type { AlternatorLogger, AlternatorNode } from "../src/types.js";
 import { commandRequests, RecordingHandler } from "./helpers.js";
 
 describe("key route affinity", () => {
@@ -82,4 +84,145 @@ describe("key route affinity", () => {
 
     expect(["node-a", "node-b", "node-c"]).toContain(commandRequests(handler)[0]?.hostname);
   });
+
+  it("orders BatchWrite voted nodes before zero-vote nodes", () => {
+    const nodes = testNodes(["node-c", "node-a", "node-b"]);
+    const target = nodeByHost(nodes, "node-b");
+    const other = nodeByHost(nodes, "node-c");
+    const targetKeys = partitionKeyValuesForNode(nodes, target, "target", 2);
+    const otherKey = partitionKeyValuesForNode(nodes, other, "other", 1)[0];
+    const planner = keyRouteAffinityPlanner({ users: "id" });
+
+    const plan = planner.queryPlanForInput(
+      {
+        RequestItems: {
+          users: [
+            { PutRequest: { Item: { id: { S: targetKeys[0] } } } },
+            { DeleteRequest: { Key: { id: { S: otherKey } } } },
+            { PutRequest: { Item: { id: { S: targetKeys[1] } } } },
+          ],
+        },
+      },
+      nodes,
+      "BatchWriteItemCommand",
+    );
+
+    expect(plan).toBeDefined();
+    expect(takeHosts(plan!, nodes.length)).toEqual(["node-b", "node-c", "node-a"]);
+  });
+
+  it("uses node URL order for tied BatchWrite votes", () => {
+    const nodes = testNodes(["node-c", "node-b", "node-a"]);
+    const left = nodeByHost(nodes, "node-a");
+    const right = nodeByHost(nodes, "node-b");
+    const leftKey = partitionKeyValuesForNode(nodes, left, "left", 1)[0];
+    const rightKey = partitionKeyValuesForNode(nodes, right, "right", 1)[0];
+    const planner = keyRouteAffinityPlanner({ users: "id" });
+
+    const plan = planner.queryPlanForInput(
+      {
+        RequestItems: {
+          users: [
+            { PutRequest: { Item: { id: { S: rightKey } } } },
+            { DeleteRequest: { Key: { id: { S: leftKey } } } },
+          ],
+        },
+      },
+      nodes,
+      "BatchWriteItemCommand",
+    );
+
+    expect(plan).toBeDefined();
+    expect(takeHosts(plan!, nodes.length)).toEqual(["node-a", "node-b", "node-c"]);
+  });
+
+  it("ignores malformed BatchWrite union writes", () => {
+    const nodes = testNodes(["node-a", "node-b", "node-c"]);
+    const valid = nodeByHost(nodes, "node-a");
+    const invalid = nodeByHost(nodes, "node-b");
+    const validKey = partitionKeyValuesForNode(nodes, valid, "valid", 1)[0];
+    const invalidKeys = partitionKeyValuesForNode(nodes, invalid, "invalid", 2);
+    const planner = keyRouteAffinityPlanner({ users: "id" });
+
+    const plan = planner.queryPlanForInput(
+      {
+        RequestItems: {
+          users: [
+            { PutRequest: { Item: { id: { S: validKey } } } },
+            {
+              PutRequest: { Item: { id: { S: invalidKeys[0] } } },
+              DeleteRequest: { Key: { id: { S: invalidKeys[1] } } },
+            },
+          ],
+        },
+      },
+      nodes,
+      "BatchWriteItemCommand",
+    );
+
+    expect(plan).toBeDefined();
+    expect(plan!.next()?.host).toBe("node-a");
+  });
 });
+
+function keyRouteAffinityPlanner(partitionKeys: Record<string, string>): KeyRouteAffinityPlanner {
+  const logger: AlternatorLogger = {};
+  return new KeyRouteAffinityPlanner(
+    {
+      enabled: true,
+      mode: "any-write",
+      partitionKeys: new Map(Object.entries(partitionKeys)),
+      autoDiscoverPartitionKeys: false,
+    },
+    () => undefined,
+    logger,
+  );
+}
+
+function testNodes(hosts: readonly string[], port = 8000): AlternatorNode[] {
+  return hosts.map((host) => ({
+    host,
+    scheme: "http",
+    port,
+    url: `http://${host}:${port}`,
+  }));
+}
+
+function nodeByHost(nodes: readonly AlternatorNode[], host: string): AlternatorNode {
+  const node = nodes.find((candidate) => candidate.host === host);
+  if (!node) {
+    throw new Error(`missing test node ${host}`);
+  }
+  return node;
+}
+
+function partitionKeyValuesForNode(
+  nodes: readonly AlternatorNode[],
+  target: AlternatorNode,
+  prefix: string,
+  count: number,
+): string[] {
+  const values: string[] = [];
+  for (let index = 0; index < 10_000 && values.length < count; index += 1) {
+    const value = `${prefix}-${target.host}-${index}`;
+    const node = firstNodeWithSeed(nodes, hashAttributeValue({ S: value }));
+    if (node?.url === target.url) {
+      values.push(value);
+    }
+  }
+
+  expect(values).toHaveLength(count);
+  return values;
+}
+
+function takeHosts(plan: { next(): AlternatorNode | undefined }, count: number): string[] {
+  const hosts: string[] = [];
+  for (let index = 0; index < count; index += 1) {
+    const node = plan.next();
+    if (!node) {
+      break;
+    }
+    hosts.push(node.host);
+  }
+  return hosts;
+}

test/query-plan.test.ts

@@ -30,6 +30,18 @@ describe("AlternatorQueryPlan", () => {
     expect(plan.next()).toBeUndefined();
   });
 
+  it("tries preferred nodes in order, then remaining nodes in sorted order", () => {
+    const nodes = testNodes(["node-c", "node-a", "node-d", "node-b"]);
+    const preferred = [
+      nodes.find((node) => node.host === "node-d"),
+      nodes.find((node) => node.host === "node-b"),
+    ].filter((node): node is AlternatorNode => node !== undefined);
+    const plan = new AlternatorQueryPlan(nodes, [], preferred, true);
+
+    expect(takeHosts(plan, nodes.length)).toEqual(["node-d", "node-b", "node-a", "node-c"]);
+    expect(plan.next()).toBeUndefined();
+  });
+
   it("matches seeded raw query plan vectors", () => {
     const hosts = Array.from({ length: 10 }, (_, index) => `node${index + 1}.example.com:8043`);