Wrapped statsd

lib/kubernetes-deploy/deploy_task.rb

@@ -150,7 +150,7 @@ def run!(verify_result: true, allow_protected_ns: false, prune: true)
         @logger.phase_heading("Predeploying priority resources")
         start_priority_resource = Time.now.utc
         predeploy_priority_resources(resources)
-        ::StatsD.measure('priority_resources.duration', StatsD.duration(start_priority_resource), tags: statsd_tags)
+        StatsD.measure('priority_resources.duration', StatsD.duration(start_priority_resource), tags: statsd_tags)
       end
 
       @logger.phase_heading("Deploying all resources")
@@ -161,7 +161,7 @@ def run!(verify_result: true, allow_protected_ns: false, prune: true)
       if verify_result
         start_normal_resource = Time.now.utc
         deploy_resources(resources, prune: prune, verify: true)
-        ::StatsD.measure('normal_resources.duration', StatsD.duration(start_normal_resource), tags: statsd_tags)
+        StatsD.measure('normal_resources.duration', StatsD.duration(start_normal_resource), tags: statsd_tags)
         failed_resources = resources.reject(&:deploy_succeeded?)
         success = failed_resources.empty?
         if !success && failed_resources.all?(&:deploy_timed_out?)
@@ -177,25 +177,25 @@ def run!(verify_result: true, allow_protected_ns: false, prune: true)
         MSG
         @logger.summary.add_paragraph(ColorizedString.new(warning).yellow)
       end
