github.com/rclone/rclone@v1.66.1-0.20240517100346-7b89735ae726/backend/quatrix/upload_memory.go (about)

     1  package quatrix
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/rclone/rclone/fs"
     8  )
     9  
    10  // UploadMemoryManager dynamically calculates every chunk size for the transfer and increases or decreases it
    11  // depending on the upload speed. This makes general upload time smaller, because transfers that are faster
    12  // does not have to wait for the slower ones until they finish upload.
    13  type UploadMemoryManager struct {
    14  	m              sync.Mutex
    15  	useDynamicSize bool
    16  	shared         int64
    17  	reserved       int64
    18  	effectiveTime  time.Duration
    19  	fileUsage      map[string]int64
    20  }
    21  
    22  // NewUploadMemoryManager is a constructor for UploadMemoryManager
    23  func NewUploadMemoryManager(ci *fs.ConfigInfo, opt *Options) *UploadMemoryManager {
    24  	useDynamicSize := true
    25  
    26  	sharedMemory := int64(opt.MaximalSummaryChunkSize) - int64(opt.MinimalChunkSize)*int64(ci.Transfers)
    27  	if sharedMemory <= 0 {
    28  		sharedMemory = 0
    29  		useDynamicSize = false
    30  	}
    31  
    32  	return &UploadMemoryManager{
    33  		useDynamicSize: useDynamicSize,
    34  		shared:         sharedMemory,
    35  		reserved:       int64(opt.MinimalChunkSize),
    36  		effectiveTime:  time.Duration(opt.EffectiveUploadTime),
    37  		fileUsage:      map[string]int64{},
    38  	}
    39  }
    40  
    41  // Consume -- decide amount of memory to consume
    42  func (u *UploadMemoryManager) Consume(fileID string, neededMemory int64, speed float64) int64 {
    43  	if !u.useDynamicSize {
    44  		if neededMemory < u.reserved {
    45  			return neededMemory
    46  		}
    47  
    48  		return u.reserved
    49  	}
    50  
    51  	u.m.Lock()
    52  	defer u.m.Unlock()
    53  
    54  	borrowed, found := u.fileUsage[fileID]
    55  	if found {
    56  		u.shared += borrowed
    57  		borrowed = 0
    58  	}
    59  
    60  	defer func() { u.fileUsage[fileID] = borrowed }()
    61  
    62  	effectiveChunkSize := int64(speed * u.effectiveTime.Seconds())
    63  
    64  	if effectiveChunkSize < u.reserved {
    65  		effectiveChunkSize = u.reserved
    66  	}
    67  
    68  	if neededMemory < effectiveChunkSize {
    69  		effectiveChunkSize = neededMemory
    70  	}
    71  
    72  	if effectiveChunkSize <= u.reserved {
    73  		return effectiveChunkSize
    74  	}
    75  
    76  	toBorrow := effectiveChunkSize - u.reserved
    77  
    78  	if toBorrow <= u.shared {
    79  		u.shared -= toBorrow
    80  		borrowed = toBorrow
    81  
    82  		return effectiveChunkSize
    83  	}
    84  
    85  	borrowed = u.shared
    86  	u.shared = 0
    87  
    88  	return borrowed + u.reserved
    89  }
    90  
    91  // Return returns consumed memory for the previous chunk upload to the memory pool
    92  func (u *UploadMemoryManager) Return(fileID string) {
    93  	if !u.useDynamicSize {
    94  		return
    95  	}
    96  
    97  	u.m.Lock()
    98  	defer u.m.Unlock()
    99  
   100  	borrowed, found := u.fileUsage[fileID]
   101  	if !found {
   102  		return
   103  	}
   104  
   105  	u.shared += borrowed
   106  
   107  	delete(u.fileUsage, fileID)
   108  }