Refactor Endpoint API

- Allow reads to complete immediately
- Unify read/write signatures
- Simplify memory management to allow future optimization work

2 files changed, 174 insertions, 133 deletions

diff --git a/src/core/security/secure_endpoint.c b/src/core/security/secure_endpoint.c
index 81b3e33cb2..4206c39318 100644
--- a/src/core/security/secure_endpoint.c
+++ b/src/core/security/secure_endpoint.c
@@ -49,15 +49,15 @@ typedef struct {
   struct tsi_frame_protector *protector;
   gpr_mu protector_mu;
   /* saved upper level callbacks and user_data. */
-  grpc_endpoint_read_cb read_cb;
-  void *read_user_data;
-  grpc_endpoint_write_cb write_cb;
-  void *write_user_data;
+  grpc_iomgr_closure *read_cb;
+  grpc_iomgr_closure *write_cb;
+  grpc_iomgr_closure on_read;
+  gpr_slice_buffer *read_buffer;
+  gpr_slice_buffer source_buffer;
   /* saved handshaker leftover data to unprotect. */
   gpr_slice_buffer leftover_bytes;
   /* buffers for read and write */
   gpr_slice read_staging_buffer;
-  gpr_slice_buffer input_buffer;
 
   gpr_slice write_staging_buffer;
   gpr_slice_buffer output_buffer;
@@ -67,62 +67,91 @@ typedef struct {
 
 int grpc_trace_secure_endpoint = 0;
 
-static void secure_endpoint_ref(secure_endpoint *ep) { gpr_ref(&ep->ref); }
-
 static void destroy(secure_endpoint *secure_ep) {
   secure_endpoint *ep = secure_ep;
   grpc_endpoint_destroy(ep->wrapped_ep);
   tsi_frame_protector_destroy(ep->protector);
   gpr_slice_buffer_destroy(&ep->leftover_bytes);
   gpr_slice_unref(ep->read_staging_buffer);
-  gpr_slice_buffer_destroy(&ep->input_buffer);
   gpr_slice_unref(ep->write_staging_buffer);
   gpr_slice_buffer_destroy(&ep->output_buffer);
+  gpr_slice_buffer_destroy(&ep->source_buffer);
   gpr_mu_destroy(&ep->protector_mu);
   gpr_free(ep);
 }
 
+#define GRPC_SECURE_ENDPOINT_REFCOUNT_DEBUG
+#define SECURE_ENDPOINT_UNREF(ep, reason) \
+  secure_endpoint_unref((ep), (reason), __FILE__, __LINE__)
+#define SECURE_ENDPOINT_REF(ep, reason) \
+  secure_endpoint_ref((ep), (reason), __FILE__, __LINE__)
+static void secure_endpoint_unref(secure_endpoint *ep, const char *reason,
+                                  const char *file, int line) {
+  gpr_log(file, line, GPR_LOG_SEVERITY_DEBUG, "SECENDP unref %p : %s %d -> %d",
+          ep, reason, ep->ref.count, ep->ref.count - 1);
+  if (gpr_unref(&ep->ref)) {
+    destroy(ep);
+  }
+}
+
+static void secure_endpoint_ref(secure_endpoint *ep, const char *reason,
+                                const char *file, int line) {
+  gpr_log(file, line, GPR_LOG_SEVERITY_DEBUG, "SECENDP   ref %p : %s %d -> %d",
+          ep, reason, ep->ref.count, ep->ref.count + 1);
+  gpr_ref(&ep->ref);
+}
+#ifdef GRPC_SECURE_ENDPOINT_REFCOUNT_DEBUG
+#else
+#define SECURE_ENDPOINT_UNREF(ep, reason) secure_endpoint_unref((ep))
+#define SECURE_ENDPOINT_REF(ep, reason) secure_endpoint_ref((ep))
 static void secure_endpoint_unref(secure_endpoint *ep) {
   if (gpr_unref(&ep->ref)) {
     destroy(ep);
   }
 }
 
+static void secure_endpoint_ref(secure_endpoint *ep) { gpr_ref(&ep->ref); }
+#endif
+
 static void flush_read_staging_buffer(secure_endpoint *ep, gpr_uint8 **cur,
                                       gpr_uint8 **end) {
-  gpr_slice_buffer_add(&ep->input_buffer, ep->read_staging_buffer);
+  gpr_slice_buffer_add(ep->read_buffer, ep->read_staging_buffer);
   ep->read_staging_buffer = gpr_slice_malloc(STAGING_BUFFER_SIZE);
   *cur = GPR_SLICE_START_PTR(ep->read_staging_buffer);
   *end = GPR_SLICE_END_PTR(ep->read_staging_buffer);
 }
 
-static void call_read_cb(secure_endpoint *ep, gpr_slice *slices, size_t nslices,
-                         grpc_endpoint_cb_status error) {
+static void call_read_cb(secure_endpoint *ep, int success) {
   if (grpc_trace_secure_endpoint) {
     size_t i;
-    for (i = 0; i < nslices; i++) {
-      char *data = gpr_dump_slice(slices[i], GPR_DUMP_HEX | GPR_DUMP_ASCII);
+    for (i = 0; i < ep->read_buffer->count; i++) {
+      char *data = gpr_dump_slice(ep->read_buffer->slices[i],
+                                  GPR_DUMP_HEX | GPR_DUMP_ASCII);
       gpr_log(GPR_DEBUG, "READ %p: %s", ep, data);
       gpr_free(data);
     }
   }
-  ep->read_cb(ep->read_user_data, slices, nslices, error);
-  secure_endpoint_unref(ep);
+  ep->read_buffer = NULL;
+  ep->read_cb->cb(ep->read_cb->cb_arg, success);
+  SECURE_ENDPOINT_UNREF(ep, "read");
 }
 
-static void on_read(void *user_data, gpr_slice *slices, size_t nslices,
-                    grpc_endpoint_cb_status error) {
+static int on_read(void *user_data, int success) {
   unsigned i;
   gpr_uint8 keep_looping = 0;
-  size_t input_buffer_count = 0;
   tsi_result result = TSI_OK;
   secure_endpoint *ep = (secure_endpoint *)user_data;
   gpr_uint8 *cur = GPR_SLICE_START_PTR(ep->read_staging_buffer);
   gpr_uint8 *end = GPR_SLICE_END_PTR(ep->read_staging_buffer);
 
+  if (!success) {
+    gpr_slice_buffer_reset_and_unref(ep->read_buffer);
+    return 0;
+  }
+
   /* TODO(yangg) check error, maybe bail out early */
-  for (i = 0; i < nslices; i++) {
-    gpr_slice encrypted = slices[i];
+  for (i = 0; i < ep->source_buffer.count; i++) {
+    gpr_slice encrypted = ep->source_buffer.slices[i];
     gpr_uint8 *message_bytes = GPR_SLICE_START_PTR(encrypted);
     size_t message_size = GPR_SLICE_LENGTH(encrypted);
 
@@ -161,7 +190,7 @@ static void on_read(void *user_data, gpr_slice *slices, size_t nslices,
 
   if (cur != GPR_SLICE_START_PTR(ep->read_staging_buffer)) {
     gpr_slice_buffer_add(
-        &ep->input_buffer,
+        ep->read_buffer,
         gpr_slice_split_head(
             &ep->read_staging_buffer,
             (size_t)(cur - GPR_SLICE_START_PTR(ep->read_staging_buffer))));
@@ -169,38 +198,53 @@ static void on_read(void *user_data, gpr_slice *slices, size_t nslices,
 
   /* TODO(yangg) experiment with moving this block after read_cb to see if it
      helps latency */
-  for (i = 0; i < nslices; i++) {
-    gpr_slice_unref(slices[i]);
-  }
+  gpr_slice_buffer_reset_and_unref(&ep->source_buffer);
 
   if (result != TSI_OK) {
-    gpr_slice_buffer_reset_and_unref(&ep->input_buffer);
-    call_read_cb(ep, NULL, 0, GRPC_ENDPOINT_CB_ERROR);
-    return;
+    gpr_slice_buffer_reset_and_unref(ep->read_buffer);
+    return 0;
   }
-  /* The upper level will unref the slices. */
-  input_buffer_count = ep->input_buffer.count;
-  ep->input_buffer.count = 0;
-  call_read_cb(ep, ep->input_buffer.slices, input_buffer_count, error);
+
+  return 1;
+}
+
+static void on_read_cb(void *user_data, int success) {
+  call_read_cb(user_data, on_read(user_data, success));
 }
 
-static void endpoint_notify_on_read(grpc_endpoint *secure_ep,
-                                    grpc_endpoint_read_cb cb, void *user_data) {
+static grpc_endpoint_op_status endpoint_read(grpc_endpoint *secure_ep,
+                                             gpr_slice_buffer *slices,
+                                             grpc_iomgr_closure *cb) {
   secure_endpoint *ep = (secure_endpoint *)secure_ep;
+  int immediate_read_success = -1;
   ep->read_cb = cb;
-  ep->read_user_data = user_data;
-
-  secure_endpoint_ref(ep);
+  ep->read_buffer = slices;
+  gpr_slice_buffer_reset_and_unref(ep->read_buffer);
 
   if (ep->leftover_bytes.count) {
-    size_t leftover_nslices = ep->leftover_bytes.count;
-    ep->leftover_bytes.count = 0;
-    on_read(ep, ep->leftover_bytes.slices, leftover_nslices,
-            GRPC_ENDPOINT_CB_OK);
-    return;
+    gpr_slice_buffer_swap(&ep->leftover_bytes, &ep->source_buffer);
+    GPR_ASSERT(ep->leftover_bytes.count == 0);
+    return on_read(ep, 1) ? GRPC_ENDPOINT_DONE : GRPC_ENDPOINT_ERROR;
   }
 
-  grpc_endpoint_notify_on_read(ep->wrapped_ep, on_read, ep);
+  SECURE_ENDPOINT_REF(ep, "read");
+
+  switch (
+      grpc_endpoint_read(ep->wrapped_ep, &ep->source_buffer, &ep->on_read)) {
+    case GRPC_ENDPOINT_DONE:
+      immediate_read_success = on_read(ep, 1);
+      break;
+    case GRPC_ENDPOINT_PENDING:
+      return GRPC_ENDPOINT_PENDING;
+    case GRPC_ENDPOINT_ERROR:
+      immediate_read_success = on_read(ep, 0);
+      break;
+  }
+
+  GPR_ASSERT(immediate_read_success != -1);
+  SECURE_ENDPOINT_UNREF(ep, "read");
+
+  return immediate_read_success ? GRPC_ENDPOINT_DONE : GRPC_ENDPOINT_ERROR;
 }
 
 static void flush_write_staging_buffer(secure_endpoint *ep, gpr_uint8 **cur,
@@ -211,36 +255,28 @@ static void flush_write_staging_buffer(secure_endpoint *ep, gpr_uint8 **cur,
   *end = GPR_SLICE_END_PTR(ep->write_staging_buffer);
 }
 
-static void on_write(void *data, grpc_endpoint_cb_status error) {
-  secure_endpoint *ep = data;
-  ep->write_cb(ep->write_user_data, error);
-  secure_endpoint_unref(ep);
-}
-
-static grpc_endpoint_write_status endpoint_write(grpc_endpoint *secure_ep,
-                                                 gpr_slice *slices,
-                                                 size_t nslices,
-                                                 grpc_endpoint_write_cb cb,
-                                                 void *user_data) {
+static grpc_endpoint_op_status endpoint_write(grpc_endpoint *secure_ep,
+                                              gpr_slice_buffer *slices,
+                                              grpc_iomgr_closure *cb) {
   unsigned i;
-  size_t output_buffer_count = 0;
   tsi_result result = TSI_OK;
   secure_endpoint *ep = (secure_endpoint *)secure_ep;
   gpr_uint8 *cur = GPR_SLICE_START_PTR(ep->write_staging_buffer);
   gpr_uint8 *end = GPR_SLICE_END_PTR(ep->write_staging_buffer);
-  grpc_endpoint_write_status status;
-  GPR_ASSERT(ep->output_buffer.count == 0);
+
+  gpr_slice_buffer_reset_and_unref(&ep->output_buffer);
 
   if (grpc_trace_secure_endpoint) {
-    for (i = 0; i < nslices; i++) {
-      char *data = gpr_dump_slice(slices[i], GPR_DUMP_HEX | GPR_DUMP_ASCII);
+    for (i = 0; i < slices->count; i++) {
+      char *data =
+          gpr_dump_slice(slices->slices[i], GPR_DUMP_HEX | GPR_DUMP_ASCII);
       gpr_log(GPR_DEBUG, "WRITE %p: %s", ep, data);
       gpr_free(data);
     }
   }
 
-  for (i = 0; i < nslices; i++) {
-    gpr_slice plain = slices[i];
+  for (i = 0; i < slices->count; i++) {
+    gpr_slice plain = slices->slices[i];
     gpr_uint8 *message_bytes = GPR_SLICE_START_PTR(plain);
     size_t message_size = GPR_SLICE_LENGTH(plain);
     while (message_size > 0) {
@@ -290,29 +326,13 @@ static grpc_endpoint_write_status endpoint_write(grpc_endpoint *secure_ep,
     }
   }
 
-  for (i = 0; i < nslices; i++) {
-    gpr_slice_unref(slices[i]);
-  }
-
   if (result != TSI_OK) {
     /* TODO(yangg) do different things according to the error type? */
     gpr_slice_buffer_reset_and_unref(&ep->output_buffer);
-    return GRPC_ENDPOINT_WRITE_ERROR;
+    return GRPC_ENDPOINT_ERROR;
   }
 
-  /* clear output_buffer and let the lower level handle its slices. */
-  output_buffer_count = ep->output_buffer.count;
-  ep->output_buffer.count = 0;
-  ep->write_cb = cb;
-  ep->write_user_data = user_data;
-  /* Need to keep the endpoint alive across a transport */
-  secure_endpoint_ref(ep);
-  status = grpc_endpoint_write(ep->wrapped_ep, ep->output_buffer.slices,
-                               output_buffer_count, on_write, ep);
-  if (status != GRPC_ENDPOINT_WRITE_PENDING) {
-    secure_endpoint_unref(ep);
-  }
-  return status;
+  return grpc_endpoint_write(ep->wrapped_ep, &ep->output_buffer, cb);
 }
 
 static void endpoint_shutdown(grpc_endpoint *secure_ep) {
@@ -320,9 +340,9 @@ static void endpoint_shutdown(grpc_endpoint *secure_ep) {
   grpc_endpoint_shutdown(ep->wrapped_ep);
 }
 
-static void endpoint_unref(grpc_endpoint *secure_ep) {
+static void endpoint_destroy(grpc_endpoint *secure_ep) {
   secure_endpoint *ep = (secure_endpoint *)secure_ep;
-  secure_endpoint_unref(ep);
+  SECURE_ENDPOINT_UNREF(ep, "destroy");
 }
 
 static void endpoint_add_to_pollset(grpc_endpoint *secure_ep,
@@ -343,9 +363,9 @@ static char *endpoint_get_peer(grpc_endpoint *secure_ep) {
 }
 
 static const grpc_endpoint_vtable vtable = {
-    endpoint_notify_on_read, endpoint_write,
+    endpoint_read,           endpoint_write,
     endpoint_add_to_pollset, endpoint_add_to_pollset_set,
-    endpoint_shutdown,       endpoint_unref,
+    endpoint_shutdown,       endpoint_destroy,
     endpoint_get_peer};
 
 grpc_endpoint *grpc_secure_endpoint_create(
@@ -363,8 +383,10 @@ grpc_endpoint *grpc_secure_endpoint_create(
   }
   ep->write_staging_buffer = gpr_slice_malloc(STAGING_BUFFER_SIZE);
   ep->read_staging_buffer = gpr_slice_malloc(STAGING_BUFFER_SIZE);
-  gpr_slice_buffer_init(&ep->input_buffer);
   gpr_slice_buffer_init(&ep->output_buffer);
+  gpr_slice_buffer_init(&ep->source_buffer);
+  ep->read_buffer = NULL;
+  grpc_iomgr_closure_init(&ep->on_read, on_read_cb, ep);
   gpr_mu_init(&ep->protector_mu);
   gpr_ref_init(&ep->ref, 1);
   return &ep->base;
diff --git a/src/core/security/secure_transport_setup.c b/src/core/security/secure_transport_setup.c
index 0c3572b53c..bf0079577e 100644
--- a/src/core/security/secure_transport_setup.c
+++ b/src/core/security/secure_transport_setup.c
@@ -50,16 +50,17 @@ typedef struct {
   grpc_endpoint *wrapped_endpoint;
   grpc_endpoint *secure_endpoint;
   gpr_slice_buffer left_overs;
+  gpr_slice_buffer incoming;
+  gpr_slice_buffer outgoing;
   grpc_secure_transport_setup_done_cb cb;
   void *user_data;
+  grpc_iomgr_closure on_handshake_data_sent_to_peer;
+  grpc_iomgr_closure on_handshake_data_received_from_peer;
 } grpc_secure_transport_setup;
 
-static void on_handshake_data_received_from_peer(void *setup, gpr_slice *slices,
-                                                 size_t nslices,
-                                                 grpc_endpoint_cb_status error);
+static void on_handshake_data_received_from_peer(void *setup, int success);
 
-static void on_handshake_data_sent_to_peer(void *setup,
-                                           grpc_endpoint_cb_status error);
+static void on_handshake_data_sent_to_peer(void *setup, int success);
 
 static void secure_transport_setup_done(grpc_secure_transport_setup *s,
                                         int is_success) {
@@ -78,6 +79,8 @@ static void secure_transport_setup_done(grpc_secure_transport_setup *s,
   if (s->handshaker != NULL) tsi_handshaker_destroy(s->handshaker);
   if (s->handshake_buffer != NULL) gpr_free(s->handshake_buffer);
   gpr_slice_buffer_destroy(&s->left_overs);
+  gpr_slice_buffer_destroy(&s->outgoing);
+  gpr_slice_buffer_destroy(&s->incoming);
   GRPC_SECURITY_CONNECTOR_UNREF(s->connector, "secure_transport_setup");
   gpr_free(s);
 }
@@ -102,6 +105,8 @@ static void on_peer_checked(void *user_data, grpc_security_status status) {
   s->secure_endpoint =
       grpc_secure_endpoint_create(protector, s->wrapped_endpoint,
                                   s->left_overs.slices, s->left_overs.count);
+  s->left_overs.count = 0;
+  s->left_overs.length = 0;
   secure_transport_setup_done(s, 1);
   return;
 }
@@ -132,7 +137,6 @@ static void send_handshake_bytes_to_peer(grpc_secure_transport_setup *s) {
   size_t offset = 0;
   tsi_result result = TSI_OK;
   gpr_slice to_send;
-  grpc_endpoint_write_status write_status;
 
   do {
     size_t to_send_size = s->handshake_buffer_size - offset;
@@ -155,28 +159,25 @@ static void send_handshake_bytes_to_peer(grpc_secure_transport_setup *s) {
 
   to_send =
       gpr_slice_from_copied_buffer((const char *)s->handshake_buffer, offset);
+  gpr_slice_buffer_reset_and_unref(&s->outgoing);
+  gpr_slice_buffer_add(&s->outgoing, to_send);
   /* TODO(klempner,jboeuf): This should probably use the client setup
          deadline */
-  write_status = grpc_endpoint_write(s->wrapped_endpoint, &to_send, 1,
-                                     on_handshake_data_sent_to_peer, s);
-  if (write_status == GRPC_ENDPOINT_WRITE_ERROR) {
-    gpr_log(GPR_ERROR, "Could not send handshake data to peer.");
-    secure_transport_setup_done(s, 0);
-  } else if (write_status == GRPC_ENDPOINT_WRITE_DONE) {
-    on_handshake_data_sent_to_peer(s, GRPC_ENDPOINT_CB_OK);
-  }
-}
-
-static void cleanup_slices(gpr_slice *slices, size_t num_slices) {
-  size_t i;
-  for (i = 0; i < num_slices; i++) {
-    gpr_slice_unref(slices[i]);
+  switch (grpc_endpoint_write(s->wrapped_endpoint, &s->outgoing,
+                              &s->on_handshake_data_sent_to_peer)) {
+    case GRPC_ENDPOINT_ERROR:
+      gpr_log(GPR_ERROR, "Could not send handshake data to peer.");
+      secure_transport_setup_done(s, 0);
+      break;
+    case GRPC_ENDPOINT_DONE:
+      on_handshake_data_sent_to_peer(s, 1);
+      break;
+    case GRPC_ENDPOINT_PENDING:
+      break;
   }
 }
 
-static void on_handshake_data_received_from_peer(
-    void *setup, gpr_slice *slices, size_t nslices,
-    grpc_endpoint_cb_status error) {
+static void on_handshake_data_received_from_peer(void *setup, int success) {
   grpc_secure_transport_setup *s = setup;
   size_t consumed_slice_size = 0;
   tsi_result result = TSI_OK;
@@ -184,32 +185,37 @@ static void on_handshake_data_received_from_peer(
   size_t num_left_overs;
   int has_left_overs_in_current_slice = 0;
 
-  if (error != GRPC_ENDPOINT_CB_OK) {
+  if (!success) {
     gpr_log(GPR_ERROR, "Read failed.");
-    cleanup_slices(slices, nslices);
     secure_transport_setup_done(s, 0);
     return;
   }
 
-  for (i = 0; i < nslices; i++) {
-    consumed_slice_size = GPR_SLICE_LENGTH(slices[i]);
+  for (i = 0; i < s->incoming.count; i++) {
+    consumed_slice_size = GPR_SLICE_LENGTH(s->incoming.slices[i]);
     result = tsi_handshaker_process_bytes_from_peer(
-        s->handshaker, GPR_SLICE_START_PTR(slices[i]), &consumed_slice_size);
+        s->handshaker, GPR_SLICE_START_PTR(s->incoming.slices[i]),
+        &consumed_slice_size);
     if (!tsi_handshaker_is_in_progress(s->handshaker)) break;
   }
 
   if (tsi_handshaker_is_in_progress(s->handshaker)) {
     /* We may need more data. */
     if (result == TSI_INCOMPLETE_DATA) {
-      /* TODO(klempner,jboeuf): This should probably use the client setup
-         deadline */
-      grpc_endpoint_notify_on_read(s->wrapped_endpoint,
-                                   on_handshake_data_received_from_peer, setup);
-      cleanup_slices(slices, nslices);
+      switch (grpc_endpoint_read(s->wrapped_endpoint, &s->incoming,
+                                 &s->on_handshake_data_received_from_peer)) {
+        case GRPC_ENDPOINT_DONE:
+          on_handshake_data_received_from_peer(s, 1);
+          break;
+        case GRPC_ENDPOINT_ERROR:
+          on_handshake_data_received_from_peer(s, 0);
+          break;
+        case GRPC_ENDPOINT_PENDING:
+          break;
+      }
       return;
     } else {
       send_handshake_bytes_to_peer(s);
-      cleanup_slices(slices, nslices);
       return;
     }
   }
@@ -217,42 +223,40 @@ static void on_handshake_data_received_from_peer(
   if (result != TSI_OK) {
     gpr_log(GPR_ERROR, "Handshake failed with error %s",
             tsi_result_to_string(result));
-    cleanup_slices(slices, nslices);
     secure_transport_setup_done(s, 0);
     return;
   }
 
   /* Handshake is done and successful this point. */
   has_left_overs_in_current_slice =
-      (consumed_slice_size < GPR_SLICE_LENGTH(slices[i]));
-  num_left_overs = (has_left_overs_in_current_slice ? 1 : 0) + nslices - i - 1;
+      (consumed_slice_size < GPR_SLICE_LENGTH(s->incoming.slices[i]));
+  num_left_overs =
+      (has_left_overs_in_current_slice ? 1 : 0) + s->incoming.count - i - 1;
   if (num_left_overs == 0) {
-    cleanup_slices(slices, nslices);
     check_peer(s);
     return;
   }
-  cleanup_slices(slices, nslices - num_left_overs);
-
   /* Put the leftovers in our buffer (ownership transfered). */
   if (has_left_overs_in_current_slice) {
-    gpr_slice_buffer_add(&s->left_overs,
-                         gpr_slice_split_tail(&slices[i], consumed_slice_size));
-    gpr_slice_unref(slices[i]); /* split_tail above increments refcount. */
+    gpr_slice_buffer_add(
+        &s->left_overs,
+        gpr_slice_split_tail(&s->incoming.slices[i], consumed_slice_size));
+    gpr_slice_unref(
+        s->incoming.slices[i]); /* split_tail above increments refcount. */
   }
   gpr_slice_buffer_addn(
-      &s->left_overs, &slices[i + 1],
+      &s->left_overs, &s->incoming.slices[i + 1],
       num_left_overs - (size_t)has_left_overs_in_current_slice);
   check_peer(s);
 }
 
 /* If setup is NULL, the setup is done. */
-static void on_handshake_data_sent_to_peer(void *setup,
-                                           grpc_endpoint_cb_status error) {
+static void on_handshake_data_sent_to_peer(void *setup, int success) {
   grpc_secure_transport_setup *s = setup;
 
   /* Make sure that write is OK. */
-  if (error != GRPC_ENDPOINT_CB_OK) {
-    gpr_log(GPR_ERROR, "Write failed with error %d.", error);
+  if (!success) {
+    gpr_log(GPR_ERROR, "Write failed.");
     if (setup != NULL) secure_transport_setup_done(s, 0);
     return;
   }
@@ -261,8 +265,17 @@ static void on_handshake_data_sent_to_peer(void *setup,
   if (tsi_handshaker_is_in_progress(s->handshaker)) {
     /* TODO(klempner,jboeuf): This should probably use the client setup
        deadline */
-    grpc_endpoint_notify_on_read(s->wrapped_endpoint,
-                                 on_handshake_data_received_from_peer, setup);
+    switch (grpc_endpoint_read(s->wrapped_endpoint, &s->incoming,
+                               &s->on_handshake_data_received_from_peer)) {
+      case GRPC_ENDPOINT_ERROR:
+        on_handshake_data_received_from_peer(s, 0);
+        break;
+      case GRPC_ENDPOINT_PENDING:
+        break;
+      case GRPC_ENDPOINT_DONE:
+        on_handshake_data_received_from_peer(s, 1);
+        break;
+    }
   } else {
     check_peer(s);
   }
@@ -288,6 +301,12 @@ void grpc_setup_secure_transport(grpc_security_connector *connector,
   s->wrapped_endpoint = nonsecure_endpoint;
   s->user_data = user_data;
   s->cb = cb;
+  grpc_iomgr_closure_init(&s->on_handshake_data_sent_to_peer,
+                          on_handshake_data_sent_to_peer, s);
+  grpc_iomgr_closure_init(&s->on_handshake_data_received_from_peer,
+                          on_handshake_data_received_from_peer, s);
   gpr_slice_buffer_init(&s->left_overs);
+  gpr_slice_buffer_init(&s->outgoing);
+  gpr_slice_buffer_init(&s->incoming);
   send_handshake_bytes_to_peer(s);
 }