diff --git a/backend/src/app/http/websocket.clj b/backend/src/app/http/websocket.clj index 50bd64f115..ddfd50455b 100644 --- a/backend/src/app/http/websocket.clj +++ b/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 diff --git a/backend/test/backend_tests/http_websocket_test.clj b/backend/test/backend_tests/http_websocket_test.clj index 3c667b41b5..74782c1054 100644 --- a/backend/test/backend_tests/http_websocket_test.clj +++ b/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})