buffer.go raw

   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  
   5  package gensupport
   6  
   7  import (
   8  	"bytes"
   9  	"hash/crc32"
  10  	"io"
  11  
  12  	"google.golang.org/api/googleapi"
  13  )
  14  
  15  // MediaBuffer buffers data from an io.Reader to support uploading media in
  16  // retryable chunks. It should be created with NewMediaBuffer.
  17  type MediaBuffer struct {
  18  	media io.Reader
  19  
  20  	chunk []byte // The current chunk which is pending upload.  The capacity is the chunk size.
  21  	err   error  // Any error generated when populating chunk by reading media.
  22  
  23  	// The absolute position of chunk in the underlying media.
  24  	off int64
  25  
  26  	// fullObjectChecksum holds the running checksum of streamed media chunks when automatic checksum
  27  	// calculation is enabled via enableAutoChecksum.
  28  	fullObjectChecksum uint32
  29  	enableAutoChecksum bool
  30  }
  31  
  32  var (
  33  	crc32cTable = crc32.MakeTable(crc32.Castagnoli)
  34  )
  35  
  36  // NewMediaBuffer initializes a MediaBuffer.
  37  func NewMediaBuffer(media io.Reader, chunkSize int) *MediaBuffer {
  38  	return &MediaBuffer{media: media, chunk: make([]byte, 0, chunkSize)}
  39  }
  40  
  41  // Chunk returns the current buffered chunk, the offset in the underlying media
  42  // from which the chunk is drawn, and the size of the chunk.
  43  // Successive calls to Chunk return the same chunk between calls to Next.
  44  func (mb *MediaBuffer) Chunk() (chunk io.Reader, off int64, size int, err error) {
  45  	// There may already be data in chunk if Next has not been called since the previous call to Chunk.
  46  	if mb.err == nil && len(mb.chunk) == 0 {
  47  		mb.err = mb.loadChunk()
  48  	}
  49  	return bytes.NewReader(mb.chunk), mb.off, len(mb.chunk), mb.err
  50  }
  51  
  52  // loadChunk will read from media into chunk, up to the capacity of chunk.
  53  func (mb *MediaBuffer) loadChunk() error {
  54  	bufSize := cap(mb.chunk)
  55  	mb.chunk = mb.chunk[:bufSize]
  56  
  57  	read := 0
  58  	var err error
  59  	for err == nil && read < bufSize {
  60  		var n int
  61  		n, err = mb.media.Read(mb.chunk[read:])
  62  		read += n
  63  	}
  64  	mb.chunk = mb.chunk[:read]
  65  	if mb.enableAutoChecksum {
  66  		mb.fullObjectChecksum = crc32.Update(mb.fullObjectChecksum, crc32cTable, mb.chunk)
  67  	}
  68  	return err
  69  }
  70  
  71  // Next advances to the next chunk, which will be returned by the next call to Chunk.
  72  // Calls to Next without a corresponding prior call to Chunk will have no effect.
  73  func (mb *MediaBuffer) Next() {
  74  	mb.off += int64(len(mb.chunk))
  75  	mb.chunk = mb.chunk[0:0]
  76  }
  77  
  78  type readerTyper struct {
  79  	io.Reader
  80  	googleapi.ContentTyper
  81  }
  82  
  83  // ReaderAtToReader adapts a ReaderAt to be used as a Reader.
  84  // If ra implements googleapi.ContentTyper, then the returned reader
  85  // will also implement googleapi.ContentTyper, delegating to ra.
  86  func ReaderAtToReader(ra io.ReaderAt, size int64) io.Reader {
  87  	r := io.NewSectionReader(ra, 0, size)
  88  	if typer, ok := ra.(googleapi.ContentTyper); ok {
  89  		return readerTyper{r, typer}
  90  	}
  91  	return r
  92  }
  93