1# frozen_string_literal: true 2 3require 'securerandom' 4 5module Gitlab 6 # This class implements an 'exclusive lease'. We call it a 'lease' 7 # because it has a set expiry time. We call it 'exclusive' because only 8 # one caller may obtain a lease for a given key at a time. The 9 # implementation is intended to work across GitLab processes and across 10 # servers. It is a cheap alternative to using SQL queries and updates: 11 # you do not need to change the SQL schema to start using 12 # ExclusiveLease. 13 # 14 class ExclusiveLease 15 PREFIX = 'gitlab:exclusive_lease' 16 NoKey = Class.new(ArgumentError) 17 18 LUA_CANCEL_SCRIPT = <<~EOS 19 local key, uuid = KEYS[1], ARGV[1] 20 if redis.call("get", key) == uuid then 21 redis.call("del", key) 22 end 23 EOS 24 25 LUA_RENEW_SCRIPT = <<~EOS 26 local key, uuid, ttl = KEYS[1], ARGV[1], ARGV[2] 27 if redis.call("get", key) == uuid then 28 redis.call("expire", key, ttl) 29 return uuid 30 end 31 EOS 32 33 def self.get_uuid(key) 34 Gitlab::Redis::SharedState.with do |redis| 35 redis.get(redis_shared_state_key(key)) || false 36 end 37 end 38 39 # yield to the {block} at most {count} times per {period} 40 # 41 # Defaults to once per hour. 42 # 43 # For example: 44 # 45 # # toot the train horn at most every 20min: 46 # throttle(locomotive.id, count: 3, period: 1.hour) { toot_train_horn } 47 # # Brake suddenly at most once every minute: 48 # throttle(locomotive.id, period: 1.minute) { brake_suddenly } 49 # # Specify a uniqueness group: 50 # throttle(locomotive.id, group: :locomotive_brake) { brake_suddenly } 51 # 52 # If a group is not specified, each block will get a separate group to itself. 53 def self.throttle(key, group: nil, period: 1.hour, count: 1, &block) 54 group ||= block.source_location.join(':') 55 56 return if new("el:throttle:#{group}:#{key}", timeout: period.to_i / count).waiting? 57 58 yield 59 end 60 61 def self.cancel(key, uuid) 62 return unless key.present? 63 64 Gitlab::Redis::SharedState.with do |redis| 65 redis.eval(LUA_CANCEL_SCRIPT, keys: [ensure_prefixed_key(key)], argv: [uuid]) 66 end 67 end 68 69 def self.redis_shared_state_key(key) 70 "#{PREFIX}:#{key}" 71 end 72 73 def self.ensure_prefixed_key(key) 74 raise NoKey unless key.present? 75 76 key.start_with?(PREFIX) ? key : redis_shared_state_key(key) 77 end 78 79 # Removes any existing exclusive_lease from redis 80 # Don't run this in a live system without making sure no one is using the leases 81 def self.reset_all!(scope = '*') 82 Gitlab::Redis::SharedState.with do |redis| 83 redis.scan_each(match: redis_shared_state_key(scope)).each do |key| 84 redis.del(key) 85 end 86 end 87 end 88 89 def initialize(key, uuid: nil, timeout:) 90 @redis_shared_state_key = self.class.redis_shared_state_key(key) 91 @timeout = timeout 92 @uuid = uuid || SecureRandom.uuid 93 end 94 95 # Try to obtain the lease. Return lease UUID on success, 96 # false if the lease is already taken. 97 def try_obtain 98 # Performing a single SET is atomic 99 Gitlab::Redis::SharedState.with do |redis| 100 redis.set(@redis_shared_state_key, @uuid, nx: true, ex: @timeout) && @uuid 101 end 102 end 103 104 # This lease is waiting to obtain 105 def waiting? 106 !try_obtain 107 end 108 109 # Try to renew an existing lease. Return lease UUID on success, 110 # false if the lease is taken by a different UUID or inexistent. 111 def renew 112 Gitlab::Redis::SharedState.with do |redis| 113 result = redis.eval(LUA_RENEW_SCRIPT, keys: [@redis_shared_state_key], argv: [@uuid, @timeout]) 114 result == @uuid 115 end 116 end 117 118 # Returns true if the key for this lease is set. 119 def exists? 120 Gitlab::Redis::SharedState.with do |redis| 121 redis.exists(@redis_shared_state_key) 122 end 123 end 124 125 # Returns the TTL of the Redis key. 126 # 127 # This method will return `nil` if no TTL could be obtained. 128 def ttl 129 Gitlab::Redis::SharedState.with do |redis| 130 ttl = redis.ttl(@redis_shared_state_key) 131 132 ttl if ttl > 0 133 end 134 end 135 136 # Gives up this lease, allowing it to be obtained by others. 137 def cancel 138 self.class.cancel(@redis_shared_state_key, @uuid) 139 end 140 end 141end 142 143Gitlab::ExclusiveLease.prepend_mod_with('Gitlab::ExclusiveLease') 144