Add mutex for stream in ProtoReflectionDescriptorDatabase, fix headers check

5 files changed, 22 insertions, 21 deletions

diff --git a/build.yaml b/build.yaml
index 5465c3b882..420a64bcf7 100644
--- a/build.yaml
+++ b/build.yaml
@@ -914,6 +914,8 @@ libs:
   - extensions/reflection/proto_server_reflection_plugin.cc
   - extensions/reflection/reflection.grpc.pb.cc
   - extensions/reflection/reflection.pb.cc
+  uses:
+  - grpc++_base
 - name: grpc++_test_config
   build: private
   language: c++
diff --git a/test/cpp/util/proto_reflection_descriptor_database.cc b/test/cpp/util/proto_reflection_descriptor_database.cc
index 3963b0c093..6513a14992 100644
--- a/test/cpp/util/proto_reflection_descriptor_database.cc
+++ b/test/cpp/util/proto_reflection_descriptor_database.cc
@@ -31,7 +31,7 @@
  *
  */
 
-#include "proto_reflection_descriptor_database.h"
+#include "test/cpp/util/proto_reflection_descriptor_database.h"
 
 #include <vector>
 
@@ -69,16 +69,14 @@ bool ProtoReflectionDescriptorDatabase::FindFileByName(
   request.set_file_by_filename(filename);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kFileDescriptorResponse) {
     AddFileFromResponse(response.file_descriptor_response());
-    // const google::protobuf::FileDescriptorProto file_proto =
-    //     ParseFileDescriptorProtoResponse(response.file_descriptor_response());
-    // known_files_.insert(file_proto.name());
-    // cached_db_.Add(file_proto);
   } else if (response.message_response_case() ==
              ServerReflectionResponse::MessageResponseCase::kErrorResponse) {
     const ErrorResponse error = response.error_response();
@@ -119,19 +117,14 @@ bool ProtoReflectionDescriptorDatabase::FindFileContainingSymbol(
   request.set_file_containing_symbol(symbol_name);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
-  // Status status = stub_->GetFileContainingSymbol(&ctx, request, &response);
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kFileDescriptorResponse) {
     AddFileFromResponse(response.file_descriptor_response());
-    // const google::protobuf::FileDescriptorProto file_proto =
-    //     ParseFileDescriptorProtoResponse(response.file_descriptor_response());
-    // if (known_files_.find(file_proto.name()) == known_files_.end()) {
-    //   known_files_.insert(file_proto.name());
-    //   cached_db_.Add(file_proto);
-    // }
   } else if (response.message_response_case() ==
              ServerReflectionResponse::MessageResponseCase::kErrorResponse) {
     const ErrorResponse error = response.error_response();
@@ -181,20 +174,14 @@ bool ProtoReflectionDescriptorDatabase::FindFileContainingExtension(
       field_number);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
-  // Status status = stub_->GetFileContainingExtension(&ctx, request,
-  // &response);
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kFileDescriptorResponse) {
     AddFileFromResponse(response.file_descriptor_response());
-    // const google::protobuf::FileDescriptorProto file_proto =
-    //     ParseFileDescriptorProtoResponse(response.file_descriptor_response());
-    // if (known_files_.find(file_proto.name()) == known_files_.end()) {
-    //   known_files_.insert(file_proto.name());
-    //   cached_db_.Add(file_proto);
-    // }
   } else if (response.message_response_case() ==
              ServerReflectionResponse::MessageResponseCase::kErrorResponse) {
     const ErrorResponse error = response.error_response();
@@ -240,8 +227,10 @@ bool ProtoReflectionDescriptorDatabase::FindAllExtensionNumbers(
   request.set_all_extension_numbers_of_type(extendee_type);
   ServerReflectionResponse response;
 
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::
@@ -272,8 +261,11 @@ bool ProtoReflectionDescriptorDatabase::GetServices(
   ServerReflectionRequest request;
   request.set_list_services("");
   ServerReflectionResponse response;
+
+  stream_mutex_.lock();
   GetStream()->Write(request);
   GetStream()->Read(&response);
+  stream_mutex_.unlock();
 
   if (response.message_response_case() ==
       ServerReflectionResponse::MessageResponseCase::kListServicesResponse) {
diff --git a/test/cpp/util/proto_reflection_descriptor_database.h b/test/cpp/util/proto_reflection_descriptor_database.h
index 4bb9c21a92..bc25fb0f5c 100644
--- a/test/cpp/util/proto_reflection_descriptor_database.h
+++ b/test/cpp/util/proto_reflection_descriptor_database.h
@@ -32,6 +32,7 @@
  */
 
 #include <memory>
+#include <mutex>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -98,6 +99,7 @@ class ProtoReflectionDescriptorDatabase
   std::unordered_set<string> missing_symbols_;
   std::unordered_map<string, std::unordered_set<int>> missing_extensions_;
   std::unordered_map<string, std::vector<int>> cached_extension_numbers_;
+  std::mutex stream_mutex_;
 
   google::protobuf::SimpleDescriptorDatabase cached_db_;
 };
diff --git a/tools/run_tests/sanity/check_sources_and_headers.py b/tools/run_tests/sanity/check_sources_and_headers.py
index c028499ca6..0eb804eb5b 100755
--- a/tools/run_tests/sanity/check_sources_and_headers.py
+++ b/tools/run_tests/sanity/check_sources_and_headers.py
@@ -57,6 +57,9 @@ def target_has_header(target, name):
       return True
   if name == 'src/core/lib/profiling/stap_probes.h':
     return True
+  if not name.startswith('extensions') \
+     and target_has_header(target, 'extensions/' + name):
+    return True
   return False
 
 def produces_object(name):
diff --git a/tools/run_tests/sources_and_headers.json b/tools/run_tests/sources_and_headers.json
index f1d344b613..7ad61845d5 100644
--- a/tools/run_tests/sources_and_headers.json
+++ b/tools/run_tests/sources_and_headers.json
@@ -4371,7 +4371,9 @@
     "type": "lib"
   }, 
   {
-    "deps": [], 
+    "deps": [
+      "grpc++_base"
+    ], 
     "headers": [
       "extensions/include/grpc++/impl/proto_server_reflection_plugin.h", 
       "extensions/include/grpc++/impl/reflection.grpc.pb.h",