diff --git a/cpp/src/remote_handle.cpp b/cpp/src/remote_handle.cpp index cffb582ed4..41f0498e7f 100644 --- a/cpp/src/remote_handle.cpp +++ b/cpp/src/remote_handle.cpp @@ -625,6 +625,7 @@ RemoteHandle RemoteHandle::open(std::string url, }; std::unique_ptr endpoint; + std::optional probed_nbytes; if (remote_endpoint_type == RemoteEndpointType::AUTO) { // Try each allowed type in the order of allowlist @@ -634,8 +635,8 @@ RemoteHandle RemoteHandle::open(std::string url, if (endpoint == nullptr) { continue; } if (type == RemoteEndpointType::S3) { // Check connectivity for the credential-based S3 endpoint, and throw an exception if - // failed - endpoint->get_file_size(); + // failed. Reuse this size when constructing the handle to avoid a second HEAD request. + probed_nbytes = endpoint->get_file_size(); } } catch (...) { // If the credential-based S3 endpoint cannot be used to access the URL, try using S3 public @@ -643,7 +644,8 @@ RemoteHandle RemoteHandle::open(std::string url, if (type == RemoteEndpointType::S3 && std::find(allow_list->begin(), allow_list->end(), RemoteEndpointType::S3_PUBLIC) != allow_list->end()) { - endpoint = std::make_unique(url); + endpoint = std::make_unique(url); + probed_nbytes = std::nullopt; } else { throw; } @@ -669,8 +671,11 @@ RemoteHandle RemoteHandle::open(std::string url, std::runtime_error); } - return nbytes.has_value() ? RemoteHandle(std::move(endpoint), nbytes.value()) - : RemoteHandle(std::move(endpoint)); + if (nbytes.has_value()) { return RemoteHandle(std::move(endpoint), nbytes.value()); } + if (probed_nbytes.has_value()) { + return RemoteHandle(std::move(endpoint), probed_nbytes.value()); + } + return RemoteHandle(std::move(endpoint)); } RemoteHandle::RemoteHandle(std::unique_ptr endpoint, std::size_t nbytes) diff --git a/cpp/tests/test_remote_handle.cpp b/cpp/tests/test_remote_handle.cpp index 41d975fd00..38743a9b68 100644 --- a/cpp/tests/test_remote_handle.cpp +++ b/cpp/tests/test_remote_handle.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -67,12 +67,12 @@ class RemoteHandleTest : public testing::Test { // Test unified interface { // Here we pass the 1-byte argument to RemoteHandle::open. For all endpoints except - // kvikio::RemoteEndpointType::S3, this prevents the endpoint constructor from querying - // the file size and sending requests to the server, thus allowing us to use dummy URLs - // for testing purpose. - // For kvikio::RemoteEndpointType::S3, RemoteHandle::open sends HEAD request as a - // connectivity check and will fail on the syntactically valid dummy URL. The - // kvikio::RemoteEndpointType::S3_PUBLIC will then be used as the endpoint. + // kvikio::RemoteEndpointType::S3 in AUTO mode, this prevents querying the file size and + // sending requests to the server, thus allowing us to use dummy URLs for testing. + // For kvikio::RemoteEndpointType::S3 with AUTO, RemoteHandle::open sends a HEAD request + // as a connectivity check (and reuses that size when nbytes is not provided). It will + // fail on the syntactically valid dummy URL, and kvikio::RemoteEndpointType::S3_PUBLIC + // will then be used as the endpoint. auto remote_handle = kvikio::RemoteHandle::open(url, kvikio::RemoteEndpointType::AUTO, std::nullopt, 1); EXPECT_EQ(remote_handle.remote_endpoint_type(), expected_endpoint_type);