♻️ Index websocket connections by profile

The connection registry was a flat map from connection id to
connection data, so finding the connections of one profile meant
walking every connection. `on-connect` and `on-disconnect` also
`swap!`ed a whole-map `assoc`/`dissoc`, which is O(n) per connection
churn.

The registry now holds `:connections` and `:by-profile` in the same
atom, so one `swap!` keeps both consistent and a profile lookup is a
single map access. Both maps are updated together by
`register-connection` / `unregister-connection`; a profile whose last
connection leaves is dropped from the index instead of lingering with
an empty set.

Also extracts `close-file-subscription` out of the
`:unsubscribe-file` handler so a revocation can tear a subscription
down through exactly the same path the client uses: announce
`:leave-file`, close the relay channel, purge the bus. Closing the
channel is what stops the `:subscribe-file` go-loop, since `take!` on
a closed channel returns nil.

Adds `connections-for-profile` and `repl-get-connections-for-profile`
so an operator can inspect a profile's connections from the REPL the
same way as for a file or a team.

The new helpers take no `!` suffix: per `mem:clojure/design-rules` the
default is no `!`, and these are neither `app.db` API nor transient
helpers. Existing `!` names such as `mbus/pub!` are legacy and stay
until their own refactor.

The registry stays local to one backend instance. That is documented
on the var: reaching connections owned by another instance has to go
through the message bus, which is what the revocation watcher will do.

Closes #12155

AI-assisted-by: space-bunny-free

backend/src/app/http/websocket.clj

