fix(#337): Replace blocking KEYS command with SCAN in Valkey client
- Use SCAN with cursor for non-blocking iteration - Prevents Redis DoS under high key counts - Same API, safer implementation Refs #337 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -12,7 +12,7 @@ const mockRedisInstance = {
|
||||
on: vi.fn(),
|
||||
quit: vi.fn(),
|
||||
duplicate: vi.fn(),
|
||||
keys: vi.fn(),
|
||||
scan: vi.fn(),
|
||||
};
|
||||
|
||||
// Mock ioredis
|
||||
@@ -153,15 +153,25 @@ describe("ValkeyClient", () => {
|
||||
);
|
||||
});
|
||||
|
||||
it("should list all task states", async () => {
|
||||
mockRedis.keys.mockResolvedValue(["orchestrator:task:task-1", "orchestrator:task:task-2"]);
|
||||
it("should list all task states using SCAN", async () => {
|
||||
// SCAN returns [cursor, keys] - cursor "0" means complete
|
||||
mockRedis.scan.mockResolvedValue([
|
||||
"0",
|
||||
["orchestrator:task:task-1", "orchestrator:task:task-2"],
|
||||
]);
|
||||
mockRedis.get
|
||||
.mockResolvedValueOnce(JSON.stringify({ ...mockTaskState, taskId: "task-1" }))
|
||||
.mockResolvedValueOnce(JSON.stringify({ ...mockTaskState, taskId: "task-2" }));
|
||||
|
||||
const result = await client.listTasks();
|
||||
|
||||
expect(mockRedis.keys).toHaveBeenCalledWith("orchestrator:task:*");
|
||||
expect(mockRedis.scan).toHaveBeenCalledWith(
|
||||
"0",
|
||||
"MATCH",
|
||||
"orchestrator:task:*",
|
||||
"COUNT",
|
||||
100
|
||||
);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0].taskId).toBe("task-1");
|
||||
expect(result[1].taskId).toBe("task-2");
|
||||
@@ -251,10 +261,11 @@ describe("ValkeyClient", () => {
|
||||
);
|
||||
});
|
||||
|
||||
it("should list all agent states", async () => {
|
||||
mockRedis.keys.mockResolvedValue([
|
||||
"orchestrator:agent:agent-1",
|
||||
"orchestrator:agent:agent-2",
|
||||
it("should list all agent states using SCAN", async () => {
|
||||
// SCAN returns [cursor, keys] - cursor "0" means complete
|
||||
mockRedis.scan.mockResolvedValue([
|
||||
"0",
|
||||
["orchestrator:agent:agent-1", "orchestrator:agent:agent-2"],
|
||||
]);
|
||||
mockRedis.get
|
||||
.mockResolvedValueOnce(JSON.stringify({ ...mockAgentState, agentId: "agent-1" }))
|
||||
@@ -262,7 +273,13 @@ describe("ValkeyClient", () => {
|
||||
|
||||
const result = await client.listAgents();
|
||||
|
||||
expect(mockRedis.keys).toHaveBeenCalledWith("orchestrator:agent:*");
|
||||
expect(mockRedis.scan).toHaveBeenCalledWith(
|
||||
"0",
|
||||
"MATCH",
|
||||
"orchestrator:agent:*",
|
||||
"COUNT",
|
||||
100
|
||||
);
|
||||
expect(result).toHaveLength(2);
|
||||
expect(result[0].agentId).toBe("agent-1");
|
||||
expect(result[1].agentId).toBe("agent-2");
|
||||
@@ -462,7 +479,10 @@ describe("ValkeyClient", () => {
|
||||
});
|
||||
|
||||
it("should filter out null values in listTasks", async () => {
|
||||
mockRedis.keys.mockResolvedValue(["orchestrator:task:task-1", "orchestrator:task:task-2"]);
|
||||
mockRedis.scan.mockResolvedValue([
|
||||
"0",
|
||||
["orchestrator:task:task-1", "orchestrator:task:task-2"],
|
||||
]);
|
||||
mockRedis.get
|
||||
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-1", status: "pending" }))
|
||||
.mockResolvedValueOnce(null); // Simulate deleted task
|
||||
@@ -474,9 +494,9 @@ describe("ValkeyClient", () => {
|
||||
});
|
||||
|
||||
it("should filter out null values in listAgents", async () => {
|
||||
mockRedis.keys.mockResolvedValue([
|
||||
"orchestrator:agent:agent-1",
|
||||
"orchestrator:agent:agent-2",
|
||||
mockRedis.scan.mockResolvedValue([
|
||||
"0",
|
||||
["orchestrator:agent:agent-1", "orchestrator:agent:agent-2"],
|
||||
]);
|
||||
mockRedis.get
|
||||
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-1", status: "running" }))
|
||||
@@ -488,4 +508,69 @@ describe("ValkeyClient", () => {
|
||||
expect(result[0].agentId).toBe("agent-1");
|
||||
});
|
||||
});
|
||||
|
||||
describe("SCAN-based iteration (large key sets)", () => {
|
||||
it("should handle multiple SCAN iterations for tasks", async () => {
|
||||
// Simulate SCAN returning multiple batches with cursor pagination
|
||||
mockRedis.scan
|
||||
.mockResolvedValueOnce(["42", ["orchestrator:task:task-1", "orchestrator:task:task-2"]]) // First batch, cursor 42
|
||||
.mockResolvedValueOnce(["0", ["orchestrator:task:task-3"]]); // Second batch, cursor 0 = done
|
||||
|
||||
mockRedis.get
|
||||
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-1", status: "pending" }))
|
||||
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-2", status: "pending" }))
|
||||
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-3", status: "pending" }));
|
||||
|
||||
const result = await client.listTasks();
|
||||
|
||||
expect(mockRedis.scan).toHaveBeenCalledTimes(2);
|
||||
expect(mockRedis.scan).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
"0",
|
||||
"MATCH",
|
||||
"orchestrator:task:*",
|
||||
"COUNT",
|
||||
100
|
||||
);
|
||||
expect(mockRedis.scan).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
"42",
|
||||
"MATCH",
|
||||
"orchestrator:task:*",
|
||||
"COUNT",
|
||||
100
|
||||
);
|
||||
expect(result).toHaveLength(3);
|
||||
expect(result.map((t) => t.taskId)).toEqual(["task-1", "task-2", "task-3"]);
|
||||
});
|
||||
|
||||
it("should handle multiple SCAN iterations for agents", async () => {
|
||||
// Simulate SCAN returning multiple batches with cursor pagination
|
||||
mockRedis.scan
|
||||
.mockResolvedValueOnce(["99", ["orchestrator:agent:agent-1", "orchestrator:agent:agent-2"]]) // First batch
|
||||
.mockResolvedValueOnce(["50", ["orchestrator:agent:agent-3"]]) // Second batch
|
||||
.mockResolvedValueOnce(["0", ["orchestrator:agent:agent-4"]]); // Third batch, done
|
||||
|
||||
mockRedis.get
|
||||
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-1", status: "running" }))
|
||||
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-2", status: "running" }))
|
||||
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-3", status: "running" }))
|
||||
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-4", status: "running" }));
|
||||
|
||||
const result = await client.listAgents();
|
||||
|
||||
expect(mockRedis.scan).toHaveBeenCalledTimes(3);
|
||||
expect(result).toHaveLength(4);
|
||||
expect(result.map((a) => a.agentId)).toEqual(["agent-1", "agent-2", "agent-3", "agent-4"]);
|
||||
});
|
||||
|
||||
it("should handle empty result from SCAN", async () => {
|
||||
mockRedis.scan.mockResolvedValue(["0", []]);
|
||||
|
||||
const result = await client.listTasks();
|
||||
|
||||
expect(mockRedis.scan).toHaveBeenCalledTimes(1);
|
||||
expect(result).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user