Stop all channels to allow ZMQContext to properly cleanup between cal…

…ls (#183)
jupyter-server · Mar 21, 2023 · 1058a86 · 1058a86
1 parent 215950e
commit 1058a86
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/jupyter_resource_usage/api.py b/jupyter_resource_usage/api.py
@@ -122,14 +122,13 @@ async def get(self, matched_part=None, *args, **kwargs):
         timeout_ms = 10_000
         events = dict(await poller.poll(timeout_ms))
         if control_socket not in events:
-            self.write(
-                json.dumps(
-                    {
-                        "content": {"reason": "timeout", "timeout_ms": timeout_ms},
-                        "kernel_id": kernel_id,
-                    }
-                )
+            out = json.dumps(
+                {
+                    "content": {"reason": "timeout", "timeout_ms": timeout_ms},
+                    "kernel_id": kernel_id,
+                }
             )
+
         else:
             res = client.control_channel.get_msg(timeout=0)
             if isawaitable(res):
@@ -138,4 +137,6 @@ async def get(self, matched_part=None, *args, **kwargs):
                 res = await res
             if res:
                 res["kernel_id"] = kernel_id
-            self.write(json.dumps(res, default=date_default))
+            out = json.dumps(res, default=date_default)
+        client.stop_channels()
+        self.write(out)