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

     1  package gorill
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  	"time"
     8  )
     9  
    10  // ErrTimeout error is returned whenever a Write operation exceeds the preset timeout period. Even
    11  // after a timeout takes place, the write may still independantly complete.
    12  type ErrTimeout time.Duration
    13  
    14  // Error returns a string representing the ErrTimeout.
    15  func (e ErrTimeout) Error() string {
    16  	return fmt.Sprintf("timeout after %s", time.Duration(e))
    17  }
    18  
    19  // TimedWriteCloser is an io.Writer that enforces a preset timeout period on every Write operation.
    20  type TimedWriteCloser struct {
    21  	halted   bool
    22  	iowc     io.WriteCloser
    23  	jobs     chan *rillJob
    24  	jobsDone sync.WaitGroup
    25  	lock     sync.RWMutex
    26  	timeout  time.Duration
    27  }
    28  
    29  // NewTimedWriteCloser returns a TimedWriteCloser that enforces a preset timeout period on every Write
    30  // operation.  It panics when timeout is less than or equal to 0.
    31  func NewTimedWriteCloser(iowc io.WriteCloser, timeout time.Duration) *TimedWriteCloser {
    32  	if timeout <= 0 {
    33  		panic(fmt.Errorf("timeout must be greater than 0: %s", timeout))
    34  	}
    35  	wc := &TimedWriteCloser{
    36  		iowc:    iowc,
    37  		jobs:    make(chan *rillJob, 1),
    38  		timeout: timeout,
    39  	}
    40  	wc.jobsDone.Add(1)
    41  	go func() {
    42  		for job := range wc.jobs {
    43  			n, err := wc.iowc.Write(job.data)
    44  			job.results <- rillResult{n, err}
    45  		}
    46  		wc.jobsDone.Done()
    47  	}()
    48  	return wc
    49  }
    50  
    51  // Write writes data to the underlying io.Writer, but returns ErrTimeout if the Write
    52  // operation exceeds a preset timeout duration.  Even after a timeout takes place, the write may
    53  // still independantly complete as writes are queued from a different go routine.
    54  func (wc *TimedWriteCloser) Write(data []byte) (int, error) {
    55  	wc.lock.RLock()
    56  	defer wc.lock.RUnlock()
    57  
    58  	if wc.halted {
    59  		return 0, ErrWriteAfterClose{}
    60  	}
    61  
    62  	job := newRillJob(_write, data)
    63  	wc.jobs <- job
    64  
    65  	// wait for result or timeout
    66  	select {
    67  	case result := <-job.results:
    68  		return result.n, result.err
    69  	case <-time.After(wc.timeout):
    70  		return 0, ErrTimeout(wc.timeout)
    71  	}
    72  }
    73  
    74  // Close frees resources when a SpooledWriteCloser is no longer needed.
    75  func (wc *TimedWriteCloser) Close() error {
    76  	wc.lock.Lock()
    77  	defer wc.lock.Unlock()
    78  
    79  	close(wc.jobs)
    80  	wc.jobsDone.Wait()
    81  	wc.halted = true
    82  	return wc.iowc.Close()
    83  }