1// Copyright 2016 The Go Authors. All rights reserved. 2// Use of this source code is governed by a BSD-style 3// license that can be found in the LICENSE file. 4 5package gensupport 6 7import ( 8 "bytes" 9 "io" 10 11 "google.golang.org/api/googleapi" 12) 13 14// MediaBuffer buffers data from an io.Reader to support uploading media in 15// retryable chunks. It should be created with NewMediaBuffer. 16type MediaBuffer struct { 17 media io.Reader 18 19 chunk []byte // The current chunk which is pending upload. The capacity is the chunk size. 20 err error // Any error generated when populating chunk by reading media. 21 22 // The absolute position of chunk in the underlying media. 23 off int64 24} 25 26// NewMediaBuffer initializes a MediaBuffer. 27func NewMediaBuffer(media io.Reader, chunkSize int) *MediaBuffer { 28 return &MediaBuffer{media: media, chunk: make([]byte, 0, chunkSize)} 29} 30 31// Chunk returns the current buffered chunk, the offset in the underlying media 32// from which the chunk is drawn, and the size of the chunk. 33// Successive calls to Chunk return the same chunk between calls to Next. 34func (mb *MediaBuffer) Chunk() (chunk io.Reader, off int64, size int, err error) { 35 // There may already be data in chunk if Next has not been called since the previous call to Chunk. 36 if mb.err == nil && len(mb.chunk) == 0 { 37 mb.err = mb.loadChunk() 38 } 39 return bytes.NewReader(mb.chunk), mb.off, len(mb.chunk), mb.err 40} 41 42// loadChunk will read from media into chunk, up to the capacity of chunk. 43func (mb *MediaBuffer) loadChunk() error { 44 bufSize := cap(mb.chunk) 45 mb.chunk = mb.chunk[:bufSize] 46 47 read := 0 48 var err error 49 for err == nil && read < bufSize { 50 var n int 51 n, err = mb.media.Read(mb.chunk[read:]) 52 read += n 53 } 54 mb.chunk = mb.chunk[:read] 55 return err 56} 57 58// Next advances to the next chunk, which will be returned by the next call to Chunk. 59// Calls to Next without a corresponding prior call to Chunk will have no effect. 60func (mb *MediaBuffer) Next() { 61 mb.off += int64(len(mb.chunk)) 62 mb.chunk = mb.chunk[0:0] 63} 64 65type readerTyper struct { 66 io.Reader 67 googleapi.ContentTyper 68} 69 70// ReaderAtToReader adapts a ReaderAt to be used as a Reader. 71// If ra implements googleapi.ContentTyper, then the returned reader 72// will also implement googleapi.ContentTyper, delegating to ra. 73func ReaderAtToReader(ra io.ReaderAt, size int64) io.Reader { 74 r := io.NewSectionReader(ra, 0, size) 75 if typer, ok := ra.(googleapi.ContentTyper); ok { 76 return readerTyper{r, typer} 77 } 78 return r 79} 80