Ensure events send out after being forked

lib/segment/analytics.rb

@@ -2,7 +2,7 @@
 require 'segment/analytics/utils'
 require 'segment/analytics/version'
 require 'segment/analytics/client'
-require 'segment/analytics/consumer'
+require 'segment/analytics/worker'
 require 'segment/analytics/request'
 require 'segment/analytics/response'
 require 'segment/analytics/logging'

lib/segment/analytics/client.rb

@@ -1,7 +1,7 @@
 require 'thread'
 require 'time'
 require 'segment/analytics/utils'
-require 'segment/analytics/consumer'
+require 'segment/analytics/worker'
 require 'segment/analytics/defaults'
 
 module Segment
@@ -21,27 +21,19 @@ def initialize options = {}
         @queue = Queue.new
         @write_key = options[:write_key]
         @max_queue_size = options[:max_queue_size] || Defaults::Queue::MAX_SIZE
-        @consumer = Consumer.new @queue, @write_key, options
-        @thread = ConsumerThread.new { @consumer.run }
+        @options = options
+        @worker_mutex = Mutex.new
 
         check_write_key!
-
-        at_exit do
-          # Let the consumer thread know it should exit.
-          @thread[:should_exit] = true
-
-          # Push a flag value to the consumer queue in case it's blocked waiting for a value.  This will allow it
-          # to continue its normal chain of processing, giving it a chance to exit.
-          @queue << nil
-        end
       end
 
-      # public: Synchronously waits until the consumer has flushed the queue.
+      # public: Synchronously waits until the worker has flushed the queue.
       #         Use only for scripts which are not long-running, and will
       #         specifically exit
       #
       def flush
-        while !@queue.empty? || @consumer.is_requesting?
+        while !@queue.empty? || @worker.is_requesting?
+          ensure_worker_running
           sleep(0.1)
         end
       end
@@ -263,10 +255,10 @@ def queued_messages
       def enqueue(action)
         # add our request id for tracing purposes
         action[:messageId] = uid
-
-        queue_full = @queue.length >= @max_queue_size
-        @queue << action unless queue_full
-
+        unless queue_full = @queue.length >= @max_queue_size
+          ensure_worker_running
+          @queue << action
+        end
         !queue_full
       end
 
@@ -313,8 +305,19 @@ def check_user_id! options
         fail ArgumentError, 'Must supply either user_id or anonymous_id' unless options[:user_id] || options[:anonymous_id]
       end
 
-      # Sub-class thread so we have a named thread (useful for debugging in Thread.list).
-      class ConsumerThread < Thread
+      def ensure_worker_running
+        return if worker_running?
+        @worker_mutex.synchronize do
+          return if worker_running?
+          @worker_thread = Thread.new do
+            @worker = Worker.new @queue, @write_key, @options
+            @worker.run
+          end
+        end
+      end
+
+      def worker_running?
+        @worker_thread && @worker_thread.alive?
       end
     end
   end

lib/segment/analytics/worker.rb

@@ -0,0 +1,59 @@
+require 'segment/analytics/defaults'
+require 'segment/analytics/utils'
+require 'segment/analytics/defaults'
+require 'segment/analytics/request'
+
+module Segment
+  class Analytics
+    class Worker
+      include Segment::Analytics::Utils
+      include Segment::Analytics::Defaults
+
+      # public: Creates a new worker
+      #
+      # The worker continuously takes messages off the queue
+      # and makes requests to the segment.io api
+      #
+      # queue   - Queue synchronized between client and worker
+      # write_key  - String of the project's Write key
+      # options - Hash of worker options
+      #           batch_size - Fixnum of how many items to send in a batch
+      #           on_error   - Proc of what to do on an error
+      #
+      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.new { |status, error| }
+        @batch = []
+        @lock = Mutex.new
+      end
+
+      # public: Continuously runs the loop to check for new events
+      #
+      def run
+        loop do
+          return if @queue.empty?
+
+          @lock.synchronize do
+            until @batch.length >= @batch_size || @queue.empty?
+              @batch << @queue.pop
+            end
+          end
+
+          res = Request.new.post @write_key, @batch
+          @on_error.call res.status, res.error unless res.status == 200
+
+          @lock.synchronize { @batch.clear }
+        end
+      end
+
+      # public: Check whether we have outstanding requests.
+      #
+      def is_requesting?
+        @lock.synchronize { !@batch.empty? }
+      end
+    end
+  end
+end

