Cache diff highlighting upon Merge Request creation (refactors diff caching)

3 files changed, 6 insertions, 2 deletions

diff --git a/app/controllers/projects/merge_requests/diffs_controller.rb b/app/controllers/projects/merge_requests/diffs_controller.rb
index 48e02581d54..34de554212f 100644
--- a/app/controllers/projects/merge_requests/diffs_controller.rb
+++ b/app/controllers/projects/merge_requests/diffs_controller.rb
@@ -21,6 +21,8 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
   def render_diffs
     @environment = @merge_request.environments_for(current_user).last
 
+    @diffs.write_cache
+
     render json: DiffsSerializer.new(current_user: current_user).represent(@diffs, additional_attributes)
   end
 
diff --git a/app/services/merge_requests/reload_diffs_service.rb b/app/services/merge_requests/reload_diffs_service.rb
index 8d85dc9eb5f..1390ae0e199 100644
--- a/app/services/merge_requests/reload_diffs_service.rb
+++ b/app/services/merge_requests/reload_diffs_service.rb
@@ -30,7 +30,7 @@ module MergeRequests
     def clear_cache(new_diff)
       # Executing the iteration we cache highlighted diffs for each diff file of
       # MergeRequestDiff.
-      new_diff.diffs_collection.diff_files.to_a
+      new_diff.diffs_collection.write_cache
 
       # Remove cache for all diffs on this MR. Do not use the association on the
       # model, as that will interfere with other actions happening when
@@ -38,7 +38,7 @@ module MergeRequests
       MergeRequestDiff.where(merge_request: merge_request).each do |merge_request_diff|
         next if merge_request_diff == new_diff
 
-        merge_request_diff.diffs_collection.clear_cache!
+        merge_request_diff.diffs_collection.clear_cache
       end
     end
   end
diff --git a/app/workers/new_merge_request_worker.rb b/app/workers/new_merge_request_worker.rb
index 5d8b8904502..62f9d9b6f57 100644
--- a/app/workers/new_merge_request_worker.rb
+++ b/app/workers/new_merge_request_worker.rb
@@ -9,6 +9,8 @@ class NewMergeRequestWorker
 
     EventCreateService.new.open_mr(issuable, user)
     NotificationService.new.new_merge_request(issuable, user)
+
+    issuable.diffs.write_cache
     issuable.create_cross_references!(user)
   end