github.com/geph-official/geph2@v0.22.6-0.20210211030601-f527cb59b0df/libs/cwl/limiter.go (about)

     1  package cwl
     2  
     3  import (
     4  	"context"
     5  	. "io"
     6  	"net"
     7  	"time"
     8  
     9  	"golang.org/x/time/rate"
    10  )
    11  
    12  // CopyWithLimit is like io.Copy but subject to a rate limit and calling a callback.
    13  func CopyWithLimit(dst Writer, src net.Conn, limiter *rate.Limiter, callback func(int), idleTimeout time.Duration) (n int, err error) {
    14  	var buf []byte
    15  	if buf == nil {
    16  		size := 16 * 1024
    17  		buf = make([]byte, size)
    18  	}
    19  	for {
    20  		if idleTimeout > 0 {
    21  			src.SetReadDeadline(time.Now().Add(idleTimeout))
    22  		}
    23  		nr, er := src.Read(buf)
    24  		if nr > 0 {
    25  			if callback != nil {
    26  				callback(nr)
    27  			}
    28  			if limiter != nil {
    29  				limiter.WaitN(context.Background(), nr)
    30  			}
    31  			nw, ew := dst.Write(buf[0:nr])
    32  			if nw > 0 {
    33  				n += nw
    34  			}
    35  			if ew != nil {
    36  				err = ew
    37  				break
    38  			}
    39  			if nr != nw {
    40  				err = ErrShortWrite
    41  				break
    42  			}
    43  		}
    44  		if er != nil {
    45  			if er != EOF {
    46  				err = er
    47  			}
    48  			break
    49  		}
    50  	}
    51  	return
    52  }