Skip to content

Commit a886bd9

Browse files
JaySon-HuangJaySon-Huang
authored
disagg: honor pager early-stop on truncated list (#10762)
close #10761 fix(s3): honor pager early-stop on truncated list Treat pager `more=false` as an explicit termination condition for both `listPrefix` and `listPrefixWithDelimiter`, so truncated listings do not re-fetch the same page indefinitely. Add a regression unit test that creates a multi-page prefix and requests early stop on the first callback; verify listing exits immediately. Signed-off-by: JaySon-Huang <tshent@qq.com>
1 parent 7fe4b3b commit a886bd9

3 files changed

Lines changed: 139 additions & 19 deletions

File tree

‎dbms/src/Storages/S3/MockS3Client.cpp‎

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,8 @@ Model::DeleteObjectOutcome MockS3Client::DeleteObject(const Model::DeleteObjectR
195195

196196
Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Model::ListObjectsV2Request & request) const
197197
{
198+
constexpr int default_max_keys = 1000;
199+
198200
std::lock_guard lock(mtx);
199201
auto itr = storage.find(request.GetBucket());
200202
if (itr == storage.end())
@@ -203,25 +205,60 @@ Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Model::ListObjects
203205
}
204206
const auto & bucket_storage = itr->second;
205207
Model::ListObjectsV2Result result;
206-
RUNTIME_CHECK(!request.DelimiterHasBeenSet() || request.GetDelimiter() == "/", request.GetDelimiter());
208+
if (request.DelimiterHasBeenSet())
209+
RUNTIME_CHECK(request.GetDelimiter() == "/", request.GetDelimiter());
210+
211+
const auto max_keys = request.MaxKeysHasBeenSet() ? request.GetMaxKeys() : default_max_keys;
212+
RUNTIME_CHECK(max_keys > 0, max_keys);
213+
const auto continuation_token = request.ContinuationTokenHasBeenSet() ? request.GetContinuationToken() : String{};
214+
215+
auto finalize_page = [&](const auto & sorted_keys, auto append_result) {
216+
auto begin = sorted_keys.begin();
217+
if (!continuation_token.empty())
218+
begin = sorted_keys.lower_bound(continuation_token);
219+
220+
int key_count = 0;
221+
for (auto iter = begin; iter != sorted_keys.end() && key_count < max_keys; ++iter, ++key_count)
222+
append_result(*iter);
223+
224+
result.SetKeyCount(key_count);
225+
226+
auto next_iter = begin;
227+
for (int i = 0; i < key_count && next_iter != sorted_keys.end(); ++i)
228+
++next_iter;
229+
if (next_iter != sorted_keys.end())
230+
{
231+
result.SetIsTruncated(true);
232+
result.SetNextContinuationToken(*next_iter);
233+
}
234+
else
235+
{
236+
result.SetIsTruncated(false);
237+
}
238+
};
207239

208240
auto normalized_prefix = normalizedKey(request.GetPrefix());
209241
if (!request.DelimiterHasBeenSet())
210242
{
243+
std::set<String> matched_keys;
211244
for (auto itr_obj = bucket_storage.lower_bound(normalized_prefix); itr_obj != bucket_storage.end(); ++itr_obj)
212245
{
213246
if (startsWith(itr_obj->first, normalized_prefix))
214247
{
215-
Model::Object obj;
216-
obj.SetKey(itr_obj->first);
217-
obj.SetSize(itr_obj->second.size());
218-
result.AddContents(std::move(obj));
248+
matched_keys.insert(itr_obj->first);
219249
}
220250
else
221251
{
222252
break;
223253
}
224254
}
255+
256+
finalize_page(matched_keys, [&](const auto & key) {
257+
Model::Object obj;
258+
obj.SetKey(key);
259+
obj.SetSize(bucket_storage.at(key).size());
260+
result.AddContents(std::move(obj));
261+
});
225262
}
226263
else
227264
{
@@ -237,10 +274,9 @@ Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Model::ListObjects
237274
continue;
238275
common_prefix.insert(key.substr(0, pos + delimiter.size()));
239276
}
240-
for (const auto & p : common_prefix)
241-
{
242-
result.AddCommonPrefixes(Model::CommonPrefix().WithPrefix(p));
243-
}
277+
finalize_page(common_prefix, [&](const auto & key) {
278+
result.AddCommonPrefixes(Model::CommonPrefix().WithPrefix(key));
279+
});
244280
}
245281
return result;
246282
}

‎dbms/src/Storages/S3/S3Common.cpp‎

Lines changed: 20 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -859,7 +859,7 @@ void listPrefix(
859859
}
860860
GET_METRIC(tiflash_storage_s3_request_seconds, type_list_objects).Observe(sw_list.elapsedSeconds());
861861

