ddtrace 0.41.0 → 0.52.0

data/lib/ddtrace/transport/io.rb

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/encoding'
 require 'ddtrace/transport/io/client'
 require 'ddtrace/transport/io/traces'
@@ -17,7 +18,7 @@ module Datadog
       # Pass options to override any settings.
       def default(options = {})
         new(
-          options.fetch(:out, STDOUT),
+          options.fetch(:out, $stdout),
           options.fetch(:encoder, Encoding::JSONEncoder)
         )
       end

data/lib/ddtrace/transport/io/client.rb

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/transport/statistics'
 require 'ddtrace/transport/io/response'
 
@@ -12,19 +13,24 @@ module Datadog
           :encoder,
           :out
 
-        def initialize(out, encoder)
+        def initialize(out, encoder, options = {})
           @out = out
           @encoder = encoder
+
+          @request_block = options.fetch(:request, method(:send_default_request))
+          @encode_block = options.fetch(:encode, method(:encode_data))
+          @write_block = options.fetch(:write, method(:write_data))
+          @response_block = options.fetch(:response, method(:build_response))
         end
 
         def send_request(request)
           # Write data to IO
           # If block is given, allow it to handle writing
-          # Otherwise use default encoding.
+          # Otherwise do a standard encode/write/response.
           response = if block_given?
                        yield(out, request)
                      else
-                       send_default_request(out, request)
+                       @request_block.call(out, request)
                      end
 
           # Update statistics
@@ -33,7 +39,7 @@ module Datadog
           # Return response
           response
         rescue StandardError => e
-          message = "Internal error during IO transport request. Cause: #{e.message} Location: #{e.backtrace.first}"
+          message = "Internal error during IO transport request. Cause: #{e.message} Location: #{Array(e.backtrace).first}"
 
           # Log error
           if stats.consecutive_errors > 0
@@ -48,8 +54,6 @@ module Datadog
           InternalErrorResponse.new(e)
         end
 
-        protected
-
         def encode_data(encoder, request)
           request.parcel.encode_with(encoder)
         end
@@ -58,17 +62,21 @@ module Datadog
           out.puts(data)
         end
 
+        def build_response(_request, _data, result)
+          IO::Response.new(result)
+        end
+
         private
 
         def send_default_request(out, request)
           # Encode data
-          data = encode_data(encoder, request)
+          data = @encode_block.call(encoder, request)
 
           # Write to IO
-          result = write_data(out, data)
+          result = @write_block.call(out, data)
 
           # Generate a response
-          IO::Response.new(result)
+          @response_block.call(request, data, result)
         end
       end
     end

data/lib/ddtrace/transport/io/response.rb

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/transport/response'
 
 module Datadog
@@ -6,14 +7,12 @@ module Datadog
       # Response from HTTP transport for traces
       class Response
         include Transport::Response
-        include Transport::Traces::Response
 
         attr_reader \
           :result
 
-        def initialize(result, trace_count = 1)
+        def initialize(result)
           @result = result
-          @trace_count = trace_count
         end
 
         def ok?

data/lib/ddtrace/transport/io/traces.rb

@@ -1,3 +1,4 @@
+# typed: false
 require 'ddtrace/transport/traces'
 
 require 'ddtrace/transport/io/response'
@@ -10,10 +11,18 @@ module Datadog
       module Traces
         # Response from HTTP transport for traces
         class Response < IO::Response
+          include Transport::Traces::Response
+
+          def initialize(result, trace_count = 1)
+            super(result)
+            @trace_count = trace_count
+          end
         end
 
         # Extensions for HTTP client
         module Client
+          include Kernel # Ensure that kernel methods are always available (https://sorbet.org/docs/error-reference#7003)
+
           def send_traces(traces)
             # Build a request
             req = Transport::Traces::Request.new(Parcel.new(traces))
@@ -84,7 +93,7 @@ module Datadog
         end
 
         # Add traces behavior to transport components
-        IO::Client.send(:include, Traces::Client)
+        IO::Client.include(Traces::Client)
       end
     end
   end

data/lib/ddtrace/transport/parcel.rb

@@ -1,13 +1,20 @@
+# typed: true
 module Datadog
   module Transport
     # Data transfer object for generic data
     module Parcel
+      include Kernel # Ensure that kernel methods are always available (https://sorbet.org/docs/error-reference#7003)
+
       attr_reader \
         :data
 
       def initialize(data)
         @data = data
       end
+
+      def encode_with(encoder)
+        raise NotImplementedError
+      end
     end
   end
 end

data/lib/ddtrace/transport/request.rb

