Dalli quantize optimization

lib/datadog/tracing/contrib/dalli/quantize.rb

@@ -10,7 +10,21 @@ module Quantize
 
           def format_command(operation, args)
             placeholder = "#{operation} BLOB (OMITTED)"
-            command = [operation, *args].join(' ').strip
+            command = +operation.to_s
+
+            args.each do |arg|
+              str = arg.to_s
+
+              if str.bytesize >= Ext::QUANTIZE_MAX_CMD_LENGTH
+                command << ' ' << Core::Utils.truncate(str, Ext::QUANTIZE_MAX_CMD_LENGTH)
+                break
+              elsif !str.empty?
+                command << ' ' << str
+              end
+
+              break if command.length >= Ext::QUANTIZE_MAX_CMD_LENGTH
+            end
+
             command = Core::Utils.utf8_encode(command, binary: true, placeholder: placeholder)
             Core::Utils.truncate(command, Ext::QUANTIZE_MAX_CMD_LENGTH)
           rescue => e

spec/datadog/tracing/contrib/dalli/quantize_spec.rb

@@ -7,25 +7,40 @@
 RSpec.describe Datadog::Tracing::Contrib::Dalli::Quantize do
   describe '#format_command' do
     subject(:formatted_command) { described_class.format_command(op, args) }
+    let(:op) { :set }
 
     context 'output' do
-      let(:op) { :set }
-      let(:args) { [123, 'foo', nil] }
+      context 'given `nil` as last element' do
+        let(:args) { [123, 'foo', nil] }
 
-      it { is_expected.to eq('set 123 foo') }
+        it { is_expected.to eq('set 123 foo') }
+      end
+
+      context 'given `nil` within an array' do
+        let(:op) { :set }
+        let(:args) { [123, nil, 'foo'] }
+
+        it { is_expected.to eq('set 123 foo') }
+      end
+
+      context 'given a large binary as key' do
+        let(:bytes) { Random.bytes(100) }
+        let(:args) { [bytes, 'foo'] }
+
+        it {
+          is_expected.to eq('set 123 foo')
+        }
+      end
     end
 
     context 'truncation' do
-      let(:op) { :set }
-      let(:args) { ['foo', 'A' * 100] }
+      let(:args) { ['foo', 'A' * 100, 'B' * 100] }
 
       it { expect(formatted_command.size).to eq(Datadog::Tracing::Contrib::Dalli::Ext::QUANTIZE_MAX_CMD_LENGTH) }
-      it { is_expected.to end_with('...') }
       it { is_expected.to eq("set foo #{'A' * 89}...") }
     end
 
     context 'different encodings' do
-      let(:op) { :set }
       let(:args) { ["\xa1".force_encoding('iso-8859-1'), "\xa1\xa1".force_encoding('euc-jp')] }
 
       it { is_expected.to match(/BLOB \(OMITTED\)/) }

spec/ddtrace/benchmark/dalli_quantization_spec.rb

@@ -0,0 +1,85 @@
+# require 'spec_helper'
+
+require 'benchmark/ips'
+require 'ddtrace'
+require 'datadog/tracing/contrib/dalli/quantize'
+
+operation = :set
+args = ['foo', Random.bytes(1_000_000), *Array.new(100, 1)].freeze
+
+module Original
+  module_function
+
+  def format_command(operation, args)
+    placeholder = "#{operation} BLOB (OMITTED)"
+    command = [operation, *args].join(' ').strip
+    command = Datadog::Core::Utils.utf8_encode(command, binary: true, placeholder: placeholder)
+    Datadog::Core::Utils.truncate(command, 100)
+  rescue => e
+    puts("Error sanitizing Dalli operation: #{e}")
+    placeholder
+  end
+end
+
+module StaticObfuscation
+  module_function
+
+  def format_command(operation, args)
+    placeholder = "#{operation} BLOB (OMITTED)"
+    command = +operation.to_s
+
+    args.each do |arg|
+      str = arg.to_s
+
+      if str.bytesize >= 100
+        command << ' (TOO LARGE OMITTED)'
+        break
+      elsif !str.empty?
+        command << ' ' << str
+      end
+
+      break if command.length >= 100
+    end
+
+    command = Datadog::Core::Utils.utf8_encode(command, binary: true, placeholder: placeholder)
+    Datadog::Core::Utils.truncate(command, 100)
+  rescue => e
+    puts("Error sanitizing Dalli operation: #{e}")
+    placeholder
+  end
+end
+
+module ParsedObfuscation
+  module_function
+
+  def format_command(operation, args)
+    placeholder = "#{operation} BLOB (OMITTED)"
+    command = +operation.to_s
+
+    args.each do |arg|
+      str = arg.to_s
+
+      if str.bytesize >= 100
+        command << ' ' << Datadog::Core::Utils.truncate(str, 100)
+        break
+      elsif !str.empty?
+        command << ' ' << str
+      end
+
+      break if command.length >= 100
+    end
+
+    command = Datadog::Core::Utils.utf8_encode(command, binary: true, placeholder: placeholder)
+    Datadog::Core::Utils.truncate(command, 100)
+  rescue => e
+    puts("Error sanitizing Dalli operation: #{e}")
+    placeholder
+  end
+end
+
+Benchmark.ips do |x|
+  x.report('original') { Original.format_command(operation, args) }
+  x.report('parsed') { ParsedObfuscation.format_command(operation, args) }
+  x.report('static') { StaticObfuscation.format_command(operation, args) }
+  x.compare!(order: :baseline)
+end