Log DB query count per request

philippthun · philippthun · commit ce8cd7f1af0e · 2025-07-10T11:22:58.000+02:00
diff --git a/lib/cloud_controller/db.rb b/lib/cloud_controller/db.rb
@@ -26,6 +26,7 @@ def self.connect(opts, logger)
       connection_options = VCAP::CloudController::DbConnection::OptionsFactory.build(opts)
 
       db = get_connection(opts, connection_options)
+      db.extension(:request_query_counter)
 
       if opts[:log_db_queries]
         db.logger = logger
@@ -38,7 +39,6 @@ def self.connect(opts, logger)
 
         db.extension(:query_length_logging)
         db.opts[:query_size_log_threshold] = opts[:query_size_log_threshold]
-        db.extension(:request_query_counter)
       end
       db.default_collate = 'utf8_bin' if db.database_type == :mysql
       add_connection_expiration_extension(db, opts)
diff --git a/lib/cloud_controller/logs/request_logs.rb b/lib/cloud_controller/logs/request_logs.rb
@@ -26,11 +26,10 @@ def start_request(request_id, env)
         @incomplete_requests.store(request_id, env)
       end
 
-      def complete_request(request_id, status, env, time_taken)
+      def complete_request(request_id, status, env, time_taken, query_count)
         return if @incomplete_requests.delete(request_id).nil?
 
         request = ActionDispatch::Request.new(env)
