aboutsummaryrefslogtreecommitdiffstats
path: root/actioncable/lib/action_cable/subscription_adapter/postgresql.rb
blob: afa99355e89edd2a3888b2574b1adab04fdbf8cf (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
require 'thread'

begin
  require 'pg'
rescue Gem::LoadError => e
  raise Gem::LoadError, "You are trying to use the PostgreSQL ActionCable adapter, but do not have the proper gems installed. Add `gem 'pg'` to your Gemfile (and ensure its version is at the minimum required by ActionCable)."
end

module ActionCable
  module SubscriptionAdapter
    class PostgreSQL < Base # :nodoc:
      # The storage instance used for broadcasting. Not intended for direct user use.
      def broadcast(channel, payload)
        with_connection do |pg_conn|
          pg_conn.exec("NOTIFY #{pg_conn.escape_identifier(channel)}, '#{pg_conn.escape_string(payload)}'")
        end
      end

      def subscribe(channel, callback, success_callback = nil)
        listener.subscribe_to(channel, callback, success_callback)
      end

      def unsubscribe(channel, callback)
        listener.unsubscribe_from(channel, callback)
      end

      def with_connection(&block) # :nodoc:
        ActiveRecord::Base.connection_pool.with_connection do |ar_conn|
          pg_conn = ar_conn.raw_connection

          unless pg_conn.is_a?(PG::Connection)
            raise 'ActiveRecord database must be Postgres in order to use the Postgres ActionCable storage adapter'
          end

          yield pg_conn
        end
      end

      private
        def listener
          @listener ||= Listener.new(self)
        end

        class Listener
          def initialize(adapter)
            @adapter = adapter
            @subscribers = Hash.new { |h,k| h[k] = [] }
            @sync = Mutex.new
            @queue = Queue.new

            Thread.new do
              Thread.current.abort_on_exception = true
              listen
            end
          end

          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}")
                  end
                end

                pg_conn.wait_for_notify(1) do |chan, pid, message|
                  @subscribers[chan].each do |callback|
                    ::EM.next_tick { callback.call(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

              @subscribers[channel] << callback
            end
          end

          def unsubscribe_from(channel, callback)
            @sync.synchronize do
              @subscribers[channel].delete(callback)

              if @subscribers[channel].empty?
                @queue.push([:unlisten, channel])
              end
            end
          end
        end
    end
  end
end