db: fix plumbing of CompactionGrantHandle

sumeerbhola · sumeerbhola · commit c0b18bc8970e · 2025-04-02T14:24:38.000-04:00
CompactionGrantHandle now includes the CPUMeasurer interface, that has
a separate method for each kind of goroutine, instead of identifying the
goroutine kind by an index. This includes the secondary goroutine for
writing to blob files. This interface is plumbed to RawColumnWriter and
blob.FileWriter, so that the secondary writing goroutines can report
CPU after writing each block.
diff --git a/compaction.go b/compaction.go
@@ -3305,7 +3305,7 @@ func (d *DB) newCompactionOutput(
 		},
 	})
 
-	tw := sstable.NewRawWriter(writable, writerOpts)
+	tw := sstable.NewRawWriterWithCPUMeasurer(writable, writerOpts, c.grantHandle)
 	return objMeta, tw, nil
 }
 
diff --git a/compaction_scheduler.go b/compaction_scheduler.go
@@ -14,6 +14,13 @@ import (
 
 type CompactionGrantHandle = base.CompactionGrantHandle
 type CompactionGrantHandleStats = base.CompactionGrantHandleStats
+type CompactionGoroutineKind = base.CompactionGoroutineKind
+
+const (
+	CompactionGoroutinePrimary           = base.CompactionGoroutinePrimary
+	CompactionGoroutineSSTableSecondary  = base.CompactionGoroutineSSTableSecondary
+	CompactionGoroutineBlobFileSecondary = base.CompactionGoroutineBlobFileSecondary
+)
 
 // NB: This interface is experimental and subject to change.
 //
@@ -157,7 +164,7 @@ type noopGrantHandle struct{}
 var _ CompactionGrantHandle = noopGrantHandle{}
 
 func (h noopGrantHandle) Started()                                              {}
-func (h noopGrantHandle) MeasureCPU(g int)                                      {}
+func (h noopGrantHandle) MeasureCPU(CompactionGoroutineKind)                    {}
 func (h noopGrantHandle) CumulativeStats(stats base.CompactionGrantHandleStats) {}
 func (h noopGrantHandle) Done()                                                 {}
 
@@ -333,10 +340,8 @@ func (s *ConcurrencyLimitScheduler) TrySchedule() (bool, CompactionGrantHandle)
 	return false, nil
 }
 
