Add Feature#around_request hook for wrapping the HTTP exchange lifecycle

Add an around-style hook that wraps the per-attempt send/receive cycle,
with the request flowing through the feature chain. This enables
features like instrumentation to use ensure-based cleanup for guaranteed
span closing on both success and failure.

Rewrite the instrumentation feature to use around_request, replacing the
separate on_request/wrap_response pair with a single call to
instrumenter.instrument, which naturally provides paired start/finish
via its own ensure block.

Author: sferik

CHANGELOG.md

@@ -9,8 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- `Feature#on_request` lifecycle hook, called before each request attempt
-  (including retries), for per-attempt side effects like instrumentation (#826)
+- `Feature#on_request` and `Feature#around_request` lifecycle hooks, called
+  before/around each request attempt (including retries), for per-attempt side
+  effects like instrumentation spans and circuit breakers (#826)
 - Pattern matching support (`deconstruct_keys`) for Response, Response::Status,
   Headers, ContentType, and URI (#642)
 - Combined global and per-operation timeouts: global and per-operation timeouts

lib/http/client.rb

@@ -136,13 +136,9 @@ def close
     # @return [HTTP::Response] the response
     # @api private
     def perform_once(req, options)
-      verify_connection!(req.uri)
-      @state = :dirty
+      res = perform_exchange(req, options)
 
-      send_request(req, options)
-      res = build_wrapped_response(req, options)
-
-      @connection.finish_response if req.verb == :head
+      @connection.finish_response if res.request.verb == :head
       @state = :clean
 
       res
@@ -198,6 +194,27 @@ def notify_features(req, options)
       options.features.each_value { |feature| feature.on_request(req) }
     end
 
+    # Execute the HTTP exchange wrapped by feature around_request hooks
+    # @return [HTTP::Response] the response
+    # @api private
+    def perform_exchange(req, options)
+      around_request(req, options) do |request|
+        verify_connection!(request.uri)
+        @state = :dirty
+        send_request(request, options)
+        build_wrapped_response(request, options)
+      end
+    end
+
+    # Compose around_request chains from all features
+    # @return [HTTP::Response] the response
+    # @api private
+    def around_request(request, options, &block)
+      options.features.values.reverse.reduce(block) do |inner, feature|
+        ->(req) { feature.around_request(req) { |r| inner.call(r) } }
+      end.call(request)
+    end
+
     # Wrap request through feature middleware
     # @return [HTTP::Request] the wrapped request
     # @api private

lib/http/feature.rb

@@ -41,6 +41,28 @@ def wrap_response(response)
     # @api public
     def on_request(_request); end
 
+    # Wraps the HTTP exchange for a single request attempt
+    #
+    # Called once per attempt (including retries), wrapping the send and
+    # receive cycle. The block performs the I/O and returns the response.
+    # Override this to add behavior that must span the entire exchange,
+    # such as instrumentation spans or circuit breakers.
+    #
+    # @example Timing a request
+    #   def around_request(request)
+    #     start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    #     yield(request).tap { log_duration(Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) }
+    #   end
+    #
+    # @param request [HTTP::Request]
+    # @yield [HTTP::Request] the request to perform
+    # @yieldreturn [HTTP::Response]
+    # @return [HTTP::Response] must return the response from yield
+    # @api public
+    def around_request(request)
+      yield request
+    end
+
     # Callback for request errors
     #
     # @example

lib/http/features/instrumentation.rb

@@ -62,37 +62,29 @@ def initialize(instrumenter: NullInstrumenter.new, namespace: "http")
         @error_name = "error.#{namespace}"
       end
 
-      # Starts instrumentation for a request attempt
+      # Wraps the HTTP exchange with instrumentation
       #
-      # Called before each request attempt, including retries. Emits a
-      # separate "start" event so a logger can print the request being run
-      # without waiting for a response, then starts the main request span.
+      # Emits a `"start_request.http"` event before the request, then wraps
+      # the exchange in a `"request.http"` span that is guaranteed to close
+      # on both success and failure (via the instrumenter's ensure block).
       #
       # @example
-      #   feature.on_request(request)
+      #   feature.around_request(request) { perform_io }
       #
       # @param request [HTTP::Request]
-      # @return [nil]
+      # @yield Executes the HTTP exchange
+      # @yieldreturn [HTTP::Response]
+      # @return [HTTP::Response]
       # @api public
-      def on_request(request)
+      def around_request(request)
         # Emit a separate "start" event, so a logger can print the request
         # being run without waiting for a response
         instrumenter.instrument("start_#{name}", request: request) {} # rubocop:disable Lint/EmptyBlock
-        instrumenter.start(name, request: request)
-        nil
-      end
-
-      # Wraps a response with instrumentation events
-      #
-      # @example
-      #   feature.wrap_response(response)
-      #
-      # @param response [HTTP::Response]
-      # @return [HTTP::Response]
-      # @api public
-      def wrap_response(response)
-        instrumenter.finish(name, response: response)
-        response
+        instrumenter.instrument(name, request: request) do |payload|
+          response = yield request
+          payload[:response] = response if payload
+          response
+        end
       end
 
       # Instruments a request error

sig/http.rbs

@@ -86,6 +86,8 @@ module HTTP
     def perform_once: (Request req, Options options) -> Response
     def perform_with_retry: (Request req, Options options) -> Response
     def notify_features: (Request req, Options options) -> void
+    def perform_exchange: (Request req, Options options) -> Response
+    def around_request: (Request request, Options options) { (Request) -> Response } -> Response
     def wrap_request: (Request req, Options opts) -> Request
     def build_response: (Request req, Options options) -> Response
     def build_wrapped_response: (Request req, Options options) -> Response
@@ -246,6 +248,7 @@ module HTTP
     def wrap_request: (untyped request) -> untyped
     def wrap_response: (untyped response) -> untyped
     def on_request: (untyped _request) -> nil
+    def around_request: (untyped request) { (untyped) -> Response } -> Response
     def on_error: (untyped _request, untyped _error) -> nil
   end
 
@@ -314,8 +317,7 @@ module HTTP
       attr_reader error_name: String
 
       def initialize: (?instrumenter: untyped, ?namespace: String) -> void
-      def on_request: (untyped request) -> nil
-      def wrap_response: (untyped response) -> untyped
+      def around_request: (untyped request) { (untyped) -> Response } -> Response
       def on_error: (untyped request, untyped error) -> untyped
 
       class NullInstrumenter

test/http/client_test.rb

@@ -441,6 +441,92 @@ def on_request(_request)
 
         assert_equal 1, feature_instance.call_count
       end
+
+      it "calls on_request once per retry attempt" do
+        feature_class_on_request =
+          Class.new(HTTP::Feature) do
+            attr_reader :call_count
+
+            def initialize
+              super
+              @call_count = 0
+            end
+
+            def on_request(_request)
+              @call_count += 1
+            end
+          end
+        feature_instance = feature_class_on_request.new
+
+        client.use(test_feature: feature_instance)
+              .retriable(delay: 0, retry_statuses: [500])
+              .request(:get, "#{dummy.endpoint}/retry-2")
+
+        assert_equal 2, feature_instance.call_count
+      end
+
+      it "wraps each retry attempt with around_request" do
+        feature_class_around =
+          Class.new(HTTP::Feature) do
+            attr_reader :events
+
+            def initialize
+              super
+              @events = []
+            end
+
+            def around_request(request)
+              @events << :before
+              yield(request).tap do
+                @events << :after
+              end
+            end
+          end
+        feature_instance = feature_class_around.new
+
+        client.use(test_feature: feature_instance)
+              .retriable(delay: 0, retry_statuses: [500])
+              .request(:get, "#{dummy.endpoint}/retry-2")
+
+        assert_equal %i[before after before after], feature_instance.events
+      end
+
+      it "wraps the exchange with around_request in feature order" do
+        feature_class_around =
+          Class.new(HTTP::Feature) do
+            @order = []
+
+            class << self
+              attr_reader :order
+            end
+
+            def initialize(id:)
+              super()
+              @id = id
+            end
+
+            def around_request(request)
+              self.class.order << "before.#{@id}"
+              yield(request).tap do
+                self.class.order << "after.#{@id}"
+              end
+            end
+          end
+        feature_instance_a = feature_class_around.new(id: "a")
+        feature_instance_b = feature_class_around.new(id: "b")
+        feature_instance_c = feature_class_around.new(id: "c")
+
+        client.use(
+          test_feature_a: feature_instance_a,
+          test_feature_b: feature_instance_b,
+          test_feature_c: feature_instance_c
+        ).request(:get, dummy.endpoint)
+
+        assert_equal(
+          ["before.a", "before.b", "before.c", "after.c", "after.b", "after.a"],
+          feature_class_around.order
+        )
+      end
     end
   end
 

test/http/features/instrumentation_test.rb

@@ -14,19 +14,19 @@ def initialize
       end
 
       def start(_name, payload)
-        output[:start] = payload
+        output[:start] = payload.dup
       end
 
       def finish(_name, payload)
-        output[:finish] = payload
+        output[:finish] = payload.dup
       end
     end
   end
 
   let(:instrumenter) { instrumenter_class.new }
   let(:feature) { HTTP::Features::Instrumentation.new(instrumenter: instrumenter) }
 
-  describe "starting instrumentation" do
+  describe "around_request" do
     let(:request) do
       HTTP::Request.new(
         verb:    :post,
@@ -36,28 +36,40 @@ def finish(_name, payload)
       )
     end
 
-    it "starts the instrumentation span" do
-      feature.on_request(request)
-
-      assert_equal({ request: request }, instrumenter.output[:start])
-    end
-  end
-
-  describe "logging the response" do
     let(:response) do
       HTTP::Response.new(
         version: "1.1",
         status:  200,
         headers: { content_type: "application/json" },
         body:    '{"success": true}',
-        request: HTTP::Request.new(verb: :get, uri: "https://example.com")
+        request: request
       )
     end
 
-    it "logs the response" do
-      feature.wrap_response(response)
+    it "starts the instrumentation span" do
+      feature.around_request(request) { response }
 
-      assert_equal({ response: response }, instrumenter.output[:finish])
+      assert_equal({ request: request }, instrumenter.output[:start])
+    end
+
+    it "finishes the instrumentation span with the response" do
+      feature.around_request(request) { response }
+
+      assert_equal({ request: request, response: response }, instrumenter.output[:finish])
+    end
+
+    it "returns the response from the block" do
+      result = feature.around_request(request) { response }
+
+      assert_same response, result
+    end
+
+    it "finishes the span even when the block raises" do
+      assert_raises(RuntimeError) do
+        feature.around_request(request) { raise "boom" }
+      end
+
+      assert_equal({ request: request }, instrumenter.output[:finish])
     end
   end
 
@@ -100,8 +112,17 @@ def finish(_name, _payload = {}); end
       HTTP::Request.new(verb: :get, uri: "https://example.com/", headers: {})
     end
 
-    it "does not raise LocalJumpError in on_request" do
-      feature.on_request(request)
+    let(:response) do
+      HTTP::Response.new(
+        version: "1.1",
+        status:  200,
+        body:    "",
+        request: request
+      )
+    end
+
+    it "does not raise LocalJumpError in around_request" do
+      feature.around_request(request) { response }
     end
 
     it "does not raise LocalJumpError in on_error" do