sstable: free physical buffers on write

This change moves the responsibility of releasing the `PhysicalBlock`s
to the code that writes out the block. This makes the higher level
code simpler and more robust.

To avoid misuse, we add an `OwnedPhysicalBlock` type which is used to
indicate that once passed, the function takes responsibility for
releasing the block. The only way to create an `OwnedPhysicalBlock`
clears the `PhysicalBlock` so it can't be incorrectly used again.

Author: RaduBerinde

sstable/blob/blob.go

@@ -229,7 +229,7 @@ func (w *FileWriter) drainWriteQueue() {
 	// Call once to initialize the CPU measurer.
 	w.cpuMeasurer.MeasureCPU(base.CompactionGoroutineBlobFileSecondary)
 	for cb := range w.writeQueue.ch {
-		_, err := cb.pb.WriteTo(w.w)
+		length, err := block.WriteAndReleasePhysicalBlock(cb.pb.Take(), w.w)
 		// Report to the CPU measurer immediately after writing (note that there
 		// may be a time lag until the next block is available to write).
 		w.cpuMeasurer.MeasureCPU(base.CompactionGoroutineBlobFileSecondary)
@@ -239,10 +239,8 @@ func (w *FileWriter) drainWriteQueue() {
 		}
 		w.indexEncoder.AddBlockHandle(block.Handle{
 			Offset: cb.off,
-			Length: uint64(cb.pb.LengthWithoutTrailer()),
+			Length: uint64(length.WithoutTrailer()),
 		})
-		// We're done with this physical block.
-		cb.pb.Release()
 	}
 }
 
