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

3 files changed, 76 insertions, 3 deletions

diff --git a/lib/gitlab_post_receive.rb b/lib/gitlab_post_receive.rb
index ebcf317..7874c85 100644
--- a/lib/gitlab_post_receive.rb
+++ b/lib/gitlab_post_receive.rb
@@ -1,5 +1,6 @@
 require_relative 'gitlab_init'
 require_relative 'gitlab_net'
+require_relative 'gitlab_reference_counter'
 require 'json'
 require 'base64'
 require 'securerandom'
@@ -28,7 +29,7 @@ class GitlabPostReceive
       nil
     end
 
-    result
+    result && GitlabReferenceCounter.new(repo_path).decrease
   end
 
   protected
diff --git a/lib/gitlab_projects.rb b/lib/gitlab_projects.rb
index 0f643b0..a40cc4f 100644
--- a/lib/gitlab_projects.rb
+++ b/lib/gitlab_projects.rb
@@ -4,6 +4,7 @@ require 'open3'
 
 require_relative 'gitlab_config'
 require_relative 'gitlab_logger'
+require_relative 'gitlab_reference_counter'
 
 class GitlabProjects
   GLOBAL_HOOKS_DIRECTORY = File.join(ROOT_PATH, 'hooks')
@@ -313,8 +314,13 @@ class GitlabProjects
     # contents of the directory, as opposed to copying the directory by name
     source_path = File.join(full_path, '')
 
-    $logger.info "Syncing project #{@project_name} from <#{full_path}> to <#{new_full_path}>."
-    system(*%W(rsync -a #{source_path} #{new_full_path}))
+    if wait_for_pushes
+      $logger.info "Syncing project #{@project_name} from <#{full_path}> to <#{new_full_path}>."
+      system(*%W(rsync -a --delete #{source_path} #{new_full_path}))
+    else
+      $logger.error "mv-storage failed: source path <#{full_path}> is waiting for pushes to finish."
+      false
+    end
   end
 
   def fork_project
@@ -361,4 +367,18 @@ class GitlabProjects
     cmd = %W(git --git-dir=#{full_path} gc)
     system(*cmd)
   end
+
+  def wait_for_pushes
+    # Try for 30 seconds, polling every 10
+    3.times do
+      return true if gitlab_reference_counter.value == 0
+      sleep 10
+    end
+
+    false
+  end
+
+  def gitlab_reference_counter
+    @gitlab_reference_counter ||= GitlabReferenceCounter.new(full_path)
+  end
 end
diff --git a/lib/gitlab_reference_counter.rb b/lib/gitlab_reference_counter.rb
new file mode 100644
index 0000000..dc08f42
--- /dev/null
+++ b/lib/gitlab_reference_counter.rb
@@ -0,0 +1,52 @@
+require_relative 'gitlab_init'
+require_relative 'gitlab_net'
+
+class GitlabReferenceCounter
+  REFERENCE_EXPIRE_TIME = 600
+
+  attr_reader :path, :key
+
+  def initialize(path)
+    @path = path
+    @key = "git-receive-pack-reference-counter:#{path}"
+  end
+
+  def value
+    (redis_client.get(key) || 0).to_i
+  end
+
+  def increase
+    redis_cmd do
+      redis_client.incr(key)
+      redis_client.expire(key, REFERENCE_EXPIRE_TIME)
+    end
+  end
+
+  def decrease
+    redis_cmd do
+      current_value = redis_client.decr(key)
+      if current_value < 0
+        $logger.warn "Reference counter for #{path} decreased when its value was less than 1. Reseting the counter."
+        redis_client.del(key)
+      end
+    end
+  end
+
+  private
+
+  def redis_client
+    @redis_client ||= GitlabNet.new.redis_client
+  end
+
+  def redis_cmd
+    begin
+      yield
+      true
+    rescue => e
+      message = "GitLab: An unexpected error occurred in writing to Redis: #{e}"
+      $stderr.puts message
+      $logger.error message
+      false
+    end
+  end
+end