rack-mini-profiler 2.3.4 → 3.1.0

data/lib/mini_profiler/storage/redis_store.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'digest'
+require 'securerandom'
 
 module Rack
   class MiniProfiler
@@ -133,85 +134,127 @@ unviewed_ids: #{get_unviewed_ids(user)}
         )
       end
 
-      def push_snapshot(page_struct, config)
-        zset_key = snapshot_zset_key()
-        hash_key = snapshot_hash_key()
+      def push_snapshot(page_struct, group_name, config)
+        group_zset_key = group_snapshot_zset_key(group_name)
+        group_hash_key = group_snapshot_hash_key(group_name)
+        overview_zset_key = snapshot_overview_zset_key
 
         id = page_struct[:id]
-        score = page_struct.duration_ms
-        limit = config.snapshots_limit
+        score = page_struct.duration_ms.to_s
+
+        per_group_limit = config.max_snapshots_per_group.to_s
+        groups_limit = config.max_snapshot_groups.to_s
         bytes = Marshal.dump(page_struct)
 
         lua = <<~LUA
-          local zset_key = KEYS[1]
-          local hash_key = KEYS[2]
+          local group_zset_key = KEYS[1]
+          local group_hash_key = KEYS[2]
+          local overview_zset_key = KEYS[3]
+
           local id = ARGV[1]
           local score = tonumber(ARGV[2])
-          local bytes = ARGV[3]
-          local limit = tonumber(ARGV[4])
-          redis.call("ZADD", zset_key, score, id)
-          redis.call("HSET", hash_key, id, bytes)
-          if redis.call("ZCARD", zset_key) > limit then
-            local lowest_snapshot_id = redis.call("ZRANGE", zset_key, 0, 0)[1]
-            redis.call("ZREM", zset_key, lowest_snapshot_id)
-            redis.call("HDEL", hash_key, lowest_snapshot_id)
+          local group_name = ARGV[3]
+          local per_group_limit = tonumber(ARGV[4])
+          local groups_limit = tonumber(ARGV[5])
+          local prefix = ARGV[6]
+          local bytes = ARGV[7]
+
+          local current_group_score = redis.call("ZSCORE", overview_zset_key, group_name)
+          if current_group_score == false or score > tonumber(current_group_score) then
+            redis.call("ZADD", overview_zset_key, score, group_name)
+          end
+
+          local do_save = true
+          local overview_size = redis.call("ZCARD", overview_zset_key)
+          while (overview_size > groups_limit) do
+            local lowest_group = redis.call("ZRANGE", overview_zset_key, 0, 0)[1]
+            redis.call("ZREM", overview_zset_key, lowest_group)
+            if lowest_group == group_name then
+              do_save = false
+            else
+              local lowest_group_zset_key = prefix .. "-mp-group-snapshot-zset-key-" .. lowest_group
+              local lowest_group_hash_key = prefix .. "-mp-group-snapshot-hash-key-" .. lowest_group
+              redis.call("DEL", lowest_group_zset_key, lowest_group_hash_key)
+            end
+            overview_size = overview_size - 1
+          end
+
+          if do_save then
+            redis.call("ZADD", group_zset_key, score, id)
+            local group_size = redis.call("ZCARD", group_zset_key)
+            while (group_size > per_group_limit) do
+              local lowest_snapshot_id = redis.call("ZRANGE", group_zset_key, 0, 0)[1]
+              redis.call("ZREM", group_zset_key, lowest_snapshot_id)
+              if lowest_snapshot_id == id then
+                do_save = false
+              else
+                redis.call("HDEL", group_hash_key, lowest_snapshot_id)
+              end
+              group_size = group_size - 1
+            end
+            if do_save then
+              redis.call("HSET", group_hash_key, id, bytes)
+            end
           end
         LUA
         redis.eval(
           lua,
-          keys: [zset_key, hash_key],
-          argv: [id, score, bytes, limit]
+          keys: [group_zset_key, group_hash_key, overview_zset_key],
+          argv: [id, score, group_name, per_group_limit, groups_limit, @prefix, bytes]
         )
       end
 
