Add support for more environment variables

Gemfile

@@ -6,6 +6,7 @@ gem 'minitest'
 gem 'minitest-matchers'
 gem 'yard', '~> 0.9.20'
 gem 'single_cov'
+gem 'climate_control'
 
 if RUBY_VERSION >= '2.0.0'
   gem 'rubocop', '~> 0.50.0' # bump this and TargetRubyVersion once we drop ruby 2.0

lib/datadog/statsd/connection.rb

@@ -10,7 +10,7 @@ def initialize(telemetry)
       # Close the underlying socket
       def close
         begin
-          @socket && @socket.close
+          @socket && @socket.close if instance_variable_defined?(:@socket)
         rescue StandardError => boom
           logger.error { "Statsd: #{boom.class} #{boom}" } if logger
         end

lib/datadog/statsd/serialization/tag_serializer.rb

@@ -5,12 +5,14 @@ class Statsd
     module Serialization
       class TagSerializer
         def initialize(global_tags = [], env = ENV)
-          @global_tags = to_tags_list(global_tags)
+          # Convert to hash
+          global_tags = to_tags_hash(global_tags)
 
-          # append the entity id to tags if DD_ENTITY_ID env var is set
-          if dd_entity = env.fetch('DD_ENTITY_ID', nil)
-            @global_tags << to_tags_list('dd.internal.entity_id' => dd_entity).first
-          end
+          # Merge with default tags
+          global_tags = default_tags(env).merge(global_tags)
+
+          # Convert to tag list and set
+          @global_tags = to_tags_list(global_tags)
         end
 
         def format(message_tags)
@@ -28,11 +30,33 @@ def format(message_tags)
         attr_reader :global_tags
 
         private
+
+        def to_tags_hash(tags)
+          case tags
+          when Hash
+            tags.dup
+          when Array
+            Hash[
+              tags.collect do |string|
+                string.split(':').tap do |tokens|
+                  tokens << nil if tokens.length == 1
+                end
+              end
+            ]
+          else
+            {}
+          end
+        end
+
         def to_tags_list(tags)
           case tags
           when Hash
-            tags.each_with_object([]) do |tag_pair, formated_tags|
-              formated_tags << "#{tag_pair.first}:#{tag_pair.last}"
+            tags.each_with_object([]) do |tag_pair, formatted_tags|
+              if tag_pair.last.nil?
+                formatted_tags << "#{tag_pair.first}"
+              else
+                formatted_tags << "#{tag_pair.first}:#{tag_pair.last}"
+              end
             end
           when Array
             tags.dup
@@ -46,6 +70,21 @@ def to_tags_list(tags)
         def escape_tag_content(tag)
           tag.to_s.delete('|,')
         end
+
+        def dd_tags(env = ENV)
+          return {} unless dd_tags = env['DD_TAGS']
+
+          to_tags_hash(dd_tags.split(','))
+        end
+
+        def default_tags(env = ENV)
+          dd_tags(env).tap do |tags|
+            tags['dd.internal.entity_id'] = env['DD_ENTITY_ID'] if env.key?('DD_ENTITY_ID')
+            tags['env'] = env['DD_ENV'] if env.key?('DD_ENV')
+            tags['service'] = env['DD_SERVICE'] if env.key?('DD_SERVICE')
+            tags['version'] = env['DD_VERSION'] if env.key?('DD_VERSION')
+          end
+        end
       end
     end
   end

lib/datadog/statsd/udp_connection.rb

@@ -17,7 +17,7 @@ class UDPConnection < Connection
       def initialize(host, port, logger, telemetry)
         super(telemetry)
         @host = host || ENV.fetch('DD_AGENT_HOST', DEFAULT_HOST)
-        @port = port || ENV.fetch('DD_DOGSTATSD_PORT', DEFAULT_PORT)
+        @port = port || ENV.fetch('DD_DOGSTATSD_PORT', DEFAULT_PORT).to_i
         @logger = logger
       end
 

spec/spec_helper.rb

@@ -8,6 +8,7 @@
 require 'logger'
 require 'faker'
 require 'allocation_stats' if RUBY_VERSION >= '2.3.0'
