Skip to content

Commit 3170714

Browse files
authored
fix(db): wave 1 performance follow-ups from re-audit (#8984)
* fix: wave 1 performance follow-ups from re-audit - chat refresh coalescer cancels a slow first load at most once per chain, and a reload that joins the load in flight reads again after it - Better Auth user delete.after hook discards reporting usage prefixes, covering the admin plugin's remove-user path - a reporting prefix epoch bump that finds Redis configured but not ready is logged like a failed bump; payer bump logs carry the payer type - unique probes inside a write transaction also turn off parallel plans, restored with the seqscan flag - integration cleanup in finally, drop the stale GIN pending-list flush, fold the keyword-ranking seed check into beforeAll - TSDoc corrections in storage tracking and log stats queries * fix: restore the two-step unauthenticated keyword ranking, join cached slow reads, drop the probe parallel toggle
1 parent d8f91a3 commit 3170714

9 files changed

Lines changed: 181 additions & 78 deletions

File tree

‎apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-resource-event.test.ts‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { flushMacrotask } from '@sim/testing/helpers/async'
12
import { createDeferred } from '@sim/testing/helpers/deferred'
23
import { QueryClient, QueryObserver } from '@tanstack/react-query'
34
import { beforeEach, describe, expect, it, vi } from 'vitest'
@@ -330,6 +331,52 @@ describe('recovered resource events', () => {
330331
}
331332
}
332333
)
334+
335+
it.each([
336+
{ name: 'a loaded transcript', cached: true },
337+
{ name: 'a transcript still on its first load', cached: false },
338+
])(
339+
'cancels a stuck read of $name at most once, so a slow replacement finishes and is read again',
340+
async ({ cached }) => {
341+
const queryClient = new QueryClient()
342+
const { server, reads, unsubscribe } = mountedChat(queryClient, {
343+
cached,
344+
hangFirstRead: true,
345+
})
346+
const slowReadReleased = createDeferred<void>()
347+
reads.mockImplementationOnce(async () => {
348+
const resources = [...server.stored]
349+
await slowReadReleased.promise
350+
return { id: 'chat-1', title: null, messages: [], activeStreamId: 'stream-1', resources }
351+
})
352+
const ctx = createStreamLoopContext(
353+
makeStreamLoopDeps({ queryClient, chatIdRef: ref<string | undefined>('chat-1') })
354+
)
355+
vi.useFakeTimers({ toFake: ['Date'] })
356+
try {
357+
handleResourceEvent(ctx, replayed(tables[0], 1))
358+
await vi.waitFor(() => expect(reads).toHaveBeenCalledOnce())
359+
360+
vi.setSystemTime(Date.now() + 30_000)
361+
handleResourceEvent(ctx, replayed(tables[1], 2))
362+
await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(2))
363+
364+
server.stored = tables
365+
vi.setSystemTime(Date.now() + 30_000)
366+
handleResourceEvent(ctx, replayed(tables[2], 3))
367+
await flushMacrotask()
368+
expect(reads).toHaveBeenCalledTimes(2)
369+
370+
slowReadReleased.resolve()
371+
await vi.waitFor(() => expect(reads).toHaveBeenCalledTimes(3))
372+
expect(await settledResources(queryClient)).toEqual(tables)
373+
expect(reads).toHaveBeenCalledTimes(3)
374+
} finally {
375+
vi.useRealTimers()
376+
unsubscribe()
377+
}
378+
}
379+
)
333380
})
334381

