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