From 393b2368075789d25181bd45c5007806074dfbd2 Mon Sep 17 00:00:00 2001
From: Zach McCormick <zachary.tyler.mccormick@gmail.com>
Date: Mon, 15 Apr 2024 16:57:07 -0400
Subject: [PATCH] Improve key quantization and add command length

---
 lib/datadog/tracing/contrib/dalli/ext.rb      |  3 +++
 .../tracing/contrib/dalli/instrumentation.rb  |  3 +++
 lib/datadog/tracing/contrib/dalli/quantize.rb | 22 ++++++++++++++++++-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/lib/datadog/tracing/contrib/dalli/ext.rb b/lib/datadog/tracing/contrib/dalli/ext.rb
index 405956bc7ae..299c08d3b54 100644
--- a/lib/datadog/tracing/contrib/dalli/ext.rb
+++ b/lib/datadog/tracing/contrib/dalli/ext.rb
@@ -26,6 +26,9 @@ module Ext
           SPAN_COMMAND = 'memcached.command'
           SPAN_TYPE_COMMAND = 'memcached'
           TAG_COMMAND = 'memcached.command'
+          # BEGIN BRAZE MODIFICATION
+          TAG_LENGTH = 'memcached.length'
+          # END BRAZE MODIFICATION
           TAG_COMPONENT = 'dalli'
           TAG_OPERATION_COMMAND = 'command'
           TAG_SYSTEM = 'memcached'
diff --git a/lib/datadog/tracing/contrib/dalli/instrumentation.rb b/lib/datadog/tracing/contrib/dalli/instrumentation.rb
index 332ddbc4d1c..0bffc2dc630 100644
--- a/lib/datadog/tracing/contrib/dalli/instrumentation.rb
+++ b/lib/datadog/tracing/contrib/dalli/instrumentation.rb
@@ -54,6 +54,9 @@ def request(op, *args)
                   cmd = Quantize.format_command(op, args)
                   span.set_tag(Ext::TAG_COMMAND, cmd)
                 end
+                # BEGIN BRAZE MODIFICATION
+                span.set_tag(Ext::TAG_LENGTH, Core::Utils.utf8_encode([op, *args].join(' ').strip, binary: true).length)
+                # END BRAZE MODIFICATION
 
                 Contrib::SpanAttributeSchema.set_peer_service!(span, Ext::PEER_SERVICE_SOURCES)
                 super
diff --git a/lib/datadog/tracing/contrib/dalli/quantize.rb b/lib/datadog/tracing/contrib/dalli/quantize.rb
index 6df12264f40..8e058946b50 100644
--- a/lib/datadog/tracing/contrib/dalli/quantize.rb
+++ b/lib/datadog/tracing/contrib/dalli/quantize.rb
@@ -6,12 +6,32 @@ module Contrib
       module Dalli
         # Quantize contains dalli-specic quantization tools.
         module Quantize
+          # BEGIN BRAZE MODIFICATION
+          GUID_ID_REGEX = /[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}/i
+          BSON_ID_REGEX = /[0-9a-f]{24}/i
+          XXHASH_REGEX = /[0-9a-f]{16}/
+          INTEGER_ID_REGEX = /\d+/
+          # END BRAZE MODIFICATION
           module_function
 
           def format_command(operation, args)
             placeholder = "#{operation} BLOB (OMITTED)"
-            command = [operation, *args].join(' ').strip
+            # BEGIN BRAZE MODIFICATION
+            if operation == :send_multiget
+              command = [operation, *args].join(' ').strip
+            else
+              # all operations except multiget have the key as the first arg
+              command = [operation, args[0]].join(' ').strip
+            end
+            # END BRAZE MODIFICATION
+
             command = Core::Utils.utf8_encode(command, binary: true, placeholder: placeholder)
+            # BEGIN BRAZE MODIFICATION
+            command = command.gsub(GUID_ID_REGEX, "GUID")
+            command = command.gsub(BSON_ID_REGEX, "BSON")
+            command = command.gsub(XXHASH_REGEX, "XXHASH")
+            command = command.gsub(INTEGER_ID_REGEX, "INTEGER")
+            # END BRAZE MODIFICATION
             Core::Utils.truncate(command, Ext::QUANTIZE_MAX_CMD_LENGTH)
           rescue => e
             Datadog.logger.debug("Error sanitizing Dalli operation: #{e}")