RUBY-3803 Fix retries when an overload error (#3021)

Author: comandeo-mongo

lib/mongo/retryable/read_worker.rb

@@ -332,13 +332,19 @@ def retry_read(original_error, session, server_selector, context: nil, failed_se
       # Each retry sleeps with jittered backoff, respects MAX_RETRIES,
       # and consumes a token from the bucket when adaptive retries
       # are enabled.
+      #
+      # Per the client-backpressure spec, backoff is applied if and only
+      # if the error triggering the retry is an overload error. Non-overload
+      # retryable errors that occur within this loop are retried immediately
+      # (without backoff) but still count toward MAX_RETRIES.
       def overload_read_retry(last_error, session, server_selector, context, failed_server, error_count:)
+        last_was_overload = true
         loop do
-          delay = retry_policy.backoff_delay(error_count)
+          delay = last_was_overload ? retry_policy.backoff_delay(error_count) : 0
           raise last_error unless retry_policy.should_retry_overload?(error_count, delay, context: context)
 
           log_retry(last_error, message: 'Read retry (overload backoff)')
-          sleep(delay)
+          sleep(delay) if last_was_overload
 
           begin
             server = select_server(
@@ -363,6 +369,7 @@ def overload_read_retry(last_error, session, server_selector, context, failed_se
             is_overload = retryable_overload_error?(e)
             raise e unless is_overload || is_retryable_exception?(e) || e.write_retryable?
 
+            last_was_overload = is_overload
             failed_server = server
             last_error = e
           end

lib/mongo/retryable/write_worker.rb

@@ -372,19 +372,25 @@ def retry_write(original_error, txn_num, context:, failed_server: nil, &block)
       end
 
       # Retry loop for overload write errors with exponential backoff.
+      #
+      # Per the client-backpressure spec, backoff is applied if and only
+      # if the error triggering the retry is an overload error. Non-overload
+      # retryable errors that occur within this loop are retried immediately
+      # (without backoff) but still count toward MAX_RETRIES.
       def overload_write_retry(last_error, session, txn_num, context:, failed_server:, error_count:,
                                was_starting_transaction: false)
         # Track the error to return per the NoWritesPerformed spec rules:
         # - first error is always saved
         # - only update when a new error does NOT have NoWritesPerformed
         error_to_raise = last_error
+        last_was_overload = true
 
         loop do
-          delay = retry_policy.backoff_delay(error_count)
+          delay = last_was_overload ? retry_policy.backoff_delay(error_count) : 0
           raise error_to_raise unless retry_policy.should_retry_overload?(error_count, delay, context: context)
 
           log_retry(last_error, message: 'Write retry (overload backoff)')
-          sleep(delay)
+          sleep(delay) if last_was_overload
 
           begin
             server = select_server(
@@ -423,6 +429,7 @@ def overload_write_retry(last_error, session, txn_num, context:, failed_server:,
             unless e.respond_to?(:label?) && e.label?('NoWritesPerformed')
               error_to_raise = e
             end
+            last_was_overload = is_overload
             context = context.with(overload_only_retry: false) unless is_overload
             failed_server = server
             last_error = e

spec/mongo/retryable/client_backpressure_no_backoff_prose_spec.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# Client Backpressure Prose Tests for backoff behavior.
+#
+# Spec reference:
+#   specifications/source/client-backpressure/tests/README.md
+describe 'Client Backpressure backoff prose tests' do
+  require_topology :replica_set
+  min_server_version '4.4'
+
+  let(:client) do
+    authorized_client.with(retry_reads: true)
+  end
+
+  let(:admin_client) { client.use(:admin) }
+
+  let(:collection) { client['backoff-prose-test'] }
+
+  let(:subscriber) { Mrss::EventSubscriber.new }
+
+  before do
+    # Inflate BASE_BACKOFF so any accidental backoff is clearly visible
+    # through timing. Without backoff the operation completes in
+    # milliseconds; with backoff it would take at least 5 seconds.
+    stub_const('Mongo::Retryable::Backpressure::BASE_BACKOFF', 5.0)
+  end
+
+  after do
+    admin_client.command(configureFailPoint: 'failCommand', mode: 'off')
+  rescue Mongo::Error
+    # Ignore cleanup failures.
+  end
+
+  # -------------------------------------------------------------------------
+  # Test 4: Backoff is applied if and only if the error is an
+  # overload error (mixed overload + non-overload in the overload loop)
+  # -------------------------------------------------------------------------
+  describe 'Test 4: backoff applied only for overload errors in overload retry loop' do
+    it 'applies backoff for the overload error but not for subsequent non-overload errors' do
+      # Configure first fail point: overload error, fires once.
+      admin_client.command(
+        configureFailPoint: 'failCommand',
+        mode: { times: 1 },
+        data: {
+          failCommands: %w[find],
+          errorCode: 91,
+          errorLabels: %w[RetryableError SystemOverloadedError]
+        }
+      )
+
+      # Via CommandFailedEvent, switch to a non-overload retryable error.
+      failpoint_set = false
+      client.subscribe(Mongo::Monitoring::COMMAND, subscriber)
+
+      allow(subscriber).to receive(:failed).and_wrap_original do |m, event|
+        m.call(event)
+        if !failpoint_set && event.command_name == 'find'
+          failpoint_set = true
+          admin_client.command(
+            configureFailPoint: 'failCommand',
+            mode: 'alwaysOn',
+            data: {
+              failCommands: %w[find],
+              errorCode: 91,
+              errorLabels: %w[RetryableError]
+            }
+          )
+        end
+      end
+
+      subscriber.clear_events!
+
+      start_time = Mongo::Utils.monotonic_time
+      expect do
+        collection.find.first
+      end.to raise_error(Mongo::Error::OperationFailure)
+      elapsed = Mongo::Utils.monotonic_time - start_time
+
+      # With BASE_BACKOFF=5s, correct behavior applies one backoff
+      # (bounded by BASE_BACKOFF) for the overload error, then retries
+      # non-overload errors immediately. The elapsed time should stay
+      # under BASE_BACKOFF plus a small margin for network overhead.
+      expect(elapsed).to be < Mongo::Retryable::Backpressure::BASE_BACKOFF + 2
+    end
+  end
+end

spec/mongo/retryable/retryable_reads_overload_retry_count_prose_spec.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# Retryable Reads Prose Test 1: Test that drivers set the maximum number
+# of retries for all retryable read errors when an overload error is
+# encountered.
+#
+# Spec reference:
+#   specifications/source/retryable-reads/tests/README.md
+#   "1: Test that drivers set the maximum number of retries for all
+#    retryable read errors when an overload error is encountered"
+describe 'Retryable reads prose test 1: overload retry count' do
+  require_topology :replica_set
+  min_server_version '6.0'
+
+  let(:client) do
+    authorized_client.with(retry_reads: true)
+  end
+
+  let(:admin_client) { client.use(:admin) }
+
+  let(:collection) { client['overload-retry-count-reads-prose-test'] }
+
+  let(:subscriber) { Mrss::EventSubscriber.new }
+
+  # MAX_ADAPTIVE_RETRIES in spec terminology; the configured max overload
+  # retries for this client (defaults to DEFAULT_MAX_RETRIES).
+  let(:max_adaptive_retries) { client.retry_policy.max_retries }
+
+  let(:find_started_events) do
+    subscriber.started_events.select { |e| e.command_name == 'find' }
+  end
+
+  after do
+    admin_client.command(configureFailPoint: 'failCommand', mode: 'off')
+  rescue Mongo::Error
+    # Ignore cleanup failures.
+  end
+
+  it 'makes MAX_ADAPTIVE_RETRIES + 1 total attempts' do
+    # Step 2: Configure a fail point for find that fires once with an
+    # overload error (code 91, labels RetryableError + SystemOverloadedError).
+    admin_client.command(
+      configureFailPoint: 'failCommand',
+      mode: { times: 1 },
+      data: {
+        failCommands: %w[find],
+        errorCode: 91,
+        errorLabels: %w[RetryableError SystemOverloadedError]
+      }
+    )
+
+    # Step 3: Via CommandFailedEvent, when the first find error fires,
+    # configure a second fail point for find with a non-overload retryable
+    # error (code 91, label RetryableError only), set to alwaysOn.
+    failpoint_set = false
+    client.subscribe(Mongo::Monitoring::COMMAND, subscriber)
+
+    allow(subscriber).to receive(:failed).and_wrap_original do |m, event|
+      m.call(event)
+      if !failpoint_set && event.command_name == 'find'
+        failpoint_set = true
+        admin_client.command(
+          configureFailPoint: 'failCommand',
+          mode: 'alwaysOn',
+          data: {
+            failCommands: %w[find],
+            errorCode: 91,
+            errorLabels: %w[RetryableError]
+          }
+        )
+      end
+    end
+
+    # Step 4: Attempt a find. Expect it to fail.
+    subscriber.clear_events!
+    expect do
+      collection.find.first
+    end.to raise_error(Mongo::Error::OperationFailure)
+
+    # Step 5: Assert that MAX_ADAPTIVE_RETRIES + 1 total find commands
+    # were sent (1 initial attempt + MAX_ADAPTIVE_RETRIES retries).
+    expect(find_started_events.length).to eq(max_adaptive_retries + 1)
+  end
+end

spec/mongo/retryable/retryable_writes_overload_retry_count_prose_spec.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# Retryable Writes Prose Test Case 4: Test that drivers set the maximum
+# number of retries for all retryable write errors when an overload error
+# is encountered.
+#
+# Spec reference:
+#   specifications/source/retryable-writes/tests/README.md
+#   "Case 4: Test that drivers set the maximum number of retries for all
+#    retryable write errors when an overload error is encountered"
+describe 'Retryable writes prose test case 4: overload retry count' do
+  require_topology :replica_set
+  min_server_version '6.0'
+
+  let(:client) do
+    authorized_client.with(retry_writes: true)
+  end
+
+  let(:admin_client) { client.use(:admin) }
+
+  let(:collection) { client['overload-retry-count-writes-prose-test'] }
+
+  let(:subscriber) { Mrss::EventSubscriber.new }
+
+  # MAX_ADAPTIVE_RETRIES in spec terminology; the configured max overload
+  # retries for this client (defaults to DEFAULT_MAX_RETRIES).
+  let(:max_adaptive_retries) { client.retry_policy.max_retries }
+
+  let(:insert_started_events) do
+    subscriber.started_events.select { |e| e.command_name == 'insert' }
+  end
+
+  after do
+    admin_client.command(configureFailPoint: 'failCommand', mode: 'off')
+  rescue Mongo::Error
+    # Ignore cleanup failures.
+  end
+
+  it 'makes MAX_ADAPTIVE_RETRIES + 1 total attempts' do
+    # Step 2: Configure a fail point for insert that fires once with an
+    # overload error (code 91, labels RetryableError + SystemOverloadedError).
+    admin_client.command(
+      configureFailPoint: 'failCommand',
+      mode: { times: 1 },
+      data: {
+        failCommands: %w[insert],
+        errorCode: 91,
+        errorLabels: %w[RetryableError SystemOverloadedError]
+      }
+    )
+
+    # Step 3: Via CommandFailedEvent, when the first insert error fires,
+    # configure a second fail point for insert with a non-overload retryable
+    # write error (code 91, labels RetryableError + RetryableWriteError),
+    # set to alwaysOn.
+    failpoint_set = false
+    client.subscribe(Mongo::Monitoring::COMMAND, subscriber)
+
+    allow(subscriber).to receive(:failed).and_wrap_original do |m, event|
+      m.call(event)
+      if !failpoint_set && event.command_name == 'insert'
+        failpoint_set = true
+        admin_client.command(
+          configureFailPoint: 'failCommand',
+          mode: 'alwaysOn',
+          data: {
+            failCommands: %w[insert],
+            errorCode: 91,
+            errorLabels: %w[RetryableError RetryableWriteError]
+          }
+        )
+      end
+    end
+
+    # Step 4: Attempt an insertOne. Expect it to fail.
+    subscriber.clear_events!
+    expect do
+      collection.insert_one(x: 1)
+    end.to raise_error(Mongo::Error::OperationFailure)
+
+    # Step 5: Assert that MAX_ADAPTIVE_RETRIES + 1 total insert commands
+    # were sent (1 initial attempt + MAX_ADAPTIVE_RETRIES retries).
+    expect(insert_started_events.length).to eq(max_adaptive_retries + 1)
+  end
+end