diff --git a/ioutil/cancel_reader.go b/ioutil/cancel_reader.go index 4c1b661..4825b70 100644 --- a/ioutil/cancel_reader.go +++ b/ioutil/cancel_reader.go @@ -21,19 +21,25 @@ type CancelableReader struct { func (c *CancelableReader) begin() { for { - sizewant := <-c.sizeWant - buf := make([]byte, sizewant) - // readFull, cause some readers like net.TCPConn returns size smaller than buf size - n, err := io.ReadFull(c.r, buf) - if err != nil { - c.err = err + select { + case sizeWant := <-c.sizeWant: + buf := make([]byte, sizeWant) + // readFull, cause some readers like net.TCPConn returns size smaller than buf size + n, err := io.ReadFull(c.r, buf) + if err != nil { + c.err = err + close(c.data) + return + } + if n != sizeWant { + panic("read " + strconv.Itoa(n) + ", want " + strconv.Itoa(sizeWant)) + } + c.data <- buf + case <-c.ctx.Done(): close(c.data) + close(c.sizeWant) return } - if n != sizewant { - panic("read " + strconv.Itoa(n) + ", want " + strconv.Itoa(sizewant)) - } - c.data <- buf } }