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