Refactor connection channel for delivering notifications

src/invidious.cr

@@ -186,10 +186,21 @@ spawn do
   end
 end
 
-notification_channels = [] of Channel(PQ::Notification)
-PG.connect_listen(PG_URL, "notifications") do |event|
-  notification_channels.each do |channel|
-    channel.send(event)
+connection_channel = Channel({Bool, Channel(PQ::Notification)}).new
+spawn do
+  connections = [] of Channel(PQ::Notification)
+
+  PG.connect_listen(PG_URL, "notifications") { |event| connections.each { |connection| connection.send(event) } }
+
+  loop do
+    action, connection = connection_channel.receive
+
+    case action
+    when true
+      connections << connection
+    when false
+      connections.delete(connection)
+    end
   end
 end
 
@@ -4469,15 +4480,7 @@ get "/api/v1/auth/notifications" do |env|
   topics = env.params.query["topics"]?.try &.split(",").uniq.first(1000)
   topics ||= [] of String
 
-  notification_channel = Channel(PQ::Notification).new
-  notification_channels << notification_channel
-
-  begin
-    create_notification_stream(env, proxies, config, Kemal.config, decrypt_function, topics, notification_channel)
-  rescue ex
-  ensure
-    notification_channels.delete(notification_channel)
-  end
+  create_notification_stream(env, proxies, config, Kemal.config, decrypt_function, topics, connection_channel)
 end
 
 post "/api/v1/auth/notifications" do |env|
@@ -4486,15 +4489,7 @@ post "/api/v1/auth/notifications" do |env|
   topics = env.params.body["topics"]?.try &.split(",").uniq.first(1000)
   topics ||= [] of String
 
-  notification_channel = Channel(PQ::Notification).new
-  notification_channels << notification_channel
-
-  begin
-    create_notification_stream(env, proxies, config, Kemal.config, decrypt_function, topics, notification_channel)
-  rescue ex
-  ensure
-    notification_channels.delete(notification_channel)
-  end
+  create_notification_stream(env, proxies, config, Kemal.config, decrypt_function, topics, connection_channel)
 end
 
 get "/api/v1/auth/preferences" do |env|

src/invidious/helpers/helpers.cr

@@ -661,7 +661,10 @@ def copy_in_chunks(input, output, chunk_size = 4096)
   end
 end
 
-def create_notification_stream(env, proxies, config, kemal_config, decrypt_function, topics, notification_channel)
+def create_notification_stream(env, proxies, config, kemal_config, decrypt_function, topics, connection_channel)
+  connection = Channel(PQ::Notification).new
+  connection_channel.send({true, connection})
+
   locale = LOCALES[env.get("preferences").as(Preferences).locale]?
 
   since = env.params.query["since"]?.try &.to_i?
@@ -669,6 +672,7 @@ def create_notification_stream(env, proxies, config, kemal_config, decrypt_funct
 
   if topics.includes? "debug"
     spawn do
+      begin
         loop do
           time_span = [0, 0, 0, 0]
           time_span[rand(4)] = rand(30) + 5
@@ -697,10 +701,13 @@ def create_notification_stream(env, proxies, config, kemal_config, decrypt_funct
 
           sleep 1.minute
         end
+      rescue ex
+      end
     end
   end
 
   spawn do
+    begin
       if since
         topics.try &.each do |topic|
           case topic
@@ -731,10 +738,12 @@ def create_notification_stream(env, proxies, config, kemal_config, decrypt_funct
         end
       end
     end
+  end
 
   spawn do
+    begin
       loop do
-      event = notification_channel.receive
+        event = connection.receive
 
         notification = JSON.parse(event.payload)
         topic = notification["topic"].as_s
@@ -763,8 +772,13 @@ def create_notification_stream(env, proxies, config, kemal_config, decrypt_funct
           id += 1
         end
       end
+    rescue ex
+    ensure
+      connection_channel.send({false, connection})
+    end
   end
 
+  begin
     # Send heartbeat
     loop do
       env.response.puts ":keepalive #{Time.now.to_unix}"
@@ -772,4 +786,8 @@ def create_notification_stream(env, proxies, config, kemal_config, decrypt_funct
       env.response.flush
       sleep (20 + rand(11)).seconds
     end
+  rescue ex
+  ensure
+    connection_channel.send({false, connection})
+  end
 end