Grant the configured number of retries, and drop the at_exit wakeup

Two fixes.

RetryBudget decremented @retries_remaining before testing it, so one retry
was spent on the exhaustion check itself: a configured N performed N-1, and
retries: 1 and retries: 0 both performed none. go, python and java all
grant N. The existing spec asserted `.exactly(retries - 1).times`, so it
codified the bug and would have passed either way; it now asserts N, and
the new retry_budget_spec covers the 1-and-0 cases that were previously
indistinguishable. Both fail without this change.

The at_exit block called Thread#wakeup on the worker. That raises
ThreadError if the thread finished between the alive? check and the call,
and an unrescued raise inside at_exit forces the process to exit 1 — a
spurious failure for any script or CI step using this client. Reproduced
under the pinned ruby 3.2.

Rescuing the error would have left the other half broken: wakeup only cuts
short a sleep already in progress, so one arriving while the worker is
mid-request is lost and the next sleep runs in full, which is the hang the
wakeup was added to prevent. The retry wait is now sliced and checks
should_exit between slices, so shutdown is noticed within a second without
needing wakeup at all. Same shape as python's interruptible wait.

Specs stub the new interruptible_sleep seam instead of sleep, and must
return truthy — a nil return reads as "shutting down" and stops the retry.

Metrics/ClassLength for transport.rb goes 155 -> 161. Regenerating
.rubocop_todo.yml as the repo's CLAUDE.md suggests rewrote ~100 unrelated
lines (rubocop 1.90 against a file generated by 1.44, plus e2e-cli files
that were not previously inspected), so this is the minimal edit instead.

210 examples pass, rubocop is clean on lib and spec, and the 61-test e2e
suite passes.

Author: MichaelGHSeg

.rubocop_todo.yml

@@ -73,7 +73,7 @@ Metrics/BlockLength:
 # Offense count: 1
 # Configuration parameters: CountComments, CountAsOne.
 Metrics/ClassLength:
-  Max: 155
+  Max: 161
 
 # Offense count: 2
 # Configuration parameters: AllowedMethods, AllowedPatterns, IgnoredMethods.

History.md

@@ -19,6 +19,7 @@ sent the write key as HTTP Basic credentials.
 * Backoff waits no longer block shutdown for the full delay.
 * Retry timing uses a monotonic clock, so a system clock change cannot stretch or collapse a backoff.
 * Backoff intervals are now jittered at the ceiling as well, so clients that back off together do not retry in lockstep.
+* Fix `retries` granting one fewer attempt than configured. A configured 10 performed 9, and `retries: 1` performed none at all.
 
 2.5.0 / 2024-07-17
 ==================

lib/segment/analytics/client.rb

@@ -32,14 +32,11 @@ def initialize(opts = {})
 
         check_write_key!
 
-        at_exit do
-          if @worker_thread
-            @worker_thread[:should_exit] = true
-            # Break any Retry-After or backoff sleep so shutdown is not held for
-            # up to rate_limit_retry_after_cap seconds.
-            @worker_thread.wakeup if @worker_thread.alive?
-          end
-        end
+        # The worker checks this between sleep slices, so a Retry-After or backoff
+        # wait is abandoned within a second rather than holding shutdown for up to
+        # rate_limit_retry_after_cap seconds. Assigning to a dead thread is safe;
+        # Thread#wakeup is not, and raising here would force a non-zero exit status.
+        at_exit { @worker_thread[:should_exit] = true if @worker_thread }
       end
 
       # Synchronously waits until the worker has flushed the queue.

lib/segment/analytics/retry_budget.rb

@@ -30,9 +30,13 @@ def initialize(options = {})
       end
 
       def next_backoff_delay
-        @retries_remaining -= 1
+        # Checked before the decrement: decrementing first spent one retry on the
+        # exhaustion test itself, so a configured N only ever performed N-1, and
+        # retries: 1 and retries: 0 were indistinguishable.
         return spent('Retries exhausted for batch') if @retries_remaining <= 0
 
+        @retries_remaining -= 1
+
         @backoff_start_time ||= monotonic_now
         return spent('Max total backoff duration exceeded for batch') if elapsed?(@backoff_start_time, @max_total_backoff_duration)
 

lib/segment/analytics/transport.rb

@@ -120,6 +120,9 @@ def build_http(options)
         http
       end
 
