Tweaked grpc_handshaker API and added grpc_handshake_manager API.

2 files changed, 150 insertions, 6 deletions

diff --git a/src/core/lib/channel/handshaker.c b/src/core/lib/channel/handshaker.c
index 00d810e63d..3018b95464 100644
--- a/src/core/lib/channel/handshaker.c
+++ b/src/core/lib/channel/handshaker.c
@@ -33,6 +33,10 @@
 
 #include "src/core/lib/channel/handshaker.h"
 
+//
+// grpc_handshaker
+//
+
 void grpc_handshaker_init(const struct grpc_handshaker_vtable* vtable,
                           grpc_handshaker** handshaker) {
   handshaker->vtable = vtable;
@@ -43,11 +47,116 @@ void grpc_handshaker_destroy(grpc_exec_ctx* exec_ctx,
   handshaker->vtable->destroy(exec_ctx, handshaker);
 }
 
+void grpc_handshaker_shutdown(grpc_exec_ctx* exec_ctx,
+                              grpc_handshaker* handshaker) {
+  handshaker->vtable->shutdown(exec_ctx, handshaker);
+}
+
 void grpc_handshaker_do_handshake(
     grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker,
-    grpc_endpoint* endpoint_in, gpr_timespec deadline, grpc_closure *on_done) {
-  handshaker->vtable->do_handshake(exec_ctx, handshaker, endpoint_in,
-                                   deadline, on_done);
+    grpc_endpoint* endpoint, gpr_timespec deadline,
+    grpc_handshaker_done_cb cb, void* arg) {
+  handshaker->vtable->do_handshake(exec_ctx, handshaker, endpoint,
+                                   deadline, cb, arg);
+}
+
+//
+// grpc_handshake_manager
+//
+
+// State used while chaining handshakers.
+struct grpc_handshaker_state {
+  // The index of the handshaker to invoke next.
+  size_t index;
+  // The deadline for all handshakers.
+  gpr_timespec deadline;
+  // The final callback and arg to invoke after the last handshaker.
+  grpc_handshaker_done_cb final_cb;
+  void* final_arg;
+};
+
+struct grpc_handshake_manager {
+  // An array of handshakers added via grpc_handshake_manager_add().
+  size_t count;
+  grpc_handshaker* handshakers;
+  // State used while chaining handshakers.
+  grpc_handshaker_state* state;
+};
+
+grpc_handshake_manager* grpc_handshake_manager_create() {
+  grpc_handshake_manager* mgr = gpr_malloc(sizeof(grpc_handshake_manager));
+  memset(mgr, 0, sizeof(*mgr));
+  return mgr;
+}
+
+void grpc_handshake_manager_add(grpc_handshaker* handshaker,
+                                grpc_handshake_manager* mgr) {
+  mgr->handshakers = gpr_realloc(mgr->handshakers,
+                                 (mgr->count + 1) * sizeof(grpc_handshaker*));
+  mgr->handshakers[mgr->count++] = handshaker;
+}
+
+void grpc_handshake_manager_destroy(
+    grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr) {
+  for (int i = 0; i < mgr->count; ++i) {
+    grpc_handshaker_destroy(exec_ctx, mgr->handshakers[i]);
+  }
+  gpr_free(mgr);
+}
+
+void grpc_handshake_manager_shutdown(
+    grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr) {
+  // FIXME: maybe check which handshaker is currently in progress, and
+  // only shut down that one?
+  for (int i = 0; i < mgr->count; ++i) {
+    grpc_handshaker_shutdown(exec_ctx, mgr->handshakers[i]);
+  }
+  if (mgr->state != NULL) {
+    gpr_free(mgr->state);
+    mgr->state = NULL;
+  }
+}
+
+// A function used as the handshaker-done callback when chaining
+// handshakers together.
+static void call_next_handshaker(grpc_exec_ctx* exec_ctx,
+                                 grpc_endpoint* endpoint, void* arg) {
+  grpc_handshake_manager* mgr = arg;
+  GPR_ASSERT(mgr->state != NULL);
+  GPR_ASSERT(mgr->state->index < mgr->count);
+  grpc_handshaker_done_cb cb = call_next_handshaker;
+  // If this is the last handshaker, use the caller-supplied callback
+  // and arg instead of chaining back to this function again.
+  if (mgr->state->index == mgr->count - 1) {
+    cb = mgr->state->final_cb;
+    arg = mgr->state->final_arg;
+  }
+  // Invoke handshaker.
+  grpc_handshaker_do_handshake(exec_ctx, mgr->handshakers[mgr->state->index],
+                               endpoint, mgr->state->deadline, cb, arg);
+  ++mgr->state->index;
+  // If this is the last handshaker, clean up state.
+  if (mgr->state->index == mgr->count) {
+    gpr_free(mgr->state);
+    mgr->state = NULL;
+  }
+}
+
+void grpc_handshake_manager_do_handshake(
+    grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr,
+    grpc_endpoint* endpoint, gpr_timespec deadline,
+    grpc_handshaker_done_cb cb, void* arg) {
+  if (mgr->count == 0) {
+    // No handshakers registered, so we just immediately call the done
+    // callback with the passed-in endpoint.
+    cb(exec_ctx, endpoint, arg);
+  } else {
+    GPR_ASSERT(mgr->state == NULL);
+    mgr->state = gpr_malloc(sizeof(grpc_handshaker_state));
+    memset(mgr->state, 0, sizeof(grpc_handshaker_state));
+    *mgr->state = {0, deadline, cb, arg};
+    call_next_handshaker(exec_ctx, endpoint, mgr);
+  }
 }
 
 #endif /* GRPC_CORE_LIB_CHANNEL_HANDSHAKER_H */
diff --git a/src/core/lib/channel/handshaker.h b/src/core/lib/channel/handshaker.h
index 1ca85fd008..c4e3bef5a5 100644
--- a/src/core/lib/channel/handshaker.h
+++ b/src/core/lib/channel/handshaker.h
@@ -40,17 +40,26 @@
 #include "src/core/lib/iomgr/endpoint.h"
 #include "src/core/lib/iomgr/exec_ctx.h"
 
+//
+// grpc_handshaker -- API for initial handshaking for a new connection
+//
+
 // FIXME: document
 
 typedef struct grpc_handshaker grpc_handshaker;
 
+typedef void (*grpc_handshaker_done_cb)(
+    grpc_exec_ctx* exec_ctx, grpc_endpoint* endpoint, void* arg);
+
 struct grpc_handshaker_vtable {
   void (*destroy)(grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker);
 
+  void (*shutdown)(grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker);
+
   void (*do_handshake)(
       grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker,
-      grpc_endpoint* endpoint_in, gpr_timespec deadline,
-      grpc_closure *on_done);
+      grpc_endpoint* endpoint, gpr_timespec deadline,
+      grpc_handshaker_done_cb cb, void* arg);
 };
 
 struct grpc_handshaker {
@@ -64,8 +73,34 @@ void grpc_handshaker_init(const struct grpc_handshaker_vtable* vtable,
 // Convenient wrappers for invoking methods via the vtable.
 void grpc_handshaker_destroy(grpc_exec_ctx* exec_ctx,
                              grpc_handshaker* handshaker);
+void grpc_handshaker_shutdown(grpc_exec_ctx* exec_ctx,
+                              grpc_handshaker* handshaker);
 void grpc_handshaker_do_handshake(
     grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker,
-    grpc_endpoint* in, gpr_timespec deadline, grpc_closure *on_done);
+    grpc_endpoint* endpoint, gpr_timespec deadline,
+    grpc_handshaker_done_cb cb, void* arg);
+
+//
+// grpc_handshake_manager -- manages a set of handshakers
+//
+
+typedef struct grpc_handshake_manager grpc_handshake_manager;
+
+grpc_handshake_manager* grpc_handshake_manager_create();
+
+// Handshakers will be invoked in the order added.
+void grpc_handshake_manager_add(grpc_handshaker* handshaker,
+                                grpc_handshake_manager* mgr);
+
+void grpc_handshake_manager_destroy(
+    grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr);
+
+void grpc_handshake_manager_shutdown(
+    grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr);
+
+void grpc_handshake_manager_do_handshake(
+    grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr,
+    grpc_endpoint* endpoint, gpr_timespec deadline,
+    grpc_handshaker_done_cb cb, void* arg);
 
 #endif /* GRPC_CORE_LIB_CHANNEL_HANDSHAKER_H */