Add Redis backend for Perfetto traces

lib/graphql/schema.rb

@@ -1366,6 +1366,9 @@ def default_directives
         }.freeze
       end
 
+      # @return [nil, GraphQL::Tracing::PerfettoSampler]
+      attr_accessor :perfetto_sampler
+
       def tracer(new_tracer, silence_deprecation_warning: false)
         if !silence_deprecation_warning
           warn("`Schema.tracer(#{new_tracer.inspect})` is deprecated; use module-based `trace_with` instead. See: https://graphql-ruby.org/queries/tracing.html")

lib/graphql/tracing.rb

@@ -32,6 +32,7 @@ module Tracing
     autoload :StatsdTrace, "graphql/tracing/statsd_trace"
     autoload :PrometheusTrace, "graphql/tracing/prometheus_trace"
     autoload :PerfettoTrace, "graphql/tracing/perfetto_trace"
+    autoload :PerfettoSampler, "graphql/tracing/perfetto_sampler"
 
     # Objects may include traceable to gain a `.trace(...)` method.
     # The object must have a `@tracers` ivar of type `Array<<#trace(k, d, &b)>>`.

lib/graphql/tracing/perfetto_sampler.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+require "graphql/tracing/perfetto_sampler/memory_backend"
+
+module GraphQL
+  module Tracing
+    class PerfettoSampler
+      def self.use(schema, trace_mode: :perfetto_sample, memory: false, active_record: true)
+        storage = if memory
+          MemoryBackend.new
+        else
+          raise ArgumentError, "A storage option must be chosen"
+        end
+        schema.perfetto_sampler = self.new(storage: storage)
+        schema.trace_with(PerfettoTrace, mode: trace_mode, save_trace_mode: trace_mode)
+      end
+
+      def initialize(storage:)
+        @storage = storage
+      end
+
+      def save_trace(operation_name, duration_ms, timestamp, trace_data)
+        @storage.save_trace(operation_name, duration_ms, timestamp, trace_data)
+      end
+
+      def traces
+        @storage.traces
+      end
+
+      def find_trace(id)
+        @storage.find_trace(id)
+      end
+
+      def delete_trace(id)
+        @storage.delete_trace(id)
+      end
+
+      def delete_all_traces
+        @storage.delete_all_traces
+      end
+
+      class StoredTrace
+        def initialize(id:, operation_name:, duration_ms:, timestamp:, trace_data:)
+          @id = id
+          @operation_name = operation_name
+          @duration_ms = duration_ms
+          @timestamp = timestamp
+          @trace_data = trace_data
+        end
+
+        attr_reader :id, :operation_name, :duration_ms, :timestamp, :trace_data
+      end
+    end
+  end
+end

lib/graphql/tracing/perfetto_sampler/memory_backend.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module Tracing
+    class PerfettoSampler
+      # An in-memory trace storage backend. Suitable for testing and development only.
+      # It won't work for multi-process deployments and everything is erased when the app is restarted.
+      class MemoryBackend
+        def initialize
+          @traces = {}
+        end
+
+        def traces
+          @traces.values
+        end
+
+        def find_trace(id)
+          @traces[id]
+        end
+
+        def delete_trace(id)
+          @traces.delete(id)
+          nil
+        end
+
+        def delete_all_traces
+          @traces.clear
+          nil
+        end
+
+        def save_trace(operation_name, duration, timestamp, trace_data)
+          id = @traces.size
+          @traces[id] = PerfettoSampler::StoredTrace.new(
+            id: id,
+            operation_name: operation_name,
+            duration_ms: duration,
+            timestamp: timestamp,
+            trace_data: trace_data
+          )
+          id
+        end
+      end
+    end
+  end
+end

lib/graphql/tracing/perfetto_trace.rb

@@ -61,8 +61,9 @@ def self.included(_trace_class)
       DA_STR_VAL_NIL_IID = 14
 
       # @param active_support_notifications_pattern [String, RegExp, false] A filter for `ActiveSupport::Notifications`, if it's present. Or `false` to skip subscribing.
-      def initialize(active_support_notifications_pattern: nil, **_rest)
+      def initialize(active_support_notifications_pattern: nil, save_trace_mode: nil, **_rest)
         super
+        @save_trace_mode = save_trace_mode
         @sequence_id = object_id
         @pid = Process.pid
         @flow_ids = Hash.new { |h, source_inst| h[source_inst] = [] }.compare_by_identity
@@ -108,6 +109,7 @@ def initialize(active_support_notifications_pattern: nil, **_rest)
         @fibers_counter_id = :fibers_counter.object_id
         @fields_counter_id = :fields_counter.object_id
         @begin_validate = nil
