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

     1  package gorill
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  	"testing"
     7  )
     8  
     9  // channelWriter provided to benchmark against LockingWriter and TimedWriteCloser.
    10  type channelWriter struct {
    11  	halted   bool
    12  	jobsDone sync.WaitGroup
    13  	iowc     io.WriteCloser
    14  	jobs     chan rillJob
    15  	lock     sync.RWMutex
    16  }
    17  
    18  func newChannelWriter(iowc io.WriteCloser) *channelWriter {
    19  	w := &channelWriter{
    20  		iowc: iowc,
    21  		jobs: make(chan rillJob, 1),
    22  	}
    23  	go func(w *channelWriter) {
    24  		w.jobsDone.Add(1)
    25  		for job := range w.jobs {
    26  			n, err := w.iowc.Write(job.data)
    27  			job.results <- rillResult{n, err}
    28  		}
    29  		w.jobsDone.Done()
    30  	}(w)
    31  	return w
    32  }
    33  
    34  func (w *channelWriter) Write(data []byte) (int, error) {
    35  	w.lock.RLock()
    36  	defer w.lock.RUnlock()
    37  
    38  	if w.halted {
    39  		return 0, ErrWriteAfterClose{}
    40  	}
    41  
    42  	job := rillJob{data: data, results: make(chan rillResult, 1)}
    43  	w.jobs <- job
    44  	// wait for result
    45  	result := <-job.results
    46  	return result.n, result.err
    47  }
    48  
    49  func (w *channelWriter) Close() error {
    50  	w.lock.Lock()
    51  	defer w.lock.Unlock()
    52  
    53  	close(w.jobs)
    54  	w.jobsDone.Wait()
    55  	w.halted = true
    56  	return w.iowc.Close()
    57  }
    58  
    59  func BenchmarkWriterChannelWriter(b *testing.B) {
    60  	consumers := make([]io.WriteCloser, consumerCount)
    61  	for i := 0; i < len(consumers); i++ {
    62  		consumers[i] = newChannelWriter(NewNopCloseBuffer())
    63  	}
    64  	benchmarkWriter(b, b.N, consumers)
    65  }