@@ -1,3 +1,4 @@
+# typed: true
 module Datadog
   module Transport
     # Defines request for transport operations

data/lib/ddtrace/transport/response.rb

@@ -1,3 +1,4 @@
+# typed: false
 module Datadog
   module Transport
     # Defines abstract response for transport operations

data/lib/ddtrace/transport/statistics.rb

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/diagnostics/health'
 
 module Datadog

data/lib/ddtrace/transport/traces.rb

@@ -1,3 +1,4 @@
+# typed: true
 require 'ddtrace/transport/parcel'
 require 'ddtrace/transport/request'
 require 'ddtrace/chunker'
@@ -7,7 +8,7 @@ module Datadog
     module Traces
       # Data transfer object for encoded traces
       class EncodedParcel
-        include Transport::Parcel
+        include Datadog::Transport::Parcel
 
         attr_reader :trace_count
 
@@ -22,7 +23,7 @@ module Datadog
       end
 
       # Traces request
-      class Request < Transport::Request
+      class Request < Datadog::Transport::Request
       end
 
       # Traces response
@@ -123,7 +124,21 @@ module Datadog
                 return send_traces(traces)
               end
             end
-          end.force
+          end
+
+          # Force resolution of lazy enumerator.
+          #
+          # The "correct" method to call here would be `#force`,
+          # as this method was created to force the eager loading
+          # of a lazy enumerator.
+          #
+          # Unfortunately, JRuby < 9.2.9.0 erroneously eagerly loads
+          # the lazy Enumerator during intermediate steps.
+          # This forces us to use `#to_a`, as this method works for both
+          # lazy and regular Enumerators.
+          # Using `#to_a` can mask the fact that we expect a lazy
+          # Enumerator.
+          responses = responses.to_a
 
           Datadog.health_metrics.transport_chunked(responses.size)
 
@@ -142,17 +157,20 @@ module Datadog
 
         def downgrade?(response)
           return false unless apis.fallbacks.key?(@current_api_id)
+
           response.not_found? || response.unsupported?
         end
 
         def downgrade!
           downgrade_api_id = apis.fallbacks[@current_api_id]
           raise NoDowngradeAvailableError, @current_api_id if downgrade_api_id.nil?
+
           change_api!(downgrade_api_id)
         end
 
         def change_api!(api_id)
           raise UnknownApiVersionError, api_id unless apis.key?(api_id)
+
           @current_api_id = api_id
           @client = HTTP::Client.new(current_api)
         end

data/lib/ddtrace/utils.rb

@@ -1,33 +1,42 @@
+# typed: true
 require 'ddtrace/utils/database'
+require 'ddtrace/utils/forking'
 
 module Datadog
   # Utils contains low-level utilities, typically to provide pseudo-random trace IDs.
   module Utils
+    extend Utils::Forking
+
     EMPTY_STRING = ''.encode(::Encoding::UTF_8).freeze
     # We use a custom random number generator because we want no interference
     # with the default one. Using the default prng, we could break code that
     # would rely on srand/rand sequences.
 
-    # Return a span id
+    # Return a randomly generated integer, valid as a Span ID or Trace ID.
+    # This method is thread-safe and fork-safe.
     def self.next_id
-      reset! if was_forked?
-
-      @rnd.rand(Datadog::Span::RUBY_MAX_ID)
+      after_fork! { reset! }
+      id_rng.rand(Datadog::Span::RUBY_MAX_ID) # TODO: This should never return zero
     end
 
-    def self.reset!
-      @pid = Process.pid
-      @rnd = Random.new
+    def self.id_rng
+      @id_rng ||= Random.new
     end
 
-    def self.was_forked?
-      Process.pid != @pid
+    def self.reset!
+      @id_rng = Random.new
     end
 
-    private_class_method :reset!, :was_forked?
-
-    reset!
+    private_class_method :id_rng, :reset!
 
+    # Stringifies `value` and ensures the outcome is
+    # string is no longer than `size`.
+    # `omission` replaces the end of the output if
+    # `value.to_s` does not fit in `size`, to signify
+    # truncation.
+    #
+    # If `omission.size` is larger than `size`, the output
+    # will still be `omission.size` in length.
     def self.truncate(value, size, omission = '...'.freeze)
       string = value.to_s
 
@@ -44,6 +53,8 @@ module Datadog
       string
     end
 
+    # Ensure `str` is a valid UTF-8, ready to be
+    # sent through the tracer transport.
     def self.utf8_encode(str, options = {})
       str = str.to_s
 

data/lib/ddtrace/utils/compression.rb

