ensure request completion in lease keep alive to prevent orphaned threads

Author: braintreeps

etcdv3.gemspec

@@ -20,4 +20,5 @@ Gem::Specification.new do |s|
   end
   s.add_dependency("grpc", "~> 1.6")
   s.add_development_dependency("rspec", "~> 3.6.0")
+  s.add_development_dependency("pry")
 end

lib/etcdv3/lease.rb

@@ -1,24 +1,29 @@
 class Etcdv3
   class Lease
-    class Enum < Array
-      def proceed
+
+    # this is used for gRPC proxy compatibility so that we do not
+    # mark as finished writing until we've received a response
+    class BlockingRequest
+      def initialize(request_op)
+        @blocked = false
+        @request_op = request_op
+      end
+
+      def read_done!
         @proceed = true
       end
 
-      def completed
-        @completed
+      def blocked?
+        @blocked
       end
 
       def each
-        @completed = false
-        i = 0
-        while i < self.size do
-          @proceed = false
-          yield self[i]
-          i+=1
-          sleep 0.001 until @proceed == true
-        end
-        @completed = true
+        @blocked = true
+
+        yield @request_op
+
+        sleep 0.001 until @proceed == true
+        @blocked = false
       end
     end
 
@@ -44,16 +49,20 @@ def lease_ttl(id, timeout: nil)
     end
 
     def lease_keep_alive_once(id, timeout: nil)
-      enum = Enum.new
-      enum << Etcdserverpb::LeaseKeepAliveRequest.new(ID: id)
+      request = BlockingRequest.new Etcdserverpb::LeaseKeepAliveRequest.new(ID: id)
       response = nil
-      @stub.lease_keep_alive(enum, metadata: @metadata, deadline: deadline(timeout)).each do |resp|
-        response = resp
-        break;
+      begin
+        @stub.lease_keep_alive(request, metadata: @metadata, deadline: deadline(timeout)).each do |resp|
+          response = resp
+          break;
+        end
+      ensure
+        request.read_done!
+        while request.blocked?
+          sleep 0.001
+        end
       end
-      enum.proceed
-      sleep 0.001 until enum.completed
-      response
+      return response
     end
 
     private

spec/etcdv3/lease_spec.rb

@@ -41,6 +41,12 @@
       stub = local_stub(Etcdv3::Lease, -1)
       expect { stub.lease_keep_alive_once(id) }.to raise_error(GRPC::DeadlineExceeded)
     end
+
+    it "doesn't orphan threads if there is a server error" do
+      expect_any_instance_of(GRPC::BidiCall).to receive(:read_loop).and_raise(GRPC::DeadlineExceeded)
+      stub = local_stub(Etcdv3::Lease, 2)
+      expect { stub.lease_keep_alive_once(314159) rescue nil; sleep 0.5}.to_not change { Thread.list.size }
+    end
   end
 
   describe '#lease_ttl' do