brown 2.2.2 → 2.2.2.25.g85ddf08

data/lib/brown/agent/amqp.rb

@@ -0,0 +1,20 @@
+# AMQP support for Brown agents.
+#
+# Including this module in your agent provides support for publishing and
+# receiving AMQP messages from an AMQP broker, such as RabbitMQ.
+#
+# The methods in this module itself aren't particularly interesting to end-users;
+# the good stuff is in {Brown::Agent::AMQP::ClassMethods}.
+#
+module Brown::Agent::AMQP
+	private
+
+	def self.included(mod)
+		mod.url :AMQP_URL
+		mod.prepend(Brown::Agent::AMQP::Initializer)
+		mod.extend(Brown::Agent::AMQP::ClassMethods)
+	end
+end
+
+require_relative "amqp/initializer"
+require_relative "amqp/class_methods"

data/lib/brown/agent/amqp/class_methods.rb

@@ -0,0 +1,147 @@
+require 'logger'
+require 'securerandom'
+
+# Class-level AMQP support for Brown agents.
+#
+# These methods are intended to be applied to a `Brown::Agent` subclass, so you
+# can use them to define new AMQP listeners and publishers in your agent classes.
+# You should not attempt to extend your classes directly with this module; the
+# {Brown::Agent::AMQP} module should handle that for you automatically.
+#
+module Brown::Agent::AMQP::ClassMethods
+	attr_reader :amqp_publishers, :amqp_listeners
+
+	# Declare an AMQP publisher, and create an AMQP exchange to publish to.
+	#
+	# On the assumption that you already know [how exchanges
+	# work](http://www.rabbitmq.com/tutorials/amqp-concepts.html), lets
+	# just dive right in.
+	#
+	# This method creates an accessor method on your agent class, named after the
+	# symbol you pass in as `name`, which returns an instance of
+	# `Brown::Agent::AMQPPublisher`.  This object, in turn, defines an
+	# AMQP exchange when it is created, and has a
+	# `publish` method on it (see {Brown::Agent::AMQPPublisher#publish}) which
+	# sends arbitrary messages to the exchange.
+	#
+	# @param name [Symbol] the name of the accessor method to call when
+	#   you want to reference this publisher in your agent code.
+	#
+	# @param publisher_opts [Hash] options which are passed to
+	#   {Brown::Agent::AMQPPublisher#initialize}.
+	#
+	# This method is a thin shim around {Brown::Agent::AMQPPublisher#initialize};
+	# you should read that method's documentation for details of what
+	# constitutes valid `publisher_opts`, and also what exceptions can be
+	# raised.
+	#
+	# @see Brown::Agent::AMQPPublisher#initialize
+	# @see Brown::Agent::AMQPPublisher#publish
+	#
+	def amqp_publisher(name, publisher_opts = {})
+		@amqp_publishers ||= []
+		@amqp_publishers << { name: name, opts: publisher_opts }
+	end
+
+	# Declare a stimulus to listen for messages from an AMQP broker.
+	#
+	# When the agent is started, we will setup a queue, bound to the exchange
+	# specified by the `exchange_name` argument, and then proceed to hoover up
+	# all the messages we can.
+	#
+	# The name of the queue that is created by default is derived from the
+	# agent class name and the exchange name being bound to.  This allows
+	# for multiple instances of the same agent, running in separate
+	# processes or machines, to share the same queue of messages to
+	# process, for throughput or redundancy reasons.
+	#
+	# @param exchange_name [#to_s, Array<#to_s>] the name of the exchange
+	#   to bind to.  You can also specify an array of exchange names, to
+	#   have all of them put their messages into the one queue.  This can
+	#   be dangerous, because you need to make sure that your message
+	#   handler can process the different types of messages that might be
+	#   sent to the different exchangs.
+	#
+	# @param queue_name [#to_s] the name of the queue to create, if you
+	#   don't want to use the class-derived default for some reason.
+	#
+	# @param concurrency [Integer] how many messages to process in parallel.
+	#   The default, `1`, means that a message will need to be acknowledged
+	#   (by calling `message.ack`) in your worker `blk` before the broker
+	#   will consider sending another.
+	#
+	#   If your agent is capable of processing more than one message in
+	#   parallel (because the agent spends a lot of its time waiting for
+	#   databases or HTTP requests, for example, or perhaps you're running
+	#   your agents in a Ruby VM which has no GIL) you should increase
+	#   this value to improve performance.  Alternately, if you want/need
+	#   to batch processing (say, you insert 100 records into a database
+	#   in a single query) you'll need to increase this to get multiple
+	#   records at once.
+	#
+	#   Setting this to `0` is only for the adventurous.  It tells the
+	#   broker to send your agent messages as fast as it can.  You still
+	#   need to acknowledge the messages as you finish processing them
+	#   (otherwise the broker will not consider them "delivered") but you
+	#   will always be sent more messages if there are more to send, even
+	#   if you never acknowledge any of them.  This *can* get you into an
+	#   awful lot of trouble if you're not careful, so don't do it just
+	#   because you can.
+	#
+	# @param autoparse [Boolean] if your messages are all of a single
+	#   content type, and the first thing you do on receiving a message is to
+	#   call `JSON.parse(msg.payload)`, then you can turn on `autoparse`, and
+	#   *as long as the message content-type is set correctly*, your messages
+	#   will be parsed into native data structures before being turned into the
+	#   `payload`.  This also enables automatic deserialization of object
+	#   messages.
+	#
+	# @param allowed_classes [Array<Class>] a list of classes which are
+	#   permitted to be instantiated when a message's `content_type` attribute
+	#   indicates it is a serialized object.  If a message is received that
+	#   contains a class not in the list (even if that object is embedded inside
+	#   the message itself), the message will be rejected.
+	#
+	# @param routing_key [#to_s] if specified, then the queue that is created
+	#   will be bound to the exchange with a defined routing key.
+	#
+	# @param predeclared [Boolean] indicates that the necessary queues and
+	#   bindings are already in place, and thus no declarations should be
+	#   made to the AMQP broker.
+	#
+	# @param blk [Proc] is called every time a message is received from
+	#   the queue, and an instance of {Brown::Agent::AMQPMessage} will
+	#   be passed as the sole argument.
+	#
+	# @yieldparam message [Brown::Agent::AMQPMessage] is passed to `blk`
+	#   each time a message is received from the queue.
+	#
+	def amqp_listener(exchange_name = "",
+	                  queue_name:      nil,
+	                  concurrency:     1,
+	                  autoparse:       false,
+	                  allowed_classes: nil,
+	                  routing_key:     nil,
+	                  predeclared:     false,
+	                  &blk
+	                 )
+		exchange_list = (Array === exchange_name ? exchange_name : [exchange_name]).map(&:to_s)
+
+		if queue_name.nil?
+			munged_exchange_list = exchange_list.map { |n| n.to_s == "" ? "" : "-#{n.to_s}" }.join
+			queue_name = self.name.to_s + munged_exchange_list + (routing_key ? ":#{routing_key}" : "")
+		end
+
+		@amqp_listeners ||= []
+		@amqp_listeners << {
+			exchange_list:   exchange_list,
+			queue_name:      queue_name,
+			concurrency:     concurrency,
+			autoparse:       autoparse,
+			allowed_classes: allowed_classes,
+			routing_key:     routing_key,
+			callback:        blk,
+			predeclared:     predeclared,
+		}
+	end
+end

