Fix too frequent queries from frontend and too long timeout on backend

jupyterlab_kernel_usage/__init__.py

@@ -1,4 +1,3 @@
-
 import json
 from pathlib import Path
 

jupyterlab_kernel_usage/handlers.py

@@ -1,20 +1,19 @@
-import asyncio
 import ipykernel
 import json
 import tornado
 import zmq
 
-from functools import partial
-
 from jupyter_server.base.handlers import APIHandler
-from jupyter_server.utils import url_path_join, ensure_async
+from jupyter_server.utils import url_path_join
 from jupyter_client.jsonutil import date_default
 
 from packaging import version
 
 
 USAGE_IS_SUPPORTED = version.parse("6.9.0") <= version.parse(ipykernel.__version__)
 
+MAX_RETRIES = 3
+
 
 class RouteHandler(APIHandler):
 
@@ -38,9 +37,8 @@ async def get(self, matched_part=None, *args, **kwargs):
         poller = zmq.Poller()
         control_socket = control_channel.socket
         poller.register(control_socket, zmq.POLLIN)
-        while True:
-            timeout = 100
-            timeout_ms = int(1000 * timeout)
+        for i in range(1, MAX_RETRIES + 1):
+            timeout_ms = 1000 * i
             events = dict(poller.poll(timeout_ms))
             if not events:
                 self.write(json.dumps({}))

src/index.ts

