Skip to content

Commit 2c2e58e

Browse files
DonJayamanneDonJayamanne
authored
Do not await before queuing cells for execution (#15841)
1 parent b4af4bf commit 2c2e58e

File tree

8 files changed

+287
-88
lines changed

8 files changed

+287
-88
lines changed

‎src/kernels/kernelExecution.ts

+6-22
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ import { CellExecutionMessageHandlerService } from './execution/cellExecutionMes
2121
import { CellExecutionQueue } from './execution/cellExecutionQueue';
2222
import { cellOutputToVSCCellOutput, traceCellMessage } from './execution/helpers';
2323
import { executeSilently } from './helpers';
24-
import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from './telemetry/helper';
2524
import { sendKernelTelemetryEvent } from './telemetry/sendKernelTelemetryEvent';
2625
import {
2726
IKernel,
@@ -140,10 +139,6 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
140139
return;
141140
}
142141

143-
await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
144-
this.kernel.resourceUri,
145-
this.kernel.kernelConnectionMetadata
146-
);
147142
sendKernelTelemetryEvent(this.kernel.resourceUri, Telemetry.ResumeCellExecution);
148143
const sessionPromise = this.kernel.start(new DisplayOptions(false));
149144
const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, sessionPromise);
@@ -166,29 +161,24 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
166161
return;
167162
}
168163

169-
traceCellMessage(cell, `NotebookKernelExecution.executeCell, ${getDisplayPath(cell.notebook.uri)}`);
170-
await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
171-
this.kernel.resourceUri,
172-
this.kernel.kernelConnectionMetadata
173-
);
174-
const sessionPromise = this.kernel.start(new DisplayOptions(false));
175-
164+
traceCellMessage(cell, `NotebookKernelExecution.executeCell (2), ${getDisplayPath(cell.notebook.uri)}`);
176165
// If we're restarting, wait for it to finish
177-
await this.kernel.restarting;
166+
const sessionPromise = this.kernel.restarting.then(() => this.kernel.start(new DisplayOptions(false)));
178167