-      def fetch_snapshots(batch_size: 200, &blk)
-        zset_key = snapshot_zset_key()
-        hash_key = snapshot_hash_key()
-        iteration = 0
-        corrupt_snapshots = []
-        while true
-          ids = redis.zrange(
-            zset_key,
-            batch_size * iteration,
-            batch_size * iteration + batch_size - 1
-          )
-          break if ids.size == 0
-          batch = redis.mapped_hmget(hash_key, *ids).to_a
-          batch.map! do |id, bytes|
-            begin
-              # rubocop:disable Security/MarshalLoad
-              Marshal.load(bytes)
-              # rubocop:enable Security/MarshalLoad
-            rescue
-              corrupt_snapshots << id
-              nil
-            end
+      def fetch_snapshots_overview
+        overview_zset_key = snapshot_overview_zset_key
+        groups = redis
+          .zrange(overview_zset_key, 0, -1, withscores: true)
+          .map { |(name, worst_score)| [name, { worst_score: worst_score }] }
+
+        prefixed_group_names = groups.map { |(group_name, _)| group_snapshot_zset_key(group_name) }
+        metadata = redis.eval(<<~LUA, keys: prefixed_group_names)
+          local metadata = {}
+          for i, k in ipairs(KEYS) do
+            local best = redis.call("ZRANGE", k, 0, 0, "WITHSCORES")[2]
+            local count = redis.call("ZCARD", k)
+            metadata[i] = {best, count}
           end
-          batch.compact!
-          blk.call(batch) if batch.size != 0
-          break if ids.size < batch_size
-          iteration += 1
+          return metadata
+        LUA
+        groups.each.with_index do |(_, hash), index|
+          best, count = metadata[index]
+          hash[:best_score] = best.to_f
+          hash[:snapshots_count] = count.to_i
+        end
+        groups.to_h
+      end
+
+      def fetch_snapshots_group(group_name)
+        group_hash_key = group_snapshot_hash_key(group_name)
+        snapshots = []
+        corrupt_snapshots = []
+        redis.hgetall(group_hash_key).each do |id, bytes|
+          # rubocop:disable Security/MarshalLoad
+          snapshots << Marshal.load(bytes)
+          # rubocop:enable Security/MarshalLoad
+        rescue
+          corrupt_snapshots << id
         end
         if corrupt_snapshots.size > 0
-          redis.pipelined do |pipeline|
-            pipeline.zrem(zset_key, corrupt_snapshots)
-            pipeline.hdel(hash_key, corrupt_snapshots)
-          end
+          cleanup_corrupt_snapshots(corrupt_snapshots, group_name)
         end
+        snapshots
       end
 
-      def load_snapshot(id)
-        hash_key = snapshot_hash_key()
-        bytes = redis.hget(hash_key, id)
+      def load_snapshot(id, group_name)
+        group_hash_key = group_snapshot_hash_key(group_name)
+        bytes = redis.hget(group_hash_key, id)
+        return if !bytes
         begin
           # rubocop:disable Security/MarshalLoad
           Marshal.load(bytes)
           # rubocop:enable Security/MarshalLoad
         rescue
-          redis.pipelined do |pipeline|
-            pipeline.zrem(snapshot_zset_key(), id)
-            pipeline.hdel(hash_key, id)
-          end
+          cleanup_corrupt_snapshots([id], group_name)
           nil
         end
       end
@@ -237,12 +280,20 @@ unviewed_ids: #{get_unviewed_ids(user)}
         @snapshot_counter_key ||= "#{@prefix}-mini-profiler-snapshots-counter"
       end
 
