Merge branch 'master' into failhijackedrecv

6 files changed, 80 insertions, 35 deletions

diff --git a/src/cpp/server/channelz/channelz_service.cc b/src/cpp/server/channelz/channelz_service.cc
index 9ecb9de7e4..c44a9ac0de 100644
--- a/src/cpp/server/channelz/channelz_service.cc
+++ b/src/cpp/server/channelz/channelz_service.cc
@@ -79,8 +79,8 @@ Status ChannelzService::GetServer(ServerContext* unused,
 Status ChannelzService::GetServerSockets(
     ServerContext* unused, const channelz::v1::GetServerSocketsRequest* request,
     channelz::v1::GetServerSocketsResponse* response) {
-  char* json_str = grpc_channelz_get_server_sockets(request->server_id(),
-                                                    request->start_socket_id());
+  char* json_str = grpc_channelz_get_server_sockets(
+      request->server_id(), request->start_socket_id(), request->max_results());
   if (json_str == nullptr) {
     return Status(StatusCode::INTERNAL,
                   "grpc_channelz_get_server_sockets returned null");
diff --git a/src/cpp/server/health/default_health_check_service.cc b/src/cpp/server/health/default_health_check_service.cc
index c951c69d51..44aebd2f9d 100644
--- a/src/cpp/server/health/default_health_check_service.cc
+++ b/src/cpp/server/health/default_health_check_service.cc
@@ -42,18 +42,37 @@ DefaultHealthCheckService::DefaultHealthCheckService() {
 void DefaultHealthCheckService::SetServingStatus(
     const grpc::string& service_name, bool serving) {
   std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    // Set to NOT_SERVING in case service_name is not in the map.
+    serving = false;
+  }
   services_map_[service_name].SetServingStatus(serving ? SERVING : NOT_SERVING);
 }
 
 void DefaultHealthCheckService::SetServingStatus(bool serving) {
   const ServingStatus status = serving ? SERVING : NOT_SERVING;
   std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    return;
+  }
   for (auto& p : services_map_) {
     ServiceData& service_data = p.second;
     service_data.SetServingStatus(status);
   }
 }
 
+void DefaultHealthCheckService::Shutdown() {
+  std::unique_lock<std::mutex> lock(mu_);
+  if (shutdown_) {
+    return;
+  }
+  shutdown_ = true;
+  for (auto& p : services_map_) {
+    ServiceData& service_data = p.second;
+    service_data.SetServingStatus(NOT_SERVING);
+  }
+}
+
 DefaultHealthCheckService::ServingStatus
 DefaultHealthCheckService::GetServingStatus(
     const grpc::string& service_name) const {
diff --git a/src/cpp/server/health/default_health_check_service.h b/src/cpp/server/health/default_health_check_service.h
index 450bd543f5..9551cd2e2c 100644
--- a/src/cpp/server/health/default_health_check_service.h
+++ b/src/cpp/server/health/default_health_check_service.h
@@ -237,6 +237,8 @@ class DefaultHealthCheckService final : public HealthCheckServiceInterface {
                         bool serving) override;
   void SetServingStatus(bool serving) override;
 
+  void Shutdown() override;
+
   ServingStatus GetServingStatus(const grpc::string& service_name) const;
 
   HealthCheckServiceImpl* GetHealthCheckService(
@@ -272,6 +274,7 @@ class DefaultHealthCheckService final : public HealthCheckServiceInterface {
       const std::shared_ptr<HealthCheckServiceImpl::CallHandler>& handler);
 
   mutable std::mutex mu_;
+  bool shutdown_ = false;                             // Guarded by mu_.
   std::map<grpc::string, ServiceData> services_map_;  // Guarded by mu_.
   std::unique_ptr<HealthCheckServiceImpl> impl_;
 };
diff --git a/src/cpp/server/secure_server_credentials.cc b/src/cpp/server/secure_server_credentials.cc
index ebb17def32..453e76eb25 100644
--- a/src/cpp/server/secure_server_credentials.cc
+++ b/src/cpp/server/secure_server_credentials.cc
@@ -61,7 +61,7 @@ void AuthMetadataProcessorAyncWrapper::InvokeProcessor(
     metadata.insert(std::make_pair(StringRefFromSlice(&md[i].key),
                                    StringRefFromSlice(&md[i].value)));
   }
-  SecureAuthContext context(ctx, false);
+  SecureAuthContext context(ctx);
   AuthMetadataProcessor::OutputMetadata consumed_metadata;
   AuthMetadataProcessor::OutputMetadata response_metadata;
 
diff --git a/src/cpp/server/server_cc.cc b/src/cpp/server/server_cc.cc
index 8b1658dd27..13741ce7aa 100644
--- a/src/cpp/server/server_cc.cc
+++ b/src/cpp/server/server_cc.cc
@@ -84,10 +84,7 @@ class ShutdownTag : public internal::CompletionQueueTag {
 
 class DummyTag : public internal::CompletionQueueTag {
  public:
-  bool FinalizeResult(void** tag, bool* status) {
-    *status = true;
-    return true;
-  }
+  bool FinalizeResult(void** tag, bool* status) { return true; }
 };
 
 class UnimplementedAsyncRequestContext {
@@ -239,9 +236,10 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
                                                 : nullptr),
           request_(nullptr),
           method_(mrd->method_),
-          call_(mrd->call_, server, &cq_, server->max_receive_message_size(),
-                ctx_.set_server_rpc_info(method_->name(),
-                                         server->interceptor_creators_)),
+          call_(
+              mrd->call_, server, &cq_, server->max_receive_message_size(),
+              ctx_.set_server_rpc_info(method_->name(), method_->method_type(),
+                                       server->interceptor_creators_)),
           server_(server),
           global_callbacks_(nullptr),
           resources_(false) {
@@ -294,7 +292,7 @@ class Server::SyncRequest final : public internal::CompletionQueueTag {
 
     void ContinueRunAfterInterception() {
       {
-        ctx_.BeginCompletionOp(&call_, false);
+        ctx_.BeginCompletionOp(&call_, nullptr, nullptr);
         global_callbacks_->PreSynchronousRequest(&ctx_);
         auto* handler = resources_ ? method_->handler()
                                    : server_->resource_exhausted_handler_.get();
@@ -430,7 +428,8 @@ class Server::CallbackRequest final : public internal::CompletionQueueTag {
               req_->call_, req_->server_, req_->cq_,
               req_->server_->max_receive_message_size(),
               req_->ctx_.set_server_rpc_info(
-                  req_->method_->name(), req_->server_->interceptor_creators_));
+                  req_->method_->name(), req_->method_->method_type(),
+                  req_->server_->interceptor_creators_));
 
       req_->interceptor_methods_.SetCall(call_);
       req_->interceptor_methods_.SetReverse();
@@ -459,7 +458,6 @@ class Server::CallbackRequest final : public internal::CompletionQueueTag {
       }
     }
     void ContinueRunAfterInterception() {
-      req_->ctx_.BeginCompletionOp(call_, true);
       req_->method_->handler()->RunHandler(
           internal::MethodHandler::HandlerParameter(
               call_, &req_->ctx_, req_->request_, req_->request_status_,
@@ -732,14 +730,15 @@ std::shared_ptr<Channel> Server::InProcessChannel(
   grpc_channel_args channel_args = args.c_channel_args();
   return CreateChannelInternal(
       "inproc", grpc_inproc_channel_create(server_, &channel_args, nullptr),
-      nullptr);
+      std::vector<
+          std::unique_ptr<experimental::ClientInterceptorFactoryInterface>>());
 }
 
 std::shared_ptr<Channel>
 Server::experimental_type::InProcessChannelWithInterceptors(
     const ChannelArguments& args,
-    std::unique_ptr<std::vector<
-        std::unique_ptr<experimental::ClientInterceptorFactoryInterface>>>
+    std::vector<
+        std::unique_ptr<experimental::ClientInterceptorFactoryInterface>>
         interceptor_creators) {
   grpc_channel_args channel_args = args.c_channel_args();
   return CreateChannelInternal(
@@ -1020,7 +1019,7 @@ bool ServerInterface::BaseAsyncRequest::FinalizeResult(void** tag,
     }
   }
   if (*status && call_) {
-    context_->BeginCompletionOp(&call_wrapper_, false);
+    context_->BeginCompletionOp(&call_wrapper_, nullptr, nullptr);
   }
   *tag = tag_;
   if (delete_on_finalize_) {
@@ -1031,7 +1030,7 @@ bool ServerInterface::BaseAsyncRequest::FinalizeResult(void** tag,
 
 void ServerInterface::BaseAsyncRequest::
     ContinueFinalizeResultAfterInterception() {
-  context_->BeginCompletionOp(&call_wrapper_, false);
+  context_->BeginCompletionOp(&call_wrapper_, nullptr, nullptr);
   // Queue a tag which will be returned immediately
   grpc_core::ExecCtx exec_ctx;
   grpc_cq_begin_op(notification_cq_->cq(), this);
@@ -1044,10 +1043,12 @@ void ServerInterface::BaseAsyncRequest::
 ServerInterface::RegisteredAsyncRequest::RegisteredAsyncRequest(
     ServerInterface* server, ServerContext* context,
     internal::ServerAsyncStreamingInterface* stream, CompletionQueue* call_cq,
-    ServerCompletionQueue* notification_cq, void* tag, const char* name)
+    ServerCompletionQueue* notification_cq, void* tag, const char* name,
+    internal::RpcMethod::RpcType type)
     : BaseAsyncRequest(server, context, stream, call_cq, notification_cq, tag,
                        true),
-      name_(name) {}
+      name_(name),
+      type_(type) {}
 
 void ServerInterface::RegisteredAsyncRequest::IssueRequest(
     void* registered_method, grpc_byte_buffer** payload,
@@ -1094,6 +1095,7 @@ bool ServerInterface::GenericAsyncRequest::FinalizeResult(void** tag,
       call_, server_, call_cq_, server_->max_receive_message_size(),
       context_->set_server_rpc_info(
           static_cast<GenericServerContext*>(context_)->method_.c_str(),
+          internal::RpcMethod::BIDI_STREAMING,
           *server_->interceptor_creators()));
   return BaseAsyncRequest::FinalizeResult(tag, status);
 }
diff --git a/src/cpp/server/server_context.cc b/src/cpp/server/server_context.cc
index 396996e5bc..1b524bc3e8 100644
--- a/src/cpp/server/server_context.cc
+++ b/src/cpp/server/server_context.cc
@@ -17,6 +17,7 @@
  */
 
 #include <grpcpp/server_context.h>
+#include <grpcpp/support/server_callback.h>
 
 #include <algorithm>
 #include <mutex>
@@ -41,8 +42,9 @@ class ServerContext::CompletionOp final : public internal::CallOpSetInterface {
  public:
   // initial refs: one in the server context, one in the cq
   // must ref the call before calling constructor and after deleting this
-  CompletionOp(internal::Call* call)
+  CompletionOp(internal::Call* call, internal::ServerReactor* reactor)
       : call_(*call),
+        reactor_(reactor),
         has_tag_(false),
         tag_(nullptr),
         core_cq_tag_(this),
@@ -124,9 +126,9 @@ class ServerContext::CompletionOp final : public internal::CallOpSetInterface {
       return;
     }
     /* Start a dummy op so that we can return the tag */
-    GPR_CODEGEN_ASSERT(GRPC_CALL_OK ==
-                       g_core_codegen_interface->grpc_call_start_batch(
-                           call_.call(), nullptr, 0, this, nullptr));
+    GPR_CODEGEN_ASSERT(
+        GRPC_CALL_OK ==
+        grpc_call_start_batch(call_.call(), nullptr, 0, core_cq_tag_, nullptr));
   }
 
  private:
@@ -136,13 +138,14 @@ class ServerContext::CompletionOp final : public internal::CallOpSetInterface {
   }
 
   internal::Call call_;
+  internal::ServerReactor* reactor_;
   bool has_tag_;
   void* tag_;
   void* core_cq_tag_;
   std::mutex mu_;
   int refs_;
   bool finalized_;
-  int cancelled_;
+  int cancelled_;  // This is an int (not bool) because it is passed to core
   bool done_intercepting_;
   internal::InterceptorBatchMethodsImpl interceptor_methods_;
 };
@@ -190,7 +193,16 @@ bool ServerContext::CompletionOp::FinalizeResult(void** tag, bool* status) {
   }
   finalized_ = true;
 
-  if (!*status) cancelled_ = 1;
+  // If for some reason the incoming status is false, mark that as a
+  // cancellation.
+  // TODO(vjpai): does this ever happen?
+  if (!*status) {
+    cancelled_ = 1;
+  }
+
+  if (cancelled_ && (reactor_ != nullptr)) {
+    reactor_->OnCancel();
+  }
   /* Release the lock since we are going to be running through interceptors now
    */
   lock.unlock();
@@ -247,21 +259,29 @@ void ServerContext::BindDeadlineAndMetadata(gpr_timespec deadline,
 ServerContext::~ServerContext() { Clear(); }
 
 void ServerContext::Clear() {
-  if (call_) {
-    grpc_call_unref(call_);
-  }
+  auth_context_.reset();
+  initial_metadata_.clear();
+  trailing_metadata_.clear();
+  client_metadata_.Reset();
   if (completion_op_) {
     completion_op_->Unref();
+    completion_op_ = nullptr;
     completion_tag_.Clear();
   }
   if (rpc_info_) {
     rpc_info_->Unref();
+    rpc_info_ = nullptr;
+  }
+  if (call_) {
+    auto* call = call_;
+    call_ = nullptr;
+    grpc_call_unref(call);
   }
-  // Don't need to clear out call_, completion_op_, or rpc_info_ because this is
-  // either called from destructor or just before Setup
 }
 
-void ServerContext::BeginCompletionOp(internal::Call* call, bool callback) {
+void ServerContext::BeginCompletionOp(internal::Call* call,
+                                      std::function<void(bool)> callback,
+                                      internal::ServerReactor* reactor) {
   GPR_ASSERT(!completion_op_);
   if (rpc_info_) {
     rpc_info_->Ref();
@@ -269,10 +289,11 @@ void ServerContext::BeginCompletionOp(internal::Call* call, bool callback) {
   grpc_call_ref(call->call());
   completion_op_ =
       new (grpc_call_arena_alloc(call->call(), sizeof(CompletionOp)))
-          CompletionOp(call);
-  if (callback) {
-    completion_tag_.Set(call->call(), nullptr, completion_op_);
+          CompletionOp(call, reactor);
+  if (callback != nullptr) {
+    completion_tag_.Set(call->call(), std::move(callback), completion_op_);
     completion_op_->set_core_cq_tag(&completion_tag_);
+    completion_op_->set_tag(completion_op_);
   } else if (has_notify_when_done_tag_) {
     completion_op_->set_tag(async_notify_when_done_tag_);
   }