Persist truncated note diffs on a new table

We request Gitaly in a N+1 manner to build discussion diffs. Once the diffs are from different revisions, it's hard to make a single request to the service in order to build the whole response.
With this change we solve this problem and simplify a lot fetching this piece of info.

app/models/concerns/diff_file.rb

+module DiffFile
+  extend ActiveSupport::Concern
+
+  def to_hash
+    keys = Gitlab::Git::Diff::SERIALIZE_KEYS - [:diff]
+
+    as_json(only: keys).merge(diff: diff).with_indifferent_access
+  end
+end

app/models/diff_note.rb

@@ -3,6 +3,7 @@
 # A note of this type can be resolvable.
 class DiffNote < Note
   include NoteOnDiff
+  include Gitlab::Utils::StrongMemoize
 
   NOTEABLE_TYPES = %w(MergeRequest Commit).freeze
 
@@ -12,7 +13,6 @@ class DiffNote < Note
 
   validates :original_position, presence: true
   validates :position, presence: true
-  validates :diff_line, presence: true, if: :on_text?
   validates :line_code, presence: true, line_code: true, if: :on_text?
   validates :noteable_type, inclusion: { in: NOTEABLE_TYPES }
   validate :positions_complete
@@ -23,6 +23,7 @@ class DiffNote < Note
   before_validation :update_position, on: :create, if: :on_text?
   before_validation :set_line_code, if: :on_text?
   after_save :keep_around_commits
+  after_commit :create_diff_file, on: :create
 
   def discussion_class(*)
     DiffDiscussion
@@ -53,21 +54,25 @@ class DiffNote < Note
     position.position_type == "image"
   end
 
+  def create_diff_file
+    return unless should_create_diff_file?
+
+    diff_file = fetch_diff_file
+    diff_line = diff_file.line_for_position(self.original_position)
+
+    creation_params = diff_file.diff.to_hash
+      .except(:too_large)
+      .merge(diff: diff_file.diff_hunk(diff_line))
+
+    create_note_diff_file(creation_params)
+  end
+
   def diff_file
-    @diff_file ||=
-      begin
-        if created_at_diff?(noteable.diff_refs)
-          # We're able to use the already persisted diffs (Postgres) if we're
-          # presenting a "current version" of the MR discussion diff.
-          # So no need to make an extra Gitaly diff request for it.
-          # As an extra benefit, the returned `diff_file` already
-          # has `highlighted_diff_lines` data set from Redis on
-          # `Diff::FileCollection::MergeRequestDiff`.
-          noteable.diffs(paths: original_position.paths, expanded: true).diff_files.first
-        else
-          original_position.diff_file(self.project.repository)
-        end
-      end
+    strong_memoize(:diff_file) do
+      enqueue_diff_file_creation_job if should_create_diff_file?
+
+      fetch_diff_file
+    end
   end
 
   def diff_line
@@ -98,6 +103,38 @@ class DiffNote < Note
 
   private
 
+  def enqueue_diff_file_creation_job
+    # Avoid enqueuing multiple file creation jobs at once for a note (i.e.
+    # parallel calls to `DiffNote#diff_file`).
+    lease = Gitlab::ExclusiveLease.new("note_diff_file_creation:#{id}", timeout: 1.hour.to_i)
+    return unless lease.try_obtain
+
+    CreateNoteDiffFileWorker.perform_async(id)
+  end
+
+  def should_create_diff_file?
+    on_text? && note_diff_file.nil? && self == discussion.first_note
+  end
+
+  def fetch_diff_file
+    if note_diff_file
+      diff = Gitlab::Git::Diff.new(note_diff_file.to_hash)
+      Gitlab::Diff::File.new(diff,
+                             repository: project.repository,
+                             diff_refs: original_position.diff_refs)
+    elsif created_at_diff?(noteable.diff_refs)
+      # We're able to use the already persisted diffs (Postgres) if we're
+      # presenting a "current version" of the MR discussion diff.
+      # So no need to make an extra Gitaly diff request for it.
+      # As an extra benefit, the returned `diff_file` already
+      # has `highlighted_diff_lines` data set from Redis on
+      # `Diff::FileCollection::MergeRequestDiff`.
+      noteable.diffs(paths: original_position.paths, expanded: true).diff_files.first
+    else
+      original_position.diff_file(self.project.repository)
+    end
+  end
+
   def supported?
     for_commit? || self.noteable.has_complete_diff_refs?
   end

app/models/merge_request_diff_file.rb

 class MergeRequestDiffFile < ActiveRecord::Base
   include Gitlab::EncodingHelper
+  include DiffFile
 
   belongs_to :merge_request_diff
 
@@ -12,10 +13,4 @@ class MergeRequestDiffFile < ActiveRecord::Base
   def diff
     binary? ? super.unpack('m0').first : super
   end
