Merge pull request #16983 from vjpai/compl_op

Arena-allocate the ServerContext::CompletionOp

1 files changed, 29 insertions, 3 deletions

diff --git a/src/cpp/server/server_context.cc b/src/cpp/server/server_context.cc
index b7254b6bb9..bd532a968d 100644
--- a/src/cpp/server/server_context.cc
+++ b/src/cpp/server/server_context.cc
@@ -40,13 +40,29 @@ namespace grpc {
 class ServerContext::CompletionOp final : public internal::CallOpSetInterface {
  public:
   // initial refs: one in the server context, one in the cq
-  CompletionOp()
-      : has_tag_(false),
+  // must ref the call before calling constructor and after deleting this
+  CompletionOp(grpc_call* call)
+      : call_(call),
+        has_tag_(false),
         tag_(nullptr),
         refs_(2),
         finalized_(false),
         cancelled_(0) {}
 
+  // This should always be arena allocated in the call, so override delete.
+  // But this class is not trivially destructible, so must actually call delete
+  // before allowing the arena to be freed
+  static void operator delete(void* ptr, std::size_t size) {
+    assert(size == sizeof(CompletionOp));
+  }
+
+  // This operator should never be called as the memory should be freed as part
+  // of the arena destruction. It only exists to provide a matching operator
+  // delete to the operator new so that some compilers will not complain (see
+  // https://github.com/grpc/grpc/issues/11301) Note at the time of adding this
+  // there are no tests catching the compiler warning.
+  static void operator delete(void*, void*) { assert(0); }
+
   void FillOps(grpc_call* call, grpc_op* ops, size_t* nops) override;
   bool FinalizeResult(void** tag, bool* status) override;
 
@@ -72,6 +88,7 @@ class ServerContext::CompletionOp final : public internal::CallOpSetInterface {
     return finalized_ ? (cancelled_ != 0) : false;
   }
 
+  grpc_call* call_;
   bool has_tag_;
   void* tag_;
   std::mutex mu_;
@@ -84,7 +101,10 @@ void ServerContext::CompletionOp::Unref() {
   std::unique_lock<std::mutex> lock(mu_);
   if (--refs_ == 0) {
     lock.unlock();
+    // Save aside the call pointer before deleting for later unref
+    grpc_call* call = call_;
     delete this;
+    grpc_call_unref(call);
   }
 }
 
@@ -108,7 +128,10 @@ bool ServerContext::CompletionOp::FinalizeResult(void** tag, bool* status) {
   if (!*status) cancelled_ = 1;
   if (--refs_ == 0) {
     lock.unlock();
+    // Save aside the call pointer before deleting for later unref
+    grpc_call* call = call_;
     delete this;
+    grpc_call_unref(call);
   }
   return ret;
 }
@@ -150,7 +173,10 @@ ServerContext::~ServerContext() {
 
 void ServerContext::BeginCompletionOp(internal::Call* call) {
   GPR_ASSERT(!completion_op_);
-  completion_op_ = new CompletionOp();
+  grpc_call_ref(call->call());
+  completion_op_ =
+      new (grpc_call_arena_alloc(call->call(), sizeof(CompletionOp)))
+          CompletionOp(call->call());
   if (has_notify_when_done_tag_) {
     completion_op_->set_tag(async_notify_when_done_tag_);
   }