[EXPERIMENTAL] allocate unary response writer against call arena

2 files changed, 45 insertions, 12 deletions

diff --git a/include/grpc++/impl/codegen/async_unary_call.h b/include/grpc++/impl/codegen/async_unary_call.h
index b77a16b699..d11986b6ec 100644
--- a/include/grpc++/impl/codegen/async_unary_call.h
+++ b/include/grpc++/impl/codegen/async_unary_call.h
@@ -34,6 +34,7 @@
 #ifndef GRPCXX_IMPL_CODEGEN_ASYNC_UNARY_CALL_H
 #define GRPCXX_IMPL_CODEGEN_ASYNC_UNARY_CALL_H
 
+#include <assert.h>
 #include <grpc++/impl/codegen/call.h>
 #include <grpc++/impl/codegen/channel_interface.h>
 #include <grpc++/impl/codegen/client_context.h>
@@ -41,6 +42,8 @@
 #include <grpc++/impl/codegen/service_type.h>
 #include <grpc++/impl/codegen/status.h>
 
+extern "C" void* grpc_call_arena_alloc(grpc_call* call, size_t size);
+
 namespace grpc {
 
 class CompletionQueue;
@@ -59,19 +62,33 @@ class ClientAsyncResponseReader final
     : public ClientAsyncResponseReaderInterface<R> {
  public:
   template <class W>
-  ClientAsyncResponseReader(ChannelInterface* channel, CompletionQueue* cq,
-                            const RpcMethod& method, ClientContext* context,
-                            const W& request)
-      : context_(context),
-        call_(channel->CreateCall(method, context, cq)),
-        collection_(std::make_shared<CallOpSetCollection>()) {
-    collection_->init_buf_.SetCollection(collection_);
-    collection_->init_buf_.SendInitialMetadata(
+  static ClientAsyncResponseReader* Create(ChannelInterface* channel,
+                                           CompletionQueue* cq,
+                                           const RpcMethod& method,
+                                           ClientContext* context,
+                                           const W& request) {
+    Call call = channel->CreateCall(method, context, cq);
+    ClientAsyncResponseReader* reader = static_cast<ClientAsyncResponseReader*>(
+        grpc_call_arena_alloc(call.call(), sizeof(*reader)));
+    new (&reader->call_) Call(std::move(call));
+    reader->context_ = context;
+    new (&reader->collection_) std::shared_ptr<CallOpSetCollection>(
+        new (grpc_call_arena_alloc(call.call(), sizeof(CallOpSetCollection)))
+            CallOpSetCollection());
+    reader->collection_->init_buf_.SetCollection(reader->collection_);
+    reader->collection_->init_buf_.SendInitialMetadata(
         context->send_initial_metadata_, context->initial_metadata_flags());
     // TODO(ctiller): don't assert
-    GPR_CODEGEN_ASSERT(collection_->init_buf_.SendMessage(request).ok());
-    collection_->init_buf_.ClientSendClose();
-    call_.PerformOps(&collection_->init_buf_);
+    GPR_CODEGEN_ASSERT(
+        reader->collection_->init_buf_.SendMessage(request).ok());
+    reader->collection_->init_buf_.ClientSendClose();
+    reader->call_.PerformOps(&reader->collection_->init_buf_);
+    return reader;
+  }
+
+  // always allocated against a call arena, no memory free required
+  static void operator delete(void* ptr, std::size_t size) {
+    assert(size == sizeof(ClientAsyncResponseReader));
   }
 
   void ReadInitialMetadata(void* tag) {
@@ -99,7 +116,10 @@ class ClientAsyncResponseReader final
   ClientContext* context_;
   Call call_;
 
-  class CallOpSetCollection : public CallOpSetCollectionInterface {
+  // disable operator new
+  static void* operator new(std::size_t size);
+
+  class CallOpSetCollection final : public CallOpSetCollectionInterface {
    public:
     SneakyCallOpSet<CallOpSendInitialMetadata, CallOpSendMessage,
                     CallOpClientSendClose>
@@ -108,6 +128,15 @@ class ClientAsyncResponseReader final
     CallOpSet<CallOpRecvInitialMetadata, CallOpRecvMessage<R>,
               CallOpClientRecvStatus>
         finish_buf_;
+
+    static void* operator new(std::size_t size, void* p) { return p; }
+    static void operator delete(void* ptr, std::size_t size) {
+      assert(size == sizeof(CallOpSetCollection));
+    }
+
+   private:
+    // disable operator new
+    static void* operator new(std::size_t size);
   };
   std::shared_ptr<CallOpSetCollection> collection_;
 };
diff --git a/include/grpc/grpc.h b/include/grpc/grpc.h
index 1b33d48c02..d603f89c28 100644
--- a/include/grpc/grpc.h
+++ b/include/grpc/grpc.h
@@ -198,6 +198,10 @@ GRPCAPI grpc_call *grpc_channel_create_registered_call(
     grpc_completion_queue *completion_queue, void *registered_call_handle,
     gpr_timespec deadline, void *reserved);
 
+/** Allocate memory in the grpc_call arena: this memory is automatically
+    discarded at call completion */
+GRPCAPI void *grpc_call_arena_alloc(grpc_call *call, size_t size);
+
 /** Start a batch of operations defined in the array ops; when complete, post a
     completion of type 'tag' to the completion queue bound to the call.
     The order of ops specified in the batch has no significance.