droid 0.9.0

data/bin/conf

@@ -0,0 +1,31 @@
+#!/usr/bin/env ruby
+require 'socket'
+
+def wait_for_tcp_port(port, options={})
+  options[:timeout] ||= 30
+  options[:host]    ||= "localhost"
+  
+  options[:timeout].times do
+    begin
+      x = TCPSocket.new(options[:host], port)
+      x.close
+      return
+    rescue
+      sleep 1
+      next
+    end
+  end
+
+  raise "#{options[:host]}:#{port} did not come up after #{options[:timeout]} seconds"
+end
+
+puts "Waiting for RabbitMQ"
+wait_for_tcp_port(5672)
+
+puts `rabbitmqctl add_vhost /heroku`
+
+%w[hermes core].each do |agent|
+  puts `rabbitmqctl add_user #{agent} thisisheroku!`
+  puts `rabbitmqctl set_permissions -p /heroku #{agent} \".*\" \".*\" \".*\"`
+end
+

data/lib/droid.rb

@@ -0,0 +1,531 @@
+require 'socket'
+require 'digest/md5'
+require File.dirname(__FILE__) + '/../vendor/logger_client/init'
+
+$:.unshift *Dir[File.dirname(__FILE__) + '/../vendor/*/lib']
+require 'json'
+require 'mq'
+require 'time'
+require 'bunny'
+
+require File.dirname(__FILE__) + '/utilization'
+
+class Droid
+
+	class BadPayload < RuntimeError; end
+
+## basic ops
+## ######
+## publish / broadcast
+## listen / subscribe
+
+	def self.con_type
+		Thread.current['con_type'] ||= :sync
+	end
+
+	def self.con_type=(type)
+		Thread.current['con_type'] = type
+	end
+
+	def self.async?
+		return @@async == true rescue false
+	end
+
+	def self.async(&blk)
+		ensure_con_type(:async, &blk)
+	end
+
+	def self.sync(&blk)
+		ensure_con_type(:sync, &blk)
+	end
+
+	def self.ensure_con_type(type, &blk)
+		old_type = con_type
+		self.con_type = type
+		begin
+			blk.call
+		ensure
+			self.con_type = old_type
+		end
+	end
+
+	def self.new_event_hash
+		s = Time.now.to_s + self.object_id.to_s + rand(100).to_s
+		Digest::MD5.hexdigest(s)
+	end
+
+	def self.queue(name, options = {})
+		reconnect_on_error do
+			if con_type == :async
+				MQ.queue(name, options)
+			else
+					bunny.queue(name, options)
+			end
+		end
+	end
+
+	if defined? Bunny::Client   # the blog below doesn't work with newer versions of bunny
+		# Disable bunny's 1 second socket connect timeout since the
+		# reconnect_on_error method sets up a separate 10 second timeout.
+		# Using a timeout of zero sets an "infinite" timeout and has the nice
+		# benefit of not starting up another thread.
+		::Bunny::Client.send(:remove_const, :CONNECT_TIMEOUT)
+		::Bunny::Client::CONNECT_TIMEOUT = 0
+	end
+
+	def self.reconnect_on_error
+		Timeout::timeout(20) do
+			begin
+				yield
+			rescue Bunny::ProtocolError
+				sleep 0.5
+				retry
+			rescue Bunny::ConnectionError
+				sleep 0.5
+				@@bunny = nil
+				retry
+			rescue Bunny::ServerDownError
+				sleep 0.5
+				@@bunny = nil
+				retry
+			end
+		end
+	end
+
+	def self.call(key, payload, options={})
+		sync do
+			reply_to = key + '.reply.' + Droid.gensym
+			@q = nil
+			begin
+				reconnect_on_error do
+
+					## this is retarded - i shouldn't be binding here - just popping the queue - need to teach hermes/em
+					@q = queue(reply_to, :auto_delete => true)
+					@q.bind(exchange, :key => reply_to)
+
+					payload[:reply_to] = reply_to
+					publish(key, payload, options)
+
+					pop(reply_to)
+				end
+			ensure
+				# for some reason the auto_delete flag is not working correctly with Bunny
+				# so we're deleting the queue manually here
+				@q.delete if @q
+			end
+		end
+	end
+
+	def self.pop(queue)
+		loop do
+			raise "POP must be sync" unless con_type == :sync
+			result = queue(queue).pop
+			result = result[:payload] if result.is_a?(Hash)
+			return JSON.parse(result) unless result == :queue_empty
+			sleep 0.1
+		end
+	end
+
+	def self.push(queue_name, payload, options={})
+		reconnect_on_error do
+			queue(queue_name).publish(payload_to_data(payload, options))
+		end
+	end
+
+	def self.payload_to_data(payload, options)
+		raise BadPayload unless payload.is_a?(Hash)
+
+		payload[:event_hash] ||= new_event_hash
+		payload[:published_on] = options[:published_on] || Time.now.getgm.to_i
+		payload[:ttl] ||= (options[:ttl].to_i || DEFAULT_TTL)
+
+		payload.to_json
+	end
+
+	def self.publish(key, payload, options={})
+		res =
+			reconnect_on_error do
+				exchange.publish(payload_to_data(payload, options), :key => key, :immediate => options[:immediate])
+			end
+
+		unless options[:log] == false
+			Log.notice "#{key} published #{payload_summary(payload)}",
+				:event_hash => payload[:event_hash],
+				:addendum => payload.inspect
+		end
+
+		res
+	end
+
+	def self.header_keys
+		@header_keys ||= [:exchange, :delivery_mode, :delivery_tag, :redelivered, :consumer_tag, :content_type, :key, :priority]
+	end
+
+	def self.payload_summary(payload)
+		payload = payload.select do |k, v|
+			!header_keys.include?(k.to_sym)
+		end
+		return ' -> (empty payload)' if payload.empty?
+		resume = payload.map do |k, v|
+			v = v.to_s
+			v = v[0..17] + '...' if v.size > 20
+			"#{k}=#{v}"
+		end[0,15].join(', ')
+		resume << " (and #{payload.size - 15} more)" if payload.size > 15
+		" -> #{resume}"
+	end
+
+	def self.exchange
+		if con_type == :async
+			MQ.topic('amq.topic')
+		else
+			bunny.exchange("amq.topic")
+		end
+	end
+
+	def self.default_options
+      uri = URI.parse(ENV["AMQP_URI"] || 'rabbit://guest:guest@localhost:5672/')
+      raise "invalid AMQP_URI [#{uri.to_s}]" unless uri.scheme == "rabbit"
+      {
+         :vhost => uri.path,
+         :host => uri.host,
+         :user => uri.user,
+         :port => uri.port,
+         :pass => uri.password
+      }
+   end
+
+	def self.default_config
+		default_options
+	end
+
+	def self.new_bunny
+		b = Bunny.new(default_options)
+		b.start
+		b
+	end
+
+	def self.bunny
+		@@bunny ||= new_bunny
+	end
+
+	def self.start(options = nil, &blk)
+		async do
+			begin
+				Signal.trap('INT') { AMQP.stop{ EM.stop } }
+				Signal.trap('TERM'){ AMQP.stop{ EM.stop } }
+				EM.run do
+					AMQP.start(options || default_options)
+					blk.call if blk
+				end
+			rescue AMQP::Error => e
+				STDERR.puts "Caught #{e.class}, sleeping to avoid inittab thrashing"
+				sleep 5
+				STDERR.puts "Done."
+				raise
+			end
+		end
+	end
+
+	def self.stop_safe
+		EM.add_timer(1) { AMQP.stop { EM.stop }}
+	end
+
+	def self.gen_queue(droid, key)
+		dn = droid
+		dn = dn.name if dn.respond_to?(:name)
+		dn ||= "d"
+		dn.gsub!(" ", "")
+		"#{self.gen_instance_queue(key)}.#{dn}"
+	end
+
+	def self.gen_instance_queue(key)
+		"#{key}.#{LocalStats.slot}.#{LocalStats.ion_instance_id}"
+	end
+
+	def self.gensym
+		values = [
+			rand(0x0010000),
+			rand(0x0010000),
+			rand(0x0010000),
+			rand(0x0010000),
+			rand(0x0010000),
+			rand(0x1000000),
+			rand(0x1000000),
+		]
+		"%04x%04x%04x%04x%04x%06x%06x" % values
+	end
+
+	def self.wait_for_tcp_port(host, port, options={:retries => 5, :timeout => 5})
+		require 'timeout'
+		options[:retries].times do
+			begin
+				Timeout::timeout(options[:timeout]) {
+					TCPSocket.new(host.to_s, port).close
+				}
+				return
+			rescue Object
+				Log.notice "#{host}:#{port} not available, waiting..."
+				sleep 1
+			end
+		end
+
+		raise "#{host}:#{port} did not come up after #{options[:retries]} retries"
+	end
+
+	# Trap exceptions leaving the block and log them. Do not re-raise.
+	def self.trap_exceptions
+		yield
+	rescue => boom
+		Log.default_error boom
+	end
+
+	# Add a one-shot timer.
+	def self.timer(duration, &bk)
+		EM.add_timer(duration) { trap_exceptions(&bk) }
+	end
+
+	# Add a periodic timer. If the now argument is true, run the block
+	# immediately in addition to scheduling the periodic timer.
+	def self.periodic_timer(duration, now=false, &bk)
+		timer(1, &bk) if now
+		EM.add_periodic_timer(duration) { trap_exceptions(&bk) }
+	end
+
+	def timer(duration, &bk) ; self.class.timer(duration, &bk) ; end
+	def periodic_timer(duration, now=false, &bk) ; self.class.periodic_timer(duration, now, &bk) ; end
+
+	class Basic
+		DEFAULT_TTL = 300
+
+		def initialize(droid, options={})
+			@droid = droid
+		end
+
+		def exchange
+			Droid.exchange
+		end
+
+		def headers
+			@headers ||= { :event_hash => self.event_hash }
+		end
+
+		def event_hash
+			@event_hash ||= Droid.new_event_hash
+		end
+
+		def publish(key, payload, options={}, &blk)
+			raise BadPayload unless payload.is_a?(Hash)
+
+			result =if blk
+				headers[:reply_to] = key + '.reply.' + Droid.gensym
+				@droid.listen4(headers[:reply_to], { :temp => true }, &blk)
+			end
+
+			Droid.publish(key, headers.merge(payload), {:log => true}.merge(options))
+
+			result
+		end
+
+		def payload_summary(payload)
+			Droid.payload_summary(payload)
+		end
+	end
+
+	class Listener < Basic
+		attr_accessor :params
+
+		def initialize(droid, key, options={})
+			@key     = key
+			@options = options
+			queue    = @options.delete(:queue) || Droid.gen_queue(droid, key)
+			auto_delete = @options.has_key?(:auto_delete) ? !!@options.delete(:auto_delete) : true
+			@mq      = MQ.new
+			@prefetch = !!@options[:prefetch]
+			@mq.prefetch(@options[:prefetch]) if @prefetch
+			@q       = @mq.queue(queue, :auto_delete => auto_delete)
+			super(droid, options)
+		end
+
+		def destroy
+			@q.unsubscribe
+			@mq.close
+		end
+
+		def mq
+			@mq
+		end
+
+		def exchange
+			Droid.exchange
+		end
+
+		def error(e)
+			begin
+				publish("event.error", :event_hash => headers[:event_hash])	# hermes fail whale
+				msg = "#{e.class}: #{e.message}\n	#{e.backtrace.join("\n	")}\n"
+				stderr_puts "About to log error #{headers[:event_hash]}"
+				stderr_puts e.message
+				stderr_puts msg
+				Log.error e.message, :addendum => msg, :event_hash => headers[:event_hash], :exception => e
+				@droid.error_handler.call(@message, e, self) if @droid.error_handler
+			rescue Exception => e
+				stderr_puts "error handling error! #{e.inspect}"
+			end
+		end
+
+		def stderr_puts(msg)
+			STDERR.puts msg
+		end
+
+		def defer(&blk)
+			EM.defer(lambda do
+				begin
+					blk.call
+				 rescue => e
+					 error(e)
+				 end
+			 end)
+		end
+
+		def reply(payload, options={})
+			publish(headers[:reply_to], payload, options)
+		end
+
+		def listen(opts={}, &blk)
+			opts[:temp] = opts[:temp] === true
+			opts[:ack] = opts[:ack] === true
+
+			if @prefetch
+				opts[:ack] = true   # we must ack messages received in order for prefetch to work
+				opts[:temp] = false # doesn't make sense for it to be temporary if we're setting prefetch
+			end
+			@q.bind(exchange, :key => @key ).subscribe(:ack => opts[:ack]) do |info, data|
+				Utilization.monitor(@key, :temp => opts[:temp]) do
+					begin
+						parse(info, data)
+
+						if opts[:detail]
+							callargs = [self.dup, info, data]
+						else
+							callargs = [self.dup]
+						end
+
+						ttl = headers[:ttl]
+						now = Time.now.getgm.to_i
+						published_on = headers[:published_on]
+
+						Log.notice "#{@key} received #{payload_summary(params)}"
+
+						if ttl == -1 or now <= (published_on + ttl)
+							@droid.before_filter.call(*callargs) if @droid.before_filter
+							blk.call(*callargs)
+						else
+							Log.error "#{@key} message timeout", :addendum => payload_summary(params)
+							info.ack if opts[:ack]
+						end
+
+						duration = Time.now.getgm.to_i - now
+
+						Log.notice "(#{duration} seconds) SLOW AMQP #{@key} #{payload_summary(params)}" if duration > 3
+					rescue => e
+						error(e)
+					ensure
+						if opts[:temp]
+							@q.unbind(exchange)
+							@q.delete
+						end
+					end
+				end
+			end
+
+			self
+		end
+
+		def requeue(opts={})
+			opts[:ttl] ||= 10
+
+			now = Time.now.getgm.to_i
+
+			payload = @params.merge(extra_headers)
+			payload.delete('ttl')
+			payload[:ttl] = opts[:ttl]
+
+			newpayload = Droid.payload_to_data(payload, :published_on => (headers[:published_on] || now))
+			@q.publish(newpayload)
+		end
+
+		def parse(info, data)
+			@headers = nil
+			@params = JSON.parse(data)
+
+			headers[:event_hash]   = @params.delete('event_hash') if @params['event_hash']
+			headers[:reply_to]     = @params.delete('reply_to') if @params['reply_to']
+			headers[:published_on] = @params.delete('published_on').to_i rescue 0
+			headers[:ttl]          = @params.delete('ttl').to_i rescue -1
+			headers[:ttl] = -1 if headers[:ttl] == 0
+			headers.merge!(info.properties) # add protocol headers
+		end
+
+		def extra_headers
+			extra = {}
+			[:event_hash, :reply_to, :published_on].each do |key|
+				extra[key] = headers[key]
+			end
+			extra
+		end
+
+		def [](key)
+			@params[key.to_s]
+		end
+
+		def unsubscribe
+			@q.unsubscribe
+		end
+	end
+
+	attr_accessor :name
+	attr_reader	 :error_handler, :before_filter
+
+	def initialize(name, credentials, &blk)
+		Log.notice "=== #{name} droid initializing"
+		credentials[:port] ||= 5672
+		self.class.wait_for_tcp_port(credentials[:host], credentials[:port], :retries => 6) # retry for 30s before giving up
+
+		@name = name
+		Log.notice "=== #{name} droid starting"
+		self.class.start(credentials) do
+			blk.call(self)
+		end
+		self
+	end
+
+	def publish(key, payload={}, options={}, &blk)
+		Basic.new(self, options).publish(key, payload, options, &blk)
+	end
+
+	def listen4(key, options={}, &blk)
+		Listener.new(self, key, options).listen({
+			:temp => options.delete(:temp),
+			:detail => options.delete(:detail),
+			:ack => options.delete(:ack),
+			}, &blk)
+	end
+
+	def on_error(&blk)
+		@error_handler = blk
+	end
+
+	def before_filter(&blk)
+		blk ? @before_filter = blk : @before_filter
+	end
+
+	def stats(&blk)
+		@stats = blk
+		Log.notice call_stats
+	end
+
+	def call_stats
+		@stats ? @stats.call : nil
+	end
+end