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,59 @@
+# frozen_string_literal: true
+require "graphql/tracing/perfetto_sampler/memory_backend"
+require "graphql/tracing/perfetto_sampler/redis_backend"
+
+module GraphQL
+  module Tracing
+    class PerfettoSampler
+      def self.use(schema, trace_mode: :perfetto_sample, memory: false, redis: nil, active_record: true)
+        storage = if redis
+          RedisBackend.new(redis: redis)
+        elsif memory
+          MemoryBackend.new
+        elsif active_record != false
+          ActiveRecordBackend.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_sampler/redis_backend.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module Tracing
+    class PerfettoSampler
+      class RedisBackend
+        KEY_PREFIX = "gql:trace:"
+        def initialize(redis:)
+          @redis = redis
+        end
+
+        def traces
+          keys = @redis.scan_each(match: "#{KEY_PREFIX}*").to_a
+          keys.sort!
+          keys.map do |k|
+            h = @redis.hgetall(k)
+            StoredTrace.new(
+              id: k.sub(KEY_PREFIX, ""),
+              operation_name: h["operation_name"],
+              duration_ms: h["duration_ms"].to_f,
+              timestamp: Time.at(h["timestamp"].to_i),
+              trace_data: h["trace_data"],
+            )
+          end
+        end
+
+        def delete_trace(id)
+          @redis.del("#{KEY_PREFIX}#{id}")
+          nil
+        end
+
+        def delete_all_traces
+          keys = @redis.scan_each(match: "#{KEY_PREFIX}*")
+          @redis.del(*keys)
+        end
+
+        def find_trace(id)
+          redis_h = @redis.hgetall("#{KEY_PREFIX}#{id}")
+          if redis_h.empty?
+            nil
+          else
+            StoredTrace.new(
+              id: id,
+              operation_name: redis_h["operation_name"],
+              duration_ms: redis_h["duration_ms"].to_f,
+              timestamp: Time.at(redis_h["timestamp"].to_i),
+              trace_data: redis_h["trace_data"],
+            )
+          end
+        end
+
+        def save_trace(operation_name, duration_ms, timestamp, trace_data)
+          id = (timestamp.to_i * 1000) + rand(1000)
+          @redis.hmset("#{KEY_PREFIX}#{id}",
+            "operation_name", operation_name,
+            "duration_ms", duration_ms,
+            "timestamp", timestamp.to_i,
+            "trace_data", trace_data,
+          )
+          id
+        end
+      end
+    end
+  end
+end

lib/graphql/tracing/perfetto_trace.rb

@@ -61,9 +61,12 @@ 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
+        puts "Starting perfetto trace on Fiber #{Fiber.current.object_id}"
+        Fiber[:graphql_flow_stack] = nil
         @pid = Process.pid
         @flow_ids = Hash.new { |h, source_inst| h[source_inst] = [] }.compare_by_identity
         @new_interned_event_names = {}
@@ -108,6 +111,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 +176,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 +195,11 @@ 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
+          data = Trace.encode(Trace.new(packet: @packets))
+          m.schema.perfetto_sampler.save_trace(@operation_name, duration_ms, @begin_time, data)
+        end
         super
       end
 
@@ -611,7 +622,10 @@ def dup_with(message, attrs, delete_counters: false)
       end
 
       def fiber_flow_stack
-        Fiber[:graphql_flow_stack] ||= []
+        Fiber[:graphql_flow_stack] ||= begin
+          puts "Creating new stack for Fiber ##{Fiber.current.object_id}"
+          []
+        end
       end
 
       def trace_packet(event_attrs)

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/dataloader/async_dataloader_spec.rb

@@ -343,15 +343,7 @@ class TraceAsyncSchema < AsyncSchema
         }
         GRAPHQL
         res = @schema.execute(query_str)
-        if ENV["DUMP_PERFETTO"]
-          res.context.query.current_trace.write(file: "perfetto.dump")
-        end
-
-        json = res.context.query.current_trace.write(file: nil, debug_json: true)
-        data = JSON.parse(json)
-
-
-        check_snapshot(data, "example.json")
+        check_snapshot(res, "example.json")
       end
     end
   end

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/redis_backend_spec.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+require "spec_helper"
+require_relative "./backend_assertions"
+
+if testing_redis?
+  describe GraphQL::Tracing::PerfettoSampler::RedisBackend do
+    include GraphQLTracingPerfettoSamplerBackendAssertions
+
+    before do
+      @backend = GraphQL::Tracing::PerfettoSampler::RedisBackend.new(redis: Redis.new)
+    end
+  end
+end

spec/graphql/tracing/perfetto_sampler_spec.rb

@@ -0,0 +1,48 @@
+# 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
+    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
+    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