-      def snapshot_zset_key
-        @snapshot_zset_key ||= "#{@prefix}-mini-profiler-snapshots-zset"
+      def group_snapshot_zset_key(group_name)
+        # if you change this key, remember to change it in the LUA script in
+        # the push_snapshot method as well
+        "#{@prefix}-mp-group-snapshot-zset-key-#{group_name}"
       end
 
-      def snapshot_hash_key
-        @snapshot_hash_key ||= "#{@prefix}-mini-profiler-snapshots-hash"
+      def group_snapshot_hash_key(group_name)
+        # if you change this key, remember to change it in the LUA script in
+        # the push_snapshot method as well
+        "#{@prefix}-mp-group-snapshot-hash-key-#{group_name}"
+      end
+
+      def snapshot_overview_zset_key
+        "#{@prefix}-mp-overviewgroup-snapshot-zset-key"
       end
 
       def cached_redis_eval(script, script_sha, reraise: true, argv: [], keys: [])
@@ -257,12 +308,44 @@ unviewed_ids: #{get_unviewed_ids(user)}
         end
       end
 
+      def cleanup_corrupt_snapshots(corrupt_snapshots_ids, group_name)
+        group_hash_key = group_snapshot_hash_key(group_name)
+        group_zset_key = group_snapshot_zset_key(group_name)
+        overview_zset_key = snapshot_overview_zset_key
+        lua = <<~LUA
+          local group_hash_key = KEYS[1]
+          local group_zset_key = KEYS[2]
+          local overview_zset_key = KEYS[3]
+          local group_name = ARGV[1]
+          for i, k in ipairs(ARGV) do
+            if k ~= group_name then
+              redis.call("HDEL", group_hash_key, k)
+              redis.call("ZREM", group_zset_key, k)
+            end
+          end
+          if redis.call("ZCARD", group_zset_key) == 0 then
+            redis.call("ZREM", overview_zset_key, group_name)
+            redis.call("DEL", group_hash_key, group_zset_key)
+          else
+            local worst_score = tonumber(redis.call("ZRANGE", group_zset_key, -1, -1, "WITHSCORES")[2])
+            redis.call("ZADD", overview_zset_key, worst_score, group_name)
+          end
+        LUA
+        redis.eval(
+          lua,
+          keys: [group_hash_key, group_zset_key, overview_zset_key],
+          argv: [group_name, *corrupt_snapshots_ids]
+        )
+      end
+
       # only used in tests
       def wipe_snapshots_data
+        keys = redis.keys(group_snapshot_hash_key('*'))
+        keys += redis.keys(group_snapshot_zset_key('*'))
         redis.del(
-          snapshot_counter_key(),
-          snapshot_zset_key(),
-          snapshot_hash_key(),
+          keys,
+          snapshot_overview_zset_key,
+          snapshot_counter_key
         )
       end
     end

data/lib/mini_profiler/storage.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+require 'mini_profiler/storage/abstract_store'
+require 'mini_profiler/storage/memcache_store'
+require 'mini_profiler/storage/memory_store'
+require 'mini_profiler/storage/redis_store'
+require 'mini_profiler/storage/file_store'

data/lib/mini_profiler/timer_struct/base.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'json'
+
 module Rack
   class MiniProfiler
     module TimerStruct

data/lib/mini_profiler/timer_struct/sql.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'erb'
+
 module Rack
   class MiniProfiler
 

data/lib/mini_profiler/timer_struct.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require 'mini_profiler/timer_struct/base'
+require 'mini_profiler/timer_struct/page'
+require 'mini_profiler/timer_struct/sql'
+require 'mini_profiler/timer_struct/custom'
+require 'mini_profiler/timer_struct/client'
+require 'mini_profiler/timer_struct/request'

data/lib/mini_profiler/version.rb

@@ -2,7 +2,7 @@
 
 module Rack
   class MiniProfiler
-    VERSION = '2.3.4'
+    VERSION = '3.1.0'
     SOURCE_CODE_URI = 'https://github.com/MiniProfiler/rack-mini-profiler'
   end
 end