|
27 | 27 | #include "remote_mock.h" |
28 | 28 | #include "storage/record_batch_builder.h" |
29 | 29 | #include "test_utils/random_bytes.h" |
| 30 | +#include "test_utils/scoped_config.h" |
30 | 31 | #include "test_utils/test.h" |
31 | 32 |
|
32 | 33 | #include <seastar/core/abort_source.hh> |
@@ -352,6 +353,80 @@ TEST_CORO(batcher_test, expired_write_request) { |
352 | 353 | } |
353 | 354 | } |
354 | 355 |
|
355 | | -// TODO: add more tests |
356 | | -// - behaviour in case if pending write request sizes exceed L0 object size |
357 | | -// limit |
| 356 | +TEST_CORO(batcher_test, chunk_splitting_balances_upload_sizes) { |
| 357 | + scoped_config cfg; |
| 358 | + // Use a small threshold so test data splits into multiple chunks. |
| 359 | + cfg.get("cloud_topics_produce_batching_size_threshold") |
| 360 | + .set_value(size_t{4096}); |
| 361 | + |
| 362 | + remote_mock mock; |
| 363 | + mock.expect_upload_object_repeatedly(); |
| 364 | + |
| 365 | + cloud_storage_clients::bucket_name bucket("foo"); |
| 366 | + cloud_topics::l0::write_pipeline<ss::manual_clock> pipeline; |
| 367 | + static_cluster_services cluster_services; |
| 368 | + cloud_topics::l0::batcher<ss::manual_clock> batcher( |
| 369 | + pipeline.register_write_pipeline_stage(), |
| 370 | + bucket, |
| 371 | + mock, |
| 372 | + &cluster_services); |
| 373 | + cloud_topics::l0::write_pipeline_accessor pipeline_accessor{ |
| 374 | + .pipeline = &pipeline, |
| 375 | + }; |
| 376 | + |
| 377 | + // Push several write requests. Each has 1 batch with 10 records |
| 378 | + // (~3KB serialized), so 6 requests total ~18KB. With threshold=4096 |
| 379 | + // this should produce multiple balanced chunks. |
| 380 | + const int num_requests = 6; |
| 381 | + std::vector<ss::future<std::expected< |
| 382 | + chunked_vector<cloud_topics::extent_meta>, |
| 383 | + std::error_code>>> |
| 384 | + futures; |
| 385 | + |
| 386 | + const auto timeout = 10s; |
| 387 | + auto deadline = ss::manual_clock::now() + timeout; |
| 388 | + |
| 389 | + for (int i = 0; i < num_requests; i++) { |
| 390 | + auto [_, records, batches] = get_random_batches(1, 10); |
| 391 | + futures.push_back(pipeline.write_and_debounce( |
| 392 | + model::controller_ntp, min_epoch, std::move(batches), deadline)); |
| 393 | + } |
| 394 | + |
| 395 | + // Wait for all write requests to be staged in the pipeline |
| 396 | + // before starting the batcher. subscribe() checks pre-existing |
| 397 | + // pending data, so bg_controller_loop's wait_next will return |
| 398 | + // immediately seeing all requests at once. |
| 399 | + co_await sleep_until(10ms, [&] { |
| 400 | + return pipeline_accessor.write_requests_pending(num_requests); |
| 401 | + }); |
| 402 | + |
| 403 | + // Start the batcher — bg_controller_loop will pull all 6 requests |
| 404 | + // in one batch and split them into balanced chunks. |
| 405 | + co_await batcher.start(); |
| 406 | + |
| 407 | + // Wait for all write request futures to resolve (the batcher |
| 408 | + // loop processes chunks via spawn_with_gate, which sets the |
| 409 | + // promises on each write request). |
| 410 | + auto results = co_await ss::when_all_succeed(std::move(futures)); |
| 411 | + for (auto& res : results) { |
| 412 | + ASSERT_TRUE_CORO(res.has_value()); |
| 413 | + } |
| 414 | + |
| 415 | + co_await batcher.stop(); |
| 416 | + |
| 417 | + // Multiple uploads should happen since total data exceeds threshold. |
| 418 | + ASSERT_GT_CORO(mock.payloads.size(), size_t{1}); |
| 419 | + |
| 420 | + // Verify uploads are balanced: no upload should be excessively small |
| 421 | + // compared to the average. |
| 422 | + size_t total_payload = 0; |
| 423 | + for (const auto& p : mock.payloads) { |
| 424 | + total_payload += p.size(); |
| 425 | + } |
| 426 | + size_t avg_size = total_payload / mock.payloads.size(); |
| 427 | + for (size_t i = 0; i < mock.payloads.size(); i++) { |
| 428 | + EXPECT_GE(mock.payloads[i].size(), avg_size / 3) |
| 429 | + << "Upload " << i << " size " << mock.payloads[i].size() |
| 430 | + << " is too small relative to average " << avg_size; |
| 431 | + } |
| 432 | +} |
0 commit comments