services/api/lib/eventbus.rb

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