+      # How long a sliced retry wait sleeps before re-checking for shutdown.
+      SHUTDOWN_CHECK_INTERVAL = 1
+
       def new_retry_budget
         RetryBudget.new(
           :retries => @retries,
@@ -135,15 +138,31 @@ def new_retry_budget
       # RetryBudget keeps the wait on Transport, where callers stub it.
       def wait_to_retry(delay, budget)
         return false if delay.nil?
-
-        sleep(delay)
-        # Client#shutdown wakes this thread, so the sleep above returns early.
-        return false if Thread.current[:should_exit]
+        return false unless interruptible_sleep(delay)
 
         budget.record_retry
         true
       end
 
+      # Sleeps in slices so shutdown is noticed within SHUTDOWN_CHECK_INTERVAL
+      # rather than after the whole delay, which can be rate_limit_retry_after_cap
+      # seconds. Returns false if shutdown was requested.
+      #
+      # Thread#wakeup is deliberately not used for this: it only cuts short a sleep
+      # already in progress, so a wakeup arriving while the worker is mid-request is
+      # lost and the next sleep still runs in full.
+      def interruptible_sleep(seconds)
+        remaining = seconds
+        while remaining > 0
+          return false if Thread.current[:should_exit]
+
+          slice = [remaining, SHUTDOWN_CHECK_INTERVAL].min
+          sleep(slice)
+          remaining -= slice
+        end
+        !Thread.current[:should_exit]
+      end
+
       def parse_error(body)
         JSON.parse(body)['error']
       rescue StandardError

spec/segment/analytics/retry_budget_spec.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+module Segment
+  class Analytics
+    describe RetryBudget do
+      let(:logger) { Logger.new(File::NULL) }
+
+      def budget(retries, intervals = nil)
+        described_class.new(
+          :retries => retries,
+          :backoff_policy => FakeBackoffPolicy.new(intervals || Array.new(retries, 1000)),
+          :max_total_backoff_duration => 43_200,
+          :max_rate_limit_duration => 43_200,
+          :rate_limit_retry_after_cap => 300,
+          :logger => logger
+        )
+      end
+
+      describe '#next_backoff_delay' do
+        it 'grants exactly as many retries as configured' do
+          # The count used to be decremented before the exhaustion check, so a
+          # configured N yielded N-1. go, python and java all grant N.
+          subject = budget(3)
+
+          expect(subject.next_backoff_delay).to eq(1.0)
+          expect(subject.next_backoff_delay).to eq(1.0)
+          expect(subject.next_backoff_delay).to eq(1.0)
+          expect(subject.next_backoff_delay).to be_nil
+        end
+
+        it 'grants one retry for retries: 1' do
+          subject = budget(1)
+
+          expect(subject.next_backoff_delay).to eq(1.0)
+          expect(subject.next_backoff_delay).to be_nil
+        end
+
+        it 'grants no retries for retries: 0' do
+          # retries: 0 and retries: 1 were previously indistinguishable.
+          subject = budget(0, [1000])
+
+          expect(subject.next_backoff_delay).to be_nil
+        end
+      end
+    end
+  end
+end

spec/segment/analytics/transport_spec.rb

@@ -163,18 +163,18 @@ class Analytics
             let(:status_code) { status_code }
             let(:body) { body }
             let(:retries) { 4 }
-            let(:backoff_policy) { FakeBackoffPolicy.new([1000, 1000, 1000]) }
+            let(:backoff_policy) { FakeBackoffPolicy.new([1000, 1000, 1000, 1000]) }
             subject {
               described_class.new(retries: retries,
                                   backoff_policy: backoff_policy)
             }
 
             it 'retries the request' do
               expect(subject)
-                .to receive(:sleep)
-                .exactly(retries - 1).times
+                .to receive(:interruptible_sleep)
+                .exactly(retries).times
                 .with(1)
-                .and_return(nil)
+                .and_return(true)
               subject.send(write_key, batch)
             end
           end
@@ -188,7 +188,7 @@ class Analytics
 
             it 'does not retry the request' do
               expect(subject)
-                .to receive(:sleep)
+                .to receive(:interruptible_sleep)
                 .never
               subject.send(write_key, batch)
             end
@@ -208,7 +208,7 @@ class Analytics
           context '3xx is not retried and is not success' do
             let(:status_code) { 301 }
             it 'returns the status without retrying, and does not report success' do
-              expect(subject).not_to receive(:sleep)
+              expect(subject).not_to receive(:interruptible_sleep)
               response = subject.send(write_key, batch)
               expect(response.status).to eq(301)
               expect(response.success?).to be false
@@ -261,18 +261,18 @@ class Analytics
             end
 
             it 'sleeps for the Retry-After duration' do
-              expect(subject).to receive(:sleep).with(2).once
+              expect(subject).to receive(:interruptible_sleep).with(2).once.and_return(true)
               subject.send(write_key, batch)
             end
 
             it 'caps Retry-After at RATE_LIMIT_RETRY_AFTER_CAP' do
               allow(response).to receive(:to_hash) { { 'retry-after' => ['9999'] } }
-              expect(subject).to receive(:sleep).with(described_class::RATE_LIMIT_RETRY_AFTER_CAP).once
+              expect(subject).to receive(:interruptible_sleep).with(described_class::RATE_LIMIT_RETRY_AFTER_CAP).once.and_return(true)
               subject.send(write_key, batch)
             end
 
             it 'returns success after retry' do
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
               expect(subject.send(write_key, batch).success?).to be true
             end
           end
@@ -291,12 +291,12 @@ class Analytics
             end
 
             it 'sleeps for the Retry-After duration' do
-              expect(subject).to receive(:sleep).with(2).once
+              expect(subject).to receive(:interruptible_sleep).with(2).once.and_return(true)
               subject.send(write_key, batch)
             end
 
             it 'does not decrement retries_remaining (uses rate-limit path)' do
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
               # With retries: 1, a 503+Retry-After should NOT exhaust retries because
               # it uses the rate-limit path (no retry budget cost)
               transport = described_class.new(retries: 1, backoff_policy: FakeBackoffPolicy.new([1000]))
@@ -306,7 +306,7 @@ class Analytics
               allow(success_response).to receive(:body) { '{}' }
               allow(success_response).to receive(:to_hash) { {} }
               allow(http).to receive(:request).and_return(response, success_response)
-              allow(transport).to receive(:sleep)
+              allow(transport).to receive(:interruptible_sleep).and_return(true)
               result = transport.send(write_key, batch)
               expect(result.status).to eq(200)
             end
@@ -326,12 +326,12 @@ class Analytics
             end
 
             it 'sleeps for the Retry-After duration' do
-              expect(subject).to receive(:sleep).with(1).once
+              expect(subject).to receive(:interruptible_sleep).with(1).once.and_return(true)
               subject.send(write_key, batch)
             end
 
             it 'returns success after retry' do
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
               expect(subject.send(write_key, batch).success?).to be true
             end
 
@@ -344,19 +344,19 @@ class Analytics
               allow(success_response).to receive(:body) { '{}' }
               allow(success_response).to receive(:to_hash) { {} }
               allow(http).to receive(:request).and_return(response, success_response)
-              allow(transport).to receive(:sleep)
+              allow(transport).to receive(:interruptible_sleep).and_return(true)
               result = transport.send(write_key, batch)
               expect(result.status).to eq(200)
             end
           end
 
           context 'X-Retry-Count header' do
             let(:status_code) { 500 }
-            let(:backoff_policy) { FakeBackoffPolicy.new([1, 1]) }
+            let(:backoff_policy) { FakeBackoffPolicy.new([1, 1, 1]) }
             subject { described_class.new(retries: 3, backoff_policy: backoff_policy) }
 
             it 'does not send X-Retry-Count on first attempt' do
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
               first_request = nil
               http = subject.instance_variable_get(:@http)
               allow(http).to receive(:request) do |req, _|
@@ -368,7 +368,7 @@ class Analytics
             end
 
             it 'sends X-Retry-Count incrementing on retries' do
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
               requests = []
               http = subject.instance_variable_get(:@http)
               allow(http).to receive(:request) do |req, _|
@@ -378,6 +378,7 @@ class Analytics
               subject.send(write_key, batch)
               expect(requests[1]['X-Retry-Count']).to eq('1')
               expect(requests[2]['X-Retry-Count']).to eq('2')
+              expect(requests[3]['X-Retry-Count']).to eq('3')
             end
           end
 
@@ -395,7 +396,7 @@ class Analytics
 
                 success_response
               end
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
 
               response = subject.send(write_key, batch)
 
@@ -406,7 +407,7 @@ class Analytics
             it 'gives up once the retry budget is spent' do
               http = subject.instance_variable_get(:@http)
               allow(http).to receive(:request).and_raise(Errno::ECONNRESET, 'reset')
-              allow(subject).to receive(:sleep)
+              allow(subject).to receive(:interruptible_sleep).and_return(true)
 
               response = subject.send(write_key, batch)
 
@@ -480,6 +481,34 @@ class Analytics
           end
         end
       end
+
+      describe '#interruptible_sleep' do
+        subject { described_class.new }
+
+        it 'abandons the wait when shutdown is requested instead of sleeping it out' do
+          # The wait used to be a single sleep broken by Thread#wakeup, which only
+          # interrupts a sleep already in progress and raises ThreadError if the
+          # thread has finished. Slicing removes the need for it.
+          elapsed = nil
+
+          worker = Thread.new do
+            started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            result = subject.__send__(:interruptible_sleep, 30)
+            elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started
+            result
+          end
+
+          sleep 0.1
+          worker[:should_exit] = true
+
+          expect(worker.value).to be false
+          expect(elapsed).to be < 5
+        end
+
+        it 'reports completion when the delay elapses' do
+          expect(Thread.new { subject.__send__(:interruptible_sleep, 0) }.value).to be true
+        end
+      end
     end
   end
 end