class Rack::MiniProfiler::RedisStore

Constants

COUNTER_LUA
COUNTER_LUA_SHA
EXPIRES_IN_SECONDS

Attributes

prefix[R]

Public Class Methods

new(args = nil) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 14
def initialize(args = nil)
  @args               = args || {}
  @prefix             = @args.delete(:prefix) || 'MPRedisStore'
  @redis_connection   = @args.delete(:connection)
  @expires_in_seconds = @args.delete(:expires_in) || EXPIRES_IN_SECONDS
end

Public Instance Methods

allowed_tokens() click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 88
def allowed_tokens
  key1, key1_old, key2 = redis.mget("#{@prefix}-key1", "#{@prefix}-key1_old", "#{@prefix}-key2")

  if key1 && (key1.length == 32)
    return [key1, key2].compact
  end

  timeout = Rack::MiniProfiler::AbstractStore::MAX_TOKEN_AGE

  # TODO  this could be moved to lua to correct a concurrency flaw
  # it is not critical cause worse case some requests will miss profiling info

  # no key so go ahead and set it
  key1 = SecureRandom.hex

  if key1_old && (key1_old.length == 32)
    key2 = key1_old
    redis.setex "#{@prefix}-key2", timeout, key2
  else
    key2 = nil
  end

  redis.setex "#{@prefix}-key1", timeout, key1
  redis.setex "#{@prefix}-key1_old", timeout * 2, key1

  [key1, key2].compact
end
diagnostics(user) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 71
      def diagnostics(user)
        client = (redis.respond_to? :_client) ? redis._client : redis.client
"Redis prefix: #{@prefix}
Redis location: #{client.host}:#{client.port} db: #{client.db}
unviewed_ids: #{get_unviewed_ids(user)}
"
      end
fetch_snapshots_group(group_name) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 231
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
    cleanup_corrupt_snapshots(corrupt_snapshots, group_name)
  end
  snapshots
end
fetch_snapshots_overview() click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 207
      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
          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
flush_tokens() click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 79
def flush_tokens
  redis.del("#{@prefix}-key1", "#{@prefix}-key1_old", "#{@prefix}-key2")
end
get_unviewed_ids(user) click to toggle source

Remove expired ids from the unviewed sorted set and return the remaining ids

# File lib/mini_profiler/storage/redis_store.rb, line 65
def get_unviewed_ids(user)
  key = user_key(user)
  redis.zremrangebyscore(key, '-inf', Process.clock_gettime(Process::CLOCK_MONOTONIC).to_i)
  redis.zrevrangebyscore(key, '+inf', '-inf')
end
load(id) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 25
def load(id)
  key = prefixed_id(id)
  raw = redis.get key
  begin
    # rubocop:disable Security/MarshalLoad
    Marshal.load(raw) if raw
    # rubocop:enable Security/MarshalLoad
  rescue
    # bad format, junk old data
    redis.del key
    nil
  end
end
load_snapshot(id, group_name) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 248
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
    cleanup_corrupt_snapshots([id], group_name)
    nil
  end
end
push_snapshot(page_struct, group_name, config) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 137
      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.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 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 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: [group_zset_key, group_hash_key, overview_zset_key],
          argv: [id, score, group_name, per_group_limit, groups_limit, @prefix, bytes]
        )
      end
save(page_struct) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 21
def save(page_struct)
  redis.setex prefixed_id(page_struct[:id]), @expires_in_seconds, Marshal::dump(page_struct)
end
set_all_unviewed(user, ids) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 48
def set_all_unviewed(user, ids)
  key = user_key(user)
  redis.del(key)
  ids.each do |id|
    if redis.call([:exists, prefixed_id(id)]) == 1
      expire_at = Process.clock_gettime(Process::CLOCK_MONOTONIC).to_i + redis.ttl(prefixed_id(id))
      redis.zadd(key, expire_at, id)
    end
  end
  redis.expire(key, @expires_in_seconds)
end
set_unviewed(user, id) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 39
def set_unviewed(user, id)
  key = user_key(user)
  if redis.call([:exists, prefixed_id(id)]) == 1
    expire_at = Process.clock_gettime(Process::CLOCK_MONOTONIC).to_i + redis.ttl(prefixed_id(id))
    redis.zadd(key, expire_at, id)
  end
  redis.expire(key, @expires_in_seconds)
end
set_viewed(user, id) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 60
def set_viewed(user, id)
  redis.zrem(user_key(user), id)
end
should_take_snapshot?(period) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 127
def should_take_snapshot?(period)
  1 == cached_redis_eval(
    COUNTER_LUA,
    COUNTER_LUA_SHA,
    reraise: false,
    keys: [snapshot_counter_key()],
    argv: [period]
  )
end
simulate_expire() click to toggle source

Only used for testing

# File lib/mini_profiler/storage/redis_store.rb, line 84
def simulate_expire
  redis.del("#{@prefix}-key1")
end

Private Instance Methods

cached_redis_eval(script, script_sha, reraise: true, argv: [], keys: []) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 299
def cached_redis_eval(script, script_sha, reraise: true, argv: [], keys: [])
  begin
    redis.evalsha(script_sha, argv: argv, keys: keys)
  rescue ::Redis::CommandError => e
    if e.message.start_with?('NOSCRIPT')
      redis.eval(script, argv: argv, keys: keys)
    else
      raise e if reraise
    end
  end
end
cleanup_corrupt_snapshots(corrupt_snapshots_ids, group_name) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 311
      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
group_snapshot_hash_key(group_name) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 289
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
group_snapshot_zset_key(group_name) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 283
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
prefixed_id(id) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 268
def prefixed_id(id)
  "#{@prefix}#{id}"
end
redis() click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 272
def redis
  @redis_connection ||= begin
    require 'redis' unless defined? Redis
    Redis.new(@args)
  end
end
snapshot_counter_key() click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 279
def snapshot_counter_key
  @snapshot_counter_key ||= "#{@prefix}-mini-profiler-snapshots-counter"
end
snapshot_overview_zset_key() click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 295
def snapshot_overview_zset_key
  "#{@prefix}-mp-overviewgroup-snapshot-zset-key"
end
user_key(user) click to toggle source
# File lib/mini_profiler/storage/redis_store.rb, line 264
def user_key(user)
  "#{@prefix}-#{user}-v1"
end
wipe_snapshots_data() click to toggle source

only used in tests

# File lib/mini_profiler/storage/redis_store.rb, line 342
def wipe_snapshots_data
  keys = redis.keys(group_snapshot_hash_key('*'))
  keys += redis.keys(group_snapshot_zset_key('*'))
  redis.del(
    keys,
    snapshot_overview_zset_key,
    snapshot_counter_key
  )
end