github.com/karrick/gorill@v1.10.3/timedReadCloser.go (about)

     1  package gorill
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  	"time"
     8  )
     9  
    10  // TimedReadCloser is an io.Reader that enforces a preset timeout period on every Read operation.
    11  type TimedReadCloser struct {
    12  	halted   bool
    13  	iorc     io.ReadCloser
    14  	jobs     chan *rillJob
    15  	jobsDone sync.WaitGroup
    16  	lock     sync.RWMutex
    17  	timeout  time.Duration
    18  }
    19  
    20  // NewTimedReadCloser returns a TimedReadCloser that enforces a preset timeout period on every Read
    21  // operation.  It panics when timeout is less than or equal to 0.
    22  func NewTimedReadCloser(iowc io.ReadCloser, timeout time.Duration) *TimedReadCloser {
    23  	if timeout <= 0 {
    24  		panic(fmt.Errorf("timeout must be greater than 0: %s", timeout))
    25  	}
    26  	rc := &TimedReadCloser{
    27  		iorc:    iowc,
    28  		jobs:    make(chan *rillJob, 1),
    29  		timeout: timeout,
    30  	}
    31  	rc.jobsDone.Add(1)
    32  	go func() {
    33  		for job := range rc.jobs {
    34  			n, err := rc.iorc.Read(job.data)
    35  			job.results <- rillResult{n, err}
    36  		}
    37  		rc.jobsDone.Done()
    38  	}()
    39  	return rc
    40  }
    41  
    42  // Read reads data to the underlying io.Reader, but returns ErrTimeout if the Read operation exceeds
    43  // a preset timeout duration.
    44  //
    45  // Even after a timeout takes place, the read may still independently complete as reads are queued
    46  // from a different go-routine.  Race condition for the data slice is prevented by reading into a
    47  // temporary byte slice, and copying the results to the client's slice when the actual read returns.
    48  func (rc *TimedReadCloser) Read(data []byte) (int, error) {
    49  	rc.lock.RLock()
    50  	defer rc.lock.RUnlock()
    51  
    52  	if rc.halted {
    53  		return 0, ErrReadAfterClose{}
    54  	}
    55  
    56  	job := newRillJob(_read, make([]byte, len(data)))
    57  	rc.jobs <- job
    58  
    59  	// wait for result or timeout
    60  	select {
    61  	case result := <-job.results:
    62  		copy(data, job.data)
    63  		return result.n, result.err
    64  	case <-time.After(rc.timeout):
    65  		return 0, ErrTimeout(rc.timeout)
    66  	}
    67  }
    68  
    69  // Close frees resources when a SpooledReadCloser is no longer needed.
    70  func (rc *TimedReadCloser) Close() error {
    71  	rc.lock.Lock()
    72  	defer rc.lock.Unlock()
    73  
    74  	close(rc.jobs)
    75  	rc.jobsDone.Wait()
    76  	rc.halted = true
    77  	return rc.iorc.Close()
    78  }