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,
|
||||
getPendingConversationKeys,
|
||||
getPendingMessagesByConversation,
|
||||
revertStuckProcessingMessages,
|
||||
updateMessageAIAnalysis,
|
||||
updateMessagesAIAnalysisBulk,
|
||||
} from "../message-capture/messageStore.js";
|
||||
@@ -661,17 +662,24 @@ function enqueueIndividualFallbacks(messages: MessageRecord[]): void {
|
||||
async function processBatch(
|
||||
conversationKey: string,
|
||||
messages: MessageRecord[],
|
||||
processingStartedAt: number,
|
||||
): 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;
|
||||
if (Date.now() < cooldownUntil) {
|
||||
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||
conversationProcessing.delete(conversationKey);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
activeRequests++;
|
||||
let shouldScheduleNext = false;
|
||||
const processingStartedAt = Date.now();
|
||||
conversationProcessing.set(conversationKey, processingStartedAt);
|
||||
try {
|
||||
const result = (await workerPool.run({
|
||||
type: "batch",
|
||||
@@ -936,16 +944,33 @@ function scheduleConversationAnalysis(conversationKey: string): void {
|
||||
const timer = setTimeout(() => {
|
||||
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.
|
||||
getPendingMessagesByConversation(
|
||||
conversationKey,
|
||||
config.AI_ANALYSIS_MAX_BATCH_SIZE,
|
||||
)
|
||||
.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);
|
||||
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.
|
||||
// 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) => {
|
||||
if (conversationProcessing.get(conversationKey) === processingStartedAt) {
|
||||
conversationProcessing.delete(conversationKey);
|
||||
}
|
||||
logger.error(
|
||||
{
|
||||
conversationKey,
|
||||
@@ -1070,6 +1098,13 @@ export function startPendingAIAnalysisWorker(
|
||||
if (!config.AI_ANALYSIS_ENABLED) return;
|
||||
|
||||
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.
|
||||
Promise.all([
|
||||
getPendingConversationKeys(500),
|
||||
|
||||
@@ -651,8 +651,8 @@ export async function getPendingMessagesByConversation(
|
||||
|
||||
// conversationKey is either thread_id or channel_id
|
||||
// Query both to safely handle the key
|
||||
const rows = await database
|
||||
.select()
|
||||
const sq = database
|
||||
.select({ id: messagesTable.id })
|
||||
.from(messagesTable)
|
||||
.where(
|
||||
and(
|
||||
@@ -665,7 +665,14 @@ export async function getPendingMessagesByConversation(
|
||||
),
|
||||
)
|
||||
.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[];
|
||||
} catch (error) {
|
||||
@@ -856,8 +863,8 @@ export async function getIncompleteMessagesByConversation(
|
||||
): Promise<MessageRecord[]> {
|
||||
try {
|
||||
const database = db();
|
||||
const rows = await database
|
||||
.select()
|
||||
const sq = database
|
||||
.select({ id: messagesTable.id })
|
||||
.from(messagesTable)
|
||||
.where(
|
||||
and(
|
||||
@@ -874,7 +881,14 @@ export async function getIncompleteMessagesByConversation(
|
||||
),
|
||||
)
|
||||
.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[];
|
||||
} catch (error) {
|
||||
@@ -1247,3 +1261,38 @@ export async function getExpiredMessages(
|
||||
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 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 AIRecommendedAction =
|
||||
| "none"
|
||||
|
||||
@@ -62,7 +62,7 @@ export const pgMessagesTable = pgTable(
|
||||
.default("text"),
|
||||
metadata: pgText("metadata"),
|
||||
ai_status: pgText("ai_status", {
|
||||
enum: ["pending", "clean", "warn", "flagged", "error"],
|
||||
enum: ["pending", "processing", "clean", "warn", "flagged", "error"],
|
||||
})
|
||||
.notNull()
|
||||
.default("pending"),
|
||||
|
||||
Reference in New Issue
Block a user