Add latest changes from gitlab-org/gitlab@master

1 files changed, 33 insertions, 135 deletions

diff --git a/app/models/concerns/counter_attribute.rb b/app/models/concerns/counter_attribute.rb
index 4dfaae674bf..f1efbba67e1 100644
--- a/app/models/concerns/counter_attribute.rb
+++ b/app/models/concerns/counter_attribute.rb
@@ -30,13 +30,16 @@
 #   end
 #
 # To increment the counter we can use the method:
-#   delayed_increment_counter(:commit_count, 3)
+#   increment_counter(:commit_count, 3)
+#
+# This method would determine whether it would increment the counter using Redis,
+# or fallback to legacy increment on ActiveRecord counters.
 #
 # It is possible to register callbacks to be executed after increments have
 # been flushed to the database. Callbacks are not executed if there are no increments
 # to flush.
 #
-#  counter_attribute_after_flush do |statistic|
+#  counter_attribute_after_commit do |statistic|
 #    Namespaces::ScheduleAggregationWorker.perform_async(statistic.namespace_id)
 #  end
 #
@@ -44,21 +47,7 @@ module CounterAttribute
   extend ActiveSupport::Concern
   extend AfterCommitQueue
   include Gitlab::ExclusiveLeaseHelpers
-
-  LUA_STEAL_INCREMENT_SCRIPT = <<~EOS
-    local increment_key, flushed_key = KEYS[1], KEYS[2]
-    local increment_value = redis.call("get", increment_key) or 0
-    local flushed_value = redis.call("incrby", flushed_key, increment_value)
-    if flushed_value == 0 then
-      redis.call("del", increment_key, flushed_key)
-    else
-      redis.call("del", increment_key)
-    end
-    return flushed_value
-  EOS
-
-  WORKER_DELAY = 10.minutes
-  WORKER_LOCK_TTL = 10.minutes
+  include Gitlab::Utils::StrongMemoize
 
   class_methods do
     def counter_attribute(attribute, if: nil)
@@ -72,13 +61,13 @@ module CounterAttribute
       @counter_attributes ||= []
     end
 
-    def after_flush_callbacks
-      @after_flush_callbacks ||= []
+    def after_commit_callbacks
+      @after_commit_callbacks ||= []
     end
 
-    # perform registered callbacks after increments have been flushed to the database
-    def counter_attribute_after_flush(&callback)
-      after_flush_callbacks << callback
+    # perform registered callbacks after increments have been committed to the database
+    def counter_attribute_after_commit(&callback)
+      after_commit_callbacks << callback
     end
   end
 
@@ -90,60 +79,19 @@ module CounterAttribute
     counter_attribute[:if_proc].call(self)
   end
 
-  # This method must only be called by FlushCounterIncrementsWorker
-  # because it should run asynchronously and with exclusive lease.
-  # This will
-  #  1. temporarily move the pending increment for a given attribute
-  #     to a relative "flushed" Redis key, delete the increment key and return
-  #     the value. If new increments are performed at this point, the increment
-  #     key is recreated as part of `delayed_increment_counter`.
-  #     The "flushed" key is used to ensure that we can keep incrementing
-  #     counters in Redis while flushing existing values.
-  #  2. then the value is used to update the counter in the database.
-  #  3. finally the "flushed" key is deleted.
-  def flush_increments_to_database!(attribute)
-    lock_key = counter_lock_key(attribute)
-
-    with_exclusive_lease(lock_key) do
-      previous_db_value = read_attribute(attribute)
-      increment_key = counter_key(attribute)
-      flushed_key = counter_flushed_key(attribute)
-      increment_value = steal_increments(increment_key, flushed_key)
-      new_db_value = nil
-
-      next if increment_value == 0
-
-      transaction do
-        update_counters_with_lease({ attribute => increment_value })
-        redis_state { |redis| redis.del(flushed_key) }
-        new_db_value = reset.read_attribute(attribute)
-      end
-
-      execute_after_flush_callbacks
-
-      log_flush_counter(attribute, increment_value, previous_db_value, new_db_value)
+  def counter(attribute)
+    strong_memoize_with(:counter, attribute) do
+      # This needs #to_sym because attribute could come from a Sidekiq param,
+      # which would be a string.
+      build_counter_for(attribute.to_sym)
     end
   end
 