-      ::StatsD.event("Deployment of #{@namespace} succeeded",
+      StatsD.event("Deployment of #{@namespace} succeeded",
         "Successfully deployed all #{@namespace} resources to #{@context}",
         alert_type: "success", tags: statsd_tags << "status:success")
-      ::StatsD.measure('all_resources.duration', StatsD.duration(start), tags: statsd_tags << "status:success")
+      StatsD.measure('all_resources.duration', StatsD.duration(start), tags: statsd_tags << "status:success")
       @logger.print_summary(:success)
     rescue DeploymentTimeoutError
       @logger.print_summary(:timed_out)
-      ::StatsD.event("Deployment of #{@namespace} timed out",
+      StatsD.event("Deployment of #{@namespace} timed out",
         "One or more #{@namespace} resources failed to deploy to #{@context} in time",
         alert_type: "error", tags: statsd_tags << "status:timeout")
-      ::StatsD.measure('all_resources.duration', StatsD.duration(start), tags: statsd_tags << "status:timeout")
+      StatsD.measure('all_resources.duration', StatsD.duration(start), tags: statsd_tags << "status:timeout")
       raise
     rescue FatalDeploymentError => error
       @logger.summary.add_action(error.message) if error.message != error.class.to_s
       @logger.print_summary(:failure)
-      ::StatsD.event("Deployment of #{@namespace} failed",
+      StatsD.event("Deployment of #{@namespace} failed",
         "One or more #{@namespace} resources failed to deploy to #{@context}",
         alert_type: "error", tags: statsd_tags << "status:failed")
-      ::StatsD.measure('all_resources.duration', StatsD.duration(start), tags: statsd_tags << "status:failed")
+      StatsD.measure('all_resources.duration', StatsD.duration(start), tags: statsd_tags << "status:failed")
       raise
     end
 

lib/kubernetes-deploy/kubernetes_resource.rb

@@ -293,7 +293,7 @@ def pretty_status
 
     def report_status_to_statsd(watch_time)
       unless @statsd_report_done
-        ::StatsD.measure('resource.duration', watch_time, tags: statsd_tags)
+        StatsD.measure('resource.duration', watch_time, tags: statsd_tags)
         @statsd_report_done = true
       end
     end

lib/kubernetes-deploy/restart_task.rb

@@ -58,14 +58,14 @@ def perform!(deployments_names = nil)
         raise DeploymentTimeoutError
       end
       raise FatalDeploymentError unless success
-      ::StatsD.distribution('restart.duration', StatsD.duration(start), tags: tags('success', deployments))
+      StatsD.distribution('restart.duration', StatsD.duration(start), tags: tags('success', deployments))
       @logger.print_summary(:success)
     rescue DeploymentTimeoutError
-      ::StatsD.distribution('restart.duration', StatsD.duration(start), tags: tags('timeout', deployments))
+      StatsD.distribution('restart.duration', StatsD.duration(start), tags: tags('timeout', deployments))
       @logger.print_summary(:timed_out)
       raise
     rescue FatalDeploymentError => error
-      ::StatsD.distribution('restart.duration', StatsD.duration(start), tags: tags('failure', deployments))
+      StatsD.distribution('restart.duration', StatsD.duration(start), tags: tags('failure', deployments))
       @logger.summary.add_action(error.message) if error.message != error.class.to_s
       @logger.print_summary(:failure)
       raise

lib/kubernetes-deploy/runner_task.rb

@@ -45,14 +45,14 @@ def run!(task_template:, entrypoint:, args:, env_vars: [], verify_result: true)
       else
         record_status_once(pod)
       end
-      ::StatsD.distribution('task_runner.duration', StatsD.duration(start), tags: statsd_tags('success'))
+      StatsD.distribution('task_runner.duration', StatsD.duration(start), tags: statsd_tags('success'))
       @logger.print_summary(:success)
     rescue DeploymentTimeoutError
-      ::StatsD.distribution('task_runner.duration', StatsD.duration(start), tags: statsd_tags('timeout'))
+      StatsD.distribution('task_runner.duration', StatsD.duration(start), tags: statsd_tags('timeout'))
       @logger.print_summary(:timed_out)
       raise
     rescue FatalDeploymentError
-      ::StatsD.distribution('task_runner.duration', StatsD.duration(start), tags: statsd_tags('failure'))
+      StatsD.distribution('task_runner.duration', StatsD.duration(start), tags: statsd_tags('failure'))
       @logger.print_summary(:failure)
       raise
     end

lib/kubernetes-deploy/statsd.rb

@@ -4,23 +4,36 @@
 
 module KubernetesDeploy
   class StatsD
+    extend ::StatsD
     def self.duration(start_time)
       (Time.now.utc - start_time).round(1)
     end
 
     def self.build
-      ::StatsD.default_sample_rate = 1.0
-      ::StatsD.prefix = "KubernetesDeploy"
+      self.default_sample_rate = 1.0
+      self.prefix = "KubernetesDeploy"
 
       if ENV['STATSD_DEV'].present?
-        ::StatsD.backend = ::StatsD::Instrument::Backends::LoggerBackend.new(Logger.new($stderr))
+        self.backend = ::StatsD::Instrument::Backends::LoggerBackend.new(Logger.new($stderr))
       elsif ENV['STATSD_ADDR'].present?
         statsd_impl = ENV['STATSD_IMPLEMENTATION'].present? ? ENV['STATSD_IMPLEMENTATION'] : "datadog"
-        ::StatsD.backend = ::StatsD::Instrument::Backends::UDPBackend.new(ENV['STATSD_ADDR'], statsd_impl)
+        self.backend = ::StatsD::Instrument::Backends::UDPBackend.new(ENV['STATSD_ADDR'], statsd_impl)
       else
-        ::StatsD.backend = ::StatsD::Instrument::Backends::NullBackend.new
+        self.backend = ::StatsD::Instrument::Backends::NullBackend.new
       end
-      ::StatsD.backend
+    end
+
+    def self.capture_statsd_calls(&block)
+      mock_backend = ::StatsD::Instrument::Backends::CaptureBackend.new
+      old_backend, self.backend = self.backend, mock_backend
+      block.call
+      mock_backend.collected_metrics
+    ensure
+      if old_backend.kind_of?(::StatsD::Instrument::Backends::CaptureBackend)
+        old_backend.collected_metrics.concat(mock_backend.collected_metrics)
+      end
+
+      self.backend = old_backend
     end
   end
 end

test/integration-serial/serial_task_run_test.rb

@@ -29,14 +29,14 @@ def test_run_without_verify_result_fails_if_pod_was_not_created
     ], in_order: true)
   end
 
-  # Run statsd tests in serial because capture_statsd_calls modifies global state in a way
+  # Run statsd tests in serial because KubernetesDeploy::StatsD.capture_statsd_calls modifies global state in a way
   # that makes capturing metrics across parrallel runs unreliable
   def test_failure_statsd_metric_emitted
     bad_ns = "missing"
     task_runner = build_task_runner(ns: bad_ns)
 
     result = false
-    metrics = capture_statsd_calls do
+    metrics = KubernetesDeploy::StatsD.capture_statsd_calls do
       result = task_runner.run(run_params)
     end
 
@@ -54,7 +54,7 @@ def test_success_statsd_metric_emitted
     task_runner = build_task_runner
 
     result = false
-    metrics = capture_statsd_calls do
+    metrics = KubernetesDeploy::StatsD.capture_statsd_calls do
       result = task_runner.run(run_params.merge(verify_result: false))
     end
 
@@ -72,7 +72,7 @@ def test_timedout_statsd_metric_emitted
     task_runner = build_task_runner(max_watch_seconds: 0)
 
     result = false
-    metrics = capture_statsd_calls do
+    metrics = KubernetesDeploy::StatsD.capture_statsd_calls do
       result = task_runner.run(run_params.merge(args: ["sleep 5"]))
     end
 

test/integration/kubernetes_deploy_test.rb

@@ -1023,7 +1023,7 @@ def test_adds_namespace_labels_to_statsd_tags
     desired_tags = %W(context:#{KubeclientHelper::TEST_CONTEXT} namespace:#{@namespace} foo:bar)
     hello_cloud = FixtureSetAssertions::HelloCloud.new(@namespace)
     kubeclient.patch_namespace(hello_cloud.namespace, metadata: { labels: { foo: 'bar' } })
-    metrics = capture_statsd_calls do
+    metrics = KubernetesDeploy::StatsD.capture_statsd_calls do
       assert_deploy_success deploy_fixtures("hello-cloud", subset: ["configmap-data.yml"])
     end
 

test/unit/kubernetes-deploy/statsd_test.rb

@@ -13,7 +13,7 @@ def test_build_when_statsd_addr_env_present_but_statsd_implementation_is_not
 
     KubernetesDeploy::StatsD.build
 
-    assert_equal :datadog, StatsD.backend.implementation
+    assert_equal :datadog, KubernetesDeploy::StatsD.backend.implementation
   ensure
     ENV['STATSD_ADDR'] = original_addr
     ENV['STATSD_IMPLEMENTATION'] = original_impl