1 files changed, 27 insertions, 10 deletions

diff --git a/grpc/include/grpcpp/server.h b/grpc/include/grpcpp/server.h
index 5dc73a19..a69e64b4 100644
--- a/grpc/include/grpcpp/server.h
+++ b/grpc/include/grpcpp/server.h
@@ -58,13 +58,13 @@ class ExternalConnectionAcceptorImpl;
 /// \a Server instances.
 class Server : public ServerInterface, private GrpcLibraryCodegen {
  public:
-  ~Server() override;
+  ~Server() ABSL_LOCKS_EXCLUDED(mu_) override;
 
   /// Block until the server shuts down.
   ///
   /// \warning The server must be either shutting down or some other thread must
   /// call \a Shutdown for this function to ever return.
-  void Wait() override;
+  void Wait() ABSL_LOCKS_EXCLUDED(mu_) override;
 
   /// Global callbacks are a set of hooks that are called when server
   /// events occur.  \a SetGlobalCallbacks method is used to register
@@ -203,6 +203,8 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
     health_check_service_ = std::move(service);
   }
 
+  ContextAllocator* context_allocator() { return context_allocator_.get(); }
+
   /// NOTE: This method is not part of the public API for this class.
   bool health_check_service_disabled() const {
     return health_check_service_disabled_;
@@ -240,6 +242,12 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
   /// ownership of theservice. The service must exist for the lifetime of the
   /// Server instance.
   void RegisterCallbackGenericService(CallbackGenericService* service) override;
+
+  void RegisterContextAllocator(
+      std::unique_ptr<ContextAllocator> context_allocator) {
+    context_allocator_ = std::move(context_allocator);
+  }
+
 #else
   /// NOTE: class experimental_registration_type is not part of the public API
   /// of this class
@@ -254,6 +262,11 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
       server_->RegisterCallbackGenericService(service);
     }
 
+    void RegisterContextAllocator(
+        std::unique_ptr<ContextAllocator> context_allocator) override {
+      server_->context_allocator_ = std::move(context_allocator);
+    }
+
    private:
     Server* server_;
   };
@@ -273,13 +286,14 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
   void PerformOpsOnCall(internal::CallOpSetInterface* ops,
                         internal::Call* call) override;
 
-  void ShutdownInternal(gpr_timespec deadline) override;
+  void ShutdownInternal(gpr_timespec deadline)
+      ABSL_LOCKS_EXCLUDED(mu_) override;
 
   int max_receive_message_size() const override {
     return max_receive_message_size_;
   }
 
-  CompletionQueue* CallbackCQ() override;
+  CompletionQueue* CallbackCQ() ABSL_LOCKS_EXCLUDED(mu_) override;
 
   ServerInitializer* initializer();
 
@@ -287,8 +301,8 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
   // the ref count are the running state of the server (take a ref at start and
   // drop it at shutdown) and each running callback RPC.
   void Ref();
-  void UnrefWithPossibleNotify() /* LOCKS_EXCLUDED(mu_) */;
-  void UnrefAndWaitLocked() /* EXCLUSIVE_LOCKS_REQUIRED(mu_) */;
+  void UnrefWithPossibleNotify() ABSL_LOCKS_EXCLUDED(mu_);
+  void UnrefAndWaitLocked() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
   std::vector<std::shared_ptr<internal::ExternalConnectionAcceptorImpl>>
       acceptors_;
@@ -322,10 +336,11 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
   // Server status
   internal::Mutex mu_;
   bool started_;
-  bool shutdown_;
-  bool shutdown_notified_;  // Was notify called on the shutdown_cv_
+  bool shutdown_ ABSL_GUARDED_BY(mu_);
+  bool shutdown_notified_
+      ABSL_GUARDED_BY(mu_);  // Was notify called on the shutdown_cv_
   internal::CondVar shutdown_done_cv_;
-  bool shutdown_done_ = false;
+  bool shutdown_done_ ABSL_GUARDED_BY(mu_) = false;
   std::atomic_int shutdown_refs_outstanding_{1};
 
   internal::CondVar shutdown_cv_;
@@ -342,6 +357,8 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
 
   std::unique_ptr<ServerInitializer> server_initializer_;
 
+  std::unique_ptr<ContextAllocator> context_allocator_;
+
   std::unique_ptr<HealthCheckServiceInterface> health_check_service_;
   bool health_check_service_disabled_;
 
@@ -363,7 +380,7 @@ class Server : public ServerInterface, private GrpcLibraryCodegen {
   // with this server (if any). It is set on the first call to CallbackCQ().
   // It is _not owned_ by the server; ownership belongs with its internal
   // shutdown callback tag (invoked when the CQ is fully shutdown).
-  CompletionQueue* callback_cq_ /* GUARDED_BY(mu_) */ = nullptr;
+  std::atomic<CompletionQueue*> callback_cq_{nullptr};
 
   // List of CQs passed in by user that must be Shutdown only after Server is
   // Shutdown.  Even though this is only used with NDEBUG, instantiate it in all