├── .ruby-version ├── .ruby-gemset ├── .rspec ├── lib ├── nsq │ ├── frames │ │ ├── error.rb │ │ ├── response.rb │ │ ├── frame.rb │ │ └── message.rb │ ├── logger.rb │ ├── producer.rb │ ├── discovery.rb │ ├── client_base.rb │ ├── consumer.rb │ └── connection.rb ├── version.rb └── nsq.rb ├── Gemfile ├── .gitignore ├── LICENSE.txt ├── Rakefile ├── Gemfile.lock ├── nsq-ruby.gemspec ├── spec ├── spec_helper.rb └── lib │ └── nsq │ ├── discovery_spec.rb │ ├── connection_spec.rb │ ├── consumer_with_shaky_connections_spec.rb │ ├── producer_spec.rb │ └── consumer_spec.rb └── README.md /.ruby-version: -------------------------------------------------------------------------------- 1 | 2.1.2 2 | -------------------------------------------------------------------------------- /.ruby-gemset: -------------------------------------------------------------------------------- 1 | nsq-ruby 2 | -------------------------------------------------------------------------------- /.rspec: -------------------------------------------------------------------------------- 1 | +--require rspec/legacy_formatters 2 | +--format Rainbow 3 | --order rand 4 | -------------------------------------------------------------------------------- /lib/nsq/frames/error.rb: -------------------------------------------------------------------------------- 1 | require_relative 'frame' 2 | 3 | module Nsq 4 | class Error < Frame 5 | end 6 | end 7 | -------------------------------------------------------------------------------- /lib/nsq/frames/response.rb: -------------------------------------------------------------------------------- 1 | require_relative 'frame' 2 | 3 | module Nsq 4 | class Response < Frame 5 | end 6 | end 7 | -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | source 'http://rubygems.org' 2 | 3 | group :development do 4 | gem 'bundler', '~> 1.0' 5 | gem 'jeweler', '~> 2.0.1' 6 | gem 'nsq-cluster', '~> 1.1.0' 7 | gem 'rspec', '~> 3.0.0' 8 | end 9 | -------------------------------------------------------------------------------- /lib/version.rb: -------------------------------------------------------------------------------- 1 | module Nsq 2 | module Version 3 | MAJOR = 0 4 | MINOR = 1 5 | PATCH = 0 6 | BUILD = nil 7 | STRING = [MAJOR, MINOR, PATCH, BUILD].compact.join('.') 8 | end 9 | end 10 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # rcov generated 2 | coverage 3 | coverage.data 4 | 5 | # rdoc generated 6 | rdoc 7 | 8 | # yard generated 9 | doc 10 | .yardoc 11 | 12 | # bundler 13 | .bundle 14 | 15 | # jeweler generated 16 | pkg 17 | 18 | # for Mac 19 | .DS_Store 20 | -------------------------------------------------------------------------------- /lib/nsq.rb: -------------------------------------------------------------------------------- 1 | require_relative 'version' 2 | 3 | require_relative 'nsq/logger' 4 | 5 | require_relative 'nsq/frames/frame' 6 | require_relative 'nsq/frames/error' 7 | require_relative 'nsq/frames/response' 8 | require_relative 'nsq/frames/message' 9 | 10 | require_relative 'nsq/consumer' 11 | require_relative 'nsq/producer' 12 | -------------------------------------------------------------------------------- /lib/nsq/frames/frame.rb: -------------------------------------------------------------------------------- 1 | require_relative '../logger' 2 | 3 | module Nsq 4 | class Frame 5 | include Nsq::AttributeLogger 6 | @@log_attributes = [:connection] 7 | 8 | attr_reader :data 9 | attr_reader :connection 10 | 11 | def initialize(data, connection) 12 | @data = data 13 | @connection = connection 14 | end 15 | end 16 | end 17 | -------------------------------------------------------------------------------- /lib/nsq/frames/message.rb: -------------------------------------------------------------------------------- 1 | require_relative 'frame' 2 | 3 | module Nsq 4 | class Message < Frame 5 | 6 | attr_reader :timestamp 7 | attr_reader :attempts 8 | attr_reader :id 9 | attr_reader :body 10 | 11 | def initialize(data, connection) 12 | super 13 | @timestamp, @attempts, @id, @body = @data.unpack('Q>s>a16a*') 14 | @body.force_encoding('UTF-8') 15 | end 16 | 17 | def finish 18 | connection.fin(id) 19 | end 20 | 21 | def requeue(timeout = 0) 22 | connection.req(id, timeout) 23 | end 24 | 25 | def touch 26 | connection.touch(id) 27 | end 28 | 29 | end 30 | end 31 | -------------------------------------------------------------------------------- /lib/nsq/logger.rb: -------------------------------------------------------------------------------- 1 | require 'logger' 2 | module Nsq 3 | @@logger = Logger.new(nil) 4 | 5 | 6 | def self.logger 7 | @@logger 8 | end 9 | 10 | 11 | def self.logger=(new_logger) 12 | @@logger = new_logger 13 | end 14 | 15 | 16 | module AttributeLogger 17 | def self.included(klass) 18 | klass.send :class_variable_set, :@@log_attributes, [] 19 | end 20 | 21 | %w(fatal error warn info debug).map{|m| m.to_sym}.each do |level| 22 | define_method level do |msg| 23 | Nsq.logger.send(level, "#{prefix} #{msg}") 24 | end 25 | end 26 | 27 | 28 | private 29 | def prefix 30 | attrs = self.class.send(:class_variable_get, :@@log_attributes) 31 | if attrs.count > 0 32 | "[#{attrs.map{|a| "#{a.to_s}: #{self.send(a)}"}.join(' ')}] " 33 | else 34 | '' 35 | end 36 | end 37 | end 38 | end -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2014 Wistia, Inc. 2 | 3 | MIT License 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining 6 | a copy of this software and associated documentation files (the 7 | "Software"), to deal in the Software without restriction, including 8 | without limitation the rights to use, copy, modify, merge, publish, 9 | distribute, sublicense, and/or sell copies of the Software, and to 10 | permit persons to whom the Software is furnished to do so, subject to 11 | the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be 14 | included in all copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 19 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 20 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 22 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 23 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require 'rubygems' 2 | require 'bundler' 3 | begin 4 | Bundler.setup(:default, :development) 5 | rescue Bundler::BundlerError => e 6 | $stderr.puts e.message 7 | $stderr.puts 'Run `bundle install` to install missing gems' 8 | exit e.status_code 9 | end 10 | 11 | require 'rake' 12 | 13 | ########### 14 | # Jeweler # 15 | ########### 16 | require 'jeweler' 17 | require_relative 'lib/version' 18 | Jeweler::Tasks.new do |gem| 19 | # gem is a Gem::Specification... see http://guides.rubygems.org/specification-reference/ for more options 20 | gem.name = "nsq-ruby" 21 | gem.version = Nsq::Version::STRING 22 | gem.homepage = "http://github.com/wistia/nsq-ruby" 23 | gem.license = "MIT" 24 | gem.summary = %Q{Ruby client library for NSQ} 25 | gem.description = %Q{} 26 | gem.email = "dev@wistia.com" 27 | gem.authors = ["Wistia"] 28 | gem.files = Dir.glob('lib/**/*.rb') + ['README.md'] 29 | # dependencies defined in Gemfile 30 | end 31 | Jeweler::RubygemsDotOrgTasks.new 32 | 33 | ######### 34 | # Rspec # 35 | ######### 36 | require 'rspec/core' 37 | require 'rspec/core/rake_task' 38 | RSpec::Core::RakeTask.new(:spec) do |spec| 39 | spec.pattern = FileList['spec/**/*_spec.rb'] 40 | end 41 | -------------------------------------------------------------------------------- /Gemfile.lock: -------------------------------------------------------------------------------- 1 | GEM 2 | remote: http://rubygems.org/ 3 | specs: 4 | addressable (2.3.6) 5 | builder (3.2.2) 6 | descendants_tracker (0.0.4) 7 | thread_safe (~> 0.3, >= 0.3.1) 8 | diff-lcs (1.2.5) 9 | faraday (0.9.0) 10 | multipart-post (>= 1.2, < 3) 11 | git (1.2.7) 12 | github_api (0.11.3) 13 | addressable (~> 2.3) 14 | descendants_tracker (~> 0.0.1) 15 | faraday (~> 0.8, < 0.10) 16 | hashie (>= 1.2) 17 | multi_json (>= 1.7.5, < 2.0) 18 | nokogiri (~> 1.6.0) 19 | oauth2 20 | hashie (3.2.0) 21 | highline (1.6.21) 22 | jeweler (2.0.1) 23 | builder 24 | bundler (>= 1.0) 25 | git (>= 1.2.5) 26 | github_api 27 | highline (>= 1.6.15) 28 | nokogiri (>= 1.5.10) 29 | rake 30 | rdoc 31 | json (1.8.1) 32 | jwt (1.0.0) 33 | mini_portile (0.6.0) 34 | multi_json (1.10.1) 35 | multi_xml (0.5.5) 36 | multipart-post (2.0.0) 37 | nokogiri (1.6.3.1) 38 | mini_portile (= 0.6.0) 39 | nsq-cluster (1.1.0) 40 | sys-proctable 41 | oauth2 (1.0.0) 42 | faraday (>= 0.8, < 0.10) 43 | jwt (~> 1.0) 44 | multi_json (~> 1.3) 45 | multi_xml (~> 0.5) 46 | rack (~> 1.2) 47 | rack (1.5.2) 48 | rake (10.3.2) 49 | rdoc (4.1.1) 50 | json (~> 1.4) 51 | rspec (3.0.0) 52 | rspec-core (~> 3.0.0) 53 | rspec-expectations (~> 3.0.0) 54 | rspec-mocks (~> 3.0.0) 55 | rspec-core (3.0.3) 56 | rspec-support (~> 3.0.0) 57 | rspec-expectations (3.0.3) 58 | diff-lcs (>= 1.2.0, < 2.0) 59 | rspec-support (~> 3.0.0) 60 | rspec-mocks (3.0.3) 61 | rspec-support (~> 3.0.0) 62 | rspec-support (3.0.3) 63 | sys-proctable (0.9.4) 64 | thread_safe (0.3.4) 65 | 66 | PLATFORMS 67 | ruby 68 | 69 | DEPENDENCIES 70 | bundler (~> 1.0) 71 | jeweler (~> 2.0.1) 72 | nsq-cluster (~> 1.1.0) 73 | rspec (~> 3.0.0) 74 | -------------------------------------------------------------------------------- /lib/nsq/producer.rb: -------------------------------------------------------------------------------- 1 | require_relative 'client_base' 2 | require_relative 'connection' 3 | require_relative 'logger' 4 | 5 | module Nsq 6 | class Producer < ClientBase 7 | include Nsq::AttributeLogger 8 | @@log_attributes = [:host, :port, :topic] 9 | 10 | attr_reader :host 11 | attr_reader :port 12 | attr_reader :topic 13 | 14 | def initialize(opts = {}) 15 | @connections = {} 16 | @topic = opts[:topic] || raise(ArgumentError, 'topic is required') 17 | @discovery_interval = opts[:discovery_interval] || 60 18 | 19 | nsqlookupds = [] 20 | if opts[:nsqlookupd] 21 | nsqlookupds = [opts[:nsqlookupd]].flatten 22 | @discovery = Discovery.new(nsqlookupds) 23 | discover_repeatedly(discover_by_topic: false) 24 | 25 | elsif opts[:nsqd] 26 | nsqds = [opts[:nsqd]].flatten 27 | nsqds.each{|d| add_connection(d)} 28 | 29 | else 30 | add_connection('127.0.0.1:4150') 31 | end 32 | 33 | at_exit{terminate} 34 | end 35 | 36 | 37 | def write(*raw_messages) 38 | # stringify the messages 39 | messages = raw_messages.map(&:to_s) 40 | 41 | # get a suitable connection to write to 42 | connection = connection_for_write 43 | 44 | if messages.length > 1 45 | connection.mpub(@topic, messages) 46 | else 47 | connection.pub(@topic, messages.first) 48 | end 49 | end 50 | 51 | 52 | private 53 | def connection_for_write 54 | # Choose a random Connection that's currently connected 55 | # Or, if there's nothing connected, just take any random one 56 | connections_currently_connected = connections.select{|_,c| c.connected?} 57 | connection = connections_currently_connected.values.sample || connections.values.sample 58 | 59 | # Raise an exception if there's no connection available 60 | unless connection 61 | raise 'No connections available' 62 | end 63 | 64 | connection 65 | end 66 | 67 | end 68 | end 69 | -------------------------------------------------------------------------------- /nsq-ruby.gemspec: -------------------------------------------------------------------------------- 1 | # Generated by jeweler 2 | # DO NOT EDIT THIS FILE DIRECTLY 3 | # Instead, edit Jeweler::Tasks in Rakefile, and run 'rake gemspec' 4 | # -*- encoding: utf-8 -*- 5 | # stub: nsq-ruby 0.1.0 ruby lib 6 | 7 | Gem::Specification.new do |s| 8 | s.name = "nsq-ruby" 9 | s.version = "0.1.0" 10 | 11 | s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version= 12 | s.require_paths = ["lib"] 13 | s.authors = ["Wistia"] 14 | s.date = "2014-07-28" 15 | s.description = "" 16 | s.email = "dev@wistia.com" 17 | s.extra_rdoc_files = [ 18 | "LICENSE.txt", 19 | "README.md" 20 | ] 21 | s.files = [ 22 | "README.md", 23 | "lib/nsq.rb", 24 | "lib/nsq/connection.rb", 25 | "lib/nsq/consumer.rb", 26 | "lib/nsq/discovery.rb", 27 | "lib/nsq/frames/error.rb", 28 | "lib/nsq/frames/frame.rb", 29 | "lib/nsq/frames/message.rb", 30 | "lib/nsq/frames/response.rb", 31 | "lib/nsq/logger.rb", 32 | "lib/nsq/producer.rb", 33 | "lib/version.rb" 34 | ] 35 | s.homepage = "http://github.com/wistia/nsq-ruby" 36 | s.licenses = ["MIT"] 37 | s.rubygems_version = "2.2.2" 38 | s.summary = "Ruby client library for NSQ" 39 | 40 | if s.respond_to? :specification_version then 41 | s.specification_version = 4 42 | 43 | if Gem::Version.new(Gem::VERSION) >= Gem::Version.new('1.2.0') then 44 | s.add_development_dependency(%q, ["~> 1.0"]) 45 | s.add_development_dependency(%q, ["~> 2.0.1"]) 46 | s.add_development_dependency(%q, ["~> 0.2.7"]) 47 | s.add_development_dependency(%q, ["~> 3.0.0"]) 48 | else 49 | s.add_dependency(%q, ["~> 1.0"]) 50 | s.add_dependency(%q, ["~> 2.0.1"]) 51 | s.add_dependency(%q, ["~> 0.2.7"]) 52 | s.add_dependency(%q, ["~> 3.0.0"]) 53 | end 54 | else 55 | s.add_dependency(%q, ["~> 1.0"]) 56 | s.add_dependency(%q, ["~> 2.0.1"]) 57 | s.add_dependency(%q, ["~> 0.2.7"]) 58 | s.add_dependency(%q, ["~> 3.0.0"]) 59 | end 60 | end 61 | 62 | -------------------------------------------------------------------------------- /lib/nsq/discovery.rb: -------------------------------------------------------------------------------- 1 | require 'json' 2 | require 'net/http' 3 | require 'uri' 4 | 5 | require_relative 'logger' 6 | 7 | # Connects to nsqlookup's to find the nsqd instances for a given topic 8 | module Nsq 9 | class Discovery 10 | include Nsq::AttributeLogger 11 | 12 | # lookupd addresses must be formatted like so: ':' 13 | def initialize(lookupds) 14 | @lookupds = lookupds 15 | end 16 | 17 | # Returns an array of nsqds instances 18 | # 19 | # nsqd instances returned are strings in this format: ':' 20 | # 21 | # discovery.nsqds 22 | # #=> ['127.0.0.1:4150', '127.0.0.1:4152'] 23 | # 24 | def nsqds 25 | @lookupds.map do |lookupd| 26 | get_nsqds(lookupd) 27 | end.flatten.uniq 28 | end 29 | 30 | # Returns an array of nsqds instances that have messages for 31 | # that topic. 32 | # 33 | # nsqd instances returned are strings in this format: ':' 34 | # 35 | # discovery.nsqds_for_topic('a-topic') 36 | # #=> ['127.0.0.1:4150', '127.0.0.1:4152'] 37 | # 38 | def nsqds_for_topic(topic) 39 | @lookupds.map do |lookupd| 40 | get_nsqds(lookupd, topic) 41 | end.flatten.uniq 42 | end 43 | 44 | private 45 | 46 | def get_nsqds(lookupd, topic = nil) 47 | uri_scheme = 'http://' unless lookupd.match(%r(https?://)) 48 | uri = URI.parse("#{uri_scheme}#{lookupd}") 49 | 50 | uri.query = "ts=#{Time.now.to_i}" 51 | if topic 52 | uri.path = '/lookup' 53 | uri.query += "&topic=#{topic}" 54 | else 55 | uri.path = '/nodes' 56 | end 57 | 58 | begin 59 | body = Net::HTTP.get(uri) 60 | data = JSON.parse(body) 61 | 62 | if data['data'] && data['data']['producers'] 63 | data['data']['producers'].map do |producer| 64 | "#{producer['broadcast_address']}:#{producer['tcp_port']}" 65 | end 66 | else 67 | [] 68 | end 69 | rescue Exception => e 70 | error "Error during discovery for #{lookupd}: #{e}" 71 | [] 72 | end 73 | end 74 | 75 | end 76 | end 77 | -------------------------------------------------------------------------------- /lib/nsq/client_base.rb: -------------------------------------------------------------------------------- 1 | module Nsq 2 | class ClientBase 3 | 4 | attr_reader :connections 5 | 6 | def connected? 7 | @connections.values.any?(&:connected?) 8 | end 9 | 10 | 11 | def terminate 12 | @discovery_thread.kill if @discovery_thread 13 | drop_all_connections 14 | end 15 | 16 | private 17 | 18 | # discovers nsqds from an nsqlookupd repeatedly 19 | # 20 | # opts: 21 | # discover_by_topic: true 22 | # 23 | def discover_repeatedly(opts = {}) 24 | @discovery_thread = Thread.new do 25 | loop do 26 | discover opts 27 | sleep @discovery_interval 28 | end 29 | end 30 | @discovery_thread.abort_on_exception = true 31 | end 32 | 33 | 34 | def discover(opts) 35 | nsqds = nil 36 | if opts[:discover_by_topic] 37 | nsqds = @discovery.nsqds_for_topic(@topic) 38 | else 39 | nsqds = @discovery.nsqds 40 | end 41 | 42 | # drop nsqd connections that are no longer in lookupd 43 | missing_nsqds = @connections.keys - nsqds 44 | missing_nsqds.each do |nsqd| 45 | drop_connection(nsqd) 46 | end 47 | 48 | # add new ones 49 | new_nsqds = nsqds - @connections.keys 50 | new_nsqds.each do |nsqd| 51 | add_connection(nsqd) 52 | end 53 | 54 | # balance RDY state amongst the connections 55 | connections_changed 56 | end 57 | 58 | 59 | def add_connection(nsqd) 60 | info "+ Adding connection #{nsqd}" 61 | host, port = nsqd.split(':') 62 | connection = Connection.new( 63 | host: host, 64 | port: port 65 | ) 66 | @connections[nsqd] = connection 67 | end 68 | 69 | 70 | def drop_connection(nsqd) 71 | info "- Dropping connection #{nsqd}" 72 | connection = @connections.delete(nsqd) 73 | connection.close if connection 74 | connections_changed 75 | end 76 | 77 | 78 | def drop_all_connections 79 | @connections.keys.each do |nsqd| 80 | drop_connection(nsqd) 81 | end 82 | end 83 | 84 | # optional subclass hook 85 | def connections_changed 86 | end 87 | end 88 | end 89 | -------------------------------------------------------------------------------- /spec/spec_helper.rb: -------------------------------------------------------------------------------- 1 | require 'nsq-cluster' 2 | 3 | $LOAD_PATH.unshift(File.join(File.dirname(__FILE__), '..', 'lib')) 4 | $LOAD_PATH.unshift(File.dirname(__FILE__)) 5 | 6 | require 'rspec' 7 | require 'nsq' 8 | 9 | # Requires supporting files with custom matchers and macros, etc, 10 | # in ./support/ and its subdirectories. 11 | Dir["#{File.dirname(__FILE__)}/support/**/*.rb"].each {|f| require f} 12 | 13 | RSpec.configure do |config| 14 | config.before(:suite) do 15 | Nsq.logger = Logger.new(STDOUT) if ENV['VERBOSE'] 16 | end 17 | end 18 | 19 | 20 | require 'timeout' 21 | def assert_no_timeout(time = 1, &block) 22 | expect{ 23 | Timeout::timeout(time) do 24 | yield 25 | end 26 | }.not_to raise_error 27 | end 28 | 29 | def assert_timeout(time = 1, &block) 30 | expect{ 31 | Timeout::timeout(time) do 32 | yield 33 | end 34 | }.to raise_error(Timeout::Error) 35 | end 36 | 37 | # Block execution until a condition is met 38 | # Times out after 5 seconds by default 39 | # 40 | # example: 41 | # wait_for { @consumer.queue.length > 0 } 42 | # 43 | def wait_for(timeout = 5, &block) 44 | Timeout::timeout(timeout) do 45 | loop do 46 | break if yield 47 | sleep(0.1) 48 | end 49 | end 50 | end 51 | 52 | TOPIC = 'some-topic' 53 | CHANNEL = 'some-channel' 54 | 55 | def new_consumer(opts = {}) 56 | lookupd = @cluster.nsqlookupd.map{|l| "#{l.host}:#{l.http_port}"} 57 | Nsq::Consumer.new({ 58 | topic: TOPIC, 59 | channel: CHANNEL, 60 | nsqlookupd: lookupd, 61 | max_in_flight: 1 62 | }.merge(opts)) 63 | end 64 | 65 | 66 | def new_producer(nsqd, opts = {}) 67 | Nsq::Producer.new({ 68 | topic: TOPIC, 69 | nsqd: "#{nsqd.host}:#{nsqd.tcp_port}", 70 | discovery_interval: 1 71 | }.merge(opts)) 72 | end 73 | 74 | def new_lookupd_producer(opts = {}) 75 | lookupd = @cluster.nsqlookupd.map{|l| "#{l.host}:#{l.http_port}"} 76 | Nsq::Producer.new({ 77 | topic: TOPIC, 78 | nsqlookupd: lookupd, 79 | discovery_interval: 1 80 | }.merge(opts)) 81 | end 82 | 83 | # This is for certain spots where we're testing connections going up and down. 84 | # Don't want these tests to take forever to run! 85 | def set_speedy_connection_timeouts! 86 | allow_any_instance_of(Nsq::Connection).to receive(:snooze).and_return(nil) 87 | end 88 | -------------------------------------------------------------------------------- /spec/lib/nsq/discovery_spec.rb: -------------------------------------------------------------------------------- 1 | require_relative '../../spec_helper' 2 | 3 | NSQD_COUNT = 5 4 | 5 | describe Nsq::Discovery do 6 | before do 7 | @cluster = NsqCluster.new(nsqd_count: NSQD_COUNT, nsqlookupd_count: 2) 8 | @topic = 'some-topic' 9 | 10 | # make sure each nsqd has a message for this topic 11 | # leave the last nsqd without this topic for testing 12 | @cluster.nsqd.take(NSQD_COUNT-1).each do |nsqd| 13 | nsqd.pub(@topic, 'some-message') 14 | end 15 | @cluster.nsqd.last.pub('some-other-topic', 'some-message') 16 | 17 | @expected_topic_lookup_nsqds = @cluster.nsqd.take(NSQD_COUNT-1).map{|d|"#{d.host}:#{d.tcp_port}"}.sort 18 | @expected_all_nsqds = @cluster.nsqd.map{|d|"#{d.host}:#{d.tcp_port}"}.sort 19 | end 20 | 21 | after do 22 | @cluster.destroy 23 | end 24 | 25 | 26 | def new_discovery(cluster_lookupds) 27 | lookupds = cluster_lookupds.map do |lookupd| 28 | "#{lookupd.host}:#{lookupd.http_port}" 29 | end 30 | 31 | # one lookupd has scheme and one does not 32 | lookupds.last.prepend 'http://' 33 | 34 | Nsq::Discovery.new(lookupds) 35 | end 36 | 37 | 38 | describe 'a single nsqlookupd' do 39 | before do 40 | @discovery = new_discovery([@cluster.nsqlookupd.first]) 41 | end 42 | 43 | describe '#nsqds' do 44 | it 'returns all nsqds' do 45 | nsqds = @discovery.nsqds 46 | expect(nsqds.sort).to eq(@expected_all_nsqds) 47 | end 48 | end 49 | 50 | describe '#nsqds_for_topic' do 51 | it 'returns [] for a topic that doesn\'t exist' do 52 | nsqds = @discovery.nsqds_for_topic('topic-that-does-not-exists') 53 | expect(nsqds).to eq([]) 54 | end 55 | 56 | it 'returns all nsqds' do 57 | nsqds = @discovery.nsqds_for_topic(@topic) 58 | expect(nsqds.sort).to eq(@expected_topic_lookup_nsqds) 59 | end 60 | end 61 | end 62 | 63 | 64 | describe 'multiple nsqlookupds' do 65 | before do 66 | @discovery = new_discovery(@cluster.nsqlookupd) 67 | end 68 | 69 | describe '#nsqds_for_topic' do 70 | it 'returns all nsqds' do 71 | nsqds = @discovery.nsqds_for_topic(@topic) 72 | expect(nsqds.sort).to eq(@expected_topic_lookup_nsqds) 73 | end 74 | end 75 | end 76 | 77 | 78 | describe 'multiple nsqlookupds, but one is down' do 79 | before do 80 | @downed_nsqlookupd = @cluster.nsqlookupd.first 81 | @downed_nsqlookupd.stop 82 | 83 | @discovery = new_discovery(@cluster.nsqlookupd) 84 | end 85 | 86 | describe '#nsqds_for_topic' do 87 | it 'returns all nsqds' do 88 | nsqds = @discovery.nsqds_for_topic(@topic) 89 | expect(nsqds.sort).to eq(@expected_topic_lookup_nsqds) 90 | end 91 | end 92 | end 93 | end 94 | -------------------------------------------------------------------------------- /spec/lib/nsq/connection_spec.rb: -------------------------------------------------------------------------------- 1 | require_relative '../../spec_helper' 2 | 3 | describe Nsq::Connection do 4 | before do 5 | @cluster = NsqCluster.new(nsqd_count: 1) 6 | @nsqd = @cluster.nsqd.first 7 | @connection = Nsq::Connection.new(host: @cluster.nsqd[0].host, port: @cluster.nsqd[0].tcp_port) 8 | end 9 | after do 10 | @connection.close 11 | @cluster.destroy 12 | end 13 | 14 | 15 | describe '::new' do 16 | it 'should raise an exception if it cannot connect to nsqd' do 17 | @nsqd.stop 18 | 19 | expect{ 20 | Nsq::Connection.new(host: @nsqd.host, port: @nsqd.tcp_port) 21 | }.to raise_error 22 | end 23 | 24 | it 'should raise an exception if it connects to something that isn\'t nsqd' do 25 | expect{ 26 | # try to connect to the HTTP port instead of TCP 27 | Nsq::Connection.new(host: @nsqd.host, port: @nsqd.http_port) 28 | }.to raise_error 29 | end 30 | end 31 | 32 | 33 | describe '#close' do 34 | it 'can be called multiple times, without issue' do 35 | expect{ 36 | 10.times{@connection.close} 37 | }.not_to raise_error 38 | end 39 | end 40 | 41 | 42 | # This is really testing the ability for Connection to reconnect 43 | describe '#connected?' do 44 | before do 45 | # For speedier timeouts 46 | set_speedy_connection_timeouts! 47 | end 48 | 49 | it 'should return true when nsqd is up and false when nsqd is down' do 50 | wait_for{@connection.connected?} 51 | expect(@connection.connected?).to eq(true) 52 | @nsqd.stop 53 | wait_for{!@connection.connected?} 54 | expect(@connection.connected?).to eq(false) 55 | @nsqd.start 56 | wait_for{@connection.connected?} 57 | expect(@connection.connected?).to eq(true) 58 | end 59 | 60 | end 61 | 62 | 63 | describe 'private methods' do 64 | describe '#frame_class_for_type' do 65 | MAX_VALID_TYPE = described_class::FRAME_CLASSES.length - 1 66 | it "returns a frame class for types 0-#{MAX_VALID_TYPE}" do 67 | (0..MAX_VALID_TYPE).each do |type| 68 | expect( 69 | described_class::FRAME_CLASSES.include?( 70 | @connection.send(:frame_class_for_type, type) 71 | ) 72 | ).to be_truthy 73 | end 74 | end 75 | it "raises an error if invalid type > #{MAX_VALID_TYPE} specified" do 76 | expect { 77 | @connection.send(:frame_class_for_type, 3) 78 | }.to raise_error(RuntimeError) 79 | end 80 | end 81 | 82 | 83 | describe '#handle_response' do 84 | it 'responds to heartbeat with NOP' do 85 | frame = Nsq::Response.new(described_class::RESPONSE_HEARTBEAT, @connection) 86 | expect(@connection).to receive(:nop) 87 | @connection.send(:handle_response, frame) 88 | end 89 | end 90 | end 91 | end 92 | -------------------------------------------------------------------------------- /lib/nsq/consumer.rb: -------------------------------------------------------------------------------- 1 | require_relative 'client_base' 2 | require_relative 'connection' 3 | require_relative 'discovery' 4 | require_relative 'logger' 5 | 6 | module Nsq 7 | class Consumer < ClientBase 8 | include Nsq::AttributeLogger 9 | @@log_attributes = [:topic] 10 | 11 | attr_reader :topic 12 | attr_reader :max_in_flight 13 | attr_reader :discovery_interval 14 | 15 | def initialize(opts = {}) 16 | if opts[:nsqlookupd] 17 | @nsqlookupds = [opts[:nsqlookupd]].flatten 18 | else 19 | @nsqlookupds = [] 20 | end 21 | 22 | @topic = opts[:topic] || raise(ArgumentError, 'topic is required') 23 | @channel = opts[:channel] || raise(ArgumentError, 'channel is required') 24 | @max_in_flight = opts[:max_in_flight] || 1 25 | @discovery_interval = opts[:discovery_interval] || 60 26 | @msg_timeout = opts[:msg_timeout] 27 | 28 | # This is where we queue up the messages we receive from each connection 29 | @messages = Queue.new 30 | 31 | # This is where we keep a record of our active nsqd connections 32 | # The key is a string with the host and port of the instance (e.g. 33 | # '127.0.0.1:4150') and the key is the Connection instance. 34 | @connections = {} 35 | 36 | if !@nsqlookupds.empty? 37 | @discovery = Discovery.new(@nsqlookupds) 38 | discover_repeatedly(discover_by_topic: true) 39 | else 40 | # normally, we find nsqd instances to connect to via nsqlookupd(s) 41 | # in this case let's connect to an nsqd instance directly 42 | add_connection(opts[:nsqd] || '127.0.0.1:4150', @max_in_flight) 43 | end 44 | 45 | at_exit{terminate} 46 | end 47 | 48 | 49 | # pop the next message off the queue 50 | def pop 51 | @messages.pop 52 | end 53 | 54 | 55 | # returns the number of messages we have locally in the queue 56 | def size 57 | @messages.size 58 | end 59 | 60 | 61 | private 62 | def add_connection(nsqd, max_in_flight = 1) 63 | info "+ Adding connection #{nsqd}" 64 | host, port = nsqd.split(':') 65 | connection = Connection.new( 66 | host: host, 67 | port: port, 68 | topic: @topic, 69 | channel: @channel, 70 | queue: @messages, 71 | msg_timeout: @msg_timeout, 72 | max_in_flight: max_in_flight 73 | ) 74 | @connections[nsqd] = connection 75 | end 76 | 77 | # Be conservative, but don't set a connection's max_in_flight below 1 78 | def max_in_flight_per_connection(number_of_connections = @connections.length) 79 | [@max_in_flight / number_of_connections, 1].max 80 | end 81 | 82 | def connections_changed 83 | redistribute_ready 84 | end 85 | 86 | def redistribute_ready 87 | @connections.values.each do |connection| 88 | connection.max_in_flight = max_in_flight_per_connection 89 | connection.re_up_ready 90 | end 91 | end 92 | end 93 | end 94 | -------------------------------------------------------------------------------- /spec/lib/nsq/consumer_with_shaky_connections_spec.rb: -------------------------------------------------------------------------------- 1 | require_relative '../../spec_helper' 2 | 3 | describe Nsq::Consumer do 4 | 5 | before do 6 | set_speedy_connection_timeouts! 7 | 8 | @nsqd_count = 3 9 | @cluster = NsqCluster.new(nsqlookupd_count: 2, nsqd_count: @nsqd_count) 10 | 11 | # publish a message to each queue 12 | # so that when the consumer starts up, it will connect to all of them 13 | @cluster.nsqd.each do |nsqd| 14 | nsqd.pub(TOPIC, 'hi') 15 | end 16 | 17 | @consumer = new_consumer(max_in_flight: 20, discovery_interval: 0.1) 18 | wait_for{@consumer.connections.length == @nsqd_count} 19 | end 20 | 21 | after do 22 | @consumer.terminate 23 | @cluster.destroy 24 | end 25 | 26 | 27 | # This is really testing that the discovery loop works as expected. 28 | # 29 | # The consumer won't evict connections if they go down, the connection itself 30 | # will try to reconnect. 31 | # 32 | # But, when nsqd goes down, nsqlookupd will see that its gone and unregister 33 | # it. So when the next time the discovery loop runs, that nsqd will no longer 34 | # be listed. 35 | it 'should drop a connection when an nsqd goes down and add one when it comes back' do 36 | @cluster.nsqd.last.stop 37 | wait_for{@consumer.connections.length == @nsqd_count - 1} 38 | 39 | @cluster.nsqd.last.start 40 | wait_for{@consumer.connections.length == @nsqd_count} 41 | end 42 | 43 | 44 | it 'should continue processing messages from live queues when one queue is down' do 45 | # shut down the last nsqd 46 | @cluster.nsqd.last.stop 47 | 48 | # make sure there are more messages on each queue than max in flight 49 | 50.times{@cluster.nsqd[0].pub(TOPIC, 'hay')} 50 | 50.times{@cluster.nsqd[1].pub(TOPIC, 'hay')} 51 | 52 | assert_no_timeout(5) do 53 | 100.times{@consumer.pop.finish} 54 | end 55 | end 56 | 57 | 58 | it 'should process messages from a new queue when it comes online' do 59 | nsqd = @cluster.nsqd.last 60 | nsqd.stop 61 | 62 | thread = Thread.new do 63 | nsqd.start 64 | nsqd.pub(TOPIC, 'needle') 65 | end 66 | 67 | assert_no_timeout(5) do 68 | string = nil 69 | until string == 'needle' 70 | msg = @consumer.pop 71 | string = msg.body 72 | msg.finish 73 | end 74 | true 75 | end 76 | 77 | thread.join 78 | end 79 | 80 | 81 | it 'should be able to rely on the second nsqlookupd if the first dies' do 82 | bad_lookupd = @cluster.nsqlookupd.first 83 | bad_lookupd.stop 84 | 85 | @cluster.nsqd.first.pub('new-topic', 'new message on new topic') 86 | consumer = new_consumer(topic: 'new-topic') 87 | 88 | assert_no_timeout do 89 | msg = consumer.pop 90 | expect(msg.body).to eq('new message on new topic') 91 | msg.finish 92 | end 93 | 94 | consumer.terminate 95 | end 96 | 97 | 98 | it 'should be able to handle all queues going offline and coming back' do 99 | expected_messages = @cluster.nsqd.map{|nsqd| nsqd.tcp_port.to_s} 100 | 101 | @cluster.nsqd.each{|q| q.stop} 102 | @cluster.nsqd.each{|q| q.start} 103 | 104 | @cluster.nsqd.each_with_index do |nsqd, idx| 105 | nsqd.pub(TOPIC, expected_messages[idx]) 106 | end 107 | 108 | assert_no_timeout(10) do 109 | received_messages = [] 110 | 111 | while (expected_messages & received_messages).length < expected_messages.length do 112 | msg = @consumer.pop 113 | received_messages << msg.body 114 | msg.finish 115 | end 116 | 117 | # ladies and gentlemen, we got 'em 118 | end 119 | end 120 | 121 | end 122 | 123 | -------------------------------------------------------------------------------- /spec/lib/nsq/producer_spec.rb: -------------------------------------------------------------------------------- 1 | require_relative '../../spec_helper' 2 | require 'json' 3 | 4 | describe Nsq::Producer do 5 | 6 | def message_count 7 | topics_info = JSON.parse(@nsqd.stats.body)['data']['topics'] 8 | topic_info = topics_info.select{|t| t['topic_name'] == @producer.topic }.first 9 | if topic_info 10 | topic_info['message_count'] 11 | else 12 | 0 13 | end 14 | end 15 | 16 | context 'connecting directly to a single nsqd' do 17 | 18 | def new_consumer 19 | Nsq::Consumer.new( 20 | topic: TOPIC, 21 | channel: CHANNEL, 22 | nsqd: "#{@nsqd.host}:#{@nsqd.tcp_port}", 23 | max_in_flight: 1 24 | ) 25 | end 26 | 27 | before do 28 | @cluster = NsqCluster.new(nsqd_count: 1) 29 | @nsqd = @cluster.nsqd.first 30 | @producer = new_producer(@nsqd) 31 | end 32 | 33 | after do 34 | @producer.terminate if @producer 35 | @cluster.destroy 36 | end 37 | 38 | describe '::new' do 39 | it 'should throw an exception when trying to connect to a server that\'s down' do 40 | @nsqd.stop 41 | 42 | expect{ 43 | new_producer(@nsqd) 44 | }.to raise_error 45 | end 46 | end 47 | 48 | describe '#connected?' do 49 | it 'should return true when it is connected' do 50 | expect(@producer.connected?).to eq(true) 51 | end 52 | 53 | it 'should return false when nsqd is down' do 54 | @nsqd.stop 55 | wait_for{!@producer.connected?} 56 | expect(@producer.connected?).to eq(false) 57 | end 58 | end 59 | 60 | describe '#write' do 61 | 62 | it 'can queue a message' do 63 | @producer.write('some-message') 64 | wait_for{message_count==1} 65 | expect(message_count).to eq(1) 66 | end 67 | 68 | it 'can queue multiple messages at once' do 69 | @producer.write(1, 2, 3, 4, 5, 6, 7, 8, 9, 10) 70 | wait_for{message_count==10} 71 | expect(message_count).to eq(10) 72 | end 73 | 74 | it 'shouldn\'t raise an error when nsqd is down' do 75 | @nsqd.stop 76 | 77 | expect{ 78 | 10.times{@producer.write('fail')} 79 | }.to_not raise_error 80 | end 81 | 82 | it 'will attempt to resend messages when it reconnects to nsqd' do 83 | @nsqd.stop 84 | 85 | # Write 10 messages while nsqd is down 86 | 10.times{|i| @producer.write(i)} 87 | 88 | @nsqd.start 89 | 90 | messages_received = [] 91 | 92 | begin 93 | consumer = new_consumer 94 | assert_no_timeout(5) do 95 | # TODO: make the socket fail faster 96 | # We only get 8 or 9 of the 10 we send. The first few can be lost 97 | # because we can't detect that they didn't make it. 98 | 8.times do |i| 99 | msg = consumer.pop 100 | messages_received << msg.body 101 | msg.finish 102 | end 103 | end 104 | ensure 105 | consumer.terminate 106 | end 107 | 108 | expect(messages_received.uniq.length).to eq(8) 109 | end 110 | 111 | # Test PUB 112 | it 'can send a single message with unicode characters' do 113 | @producer.write('☺') 114 | consumer = new_consumer 115 | assert_no_timeout do 116 | expect(consumer.pop.body).to eq('☺') 117 | end 118 | consumer.terminate 119 | end 120 | 121 | # Test MPUB as well 122 | it 'can send multiple message with unicode characters' do 123 | @producer.write('☺', '☺', '☺') 124 | consumer = new_consumer 125 | assert_no_timeout do 126 | 3.times do 127 | msg = consumer.pop 128 | expect(msg.body).to eq('☺') 129 | msg.finish 130 | end 131 | end 132 | consumer.terminate 133 | end 134 | end 135 | 136 | end 137 | 138 | 139 | context 'connecting via nsqlookupd' do 140 | 141 | before do 142 | @cluster = NsqCluster.new(nsqd_count: 2, nsqlookupd_count: 1) 143 | @producer = new_lookupd_producer 144 | 145 | # wait for it to connect to all nsqds 146 | wait_for{ @producer.connections.length == @cluster.nsqd.length } 147 | end 148 | 149 | after do 150 | @producer.terminate if @producer 151 | @cluster.destroy 152 | end 153 | 154 | 155 | describe '#connections' do 156 | it 'should be connected to all nsqds' do 157 | expect(@producer.connections.length).to eq(@cluster.nsqd.length) 158 | end 159 | 160 | it 'should drop a connection when an nsqd goes offline' do 161 | @cluster.nsqd.first.stop 162 | wait_for{ @producer.connections.length == @cluster.nsqd.length - 1 } 163 | expect(@producer.connections.length).to eq(@cluster.nsqd.length - 1) 164 | end 165 | end 166 | 167 | 168 | describe '#connected?' do 169 | it 'should return true if it\'s connected to at least one nsqd' do 170 | expect(@producer.connected?).to eq(true) 171 | end 172 | 173 | it 'should return false when it\'s not connected to any nsqds' do 174 | @cluster.nsqd.each{|nsqd| nsqd.stop} 175 | wait_for{ !@producer.connected? } 176 | expect(@producer.connected?).to eq(false) 177 | end 178 | end 179 | 180 | 181 | describe '#write' do 182 | it 'writes to a random connection' do 183 | expect_any_instance_of(Nsq::Connection).to receive(:pub) 184 | @producer.write('howdy!') 185 | end 186 | 187 | it 'raises an error if there are no connections to write to' do 188 | @cluster.nsqd.each{|nsqd| nsqd.stop} 189 | wait_for{ @producer.connections.length == 0 } 190 | expect { 191 | @producer.write('die') 192 | }.to raise_error 193 | end 194 | end 195 | 196 | end 197 | end 198 | -------------------------------------------------------------------------------- /spec/lib/nsq/consumer_spec.rb: -------------------------------------------------------------------------------- 1 | require_relative '../../spec_helper' 2 | require 'json' 3 | require 'timeout' 4 | 5 | describe Nsq::Consumer do 6 | before do 7 | @cluster = NsqCluster.new(nsqd_count: 2, nsqlookupd_count: 1) 8 | end 9 | 10 | after do 11 | @cluster.destroy 12 | end 13 | 14 | 15 | describe 'when connecting to nsqd directly' do 16 | before do 17 | @nsqd = @cluster.nsqd.first 18 | @consumer = new_consumer(nsqlookupd: nil, nsqd: "#{@nsqd.host}:#{@nsqd.tcp_port}", max_in_flight: 10) 19 | end 20 | after do 21 | @consumer.terminate 22 | end 23 | 24 | 25 | describe '::new' do 26 | it 'should throw an exception when trying to connect to a server that\'s down' do 27 | @nsqd.stop 28 | 29 | expect{ 30 | new_consumer(nsqlookupd: nil, nsqd: "#{@nsqd.host}:#{@nsqd.tcp_port}") 31 | }.to raise_error 32 | end 33 | end 34 | 35 | 36 | # This is testing the behavior of the consumer, rather than the size method itself 37 | describe '#size' do 38 | it 'doesn\'t exceed max_in_flight for the consumer' do 39 | # publish a bunch of messages 40 | (@consumer.max_in_flight * 2).times do 41 | @nsqd.pub(@consumer.topic, 'some-message') 42 | end 43 | 44 | wait_for{@consumer.size >= @consumer.max_in_flight} 45 | expect(@consumer.size).to eq(@consumer.max_in_flight) 46 | end 47 | end 48 | 49 | 50 | describe '#pop' do 51 | it 'can pop off a message' do 52 | @nsqd.pub(@consumer.topic, 'some-message') 53 | assert_no_timeout(1) do 54 | msg = @consumer.pop 55 | expect(msg.body).to eq('some-message') 56 | msg.finish 57 | end 58 | end 59 | 60 | it 'can pop off many messages' do 61 | 10.times{@nsqd.pub(@consumer.topic, 'some-message')} 62 | assert_no_timeout(1) do 63 | 10.times{@consumer.pop.finish} 64 | end 65 | end 66 | 67 | it 'can receive messages with unicode characters' do 68 | @nsqd.pub(@consumer.topic, '☺') 69 | expect(@consumer.pop.body).to eq('☺') 70 | end 71 | end 72 | 73 | 74 | describe '#req' do 75 | it 'can successfully requeue a message' do 76 | # queue a message 77 | @nsqd.pub(TOPIC, 'twice') 78 | 79 | msg = @consumer.pop 80 | 81 | expect(msg.body).to eq('twice') 82 | 83 | # requeue it 84 | msg.requeue 85 | 86 | req_msg = @consumer.pop 87 | expect(req_msg.body).to eq('twice') 88 | expect(req_msg.attempts).to eq(2) 89 | end 90 | end 91 | end 92 | 93 | 94 | describe 'when using lookupd' do 95 | before do 96 | @expected_messages = (1..20).to_a.map(&:to_s) 97 | @expected_messages.each_with_index do |message, idx| 98 | @cluster.nsqd[idx % @cluster.nsqd.length].pub(TOPIC, message) 99 | end 100 | 101 | @consumer = new_consumer(max_in_flight: 10) 102 | end 103 | 104 | after do 105 | @consumer.terminate 106 | end 107 | 108 | describe '#pop' do 109 | it 'receives messages from both queues' do 110 | received_messages = [] 111 | 112 | # gather all the messages 113 | assert_no_timeout(2) do 114 | @expected_messages.length.times do 115 | msg = @consumer.pop 116 | received_messages << msg.body 117 | msg.finish 118 | end 119 | end 120 | 121 | expect(received_messages.sort).to eq(@expected_messages.sort) 122 | end 123 | end 124 | 125 | # This is testing the behavior of the consumer, rather than the size method itself 126 | describe '#size' do 127 | it 'doesn\'t exceed max_in_flight for the consumer' do 128 | wait_for{@consumer.size >= @consumer.max_in_flight} 129 | expect(@consumer.size).to eq(@consumer.max_in_flight) 130 | end 131 | end 132 | end 133 | 134 | 135 | describe 'with a low message timeout' do 136 | before do 137 | @nsqd = @cluster.nsqd.first 138 | @msg_timeout = 1 139 | @consumer = new_consumer( 140 | nsqlookupd: nil, 141 | nsqd: "#{@nsqd.host}:#{@nsqd.tcp_port}", 142 | msg_timeout: @msg_timeout * 1000 # in milliseconds 143 | ) 144 | end 145 | after do 146 | @consumer.terminate 147 | end 148 | 149 | 150 | # This testing that our msg_timeout is being honored 151 | it 'should give us the same message over and over' do 152 | @nsqd.pub(TOPIC, 'slow') 153 | 154 | msg1 = @consumer.pop 155 | expect(msg1.body).to eq('slow') 156 | expect(msg1.attempts).to eq(1) 157 | 158 | # wait for it to be reclaimed by nsqd and then finish it so we can get 159 | # another. this fin won't actually succeed, because the message is no 160 | # longer in flight 161 | sleep(@msg_timeout + 0.1) 162 | msg1.finish 163 | 164 | assert_no_timeout do 165 | msg2 = @consumer.pop 166 | expect(msg2.body).to eq('slow') 167 | expect(msg2.attempts).to eq(2) 168 | end 169 | end 170 | 171 | 172 | # This is like the test above, except we touch the message to reset its 173 | # timeout 174 | it 'should be able to touch a message to reset its timeout' do 175 | @nsqd.pub(TOPIC, 'slow') 176 | 177 | msg1 = @consumer.pop 178 | expect(msg1.body).to eq('slow') 179 | 180 | # touch the message in the middle of a sleep session whose total just 181 | # exceeds the msg_timeout 182 | sleep(@msg_timeout / 2.0 + 0.1) 183 | msg1.touch 184 | sleep(@msg_timeout / 2.0 + 0.1) 185 | msg1.finish 186 | 187 | # if our touch didn't work, we should receive a message 188 | assert_timeout do 189 | @consumer.pop 190 | end 191 | end 192 | end 193 | 194 | 195 | describe 'with a high max_in_flight and tons of messages' do 196 | it 'should receive all messages in a reasonable amount of time' do 197 | expected_messages = (1..10_000).to_a.map(&:to_s) 198 | expected_messages.each_slice(100) do |slice| 199 | @cluster.nsqd.sample.mpub(TOPIC, *slice) 200 | end 201 | 202 | consumer = new_consumer(max_in_flight: 1000) 203 | received_messages = [] 204 | 205 | assert_no_timeout(5) do 206 | expected_messages.length.times do 207 | msg = consumer.pop 208 | received_messages << msg.body 209 | msg.finish 210 | end 211 | end 212 | 213 | consumer.terminate 214 | 215 | expect(received_messages.sort).to eq(expected_messages.sort) 216 | end 217 | end 218 | 219 | end 220 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # nsq-ruby 2 | 3 | nsq-ruby is a simple NSQ client library written in Ruby. 4 | 5 | - The code is straightforward. 6 | - It has no dependencies. 7 | - It's well tested. 8 | 9 | 10 | ## Quick start 11 | 12 | ### Publish messages 13 | 14 | ```Ruby 15 | require 'nsq' 16 | producer = Nsq::Producer.new( 17 | nsqd: '127.0.0.1:4150', 18 | topic: 'some-topic' 19 | ) 20 | 21 | # Write a message to NSQ 22 | producer.write('some-message') 23 | 24 | # Write a bunch of messages to NSQ (uses mpub) 25 | producer.write('one', 'two', 'three', 'four', 'five') 26 | 27 | # Close the connection 28 | producer.terminate 29 | ``` 30 | 31 | ### Consume messages 32 | 33 | ```Ruby 34 | require 'nsq' 35 | consumer = Nsq::Consumer.new( 36 | nsqlookupd: '127.0.0.1:4161', 37 | topic: 'some-topic', 38 | channel: 'some-channel' 39 | ) 40 | 41 | # Pop a message off the queue 42 | msg = consumer.pop 43 | puts msg.body 44 | msg.finish 45 | 46 | # Close the connections 47 | consumer.terminate 48 | ``` 49 | 50 | 51 | ## Producer 52 | 53 | ### Initialization 54 | 55 | The Nsq::Producer constructor takes the following options: 56 | 57 | | Option | Description | Default | 58 | |---------------|----------------------------------------|--------------------| 59 | | `topic` | Topic to which to publish messages | | 60 | | `nsqd` | Host and port of the nsqd instance | '127.0.0.1:4150' | 61 | 62 | For example: 63 | 64 | ```Ruby 65 | producer = Nsq::Producer.new( 66 | nsqd: '6.7.8.9:4150', 67 | topic: 'topic-of-great-esteem' 68 | ) 69 | ``` 70 | 71 | ### `#write` 72 | 73 | Publishes one or more message to nsqd. If you give it a single argument, it will 74 | send it to nsqd via `PUB`. If you give it multiple arguments, it will send all 75 | those messages to nsqd via `MPUB`. It will automatically call `to_s` on any 76 | arguments you give it. 77 | 78 | ```Ruby 79 | # Send a single message via PUB 80 | producer.write(123) 81 | 82 | # Send three messages via MPUB 83 | producer.write(456, 'another-message', { key: 'value' }.to_json) 84 | ``` 85 | 86 | If it's connection to nsqd fails, it will automatically try to reconnect with 87 | exponential backoff. Any messages that were sent to `#write` will be queued 88 | and transmitted after reconnecting. 89 | 90 | **Note** we don't wait for nsqd to acknowledge our writes. As a result, if the 91 | connection to nsqd fails, you can lose messages. This is acceptable for our use 92 | cases, mostly because we are sending messages to a local nsqd instance and 93 | failure is very rare. 94 | 95 | ### `#connected?` 96 | 97 | Returns true if it's currently connected to nsqd and false if not. 98 | 99 | ### `#terminate` 100 | 101 | Closes the connection to nsqd and stops it from trying to automatically 102 | reconnect. 103 | 104 | This is automatically called `at_exit`, but it's good practice to close your 105 | producers when you're done with them. 106 | 107 | 108 | ## Consumer 109 | 110 | ### Initialization 111 | 112 | | Option | Description | Default | 113 | |----------------------|-----------------------------------------------|--------------------| 114 | | `topic` | Topic to consume messages from | | 115 | | `channel` | Channel name for this consumer | | 116 | | `nsqlookupd` | Use lookupd to automatically discover nsqds | | 117 | | `nsqd` | Connect directly to a single nsqd instance | '127.0.0.1:4150' | 118 | | `max_in_flight` | Max number of messages for this consumer to have in flight at a time | 1 | 119 | | `discovery_interval` | Seconds between queue discovery via nsqlookupd | 60.0 | 120 | | `msg_timeout` | Milliseconds before nsqd will timeout a message | 60000 | 121 | 122 | 123 | For example: 124 | 125 | ```Ruby 126 | consumer = Nsq::Consumer.new( 127 | topic: 'the-topic', 128 | channel: 'my-channel', 129 | nsqlookupd: ['127.0.0.1:4161', '4.5.6.7:4161'], 130 | max_in_flight: 100, 131 | discovery_interval: 30, 132 | msq_timeout: 120_000 133 | ) 134 | ``` 135 | 136 | Notes: 137 | 138 | - `nsqlookupd` can be a string or array of strings for each nsqlookupd service 139 | you'd like to use. The format is `":"`. If you specify 140 | `nsqlookupd`, it ignores the `nsqd` option. 141 | - `max_in_flight` is for the total max in flight across all the connections, 142 | but to make the implementation of `nsq-ruby` as simple as possible, the minimum 143 | `max_in_flight` _per_ connection is 1. So if you set `max_in_flight` to 1 and 144 | are connected to 3 nsqds, you may have up to 3 messages in flight at a time. 145 | 146 | 147 | ### `#pop` 148 | 149 | `nsq-ruby` works by maintaining a local queue of in flight messages from NSQ. 150 | To get at these messages, just call pop. 151 | 152 | ```Ruby 153 | message = consumer.pop 154 | ``` 155 | 156 | If there are messages on the queue, `pop` will return one immediately. If there 157 | are no messages on the queue, `pop` will block execution until one arrives. 158 | 159 | 160 | ### `#size` 161 | 162 | `size` returns the size of the local message queue. 163 | 164 | 165 | ### `#terminate` 166 | 167 | Gracefully closes all connections and stops the consumer. You should call this 168 | when you're finished with a consumer object. 169 | 170 | 171 | ## Message 172 | 173 | The `Message` object is what you get when you call `pop` on a consumer. 174 | Once you have a message, you'll likely want to get its contents using the `#body` 175 | method, and then call `#finish` once you're done with it. 176 | 177 | ### `body` 178 | 179 | Returns the body of the message as a UTF-8 encoded string. 180 | 181 | ### `attempts` 182 | 183 | Returns the number of times this message was attempted to be processed. For 184 | most messages this should be 1 (since it will be your first attempt processing 185 | them). If it's more than 1, that means that you requeued the message or it 186 | timed out in flight. 187 | 188 | ### `#finish` 189 | 190 | Notify NSQ that you've completed processing of this message. 191 | 192 | ### `#touch` 193 | 194 | Tells NSQ to reset the message timeout for this message so you have more time 195 | to process it. 196 | 197 | ### `#requeue(timeout = 0)` 198 | 199 | Tells NSQ to requeue this message. Called with no arguments, this will requeue 200 | the message and it will be available to be received immediately. 201 | 202 | Optionally you can pass a number of milliseconds as an argument. This tells 203 | NSQ to delay its requeueing by that number of milliseconds. 204 | 205 | 206 | ## Logging 207 | 208 | By default, `nsq-ruby` doesn't log anything. To enable logging, use 209 | `Nsq.logger=` and point it at a Ruby Logger instance. Like this: 210 | 211 | ```Ruby 212 | Nsq.logger = Logger.new(STDOUT) 213 | ``` 214 | 215 | 216 | ## Requirements 217 | 218 | NSQ v0.2.29 or later due for IDENTITY metadata specification (0.2.28) and per- 219 | connection timeout support (0.2.29). 220 | 221 | 222 | ### Supports 223 | 224 | - Discovery via nsqlookupd 225 | - Automatic reconnection to nsqd 226 | 227 | ### Does not support 228 | 229 | - TLS 230 | - Compression 231 | - Backoff 232 | - Authentication 233 | 234 | If you need more advanced features, like these, you should check out 235 | [Krakow](https://github.com/chrisroberts/krakow), a more fully featured NSQ 236 | client for Ruby. 237 | 238 | 239 | ## Testing 240 | 241 | Run the tests like this: 242 | 243 | ``` 244 | rake spec 245 | ``` 246 | 247 | Want a deluge of logging while running the specs to help determine what is 248 | going on? 249 | 250 | ``` 251 | VERBOSE=true rake spec 252 | ``` 253 | 254 | 255 | ## MIT License 256 | 257 | Copyright (C) 2014 Wistia, Inc. 258 | 259 | Permission is hereby granted, free of charge, to any person obtaining a copy of 260 | this software and associated documentation files (the "Software"), to deal in 261 | the Software without restriction, including without limitation the rights to 262 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 263 | of the Software, and to permit persons to whom the Software is furnished to do 264 | so, subject to the following conditions: 265 | 266 | The above copyright notice and this permission notice shall be included in all 267 | copies or substantial portions of the Software. 268 | 269 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 270 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 271 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 272 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 273 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 274 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 275 | SOFTWARE. 276 | 277 | -------------------------------------------------------------------------------- /lib/nsq/connection.rb: -------------------------------------------------------------------------------- 1 | require 'json' 2 | require 'socket' 3 | require 'timeout' 4 | 5 | require_relative 'frames/error' 6 | require_relative 'frames/message' 7 | require_relative 'frames/response' 8 | require_relative 'logger' 9 | 10 | module Nsq 11 | class Connection 12 | include Nsq::AttributeLogger 13 | @@log_attributes = [:host, :port] 14 | 15 | attr_reader :host 16 | attr_reader :port 17 | attr_accessor :max_in_flight 18 | attr_reader :presumed_in_flight 19 | 20 | USER_AGENT = "nsq-ruby/#{Nsq::Version::STRING}" 21 | RESPONSE_HEARTBEAT = '_heartbeat_' 22 | RESPONSE_OK = 'OK' 23 | 24 | 25 | def initialize(opts = {}) 26 | @host = opts[:host] || (raise ArgumentError, 'host is required') 27 | @port = opts[:port] || (raise ArgumentError, 'host is required') 28 | @queue = opts[:queue] 29 | @topic = opts[:topic] 30 | @channel = opts[:channel] 31 | @msg_timeout = opts[:msg_timeout] || 60_000 # 60s 32 | @max_in_flight = opts[:max_in_flight] || 1 33 | 34 | if @msg_timeout < 1000 35 | raise ArgumentError, 'msg_timeout cannot be less than 1000. it\'s in milliseconds.' 36 | end 37 | 38 | # for outgoing communication 39 | @write_queue = Queue.new 40 | 41 | # For indicating that the connection has died. 42 | # We use a Queue so we don't have to poll. Used to communicate across 43 | # threads (from write_loop and read_loop to connect_and_monitor). 44 | @death_queue = Queue.new 45 | 46 | @connected = false 47 | @presumed_in_flight = 0 48 | 49 | open_connection 50 | start_monitoring_connection 51 | end 52 | 53 | 54 | def connected? 55 | @connected 56 | end 57 | 58 | 59 | # close the connection and don't try to re-open it 60 | def close 61 | stop_monitoring_connection 62 | close_connection 63 | end 64 | 65 | 66 | def sub(topic, channel) 67 | write "SUB #{topic} #{channel}\n" 68 | end 69 | 70 | 71 | def rdy(count) 72 | write "RDY #{count}\n" 73 | end 74 | 75 | 76 | def fin(message_id) 77 | write "FIN #{message_id}\n" 78 | decrement_in_flight 79 | end 80 | 81 | 82 | def req(message_id, timeout) 83 | write "REQ #{message_id} #{timeout}\n" 84 | decrement_in_flight 85 | end 86 | 87 | 88 | def touch(message_id) 89 | write "TOUCH #{message_id}\n" 90 | end 91 | 92 | 93 | def pub(topic, message) 94 | write ["PUB #{topic}\n", message.bytesize, message].pack('a*l>a*') 95 | end 96 | 97 | 98 | def mpub(topic, messages) 99 | body = messages.map do |message| 100 | [message.bytesize, message].pack('l>a*') 101 | end.join 102 | 103 | write ["MPUB #{topic}\n", body.bytesize, messages.size, body].pack('a*l>l>a*') 104 | end 105 | 106 | 107 | # Tell the server we are ready for more messages! 108 | def re_up_ready 109 | rdy(@max_in_flight) 110 | # assume these messages are coming our way. yes, this might not be the 111 | # case, but it's much easier to manage our RDY state with the server if 112 | # we treat things this way. 113 | @presumed_in_flight = @max_in_flight 114 | end 115 | 116 | 117 | private 118 | 119 | def cls 120 | write "CLS\n" 121 | end 122 | 123 | 124 | def nop 125 | write "NOP\n" 126 | end 127 | 128 | 129 | def write(raw) 130 | @write_queue.push(raw) 131 | end 132 | 133 | 134 | def write_to_socket(raw) 135 | debug ">>> #{raw.inspect}" 136 | @socket.write(raw) 137 | end 138 | 139 | 140 | # Block until we get an OK from nsqd 141 | def wait_for_ok 142 | frame = receive_frame 143 | unless frame.is_a?(Response) && frame.data == RESPONSE_OK 144 | raise "Received non-OK response while IDENTIFYing: #{frame.data}" 145 | end 146 | end 147 | 148 | 149 | def identify 150 | hostname = Socket.gethostname 151 | metadata = { 152 | client_id: Socket.gethostbyname(hostname).flatten.compact.first, 153 | hostname: hostname, 154 | feature_negotiation: false, 155 | heartbeat_interval: 30_000, # 30 seconds 156 | output_buffer: 16_000, # 16kb 157 | output_buffer_timeout: 250, # 250ms 158 | tls_v1: false, 159 | snappy: false, 160 | deflate: false, 161 | sample_rate: 0, # disable sampling 162 | user_agent: USER_AGENT, 163 | msg_timeout: @msg_timeout 164 | }.to_json 165 | write_to_socket ["IDENTIFY\n", metadata.length, metadata].pack('a*l>a*') 166 | end 167 | 168 | 169 | def handle_response(frame) 170 | if frame.data == RESPONSE_HEARTBEAT 171 | debug 'Received heartbeat' 172 | nop 173 | elsif frame.data == RESPONSE_OK 174 | debug 'Received OK' 175 | else 176 | die "Received response we don't know how to handle: #{frame.data}" 177 | end 178 | end 179 | 180 | 181 | def receive_frame 182 | if buffer = @socket.read(8) 183 | size, type = buffer.unpack('l>l>') 184 | size -= 4 # we want the size of the data part and type already took up 4 bytes 185 | data = @socket.read(size) 186 | frame_class = frame_class_for_type(type) 187 | return frame_class.new(data, self) 188 | end 189 | end 190 | 191 | 192 | FRAME_CLASSES = [Response, Error, Message] 193 | def frame_class_for_type(type) 194 | raise "Bad frame type specified: #{type}" if type > FRAME_CLASSES.length - 1 195 | [Response, Error, Message][type] 196 | end 197 | 198 | 199 | def decrement_in_flight 200 | @presumed_in_flight -= 1 201 | 202 | # now that we're less than @max_in_flight we might need to re-up our RDY 203 | # state 204 | threshold = (@max_in_flight * 0.2).ceil 205 | re_up_ready if @presumed_in_flight <= threshold 206 | end 207 | 208 | 209 | def start_read_loop 210 | @read_loop_thread ||= Thread.new{read_loop} 211 | end 212 | 213 | 214 | def stop_read_loop 215 | @read_loop_thread.kill if @read_loop_thread 216 | @read_loop_thread = nil 217 | end 218 | 219 | 220 | def read_loop 221 | loop do 222 | frame = receive_frame 223 | if frame.is_a?(Response) 224 | handle_response(frame) 225 | elsif frame.is_a?(Error) 226 | error "Error received: #{frame.data}" 227 | elsif frame.is_a?(Message) 228 | debug "<<< #{frame.body}" 229 | @queue.push(frame) if @queue 230 | else 231 | raise 'No data from socket' 232 | end 233 | end 234 | rescue Exception => ex 235 | die(ex) 236 | end 237 | 238 | 239 | def start_write_loop 240 | @write_loop_thread ||= Thread.new{write_loop} 241 | end 242 | 243 | 244 | def stop_write_loop 245 | @stop_write_loop = true 246 | @write_loop_thread.join(1) if @write_loop_thread 247 | @write_loop_thread = nil 248 | end 249 | 250 | 251 | def write_loop 252 | @stop_write_loop = false 253 | data = nil 254 | loop do 255 | data = @write_queue.pop 256 | write_to_socket(data) 257 | break if @stop_write_loop && @write_queue.size == 0 258 | end 259 | rescue Exception => ex 260 | # requeue PUB and MPUB commands 261 | if data =~ /^M?PUB/ 262 | debug "Requeueing to write_queue: #{data.inspect}" 263 | @write_queue.push(data) 264 | end 265 | die(ex) 266 | end 267 | 268 | 269 | # Waits for death of connection 270 | def start_monitoring_connection 271 | @connection_monitor_thread ||= Thread.new{monitor_connection} 272 | @connection_monitor_thread.abort_on_exception = true 273 | end 274 | 275 | 276 | def stop_monitoring_connection 277 | @connection_monitor_thread.kill if @connection_monitor_thread 278 | @connection_monitor = nil 279 | end 280 | 281 | 282 | def monitor_connection 283 | loop do 284 | # wait for death, hopefully it never comes 285 | cause_of_death = @death_queue.pop 286 | warn "Died from: #{cause_of_death}" 287 | 288 | debug 'Reconnecting...' 289 | reconnect 290 | debug 'Reconnected!' 291 | 292 | # clear all death messages, since we're now reconnected. 293 | # we don't want to complete this loop and immediately reconnect again. 294 | @death_queue.clear 295 | end 296 | end 297 | 298 | 299 | # close the connection if it's not already closed and try to reconnect 300 | # over and over until we succeed! 301 | def reconnect 302 | close_connection 303 | with_retries do 304 | open_connection 305 | end 306 | end 307 | 308 | 309 | def open_connection 310 | @socket = TCPSocket.new(@host, @port) 311 | # write the version and IDENTIFY directly to the socket to make sure 312 | # it gets to nsqd ahead of anything in the `@write_queue` 313 | write_to_socket ' V2' 314 | identify 315 | wait_for_ok 316 | 317 | start_read_loop 318 | start_write_loop 319 | @connected = true 320 | 321 | # we need to re-subscribe if there's a topic specified 322 | if @topic 323 | debug "Subscribing to #{@topic}" 324 | sub(@topic, @channel) 325 | re_up_ready 326 | end 327 | end 328 | 329 | 330 | # closes the connection and stops listening for messages 331 | def close_connection 332 | cls if connected? 333 | stop_read_loop 334 | stop_write_loop 335 | @socket = nil 336 | @connected = false 337 | end 338 | 339 | 340 | # this is called when there's a connection error in the read or write loop 341 | # it triggers `connect_and_monitor` to try to reconnect 342 | def die(reason) 343 | @connected = false 344 | @death_queue.push(reason) 345 | end 346 | 347 | 348 | # Retry the supplied block with exponential backoff. 349 | # 350 | # Borrowed liberally from: 351 | # https://github.com/ooyala/retries/blob/master/lib/retries.rb 352 | def with_retries(&block) 353 | base_sleep_seconds = 0.5 354 | max_sleep_seconds = 300 # 5 minutes 355 | 356 | # Let's do this thing 357 | attempts = 0 358 | start_time = Time.now 359 | 360 | begin 361 | attempts += 1 362 | return block.call(attempts) 363 | 364 | rescue Errno::ECONNREFUSED, Errno::ECONNRESET, Errno::EHOSTUNREACH, 365 | Errno::ENETDOWN, Errno::ENETUNREACH, Errno::ETIMEDOUT, Timeout::Error => ex 366 | 367 | raise ex if attempts >= 100 368 | 369 | # The sleep time is an exponentially-increasing function of base_sleep_seconds. 370 | # But, it never exceeds max_sleep_seconds. 371 | sleep_seconds = [base_sleep_seconds * (2 ** (attempts - 1)), max_sleep_seconds].min 372 | # Randomize to a random value in the range sleep_seconds/2 .. sleep_seconds 373 | sleep_seconds = sleep_seconds * (0.5 * (1 + rand())) 374 | # But never sleep less than base_sleep_seconds 375 | sleep_seconds = [base_sleep_seconds, sleep_seconds].max 376 | 377 | warn "Failed to connect: #{ex}. Retrying in #{sleep_seconds.round(1)} seconds." 378 | 379 | snooze(sleep_seconds) 380 | 381 | retry 382 | end 383 | end 384 | 385 | 386 | # Se we can stub for testing and reconnect in a tight loop 387 | def snooze(t) 388 | sleep(t) 389 | end 390 | end 391 | end 392 | --------------------------------------------------------------------------------