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