Use the Gitaly Snapshot RPCs in Geo synchronization

ee/app/models/geo_node.rb

@@ -105,6 +105,13 @@ class GeoNode < ActiveRecord::Base
     geo_api_url('status')
   end
 
+  def snapshot_url(repository)
+    url = api_url("projects/#{repository.project.id}/snapshot")
+    url += "?wiki=1" if repository.is_wiki
+
+    url
+  end
+
   def oauth_callback_url
     Gitlab::Routing.url_helpers.oauth_geo_callback_url(url_helper_args)
   end
@@ -185,7 +192,11 @@ class GeoNode < ActiveRecord::Base
   private
 
   def geo_api_url(suffix)
-    URI.join(uri, "#{uri.path}", "api/#{API::API.version}/geo/#{suffix}").to_s
+    api_url("geo/#{suffix}")
+  end
+
+  def api_url(suffix)
+    URI.join(uri, "#{uri.path}", "api/#{API::API.version}/#{suffix}").to_s
   end
 
   def ensure_access_keys!

ee/app/services/geo/base_sync_service.rb

@@ -59,7 +59,22 @@ module Geo
 
       if redownload
         log_info("Redownloading #{type}")
-        fetch_geo_mirror(build_temporary_repository)
+        temp_repo = build_temporary_repository
+
+        if Feature.enabled?(:geo_redownload_with_snapshot)
+          begin
+            fetch_snapshot(temp_repo)
+          rescue => err
+            log_error('Snapshot attempt failed', err)
+          end
+        end
+
+        # A git fetch should be attempted, regardless of whether a snapshot was
+        # performed. A git fetch *may* succeed where a snapshot has failed. If
+        # the snapshot succeeded, it may not be in a consistent state - the git
+        # git fetch may repair it.
+        fetch_geo_mirror(temp_repo)
+
         set_temp_repository_as_main
       else
         ensure_repository
@@ -99,6 +114,13 @@ module Geo
       end
     end
 
+    def fetch_snapshot(repository)
+      repository.create_from_snapshot(
+        ::Gitlab::Geo.primary_node.snapshot_url(repository),
+        ::Gitlab::Geo::RepoSyncRequest.new.authorization
+      )
+    end
+
     def registry
       @registry ||= Geo::ProjectRegistry.find_or_initialize_by(project_id: project.id)
     end

ee/lib/api/geo.rb

@@ -38,23 +38,5 @@ module API
         present status, with: EE::API::Entities::GeoNodeStatus
       end
     end
-
-    helpers do
-      def authenticate_by_gitlab_geo_node_token!
-        auth_header = headers['Authorization']
-
-        begin
-          unless auth_header && Gitlab::Geo::JwtRequestDecoder.new(auth_header).decode
-            unauthorized!
-          end
-        rescue Gitlab::Geo::InvalidDecryptionKeyError, Gitlab::Geo::SignatureTimeInvalidError => e
-          render_api_error!(e.to_s, 401)
-        end
-      end
-
-      def require_node_to_be_enabled!
-        forbidden! 'Geo node is disabled.' unless Gitlab::Geo.current_node&.enabled?
-      end
-    end
   end
 end

ee/lib/ee/api/helpers.rb

@@ -3,6 +3,26 @@ module EE
     module Helpers
       extend ::Gitlab::Utils::Override
 
+      def require_node_to_be_enabled!
+        forbidden! 'Geo node is disabled.' unless ::Gitlab::Geo.current_node&.enabled?
+      end
+
+      def gitlab_geo_node_token?
+        headers['Authorization']&.start_with?(::Gitlab::Geo::BaseRequest::GITLAB_GEO_AUTH_TOKEN_TYPE)
+      end
+
+      def authenticate_by_gitlab_geo_node_token!
+        auth_header = headers['Authorization']
+
+        begin
+          unless auth_header && ::Gitlab::Geo::JwtRequestDecoder.new(auth_header).decode
+            unauthorized!
+          end
+        rescue ::Gitlab::Geo::InvalidDecryptionKeyError, ::Gitlab::Geo::SignatureTimeInvalidError => e
+          render_api_error!(e.to_s, 401)
+        end
+      end
+
       override :current_user
       def current_user
         strong_memoize(:current_user) do

ee/lib/ee/api/helpers/project_snapshots_helpers.rb

+module EE
+  module API
+    module Helpers
+      module ProjectSnapshotsHelpers
+        extend ::Gitlab::Utils::Override
+
+        # Allow Geo nodes to access snapshots by presenting a valid JWT
+        override :authorize_read_git_snapshot!
+        def authorize_read_git_snapshot!
+          if gitlab_geo_node_token?
+            require_node_to_be_enabled!
+            authenticate_by_gitlab_geo_node_token!
+          else
+            super
+          end
+        end
+
+        # Skip checking authorization of current_user if authenticated via Geo
+        override :snapshot_project
+        def snapshot_project
+          if gitlab_geo_node_token?
+            project = find_project(params[:id])
+            not_found!('Project') if project.nil?
+
+            project
+          else
+            super
+          end
+        end
+      end
+    end
+  end
+end

