Function DoSinglePerfRun

cpp/src/arrow/flight/flight_benchmark.cc:273–346

}

Status DoSinglePerfRun(FlightClient* client, const FlightClientOptions client_options,
                       const FlightCallOptions& call_options, bool test_put,
                       PerformanceStats* stats) {
  // schema not needed
  perf::Perf perf;
  perf.set_stream_count(FLAGS_num_streams);
  perf.set_records_per_stream(FLAGS_records_per_stream);
  perf.set_records_per_batch(FLAGS_records_per_batch);

  // Plan the query
  FlightDescriptor descriptor;
  descriptor.type = FlightDescriptor::CMD;
  perf.SerializeToString(&descriptor.cmd);

  ARROW_ASSIGN_OR_RAISE(auto plan, client->GetFlightInfo(call_options, descriptor));

  // Read the streams in parallel
  ipc::DictionaryMemo dict_memo;
  ARROW_ASSIGN_OR_RAISE(auto schema, plan->GetSchema(&dict_memo));

  int64_t start_total_records = stats->total_records;

  auto test_loop = test_put ? &RunDoPutTest : &RunDoGetTest;
  auto ConsumeStream = [&client, &stats, &test_loop, &client_options,
                        &call_options](const FlightEndpoint& endpoint) {
    std::unique_ptr<FlightClient> local_client;
    FlightClient* data_client;
    if (endpoint.locations.empty()) {
      data_client = client;
    } else {
      ARROW_ASSIGN_OR_RAISE(
          local_client,
          FlightClient::Connect(endpoint.locations.front(), client_options));
      data_client = local_client.get();
    }

    perf::Token token;
    token.ParseFromString(endpoint.ticket.ticket);

    const auto& result = test_loop(data_client, call_options, token, endpoint, stats);
    if (result.ok()) {
      const PerformanceResult& perf = result.ValueOrDie();
      stats->Update(perf.num_batches, perf.num_records, perf.num_bytes);
    }
    return result.status();
  };

  // XXX(wesm): Serial version for debugging
  // for (const auto& endpoint : plan->endpoints()) {
  //   RETURN_NOT_OK(ConsumeStream(endpoint));
  // }

  ARROW_ASSIGN_OR_RAISE(auto pool, ThreadPool::Make(FLAGS_num_threads));
  std::vector<Future<>> tasks;
  for (const auto& endpoint : plan->endpoints()) {
    ARROW_ASSIGN_OR_RAISE(auto task, pool->Submit(ConsumeStream, endpoint));
    tasks.push_back(std::move(task));
  }