1 files changed, 20 insertions, 13 deletions
diff --git a/lib/gitlab/gitaly_client/cleanup_service.rb b/lib/gitlab/gitaly_client/cleanup_service.rb
index 3e8d6a773ca..a56bc35f6d7 100644
--- a/lib/gitlab/gitaly_client/cleanup_service.rb
+++ b/lib/gitlab/gitaly_client/cleanup_service.rb
@@ -12,25 +12,32 @@ module Gitlab
         @storage = repository.storage
       end
 
-      def apply_bfg_object_map(io)
-        first_request = Gitaly::ApplyBfgObjectMapRequest.new(repository: gitaly_repo)
+      def apply_bfg_object_map_stream(io, &blk)
+        responses = GitalyClient.call(
+          storage,
+          :cleanup_service,
+          :apply_bfg_object_map_stream,
+          build_object_map_enum(io),
+          timeout: GitalyClient.no_timeout
+        )
+
+        responses.each(&blk)
+      end
+
+      private
 
-        enum = Enumerator.new do |y|
-          y.yield first_request
+      def build_object_map_enum(io)
+        Enumerator.new do |y|
+          # First request. For simplicity, doesn't include any object map data
+          y << Gitaly::ApplyBfgObjectMapStreamRequest.new(repository: gitaly_repo)
 
+          # Now stream the BFG object map file to gitaly in chunks
           while data = io.read(RepositoryService::MAX_MSG_SIZE)
-            y.yield Gitaly::ApplyBfgObjectMapRequest.new(object_map: data)
+            y << Gitaly::ApplyBfgObjectMapStreamRequest.new(object_map: data)
+
             break if io&.eof?
           end
         end
-
-        GitalyClient.call(
-          storage,
-          :cleanup_service,
-          :apply_bfg_object_map,
-          enum,
-          timeout: GitalyClient.no_timeout
-        )
       end
     end
   end