4 files changed, 220 insertions, 40 deletions

diff --git a/src/core/lib/transport/byte_stream.c b/src/core/lib/transport/byte_stream.c
index 3355814017..fb03a10315 100644
--- a/src/core/lib/transport/byte_stream.c
+++ b/src/core/lib/transport/byte_stream.c
@@ -19,29 +19,37 @@
 #include "src/core/lib/transport/byte_stream.h"
 
 #include <stdlib.h>
+#include <string.h>
 
 #include <grpc/support/log.h>
 
 #include "src/core/lib/slice/slice_internal.h"
 
-int grpc_byte_stream_next(grpc_exec_ctx *exec_ctx,
-                          grpc_byte_stream *byte_stream, size_t max_size_hint,
-                          grpc_closure *on_complete) {
-  return byte_stream->next(exec_ctx, byte_stream, max_size_hint, on_complete);
+bool grpc_byte_stream_next(grpc_exec_ctx *exec_ctx,
+                           grpc_byte_stream *byte_stream, size_t max_size_hint,
+                           grpc_closure *on_complete) {
+  return byte_stream->vtable->next(exec_ctx, byte_stream, max_size_hint,
+                                   on_complete);
 }
 
 grpc_error *grpc_byte_stream_pull(grpc_exec_ctx *exec_ctx,
                                   grpc_byte_stream *byte_stream,
                                   grpc_slice *slice) {
-  return byte_stream->pull(exec_ctx, byte_stream, slice);
+  return byte_stream->vtable->pull(exec_ctx, byte_stream, slice);
+}
+
+void grpc_byte_stream_shutdown(grpc_exec_ctx *exec_ctx,
+                               grpc_byte_stream *byte_stream,
+                               grpc_error *error) {
+  byte_stream->vtable->shutdown(exec_ctx, byte_stream, error);
 }
 
 void grpc_byte_stream_destroy(grpc_exec_ctx *exec_ctx,
                               grpc_byte_stream *byte_stream) {
-  byte_stream->destroy(exec_ctx, byte_stream);
+  byte_stream->vtable->destroy(exec_ctx, byte_stream);
 }
 
-/* slice_buffer_stream */
+// grpc_slice_buffer_stream
 
 static bool slice_buffer_stream_next(grpc_exec_ctx *exec_ctx,
                                      grpc_byte_stream *byte_stream,
@@ -56,6 +64,9 @@ static grpc_error *slice_buffer_stream_pull(grpc_exec_ctx *exec_ctx,
                                             grpc_byte_stream *byte_stream,
                                             grpc_slice *slice) {
   grpc_slice_buffer_stream *stream = (grpc_slice_buffer_stream *)byte_stream;
+  if (stream->shutdown_error != GRPC_ERROR_NONE) {
+    return GRPC_ERROR_REF(stream->shutdown_error);
+  }
   GPR_ASSERT(stream->cursor < stream->backing_buffer->count);
   *slice =
       grpc_slice_ref_internal(stream->backing_buffer->slices[stream->cursor]);
@@ -63,8 +74,23 @@ static grpc_error *slice_buffer_stream_pull(grpc_exec_ctx *exec_ctx,
   return GRPC_ERROR_NONE;
 }
 
+static void slice_buffer_stream_shutdown(grpc_exec_ctx *exec_ctx,
+                                         grpc_byte_stream *byte_stream,
+                                         grpc_error *error) {
+  grpc_slice_buffer_stream *stream = (grpc_slice_buffer_stream *)byte_stream;
+  GRPC_ERROR_UNREF(stream->shutdown_error);
+  stream->shutdown_error = error;
+}
+
 static void slice_buffer_stream_destroy(grpc_exec_ctx *exec_ctx,
-                                        grpc_byte_stream *byte_stream) {}
+                                        grpc_byte_stream *byte_stream) {
+  grpc_slice_buffer_stream *stream = (grpc_slice_buffer_stream *)byte_stream;
+  GRPC_ERROR_UNREF(stream->shutdown_error);
+}
+
+static const grpc_byte_stream_vtable slice_buffer_stream_vtable = {
+    slice_buffer_stream_next, slice_buffer_stream_pull,
+    slice_buffer_stream_shutdown, slice_buffer_stream_destroy};
 
 void grpc_slice_buffer_stream_init(grpc_slice_buffer_stream *stream,
                                    grpc_slice_buffer *slice_buffer,
@@ -72,9 +98,89 @@ void grpc_slice_buffer_stream_init(grpc_slice_buffer_stream *stream,
   GPR_ASSERT(slice_buffer->length <= UINT32_MAX);
   stream->base.length = (uint32_t)slice_buffer->length;
   stream->base.flags = flags;
-  stream->base.next = slice_buffer_stream_next;
-  stream->base.pull = slice_buffer_stream_pull;
-  stream->base.destroy = slice_buffer_stream_destroy;
+  stream->base.vtable = &slice_buffer_stream_vtable;
   stream->backing_buffer = slice_buffer;
   stream->cursor = 0;
+  stream->shutdown_error = GRPC_ERROR_NONE;
+}
+
+// grpc_caching_byte_stream
+
+void grpc_byte_stream_cache_init(grpc_byte_stream_cache *cache,
+                                 grpc_byte_stream *underlying_stream) {
+  cache->underlying_stream = underlying_stream;
+  grpc_slice_buffer_init(&cache->cache_buffer);
+}
+
+void grpc_byte_stream_cache_destroy(grpc_exec_ctx *exec_ctx,
+                                    grpc_byte_stream_cache *cache) {
+  grpc_byte_stream_destroy(exec_ctx, cache->underlying_stream);
+  grpc_slice_buffer_destroy_internal(exec_ctx, &cache->cache_buffer);
+}
+
+static bool caching_byte_stream_next(grpc_exec_ctx *exec_ctx,
+                                     grpc_byte_stream *byte_stream,
+                                     size_t max_size_hint,
+                                     grpc_closure *on_complete) {
+  grpc_caching_byte_stream *stream = (grpc_caching_byte_stream *)byte_stream;
+  if (stream->shutdown_error != GRPC_ERROR_NONE) return true;
+  if (stream->cursor < stream->cache->cache_buffer.count) return true;
+  return grpc_byte_stream_next(exec_ctx, stream->cache->underlying_stream,
+                               max_size_hint, on_complete);
+}
+
+static grpc_error *caching_byte_stream_pull(grpc_exec_ctx *exec_ctx,
+                                            grpc_byte_stream *byte_stream,
+                                            grpc_slice *slice) {
+  grpc_caching_byte_stream *stream = (grpc_caching_byte_stream *)byte_stream;
+  if (stream->shutdown_error != GRPC_ERROR_NONE) {
+    return GRPC_ERROR_REF(stream->shutdown_error);
+  }
+  if (stream->cursor < stream->cache->cache_buffer.count) {
+    *slice = grpc_slice_ref_internal(
+        stream->cache->cache_buffer.slices[stream->cursor]);
+    ++stream->cursor;
+    return GRPC_ERROR_NONE;
+  }
+  grpc_error *error =
+      grpc_byte_stream_pull(exec_ctx, stream->cache->underlying_stream, slice);
+  if (error == GRPC_ERROR_NONE) {
+    ++stream->cursor;
+    grpc_slice_buffer_add(&stream->cache->cache_buffer,
+                          grpc_slice_ref_internal(*slice));
+  }
+  return error;
+}
+
+static void caching_byte_stream_shutdown(grpc_exec_ctx *exec_ctx,
+                                         grpc_byte_stream *byte_stream,
+                                         grpc_error *error) {
+  grpc_caching_byte_stream *stream = (grpc_caching_byte_stream *)byte_stream;
+  GRPC_ERROR_UNREF(stream->shutdown_error);
+  stream->shutdown_error = GRPC_ERROR_REF(error);
+  grpc_byte_stream_shutdown(exec_ctx, stream->cache->underlying_stream, error);
+}
+
+static void caching_byte_stream_destroy(grpc_exec_ctx *exec_ctx,
+                                        grpc_byte_stream *byte_stream) {
+  grpc_caching_byte_stream *stream = (grpc_caching_byte_stream *)byte_stream;
+  GRPC_ERROR_UNREF(stream->shutdown_error);
+}
+
+static const grpc_byte_stream_vtable caching_byte_stream_vtable = {
+    caching_byte_stream_next, caching_byte_stream_pull,
+    caching_byte_stream_shutdown, caching_byte_stream_destroy};
+
+void grpc_caching_byte_stream_init(grpc_caching_byte_stream *stream,
+                                   grpc_byte_stream_cache *cache) {
+  memset(stream, 0, sizeof(*stream));
+  stream->base.length = cache->underlying_stream->length;
+  stream->base.flags = cache->underlying_stream->flags;
+  stream->base.vtable = &caching_byte_stream_vtable;
+  stream->cache = cache;
+  stream->shutdown_error = GRPC_ERROR_NONE;
+}
+
+void grpc_caching_byte_stream_reset(grpc_caching_byte_stream *stream) {
+  stream->cursor = 0;
 }
diff --git a/src/core/lib/transport/byte_stream.h b/src/core/lib/transport/byte_stream.h
index f172296e4b..1e1e8310b8 100644
--- a/src/core/lib/transport/byte_stream.h
+++ b/src/core/lib/transport/byte_stream.h
@@ -28,52 +28,109 @@
 /** Mask of all valid internal flags. */
 #define GRPC_WRITE_INTERNAL_USED_MASK (GRPC_WRITE_INTERNAL_COMPRESS)
 
-struct grpc_byte_stream;
 typedef struct grpc_byte_stream grpc_byte_stream;
 
-struct grpc_byte_stream {
-  uint32_t length;
-  uint32_t flags;
+typedef struct {
   bool (*next)(grpc_exec_ctx *exec_ctx, grpc_byte_stream *byte_stream,
                size_t max_size_hint, grpc_closure *on_complete);
   grpc_error *(*pull)(grpc_exec_ctx *exec_ctx, grpc_byte_stream *byte_stream,
                       grpc_slice *slice);
+  void (*shutdown)(grpc_exec_ctx *exec_ctx, grpc_byte_stream *byte_stream,
+                   grpc_error *error);
   void (*destroy)(grpc_exec_ctx *exec_ctx, grpc_byte_stream *byte_stream);
+} grpc_byte_stream_vtable;
+
+struct grpc_byte_stream {
+  uint32_t length;
+  uint32_t flags;
+  const grpc_byte_stream_vtable *vtable;
 };
 
-/* returns 1 if the bytes are available immediately (in which case
- * on_complete will not be called), 0 if the bytes will be available
- * asynchronously.
- *
- * max_size_hint can be set as a hint as to the maximum number
- * of bytes that would be acceptable to read.
- */
-int grpc_byte_stream_next(grpc_exec_ctx *exec_ctx,
-                          grpc_byte_stream *byte_stream, size_t max_size_hint,
-                          grpc_closure *on_complete);
+// Returns true if the bytes are available immediately (in which case
+// on_complete will not be called), false if the bytes will be available
+// asynchronously.
+//
+// max_size_hint can be set as a hint as to the maximum number
+// of bytes that would be acceptable to read.
+bool grpc_byte_stream_next(grpc_exec_ctx *exec_ctx,
+                           grpc_byte_stream *byte_stream, size_t max_size_hint,
+                           grpc_closure *on_complete);
 
-/* returns the next slice in the byte stream when it is ready (indicated by
- * either grpc_byte_stream_next returning 1 or on_complete passed to
- * grpc_byte_stream_next is called).
- *
- * once a slice is returned into *slice, it is owned by the caller.
- */
+// Returns the next slice in the byte stream when it is ready (indicated by
+// either grpc_byte_stream_next returning true or on_complete passed to
+// grpc_byte_stream_next is called).
+//
+// Once a slice is returned into *slice, it is owned by the caller.
 grpc_error *grpc_byte_stream_pull(grpc_exec_ctx *exec_ctx,
                                   grpc_byte_stream *byte_stream,
                                   grpc_slice *slice);
 
+// Shuts down the byte stream.
+//
+// If there is a pending call to on_complete from grpc_byte_stream_next(),
+// it will be invoked with the error passed to grpc_byte_stream_shutdown().
+//
+// The next call to grpc_byte_stream_pull() (if any) will return the error
+// passed to grpc_byte_stream_shutdown().
+void grpc_byte_stream_shutdown(grpc_exec_ctx *exec_ctx,
+                               grpc_byte_stream *byte_stream,
+                               grpc_error *error);
+
 void grpc_byte_stream_destroy(grpc_exec_ctx *exec_ctx,
                               grpc_byte_stream *byte_stream);
 
-/* grpc_byte_stream that wraps a slice buffer */
+// grpc_slice_buffer_stream
+//
+// A grpc_byte_stream that wraps a slice buffer.
+
 typedef struct grpc_slice_buffer_stream {
   grpc_byte_stream base;
   grpc_slice_buffer *backing_buffer;
   size_t cursor;
+  grpc_error *shutdown_error;
 } grpc_slice_buffer_stream;
 
 void grpc_slice_buffer_stream_init(grpc_slice_buffer_stream *stream,
                                    grpc_slice_buffer *slice_buffer,
                                    uint32_t flags);
 
+// grpc_caching_byte_stream
+//
+// A grpc_byte_stream that that wraps an underlying byte stream but caches
+// the resulting slices in a slice buffer.  If an initial attempt fails
+// without fully draining the underlying stream, a new caching stream
+// can be created from the same underlying cache, in which case it will
+// return whatever is in the backing buffer before continuing to read the
+// underlying stream.
+//
+// NOTE: No synchronization is done, so it is not safe to have multiple
+// grpc_caching_byte_streams simultaneously drawing from the same underlying
+// grpc_byte_stream_cache at the same time.
+
+typedef struct {
+  grpc_byte_stream *underlying_stream;
+  grpc_slice_buffer cache_buffer;
+} grpc_byte_stream_cache;
+
+// Takes ownership of underlying_stream.
+void grpc_byte_stream_cache_init(grpc_byte_stream_cache *cache,
+                                 grpc_byte_stream *underlying_stream);
+
+// Must not be called while still in use by a grpc_caching_byte_stream.
+void grpc_byte_stream_cache_destroy(grpc_exec_ctx *exec_ctx,
+                                    grpc_byte_stream_cache *cache);
+
+typedef struct {
+  grpc_byte_stream base;
+  grpc_byte_stream_cache *cache;
+  size_t cursor;
+  grpc_error *shutdown_error;
+} grpc_caching_byte_stream;
+
+void grpc_caching_byte_stream_init(grpc_caching_byte_stream *stream,
+                                   grpc_byte_stream_cache *cache);
+
+// Resets the byte stream to the start of the underlying stream.
+void grpc_caching_byte_stream_reset(grpc_caching_byte_stream *stream);
+
 #endif /* GRPC_CORE_LIB_TRANSPORT_BYTE_STREAM_H */
diff --git a/src/core/lib/transport/transport.c b/src/core/lib/transport/transport.c
index 7281602d66..6c61f4b8d9 100644
--- a/src/core/lib/transport/transport.c
+++ b/src/core/lib/transport/transport.c
@@ -207,27 +207,35 @@ grpc_endpoint *grpc_transport_get_endpoint(grpc_exec_ctx *exec_ctx,
   return transport->vtable->get_endpoint(exec_ctx, transport);
 }
 
+// This comment should be sung to the tune of
+// "Supercalifragilisticexpialidocious":
+//
 // grpc_transport_stream_op_batch_finish_with_failure
 // is a function that must always unref cancel_error
 // though it lives in lib, it handles transport stream ops sure
 // it's grpc_transport_stream_op_batch_finish_with_failure
 
 void grpc_transport_stream_op_batch_finish_with_failure(
-    grpc_exec_ctx *exec_ctx, grpc_transport_stream_op_batch *op,
+    grpc_exec_ctx *exec_ctx, grpc_transport_stream_op_batch *batch,
     grpc_error *error) {
-  if (op->recv_message) {
-    GRPC_CLOSURE_SCHED(exec_ctx, op->payload->recv_message.recv_message_ready,
+  if (batch->send_message) {
+    grpc_byte_stream_destroy(exec_ctx,
+                             batch->payload->send_message.send_message);
+  }
+  if (batch->recv_message) {
+    GRPC_CLOSURE_SCHED(exec_ctx,
+                       batch->payload->recv_message.recv_message_ready,
                        GRPC_ERROR_REF(error));
   }
-  if (op->recv_initial_metadata) {
+  if (batch->recv_initial_metadata) {
     GRPC_CLOSURE_SCHED(
         exec_ctx,
-        op->payload->recv_initial_metadata.recv_initial_metadata_ready,
+        batch->payload->recv_initial_metadata.recv_initial_metadata_ready,
         GRPC_ERROR_REF(error));
   }
-  GRPC_CLOSURE_SCHED(exec_ctx, op->on_complete, error);
-  if (op->cancel_stream) {
-    GRPC_ERROR_UNREF(op->payload->cancel_stream.cancel_error);
+  GRPC_CLOSURE_SCHED(exec_ctx, batch->on_complete, error);
+  if (batch->cancel_stream) {
+    GRPC_ERROR_UNREF(batch->payload->cancel_stream.cancel_error);
   }
 }
 
diff --git a/src/core/lib/transport/transport.h b/src/core/lib/transport/transport.h
index 84e53e683a..099138ea14 100644
--- a/src/core/lib/transport/transport.h
+++ b/src/core/lib/transport/transport.h
@@ -159,6 +159,11 @@ struct grpc_transport_stream_op_batch_payload {
   } send_trailing_metadata;
 
   struct {
+    // The transport (or a filter that decides to return a failure before
+    // the op gets down to the transport) is responsible for calling
+    // grpc_byte_stream_destroy() on this.
+    // The batch's on_complete will not be called until after the byte
+    // stream is destroyed.
     grpc_byte_stream *send_message;
   } send_message;
 
@@ -174,6 +179,10 @@ struct grpc_transport_stream_op_batch_payload {
   } recv_initial_metadata;
 
   struct {
+    // Will be set by the transport to point to the byte stream
+    // containing a received message.
+    // The caller is responsible for calling grpc_byte_stream_destroy()
+    // on this byte stream.
     grpc_byte_stream **recv_message;
     /** Should be enqueued when one message is ready to be processed. */
     grpc_closure *recv_message_ready;