ee/spec/models/geo_node_spec.rb

@@ -200,6 +200,19 @@ describe GeoNode, type: :model do
     end
   end
 
+  describe '#snapshot_url' do
+    let(:project) { create(:project) }
+    let(:snapshot_url) { "https://localhost:3000/gitlab/api/#{api_version}/projects/#{project.id}/snapshot" }
+
+    it 'returns snapshot URL based on node URI' do
+      expect(new_node.snapshot_url(project.repository)).to eq(snapshot_url)
+    end
+
+    it 'adds ?wiki=1 to the snapshot URL when the repository is a wiki' do
+      expect(new_node.snapshot_url(project.wiki.repository)).to eq(snapshot_url + "?wiki=1")
+    end
+  end
+
   describe '#find_or_build_status' do
     it 'returns a new status' do
       status = new_node.find_or_build_status

ee/spec/requests/api/project_snapshots_spec.rb

+require 'spec_helper'
+
+describe API::ProjectSnapshots do
+  include ::EE::GeoHelpers
+
+  let(:project) { create(:project) }
+
+  describe 'GET /projects/:id/snapshot' do
+    let(:primary) { create(:geo_node, :primary) }
+    let(:secondary) { create(:geo_node) }
+
+    before do
+      stub_current_geo_node(primary)
+    end
+
+    it 'requests project repository raw archive from Geo primary as Geo secondary' do
+      req = Gitlab::Geo::BaseRequest.new
+      allow(req).to receive(:requesting_node) { secondary }
+
+      get api("/projects/#{project.id}/snapshot", nil), {}, req.headers
+
+      expect(response).to have_gitlab_http_status(200)
+    end
+  end
+end

ee/spec/services/geo/repository_sync_service_spec.rb

@@ -335,5 +335,9 @@ describe Geo::RepositorySyncService do
         end
       end
     end
+
+    it_behaves_like 'sync retries use the snapshot RPC' do
+      let(:repository) { project.repository }
+    end
   end
 end

ee/spec/services/geo/wiki_sync_service_spec.rb

@@ -200,5 +200,9 @@ RSpec.describe Geo::WikiSyncService do
         end
       end
     end
+
+    it_behaves_like 'sync retries use the snapshot RPC' do
+      let(:repository) { project.wiki.repository }
+    end
   end
 end

ee/spec/support/services/base_sync_service_shared_examples.rb

@@ -47,3 +47,60 @@ shared_examples 'cleans temporary repositories' do
     end
   end
 end
+
+shared_examples 'sync retries use the snapshot RPC' do
+  let(:retry_count) { Geo::BaseSyncService::RETRY_BEFORE_REDOWNLOAD }
+
+  context 'snapshot synchronization method' do
+    before do
+      allow(subject).to receive(:build_temporary_repository) { repository }
+    end
+
+    def receive_create_from_snapshot
+      receive(:create_from_snapshot).with(primary.snapshot_url(repository), match(/^GL-Geo/))
+    end
+
+    it 'does not attempt to snapshot for initial sync' do
+      expect(repository).not_to receive_create_from_snapshot
+      expect(subject).to receive(:fetch_geo_mirror).with(repository)
+
+      subject.execute
+    end
+
+    it 'does not attempt to snapshot for ordinary retries' do
+      create(:geo_project_registry, project: project, repository_retry_count: retry_count - 1, wiki_retry_count: retry_count - 1)
+
+      expect(repository).not_to receive_create_from_snapshot
+      expect(subject).to receive(:fetch_geo_mirror).with(repository)
+
+      subject.execute
+    end
+
+    context 'registry is ready to be snapshotted' do
+      let!(:registry) { create(:geo_project_registry, project: project, repository_retry_count: retry_count + 1, wiki_retry_count: retry_count + 1) }
+
+      it 'attempts to snapshot + fetch' do
+        expect(repository).to receive_create_from_snapshot
+        expect(subject).to receive(:fetch_geo_mirror).with(repository)
+
+        subject.execute
+      end
+
+      it 'attempts to fetch if snapshotting raises an exception' do
+        expect(repository).to receive_create_from_snapshot.and_raise(ArgumentError)
+        expect(subject).to receive(:fetch_geo_mirror).with(repository)
+
+        subject.execute
+      end
+
+      it 'does not attempt to snapshot if the feature flag is disabled' do
+        stub_feature_flags(geo_redownload_with_snapshot: false)
+
+        expect(repository).not_to receive_create_from_snapshot
+        expect(subject).to receive(:fetch_geo_mirror).with(repository)
+
+        subject.execute
+      end
+    end
+  end
+end

lib/api/helpers/project_snapshots_helpers.rb

 module API
   module Helpers
     module ProjectSnapshotsHelpers
+      prepend ::EE::API::Helpers::ProjectSnapshotsHelpers
+
       def authorize_read_git_snapshot!
         authenticated_with_full_private_access!
       end