Revert "Revert "Refactor Endpoint API""

1 files changed, 105 insertions, 83 deletions

diff --git a/src/core/security/secure_endpoint.c b/src/core/security/secure_endpoint.c
index 81b3e33cb2..b696e384fc 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*/
+#ifdef 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);
+}
+#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;