github.com/10XDev/rclone@v1.52.3-0.20200626220027-16af9ab76b2a/fs/operations/multithread.go (about)

     1  package operations
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/pkg/errors"
     8  	"github.com/rclone/rclone/fs"
     9  	"github.com/rclone/rclone/fs/accounting"
    10  	"golang.org/x/sync/errgroup"
    11  )
    12  
    13  const (
    14  	multithreadChunkSize     = 64 << 10
    15  	multithreadChunkSizeMask = multithreadChunkSize - 1
    16  	multithreadBufferSize    = 32 * 1024
    17  )
    18  
    19  // Return a boolean as to whether we should use multi thread copy for
    20  // this transfer
    21  func doMultiThreadCopy(f fs.Fs, src fs.Object) bool {
    22  	// Disable multi thread if...
    23  
    24  	// ...it isn't configured
    25  	if fs.Config.MultiThreadStreams <= 1 {
    26  		return false
    27  	}
    28  	// ...size of object is less than cutoff
    29  	if src.Size() < int64(fs.Config.MultiThreadCutoff) {
    30  		return false
    31  	}
    32  	// ...source doesn't support it
    33  	dstFeatures := f.Features()
    34  	if dstFeatures.OpenWriterAt == nil {
    35  		return false
    36  	}
    37  	// ...if --multi-thread-streams not in use and source and
    38  	// destination are both local
    39  	if !fs.Config.MultiThreadSet && dstFeatures.IsLocal && src.Fs().Features().IsLocal {
    40  		return false
    41  	}
    42  	return true
    43  }
    44  
    45  // state for a multi-thread copy
    46  type multiThreadCopyState struct {
    47  	ctx      context.Context
    48  	partSize int64
    49  	size     int64
    50  	wc       fs.WriterAtCloser
    51  	src      fs.Object
    52  	acc      *accounting.Account
    53  	streams  int
    54  }
    55  
    56  // Copy a single stream into place
    57  func (mc *multiThreadCopyState) copyStream(ctx context.Context, stream int) (err error) {
    58  	defer func() {
    59  		if err != nil {
    60  			fs.Debugf(mc.src, "multi-thread copy: stream %d/%d failed: %v", stream+1, mc.streams, err)
    61  		}
    62  	}()
    63  	start := int64(stream) * mc.partSize
    64  	if start >= mc.size {
    65  		return nil
    66  	}
    67  	end := start + mc.partSize
    68  	if end > mc.size {
    69  		end = mc.size
    70  	}
    71  
    72  	fs.Debugf(mc.src, "multi-thread copy: stream %d/%d (%d-%d) size %v starting", stream+1, mc.streams, start, end, fs.SizeSuffix(end-start))
    73  
    74  	rc, err := NewReOpen(ctx, mc.src, fs.Config.LowLevelRetries, &fs.RangeOption{Start: start, End: end - 1})
    75  	if err != nil {
    76  		return errors.Wrap(err, "multpart copy: failed to open source")
    77  	}
    78  	defer fs.CheckClose(rc, &err)
    79  
    80  	// Copy the data
    81  	buf := make([]byte, multithreadBufferSize)
    82  	offset := start
    83  	for {
    84  		// Check if context cancelled and exit if so
    85  		if mc.ctx.Err() != nil {
    86  			return mc.ctx.Err()
    87  		}
    88  		nr, er := rc.Read(buf)
    89  		if nr > 0 {
    90  			err = mc.acc.AccountRead(nr)
    91  			if err != nil {
    92  				return errors.Wrap(err, "multpart copy: accounting failed")
    93  			}
    94  			nw, ew := mc.wc.WriteAt(buf[0:nr], offset)
    95  			if nw > 0 {
    96  				offset += int64(nw)
    97  			}
    98  			if ew != nil {
    99  				return errors.Wrap(ew, "multpart copy: write failed")
   100  			}
   101  			if nr != nw {
   102  				return errors.Wrap(io.ErrShortWrite, "multpart copy")
   103  			}
   104  		}
   105  		if er != nil {
   106  			if er != io.EOF {
   107  				return errors.Wrap(er, "multpart copy: read failed")
   108  			}
   109  			break
   110  		}
   111  	}
   112  
   113  	if offset != end {
   114  		return errors.Errorf("multpart copy: wrote %d bytes but expected to write %d", offset-start, end-start)
   115  	}
   116  
   117  	fs.Debugf(mc.src, "multi-thread copy: stream %d/%d (%d-%d) size %v finished", stream+1, mc.streams, start, end, fs.SizeSuffix(end-start))
   118  	return nil
   119  }
   120  
   121  // Calculate the chunk sizes and updated number of streams
   122  func (mc *multiThreadCopyState) calculateChunks() {
   123  	partSize := mc.size / int64(mc.streams)
   124  	// Round partition size up so partSize * streams >= size
   125  	if (mc.size % int64(mc.streams)) != 0 {
   126  		partSize++
   127  	}
   128  	// round partSize up to nearest multithreadChunkSize boundary
   129  	mc.partSize = (partSize + multithreadChunkSizeMask) &^ multithreadChunkSizeMask
   130  	// recalculate number of streams
   131  	mc.streams = int(mc.size / mc.partSize)
   132  	// round streams up so partSize * streams >= size
   133  	if (mc.size % mc.partSize) != 0 {
   134  		mc.streams++
   135  	}
   136  }
   137  
   138  // Copy src to (f, remote) using streams download threads and the OpenWriterAt feature
   139  func multiThreadCopy(ctx context.Context, f fs.Fs, remote string, src fs.Object, streams int, tr *accounting.Transfer) (newDst fs.Object, err error) {
   140  	openWriterAt := f.Features().OpenWriterAt
   141  	if openWriterAt == nil {
   142  		return nil, errors.New("multi-thread copy: OpenWriterAt not supported")
   143  	}
   144  	if src.Size() < 0 {
   145  		return nil, errors.New("multi-thread copy: can't copy unknown sized file")
   146  	}
   147  	if src.Size() == 0 {
   148  		return nil, errors.New("multi-thread copy: can't copy zero sized file")
   149  	}
   150  
   151  	g, gCtx := errgroup.WithContext(ctx)
   152  	mc := &multiThreadCopyState{
   153  		ctx:     gCtx,
   154  		size:    src.Size(),
   155  		src:     src,
   156  		streams: streams,
   157  	}
   158  	mc.calculateChunks()
   159  
   160  	// Make accounting
   161  	mc.acc = tr.Account(nil)
   162  
   163  	// create write file handle
   164  	mc.wc, err = openWriterAt(gCtx, remote, mc.size)
   165  	if err != nil {
   166  		return nil, errors.Wrap(err, "multpart copy: failed to open destination")
   167  	}
   168  
   169  	fs.Debugf(src, "Starting multi-thread copy with %d parts of size %v", mc.streams, fs.SizeSuffix(mc.partSize))
   170  	for stream := 0; stream < mc.streams; stream++ {
   171  		stream := stream
   172  		g.Go(func() (err error) {
   173  			return mc.copyStream(gCtx, stream)
   174  		})
   175  	}
   176  	err = g.Wait()
   177  	closeErr := mc.wc.Close()
   178  	if err != nil {
   179  		return nil, err
   180  	}
   181  	if closeErr != nil {
   182  		return nil, errors.Wrap(closeErr, "multi-thread copy: failed to close object after copy")
   183  	}
   184  
   185  	obj, err := f.NewObject(ctx, remote)
   186  	if err != nil {
   187  		return nil, errors.Wrap(err, "multi-thread copy: failed to find object after copy")
   188  	}
   189  
   190  	err = obj.SetModTime(ctx, src.ModTime(ctx))
   191  	switch err {
   192  	case nil, fs.ErrorCantSetModTime, fs.ErrorCantSetModTimeWithoutDelete:
   193  	default:
   194  		return nil, errors.Wrap(err, "multi-thread copy: failed to set modification time")
   195  	}
   196  
   197  	fs.Debugf(src, "Finished multi-thread copy with %d parts of size %v", mc.streams, fs.SizeSuffix(mc.partSize))
   198  	return obj, nil
   199  }