Refactor big methods in Sidekiq size limiter

Issue https://gitlab.com/gitlab-com/gl-infra/scalability/-/issues/1054

lib/gitlab/sidekiq_middleware/size_limiter/compressor.rb

@@ -30,18 +30,21 @@ module Gitlab
         def self.decompress(job)
           return unless compressed?(job)
 
+          validate_args!(job)
+
+          job.except!(ORIGINAL_SIZE_KEY, COMPRESSED_KEY)
+          job['args'] = Sidekiq.load_json(Zlib::Inflate.inflate(Base64.strict_decode64(job['args'].first)))
+        rescue Zlib::Error
+          raise PayloadDecompressionError, 'Fail to decompress Sidekiq job payload'
+        end
+
+        def self.validate_args!(job)
           if job['args'] && job['args'].length != 1
             exception = PayloadDecompressionConflictError.new('Sidekiq argument list should include 1 argument.\
                                                               There should be a middleware interferes the job payload.\
                                                               That conflicts with the payload compressor')
             ::Gitlab::ErrorTracking.track_and_raise_for_dev_exception(exception)
           end
-
-          job.delete(ORIGINAL_SIZE_KEY)
-          job.delete(COMPRESSED_KEY)
-          job['args'] = Sidekiq.load_json(Zlib::Inflate.inflate(Base64.strict_decode64(job['args'].first)))
-        rescue Zlib::Error
-          raise PayloadDecompressionError, 'Fail to decompress Sidekiq job payload'
         end
       end
     end

lib/gitlab/sidekiq_middleware/size_limiter/validator.rb

@@ -3,16 +3,20 @@
 module Gitlab
   module SidekiqMiddleware
     module SizeLimiter
-      # Validate a Sidekiq job payload limit based on current configuration.
+      # Handle a Sidekiq job payload limit based on current configuration.
       # This validator pulls the configuration from the environment variables:
-      #
       # - GITLAB_SIDEKIQ_SIZE_LIMITER_MODE: the current mode of the size
-      # limiter. This must be either `track` or `raise`.
-      #
+      # limiter. This must be either `track` or `compress`.
+      # - GITLAB_SIDEKIQ_SIZE_LIMITER_COMPRESSION_THRESHOLD_BYTES: the
+      # threshold before the input job payload is compressed.
       # - GITLAB_SIDEKIQ_SIZE_LIMITER_LIMIT_BYTES: the size limit in bytes.
       #
-      # If the size of job payload after serialization exceeds the limit, an
-      # error is tracked raised adhering to the mode.
+      # In track mode, if a job payload limit exceeds the size limit, an
+      # event is sent to Sentry and the job is scheduled like normal.
+      #
+      # In compress mode, if a job payload limit exceeds the threshold, it is
+      # then compressed. If the compressed payload still exceeds the limit, the
+      # job is discarded, and a ExceedLimitError exception is raised.
       class Validator
         def self.validate!(worker_class, job)
           new(worker_class, job).validate!
@@ -37,45 +41,59 @@ module Gitlab
           @worker_class = worker_class
           @job = job
 
+          set_mode(mode)
+          set_compression_threshold(compression_threshold)
+          set_size_limit(size_limit)
+        end
+
+        def validate!
+          return unless @size_limit > 0
+          return if allow_big_payload?
+
+          job_args = compress_if_necessary(::Sidekiq.dump_json(@job['args']))
+          return if job_args.bytesize <= @size_limit
+
+          exception = exceed_limit_error(job_args)
+          if compress_mode?
+            raise exception
+          else
+            track(exception)
+          end
+        end
+
+        private
+
+        def set_mode(mode)
           @mode = (mode || TRACK_MODE).to_s.strip
           unless MODES.include?(@mode)
             ::Sidekiq.logger.warn "Invalid Sidekiq size limiter mode: #{@mode}. Fallback to #{TRACK_MODE} mode."
             @mode = TRACK_MODE
           end
+        end
 
+        def set_compression_threshold(compression_threshold)
           @compression_threshold = (compression_threshold || DEFAULT_COMPRESION_THRESHOLD_BYTES).to_i
           if @compression_threshold < 0
             ::Sidekiq.logger.warn "Invalid Sidekiq size limiter compression threshold: #{@compression_threshold}"
           end
+        end
 
+        def set_size_limit(size_limit)
           @size_limit = (size_limit || DEFAULT_SIZE_LIMIT).to_i
           if @size_limit < 0
             ::Sidekiq.logger.warn "Invalid Sidekiq size limiter limit: #{@size_limit}"
           end
         end
 
-        def validate!
-          return unless @size_limit > 0
-          return if allow_big_payload?
-
-          job_args = compress_if_necessary(::Sidekiq.dump_json(@job['args']))
-          return if job_args.bytesize <= @size_limit
-
-          exception = ExceedLimitError.new(@worker_class, job_args.bytesize, @size_limit)
-          # This should belong to Gitlab::ErrorTracking. We'll remove this
-          # after this epic is done:
-          # https://gitlab.com/groups/gitlab-com/gl-infra/-/epics/396
-          exception.set_backtrace(backtrace)
-
-          if compress_mode?
-            raise exception
-          else
-            track(exception)
+        def exceed_limit_error(job_args)
+          ExceedLimitError.new(@worker_class, job_args.bytesize, @size_limit).tap do |exception|
+            # This should belong to Gitlab::ErrorTracking. We'll remove this
+            # after this epic is done:
+            # https://gitlab.com/groups/gitlab-com/gl-infra/-/epics/396
+            exception.set_backtrace(backtrace)
           end
         end
 
-        private
-
         def compress_if_necessary(job_args)
           return job_args if !compress_mode? || job_args.bytesize < @compression_threshold