From acc77c3836974473e7c6a423cbd1138479ae197a Mon Sep 17 00:00:00 2001
From: Renaud Chaput <renchap@gmail.com>
Date: Fri, 24 May 2024 15:13:23 +0200
Subject: [PATCH] Add instrumentation to the search services (#30350)

---
 Gemfile                                 |  2 ++
 Gemfile.lock                            |  1 +
 app/services/account_search_service.rb  | 22 ++++++++++++++++------
 app/services/statuses_search_service.rb | 24 +++++++++++++++++-------
 app/services/tag_search_service.rb      | 24 +++++++++++++++++-------
 config/initializers/opentelemetry.rb    |  2 ++
 6 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/Gemfile b/Gemfile
index 240dcce95..d9de33182 100644
--- a/Gemfile
+++ b/Gemfile
@@ -103,6 +103,8 @@ gem 'rdf-normalize', '~> 0.5'
 
 gem 'private_address_check', '~> 0.5'
 
+gem 'opentelemetry-api', '~> 1.2.5'
+
 group :opentelemetry do
   gem 'opentelemetry-exporter-otlp', '~> 0.26.3', require: false
   gem 'opentelemetry-instrumentation-active_job', '~> 0.7.1', require: false
diff --git a/Gemfile.lock b/Gemfile.lock
index 78023bd7f..ad60eb8be 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -975,6 +975,7 @@ DEPENDENCIES
   omniauth-rails_csrf_protection (~> 1.0)
   omniauth-saml (~> 2.0)
   omniauth_openid_connect (~> 0.6.1)
+  opentelemetry-api (~> 1.2.5)
   opentelemetry-exporter-otlp (~> 0.26.3)
   opentelemetry-instrumentation-active_job (~> 0.7.1)
   opentelemetry-instrumentation-active_model_serializers (~> 0.20.1)
diff --git a/app/services/account_search_service.rb b/app/services/account_search_service.rb
index 571a0fa57..b86c9b9e7 100644
--- a/app/services/account_search_service.rb
+++ b/app/services/account_search_service.rb
@@ -151,13 +151,23 @@ class AccountSearchService < BaseService
   end
 
   def call(query, account = nil, options = {})
-    @query   = query&.strip&.gsub(/\A@/, '')
-    @limit   = options[:limit].to_i
-    @offset  = options[:offset].to_i
-    @options = options
-    @account = account
+    MastodonOTELTracer.in_span('AccountSearchService#call') do |span|
+      @query   = query&.strip&.gsub(/\A@/, '')
+      @limit   = options[:limit].to_i
+      @offset  = options[:offset].to_i
+      @options = options
+      @account = account
 
-    search_service_results.compact.uniq
+      span.add_attributes(
+        'search.offset' => @offset,
+        'search.limit' => @limit,
+        'search.backend' => Chewy.enabled? ? 'elasticsearch' : 'database'
+      )
+
+      search_service_results.compact.uniq.tap do |results|
+        span.set_attribute('search.results.count', results.size)
+      end
+    end
   end
 
   private
diff --git a/app/services/statuses_search_service.rb b/app/services/statuses_search_service.rb
index 7d5b0203a..ab8e28f61 100644
--- a/app/services/statuses_search_service.rb
+++ b/app/services/statuses_search_service.rb
@@ -2,14 +2,24 @@
 
 class StatusesSearchService < BaseService
   def call(query, account = nil, options = {})
-    @query   = query&.strip
-    @account = account
-    @options = options
-    @limit   = options[:limit].to_i
-    @offset  = options[:offset].to_i
+    MastodonOTELTracer.in_span('StatusesSearchService#call') do |span|
+      @query   = query&.strip
+      @account = account
+      @options = options
+      @limit   = options[:limit].to_i
+      @offset  = options[:offset].to_i
+      convert_deprecated_options!
 
-    convert_deprecated_options!
-    status_search_results
+      span.add_attributes(
+        'search.offset' => @offset,
+        'search.limit' => @limit,
+        'search.backend' => Chewy.enabled? ? 'elasticsearch' : 'database'
+      )
+
+      status_search_results.tap do |results|
+        span.set_attribute('search.results.count', results.size)
+      end
+    end
   end
 
   private
diff --git a/app/services/tag_search_service.rb b/app/services/tag_search_service.rb
index 929cfd884..57400b76a 100644
--- a/app/services/tag_search_service.rb
+++ b/app/services/tag_search_service.rb
@@ -2,15 +2,25 @@
 
 class TagSearchService < BaseService
   def call(query, options = {})
-    @query   = query.strip.delete_prefix('#')
-    @offset  = options.delete(:offset).to_i
-    @limit   = options.delete(:limit).to_i
-    @options = options
+    MastodonOTELTracer.in_span('TagSearchService#call') do |span|
+      @query   = query.strip.delete_prefix('#')
+      @offset  = options.delete(:offset).to_i
+      @limit   = options.delete(:limit).to_i
+      @options = options
 
-    results   = from_elasticsearch if Chewy.enabled?
-    results ||= from_database
+      span.add_attributes(
+        'search.offset' => @offset,
+        'search.limit' => @limit,
+        'search.backend' => Chewy.enabled? ? 'elasticsearch' : 'database'
+      )
 
-    results
+      results   = from_elasticsearch if Chewy.enabled?
+      results ||= from_database
+
+      span.set_attribute('search.results.count', results.size)
+
+      results
+    end
   end
 
   private
diff --git a/config/initializers/opentelemetry.rb b/config/initializers/opentelemetry.rb
index cf9f0b96f..d121a95a3 100644
--- a/config/initializers/opentelemetry.rb
+++ b/config/initializers/opentelemetry.rb
@@ -66,3 +66,5 @@ if ENV.keys.any? { |name| name.match?(/OTEL_.*_ENDPOINT/) }
     c.service_version = Mastodon::Version.to_s
   end
 end
+
+MastodonOTELTracer = OpenTelemetry.tracer_provider.tracer('mastodon')