a datadog agent fake server to intercept ddtrace requests. NOT compatible with ddtrace v1

datadog-mock.rb

# frozen_string_literal: true

require 'roda'
require 'pry'
require 'msgpack'
require 'json'
require 'zlib'
require 'multipart_parser/reader'
require 'google/protobuf'

Google::Protobuf::DescriptorPool.generated_pool.build do
  add_message "perftools.profiles.Profile" do
    repeated :sample_type, :message, 1, "perftools.profiles.ValueType"
    repeated :sample, :message, 2, "perftools.profiles.Sample"
    repeated :mapping, :message, 3, "perftools.profiles.Mapping"
    repeated :location, :message, 4, "perftools.profiles.Location"
    repeated :function, :message, 5, "perftools.profiles.Function"
    repeated :string_table, :string, 6
    optional :drop_frames, :int64, 7
    optional :keep_frames, :int64, 8
    optional :time_nanos, :int64, 9
    optional :duration_nanos, :int64, 10
    optional :period_type, :message, 11, "perftools.profiles.ValueType"
    optional :period, :int64, 12
    repeated :comment, :int64, 13
    optional :default_sample_type, :int64, 14
  end
  add_message "perftools.profiles.ValueType" do
    optional :type, :int64, 1
    optional :unit, :int64, 2
  end
  add_message "perftools.profiles.Sample" do
    repeated :location_id, :uint64, 1
    repeated :value, :int64, 2
    repeated :label, :message, 3, "perftools.profiles.Label"
  end
  add_message "perftools.profiles.Label" do
    optional :key, :int64, 1
    optional :str, :int64, 2
    optional :num, :int64, 3
    optional :num_unit, :int64, 4
  end
  add_message "perftools.profiles.Mapping" do
    optional :id, :uint64, 1
    optional :memory_start, :uint64, 2
    optional :memory_limit, :uint64, 3
    optional :file_offset, :uint64, 4
    optional :filename, :int64, 5
    optional :build_id, :int64, 6
    optional :has_functions, :bool, 7
    optional :has_filenames, :bool, 8
    optional :has_line_numbers, :bool, 9
    optional :has_inline_frames, :bool, 10
  end
  add_message "perftools.profiles.Location" do
    optional :id, :uint64, 1
    optional :mapping_id, :uint64, 2
    optional :address, :uint64, 3
    repeated :line, :message, 4, "perftools.profiles.Line"
    optional :is_folded, :bool, 5
  end
  add_message "perftools.profiles.Line" do
    optional :function_id, :uint64, 1
    optional :line, :int64, 2
  end
  add_message "perftools.profiles.Function" do
    optional :id, :uint64, 1
    optional :name, :int64, 2
    optional :system_name, :int64, 3
    optional :filename, :int64, 4
    optional :start_line, :int64, 5
  end
end

module Perftools
  module Profiles
    Profile = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Profile").msgclass
    ValueType = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.ValueType").msgclass
    Sample = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Sample").msgclass
    Label = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Label").msgclass
    Mapping = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Mapping").msgclass
    Location = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Location").msgclass
    Line = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Line").msgclass
    Function = Google::Protobuf::DescriptorPool.generated_pool.lookup("perftools.profiles.Function").msgclass
  end
end

module Utils
  extend self

  def parse_multipart(request)
    boundary = request.content_type[/boundary=(.*)$/, 1]
    body = request.body.read
    request.body.rewind
    reader = MultipartParser::Reader.new(boundary)
    result = { errors: [], parts: [] }
    def result.part(name)
      hash = self[:parts].detect { |h| h[:part].name == name }
      [hash[:part], hash[:body].join]
    end

    reader.on_part do |part|
      result[:parts] << thispart = {
        part: part,
        body: []
      }
      part.on_data do |chunk|
        thispart[:body] << chunk
      end
    end
    reader.on_error do |msg|
      result[:errors] << msg
    end

    reader.write(body)
    result
  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

class App < Roda
  plugin :json
  @@raw_traces = []
  @@raw_profiles = []
  @@results = []

  route do |r|
    r.get 'results' do |r|
      total = @@results.reduce(0) { |sum, result| sum += result[:duration_s] }
      [{total: total}, *@@results]
    end
    r.get 'raw_traces' do |r|
      @@raw_traces
    end
    r.get 'raw_profiles' do |r|
      binding.pry
      @@raw_profiles.map do |profile|
        profile[:string_table].select {|_| _.match?(/\.rb$/) }
      end
    end
    r.post 'v0.4/traces' do |r|
      traces = MessagePack.unpack(request.body.read)
      traces.each do |traces|
        traces.each do |trace|
          @@raw_traces << trace
          @@results.append({
            name: trace["meta"]["test.name"],
            suite: trace["meta"]["test.suite"],
            duration_s: trace["duration"] / 1000000000.to_f
          })

          puts @@results.size
        end
      end
      {success: true}
    end
    r.post 'profiling/v1/input' do |r|
      result = Utils.parse_multipart(request)
      profiling_part = result.part('data[rubyprofile.pprof]')
      gzipped_content = profiling_part[1]
      encoded_content = Utils.gunzip(gzipped_content)
      @@raw_profiles << Perftools::Profiles::Profile.decode(encoded_content).to_h
      #@@raw_profiles << [Perftools::Profiles::Profile.decode(encoded_content).to_h[:string_table].select{_1.match?(/\.rb$/)}]
      puts 'profile ingested'
      {success: true}
    end
    r.get 'reset' do |r|
      @@raw_traces = @@raw_profiles = @@results = []
    end
    r.is /.*/ do |r|
      binding.pry
      puts 'hello u there'
      {success: true}
    end
  end
end

run App.freeze.app