diff --git a/lib/segment/analytics/defaults.rb b/lib/segment/analytics/defaults.rb index 18cf861..8562346 100644 --- a/lib/segment/analytics/defaults.rb +++ b/lib/segment/analytics/defaults.rb @@ -13,7 +13,6 @@ module Request end module Queue - BATCH_SIZE = 100 MAX_SIZE = 10000 end @@ -21,6 +20,11 @@ module Message MAX_BYTES = 32768 # 32Kb end + module MessageBatch + MAX_BYTES = 512_000 # 500Kb + MAX_SIZE = 100 + end + module BackoffPolicy MIN_TIMEOUT_MS = 100 MAX_TIMEOUT_MS = 10000 diff --git a/lib/segment/analytics/message.rb b/lib/segment/analytics/message.rb index b11b121..e00a74d 100644 --- a/lib/segment/analytics/message.rb +++ b/lib/segment/analytics/message.rb @@ -9,7 +9,11 @@ def initialize(hash) end def too_big? - to_json.bytesize > Defaults::Message::MAX_BYTES + json_size > Defaults::Message::MAX_BYTES + end + + def json_size + to_json.bytesize end # Since the hash is expected to not be modified (set at initialization), diff --git a/lib/segment/analytics/message_batch.rb b/lib/segment/analytics/message_batch.rb index 0cf1dc4..bc44971 100644 --- a/lib/segment/analytics/message_batch.rb +++ b/lib/segment/analytics/message_batch.rb @@ -6,9 +6,12 @@ class Analytics class MessageBatch extend Forwardable include Segment::Analytics::Logging + include Segment::Analytics::Defaults::MessageBatch - def initialize + def initialize(max_message_count) @messages = [] + @max_message_count = max_message_count + @json_size = 0 end def <<(message) @@ -16,13 +19,40 @@ def <<(message) logger.error('a message exceeded the maximum allowed size') else @messages << message + @json_size += message.json_size + 1 # One byte for the comma end end + def full? + item_count_exhausted? || size_exhausted? + end + + def clear + @messages.clear + @json_size = 0 + end + def_delegators :@messages, :to_json - def_delegators :@messages, :clear def_delegators :@messages, :empty? def_delegators :@messages, :length + + private + + def item_count_exhausted? + @messages.length >= @max_message_count + end + + # We consider the max size here as just enough to leave room for one more + # message of the largest size possible. This is a shortcut that allows us + # to use a native Ruby `Queue` that doesn't allow peeking. The tradeoff + # here is that we might fit in less messages than possible into a batch. + # + # The alternative is to use our own `Queue` implementation that allows + # peeking, and to consider the next message size when calculating whether + # the message can be accomodated in this batch. + def size_exhausted? + @json_size >= (MAX_BYTES - Defaults::Message::MAX_BYTES) + end end end end diff --git a/lib/segment/analytics/worker.rb b/lib/segment/analytics/worker.rb index a5fb638..f1c7ef4 100644 --- a/lib/segment/analytics/worker.rb +++ b/lib/segment/analytics/worker.rb @@ -25,9 +25,9 @@ def initialize(queue, write_key, options = {}) symbolize_keys! options @queue = queue @write_key = write_key - @batch_size = options[:batch_size] || Queue::BATCH_SIZE @on_error = options[:on_error] || proc { |status, error| } - @batch = MessageBatch.new + batch_size = options[:batch_size] || Defaults::MessageBatch::MAX_SIZE + @batch = MessageBatch.new(batch_size) @lock = Mutex.new @request = Request.new end @@ -39,7 +39,7 @@ def run return if @queue.empty? @lock.synchronize do - until @batch.length >= @batch_size || @queue.empty? + until @batch.full? || @queue.empty? @batch << Message.new(@queue.pop) end end diff --git a/spec/helpers/runscope_client.rb b/spec/helpers/runscope_client.rb index 7d0e7f8..a68be26 100644 --- a/spec/helpers/runscope_client.rb +++ b/spec/helpers/runscope_client.rb @@ -9,7 +9,7 @@ def initialize(api_token) def requests(bucket_key) with_retries(3) do - response = @conn.get("/buckets/#{bucket_key}/messages", count: 10) + response = @conn.get("/buckets/#{bucket_key}/messages", count: 20) raise "Runscope error. #{response.body}" unless response.status == 200 diff --git a/spec/segment/analytics/message_batch_spec.rb b/spec/segment/analytics/message_batch_spec.rb index e50f4d0..6e014c8 100644 --- a/spec/segment/analytics/message_batch_spec.rb +++ b/spec/segment/analytics/message_batch_spec.rb @@ -3,16 +3,16 @@ module Segment class Analytics describe MessageBatch do - describe '#<<' do - subject { described_class.new } + subject { described_class.new(100) } + describe '#<<' do it 'appends messages' do subject << Message.new('a' => 'b') expect(subject.length).to eq(1) end it 'rejects messages that exceed the maximum allowed size' do - max_bytes = Segment::Analytics::Defaults::Message::MAX_BYTES + max_bytes = Defaults::Message::MAX_BYTES hash = { 'a' => 'b' * max_bytes } message = Message.new(hash) @@ -20,6 +20,30 @@ class Analytics expect(subject.length).to eq(0) end end + + describe '#full?' do + it 'returns true once item count is exceeded' do + 99.times { subject << Message.new(a: 'b') } + expect(subject.full?).to be(false) + + subject << Message.new(a: 'b') + expect(subject.full?).to be(true) + end + + it 'returns true once max size is almost exceeded' do + message = Message.new(a: 'b' * (Defaults::Message::MAX_BYTES - 10)) + + # Each message is under the individual limit + expect(message.json_size).to be < Defaults::Message::MAX_BYTES + + # Size of the batch is over the limit + expect(50 * message.json_size).to be > Defaults::MessageBatch::MAX_BYTES + + expect(subject.full?).to be(false) + 50.times { subject << message } + expect(subject.full?).to be(true) + end + end end end end diff --git a/spec/segment/analytics/worker_spec.rb b/spec/segment/analytics/worker_spec.rb index 3abd4d8..e66db64 100644 --- a/spec/segment/analytics/worker_spec.rb +++ b/spec/segment/analytics/worker_spec.rb @@ -6,8 +6,11 @@ class Analytics describe '#init' do it 'accepts string keys' do queue = Queue.new - worker = Segment::Analytics::Worker.new(queue, 'secret', 'batch_size' => 100) - expect(worker.instance_variable_get(:@batch_size)).to eq(100) + worker = Segment::Analytics::Worker.new(queue, + 'secret', + 'batch_size' => 100) + batch = worker.instance_variable_get(:@batch) + expect(batch.instance_variable_get(:@max_message_count)).to eq(100) end end @@ -35,8 +38,11 @@ class Analytics end.to_not raise_error end - it 'executes the error handler, before the request phase ends, if the request is invalid' do - Segment::Analytics::Request.any_instance.stub(:post).and_return(Segment::Analytics::Response.new(400, 'Some error')) + it 'executes the error handler if the request is invalid' do + Segment::Analytics::Request + .any_instance + .stub(:post) + .and_return(Segment::Analytics::Response.new(400, 'Some error')) status = error = nil on_error = proc do |yielded_status, yielded_error| @@ -46,9 +52,10 @@ class Analytics queue = Queue.new queue << {} - worker = Segment::Analytics::Worker.new queue, 'secret', :on_error => on_error + worker = described_class.new(queue, 'secret', :on_error => on_error) - # This is to ensure that Client#flush doesn't finish before calling the error handler. + # This is to ensure that Client#flush doesn't finish before calling + # the error handler. Thread.new { worker.run } sleep 0.1 # First give thread time to spin-up. sleep 0.01 while worker.is_requesting? @@ -69,7 +76,9 @@ class Analytics queue = Queue.new queue << Requested::TRACK - worker = Segment::Analytics::Worker.new queue, 'testsecret', :on_error => on_error + worker = described_class.new(queue, + 'testsecret', + :on_error => on_error) worker.run expect(queue).to be_empty @@ -89,12 +98,11 @@ class Analytics queue << Requested::TRACK worker = Segment::Analytics::Worker.new(queue, 'testsecret') - Thread.new do - worker.run - expect(worker.is_requesting?).to eq(false) - end - + worker_thread = Thread.new { worker.run } eventually { expect(worker.is_requesting?).to eq(true) } + + worker_thread.join + expect(worker.is_requesting?).to eq(false) end end end