@@ -5,7 +5,6 @@ import {
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { LabIcon } from '@jupyterlab/ui-components';
 import { ICommandPalette } from '@jupyterlab/apputils';
-import { ILauncher } from '@jupyterlab/launcher';
 import { KernelUsagePanel } from './panel';
 import tachometer from '../style/tachometer.svg';
 
@@ -19,24 +18,25 @@ namespace CommandIDs {
 const plugin: JupyterFrontEndPlugin<void> = {
   id: 'kernelusage:plugin',
   requires: [ICommandPalette, INotebookTracker],
-  optional: [ILauncher],
   autoStart: true,
   activate: (
     app: JupyterFrontEnd,
     palette: ICommandPalette,
-    notebookTracker: INotebookTracker,
-    launcher: ILauncher | null
+    notebookTracker: INotebookTracker
   ) => {
     const { commands, shell } = app;
     const category = 'Kernel Resource';
 
-    async function createPanel(): Promise<KernelUsagePanel> {
-      const panel = new KernelUsagePanel({
-        widgetAdded: notebookTracker.widgetAdded,
-        currentNotebookChanged: notebookTracker.currentChanged
-      });
-      shell.add(panel, 'right', { rank: 200 });
-      return panel;
+    let panel: KernelUsagePanel | null = null;
+
+    function createPanel() {
+      if (!panel || panel.isDisposed) {
+        panel = new KernelUsagePanel({
+          widgetAdded: notebookTracker.widgetAdded,
+          currentNotebookChanged: notebookTracker.currentChanged
+        });
+        shell.add(panel, 'right', { rank: 200 });
+      }
     }
 
     commands.addCommand(CommandIDs.getKernelUsage, {

src/widget.tsx

@@ -1,4 +1,4 @@
-import React, { useState } from 'react';
+import React, { useState, useEffect } from 'react';
 import { ISignal } from '@lumino/signaling';
 import { ReactWidget, ISessionContext } from '@jupyterlab/apputils';
 import { IChangedArgs } from '@jupyterlab/coreutils';
@@ -32,6 +32,19 @@ type Usage = {
 
 const POLL_INTERVAL_SEC = 5;
 
+type KernelChangeCallback = (
+  _sender: ISessionContext,
+  args: IChangedArgs<
+    Kernel.IKernelConnection | null,
+    Kernel.IKernelConnection | null,
+    'kernel'
+  >
+) => void;
+let kernelChangeCallback: {
+  callback: KernelChangeCallback;
+  panel: NotebookPanel;
+} | null = null;
+
 const KernelUsage = (props: {
   widgetAdded: ISignal<INotebookTracker, NotebookPanel | null>;
   currentNotebookChanged: ISignal<INotebookTracker, NotebookPanel | null>;
@@ -44,61 +57,87 @@ const KernelUsage = (props: {
 
   useInterval(async () => {
     if (kernelId && panel.isVisible) {
-      requestUsage(kernelId).then(usage => setUsage(usage));
+      requestUsage(kernelId)
+        .then(usage => setUsage(usage))
+        .catch(() => {
+          console.warn(`Request failed for ${kernelId}. Kernel restarting?`);
+        });
     }
   }, POLL_INTERVAL_SEC * 1000);
 
-  const requestUsage = (kid: string) =>
-    requestAPI<any>(`get_usage/${kid}`).then(data => {
+  const requestUsage = (kid: string) => {
+    return requestAPI<any>(`get_usage/${kid}`).then(data => {
       const usage: Usage = {
         ...data.content,
         kernelId: kid,
         timestamp: new Date()
       };
       return usage;
     });
+  };
 
-  props.currentNotebookChanged.connect(
-    (sender: INotebookTracker, panel: NotebookPanel | null) => {
-      panel?.sessionContext.kernelChanged.connect(
-        (
-          _sender: ISessionContext,
-          args: IChangedArgs<
-            Kernel.IKernelConnection | null,
-            Kernel.IKernelConnection | null,
-            'kernel'
-          >
-        ) => {
-          /*
-          const oldKernelId = args.oldValue?.id;
-          if (oldKernelId) {
-            const poll = kernelPools.get(oldKernelId);
-            poll?.poll.dispose();
-            kernelPools.delete(oldKernelId);
-          }
-          */
-          const newKernelId = args.newValue?.id;
-          if (newKernelId) {
-            setKernelId(newKernelId);
-            const path = panel?.sessionContext.session?.model.path;
-            setPath(path);
-            requestUsage(newKernelId).then(usage => setUsage(usage));
-          }
+  useEffect(() => {
+    const createKernelChangeCallback = (panel: NotebookPanel) => {
+      return (
+        _sender: ISessionContext,
+        args: IChangedArgs<
+          Kernel.IKernelConnection | null,
+          Kernel.IKernelConnection | null,
+          'kernel'
+        >
+      ) => {
+        const newKernelId = args.newValue?.id;
+        if (newKernelId) {
+          setKernelId(newKernelId);
+          const path = panel?.sessionContext.session?.model.path;
+          setPath(path);
+          requestUsage(newKernelId).then(usage => setUsage(usage));
+        } else {
+          // Kernel was disposed
+          setKernelId(newKernelId);
         }
-      );
-      if (panel?.sessionContext.session?.id !== kernelId) {
-        if (panel?.sessionContext.session?.kernel?.id) {
-          const kernelId = panel?.sessionContext.session?.kernel?.id;
-          if (kernelId) {
-            setKernelId(kernelId);
-            const path = panel?.sessionContext.session?.model.path;
-            setPath(path);
-            requestUsage(kernelId).then(usage => setUsage(usage));
-          }
+      };
+    };
+
+    const notebookChangeCallback = (
+      sender: INotebookTracker,
+      panel: NotebookPanel | null
+    ) => {
+      if (panel === null) {
+        // Ideally we would switch to a new "select a notebook to get kernel
+        // usage" screen instead of showing outdated info.
+        return;
+      }
+      if (kernelChangeCallback) {
+        kernelChangeCallback.panel.sessionContext.kernelChanged.disconnect(
+          kernelChangeCallback.callback
+        );
+      }
+      kernelChangeCallback = {
+        callback: createKernelChangeCallback(panel),
+        panel
+      };
+      panel.sessionContext.kernelChanged.connect(kernelChangeCallback.callback);
+
+      if (panel.sessionContext.session?.kernel?.id !== kernelId) {
+        const kernelId = panel.sessionContext.session?.kernel?.id;
+        if (kernelId) {
+          setKernelId(kernelId);
+          const path = panel.sessionContext.session?.model.path;
+          setPath(path);
+          requestUsage(kernelId).then(usage => setUsage(usage));
         }
       }
-    }
-  );
+    };
+    props.currentNotebookChanged.connect(notebookChangeCallback);
+    return () => {
+      props.currentNotebookChanged.disconnect(notebookChangeCallback);
+      // In the ideal world we would disconnect kernelChangeCallback from
+      // last panel here, but this can lead to a race condition. Instead,
+      // we make sure there is ever only one callback active by holding
+      // it in a global state.
+    };
+  }, [kernelId]);
 
   if (kernelId) {
     if (usage) {