1 # If any threads raise an unhandled exception, make them all die.
2 # We trust a supervisor like runit to restart the server in this case.
3 Thread.abort_on_exception = true
7 require 'faye/websocket'
8 require 'record_filters'
13 # Patch in user, last_log_id and filters fields into the Faye::Websocket class.
17 attr_accessor :last_log_id
18 attr_accessor :filters
19 attr_accessor :sent_ids
21 attr_accessor :frame_mtx
29 alias_method :_frame, :frame
31 def frame(data, type = nil, code = nil)
32 # Most of the sending activity will be from the thread set up in
33 # on_connect. However, there is also some automatic activity in the
34 # form of ping/pong messages, so ensure that the frame method (which
35 # actually packs and sends one complete message to the underlying
36 # socket) can only be called by one thread at a time.
37 @socket.frame_mtx.synchronize do
38 _frame(data, type, code)
45 # Store the filters supplied by the user that will be applied to the logs table
46 # to determine which events to return to the listener.
50 attr_accessor :filters
62 # Manages websocket connections, accepts subscription messages and publishes
65 include CurrentApiClient
68 # used in RecordFilters
73 # Initialize EventBus. Takes no parameters.
75 @channel = EventMachine::Channel.new
81 # Push out any pending events to the connection +ws+
82 # +notify_id+ the id of the most recent row in the log table, may be nil
84 # This accepts a websocket and a notify_id (this is the row id from Postgres
85 # LISTEN/NOTIFY, it may be nil if called from somewhere else)
87 # It queries the database for log rows that are either
88 # a) greater than ws.last_log_id, which is the last log id which was a candidate to be sent out
89 # b) if ws.last_log_id is nil, then it queries the row notify_id
91 # Regular Arvados permissions are applied using readable_by() and filters using record_filters().
92 def push_events ws, notify_id
94 # Must have at least one filter set up to receive events
95 if ws.filters.length > 0
96 # Start with log rows readable by user, sorted in ascending order
97 logs = Log.readable_by(ws.user).order("id asc")
103 if not ws.last_log_id.nil?
104 # We are catching up from some starting point.
105 cond_id = "logs.id > ?"
106 param_out << ws.last_log_id
107 elsif not notify_id.nil?
108 # Get next row being notified.
109 cond_id = "logs.id = ?"
110 param_out << notify_id
112 # No log id to start from, nothing to do, return
116 # Now build filters provided by client
117 ws.filters.each do |filter|
118 ft = record_filters filter.filters, Log
119 if ft[:cond_out].any?
120 # Join the clauses within a single subscription filter with AND
121 # so it is consistent with regular queries
122 cond_out << "(#{ft[:cond_out].join ') AND ('})"
123 param_out += ft[:param_out]
127 # Add filters to query
129 # Join subscriptions with OR
130 logs = logs.where(cond_id + " AND ((#{cond_out.join ') OR ('}))", *param_out)
132 logs = logs.where(cond_id, *param_out)
135 # Execute query and actually send the matching log rows
137 if not ws.sent_ids.include?(l.id)
138 # only send if not a duplicate
139 ws.send(l.as_api_response.to_json)
141 if not ws.last_log_id.nil?
142 # record ids only when sending "catchup" messages, not notifies
148 rescue ArgumentError => e
149 # There was some kind of user error.
150 Rails.logger.warn "Error publishing event: #{$!}"
151 ws.send ({status: 500, message: $!}.to_json)
154 Rails.logger.warn "Error publishing event: #{$!}"
155 Rails.logger.warn "Backtrace:\n\t#{e.backtrace.join("\n\t")}"
156 ws.send ({status: 500, message: $!}.to_json)
158 # These exceptions typically indicate serious server trouble:
159 # out of memory issues, database connection problems, etc. Go ahead and
160 # crash; we expect that a supervisor service like runit will restart us.
165 # Handle inbound subscribe or unsubscribe message.
166 def handle_message ws, event
169 # Parse event data as JSON
170 p = (Oj.strict_load event.data).symbolize_keys
171 filter = Filter.new(p)
172 rescue Oj::Error => e
173 ws.send ({status: 400, message: "malformed request"}.to_json)
177 if p[:method] == 'subscribe'
178 # Handle subscribe event
181 # Set or reset the last_log_id. The event bus only reports events
182 # for rows that come after last_log_id.
183 ws.last_log_id = p[:last_log_id].to_i
184 # Reset sent_ids for consistency
185 # (always re-deliver all matching messages following last_log_id)
186 ws.sent_ids = Set.new
189 if ws.filters.length < MAX_FILTERS
190 # Add a filter. This gets the :filters field which is the same
191 # format as used for regular index queries.
193 ws.send ({status: 200, message: 'subscribe ok', filter: p}.to_json)
195 # Send any pending events
198 ws.send ({status: 403, message: "maximum of #{MAX_FILTERS} filters allowed per connection"}.to_json)
201 elsif p[:method] == 'unsubscribe'
202 # Handle unsubscribe event
204 len = ws.filters.length
205 ws.filters.select! { |f| not ((f.filters == p[:filters]) or (f.filters.empty? and p[:filters].nil?)) }
206 if ws.filters.length < len
207 ws.send ({status: 200, message: 'unsubscribe ok'}.to_json)
209 ws.send ({status: 404, message: 'filter not found'}.to_json)
213 ws.send ({status: 400, message: "missing or unrecognized method"}.to_json)
216 Rails.logger.warn "Error handling message: #{$!}"
217 Rails.logger.warn "Backtrace:\n\t#{e.backtrace.join("\n\t")}"
218 ws.send ({status: 500, message: 'error'}.to_json)
223 # Constant maximum number of filters, to avoid silly huge database queries.
225 MAX_NOTIFY_BACKLOG = 1000
226 MAX_CONNECTIONS = 500
230 @connection_count >= MAX_CONNECTIONS
234 # Called by RackSocket when a new websocket connection has been established.
236 # Disconnect if no valid API token.
237 # current_user is included from CurrentApiClient
239 ws.send ({status: 401, message: "Valid API token required"}.to_json)
244 # Initialize our custom fields on the websocket connection object.
245 ws.user = current_user
248 ws.sent_ids = Set.new
250 ws.frame_mtx = Mutex.new
253 @connection_count += 1
256 # Subscribe to internal postgres notifications through @channel and
257 # forward them to the thread associated with the connection.
258 sub = @channel.subscribe do |msg|
260 ws.queue << [:notify, msg]
266 # Set up callback for inbound message dispatch.
267 ws.on :message do |event|
268 ws.queue << [:message, event]
271 # Set up socket close callback
272 ws.on :close do |event|
273 @channel.unsubscribe sub
275 ws.queue << [:close, nil]
278 # Spin off a new thread to handle sending events to the client. We need a
279 # separate thread per connection so that a slow client doesn't interfere
280 # with other clients.
282 # We don't want the loop in the request thread because on a TERM signal,
283 # Puma waits for outstanding requests to complete, and long-lived websocket
284 # connections may not complete in a timely manner.
286 # Loop and react to socket events.
288 eventType, msg = ws.queue.pop
289 if ws.queue.length > MAX_NOTIFY_BACKLOG
290 ws.send ({status: 500, message: 'Notify backlog too long'}.to_json)
293 if eventType == :message
294 handle_message ws, msg
295 elsif eventType == :notify
297 elsif eventType == :close
303 @connection_count -= 1
307 # Start up thread to monitor the Postgres database, if none exists already.
312 # from http://stackoverflow.com/questions/16405520/postgres-listen-notify-rails
313 ActiveRecord::Base.connection_pool.with_connection do |connection|
314 conn = connection.instance_variable_get(:@connection)
316 conn.async_exec "LISTEN logs"
318 # wait_for_notify will block until there is a change
319 # notification from Postgres about the logs table, then push
320 # the notification into the EventMachine channel. Each
321 # websocket connection subscribes to the other end of the
322 # channel and calls #push_events to actually dispatch the
323 # events to the client.
324 conn.wait_for_notify do |channel, pid, payload|
325 @channel.push payload.to_i
329 # Don't want the connection to still be listening once we return
330 # it to the pool - could result in weird behavior for the next
331 # thread to check it out.
332 conn.async_exec "UNLISTEN *"