gopkg.in/cavaliercoder/grab.v2@v2.0.0/transfer.go (about) 1 package grab 2 3 import ( 4 "context" 5 "io" 6 "sync/atomic" 7 ) 8 9 type transfer struct { 10 n int64 // must be 64bit aligned on 386 11 ctx context.Context 12 lim RateLimiter 13 w io.Writer 14 r io.Reader 15 b []byte 16 } 17 18 func newTransfer(ctx context.Context, lim RateLimiter, dst io.Writer, src io.Reader, buf []byte) *transfer { 19 return &transfer{ 20 ctx: ctx, 21 lim: lim, 22 w: dst, 23 r: src, 24 b: buf, 25 } 26 } 27 28 // copy behaves similarly to io.CopyBuffer except that it checks for cancelation 29 // of the given context.Context and reports progress in a thread-safe manner. 30 func (c *transfer) copy() (written int64, err error) { 31 if c.b == nil { 32 c.b = make([]byte, 32*1024) 33 } 34 for { 35 select { 36 case <-c.ctx.Done(): 37 err = c.ctx.Err() 38 return 39 default: 40 // keep working 41 } 42 if c.lim != nil { 43 err = c.lim.WaitN(c.ctx, len(c.b)) 44 if err != nil { 45 return 46 } 47 } 48 nr, er := c.r.Read(c.b) 49 if nr > 0 { 50 nw, ew := c.w.Write(c.b[0:nr]) 51 if nw > 0 { 52 written += int64(nw) 53 atomic.StoreInt64(&c.n, written) 54 } 55 if ew != nil { 56 err = ew 57 break 58 } 59 if nr != nw { 60 err = io.ErrShortWrite 61 break 62 } 63 } 64 if er != nil { 65 if er != io.EOF { 66 err = er 67 } 68 break 69 } 70 } 71 return written, err 72 } 73 74 // N returns the number of bytes transferred. 75 func (c *transfer) N() (n int64) { 76 if c == nil { 77 return 0 78 } 79 n = atomic.LoadInt64(&c.n) 80 return 81 }