862-
PageResult page_res{};
862+
bool should_continue = true;
863863
const auto & result = outcome.GetResult();
864864
auto page_keys = result.GetContents().size();
865865
num_keys += page_keys;
@@ -868,22 +868,27 @@ void listPrefix(
868868
{
869869
if (!need_cut)
870870
{
871-
page_res = pager(object);
871+
should_continue = pager(object).more;
872872
}
873873
else
874874
{
875875
// Copy the `Object` to cut off the `root` from key, the cost should be acceptable :(
876876
auto object_without_root = object;
877877
object_without_root.SetKey(object.GetKey().substr(cut_size, object.GetKey().size()));
878-
page_res = pager(object_without_root);
878+
should_continue = pager(object_without_root).more;
879879
}
880-
if (!page_res.more)
880+
if (!should_continue)
881881
break;
882882
}
883883

884+
if (!should_continue)
885+
{
886+
break;
887+
}
888+
884889
// handle the result size over max size
885890
done = !result.GetIsTruncated();
886-
if (!done && page_res.more)
891+
if (!done)
887892
{
888893
const auto & next_token = result.GetNextContinuationToken();
889894
req.SetContinuationToken(next_token);
@@ -946,30 +951,35 @@ void listPrefixWithDelimiter(
946951
}
947952
GET_METRIC(tiflash_storage_s3_request_seconds, type_list_objects).Observe(sw_list.elapsedSeconds());
948953

949-
PageResult page_res{};
954+
bool should_continue = true;
950955
const auto & result = outcome.GetResult();
951956
auto page_keys = result.GetCommonPrefixes().size();
952957
num_keys += page_keys;
953958
for (const auto & prefix : result.GetCommonPrefixes())
954959
{
955960
if (!need_cut)
956961
{
957-
page_res = pager(prefix);
962+
should_continue = pager(prefix).more;
958963
}
959964
else
960965
{
961966
// Copy the `CommonPrefix` to cut off the `root`, the cost should be acceptable :(
962967
auto prefix_without_root = prefix;
963968
prefix_without_root.SetPrefix(prefix.GetPrefix().substr(cut_size, prefix.GetPrefix().size()));
964-
page_res = pager(prefix_without_root);
969+
should_continue = pager(prefix_without_root).more;
965970
}
966-
if (!page_res.more)
971+
if (!should_continue)
967972
break;
968973
}
969974

975+
if (!should_continue)
976+
{
977+
break;
978+
}
979+
970980
// handle the result size over max size
971981
done = !result.GetIsTruncated();
972-
if (!done && page_res.more)
982+
if (!done)
973983
{
974984
const auto & next_token = result.GetNextContinuationToken();
975985
req.SetContinuationToken(next_token);

‎dbms/src/Storages/S3/tests/gtest_s3client.cpp‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,4 +252,78 @@ try
252252
}
253253
CATCH
254254

255+
TEST_F(S3ClientTest, ListPrefixEarlyStopOnTruncatedResult)
256+
try
257+
{
258+
// Keep key count above S3's default one-page listing size so listing is truncated.
259+
constexpr size_t key_count = 1001;
260+
SCOPE_EXIT({
261+
for (size_t i = 0; i < key_count; ++i)
262+
deleteObject(*client, fmt::format("s999/list_prefix_early_stop/key_{}", i));
263+
});
264+
for (size_t i = 0; i < key_count; ++i)
265+
{
266+
uploadEmptyFile(*client, fmt::format("s999/list_prefix_early_stop/key_{}", i));
267+
}
268+
269+
Aws::S3::Model::ListObjectsV2Request req;
270+
req.WithBucket(client->bucket()).WithPrefix(client->root() + "s999/list_prefix_early_stop/");
271+
auto outcome = client->ListObjectsV2(req);
272+
ASSERT_TRUE(outcome.IsSuccess());
273+
const auto & result = outcome.GetResult();
274+
ASSERT_TRUE(result.GetIsTruncated());
275+
ASSERT_EQ(result.GetContents().size(), 1000);
276+
ASSERT_FALSE(result.GetNextContinuationToken().empty());
277+
278+
size_t visited = 0;
279+
listPrefix(*client, "s999/list_prefix_early_stop/", [&](const Aws::S3::Model::Object & object) {
280+
UNUSED(object);
281+
++visited;
282+
return PageResult{.num_keys = 1, .more = false};
283+
});
284+
285+
ASSERT_EQ(visited, 1);
286+
}
287+
CATCH
288+
289+
TEST_F(S3ClientTest, ListPrefixWithDelimiterEarlyStopOnTruncatedResult)
290+
try
291+
{
292+
// Keep common prefix count above S3's default one-page listing size so listing is truncated.
293+
constexpr size_t prefix_count = 1001;
294+
SCOPE_EXIT({
295+
for (size_t i = 0; i < prefix_count; ++i)
296+
deleteObject(*client, fmt::format("s999/list_prefix_with_delimiter_early_stop/dir_{}/key", i));
297+
});
298+
for (size_t i = 0; i < prefix_count; ++i)
299+
{
300+
uploadEmptyFile(*client, fmt::format("s999/list_prefix_with_delimiter_early_stop/dir_{}/key", i));
301+
}
302+
303+
Aws::S3::Model::ListObjectsV2Request req;
304+
req.WithBucket(client->bucket())
305+
.WithPrefix(client->root() + "s999/list_prefix_with_delimiter_early_stop/")
306+
.WithDelimiter("/");
307+
auto outcome = client->ListObjectsV2(req);
308+
ASSERT_TRUE(outcome.IsSuccess());
309+
const auto & result = outcome.GetResult();
310+
ASSERT_TRUE(result.GetIsTruncated());
311+
ASSERT_EQ(result.GetCommonPrefixes().size(), 1000);
312+
ASSERT_FALSE(result.GetNextContinuationToken().empty());
313+
314+
size_t visited = 0;
315+
listPrefixWithDelimiter(
316+
*client,
317+
"s999/list_prefix_with_delimiter_early_stop/",
318+
"/",
319+
[&](const Aws::S3::Model::CommonPrefix & prefix) {
320+
UNUSED(prefix);
321+
++visited;
322+
return PageResult{.num_keys = 1, .more = false};
323+
});
324+
325+
ASSERT_EQ(visited, 1);
326+
}
327+
CATCH
328+
255329
} // namespace DB::S3::tests

0 commit comments

Comments
 (0)