-func (s *ConcurrencyLimitScheduler) Started() {}
-
-func (s *ConcurrencyLimitScheduler) MeasureCPU(g int) {}
-
+func (s *ConcurrencyLimitScheduler) Started()                                              {}
+func (s *ConcurrencyLimitScheduler) MeasureCPU(CompactionGoroutineKind)                    {}
 func (s *ConcurrencyLimitScheduler) CumulativeStats(stats base.CompactionGrantHandleStats) {}
 
 func (s *ConcurrencyLimitScheduler) Done() {
diff --git a/internal/base/compaction_grant_handle.go b/internal/base/compaction_grant_handle.go
@@ -20,14 +20,7 @@ type CompactionGrantHandle interface {
 	// Started is called once and must precede calls to MeasureCPU and
 	// CumulativeStats.
 	Started()
-	// MeasureCPU is used to measure the CPU consumption of a goroutine involved
-	// in a compaction. It must be called from each of the two goroutines that
-	// consume CPU during a compaction, and the first call must be before any
-	// significant work is done, since the first call is used to initialize the
-	// measurer for the goroutine. The parameter g must be 0 or 1, to
-	// differentiate between the goroutines. If a compaction is only using one
-	// goroutine, then it can skip calling MeasureCPU(1).
-	MeasureCPU(g int)
+	CPUMeasurer
 	// CumulativeStats reports the current cumulative stats. This method may
 	// block if the scheduler wants to pace the compaction (say to moderate its
 	// consumption of disk write bandwidth).
@@ -39,3 +32,36 @@ type CompactionGrantHandle interface {
 	// been installed.
 	Done()
 }
+
+// CompactionGoroutineKind identifies the kind of compaction goroutine.
+type CompactionGoroutineKind uint8
+
+const (
+	// CompactionGoroutinePrimary is the primary compaction goroutine that
+	// iterates over key-value pairs in the input and calls the current sstable
+	// writer and blob file writer.
+	CompactionGoroutinePrimary CompactionGoroutineKind = iota
+	// CompactionGoroutineSSTableSecondary is the secondary goroutine in the
+	// current sstable writer that writes blocks to the sstable.
+	CompactionGoroutineSSTableSecondary
+	// CompactionGoroutineBlobFileSecondary is the secondary goroutine in the
+	// current blob file writer that writes blocks to the blob file.
+	CompactionGoroutineBlobFileSecondary
+)
+
+// CPUMeasurer is used to measure the CPU consumption of goroutines involved
+// in a compaction.
+type CPUMeasurer interface {
+	// MeasureCPU allows the measurer to keep track of CPU usage while a
+	// compaction is ongoing. It is to be called regularly from the compaction
+	// goroutine corresponding to the argument. The first call from a goroutine
+	// must be done before any significant CPU consumption, since it is used to
+	// initialize the measurer for the goroutine making the call. If a
+	// compaction is not using a certain kind of goroutine, it can skip calling
+	// this method with the corresponding argument.
+	MeasureCPU(CompactionGoroutineKind)
+}
+
+type NoopCPUMeasurer struct{}
+
+func (NoopCPUMeasurer) MeasureCPU(CompactionGoroutineKind) {}
diff --git a/internal/compact/run.go b/internal/compact/run.go
@@ -273,12 +273,14 @@ func (r *Runner) writeKeysToTable(tw sstable.RawWriter) (splitKey []byte, _ erro
 		iteratedKeys++
 		if iteratedKeys%updateGrantHandleEveryNKeys == 0 {
 			r.cfg.GrantHandle.CumulativeStats(base.CompactionGrantHandleStats{
+				// TODO(jackson): CumulativeWrittenSize does not include blob files
+				// already written by this compaction, and
+				// ValueSeparation.EstimatedFileSize does not either. So this
+				// CumWriteBytes is incomplete.
 				CumWriteBytes: r.stats.CumulativeWrittenSize + tw.EstimatedSize() +
 					r.cfg.ValueSeparation.EstimatedFileSize(),
 			})
-			// TODO(sumeer): give the GrantHandle to the writer so it can account on
-			// all its goroutines.
-			r.cfg.GrantHandle.MeasureCPU(0)
+			r.cfg.GrantHandle.MeasureCPU(base.CompactionGoroutinePrimary)
 		}
 		outputSize := tw.EstimatedSize()
 		outputSize += r.cfg.ValueSeparation.EstimatedReferenceSize()
@@ -336,11 +338,15 @@ func (r *Runner) writeKeysToTable(tw sstable.RawWriter) (splitKey []byte, _ erro
 	r.stats.CumulativePinnedSize += pinnedKeySize + pinnedValueSize
 
 	r.cfg.GrantHandle.CumulativeStats(base.CompactionGrantHandleStats{
+		// TODO(jackson): CumulativeWrittenSize does not include blob files
+		// already written by this compaction, and
+		// ValueSeparation.EstimatedFileSize does not either. So this
+		// CumWriteBytes is incomplete.
 		CumWriteBytes: r.stats.CumulativeWrittenSize +
 			tw.EstimatedSize() +
 			r.cfg.ValueSeparation.EstimatedFileSize(),
 	})
-	r.cfg.GrantHandle.MeasureCPU(0)
+	r.cfg.GrantHandle.MeasureCPU(base.CompactionGoroutinePrimary)
 	return splitKey, nil
 }
 
diff --git a/sstable/blob/blob.go b/sstable/blob/blob.go
@@ -52,6 +52,8 @@ type FileWriterOptions struct {
 	Compression   block.Compression
 	ChecksumType  block.ChecksumType
 	FlushGovernor block.FlushGovernor
+	// Only CPUMeasurer.MeasureCPUBlobFileSecondary is used.
+	CpuMeasurer base.CPUMeasurer
 }
 
 func (o *FileWriterOptions) ensureDefaults() {
@@ -68,6 +70,9 @@ func (o *FileWriterOptions) ensureDefaults() {
 			base.SizeClassAwareBlockSizeThreshold,
 			nil)
 	}
+	if o.CpuMeasurer == nil {
+		o.CpuMeasurer = base.NoopCPUMeasurer{}
+	}
 }
 
 // FileWriterStats aggregates statistics about a blob file written by a
@@ -98,6 +103,7 @@ type FileWriter struct {
 	blockOffsets []uint64
 	err          error
 	checksumType block.ChecksumType
+	cpuMeasurer  base.CPUMeasurer
 	writeQueue   struct {
 		wg  sync.WaitGroup
 		ch  chan compressedBlock
@@ -120,6 +126,7 @@ func NewFileWriter(fn base.DiskFileNum, w objstorage.Writable, opts FileWriterOp
 	fw.b.Init(opts.Compression, opts.ChecksumType)
 	fw.flushGov = opts.FlushGovernor
 	fw.checksumType = opts.ChecksumType
+	fw.cpuMeasurer = opts.CpuMeasurer
 	fw.writeQueue.ch = make(chan compressedBlock)
 	fw.writeQueue.wg.Add(1)
 	go fw.drainWriteQueue()
@@ -173,8 +180,13 @@ func (w *FileWriter) flush() {
 // until the channel is closed. All value blocks are written by this goroutine.
 func (w *FileWriter) drainWriteQueue() {
 	defer w.writeQueue.wg.Done()
+	// Call once to initialize the CPU measurer.
+	w.cpuMeasurer.MeasureCPU(base.CompactionGoroutineBlobFileSecondary)
 	for cb := range w.writeQueue.ch {
 		_, err := cb.pb.WriteTo(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)
 		if err != nil {
 			w.writeQueue.err = err
 			continue
diff --git a/sstable/colblk_writer.go b/sstable/colblk_writer.go
@@ -102,12 +102,17 @@ type RawColumnWriter struct {
 	previousUserKey       invariants.Value[[]byte]
 	validator             invariants.Value[*colblk.DataBlockValidator]
 	disableKeyOrderChecks bool
+	cpuMeasurer           base.CPUMeasurer
 }
 
 // Assert that *RawColumnWriter implements RawWriter.
 var _ RawWriter = (*RawColumnWriter)(nil)
 
-func newColumnarWriter(writable objstorage.Writable, o WriterOptions) *RawColumnWriter {
+// cpuMeasurer, if non-nil, is only used for calling
+// cpuMeasurer.MeasureCPUSSTableSecondary.
+func newColumnarWriter(
+	writable objstorage.Writable, o WriterOptions, cpuMeasurer base.CPUMeasurer,
+) *RawColumnWriter {
 	if writable == nil {
 		panic("pebble: nil writable")
 	}
@@ -177,6 +182,7 @@ func newColumnarWriter(writable objstorage.Writable, o WriterOptions) *RawColumn
 
 	w.writeQueue.ch = make(chan *compressedBlock)
 	w.writeQueue.wg.Add(1)
+	w.cpuMeasurer = cpuMeasurer
 	go w.drainWriteQueue()
 	return w
 }
@@ -863,10 +869,15 @@ func (w *RawColumnWriter) flushBufferedIndexBlocks() (rootIndex block.Handle, er
 // Other blocks are written directly by the client goroutine. See Close.
 func (w *RawColumnWriter) drainWriteQueue() {
 	defer w.writeQueue.wg.Done()
+	// Call once to initialize the CPU measurer.
+	w.cpuMeasurer.MeasureCPU(base.CompactionGoroutineSSTableSecondary)
 	for cb := range w.writeQueue.ch {
 		if _, err := w.layout.WritePrecompressedDataBlock(cb.physical); err != nil {
 			w.writeQueue.err = err
 		}
+		// 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)
 		cb.blockBuf.clear()
 		cb.physical = block.PhysicalBlock{}
 		compressedBlockPool.Put(cb)
diff --git a/sstable/writer.go b/sstable/writer.go
@@ -18,10 +18,22 @@ import (
 // NewRawWriter returns a new table writer for the file. Closing the writer will
 // close the file.
 func NewRawWriter(writable objstorage.Writable, o WriterOptions) RawWriter {
+	return NewRawWriterWithCPUMeasurer(writable, o, base.NoopCPUMeasurer{})
+}
+
+// NewRawWriterWithCPUMeasurer is like NewRawWriter, but additionally allows
+// the caller to specify a CPUMeasurer. Only
+// CPUMeasurer.MeasureCPU(CompactionGoroutineSSTableSecondary) is used by the
+// writer.
+func NewRawWriterWithCPUMeasurer(
+	writable objstorage.Writable, o WriterOptions, cpuMeasurer base.CPUMeasurer,
+) RawWriter {
 	if o.TableFormat <= TableFormatPebblev4 {
+		// Don't bother plumbing the cpuMeasurer to the row writer since it is not
+		// the default and will be removed.
 		return newRowWriter(writable, o)
 	}
-	return newColumnarWriter(writable, o)
+	return newColumnarWriter(writable, o, cpuMeasurer)
 }
 
 // Writer is a table writer.

Original file line number	Diff line number	Diff line change
`@@ -3305,7 +3305,7 @@ func (d *DB) newCompactionOutput(`
`3305`	`3305`	`},`
`3306`	`3306`	`})`
`3307`	`3307`
`3308`		`- tw := sstable.NewRawWriter(writable, writerOpts)`
	`3308`	`+ tw := sstable.NewRawWriterWithCPUMeasurer(writable, writerOpts, c.grantHandle)`
`3309`	`3309`	`return objMeta, tw, nil`
`3310`	`3310`	`}`
`3311`	`3311`