@@ -37,31 +37,84 @@
 ;; WEBSOCKET HOOKS
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
-(def state (atom {}))
+(def state
+  "Registry of the websocket connections held by this backend instance.
+
+  Holds `:connections` (connection id -> connection data) and
+  `:by-profile` (profile id -> set of connection ids). Both live in the
+  same atom so a single `swap!` keeps them consistent, and so the
+  profile index never has to be rebuilt by scanning every connection.
+
+  NOTE: the registry is local to one backend instance, unlike the
+  message bus, which is shared. Code that must reach connections
+  regardless of which instance owns them has to go through the bus."
+  (atom {:connections {}
+         :by-profile {}}))
+
+(defn register-connection
+  "Adds `wsp` to the registry under `id`, indexing it by its profile."
+  [id wsp]
+  (let [profile-id (::profile-id wsp)]
+    (swap! state
+           (fn [st]
+             (-> st
+                 (assoc-in [:connections id] wsp)
+                 (cond-> profile-id
+                   (update-in [:by-profile profile-id]
+                              (fnil conj #{}) id)))))))
+
+(defn unregister-connection
+  "Removes the connection `id` from the registry and from the profile
+  index."
+  [id]
+  (swap! state
+         (fn [{:keys [connections] :as st}]
+           (let [owner   (::profile-id (get connections id))
+                 indexed (get-in st [:by-profile owner])
+                 indexed (if (set? indexed) (disj indexed id) indexed)]
+             (cond-> (assoc st :connections (dissoc connections id))
+               (and owner (seq indexed))
+               (assoc-in [:by-profile owner] indexed)
+               (and owner (empty? indexed))
+               (update :by-profile dissoc owner))))))
+
+(defn get-connection
+  "Returns the connection data registered under `id`, or nil."
+  [id]
+  (get-in @state [:connections id]))
+
+(defn connections-for-profile
+  "Returns the set of connection ids currently held by `profile-id`."
+  [profile-id]
+  (get-in @state [:by-profile profile-id] #{}))
 
 ;; REPL HELPERS
 
 (defn repl-get-connections-for-file
   [file-id]
-  (->> (vals @state)
+  (->> (vals (:connections @state))
        (filter #(= file-id (-> % ::ws/state deref ::file-subscription :file-id)))
        (map ::ws/id)))
 
 (defn repl-get-connections-for-team
   [team-id]
-  (->> (vals @state)
+  (->> (vals (:connections @state))
        (filter #(= team-id (-> % ::ws/state deref ::team-subscription :team-id)))
        (map ::ws/id)))
 
+(defn repl-get-connections-for-profile
+  [profile-id]
+  (connections-for-profile profile-id))
+
 (defn repl-close-connection
   [id]
-  (when-let [{:keys [::ws/close-ch] :as wsp} (get @state id)]
+  (when-let [{:keys [::ws/close-ch]} (get-connection id)]
     (sp/put! close-ch [8899 "closed from server"])
     (sp/close! close-ch)))
 
 (defn repl-get-connection-info
   [id]
-  (when-let [wsp (get @state id)]
+  (when-let [wsp (get-connection id)]
     (let [subs (some-> wsp ::ws/state deref)]
       {:id               id
        :created-at       (::created-at wsp)
@@ -220,22 +273,35 @@
                    :profile-id profile-id}]
       (mbus/pub! (::mbus/msgbus cfg) :topic file-id :message message))))
 
-(defmethod handle-message :unsubscribe-file
-  [{:keys [::mbus/msgbus]} {:keys [::ws/id ::ws/state ::session-id ::profile-id]} {:keys [file-id] :as params}]
-  (l/trace :fn "handle-message" :event "unsubscribe-file" :file-id file-id :conn-id id)
+(defn close-file-subscription
+  "Tears down the file subscription held by the connection `wsp`, if it
+  is subscribed to `file-id`.
 
-  (let [subs    (::file-subscription @state)
-        message {:type :leave-file
-                 :file-id file-id
-                 :session-id session-id
-                 :profile-id profile-id}]
+  Announces the departure so the remaining participants drop the
+  presence of this session, closes the relay channel (which in turn
+  stops the `:subscribe-file` go-loop, because `take!` on a closed
+  channel returns nil) and removes the subscription from the bus.
 
+  Does nothing when the connection is subscribed to a different file,
+  so it is safe to call for a subscription that is already gone."
+  [{:keys [::mbus/msgbus]} {:keys [::ws/state ::session-id ::profile-id]} file-id]
+  (let [subs (::file-subscription @state)]
     (when (= (:file-id subs) file-id)
-      (mbus/pub! msgbus :topic file-id :message message)
+      (mbus/pub! msgbus
+                 :topic file-id
+                 :message {:type :leave-file
+                           :file-id file-id
+                           :session-id session-id
+                           :profile-id profile-id})
       (let [ch (:channel subs)]
         (sp/close! ch)
         (mbus/purge! msgbus [ch])))))
 
+(defmethod handle-message :unsubscribe-file
+  [cfg {:keys [::ws/id] :as wsp} {:keys [file-id] :as params}]
+  (l/trace :fn "handle-message" :event "unsubscribe-file" :file-id file-id :conn-id id)
+  (close-file-subscription cfg wsp file-id))
+
 (defmethod handle-message :keepalive
   [_ _ _]
   (l/trace :fn "handle-message" :event :keepalive))
@@ -273,7 +339,7 @@
   [{:keys [::mtx/metrics]} {:keys [::ws/id] :as wsp}]
   (let [created-at (ct/now)]
     (l/trace :fn "on-connect" :conn-id id)
-    (swap! state assoc id wsp)
+    (register-connection id wsp)
     (mtx/run! metrics
               :id :websocket-active-connections
               :inc 1)
@@ -281,7 +347,7 @@
     (assoc wsp ::ws/on-disconnect
            (fn []
              (l/trace :fn "on-disconnect" :conn-id id)
-             (swap! state dissoc id)
+             (unregister-connection id)
              (mtx/run! metrics :id :websocket-active-connections :dec 1)
              (mtx/run! metrics
                        :id :websocket-session-timing

backend/test/backend_tests/http_websocket_test.clj

@@ -198,6 +198,92 @@
     (fn [topic message]
       (mbus/pub! msgbus :topic topic :message message))))
 
+;; --- CONNECTION REGISTRY
+
+(defn- with-clean-registry
+  "Runs `f` against an empty connection registry, restoring the previous
+  contents afterwards so tests cannot leak connections into each other."
+  [f]
+  (let [saved @ws/state]
+    (reset! ws/state {:connections {} :by-profile {}})
+    (try (f)
+         (finally (reset! ws/state saved)))))
+
+(t/deftest registry-indexes-connections-by-profile
+  (let [profile-a (uuid/next)
+        profile-b (uuid/next)
+        conn-a    (uuid/next)
+        conn-b    (uuid/next)
+        conn-c    (uuid/next)]
+
+    (with-clean-registry
+      (fn []
+        (t/testing "a profile with no connections has an empty index"
+          (t/is (= #{} (ws/connections-for-profile profile-a))))
+
+        (ws/register-connection conn-a {::ws/profile-id profile-a})
+        (ws/register-connection conn-b {::ws/profile-id profile-a})
+        (ws/register-connection conn-c {::ws/profile-id profile-b})
+
+        (t/testing "connections are indexed under their profile"
+          (t/is (= #{conn-a conn-b} (ws/connections-for-profile profile-a)))
+          (t/is (= #{conn-c} (ws/connections-for-profile profile-b))))
+
+        (t/testing "connections are retrievable by id"
+          (t/is (= profile-a (::ws/profile-id (ws/get-connection conn-a))))
+          (t/is (nil? (ws/get-connection (uuid/next)))))
+
+        (t/testing "unregistering removes it from both the map and the index"
+          (ws/unregister-connection conn-b)
+          (t/is (= #{conn-a} (ws/connections-for-profile profile-a)))
+          (t/is (nil? (ws/get-connection conn-b))))
+
+        (t/testing "a profile whose last connection left is dropped from the index"
+          (ws/unregister-connection conn-a)
+          (t/is (= #{} (ws/connections-for-profile profile-a)))
+          (t/is (not (contains? (:by-profile @ws/state) profile-a)))
+          (t/is (= #{conn-c} (ws/connections-for-profile profile-b))))
+
+        (t/testing "unregistering an unknown connection is a no-op"
+          (ws/unregister-connection (uuid/next))
+          (t/is (= #{conn-c} (ws/connections-for-profile profile-b)))
+          (t/is (= conn-c (-> (:connections @ws/state) keys first))))))))
+
+(t/deftest close-file-subscription-drops-only-the-matching-file
+  (let [profile  (th/create-profile* 1 {:is-active true})
+        file     (th/create-file* 1 {:profile-id (:id profile)
+                                     :project-id (:default-project-id profile)})
+        other    (th/create-file* 2 {:profile-id (:id profile)
+                                     :project-id (:default-project-id profile)})
+        output   (sp/chan :buf (sp/dropping-buffer 64))
+        ws-state (atom {::ws/file-subscription
+                        {:file-id (:id file)
+                         :channel (sp/chan :buf (sp/dropping-buffer 64))
+                         :topic (:id file)}})
+        wsp      (assoc (make-wsp (:id profile) ws-state output)
+                        ::ws/state ws-state)]
+
+    (t/testing "closing a file the connection is not subscribed to does nothing"
+      (with-redefs [mbus/pub!  (fn [& _] (throw ::unexpected-publish))
+                    mbus/purge! (fn [& _] (throw ::unexpected-purge))]
+        (t/is (nil? (ws/close-file-subscription th/*system* wsp (:id other))))
+        (t/is (some? (::ws/file-subscription @ws-state)))))
+
+    (t/testing "closing the subscribed file purges the bus"
+      (let [purged (atom [])]
+        (with-redefs [mbus/purge! (fn [_ chans] (swap! purged conj chans))]
+          (ws/close-file-subscription th/*system* wsp (:id file))
+          (t/is (= 1 (count @purged))))))
+
+    (t/testing "the relay channel is closed so the go-loop stops"
+      (let [sub-ch (-> @ws-state ::ws/file-subscription :channel)]
+        (t/is (some? sub-ch))
+        (t/is (nil? (sp/poll! sub-ch)))
+        (t/is (true? (sp/closed? sub-ch)))))
+
+    (t/testing "closing twice is safe"
+      (t/is (nil? (ws/close-file-subscription th/*system* wsp (:id file)))))))
+
 (t/deftest revoked-member-stops-receiving-file-changes
   (let [owner   (th/create-profile* 1 {:is-active true})
         editor  (th/create-profile* 2 {:is-active true})