Merge pull request #175 from wayfair/add_shutdown_to_client

Add shutdown to client

Author: evanphx

lib/vault/client.rb

@@ -158,6 +158,12 @@ def pool
 
     private :pool
 
+    # Shutdown any open pool connections. Pool will be recreated upon next request.
+    def shutdown
+      @nhp.shutdown()
+      @nhp = nil
+    end
+
     # Creates and yields a new client object with the given token. This may be
     # used safely in a threadsafe manner because the original client remains
     # unchanged. The value of the block is returned.

spec/integration/client_spec.rb

@@ -52,5 +52,65 @@ def free_address
         }.to raise_error(MissingTokenError)
       end
     end
+
+    describe "#shutdown" do
+      it "clears the pool after calling shutdown and sets nhp to nil" do
+        TCPServer.open('localhost', 0) do |server|
+          Thread.new do
+            loop do
+              client = server.accept
+              sleep 0.25
+              client.close
+            end
+          end
+
+          address = "http://%s:%s" % ["localhost", server.addr[1]]
+
+          client = described_class.new(address: address, token: "foo")
+
+          expect { client.request(:get, "/", {}, {}) }.to raise_error(HTTPConnectionError)
+
+          pool = client.instance_variable_get(:@nhp).pool
+
+          client.shutdown()
+
+          expect(pool.available.instance_variable_get(:@enqueued)).to eq(0)
+          expect(pool.available.instance_variable_get(:@shutdown_block)).not_to be_nil
+          expect(client.instance_variable_get(:@nhp)).to be_nil
+
+          server.close
+        end
+      end
+
+      it "the pool is recreated on the following request" do
+        TCPServer.open('localhost', 0) do |server|
+          Thread.new do
+            loop do
+              client = server.accept
+              sleep 0.25
+              client.close
+            end
+          end
+
+          address = "http://%s:%s" % ["localhost", server.addr[1]]
+
+          client = described_class.new(address: address, token: "foo")
+
+          expect { client.request(:get, "/", {}, {}) }.to raise_error(HTTPConnectionError)
+
+          client.shutdown()
+
+          expect { client.request(:get, "/", {}, {}) }.to raise_error(HTTPConnectionError)
+
+          pool = client.instance_variable_get(:@nhp).pool
+
+          expect(pool.available.instance_variable_get(:@enqueued)).to eq(1)
+          expect(pool.available.instance_variable_get(:@shutdown_block)).to be_nil
+          expect(client.instance_variable_get(:@nhp)).not_to be_nil
+
+          server.close
+        end
+      end
+    end
   end
 end