+        @begin_time = nil
         @packets = []
         @packets << TracePacket.new(
           track_descriptor: TrackDescriptor.new(
@@ -172,6 +174,8 @@ def initialize(active_support_notifications_pattern: nil, **_rest)
       end
 
       def begin_execute_multiplex(m)
+        @operation_name = m.queries.map { |q| q.selected_operation_name || "anonymous" }.join(",")
+        @begin_time = Time.now
         @packets << trace_packet(
           type: TrackEvent::Type::TYPE_SLICE_BEGIN,
           track_uuid: fid,
@@ -189,6 +193,10 @@ def end_execute_multiplex(m)
           track_uuid: fid,
         )
         unsubscribe_from_active_support_notifications
+        # if @save_trace_mode && m.context[:trace_mode] == @save_trace_mode
+        #   duration_ms = (Time.now.to_f - @begin_time.to_f) * 1000
+        #   m.schema.perfetto_sampler.save_trace(@operation_name, duration_ms, @begin_time, Trace.encode(Trace.new(packet: @packets)))
+        # end
         super
       end
 

spec/graphql/autoload_spec.rb

@@ -50,7 +50,7 @@ def self.eager_load!
   describe "loading nested files in the repo" do
     it "can load them individually" do
       files_to_load = Dir.glob("lib/**/tracing/*.rb")
-      assert_equal 27, files_to_load.size, "It found all the expected files"
+      assert_equal 28, files_to_load.size, "It found all the expected files"
       files_to_load.each do |file|
         require_path = file.sub("lib/", "").sub(".rb", "")
         stderr_and_stdout, _status = Open3.capture2e("ruby -Ilib -e 'require \"#{require_path}\"'")

spec/graphql/tracing/perfetto_sampler/backend_assertions.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module GraphQLTracingPerfettoSamplerBackendAssertions
+  def self.included(child_class)
+    child_class.instance_eval do
+      describe "BackendAssertions" do
+        before do
+          @backend.delete_all_traces
+        end
+
+        it "can save, retreive, list, and delete traces" do
+          data = SecureRandom.bytes(1000)
+          trace_id = @backend.save_trace(
+            "GetStuff",
+            100.56,
+            Time.utc(2024, 01, 01, 04, 44, 33, 695),
+            data
+          )
+
+          trace = @backend.find_trace(trace_id)
+          assert_kind_of GraphQL::Tracing::PerfettoSampler::StoredTrace, trace
+          assert_equal "GetStuff", trace.operation_name
+          assert_equal 100.56, trace.duration_ms
+          assert_equal "2024-01-01 04:44:33.000", trace.timestamp.utc.strftime("%Y-%m-%d %H:%M:%S.%L")
+          assert_equal data, trace.trace_data
+
+
+          @backend.save_trace(
+            "GetOtherStuff",
+            200.16,
+            Time.utc(2024, 01, 03, 04, 44, 33, 695),
+            data
+          )
+
+          assert_equal ["GetStuff", "GetOtherStuff"], @backend.traces.map(&:operation_name)
+
+          @backend.delete_trace(trace_id)
+
+          assert_equal ["GetOtherStuff"], @backend.traces.map(&:operation_name)
+
+          @backend.delete_all_traces
+          assert_equal [], @backend.traces
+        end
+
+        it "returns nil for nonexistent IDs" do
+          assert_nil @backend.find_trace(999_999_999)
+        end
+      end
+    end
+  end
+end

spec/graphql/tracing/perfetto_sampler/memory_backend_spec.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+require "spec_helper"
+require_relative "./backend_assertions"
+
+describe GraphQL::Tracing::PerfettoSampler::MemoryBackend do
+  include GraphQLTracingPerfettoSamplerBackendAssertions
+
+  before do
+    @backend = GraphQL::Tracing::PerfettoSampler::MemoryBackend.new
+  end
+end

spec/graphql/tracing/perfetto_sampler_spec.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+require "spec_helper"
+
+describe GraphQL::Tracing::PerfettoSampler do
+  class SamplerSchema < GraphQL::Schema
+    class Query < GraphQL::Schema::Object
+      field :truthy, Boolean, fallback_value: true
+    end
+
+    query(Query)
+    use GraphQL::Tracing::PerfettoSampler, memory: true
+  end
+
+  before do
+    SamplerSchema.perfetto_sampler.delete_all_traces
+  end
+
+  it "runs when the configured trace mode is set" do
+    skip("Currently disabled")
+    assert_equal 0, SamplerSchema.perfetto_sampler.traces.size
+    res = SamplerSchema.execute("{ truthy }")
+    assert_equal true, res["data"]["truthy"]
+    assert_equal 0, SamplerSchema.perfetto_sampler.traces.size
+
+    SamplerSchema.execute("{ truthy }", context: { trace_mode: :perfetto_sample })
+    assert_equal 1, SamplerSchema.perfetto_sampler.traces.size
+  end
+
+  it "calls through to storage for access methods" do
+    skip("Currently disabled")
+    SamplerSchema.execute("{ truthy }", context: { trace_mode: :perfetto_sample })
+    id = SamplerSchema.perfetto_sampler.traces.first.id
+    assert_kind_of GraphQL::Tracing::PerfettoSampler::StoredTrace, SamplerSchema.perfetto_sampler.find_trace(id)
+    SamplerSchema.perfetto_sampler.delete_trace(id)
+    assert_equal 0, SamplerSchema.perfetto_sampler.traces.size
+
+    SamplerSchema.execute("{ truthy }", context: { trace_mode: :perfetto_sample })
+    assert_equal 1, SamplerSchema.perfetto_sampler.traces.size
+    SamplerSchema.perfetto_sampler.delete_all_traces
+  end
+
+  it "raises when no storage is configured" do
+    err = assert_raises ArgumentError do
+      Class.new(GraphQL::Schema) do
+        use GraphQL::Tracing::PerfettoSampler, active_record: false
+      end
+    end
+    assert_equal "A storage option must be chosen", err.message
+  end
+end