You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

80 lines
2.4 KiB

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