Sync release leadership on stop

spec/etcd_spec.cr

@@ -66,7 +66,7 @@ describe LavinMQ::Etcd do
         fail "should not lose the leadership to #{new}"
       when timeout(2.seconds)
       end
-      lease.close
+      lease.release
     end
   end
 
@@ -75,13 +75,9 @@ describe LavinMQ::Etcd do
     cluster.run do |etcds|
       etcd = LavinMQ::Etcd.new(cluster.endpoints)
       key = "foo/#{rand}"
-      lease = etcd.elect(key, "bar", 1)
+      lease = etcd.elect(key, "bar", ttl: 1)
       etcds.first(2).each &.terminate(graceful: false)
-      select
-      when lease.receive?
-      when timeout(15.seconds)
-        fail "should lose the leadership"
-      end
+      lease.wait(15.seconds).should be_true, "should lose the leadership"
     end
   end
 
@@ -90,13 +86,9 @@ describe LavinMQ::Etcd do
     cluster.run do |etcds|
       etcd = LavinMQ::Etcd.new(cluster.endpoints)
       key = "foo/#{rand}"
-      lease = etcd.elect(key, "bar", 1)
+      lease = etcd.elect(key, "bar", ttl: 1)
       etcds.sample.terminate(graceful: false)
-      select
-      when lease.receive?
-        fail "should not lose the leadership"
-      when timeout(6.seconds)
-      end
+      lease.wait(6.seconds).should be_false, "should not lose the leadership"
     end
   end
 

src/lavinmq/etcd.cr

@@ -1,6 +1,7 @@
 require "http/client"
 require "wait_group"
 require "json"
+require "./logger"
 
 module LavinMQ
   class Etcd
@@ -78,9 +79,9 @@ module LavinMQ
     end
 
     # Campaign for an election
-    # Returns a Channel when elected leader,
-    # when the channel is closed the leadership is lost
-    def elect(name, value, ttl = 10) : Channel(Nil)
+    # Returns when elected leader
+    # Returns a `Leadership` instance
+    def elect(name, value, ttl = 10) : Leadership
       channel = Channel(Nil).new
       lease_id, ttl = lease_grant(ttl)
       wg = WaitGroup.new(1)
@@ -103,7 +104,30 @@ module LavinMQ
       end
       election_campaign(name, value, lease_id)
       wg.wait
-      channel
+      Leadership.new(self, lease_id, channel)
+    end
+
+    # Represents a holding a Leadership
+    # Can be revoked or wait until lost
+    class Leadership
+      def initialize(@etcd : Etcd, @lease_id : Int64, @lost_leadership_channel : Channel(Nil))
+      end
+
+      # Force release leadership
+      def release
+        @etcd.lease_revoke(@lease_id)
+      end
+
+      # Wait until looses leadership
+      # Returns true when lost leadership, false when timeout occured
+      def wait(timeout : Time::Span) : Bool
+        select
+        when @lost_leadership_channel.receive?
+          return true
+        when timeout(timeout)
+          return false
+        end
+      end
     end
 
     def elect_listen(name, &)

src/lavinmq/launcher.cr

@@ -6,17 +6,18 @@ require "./server"
 require "./http/http_server"
 require "./in_memory_backend"
 require "./data_dir_lock"
+require "./etcd"
 
 module LavinMQ
   class Launcher
     Log = LavinMQ::Log.for "launcher"
     @tls_context : OpenSSL::SSL::Context::Server?
     @first_shutdown_attempt = true
     @data_dir_lock : DataDirLock?
-    @lease : Channel(Nil)?
-    @running = false
+    @closed = false
+    @leadership : Etcd::Leadership?
 
-    def initialize(@config : Config, replicator = Clustering::NoopServer.new, @lease = nil)
+    def initialize(@config : Config, replicator = Clustering::NoopServer.new, @leadership = nil)
       print_environment_info
       print_max_map_count
       fd_limit = System.maximize_fd_limit
@@ -38,18 +39,15 @@ module LavinMQ
     end
 
     def run
-      @running = true
       listen
       SystemD.notify_ready
       loop do
-        if lease = @lease
-          select
-          when lease.receive?
-            break unless @running
-            Log.warn { "Lost leadership lease" }
+        if leadership = @leadership
+          if leadership.wait(30.seconds)
+            Log.warn { "Lost leadership" }
             stop
             exit 1
-          when timeout(30.seconds)
+          else
             @data_dir_lock.try &.poll
             GC.collect
           end
@@ -62,14 +60,14 @@ module LavinMQ
     end
 
     def stop
-      return unless @running
-      @running = false
+      return if @closed
+      @closed = true
       Log.warn { "Stopping" }
       SystemD.notify_stopping
       @http_server.close rescue nil
       @amqp_server.close rescue nil
       @data_dir_lock.try &.release
-      @lease.try &.close
+      @leadership.try &.release
     end
 
     private def print_environment_info