spec/segment/analytics/client_spec.rb

@@ -20,7 +20,6 @@ class Analytics
       describe '#track' do
         before(:all) do
           @client = Client.new :write_key => WRITE_KEY
-          @client.instance_variable_get(:@thread).kill
           @queue = @client.instance_variable_get :@queue
         end
 
@@ -69,10 +68,8 @@ class Analytics
 
 
       describe '#identify' do
-
         before(:all) do
           @client = Client.new :write_key => WRITE_KEY
-          @client.instance_variable_get(:@thread).kill
           @queue = @client.instance_variable_get :@queue
         end
 
@@ -203,6 +200,18 @@ class Analytics
           @client.flush
           @client.queued_messages.should == 0
         end
+
+        it 'should complete when the process forks' do
+          @client.identify Queued::IDENTIFY
+
+          Process.fork do
+            @client.track Queued::TRACK
+            @client.flush
+            @client.queued_messages.should == 0
+          end
+
+          Process.wait
+        end unless defined? JRUBY_VERSION
       end
     end
   end

spec/segment/analytics/consumer_spec.rb → spec/segment/analytics/worker_spec.rb

@@ -2,12 +2,12 @@
 
 module Segment
   class Analytics
-    describe Consumer do
+    describe Worker do
       describe "#init" do
         it 'accepts string keys' do
           queue = Queue.new
-          consumer = Segment::Analytics::Consumer.new(queue, 'secret', 'batch_size' => 100)
-          consumer.instance_variable_get(:@batch_size).should == 100
+          worker = Segment::Analytics::Worker.new(queue, 'secret', 'batch_size' => 100)
+          worker.instance_variable_get(:@batch_size).should == 100
         end
       end
 
@@ -25,8 +25,8 @@ class Analytics
 
           queue = Queue.new
           queue << {}
-          consumer = Segment::Analytics::Consumer.new(queue, 'secret')
-          consumer.flush
+          worker = Segment::Analytics::Worker.new(queue, 'secret')
+          worker.run
 
           queue.should be_empty
 
@@ -44,8 +44,8 @@ class Analytics
 
           queue = Queue.new
           queue << {}
-          consumer = Segment::Analytics::Consumer.new queue, 'secret', :on_error => on_error
-          consumer.flush
+          worker = Segment::Analytics::Worker.new queue, 'secret', :on_error => on_error
+          worker.run
 
           Segment::Analytics::Request::any_instance.unstub(:post)
 
@@ -62,37 +62,36 @@ class Analytics
 
           queue = Queue.new
           queue << Requested::TRACK
-          consumer = Segment::Analytics::Consumer.new queue, 'testsecret', :on_error => on_error
-          consumer.flush
+          worker = Segment::Analytics::Worker.new queue, 'testsecret', :on_error => on_error
+          worker.run
 
           queue.should be_empty
         end
       end
 
       describe '#is_requesting?' do
-
         it 'should not return true if there isn\'t a current batch' do
 
           queue = Queue.new
-          consumer = Segment::Analytics::Consumer.new(queue, 'testsecret')
+          worker = Segment::Analytics::Worker.new(queue, 'testsecret')
 
-          consumer.is_requesting?.should == false
+          worker.is_requesting?.should == false
         end
 
         it 'should return true if there is a current batch' do
 
           queue = Queue.new
           queue << Requested::TRACK
-          consumer = Segment::Analytics::Consumer.new(queue, 'testsecret')
+          worker = Segment::Analytics::Worker.new(queue, 'testsecret')
 
           Thread.new {
-            consumer.flush
-            consumer.is_requesting?.should == false
+            worker.run
+            worker.is_requesting?.should == false
           }
 
           # sleep barely long enough to let thread flush the queue.
           sleep(0.001)
-          consumer.is_requesting?.should == true
+          worker.is_requesting?.should == true
         end
       end
     end

spec/spec_helper.rb

@@ -49,7 +49,7 @@ module Queued
       SCREEN = SCREEN.merge :user_id => USER_ID
     end
 
-    # Hashes which are sent from the consumer, camel_cased
+    # Hashes which are sent from the worker, camel_cased
     module Requested
       TRACK = TRACK.merge({
         :userId => USER_ID,