trace_middleware.rb in ddtrace-1.16.0

- old
+ new
@@ -7,57 +7,70 @@
   module Tracing
     module Contrib
       module Redis
         # Instrumentation for Redis 5+
         module TraceMiddleware
-          def call(commands, redis_config)
-            Tracing.trace(Contrib::Redis::Ext::SPAN_COMMAND) do |span|
-              datadog_configuration = resolve(redis_config)
-              resource = get_command(commands, datadog_configuration[:command_args])
+          # Instruments {RedisClient::ConnectionMixin#call}.
+          def call(command, redis_config)
+            config = resolve(redis_config)
+            TraceMiddleware.call(redis_config, command, config[:service_name], config[:command_args]) { super }
+          end
 
-              span.service = datadog_configuration[:service_name]
-              span.span_type = Contrib::Redis::Ext::TYPE
-              span.resource = resource
+          # Instruments {RedisClient::ConnectionMixin#call_pipelined}.
+          def call_pipelined(commands, redis_config)
+            config = resolve(redis_config)
+            TraceMiddleware.call_pipelined(redis_config, commands, config[:service_name], config[:command_args]) { super }
+          end
 
-              Contrib::Redis::Tags.set_common_tags(redis_config, span, datadog_configuration[:command_args])
+          class << self
+            def call(client, command, service_name, command_args)
+              Tracing.trace(Redis::Ext::SPAN_COMMAND, type: Redis::Ext::TYPE, service: service_name) do |span|
+                raw_command = get_command(command, true)
+                span.resource = command_args ? raw_command : get_command(command, false)
 
-              super
+                Contrib::Redis::Tags.set_common_tags(client, span, raw_command)
+
+                yield
+              end
             end
-          end
 
-          def call_pipelined(commands, redis_config)
-            Tracing.trace(Contrib::Redis::Ext::SPAN_COMMAND) do |span|
-              datadog_configuration = resolve(redis_config)
-              pipelined_commands = get_pipeline_commands(commands, datadog_configuration[:command_args])
+            def call_pipelined(client, commands, service_name, command_args)
+              Tracing.trace(Redis::Ext::SPAN_COMMAND, type: Redis::Ext::TYPE, service: service_name) do |span|
+                raw_command = get_pipeline_commands(commands, true)
+                span.resource = command_args ? raw_command : get_pipeline_commands(commands, false)
 
-              span.service = datadog_configuration[:service_name]
-              span.span_type = Contrib::Redis::Ext::TYPE
-              span.resource = pipelined_commands.join("\n")
-              span.set_metric Contrib::Redis::Ext::METRIC_PIPELINE_LEN, pipelined_commands.length
+                span.set_metric Contrib::Redis::Ext::METRIC_PIPELINE_LEN, commands.length
 
-              Contrib::Redis::Tags.set_common_tags(redis_config, span, datadog_configuration[:command_args])
+                Contrib::Redis::Tags.set_common_tags(client, span, raw_command)
 
-              super
+                yield
+              end
             end
-          end
 
-          private
+            private
 
-          def get_command(commands, boolean)
-            if boolean
-              Contrib::Redis::Quantize.format_command_args(commands)
-            else
-              Contrib::Redis::Quantize.get_verb(commands)
+            # Quantizes a single Redis command
+            def get_command(command, command_args)
+              if command_args
+                Contrib::Redis::Quantize.format_command_args(command)
+              else
+                Contrib::Redis::Quantize.get_verb(command)
+              end
             end
-          end
 
-          def get_pipeline_commands(commands, boolean)
-            if boolean
-              commands.map { |c| Contrib::Redis::Quantize.format_command_args(c) }
-            else
-              commands.map { |c| Contrib::Redis::Quantize.get_verb(c) }
+            # Quantizes a multi-command Redis pipeline execution
+            def get_pipeline_commands(commands, command_args)
+              list = if command_args
+                       commands.map { |c| Contrib::Redis::Quantize.format_command_args(c) }
+                     else
+                       commands.map { |c| Contrib::Redis::Quantize.get_verb(c) }
+                     end
+
+              list.empty? ? '(none)' : list.join("\n")
             end
           end
+
+          private
 
           def resolve(redis_config)
             custom = redis_config.custom[:datadog] || {}
 
             Datadog.configuration.tracing[:redis, redis_config.server_url].to_h.merge(custom)