stop_gracefully: Wait until all clients are disconnected

droonga · Apr 22, 2015 · 139c5e1 · 139c5e1
1 parent 5f2a81e
commit 139c5e1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/lib/droonga/command/droonga_engine_service.rb b/lib/droonga/command/droonga_engine_service.rb
@@ -213,15 +213,15 @@ def stop_gracefully
         return if @stopping
         logger.trace("stop_gracefully: start")
         @stopping = true
-        @receiver.stop_gracefully
+        @receiver.stop_gracefully do
         @engine.stop_gracefully do
           shutdown_worker_process_agent
           shutdown_internal_message_receiver
-          @receiver.shutdown_clients
           logger.trace("stop_gracefully: done",
                        :n_rest_watchers => @loop.watchers.size,
                        :rest_watchers   => @loop.watchers)
         end
+        end
       end
 
       # It may be called after stop_gracefully.

diff --git a/lib/droonga/fluent_message_receiver.rb b/lib/droonga/fluent_message_receiver.rb
@@ -31,6 +31,7 @@ def initialize(loop, options={}, &on_message)
       @server = nil
       @clients = []
       @on_message = on_message
+      @on_shutdown_ready = nil
     end
 
     def start
@@ -40,11 +41,18 @@ def start
       logger.trace("start: done")
     end
 
-    def stop_gracefully
+    def stop_gracefully(&block)
       logger.trace("stop_gracefully: start")
       shutdown_heartbeat_receiver
       logger.trace("stop_gracefully: middle")
       shutdown_server
+      if @clients.empty?
+        yield
+      elsif block_given?
+        @on_shutdown_ready = lambda do
+          yield
+        end
+      end
       logger.trace("stop_gracefully: done")
     end
 
@@ -87,6 +95,9 @@ def start_server
         end
         client.on_close = lambda do
           @clients.delete(client)
+          if @on_shutdown_ready and @clients.empty?
+            @on_shutdown_ready.call
+          end
         end
         @clients << client
       end