data/lib/brown/agent/amqp/initializer.rb

@@ -0,0 +1,144 @@
+require "securerandom"
+require "json"
+require "yaml"
+
+# Methods that have to be prepended in order to work properly.
+#
+module Brown::Agent::AMQP::Initializer
+	def initialize(*_)
+		begin
+			super
+		rescue ArgumentError => ex
+			if ex.message =~ /wrong number of arguments.*expected 0/
+				super()
+			else
+				raise
+			end
+		end
+
+		initialize_publishers
+	end
+
+	def run
+		initialize_listeners
+
+		super
+	end
+
+	def shutdown
+		amqp_session.close
+
+		super
+	end
+
+	private
+
+	def amqp_session
+		@amqp_session ||= begin
+			logger.debug(logloc) { "Initializing AMQP session" }
+			Bunny.new(config.amqp_url, recover_from_connection_close: true, logger: config.logger).tap do |session|
+				session.on_blocked { |blocked| logger.warn(logloc) { "AMQP connection has become blocked: #{blocked.reason}" } }
+				session.on_unblocked { logger.info(logloc) { "AMQP connection has unblocked" } }
+				session.start
+			end
+		end
+	end
+
+	def initialize_publishers
+		(self.class.amqp_publishers || []).each do |publisher|
+			logger.debug(logloc) { "Initializing AMQP publisher #{publisher}" }
+			opts = { exchange_name: publisher[:name] }.merge(publisher[:opts])
+
+			amqp_publisher = Brown::Agent::AMQPPublisher.new(amqp_session: amqp_session, **opts)
+
+			define_singleton_method(publisher[:name]) { amqp_publisher }
+		end
+	end
+
+	def initialize_listeners
+		(self.class.amqp_listeners || []).each do |listener|
+			logger.debug(logloc) { "Initializing AMQP listener #{listener}" }
+			worker_method = "amqp_listener_worker_#{SecureRandom.uuid}".to_sym
+			define_singleton_method(worker_method, listener[:callback])
+
+			@stimuli ||= []
+			@stimuli << {
+				name: "amqp_listener_#{listener[:exchange_list].join("_").gsub(/[^A-Za-z0-9_]/, '_').gsub(/__+/, "_")}",
+				method: method(worker_method),
+				stimuli_proc: proc do |worker|
+					consumer = queue(listener).subscribe(manual_ack: true) do |di, prop, payload|
+						if listener[:autoparse]
+							logger.debug(logloc) { "Attempting to autoparse against Content-Type: #{prop.content_type.inspect}" }
+							case prop.content_type
+							when "application/json"
+								logger.debug(logloc) { "Parsing as JSON" }
+								payload = JSON.parse(payload)
+							when "application/x.yaml"
+								logger.debug(logloc) { "Parsing as YAML" }
+								payload = YAML.load(payload)
+							when "application/vnd.brown.object.v1"
+								logger.debug(logloc) { "Parsing as Brown object, allowed classes: #{listener[:allowed_classes]}" }
+								begin
+									payload = YAML.safe_load(payload, listener[:allowed_classes])
+								rescue Psych::DisallowedClass => ex
+									logger.error(logloc) { "message rejected: #{ex.message}" }
+									di.channel.nack(di.delivery_tag, false, false)
+									next
+								end
+							end
+						end
+
+						worker.call Brown::Agent::AMQPMessage.new(di, prop, payload)
+					end
+
+					while consumer&.channel&.status == :open do
+						logger.debug(logloc) { "stimuli_proc for #{listener[:queue_name]} having a snooze" }
+						sleep
+					end
+				end
+			}
+		end
+	end
+
+	def queue(listener)
+		@queue_cache ||= {}
+		@queue_cache[listener] ||= begin
+			bind_queue(
+				queue_name:    listener[:queue_name],
+				exchange_list: listener[:exchange_list].map(&:to_s),
+				concurrency:   listener[:concurrency],
+				routing_key:   listener[:routing_key],
+				predeclared:   listener[:predeclared],
+			)
+		rescue StandardError => ex
+			log_exception(ex) { "Unknown error while binding queue #{listener[:queue_name].inspect} to exchange list #{listener[:exchange_list].inspect}" }
+			sleep 5
+			retry
+		end
+	end
+
+	def bind_queue(queue_name:, exchange_list:, concurrency:, routing_key: nil, predeclared: false)
+		ch = amqp_session.create_channel
+		ch.prefetch(concurrency)
+
+		ch.queue(queue_name, durable: true, no_declare: predeclared).tap do |q|
+			next if predeclared
+			exchange_list.each do |exchange_name|
+				if exchange_name != ""
+					begin
+						q.bind(exchange_name, routing_key: routing_key)
+					rescue Bunny::NotFound => ex
+						logger.error { "bind failed: #{ex.message}" }
+						sleep 5
+						return bind_queue(
+						       queue_name:    queue_name,
+						       exchange_list: exchange_list,
+						       concurrency:   concurrency,
+						       routing_key:   routing_key,
+						      )
+					end
+				end
+			end
+		end
+	end
+end

