github.com/onsi/gomega@v1.32.0/gbytes/io_wrappers.go (about)

     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  }