From 9e2964c15a7d387e46e25c83afa478c12a856d77 Mon Sep 17 00:00:00 2001
From: Patricio Cano <suprnova32@gmail.com>
Date: Fri, 4 Nov 2016 12:53:12 -0600
Subject: [PATCH] Allow certain Sidekiq jobs to be throttled

---
 Gemfile                                       |  1 +
 Gemfile.lock                                  |  3 ++
 .../admin/application_settings_controller.rb  |  1 +
 app/models/application_setting.rb             |  1 +
 .../application_settings/_form.html.haml      | 13 +++++++++
 config/initializers/sidekiq.rb                | 16 ++++++++++
 ...ekiq_throttling_to_application_settings.rb | 29 +++++++++++++++++++
 db/schema.rb                                  |  1 +
 lib/gitlab/current_settings.rb                |  1 +
 9 files changed, 66 insertions(+)
 create mode 100644 db/migrate/20161103191444_add_sidekiq_throttling_to_application_settings.rb

diff --git a/Gemfile b/Gemfile
index cb2a8470126..f2291568d25 100644
--- a/Gemfile
+++ b/Gemfile
@@ -137,6 +137,7 @@ gem 'acts-as-taggable-on', '~> 4.0'
 gem 'sidekiq', '~> 4.2'
 gem 'sidekiq-cron', '~> 0.4.0'
 gem 'redis-namespace', '~> 1.5.2'
+gem 'sidekiq-limit_fetch', '~> 3.4'
 
 # HTTP requests
 gem 'httparty', '~> 0.13.3'
diff --git a/Gemfile.lock b/Gemfile.lock
index 290e4c3e1b3..81b43f2238a 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -685,6 +685,8 @@ GEM
       redis-namespace (>= 1.5.2)
       rufus-scheduler (>= 2.0.24)
       sidekiq (>= 4.0.0)
+    sidekiq-limit_fetch (3.4.0)
+      sidekiq (>= 4)
     simplecov (0.12.0)
       docile (~> 1.1.0)
       json (>= 1.8, < 3)
@@ -961,6 +963,7 @@ DEPENDENCIES
   shoulda-matchers (~> 2.8.0)
   sidekiq (~> 4.2)
   sidekiq-cron (~> 0.4.0)
+  sidekiq-limit_fetch (~> 3.4)
   simplecov (= 0.12.0)
   slack-notifier (~> 1.2.0)
   spinach-rails (~> 0.2.1)
diff --git a/app/controllers/admin/application_settings_controller.rb b/app/controllers/admin/application_settings_controller.rb
index 52e0256943a..a9dcf7615c4 100644
--- a/app/controllers/admin/application_settings_controller.rb
+++ b/app/controllers/admin/application_settings_controller.rb
@@ -117,6 +117,7 @@ class Admin::ApplicationSettingsController < Admin::ApplicationController
       :send_user_confirmation_email,
       :container_registry_token_expire_delay,
       :enabled_git_access_protocol,
+      :sidekiq_throttling_enabled,
       :housekeeping_enabled,
       :housekeeping_bitmaps_enabled,
       :housekeeping_incremental_repack_period,
diff --git a/app/models/application_setting.rb b/app/models/application_setting.rb
index bb60cc8736c..b728083e91c 100644
--- a/app/models/application_setting.rb
+++ b/app/models/application_setting.rb
@@ -180,6 +180,7 @@ class ApplicationSetting < ActiveRecord::Base
       container_registry_token_expire_delay: 5,
       repository_storages: ['default'],
       user_default_external: false,
+      sidekiq_throttling_enabled: false,
       housekeeping_enabled: true,
       housekeeping_bitmaps_enabled: true,
       housekeeping_incremental_repack_period: 10,
diff --git a/app/views/admin/application_settings/_form.html.haml b/app/views/admin/application_settings/_form.html.haml
index 450ec322f2c..01a14accbba 100644
--- a/app/views/admin/application_settings/_form.html.haml
+++ b/app/views/admin/application_settings/_form.html.haml
@@ -283,6 +283,19 @@
           The amount of points to store in a single UDP packet. More points
           results in fewer but larger UDP packets being sent.
 
