Track ongoing pushes and reject mv-storage commands if there are push running (after waiting some time)mv-storage

3 files changed, 78 insertions, 0 deletions

diff --git a/spec/gitlab_post_receive_spec.rb b/spec/gitlab_post_receive_spec.rb
index 26b1037..6762d49 100644
--- a/spec/gitlab_post_receive_spec.rb
+++ b/spec/gitlab_post_receive_spec.rb
@@ -25,6 +25,11 @@ describe GitlabPostReceive do
 
     before do
       allow_any_instance_of(GitlabNet).to receive(:redis_client).and_return(redis_client)
+      allow_any_instance_of(GitlabReferenceCounter).to receive(:redis_client).and_return(redis_client)
+      allow(redis_client).to receive(:get).and_return(1)
+      allow(redis_client).to receive(:incr).and_return(true)
+      allow(redis_client).to receive(:decr).and_return(0)
+      allow(redis_client).to receive(:rpush).and_return(true)
     end
 
     it "prints the broadcast message" do
@@ -59,6 +64,34 @@ describe GitlabPostReceive do
       gitlab_post_receive.exec
     end
 
+    context 'reference counter' do
+      it 'decreases the reference counter for the project' do
+        expect_any_instance_of(GitlabReferenceCounter).to receive(:decrease).and_return(true)
+
+        gitlab_post_receive.exec
+      end
+
+      context "when the redis command succeeds" do
+        before do
+          allow(redis_client).to receive(:decr).and_return(0)
+        end
+
+        it "returns true" do
+          expect(gitlab_post_receive.exec).to eq(true)
+        end
+      end
+
+      context "when the redis command fails" do
+        before do
+          allow(redis_client).to receive(:decr).and_raise('Fail')
+        end
+
+        it "returns false" do
+          expect(gitlab_post_receive.exec).to eq(false)
+        end
+      end
+    end
+
     context "when the redis command succeeds" do
 
       before do
diff --git a/spec/gitlab_projects_spec.rb b/spec/gitlab_projects_spec.rb
index 5e50af0..19646e3 100644
--- a/spec/gitlab_projects_spec.rb
+++ b/spec/gitlab_projects_spec.rb
@@ -213,6 +213,7 @@ describe GitlabProjects do
     before do
       FileUtils.mkdir_p(tmp_repo_path)
       FileUtils.mkdir_p(alternative_storage_path)
+      allow_any_instance_of(GitlabReferenceCounter).to receive(:value).and_return(0)
     end
 
     after { FileUtils.rm_rf(alternative_storage_path) }
@@ -235,6 +236,12 @@ describe GitlabProjects do
       bad_source.exec.should be_false
     end
 
+    it "should fail if there are pushes ongoing" do
+      allow_any_instance_of(GitlabReferenceCounter).to receive(:value).and_return(1)
+      $logger.should_receive(:error).with("mv-storage failed: source path <#{tmp_repo_path}> is waiting for pushes to finish.")
+      gl_projects.exec.should be_false
+    end
+
     it "should log an mv-storage event" do
       message = "Syncing project #{repo_name} from <#{tmp_repo_path}> to <#{new_repo_path}>."
       $logger.should_receive(:info).with(message)
diff --git a/spec/gitlab_reference_counter_spec.rb b/spec/gitlab_reference_counter_spec.rb
new file mode 100644
index 0000000..5be53ff
--- /dev/null
+++ b/spec/gitlab_reference_counter_spec.rb
@@ -0,0 +1,38 @@
+# coding: utf-8
+require 'spec_helper'
+require 'gitlab_reference_counter'
+
+describe GitlabReferenceCounter do
+  let(:redis_client) { double('redis_client') }
+  let(:reference_counter_key) { "git-receive-pack-reference-counter:/test/path" }
+  let(:gitlab_reference_counter) { GitlabReferenceCounter.new('/test/path') }
+
+  before do
+    allow(gitlab_reference_counter).to receive(:redis_client).and_return(redis_client)
+    $logger = double('logger').as_null_object
+  end
+
+  it 'increases and set the expire time of a reference count for a path' do
+    expect(redis_client).to receive(:incr).with(reference_counter_key)
+    expect(redis_client).to receive(:expire).with(reference_counter_key, GitlabReferenceCounter::REFERENCE_EXPIRE_TIME)
+    expect(gitlab_reference_counter.increase).to be(true)
+  end
+
+  it 'decreases the reference count for a path' do
+    allow(redis_client).to receive(:decr).and_return(0)
+    expect(redis_client).to receive(:decr).with(reference_counter_key)
+    expect(gitlab_reference_counter.decrease).to be(true)
+  end
+
+  it 'warns if attempting to decrease a counter with a value of one or less, and resets the counter' do
+    expect(redis_client).to receive(:decr).and_return(-1)
+    expect(redis_client).to receive(:del)
+    expect($logger).to receive(:warn).with("Reference counter for /test/path decreased when its value was less than 1. Reseting the counter.")
+    expect(gitlab_reference_counter.decrease).to be(true)
+  end
+
+  it 'get the reference count for a path' do
+    allow(redis_client).to receive(:get).and_return(1)
+    expect(gitlab_reference_counter.value).to be(1)
+  end
+end