From cb86af843bdee4f782da5a5116ab70542210442a Mon Sep 17 00:00:00 2001 From: Andrew Baptist Date: Thu, 11 Jul 2024 08:30:16 -0400 Subject: [PATCH] pebble: smoothing code smoother: integrate into code --- compaction.go | 19 ++-- compaction_picker_test.go | 4 +- compaction_test.go | 4 +- db.go | 5 + download.go | 2 +- internal/compact/iterator.go | 7 +- internal/compact/run.go | 14 ++- internal/rate/smoother.go | 199 +++++++++++++++++++++++++++++++++ internal/rate/smoother_test.go | 56 ++++++++++ level_iter.go | 2 + open.go | 3 + sstable/reader.go | 11 +- sstable/reader_common.go | 2 + sstable/reader_iter.go | 10 ++ sstable/reader_iter_two_lvl.go | 4 + sstable/reader_virtual.go | 4 +- table_cache.go | 2 +- 17 files changed, 326 insertions(+), 22 deletions(-) create mode 100644 internal/rate/smoother.go create mode 100644 internal/rate/smoother_test.go diff --git a/compaction.go b/compaction.go index 71186e86d09..392be5babd3 100644 --- a/compaction.go +++ b/compaction.go @@ -20,6 +20,7 @@ import ( "github.com/cockroachdb/pebble/internal/keyspan" "github.com/cockroachdb/pebble/internal/keyspan/keyspanimpl" "github.com/cockroachdb/pebble/internal/manifest" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/internal/sstableinternal" "github.com/cockroachdb/pebble/objstorage" "github.com/cockroachdb/pebble/objstorage/objstorageprovider/objiotracing" @@ -251,6 +252,8 @@ type compaction struct { metrics map[int]*LevelMetrics pickerMetrics compactionPickerMetrics + + smoother *rate.Smoother } // inputLargestSeqNumAbsolute returns the maximum LargestSeqNumAbsolute of any @@ -315,7 +318,7 @@ func (c *compaction) userKeyBounds() base.UserKeyBounds { } func newCompaction( - pc *pickedCompaction, opts *Options, beganAt time.Time, provider objstorage.Provider, + pc *pickedCompaction, opts *Options, beganAt time.Time, provider objstorage.Provider, smoother *rate.Smoother, ) *compaction { c := &compaction{ kind: compactionKindDefault, @@ -332,6 +335,7 @@ func newCompaction( maxOutputFileSize: pc.maxOutputFileSize, maxOverlapBytes: pc.maxOverlapBytes, pickerMetrics: pc.pickerMetrics, + smoother: smoother, } c.startLevel = &c.inputs[0] if pc.startLevel.l0SublevelInfo != nil { @@ -438,8 +442,7 @@ func adjustGrandparentOverlapBytesForFlush(c *compaction, flushingBytes uint64) // maxOverlapBytes will cause splits at f10, f20,..., f990, which // means an upper bound file count of 100 files. Say the input bytes // in the flush are such that acceptableFileCount=10. We will fatten - // up maxOverlapBytes by 10x to ensure that the upper bound file count - // drops to 10. However, it is possible that in practice, even without + // iter := compact.NewIter(cfg, pointIter, rangeDelIter, rangeKeyIter) 10. However, it is possible that in practice, even without // this change, we would have produced no more than 10 files, and that // this change makes the files unnecessarily wide. Say the input bytes // are distributed such that 10% are in f0...f9, 10% in f10...f19, ... @@ -747,6 +750,7 @@ func (c *compaction) newInputIters( iterOpts, c.comparer, newIters, level.files.Iter(), l, internalIterOpts{ compaction: true, bufferPool: &c.bufferPool, + smoother: c.smoother, })) // TODO(jackson): Use keyspanimpl.LevelIter to avoid loading all the range // deletions into memory upfront. (See #2015, which reverted this.) There @@ -904,6 +908,7 @@ func (c *compaction) newRangeDelIter( internalIterOpts{ compaction: true, bufferPool: &c.bufferPool, + smoother: c.smoother, }, iterRangeDeletions) if err != nil { return nil, err @@ -1784,7 +1789,7 @@ func (d *DB) tryScheduleManualCompaction(env compactionEnv, manual *manualCompac return false } - c := newCompaction(pc, d.opts, d.timeNow(), d.ObjProvider()) + c := newCompaction(pc, d.opts, d.timeNow(), d.ObjProvider(), d.smoother) d.mu.compact.compactingCount++ d.addInProgressCompaction(c) go d.compact(c, manual.done) @@ -1810,7 +1815,7 @@ func (d *DB) tryScheduleAutoCompaction( if pc == nil { return false } - c := newCompaction(pc, d.opts, d.timeNow(), d.ObjProvider()) + c := newCompaction(pc, d.opts, d.timeNow(), d.ObjProvider(), d.smoother) d.mu.compact.compactingCount++ d.addInProgressCompaction(c) go d.compact(c, nil) @@ -2510,7 +2515,7 @@ func (d *DB) compactAndWrite( IneffectualSingleDeleteCallback: d.opts.Experimental.IneffectualSingleDeleteCallback, SingleDeleteInvariantViolationCallback: d.opts.Experimental.SingleDeleteInvariantViolationCallback, } - iter := compact.NewIter(cfg, pointIter, rangeDelIter, rangeKeyIter) + iter := compact.NewIter(cfg, pointIter, rangeDelIter, rangeKeyIter, d.smoother) runnerCfg := compact.RunnerConfig{ CompactionBounds: base.UserKeyBoundsFromInternal(c.smallest, c.largest), @@ -2519,7 +2524,7 @@ func (d *DB) compactAndWrite( MaxGrandparentOverlapBytes: c.maxOverlapBytes, TargetOutputFileSize: c.maxOutputFileSize, } - runner := compact.NewRunner(runnerCfg, iter) + runner := compact.NewRunner(runnerCfg, iter, d.smoother) for runner.MoreDataToWrite() { if c.cancel.Load() { return runner.Finish().WithError(ErrCancelledCompaction) diff --git a/compaction_picker_test.go b/compaction_picker_test.go index d305b1678f1..f628ea7ac15 100644 --- a/compaction_picker_test.go +++ b/compaction_picker_test.go @@ -539,7 +539,7 @@ func TestCompactionPickerL0(t *testing.T) { var result strings.Builder if pc != nil { checkClone(t, pc) - c := newCompaction(pc, opts, time.Now(), nil /* provider */) + c := newCompaction(pc, opts, time.Now(), nil /* provider */, nil /* smoother */) fmt.Fprintf(&result, "L%d -> L%d\n", pc.startLevel.level, pc.outputLevel.level) fmt.Fprintf(&result, "L%d: %s\n", pc.startLevel.level, fileNums(pc.startLevel.files)) if !pc.outputLevel.files.Empty() { @@ -760,7 +760,7 @@ func TestCompactionPickerConcurrency(t *testing.T) { }) var result strings.Builder if pc != nil { - c := newCompaction(pc, opts, time.Now(), nil /* provider */) + c := newCompaction(pc, opts, time.Now(), nil /* provider */, nil /* smoother */) fmt.Fprintf(&result, "L%d -> L%d\n", pc.startLevel.level, pc.outputLevel.level) fmt.Fprintf(&result, "L%d: %s\n", pc.startLevel.level, fileNums(pc.startLevel.files)) if !pc.outputLevel.files.Empty() { diff --git a/compaction_test.go b/compaction_test.go index 4a7e95ddbc8..a5afa0637a8 100644 --- a/compaction_test.go +++ b/compaction_test.go @@ -516,7 +516,7 @@ func TestPickCompaction(t *testing.T) { vs.picker = &tc.picker pc, got := vs.picker.pickAuto(compactionEnv{diskAvailBytes: math.MaxUint64}), "" if pc != nil { - c := newCompaction(pc, opts, time.Now(), nil /* provider */) + c := newCompaction(pc, opts, time.Now(), nil /* provider */, nil /* smoother */) gotStart := fileNums(c.startLevel.files) gotML := "" @@ -1297,7 +1297,7 @@ func TestCompactionOutputLevel(t *testing.T) { d.ScanArgs(t, "start", &start) d.ScanArgs(t, "base", &base) pc := newPickedCompaction(opts, version, start, defaultOutputLevel(start, base), base) - c := newCompaction(pc, opts, time.Now(), nil /* provider */) + c := newCompaction(pc, opts, time.Now(), nil /* provider */, nil /* smoother */) return fmt.Sprintf("output=%d\nmax-output-file-size=%d\n", c.outputLevel.level, c.maxOutputFileSize) diff --git a/db.go b/db.go index d879f688256..4a68504db1d 100644 --- a/db.go +++ b/db.go @@ -23,6 +23,7 @@ import ( "github.com/cockroachdb/pebble/internal/keyspan/keyspanimpl" "github.com/cockroachdb/pebble/internal/manifest" "github.com/cockroachdb/pebble/internal/manual" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/objstorage" "github.com/cockroachdb/pebble/objstorage/remote" "github.com/cockroachdb/pebble/rangekey" @@ -307,6 +308,8 @@ type DB struct { commit *commitPipeline + smoother *rate.Smoother + // readState provides access to the state needed for reading without needing // to acquire DB.mu. readState struct { @@ -1621,6 +1624,8 @@ func (d *DB) Close() error { panic(err) } + d.smoother.Stop() + // Clear the finalizer that is used to check that an unreferenced DB has been // closed. We're closing the DB here, so the check performed by that // finalizer isn't necessary. diff --git a/download.go b/download.go index 3c4a331a4e5..60ba41d4abe 100644 --- a/download.go +++ b/download.go @@ -433,7 +433,7 @@ func (d *DB) tryLaunchDownloadForFile( download.numLaunchedDownloads++ doneCh = make(chan error, 1) - c := newCompaction(pc, d.opts, d.timeNow(), d.objProvider) + c := newCompaction(pc, d.opts, d.timeNow(), d.objProvider, d.smoother) c.isDownload = true d.mu.compact.downloadingCount++ d.addInProgressCompaction(c) diff --git a/internal/compact/iterator.go b/internal/compact/iterator.go index f17e88e7be9..eb9fb57e1b8 100644 --- a/internal/compact/iterator.go +++ b/internal/compact/iterator.go @@ -15,6 +15,7 @@ import ( "github.com/cockroachdb/pebble/internal/invariants" "github.com/cockroachdb/pebble/internal/keyspan" "github.com/cockroachdb/pebble/internal/rangekey" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/redact" ) @@ -248,6 +249,8 @@ type Iter struct { span keyspan.Span stats IterStats + + smoother *rate.Smoother } // IterConfig contains the parameters necessary to create a compaction iterator. @@ -305,6 +308,7 @@ func NewIter( cfg IterConfig, pointIter base.InternalIterator, rangeDelIter, rangeKeyIter keyspan.FragmentIterator, + s *rate.Smoother, ) *Iter { cfg.ensureDefaults() i := &Iter{ @@ -312,7 +316,8 @@ func NewIter( cfg: cfg, // We don't want a nil keyBuf because if the first key we encounter is // empty, it would become nil. - keyBuf: make([]byte, 8), + keyBuf: make([]byte, 8), + smoother: s, } iter := pointIter diff --git a/internal/compact/run.go b/internal/compact/run.go index a15e380a0a5..1b7a8dbc197 100644 --- a/internal/compact/run.go +++ b/internal/compact/run.go @@ -12,6 +12,7 @@ import ( "github.com/cockroachdb/pebble/internal/base" "github.com/cockroachdb/pebble/internal/keyspan" "github.com/cockroachdb/pebble/internal/manifest" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/objstorage" "github.com/cockroachdb/pebble/sstable" ) @@ -107,14 +108,16 @@ type Runner struct { // Last range key span (or portion of it) that was not yet written to a table. lastRangeKeySpan keyspan.Span stats Stats + smoother *rate.Smoother } // NewRunner creates a new Runner. -func NewRunner(cfg RunnerConfig, iter *Iter) *Runner { +func NewRunner(cfg RunnerConfig, iter *Iter, s *rate.Smoother) *Runner { r := &Runner{ - cmp: iter.cmp, - cfg: cfg, - iter: iter, + cmp: iter.cmp, + cfg: cfg, + iter: iter, + smoother: s, } r.key, r.value = r.iter.First() return r @@ -176,10 +179,13 @@ func (r *Runner) writeKeysToTable(tw *sstable.Writer) (splitKey []byte, _ error) } var pinnedKeySize, pinnedValueSize, pinnedCount uint64 key, value := r.key, r.value + si := r.smoother.Track(tw.EstimatedSize) + defer si.Close() for ; key != nil; key, value = r.iter.Next() { if splitter.ShouldSplitBefore(key.UserKey, tw.EstimatedSize(), lastUserKeyFn) { break } + si.Tick() switch key.Kind() { case base.InternalKeyKindRangeDelete: diff --git a/internal/rate/smoother.go b/internal/rate/smoother.go new file mode 100644 index 00000000000..9dcd7e5bdd7 --- /dev/null +++ b/internal/rate/smoother.go @@ -0,0 +1,199 @@ +package rate + +import ( + "fmt" + "math" + "sync/atomic" + "time" +) + +// With a decay rate of 0.01, and a 10ms sample rate, it takes ~10 second to react to a significant change. +const decayRate float64 = 0.01 + +// minUtilization is the minimum utilization we will allow. This is to prevent +// us from quickly reacting to external changes. +const minUtilization float64 = 0.1 + +// 500 measurements at 20ms sample rate is 10 seconds. +const sampleRate = 20 * time.Millisecond +const numMeasurements = 500 + +// The measurementBuffer stores numMeasurements / sample rates in the +// measurement buffer. With these defaults it looks back 10 seconds. +type measurementBuffer struct { + measurements [numMeasurements]int32 + index int +} + +func newMeasurementBuffer() measurementBuffer { + b := measurementBuffer{} + for i := range b.measurements { + b.measurements[i] = 1 + } + return b +} + +func (b *measurementBuffer) addMeasurement(value int32) { + b.measurements[b.index] = value + b.index = (b.index + 1) % numMeasurements +} + +func (b *measurementBuffer) getAverage() float64 { + sum := 0.0 + for _, value := range b.measurements { + sum += float64(value) + } + return sum / float64(numMeasurements) +} + +// Smoother will attempt to smooth out a process that runs multiple iterations. +// The goal is not to have it run faster or slower, but simply with pace itself +// to run evenly over time rather than being bunchy. +// +// Each "thread" can be in one of three states, Running, Idle, Sleeping. +// * Running means the thread is currently working (and writing to disk). +// * Idle means this thread is not doing anything as there is no work to do. +// * Sleeping means the thread has active work to do, but is sleeping to smooth load. +// +// The goal of the smoother is to turn all Idle work into Sleeping work while +// preserving the amount of Running work that is being done. The sleeping work +// can be somewhat evenly placed between small Running tasks. The perfect +// smoother would result in the flush / compaction write bandwidth to be a +// constant. +// +// Some simplifying assumptions are made which could be improved upon. +// * All tasks runs for a similar amount of time (estimatedIterDurationNs). +// * All tasks create a similar amount of IO. +// * Different types of tasks (compaction vs flush) are similar cost. +// +// After a Smoother is created, it must be started by calling start. +type Smoother struct { + disabled bool + runningCount atomic.Int32 + sleepingCount atomic.Int32 + + // estimatedUtilization is the estimated utilization of the disks during + // compaction. If it is less than 100% then we pace compactions to prevent + // spike disk behavior. + estimatedUtilization atomic.Value + mb measurementBuffer + + stopper chan struct{} +} + +// NewSmoother creates a new Smoother instance. +func NewSmoother() *Smoother { + s := Smoother{} + s.runningCount.Store(0) + s.sleepingCount.Store(0) + s.estimatedUtilization.Store(float64(1.0)) + s.stopper = make(chan struct{}) + s.mb = newMeasurementBuffer() + return &s +} + +// Start the smoothing tracking loop. This loop runs in the background and +// checks how many jobs are running and sleeping every 10ms. +func (s *Smoother) Start() { + go func() { + ticker := time.NewTicker(sampleRate) + defer ticker.Stop() + for { + select { + case <-s.stopper: + return + + case <-ticker.C: + s.mb.addMeasurement(s.runningCount.Load() - s.sleepingCount.Load()) + s.estimatedUtilization.Store(math.Max(s.mb.getAverage(), minUtilization)) + } + } + }() +} + +// Stop is called to stop them main thread running. +func (s *Smoother) Stop() { + close(s.stopper) +} + +// Track is called before work is started so the smoother can tell if there +// is active work being done. Close must be called on the smootherInstance when +// the work is complete. The function passed in should return a monotonically +// increasing value that estimates how much data this job has done. This is used +// to determine when to inject sleep intervals. +func (s *Smoother) Track(f func() uint64) Tracked { + s.runningCount.Add(1) + now := time.Now() + return &smootherInstance{s: s, f: f, bytesWritten: f(), startWorkTime: now, startTime: now} +} + +type Tracked interface { + Tick() time.Duration + Close() +} + +const chunkSize = 4 * 1024 * 1024 + +// We want to record the amount of time we were sleeping vs running over the time window. +func (s *smootherInstance) Close() { + s.s.runningCount.Add(-1) +} + +// SmootherInstance is a handle to a single instance of work being done. It is +// not safe to pass between goroutines. +type smootherInstance struct { + s *Smoother + f func() uint64 + bytesWritten uint64 + startWorkTime time.Time + startTime time.Time + chunkEstimate time.Duration +} + +// tick is called after work is done to determine how long to sleep. +func (s *smootherInstance) Tick() time.Duration { + // Duration to sleep if the caller wants to sleep after a call. + if s.s.disabled { + return 0 + } + // If this call is expensive, we could consider calling less frequently. + curVal := s.f() + + // Most of the time we don't record our work. Wait until we have done enough + // work to have a meaningful sleep. + if curVal-s.bytesWritten < chunkSize { + return 0 + } + + // Only record work if we have written enough bytes. + sleepTime := s.recordWork(time.Since(s.startWorkTime)) + // TODO: remove + fmt.Println(sleepTime, s.s.estimatedUtilization.Load(), s.chunkEstimate, curVal, s.bytesWritten) + time.Sleep(sleepTime) + + s.bytesWritten = s.f() + s.startWorkTime = time.Now() + return sleepTime +} + +// recordWork should be called periodically while a jobs is running to +// determine how long the process should sleep for. +func (s *smootherInstance) recordWork(duration time.Duration) time.Duration { + // We just did a chunk of work, add it to the estimated duration of work. + if s.chunkEstimate == 0 { + s.chunkEstimate = time.Duration(duration) + } else { + // Exponential moving average of the chunk duration estimate. + s.chunkEstimate = time.Duration(float64(s.chunkEstimate)*(1-decayRate) + float64(duration)*decayRate) + } + + currentUtil := s.s.estimatedUtilization.Load().(float64) / 3.0 + + // If we are over 100% utilization don't sleep. + if currentUtil > 1.0 { + return 0 + } + // Compute the sleep time based on the average rather than the current + // duration to smooth out the sleep times. + return time.Duration(float64(s.chunkEstimate)*(1/currentUtil) - 1) +} diff --git a/internal/rate/smoother_test.go b/internal/rate/smoother_test.go new file mode 100644 index 00000000000..9e1007ba16f --- /dev/null +++ b/internal/rate/smoother_test.go @@ -0,0 +1,56 @@ +package rate + +import ( + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestSleepConverge(t *testing.T) { + smoother := NewSmoother() + + i := smoother.Track(func() uint64 { return 0 }) + defer i.close() + + for x := 0; x < 10000; x++ { + i.recordWork(100 * time.Millisecond) + } + smoother.estimatedUtilization.Store(0.1) + sleepTime := i.recordWork(1 * time.Millisecond) + + // Sleep time is 900ms since it is 10% utilized and 90% idle. The estimated + // time is based on the 100ms work time estimate. + require.InDelta(t, 900*time.Millisecond, sleepTime, float64(time.Microsecond)) +} + +func TestSleep(t *testing.T) { + s := NewSmoother() + s.estimatedUtilization.Store(0.5) + + i := s.Track(func() uint64 { return 0 }) + // The previous estimated time is 1ms, so the sleep time should be 1ms at 50% utilization. + sleepTime := i.recordWork(time.Millisecond) + require.InDelta(t, sleepTime, time.Duration(1*float64(time.Millisecond)), float64(time.Microsecond)) +} + +// TestAvgUtil tests that the smoother converges to the correct average utilization. +func TestAvgUtil(t *testing.T) { + smoother := NewSmoother() + smoother.Start() + defer smoother.Stop() + + tick := smoother.Track(func() uint64 { return 8 * 1024 * 1024 }) + + // For the first 100ms we are 100% utilized. + time.Sleep(100 * time.Millisecond) + tick.close() + + // Now we are not utilized at all. + time.Sleep(100 * time.Millisecond) + + // Over the last 10s we should have been utilized for 9.9s/10s = 99% of the time. + fmt.Println(smoother.estimatedUtilization.Load()) + require.InDelta(t, 0.99, smoother.estimatedUtilization.Load(), 0.01) +} diff --git a/level_iter.go b/level_iter.go index 0e0390e2b24..511d216e279 100644 --- a/level_iter.go +++ b/level_iter.go @@ -14,6 +14,7 @@ import ( "github.com/cockroachdb/pebble/internal/invariants" "github.com/cockroachdb/pebble/internal/keyspan" "github.com/cockroachdb/pebble/internal/manifest" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/internal/treeprinter" "github.com/cockroachdb/pebble/sstable" ) @@ -26,6 +27,7 @@ type internalIterOpts struct { bufferPool *sstable.BufferPool stats *base.InternalIteratorStats boundLimitedFilter sstable.BoundLimitedBlockPropertyFilter + smoother *rate.Smoother } // levelIter provides a merged view of the sstables in a level. diff --git a/open.go b/open.go index 20aa384cf66..71384f00835 100644 --- a/open.go +++ b/open.go @@ -26,6 +26,7 @@ import ( "github.com/cockroachdb/pebble/internal/invariants" "github.com/cockroachdb/pebble/internal/manifest" "github.com/cockroachdb/pebble/internal/manual" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/objstorage" "github.com/cockroachdb/pebble/objstorage/objstorageprovider" "github.com/cockroachdb/pebble/objstorage/remote" @@ -218,9 +219,11 @@ func Open(dirname string, opts *Options) (db *DB, err error) { dataDir: dataDir, closed: new(atomic.Value), closedCh: make(chan struct{}), + smoother: rate.NewSmoother(), } d.mu.versions = &versionSet{} d.diskAvailBytes.Store(math.MaxUint64) + d.smoother.Start() defer func() { // If an error or panic occurs during open, attempt to release the manually diff --git a/sstable/reader.go b/sstable/reader.go index 4acbd7f419e..2fa4cce14cf 100644 --- a/sstable/reader.go +++ b/sstable/reader.go @@ -23,6 +23,7 @@ import ( "github.com/cockroachdb/pebble/internal/crc" "github.com/cockroachdb/pebble/internal/invariants" "github.com/cockroachdb/pebble/internal/keyspan" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/internal/sstableinternal" "github.com/cockroachdb/pebble/objstorage" "github.com/cockroachdb/pebble/objstorage/objstorageprovider" @@ -258,8 +259,9 @@ func (r *Reader) NewCompactionIter( statsCollector *CategoryStatsCollector, rp ReaderProvider, bufferPool *block.BufferPool, + smoother *rate.Smoother, ) (Iterator, error) { - return r.newCompactionIter(transforms, categoryAndQoS, statsCollector, rp, nil, bufferPool) + return r.newCompactionIter(transforms, categoryAndQoS, statsCollector, rp, nil, bufferPool, smoother) } func (r *Reader) newCompactionIter( @@ -269,10 +271,13 @@ func (r *Reader) newCompactionIter( rp ReaderProvider, vState *virtualState, bufferPool *block.BufferPool, + smoother *rate.Smoother, ) (Iterator, error) { if vState != nil && vState.isSharedIngested { transforms.HideObsoletePoints = true } + pos := 0 + tracked := smoother.Track(func() uint64 { pos++; return uint64(pos * 10000) }) if r.Properties.IndexType == twoLevelIndex { i, err := newTwoLevelIterator( context.Background(), @@ -283,7 +288,7 @@ func (r *Reader) newCompactionIter( return nil, err } i.setupForCompaction() - return &twoLevelCompactionIterator{twoLevelIterator: i}, nil + return &twoLevelCompactionIterator{twoLevelIterator: i, tracked: tracked}, nil } i, err := newSingleLevelIterator( context.Background(), r, vState, transforms, nil /* lower */, nil, /* upper */ @@ -293,7 +298,7 @@ func (r *Reader) newCompactionIter( return nil, err } i.setupForCompaction() - return &compactionIterator{singleLevelIterator: i}, nil + return &compactionIterator{singleLevelIterator: i, tracked: tracked}, nil } // NewRawRangeDelIter returns an internal iterator for the contents of the diff --git a/sstable/reader_common.go b/sstable/reader_common.go index ef10cf607e1..d334638afb3 100644 --- a/sstable/reader_common.go +++ b/sstable/reader_common.go @@ -9,6 +9,7 @@ import ( "github.com/cockroachdb/pebble/internal/base" "github.com/cockroachdb/pebble/internal/keyspan" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/sstable/block" ) @@ -38,6 +39,7 @@ type CommonReader interface { statsCollector *CategoryStatsCollector, rp ReaderProvider, bufferPool *block.BufferPool, + smoother *rate.Smoother, ) (Iterator, error) EstimateDiskUsage(start, end []byte) (uint64, error) diff --git a/sstable/reader_iter.go b/sstable/reader_iter.go index 52cae69a590..1298f381821 100644 --- a/sstable/reader_iter.go +++ b/sstable/reader_iter.go @@ -11,6 +11,7 @@ import ( "github.com/cockroachdb/pebble/internal/base" "github.com/cockroachdb/pebble/internal/invariants" + "github.com/cockroachdb/pebble/internal/rate" ) // Iterator iterates over an entire table of data. @@ -154,6 +155,7 @@ func checkTwoLevelIterator(obj interface{}) { // bytes that have been iterated through. type compactionIterator struct { *singleLevelIterator + tracked rate.Tracked } // compactionIterator implements the base.InternalIterator interface. @@ -185,6 +187,11 @@ func (i *compactionIterator) First() *base.InternalKV { return i.skipForward(i.singleLevelIterator.First()) } +func (i *compactionIterator) Close() error { + i.tracked.Close() + return i.singleLevelIterator.Close() +} + func (i *compactionIterator) Last() *base.InternalKV { panic("pebble: Last unimplemented") } @@ -195,6 +202,8 @@ func (i *compactionIterator) Next() *base.InternalKV { if i.err != nil { return nil } + //bytesRead := i.stats.BlockBytes + return i.skipForward(i.data.Next()) } @@ -213,6 +222,7 @@ func (i *compactionIterator) skipForward(kv *base.InternalKV) *base.InternalKV { break } result := i.loadBlock(+1) + i.tracked.Tick() if result != loadBlockOK { if i.err != nil { break diff --git a/sstable/reader_iter_two_lvl.go b/sstable/reader_iter_two_lvl.go index cd2c955162f..3840554ad47 100644 --- a/sstable/reader_iter_two_lvl.go +++ b/sstable/reader_iter_two_lvl.go @@ -11,6 +11,7 @@ import ( "github.com/cockroachdb/pebble/internal/base" "github.com/cockroachdb/pebble/internal/invariants" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/internal/treeprinter" "github.com/cockroachdb/pebble/objstorage/objstorageprovider/objiotracing" "github.com/cockroachdb/pebble/sstable/block" @@ -969,12 +970,14 @@ func (i *twoLevelIterator) Close() error { // were separated due to performance. type twoLevelCompactionIterator struct { *twoLevelIterator + tracked rate.Tracked } // twoLevelCompactionIterator implements the base.InternalIterator interface. var _ base.InternalIterator = (*twoLevelCompactionIterator)(nil) func (i *twoLevelCompactionIterator) Close() error { + i.tracked.Close() return i.twoLevelIterator.Close() } @@ -1025,6 +1028,7 @@ func (i *twoLevelCompactionIterator) skipForward(kv *base.InternalKV) *base.Inte break } result := i.loadIndex(+1) + i.tracked.Tick() if result != loadBlockOK { if i.err != nil { break diff --git a/sstable/reader_virtual.go b/sstable/reader_virtual.go index 0788c04efa9..e8d220ba1eb 100644 --- a/sstable/reader_virtual.go +++ b/sstable/reader_virtual.go @@ -10,6 +10,7 @@ import ( "github.com/cockroachdb/pebble/internal/base" "github.com/cockroachdb/pebble/internal/keyspan" "github.com/cockroachdb/pebble/internal/rangekey" + "github.com/cockroachdb/pebble/internal/rate" "github.com/cockroachdb/pebble/sstable/block" ) @@ -96,9 +97,10 @@ func (v *VirtualReader) NewCompactionIter( statsCollector *CategoryStatsCollector, rp ReaderProvider, bufferPool *block.BufferPool, + smoother *rate.Smoother, ) (Iterator, error) { return v.reader.newCompactionIter( - transforms, categoryAndQoS, statsCollector, rp, &v.vState, bufferPool) + transforms, categoryAndQoS, statsCollector, rp, &v.vState, bufferPool, smoother) } // NewIterWithBlockPropertyFiltersAndContextEtc wraps diff --git a/table_cache.go b/table_cache.go index 3ed788961fa..9424f18c10b 100644 --- a/table_cache.go +++ b/table_cache.go @@ -590,7 +590,7 @@ func (c *tableCacheShard) newPointIter( if internalOpts.compaction { iter, err = cr.NewCompactionIter( transforms, categoryAndQoS, dbOpts.sstStatsCollector, rp, - internalOpts.bufferPool) + internalOpts.bufferPool, internalOpts.smoother) } else { iter, err = cr.NewIterWithBlockPropertyFiltersAndContextEtc( ctx, transforms, opts.GetLowerBound(), opts.GetUpperBound(), filterer, useFilter,