-        query_count = Thread.current[:db_query_count]
         @logger.info("Completed #{status} vcap-request-id: #{request_id}",
                      { status_code: status,
                        time_taken_in_ms: time_taken,
diff --git a/lib/cloud_controller/metrics/prometheus_updater.rb b/lib/cloud_controller/metrics/prometheus_updater.rb
@@ -18,7 +18,6 @@ def self.allow_pid_label
 
     DURATION_BUCKETS = [5, 10, 30, 60, 300, 600, 890].freeze
     CONNECTION_DURATION_BUCKETS = [0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 10].freeze
-    QUERY_COUNT_BUCKETS = [1, 5, 10, 25, 50, 100, 200].freeze
     DELAYED_JOB_METRIC_BUCKETS = [0.01, 0.05, 0.1, 0.5, 1, 2, 5, 10, 30, 60, 120, 300, 600].freeze
 
     METRICS = [
@@ -30,7 +29,6 @@ def self.allow_pid_label
       { type: :histogram, name: :cc_staging_failed_duration_seconds, docstring: 'Durations of failed staging events', buckets: DURATION_BUCKETS },
       { type: :gauge, name: :cc_requests_outstanding_total, docstring: 'Requests outstanding', aggregation: :sum },
       { type: :counter, name: :cc_requests_completed_total, docstring: 'Requests completed' },
-      { type: :histogram, name: :cc_db_queries_per_request, docstring: 'DB queries per request', buckets: QUERY_COUNT_BUCKETS },
       { type: :gauge, name: :cc_vitals_started_at, docstring: 'CloudController Vitals: started_at', aggregation: :most_recent },
       { type: :gauge, name: :cc_vitals_mem_bytes, docstring: 'CloudController Vitals: mem_bytes', aggregation: :most_recent },
       { type: :gauge, name: :cc_vitals_cpu_load_avg, docstring: 'CloudController Vitals: cpu_load_avg', aggregation: :most_recent },
@@ -191,10 +189,6 @@ def report_staging_failure_metrics(duration_ns)
       update_histogram_metric(:cc_staging_failed_duration_seconds, nanoseconds_to_seconds(duration_ns))
     end
 
-    def record_db_query_count(count)
-      update_histogram_metric(:cc_db_queries_per_request, count)
-    end
-
     private
 
     def register(metric)
diff --git a/lib/cloud_controller/metrics/request_metrics.rb b/lib/cloud_controller/metrics/request_metrics.rb
@@ -15,11 +15,9 @@ def start_request
 
       def complete_request(status, query_count)
         @statsd_updater.complete_request(status)
-        @statsd_updater.record_db_query_count(query_count)
 
         @prometheus_updater.increment_counter_metric(:cc_requests_completed_total)
         @prometheus_updater.decrement_gauge_metric(:cc_requests_outstanding_total)
-        @prometheus_updater.record_db_query_count(query_count)
       end
     end
   end
diff --git a/lib/cloud_controller/metrics/statsd_updater.rb b/lib/cloud_controller/metrics/statsd_updater.rb
@@ -110,10 +110,6 @@ def complete_request(status)
       end
     end
 
-    def record_db_query_count(count)
-      @statsd.gauge('cc.requests.db_query_count', count)
-    end
-
     private
 
     def nanoseconds_to_milliseconds(time_ns)
diff --git a/lib/sequel/extensions/request_query_counter.rb b/lib/sequel/extensions/request_query_counter.rb
@@ -1,8 +1,7 @@
 module Sequel
   module RequestQueryCounter
     def log_connection_yield(sql, conn, args=nil)
-      Thread.current[:db_query_count] ||= 0
-      Thread.current[:db_query_count] += 1
+      VCAP::Request.db_query_count = (VCAP::Request.db_query_count || 0) + 1
       super
     end
   end
diff --git a/lib/vcap/request.rb b/lib/vcap/request.rb
@@ -55,6 +55,14 @@ def b3_span_id=(span_id)
       def b3_span_id
         Thread.current[:b3_span_id]
       end
+
+      def db_query_count=(count)
+        Thread.current[:db_query_count] = count
+      end
+
+      def db_query_count
+        Thread.current[:db_query_count]
+      end
     end
   end
 end
diff --git a/middleware/request_logs.rb b/middleware/request_logs.rb
@@ -15,7 +15,8 @@ def call(env)
         # convert to milliseconds
         time_taken = (Time.now - start_time) * 1000
         time_taken = time_taken.to_i
-        @request_logs.complete_request(request_id, status, env, time_taken)
+        query_count = VCAP::Request.db_query_count
+        @request_logs.complete_request(request_id, status, env, time_taken, query_count)
 
         [status, headers, body]
       end
diff --git a/middleware/request_metrics.rb b/middleware/request_metrics.rb
@@ -7,7 +7,7 @@ def initialize(app, request_metrics)
       end
 
       def call(env)
-        Thread.current[:db_query_count] = 0
+        VCAP::Request.db_query_count = 0
         @request_metrics.start_request
 
         status, headers, body = @app.call(env)
@@ -18,8 +18,8 @@ def call(env)
         status = 500
         raise e
       ensure
-        query_count = Thread.current[:db_query_count]
-        Thread.current[:db_query_count] = nil
+        query_count = VCAP::Request.db_query_count
+        VCAP::Request.db_query_count = nil
         @request_metrics.complete_request(status, query_count)
       end
     end
diff --git a/spec/unit/lib/cloud_controller/logs/request_logs_spec.rb b/spec/unit/lib/cloud_controller/logs/request_logs_spec.rb
@@ -46,7 +46,7 @@ module VCAP::CloudController::Logs
           end
 
           it 'logs the completion of the request' do
-            request_logs.complete_request(request_id, status, env, time_taken)
+            request_logs.complete_request(request_id, status, env, time_taken, nil)
             expect(logger).to have_received(:info).with(
               /\ACompleted 200 vcap-request-id: ID/,
               time_taken_in_ms: 30,
@@ -60,7 +60,7 @@ module VCAP::CloudController::Logs
 
         context 'without a matching start request' do
           it 'does not log the completion of the request' do
-            request_logs.complete_request(request_id, status, env, time_taken)
+            request_logs.complete_request(request_id, status, env, time_taken, nil)
             expect(logger).not_to have_received(:info)
           end
         end
diff --git a/spec/unit/lib/cloud_controller/metrics/prometheus_updater_spec.rb b/spec/unit/lib/cloud_controller/metrics/prometheus_updater_spec.rb
@@ -279,13 +279,5 @@ module VCAP::CloudController::Metrics
       end
     end
 
-    describe '#record_db_query_count' do
-      it 'records db query count histogram' do
-        updater.record_db_query_count(3)
-
-        metric = prom_client.get :cc_db_queries_per_request
-        expect(metric.get).to eq({ '1' => 0.0, '5' => 1.0, '10' => 1.0, '25' => 1.0, '50' => 1.0, '100' => 1.0, '200' => 1.0, 'sum' => 3.0, '+Inf' => 1.0 })
-      end
-    end
   end
 end
diff --git a/spec/unit/lib/cloud_controller/metrics/request_metrics_spec.rb b/spec/unit/lib/cloud_controller/metrics/request_metrics_spec.rb
@@ -12,10 +12,8 @@ module VCAP::CloudController::Metrics
       allow(prometheus_client).to receive(:decrement_gauge_metric)
       allow(prometheus_client).to receive(:increment_gauge_metric)
       allow(prometheus_client).to receive(:increment_counter_metric)
-      allow(prometheus_client).to receive(:record_db_query_count)
       allow(statsd_updater).to receive(:start_request)
       allow(statsd_updater).to receive(:complete_request)
-      allow(statsd_updater).to receive(:record_db_query_count)
     end
 
     describe '#start_request' do
@@ -36,15 +34,13 @@ module VCAP::CloudController::Metrics
       it 'increments completed, decrements outstanding, increments status for statsd' do
         request_metrics.complete_request(status, 5)
         expect(statsd_updater).to have_received(:complete_request).with(status)
-        expect(statsd_updater).to have_received(:record_db_query_count).with(5)
       end
 
       it 'increments completed and decrements outstanding for prometheus' do
         request_metrics.complete_request(status, 5)
 
         expect(prometheus_client).to have_received(:decrement_gauge_metric).with(:cc_requests_outstanding_total)
         expect(prometheus_client).to have_received(:increment_counter_metric).with(:cc_requests_completed_total)
-        expect(prometheus_client).to have_received(:record_db_query_count).with(5)
       end
     end
   end
diff --git a/spec/unit/lib/cloud_controller/metrics/statsd_updater_spec.rb b/spec/unit/lib/cloud_controller/metrics/statsd_updater_spec.rb
@@ -335,17 +335,6 @@ module VCAP::CloudController::Metrics
       end
     end
 
-    describe '#record_db_query_count' do
-      before do
-        allow(statsd_client).to receive(:gauge)
-      end
-
-      it 'emits the query count as a gauge' do
-        updater.record_db_query_count(7)
-        expect(statsd_client).to have_received(:gauge).with('cc.requests.db_query_count', 7)
-      end
-    end
-
     describe '#store' do
       let(:config) { double(:config) }
 
diff --git a/spec/unit/lib/sequel/extensions/request_query_counter_spec.rb b/spec/unit/lib/sequel/extensions/request_query_counter_spec.rb
@@ -8,8 +8,8 @@
   end
 
   it 'increments the thread local query count on each query' do
-    Thread.current[:db_query_count] = 0
+    VCAP::Request.db_query_count = 0
     db.fetch('SELECT 1').all
-    expect(Thread.current[:db_query_count]).to eq(1)
+    expect(VCAP::Request.db_query_count).to eq(1)
   end
 end
diff --git a/spec/unit/middleware/request_logs_spec.rb b/spec/unit/middleware/request_logs_spec.rb
@@ -21,7 +21,7 @@ module Middleware
 
         it 'calls complete request on request logs after the request' do
           middleware.call(env)
-          expect(request_logs).to have_received(:complete_request).with('ID', 200, { 'cf.request_id' => 'ID' }, be_a(Numeric))
+          expect(request_logs).to have_received(:complete_request).with('ID', 200, { 'cf.request_id' => 'ID' }, be_a(Numeric), an_instance_of(Integer).or(nil))
         end
       end
     end