335382
describe('handleResourceEvent saved-view pins', () => {

‎apps/sim/hooks/queries/mothership-chats.ts‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,8 @@ const CHAT_DETAIL_REFRESH_HOLD_MS = 30 * 1000
118118
interface ChatDetailRefresh {
119119
startedAt: number
120120
requestedAgain: boolean
121+
/** Whether this reload replaced one past the hold; a chain cancels a slow load at most once. */
122+
replaced: boolean
121123
}
122124

123125
/** Per client, the chats with a reload in flight. */
@@ -139,16 +141,23 @@ export function refreshMothershipChatDetail(queryClient: QueryClient, chatId: st
139141
inFlight.requestedAgain = true
140142
return
141143
}
142-
const refresh: ChatDetailRefresh = { startedAt: Date.now(), requestedAgain: false }
144+
// A query with no data yet joins its fetch in flight instead of replacing it, so a stuck
145+
// one is cancelled first, but only once per chain: a load that is slow rather than hung
146+
// would otherwise restart every hold and never finish. Later reloads in the chain join the
147+
// read in flight, cached or not, and since it may predate this request, read again after.
148+
const refresh: ChatDetailRefresh = {
149+
startedAt: Date.now(),
150+
requestedAgain: inFlight?.replaced === true,
151+
replaced: inFlight !== undefined,
152+
}
143153
refreshes.set(chatId, refresh)
144154
const queryKey = mothershipChatKeys.detail(chatId)
145-
// A query with no data yet joins its fetch in flight instead of replacing it, so a stuck
146-
// one is cancelled first.
147-
const reload = inFlight
148-
? queryClient
149-
.cancelQueries({ queryKey, exact: true })
150-
.then(() => queryClient.invalidateQueries({ queryKey }))
151-
: queryClient.invalidateQueries({ queryKey })
155+
const reload =
156+
inFlight && !inFlight.replaced
157+
? queryClient
158+
.cancelQueries({ queryKey, exact: true })
159+
.then(() => queryClient.invalidateQueries({ queryKey }))
160+
: queryClient.invalidateQueries({ queryKey }, { cancelRefetch: !inFlight })
152161
void reload.finally(() => {
153162
if (refreshes.get(chatId) !== refresh) return
154163
refreshes.delete(chatId)

‎apps/sim/lib/auth/auth.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ import {
8686
releaseCheckoutAdmission,
8787
resolveCheckoutReferenceId,
8888
} from '@/lib/billing/checkout-admission'
89+
import { invalidateReportingUsagePrefixes } from '@/lib/billing/core/reporting-usage-epoch'
8990
import {
9091
getOrganizationIdForSubscriptionReference,
9192
syncSubscriptionPlan,
@@ -442,6 +443,16 @@ export const auth = betterAuth({
442443
}
443444
},
444445
},
446+
/**
447+
* The admin plugin's `/admin/remove-user` deletes the user row through Better Auth, so its
448+
* usage rows cascade away without `deleteUserAccount`. Better Auth runs this once the delete
449+
* commits, which is when the stored reporting prefixes that counted them must be discarded.
450+
*/
451+
delete: {
452+
after: async () => {
453+
await invalidateReportingUsagePrefixes()
454+
},
455+
},
445456
},
446457
account: {
447458
create: {

‎apps/sim/lib/billing/core/reporting-usage-epoch.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,14 @@ export function readyRedisClient(): Redis | null {
6363
return redis?.status === 'ready' ? redis : null
6464
}
6565

66+
/**
67+
* Throws when a configured Redis is not ready, since its stored prefixes then outlive the bump:
68+
* each caller logs that like a failed or timed-out SET. Without Redis there is nothing to discard.
69+
*/
6670
async function replaceEpoch(epochKey: string): Promise<void> {
67-
const redis = readyRedisClient()
71+
const redis = getRedisClient()
6872
if (!redis) return
73+
if (redis.status !== 'ready') throw new Error(`Redis not ready (${redis.status})`)
6974
await withinDeadline(
7075
() => redis.set(epochKey, generateShortId(), 'PX', PREFIX_EPOCH_RETAIN_MS),
7176
Date.now() + SHARED_READ_TIMEOUT_MS
@@ -76,8 +81,8 @@ async function replaceEpoch(epochKey: string): Promise<void> {
7681
* Discards every stored prefix, for a write that removes ledger rows — a user's deletion, whose
7782
* rows cascade away. Call it once that write has committed: a refresh that read the old epoch
7883
* stores its prefix under it, where nothing reads it, and one that reads the new epoch sums after
79-
* the deletion. Best effort: without a ready Redis there is no prefix to discard, and a failed
80-
* bump leaves the deleted usage counted until the stored prefixes expire.
84+
* the deletion. Best effort: a bump that fails, times out, or finds Redis not ready leaves the
85+
* deleted usage counted until the stored prefixes expire, and is logged.
8186
*/
8287
export async function invalidateReportingUsagePrefixes(): Promise<void> {
8388
try {
@@ -100,6 +105,7 @@ export function invalidatePayerPrefixesForTopUp(entity: BillingEntity, rowCreate
100105
if (rowCreatedAt.getTime() >= Date.now() - TOP_UP_INVALIDATES_AFTER_MS) return
101106
replaceEpoch(payerPrefixEpochKey(entity)).catch((error) => {
102107
logger.warn('Shared reporting usage payer prefix invalidation failed', {
108+
payerType: entity.type,
103109
error: getErrorMessage(error),
104110
})
105111
})

‎apps/sim/lib/billing/storage/tracking.ts‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -433,13 +433,14 @@ export async function applyStorageUsageDeltasInTx(
433433
/**
434434
* Locks the canonical workspace row before a file mutation that changes its storage ledger.
435435
*
436-
* Every transaction that writes `workspace_files` and charges storage takes its locks in one order:
437-
* the folder advisory lock (only for operations that take one), then this workspace row, then file
438-
* rows, then the file writes (whose search-index trigger upserts the one-per-workspace
439-
* `workspace_file_search_dispatch_queue` row), then the payer row. Call this before the first file insert or update: an insert that reaches
440-
* the queue row first and only then waits on the workspace row deadlocks with a content write
441-
* that holds the workspace row and waits on the queue row. Re-locking later in the same
442-
* transaction, as the storage increment does, never waits.
436+
* Every transaction that writes a `context = 'workspace'` row of `workspace_files` and charges
437+
* storage takes its locks in one order: the folder advisory lock (only for operations that take
438+
* one), then this workspace row, then the file rows and their writes (whose search-index trigger
439+
* upserts the one-per-workspace `workspace_file_search_dispatch_queue` row), then the payer row.
440+
* Call this before the first file insert or update: an insert that reaches the queue row first
441+
* and only then waits on the workspace row deadlocks with a content write that holds the
442+
* workspace row and waits on the queue row. Re-locking later in the same transaction, as the
443+
* storage increment does, never waits.
443444
*/
444445
export async function lockWorkspaceStorageForMutationInTx(
445446
tx: DbOrTx,

‎apps/sim/lib/knowledge/__integration__/keyword-ranking.integration.ts‎

Lines changed: 28 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
/**
2-
* The keyword leg of a search with no signed-in reader behind it, against real PostgreSQL: ranking
3-
* every base in one statement answers what ranking each base (or all of them) in the old two-step
4-
* pass did, tie groups compared as sets, and admits nothing the reader may not see.
2+
* The keyword leg of a search with no signed-in reader behind it, against real PostgreSQL: it
3+
* answers what ranking each base (or all of them) in the two-step pass does, tie groups compared
4+
* as sets, and admits nothing the reader may not see.
55
*
66
* Every chunk's vector points away from the query vector, so the vector leg of a hybrid search
77
* finds nothing under its distance threshold and the fused list is the keyword ranking alone.
@@ -11,7 +11,7 @@ import { db } from '@sim/db'
1111
import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema'
1212
import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
1313
import { generateId } from '@sim/utils/id'
14-
import { and, desc, eq, inArray, type SQL, sql } from 'drizzle-orm'
14+
import { and, eq, inArray, type SQL, sql } from 'drizzle-orm'
1515
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
1616
import {
1717
createKnowledgeAclFixtureIds,
@@ -34,7 +34,7 @@ const queryVector = {
3434
} as const
3535
const AWAY_FROM_QUERY = [-1, ...Array<number>(DIMENSIONS - 1).fill(0)]
3636

37-
/** More bases than the old pass ranked in one statement, so it took its per-base branch. */
37+
/** More bases than the pass ranks in one statement, so it takes its per-base branch. */
3838
const EXTRA_BASE_COUNT = 5
3939
const DOCUMENTS_PER_BASE = 3
4040
const CHUNKS_PER_DOCUMENT = 2
@@ -78,11 +78,8 @@ interface RankedKeywordRow {
7878
keywordRank: number
7979
}
8080

81-
/**
82-
* The unauthenticated keyword ranking as it stood before it moved into one statement: the oracle
83-
* each case is compared against.
84-
*/
85-
async function rankBeforeSingleStatement(input: {
81+
/** The two-step unauthenticated keyword ranking: the oracle each case is compared against. */
82+
async function rankTwoStep(input: {
8683
knowledgeBaseIds: string[]
8784
topK: number
8885
query: string
@@ -152,7 +149,7 @@ function expectSameRanking(actual: string[], reference: RankedKeywordRow[], topK
152149
expect(cursor).toBe(actual.length)
153150
}
154151

155-
describe('unauthenticated keyword ranking in one statement', () => {
152+
describe('unauthenticated keyword ranking', () => {
156153
const ids = createKnowledgeAclFixtureIds()
157154
const extraBaseIds = Array.from({ length: EXTRA_BASE_COUNT }, () => generateId())
158155
const unsearchedBaseId = generateId()
@@ -307,6 +304,23 @@ describe('unauthenticated keyword ranking in one statement', () => {
307304

308305
await db.insert(document).values([...documents, ...hiddenDocuments])
309306
await db.insert(embedding).values(chunks)
307+
308+
// The comparisons below test tie-breaking only if scores tie across bases at several levels.
309+
const scores = await db
310+
.select({
311+
rank: sql<number>`ts_rank_cd(${embedding.contentTsv}, websearch_to_tsquery(${FTS_CONFIG}, ${QUERY}))`,
312+
bases: sql<number>`count(DISTINCT ${embedding.knowledgeBaseId})::int`,
313+
})
314+
.from(embedding)
315+
.where(
316+
and(
317+
inArray(embedding.knowledgeBaseId, allBaseIds),
318+
sql`${embedding.contentTsv} @@ websearch_to_tsquery(${FTS_CONFIG}, ${QUERY})`
319+
)
320+
)
321+
.groupBy(sql`1`)
322+
expect(scores.length).toBeGreaterThan(2)
323+
expect(scores.filter((score) => score.bases > 1).length).toBeGreaterThan(1)
310324
})
311325

312326
afterAll(async () => {
@@ -334,7 +348,7 @@ describe('unauthenticated keyword ranking in one statement', () => {
334348
queryVector,
335349
})
336350
expect(retrieval).toEqual({ status: 'complete', timedOutLegs: [] })
337-
const reference = await rankBeforeSingleStatement({
351+
const reference = await rankTwoStep({
338352
knowledgeBaseIds: input.knowledgeBaseIds,
339353
topK: hybridLegTopK(input.topK),
340354
query: QUERY,
@@ -352,28 +366,9 @@ describe('unauthenticated keyword ranking in one statement', () => {
352366
return { actual, reference }
353367
}
354368

355-
it('seeds scores tied across bases at several distinct levels', async () => {
356-
const scores = await db
357-
.select({
358-
rank: sql<number>`ts_rank_cd(${embedding.contentTsv}, websearch_to_tsquery(${FTS_CONFIG}, ${QUERY}))`,
359-
bases: sql<number>`count(DISTINCT ${embedding.knowledgeBaseId})::int`,
360-
})
361-
.from(embedding)
362-
.where(
363-
and(
364-
inArray(embedding.knowledgeBaseId, allBaseIds),
365-
sql`${embedding.contentTsv} @@ websearch_to_tsquery(${FTS_CONFIG}, ${QUERY})`
366-
)
367-
)
368-
.groupBy(sql`1`)
369-
.orderBy(desc(sql`1`))
370-
expect(scores.length).toBeGreaterThan(2)
371-
expect(scores.filter((score) => score.bases > 1).length).toBeGreaterThan(1)
372-
})
373-
374369
describe.each([
375-
{ label: 'every base (old per-base branch)', baseCount: allBaseIds.length },
376-
{ label: 'two bases (old single-statement branch)', baseCount: 2 },
370+
{ label: 'every base (per-base branch)', baseCount: allBaseIds.length },
371+
{ label: 'two bases (single-statement branch)', baseCount: 2 },
377372
])('$label', ({ baseCount }) => {
378373
const knowledgeBaseIds = allBaseIds.slice(0, baseCount)
379374

0 commit comments

Comments
 (0)