Skip to content

Commit 5ab585c

Browse files
brendan-kellambrendan-kellam
committed
improved cancelation handling
1 parent c7ba32f commit 5ab585c

File tree

3 files changed

+79
-57
lines changed

3 files changed

+79
-57
lines changed

‎packages/web/src/app/[domain]/search/components/searchResultsPage.tsx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import { SearchQueryParams } from "@/lib/types";
1919
import { createPathWithQueryParams } from "@/lib/utils";
2020
import { InfoCircledIcon } from "@radix-ui/react-icons";
2121
import { useLocalStorage } from "@uidotdev/usehooks";
22-
import { AlertTriangleIcon, BugIcon, FilterIcon, RefreshCcwIcon } from "lucide-react";
22+
import { AlertTriangleIcon, BugIcon, FilterIcon, RefreshCwIcon } from "lucide-react";
2323
import { useRouter } from "next/navigation";
2424
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
2525
import { useHotkeys } from "react-hotkeys-hook";
@@ -291,7 +291,7 @@ const PanelGroup = ({
291291
<div className="py-1 px-2 flex flex-row items-center">
292292
{isStreaming ? (
293293
<>
294-
<RefreshCcwIcon className="h-4 w-4 animate-spin mr-2" />
294+
<RefreshCwIcon className="h-4 w-4 animate-spin mr-2" />
295295
<p className="text-sm font-medium mr-1">Searching...</p>
296296
{numMatches > 0 && (
297297
<p className="text-sm font-medium">{`Found ${numMatches} matches in ${fileMatches.length} ${fileMatches.length > 1 ? 'files' : 'file'}`}</p>
@@ -353,7 +353,7 @@ const PanelGroup = ({
353353
/>
354354
) : isStreaming ? (
355355
<div className="flex flex-col items-center justify-center h-full gap-2">
356-
<RefreshCcwIcon className="h-6 w-6 animate-spin" />
356+
<RefreshCwIcon className="h-6 w-6 animate-spin" />
357357
<p className="font-semibold text-center">Searching...</p>
358358
</div>
359359
) : (

‎packages/web/src/app/[domain]/search/useStreamedSearch.ts‎

Lines changed: 41 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import { RepositoryInfo, SearchRequest, SearchResponse, SearchResultFile } from '@/features/search/types';
44
import { useState, useCallback, useRef, useEffect } from 'react';
5+
import * as Sentry from '@sentry/nextjs';
56

67
interface CacheEntry {
78
files: SearchResultFile[];
@@ -155,54 +156,39 @@ export const useStreamedSearch = ({ query, matches, contextLines, whole, isRegex
155156

156157
// SSE messages start with "data: "
157158
const dataMatch = message.match(/^data: (.+)$/);
158-
if (!dataMatch) continue;
159+
if (!dataMatch) {
160+
continue;
161+
}
159162

160163
const data = dataMatch[1];
161164

162165
// Check for completion signal
163166
if (data === '[DONE]') {
164-
setState(prev => ({ ...prev, isStreaming: false }));
165-
return;
167+
break;
166168
}
167169

168-
try {
169-
const chunk: SearchResponse = JSON.parse(data);
170-
setState(prev => ({
171-
...prev,
172-
files: [
173-
...prev.files,
174-
...chunk.files
175-
],
176-
repoInfo: {
177-
...prev.repoInfo,
178-
...chunk.repositoryInfo.reduce((acc, repo) => {
179-
acc[repo.id] = repo;
180-
return acc;
181-
}, {} as Record<number, RepositoryInfo>),
182-
},
183-
numMatches: prev.numMatches + chunk.stats.actualMatchCount,
184-
}));
185-
} catch (parseError) {
186-
console.error('Error parsing chunk:', parseError);
187-
}
170+
const chunk: SearchResponse = JSON.parse(data);
171+
setState(prev => ({
172+
...prev,
173+
files: [
174+
...prev.files,
175+
...chunk.files
176+
],
177+
repoInfo: {
178+
...prev.repoInfo,
179+
...chunk.repositoryInfo.reduce((acc, repo) => {
180+
acc[repo.id] = repo;
181+
return acc;
182+
}, {} as Record<number, RepositoryInfo>),
183+
},
184+
numMatches: prev.numMatches + chunk.stats.actualMatchCount,
185+
}));
188186
}
189187
}
190188

191-
setState(prev => ({ ...prev, isStreaming: false }));
192-
} catch (error) {
193-
if ((error as Error).name === 'AbortError') {
194-
console.log('Stream aborted');
195-
} else {
196-
setState(prev => ({
197-
...prev,
198-
isStreaming: false,
199-
error: error as Error,
200-
}));
201-
}
202-
} finally {
203-
const endTime = performance.now();
204-
const durationMs = endTime - startTime;
189+
const durationMs = performance.now() - startTime;
205190
setState(prev => {
191+
// Cache the final results after the stream has completed.
206192
searchCache.set(cacheKey, {
207193
files: prev.files,
208194
repoInfo: prev.repoInfo,
@@ -213,14 +199,31 @@ export const useStreamedSearch = ({ query, matches, contextLines, whole, isRegex
213199
return {
214200
...prev,
215201
durationMs,
202+
isStreaming: false,
216203
}
217204
});
205+
206+
} catch (error) {
207+
if ((error as Error).name === 'AbortError') {
208+
return;
209+
}
210+
211+
console.error(error);
212+
Sentry.captureException(error);
213+
const durationMs = performance.now() - startTime;
214+
setState(prev => ({
215+
...prev,
216+
isStreaming: false,
217+
durationMs,
218+
error: error as Error,
219+
}));
218220
}
219221
}
220222

221223
search();
222224

223225
return () => {
226+
cancel();
224227
}
225228
}, [
226229
query,
@@ -229,6 +232,7 @@ export const useStreamedSearch = ({ query, matches, contextLines, whole, isRegex
229232
whole,
230233
isRegexEnabled,
231234
isCaseSensitivityEnabled,
235+
cancel,
232236
]);
233237

234238
return {

‎packages/web/src/app/api/(server)/stream_search/route.ts‎

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import { searchRequestSchema } from '@/features/search/schemas';
44
import { SearchResponse, SourceRange } from '@/features/search/types';
5+
import { SINGLE_TENANT_ORG_ID } from '@/lib/constants';
56
import { schemaValidationError, serviceErrorResponse } from '@/lib/serviceError';
67
import { prisma } from '@/prisma';
78
import type { ProtoGrpcType } from '@/proto/webserver';
@@ -12,13 +13,13 @@ import type { StreamSearchResponse__Output } from '@/proto/zoekt/webserver/v1/St
1213
import type { WebserverServiceClient } from '@/proto/zoekt/webserver/v1/WebserverService';
1314
import * as grpc from '@grpc/grpc-js';
1415
import * as protoLoader from '@grpc/proto-loader';
16+
import * as Sentry from '@sentry/nextjs';
1517
import { PrismaClient, Repo } from '@sourcebot/db';
18+
import { parser as _parser } from '@sourcebot/query-language';
1619
import { createLogger, env } from '@sourcebot/shared';
1720
import { NextRequest } from 'next/server';
1821
import * as path from 'path';
19-
import { parser as _parser } from '@sourcebot/query-language';
2022
import { transformToZoektQuery } from './transformer';
21-
import { SINGLE_TENANT_ORG_ID } from '@/lib/constants';
2223

2324
const logger = createLogger('streamSearchApi');
2425

@@ -87,8 +88,8 @@ export const POST = async (request: NextRequest) => {
8788
input: query,
8889
isCaseSensitivityEnabled,
8990
isRegexEnabled,
90-
onExpandSearchContext: async (contextName: string) => {
91-
const context = await prisma.searchContext.findUnique({
91+
onExpandSearchContext: async (contextName: string) => {
92+
const context = await prisma.searchContext.findUnique({
9293
where: {
9394
name_orgId: {
9495
name: contextName,
@@ -108,8 +109,6 @@ export const POST = async (request: NextRequest) => {
108109
},
109110
});
110111

111-
console.log(JSON.stringify(zoektQuery, null, 2));
112-
113112
const searchRequest: SearchRequest = {
114113
query: zoektQuery,
115114
opts: {
@@ -158,9 +157,19 @@ const createSSESearchStream = async (searchRequest: SearchRequest, prisma: Prism
158157
const client = createGrpcClient();
159158
let grpcStream: ReturnType<WebserverServiceClient['StreamSearch']> | null = null;
160159
let isStreamActive = true;
160+
let pendingChunks = 0;
161161

162162
return new ReadableStream({
163163
async start(controller) {
164+
const tryCloseController = () => {
165+
if (!isStreamActive && pendingChunks === 0) {
166+
controller.enqueue(new TextEncoder().encode('data: [DONE]\n\n'));
167+
controller.close();
168+
client.close();
169+
logger.debug('SSE stream closed');
170+
}
171+
};
172+
164173
try {
165174
// @todo: we should just disable tenant enforcement for now.
166175
const metadata = new grpc.Metadata();
@@ -190,12 +199,14 @@ const createSSESearchStream = async (searchRequest: SearchRequest, prisma: Prism
190199

191200
// Handle incoming data chunks
192201
grpcStream.on('data', async (chunk: StreamSearchResponse__Output) => {
193-
console.log('chunk');
194-
195202
if (!isStreamActive) {
203+
logger.debug('SSE stream closed, skipping chunk');
196204
return;
197205
}
198206

207+
// Track that we're processing a chunk
208+
pendingChunks++;
209+
199210
// grpcStream.on doesn't actually await on our handler, so we need to
200211
// explicitly pause the stream here to prevent the stream from completing
201212
// prior to our asynchronous work being completed.
@@ -352,7 +363,18 @@ const createSSESearchStream = async (searchRequest: SearchRequest, prisma: Prism
352363
} catch (error) {
353364
console.error('Error encoding chunk:', error);
354365
} finally {
366+
pendingChunks--;
355367
grpcStream?.resume();
368+
369+
// @note: we were hitting "Controller is already closed" errors when calling
370+
// `controller.enqueue` above for the last chunk. The reasoning was the event
371+
// handler for 'end' was being invoked prior to the completion of the last chunk,
372+
// resulting in the controller being closed prematurely. The workaround was to
373+
// keep track of the number of pending chunks and only close the controller
374+
// when there are no more chunks to process. We need to explicitly call
375+
// `tryCloseController` since there _seems_ to be no ordering guarantees between
376+
// the 'end' event handler and this callback.
377+
tryCloseController();
356378
}
357379
});
358380

@@ -362,17 +384,13 @@ const createSSESearchStream = async (searchRequest: SearchRequest, prisma: Prism
362384
return;
363385
}
364386
isStreamActive = false;
365-
366-
// Send completion signal
367-
controller.enqueue(new TextEncoder().encode('data: [DONE]\n\n'));
368-
controller.close();
369-
console.log('SSE stream completed');
370-
client.close();
387+
tryCloseController();
371388
});
372389

373390
// Handle errors
374391
grpcStream.on('error', (error: grpc.ServiceError) => {
375-
console.error('gRPC stream error:', error);
392+
logger.error('gRPC stream error:', error);
393+
Sentry.captureException(error);
376394

377395
if (!isStreamActive) {
378396
return;
@@ -392,7 +410,7 @@ const createSSESearchStream = async (searchRequest: SearchRequest, prisma: Prism
392410
client.close();
393411
});
394412
} catch (error) {
395-
console.error('Stream initialization error:', error);
413+
logger.error('Stream initialization error:', error);
396414

397415
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
398416
const errorData = `data: ${JSON.stringify({
@@ -405,7 +423,7 @@ const createSSESearchStream = async (searchRequest: SearchRequest, prisma: Prism
405423
}
406424
},
407425
cancel() {
408-
console.log('SSE stream cancelled by client');
426+
logger.warn('SSE stream cancelled by client');
409427
isStreamActive = false;
410428

411429
// Cancel the gRPC stream to stop receiving data

0 commit comments

Comments
 (0)