249 lines
7.1 KiB
Go
249 lines
7.1 KiB
Go
package chunkedreader
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"sync"
|
|
|
|
"github.com/rclone/rclone/fs"
|
|
"github.com/rclone/rclone/fs/hash"
|
|
)
|
|
|
|
// io related errors returned by ChunkedReader
|
|
var (
|
|
ErrorFileClosed = errors.New("file already closed")
|
|
ErrorInvalidSeek = errors.New("invalid seek position")
|
|
)
|
|
|
|
// ChunkedReader is a reader for a Object with the possibility
|
|
// of reading the source in chunks of given size
|
|
//
|
|
// A initialChunkSize of <= 0 will disable chunked reading.
|
|
type ChunkedReader struct {
|
|
ctx context.Context
|
|
mu sync.Mutex // protects following fields
|
|
o fs.Object // source to read from
|
|
rc io.ReadCloser // reader for the current open chunk
|
|
offset int64 // offset the next Read will start. -1 forces a reopen of o
|
|
chunkOffset int64 // beginning of the current or next chunk
|
|
chunkSize int64 // length of the current or next chunk. -1 will open o from chunkOffset to the end
|
|
initialChunkSize int64 // default chunkSize after the chunk specified by RangeSeek is complete
|
|
maxChunkSize int64 // consecutive read chunks will double in size until reached. -1 means no limit
|
|
customChunkSize bool // is the current chunkSize set by RangeSeek?
|
|
closed bool // has Close been called?
|
|
}
|
|
|
|
// New returns a ChunkedReader for the Object.
|
|
//
|
|
// A initialChunkSize of <= 0 will disable chunked reading.
|
|
// If maxChunkSize is greater than initialChunkSize, the chunk size will be
|
|
// doubled after each chunk read with a maximun of maxChunkSize.
|
|
// A Seek or RangeSeek will reset the chunk size to it's initial value
|
|
func New(ctx context.Context, o fs.Object, initialChunkSize int64, maxChunkSize int64) *ChunkedReader {
|
|
if initialChunkSize <= 0 {
|
|
initialChunkSize = -1
|
|
}
|
|
if maxChunkSize != -1 && maxChunkSize < initialChunkSize {
|
|
maxChunkSize = initialChunkSize
|
|
}
|
|
return &ChunkedReader{
|
|
ctx: ctx,
|
|
o: o,
|
|
offset: -1,
|
|
chunkSize: initialChunkSize,
|
|
initialChunkSize: initialChunkSize,
|
|
maxChunkSize: maxChunkSize,
|
|
}
|
|
}
|
|
|
|
// Read from the file - for details see io.Reader
|
|
func (cr *ChunkedReader) Read(p []byte) (n int, err error) {
|
|
cr.mu.Lock()
|
|
defer cr.mu.Unlock()
|
|
|
|
if cr.closed {
|
|
return 0, ErrorFileClosed
|
|
}
|
|
|
|
for reqSize := int64(len(p)); reqSize > 0; reqSize = int64(len(p)) {
|
|
// the current chunk boundary. valid only when chunkSize > 0
|
|
chunkEnd := cr.chunkOffset + cr.chunkSize
|
|
|
|
fs.Debugf(cr.o, "ChunkedReader.Read at %d length %d chunkOffset %d chunkSize %d", cr.offset, reqSize, cr.chunkOffset, cr.chunkSize)
|
|
|
|
switch {
|
|
case cr.chunkSize > 0 && cr.offset == chunkEnd: // last chunk read completely
|
|
cr.chunkOffset = cr.offset
|
|
if cr.customChunkSize { // last chunkSize was set by RangeSeek
|
|
cr.customChunkSize = false
|
|
cr.chunkSize = cr.initialChunkSize
|
|
} else {
|
|
cr.chunkSize *= 2
|
|
if cr.chunkSize > cr.maxChunkSize && cr.maxChunkSize != -1 {
|
|
cr.chunkSize = cr.maxChunkSize
|
|
}
|
|
}
|
|
// recalculate the chunk boundary. valid only when chunkSize > 0
|
|
chunkEnd = cr.chunkOffset + cr.chunkSize
|
|
fallthrough
|
|
case cr.offset == -1: // first Read or Read after RangeSeek
|
|
err = cr.openRange()
|
|
if err != nil {
|
|
return
|
|
}
|
|
}
|
|
|
|
var buf []byte
|
|
chunkRest := chunkEnd - cr.offset
|
|
// limit read to chunk boundaries if chunkSize > 0
|
|
if reqSize > chunkRest && cr.chunkSize > 0 {
|
|
buf, p = p[0:chunkRest], p[chunkRest:]
|
|
} else {
|
|
buf, p = p, nil
|
|
}
|
|
var rn int
|
|
rn, err = io.ReadFull(cr.rc, buf)
|
|
n += rn
|
|
cr.offset += int64(rn)
|
|
if err != nil {
|
|
if err == io.ErrUnexpectedEOF {
|
|
err = io.EOF
|
|
}
|
|
return
|
|
}
|
|
}
|
|
return n, nil
|
|
}
|
|
|
|
// Close the file - for details see io.Closer
|
|
//
|
|
// All methods on ChunkedReader will return ErrorFileClosed afterwards
|
|
func (cr *ChunkedReader) Close() error {
|
|
cr.mu.Lock()
|
|
defer cr.mu.Unlock()
|
|
|
|
if cr.closed {
|
|
return ErrorFileClosed
|
|
}
|
|
cr.closed = true
|
|
|
|
return cr.resetReader(nil, 0)
|
|
}
|
|
|
|
// Seek the file - for details see io.Seeker
|
|
func (cr *ChunkedReader) Seek(offset int64, whence int) (int64, error) {
|
|
return cr.RangeSeek(context.TODO(), offset, whence, -1)
|
|
}
|
|
|
|
// RangeSeek the file - for details see RangeSeeker
|
|
//
|
|
// The specified length will only apply to the next chunk opened.
|
|
// RangeSeek will not reopen the source until Read is called.
|
|
func (cr *ChunkedReader) RangeSeek(ctx context.Context, offset int64, whence int, length int64) (int64, error) {
|
|
cr.mu.Lock()
|
|
defer cr.mu.Unlock()
|
|
|
|
fs.Debugf(cr.o, "ChunkedReader.RangeSeek from %d to %d length %d", cr.offset, offset, length)
|
|
|
|
if cr.closed {
|
|
return 0, ErrorFileClosed
|
|
}
|
|
|
|
size := cr.o.Size()
|
|
switch whence {
|
|
case io.SeekStart:
|
|
cr.offset = 0
|
|
case io.SeekEnd:
|
|
cr.offset = size
|
|
}
|
|
// set the new chunk start
|
|
cr.chunkOffset = cr.offset + offset
|
|
// force reopen on next Read
|
|
cr.offset = -1
|
|
if length > 0 {
|
|
cr.customChunkSize = true
|
|
cr.chunkSize = length
|
|
} else {
|
|
cr.chunkSize = cr.initialChunkSize
|
|
}
|
|
if cr.chunkOffset < 0 || cr.chunkOffset >= size {
|
|
cr.chunkOffset = 0
|
|
return 0, ErrorInvalidSeek
|
|
}
|
|
return cr.chunkOffset, nil
|
|
}
|
|
|
|
// Open forces the connection to be opened
|
|
func (cr *ChunkedReader) Open() (*ChunkedReader, error) {
|
|
cr.mu.Lock()
|
|
defer cr.mu.Unlock()
|
|
|
|
if cr.rc != nil && cr.offset != -1 {
|
|
return cr, nil
|
|
}
|
|
return cr, cr.openRange()
|
|
}
|
|
|
|
// openRange will open the source Object with the current chunk range
|
|
//
|
|
// If the current open reader implements RangeSeeker, it is tried first.
|
|
// When RangeSeek fails, o.Open with a RangeOption is used.
|
|
//
|
|
// A length <= 0 will request till the end of the file
|
|
func (cr *ChunkedReader) openRange() error {
|
|
offset, length := cr.chunkOffset, cr.chunkSize
|
|
fs.Debugf(cr.o, "ChunkedReader.openRange at %d length %d", offset, length)
|
|
|
|
if cr.closed {
|
|
return ErrorFileClosed
|
|
}
|
|
|
|
if rs, ok := cr.rc.(fs.RangeSeeker); ok {
|
|
n, err := rs.RangeSeek(cr.ctx, offset, io.SeekStart, length)
|
|
if err == nil && n == offset {
|
|
cr.offset = offset
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
fs.Debugf(cr.o, "ChunkedReader.openRange seek failed (%s). Trying Open", err)
|
|
} else {
|
|
fs.Debugf(cr.o, "ChunkedReader.openRange seeked to wrong offset. Wanted %d, got %d. Trying Open", offset, n)
|
|
}
|
|
}
|
|
|
|
var rc io.ReadCloser
|
|
var err error
|
|
if length <= 0 {
|
|
if offset == 0 {
|
|
rc, err = cr.o.Open(cr.ctx, &fs.HashesOption{Hashes: hash.Set(hash.None)})
|
|
} else {
|
|
rc, err = cr.o.Open(cr.ctx, &fs.HashesOption{Hashes: hash.Set(hash.None)}, &fs.RangeOption{Start: offset, End: -1})
|
|
}
|
|
} else {
|
|
rc, err = cr.o.Open(cr.ctx, &fs.HashesOption{Hashes: hash.Set(hash.None)}, &fs.RangeOption{Start: offset, End: offset + length - 1})
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return cr.resetReader(rc, offset)
|
|
}
|
|
|
|
// resetReader switches the current reader to the given reader.
|
|
// The old reader will be Close'd before setting the new reader.
|
|
func (cr *ChunkedReader) resetReader(rc io.ReadCloser, offset int64) error {
|
|
if cr.rc != nil {
|
|
if err := cr.rc.Close(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
cr.rc = rc
|
|
cr.offset = offset
|
|
return nil
|
|
}
|
|
|
|
var (
|
|
_ io.ReadCloser = (*ChunkedReader)(nil)
|
|
_ io.Seeker = (*ChunkedReader)(nil)
|
|
_ fs.RangeSeeker = (*ChunkedReader)(nil)
|
|
)
|