Skip to content

fix(rest): preserve scan splits and honor polling deadlines #1059

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhjwpku wants to merge 1 commit into
base: main
Choose a base branch
from
+137 −13
Open

fix(rest): preserve scan splits and honor polling deadlines #1059

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 20 additions & 12 deletions src/iceberg/catalog/rest/rest_table_scan.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,9 +200,20 @@ Result<FetchPlanningResultResponse> FetchPlanningResult(
ICEBERG_ASSIGN_OR_RAISE(auto path, ctx.paths->Plan(ctx.identifier, plan_id));

auto delay_ms = kMinSleepMs;
auto start = std::chrono::steady_clock::now();
const auto start = std::chrono::steady_clock::now();
const auto remaining_wait = [&]() -> Result<std::chrono::milliseconds> {
auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
std::chrono::steady_clock::now() - start)
.count();
if (elapsed_ms >= poll_config.max_wait_ms) {
return IOError("Scan planning timed out after {}ms waiting for plan_id={}",
elapsed_ms, plan_id);
}
return std::chrono::milliseconds(poll_config.max_wait_ms - elapsed_ms);
};

for (int32_t retry = 0; retry <= poll_config.max_retries; ++retry) {
ICEBERG_RETURN_UNEXPECTED(remaining_wait());
ICEBERG_ASSIGN_OR_RAISE(auto response,
ctx.client->Get(path, /*params=*/{}, /*headers=*/{},
*PlanErrorHandler::Instance(), *ctx.session));
Expand All @@ -218,15 +229,10 @@ Result<FetchPlanningResultResponse> FetchPlanningResult(
return result;
}
case PlanStatus::kSubmitted: {
auto elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
std::chrono::steady_clock::now() - start)
.count();
if (elapsed_ms >= poll_config.max_wait_ms) {
return IOError("Scan planning timed out after {}ms waiting for plan_id={}",
elapsed_ms, plan_id);
}
ICEBERG_ASSIGN_OR_RAISE(auto remaining, remaining_wait());
if (retry == poll_config.max_retries) break;
std::this_thread::sleep_for(std::chrono::milliseconds(delay_ms));
std::this_thread::sleep_for(
std::min(std::chrono::milliseconds(delay_ms), remaining));
delay_ms = std::min(delay_ms * 2, kMaxSleepMs);
continue;
}
Expand Down Expand Up @@ -278,9 +284,11 @@ class RestFileScanTaskStream final : public FileScanTaskStream {
auto task = std::move(buffer_[buffer_pos_++]);
if (ignore_residuals_ || !task->residual_filter()) {
// The server may leave row filtering to the client.
return std::make_shared<FileScanTask>(
task->data_file(), task->delete_files(),
ignore_residuals_ ? True::Instance() : filter_);
ICEBERG_ASSIGN_OR_RAISE(
task,
FileScanTask::MakeSplit(task->data_file(), task->start(), task->length(),
task->delete_files(),
ignore_residuals_ ? True::Instance() : filter_));
}
return task;
}
Expand Down
118 changes: 117 additions & 1 deletion src/iceberg/test/rest_table_scan_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,11 @@
#include "iceberg/catalog/rest/rest_table_scan.h"

#include <array>
#include <chrono>
#include <memory>
#include <optional>
#include <string>
#include <thread>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
Expand Down Expand Up @@ -935,7 +937,7 @@ class RestScanResidualTest
: public RestIncrementalAppendScanTest,
public ::testing::WithParamInterface<std::tuple<bool, bool, bool, bool, bool>> {};

TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) {
TEST_P(RestScanResidualTest, PreservesSplitRangesFilteringAndDeleteFiles) {
const auto [incremental, submitted, fetch_token, server_residual, ignore_residuals] =
GetParam();
auto filter = Expressions::GreaterThan("id", Literal::Int(10));
Expand All @@ -944,6 +946,8 @@ TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) {
ICEBERG_UNWRAP_OR_FAIL(auto residual_json, iceberg::ToJson(*residual));

auto task = FileTaskJson("s3://bucket/data.parquet");
task["start"] = 25;
task["length"] = 50;
task["delete-file-references"] = {0};
if (server_residual) task["residual-filter"] = residual_json;
auto delete_file = FileTaskJson("s3://bucket/delete.parquet")["data-file"];
Expand Down Expand Up @@ -985,6 +989,9 @@ TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) {
}));
ASSERT_EQ(tasks.size(), 1);
EXPECT_EQ(tasks[0]->data_file()->file_path, "s3://bucket/data.parquet");
EXPECT_TRUE(tasks[0]->is_split());
EXPECT_EQ(tasks[0]->start(), 25);
EXPECT_EQ(tasks[0]->length(), 50);
ASSERT_EQ(tasks[0]->delete_files().size(), 1);
EXPECT_EQ(tasks[0]->delete_files()[0]->file_path, "s3://bucket/delete.parquet");
ASSERT_NE(tasks[0]->residual_filter(), nullptr);
Expand Down Expand Up @@ -1039,6 +1046,115 @@ TEST_P(RestScanRegressionTest, AllColumnStatsIncludeUnprojectedColumns) {
EXPECT_TRUE(tasks.empty());
}

