fix(ai-moderation): resolve race conditions and implement processing state
Introduces a `processing` state to the AI analysis lifecycle to prevent duplicate processing of the same messages. - Implements row-level locking using `FOR UPDATE SKIP LOCKED` in `messageStore.ts` to ensure atomic message acquisition. - Adds a `processing` status to the `AIStatus` type and database schema. - Fixes a TOCTOU race condition in `aiAnalyzer.ts` by synchronizing the conversation processing lock before async database operations. - Implements `revertStuckProcessingMessages` to recover messages stuck in the `processing` state due to worker crashes or timeouts. - Updates `processBatch` and scheduling logic to correctly manage and release conversation-level locks.
This commit is contained in:
@@ -14,6 +14,7 @@ import {
|
|||||||
getMessageById,
|
getMessageById,
|
||||||
getPendingConversationKeys,
|
getPendingConversationKeys,
|
||||||
getPendingMessagesByConversation,
|
getPendingMessagesByConversation,
|
||||||
|
revertStuckProcessingMessages,
|
||||||
updateMessageAIAnalysis,
|
updateMessageAIAnalysis,
|
||||||
updateMessagesAIAnalysisBulk,
|
updateMessagesAIAnalysisBulk,
|
||||||
} from "../message-capture/messageStore.js";
|
} from "../message-capture/messageStore.js";
|
||||||
@@ -661,17 +662,24 @@ function enqueueIndividualFallbacks(messages: MessageRecord[]): void {
|
|||||||
async function processBatch(
|
async function processBatch(
|
||||||
conversationKey: string,
|
conversationKey: string,
|
||||||
messages: MessageRecord[],
|
messages: MessageRecord[],
|
||||||
|
processingStartedAt: number,
|
||||||
): Promise<void> {
|
): Promise<void> {
|
||||||
if (messages.length === 0) return;
|
if (messages.length === 0) {
|
||||||
|
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||||
|
conversationProcessing.delete(conversationKey);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
const cooldownUntil = conversationErrorCooldown.get(conversationKey) ?? 0;
|
const cooldownUntil = conversationErrorCooldown.get(conversationKey) ?? 0;
|
||||||
if (Date.now() < cooldownUntil) {
|
if (Date.now() < cooldownUntil) {
|
||||||
|
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||||
|
conversationProcessing.delete(conversationKey);
|
||||||
|
}
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
activeRequests++;
|
activeRequests++;
|
||||||
let shouldScheduleNext = false;
|
let shouldScheduleNext = false;
|
||||||
const processingStartedAt = Date.now();
|
|
||||||
conversationProcessing.set(conversationKey, processingStartedAt);
|
|
||||||
try {
|
try {
|
||||||
const result = (await workerPool.run({
|
const result = (await workerPool.run({
|
||||||
type: "batch",
|
type: "batch",
|
||||||
@@ -936,16 +944,33 @@ function scheduleConversationAnalysis(conversationKey: string): void {
|
|||||||
const timer = setTimeout(() => {
|
const timer = setTimeout(() => {
|
||||||
conversationDebounceTimers.delete(conversationKey);
|
conversationDebounceTimers.delete(conversationKey);
|
||||||
|
|
||||||
|
// FIX TOCTOU: Set lock synchronously BEFORE the async DB fetch starts
|
||||||
|
if (isConversationProcessingLocked(conversationKey)) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const processingStartedAt = Date.now();
|
||||||
|
conversationProcessing.set(conversationKey, processingStartedAt);
|
||||||
|
|
||||||
// FIX #3: explicit .catch() — no async arrow function to avoid unhandled rejection.
|
// FIX #3: explicit .catch() — no async arrow function to avoid unhandled rejection.
|
||||||
getPendingMessagesByConversation(
|
getPendingMessagesByConversation(
|
||||||
conversationKey,
|
conversationKey,
|
||||||
config.AI_ANALYSIS_MAX_BATCH_SIZE,
|
config.AI_ANALYSIS_MAX_BATCH_SIZE,
|
||||||
)
|
)
|
||||||
.then(async (messages) => {
|
.then(async (messages) => {
|
||||||
if (messages.length === 0) return;
|
if (messages.length === 0) {
|
||||||
|
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||||
|
conversationProcessing.delete(conversationKey);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
const processableMessages = await skipAgeRestrictedMessages(messages);
|
const processableMessages = await skipAgeRestrictedMessages(messages);
|
||||||
if (processableMessages.length === 0) return;
|
if (processableMessages.length === 0) {
|
||||||
|
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||||
|
conversationProcessing.delete(conversationKey);
|
||||||
|
}
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
// FIX #6: trim to token budget before sending to LLM.
|
// FIX #6: trim to token budget before sending to LLM.
|
||||||
// 50 tokens overhead accounts for JSON structure + id/username fields.
|
// 50 tokens overhead accounts for JSON structure + id/username fields.
|
||||||
@@ -971,9 +996,12 @@ function scheduleConversationAnalysis(conversationKey: string): void {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
return processBatch(conversationKey, trimmed);
|
return processBatch(conversationKey, trimmed, processingStartedAt);
|
||||||
})
|
})
|
||||||
.catch((err: unknown) => {
|
.catch((err: unknown) => {
|
||||||
|
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||||
|
conversationProcessing.delete(conversationKey);
|
||||||
|
}
|
||||||
logger.error(
|
logger.error(
|
||||||
{
|
{
|
||||||
conversationKey,
|
conversationKey,
|
||||||
@@ -1070,6 +1098,13 @@ export function startPendingAIAnalysisWorker(
|
|||||||
if (!config.AI_ANALYSIS_ENABLED) return;
|
if (!config.AI_ANALYSIS_ENABLED) return;
|
||||||
|
|
||||||
setInterval(() => {
|
setInterval(() => {
|
||||||
|
revertStuckProcessingMessages(300000).catch((err: unknown) => {
|
||||||
|
logger.error(
|
||||||
|
{ error: String(err) },
|
||||||
|
"Failed to run stuck processing recovery",
|
||||||
|
);
|
||||||
|
});
|
||||||
|
|
||||||
// FIX #3 pattern: no async arrow — chain promises explicitly.
|
// FIX #3 pattern: no async arrow — chain promises explicitly.
|
||||||
Promise.all([
|
Promise.all([
|
||||||
getPendingConversationKeys(500),
|
getPendingConversationKeys(500),
|
||||||
|
|||||||
@@ -651,8 +651,8 @@ export async function getPendingMessagesByConversation(
|
|||||||
|
|
||||||
// conversationKey is either thread_id or channel_id
|
// conversationKey is either thread_id or channel_id
|
||||||
// Query both to safely handle the key
|
// Query both to safely handle the key
|
||||||
const rows = await database
|
const sq = database
|
||||||
.select()
|
.select({ id: messagesTable.id })
|
||||||
.from(messagesTable)
|
.from(messagesTable)
|
||||||
.where(
|
.where(
|
||||||
and(
|
and(
|
||||||
@@ -665,7 +665,14 @@ export async function getPendingMessagesByConversation(
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
.orderBy(asc(messagesTable.created_at))
|
.orderBy(asc(messagesTable.created_at))
|
||||||
.limit(limit);
|
.limit(limit)
|
||||||
|
.for("update", { skipLocked: true });
|
||||||
|
|
||||||
|
const rows = await database
|
||||||
|
.update(messagesTable)
|
||||||
|
.set({ ai_status: "processing", ai_analyzed_at: Date.now() })
|
||||||
|
.where(inArray(messagesTable.id, sq))
|
||||||
|
.returning();
|
||||||
|
|
||||||
return rows as MessageRecord[];
|
return rows as MessageRecord[];
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
@@ -856,8 +863,8 @@ export async function getIncompleteMessagesByConversation(
|
|||||||
): Promise<MessageRecord[]> {
|
): Promise<MessageRecord[]> {
|
||||||
try {
|
try {
|
||||||
const database = db();
|
const database = db();
|
||||||
const rows = await database
|
const sq = database
|
||||||
.select()
|
.select({ id: messagesTable.id })
|
||||||
.from(messagesTable)
|
.from(messagesTable)
|
||||||
.where(
|
.where(
|
||||||
and(
|
and(
|
||||||
@@ -874,7 +881,14 @@ export async function getIncompleteMessagesByConversation(
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
.orderBy(asc(messagesTable.created_at))
|
.orderBy(asc(messagesTable.created_at))
|
||||||
.limit(limit);
|
.limit(limit)
|
||||||
|
.for("update", { skipLocked: true });
|
||||||
|
|
||||||
|
const rows = await database
|
||||||
|
.update(messagesTable)
|
||||||
|
.set({ ai_status: "processing", ai_analyzed_at: Date.now() })
|
||||||
|
.where(inArray(messagesTable.id, sq))
|
||||||
|
.returning();
|
||||||
|
|
||||||
return rows as MessageRecord[];
|
return rows as MessageRecord[];
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
@@ -1247,3 +1261,38 @@ export async function getExpiredMessages(
|
|||||||
throw error;
|
throw error;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export async function revertStuckProcessingMessages(
|
||||||
|
timeoutMs: number = 300000,
|
||||||
|
): Promise<number> {
|
||||||
|
try {
|
||||||
|
const database = db();
|
||||||
|
const cutoffTime = Date.now() - timeoutMs;
|
||||||
|
|
||||||
|
const rows = await database
|
||||||
|
.update(messagesTable)
|
||||||
|
.set({ ai_status: "pending", ai_analyzed_at: null })
|
||||||
|
.where(
|
||||||
|
and(
|
||||||
|
eq(messagesTable.ai_status, "processing"),
|
||||||
|
sql`${messagesTable.ai_analyzed_at} < ${cutoffTime}`,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.returning({ id: messagesTable.id });
|
||||||
|
|
||||||
|
if (rows.length > 0) {
|
||||||
|
logger.warn(
|
||||||
|
{ count: rows.length, messageIds: rows.map((r) => r.id) },
|
||||||
|
"Reverted stuck processing messages to pending",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
return rows.length;
|
||||||
|
} catch (error) {
|
||||||
|
logger.error(
|
||||||
|
{ error: error instanceof Error ? error.message : String(error) },
|
||||||
|
"Failed to revert stuck processing messages",
|
||||||
|
);
|
||||||
|
return 0;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
import type fs from "node:fs";
|
import type fs from "node:fs";
|
||||||
import type prism from "prism-media";
|
import type prism from "prism-media";
|
||||||
|
|
||||||
export type AIStatus = "pending" | "clean" | "warn" | "flagged" | "error";
|
export type AIStatus = "pending" | "processing" | "clean" | "warn" | "flagged" | "error";
|
||||||
export type AISeverity = "none" | "low" | "medium" | "high" | "critical";
|
export type AISeverity = "none" | "low" | "medium" | "high" | "critical";
|
||||||
export type AIRecommendedAction =
|
export type AIRecommendedAction =
|
||||||
| "none"
|
| "none"
|
||||||
|
|||||||
@@ -62,7 +62,7 @@ export const pgMessagesTable = pgTable(
|
|||||||
.default("text"),
|
.default("text"),
|
||||||
metadata: pgText("metadata"),
|
metadata: pgText("metadata"),
|
||||||
ai_status: pgText("ai_status", {
|
ai_status: pgText("ai_status", {
|
||||||
enum: ["pending", "clean", "warn", "flagged", "error"],
|
enum: ["pending", "processing", "clean", "warn", "flagged", "error"],
|
||||||
})
|
})
|
||||||
.notNull()
|
.notNull()
|
||||||
.default("pending"),
|
.default("pending"),
|
||||||
|
|||||||
Reference in New Issue
Block a user