data/lib/brown/agent/amqp_message.rb

@@ -7,11 +7,14 @@
 # "delivered".
 #
 class Brown::Agent::AMQPMessage
-	# The raw body of the message.  No translation is done on what was sent,
-	# so any serialisation that might have been applied to the message will
-	# have to be undone manually.
+	# The body of the message.
 	#
-	# @return [String]
+	# Depending on the listener configuration, this may either be a string
+	# representing the raw message, or else a deserialized object of some
+	# sort.  See the documentation for the amqp_listener method's `autoparse`
+	# and `allowed_classes` parameters for more details.
+	#
+	# @return [Object]
 	#
 	attr_reader :payload
 
@@ -39,4 +42,23 @@ class Brown::Agent::AMQPMessage
 	def ack
 		@delivery_info.channel.ack(@delivery_info.delivery_tag)
 	end
+
+	# Decline to handle this message, and requeue for later
+	#
+	# If, for some reason, the agent can't successfully process this message
+	# right *now*, but you'd like it to be processed again later, then you
+	# can ask for it to be requeued using this handy-dandy method.
+	#
+	def requeue
+		@delivery_info.channel.nack(@delivery_info.delivery_tag, false, true)
+	end
+
+	# Decline to handle this message, now and forever
+	#
+	# If you want to dead-letter (or just disappear) a message, call this
+	# method, and it will cease to exist.
+	#
+	def reject
+		@delivery_info.channel.nack(@delivery_info.delivery_tag, false, false)
+	end
 end

