Moved message handling out to a separate method for clarity in eventbus.
[arvados.git] / services / api / lib / eventbus.rb
1 require 'eventmachine'
2 require 'oj'
3 require 'faye/websocket'
4 require 'record_filters'
5 require 'load_param'
6
7 # Patch in user, last_log_id and filters fields into the Faye::Websocket class.
8 module Faye
9   class WebSocket
10     attr_accessor :user
11     attr_accessor :last_log_id
12     attr_accessor :filters
13   end
14 end
15
16 # Store the filters supplied by the user that will be applied to the logs table
17 # to determine which events to return to the listener.
18 class Filter
19   include LoadParam
20
21   attr_accessor :filters
22
23   def initialize p
24     @params = p
25     load_filters_param
26   end
27
28   def params
29     @params
30   end
31 end
32
33 # Manages websocket connections, accepts subscription messages and publishes
34 # log table events.
35 class EventBus
36   include CurrentApiClient
37   include RecordFilters
38
39   # used in RecordFilters
40   def model_class
41     Log
42   end
43
44   # Initialize EventBus.  Takes no parameters.
45   def initialize
46     @channel = EventMachine::Channel.new
47     @mtx = Mutex.new
48     @bgthread = false
49   end
50
51   # Push out any pending events to the connection +ws+
52   # +id+  the id of the most recent row in the log table, may be nil
53   def push_events ws, id = nil
54       begin
55         # Must have at least one filter set up to receive events
56         if ws.filters.length > 0
57           # Start with log rows readable by user, sorted in ascending order
58           logs = Log.readable_by(ws.user).order("id asc")
59
60           if ws.last_log_id
61             # Client is only interested in log rows that are newer than the
62             # last log row seen by the client.
63             logs = logs.where("logs.id > ?", ws.last_log_id)
64           elsif id
65             # No last log id, so only look at the most recently changed row
66             logs = logs.where("logs.id = ?", id.to_i)
67           else
68             return
69           end
70
71           # Now process filters provided by client
72           cond_out = []
73           param_out = []
74           ws.filters.each do |filter|
75             ft = record_filters filter.filters, Log.table_name
76             cond_out += ft[:cond_out]
77             param_out += ft[:param_out]
78           end
79
80           # Add filters to query
81           if cond_out.any?
82             logs = logs.where(cond_out.join(' OR '), *param_out)
83           end
84
85           # Finally execute query and actually send the matching log rows
86           logs.each do |l|
87             ws.send(l.as_api_response.to_json)
88             ws.last_log_id = l.id
89           end
90         elsif id
91           # No filters set up, so just record the sequence number
92           ws.last_log_id = id.to_i
93         end
94       rescue Exception => e
95         puts "Error publishing event: #{$!}"
96         puts "Backtrace:\n\t#{e.backtrace.join("\n\t")}"
97         ws.send ({status: 500, message: 'error'}.to_json)
98         ws.close
99       end
100   end
101
102   # Handle inbound subscribe or unsubscribe message.
103   def handle_message ws, event
104     begin
105       # Parse event data as JSON
106       p = (Oj.load event.data).symbolize_keys
107
108       if p[:method] == 'subscribe'
109         # Handle subscribe event
110
111         if p[:last_log_id]
112           # Set or reset the last_log_id.  The event bus only reports events
113           # for rows that come after last_log_id.
114           ws.last_log_id = p[:last_log_id].to_i
115         end
116
117         if ws.filters.length < MAX_FILTERS
118           # Add a filter.  This gets the :filters field which is the same
119           # format as used for regular index queries.
120           ws.filters << Filter.new(p)
121           ws.send ({status: 200, message: 'subscribe ok'}.to_json)
122
123           # Send any pending events
124           push_events ws
125         else
126           ws.send ({status: 403, message: "maximum of #{MAX_FILTERS} filters allowed per connection"}.to_json)
127         end
128
129       elsif p[:method] == 'unsubscribe'
130         # Handle unsubscribe event
131
132         len = ws.filters.length
133         ws.filters.select! { |f| not ((f.filters == p[:filters]) or (f.filters.empty? and p[:filters].nil?)) }
134         if ws.filters.length < len
135           ws.send ({status: 200, message: 'unsubscribe ok'}.to_json)
136         else
137           ws.send ({status: 404, message: 'filter not found'}.to_json)
138         end
139
140       else
141         ws.send ({status: 400, message: "missing or unrecognized method"}.to_json)
142       end
143     rescue Oj::Error => e
144       ws.send ({status: 400, message: "malformed request"}.to_json)
145     rescue Exception => e
146       puts "Error handling message: #{$!}"
147       puts "Backtrace:\n\t#{e.backtrace.join("\n\t")}"
148       ws.send ({status: 500, message: 'error'}.to_json)
149       ws.close
150     end
151   end
152
153   # Constant maximum number of filters, to avoid silly huge database queries.
154   MAX_FILTERS = 16
155
156   # Called by RackSocket when a new websocket connection has been established.
157   def on_connect ws
158
159     # Disconnect if no valid API token.
160     # current_user is included from CurrentApiClient
161     if not current_user
162       ws.send ({status: 401, message: "Valid API token required"}.to_json)
163       ws.close
164       return
165     end
166
167     # Initialize our custom fields on the websocket connection object.
168     ws.user = current_user
169     ws.filters = []
170     ws.last_log_id = nil
171
172     # Subscribe to internal postgres notifications through @channel.  This will
173     # call push_events when a notification comes through.
174     sub = @channel.subscribe do |msg|
175       push_events ws, msg
176     end
177
178     # Set up callback for inbound message dispatch.
179     ws.on :message do |event|
180       handle_message ws, event
181     end
182
183     # Set up socket close callback
184     ws.on :close do |event|
185       @channel.unsubscribe sub
186       ws = nil
187     end
188
189     # Start up thread to monitor the Postgres database, if none exists already.
190     @mtx.synchronize do
191       unless @bgthread
192         @bgthread = true
193         Thread.new do
194           # from http://stackoverflow.com/questions/16405520/postgres-listen-notify-rails
195           ActiveRecord::Base.connection_pool.with_connection do |connection|
196             conn = connection.instance_variable_get(:@connection)
197             begin
198               conn.async_exec "LISTEN logs"
199               while true
200                 # wait_for_notify will block until there is a change
201                 # notification from Postgres about the logs table, then push
202                 # the notification into the EventMachine channel.  Each
203                 # websocket connection subscribes to the other end of the
204                 # channel and calls #push_events to actually dispatch the
205                 # events to the client.
206                 conn.wait_for_notify do |channel, pid, payload|
207                   @channel.push payload
208                 end
209               end
210             ensure
211               # Don't want the connection to still be listening once we return
212               # it to the pool - could result in weird behavior for the next
213               # thread to check it out.
214               conn.async_exec "UNLISTEN *"
215             end
216           end
217           @bgthread = false
218         end
219       end
220     end
221
222     # Since EventMachine is an asynchronous event based dispatcher, #on_connect
223     # does not block but instead returns immediately after having set up the
224     # websocket and notification channel callbacks.
225   end
226 end