+require 'climate_control'
 
 Dir[File.join(File.dirname(__FILE__), '/support/**/*.rb')].each { |f| require f }
 Dir[File.join(File.dirname(__FILE__), '/matchers/**/*.rb')].each { |f| require f }

spec/statsd/serialization/tag_serializer_spec.rb

@@ -143,5 +143,77 @@
         expect(subject.format(['name:foobarfoo'.freeze])).to eq 'name:foobarfoo'
       end
     end
+
+    context '[testing management of env vars]' do
+      context 'when testing DD_TAGS' do
+        around do |example|
+          ClimateControl.modify(
+            'DD_TAGS' => 'ghi,team:qa'
+          ) do
+            example.run
+          end
+        end
+
+        it 'correctly adds individual tags' do
+          expect(subject.format([])).to eq 'ghi,team:qa'
+        end
+      end
+
+      context 'when testing DD_ENTITY_ID' do
+        around do |example|
+          ClimateControl.modify(
+            'DD_ENTITY_ID' => '04652bb7-19b7-11e9-9cc6-42010a9c016d'
+          ) do
+            example.run
+          end
+        end
+
+        it 'correctly adds the entity_id tag' do
+          expect(subject.format([])).to eq 'dd.internal.entity_id:04652bb7-19b7-11e9-9cc6-42010a9c016d'
+        end
+      end
+
+      context 'when testing DD_ENV' do
+        around do |example|
+          ClimateControl.modify(
+            'DD_ENV' => 'staging'
+          ) do
+            example.run
+          end
+        end
+
+        it 'correctly adds the env tag' do
+          expect(subject.format([])).to eq 'env:staging'
+        end
+      end
+
+      context 'when testing DD_SERVICE' do
+        around do |example|
+          ClimateControl.modify(
+            'DD_SERVICE' => 'billing-service'
+          ) do
+            example.run
+          end
+        end
+
+        it 'correctly adds the service tag' do
+          expect(subject.format([])).to eq 'service:billing-service'
+        end
+      end
+
+      context 'when testing DD_VERSION' do
+        around do |example|
+          ClimateControl.modify(
+            'DD_VERSION' => '0.1.0-alpha'
+          ) do
+            example.run
+          end
+        end
+
+        it 'correctly adds the version tag' do
+          expect(subject.format([])).to eq 'version:0.1.0-alpha'
+        end
+      end
+    end
   end
 end

spec/statsd_spec.rb

@@ -69,11 +69,19 @@
         )
       end
 
-      before do
-        allow(ENV).to receive(:fetch).and_call_original
-        allow(ENV).to receive(:fetch).with('DD_AGENT_HOST', anything).and_return('myhost')
-        allow(ENV).to receive(:fetch).with('DD_DOGSTATSD_PORT', anything).and_return(4321)
-        allow(ENV).to receive(:fetch).with('DD_ENTITY_ID', anything).and_return('04652bb7-19b7-11e9-9cc6-42010a9c016d')
+
+      around do |example|
+        ClimateControl.modify(
+          'DD_AGENT_HOST' => 'myhost',
+          'DD_DOGSTATSD_PORT' => '4321',
+          'DD_ENTITY_ID' => '04652bb7-19b7-11e9-9cc6-42010a9c016d',
+          'DD_ENV' => 'staging',
+          'DD_SERVICE' => 'billing-service',
+          'DD_VERSION' => '0.1.0-alpha',
+          'DD_TAGS' => 'ghi,team:qa'
+        ) do
+          example.run
+        end
       end
 
       it 'sets the host using the env var DD_AGENT_HOST' do
@@ -85,9 +93,14 @@
       end
 
       it 'sets the entity tag using ' do
-        expect(subject.tags).to eq [
+        expect(subject.tags).to match_array [
           'abc',
           'def',
+          'ghi',
+          'env:staging',
+          'service:billing-service',
+          'team:qa',
+          'version:0.1.0-alpha',
           'dd.internal.entity_id:04652bb7-19b7-11e9-9cc6-42010a9c016d'
         ]
       end