[C++][FlightRPC] arrow::flight::FlightClient::AuthenticateBasicToken is blocked against Ballista

[kou]

Describe the bug, including details regarding any error messages, version, and platform.

Ballista uses admin/password for (Basic) authentication: https://arrow.apache.org/ballista/user-guide/flightsql.html#a-name-tool-use-the-driver-in-your-favorite-data-tool

The following Apache Arrow C++ Flight RPC client is blocked:

#include <iostream>

#include <arrow/flight/client.h>

static arrow::Status
run(void)
{
  ARROW_ASSIGN_OR_RAISE(auto location,
                        arrow::flight::Location::Parse("grpc://127.0.0.1:50050"));
  arrow::flight::FlightClientOptions options;
  ARROW_ASSIGN_OR_RAISE(auto client,
                        arrow::flight::FlightClient::Connect(location, options));
  arrow::flight::FlightCallOptions call_options;
  ARROW_ASSIGN_OR_RAISE(auto bearer_token,
                        client->AuthenticateBasicToken(call_options, "admin", "password"));
  return arrow::Status::OK();
}

extern "C" int
main(void)
{
  auto status = run();
  if (!status.ok()) {
    std::cout << status.ToString() << std::endl;
    return 1;
  }
  return 0;
}

stream->Finish() in GrpcClientImpl::AuthenticateBasicToken() is blocked:

https://github.com/apache/arrow/blob/main/cpp/src/arrow/flight/transport/grpc/grpc_client.cc#L762

grpc::ClientReaderWriter< W, R >::Finish document refers grpc::internal::ClientStreamingInterface::Finish() document and it says:

It is appropriate to call this method exactly once when both:

• the calling code (client-side) has no more message to send (this can be declared implicitly by calling this method, or explicitly through an earlier call to WritesDone method of the class in use, e.g. ClientWriterInterface::WritesDone or ClientReaderWriterInterface::WritesDone).
• there are no more messages to be received from the server (which can be known implicitly, or explicitly from an earlier call to ReaderInterface::Read that returned "false").

We call stream->WritesDone() explicitly for the former. But we don't call stream->Read() until it returns false. Should we call stream->Read() before we call stream->Finish() like the following?

diff --git a/cpp/src/arrow/flight/transport/grpc/grpc_client.cc b/cpp/src/arrow/flight/transport/grpc/grpc_client.cc
index 5abecd91a..b382d1e11 100644
--- a/cpp/src/arrow/flight/transport/grpc/grpc_client.cc
+++ b/cpp/src/arrow/flight/transport/grpc/grpc_client.cc
@@ -740,6 +740,15 @@ class GrpcClientImpl : public internal::ClientTransport {
     RETURN_NOT_OK(auth_handler_->Authenticate(&outgoing, &incoming));
     // Explicitly close our side of the connection
     bool finished_writes = stream->WritesDone();
+    pb::HandshakeResponse response;
+    if (!stream->Read(&response)) {
+      return MakeFlightError(FlightStatusCode::Internal,
+                             "No handshake response from server");
+    }
+    if (stream->Read(&response)) {
+      return MakeFlightError(FlightStatusCode::Internal,
+                             "Too much handshake response from server");
+    }
     RETURN_NOT_OK(FromGrpcStatus(stream->Finish(), &rpc.context));
     if (!finished_writes) {
       return MakeFlightError(FlightStatusCode::Internal,
@@ -759,6 +768,15 @@ class GrpcClientImpl : public internal::ClientTransport {
         stream = stub_->Handshake(&rpc.context);
     // Explicitly close our side of the connection.
     bool finished_writes = stream->WritesDone();
+    pb::HandshakeResponse response;
+    if (!stream->Read(&response)) {
+      return MakeFlightError(FlightStatusCode::Internal,
+                             "No handshake response from server");
+    }
+    if (stream->Read(&response)) {
+      return MakeFlightError(FlightStatusCode::Internal,
+                             "Too much handshake response from server");
+    }
     RETURN_NOT_OK(FromGrpcStatus(stream->Finish(), &rpc.context));
     if (!finished_writes) {
       return MakeFlightError(FlightStatusCode::Internal,

If we call stream->Read() before stream->Finish(), the test program isn't blocked.

@lidavidm What do you think about this?

Component(s)

C++, FlightRPC

[lidavidm]

Yes, we need to drain all reads. The rest of the client code is careful to ensure this, but not here apparently.

[lidavidm]

Thanks for digging through this and figuring it out...

[kou]

OK. I'll open a pull request for this.