github.com/sandwich-go/boost@v1.3.29/xio/xio.go (about)

     1  package xio
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"time"
     7  )
     8  
     9  type ioRet struct {
    10  	n   int
    11  	err error
    12  }
    13  
    14  type ctxReader struct {
    15  	r   io.Reader
    16  	ctx context.Context
    17  }
    18  
    19  // NewReader 返回context敏感的io reader
    20  func NewReader(ctx context.Context, r io.Reader) io.Reader {
    21  	return &ctxReader{ctx: ctx, r: r}
    22  }
    23  
    24  // block for debug
    25  var block = false
    26  
    27  // Read 由于Read是block操作，内部为每一次Read启动了独立协程协助读取,如果超时，则会返回(0，ctx.Error)
    28  func (r *ctxReader) Read(buf []byte) (int, error) {
    29  	bufReading := make([]byte, len(buf))
    30  
    31  	c := make(chan ioRet, 1)
    32  
    33  	go func() {
    34  		if block {
    35  			time.Sleep(10 * time.Second)
    36  		}
    37  		n, err := r.r.Read(bufReading)
    38  		c <- ioRet{n: n, err: err}
    39  		close(c)
    40  	}()
    41  
    42  	select {
    43  	case ret := <-c:
    44  		copy(buf, bufReading)
    45  		return ret.n, ret.err
    46  	case <-r.ctx.Done():
    47  		return 0, r.ctx.Err()
    48  	}
    49  }