github.com/puellanivis/breton@v0.2.16/lib/files/copy.go (about)

     1  package files
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"time"
     8  )
     9  
    10  const defaultBufferSize = 64 * 1024
    11  
    12  // ErrWatchdogExpired is returned by files.Copy, if the watchdog time expires during a read.
    13  var ErrWatchdogExpired error = watchdogExpiredError{}
    14  
    15  type watchdogExpiredError struct{}
    16  
    17  func (watchdogExpiredError) Error() string   { return "watchdog expired" }
    18  func (watchdogExpiredError) Timeout() bool   { return true }
    19  func (watchdogExpiredError) Temporary() bool { return true }
    20  
    21  // Copy is a context aware version of io.Copy.
    22  // Do not use to Discard a reader, as a canceled context would stop the read, and it would not be fully discarded.
    23  func Copy(ctx context.Context, dst io.Writer, src io.Reader, opts ...CopyOption) (written int64, err error) {
    24  	if dst == nil {
    25  		return 0, errors.New("nil io.Writer passed to files.Copy")
    26  	}
    27  
    28  	c := new(copyConfig)
    29  
    30  	for _, opt := range opts {
    31  		// intentionally throwing away the reverting functions.
    32  		_ = opt(c)
    33  	}
    34  
    35  	if c.buffer == nil {
    36  		// we allocate a buffer to use as a temporary buffer, rather than alloc new every time.
    37  		c.buffer = make([]byte, defaultBufferSize)
    38  	}
    39  	buflen := int64(len(c.buffer))
    40  
    41  	var keepingMetrics bool
    42  
    43  	var total func(float64)
    44  	if c.bwLifetime != nil {
    45  		total = c.bwLifetime.Observe
    46  		keepingMetrics = true
    47  	}
    48  
    49  	if c.bwScale <= 0 {
    50  		c.bwScale = 1
    51  	}
    52  
    53  	if c.bwInterval <= 0 {
    54  		c.bwInterval = 1 * time.Second
    55  	}
    56  
    57  	type bwSnippet struct {
    58  		n int64
    59  		d time.Duration
    60  	}
    61  	var bwWindow []bwSnippet
    62  
    63  	var running func(float64)
    64  	if c.bwRunning != nil {
    65  		if c.bwCount < 1 {
    66  			c.bwCount = 1
    67  		}
    68  
    69  		running = c.bwRunning.Observe
    70  		keepingMetrics = true
    71  		bwWindow = make([]bwSnippet, c.bwCount)
    72  	}
    73  
    74  	// Prevent an accidental write outside of returning from this function.
    75  	ctx, cancel := context.WithCancel(ctx)
    76  	defer cancel()
    77  
    78  	w := &deadlineWriter{
    79  		ctx: ctx,
    80  		w:   dst,
    81  	}
    82  
    83  	r := &fuzzyLimitedReader{
    84  		R: src,
    85  		N: buflen,
    86  	}
    87  
    88  	t := time.NewTimer(c.runningTimeout)
    89  	if c.runningTimeout <= 0 {
    90  		if !t.Stop() {
    91  			<-t.C
    92  		}
    93  	}
    94  
    95  	start := time.Now()
    96  
    97  	var bwAccum int64
    98  	last := start
    99  	next := last.Add(c.bwInterval)
   100  
   101  	for {
   102  		r.N = buflen // reset fuzzyLimitedReader
   103  
   104  		if c.runningTimeout > 0 {
   105  			if !t.Stop() {
   106  				<-t.C
   107  			}
   108  			t.Reset(c.runningTimeout)
   109  		}
   110  
   111  		var n int64
   112  		done := make(chan struct{})
   113  		go func() {
   114  			defer close(done)
   115  
   116  			n, err = io.CopyBuffer(w, r, c.buffer)
   117  
   118  			if n < buflen && err == nil {
   119  				err = io.EOF
   120  			}
   121  		}()
   122  
   123  		select {
   124  		case <-done:
   125  
   126  		case <-t.C:
   127  			return written, ErrWatchdogExpired
   128  
   129  		case <-ctx.Done():
   130  			return written, ctx.Err()
   131  		}
   132  
   133  		// n and err are valid here because <-done HAPPENS AFTER close(done)
   134  		written += n
   135  		bwAccum += n
   136  		if err != nil {
   137  			break
   138  		}
   139  
   140  		if keepingMetrics {
   141  			if now := time.Now(); now.After(next) {
   142  				if total != nil {
   143  					dur := now.Sub(start)
   144  					total(float64(written) * c.bwScale / dur.Seconds())
   145  				}
   146  
   147  				if running != nil {
   148  					dur := now.Sub(last)
   149  
   150  					copy(bwWindow, bwWindow[1:])
   151  					bwWindow[len(bwWindow)-1].n = bwAccum
   152  					bwWindow[len(bwWindow)-1].d = dur
   153  
   154  					var n int64
   155  					var d time.Duration
   156  					for i := range bwWindow {
   157  						n += bwWindow[i].n
   158  						d += bwWindow[i].d
   159  					}
   160  
   161  					running(float64(n) * c.bwScale / d.Seconds())
   162  				}
   163  
   164  				bwAccum = 0
   165  				last = now
   166  				next = last.Add(time.Second)
   167  			}
   168  		}
   169  	}
   170  
   171  	if err == io.EOF {
   172  		return written, nil
   173  	}
   174  
   175  	return written, err
   176  }