Use stream rpc to ensure all related requests go to a single server.

1 files changed, 165 insertions, 85 deletions

diff --git a/test/cpp/util/proto_reflection_descriptor_database.cc b/test/cpp/util/proto_reflection_descriptor_database.cc
index c2ed93c12b..4ed069ffc2 100644
--- a/test/cpp/util/proto_reflection_descriptor_database.cc
+++ b/test/cpp/util/proto_reflection_descriptor_database.cc
@@ -37,15 +37,21 @@
 
 #include <grpc/support/log.h>
 
+using grpc::reflection::v1alpha::ServerReflection;
+using grpc::reflection::v1alpha::DescriptorDatabaseRequest;
+using grpc::reflection::v1alpha::DescriptorDatabaseResponse;
+using grpc::reflection::v1alpha::ListServiceResponse;
+using grpc::reflection::v1alpha::ErrorResponse;
+
 namespace grpc {
 
 ProtoReflectionDescriptorDatabase::ProtoReflectionDescriptorDatabase(
-    std::unique_ptr<reflection::v1alpha::ServerReflection::Stub> stub)
+    std::unique_ptr<ServerReflection::Stub> stub)
     : stub_(std::move(stub)) {}
 
 ProtoReflectionDescriptorDatabase::ProtoReflectionDescriptorDatabase(
     std::shared_ptr<grpc::Channel> channel)
-    : stub_(reflection::v1alpha::ServerReflection::NewStub(channel)) {}
+    : stub_(ServerReflection::NewStub(channel)) {}
 
 ProtoReflectionDescriptorDatabase::~ProtoReflectionDescriptorDatabase() {}
 
@@ -59,28 +65,40 @@ bool ProtoReflectionDescriptorDatabase::FindFileByName(
     return false;
   }
 
-  ClientContext ctx;
-  reflection::v1alpha::FileNameRequest request;
-  request.set_filename(filename);
-  reflection::v1alpha::FileDescriptorProtoResponse response;
+  DescriptorDatabaseRequest request;
+  request.set_file_by_filename(filename);
+  DescriptorDatabaseResponse response;
+
+  GetStream()->Write(request);
+  GetStream()->Read(&response);
 
