wreq 1.2.12 → 1.2.13

data/src/rt.rs

@@ -17,29 +17,13 @@ use crate::{
 /// Initialize the global runtime lazily and preserve failures for Ruby.
 static RUNTIME: OnceLock<Result<TokioRuntime, io::Error>> = OnceLock::new();
 
-enum BlockOnError<E> {
-    Interrupted,
-    Future(E),
-}
-
-/// Register fork tracking while the native extension is being loaded.
+/// Reject a child process that inherited an initialized native runtime.
 ///
 /// # Errors
 ///
-/// Returns `Wreq::ForkError` if the platform cannot install its child-process
-/// callback.
-#[cfg(unix)]
-pub fn initialize(ruby: &Ruby) -> Result<(), magnus::Error> {
-    arch::initialize_fork_tracking().map_err(|err| fork_handler_error(ruby, &err))
-}
-
-/// Reject a child process that inherited the loaded native extension.
-///
-/// # Errors
-///
-/// Returns `Wreq::ForkError` when the extension was loaded before the current
-/// process was forked.
-pub fn ensure_current(ruby: &Ruby) -> Result<(), magnus::Error> {
+/// Returns `Wreq::ForkError` when the global runtime belongs to the parent
+/// process.
+fn ensure_runtime_owner(ruby: &Ruby) -> Result<(), magnus::Error> {
     #[cfg(unix)]
     if let Some((owner_pid, current_pid)) = arch::forked_process_ids() {
         return Err(fork_error(ruby, owner_pid, current_pid));
@@ -54,21 +38,25 @@ pub fn ensure_current(ruby: &Ruby) -> Result<(), magnus::Error> {
 /// Block on a future to completion on the current process's global Tokio runtime.
 ///
 /// The future runs without Ruby's GVL, so it must not construct Ruby objects or
-/// Ruby exceptions. Convert Rust errors back into Ruby errors after the GVL has
-/// been reacquired.
+/// Ruby exceptions. Its output is returned unchanged. If that output is a
+/// `Result`, convert its error after this function returns and reacquires the
+/// GVL.
 ///
 /// # Errors
 ///
-/// Returns `Wreq::ForkError` if the extension belongs to a parent process,
+/// Returns `Wreq::ForkError` if the runtime belongs to a parent process,
 /// `Wreq::BuilderError` if the Tokio runtime cannot be initialized,
-/// `Wreq::InterruptError` if Ruby interrupts the request, or the error produced
-/// by `map_err` if the future fails.
-pub fn try_block_on<F, T, E, M>(ruby: &Ruby, future: F, map_err: M) -> Result<T, magnus::Error>
+/// or `Wreq::InterruptError` if Ruby interrupts the operation.
+pub(crate) fn block_on<F>(ruby: &Ruby, future: F) -> Result<F::Output, magnus::Error>
 where
-    F: Future<Output = Result<T, E>>,
-    M: FnOnce(&Ruby, E) -> magnus::Error,
+    F: Future,
 {
-    ensure_current(ruby)?;
+    // Install fork tracking at the same point as the lazy runtime. Loading the
+    // extension alone must not claim the runtime for the parent process.
+    #[cfg(unix)]
+    arch::initialize_fork_tracking().map_err(|err| fork_handler_error(ruby, &err))?;
+
+    ensure_runtime_owner(ruby)?;
     let runtime = RUNTIME
         .get_or_init(|| {
             let mut builder = Builder::new_multi_thread();
@@ -80,15 +68,11 @@ where
         runtime.block_on(async move {
             tokio::select! {
                 biased;
-                _ = flag.cancelled() => Err(BlockOnError::Interrupted),
-                result = future => result.map_err(BlockOnError::Future),
+                _ = flag.cancelled() => None,
+                result = future => Some(result),
             }
         })
     });
 
-    match result {
-        Ok(value) => Ok(value),
-        Err(BlockOnError::Interrupted) => Err(interrupt_error(ruby)),
-        Err(BlockOnError::Future(err)) => Err(map_err(ruby, err)),
-    }
+    result.ok_or_else(|| interrupt_error(ruby))
 }

data/test/fork_test.rb

@@ -7,15 +7,16 @@ require "timeout"
 
 class ForkTest < Minitest::Test
   FORK_ERROR_LABELS = %w[
-    before_runtime
-    invalid_client
-    invalid_request
-    fresh_body_sender
+    module_request
+    fresh_client_request
+    fresh_body_sender_push
     inherited_body_sender_push
     inherited_body_sender_close
     inherited_body_sender_closed
-    fresh_client
     inherited_client
+    inherited_jar
+    inherited_cookie_provider
+    inherited_response_metadata
     inherited_response
     inherited_response_text
     inherited_response_chunks
@@ -26,13 +27,36 @@ class ForkTest < Minitest::Test
     assert_operator Wreq::ForkError, :<, RuntimeError
   end
 
-  def test_loaded_extension_is_rejected_after_fork
+  def test_loaded_extension_can_initialize_runtime_after_fork
+    skip "fork is not supported on this platform" unless Process.respond_to?(:fork)
+
+    stdout, stderr, status = run_fork_script("prefork_runtime.rb")
+
+    assert status.success?, "subprocess failed with #{status.inspect}: #{stderr}"
+    assert_equal "ok\n", stdout
+    assert_match(/loaded_only=ok/, stderr)
+    assert_match(/before_runtime=ok/, stderr)
+    assert_match(/parent_after_children=ok/, stderr)
+    %w[
+      inherited_client_before_runtime
+      inherited_sender_before_runtime
+      inherited_jar_before_runtime
+      inherited_cookie_provider_before_runtime
+    ].each do |label|
+      assert_match(/#{label}=Wreq::ForkError:.*cannot be used after fork/, stderr)
+    end
+    refute_match(/\[BUG\]|segmentation fault|panicked/i, stderr)
+  end
+
+  def test_initialized_runtime_is_rejected_after_fork
     skip "fork is not supported on this platform" unless Process.respond_to?(:fork)
 
     stdout, stderr, status = run_fork_script("fork_safety.rb")
 
     assert status.success?, "subprocess failed with #{status.inspect}: #{stderr}"
     assert_equal "ok\n", stdout
+    assert_match(/non_runtime_construction=ok/, stderr)
+    assert_match(/inherited_snapshots=ok/, stderr)
     FORK_ERROR_LABELS.each do |label|
       assert_match(/#{label}=Wreq::ForkError:.*cannot be used after fork/, stderr)
       assert_match(/#{label}_retry=Wreq::ForkError:.*cannot be used after fork/, stderr)

data/test/scripts/fork_safety.rb

@@ -9,36 +9,22 @@ $stdout.sync = true
 $stderr.sync = true
 
 def expect_fork_error(label)
-  child_pid = fork do
-    2.times do |attempt|
-      attempt_label = attempt.zero? ? label : "#{label}_retry"
-
-      begin
-        Timeout.timeout(5) { yield }
-      rescue Wreq::ForkError => error
-        warn "#{attempt_label}=#{error.class}: #{error.message}"
-        next
-      rescue => error
-        warn "#{attempt_label}=unexpected #{error.class}: #{error.message}"
-        exit! 2
-      end
+  2.times do |attempt|
+    attempt_label = attempt.zero? ? label : "#{label}_retry"
 
-      warn "#{attempt_label}=missing Wreq::ForkError"
-      exit! 3
+    begin
+      yield
+    rescue Wreq::ForkError => error
+      warn "#{attempt_label}=#{error.class}: #{error.message}"
+      next
+    rescue => error
+      abort "#{attempt_label}=unexpected #{error.class}: #{error.message}"
     end
 
-    exit! 0
+    abort "#{attempt_label}=missing Wreq::ForkError"
   end
-
-  _, status = Process.wait2(child_pid)
-  abort "#{label} child failed with #{status.inspect}" unless status.success?
 end
 
-expect_fork_error("before_runtime") { Wreq::Client.new }
-expect_fork_error("invalid_client") { Wreq::Client.new(unknown: true) }
-expect_fork_error("invalid_request") { Wreq.get(1) }
-expect_fork_error("fresh_body_sender") { Wreq::BodySender.new(0) }
-
 server = TCPServer.new("127.0.0.1", 0)
 port = server.addr[1]
 server_pid = fork do
@@ -51,7 +37,7 @@ server_pid = fork do
       while (line = socket.gets)
         break if line == "\r\n"
       end
-      socket.write("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
+      socket.write("HTTP/1.1 200 OK\r\nX-Fork-Test: ok\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
     ensure
       socket.close
     end
@@ -63,29 +49,66 @@ end
 server.close
 
 url = "http://127.0.0.1:#{port}/"
+
+# Start the server process before Tokio creates worker threads in the parent.
+runtime_probe = Wreq::BodySender.new(1)
+runtime_probe.push("warmup")
+
 client = Wreq::Client.new
 abort "parent warm-up failed" unless client.get(url).bytes == "ok"
 
-def build_inherited_objects(client, url)
-  [Wreq::Client.new, Wreq::BodySender.new, client.get(url)]
+inherited_objects = {
+  client: Wreq::Client.new,
+  sender: Wreq::BodySender.new,
+  response: client.get(url),
+  jar: Wreq::Jar.new
+}
+inherited_weak_refs = inherited_objects.values.map { |object| WeakRef.new(object) }
+status_snapshot = inherited_objects[:response].status
+headers_snapshot = inherited_objects[:response].headers
+
+guard_pid = fork do
+  Timeout.timeout(10) do
+    jar = Wreq::Jar.new
+    jar.add("child=1; Path=/", url)
+    abort "fresh child jar failed" unless jar.get_all.one?
+
+    Wreq::Client.new(cookie_provider: jar)
+    Wreq::BodySender.new
+    warn "non_runtime_construction=ok"
+
+    abort "status snapshot changed" unless status_snapshot.to_i == 200
+    abort "headers snapshot changed" unless headers_snapshot["X-Fork-Test"] == "ok"
+    warn "inherited_snapshots=ok"
+
+    expect_fork_error("module_request") { Wreq.get(url) }
+    expect_fork_error("fresh_client_request") { Wreq::Client.new.get(url) }
+    expect_fork_error("fresh_body_sender_push") { Wreq::BodySender.new.push("chunk") }
+    expect_fork_error("inherited_body_sender_push") do
+      inherited_objects[:sender].push("chunk")
+    end
+    expect_fork_error("inherited_body_sender_close") { inherited_objects[:sender].close }
+    expect_fork_error("inherited_body_sender_closed") { inherited_objects[:sender].closed? }
+    expect_fork_error("inherited_client") { inherited_objects[:client].get(url) }
+    expect_fork_error("inherited_jar") { inherited_objects[:jar].get_all }
+    expect_fork_error("inherited_cookie_provider") do
+      Wreq::Client.new(cookie_provider: inherited_objects[:jar])
+    end
+    expect_fork_error("inherited_response_metadata") { inherited_objects[:response].status }
+    expect_fork_error("inherited_response") { inherited_objects[:response].bytes }
+    expect_fork_error("inherited_response_text") { inherited_objects[:response].text }
+    expect_fork_error("inherited_response_chunks") { inherited_objects[:response].chunks { nil } }
+    expect_fork_error("inherited_response_close") { inherited_objects[:response].close }
+  end
+  exit! 0
+rescue => error
+  warn "guard_checks=unexpected #{error.class}: #{error.message}"
+  exit! 2
 end
+_, guard_status = Process.wait2(guard_pid)
+abort "guard checks child failed with #{guard_status.inspect}" unless guard_status.success?
 
-inherited_objects = build_inherited_objects(client, url)
-inherited_weak_refs = inherited_objects.map { |object| WeakRef.new(object) }
-
-expect_fork_error("inherited_body_sender_push") do
-  inherited_objects[1].push("chunk")
-end
-expect_fork_error("inherited_body_sender_close") { inherited_objects[1].close }
-expect_fork_error("inherited_body_sender_closed") { inherited_objects[1].closed? }
-expect_fork_error("fresh_client") { Wreq::Client.new }
-expect_fork_error("inherited_client") { client.get(url) }
-expect_fork_error("inherited_response") { inherited_objects[2].bytes }
-expect_fork_error("inherited_response_text") { inherited_objects[2].text(1) }
-expect_fork_error("inherited_response_chunks") { inherited_objects[2].chunks }
-expect_fork_error("inherited_response_close") { inherited_objects[2].close }
-
-# Release the earlier test blocks so this array is the only strong reference.
+# The hash is now the only strong reference to the inherited native objects.
 GC.start(full_mark: true, immediate_sweep: true)
 gc_pid = fork do
   inherited_objects = nil

data/test/scripts/prefork_runtime.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require "socket"
+require "timeout"
+require "wreq"
+
+$stdout.sync = true
+$stderr.sync = true
+
+def expect_fork_error(label)
+  yield
+  abort "#{label}=missing Wreq::ForkError"
+rescue Wreq::ForkError => error
+  warn "#{label}=#{error.class}: #{error.message}"
+end
+
+def run_child(label)
+  child_pid = fork do
+    Timeout.timeout(10) { yield }
+    warn "#{label}=ok"
+    exit! 0
+  rescue => error
+    warn "#{label}=unexpected #{error.class}: #{error.message}"
+    exit! 2
+  end
+
+  _, status = Process.wait2(child_pid)
+  abort "#{label} child failed with #{status.inspect}" unless status.success?
+end
+
+server = TCPServer.new("127.0.0.1", 0)
+url = "http://127.0.0.1:#{server.addr[1]}/"
+server_pid = fork do
+  3.times do
+    ready = IO.select([server], nil, nil, 10)
+    exit! 4 unless ready
+
+    socket = server.accept
+    begin
+      while (line = socket.gets)
+        break if line == "\r\n"
+      end
+      socket.write("HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
+    ensure
+      socket.close
+    end
+  end
+  exit! 0
+ensure
+  server.close
+end
+server.close
+
+# Requiring the extension is the only parent-side Wreq operation before this fork.
+run_child("loaded_only") do
+  abort "module request failed" unless Wreq.get(url).bytes == "ok"
+end
+
+inherited_client = Wreq::Client.new
+inherited_sender = Wreq::BodySender.new
+inherited_jar = Wreq::Jar.new
+
+run_child("before_runtime") do
+  expect_fork_error("inherited_client_before_runtime") do
+    inherited_client.get(url)
+  end
+  expect_fork_error("inherited_sender_before_runtime") { inherited_sender.closed? }
+  expect_fork_error("inherited_jar_before_runtime") { inherited_jar.get_all }
+  expect_fork_error("inherited_cookie_provider_before_runtime") do
+    Wreq::Client.new(cookie_provider: inherited_jar)
+  end
+
+  sender = Wreq::BodySender.new
+  sender.push("child")
+
+  jar = Wreq::Jar.new
+  jar.add("child=1; Path=/", url)
+  abort "child jar failed" unless jar.get_all.one?
+
+  client = Wreq::Client.new(cookie_provider: jar)
+  abort "client request failed" unless client.get(url).bytes == "ok"
+end
+
+Timeout.timeout(10) do
+  abort "parent client failed" unless inherited_client.get(url).bytes == "ok"
+  inherited_sender.push("parent")
+  inherited_jar.add("parent=1; Path=/", url)
+  abort "parent jar failed" unless inherited_jar.get_all.one?
+end
+warn "parent_after_children=ok"
+
+_, server_status = Process.wait2(server_pid)
+abort "server failed with #{server_status.inspect}" unless server_status.success?
+
+puts "ok"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: wreq
 version: !ruby/object:Gem::Version
-  version: 1.2.12
+  version: 1.2.13
 platform: ruby
 authors:
 - SearchApi
@@ -164,6 +164,7 @@ files:
 - test/request_test.rb
 - test/response_test.rb
 - test/scripts/fork_safety.rb
+- test/scripts/prefork_runtime.rb
 - test/stream_test.rb
 - test/support/tls_server.rb
 - test/test_helper.rb