From 128fabb3cd59afd470985f1bb10316e0593377f5 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Thu, 5 Oct 2023 10:30:40 +0100
Subject: [PATCH] Lint

---
 permchain/pubsub.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/permchain/pubsub.py b/permchain/pubsub.py
index cd1c1115b..8688467b5 100644
--- a/permchain/pubsub.py
+++ b/permchain/pubsub.py
@@ -216,6 +216,7 @@ class PubSub(Runnable[Any, Any], ABC):
                     for chunk in self.connection.observe(topic_prefix):
                         yield chunk
                         if chunk["topic"] == OUTPUT_TOPIC:
+                            # All expected output has been received, close
                             self.connection.disconnect(topic_prefix)
                             break
                 else:
@@ -225,6 +226,10 @@ class PubSub(Runnable[Any, Any], ABC):
                 while inflight:
                     inflight.pop().cancel()
 
+                # Remove namespace from inflight set
+                with self.lock:
+                    self.inflight_namespaces.remove(topic_prefix)
+
                 # Raise exceptions if any
                 if exceptions:
                     raise exceptions[0]