-  def delayed_increment_counter(attribute, increment)
-    raise ArgumentError, "#{attribute} is not a counter attribute" unless counter_attribute_enabled?(attribute)
-
+  def increment_counter(attribute, increment)
     return if increment == 0
 
     run_after_commit_or_now do
-      increment_counter(attribute, increment)
-
-      FlushCounterIncrementsWorker.perform_in(WORKER_DELAY, self.class.name, self.id, attribute)
-    end
-
-    true
-  end
-
-  def increment_counter(attribute, increment)
-    if counter_attribute_enabled?(attribute)
-      new_value = redis_state do |redis|
-        redis.incrby(counter_key(attribute), increment)
-      end
+      new_value = counter(attribute).increment(increment)
 
       log_increment_counter(attribute, increment, new_value)
     end
@@ -156,70 +104,33 @@ module CounterAttribute
   end
 
   def reset_counter!(attribute)
-    if counter_attribute_enabled?(attribute)
-      detect_race_on_record(log_fields: { caller: __method__, attributes: attribute }) do
-        update!(attribute => 0)
-        clear_counter!(attribute)
-      end
-
-      log_clear_counter(attribute)
+    detect_race_on_record(log_fields: { caller: __method__, attributes: attribute }) do
+      counter(attribute).reset!
     end
-  end
-
-  def get_counter_value(attribute)
-    if counter_attribute_enabled?(attribute)
-      redis_state do |redis|
-        redis.get(counter_key(attribute)).to_i
-      end
-    end
-  end
-
-  def counter_key(attribute)
-    "project:{#{project_id}}:counters:#{self.class}:#{id}:#{attribute}"
-  end
 
-  def counter_flushed_key(attribute)
-    counter_key(attribute) + ':flushed'
+    log_clear_counter(attribute)
   end
 
-  def counter_lock_key(attribute)
-    counter_key(attribute) + ':lock'
+  def execute_after_commit_callbacks
+    self.class.after_commit_callbacks.each do |callback|
+      callback.call(self.reset)
+    end
   end
 
   private
 
-  def database_lock_key
-    "project:{#{project_id}}:#{self.class}:#{id}"
-  end
-
-  def steal_increments(increment_key, flushed_key)
-    redis_state do |redis|
-      redis.eval(LUA_STEAL_INCREMENT_SCRIPT, keys: [increment_key, flushed_key])
-    end
-  end
-
-  def clear_counter!(attribute)
-    redis_state do |redis|
-      redis.del(counter_key(attribute))
-    end
-  end
+  def build_counter_for(attribute)
+    raise ArgumentError, %(attribute "#{attribute}" does not exist) unless has_attribute?(attribute)
 
-  def execute_after_flush_callbacks
-    self.class.after_flush_callbacks.each do |callback|
-      callback.call(self)
+    if counter_attribute_enabled?(attribute)
+      Gitlab::Counters::BufferedCounter.new(self, attribute)
+    else
+      Gitlab::Counters::LegacyCounter.new(self, attribute)
     end
   end
 
-  def redis_state(&block)
-    Gitlab::Redis::SharedState.with(&block)
-  end
-
-  def with_exclusive_lease(lock_key)
-    in_lock(lock_key, ttl: WORKER_LOCK_TTL) do
-      yield
-    end
-  rescue Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError
-    # a worker is already updating the counters
+  def database_lock_key
+    "project:{#{project_id}}:#{self.class}:#{id}"
   end
 
   # detect_race_on_record uses a lease to monitor access
@@ -273,19 +184,6 @@ module CounterAttribute
     Gitlab::AppLogger.info(payload)
   end
 
-  def log_flush_counter(attribute, increment, previous_db_value, new_db_value)
-    payload = Gitlab::ApplicationContext.current.merge(
-      message: 'Flush counter attribute to database',
-      attribute: attribute,
-      project_id: project_id,
-      increment: increment,
-      previous_db_value: previous_db_value,
-      new_db_value: new_db_value
-    )
-
-    Gitlab::AppLogger.info(payload)
-  end
-
   def log_clear_counter(attribute)
     payload = Gitlab::ApplicationContext.current.merge(
       message: 'Clear counter attribute',