data/lib/brown/agent/amqp_publisher.rb

@@ -1,8 +1,13 @@
 require 'bunny'
+require 'json'
+require 'service_skeleton/logging_helpers'
+require 'yaml'
 
 # Publish messages to an AMQP exchange.
 #
 class Brown::Agent::AMQPPublisher
+	include ServiceSkeleton::LoggingHelpers
+
 	#:nodoc:
 	# Sentinel to detect that we've been sent the "default" value,
 	# since `nil` can, sometimes, be a valid value.
@@ -34,21 +39,14 @@ class Brown::Agent::AMQPPublisher
 	# Setup an exchange in the AMQP broker, and allow the publishing of
 	# messages to that exchange.
 	#
-	# @param amqp_url [#to_s] the AMQP broker to connect to, specified as a
-	#   URL.  The scheme must be `amqp`.  Username and password should be
-	#   given in the standard fashion (`amqp://<user>:<pass>@<host>`).
-	#
-	#   The path portion of AMQP URLs is totes spesh; if you want to
-	#   connect to the default vhost (`/`) you either need to specify *no*
-	#   trailing slash (ie `amqp://hostname`) or percent-encode the `/`
-	#   vhost name (ie `amqp://hostname/%2F`).  Yes, this drives me nuts,
-	#   too.
+	# @param amqp_session [Bunny::Session] an active session with the AMQP
+	#   server.  Typically this will be the agent's `@amqp_session` variable.
 	#
 	# @param exchange_type [Symbol] the type of exchange to create or publish
 	#   to.  By default, the exchange is created as a *direct* exchange; this
 	#   routes messages to their destination queue(s) based on the
 	#   `routing_key` (set per-publisher or per-queue).  Other valid values
-	#   for this option are `:direct`, `:topic`, and `:headers`.
+	#   for this option are `:fanout`, `:topic`, and `:headers`.
 	#
 	# @param exchange_name [#to_s] the name of the exchange to create or
 	#   publish to.  If not specified, then the "default" exchange is used,
@@ -88,7 +86,7 @@ class Brown::Agent::AMQPPublisher
 	#   create the exchange fails for some reason (such as the exchange
 	#   already existing with a different configuration).
 	#
