github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/sstable/write_queue.go (about)

     1  package sstable
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/zuoyebang/bitalostable/internal/base"
     7  )
     8  
     9  type writeTask struct {
    10  	// Since writeTasks are pooled, the compressionDone channel will be re-used.
    11  	// It is necessary that any writes to the channel have already been read,
    12  	// before adding the writeTask back to the pool.
    13  	compressionDone chan bool
    14  	buf             *dataBlockBuf
    15  	// If this is not nil, then this index block will be flushed.
    16  	flushableIndexBlock *indexBlockBuf
    17  	// currIndexBlock is the index block on which indexBlock.add must be called.
    18  	currIndexBlock *indexBlockBuf
    19  	indexEntrySep  InternalKey
    20  	// inflightSize is used to decrement Writer.coordination.sizeEstimate.inflightSize.
    21  	inflightSize int
    22  	// inflightIndexEntrySize is used to decrement Writer.indexBlock.sizeEstimate.inflightSize.
    23  	indexInflightSize int
    24  	// If the index block is finished, then we set the finishedIndexProps here.
    25  	finishedIndexProps []byte
    26  }
    27  
    28  // It is not the responsibility of the writeTask to clear the
    29  // task.flushableIndexBlock, and task.buf.
    30  func (task *writeTask) clear() {
    31  	*task = writeTask{
    32  		indexEntrySep:   base.InvalidInternalKey,
    33  		compressionDone: task.compressionDone,
    34  	}
    35  }
    36  
    37  // Note that only the Writer client goroutine will be adding tasks to the writeQueue.
    38  // Both the Writer client and the compression goroutines will be able to write to
    39  // writeTask.compressionDone to indicate that the compression job associated with
    40  // a writeTask has finished.
    41  type writeQueue struct {
    42  	tasks  chan *writeTask
    43  	wg     sync.WaitGroup
    44  	writer *Writer
    45  
    46  	// err represents an error which is encountered when the write queue attempts
    47  	// to write a block to disk. The error is stored here to skip unnecessary block
    48  	// writes once the first error is encountered.
    49  	err    error
    50  	closed bool
    51  }
    52  
    53  func newWriteQueue(size int, writer *Writer) *writeQueue {
    54  	w := &writeQueue{}
    55  	w.tasks = make(chan *writeTask, size)
    56  	w.writer = writer
    57  
    58  	w.wg.Add(1)
    59  	go w.runWorker()
    60  	return w
    61  }
    62  
    63  func (w *writeQueue) performWrite(task *writeTask) error {
    64  	var bh BlockHandle
    65  	var bhp BlockHandleWithProperties
    66  
    67  	var err error
    68  	if bh, err = w.writer.writeCompressedBlock(task.buf.compressed, task.buf.tmp[:]); err != nil {
    69  		return err
    70  	}
    71  
    72  	// Update the size estimates after writing the data block to disk.
    73  	w.writer.coordination.sizeEstimate.dataBlockWritten(
    74  		w.writer.meta.Size, task.inflightSize, int(bh.Length),
    75  	)
    76  
    77  	bhp = BlockHandleWithProperties{BlockHandle: bh, Props: task.buf.dataBlockProps}
    78  	if err = w.writer.addIndexEntry(
    79  		task.indexEntrySep, bhp, task.buf.tmp[:], task.flushableIndexBlock, task.currIndexBlock,
    80  		task.indexInflightSize, task.finishedIndexProps); err != nil {
    81  		return err
    82  	}
    83  
    84  	return nil
    85  }
    86  
    87  // It is necessary to ensure that none of the buffers in the writeTask,
    88  // dataBlockBuf, indexBlockBuf, are pointed to by another struct.
    89  func (w *writeQueue) releaseBuffers(task *writeTask) {
    90  	task.buf.clear()
    91  	dataBlockBufPool.Put(task.buf)
    92  
    93  	// This index block is no longer used by the Writer, so we can add it back
    94  	// to the pool.
    95  	if task.flushableIndexBlock != nil {
    96  		task.flushableIndexBlock.clear()
    97  		indexBlockBufPool.Put(task.flushableIndexBlock)
    98  	}
    99  
   100  	task.clear()
   101  	writeTaskPool.Put(task)
   102  }
   103  
   104  func (w *writeQueue) runWorker() {
   105  	for task := range w.tasks {
   106  		<-task.compressionDone
   107  
   108  		if w.err == nil {
   109  			w.err = w.performWrite(task)
   110  		}
   111  
   112  		w.releaseBuffers(task)
   113  	}
   114  	w.wg.Done()
   115  }
   116  
   117  func (w *writeQueue) add(task *writeTask) {
   118  	w.tasks <- task
   119  }
   120  
   121  // addSync will perform the writeTask synchronously with the caller goroutine. Calls to addSync
   122  // are no longer valid once writeQueue.add has been called at least once.
   123  func (w *writeQueue) addSync(task *writeTask) error {
   124  	// This should instantly return without blocking.
   125  	<-task.compressionDone
   126  
   127  	if w.err == nil {
   128  		w.err = w.performWrite(task)
   129  	}
   130  
   131  	w.releaseBuffers(task)
   132  
   133  	return w.err
   134  }
   135  
   136  // finish should only be called once no more tasks will be added to the writeQueue.
   137  // finish will return any error which was encountered while tasks were processed.
   138  func (w *writeQueue) finish() error {
   139  	if w.closed {
   140  		return w.err
   141  	}
   142  
   143  	close(w.tasks)
   144  	w.wg.Wait()
   145  	w.closed = true
   146  	return w.err
   147  }