Add shutdown_starting callbacks to tcp_server.

tcp_server_posix_test illustrates how this can be used to implement a
weak referencing mechanism.

3 files changed, 45 insertions, 2 deletions

diff --git a/src/core/iomgr/tcp_server.h b/src/core/iomgr/tcp_server.h
index 5199c25eed..8f3184ff1e 100644
--- a/src/core/iomgr/tcp_server.h
+++ b/src/core/iomgr/tcp_server.h
@@ -34,6 +34,7 @@
 #ifndef GRPC_INTERNAL_CORE_IOMGR_TCP_SERVER_H
 #define GRPC_INTERNAL_CORE_IOMGR_TCP_SERVER_H
 
+#include "src/core/iomgr/closure.h"
 #include "src/core/iomgr/endpoint.h"
 
 /* Forward decl of grpc_tcp_server */
@@ -90,6 +91,12 @@ int grpc_tcp_server_port_fd(grpc_tcp_server *s, unsigned port_index,
 /* Ref s and return s. */
 grpc_tcp_server *grpc_tcp_server_ref(grpc_tcp_server *s);
 
+/* shutdown_starting is called when ref count has reached zero and the server is
+   about to be destroyed. The server will be deleted after it returns. Calling
+   grpc_tcp_server_ref() from it has no effect. */
+void grpc_tcp_server_shutdown_starting_add(grpc_tcp_server *s,
+                                           grpc_closure *shutdown_starting);
+
 /* If the recount drops to zero, delete s, and call (exec_ctx==NULL) or enqueue
    a call (exec_ctx!=NULL) to shutdown_complete. */
 void grpc_tcp_server_unref(grpc_exec_ctx *exec_ctx, grpc_tcp_server *s);
diff --git a/src/core/iomgr/tcp_server_posix.c b/src/core/iomgr/tcp_server_posix.c
index 6a22c08edb..adf14aeb59 100644
--- a/src/core/iomgr/tcp_server_posix.c
+++ b/src/core/iomgr/tcp_server_posix.c
@@ -128,6 +128,9 @@ struct grpc_tcp_server {
   grpc_tcp_listener *tail;
   unsigned nports;
 
+  /* List of closures passed to shutdown_starting_add(). */
+  grpc_closure_list shutdown_starting;
+
   /* shutdown callback */
   grpc_closure *shutdown_complete;
 
@@ -144,6 +147,8 @@ grpc_tcp_server *grpc_tcp_server_create(grpc_closure *shutdown_complete) {
   s->active_ports = 0;
   s->destroyed_ports = 0;
   s->shutdown = 0;
+  s->shutdown_starting.head = NULL;
+  s->shutdown_starting.tail = NULL;
   s->shutdown_complete = shutdown_complete;
   s->on_accept_cb = NULL;
   s->on_accept_cb_arg = NULL;
@@ -585,13 +590,26 @@ grpc_tcp_server *grpc_tcp_server_ref(grpc_tcp_server *s) {
   return s;
 }
 
+void grpc_tcp_server_shutdown_starting_add(grpc_tcp_server *s,
+                                           grpc_closure *shutdown_starting) {
+  gpr_mu_lock(&s->mu);
+  grpc_closure_list_add(&s->shutdown_starting, shutdown_starting, 1);
+  gpr_mu_unlock(&s->mu);
+}
+
 void grpc_tcp_server_unref(grpc_exec_ctx *exec_ctx, grpc_tcp_server *s) {
   if (gpr_unref(&s->refs)) {
+    /* Complete shutdown_starting work before destroying. */
+    grpc_exec_ctx local_exec_ctx = GRPC_EXEC_CTX_INIT;
+    gpr_mu_lock(&s->mu);
+    grpc_exec_ctx_enqueue_list(&local_exec_ctx, &s->shutdown_starting);
+    gpr_mu_unlock(&s->mu);
     if (exec_ctx == NULL) {
-      grpc_exec_ctx local_exec_ctx = GRPC_EXEC_CTX_INIT;
+      grpc_exec_ctx_flush(&local_exec_ctx);
       tcp_server_destroy(&local_exec_ctx, s);
       grpc_exec_ctx_finish(&local_exec_ctx);
     } else {
+      grpc_exec_ctx_finish(&local_exec_ctx);
       tcp_server_destroy(exec_ctx, s);
     }
   }
diff --git a/src/core/iomgr/tcp_server_windows.c b/src/core/iomgr/tcp_server_windows.c
index 46617a6c65..8ee8149f25 100644
--- a/src/core/iomgr/tcp_server_windows.c
+++ b/src/core/iomgr/tcp_server_windows.c
@@ -93,6 +93,9 @@ struct grpc_tcp_server {
   grpc_tcp_listener *head;
   grpc_tcp_listener *tail;
 
+  /* List of closures passed to shutdown_starting_add(). */
+  grpc_closure_list shutdown_starting;
+
   /* shutdown callback */
   grpc_closure *shutdown_complete;
 };
@@ -108,6 +111,8 @@ grpc_tcp_server *grpc_tcp_server_create(grpc_closure *shutdown_complete) {
   s->on_accept_cb_arg = NULL;
   s->head = NULL;
   s->tail = NULL;
+  s->shutdown_starting.head = NULL;
+  s->shutdown_starting.tail = NULL;
   s->shutdown_complete = shutdown_complete;
   return s;
 }
@@ -135,6 +140,13 @@ grpc_tcp_server *grpc_tcp_server_ref(grpc_tcp_server *s) {
   return s;
 }
 
+void grpc_tcp_server_shutdown_starting_add(grpc_tcp_server *s,
+                                           grpc_closure *shutdown_starting) {
+  gpr_mu_lock(&s->mu);
+  grpc_closure_list_add(&s->shutdown_starting, shutdown_starting, 1);
+  gpr_mu_unlock(&s->mu);
+}
+
 static void tcp_server_destroy(grpc_exec_ctx *exec_ctx, grpc_tcp_server *s) {
   int immediately_done = 0;
   grpc_tcp_listener *sp;
@@ -158,11 +170,17 @@ static void tcp_server_destroy(grpc_exec_ctx *exec_ctx, grpc_tcp_server *s) {
 
 void grpc_tcp_server_unref(grpc_exec_ctx *exec_ctx, grpc_tcp_server *s) {
   if (gpr_unref(&s->refs)) {
+    /* Complete shutdown_starting work before destroying. */
+    grpc_exec_ctx local_exec_ctx = GRPC_EXEC_CTX_INIT;
+    gpr_mu_lock(&s->mu);
+    grpc_exec_ctx_enqueue_list(&local_exec_ctx, &s->shutdown_starting);
+    gpr_mu_unlock(&s->mu);
     if (exec_ctx == NULL) {
-      grpc_exec_ctx local_exec_ctx = GRPC_EXEC_CTX_INIT;
+      grpc_exec_ctx_flush(&local_exec_ctx);
       tcp_server_destroy(&local_exec_ctx, s);
       grpc_exec_ctx_finish(&local_exec_ctx);
     } else {
+      grpc_exec_ctx_finish(&local_exec_ctx);
       tcp_server_destroy(exec_ctx, s);
     }
   }