1 files changed, 45 insertions, 36 deletions

diff --git a/actioncable/lib/action_cable/subscription_adapter/postgresql.rb b/actioncable/lib/action_cable/subscription_adapter/postgresql.rb
index 6465663c97..66c7852f6e 100644
--- a/actioncable/lib/action_cable/subscription_adapter/postgresql.rb
+++ b/actioncable/lib/action_cable/subscription_adapter/postgresql.rb
@@ -5,6 +5,11 @@ require 'thread'
 module ActionCable
   module SubscriptionAdapter
     class PostgreSQL < Base # :nodoc:
+      def initialize(*)
+        super
+        @listener = nil
+      end
+
       def broadcast(channel, payload)
         with_connection do |pg_conn|
           pg_conn.exec("NOTIFY #{pg_conn.escape_identifier(channel)}, '#{pg_conn.escape_string(payload)}'")
@@ -12,11 +17,15 @@ module ActionCable
       end
 
       def subscribe(channel, callback, success_callback = nil)
-        listener.subscribe_to(channel, callback, success_callback)
+        listener.add_subscriber(channel, callback, success_callback)
       end
 
       def unsubscribe(channel, callback)
-        listener.unsubscribe_from(channel, callback)
+        listener.remove_subscriber(channel, callback)
+      end
+
+      def shutdown
+        listener.shutdown
       end
 
       def with_connection(&block) # :nodoc:
@@ -33,17 +42,18 @@ module ActionCable
 
       private
         def listener
-          @listener ||= Listener.new(self)
+          @listener || @server.mutex.synchronize { @listener ||= Listener.new(self, @server.event_loop) }
         end
 
-        class Listener
-          def initialize(adapter)
+        class Listener < SubscriberMap
+          def initialize(adapter, event_loop)
+            super()
+
             @adapter = adapter
-            @subscribers = Hash.new { |h,k| h[k] = [] }
-            @sync = Mutex.new
+            @event_loop = event_loop
             @queue = Queue.new
 
-            Thread.new do
+            @thread = Thread.new do
               Thread.current.abort_on_exception = true
               listen
             end
@@ -51,46 +61,45 @@ module ActionCable
 
           def listen
             @adapter.with_connection do |pg_conn|
-              loop do
-                until @queue.empty?
-                  action, channel, callback = @queue.pop(true)
-                  escaped_channel = pg_conn.escape_identifier(channel)
-
-                  if action == :listen
-                    pg_conn.exec("LISTEN #{escaped_channel}")
-                    ::EM.next_tick(&callback) if callback
-                  elsif action == :unlisten
-                    pg_conn.exec("UNLISTEN #{escaped_channel}")
+              catch :shutdown do
+                loop do
+                  until @queue.empty?
+                    action, channel, callback = @queue.pop(true)
+
+                    case action
+                    when :listen
+                      pg_conn.exec("LISTEN #{pg_conn.escape_identifier channel}")
+                      @event_loop.post(&callback) if callback
+                    when :unlisten
+                      pg_conn.exec("UNLISTEN #{pg_conn.escape_identifier channel}")
+                    when :shutdown
+                      throw :shutdown
+                    end
                   end
-                end
 
-                pg_conn.wait_for_notify(1) do |chan, pid, message|
-                  @subscribers[chan].each do |callback|
-                    ::EM.next_tick { callback.call(message) }
+                  pg_conn.wait_for_notify(1) do |chan, pid, message|
+                    broadcast(chan, message)
                   end
                 end
               end
             end
           end
 
-          def subscribe_to(channel, callback, success_callback)
-            @sync.synchronize do
-              if @subscribers[channel].empty?
-                @queue.push([:listen, channel, success_callback])
-              end
+          def shutdown
+            @queue.push([:shutdown])
+            Thread.pass while @thread.alive?
+          end
 
-              @subscribers[channel] << callback
-            end
+          def add_channel(channel, on_success)
+            @queue.push([:listen, channel, on_success])
           end
 
-          def unsubscribe_from(channel, callback)
-            @sync.synchronize do
-              @subscribers[channel].delete(callback)
+          def remove_channel(channel)
+            @queue.push([:unlisten, channel])
+          end
 
-              if @subscribers[channel].empty?
-                @queue.push([:unlisten, channel])
-              end
-            end
+          def invoke_callback(*)
+            @event_loop.post { super }
           end
         end
     end