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   # +notify_id+  the id of the most recent row in the log table, may be nil
  53   #
  54   # This accepts a websocket and a notify_id (this is the row id from Postgres LISTEN/NOTIFY, it may nil)
  55   # It queries the database for log rows that are either
  56   #  a) greater than ws.last_log_id, which is the last log id which was a candidate to be sent out
  57   #  b) if ws.last_log_id is nil, then it queries rows starting with notify_id
  58   #
  59   # Regular Arvados permissions are applied using readable_by() and filters using record_filters()
  60   # To avoid clogging up the database, queries are limited to batches of 100.  It will schedule a new
  61   # push_events call if there are more log rows to send.
  62   def push_events ws, notify_id
  63     begin
  64       if !notify_id.nil? and !ws.last_log_id.nil? and notify_id <= ws.last_log_id
  65         # This notify is for a row we've handled already.
  66         return
  67       end
  68
  69       # Must have at least one filter set up to receive events
  70       if ws.filters.length > 0
  71         # Start with log rows readable by user, sorted in ascending order
  72         logs = Log.readable_by(ws.user).order("id asc")
  73
  74         cond_id = nil
  75         cond_out = []
  76         param_out = []
  77
  78         if !ws.last_log_id.nil?
  79           # Client is only interested in log rows that are newer than the
  80           # last log row seen by the client.
  81           cond_id = "logs.id > ?"
  82           param_out << ws.last_log_id
  83         elsif !notify_id.nil?
  84           # No last log id, so look at rows starting with notify id
  85           cond_id = "logs.id >= ?"
  86           param_out << notify_id
  87         else
  88           # No log id to start from, nothing to do, return
  89           return
  90         end
  91
  92         # Now build filters provided by client
  93         ws.filters.each do |filter|
  94           ft = record_filters filter.filters, Log
  95           if ft[:cond_out].any?
  96             # Join the clauses within a single subscription filter with AND
  97             # so it is consistent with regular queries
  98             cond_out << "(#{ft[:cond_out].join ') AND ('})"
  99             param_out += ft[:param_out]
 100           end
 101         end
 102
 103         # Add filters to query
 104         if cond_out.any?
 105           # Join subscriptions with OR
 106           logs = logs.where(cond_id + " AND ((#{cond_out.join ') OR ('}))", *param_out)
 107         else
 108           logs = logs.where(cond_id, *param_out)
 109         end
 110
 111         # Execute query and actually send the matching log rows
 112         count = 0
 113         limit = 100
 114
 115         logs.limit(limit).each do |l|
 116           ws.send(l.as_api_response.to_json)
 117           ws.last_log_id = l.id
 118           count += 1
 119         end
 120
 121         if count == limit
 122           # Number of rows returned was capped by limit(), we need to schedule
 123           # another query to get more logs (will start from last_log_id
 124           # reported by current query)
 125           @channel.push nil
 126         elsif !notify_id.nil? and (ws.last_log_id.nil? or notify_id > ws.last_log_id)
 127           # Number of rows returned was less than cap, but the notify id is
 128           # higher than the last id visible to the client, so update last_log_id
 129           ws.last_log_id = notify_id
 130         end
 131       elsif !notify_id.nil?
 132         # No filters set up, so just record the sequence number
 133         ws.last_log_id = notify_id
 134       end
 135     rescue => e
 136       Rails.logger.warn "Error publishing event: #{$!}"
 137       Rails.logger.warn "Backtrace:\n\t#{e.backtrace.join("\n\t")}"
 138       ws.send ({status: 500, message: 'error'}.to_json)
 139       ws.close
 140     end
 141   end
 142
 143   # Handle inbound subscribe or unsubscribe message.
 144   def handle_message ws, event
 145     begin
 146       # Parse event data as JSON
 147       p = (Oj.load event.data).symbolize_keys
 148
 149       if p[:method] == 'subscribe'
 150         # Handle subscribe event
 151
 152         if p[:last_log_id]
 153           # Set or reset the last_log_id.  The event bus only reports events
 154           # for rows that come after last_log_id.
 155           ws.last_log_id = p[:last_log_id].to_i
 156         end
 157
 158         if ws.filters.length < MAX_FILTERS
 159           # Add a filter.  This gets the :filters field which is the same
 160           # format as used for regular index queries.
 161           ws.filters << Filter.new(p)
 162           ws.send ({status: 200, message: 'subscribe ok', filter: p}.to_json)
 163
 164           # Send any pending events
 165           push_events ws, nil
 166         else
 167           ws.send ({status: 403, message: "maximum of #{MAX_FILTERS} filters allowed per connection"}.to_json)
 168         end
 169
 170       elsif p[:method] == 'unsubscribe'
 171         # Handle unsubscribe event
 172
 173         len = ws.filters.length
 174         ws.filters.select! { |f| not ((f.filters == p[:filters]) or (f.filters.empty? and p[:filters].nil?)) }
 175         if ws.filters.length < len
 176           ws.send ({status: 200, message: 'unsubscribe ok'}.to_json)
 177         else
 178           ws.send ({status: 404, message: 'filter not found'}.to_json)
 179         end
 180
 181       else
 182         ws.send ({status: 400, message: "missing or unrecognized method"}.to_json)
 183       end
 184     rescue Oj::Error => e
 185       ws.send ({status: 400, message: "malformed request"}.to_json)
 186     rescue => e
 187       Rails.logger.warn "Error handling message: #{$!}"
 188       Rails.logger.warn "Backtrace:\n\t#{e.backtrace.join("\n\t")}"
 189       ws.send ({status: 500, message: 'error'}.to_json)
 190       ws.close
 191     end
 192   end
 193
 194   # Constant maximum number of filters, to avoid silly huge database queries.
 195   MAX_FILTERS = 16
 196
 197   # Called by RackSocket when a new websocket connection has been established.
 198   def on_connect ws
 199
 200     # Disconnect if no valid API token.
 201     # current_user is included from CurrentApiClient
 202     if not current_user
 203       ws.send ({status: 401, message: "Valid API token required"}.to_json)
 204       ws.close
 205       return
 206     end
 207
 208     # Initialize our custom fields on the websocket connection object.
 209     ws.user = current_user
 210     ws.filters = []
 211     ws.last_log_id = nil
 212
 213     # Subscribe to internal postgres notifications through @channel.  This will
 214     # call push_events when a notification comes through.
 215     sub = @channel.subscribe do |msg|
 216       push_events ws, msg
 217     end
 218
 219     # Set up callback for inbound message dispatch.
 220     ws.on :message do |event|
 221       handle_message ws, event
 222     end
 223
 224     # Set up socket close callback
 225     ws.on :close do |event|
 226       @channel.unsubscribe sub
 227       ws = nil
 228     end
 229
 230     # Start up thread to monitor the Postgres database, if none exists already.
 231     @mtx.synchronize do
 232       unless @bgthread
 233         @bgthread = true
 234         Thread.new do
 235           # from http://stackoverflow.com/questions/16405520/postgres-listen-notify-rails
 236           ActiveRecord::Base.connection_pool.with_connection do |connection|
 237             conn = connection.instance_variable_get(:@connection)
 238             begin
 239               conn.async_exec "LISTEN logs"
 240               while true
 241                 # wait_for_notify will block until there is a change
 242                 # notification from Postgres about the logs table, then push
 243                 # the notification into the EventMachine channel.  Each
 244                 # websocket connection subscribes to the other end of the
 245                 # channel and calls #push_events to actually dispatch the
 246                 # events to the client.
 247                 conn.wait_for_notify do |channel, pid, payload|
 248                   @channel.push payload.to_i
 249                 end
 250               end
 251             ensure
 252               # Don't want the connection to still be listening once we return
 253               # it to the pool - could result in weird behavior for the next
 254               # thread to check it out.
 255               conn.async_exec "UNLISTEN *"
 256             end
 257           end
 258           @bgthread = false
 259         end
 260       end
 261     end
 262
 263     # Since EventMachine is an asynchronous event based dispatcher, #on_connect
 264     # does not block but instead returns immediately after having set up the
 265     # websocket and notification channel callbacks.
 266   end
 267 end