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 }