From 824cca5f9135c6f5ae078b213f16afda7b7c1575 Mon Sep 17 00:00:00 2001
From: "Jason R. Clark" <jclark@newrelic.com>
Date: Thu, 12 Mar 2015 10:32:35 -0700
Subject: [PATCH] RUBY-1424 Close connections on ForceRestartException

If a shared connection was being held when a ForceRestartException was
received, this wouldn't close the connection. While not directly a
problem since the connection is likely to die and be re-established next
time around, it's best if we clean things up when we're restarting.
---
 lib/new_relic/agent/agent.rb                         |  2 +-
 lib/new_relic/agent/new_relic_service.rb             |  9 +++++++++
 .../agent/agent/start_worker_thread_test.rb          |  2 ++
 test/new_relic/agent/new_relic_service_test.rb       | 12 ++++++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/lib/new_relic/agent/agent.rb b/lib/new_relic/agent/agent.rb
index dfd25c8f6e..0b149c90a6 100644
--- a/lib/new_relic/agent/agent.rb
+++ b/lib/new_relic/agent/agent.rb
@@ -651,7 +651,7 @@ def create_and_run_event_loop
           def handle_force_restart(error)
             ::NewRelic::Agent.logger.debug error.message
             drop_buffered_data
-            @service.reset_metric_id_cache if @service
+            @service.force_restart if @service
             @connect_state = :pending
             sleep 30
           end
diff --git a/lib/new_relic/agent/new_relic_service.rb b/lib/new_relic/agent/new_relic_service.rb
index 9cfdfd7c95..a35a06beba 100644
--- a/lib/new_relic/agent/new_relic_service.rb
+++ b/lib/new_relic/agent/new_relic_service.rb
@@ -86,6 +86,11 @@ def reset_metric_id_cache
         @metric_id_cache = {}
       end
 
+      def force_restart
+        reset_metric_id_cache
+        close_shared_connection
+      end
+
       # takes an array of arrays of spec and id, adds it into the
       # metric cache so we can save the collector some work by
       # sending integers instead of strings the next time around
@@ -244,6 +249,10 @@ def close_shared_connection
         end
       end
 
+      def has_shared_connection?
+        !@shared_tcp_connection.nil?
+      end
+
       def ssl_cert_store
         path = cert_file_path
         if !@ssl_cert_store || path != @cached_cert_store_path
diff --git a/test/new_relic/agent/agent/start_worker_thread_test.rb b/test/new_relic/agent/agent/start_worker_thread_test.rb
index 6bbbbc6a85..62644312bd 100644
--- a/test/new_relic/agent/agent/start_worker_thread_test.rb
+++ b/test/new_relic/agent/agent/start_worker_thread_test.rb
@@ -28,7 +28,9 @@ def test_handle_force_restart
 
     self.expects(:drop_buffered_data)
     self.expects(:sleep).with(30)
+
     @connected = true
+    @service = mock('service', :force_restart => nil)
 
     handle_force_restart(error)
 
diff --git a/test/new_relic/agent/new_relic_service_test.rb b/test/new_relic/agent/new_relic_service_test.rb
index 2f9d8b683b..751f34baff 100644
--- a/test/new_relic/agent/new_relic_service_test.rb
+++ b/test/new_relic/agent/new_relic_service_test.rb
@@ -825,6 +825,18 @@ def test_supportability_metrics_with_serialization_failure
     ])
   end
 
+  def test_force_restart_clears_metric_cache
+    @service.metric_id_cache[1] = "boo"
+    @service.force_restart
+    assert_empty @service.metric_id_cache
+  end
+
+  def test_force_restart_closes_shared_connections
+    @service.establish_shared_connection
+    @service.force_restart
+    refute @service.has_shared_connection?
+  end
+
   def build_stats_hash(items={})
     hash = NewRelic::Agent::StatsHash.new
     items.each do |key, value|