-
-  def to_hash
-    keys = Gitlab::Git::Diff::SERIALIZE_KEYS - [:diff]
-
-    as_json(only: keys).merge(diff: diff).with_indifferent_access
-  end
 end

app/models/note.rb

@@ -63,6 +63,7 @@ class Note < ActiveRecord::Base
   has_many :todos
   has_many :events, as: :target, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
   has_one :system_note_metadata
+  has_one :note_diff_file, inverse_of: :diff_note, foreign_key: :diff_note_id
 
   delegate :gfm_reference, :local_reference, to: :noteable
   delegate :name, to: :project, prefix: true
@@ -100,7 +101,8 @@ class Note < ActiveRecord::Base
   scope :inc_author_project, -> { includes(:project, :author) }
   scope :inc_author, -> { includes(:author) }
   scope :inc_relations_for_view, -> do
-    includes(:project, :author, :updated_by, :resolved_by, :award_emoji, :system_note_metadata)
+    includes(:project, :author, :updated_by, :resolved_by, :award_emoji,
+             :system_note_metadata, :note_diff_file)
   end
 
   scope :diff_notes, -> { where(type: %w(LegacyDiffNote DiffNote)) }

app/models/note_diff_file.rb

+class NoteDiffFile < ActiveRecord::Base
+  include DiffFile
+
+  belongs_to :diff_note, inverse_of: :note_diff_file
+
+  validates :diff_note, presence: true
+end

app/workers/all_queues.yml

@@ -115,3 +115,4 @@
 - upload_checksum
 - web_hook
 - repository_update_remote_mirror
+- create_note_diff_file

app/workers/create_note_diff_file_worker.rb

+class CreateNoteDiffFileWorker
+  include ApplicationWorker
+
+  def perform(diff_note_id)
+    diff_note = DiffNote.find(diff_note_id)
+
+    diff_note.create_diff_file
+  end
+end

changelogs/unreleased/45190-create-notes-diff-files.yml

+---
+title: Persist truncated note diffs on a new table
+merge_request:
+author:
+type: performance

config/sidekiq_queues.yml

@@ -75,4 +75,5 @@
   - [pipeline_background, 1]
   - [repository_update_remote_mirror, 1]
   - [repository_remove_remote, 1]
+  - [create_note_diff_file, 1]
  

db/migrate/20180515121227_create_notes_diff_files.rb

+class CreateNotesDiffFiles < ActiveRecord::Migration
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  def change
+    create_table :note_diff_files do |t|
+      t.references :diff_note, references: :notes, null: false, index: { unique: true }
+      t.text :diff, null: false
+      t.boolean :new_file, null: false
+      t.boolean :renamed_file, null: false
+      t.boolean :deleted_file, null: false
+      t.string :a_mode, null: false
+      t.string :b_mode, null: false
+      t.text :new_path, null: false
+      t.text :old_path, null: false
+    end
+
+    add_foreign_key :note_diff_files, :notes, column: :diff_note_id, on_delete: :cascade
+  end
+end

db/schema.rb

@@ -1302,6 +1302,20 @@ ActiveRecord::Schema.define(version: 20180521171529) do
   add_index "namespaces", ["runners_token"], name: "index_namespaces_on_runners_token", unique: true, using: :btree
   add_index "namespaces", ["type"], name: "index_namespaces_on_type", using: :btree
 
+  create_table "note_diff_files", force: :cascade do |t|
+    t.integer "diff_note_id", null: false
+    t.text "diff", null: false
+    t.boolean "new_file", null: false
+    t.boolean "renamed_file", null: false
+    t.boolean "deleted_file", null: false
+    t.string "a_mode", null: false
+    t.string "b_mode", null: false
+    t.text "new_path", null: false
+    t.text "old_path", null: false
+  end
+
+  add_index "note_diff_files", ["diff_note_id"], name: "index_note_diff_files_on_diff_note_id", unique: true, using: :btree
+
   create_table "notes", force: :cascade do |t|
     t.text "note"
     t.string "noteable_type"
@@ -2243,6 +2257,7 @@ ActiveRecord::Schema.define(version: 20180521171529) do
   add_foreign_key "merge_requests_closing_issues", "merge_requests", on_delete: :cascade
   add_foreign_key "milestones", "namespaces", column: "group_id", name: "fk_95650a40d4", on_delete: :cascade
   add_foreign_key "milestones", "projects", name: "fk_9bd0a0c791", on_delete: :cascade
+  add_foreign_key "note_diff_files", "notes", column: "diff_note_id", on_delete: :cascade
   add_foreign_key "notes", "projects", name: "fk_99e097b079", on_delete: :cascade
   add_foreign_key "oauth_openid_requests", "oauth_access_grants", column: "access_grant_id", name: "fk_oauth_openid_requests_oauth_access_grants_access_grant_id"
   add_foreign_key "pages_domains", "projects", name: "fk_ea2f6dfc6f", on_delete: :cascade

