diff --git a/src/affinity.ts b/src/affinity.ts index de00b5b..60d5c2d 100644 --- a/src/affinity.ts +++ b/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): 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): 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 { diff --git a/src/query-plan.ts b/src/query-plan.ts index ba9e205..6eaca3f 100644 --- a/src/query-plan.ts +++ b/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); +} diff --git a/test/affinity.test.ts b/test/affinity.test.ts index 1fe8eac..9181146 100644 --- a/test/affinity.test.ts +++ b/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): 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; +} diff --git a/test/query-plan.test.ts b/test/query-plan.test.ts index a36392b..27f19c6 100644 --- a/test/query-plan.test.ts +++ b/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`);