Add limit: option

Author: rmosolgo

lib/graphql/tracing/perfetto_sampler.rb

@@ -5,15 +5,16 @@
 module GraphQL
   module Tracing
     class PerfettoSampler
-      def self.use(schema, trace_mode: :perfetto_sample, memory: false, redis: nil, active_record: true)
+      # @param trace_mode [Symbol] When this is set as `context[:trace_mode]`, a profile will be taken
+      # @param redis [Redis] If provided, profiles will be stored in Redis for later review
+      # @param limit [Integer] A maximum number of profiles to store
+      def self.use(schema, trace_mode: :perfetto_sample, memory: false, redis: nil, limit: nil)
         storage = if redis
-          RedisBackend.new(redis: redis)
+          RedisBackend.new(redis: redis, limit: limit)
         elsif memory
-          MemoryBackend.new
-        elsif active_record != false
-          ActiveRecordBackend.new
+          MemoryBackend.new(limit: limit)
         else
-          raise ArgumentError, "A storage option must be chosen"
+          raise ArgumentError, "Pass `redis: ...` to store traces in Redis for later review"
         end
         schema.perfetto_sampler = self.new(storage: storage)
         schema.trace_with(PerfettoTrace, mode: trace_mode, save_trace_mode: trace_mode)

lib/graphql/tracing/perfetto_sampler/memory_backend.rb

@@ -6,8 +6,10 @@ 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
+        def initialize(limit: nil)
+          @limit = limit
           @traces = {}
+          @next_id = 0
         end
 
         def traces(last:, before:)
@@ -27,7 +29,7 @@ def find_trace(id)
         end
 
         def delete_trace(id)
-          @traces.delete(id.to_i)
+          @traces.delete(id)
           nil
         end
 
@@ -37,14 +39,19 @@ def delete_all_traces
         end
 
         def save_trace(operation_name, duration, timestamp, trace_data)
-          id = @traces.size
+          id = @next_id
+          @next_id += 1
           @traces[id] = PerfettoSampler::StoredTrace.new(
             id: id,
             operation_name: operation_name,
             duration_ms: duration,
             timestamp: timestamp,
             trace_data: trace_data
           )
+          if @limit && @traces.size > @limit
+            del_keys = @traces.keys[0...-@limit]
+            del_keys.each { |k| @traces.delete(k) }
+          end
           id
         end
       end

lib/graphql/tracing/perfetto_sampler/redis_backend.rb

@@ -5,9 +5,10 @@ module Tracing
     class PerfettoSampler
       class RedisBackend
         KEY_PREFIX = "gql:trace:"
-        def initialize(redis:)
+        def initialize(redis:, limit: nil)
           @redis = redis
           @key = KEY_PREFIX + "traces"
+          @remrangebyrank_limit = limit ? -limit - 1 : nil
         end
 
         def traces(last:, before:)
@@ -44,7 +45,12 @@ def find_trace(id)
         def save_trace(operation_name, duration_ms, begin_ms, trace_data)
           id = begin_ms
           data = JSON.dump({ "o" => operation_name, "d" => duration_ms, "b" => begin_ms, "t" => Base64.encode64(trace_data) })
-          @redis.zadd(@key, id, data)
+          @redis.pipelined do |pipeline|
+            pipeline.zadd(@key, id, data)
+            if @remrangebyrank_limit
+              pipeline.zremrangebyrank(@key, 0, @remrangebyrank_limit)
+            end
+          end
           id
         end
 

spec/graphql/tracing/perfetto_sampler/backend_assertions.rb

@@ -5,6 +5,7 @@ def self.included(child_class)
     child_class.instance_eval do
       describe "BackendAssertions" do
         before do
+          @backend = new_backend
           @backend.delete_all_traces
         end
 
@@ -57,6 +58,38 @@ def self.included(child_class)
         it "returns nil for nonexistent IDs" do
           assert_nil @backend.find_trace(999_999_999)
         end
+
+        it "implements a limit" do
+          limit_backend = new_backend(limit: 10)
+
+          10.times do |n|
+            limit_backend.save_trace(
+              "Trace#{n}",
+              1.5,
+              5000 + n,
+              "some-data"
+            )
+          end
+
+          all_traces = limit_backend.traces(last: nil, before: nil)
+          assert_equal 10, all_traces.size
+          assert_equal "Trace9", all_traces.first.operation_name
+          assert_equal "Trace0", all_traces.last.operation_name
+
+          3.times do |n|
+            limit_backend.save_trace(
+              "Trace 2-#{n}",
+              1.5,
+              5020 + n,
+              "some-data"
+            )
+          end
+
+          all_traces = limit_backend.traces(last: nil, before: nil)
+          assert_equal 10, all_traces.size
+          assert_equal "Trace 2-2", all_traces.first.operation_name
+          assert_equal "Trace3", all_traces.last.operation_name
+        end
       end
     end
   end

spec/graphql/tracing/perfetto_sampler/memory_backend_spec.rb

@@ -4,8 +4,7 @@
 
 describe GraphQL::Tracing::PerfettoSampler::MemoryBackend do
   include GraphQLTracingPerfettoSamplerBackendAssertions
-
-  before do
-    @backend = GraphQL::Tracing::PerfettoSampler::MemoryBackend.new
+  def new_backend(**kwargs)
+    GraphQL::Tracing::PerfettoSampler::MemoryBackend.new(**kwargs)
   end
 end

spec/graphql/tracing/perfetto_sampler/redis_backend_spec.rb

@@ -5,9 +5,8 @@
 if testing_redis?
   describe GraphQL::Tracing::PerfettoSampler::RedisBackend do
     include GraphQLTracingPerfettoSamplerBackendAssertions
-
-    before do
-      @backend = GraphQL::Tracing::PerfettoSampler::RedisBackend.new(redis: Redis.new)
+    def new_backend(**kwargs)
+      GraphQL::Tracing::PerfettoSampler::RedisBackend.new(redis: Redis.new, **kwargs)
     end
   end
 end