-	def initialize(amqp_url:      "amqp://localhost",
+	def initialize(amqp_session:,
 	               exchange_type: :direct,
 	               exchange_name: "",
 	               routing_key:   nil,
@@ -96,23 +94,11 @@ class Brown::Agent::AMQPPublisher
 	               logger:        Logger.new("/dev/null"),
 	               **amqp_opts
 	              )
-		begin
-			@amqp_session = Bunny.new(amqp_url, logger: logger)
-			@amqp_session.start
-		rescue Bunny::TCPConnectionFailed
-			raise BrokerError,
-			      "Failed to connect to #{amqp_url}"
-		rescue Bunny::PossibleAuthenticationFailureError
-			raise BrokerError,
-			      "Authentication failed for #{amqp_url}"
-		rescue StandardError => ex
-			raise Error,
-			      "Unknown error occured: #{ex.message} (#{ex.class})"
-		end
-
-		@amqp_channel = @amqp_session.create_channel
+		@logger = logger
+		@amqp_channel = amqp_session.create_channel
 
 		begin
+			@logger.debug(logloc) { "Initializing exchange #{exchange_name.inspect}" }
 			@amqp_exchange = @amqp_channel.exchange(
 			                   exchange_name,
 			                   type: exchange_type,
@@ -133,7 +119,35 @@ class Brown::Agent::AMQPPublisher
 
 	# Publish a message to the exchange.
 	#
-	# @param payload [#to_s] the "body" of the message to send.
+	# A message can be provided in the `payload` parameter in one of several
+	# ways:
+	#
+	# * as a string, in which case the provided string is used as the message
+	#   payload as-is, with no further processing or added metadata.
+	#
+	# * as a single-element hash, in which case the key is used as the identifier
+	#   of the serialization format (either `:yaml` or `:json`), while the value
+	#   is serialized by calling the appropriate conversion method (`.to_yaml` or
+	#   `.to_json`, as appropriate) and the resulting string is used as the message
+	#   payload.  The message's `content_type` attribute is set appropriately, also,
+	#   allowing for automatic deserialization at the receiver.
+	#
+	#   This payload format is best used where you have complex structured data to
+	#   pass around between system components with very diverse implementation
+	#   environments, or where very loose coupling is desired.
+	#
+	# * as an arbitrary object, in which case the object is serialized in a form
+	#   which allows it to be transparently deserialized again into a native
+	#   Ruby object at the other end.
+	#
+	#   This payload format is most appropriate when you're deploying a
+	#   tightly-coupled system where all the consumers are Brown agents, or at
+	#   least all Ruby-based.
+	#
+	# @param payload [String, Hash<Symbol, Object>, Object] the content of the
+	#   message to send.  There are several possibilities for what can be passed
+	#   here, and how it is encoded into an AMQP message, as per the description
+	#   above.
 	#
 	# @param type [#to_s] override the default message type set in the
 	#   publisher, just for this one message.
@@ -150,6 +164,8 @@ class Brown::Agent::AMQPPublisher
 	#   for full details of every possible permutation.
 	#
 	def publish(payload, type: NoValue, routing_key: NoValue, **amqp_opts)
+		@logger.debug(logloc) { "Publishing message #{payload.inspect}, type: #{type.inspect}, routing_key: #{routing_key.inspect}, options: #{amqp_opts.inspect}" }
+
 		opts = @message_defaults.merge(
 		          {
 		            type:        type,
@@ -157,6 +173,30 @@ class Brown::Agent::AMQPPublisher
 		          }.delete_if { |_,v| v == NoValue }
 		       ).delete_if { |_,v| v.nil? }.merge(amqp_opts)
 
+		if payload.is_a?(Hash)
+			if payload.length != 1
+				raise ArgumentError,
+					   "Payload hash must have exactly one element"
+			end
+
+			case payload.keys.first
+			when :json
+				@logger.debug(logloc) { "JSON serialisation activated" }
+				opts[:content_type] = "application/json"
+				payload = payload.values.first.to_json
+			when :yaml
+				@logger.debug(logloc) { "YAML serialisation activated" }
+				opts[:content_type] = "application/x.yaml"
+				payload = payload.values.first.to_yaml
+			else
+				raise ArgumentError,
+				      "Unknown format type: #{payload.keys.first.inspect} (must be :json or :yaml)"
+			end
+		elsif !payload.is_a?(String)
+			payload = payload.to_yaml
+			opts[:content_type] = "application/vnd.brown.object.v1"
+		end
+
 		if @amqp_exchange.name == "" and opts[:routing_key].nil?
 			raise ExchangeError,
 			      "Cannot send a message to the default exchange without a routing key"
@@ -164,6 +204,7 @@ class Brown::Agent::AMQPPublisher
 
 		@channel_mutex.synchronize do
 			@amqp_exchange.publish(payload, opts)
+			@logger.debug(logloc) { "... and it's gone" }
 		end
 	end
 end