| 1433 | |
| 1434 | template <typename BatchesMaker> |
| 1435 | void TestBackpressure(BatchesMaker maker, int batch_size, int num_l_batches, |
| 1436 | int num_r0_batches, int num_r1_batches, bool slow_r0) { |
| 1437 | auto l_schema = |
| 1438 | schema({field("time", int32()), field("key", int32()), field("l_value", int32())}); |
| 1439 | auto r0_schema = |
| 1440 | schema({field("time", int32()), field("key", int32()), field("r0_value", int32())}); |
| 1441 | auto r1_schema = |
| 1442 | schema({field("time", int32()), field("key", int32()), field("r1_value", int32())}); |
| 1443 | |
| 1444 | auto make_shift = [&maker, batch_size](int num_batches, |
| 1445 | const std::shared_ptr<Schema>& schema, |
| 1446 | int shift) { |
| 1447 | return maker({[](int row) -> int64_t { return row; }, |
| 1448 | [num_batches](int row) -> int64_t { return row / num_batches; }, |
| 1449 | [shift](int row) -> int64_t { return row * 10 + shift; }}, |
| 1450 | schema, num_batches, batch_size); |
| 1451 | }; |
| 1452 | ASSERT_OK_AND_ASSIGN(auto l_batches, make_shift(num_l_batches, l_schema, 0)); |
| 1453 | ASSERT_OK_AND_ASSIGN(auto r0_batches, make_shift(num_r0_batches, r0_schema, 1)); |
| 1454 | ASSERT_OK_AND_ASSIGN(auto r1_batches, make_shift(num_r1_batches, r1_schema, 2)); |
| 1455 | |
| 1456 | BackpressureCountingNode::Register(); |
| 1457 | RegisterTestNodes(); // for GatedNode |
| 1458 | |
| 1459 | struct BackpressureSourceConfig { |
| 1460 | std::string name_prefix; |
| 1461 | bool is_gated; |
| 1462 | bool is_delayed; |
| 1463 | std::shared_ptr<Schema> schema; |
| 1464 | decltype(l_batches) batches; |
| 1465 | |
| 1466 | std::string name() const { |
| 1467 | return name_prefix + ";" + (is_gated ? "gated" : "ungated"); |
| 1468 | } |
| 1469 | }; |
| 1470 | |
| 1471 | auto gate_ptr = Gate::Make(); |
| 1472 | auto& gate = *gate_ptr; |
| 1473 | GatedNodeOptions gate_options(gate_ptr.get()); |
| 1474 | |
| 1475 | // Two ungated and one gated |
| 1476 | std::vector<BackpressureSourceConfig> source_configs = { |
| 1477 | {"0", false, false, l_schema, l_batches}, |
| 1478 | {"1", true, slow_r0, r0_schema, r0_batches}, |
| 1479 | {"2", false, false, r1_schema, r1_batches}, |
| 1480 | }; |
| 1481 | |
| 1482 | std::vector<BackpressureCounters> bp_counters(source_configs.size()); |
| 1483 | std::vector<Declaration> src_decls; |
| 1484 | std::vector<std::shared_ptr<BackpressureCountingNodeOptions>> bp_options; |
| 1485 | std::vector<Declaration::Input> bp_decls; |
| 1486 | for (size_t i = 0; i < source_configs.size(); i++) { |
| 1487 | const auto& config = source_configs[i]; |
| 1488 | if (config.is_delayed) { |
| 1489 | src_decls.emplace_back( |
| 1490 | "source", |
| 1491 | SourceNodeOptions(config.schema, MakeDelayedGen(config.batches, "slow_source", |
| 1492 | /*delay_sec=*/0.5, |
no test coverage detected