Capture Gitaly calls that don't go through `migrate`

2 files changed, 23 insertions, 13 deletions

diff --git a/lib/gitlab/gitaly_client.rb b/lib/gitlab/gitaly_client.rb
index c3aa3ccd7b1..8ca30ffc232 100644
--- a/lib/gitlab/gitaly_client.rb
+++ b/lib/gitlab/gitaly_client.rb
@@ -119,7 +119,8 @@ module Gitlab
     #
     def self.call(storage, service, rpc, request, remote_storage: nil, timeout: nil)
       start = Gitlab::Metrics::System.monotonic_time
-      @last_request = request.is_a?(Google::Protobuf::MessageExts) ? request.to_h : nil
+      request_hash = request.is_a?(Google::Protobuf::MessageExts) ? request.to_h : {}
+      @current_call_id ||= SecureRandom.uuid
 
       enforce_gitaly_request_limits(:call)
 
@@ -137,6 +138,10 @@ module Gitlab
       gitaly_controller_action_duration_seconds.observe(
         current_transaction_labels.merge(gitaly_service: service.to_s, rpc: rpc.to_s),
         duration)
+
+      add_call_details(id: @current_call_id, feature: service, duration: duration, request: request_hash)
+
+      @current_call_id = nil
     end
 
     def self.handle_grpc_unavailable!(ex)
@@ -254,15 +259,16 @@ module Gitlab
           feature_stack.unshift(feature)
           begin
             start = Gitlab::Metrics::System.monotonic_time
+            @current_call_id = SecureRandom.uuid
+            call_details = { id: @current_call_id }
             yield is_enabled
           ensure
             total_time = Gitlab::Metrics::System.monotonic_time - start
             gitaly_migrate_call_duration_seconds.observe({ gitaly_enabled: is_enabled, feature: feature }, total_time)
             feature_stack.shift
             Thread.current[:gitaly_feature_stack] = nil if feature_stack.empty?
-            add_call_details(feature: feature,
-                             duration: total_time,
-                             request: is_enabled ? @last_request : {})
+
+            add_call_details(call_details.merge(feature: feature, duration: total_time))
           end
         end
       end
@@ -350,16 +356,19 @@ module Gitlab
     end
 
     def self.add_call_details(details)
-      return unless RequestStore.active? && RequestStore.store[:peek_enabled]
+      id = details.delete(:id)
+
+      return unless id && RequestStore.active? && RequestStore.store[:peek_enabled]
 
-      RequestStore.store['gitaly_call_details'] ||= []
-      RequestStore.store['gitaly_call_details'] << details
+      RequestStore.store['gitaly_call_details'] ||= {}
+      RequestStore.store['gitaly_call_details'][id] ||= {}
+      RequestStore.store['gitaly_call_details'][id].merge!(details)
     end
 
-    def self.call_details
-      return [] unless RequestStore.active? && RequestStore.store[:peek_enabled]
+    def self.list_call_details
+      return {} unless RequestStore.active? && RequestStore.store[:peek_enabled]
 
-      RequestStore.store['gitaly_call_details'] || []
+      RequestStore.store['gitaly_call_details'] || {}
     end
 
     def self.expected_server_version
diff --git a/lib/peek/views/gitaly.rb b/lib/peek/views/gitaly.rb
index 79851640318..ab35f7a2258 100644
--- a/lib/peek/views/gitaly.rb
+++ b/lib/peek/views/gitaly.rb
@@ -20,16 +20,17 @@ module Peek
       private
 
       def details
-        ::Gitlab::GitalyClient.call_details
+        ::Gitlab::GitalyClient.list_call_details
+          .values
           .sort { |a, b| b[:duration] <=> a[:duration] }
           .map(&method(:format_call_details))
       end
 
       def format_call_details(call)
-        pretty_request = call[:request].reject { |k, v| v.blank? }.to_h.pretty_inspect
+        pretty_request = call[:request]&.reject { |k, v| v.blank? }.to_h.pretty_inspect
 
         call.merge(duration: (call[:duration] * 1000).round(3),
-                   request: pretty_request)
+                   request: pretty_request || {})
       end
 
       def formatted_duration