github.com/peterdeka/grab@v2.0.0+incompatible/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  }