From ae27a47b6e5a27bb16d8f1555207784e34ec2ec1 Mon Sep 17 00:00:00 2001
From: Alessio Caiazza <acaiazza@gitlab.com>
Date: Wed, 28 Mar 2018 14:40:47 +0200
Subject: [PATCH] Add Gitlab::ExclusiveLease to ObjectStorage#use_file

---
 app/uploaders/object_storage.rb               | 45 ++++++++++++-------
 .../unreleased/ac-fix-use_file-race.yml       |  5 +++
 .../object_storage_shared_examples.rb         | 16 ++++++-
 spec/uploaders/object_storage_spec.rb         | 24 ++++++++++
 4 files changed, 72 insertions(+), 18 deletions(-)
 create mode 100644 changelogs/unreleased/ac-fix-use_file-race.yml

diff --git a/app/uploaders/object_storage.rb b/app/uploaders/object_storage.rb
index 30cc4425ae4..4028b052768 100644
--- a/app/uploaders/object_storage.rb
+++ b/app/uploaders/object_storage.rb
@@ -228,16 +228,9 @@ module ObjectStorage
       raise 'Failed to update object store' unless updated
     end
 
-    def use_file
-      if file_storage?
-        return yield path
-      end
-
-      begin
-        cache_stored_file!
-        yield cache_path
-      ensure
-        cache_storage.delete_dir!(cache_path(nil))
+    def use_file(&blk)
+      with_exclusive_lease do
+        unsafe_use_file(&blk)
       end
     end
 
@@ -247,12 +240,9 @@ module ObjectStorage
     #   new_store: Enum (Store::LOCAL, Store::REMOTE)
     #
     def migrate!(new_store)
-      uuid = Gitlab::ExclusiveLease.new(exclusive_lease_key, timeout: 1.hour.to_i).try_obtain
-      raise 'Already running' unless uuid
-
-      unsafe_migrate!(new_store)
-    ensure
-      Gitlab::ExclusiveLease.cancel(exclusive_lease_key, uuid)
+      with_exclusive_lease do
+        unsafe_migrate!(new_store)
+      end
     end
 
     def schedule_background_upload(*args)
@@ -384,6 +374,15 @@ module ObjectStorage
       "object_storage_migrate:#{model.class}:#{model.id}"
     end
 
+    def with_exclusive_lease
+      uuid = Gitlab::ExclusiveLease.new(exclusive_lease_key, timeout: 1.hour.to_i).try_obtain
+      raise 'exclusive lease already taken' unless uuid
+
+      yield uuid
+    ensure
+      Gitlab::ExclusiveLease.cancel(exclusive_lease_key, uuid)
+    end
+
     #
     # Move the file to another store
     #
@@ -418,4 +417,18 @@ module ObjectStorage
       raise e
     end
   end
+
+  def unsafe_use_file
+    if file_storage?
+      return yield path
+    end
+
+    begin
+      cache_stored_file!
+      yield cache_path
+    ensure
+      FileUtils.rm_f(cache_path)
+      cache_storage.delete_dir!(cache_path(nil))
+    end
+  end
 end
diff --git a/changelogs/unreleased/ac-fix-use_file-race.yml b/changelogs/unreleased/ac-fix-use_file-race.yml
new file mode 100644
index 00000000000..f1315d5d50e
--- /dev/null
+++ b/changelogs/unreleased/ac-fix-use_file-race.yml
@@ -0,0 +1,5 @@
+---
+title: Fix data race between ObjectStorage background_upload and Pages publishing
+merge_request:
+author:
+type: fixed
diff --git a/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb b/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb
index cd9974cd6e2..6352f1527cd 100644
--- a/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb
+++ b/spec/support/shared_examples/uploaders/object_storage_shared_examples.rb
@@ -61,12 +61,18 @@ shared_examples "migrates" do |to_store:, from_store: nil|
     expect { migrate(to) }.not_to change { file.exists? }
   end
 
-  context 'when migrate! is not oqqupied by another process' do
+  context 'when migrate! is not occupied by another process' do
     it 'executes migrate!' do
       expect(subject).to receive(:object_store=).at_least(1)
 
       migrate(to)
     end
+
+    it 'executes use_file' do
+      expect(subject).to receive(:unsafe_use_file).once
+
+      subject.use_file
+    end
   end
 
   context 'when migrate! is occupied by another process' do
@@ -79,7 +85,13 @@ shared_examples "migrates" do |to_store:, from_store: nil|
     it 'does not execute migrate!' do
       expect(subject).not_to receive(:unsafe_migrate!)
 
-      expect { migrate(to) }.to raise_error('Already running')
+      expect { migrate(to) }.to raise_error('exclusive lease already taken')
+    end
+
+    it 'does not execute use_file' do
+      expect(subject).not_to receive(:unsafe_use_file)
+
+      expect { subject.use_file }.to raise_error('exclusive lease already taken')
     end
 
     after do
diff --git a/spec/uploaders/object_storage_spec.rb b/spec/uploaders/object_storage_spec.rb
index 1d406c71955..59e02fecbce 100644
--- a/spec/uploaders/object_storage_spec.rb
+++ b/spec/uploaders/object_storage_spec.rb
@@ -308,6 +308,30 @@ describe ObjectStorage do
     it { is_expected.to eq(remote_directory) }
   end
 
+  context 'when file is in use' do
+    def when_file_is_in_use
+      uploader.use_file do
+        yield
+      end
+    end
+
+    it 'cannot migrate' do
+      when_file_is_in_use do
+        expect(uploader).not_to receive(:unsafe_migrate!)
+
+        expect { uploader.migrate!(described_class::Store::REMOTE) }.to raise_error('exclusive lease already taken')
+      end
+    end
+
+    it 'cannot use_file' do
+      when_file_is_in_use do
+        expect(uploader).not_to receive(:unsafe_use_file)
+
+        expect { uploader.use_file }.to raise_error('exclusive lease already taken')
+      end
+    end
+  end
+
   describe '#fog_credentials' do
     let(:connection) { Settingslogic.new("provider" => "AWS") }