From 74d2d4f658171bcab41ccd7257495e7e39fb45e4 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Thu, 14 Feb 2019 13:33:58 +1100
Subject: [PATCH] FEATURE: add APIS for unpausing all sites

This adjusts 53d592ad by @tgxworld

- Adds Sidekiq.upause_all! to unpause all sites
- Adds Sidekiq.paused_dbs to list dbs that are currently paused
- Handles some edge cases where unpause thread could extend expiry on
sites that were unpaused from a different process
- Ensures tests always terminates background thread used for pause
keepalive
---
 lib/sidekiq/pausable.rb         | 54 +++++++++++++++++++++++++++++----
 spec/multisite/pausable_spec.rb | 13 +++++---
 2 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/lib/sidekiq/pausable.rb b/lib/sidekiq/pausable.rb
index 16d4a714623..48106478568 100644
--- a/lib/sidekiq/pausable.rb
+++ b/lib/sidekiq/pausable.rb
@@ -14,7 +14,6 @@ class SidekiqPauser
 
     @mutex.synchronize do
       extend_lease_thread
-      sleep 0.001 while !paused?
     end
 
     true
@@ -24,10 +23,29 @@ class SidekiqPauser
     !!$redis.get(PAUSED_KEY)
   end
 
+  def unpause_all!
+    @mutex.synchronize do
+      @dbs = Set.new
+      stop_extend_lease_thread
+    end
+
+    RailsMultisite::ConnectionManagement.each_connection do
+      unpause! if paused?
+    end
+  end
+
+  def paused_dbs
+    dbs = []
+    RailsMultisite::ConnectionManagement.each_connection do
+      dbs << RailsMultisite::ConnectionManagement.current_db if paused?
+    end
+    dbs
+  end
+
   def unpause!
     @mutex.synchronize do
       @dbs.delete(RailsMultisite::ConnectionManagement.current_db)
-      @extend_lease_thread = nil if @dbs.size == 0
+      stop_extend_lease_thread if @dbs.size == 0
     end
 
     $redis.del(PAUSED_KEY)
@@ -36,16 +54,32 @@ class SidekiqPauser
 
   private
 
-  def extend_lease_thread
-    @dbs << RailsMultisite::ConnectionManagement.current_db
+  def stop_extend_lease_thread
+    # should always be called from a mutex
+    if t = @extend_lease_thread
+      @extend_lease_thread = nil
+      while t.alive?
+        t.wakeup
+        sleep 0
+      end
+    end
+  end
 
+  def extend_lease_thread
+    # should always be called from a mutex
+    @dbs << RailsMultisite::ConnectionManagement.current_db
     @extend_lease_thread ||= Thread.new do
       while true do
-        break unless @mutex.synchronize { @extend_lease_thread }
+
+        break if !@extend_lease_thread
 
         @dbs.each do |db|
           RailsMultisite::ConnectionManagement.with_connection(db) do
-            $redis.expire PAUSED_KEY, TTL
+            if !$redis.expire(PAUSED_KEY, TTL)
+              # if it was unpaused in another process we got to remove the
+              # bad key
+              @dbs.delete(db)
+            end
           end
         end
 
@@ -68,6 +102,14 @@ module Sidekiq
   def self.unpause!
     @pauser.unpause!
   end
+
+  def self.unpause_all!
+    @pauser.unpause_all!
+  end
+
+  def self.paused_dbs
+    @pauser.paused_dbs
+  end
 end
 
 # server middleware that will reschedule work whenever Sidekiq is paused
diff --git a/spec/multisite/pausable_spec.rb b/spec/multisite/pausable_spec.rb
index 57297dd4f7d..dc0e0f9ee30 100644
--- a/spec/multisite/pausable_spec.rb
+++ b/spec/multisite/pausable_spec.rb
@@ -2,9 +2,6 @@ require 'rails_helper'
 require_dependency 'sidekiq/pausable'
 
 RSpec.describe "Pausing/Unpausing Sidekiq", type: :multisite do
-  after do
-    $redis.flushall
-  end
 
   describe '#pause!, #unpause! and #paused?' do
     it "can pause and unpause" do
@@ -23,13 +20,21 @@ RSpec.describe "Pausing/Unpausing Sidekiq", type: :multisite do
         Sidekiq.pause!
         expect(Sidekiq.paused?).to eq(true)
       end
+
+      expect(Sidekiq.paused_dbs).to eq(["second"])
+
+      Sidekiq.unpause_all!
+
+      RailsMultisite::ConnectionManagement.each_connection do
+        expect(Sidekiq.paused?).to eq(false)
+      end
     end
   end
 end
 
 RSpec.describe Sidekiq::Pausable do
   after do
-    $redis.flushall
+    Sidekiq.unpause_all!
   end
 
   it "can still run heartbeats when paused" do