Add rake task for cleaning up sidekiq batch keys (#3912)

* Add rake task for cleaning up sidekiq batch keys

Apply suggestions from code review

Co-authored-by: Aleksandar N. Kostadinov <[email protected]>

* Remove unnecessary logs

* Refactoring

---------

Co-authored-by: Aleksandar N. Kostadinov <[email protected]>

app/lib/backend/storage_rewrite.rb

@@ -130,14 +130,8 @@ def process(collection)
         rewrite(scope: collection)
       end
 
-      # This logger just prints out a message to STDOUT, with new line before and after.
-      # New line before is to make progress log look better
       def logger
-        @logger ||= begin
-          log = ActiveSupport::Logger.new($stdout)
-          log.formatter = ->(_, _, _, msg) { "\n#{msg.is_a?(String) ? msg : msg.inspect}\n" }
-          log
-        end
+        @logger ||= ProgressCounter.stdout_logger
       end
 
       # All accounts eligible for backend sync

config/application.rb

@@ -267,6 +267,7 @@ def cache_store_config
     require 'three_scale/domain_substitution'
     require 'three_scale/middleware/multitenant'
     require 'three_scale/middleware/cors'
+    require 'three_scale/patterns/service'
 
     config.middleware.use ThreeScale::Middleware::Multitenant, :tenant_id unless ENV["DEBUG_DISABLE_TENANT_CHECK"] == "1"
     config.middleware.insert_before Rack::Runtime, Rack::UTF8Sanitizer

lib/progress_counter.rb

@@ -3,6 +3,16 @@
 class ProgressCounter
   attr_reader :total, :index
 
+  class << self
+    # This logger just prints out a message to STDOUT, with new line before and after.
+    # New line before is to make progress log look better
+    def stdout_logger
+      log = ActiveSupport::Logger.new($stdout)
+      log.formatter = ->(_, _, _, msg) { "\n#{msg.is_a?(String) ? msg : msg.inspect}\n" }
+      log
+    end
+  end
+
   def initialize(total)
     @index = 0
     @total = total

lib/tasks/sidekiq/cleanup.rake

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require 'three_scale/sidekiq_batch_cleanup_service'
+
+namespace :sidekiq do
+  desc 'cleanup BID-* keys from sidekiq-batch, specify the max age in seconds as an argument'
+  task :cleanup_batches, [:max_age_seconds] => :environment do |task, args|
+    params = args[:max_age_seconds] ? { max_age_seconds: Integer(args[:max_age_seconds]) } : {} 
+
+    ThreeScale::SidekiqBatchCleanupService.call(**params)
+  end
+end

lib/three_scale/sidekiq_batch_cleanup_service.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require 'progress_counter'
+
+module ThreeScale
+  class SidekiqBatchCleanupService < ThreeScale::Patterns::Service
+    MAX_FETCH_COUNT = 1000
+
+    BID_EXPIRE_TTL = 30.days.seconds
+    DEFAULT_MAX_AGE_SECONDS = 3.hours.seconds
+
+    # `max_age_seconds` specifies the maximum age of the keys (in seconds)
+    # all keys that are older will be deleted, calculated by the TTL that is still left, compared with the default expire value
+    def initialize(max_age_seconds: DEFAULT_MAX_AGE_SECONDS)
+      @now = Time.zone.now
+      @redis = System.redis
+
+      @bid_max_ttl = BID_EXPIRE_TTL - max_age_seconds
+      super()
+    end
+
+    attr_reader :now, :redis, :bid_max_ttl
+
+    def call
+      total = redis.dbsize
+      logger.info "Total number of keys: #{total}, will delete BID-* keys with TTL less than #{bid_max_ttl.seconds.in_hours} hours"
+
+      scan_enum = System.redis.scan_each(match: 'BID-*', type: 'hash', count: MAX_FETCH_COUNT)
+
+      each_with_progress_counter(scan_enum, total) do |key|
+        next if /-(success|complete|failed|jids)$/.match?(key)
+
+        bid = key.remove(/^BID-/)
+        perform(bid)
+      end
+    end
+
+    def perform(bid)
+      # TODO: ensure there is no task running still
+      batch_key = "BID-#{bid}"
+      ttl = redis.ttl(batch_key)
+
+      Sidekiq::Batch.cleanup_redis(bid) if ttl <= bid_max_ttl
+    end
+
+    private
+
+    def each_with_progress_counter(enumerable, count)
+      progress = ProgressCounter.new(count)
+      enumerable.each do |element|
+        yield element
+        progress.call
+      end
+    end
+
+    def logger
+      @logger ||= ProgressCounter.stdout_logger
+    end
+  end
+end