1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
# frozen_string_literal: true
module Gitlab
module SidekiqMiddleware
module SizeLimiter
# Validate 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`.
#
# - 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.
class Validator
def self.validate!(worker_class, job)
new(worker_class, job).validate!
end
DEFAULT_SIZE_LIMIT = 0
MODES = [
TRACK_MODE = 'track',
RAISE_MODE = 'raise'
].freeze
attr_reader :mode, :size_limit
def initialize(
worker_class, job,
mode: ENV['GITLAB_SIDEKIQ_SIZE_LIMITER_MODE'],
size_limit: ENV['GITLAB_SIDEKIQ_SIZE_LIMITER_LIMIT_BYTES']
)
@worker_class = worker_class
@job = job
@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
@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?
return if job_size <= @size_limit
exception = ExceedLimitError.new(@worker_class, job_size, @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 raise_mode?
raise exception
else
track(exception)
end
end
private
def job_size
# This maynot be the optimal solution, but can be acceptable solution
# for now. Internally, Sidekiq calls Sidekiq.dump_json everywhere.
# There is no clean way to intefere to prevent double serialization.
@job_size ||= ::Sidekiq.dump_json(@job).bytesize
end
def allow_big_payload?
worker_class = @worker_class.to_s.safe_constantize
worker_class.respond_to?(:big_payload?) && worker_class.big_payload?
end
def raise_mode?
@mode == RAISE_MODE
end
def track(exception)
Gitlab::ErrorTracking.track_exception(exception)
end
def backtrace
Gitlab::BacktraceCleaner.clean_backtrace(caller)
end
end
end
end
end
|
