From 7396610151e366e962183ed368a391c7d07cb1e8 Mon Sep 17 00:00:00 2001
From: David Li <li.davidm96@gmail.com>
Date: Fri, 19 Mar 2021 12:00:16 -0400
Subject: [PATCH] ARROW-11908: [C++] Work around flaky integration by retrying

---
 .../arrow/flight/test_integration_client.cc   | 34 +++++++++++++------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/cpp/src/arrow/flight/test_integration_client.cc b/cpp/src/arrow/flight/test_integration_client.cc
index 64da66564bc..6c1d6904603 100644
--- a/cpp/src/arrow/flight/test_integration_client.cc
+++ b/cpp/src/arrow/flight/test_integration_client.cc
@@ -21,9 +21,11 @@
 // client then requests the data from the server and compares it to
 // the data originally uploaded.
 
+#include <chrono>
 #include <iostream>
 #include <memory>
 #include <string>
+#include <thread>
 
 #include <gflags/gflags.h>
 
@@ -202,6 +204,19 @@ class IntegrationTestScenario : public flight::Scenario {
 }  // namespace flight
 }  // namespace arrow
 
+constexpr int kRetries = 3;
+
+arrow::Status RunScenario(arrow::flight::Scenario* scenario) {
+  auto options = arrow::flight::FlightClientOptions::Defaults();
+  std::unique_ptr<arrow::flight::FlightClient> client;
+
+  RETURN_NOT_OK(scenario->MakeClient(&options));
+  arrow::flight::Location location;
+  RETURN_NOT_OK(arrow::flight::Location::ForGrpcTcp(FLAGS_host, FLAGS_port, &location));
+  RETURN_NOT_OK(arrow::flight::FlightClient::Connect(location, options, &client));
+  return scenario->RunClient(std::move(client));
+}
+
 int main(int argc, char** argv) {
   arrow::util::ArrowLog::InstallFailureSignalHandler();
 
@@ -214,16 +229,15 @@ int main(int argc, char** argv) {
     scenario = std::make_shared<arrow::flight::IntegrationTestScenario>();
   }
 
-  arrow::flight::FlightClientOptions options =
-      arrow::flight::FlightClientOptions::Defaults();
-  std::unique_ptr<arrow::flight::FlightClient> client;
-
-  ABORT_NOT_OK(scenario->MakeClient(&options));
-
-  arrow::flight::Location location;
-  ABORT_NOT_OK(arrow::flight::Location::ForGrpcTcp(FLAGS_host, FLAGS_port, &location));
-  ABORT_NOT_OK(arrow::flight::FlightClient::Connect(location, options, &client));
-  ABORT_NOT_OK(scenario->RunClient(std::move(client)));
+  // ARROW-11908: retry a few times in case a client is slow to bring up the server
+  auto status = arrow::Status::OK();
+  for (int i = 0; i < kRetries; i++) {
+    status = RunScenario(scenario.get());
+    if (status.ok()) break;
+    // Failed, wait a bit and try again
+    std::this_thread::sleep_for(std::chrono::milliseconds((i + 1) * 500));
+  }
+  ABORT_NOT_OK(status);
 
   arrow::util::ArrowLog::UninstallSignalAction();
   return 0;