add limit option to OnlyOnceSuccessful util

lib/datadog/core/utils/only_once_successful.rb

@@ -18,16 +18,50 @@ module Utils
       # In https://github.com/DataDog/dd-trace-rb/pull/1398#issuecomment-797378810 we have a discussion of alternatives,
       # including an alternative implementation that is Ractor-safe once spent.
       class OnlyOnceSuccessful < OnlyOnce
+        def initialize(limit = 0)
+          super()
+
+          @limit = limit
+          @failed = false
+          @retries = 0
+        end
+
         def run
           @mutex.synchronize do
             return if @ran_once
 
             result = yield
             @ran_once = !!result
 
+            if !@ran_once && limited?
+              @retries += 1
+              check_limit!
+            end
+
             result
           end
         end
+
+        def success?
+          @mutex.synchronize { @ran_once && !@failed }
+        end
+
+        def failed?
+          @mutex.synchronize { @ran_once && @failed }
+        end
+
+        private
+
+        def check_limit!
+          if @retries >= @limit
+            @failed = true
+            @ran_once = true
+          end
+        end
+
+        def limited?
+          !@limit.nil? && @limit.positive?
+        end
       end
     end
   end

sig/datadog/core/utils/only_once_successful.rbs

@@ -2,6 +2,21 @@ module Datadog
   module Core
     module Utils
       class OnlyOnceSuccessful < Datadog::Core::Utils::OnlyOnce
+        @limit: Integer
+        @retries: Integer
+        @failed: bool
+
+        def initialize: (?Integer limit) -> void
+
+        def success?: () -> bool
+
+        def failed?: () -> bool
+
+        private
+
+        def check_limit!: () -> void
+
+        def limited?: () -> bool
       end
     end
   end

spec/datadog/core/utils/only_once_successful_spec.rb

@@ -1,28 +1,62 @@
 require 'datadog/core/utils/only_once_successful'
 
 RSpec.describe Datadog::Core::Utils::OnlyOnceSuccessful do
-  subject(:only_once_successful) { described_class.new }
+  subject(:only_once_successful) { described_class.new(limit) }
+
+  let(:limit) { 0 }
 
   describe '#run' do
-    context 'before running once' do
-      it do
-        expect { |block| only_once_successful.run(&block) }.to yield_control
-      end
+    context 'when limitless' do
+      context 'before running once' do
+        it do
+          expect { |block| only_once_successful.run(&block) }.to yield_control
+        end
 
-      it 'returns the result of the block ran' do
-        expect(only_once_successful.run { :result }).to be :result
+        it 'returns the result of the block ran' do
+          expect(only_once_successful.run { :result }).to be :result
+        end
       end
-    end
 
-    context 'after running once' do
-      let(:result) { nil }
+      context 'after running once' do
+        let(:result) { nil }
 
-      before do
-        only_once_successful.run { result }
+        before do
+          only_once_successful.run { result }
+        end
+
+        context 'when block returns truthy value' do
+          let(:result) { true }
+
+          it do
+            expect { |block| only_once_successful.run(&block) }.to_not yield_control
+          end
+
+          it do
+            expect(only_once_successful.run { :result }).to be nil
+          end
+        end
+
+        context 'when block returns falsey value' do
+          let(:result) { false }
+
+          it do
+            expect { |block| only_once_successful.run(&block) }.to yield_control
+          end
+
+          it 'runs again until block returns truthy value' do
+            expect(only_once_successful.run { :result }).to be :result
+
+            expect(only_once_successful.run { :result }).to be nil
+          end
+        end
       end
+    end
+
+    context 'when limited' do
+      let(:limit) { 2 }
 
       context 'when block returns truthy value' do
-        let(:result) { true }
+        before { only_once_successful.run { true } }
 
         it do
           expect { |block| only_once_successful.run(&block) }.to_not yield_control
@@ -33,16 +67,18 @@
         end
       end
 
-      context 'when block returns falsey value' do
-        let(:result) { false }
+      context 'when block returns falsey value "limit" times' do
+        before do
+          limit.times do
+            only_once_successful.run { false }
+          end
+        end
 
         it do
-          expect { |block| only_once_successful.run(&block) }.to yield_control
+          expect { |block| only_once_successful.run(&block) }.to_not yield_control
         end
 
-        it 'runs again until block returns truthy value' do
-          expect(only_once_successful.run { :result }).to be :result
-
+        it do
           expect(only_once_successful.run { :result }).to be nil
         end
       end
@@ -91,5 +127,107 @@
         end
       end
     end
+
+    context 'when limited and ran "limit" times' do
+      let(:limit) { 2 }
+
+      before do
+        limit.times do
+          only_once_successful.run { false }
+        end
+      end
+
+      it do
+        expect(only_once_successful.ran?).to be true
+      end
+    end
+  end
+
+  describe '#success?' do
+    context 'before running once' do
+      it do
+        expect(only_once_successful.success?).to be false
+      end
+    end
+
+    context 'after running once' do
+      let(:result) { nil }
+
+      before do
+        only_once_successful.run { result }
+      end
+
+      context 'when block returns truthy value' do
+        let(:result) { true }
+
+        it do
+          expect(only_once_successful.success?).to be true
+        end
+      end
+
+      context 'when block returns falsey value' do
+        it do
+          expect(only_once_successful.success?).to be false
+        end
+      end
+    end
+
+    context 'when limited and ran "limit" times' do
+      let(:limit) { 2 }
+
+      before do
+        limit.times do
+          only_once_successful.run { false }
+        end
+      end
+
+      it do
+        expect(only_once_successful.success?).to be false
+      end
+    end
+  end
+
+  describe '#failed?' do
+    context 'before running once' do
+      it do
+        expect(only_once_successful.failed?).to be false
+      end
+    end
+
+    context 'after running once' do
+      let(:result) { nil }
+
+      before do
+        only_once_successful.run { result }
+      end
+
+      context 'when block returns truthy value' do
+        let(:result) { true }
+
+        it do
+          expect(only_once_successful.failed?).to be false
+        end
+      end
+
+      context 'when block returns falsey value' do
+        it do
+          expect(only_once_successful.failed?).to be false
+        end
+      end
+    end
+
+    context 'when limited and ran "limit" times' do
+      let(:limit) { 2 }
+
+      before do
+        limit.times do
+          only_once_successful.run { false }
+        end
+      end
+
+      it do
+        expect(only_once_successful.failed?).to be true
+      end
+    end
   end
 end