ckpt

blob/blob.go

@@ -237,11 +237,33 @@ func (r *Reader) WriteTo(w io.Writer) (int64, error) {
 	return nw, err
 }
 
+// downloadAndClose is similar to WriteTo, but ensures it's the only read.
+// This pattern is more optimal for some drivers.
+func (r *Reader) downloadAndClose(w io.Writer) (err error) {
+	if r.bytesRead != 0 {
+		// Shouldn't happen.
+		return gcerr.Newf(gcerr.Internal, nil, "blob: downloadAndClose isn't the first read")
+	}
+	driverDownloader, ok := r.r.(driver.Downloader)
+	if ok {
+		err = driverDownloader.Download(w)
+	} else {
+		_, err = r.WriteTo(w)
+	}
+	cerr := r.Close()
+	if err == nil && cerr != nil {
+		err = cerr
+	}
+	return err
+}
+
 // readFromWriteTo is a helper for ReadFrom and WriteTo.
 // It reads data from r and writes to w, until EOF or a read/write error.
 // It returns the number of bytes read from r and the number of bytes
 // written to w.
 func readFromWriteTo(r io.Reader, w io.Writer) (int64, int64, error) {
+	// Note: can't use io.Copy because it will try to use r.WriteTo
+	// or w.WriteTo, which is recursive in this context.
 	buf := make([]byte, 1024)
 	var totalRead, totalWritten int64
 	for {
@@ -458,6 +480,26 @@ func (w *Writer) ReadFrom(r io.Reader) (int64, error) {
 	return nr, err
 }
 
+// uploadAndClose is similar to ReadFrom, but ensures it's the only write.
+// This pattern is more optimal for some drivers.
+func (w *Writer) uploadAndClose(r io.Reader) (err error) {
+	if w.bytesWritten != 0 {
+		// Shouldn't happen.
+		return gcerr.Newf(gcerr.Internal, nil, "blob: uploadAndClose must be the first write")
+	}
+	driverUploader, ok := w.w.(driver.Uploader)
+	if ok {
+		err = driverUploader.Upload(r)
+	} else {
+		_, err = w.ReadFrom(r)
+	}
+	cerr := w.Close()
+	if err == nil && cerr != nil {
+		err = cerr
+	}
+	return err
+}
+
 // ListOptions sets options for listing blobs via Bucket.List.
 type ListOptions struct {
 	// Prefix indicates that only blobs with a key starting with this prefix
@@ -644,6 +686,8 @@ func (b *Bucket) ErrorAs(err error, i interface{}) bool {
 
 // ReadAll is a shortcut for creating a Reader via NewReader with nil
 // ReaderOptions, and reading the entire blob.
+//
+// Using Download may be more efficient.
 func (b *Bucket) ReadAll(ctx context.Context, key string) (_ []byte, err error) {
 	b.mu.RLock()
 	defer b.mu.RUnlock()
@@ -658,6 +702,20 @@ func (b *Bucket) ReadAll(ctx context.Context, key string) (_ []byte, err error)
 	return ioutil.ReadAll(r)
 }
 
+// Download writes the content of a blob into an io.Writer w.
+func (b *Bucket) Download(ctx context.Context, key string, w io.Writer, opts *ReaderOptions) error {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+	if b.closed {
+		return errClosed
+	}
+	r, err := b.NewReader(ctx, key, opts)
+	if err != nil {
+		return err
+	}
+	return r.downloadAndClose(w)
+}
+
 // List returns a ListIterator that can be used to iterate over blobs in a
 // bucket, in lexicographical order of UTF-8 encoded keys. The underlying
 // implementation fetches results in pages.
@@ -934,6 +992,8 @@ func (b *Bucket) newRangeReader(ctx context.Context, key string, offset, length
 //
 // If opts.ContentMD5 is not set, WriteAll will compute the MD5 of p and use it
 // as the ContentMD5 option for the Writer it creates.
+//
+// Using Upload may be more efficient.
 func (b *Bucket) WriteAll(ctx context.Context, key string, p []byte, opts *WriterOptions) (err error) {
 	realOpts := new(WriterOptions)
 	if opts != nil {
@@ -954,6 +1014,20 @@ func (b *Bucket) WriteAll(ctx context.Context, key string, p []byte, opts *Write
 	return w.Close()
 }
 
+// Upload reads from an io.Reader r and writes into a blob.
+//
+// opts.ContentType is required.
+func (b *Bucket) Upload(ctx context.Context, key string, r io.Reader, opts *WriterOptions) error {
+	if opts == nil || opts.ContentType == "" {
+		return gcerr.Newf(gcerr.InvalidArgument, nil, "blob: Upload requires WriterOptions.ContentType")
+	}
+	w, err := b.NewWriter(ctx, key, opts)
+	if err != nil {
+		return err
+	}
+	return w.uploadAndClose(r)
+}
+
 // NewWriter returns a Writer that writes to the blob stored at key.
 // A nil WriterOptions is treated the same as the zero value.
 //

blob/driver/driver.go

@@ -50,11 +50,24 @@ type Reader interface {
 	As(interface{}) bool
 }
 
+// Downloader has an optional extra method for readers.
+// It is similar to io.WriteTo, but without the count of bytes returned.
+type Downloader interface {
+	// Download is similar to io.WriteTo, but without the count of bytes returned.
+	Download(w io.Writer) error
+}
+
 // Writer writes an object to the blob.
 type Writer interface {
 	io.WriteCloser
 }
 
+// Uploader has an optional extra method for writers.
+type Uploader interface {
+	// Upload is similar to io.ReadFrom, but without the count of bytes returned.
+	Upload(r io.Reader) error
+}
+
 // WriterOptions controls behaviors of Writer.
 type WriterOptions struct {
 	// BufferSize changes the default size in byte of the maximum part Writer can
@@ -263,6 +276,11 @@ type Bucket interface {
 	// exist, NewRangeReader must return an error for which ErrorCode returns
 	// gcerrors.NotFound.
 	// opts is guaranteed to be non-nil.
+	//
+	// The returned Reader *may* also implement Downloader if the underlying
+	// implementation can take advantage of that. The Download call is guaranteed
+	// to be the only call to the Reader. For such readers, offset will always
+	// be 0 and length will always be -1.
 	NewRangeReader(ctx context.Context, key string, offset, length int64, opts *ReaderOptions) (Reader, error)
 
 	// NewTypedWriter returns Writer that writes to an object associated with key.
@@ -279,6 +297,10 @@ type Bucket interface {
 	//
 	// Implementations should abort an ongoing write if ctx is later canceled,
 	// and do any necessary cleanup in Close. Close should then return ctx.Err().
+	//
+	// The returned Writer *may* also implement Uploader if the underlying
+	// implementation can take advantage of that. The Upload call is guaranteed
+	// to be the only non-Close call to the Writer..
 	NewTypedWriter(ctx context.Context, key, contentType string, opts *WriterOptions) (Writer, error)
 
 	// Copy copies the object associated with srcKey to dstKey.

blob/fileblob/fileblob.go

@@ -622,6 +622,14 @@ func (r *reader) Read(p []byte) (int, error) {
 	return r.r.Read(p)
 }
 
+func (r *reader) Download(w io.Writer) error {
+	// This should always work because r.r was created from a File.
+	// It's only not a WriterTo when we wrap it with a LimitReader,
+	// which is guaranteed not to happen by the driver interface.
+	_, err := r.r.(io.WriterTo).WriteTo(w)
+	return err
+}
+
 func (r *reader) Close() error {
 	if r.c == nil {
 		return nil

blob/memblob/memblob.go

@@ -256,6 +256,14 @@ func (r *reader) Read(p []byte) (int, error) {
 	return r.r.Read(p)
 }
 
+func (r *reader) WriteTo(w io.Writer) error {
+	// This should always work because r.r was created from a bytes.Reader.
+	// It's only not a WriterTo when we wrap it with a LimitReader,
+	// which is guaranteed not to happen by the driver interface.
+	_, err := r.r.(io.WriterTo).WriteTo(w)
+	return err
+}
+
 func (r *reader) Close() error {
 	return nil
 }

blob/s3blob/s3blob.go

@@ -264,7 +264,9 @@ func (r *reader) Attributes() *driver.ReaderAttributes {
 
 // writer writes an S3 object, it implements io.WriteCloser.
 type writer struct {
-	w *io.PipeWriter // created when the first byte is written
+	pw     *io.PipeWriter // created when the first byte is written
+	pr     *io.PipeReader
+	upload bool
 
 	ctx   context.Context
 	useV2 bool
@@ -287,12 +289,11 @@ func (w *writer) Write(p []byte) (int, error) {
 	if len(p) == 0 {
 		return 0, nil
 	}
-	if w.w == nil {
+	if w.pw == nil {
 		// We'll write into pw and use pr as an io.Reader for the
 		// Upload call to S3.
-		pr, pw := io.Pipe()
-		w.w = pw
-		if err := w.open(pr); err != nil {
+		w.pr, w.pw = io.Pipe()
+		if err := w.open(w.pr, true); err != nil {
 			return 0, err
 		}
 	}
@@ -301,33 +302,46 @@ func (w *writer) Write(p []byte) (int, error) {
 		return 0, w.err
 	default:
 	}
-	return w.w.Write(p)
+	return w.pw.Write(p)
 }
 
-// pr may be nil if we're Closing and no data was written.
-func (w *writer) open(pr *io.PipeReader) error {
+// ReadFrom reads from r. Per the driver, it is guaranteed to be the only
+// write call for this writer.
+func (w *writer) Upload(r io.Reader) error {
+	w.upload = true
+	if err := w.open(r, false); err != nil {
+		return err
+	}
+	select {
+	case <-w.donec:
+		return w.err
+	default:
+	}
+	return nil
+}
 
+// pr may be nil if we're Closing and no data was written.
+func (w *writer) open(r io.Reader, closeReaderOnError bool) error {
 	go func() {
 		defer close(w.donec)
 
-		body := io.Reader(pr)
-		if pr == nil {
+		if r == nil {
 			// AWS doesn't like a nil Body.
-			body = http.NoBody
+			r = http.NoBody
 		}
 		var err error
 		if w.useV2 {
-			w.reqV2.Body = body
+			w.reqV2.Body = r
 			_, err = w.uploaderV2.Upload(w.ctx, w.reqV2)
 		} else {
-			w.req.Body = body
+			w.req.Body = r
 			_, err = w.uploader.UploadWithContext(w.ctx, w.req)
 		}
 		if err != nil {
-			w.err = err
-			if pr != nil {
-				pr.CloseWithError(err)
+			if closeReaderOnError {
+				w.pr.CloseWithError(err)
 			}
+			w.err = err
 			return
 		}
 	}()
@@ -338,11 +352,16 @@ func (w *writer) open(pr *io.PipeReader) error {
 // will be returned. If a writer is closed before any Write is called, Close
 // will create an empty file at the given key.
 func (w *writer) Close() error {
-	if w.w == nil {
-		// We never got any bytes written. We'll write an http.NoBody.
-		w.open(nil)
-	} else if err := w.w.Close(); err != nil {
-		return err
+	if !w.upload {
+		if w.pr != nil {
+			// defer w.pr.Close()
+		}
+		if w.pw == nil {
+			// We never got any bytes written. We'll write an http.NoBody.
+			w.open(nil, false)
+		} else if err := w.pw.Close(); err != nil {
+			return err
+		}
 	}
 	<-w.donec
 	return w.err