fix(code-suggestions): detect and cleanup detached streams (b68ab533) · Commits · GitLab.org / gitlab-vscode-extension

src/common/language_server/completion_stream.test.ts

0 → 100644

+28 −0

Original line number	Diff line number	Diff line
		import { BaseLanguageClient } from 'vscode-languageclient';
		import * as vscode from 'vscode';
		import { createFakePartial } from '../test_utils/create_fake_partial';
		import { CompletionStream } from './completion_stream';
		import { createStreamIterator } from './create_stream_iterator';

		jest.mock('./create_stream_iterator', () => ({
		createStreamIterator: jest.fn(),
		}));

		const mockCleanupFn = jest.fn();

		describe('CompletionStream', () => {
		it('binds the CompletionStream instance to the cleanup function ', () => {
		const stream = new CompletionStream(
		createFakePartial<BaseLanguageClient>({}),
		createFakePartial<vscode.TextDocument>({}),
		createFakePartial<vscode.Position>({}),
		'stream-1',
		mockCleanupFn,
		);

		const cleanupFn = (createStreamIterator as jest.Mock).mock.calls[0][3];
		cleanupFn();

		expect(mockCleanupFn).toHaveBeenCalledWith(stream);
		});
		});

src/common/language_server/completion_stream.ts

+2 −0

Original line number	Diff line number	Diff line
		@@ -23,6 +23,7 @@ export class CompletionStream {
		document: vscode.TextDocument,
		position: vscode.Position,
		streamId: string,
		onCancelDetached: (stream: CompletionStream) => void,
		) {
		this.#id = streamId;
		this.#client = client;
		@@ -31,6 +32,7 @@ export class CompletionStream {
		this.#client,
		this.#id,
		this.#cancellationTokenSource.token,
		() => onCancelDetached(this),
		);
		log.debug(`Listening to stream ${this.#id}`);
		}

src/common/language_server/create_stream_iterator.test.ts

0 → 100644

+109 −0

Original line number	Diff line number	Diff line
		import { BaseLanguageClient } from 'vscode-languageclient';
		import { CancellationToken } from 'vscode';
		import { createStreamIterator, STREAM_QUEUE_TIMEOUT_MS } from './create_stream_iterator';
		import { createFakePartial } from '../test_utils/create_fake_partial';

		const mockCleanupFn = jest.fn();
		const mockOnNotification = jest.fn();
		const mockSendNotification = jest.fn();
		jest.useFakeTimers();

		const streamId = 'stream-1';

		const createIterator = () =>
		createStreamIterator(
		createFakePartial<BaseLanguageClient>({
		onNotification: mockOnNotification,
		sendNotification: mockSendNotification,
		}),
		streamId,
		createFakePartial<CancellationToken>({
		onCancellationRequested: () => ({
		dispose: () => {},
		}),
		}),
		mockCleanupFn,
		);

		describe('Create Stream Iterator', () => {
		afterEach(() => {
		jest.resetAllMocks();
		jest.clearAllTimers();
		});

		it('creates a stream iterator', () => {
		const iterator = createIterator();
		expect(iterator.next).toBeDefined();
		});

		describe('stream cleanup', () => {
		const notificationData = { id: streamId, completion: 'hello', done: false };
		let sendNotification: (_args: typeof notificationData) => void \| undefined;

		beforeEach(() => {
		mockOnNotification.mockImplementation((_, _callback) => {
		sendNotification = _callback;
		return { dispose: () => {} };
		});
		});

		it('does not run cleanup when Part is requested before data is added to the queue', async () => {
		const iterator = createIterator();

		// Request the next Part
		const next = iterator.next();
		// Send new data to the iterator queue
		sendNotification(notificationData);
		await next;
		// Ensure the cleanup delay is complete
		jest.advanceTimersByTime(STREAM_QUEUE_TIMEOUT_MS);

		expect(mockCleanupFn).not.toHaveBeenCalled();
		expect(mockSendNotification).not.toHaveBeenCalled();
		});

		it('does not run cleanup when Part is requested within delay timeout', async () => {
		const iterator = createIterator();

		// Send new data to the iterator queue
		sendNotification(notificationData);
		jest.advanceTimersByTime(100);
		// Request the next Part
		await iterator.next();
		// Ensure the cleanup delay is complete
		jest.advanceTimersByTime(STREAM_QUEUE_TIMEOUT_MS);

		expect(mockCleanupFn).not.toHaveBeenCalled();
		expect(mockSendNotification).not.toHaveBeenCalled();
		});

		it('calls the cleanup function when no Part is requested', () => {
		createIterator();

		// Send new data to the iterator queue
		sendNotification(notificationData);
		expect(mockCleanupFn).not.toHaveBeenCalled();
		// Ensure the cleanup delay is complete
		jest.advanceTimersByTime(STREAM_QUEUE_TIMEOUT_MS);
		// Note the lack of requests to the iterator for the next Part

		expect(mockCleanupFn).toHaveBeenCalled();
		expect(mockSendNotification).toHaveBeenCalled();
		});

		it('calls the cleanup function if a Part is requested too late', async () => {
		const iterator = createIterator();

		// Send new data to the iterator queue
		sendNotification(notificationData);
		expect(mockCleanupFn).not.toHaveBeenCalled();
		// Ensure the cleanup delay is complete
		jest.advanceTimersByTime(STREAM_QUEUE_TIMEOUT_MS);
		// Request the next Part
		await iterator.next();

		expect(mockCleanupFn).toHaveBeenCalled();
		expect(mockSendNotification).toHaveBeenCalled();
		});
		});
		});

