Merge remote-tracking branch 'upstream/master' into core-end2end-tests-cronet-update

1 files changed, 15 insertions, 4 deletions

diff --git a/test/cpp/qps/server_async.cc b/test/cpp/qps/server_async.cc
index 082b4bc72f..bc4c896d83 100644
--- a/test/cpp/qps/server_async.cc
+++ b/test/cpp/qps/server_async.cc
@@ -38,6 +38,7 @@
 #include <thread>
 
 #include <grpc++/generic/async_generic_service.h>
+#include <grpc++/resource_quota.h>
 #include <grpc++/security/server_credentials.h>
 #include <grpc++/server.h>
 #include <grpc++/server_builder.h>
@@ -57,7 +58,7 @@ namespace testing {
 
 template <class RequestType, class ResponseType, class ServiceType,
           class ServerContextType>
-class AsyncQpsServerTest : public Server {
+class AsyncQpsServerTest GRPC_FINAL : public grpc::testing::Server {
  public:
   AsyncQpsServerTest(
       const ServerConfig &config,
@@ -95,6 +96,11 @@ class AsyncQpsServerTest : public Server {
       srv_cqs_.emplace_back(builder.AddCompletionQueue());
     }
 
+    if (config.resource_quota_size() > 0) {
+      builder.SetResourceQuota(ResourceQuota("AsyncQpsServerTest")
+                                   .Resize(config.resource_quota_size()));
+    }
+
     server_ = builder.BuildAndStart();
 
     using namespace std::placeholders;
@@ -131,9 +137,7 @@ class AsyncQpsServerTest : public Server {
       std::lock_guard<std::mutex> lock((*ss)->mutex);
       (*ss)->shutdown = true;
     }
-    // TODO (vpai): Remove this deadline and allow Shutdown to finish properly
-    auto deadline = std::chrono::system_clock::now() + std::chrono::seconds(3);
-    server_->Shutdown(deadline);
+    std::thread shutdown_thread(&AsyncQpsServerTest::ShutdownThreadFunc, this);
     for (auto cq = srv_cqs_.begin(); cq != srv_cqs_.end(); ++cq) {
       (*cq)->Shutdown();
     }
@@ -146,9 +150,16 @@ class AsyncQpsServerTest : public Server {
       while ((*cq)->Next(&got_tag, &ok))
         ;
     }
+    shutdown_thread.join();
   }
 
  private:
+  void ShutdownThreadFunc() {
+    // TODO (vpai): Remove this deadline and allow Shutdown to finish properly
+    auto deadline = std::chrono::system_clock::now() + std::chrono::seconds(3);
+    server_->Shutdown(deadline);
+  }
+
   void ThreadFunc(int thread_idx) {
     // Wait until work is available or we are shutting down
     bool ok;