179-
traceCellMessage(cell, `NotebookKernelExecution.executeCell (2), ${getDisplayPath(cell.notebook.uri)}`);
168+
traceCellMessage(cell, `NotebookKernelExecution.executeCell (3), ${getDisplayPath(cell.notebook.uri)}`);
180169
const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, sessionPromise);
181170
executionQueue.queueCell(cell, codeOverride);
182171
let success = true;
183172
try {
173+
traceCellMessage(cell, `NotebookKernelExecution.executeCell (4), ${getDisplayPath(cell.notebook.uri)}`);
184174
await executionQueue.waitForCompletion(cell);
185175
} catch (ex) {
186176
success = false;
187177
throw ex;
188178
} finally {
189179
traceCellMessage(
190180
cell,
191-
`NotebookKernelExecution.executeCell completed (3), ${getDisplayPath(cell.notebook.uri)}`
181+
`NotebookKernelExecution.executeCell completed (5), ${getDisplayPath(cell.notebook.uri)}`
192182
);
193183
logger.trace(`Cell ${cell.index} executed ${success ? 'successfully' : 'with an error'}`);
194184
sendKernelTelemetryEvent(this.kernel.resourceUri, Telemetry.ExecuteCell, {
@@ -206,14 +196,8 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
206196
token: CancellationToken
207197
): AsyncGenerator<NotebookCellOutput, void, unknown> {
208198
const stopWatch = new StopWatch();
209-
await initializeInteractiveOrNotebookTelemetryBasedOnUserAction(
210-
this.kernel.resourceUri,
211-
this.kernel.kernelConnectionMetadata
212-
);
213-
const sessionPromise = this.kernel.start(new DisplayOptions(false));
214-
215199
// If we're restarting, wait for it to finish
216-
await this.kernel.restarting;
200+
const sessionPromise = this.kernel.restarting.then(() => this.kernel.start(new DisplayOptions(false)));
217201

218202
const executionQueue = this.getOrCreateCellExecutionQueue(this.notebook, sessionPromise);
219203

‎src/notebooks/controllers/kernelConnector.ts

+71-17
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,10 @@ export class KernelConnector {
171171
{
172172
kernel: Deferred<{
173173
kernel: IKernel;
174-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
174+
deadKernelAction?:
175+
| 'deadKernelNeedsToBeRestarted'
176+
| 'deadKernelWasNoRestarted'
177+
| 'deadKernelWasRestarted';
175178
}>;
176179
options: IDisplayOptions;
177180
}
@@ -181,7 +184,10 @@ export class KernelConnector {
181184
{
182185
kernel: Deferred<{
183186
kernel: IBaseKernel;
184-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
187+
deadKernelAction?:
188+
| 'deadKernelNeedsToBeRestarted'
189+
| 'deadKernelWasNoRestarted'
190+
| 'deadKernelWasRestarted';
185191
}>;
186192
options: IDisplayOptions;
187193
}
@@ -194,17 +200,25 @@ export class KernelConnector {
194200
promise:
195201
| Promise<{
196202
kernel: IBaseKernel;
197-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
203+
deadKernelAction?:
204+
| 'deadKernelNeedsToBeRestarted'
205+
| 'deadKernelWasNoRestarted'
206+
| 'deadKernelWasRestarted';
198207
}>
199208
| Promise<{
200209
kernel: IKernel;
201-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
210+
deadKernelAction?:
211+
| 'deadKernelNeedsToBeRestarted'
212+
| 'deadKernelWasNoRestarted'
213+
| 'deadKernelWasRestarted';
202214
}>,
203215
actionSource: KernelActionSource,
204216
onAction: (action: KernelAction, kernel: IBaseKernel) => void,
205-
disposables: IDisposable[]
217+
disposables: IDisposable[],
218+
kernelWasRestartedDueToDeadKernel?: { kernelWasRestartedDueToDeadKernel: Deferred<boolean> }
206219
): Promise<IKernel | IBaseKernel> {
207-
const { kernel, deadKernelAction } = await promise;
220+
const info = await promise;
221+
const { kernel, deadKernelAction } = info;
208222
// Before returning, but without disposing the kernel, double check it's still valid
209223
// If a restart didn't happen, then we can't connect. Throw an error.
210224
// Do this outside of the loop so that subsequent calls will still ask because the kernel isn't disposed
@@ -215,9 +229,17 @@ export class KernelConnector {
215229
if (deadKernelAction === 'deadKernelWasNoRestarted') {
216230
throw new KernelDeadError(kernel.kernelConnectionMetadata);
217231
} else if (deadKernelAction === 'deadKernelWasRestarted') {
218-
return kernel;
232+
throw new KernelDeadError(kernel.kernelConnectionMetadata);
233+
} else if (kernelWasRestartedDueToDeadKernel?.kernelWasRestartedDueToDeadKernel.value === true) {
234+
throw new KernelDeadError(kernel.kernelConnectionMetadata);
219235
}
236+
info.deadKernelAction = 'deadKernelWasRestarted';
220237
// Kernel is dead and we didn't prompt the user to restart it, hence re-run the code that will prompt the user for a restart.
238+
const deferred = createDeferred<boolean>();
239+
deferred.resolve(true);
240+
kernelWasRestartedDueToDeadKernel = kernelWasRestartedDueToDeadKernel || {
241+
kernelWasRestartedDueToDeadKernel: deferred
242+
};
221243
return KernelConnector.wrapKernelMethod(
222244
kernel.kernelConnectionMetadata,
223245
'start',
@@ -226,7 +248,8 @@ export class KernelConnector {
226248
notebookResource,
227249
options,
228250
disposables,
229-
onAction
251+
onAction,
252+
kernelWasRestartedDueToDeadKernel
230253
);
231254
}
232255
return kernel;
@@ -240,7 +263,8 @@ export class KernelConnector {
240263
notebookResource: NotebookResource,
241264
options: IDisplayOptions,
242265
disposables: IDisposable[],
243-
onAction: (action: KernelAction, kernel: IBaseKernel | IKernel) => void = () => noop()
266+
onAction: (action: KernelAction, kernel: IBaseKernel | IKernel) => void = () => noop(),
267+
kernelWasRestartedDueToDeadKernel?: { kernelWasRestartedDueToDeadKernel: Deferred<boolean> }
244268
): Promise<IBaseKernel | IKernel> {
245269
let currentPromise = this.getKernelInfo(notebookResource);
246270
if (!options.disableUI && currentPromise?.options.disableUI) {
@@ -276,15 +300,22 @@ export class KernelConnector {
276300
)}`
277301
);
278302

303+
kernelWasRestartedDueToDeadKernel = kernelWasRestartedDueToDeadKernel || {
304+
kernelWasRestartedDueToDeadKernel: createDeferred<boolean>()
305+
};
279306
const promise = KernelConnector.wrapKernelMethodImpl(
280307
metadata,
281308
initialContext,
282309
serviceContainer,
283310
notebookResource,
284311
options,
285312
actionSource,
286-
onAction
313+
onAction,
314+
kernelWasRestartedDueToDeadKernel
287315
);
316+
if (kernelWasRestartedDueToDeadKernel?.kernelWasRestartedDueToDeadKernel.resolved) {
317+
kernelWasRestartedDueToDeadKernel.kernelWasRestartedDueToDeadKernel.resolve(true);
318+
}
288319
const deferred = createDeferredFromPromise(promise);
289320
deferred.promise.catch(noop);
290321
// If the kernel gets disposed or we fail to create the kernel, then ensure we remove the cached result.
@@ -310,7 +341,8 @@ export class KernelConnector {
310341
deferred.promise,
311342
actionSource,
312343
onAction,
313-
disposables
344+
disposables,
345+
kernelWasRestartedDueToDeadKernel
314346
);
315347
}
316348
private static getKernelInfo(notebookResource: NotebookResource) {
@@ -322,15 +354,23 @@ export class KernelConnector {
322354
notebookResource: NotebookResource,
323355
deferred: Deferred<{
324356
kernel: IBaseKernel;
325-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted' | undefined;
357+
deadKernelAction?:
358+
| 'deadKernelNeedsToBeRestarted'
359+
| 'deadKernelWasNoRestarted'
360+
| 'deadKernelWasRestarted'
361+
| undefined;
326362
}>,
327363
options: IDisplayOptions
328364
) {
329365
if ('notebook' in notebookResource) {
330366
KernelConnector.connectionsByNotebook.set(notebookResource.notebook, {
331367
kernel: deferred as Deferred<{
332368
kernel: IKernel;
333-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted' | undefined;
369+
deadKernelAction?:
370+
| 'deadKernelNeedsToBeRestarted'
371+
| 'deadKernelWasNoRestarted'
372+
| 'deadKernelWasRestarted'
373+
| undefined;
334374
}>,
335375
options
336376
});
@@ -342,7 +382,11 @@ export class KernelConnector {
342382
notebookResource: NotebookResource,
343383
matchingKernelPromise?: Promise<{
344384
kernel: IBaseKernel;
345-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted' | undefined;
385+
deadKernelAction?:
386+
| 'deadKernelNeedsToBeRestarted'
387+
| 'deadKernelWasNoRestarted'
388+
| 'deadKernelWasRestarted'
389+
| undefined;
346390
}>
347391
) {
348392
if (!matchingKernelPromise) {
@@ -385,10 +429,11 @@ export class KernelConnector {
385429
notebookResource: NotebookResource,
386430
options: IDisplayOptions,
387431
actionSource: KernelActionSource,
388-
onAction: (action: KernelAction, kernel: IBaseKernel) => void
432+
onAction: (action: KernelAction, kernel: IBaseKernel) => void,
433+
kernelWasRestartedDueToDeadKernel?: { kernelWasRestartedDueToDeadKernel: Deferred<boolean> }
389434
): Promise<{
390435
kernel: IBaseKernel | IKernel;
391-
deadKernelAction?: 'deadKernelWasRestarted' | 'deadKernelWasNoRestarted';
436+
deadKernelAction?: 'deadKernelNeedsToBeRestarted' | 'deadKernelWasNoRestarted';
392437
}> {
393438
const kernelProvider = serviceContainer.get<IKernelProvider>(IKernelProvider);
394439
const thirdPartyKernelProvider = serviceContainer.get<IThirdPartyKernelProvider>(IThirdPartyKernelProvider);
@@ -416,15 +461,20 @@ export class KernelConnector {
416461
resourceUri: notebookResource.resource
417462
});
418463

464+
let attemptedToRestart = false;
419465
try {
420466
// If the kernel is dead, ask the user if they want to restart.
421467
// We need to perform this check first, as its possible we'd call this method for dead kernels.
422468
// & if the kernel is dead, prompt to restart.
423469
if (initialContext !== 'restart' && isKernelDead(kernel) && !options.disableUI) {
470+
attemptedToRestart = true;
424471
const restarted = await KernelConnector.notifyAndRestartDeadKernel(kernel);
472+
if (restarted && kernelWasRestartedDueToDeadKernel) {
473+
kernelWasRestartedDueToDeadKernel.kernelWasRestartedDueToDeadKernel.resolve(true);
474+
}
425475
return {
426476
kernel,
427-
deadKernelAction: restarted ? 'deadKernelWasRestarted' : 'deadKernelWasNoRestarted'
477+
deadKernelAction: restarted ? 'deadKernelNeedsToBeRestarted' : 'deadKernelWasNoRestarted'
428478
};
429479
} else {
430480
onAction(currentContext, kernel);
@@ -440,6 +490,10 @@ export class KernelConnector {
440490
}
441491
}
442492
} catch (error) {
493+
if (attemptedToRestart && kernelWasRestartedDueToDeadKernel) {
494+
kernelWasRestartedDueToDeadKernel.kernelWasRestartedDueToDeadKernel.resolve(true);
495+
}
496+
443497
if (!isCancellationError(error)) {
444498
logger.warn(
445499
`Error occurred while trying to ${currentContext} the kernel, options.disableUI=${options.disableUI}`,

‎src/notebooks/controllers/kernelConnector.unit.test.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import { Disposable } from 'vscode';
2929
import { PythonExtension } from '@vscode/python-extension';
3030
import { setPythonApi } from '../../platform/interpreter/helpers';
3131
import { resolvableInstance } from '../../test/datascience/helpers';
32+
import { noop } from '../../test/core';
3233

3334
suite('Kernel Connector', () => {
3435
const pythonConnection = PythonKernelConnectionMetadata.create({
@@ -163,7 +164,7 @@ suite('Kernel Connector', () => {
163164
new DisplayOptions(false),
164165
disposables,
165166
'jupyterExtension'
166-
);
167+
).catch(noop);
167168

168169
verify(kernel.restart()).once();
169170
verify(

‎src/notebooks/controllers/types.ts

-1
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ export interface IVSCodeNotebookController extends IDisposable {
3838
postMessage(message: any, editor?: vscode.NotebookEditor): Thenable<boolean>;
3939
asWebviewUri(localResource: vscode.Uri): vscode.Uri;
4040
isAssociatedWithDocument(notebook: vscode.NotebookDocument): boolean;
41-
restoreConnection(notebook: vscode.NotebookDocument): Promise<void>;
4241
updateConnection(connection: KernelConnectionMetadata): void;
4342
setPendingCellAddition(notebook: vscode.NotebookDocument, promise: Promise<void>): void;
4443
}

0 commit comments

Comments
 (0)