src/common/language_server/create_stream_iterator.ts

+50 −5

Original line number	Diff line number	Diff line
		@@ -9,17 +9,63 @@ interface CompletionPart {
		done: boolean;
		}

		/*
		* Length of time to wait before timing-out a stream.
		*
		* Designates the maximum allowed amount of time in milliseconds from:
		* a notification arriving from the LS with a CompletionPart, to:
		* the client requesting that part from the queue.
		*
		* See https://gitlab.com/gitlab-org/gitlab-vscode-extension/-/issues/1190 for more info
		*/
		export const STREAM_QUEUE_TIMEOUT_MS = 200;

		export const createStreamIterator = (
		client: BaseLanguageClient,
		streamId: string,
		cancellationToken?: CancellationToken,
		cancellationToken: CancellationToken,
		onCancelDetached: () => void,
		): AsyncIterator<{ completion: string }, { completion: string; canceled?: boolean }> => {
		let resolveWait: ((val: CompletionPart) => void) \| null = null;
		let lastCompletion = '';
		let detachedStreamCheck: ReturnType<typeof setTimeout> \| undefined;

		const queue: CompletionPart[] = [];
		const subscriptions: Disposable[] = [];

		const onComplete = () => {
		subscriptions.forEach(x => x.dispose());
		};

		const resetCleanupCountdown = () => {
		if (detachedStreamCheck) {
		clearTimeout(detachedStreamCheck);
		}
		detachedStreamCheck = undefined;
		};

		const initiateCleanupCountdown = () => {
		if (detachedStreamCheck) {
		// Prevent concurrent countdowns
		return;
		}
		log.debug(`Possible detached stream, initiating check`);
		detachedStreamCheck = setTimeout(() => {
		// Will only run if not cancelled by resetCleanupCountdown
		log.debug(`Detached stream detected, performing cleanup`);

		// Cancel the Language Server stream
		// eslint-disable-next-line @typescript-eslint/no-floating-promises
		client.sendNotification(CancelStreaming, { id: streamId });

		// Unset the loading icon
		onCancelDetached();

		// Dispose of stream iterator listeners
		onComplete();
		}, STREAM_QUEUE_TIMEOUT_MS);
		};

		const waitForNext = (): Promise<CompletionPart> => {
		if (queue.length) {
		return Promise.resolve(queue.shift()!);
		@@ -30,10 +76,6 @@ export const createStreamIterator = (
		});
		};

		const onComplete = () => {
		subscriptions.forEach(x => x.dispose());
		};

		const onNext = (part: CompletionPart): void => {
		lastCompletion = part.completion;

		@@ -42,6 +84,7 @@ export const createStreamIterator = (
		resolveWait = null;
		} else {
		queue.push(part);
		initiateCleanupCountdown();
		}
		};

		@@ -75,6 +118,8 @@ export const createStreamIterator = (

		return {
		async next() {
		resetCleanupCountdown();

		const { completion, done, canceled = false } = await waitForNext();

		if (canceled) {

src/common/language_server/language_client_middleware.ts

+3 −1

Original line number	Diff line number	Diff line
		@@ -104,7 +104,9 @@ export class LanguageClientMiddleware implements Middleware {
		return;
		}

		const stream = new CompletionStream(this.#client, document, position, streamId);
		const stream = new CompletionStream(this.#client, document, position, streamId, _stream =>
		this.#stateManager.setLoadingResource(_stream, false),
		);

		try {
		this.#stateManager.setLoadingResource(stream, true);