TEST_P(RestScanRegressionTest, PollingSleepRespectsTimeout) {
constexpr std::string_view kSubmitted =
R"({"status":"submitted","plan-id":"timeout-plan"})";
EXPECT_CALL(*mock_client_, Post(_, _, _, _, _))
.WillOnce(Return(HttpResponse(200, std::string(kSubmitted))));
int poll_count = 0;
EXPECT_CALL(*mock_client_, Get(_, _, _, _, _))
.WillRepeatedly([&](auto&&...) -> Result<HttpResponse> {
++poll_count;
return HttpResponse(
200, poll_count == 1 ? std::string(kSubmitted) : R"({"status":"completed"})");
});
EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _))
.WillOnce(Return(HttpResponse(200, "{}")));

auto ctx = MakeContext();
ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}};
auto start = std::chrono::steady_clock::now();
EXPECT_THAT(
PlanScan(GetParam(), std::move(ctx), [](auto&) {}),
::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("timed out")));
auto elapsed = std::chrono::steady_clock::now() - start;
EXPECT_EQ(poll_count, 1);
EXPECT_GE(elapsed, std::chrono::milliseconds(10));
EXPECT_LT(elapsed, std::chrono::milliseconds(500));
}

TEST_P(RestScanRegressionTest, PollingCompletionAfterTimeoutIsPreserved) {
nlohmann::json completed = {
{"status", "completed"},
{"file-scan-tasks",
nlohmann::json::array({FileTaskJson("s3://bucket/completed.parquet")})}};
EXPECT_CALL(*mock_client_, Post(_, _, _, _, _))
.WillOnce(
Return(HttpResponse(200, R"({"status":"submitted","plan-id":"slow-plan"})")));
EXPECT_CALL(*mock_client_, Get(_, _, _, _, _))
.WillOnce([&](auto&&...) -> Result<HttpResponse> {
std::this_thread::sleep_for(std::chrono::milliseconds(20));
return HttpResponse(200, completed.dump());
});
EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0);

auto ctx = MakeContext();
ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}};
ICEBERG_UNWRAP_OR_FAIL(auto tasks, PlanScan(GetParam(), std::move(ctx), [](auto&) {}));
ASSERT_EQ(tasks.size(), 1);
EXPECT_EQ(tasks[0]->data_file()->file_path, "s3://bucket/completed.parquet");
}

TEST_P(RestScanRegressionTest, PollingFailureAfterTimeoutPreservesServerError) {
EXPECT_CALL(*mock_client_, Post(_, _, _, _, _))
.WillOnce(
Return(HttpResponse(200, R"({"status":"submitted","plan-id":"slow-plan"})")));
EXPECT_CALL(*mock_client_, Get(_, _, _, _, _))
.WillOnce([](auto&&...) -> Result<HttpResponse> {
std::this_thread::sleep_for(std::chrono::milliseconds(20));
return HttpResponse(
200,
R"({"status":"failed","error":{"message":"server error","type":"ServerError","code":500}})");
});
EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _))
.WillOnce(Return(HttpResponse(200, "{}")));

auto ctx = MakeContext();
ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}};
EXPECT_THAT(
PlanScan(GetParam(), std::move(ctx), [](auto&) {}),
::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("server error")));
}

TEST_P(RestScanRegressionTest, PollingCancellationAfterTimeoutDoesNotCancelAgain) {
EXPECT_CALL(*mock_client_, Post(_, _, _, _, _))
.WillOnce(
Return(HttpResponse(200, R"({"status":"submitted","plan-id":"slow-plan"})")));
EXPECT_CALL(*mock_client_, Get(_, _, _, _, _))
.WillOnce([](auto&&...) -> Result<HttpResponse> {
std::this_thread::sleep_for(std::chrono::milliseconds(20));
return HttpResponse(200, R"({"status":"cancelled"})");
});
EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0);

auto ctx = MakeContext();
ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}};
EXPECT_THAT(
PlanScan(GetParam(), std::move(ctx), [](auto&) {}),
::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("was cancelled")));
}

TEST_P(RestScanRegressionTest, PollingSubmittedAfterTimeoutReturnsTimeout) {
constexpr std::string_view kSubmitted =
R"({"status":"submitted","plan-id":"slow-plan"})";
EXPECT_CALL(*mock_client_, Post(_, _, _, _, _))
.WillOnce(Return(HttpResponse(200, std::string(kSubmitted))));
EXPECT_CALL(*mock_client_, Get(_, _, _, _, _))
.WillOnce([&](auto&&...) -> Result<HttpResponse> {
std::this_thread::sleep_for(std::chrono::milliseconds(20));
return HttpResponse(200, std::string(kSubmitted));
});
EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _))
.WillOnce(Return(HttpResponse(200, "{}")));

auto ctx = MakeContext();
ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"},
{"rest-scan-planning.poll-num-retries", "0"}};
EXPECT_THAT(
PlanScan(GetParam(), std::move(ctx), [](auto&) {}),
::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("timed out")));
}

TEST_P(RestScanRegressionTest, InvalidPollingConfigFailsBeforePlanning) {
struct PropertyCase {
std::string_view key;
Expand Down
Loading