github.com/divyam234/rclone@v1.64.1/fs/accounting/accounting.go (about) 1 // Package accounting providers an accounting and limiting reader 2 package accounting 3 4 import ( 5 "context" 6 "errors" 7 "fmt" 8 "io" 9 "sync" 10 "time" 11 "unicode/utf8" 12 13 "github.com/divyam234/rclone/fs/rc" 14 15 "github.com/divyam234/rclone/fs" 16 "github.com/divyam234/rclone/fs/asyncreader" 17 "github.com/divyam234/rclone/fs/fserrors" 18 ) 19 20 // ErrorMaxTransferLimitReached defines error when transfer limit is reached. 21 // Used for checking on exit and matching to correct exit code. 22 var ErrorMaxTransferLimitReached = errors.New("max transfer limit reached as set by --max-transfer") 23 24 // ErrorMaxTransferLimitReachedFatal is returned from Read when the max 25 // transfer limit is reached. 26 var ErrorMaxTransferLimitReachedFatal = fserrors.FatalError(ErrorMaxTransferLimitReached) 27 28 // ErrorMaxTransferLimitReachedGraceful is returned from operations.Copy when the max 29 // transfer limit is reached and a graceful stop is required. 30 var ErrorMaxTransferLimitReachedGraceful = fserrors.NoRetryError(ErrorMaxTransferLimitReached) 31 32 // Start sets up the accounting, in particular the bandwidth limiting 33 func Start(ctx context.Context) { 34 // Start the token bucket limiter 35 TokenBucket.StartTokenBucket(ctx) 36 37 // Start the bandwidth update ticker 38 TokenBucket.StartTokenTicker(ctx) 39 40 // Start the transactions per second limiter 41 StartLimitTPS(ctx) 42 } 43 44 // Account limits and accounts for one transfer 45 type Account struct { 46 stats *StatsInfo 47 // The mutex is to make sure Read() and Close() aren't called 48 // concurrently. Unfortunately the persistent connection loop 49 // in http transport calls Read() after Do() returns on 50 // CancelRequest so this race can happen when it apparently 51 // shouldn't. 52 mu sync.Mutex // mutex protects these values 53 in io.Reader 54 ctx context.Context // current context for transfer - may change 55 ci *fs.ConfigInfo 56 origIn io.ReadCloser 57 close io.Closer 58 size int64 59 name string 60 closed bool // set if the file is closed 61 exit chan struct{} // channel that will be closed when transfer is finished 62 withBuf bool // is using a buffered in 63 64 tokenBucket buckets // per file bandwidth limiter (may be nil) 65 66 values accountValues 67 } 68 69 // accountValues holds statistics for this Account 70 type accountValues struct { 71 mu sync.Mutex // Mutex for stat values. 72 bytes int64 // Total number of bytes read 73 max int64 // if >=0 the max number of bytes to transfer 74 start time.Time // Start time of first read 75 lpTime time.Time // Time of last average measurement 76 lpBytes int // Number of bytes read since last measurement 77 avg float64 // Moving average of last few measurements in Byte/s 78 } 79 80 const averagePeriod = 16 // period to do exponentially weighted averages over 81 82 // newAccountSizeName makes an Account reader for an io.ReadCloser of 83 // the given size and name 84 func newAccountSizeName(ctx context.Context, stats *StatsInfo, in io.ReadCloser, size int64, name string) *Account { 85 acc := &Account{ 86 stats: stats, 87 in: in, 88 ctx: ctx, 89 ci: fs.GetConfig(ctx), 90 close: in, 91 origIn: in, 92 size: size, 93 name: name, 94 exit: make(chan struct{}), 95 values: accountValues{ 96 avg: 0, 97 lpTime: time.Now(), 98 max: -1, 99 }, 100 } 101 if acc.ci.CutoffMode == fs.CutoffModeHard { 102 acc.values.max = int64((acc.ci.MaxTransfer)) 103 } 104 currLimit := acc.ci.BwLimitFile.LimitAt(time.Now()) 105 if currLimit.Bandwidth.IsSet() { 106 fs.Debugf(acc.name, "Limiting file transfer to %v", currLimit.Bandwidth) 107 acc.tokenBucket = newTokenBucket(currLimit.Bandwidth) 108 } 109 110 go acc.averageLoop() 111 stats.inProgress.set(acc.name, acc) 112 return acc 113 } 114 115 // WithBuffer - If the file is above a certain size it adds an Async reader 116 func (acc *Account) WithBuffer() *Account { 117 // if already have a buffer then just return 118 if acc.withBuf { 119 return acc 120 } 121 acc.withBuf = true 122 var buffers int 123 if acc.size >= int64(acc.ci.BufferSize) || acc.size == -1 { 124 buffers = int(int64(acc.ci.BufferSize) / asyncreader.BufferSize) 125 } else { 126 buffers = int(acc.size / asyncreader.BufferSize) 127 } 128 // On big files add a buffer 129 if buffers > 0 { 130 rc, err := asyncreader.New(acc.ctx, acc.origIn, buffers) 131 if err != nil { 132 fs.Errorf(acc.name, "Failed to make buffer: %v", err) 133 } else { 134 acc.in = rc 135 acc.close = rc 136 } 137 } 138 return acc 139 } 140 141 // HasBuffer - returns true if this Account has an AsyncReader with a buffer 142 func (acc *Account) HasBuffer() bool { 143 acc.mu.Lock() 144 defer acc.mu.Unlock() 145 _, ok := acc.in.(*asyncreader.AsyncReader) 146 return ok 147 } 148 149 // GetReader returns the underlying io.ReadCloser under any Buffer 150 func (acc *Account) GetReader() io.ReadCloser { 151 acc.mu.Lock() 152 defer acc.mu.Unlock() 153 return acc.origIn 154 } 155 156 // GetAsyncReader returns the current AsyncReader or nil if Account is unbuffered 157 func (acc *Account) GetAsyncReader() *asyncreader.AsyncReader { 158 acc.mu.Lock() 159 defer acc.mu.Unlock() 160 if asyncIn, ok := acc.in.(*asyncreader.AsyncReader); ok { 161 return asyncIn 162 } 163 return nil 164 } 165 166 // StopBuffering stops the async buffer doing any more buffering 167 func (acc *Account) StopBuffering() { 168 if asyncIn, ok := acc.in.(*asyncreader.AsyncReader); ok { 169 asyncIn.StopBuffering() 170 } 171 } 172 173 // Abandon stops the async buffer doing any more buffering 174 func (acc *Account) Abandon() { 175 if asyncIn, ok := acc.in.(*asyncreader.AsyncReader); ok { 176 asyncIn.Abandon() 177 } 178 } 179 180 // UpdateReader updates the underlying io.ReadCloser stopping the 181 // async buffer (if any) and re-adding it 182 func (acc *Account) UpdateReader(ctx context.Context, in io.ReadCloser) { 183 acc.mu.Lock() 184 withBuf := acc.withBuf 185 if withBuf { 186 acc.Abandon() 187 acc.withBuf = false 188 } 189 acc.in = in 190 acc.ctx = ctx 191 acc.close = in 192 acc.origIn = in 193 acc.closed = false 194 if withBuf { 195 acc.WithBuffer() 196 } 197 acc.mu.Unlock() 198 199 // Reset counter to stop percentage going over 100% 200 acc.values.mu.Lock() 201 acc.values.lpBytes = 0 202 acc.values.bytes = 0 203 acc.values.mu.Unlock() 204 } 205 206 // averageLoop calculates averages for the stats in the background 207 func (acc *Account) averageLoop() { 208 tick := time.NewTicker(time.Second) 209 var period float64 210 defer tick.Stop() 211 for { 212 select { 213 case now := <-tick.C: 214 acc.values.mu.Lock() 215 // Add average of last second. 216 elapsed := now.Sub(acc.values.lpTime).Seconds() 217 avg := 0.0 218 if elapsed > 0 { 219 avg = float64(acc.values.lpBytes) / elapsed 220 } 221 // Soft start the moving average 222 if period < averagePeriod { 223 period++ 224 } 225 acc.values.avg = (avg + (period-1)*acc.values.avg) / period 226 acc.values.lpBytes = 0 227 acc.values.lpTime = now 228 // Unlock stats 229 acc.values.mu.Unlock() 230 case <-acc.exit: 231 return 232 } 233 } 234 } 235 236 // Check the read before it has happened is valid returning the number 237 // of bytes remaining to read. 238 func (acc *Account) checkReadBefore() (bytesUntilLimit int64, err error) { 239 // Check to see if context is cancelled 240 if err = acc.ctx.Err(); err != nil { 241 return 0, err 242 } 243 acc.values.mu.Lock() 244 if acc.values.max >= 0 { 245 bytesUntilLimit = acc.values.max - acc.stats.GetBytes() 246 if bytesUntilLimit < 0 { 247 acc.values.mu.Unlock() 248 return bytesUntilLimit, ErrorMaxTransferLimitReachedFatal 249 } 250 } else { 251 bytesUntilLimit = 1 << 62 252 } 253 // Set start time. 254 if acc.values.start.IsZero() { 255 acc.values.start = time.Now() 256 } 257 acc.values.mu.Unlock() 258 return bytesUntilLimit, nil 259 } 260 261 // Check the read call after the read has happened 262 func (acc *Account) checkReadAfter(bytesUntilLimit int64, n int, err error) (outN int, outErr error) { 263 bytesUntilLimit -= int64(n) 264 if bytesUntilLimit < 0 { 265 // chop the overage off 266 n += int(bytesUntilLimit) 267 if n < 0 { 268 n = 0 269 } 270 err = ErrorMaxTransferLimitReachedFatal 271 } 272 return n, err 273 } 274 275 // ServerSideTransferStart should be called at the start of a server-side transfer 276 // 277 // This pretends a transfer has started 278 func (acc *Account) ServerSideTransferStart() { 279 acc.values.mu.Lock() 280 // Set start time. 281 if acc.values.start.IsZero() { 282 acc.values.start = time.Now() 283 } 284 acc.values.mu.Unlock() 285 } 286 287 // ServerSideTransferEnd accounts for a read of n bytes in a sever 288 // side transfer to be treated as a normal transfer. 289 func (acc *Account) ServerSideTransferEnd(n int64) { 290 // Update Stats 291 acc.values.mu.Lock() 292 acc.values.bytes += n 293 acc.values.mu.Unlock() 294 295 acc.stats.Bytes(n) 296 } 297 298 // ServerSideCopyEnd accounts for a read of n bytes in a sever side copy 299 func (acc *Account) ServerSideCopyEnd(n int64) { 300 acc.stats.AddServerSideCopy(n) 301 } 302 303 // ServerSideMoveEnd accounts for a read of n bytes in a sever side move 304 func (acc *Account) ServerSideMoveEnd(n int64) { 305 acc.stats.AddServerSideMove(n) 306 } 307 308 // DryRun accounts for statistics without running the operation 309 func (acc *Account) DryRun(n int64) { 310 acc.ServerSideTransferStart() 311 acc.ServerSideTransferEnd(n) 312 } 313 314 // Account for n bytes from the current file bandwidth limit (if any) 315 func (acc *Account) limitPerFileBandwidth(n int) { 316 acc.values.mu.Lock() 317 tokenBucket := acc.tokenBucket[TokenBucketSlotAccounting] 318 acc.values.mu.Unlock() 319 320 if tokenBucket != nil { 321 err := tokenBucket.WaitN(context.Background(), n) 322 if err != nil { 323 fs.Errorf(nil, "Token bucket error: %v", err) 324 } 325 } 326 } 327 328 // Account the read and limit bandwidth 329 func (acc *Account) accountRead(n int) { 330 // Update Stats 331 acc.values.mu.Lock() 332 acc.values.lpBytes += n 333 acc.values.bytes += int64(n) 334 acc.values.mu.Unlock() 335 336 acc.stats.Bytes(int64(n)) 337 338 TokenBucket.LimitBandwidth(TokenBucketSlotAccounting, n) 339 acc.limitPerFileBandwidth(n) 340 } 341 342 // read bytes from the io.Reader passed in and account them 343 func (acc *Account) read(in io.Reader, p []byte) (n int, err error) { 344 bytesUntilLimit, err := acc.checkReadBefore() 345 if err == nil { 346 n, err = in.Read(p) 347 acc.accountRead(n) 348 n, err = acc.checkReadAfter(bytesUntilLimit, n, err) 349 } 350 return n, err 351 } 352 353 // Read bytes from the object - see io.Reader 354 func (acc *Account) Read(p []byte) (n int, err error) { 355 acc.mu.Lock() 356 defer acc.mu.Unlock() 357 return acc.read(acc.in, p) 358 } 359 360 // Thin wrapper for w 361 type accountWriteTo struct { 362 w io.Writer 363 acc *Account 364 } 365 366 // Write writes len(p) bytes from p to the underlying data stream. It 367 // returns the number of bytes written from p (0 <= n <= len(p)) and 368 // any error encountered that caused the write to stop early. Write 369 // must return a non-nil error if it returns n < len(p). Write must 370 // not modify the slice data, even temporarily. 371 // 372 // Implementations must not retain p. 373 func (awt *accountWriteTo) Write(p []byte) (n int, err error) { 374 bytesUntilLimit, err := awt.acc.checkReadBefore() 375 if err == nil { 376 n, err = awt.w.Write(p) 377 n, err = awt.acc.checkReadAfter(bytesUntilLimit, n, err) 378 awt.acc.accountRead(n) 379 } 380 return n, err 381 } 382 383 // WriteTo writes data to w until there's no more data to write or 384 // when an error occurs. The return value n is the number of bytes 385 // written. Any error encountered during the write is also returned. 386 func (acc *Account) WriteTo(w io.Writer) (n int64, err error) { 387 acc.mu.Lock() 388 in := acc.in 389 acc.mu.Unlock() 390 wrappedWriter := accountWriteTo{w: w, acc: acc} 391 if do, ok := in.(io.WriterTo); ok { 392 n, err = do.WriteTo(&wrappedWriter) 393 } else { 394 n, err = io.Copy(&wrappedWriter, in) 395 } 396 return 397 } 398 399 // AccountRead account having read n bytes 400 func (acc *Account) AccountRead(n int) (err error) { 401 acc.mu.Lock() 402 defer acc.mu.Unlock() 403 bytesUntilLimit, err := acc.checkReadBefore() 404 if err == nil { 405 n, err = acc.checkReadAfter(bytesUntilLimit, n, err) 406 acc.accountRead(n) 407 } 408 return err 409 } 410 411 // Close the object 412 func (acc *Account) Close() error { 413 acc.mu.Lock() 414 defer acc.mu.Unlock() 415 if acc.closed { 416 return nil 417 } 418 acc.closed = true 419 if acc.close == nil { 420 return nil 421 } 422 return acc.close.Close() 423 } 424 425 // Done with accounting - must be called to free accounting goroutine 426 func (acc *Account) Done() { 427 acc.mu.Lock() 428 defer acc.mu.Unlock() 429 close(acc.exit) 430 acc.stats.inProgress.clear(acc.name) 431 } 432 433 // progress returns bytes read as well as the size. 434 // Size can be <= 0 if the size is unknown. 435 func (acc *Account) progress() (bytes, size int64) { 436 if acc == nil { 437 return 0, 0 438 } 439 acc.values.mu.Lock() 440 bytes, size = acc.values.bytes, acc.size 441 acc.values.mu.Unlock() 442 return bytes, size 443 } 444 445 // speed returns the speed of the current file transfer 446 // in bytes per second, as well an exponentially weighted moving average 447 // If no read has completed yet, 0 is returned for both values. 448 func (acc *Account) speed() (bps, current float64) { 449 if acc == nil { 450 return 0, 0 451 } 452 acc.values.mu.Lock() 453 defer acc.values.mu.Unlock() 454 if acc.values.bytes == 0 { 455 return 0, 0 456 } 457 // Calculate speed from first read. 458 total := float64(time.Since(acc.values.start)) / float64(time.Second) 459 if total > 0 { 460 bps = float64(acc.values.bytes) / total 461 } else { 462 bps = 0.0 463 } 464 current = acc.values.avg 465 return 466 } 467 468 // eta returns the ETA of the current operation, 469 // rounded to full seconds. 470 // If the ETA cannot be determined 'ok' returns false. 471 func (acc *Account) eta() (etaDuration time.Duration, ok bool) { 472 if acc == nil { 473 return 0, false 474 } 475 acc.values.mu.Lock() 476 defer acc.values.mu.Unlock() 477 return eta(acc.values.bytes, acc.size, acc.values.avg) 478 } 479 480 // shortenName shortens in to size runes long 481 // If size <= 0 then in is left untouched 482 func shortenName(in string, size int) string { 483 if size <= 0 { 484 return in 485 } 486 if utf8.RuneCountInString(in) <= size { 487 return in 488 } 489 name := []rune(in) 490 size-- // don't count ellipsis rune 491 suffixLength := size / 2 492 prefixLength := size - suffixLength 493 suffixStart := len(name) - suffixLength 494 name = append(append(name[:prefixLength], '…'), name[suffixStart:]...) 495 return string(name) 496 } 497 498 // String produces stats for this file 499 func (acc *Account) String() string { 500 a, b := acc.progress() 501 _, cur := acc.speed() 502 eta, etaok := acc.eta() 503 etas := "-" 504 if etaok { 505 if eta > 0 { 506 etas = fmt.Sprintf("%v", eta) 507 } else { 508 etas = "0s" 509 } 510 } 511 512 if acc.ci.DataRateUnit == "bits" { 513 cur = cur * 8 514 } 515 516 percentageDone := 0 517 if b > 0 { 518 percentageDone = int(100 * float64(a) / float64(b)) 519 } 520 521 return fmt.Sprintf("%*s:%3d%% /%s, %s/s, %s", 522 acc.ci.StatsFileNameLength, 523 shortenName(acc.name, acc.ci.StatsFileNameLength), 524 percentageDone, 525 fs.SizeSuffix(b), 526 fs.SizeSuffix(cur), 527 etas, 528 ) 529 } 530 531 // rcStats produces remote control stats for this file 532 func (acc *Account) rcStats() (out rc.Params) { 533 out = make(rc.Params) 534 a, b := acc.progress() 535 out["bytes"] = a 536 out["size"] = b 537 spd, cur := acc.speed() 538 out["speed"] = spd 539 out["speedAvg"] = cur 540 541 eta, etaOK := acc.eta() 542 if etaOK { 543 out["eta"] = eta.Seconds() 544 } else { 545 out["eta"] = nil 546 } 547 out["name"] = acc.name 548 549 percentageDone := 0 550 if b > 0 { 551 percentageDone = int(100 * float64(a) / float64(b)) 552 } 553 out["percentage"] = percentageDone 554 out["group"] = acc.stats.group 555 556 return out 557 } 558 559 // OldStream returns the top io.Reader 560 func (acc *Account) OldStream() io.Reader { 561 acc.mu.Lock() 562 defer acc.mu.Unlock() 563 return acc.in 564 } 565 566 // SetStream updates the top io.Reader 567 func (acc *Account) SetStream(in io.Reader) { 568 acc.mu.Lock() 569 acc.in = in 570 acc.mu.Unlock() 571 } 572 573 // WrapStream wraps an io Reader so it will be accounted in the same 574 // way as account 575 func (acc *Account) WrapStream(in io.Reader) io.Reader { 576 return &accountStream{ 577 acc: acc, 578 in: in, 579 } 580 } 581 582 // accountStream accounts a single io.Reader into a parent *Account 583 type accountStream struct { 584 acc *Account 585 in io.Reader 586 } 587 588 // OldStream return the underlying stream 589 func (a *accountStream) OldStream() io.Reader { 590 return a.in 591 } 592 593 // SetStream set the underlying stream 594 func (a *accountStream) SetStream(in io.Reader) { 595 a.in = in 596 } 597 598 // WrapStream wrap in in an accounter 599 func (a *accountStream) WrapStream(in io.Reader) io.Reader { 600 return a.acc.WrapStream(in) 601 } 602 603 // Read bytes from the object - see io.Reader 604 func (a *accountStream) Read(p []byte) (n int, err error) { 605 return a.acc.read(a.in, p) 606 } 607 608 // Accounter accounts a stream allowing the accounting to be removed and re-added 609 type Accounter interface { 610 io.Reader 611 OldStream() io.Reader 612 SetStream(io.Reader) 613 WrapStream(io.Reader) io.Reader 614 } 615 616 // WrapFn wraps an io.Reader (for accounting purposes usually) 617 type WrapFn func(io.Reader) io.Reader 618 619 // UnWrap unwraps a reader returning unwrapped and wrap, a function to 620 // wrap it back up again. If `in` is an Accounter then this function 621 // will take the accounting unwrapped and wrap will put it back on 622 // again the new Reader passed in. 623 // 624 // This allows functions which wrap io.Readers to move the accounting 625 // to the end of the wrapped chain of readers. This is very important 626 // if buffering is being introduced and if the Reader might be wrapped 627 // again. 628 func UnWrap(in io.Reader) (unwrapped io.Reader, wrap WrapFn) { 629 acc, ok := in.(Accounter) 630 if !ok { 631 return in, func(r io.Reader) io.Reader { return r } 632 } 633 return acc.OldStream(), acc.WrapStream 634 } 635 636 // UnWrapAccounting unwraps a reader returning unwrapped and acc a 637 // pointer to the accounting. 638 // 639 // The caller is expected to manage the accounting at this point. 640 func UnWrapAccounting(in io.Reader) (unwrapped io.Reader, acc *Account) { 641 a, ok := in.(*accountStream) 642 if !ok { 643 return in, nil 644 } 645 return a.in, a.acc 646 }