Adds timeout to notebook rendering

In some cases, creating the rendering for notebook is taking too long.
There are even a few bugs
(https://gitlab.com/gitlab-org/incubation-engineering/mlops/rb-ipynbdiff/-/issues/9)
that might lead to an infinite loop. Although adding this timeout won't
solve the bug, it will avoid disrupting users.

Changelog: fixed

lib/gitlab/diff/rendered/notebook/diff_file.rb

@@ -6,6 +6,14 @@ module Gitlab
         include Gitlab::Utils::StrongMemoize
 
         class DiffFile < Gitlab::Diff::File
+          RENDERED_TIMEOUT_BACKGROUND = 10.seconds
+          RENDERED_TIMEOUT_FOREGROUND = 1.5.seconds
+          BACKGROUND_EXECUTION = 'background'
+          FOREGROUND_EXECUTION = 'foreground'
+          LOG_IPYNBDIFF_GENERATED = 'IPYNB_DIFF_GENERATED'
+          LOG_IPYNBDIFF_TIMEOUT = 'IPYNB_DIFF_TIMEOUT'
+          LOG_IPYNBDIFF_INVALID = 'IPYNB_DIFF_INVALID'
+
           attr_reader :source_diff
 
           delegate :repository, :diff_refs, :fallback_diff_refs, :unfolded, :unique_identifier,
@@ -52,14 +60,17 @@ module Gitlab
 
           def notebook_diff
             strong_memoize(:notebook_diff) do
-              Gitlab::AppLogger.info({ message: 'IPYNB_DIFF_GENERATED' })
-
+              Timeout.timeout(timeout_time) do
                 IpynbDiff.diff(source_diff.old_blob&.data, source_diff.new_blob&.data,
-                             raise_if_invalid_nb: true,
-                             diffy_opts: { include_diff_info: true })
+                               raise_if_invalid_nb: true, diffy_opts: { include_diff_info: true })&.tap do
+                  log_event(LOG_IPYNBDIFF_GENERATED)
+                end
+              end
+            rescue Timeout::Error => e
+              rendered_timeout.increment(source: Gitlab::Runtime.sidekiq? ? BACKGROUND_EXECUTION : FOREGROUND_EXECUTION)
+              log_event(LOG_IPYNBDIFF_TIMEOUT, e)
             rescue IpynbDiff::InvalidNotebookError, IpynbDiff::InvalidTokenError => e
-              Gitlab::ErrorTracking.log_exception(e)
-              nil
+              log_event(LOG_IPYNBDIFF_INVALID, e)
             end
           end
 
@@ -116,6 +127,23 @@ module Gitlab
 
             [removals, additions]
           end
+
+          def rendered_timeout
+            @rendered_timeout ||= Gitlab::Metrics.counter(
+              :ipynb_semantic_diff_timeouts_total,
+              'Counts the times notebook diff rendering timed out'
+            )
+          end
+
+          def timeout_time
+            Gitlab::Runtime.sidekiq? ? RENDERED_TIMEOUT_BACKGROUND : RENDERED_TIMEOUT_FOREGROUND
+          end
+
+          def log_event(message, error = nil)
+            Gitlab::AppLogger.info({ message: message })
+            Gitlab::ErrorTracking.track_exception(error) if error
+            nil
+          end
         end
       end
     end

spec/lib/gitlab/diff/rendered/notebook/diff_file_spec.rb

@@ -63,6 +63,28 @@ RSpec.describe Gitlab::Diff::Rendered::Notebook::DiffFile do
         expect(nb_file.diff).to be_nil
       end
     end
+
+    context 'timeout' do
+      it 'utilizes timeout for web' do
+        expect(Timeout).to receive(:timeout).with(described_class::RENDERED_TIMEOUT_FOREGROUND).and_call_original
+
+        nb_file.diff
+      end
+
+      it 'falls back to nil on timeout' do
+        allow(Gitlab::ErrorTracking).to receive(:track_and_raise_for_dev_exception)
+        expect(Timeout).to receive(:timeout).and_raise(Timeout::Error)
+
+        expect(nb_file.diff).to be_nil
+      end
+
+      it 'utilizes longer timeout for sidekiq' do
+        allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
+        expect(Timeout).to receive(:timeout).with(described_class::RENDERED_TIMEOUT_BACKGROUND).and_call_original
+
+        nb_file.diff
+      end
+    end
   end
 
   describe '#has_renderable?' do