lib/gitlab/diff/file.rb

@@ -76,6 +76,13 @@ module Gitlab
         line_code(line) if line
       end
 
+      # Returns the raw diff content up to the given line index
+      def diff_hunk(diff_line)
+        # Adding 2 because of the @@ diff header and Enum#take should consider
+        # an extra line, because we're passing an index.
+        raw_diff.each_line.take(diff_line.index + 2).join
+      end
+
       def old_sha
         diff_refs&.base_sha
       end

spec/lib/gitlab/diff/file_spec.rb

@@ -468,4 +468,58 @@ describe Gitlab::Diff::File do
       end
     end
   end
+
+  describe '#diff_hunk' do
+    let(:raw_diff) do
+      <<EOS
+@@ -6,12 +6,18 @@ module Popen
+
+   def popen(cmd, path=nil)
+     unless cmd.is_a?(Array)
+-      raise "System commands must be given as an array of strings"
++      raise RuntimeError, "System commands must be given as an array of strings"
+     end
+
+     path ||= Dir.pwd
+-    vars = { "PWD" => path }
+-    options = { chdir: path }
++
++    vars = {
++      "PWD" => path
++    }
++
++    options = {
++      chdir: path
++    }
+
+     unless File.directory?(path)
+       FileUtils.mkdir_p(path)
+@@ -19,6 +25,7 @@ module Popen
+
+     @cmd_output = ""
+     @cmd_status = 0
++
+     Open3.popen3(vars, *cmd, options) do |stdin, stdout, stderr, wait_thr|
+       @cmd_output << stdout.read
+       @cmd_output << stderr.read
+EOS
+    end
+
+    it 'returns raw diff up to given line index' do
+      allow(diff_file).to receive(:raw_diff) { raw_diff }
+      diff_line = instance_double(Gitlab::Diff::Line, index: 5)
+
+      diff_hunk = <<EOS
+@@ -6,12 +6,18 @@ module Popen
+
+   def popen(cmd, path=nil)
+     unless cmd.is_a?(Array)
+-      raise "System commands must be given as an array of strings"
++      raise RuntimeError, "System commands must be given as an array of strings"
+     end
+EOS
+
+      expect(diff_file.diff_hunk(diff_line)).to eq(diff_hunk)
+    end
+  end
 end

spec/lib/gitlab/import_export/all_models.yml

@@ -35,6 +35,7 @@ notes:
 - todos
 - events
 - system_note_metadata
+- note_diff_file
 label_links:
 - target
 - label

spec/models/concerns/discussion_on_diff_spec.rb

 require 'spec_helper'
 
 describe DiscussionOnDiff do
-  subject { create(:diff_note_on_merge_request).to_discussion }
+  subject { create(:diff_note_on_merge_request, line_number: 18).to_discussion }
 
   describe "#truncated_diff_lines" do
     let(:truncated_lines) { subject.truncated_diff_lines }

spec/models/diff_note_spec.rb

@@ -84,7 +84,47 @@ describe DiffNote do
     end
   end
 
-  describe "#diff_file" do
+  describe '#create_diff_file callback' do
+    let(:noteable) { create(:merge_request) }
+    let(:project) { noteable.project }
+
+    context 'merge request' do
+      let!(:diff_note) { create(:diff_note_on_merge_request, project: project, noteable: noteable) }
+
+      it 'creates a diff note file' do
+        expect(diff_note.reload.note_diff_file).to be_present
+      end
+
+      it 'does not create diff note file if it is a reply' do
+        expect { create(:diff_note_on_merge_request, noteable: noteable, in_reply_to: diff_note) }
+          .not_to change(NoteDiffFile, :count)
+      end
+    end
+
+    context 'commit' do
+      let!(:diff_note) { create(:diff_note_on_commit, project: project) }
+
+      it 'creates a diff note file' do
+        expect(diff_note.reload.note_diff_file).to be_present
+      end
+
+      it 'does not create diff note file if it is a reply' do
+        expect { create(:diff_note_on_commit, in_reply_to: diff_note) }
+          .not_to change(NoteDiffFile, :count)
+      end
+    end
+  end
+
+  describe '#diff_file', :clean_gitlab_redis_shared_state do
+    context 'when note_diff_file association exists' do
+      it 'returns persisted diff file data' do
+        diff_file = subject.diff_file
+
+        expect(diff_file.diff.to_hash.with_indifferent_access)
+          .to include(subject.note_diff_file.to_hash)
+      end
+    end
+
     context 'when the discussion was created in the diff' do
       it 'returns correct diff file' do
         diff_file = subject.diff_file
