github.com/decred/dcrlnd@v0.7.6/pool/write.go (about)

     1  package pool
     2  
     3  import (
     4  	"bytes"
     5  	"time"
     6  
     7  	"github.com/decred/dcrlnd/buffer"
     8  )
     9  
    10  // Write is a worker pool specifically designed for sharing access to
    11  // buffer.Write objects amongst a set of worker goroutines. This enables an
    12  // application to limit the total number of buffer.Write objects allocated at
    13  // any given time.
    14  type Write struct {
    15  	workerPool *Worker
    16  	bufferPool *WriteBuffer
    17  }
    18  
    19  // NewWrite creates a Write pool, using an underlying Writebuffer pool to
    20  // recycle buffer.Write objects across the lifetime of the Write pool's
    21  // workers.
    22  func NewWrite(writeBufferPool *WriteBuffer, numWorkers int,
    23  	workerTimeout time.Duration) *Write {
    24  
    25  	w := &Write{
    26  		bufferPool: writeBufferPool,
    27  	}
    28  	w.workerPool = NewWorker(&WorkerConfig{
    29  		NewWorkerState: w.newWorkerState,
    30  		NumWorkers:     numWorkers,
    31  		WorkerTimeout:  workerTimeout,
    32  	})
    33  
    34  	return w
    35  }
    36  
    37  // Start safely spins up the Write pool.
    38  func (w *Write) Start() error {
    39  	return w.workerPool.Start()
    40  }
    41  
    42  // Stop safely shuts down the Write pool.
    43  func (w *Write) Stop() error {
    44  	return w.workerPool.Stop()
    45  }
    46  
    47  // Submit accepts a function closure that provides access to a fresh
    48  // bytes.Buffer backed by a buffer.Write object. The function's execution will
    49  // be allocated to one of the underlying Worker pool's goroutines.
    50  func (w *Write) Submit(inner func(*bytes.Buffer) error) error {
    51  	return w.workerPool.Submit(func(s WorkerState) error {
    52  		state := s.(*writeWorkerState)
    53  		return inner(state.buf)
    54  	})
    55  }
    56  
    57  // writeWorkerState is the per-goroutine state maintained by a Write pool's
    58  // goroutines.
    59  type writeWorkerState struct {
    60  	// bufferPool is the pool to which the writeBuf will be returned when
    61  	// the goroutine exits.
    62  	bufferPool *WriteBuffer
    63  
    64  	// writeBuf is the buffer taken from the bufferPool on initialization,
    65  	// which will be used to back the buf object provided to any tasks that
    66  	// the goroutine processes before exiting.
    67  	writeBuf *buffer.Write
    68  
    69  	// buf is a buffer backed by writeBuf, that can be written to by tasks
    70  	// submitted to the Write pool. The buf will be reset between each task
    71  	// processed by a goroutine before exiting, and allows the task
    72  	// submitters to interact with the writeBuf as if it were an io.Writer.
    73  	buf *bytes.Buffer
    74  }
    75  
    76  // newWorkerState initializes a new writeWorkerState, which will be called
    77  // whenever a new goroutine is allocated to begin processing write tasks.
    78  func (w *Write) newWorkerState() WorkerState {
    79  	writeBuf := w.bufferPool.Take()
    80  
    81  	return &writeWorkerState{
    82  		bufferPool: w.bufferPool,
    83  		writeBuf:   writeBuf,
    84  		buf:        bytes.NewBuffer(writeBuf[0:0:len(writeBuf)]),
    85  	}
    86  }
    87  
    88  // Cleanup returns the writeBuf to the underlying buffer pool, and removes the
    89  // goroutine's reference to the writeBuf and encapsulating buf.
    90  func (w *writeWorkerState) Cleanup() {
    91  	w.bufferPool.Return(w.writeBuf)
    92  	w.writeBuf = nil
    93  	w.buf = nil
    94  }
    95  
    96  // Reset resets the bytes.Buffer so that it is zero-length and has the capacity
    97  // of the underlying buffer.Write.k
    98  func (w *writeWorkerState) Reset() {
    99  	w.buf.Reset()
   100  }