1
package gbytes
2

3
import (
4
	"errors"
5
	"io"
6
	"time"
7
)
8

9
// ErrTimeout is returned by TimeoutCloser, TimeoutReader, and TimeoutWriter when the underlying Closer/Reader/Writer does not return within the specified timeout
10
var ErrTimeout = errors.New("timeout occurred")
11

12
// TimeoutCloser returns an io.Closer that wraps the passed-in io.Closer.  If the underlying Closer fails to close within the allotted timeout ErrTimeout is returned.
13
func TimeoutCloser(c io.Closer, timeout time.Duration) io.Closer {
14
	return timeoutReaderWriterCloser{c: c, d: timeout}
15
}
16

17
// TimeoutReader returns an io.Reader that wraps the passed-in io.Reader.  If the underlying Reader fails to read within the allotted timeout ErrTimeout is returned.
18
func TimeoutReader(r io.Reader, timeout time.Duration) io.Reader {
19
	return timeoutReaderWriterCloser{r: r, d: timeout}
20
}
21

22
// TimeoutWriter returns an io.Writer that wraps the passed-in io.Writer.  If the underlying Writer fails to write within the allotted timeout ErrTimeout is returned.
23
func TimeoutWriter(w io.Writer, timeout time.Duration) io.Writer {
24
	return timeoutReaderWriterCloser{w: w, d: timeout}
25
}
26

27
type timeoutReaderWriterCloser struct {
28
	c io.Closer
29
	w io.Writer
30
	r io.Reader
31
	d time.Duration
32
}
33

34
func (t timeoutReaderWriterCloser) Close() error {
35
	done := make(chan struct{})
36
	var err error
37

38
	go func() {
39
		err = t.c.Close()
40
		close(done)
41
	}()
42

43
	select {
44
	case <-done:
45
		return err
46
	case <-time.After(t.d):
47
		return ErrTimeout
48
	}
49
}
50

51
func (t timeoutReaderWriterCloser) Read(p []byte) (int, error) {
52
	done := make(chan struct{})
53
	var n int
54
	var err error
55

56
	go func() {
57
		n, err = t.r.Read(p)
58
		close(done)
59
	}()
60

61
	select {
62
	case <-done:
63
		return n, err
64
	case <-time.After(t.d):
65
		return 0, ErrTimeout
66
	}
67
}
68

69
func (t timeoutReaderWriterCloser) Write(p []byte) (int, error) {
70
	done := make(chan struct{})
71
	var n int
72
	var err error
73

74
	go func() {
75
		n, err = t.w.Write(p)
76
		close(done)
77
	}()
78

79
	select {
80
	case <-done:
81
		return n, err
82
	case <-time.After(t.d):
83
		return 0, ErrTimeout
84
	}
85
}
86

87