@@ -0,0 +1,28 @@
+# typed: true
+require 'zlib'
+
+module Datadog
+  module Utils
+    # Common database-related utility functions.
+    module Compression
+      module_function
+
+      def gzip(string, level: nil, strategy: nil)
+        sio = StringIO.new
+        sio.binmode
+        gz = Zlib::GzipWriter.new(sio, level, strategy)
+        gz.write(string)
+        gz.close
+        sio.string
+      end
+
+      def gunzip(string, encoding = ::Encoding::ASCII_8BIT)
+        sio = StringIO.new(string)
+        gz = Zlib::GzipReader.new(sio, encoding: encoding)
+        gz.read
+      ensure
+        gz && gz.close
+      end
+    end
+  end
+end

data/lib/ddtrace/utils/database.rb

@@ -1,3 +1,4 @@
+# typed: true
 module Datadog
   module Utils
     # Common database-related utility functions.

data/lib/ddtrace/utils/forking.rb

@@ -0,0 +1,53 @@
+# typed: false
+module Datadog
+  module Utils
+    # Helper methods for managing forking behavior
+    module Forking
+      def self.included(base)
+        base.prepend(ClassExtensions) if base.is_a?(Class)
+      end
+
+      def self.extended(base)
+        # Explicitly update PID here because there's a case where
+        # the code path that lazily updates the PID may not be exercised
+        # until after a fork occurs, thus causing the event to be missed.
+        # By eagerly setting this, we avoid this scenario.
+        base.update_fork_pid!
+      end
+
+      def after_fork!
+        if forked?
+          yield
+          update_fork_pid!
+          true
+        else
+          false
+        end
+      end
+
+      def forked?
+        Process.pid != fork_pid
+      end
+
+      def update_fork_pid!
+        @fork_pid = Process.pid
+      end
+
+      def fork_pid
+        @fork_pid ||= Process.pid
+      end
+
+      # Adds additional functionality for Classes that implement Forking
+      module ClassExtensions
+        # Addresses an edge case where forking before invoking #update_fork_pid! on the
+        # object will cause forking to not be detected in the fork when it should have.
+        #
+        # This wrapper prevents this by initializing the fork PID when the object is created.
+        def initialize(*args, &block)
+          super
+          update_fork_pid!
+        end
+      end
+    end
+  end
+end

data/lib/ddtrace/utils/object_set.rb

@@ -0,0 +1,40 @@
+# typed: true
+require 'ddtrace/utils/sequence'
+
+module Datadog
+  module Utils
+    # Acts as a unique dictionary of objects
+    class ObjectSet
+      # You can provide a block that defines how the key
+      # for this message type is resolved.
+      def initialize(seed = 0, &block)
+        @sequence = Utils::Sequence.new(seed)
+        @items = {}
+        @key_block = block
+      end
+
+      # Submit an array of arguments that define the message.
+      # If they match an existing message, it will return the
+      # matching object. If it doesn't match, it will yield to
+      # the block with the next ID & args given.
+      def fetch(*args)
+        # TODO: Array hashing is **really** expensive, we probably want to get rid of it in the future
+        key = @key_block ? @key_block.call(*args) : args.hash
+        @items[key] ||= yield(@sequence.next, *args)
+      end
+
+      def length
+        @items.length
+      end
+
+      def objects
+        @items.values
+      end
+
+      def freeze
+        super
+        @items.freeze
+      end
+    end
+  end
+end

data/lib/ddtrace/utils/only_once.rb

@@ -0,0 +1,41 @@
+# typed: true
+# frozen_string_literal: true
+
+module Datadog
+  module Utils
+    # Helper class to execute something only once such as not repeating warning logs, and instrumenting classes
+    # only once.
+    #
+    # Thread-safe when used correctly (e.g. be careful of races when lazily initializing instances of this class).
+    #
+    # Note: In its current state, this class is not Ractor-safe.
+    # In https://github.com/DataDog/dd-trace-rb/pull/1398#issuecomment-797378810 we have a discussion of alternatives,
+    # including an alternative implementation that is Ractor-safe once spent.
+    class OnlyOnce
+      def initialize
+        @mutex = Mutex.new
+        @ran_once = false
+      end
+
+      def run
+        @mutex.synchronize do
+          return if @ran_once
+
+          @ran_once = true
+
+          yield
+        end
+      end
+
+      def ran?
+        @mutex.synchronize { @ran_once }
+      end
+
+      private
+
+      def reset_ran_once_state_for_tests
+        @mutex.synchronize { @ran_once = false }
+      end
+    end
+  end
+end