Adds timeout for notebook rendering on CustomDiff

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/custom_diff.rb

@@ -2,17 +2,29 @@
 module Gitlab
   module Diff
     module CustomDiff
+      RENDERED_TIMEOUT_BACKGROUND = 20.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'
+
       class << self
         def preprocess_before_diff(path, old_blob, new_blob)
           return unless path.ends_with? '.ipynb'
 
-          transformed_diff(old_blob&.data, new_blob&.data)&.tap do
-            transformed_for_diff(new_blob, old_blob)
-            Gitlab::AppLogger.info({ message: 'IPYNB_DIFF_GENERATED' })
+          Timeout.timeout(timeout_time) do
+            transformed_diff(old_blob&.data, new_blob&.data)&.tap do
+              transformed_for_diff(new_blob, old_blob)
+              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
 
         def transformed_diff(before, after)
@@ -50,6 +62,23 @@ module Gitlab
             blobs_with_transformed_diffs[b] = true if b
           end
         end
+
+        def rendered_timeout
+          @rendered_timeout ||= Gitlab::Metrics.counter(
+            :rendered_timeout,
+            'Counts the times notebook 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.log_exception(error) if error
+          nil
+        end
       end
     end
   end

spec/lib/gitlab/diff/custom_diff_spec.rb

@@ -34,6 +34,28 @@ RSpec.describe Gitlab::Diff::CustomDiff do
         expect(described_class.transformed_for_diff?(blob)).to be_falsey
       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
+
+        expect(described_class.preprocess_before_diff(ipynb_blob.path, nil, ipynb_blob)).not_to include('cells')
+      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(described_class.preprocess_before_diff(ipynb_blob.path, nil, ipynb_blob)).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
+
+        described_class.preprocess_before_diff(ipynb_blob.path, nil, ipynb_blob)
+      end
+    end
   end
 
   describe '#transformed_blob_data' do