1 files changed, 32 insertions, 2 deletions
diff --git a/src/cpp/server/server_context.cc b/src/cpp/server/server_context.cc
index f6c073040b..04373397f9 100644
--- a/src/cpp/server/server_context.cc
+++ b/src/cpp/server/server_context.cc
@@ -34,6 +34,7 @@
 #include <grpc++/server_context.h>
 
 #include <grpc/grpc.h>
+#include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
 #include <grpc++/impl/call.h>
 #include <grpc++/impl/sync.h>
@@ -49,16 +50,23 @@ namespace grpc {
 class ServerContext::CompletionOp GRPC_FINAL : public CallOpSetInterface {
  public:
   // initial refs: one in the server context, one in the cq
-  CompletionOp() : refs_(2), finalized_(false), cancelled_(0) {}
+  CompletionOp() : has_tag_(false), tag_(nullptr), refs_(2), finalized_(false), cancelled_(0) {}
 
   void FillOps(grpc_op* ops, size_t* nops) GRPC_OVERRIDE;
   bool FinalizeResult(void** tag, bool* status) GRPC_OVERRIDE;
 
   bool CheckCancelled(CompletionQueue* cq);
 
+  void set_tag(void* tag) {
+    has_tag_ = true;
+    tag_ = tag;
+  }
+
   void Unref();
 
  private:
+  bool has_tag_;
+  void* tag_;
   grpc::mutex mu_;
   int refs_;
   bool finalized_;
@@ -89,18 +97,25 @@ void ServerContext::CompletionOp::FillOps(grpc_op* ops, size_t* nops) {
 bool ServerContext::CompletionOp::FinalizeResult(void** tag, bool* status) {
   grpc::unique_lock<grpc::mutex> lock(mu_);
   finalized_ = true;
+  bool ret = false;
+  if (has_tag_) {
+    *tag = tag_;
+    ret = true;
+  }
   if (!*status) cancelled_ = 1;
   if (--refs_ == 0) {
     lock.unlock();
     delete this;
   }
-  return false;
+  return ret;
 }
 
 // ServerContext body
 
 ServerContext::ServerContext()
     : completion_op_(nullptr),
+      has_notify_when_done_tag_(false),
+      async_notify_when_done_tag_(nullptr),
       call_(nullptr),
       cq_(nullptr),
       sent_initial_metadata_(false) {}
@@ -108,6 +123,8 @@ ServerContext::ServerContext()
 ServerContext::ServerContext(gpr_timespec deadline, grpc_metadata* metadata,
                              size_t metadata_count)
     : completion_op_(nullptr),
+      has_notify_when_done_tag_(false),
+      async_notify_when_done_tag_(nullptr),
       deadline_(deadline),
       call_(nullptr),
       cq_(nullptr),
@@ -132,6 +149,9 @@ ServerContext::~ServerContext() {
 void ServerContext::BeginCompletionOp(Call* call) {
   GPR_ASSERT(!completion_op_);
   completion_op_ = new CompletionOp();
+  if (has_notify_when_done_tag_) {
+    completion_op_->set_tag(async_notify_when_done_tag_);
+  }
   call->PerformOps(completion_op_);
 }
 
@@ -179,6 +199,16 @@ std::shared_ptr<const AuthContext> ServerContext::auth_context() const {
   return auth_context_;
 }
 
+grpc::string ServerContext::peer() const {
+  grpc::string peer;
+  if (call_) {
+    char* c_peer = grpc_call_get_peer(call_);
+    peer = c_peer;
+    gpr_free(c_peer);
+  }
+  return peer;
+}
+
 const struct census_context* ServerContext::census_context() const {
   return grpc_census_call_get_context(call_);
 }