Merge pull request #227 from DataDog/pr163

Allow setting socket path with DD_DOGSTATSD_SOCKET

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 [//]: # (comment: Don't forget to update lib/datadog/statsd/version.rb:DogStatsd::Statsd::VERSION when releasing a new version)
 
+  * [ENHANCEMENT] The client can now be configured to use UDS via the `DD_DOGSTATSD_SOCKET` environment variable.
+    This variable does not take precedence over any explicit parameters passed to the Statsd constructor.
+
 ## 5.3.2 / 2021.11.03
 
   * [OTHER] add a warning message for the v5.x update on install #222 by @djmitche

lib/datadog/statsd.rb

@@ -5,6 +5,7 @@
 require_relative 'statsd/telemetry'
 require_relative 'statsd/udp_connection'
 require_relative 'statsd/uds_connection'
+require_relative 'statsd/connection_cfg'
 require_relative 'statsd/message_buffer'
 require_relative 'statsd/serialization'
 require_relative 'statsd/sender'
@@ -118,9 +119,11 @@ def initialize(
       end
 
       @forwarder = Forwarder.new(
-        host: host,
-        port: port,
-        socket_path: socket_path,
+        connection_cfg: ConnectionCfg.new(
+          host: host,
+          port: port,
+          socket_path: socket_path,
+        ),
 
         global_tags: tags,
         logger: logger,

lib/datadog/statsd/connection_cfg.rb

@@ -0,0 +1,76 @@
+module Datadog
+  class Statsd
+    class ConnectionCfg
+      attr_reader :host
+      attr_reader :port
+      attr_reader :socket_path
+      attr_reader :transport_type
+
+      def initialize(host: nil, port: nil, socket_path: nil)
+        initialize_with_constructor_args(host: host, port: port, socket_path: socket_path) ||
+          initialize_with_env_vars ||
+          initialize_with_defaults
+      end
+
+      def make_connection(**params)
+        case @transport_type
+        when :udp
+          UDPConnection.new(@host, @port, **params)
+        when :uds
+          UDSConnection.new(@socket_path, **params)
+        end
+      end
+
+      private
+
+      DEFAULT_HOST = '127.0.0.1'
+      DEFAULT_PORT = 8125
+
+      def initialize_with_constructor_args(host: nil, port: nil, socket_path: nil)
+        try_initialize_with(host: host, port: port, socket_path: socket_path,
+          not_both_error_message: 
+            "Both UDP: (host/port #{host}:#{port}) and UDS (socket_path #{socket_path}) " +
+            "constructor arguments were given. Use only one or the other.",
+          )
+      end
+
+      def initialize_with_env_vars()
+        try_initialize_with(
+          host: ENV['DD_AGENT_HOST'],
+          port: ENV['DD_DOGSTATSD_PORT'] && ENV['DD_DOGSTATSD_PORT'].to_i,
+          socket_path: ENV['DD_DOGSTATSD_SOCKET'],
+          not_both_error_message:
+            "Both UDP (DD_AGENT_HOST/DD_DOGSTATSD_PORT #{ENV['DD_AGENT_HOST']}:#{ENV['DD_DOGSTATSD_PORT']}) " +
+            "and UDS (DD_DOGSTATSD_SOCKET #{ENV['DD_DOGSTATSD_SOCKET']}) environment variables are set. " +
+            "Set only one or the other." %
+            [ENV['DD_AGENT_HOST'], ENV['DD_DOGSTATSD_PORT'], ENV['DD_DOGSTATSD_SOCKET']])
+      end
+
+      def initialize_with_defaults()
+        try_initialize_with(host: DEFAULT_HOST, port: DEFAULT_PORT)
+      end
+
+      def try_initialize_with(host: nil, port: nil, socket_path: nil, not_both_error_message: "")
+        if (host || port) && socket_path
+          raise ArgumentError, not_both_error_message
+        end
+
+        if host || port 
+          @host = host || DEFAULT_HOST
+          @port = port || DEFAULT_PORT
+          @socket_path = nil
+          @transport_type = :udp
+          return true
+        elsif socket_path
+          @host = nil
+          @port = nil
+          @socket_path = socket_path
+          @transport_type = :uds
+          return true
+        end
+
+        return false
+      end
+    end
+  end
+end

lib/datadog/statsd/forwarder.rb

@@ -7,9 +7,7 @@ class Forwarder
       attr_reader :transport_type
 
       def initialize(
-        host: nil,
-        port: nil,
-        socket_path: nil,
+        connection_cfg: nil,
 
         buffer_max_payload_size: nil,
         buffer_max_pool_size: nil,
@@ -22,24 +20,20 @@ def initialize(
 
         logger: nil
       )
-        @transport_type = socket_path.nil? ? :udp : :uds
+        @transport_type = connection_cfg.transport_type
 
         if telemetry_flush_interval
           @telemetry = Telemetry.new(telemetry_flush_interval,
             global_tags: global_tags,
-            transport_type: transport_type
+            transport_type: @transport_type
           )
         end
 
-        @connection = case transport_type
-                      when :udp
-                        UDPConnection.new(host, port, logger: logger, telemetry: telemetry)
-                      when :uds
-                        UDSConnection.new(socket_path, logger: logger, telemetry: telemetry)
-                      end
+        @connection = connection_cfg.make_connection(logger: logger, telemetry: telemetry)
 
         # Initialize buffer
-        buffer_max_payload_size ||= (transport_type == :udp ? UDP_DEFAULT_BUFFER_SIZE : UDS_DEFAULT_BUFFER_SIZE)
+        buffer_max_payload_size ||= (@transport_type == :udp ?
+                                     UDP_DEFAULT_BUFFER_SIZE : UDS_DEFAULT_BUFFER_SIZE)
 
         if buffer_max_payload_size <= 0
           raise ArgumentError, 'buffer_max_payload_size cannot be <= 0'

lib/datadog/statsd/udp_connection.rb

@@ -5,20 +5,17 @@
 module Datadog
   class Statsd
     class UDPConnection < Connection
-      DEFAULT_HOST = '127.0.0.1'
-      DEFAULT_PORT = 8125
-
-      # StatsD host. Defaults to 127.0.0.1.
+      # StatsD host.
       attr_reader :host
 
-      # StatsD port. Defaults to 8125.
+      # StatsD port.
       attr_reader :port
 
       def initialize(host, port, **kwargs)
         super(**kwargs)
 
-        @host = host || ENV.fetch('DD_AGENT_HOST', DEFAULT_HOST)
-        @port = port || ENV.fetch('DD_DOGSTATSD_PORT', DEFAULT_PORT).to_i
+        @host = host
+        @port = port
         @socket = nil
       end
 

spec/statsd/connection_cfg_spec.rb

@@ -0,0 +1,219 @@
+require 'spec_helper'
+
+describe Datadog::Statsd::ConnectionCfg do
+  subject do
+    described_class.new(host: host, port: port, socket_path: socket_path)
+  end
+
+  around do |example|
+    ClimateControl.modify(
+      'DD_AGENT_HOST' => dd_agent_host,
+      'DD_DOGSTATSD_PORT' => dd_dogstatsd_port,
+      'DD_DOGSTATSD_SOCKET' => dd_dogstatsd_socket,
+    ) do
+      example.run
+    end
+  end
+
+  let(:host) { nil }
+  let(:port) { nil }
+  let(:socket_path) { nil }
+  let(:dd_agent_host) { nil }
+  let(:dd_dogstatsd_port) { nil }
+  let(:dd_dogstatsd_socket) { nil }
+
+  describe '#initialize' do
+    context 'with host/port args and env vars set' do
+      let(:host) { 'my-agent' }
+      let(:port) { 1234 }
+      let(:dd_agent_host) { 'unused' }
+      let(:dd_dogstatsd_port) { '999' }
+      let(:dd_dogstatsd_socket) { '/un/used' }
+
+      it 'creates a UDP connection' do
+        expect(subject.transport_type).to eq :udp
+      end
+
+      it 'uses the agent name from the args' do
+        expect(subject.host).to eq 'my-agent'
+      end
+
+      it 'uses the port name from the args' do
+        expect(subject.port).to eq 1234
+      end
+
+      it 'sets socket_path to nil' do
+        expect(subject.socket_path).to eq nil
+      end
+    end
+
+    context 'with both host/port and socket args' do
+      let(:host) { 'my-agent' }
+      let(:port) { 1234 }
+      let(:socket_path) { '/some/socket' }
+
+      it 'raises an exception' do
+        expect do
+          subject.new(host: host, port: port, socket_path: socket_path)
+        end.to raise_error(
+          ArgumentError,
+          "Both UDP: (host/port my-agent:1234) and UDS (socket_path /some/socket) constructor arguments were given. Use only one or the other.")
+      end
+    end
+
+    context 'with socket_path arg and env vars' do
+      let(:socket_path) { '/some/socket' }
+      let(:dd_agent_host) { 'unused' }
+      let(:dd_dogstatsd_port) { '999' }
+      let(:dd_dogstatsd_socket) { '/un/used' }
+
+      it 'creates a UDS connection' do
+        expect(subject.transport_type).to eq :uds
+      end
+
+      it 'sets host to nil' do
+        expect(subject.host).to eq nil
+      end
+
+      it 'sets port to nil' do
+        expect(subject.port).to eq nil
+      end
+
+      it 'sets socket_path to path in the arg' do
+        expect(subject.socket_path).to eq '/some/socket'
+      end
+    end
+
+    context 'with no args and DD_AGENT_HOST set' do
+      let(:dd_agent_host) { 'some-host' }
+
+      it 'creates a UDP connection' do
+        expect(subject.transport_type).to eq :udp
+      end
+
+      it 'sets host to DD_AGENT_HOST' do
+        expect(subject.host).to eq 'some-host'
+      end
+
+      it 'sets port to 8125 (default)' do
+        expect(subject.port).to eq 8125
+      end
+
+      it 'sets socket_path to nil' do
+        expect(subject.socket_path).to eq nil
+      end
+
+      context 'and DD_DOGSTATSD_PORT set' do
+        let(:dd_dogstatsd_port) { '1234' }
+
+        it 'creates a UDP connection' do
+          expect(subject.transport_type).to eq :udp
+        end
+
+        it 'sets host to DD_AGENT_HOST' do
+          expect(subject.host).to eq 'some-host'
+        end
+
+        it 'sets port to DD_DOGSTATSD_PORT' do
+          expect(subject.port).to eq 1234
+        end
+
+        it 'sets socket_path to nil' do
+          expect(subject.socket_path).to eq nil
+        end
+      end
+    end
+
+    context 'with no args and DD_DOGSTATSD_SOCKET set' do
+      let(:dd_dogstatsd_socket) { '/some/socket' }
+
+      it 'creates a UDS connection' do
+        expect(subject.transport_type).to eq :uds
+      end
+
+      it 'sets host to nil' do
+        expect(subject.host).to eq nil
+      end
+
+      it 'sets port to nil' do
+        expect(subject.port).to eq nil
+      end
+
+      it 'sets socket_path to DD_DOGSTATSD_SOCKET' do
+        expect(subject.socket_path).to eq '/some/socket'
+      end
+    end
+
+    context 'with both DD_AGENT_HOST and DD_DOGSTATSD_SOCKET set' do
+      let(:dd_agent_host) { 'some-host' }
+      let(:dd_dogstatsd_socket) { '/some/socket' }
+
+      it 'raises an exception' do
+        expect do
+          subject.new(host: host, port: port, socket_path: socket_path)
+        end.to raise_error(
+          ArgumentError,
+          'Both UDP (DD_AGENT_HOST/DD_DOGSTATSD_PORT some-host:) and UDS (DD_DOGSTATSD_SOCKET /some/socket) environment variables are set. Set only one or the other.')
+      end
+    end
+
+    context 'with no args and no env vars set' do
+      it 'creates a UDP connection' do
+        expect(subject.transport_type).to eq :udp
+      end
+
+      it 'sets host to 127.0.0.1 (default)' do
+        expect(subject.host).to eq '127.0.0.1'
+      end
+
+      it 'sets port to 8125 (default)' do
+        expect(subject.port).to eq 8125
+      end
+
+      it 'sets socket_path to nil' do
+        expect(subject.socket_path).to eq nil
+      end
+    end
+  end
+
+  describe '#make_connection' do
+    context 'for a UDP connection' do
+      before do
+        allow(Datadog::Statsd::UDPConnection)
+          .to receive(:new)
+          .with(host, port, param: 'param')
+          .and_return(udp_connection)
+      end
+
+      let(:host) { 'my-agent' }
+      let(:port) { 1234 }
+      let(:udp_connection) do
+        instance_double(Datadog::Statsd::UDPConnection)
+      end
+
+      it 'creates a UDP connection, passing along params' do
+        expect(subject.make_connection(param: 'param')).to eq udp_connection
+      end
+    end
+
+    context 'for a UDS connection' do
+      before do
+        allow(Datadog::Statsd::UDSConnection)
+          .to receive(:new)
+          .with(socket_path, param: 'param')
+          .and_return(uds_connection)
+      end
+
+      let(:socket_path) { '/tmp/dd-socket' }
+      let(:uds_connection) do
+        instance_double(Datadog::Statsd::UDSConnection,
+          socket_path: socket_path
+        )
+      end
+
+      it 'creates a UDS connection, passing along params' do
+        expect(subject.make_connection(param: 'param')).to eq uds_connection
+      end
+    end
+  end
+end