Merge pull request #145 from brianstien/brian/check_your_self

Check workers before broadcasting heartbeats.

lib/protobuf/rpc/servers/zmq/server.rb

@@ -12,7 +12,8 @@ class Server
 
         DEFAULT_OPTIONS = {
           :beacon_interval => 5,
-          :broadcast_beacons => false
+          :broadcast_beacons => false,
+          :broadcast_busy => false
         }
 
         attr_accessor :options, :workers
@@ -33,6 +34,10 @@ def add_worker
           @total_workers = total_workers + 1
         end
 
+        def all_workers_busy?
+          workers.all? { |thread| !!thread.thread_variable_get(:busy) }
+        end
+
         def backend_port
           options[:worker_port] || frontend_port + 1
         end
@@ -99,6 +104,10 @@ def brokerless?
           !!options[:workers_only]
         end
 
+        def busy_worker_count
+          workers.count { |thread| !!thread.thread_variable_get(:busy) }
+        end
+
         def frontend_ip
           @frontend_ip ||= resolve_ip(options[:host])
         end
@@ -236,7 +245,14 @@ def wait_for_shutdown_signal
               start_missing_workers
             end
 
-            broadcast_heartbeat if broadcast_heartbeat?
+            if broadcast_heartbeat?
+              if all_workers_busy? && options[:broadcast_busy]
+                broadcast_flatline
+              else
+                broadcast_heartbeat
+              end
+            end
+
           end
         end
 

lib/protobuf/rpc/servers/zmq/worker.rb

@@ -51,8 +51,10 @@ def run
             break if rc == -1
 
             if rc > 0
+              ::Thread.current.thread_variable_set(:busy, true)
               initialize_request!
               process_request
+              ::Thread.current.thread_variable_set(:busy, false)
             end
           end
         ensure

lib/protobuf/rpc/servers/zmq_runner.rb

@@ -43,6 +43,7 @@ def register_signals
 
         trap(:TTOU) do
           log_info { "Current worker size: #{@server.workers.size}" }
+          log_info { "Current worker size: #{@server.busy_worker_count}" }
         end
       end
     end