@@ -281,16 +279,17 @@ func (w *FileWriter) Close() (FileWriterStats, error) {
 	{
 		indexBlock := w.indexEncoder.Finish()
 		pb := w.physBlockMaker.Make(indexBlock, blockkind.Metadata, block.DontCompress)
-		defer pb.Release()
-		if _, w.err = pb.WriteTo(w.w); w.err != nil {
+		length, err := block.WriteAndReleasePhysicalBlock(pb.Take(), w.w)
+		if err != nil {
+			w.err = err
 			if w.w != nil {
 				w.w.Abort()
 			}
-			return FileWriterStats{}, w.err
+			return FileWriterStats{}, err
 		}
 		indexBlockHandle.Offset = stats.FileLen
-		indexBlockHandle.Length = uint64(pb.LengthWithoutTrailer())
-		stats.FileLen += uint64(pb.LengthWithTrailer())
+		indexBlockHandle.Length = uint64(length.WithoutTrailer())
+		stats.FileLen += uint64(length.WithTrailer())
 	}
 
 	// Write the footer.

sstable/block/physical.go

@@ -36,9 +36,22 @@ func MakeTrailer(blockType byte, checksum uint32) (t Trailer) {
 	return t
 }
 
-// LengthWithTrailer returns the length of the data block, including the trailer.
-func (b PhysicalBlock) LengthWithTrailer() int {
-	return b.tb.Size()
+// PhysicalBlockLength represents the length of a physical block.
+type PhysicalBlockLength struct {
+	lenWithTrailer int
+}
+
+func (l PhysicalBlockLength) WithTrailer() int {
+	return l.lenWithTrailer
+}
+
+func (l PhysicalBlockLength) WithoutTrailer() int {
+	return invariants.SafeSub(l.lenWithTrailer, TrailerLen)
+}
+
+// Length returns the length of the data block.
+func (b PhysicalBlock) Length() PhysicalBlockLength {
+	return PhysicalBlockLength{lenWithTrailer: b.tb.Size()}
 }
 
 // LengthWithoutTrailer returns the length of the data block, excluding the trailer.
@@ -48,8 +61,10 @@ func (b PhysicalBlock) LengthWithoutTrailer() int {
 
 // Release the underlying TempBuffer. The PhysicalBlock should not be used again.
 func (b *PhysicalBlock) Release() {
-	b.tb.Release()
-	b.tb = nil
+	if b.tb != nil {
+		b.tb.Release()
+		b.tb = nil
+	}
 }
 
 // AlreadyEncodedPhysicalBlock creates a PhysicalBlock from the provided on-disk
@@ -61,21 +76,37 @@ func AlreadyEncodedPhysicalBlock(dataWithTrailer []byte) PhysicalBlock {
 	return PhysicalBlock{tb: tb}
 }
 
-// WriteTo writes the block (including its trailer) to the provided Writable. If
-// err == nil, n is the number of bytes successfully written to the Writable.
-//
-// WriteTo might mangle the block data.
-func (b PhysicalBlock) WriteTo(w objstorage.Writable) (n int, err error) {
-	if err := w.Write(b.tb.Data()); err != nil {
-		return 0, err
-	}
+// OwnedPhysicalBlock is a wrapper around PhysicalBlock which indicates that
+// whoever is holding it is responsible for releasing the buffer.
+type OwnedPhysicalBlock struct {
+	pb PhysicalBlock
+}
 
-	// WriteTo is allowed to mangle the data. Mangle it ourselves some of the time
-	// in invariant builds to catch callers that don't handle this.
-	if invariants.Enabled && invariants.Sometimes(1) {
-		invariants.Mangle(b.tb.Data())
+// Length returns the length of the data block.
+func (b OwnedPhysicalBlock) Length() PhysicalBlockLength {
+	return PhysicalBlockLength{lenWithTrailer: b.pb.tb.Size()}
+}
+
+// Take moves the physical block to an OwnedPhysicalBlock. The receiver must not
+// be used again. This is used to make the ownership transfer more explicit.
+func (b *PhysicalBlock) Take() OwnedPhysicalBlock {
+	pb := *b
+	b.tb = nil
+	return OwnedPhysicalBlock{pb: pb}
+}
+
+// WriteAndReleasePhysicalBlock writes the block (including its trailer) to the
+// provided Writable and releases the block. The block is released even in error
+// cases.
+func WriteAndReleasePhysicalBlock(
+	b OwnedPhysicalBlock, w objstorage.Writable,
+) (length PhysicalBlockLength, err error) {
+	defer b.pb.Release()
+	// Note that Write can mangle the buffer, but we're releasing it anyway.
+	if err := w.Write(b.pb.tb.Data()); err != nil {
+		return PhysicalBlockLength{}, err
 	}
-	return b.tb.Size(), nil
+	return b.Length(), nil
 }
 
 // PhysicalBlockMaker is used to create physical blocks from logical block data.

sstable/colblk_writer.go

@@ -99,7 +99,7 @@ type RawColumnWriter struct {
 
 	writeQueue struct {
 		wg  sync.WaitGroup
-		ch  chan block.PhysicalBlock
+		ch  chan block.OwnedPhysicalBlock
 		err error
 	}
 	layout layoutWriter
@@ -189,7 +189,7 @@ func newColumnarWriter(
 	w.props.KeySchemaName = o.KeySchema.Name
 	w.props.MergerName = o.MergerName
 
-	w.writeQueue.ch = make(chan block.PhysicalBlock)
+	w.writeQueue.ch = make(chan block.OwnedPhysicalBlock)
 	w.writeQueue.wg.Add(1)
 	w.cpuMeasurer = cpuMeasurer
 	go w.drainWriteQueue()
@@ -710,15 +710,17 @@ func (w *RawColumnWriter) enqueueDataBlock(
 
 	// Compress and checksum the data block and send it to the write queue.
 	pb := w.layout.physBlockMaker.Make(serializedBlock, blockkind.SSTableData, block.NoFlags)
-	return w.enqueuePhysicalBlock(pb, separator)
+	return w.enqueuePhysicalBlock(pb.Take(), separator)
 }
 
 // enqueuePhysicalBlock enqueues a physical block to the write queue; the
 // physical block will be automatically released.
-func (w *RawColumnWriter) enqueuePhysicalBlock(pb block.PhysicalBlock, separator []byte) error {
+func (w *RawColumnWriter) enqueuePhysicalBlock(
+	pb block.OwnedPhysicalBlock, separator []byte,
+) error {
 	dataBlockHandle := block.Handle{
 		Offset: w.queuedDataSize,
-		Length: uint64(pb.LengthWithoutTrailer()),
+		Length: uint64(pb.Length().WithoutTrailer()),
 	}
 	w.queuedDataSize += dataBlockHandle.Length + block.TrailerLen
 	w.writeQueue.ch <- pb
@@ -881,7 +883,6 @@ func (w *RawColumnWriter) drainWriteQueue() {
 		if _, err := w.layout.WritePrecompressedDataBlock(pb); err != nil {
 			w.writeQueue.err = err
 		}
-		pb.Release()
 		// Report to the CPU measurer immediately after writing (note that there
 		// may be a time lag until the next block is available to write).
 		w.cpuMeasurer.MeasureCPU(base.CompactionGoroutineSSTableSecondary)
@@ -1125,7 +1126,7 @@ func (w *RawColumnWriter) rewriteSuffixes(
 			w.separatorBuf = w.comparer.Successor(w.separatorBuf[:0], blocks[i].end.UserKey)
 			separator = w.separatorBuf
 		}
-		if err := w.enqueuePhysicalBlock(blocks[i].physical, separator); err != nil {
+		if err := w.enqueuePhysicalBlock(blocks[i].physical.Take(), separator); err != nil {
 			return err
 		}
 	}
@@ -1207,7 +1208,7 @@ func (w *RawColumnWriter) copyDataBlocks(
 			offsetDiff := blocks[i].bh.Offset - blocks[firstBlockIdx].bh.Offset
 			dataWithTrailer := buf[offsetDiff : offsetDiff+blocks[i].bh.Length+block.TrailerLen]
 			pb := block.AlreadyEncodedPhysicalBlock(dataWithTrailer)
-			if err := w.enqueuePhysicalBlock(pb, blocks[i].sep); err != nil {
+			if err := w.enqueuePhysicalBlock(pb.Take(), blocks[i].sep); err != nil {
 				return err
 			}
 		}
@@ -1242,7 +1243,7 @@ func (w *RawColumnWriter) copyDataBlocks(
 func (w *RawColumnWriter) addDataBlock(b, sep []byte, bhp block.HandleWithProperties) error {
 	// Compress and checksum the data block and send it to the write queue.
 	pb := w.layout.physBlockMaker.Make(b, blockkind.SSTableData, block.NoFlags)
-	if err := w.enqueuePhysicalBlock(pb, sep); err != nil {
+	if err := w.enqueuePhysicalBlock(pb.Take(), sep); err != nil {
 		return err
 	}
 	return nil

sstable/layout.go

@@ -859,10 +859,12 @@ func (w *layoutWriter) WriteDataBlock(b []byte) (block.Handle, error) {
 	return w.writeBlock(b, blockkind.SSTableData)
 }
 
-// WritePrecompressedDataBlock writes a pre-compressed data block and its
-// pre-computed trailer to the writer, returning its block handle. It can mangle
-// the block data.
-func (w *layoutWriter) WritePrecompressedDataBlock(blk block.PhysicalBlock) (block.Handle, error) {
+// WritePrecompressedDataBlock writes a pre-compressed data block (including the
+// trailer) and releases it, returning its block handle. Releases the block even
+// in error cases.
+func (w *layoutWriter) WritePrecompressedDataBlock(
+	blk block.OwnedPhysicalBlock,
+) (block.Handle, error) {
 	return w.writePhysicalBlock(blk)
 }
 
@@ -943,15 +945,15 @@ func (w *layoutWriter) writeNamedBlockUncompressed(
 }
 
 // WriteValueBlock writes a pre-finished value block (with the trailer) to the
-// writer. It can mangle the block data.
-func (w *layoutWriter) WriteValueBlock(blk block.PhysicalBlock) (block.Handle, error) {
+// writer and releases the buffer. The buffer is releaseed even in error cases.
+func (w *layoutWriter) WriteValueBlock(blk block.OwnedPhysicalBlock) (block.Handle, error) {
 	return w.writePhysicalBlock(blk)
 }
 
-// WriteValueIndexBlock writes a value index block and adds it to the meta
-// index. It can mangle the block data.
+// WriteValueIndexBlock writes a value index block, releases it, ands updates
+// the meta index. The block is released even in error cases.
 func (w *layoutWriter) WriteValueIndexBlock(
-	blk block.PhysicalBlock, vbih valblk.IndexHandle,
+	blk block.OwnedPhysicalBlock, vbih valblk.IndexHandle,
 ) (block.Handle, error) {
 	h, err := w.writePhysicalBlock(blk)
 	if err != nil {
@@ -965,32 +967,31 @@ func (w *layoutWriter) WriteValueIndexBlock(
 // writeBlock checksums, compresses, and writes out a block.
 func (w *layoutWriter) writeBlock(b []byte, kind block.Kind) (block.Handle, error) {
 	pb := w.physBlockMaker.Make(b, kind, block.NoFlags)
-	defer pb.Release()
-	h, err := w.writePhysicalBlock(pb)
+	h, err := w.writePhysicalBlock(pb.Take())
 	return h, err
 }
 
 // writeBlock checksums and writes out a block.
 func (w *layoutWriter) writeBlockUncompressed(b []byte, kind block.Kind) (block.Handle, error) {
 	pb := w.physBlockMaker.Make(b, kind, block.DontCompress)
-	defer pb.Release()
-	h, err := w.writePhysicalBlock(pb)
+	h, err := w.writePhysicalBlock(pb.Take())
 	return h, err
 }
 
-// writePhysicalBlock writes a physical block to the writer, returning its block
-// handle. Does not release the physical block.
+// writePhysicalBlock writes a physical block to the writer and releases it,
+// returning its block handle. Releases the physical block even in error cases.
 //
 // writePhysicalBlock might mangle the block data.
-func (w *layoutWriter) writePhysicalBlock(blk block.PhysicalBlock) (block.Handle, error) {
+func (w *layoutWriter) writePhysicalBlock(blk block.OwnedPhysicalBlock) (block.Handle, error) {
 	w.clearFromCache(w.offset)
-	// Write the bytes to the file.
-	n, err := blk.WriteTo(w.writable)
+	// Write the bytes to the file. Note that this call can mangle the block data,
+	// but it should not matter since we are releasing the TempBuffer right away.
+	length, err := block.WriteAndReleasePhysicalBlock(blk, w.writable)
 	if err != nil {
 		return block.Handle{}, err
 	}
-	bh := block.Handle{Offset: w.offset, Length: uint64(blk.LengthWithoutTrailer())}
-	w.offset += uint64(n)
+	bh := block.Handle{Offset: w.offset, Length: uint64(length.WithoutTrailer())}
+	w.offset += uint64(length.WithTrailer())
 	return bh, nil
 }
 

sstable/rowblk_writer.go

@@ -1810,7 +1810,7 @@ func (w *RawRowWriter) rewriteSuffixes(
 
 	for i := range blocks {
 		// Write the rewritten block to the file.
-		bh, err := w.layout.WritePrecompressedDataBlock(blocks[i].physical)
+		bh, err := w.layout.WritePrecompressedDataBlock(blocks[i].physical.Take())
 		blocks[i].physical.Release()
 		if err != nil {
 			return err
@@ -1922,8 +1922,7 @@ func (w *RawRowWriter) addDataBlock(b, sep []byte, bhp block.HandleWithPropertie
 
 	// layout.writePhysicalBlock keeps layout.offset up-to-date for us.
 	// Note that this can mangle the pb data.
-	bh, err := w.layout.writePhysicalBlock(pb)
-	pb.Release()
+	bh, err := w.layout.writePhysicalBlock(pb.Take())
 	if err != nil {
 		return err
 	}

sstable/valblk/writer.go

@@ -91,8 +91,8 @@ func (w *Writer) Size() uint64 {
 func (w *Writer) compressAndFlush() {
 	physicalBlock := w.physBlockMaker.Make(w.buf.Data(), blockkind.SSTableValue, block.NoFlags)
 	w.buf.Reset()
-	bh := block.Handle{Offset: w.totalBlockBytes, Length: uint64(physicalBlock.LengthWithoutTrailer())}
-	w.totalBlockBytes += uint64(physicalBlock.LengthWithTrailer())
+	bh := block.Handle{Offset: w.totalBlockBytes, Length: uint64(physicalBlock.Length().WithoutTrailer())}
+	w.totalBlockBytes += uint64(physicalBlock.Length().WithTrailer())
 	// blockFinishedFunc length excludes the block trailer.
 	w.blockFinishedFunc(physicalBlock.LengthWithoutTrailer())
 	w.blocks = append(w.blocks, bufferedValueBlock{
@@ -114,7 +114,7 @@ func (w *Writer) Finish(layout LayoutWriter, fileOffset uint64) (IndexHandle, Wr
 	largestOffset := uint64(0)
 	largestLength := uint64(0)
 	for i := range w.blocks {
-		_, err := layout.WriteValueBlock(w.blocks[i].block)
+		_, err := layout.WriteValueBlock(w.blocks[i].block.Take())
 		if err != nil {
 			return IndexHandle{}, WriterStats{}, err
 		}
@@ -165,8 +165,7 @@ func (w *Writer) writeValueBlocksIndex(layout LayoutWriter, h IndexHandle) (Inde
 		panic("incorrect length calculation")
 	}
 	pb := w.physBlockMaker.Make(w.buf.Data(), blockkind.Metadata, block.DontCompress)
-	defer pb.Release()
-	if _, err := layout.WriteValueIndexBlock(pb, h); err != nil {
+	if _, err := layout.WriteValueIndexBlock(pb.Take(), h); err != nil {
 		return IndexHandle{}, err
 	}
 	return h, nil
@@ -200,10 +199,10 @@ type WriterStats struct {
 // LayoutWriter defines the interface for a writer that writes out serialized
 // value and value index blocks.
 type LayoutWriter interface {
-	// WriteValueBlock writes a pre-finished value block (with the trailer) to
-	// the writer.
-	WriteValueBlock(blk block.PhysicalBlock) (block.Handle, error)
-	// WriteValueIndexBlock writes a pre-finished value block index to the
-	// writer.
-	WriteValueIndexBlock(blk block.PhysicalBlock, vbih IndexHandle) (block.Handle, error)
+	// WriteValueBlock writes a pre-finished value block (with the trailer) to the
+	// writer and releases the block. The block is released even in error cases.
+	WriteValueBlock(blk block.OwnedPhysicalBlock) (block.Handle, error)
+	// WriteValueIndexBlock writes a pre-finished value block index to the writer
+	// and releases the block. The block is released even in error cases.
+	WriteValueIndexBlock(blk block.OwnedPhysicalBlock, vbih IndexHandle) (block.Handle, error)
 }

sstable/write_queue.go

@@ -60,8 +60,7 @@ func newWriteQueue(size int, writer *RawRowWriter) *writeQueue {
 }
 
 func (w *writeQueue) performWrite(task *writeTask) error {
-	handle, err := w.writer.layout.WritePrecompressedDataBlock(task.buf.physical)
-	task.buf.physical.Release()
+	handle, err := w.writer.layout.WritePrecompressedDataBlock(task.buf.physical.Take())
 	if err != nil {
 		return err
 	}