From 3207b16ea39d9a5b7af0d7a1ae392058d2c39656 Mon Sep 17 00:00:00 2001
From: Yury Lebedev <yury.lebedev@datadoghq.com>
Date: Fri, 13 Sep 2024 20:57:55 +0200
Subject: [PATCH] Improve rack instrumentation and specs

---
 .../action_dispatch/instrumentation.rb        |  2 --
 lib/datadog/tracing/contrib/grape/endpoint.rb |  2 +-
 .../tracing/contrib/rack/middlewares.rb       |  2 +-
 .../action_dispatch/instrumentation_spec.rb   |  8 --------
 .../action_dispatch/journey/router_spec.rb    | 20 +++++++++++++++++--
 5 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/lib/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation.rb b/lib/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation.rb
index f527bd99558..d18ff8e2ca1 100644
--- a/lib/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation.rb
+++ b/lib/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation.rb
@@ -24,8 +24,6 @@ def set_http_route_tags(route_spec, script_name)
               if script_name && !script_name.empty?
                 request_trace.set_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE_PATH, script_name)
               end
-            rescue StandardError => e
-              Datadog.logger.error(e.message)
             end
 
             def dispatcher_route?(route)
diff --git a/lib/datadog/tracing/contrib/grape/endpoint.rb b/lib/datadog/tracing/contrib/grape/endpoint.rb
index 37b8448b9a0..69e8da18413 100644
--- a/lib/datadog/tracing/contrib/grape/endpoint.rb
+++ b/lib/datadog/tracing/contrib/grape/endpoint.rb
@@ -66,7 +66,7 @@ def endpoint_start_process(_name, _start, _finish, _id, payload)
               if (grape_route = env['grape.routing_args'] && env['grape.routing_args'][:route_info])
                 trace.set_tag(
                   Tracing::Metadata::Ext::HTTP::TAG_ROUTE,
-                  grape_route.path&.gsub(/\(\.{1}:?\w+\)\z/, '')
+                  grape_route.path&.gsub(/\(\.:?\w+\)\z/, '')
                 )
 
                 trace.set_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE_PATH, env['SCRIPT_NAME'])
diff --git a/lib/datadog/tracing/contrib/rack/middlewares.rb b/lib/datadog/tracing/contrib/rack/middlewares.rb
index 68e5cf0bc3e..869093a7e06 100644
--- a/lib/datadog/tracing/contrib/rack/middlewares.rb
+++ b/lib/datadog/tracing/contrib/rack/middlewares.rb
@@ -139,7 +139,7 @@ def set_request_tags!(trace, request_span, env, status, headers, response, origi
             request_span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_SERVER)
 
             if status != 404 && (last_route = trace.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE))
-              last_script_name = trace.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE_PATH).to_s
+              last_script_name = trace.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE_PATH).to_s || ''
 
               # If the last_script_name is empty but the env['SCRIPT_NAME'] is NOT empty
               # then the current rack request was not routed and must be accounted for
diff --git a/spec/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation_spec.rb b/spec/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation_spec.rb
index 47d55570eb2..4927478b802 100644
--- a/spec/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation_spec.rb
+++ b/spec/datadog/tracing/contrib/action_pack/action_dispatch/instrumentation_spec.rb
@@ -56,13 +56,5 @@
 
       expect(traces).to be_empty
     end
-
-    it 'rescues exceptions' do
-      expect(Datadog::Tracing).to receive(:active_trace).and_raise('boom')
-
-      expect(Datadog.logger).to receive(:error).with('boom')
-
-      described_class.set_http_route_tags('/users/:id', '/auth')
-    end
   end
 end
diff --git a/spec/datadog/tracing/contrib/action_pack/action_dispatch/journey/router_spec.rb b/spec/datadog/tracing/contrib/action_pack/action_dispatch/journey/router_spec.rb
index 6fab8071ba0..1a7610ca71a 100644
--- a/spec/datadog/tracing/contrib/action_pack/action_dispatch/journey/router_spec.rb
+++ b/spec/datadog/tracing/contrib/action_pack/action_dispatch/journey/router_spec.rb
@@ -23,7 +23,7 @@
   describe '#find_routes' do
     before do
       engine.routes.append do
-        get '/sign-in/(:expires_in)' => 'tokens#create'
+        get '/sign-in' => 'tokens#create'
       end
 
       auth_engine = engine
@@ -41,6 +41,7 @@
         get '/items/:id', to: 'items#by_id', id: /\d+/
         get '/items/:slug', to: 'items#by_slug', id: /(\w-)+/
 
+        get 'books(/:category)', to: 'books#index'
         get 'books/*section/:title', to: 'books#show'
       end
     end
@@ -88,6 +89,10 @@ def by_slug
       stub_const(
         'BooksController',
         Class.new(ActionController::Base) do
+          def index
+            head :ok
+          end
+
           def show
             head :ok
           end
@@ -180,6 +185,17 @@ def create
         expect(request_span.tags).not_to have_key('http.route.path')
       end
 
+      it 'sets http.route correctly for routes with optional parameter' do
+        get 'books/some-category'
+
+        request_span = spans.first
+
+        expect(last_response).to be_ok
+        expect(request_span.name).to eq('rack.request')
+        expect(request_span.tags.fetch('http.route')).to eq('/books(/:category)')
+        expect(request_span.tags).not_to have_key('http.route.path')
+      end
+
       it 'sets http.route and http.route.path for rails engine routes' do
         get '/api/auth/sign-in'
 
@@ -187,7 +203,7 @@ def create
 
         expect(last_response).to be_ok
         expect(request_span.name).to eq('rack.request')
-        expect(request_span.tags.fetch('http.route')).to eq('/api/auth/sign-in(/:expires_in)')
+        expect(request_span.tags.fetch('http.route')).to eq('/api/auth/sign-in')
         expect(request_span.tags).not_to have_key('http.route.path')
       end