Initial commit

vendor/github.com/cockroachdb/pebble/sstable/write_queue.go

@@ -0,0 +1,148 @@
+package sstable
+
+import (
+	"sync"
+
+	"github.com/cockroachdb/pebble/internal/base"
+)
+
+type writeTask struct {
+	// Since writeTasks are pooled, the compressionDone channel will be re-used.
+	// It is necessary that any writes to the channel have already been read,
+	// before adding the writeTask back to the pool.
+	compressionDone chan bool
+	buf             *dataBlockBuf
+	// If this is not nil, then this index block will be flushed.
+	flushableIndexBlock *indexBlockBuf
+	// currIndexBlock is the index block on which indexBlock.add must be called.
+	currIndexBlock *indexBlockBuf
+	indexEntrySep  InternalKey
+	// inflightSize is used to decrement Writer.coordination.sizeEstimate.inflightSize.
+	inflightSize int
+	// inflightIndexEntrySize is used to decrement Writer.indexBlock.sizeEstimate.inflightSize.
+	indexInflightSize int
+	// If the index block is finished, then we set the finishedIndexProps here.
+	finishedIndexProps []byte
+}
+
+// It is not the responsibility of the writeTask to clear the
+// task.flushableIndexBlock, and task.buf.
+func (task *writeTask) clear() {
+	*task = writeTask{
+		indexEntrySep:   base.InvalidInternalKey,
+		compressionDone: task.compressionDone,
+	}
+}
+
+// Note that only the Writer client goroutine will be adding tasks to the writeQueue.
+// Both the Writer client and the compression goroutines will be able to write to
+// writeTask.compressionDone to indicate that the compression job associated with
+// a writeTask has finished.
+type writeQueue struct {
+	tasks  chan *writeTask
+	wg     sync.WaitGroup
+	writer *Writer
+
+	// err represents an error which is encountered when the write queue attempts
+	// to write a block to disk. The error is stored here to skip unnecessary block
+	// writes once the first error is encountered.
+	err    error
+	closed bool
+}
+
+func newWriteQueue(size int, writer *Writer) *writeQueue {
+	w := &writeQueue{}
+	w.tasks = make(chan *writeTask, size)
+	w.writer = writer
+
+	w.wg.Add(1)
+	go w.runWorker()
+	return w
+}
+
+func (w *writeQueue) performWrite(task *writeTask) error {
+	var bh BlockHandle
+	var bhp BlockHandleWithProperties
+
+	var err error
+	if bh, err = w.writer.writeCompressedBlock(task.buf.compressed, task.buf.tmp[:]); err != nil {
+		return err
+	}
+
+	// Update the size estimates after writing the data block to disk.
+	w.writer.coordination.sizeEstimate.dataBlockWritten(
+		w.writer.meta.Size, task.inflightSize, int(bh.Length),
+	)
+
+	bhp = BlockHandleWithProperties{BlockHandle: bh, Props: task.buf.dataBlockProps}
+	if err = w.writer.addIndexEntry(
+		task.indexEntrySep, bhp, task.buf.tmp[:], task.flushableIndexBlock, task.currIndexBlock,
+		task.indexInflightSize, task.finishedIndexProps); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// It is necessary to ensure that none of the buffers in the writeTask,
+// dataBlockBuf, indexBlockBuf, are pointed to by another struct.
+func (w *writeQueue) releaseBuffers(task *writeTask) {
+	task.buf.clear()
+	dataBlockBufPool.Put(task.buf)
+
+	// This index block is no longer used by the Writer, so we can add it back
+	// to the pool.
+	if task.flushableIndexBlock != nil {
+		task.flushableIndexBlock.clear()
+		indexBlockBufPool.Put(task.flushableIndexBlock)
+	}
+
+	task.clear()
+	writeTaskPool.Put(task)
+}
+
+func (w *writeQueue) runWorker() {
+	for task := range w.tasks {
+		<-task.compressionDone
+
+		if w.err == nil {
+			w.err = w.performWrite(task)
+		}
+
+		w.releaseBuffers(task)
+	}
+	w.wg.Done()
+}
+
+//lint:ignore U1000 - Will be used in a future pr.
+func (w *writeQueue) add(task *writeTask) {
+	w.tasks <- task
+}
+
+// addSync will perform the writeTask synchronously with the caller goroutine. Calls to addSync
+// are no longer valid once writeQueue.add has been called at least once.
+func (w *writeQueue) addSync(task *writeTask) error {
+	// This should instantly return without blocking.
+	<-task.compressionDone
+
+	if w.err == nil {
+		w.err = w.performWrite(task)
+	}
+
+	w.releaseBuffers(task)
+
+	return w.err
+}
+
+// finish should only be called once no more tasks will be added to the writeQueue.
+// finish will return any error which was encountered while tasks were processed.
+func (w *writeQueue) finish() error {
+	if w.closed {
+		return w.err
+	}
+
+	close(w.tasks)
+	w.wg.Wait()
+	w.closed = true
+	return w.err
+}