improve ipywidgets performance by synchronizing only output widgets state changes

Author: sdissegna-maystreet

src/client/datascience/ipywidgets/ipyWidgetMessageDispatcher.ts

@@ -54,13 +54,23 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
     private totalWaitTime: number = 0;
     private totalWaitedMessages: number = 0;
     private hookCount: number = 0;
-    private fullHandleMessage?: { id: string; promise: Deferred<void> };
     /**
-     * This will be true if user has executed something that has resulted in the use of ipywidgets.
-     * We make this determinination based on whether we see messages coming from backend kernel of a specific shape.
-     * E.g. if it contains ipywidget mime type, then widgets are in use.
+     * The Output widget's model can set up or tear down a kernel message hook on state change.
+     * We need to wait until the kernel message hook has been connected before it's safe to send
+     * more messages to the UI kernel.
+     *
+     * To do this we:
+     * - Keep track of the id of all the Output widget models in the outputWidgetIds instance variable.
+     *   We add/remove these ids by inspecting messages in onKernelSocketMessage.
+     * - When a state update message is sent to one of these widgets, we synchronize with the UI and
+     *   stop sending messages until we receive a reply indicating that the state change has been fully handled.
+     *   We keep track of the message we're waiting for in the fullHandleMessage instance variable.
+     *   We start waiting for the state change to finish processing in onKernelSocketMessage,
+     *   and we stop waiting in iopubMessageHandled.
      */
-    private isUsingIPyWidgets?: boolean;
+    private outputWidgetIds = new Set<string>();
+    private fullHandleMessage?: { id: string; promise: Deferred<void> };
+    private isUsingIPyWidgets = false;
     private readonly deserialize: (data: string | ArrayBuffer) => KernelMessage.IMessage<KernelMessage.MessageType>;
 
     constructor(private readonly kernelProvider: IKernelProvider, public readonly document: NotebookDocument) {
@@ -251,14 +261,12 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
     // fully handled on both the UI and extension side before we process the next message incoming
     private messageNeedsFullHandle(message: any) {
         // We only get a handled callback for iopub messages, so this channel must be iopub
-        if (message.channel === 'iopub') {
-            if (message.header?.msg_type === 'comm_msg') {
-                // IOPub comm messages need to be fully handled
-                return true;
-            }
-        }
-
-        return false;
+        return (
+            message.channel === 'iopub' &&
+            message.header?.msg_type === 'comm_msg' &&
+            message.content?.data?.method === 'update' &&
+            this.outputWidgetIds.has(message.content?.comm_id)
+        );
     }
 
     // Callback from the UI kernel when an iopubMessage has been fully handled
@@ -272,46 +280,11 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
     }
     private async onKernelSocketMessage(data: WebSocketData): Promise<void> {
         // Hooks expect serialized data as this normally comes from a WebSocket
-        let message: undefined | KernelMessage.ICommOpenMsg; // = this.deserialize(data as any) as any;
-        if (!this.isUsingIPyWidgets) {
-            // Lets deserialize only if we know we have a potential case
-            // where this message contains some data we're interested in.
-            let mustDeserialize = false;
-            if (typeof data === 'string') {
-                mustDeserialize = data.includes(WIDGET_MIMETYPE) || data.includes(Identifiers.DefaultCommTarget);
-            } else {
-                // Array buffers (non-plain text data) must be deserialized.
-                mustDeserialize = true;
-            }
-            if (!message && mustDeserialize) {
-                message = this.deserialize(data as any) as any;
-            }
-
-            // Check for hints that would indicate whether ipywidgest are used in outputs.
-            if (
-                message &&
-                message.content &&
-                message.content.data &&
-                (message.content.data[WIDGET_MIMETYPE] || message.content.target_name === Identifiers.DefaultCommTarget)
-            ) {
-                this.isUsingIPyWidgets = true;
-            }
-        }
 
         const msgUuid = uuid();
         const promise = createDeferred<void>();
         this.waitingMessageIds.set(msgUuid, { startTime: Date.now(), resultPromise: promise });
 
-        // Check if we need to fully handle this message on UI and Extension side before we move to the next
-        if (this.isUsingIPyWidgets) {
-            if (!message) {
-                message = this.deserialize(data as any) as any;
-            }
-            if (this.messageNeedsFullHandle(message)) {
-                this.fullHandleMessage = { id: message!.header.msg_id, promise: createDeferred<void>() };
-            }
-        }
-
         if (typeof data === 'string') {
             this.raisePostMessage(IPyWidgetMessages.IPyWidgets_msg, { id: msgUuid, data });
         } else {
@@ -321,21 +294,42 @@ export class IPyWidgetMessageDispatcher implements IIPyWidgetMessageDispatcher {
             });
         }
 
-        // There are three handling states that we have for messages here
-        // 1. If we have not detected ipywidget usage at all, we just forward messages to the kernel
-        // 2. If we have detected ipywidget usage. We wait on our message to be received, but not
-        //      possibly processed yet by the UI kernel. This make sure our ordering is in sync
-        // 3. For iopub comm messages we wait for them to be fully handled by the UI kernel
-        //      and the Extension kernel as they may be required to do things like
-        //      register message hooks on both sides before we process the nextExtension message
-
-        // If there are no ipywidgets thusfar in the notebook, then no need to synchronize messages.
-        if (this.isUsingIPyWidgets) {
-            await promise.promise;
-
-            // Comm specific iopub messages we need to wait until they are full handled
-            // by both the UI and extension side before we move forward
-            if (this.fullHandleMessage) {
+        // Lets deserialize only if we know we have a potential case
+        // where this message contains some data we're interested in.
+        const mustDeserialize =
+            typeof data !== 'string' ||
+            data.includes(WIDGET_MIMETYPE) ||
+            data.includes(Identifiers.DefaultCommTarget) ||
+            data.includes('comm_open') ||
+            data.includes('comm_close') ||
+            data.includes('comm_msg');
+        if (mustDeserialize) {
+            const message = this.deserialize(data as any) as any;
+
+            // Check for hints that would indicate whether ipywidgest are used in outputs.
+            if (
+                message &&
+                message.content &&
+                message.content.data &&
+                (message.content.data[WIDGET_MIMETYPE] || message.content.target_name === Identifiers.DefaultCommTarget)
+            ) {
+                this.isUsingIPyWidgets = true;
+            }
+
+            const isIPYWidgetOutputModelOpen =
+                message.header?.msg_type === 'comm_open' &&
+                message.content?.data?.state?._model_module === '@jupyter-widgets/output' &&
+                message.content?.data?.state?._model_name === 'OutputModel';
+            const isIPYWidgetOutputModelClose =
+                message.header?.msg_type === 'comm_close' && this.outputWidgetIds.has(message.content?.comm_id);
+
+            if (isIPYWidgetOutputModelOpen) {
+                this.outputWidgetIds.add(message.content.comm_id);
+            } else if (isIPYWidgetOutputModelClose) {
+                this.outputWidgetIds.delete(message.content.comm_id);
+            } else if (this.messageNeedsFullHandle(message)) {
+                this.fullHandleMessage = { id: message.header.msg_id, promise: createDeferred<void>() };
+                await promise.promise;
                 await this.fullHandleMessage.promise.promise;
                 this.fullHandleMessage = undefined;
             }