@@ -115,6 +155,26 @@ describe DiffNote do
         expect(diff_file.diff_refs).to eq(position.diff_refs)
       end
     end
+
+    context 'note diff file creation enqueuing' do
+      it 'enqueues CreateNoteDiffFileWorker if it is the first note of a discussion' do
+        subject.note_diff_file.destroy!
+
+        expect(CreateNoteDiffFileWorker).to receive(:perform_async).with(subject.id)
+
+        subject.reload.diff_file
+      end
+
+      it 'does not enqueues CreateNoteDiffFileWorker if not first note of a discussion' do
+        mr = create(:merge_request)
+        diff_note = create(:diff_note_on_merge_request, project: mr.project, noteable: mr)
+        reply_diff_note = create(:diff_note_on_merge_request, in_reply_to: diff_note)
+
+        expect(CreateNoteDiffFileWorker).not_to receive(:perform_async).with(reply_diff_note.id)
+
+        reply_diff_note.reload.diff_file
+      end
+    end
   end
 
   describe "#diff_line" do

spec/models/note_diff_file_spec.rb

+require 'rails_helper'
+
+describe NoteDiffFile do
+  describe 'associations' do
+    it { is_expected.to belong_to(:diff_note) }
+  end
+
+  describe 'validations' do
+    it { is_expected.to validate_presence_of(:diff_note) }
+  end
+end

spec/services/notes/create_service_spec.rb

@@ -57,6 +57,88 @@ describe Notes::CreateService do
       end
     end
 
+    context 'note diff file' do
+      let(:project_with_repo) { create(:project, :repository) }
+      let(:merge_request) do
+        create(:merge_request,
+               source_project: project_with_repo,
+               target_project: project_with_repo)
+      end
+      let(:line_number) { 14 }
+      let(:position) do
+        Gitlab::Diff::Position.new(old_path: "files/ruby/popen.rb",
+                                   new_path: "files/ruby/popen.rb",
+                                   old_line: nil,
+                                   new_line: line_number,
+                                   diff_refs: merge_request.diff_refs)
+      end
+      let(:previous_note) do
+        create(:diff_note_on_merge_request, noteable: merge_request, project: project_with_repo)
+      end
+
+      context 'when eligible to have a note diff file' do
+        let(:new_opts) do
+          opts.merge(in_reply_to_discussion_id: nil,
+                     type: 'DiffNote',
+                     noteable_type: 'MergeRequest',
+                     noteable_id: merge_request.id,
+                     position: position.to_h)
+        end
+
+        it 'note is associated with a note diff file' do
+          note = described_class.new(project_with_repo, user, new_opts).execute
+
+          expect(note).to be_persisted
+          expect(note.note_diff_file).to be_present
+        end
+      end
+
+      context 'when DiffNote is a reply' do
+        let(:new_opts) do
+          opts.merge(in_reply_to_discussion_id: previous_note.discussion_id,
+                     type: 'DiffNote',
+                     noteable_type: 'MergeRequest',
+                     noteable_id: merge_request.id,
+                     position: position.to_h)
+        end
+
+        it 'note is not associated with a note diff file' do
+          note = described_class.new(project_with_repo, user, new_opts).execute
+
+          expect(note).to be_persisted
+          expect(note.note_diff_file).to be_nil
+        end
+
+        context 'when DiffNote from an image' do
+          let(:image_position) do
+            Gitlab::Diff::Position.new(old_path: "files/images/6049019_460s.jpg",
+                                       new_path: "files/images/6049019_460s.jpg",
+                                       width: 100,
+                                       height: 100,
+                                       x: 1,
+                                       y: 100,
+                                       diff_refs: merge_request.diff_refs,
+                                       position_type: 'image')
+          end
+
+          let(:new_opts) do
+            opts.merge(in_reply_to_discussion_id: nil,
+                       type: 'DiffNote',
+                       noteable_type: 'MergeRequest',
+                       noteable_id: merge_request.id,
+                       position: image_position.to_h)
+          end
+
+          it 'note is not associated with a note diff file' do
+            note = described_class.new(project_with_repo, user, new_opts).execute
+
+            expect(note).to be_persisted
+            expect(note.note_diff_file).to be_nil
+          end
+        end
+      end
+    end
+
     context 'note with commands' do
       context 'as a user who can update the target' do
         context '/close, /label, /assign & /milestone' do

spec/workers/create_note_diff_file_worker_spec.rb

+require 'spec_helper'
+
+describe CreateNoteDiffFileWorker do
+  describe '#perform' do
+    let(:diff_note) { create(:diff_note_on_merge_request) }
+
+    it 'creates diff file' do
+      diff_note.note_diff_file.destroy!
+
+      expect_any_instance_of(DiffNote).to receive(:create_diff_file)
+        .and_call_original
+
+      described_class.new.perform(diff_note.id)
+    end
+  end
+end