-  Status status = stub_->GetFileByName(&ctx, request, &response);
-  if (status.ok()) {
-    // const google::protobuf::FileDescriptorProto* file_proto =
-    //     response.mutable_file_descriptor_proto();
+  if (response.message_response_case() ==
+      DescriptorDatabaseResponse::MessageResponseCase::kFileDescriptorProto) {
     const google::protobuf::FileDescriptorProto file_proto =
-        ParseFileDescriptorProtoResponse(&response);
+        ParseFileDescriptorProtoResponse(response.file_descriptor_proto());
     known_files_.insert(file_proto.name());
     cached_db_.Add(file_proto);
-  } else if (status.error_code() == StatusCode::NOT_FOUND) {
-    gpr_log(GPR_INFO, "NOT_FOUND from server for FindFileByName(%s)",
-            filename.c_str());
+  } else if (response.message_response_case() ==
+             DescriptorDatabaseResponse::MessageResponseCase::kErrorResponse) {
+    const ErrorResponse error = response.error_response();
+    if (error.error_code() == StatusCode::NOT_FOUND) {
+      gpr_log(GPR_INFO, "NOT_FOUND from server for FindFileByName(%s)",
+              filename.c_str());
+    } else {
+      gpr_log(GPR_INFO,
+              "Error on FindFileByName(%s)\n\tError code: %d\n"
+              "\tError Message: %s",
+              filename.c_str(), error.error_code(),
+              error.error_message().c_str());
+    }
   } else {
-    gpr_log(GPR_INFO,
-            "Error on FindFileByName(%s)\n\tError code: %d\n"
-            "\tError Message: %s",
-            filename.c_str(), status.error_code(),
-            status.error_message().c_str());
+    gpr_log(
+        GPR_INFO,
+        "Error on FindFileByName(%s) response type\n"
+        "\tExpecting: %d\n\tReceived: %d",
+        filename.c_str(),
+        DescriptorDatabaseResponse::MessageResponseCase::kFileDescriptorProto,
+        response.message_response_case());
   }
 
   return cached_db_.FindFileByName(filename, output);
@@ -96,31 +114,46 @@ bool ProtoReflectionDescriptorDatabase::FindFileContainingSymbol(
     return false;
   }
 
-  ClientContext ctx;
-  reflection::v1alpha::SymbolRequest request;
-  request.set_symbol(symbol_name);
-  reflection::v1alpha::FileDescriptorProtoResponse response;
+  DescriptorDatabaseRequest request;
+  request.set_file_containing_symbol(symbol_name);
+  DescriptorDatabaseResponse response;
+
+  GetStream()->Write(request);
+  GetStream()->Read(&response);
 
-  Status status = stub_->GetFileContainingSymbol(&ctx, request, &response);
-  if (status.ok()) {
+  // Status status = stub_->GetFileContainingSymbol(&ctx, request, &response);
+  if (response.message_response_case() ==
+      DescriptorDatabaseResponse::MessageResponseCase::kFileDescriptorProto) {
     const google::protobuf::FileDescriptorProto file_proto =
-        ParseFileDescriptorProtoResponse(&response);
+        ParseFileDescriptorProtoResponse(response.file_descriptor_proto());
     if (known_files_.find(file_proto.name()) == known_files_.end()) {
       known_files_.insert(file_proto.name());
       cached_db_.Add(file_proto);
     }
-  } else if (status.error_code() == StatusCode::NOT_FOUND) {
-    missing_symbols_.insert(symbol_name);
-    gpr_log(GPR_INFO, "NOT_FOUND from server for FindFileContainingSymbol(%s)",
-            symbol_name.c_str());
+  } else if (response.message_response_case() ==
+             DescriptorDatabaseResponse::MessageResponseCase::kErrorResponse) {
+    const ErrorResponse error = response.error_response();
+    if (error.error_code() == StatusCode::NOT_FOUND) {
+      missing_symbols_.insert(symbol_name);
+      gpr_log(GPR_INFO,
+              "NOT_FOUND from server for FindFileContainingSymbol(%s)",
+              symbol_name.c_str());
+    } else {
+      gpr_log(GPR_INFO,
+              "Error on FindFileContainingSymbol(%s)\n"
+              "\tError code: %d\n\tError Message: %s",
+              symbol_name.c_str(), error.error_code(),
+              error.error_message().c_str());
+    }
   } else {
-    gpr_log(GPR_INFO,
-            "Error on FindFileContainingSymbol(%s)\n"
-            "\tError code: %d\n\tError Message: %s",
-            symbol_name.c_str(), status.error_code(),
-            status.error_message().c_str());
+    gpr_log(
+        GPR_INFO,
+        "Error on FindFileContainingSymbol(%s) response type\n"
+        "\tExpecting: %d\n\tReceived: %d",
+        symbol_name.c_str(),
+        DescriptorDatabaseResponse::MessageResponseCase::kFileDescriptorProto,
+        response.message_response_case());
   }
-
   return cached_db_.FindFileContainingSymbol(symbol_name, output);
 }
 
@@ -139,35 +172,53 @@ bool ProtoReflectionDescriptorDatabase::FindFileContainingExtension(
     return false;
   }
 
-  ClientContext ctx;
-  reflection::v1alpha::ExtensionRequest request;
-  request.set_containing_type(containing_type);
-  request.set_extension_number(field_number);
-  reflection::v1alpha::FileDescriptorProtoResponse response;
+  DescriptorDatabaseRequest request;
+  request.mutable_file_containing_extension()->set_containing_type(
+      containing_type);
+  request.mutable_file_containing_extension()->set_extension_number(
+      field_number);
+  DescriptorDatabaseResponse response;
 
-  Status status = stub_->GetFileContainingExtension(&ctx, request, &response);
-  if (status.ok()) {
+  GetStream()->Write(request);
+  GetStream()->Read(&response);
+
+  // Status status = stub_->GetFileContainingExtension(&ctx, request,
+  // &response);
+  if (response.message_response_case() ==
+      DescriptorDatabaseResponse::MessageResponseCase::kFileDescriptorProto) {
     const google::protobuf::FileDescriptorProto file_proto =
-        ParseFileDescriptorProtoResponse(&response);
+        ParseFileDescriptorProtoResponse(response.file_descriptor_proto());
     if (known_files_.find(file_proto.name()) == known_files_.end()) {
       known_files_.insert(file_proto.name());
       cached_db_.Add(file_proto);
     }
-  } else if (status.error_code() == StatusCode::NOT_FOUND) {
-    if (missing_extensions_.find(containing_type) ==
-        missing_extensions_.end()) {
-      missing_extensions_[containing_type] = {};
+  } else if (response.message_response_case() ==
+             DescriptorDatabaseResponse::MessageResponseCase::kErrorResponse) {
+    const ErrorResponse error = response.error_response();
+    if (error.error_code() == StatusCode::NOT_FOUND) {
+      if (missing_extensions_.find(containing_type) ==
+          missing_extensions_.end()) {
+        missing_extensions_[containing_type] = {};
+      }
+      missing_extensions_[containing_type].insert(field_number);
+      gpr_log(GPR_INFO,
+              "NOT_FOUND from server for FindFileContainingExtension(%s, %d)",
+              containing_type.c_str(), field_number);
+    } else {
+      gpr_log(GPR_INFO,
+              "Error on FindFileContainingExtension(%s, %d)\n"
+              "\tError code: %d\n\tError Message: %s",
+              containing_type.c_str(), field_number, error.error_code(),
+              error.error_message().c_str());
     }
-    missing_extensions_[containing_type].insert(field_number);
-    gpr_log(GPR_INFO,
-            "NOT_FOUND from server for FindFileContainingExtension(%s, %d)",
-            containing_type.c_str(), field_number);
   } else {
-    gpr_log(GPR_INFO,
-            "Error on FindFileContainingExtension(%s, %d)\n"
-            "\tError code: %d\n\tError Message: %s",
-            containing_type.c_str(), field_number, status.error_code(),
-            status.error_message().c_str());
+    gpr_log(
+        GPR_INFO,
+        "Error on FindFileContainingExtension(%s, %d) response type\n"
+        "\tExpecting: %d\n\tReceived: %d",
+        containing_type.c_str(), field_number,
+        DescriptorDatabaseResponse::MessageResponseCase::kFileDescriptorProto,
+        response.message_response_case());
   }
 
   return cached_db_.FindFileContainingExtension(containing_type, field_number,
@@ -182,57 +233,86 @@ bool ProtoReflectionDescriptorDatabase::FindAllExtensionNumbers(
     return true;
   }
 
-  ClientContext ctx;
-  reflection::v1alpha::TypeRequest request;
-  request.set_type(extendee_type);
-  reflection::v1alpha::ExtensionNumberResponse response;
+  DescriptorDatabaseRequest request;
+  request.set_all_extension_numbers_of_type(extendee_type);
+  DescriptorDatabaseResponse response;
+
+  GetStream()->Write(request);
+  GetStream()->Read(&response);
 
-  Status status = stub_->GetAllExtensionNumbers(&ctx, request, &response);
-  if (status.ok()) {
-    auto number = response.extension_number();
+  // Status status = stub_->GetAllExtensionNumbers(&ctx, request, &response);
+  if (response.message_response_case() ==
+      DescriptorDatabaseResponse::MessageResponseCase::
+          kAllExtensionNumbersResponse) {
+    auto number = response.all_extension_numbers_response().extension_number();
     *output = std::vector<int>(number.begin(), number.end());
     cached_extension_numbers_[extendee_type] = *output;
     return true;
-  } else if (status.error_code() == StatusCode::NOT_FOUND) {
-    gpr_log(GPR_INFO, "NOT_FOUND from server for FindAllExtensionNumbers(%s)",
-            extendee_type.c_str());
-  } else {
-    gpr_log(GPR_INFO,
-            "Error on FindAllExtensionNumbersExtension(%s)\n"
-            "\tError code: %d\n\tError Message: %s",
-            extendee_type.c_str(), status.error_code(),
-            status.error_message().c_str());
+  } else if (response.message_response_case() ==
+             DescriptorDatabaseResponse::MessageResponseCase::kErrorResponse) {
+    const ErrorResponse error = response.error_response();
+    if (error.error_code() == StatusCode::NOT_FOUND) {
+      gpr_log(GPR_INFO, "NOT_FOUND from server for FindAllExtensionNumbers(%s)",
+              extendee_type.c_str());
+    } else {
+      gpr_log(GPR_INFO,
+              "Error on FindAllExtensionNumbersExtension(%s)\n"
+              "\tError code: %d\n\tError Message: %s",
+              extendee_type.c_str(), error.error_code(),
+              error.error_message().c_str());
+    }
   }
   return false;
 }
 
 bool ProtoReflectionDescriptorDatabase::GetServices(
     std::vector<std::string>* output) {
-  ClientContext ctx;
-  reflection::v1alpha::EmptyRequest request;
-  reflection::v1alpha::ListServiceResponse response;
-
-  Status status = stub_->ListService(&ctx, request, &response);
-  if (status.ok()) {
-    for (int i = 0; i < response.services_size(); ++i) {
-      (*output).push_back(response.services(i));
+  DescriptorDatabaseRequest request;
+  request.set_list_services("");
+  DescriptorDatabaseResponse response;
+  GetStream()->Write(request);
+  GetStream()->Read(&response);
+
+  // Status status = stub_->ListService(&ctx, request, &response);
+  if (response.message_response_case() ==
+      DescriptorDatabaseResponse::MessageResponseCase::kListServicesResponse) {
+    const ListServiceResponse ls_response = response.list_services_response();
+    for (int i = 0; i < ls_response.service_size(); ++i) {
+      (*output).push_back(ls_response.service(i));
     }
     return true;
-  } else {
+  } else if (response.message_response_case() ==
+             DescriptorDatabaseResponse::MessageResponseCase::kErrorResponse) {
+    const ErrorResponse error = response.error_response();
     gpr_log(GPR_INFO,
             "Error on GetServices()\n\tError code: %d\n"
             "\tError Message: %s",
-            status.error_code(), status.error_message().c_str());
+            error.error_code(), error.error_message().c_str());
+  } else {
+    gpr_log(
+        GPR_INFO,
+        "Error on GetServices() response type\n\tExpecting: %d\n\tReceived: %d",
+        DescriptorDatabaseResponse::MessageResponseCase::kListServicesResponse,
+        response.message_response_case());
   }
   return false;
 }
 
 const google::protobuf::FileDescriptorProto
 ProtoReflectionDescriptorDatabase::ParseFileDescriptorProtoResponse(
-    reflection::v1alpha::FileDescriptorProtoResponse* response) {
+    const std::string& byte_fd_proto) {
   google::protobuf::FileDescriptorProto file_desc_proto;
-  file_desc_proto.ParseFromString(response->file_descriptor_proto());
+  file_desc_proto.ParseFromString(byte_fd_proto);
   return file_desc_proto;
 }
 
+const std::shared_ptr<ProtoReflectionDescriptorDatabase::ClientStream>
+ProtoReflectionDescriptorDatabase::GetStream() {
+  if (stream_ == nullptr) {
+    stream_ = stub_->DescriptorDatabaseInfo(&ctx_);
+    // stream_.reset(std::move(stub_->DescriptorDatabaseInfo(&ctx_)));
+  }
+  return stream_;
+}
+
 }  // namespace grpc