From 7c8920c9dc6449f9d9a76dac619e6829398ec7db Mon Sep 17 00:00:00 2001
From: Gabriel Mazetto <brodock@gmail.com>
Date: Wed, 27 Feb 2019 06:08:34 +0100
Subject: [PATCH] Extract common logic to HashedStorage::BaseWorker

New class contains the ExclusiveLease specifics that is shared among
both the Migration and Rollback workers.
---
 .../hashed_storage/base_repository_service.rb |  2 ++
 app/workers/hashed_storage/base_worker.rb     | 21 ++++++++++++
 .../hashed_storage/project_migrate_worker.rb  | 34 ++++---------------
 .../hashed_storage/project_rollback_worker.rb | 33 ++++--------------
 4 files changed, 35 insertions(+), 55 deletions(-)
 create mode 100644 app/workers/hashed_storage/base_worker.rb

diff --git a/app/services/projects/hashed_storage/base_repository_service.rb b/app/services/projects/hashed_storage/base_repository_service.rb
index 00a6af5238e..a682df35d4c 100644
--- a/app/services/projects/hashed_storage/base_repository_service.rb
+++ b/app/services/projects/hashed_storage/base_repository_service.rb
@@ -39,6 +39,8 @@ module Projects
         if !from_exists && !to_exists
           logger.warn "Can't find a repository on either source or target paths for #{project.full_path} (ID=#{project.id}) ..."
 
+          # We return true so we still reflect the change in the database.
+          # Next time the repository is (re)created it will be under the new storage layout
           return true
         elsif !from_exists
           # Repository have been moved already.
diff --git a/app/workers/hashed_storage/base_worker.rb b/app/workers/hashed_storage/base_worker.rb
new file mode 100644
index 00000000000..816e0504db6
--- /dev/null
+++ b/app/workers/hashed_storage/base_worker.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module HashedStorage
+  class BaseWorker
+    include ExclusiveLeaseGuard
+
+    LEASE_TIMEOUT = 30.seconds.to_i
+    LEASE_KEY_SEGMENT = 'project_migrate_hashed_storage_worker'.freeze
+
+    protected
+
+    def lease_key
+      # we share the same lease key for both migration and rollback so they don't run simultaneously
+      "#{LEASE_KEY_SEGMENT}:#{project_id}"
+    end
+
+    def lease_timeout
+      LEASE_TIMEOUT
+    end
+  end
+end
diff --git a/app/workers/hashed_storage/project_migrate_worker.rb b/app/workers/hashed_storage/project_migrate_worker.rb
index d516929129f..f00a459a097 100644
--- a/app/workers/hashed_storage/project_migrate_worker.rb
+++ b/app/workers/hashed_storage/project_migrate_worker.rb
@@ -1,48 +1,26 @@
 # frozen_string_literal: true
 
 module HashedStorage
-  class ProjectMigrateWorker
+  class ProjectMigrateWorker < BaseWorker
     include ApplicationWorker
 
-    LEASE_TIMEOUT = 30.seconds.to_i
-    LEASE_KEY_SEGMENT = 'project_migrate_hashed_storage_worker'.freeze
-
     queue_namespace :hashed_storage
 
+    attr_reader :project_id
+
     # rubocop: disable CodeReuse/ActiveRecord
     def perform(project_id, old_disk_path = nil)
-      uuid = lease_for(project_id).try_obtain
+      @project_id = project_id # we need to set this in order to create the lease_key
 
-      if uuid
+      try_obtain_lease do
         project = Project.without_deleted.find_by(id: project_id)
-        return unless project
+        break unless project
 
         old_disk_path ||= project.disk_path
 
         ::Projects::HashedStorage::MigrationService.new(project, old_disk_path, logger: logger).execute
-      else
-        return false
       end
-
-    ensure
-      cancel_lease_for(project_id, uuid) if uuid
     end
-
     # rubocop: enable CodeReuse/ActiveRecord
-
-    def lease_for(project_id)
-      Gitlab::ExclusiveLease.new(lease_key(project_id), timeout: LEASE_TIMEOUT)
-    end
-
-    private
-
-    def lease_key(project_id)
-      # we share the same lease key for both migration and rollback so they don't run simultaneously
-      "#{LEASE_KEY_SEGMENT}:#{project_id}"
-    end
-
-    def cancel_lease_for(project_id, uuid)
-      Gitlab::ExclusiveLease.cancel(lease_key(project_id), uuid)
-    end
   end
 end
diff --git a/app/workers/hashed_storage/project_rollback_worker.rb b/app/workers/hashed_storage/project_rollback_worker.rb
index c6ac76a1674..55e1d7ab23e 100644
--- a/app/workers/hashed_storage/project_rollback_worker.rb
+++ b/app/workers/hashed_storage/project_rollback_worker.rb
@@ -1,47 +1,26 @@
 # frozen_string_literal: true
 
 module HashedStorage
-  class ProjectRollbackWorker
+  class ProjectRollbackWorker < BaseWorker
     include ApplicationWorker
 
-    LEASE_TIMEOUT = 30.seconds.to_i
-
     queue_namespace :hashed_storage
 
+    attr_reader :project_id
+
     # rubocop: disable CodeReuse/ActiveRecord
     def perform(project_id, old_disk_path = nil)
-      uuid = lease_for(project_id).try_obtain
+      @project_id = project_id # we need to set this in order to create the lease_key
 
-      if uuid
+      try_obtain_lease do
         project = Project.without_deleted.find_by(id: project_id)
-        return unless project
+        break unless project
 
         old_disk_path ||= project.disk_path
 
         ::Projects::HashedStorage::RollbackService.new(project, old_disk_path, logger: logger).execute
-      else
-        return false
       end
-
-    ensure
-      cancel_lease_for(project_id, uuid) if uuid
     end
-
     # rubocop: enable CodeReuse/ActiveRecord
-
-    def lease_for(project_id)
-      Gitlab::ExclusiveLease.new(lease_key(project_id), timeout: LEASE_TIMEOUT)
-    end
-
-    private
-
-    def lease_key(project_id)
-      # we share the same lease key for both migration and rollback so they don't run simultaneously
-      "#{ProjectMigrateWorker::LEASE_KEY_SEGMENT}:#{project_id}"
-    end
-
-    def cancel_lease_for(project_id, uuid)
-      Gitlab::ExclusiveLease.cancel(lease_key(project_id), uuid)
-    end
   end
 end