github.com/xhghs/rclone@v1.51.1-0.20200430155106-e186a28cced8/lib/oauthutil/renew.go (about)

     1  package oauthutil
     2  
     3  import (
     4  	"sync/atomic"
     5  
     6  	"github.com/rclone/rclone/fs"
     7  )
     8  
     9  // Renew allows tokens to be renewed on expiry if uploads are in progress.
    10  type Renew struct {
    11  	name    string       // name to use in logs
    12  	ts      *TokenSource // token source that needs renewing
    13  	uploads int32        // number of uploads in progress - atomic access required
    14  	run     func() error // a transaction to run to renew the token on
    15  }
    16  
    17  // NewRenew creates a new Renew struct and starts a background process
    18  // which renews the token whenever it expires.  It uses the run() call
    19  // to run a transaction to do this.
    20  //
    21  // It will only renew the token if the number of uploads > 0
    22  func NewRenew(name string, ts *TokenSource, run func() error) *Renew {
    23  	r := &Renew{
    24  		name: name,
    25  		ts:   ts,
    26  		run:  run,
    27  	}
    28  	go r.renewOnExpiry()
    29  	return r
    30  }
    31  
    32  // renewOnExpiry renews the token whenever it expires.  Useful when there
    33  // are lots of uploads in progress and the token doesn't get renewed.
    34  // Amazon seem to cancel your uploads if you don't renew your token
    35  // for 2hrs.
    36  func (r *Renew) renewOnExpiry() {
    37  	expiry := r.ts.OnExpiry()
    38  	for {
    39  		<-expiry
    40  		uploads := atomic.LoadInt32(&r.uploads)
    41  		if uploads != 0 {
    42  			fs.Debugf(r.name, "Token expired - %d uploads in progress - refreshing", uploads)
    43  			// Do a transaction
    44  			err := r.run()
    45  			if err == nil {
    46  				fs.Debugf(r.name, "Token refresh successful")
    47  			} else {
    48  				fs.Errorf(r.name, "Token refresh failed: %v", err)
    49  			}
    50  		} else {
    51  			fs.Debugf(r.name, "Token expired but no uploads in progress - doing nothing")
    52  		}
    53  	}
    54  }
    55  
    56  // Start should be called before starting an upload
    57  func (r *Renew) Start() {
    58  	atomic.AddInt32(&r.uploads, 1)
    59  }
    60  
    61  // Stop should be called after finishing an upload
    62  func (r *Renew) Stop() {
    63  	atomic.AddInt32(&r.uploads, -1)
    64  }
    65  
    66  // Invalidate invalidates the token source
    67  func (r *Renew) Invalidate() {
    68  	r.ts.Invalidate()
    69  }