+  %fieldset
+    %legend Background Jobs
+    %p
+      These settings require a restart to take effect.
+    .form-group
+      .col-sm-offset-2.col-sm-10
+        .checkbox
+          = f.label :sidekiq_throttling_enabled do
+            = f.check_box :sidekiq_throttling_enabled
+            Enable Sidekiq Job Throttling
+          .help-block
+            Limit the amount of resources slow running jobs are assigned.
+
   %fieldset
     %legend Spam and Anti-bot Protection
     .form-group
diff --git a/config/initializers/sidekiq.rb b/config/initializers/sidekiq.rb
index 023af2af23c..6e660a8c026 100644
--- a/config/initializers/sidekiq.rb
+++ b/config/initializers/sidekiq.rb
@@ -1,3 +1,6 @@
+require 'gitlab/current_settings'
+include Gitlab::CurrentSettings
+
 # Custom Redis configuration
 redis_config_hash = Gitlab::Redis.params
 redis_config_hash[:namespace] = Gitlab::Redis::SIDEKIQ_NAMESPACE
@@ -29,6 +32,19 @@ Sidekiq.configure_server do |config|
   end
   Sidekiq::Cron::Job.load_from_hash! cron_jobs
 
+  # allow it to fail: it may do so when create_from_defaults is executed before migrations are actually done
+  begin
+    throttling_enabled = current_application_settings.sidekiq_throttling_enabled
+  rescue
+    throttling_enabled = false
+  end
+
+  if throttling_enabled
+    { 'project_cache' => 0.1, 'pipeline' => 0.1 }.each do |queue, ratio|
+      Sidekiq::Queue[queue].limit = (ratio * Sidekiq.options[:concurrency]).ceil
+    end
+  end
+
   # Database pool should be at least `sidekiq_concurrency` + 2
   # For more info, see: https://github.com/mperham/sidekiq/blob/master/4.0-Upgrade.md
   config = ActiveRecord::Base.configurations[Rails.env] ||
diff --git a/db/migrate/20161103191444_add_sidekiq_throttling_to_application_settings.rb b/db/migrate/20161103191444_add_sidekiq_throttling_to_application_settings.rb
new file mode 100644
index 00000000000..e2839219fb7
--- /dev/null
+++ b/db/migrate/20161103191444_add_sidekiq_throttling_to_application_settings.rb
@@ -0,0 +1,29 @@
+# See http://doc.gitlab.com/ce/development/migration_style_guide.html
+# for more information on how to write migrations for GitLab.
+
+class AddSidekiqThrottlingToApplicationSettings < ActiveRecord::Migration
+  include Gitlab::Database::MigrationHelpers
+
+  # Set this constant to true if this migration requires downtime.
+  DOWNTIME = false
+
+  # When a migration requires downtime you **must** uncomment the following
+  # constant and define a short and easy to understand explanation as to why the
+  # migration requires downtime.
+  # DOWNTIME_REASON = ''
+
+  # When using the methods "add_concurrent_index" or "add_column_with_default"
+  # you must disable the use of transactions as these methods can not run in an
+  # existing transaction. When using "add_concurrent_index" make sure that this
+  # method is the _only_ method called in the migration, any other changes
+  # should go in a separate migration. This ensures that upon failure _only_ the
+  # index creation fails and can be retried or reverted easily.
+  #
+  # To disable transactions uncomment the following line and remove these
+  # comments:
+  # disable_ddl_transaction!
+
+  def change
+    add_column :application_settings, :sidekiq_throttling_enabled, :boolean, default: false
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index 62c325a52d7..31d01403508 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -98,6 +98,7 @@ ActiveRecord::Schema.define(version: 20161106185620) do
     t.text "help_page_text_html"
     t.text "shared_runners_text_html"
     t.text "after_sign_up_text_html"
+    t.boolean "sidekiq_throttling_enabled", default: false
     t.boolean "housekeeping_enabled", default: true, null: false
     t.boolean "housekeeping_bitmaps_enabled", default: true, null: false
     t.integer "housekeeping_incremental_repack_period", default: 10, null: false
diff --git a/lib/gitlab/current_settings.rb b/lib/gitlab/current_settings.rb
index ef9160d6437..801c2d5abcd 100644
--- a/lib/gitlab/current_settings.rb
+++ b/lib/gitlab/current_settings.rb
@@ -50,6 +50,7 @@ module Gitlab
         repository_checks_enabled: true,
         container_registry_token_expire_delay: 5,
         user_default_external: false,
+        sidekiq_throttling_enabled: false,
       )
     end
 
-- 
GitLab