Skip to content

Commit 171ac57

Browse files
HuaHuaYHuaHuaY
authored
feat: add executor pool support (#687)
1 parent ec20752 commit 171ac57

19 files changed

Lines changed: 1204 additions & 372 deletions

‎cmake_modules/IcebergThirdpartyToolchain.cmake‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ function(resolve_arrow_dependency)
105105
set(ARROW_S3 ${ICEBERG_S3})
106106
set(ARROW_JSON ON)
107107
set(ARROW_PARQUET ON)
108+
set(ARROW_ENABLE_THREADING ON)
108109
set(ARROW_SIMD_LEVEL "NONE")
109110
set(ARROW_RUNTIME_SIMD_LEVEL "NONE")
110111
set(ARROW_POSITION_INDEPENDENT_CODE ON)
@@ -167,8 +168,8 @@ function(resolve_arrow_dependency)
167168
endif()
168169

169170
# Arrow's exported static target interface may reference system libraries
170-
# (e.g. OpenSSL, CURL, ZLIB) that consumers need to find.
171-
list(APPEND ICEBERG_SYSTEM_DEPENDENCIES ZLIB)
171+
# (e.g. Threads, OpenSSL, CURL, ZLIB) that consumers need to find.
172+
list(APPEND ICEBERG_SYSTEM_DEPENDENCIES Threads ZLIB)
172173
if(ARROW_S3)
173174
list(APPEND ICEBERG_SYSTEM_DEPENDENCIES OpenSSL CURL)
174175
endif()

‎src/iceberg/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ set(ICEBERG_SOURCES
114114
util/snapshot_util.cc
115115
util/string_util.cc
116116
util/struct_like_set.cc
117+
util/task_group.cc
117118
util/temporal_util.cc
118119
util/timepoint.cc
119120
util/transform_util.cc

‎src/iceberg/meson.build‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ iceberg_sources = files(
136136
'util/snapshot_util.cc',
137137
'util/string_util.cc',
138138
'util/struct_like_set.cc',
139+
'util/task_group.cc',
139140
'util/temporal_util.cc',
140141
'util/timepoint.cc',
141142
'util/transform_util.cc',

‎src/iceberg/result.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,11 @@
1919

2020
#pragma once
2121

22+
#include <concepts>
2223
#include <expected>
2324
#include <format>
2425
#include <string>
26+
#include <type_traits>
2527

2628
#include "iceberg/iceberg_export.h"
2729

@@ -126,4 +128,11 @@ DEFINE_ERROR_FUNCTION(ValidationFailed)
126128

127129
#undef DEFINE_ERROR_FUNCTION
128130

131+
template <typename T>
132+
concept AsResult = std::derived_from<std::remove_cvref_t<T>,
133+
Result<typename std::remove_cvref_t<T>::value_type>>;
134+
135+
template <AsResult T>
136+
using ResultValueT = typename std::remove_cvref_t<T>::value_type;
137+
129138
} // namespace iceberg

‎src/iceberg/test/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ add_iceberg_test(util_test
136136
retry_util_test.cc
137137
string_util_test.cc
138138
struct_like_set_test.cc
139+
task_group_test.cc
139140
temporal_util_test.cc
140141
transform_util_test.cc
141142
truncate_util_test.cc

‎src/iceberg/test/arrow_test.cc‎

Lines changed: 69 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,28 +17,39 @@
1717
* under the License.
1818
*/
1919

20+
#include <functional>
2021
#include <memory>
22+
#include <mutex>
23+
#include <thread>
2124
#include <unordered_map>
25+
#include <utility>
26+
#include <vector>
2227

2328
#include <arrow/api.h>
2429
#include <arrow/c/bridge.h>
2530
#include <arrow/extension/uuid.h>
2631
#include <arrow/result.h>
2732
#include <arrow/type_fwd.h>
33+
#include <arrow/util/config.h>
2834
#include <arrow/util/key_value_metadata.h>
35+
#include <arrow/util/thread_pool.h>
2936
#include <gtest/gtest.h>
3037

38+
#include "iceberg/arrow/arrow_status_internal.h"
3139
#include "iceberg/constants.h"
40+
#include "iceberg/result.h"
3241
#include "iceberg/schema.h"
3342
#include "iceberg/schema_internal.h"
3443
#include "iceberg/test/matchers.h"
44+
#include "iceberg/util/executor.h"
45+
#include "iceberg/util/task_group.h"
3546

3647
namespace iceberg {
3748

3849
struct ToArrowSchemaParam {
3950
std::shared_ptr<Type> iceberg_type;
4051
bool optional = true;
41-
std::shared_ptr<arrow::DataType> arrow_type;
52+
std::shared_ptr<::arrow::DataType> arrow_type;
4253
};
4354

4455
class ToArrowSchemaTest : public ::testing::TestWithParam<ToArrowSchemaParam> {};
@@ -89,17 +100,17 @@ INSTANTIATE_TEST_SUITE_P(
89100
ToArrowSchemaParam{.iceberg_type = iceberg::date(),
90101
.arrow_type = ::arrow::date32()},
91102
ToArrowSchemaParam{.iceberg_type = iceberg::time(),
92-
.arrow_type = ::arrow::time64(arrow::TimeUnit::MICRO)},
103+
.arrow_type = ::arrow::time64(::arrow::TimeUnit::MICRO)},
93104
ToArrowSchemaParam{.iceberg_type = iceberg::timestamp(),
94-
.arrow_type = ::arrow::timestamp(arrow::TimeUnit::MICRO)},
105+
.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO)},
95106
ToArrowSchemaParam{
96107
.iceberg_type = iceberg::timestamp_tz(),
97-
.arrow_type = ::arrow::timestamp(arrow::TimeUnit::MICRO, "UTC")},
108+
.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC")},
98109
ToArrowSchemaParam{.iceberg_type = iceberg::timestamp_ns(),
99-
.arrow_type = ::arrow::timestamp(arrow::TimeUnit::NANO)},
110+
.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO)},
100111
ToArrowSchemaParam{
101112
.iceberg_type = iceberg::timestamptz_ns(),
102-
.arrow_type = ::arrow::timestamp(arrow::TimeUnit::NANO, "UTC")},
113+
.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC")},
103114
ToArrowSchemaParam{.iceberg_type = iceberg::string(),
104115
.arrow_type = ::arrow::utf8()},
105116
ToArrowSchemaParam{.iceberg_type = iceberg::binary(),
@@ -311,7 +322,7 @@ TEST(ToArrowSchemaTest, NestedUnknownFieldsRoundTrip) {
311322
}
312323

313324
struct FromArrowSchemaParam {
314-
std::shared_ptr<arrow::DataType> arrow_type;
325+
std::shared_ptr<::arrow::DataType> arrow_type;
315326
bool optional = true;
316327
std::shared_ptr<Type> iceberg_type;
317328
};
@@ -365,17 +376,17 @@ INSTANTIATE_TEST_SUITE_P(
365376
.iceberg_type = iceberg::decimal(10, 2)},
366377
FromArrowSchemaParam{.arrow_type = ::arrow::date32(),
367378
.iceberg_type = iceberg::date()},
368-
FromArrowSchemaParam{.arrow_type = ::arrow::time64(arrow::TimeUnit::MICRO),
379+
FromArrowSchemaParam{.arrow_type = ::arrow::time64(::arrow::TimeUnit::MICRO),
369380
.iceberg_type = iceberg::time()},
370-
FromArrowSchemaParam{.arrow_type = ::arrow::timestamp(arrow::TimeUnit::MICRO),
381+
FromArrowSchemaParam{.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO),
371382
.iceberg_type = iceberg::timestamp()},
372383
FromArrowSchemaParam{
373-
.arrow_type = ::arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"),
384+
.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC"),
374385
.iceberg_type = std::make_shared<TimestampTzType>()},
375-
FromArrowSchemaParam{.arrow_type = ::arrow::timestamp(arrow::TimeUnit::NANO),
386+
FromArrowSchemaParam{.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO),
376387
.iceberg_type = iceberg::timestamp_ns()},
377388
FromArrowSchemaParam{
378-
.arrow_type = ::arrow::timestamp(arrow::TimeUnit::NANO, "UTC"),
389+
.arrow_type = ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC"),
379390
.iceberg_type = iceberg::timestamptz_ns()},
380391
FromArrowSchemaParam{.arrow_type = ::arrow::utf8(),
381392
.iceberg_type = iceberg::string()},
@@ -588,4 +599,50 @@ TEST(FromArrowSchemaTest, MapType) {
588599
ASSERT_EQ(value.type()->type_id(), TypeId::kInt);
589600
}
590601

602+
TEST(ArrowExecutorAdapterTest, RunsTaskGroupOnThreadPool) {
603+
#ifndef ARROW_ENABLE_THREADING
604+
GTEST_SKIP() << "Test requires ARROW_ENABLE_THREADING=ON";
605+
#endif
606+
607+
class ArrowExecutorAdapter final : public Executor {
608+
public:
609+
explicit ArrowExecutorAdapter(::arrow::internal::Executor& executor)
610+
: executor_(executor) {}
611+
612+
Status Submit(ExecutorTask task) override {
613+
ICEBERG_ARROW_RETURN_NOT_OK(executor_.Spawn(std::move(task)));
614+
return {};
615+
}
616+
617+
private:
618+
::arrow::internal::Executor& executor_;
619+
};
620+
621+
auto thread_pool = ::arrow::internal::ThreadPool::Make(2).ValueOrDie();
622+
ArrowExecutorAdapter executor(*thread_pool);
623+
624+
std::mutex mutex;
625+
std::vector<std::thread::id> thread_ids;
626+
627+
auto status = TaskGroup<>()
628+
.SetExecutor(std::ref(executor))
629+
.Submit([&]() -> Status {
630+
std::lock_guard lock(mutex);
631+
thread_ids.push_back(std::this_thread::get_id());
632+
return {};
633+
})
634+
.Submit([&]() -> Status {
635+
std::lock_guard lock(mutex);
636+
thread_ids.push_back(std::this_thread::get_id());
637+
return {};
638+
})
639+
.Run();
640+
641+
EXPECT_THAT(status, IsOk());
642+
EXPECT_EQ(thread_ids.size(), 2);
643+
EXPECT_NE(thread_ids[0], std::this_thread::get_id());
644+
EXPECT_NE(thread_ids[1], std::this_thread::get_id());
645+
EXPECT_TRUE(thread_pool->Shutdown().ok());
646+
}
647+
591648
} // namespace iceberg

‎src/iceberg/test/executor.h‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#pragma once
21+
22+
#include <atomic>
23+
#include <thread>
24+
#include <utility>
25+
#include <vector>
26+
27+
#include "iceberg/result.h"
28+
#include "iceberg/util/executor.h"
29+
30+
namespace iceberg::test {
31+
32+
class ThreadExecutor final : public Executor {
33+
public:
34+
explicit ThreadExecutor(Status submit_status = {})
35+
: submit_status_(std::move(submit_status)) {}
36+
37+
~ThreadExecutor() override {
38+
for (auto& thread : threads_) {
39+
if (thread.joinable()) {
40+
thread.join();
41+
}
42+
}
43+
}
44+
45+
Status Submit(ExecutorTask task) override {
46+
submit_count_.fetch_add(1, std::memory_order_relaxed);
47+
if (!submit_status_.has_value()) {
48+
return std::unexpected(submit_status_.error());
49+
}
50+
threads_.emplace_back(std::move(task));
51+
return {};
52+
}
53+
54+
int submit_count() const { return submit_count_.load(std::memory_order_relaxed); }
55+
56+
private:
57+
Status submit_status_;
58+
std::atomic<int> submit_count_{0};
59+
std::vector<std::thread> threads_;
60+
};
61+
62+
} // namespace iceberg::test

‎src/iceberg/test/meson.build‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ iceberg_tests = {
100100
'roaring_position_bitmap_test.cc',
101101
'string_util_test.cc',
102102
'struct_like_set_test.cc',
103+
'task_group_test.cc',
103104
'temporal_util_test.cc',
104105
'transform_util_test.cc',
105106
'truncate_util_test.cc',

‎src/iceberg/test/retry.h‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#pragma once
21+
22+
#include <cstdint>
23+
#include <vector>
24+
25+
#include "iceberg/result.h"
26+
#include "iceberg/util/retry_util.h"
27+
#include "iceberg/util/retry_util_internal.h"
28+
29+
namespace iceberg::test {
30+
31+
using CommitFailedRetry = retry::OnlyRetryOn<ErrorKind::kCommitFailed>;
32+
33+
using TransientIORetry =
34+
retry::OnlyRetryOn<ErrorKind::kIOError, ErrorKind::kServiceUnavailable,
35+
ErrorKind::kInternalServerError>;
36+
37+
class FakeRetryEnvironment {
38+
public:
39+
using Duration = RetryTestHooks::Duration;
40+
using TimePoint = RetryTestHooks::TimePoint;
41+
42+
FakeRetryEnvironment() {
43+
hooks_.now = [this]() { return now_; };
44+
hooks_.sleep_for = [this](Duration duration) {
45+
sleep_durations_.push_back(duration);
46+
now_ += duration;
47+
};
48+
hooks_.jitter = [this](int32_t base_delay_ms) {
49+
observed_base_delays_ms_.push_back(base_delay_ms);
50+
return base_delay_ms + jitter_offset_ms_;
51+
};
52+
}
53+
54+
void Advance(Duration duration) { now_ += duration; }
55+
56+
void SetJitterOffsetMs(int32_t jitter_offset_ms) {
57+
jitter_offset_ms_ = jitter_offset_ms;
58+
}
59+
60+
const RetryTestHooks& hooks() const { return hooks_; }
61+
62+
const std::vector<Duration>& sleep_durations() const { return sleep_durations_; }
63+
64+
const std::vector<int32_t>& observed_base_delays_ms() const {
65+
return observed_base_delays_ms_;
66+
}
67+
68+
private:
69+
RetryTestHooks hooks_;
70+
TimePoint now_{};
71+
int32_t jitter_offset_ms_ = 0;
72+
std::vector<Duration> sleep_durations_;
73+
std::vector<int32_t> observed_base_delays_ms_;
74+
};
75+
76+
} // namespace iceberg::test

0 commit comments

Comments
 (0)