86 lines
2 KiB
Go
86 lines
2 KiB
Go
|
package gbytes
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
// ErrTimeout is returned by TimeoutCloser, TimeoutReader, and TimeoutWriter when the underlying Closer/Reader/Writer does not return within the specified timeout
|
||
|
|
var ErrTimeout = errors.New("timeout occurred")
|
||
|
|
|
||
|
|
// TimeoutCloser returns an io.Closer that wraps the passed-in io.Closer. If the underlying Closer fails to close within the alloted timeout ErrTimeout is returned.
|
||
|
|
func TimeoutCloser(c io.Closer, timeout time.Duration) io.Closer {
|
||
|
|
return timeoutReaderWriterCloser{c: c, d: timeout}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TimeoutReader returns an io.Reader that wraps the passed-in io.Reader. If the underlying Reader fails to read within the alloted timeout ErrTimeout is returned.
|
||
|
|
func TimeoutReader(r io.Reader, timeout time.Duration) io.Reader {
|
||
|
|
return timeoutReaderWriterCloser{r: r, d: timeout}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TimeoutWriter returns an io.Writer that wraps the passed-in io.Writer. If the underlying Writer fails to write within the alloted timeout ErrTimeout is returned.
|
||
|
|
func TimeoutWriter(w io.Writer, timeout time.Duration) io.Writer {
|
||
|
|
return timeoutReaderWriterCloser{w: w, d: timeout}
|
||
|
|
}
|
||
|
|
|
||
|
|
type timeoutReaderWriterCloser struct {
|
||
|
|
c io.Closer
|
||
|
|
w io.Writer
|
||
|
|
r io.Reader
|
||
|
|
d time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
func (t timeoutReaderWriterCloser) Close() error {
|
||
|
|
done := make(chan struct{})
|
||
|
|
var err error
|
||
|
|
|
||
|
|
go func() {
|
||
|
|
err = t.c.Close()
|
||
|
|
close(done)
|
||
|
|
}()
|
||
|
|
|
||
|
|
select {
|
||
|
|
case <-done:
|
||
|
|
return err
|
||
|
|
case <-time.After(t.d):
|
||
|
|
return ErrTimeout
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (t timeoutReaderWriterCloser) Read(p []byte) (int, error) {
|
||
|
|
done := make(chan struct{})
|
||
|
|
var n int
|
||
|
|
var err error
|
||
|
|
|
||
|
|
go func() {
|
||
|
|
n, err = t.r.Read(p)
|
||
|
|
close(done)
|
||
|
|
}()
|
||
|
|
|
||
|
|
select {
|
||
|
|
case <-done:
|
||
|
|
return n, err
|
||
|
|
case <-time.After(t.d):
|
||
|
|
return 0, ErrTimeout
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (t timeoutReaderWriterCloser) Write(p []byte) (int, error) {
|
||
|
|
done := make(chan struct{})
|
||
|
|
var n int
|
||
|
|
var err error
|
||
|
|
|
||
|
|
go func() {
|
||
|
|
n, err = t.w.Write(p)
|
||
|
|
close(done)
|
||
|
|
}()
|
||
|
|
|
||
|
|
select {
|
||
|
|
case <-done:
|
||
|
|
return n, err
|
||
|
|
case <-time.After(t.d):
|
||
|
|
return 0, ErrTimeout
|
||
|
|
}
|
||
|
|
}
|