Update Redis backend

Author: rmosolgo

lib/graphql/tracing/perfetto_sampler/redis_backend.rb

@@ -7,58 +7,59 @@ class RedisBackend
         KEY_PREFIX = "gql:trace:"
         def initialize(redis:)
           @redis = redis
+          @key = KEY_PREFIX + "traces"
         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"],
-            )
+        def traces(last:, before:)
+          before = case before
+          when Numeric
+            "(#{before}"
+          when nil
+            "+inf"
+          end
+          str_pairs = @redis.zrange(@key, before, 0, byscore: true, rev: true, limit: [0, last || 100], withscores: true)
+          str_pairs.map do |(str_data, score)|
+            entry_to_trace(score, str_data)
           end
         end
 
         def delete_trace(id)
-          @redis.del("#{KEY_PREFIX}#{id}")
+          @redis.zremrangebyscore(@key, id, id)
           nil
         end
 
         def delete_all_traces
-          keys = @redis.scan_each(match: "#{KEY_PREFIX}*")
-          @redis.del(*keys)
+          @redis.del(@key)
         end
 
         def find_trace(id)
-          redis_h = @redis.hgetall("#{KEY_PREFIX}#{id}")
-          if redis_h.empty?
+          str_data = @redis.zrange(@key, id, id, byscore: true).first
+          if str_data.nil?
             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"],
-            )
+            entry_to_trace(id, str_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,
-          )
+        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)
           id
         end
+
+        private
+
+        def entry_to_trace(id, json_str)
+          data = JSON.parse(json_str)
+          StoredTrace.new(
+            id: id,
+            operation_name: data["o"],
+            duration_ms: data["d"].to_f,
+            timestamp: data["b"].to_i,
+            trace_data: Base64.decode64(data["t"]),
+          )
+        end
       end
     end
   end

spec/graphql/tracing/perfetto_sampler/backend_assertions.rb

@@ -13,33 +13,45 @@ def self.included(child_class)
           trace_id = @backend.save_trace(
             "GetStuff",
             100.56,
-            Time.utc(2024, 01, 01, 04, 44, 33, 695),
+            (Time.utc(2024, 01, 01, 04, 44, 33, 695000).to_f * 1000).round,
             data
           )
 
           trace = @backend.find_trace(trace_id)
           assert_kind_of GraphQL::Tracing::PerfettoSampler::StoredTrace, trace
+          assert_equal trace_id, trace.id
           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 "2024-01-01 04:44:33.694", Time.at(trace.timestamp / 1000.0).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),
+            (Time.utc(2024, 01, 03, 04, 44, 33, 695000).to_f * 1000).round,
             data
           )
 
-          assert_equal ["GetStuff", "GetOtherStuff"], @backend.traces.map(&:operation_name)
+          @backend.save_trace(
+            "GetMoreOtherStuff",
+            200.16,
+            (Time.utc(2024, 01, 03, 04, 44, 33, 795000).to_f * 1000).round,
+            data
+          )
+
+          assert_equal ["GetMoreOtherStuff", "GetOtherStuff", "GetStuff" ], @backend.traces(last: 20, before: nil).map(&:operation_name)
+
+          assert_equal ["GetMoreOtherStuff"], @backend.traces(last: 1, before: nil).map(&:operation_name)
+          assert_equal ["GetOtherStuff", "GetStuff"], @backend.traces(last: 2, before: Time.utc(2024, 01, 03, 04, 44, 33, 795000).to_f * 1000).map(&:operation_name)
+
 
           @backend.delete_trace(trace_id)
 
-          assert_equal ["GetOtherStuff"], @backend.traces.map(&:operation_name)
+          assert_equal ["GetMoreOtherStuff", "GetOtherStuff"], @backend.traces(last: 20, before: nil).map(&:operation_name)
 
           @backend.delete_all_traces
-          assert_equal [], @backend.traces
+          assert_equal [], @backend.traces(last: 20, before: nil)
         end
 
         it "returns nil for nonexistent IDs" do