Further simplify the ref-count management for CallbackWrapper

src/python/grpcio/grpc/_cython/_cygrpc/aio/callback_common.pxd.pxi

@@ -33,9 +33,12 @@ cdef struct CallbackContext:
     #       invoked by Core.
     #     failure_handler: A CallbackFailureHandler object that called when Core
     #       returns 'success == 0' state.
+    #     wrapper: A self-reference to the CallbackWrapper to help life cycle
+    #       management.
     grpc_experimental_completion_queue_functor functor
     cpython.PyObject *waiter
     cpython.PyObject *failure_handler
+    cpython.PyObject *callback_wrapper
 
 
 cdef class CallbackWrapper:

src/python/grpcio/grpc/_cython/_cygrpc/aio/callback_common.pyx.pxi

@@ -36,10 +36,12 @@ cdef class CallbackWrapper:
         self.context.functor.functor_run = self.functor_run
         self.context.waiter = <cpython.PyObject*>future
         self.context.failure_handler = <cpython.PyObject*>failure_handler
+        self.context.callback_wrapper = <cpython.PyObject*>self
         # NOTE(lidiz) Not using a list here, because this class is critical in
         # data path. We should make it as efficient as possible.
         self._reference_of_future = future
         self._reference_of_failure_handler = failure_handler
+        cpython.Py_INCREF(self)
 
     @staticmethod
     cdef void functor_run(
@@ -47,12 +49,12 @@ cdef class CallbackWrapper:
             int success):
         cdef CallbackContext *context = <CallbackContext *>functor
         cdef object waiter = <object>context.waiter
-        if waiter.cancelled():
-            return
-        if success == 0:
-            (<CallbackFailureHandler>context.failure_handler).handle(waiter)
-        else:
-            waiter.set_result(None)
+        if not waiter.cancelled():
+            if success == 0:
+                (<CallbackFailureHandler>context.failure_handler).handle(waiter)
+            else:
+                waiter.set_result(None)
+        cpython.Py_DECREF(<object>context.callback_wrapper)
 
     cdef grpc_experimental_completion_queue_functor *c_functor(self):
         return &self.context.functor
@@ -99,9 +101,6 @@ async def execute_batch(GrpcCallWrapper grpc_call_wrapper,
     cdef CallbackWrapper wrapper = CallbackWrapper(
         future,
         CallbackFailureHandler('execute_batch', operations, ExecuteBatchError))
-    # NOTE(lidiz) Without Py_INCREF, the wrapper object will be destructed
-    # when calling "await". This is an over-optimization by Cython.
-    cpython.Py_INCREF(wrapper)
     cdef grpc_call_error error = grpc_call_start_batch(
         grpc_call_wrapper.call,
         batch_operation_tag.c_ops,
@@ -111,10 +110,6 @@ async def execute_batch(GrpcCallWrapper grpc_call_wrapper,
     if error != GRPC_CALL_OK:
         raise ExecuteBatchError("Failed grpc_call_start_batch: {}".format(error))
 
-    # NOTE(lidiz) Guard against CanceledError from future.
-    def dealloc_wrapper(_):
-        cpython.Py_DECREF(wrapper)
-    future.add_done_callback(dealloc_wrapper)
     await future
 
     cdef grpc_event c_event

src/python/grpcio/grpc/_cython/_cygrpc/aio/channel.pyx.pxi

@@ -73,7 +73,6 @@ cdef class AioChannel:
         cdef CallbackWrapper wrapper = CallbackWrapper(
             future,
             _WATCH_CONNECTIVITY_FAILURE_HANDLER)
-        cpython.Py_INCREF(wrapper)
         grpc_channel_watch_connectivity_state(
             self.channel,
             last_observed_state,
@@ -81,13 +80,6 @@ cdef class AioChannel:
             self.cq.c_ptr(),
             wrapper.c_functor())
 
-        # NOTE(lidiz) The callback will be invoked after the channel is closed
-        # with a failure state. We need to keep wrapper alive until then, or we
-        # will observe a segfault.
-        def dealloc_wrapper(_):
-            cpython.Py_DECREF(wrapper)
-        future.add_done_callback(dealloc_wrapper)
-
         try:
             await future
         except _WatchConnectivityFailed:

src/python/grpcio/grpc/_cython/_cygrpc/aio/server.pyx.pxi

@@ -307,9 +307,6 @@ cdef class AioServer:
         cdef CallbackWrapper wrapper = CallbackWrapper(
             future,
             REQUEST_CALL_FAILURE_HANDLER)
-        # NOTE(lidiz) Without Py_INCREF, the wrapper object will be destructed
-        # when calling "await". This is an over-optimization by Cython.
-        cpython.Py_INCREF(wrapper)
         error = grpc_server_request_call(
             self._server.c_server, &rpc_state.call, &rpc_state.details,
             &rpc_state.request_metadata,
@@ -320,7 +317,6 @@ cdef class AioServer:
             raise RuntimeError("Error in grpc_server_request_call: %s" % error)
 
         await future
-        cpython.Py_DECREF(wrapper)
         return rpc_state
 
     async def _server_main_loop(self,