pebble: separate compaction concurrency override logic

Separate the logic that allows temporarily overriding compaction
concurrency.

Epic: none
Release note: None

Author: RaduBerinde

pkg/kv/kvserver/stores_server.go

@@ -199,11 +199,11 @@ func (is Server) SetCompactionConcurrency(
 	resp := &CompactionConcurrencyResponse{}
 	err := is.execStoreCommand(ctx, req.StoreRequestHeader,
 		func(ctx context.Context, s *Store) error {
-			prevConcurrency := s.TODOEngine().SetCompactionConcurrency(req.CompactionConcurrency)
+			s.TODOEngine().SetCompactionConcurrency(req.CompactionConcurrency)
 
 			// Wait for cancellation, and once cancelled, reset the compaction concurrency.
 			<-ctx.Done()
-			s.TODOEngine().SetCompactionConcurrency(prevConcurrency)
+			s.TODOEngine().SetCompactionConcurrency(0)
 			return nil
 		})
 	return resp, err

pkg/storage/engine.go

@@ -1102,9 +1102,9 @@ type Engine interface {
 	// version that it must maintain compatibility with.
 	SetMinVersion(version roachpb.Version) error
 
-	// SetCompactionConcurrency is used to set the engine's compaction
-	// concurrency. It returns the previous compaction concurrency.
-	SetCompactionConcurrency(n uint64) uint64
+	// SetCompactionConcurrency is used to override the engine's max compaction
+	// concurrency. A value of 0 removes any existing override.
+	SetCompactionConcurrency(n uint64)
 
 	// SetStoreID informs the engine of the store ID, once it is known.
 	// Used to show the store ID in logs and to initialize the shared object

pkg/storage/pebble.go

@@ -364,7 +364,7 @@ func DefaultPebbleOptions() *pebble.Options {
 		L0StopWritesThreshold: 1000,
 		LBaseMaxBytes:         64 << 20, // 64 MB
 		Levels:                make([]pebble.LevelOptions, 7),
-		// NB: Options.MaxConcurrentCompactions may be overidden in NewPebble to
+		// NB: Options.MaxConcurrentCompactions may be "wrapped" in NewPebble to
 		// allow overriding the max at runtime through
 		// Engine.SetCompactionConcurrency.
 		MaxConcurrentCompactions:    getMaxConcurrentCompactions,
@@ -501,24 +501,15 @@ type engineConfig struct {
 
 // Pebble is a wrapper around a Pebble database instance.
 type Pebble struct {
-	atomic struct {
-		// compactionConcurrency is the current compaction concurrency set on
-		// the Pebble store. The compactionConcurrency option in the Pebble
-		// Options struct is a closure which will return
-		// Pebble.atomic.compactionConcurrency.
-		//
-		// This mechanism allows us to change the Pebble compactionConcurrency
-		// on the fly without restarting Pebble.
-		compactionConcurrency uint64
-	}
-
 	cfg         engineConfig
 	db          *pebble.DB
 	closed      bool
 	auxDir      string
 	ballastPath string
 	properties  roachpb.StoreProperties
 
+	cco compactionConcurrencyOverride
+
 	// Stats updated by pebble.EventListener invocations, and returned in
 	// GetMetrics. Updated and retrieved atomically.
 	writeStallCount                  int64
@@ -575,10 +566,10 @@ var _ Engine = &Pebble{}
 // WorkloadCollectorEnabled specifies if the workload collector will be enabled
 var WorkloadCollectorEnabled = envutil.EnvOrDefaultBool("COCKROACH_STORAGE_WORKLOAD_COLLECTOR", false)
 
-// SetCompactionConcurrency will return the previous compaction concurrency.
-func (p *Pebble) SetCompactionConcurrency(n uint64) uint64 {
-	prevConcurrency := atomic.SwapUint64(&p.atomic.compactionConcurrency, n)
-	return prevConcurrency
+// SetCompactionConcurrency is used to override the engine's max compaction
+// concurrency. A value of 0 removes any existing override.
+func (p *Pebble) SetCompactionConcurrency(n uint64) {
+	p.cco.Set(uint32(n))
 }
 
 // RegisterDiskSlowCallback registers a callback that will be run when a write
@@ -692,6 +683,8 @@ func newPebble(ctx context.Context, cfg engineConfig) (p *Pebble, err error) {
 		}
 	}
 
+	cfg.opts.MaxConcurrentCompactions = getMaxConcurrentCompactions
+
 	cfg.opts.EnsureDefaults()
 
 	// The context dance here is done so that we have a clean context without
@@ -798,15 +791,9 @@ func newPebble(ctx context.Context, cfg engineConfig) (p *Pebble, err error) {
 
 	// MaxConcurrentCompactions can be set by multiple sources, but all the
 	// sources will eventually call NewPebble. So, we override
-	// cfg.opts.MaxConcurrentCompactions to a closure which will return
-	// Pebble.atomic.compactionConcurrency. This will allow us to both honor
-	// the compactions concurrency which has already been set and allow us
-	// to update the compactionConcurrency on the fly by changing the
-	// Pebble.atomic.compactionConcurrency variable.
-	p.atomic.compactionConcurrency = uint64(cfg.opts.MaxConcurrentCompactions())
-	cfg.opts.MaxConcurrentCompactions = func() int {
-		return int(atomic.LoadUint64(&p.atomic.compactionConcurrency))
-	}
+	// cfg.opts.MaxConcurrentCompactions to a closure which allows ovderriding the
+	// value.
+	cfg.opts.MaxConcurrentCompactions = p.cco.Wrap(cfg.opts.MaxConcurrentCompactions)
 
 	// NB: The ordering of the event listeners passed to TeeEventListener is
 	// deliberate. The listener returned by makeMetricEtcEventListener is
@@ -2750,3 +2737,24 @@ var _ error = &ExceedMaxSizeError{}
 func (e *ExceedMaxSizeError) Error() string {
 	return fmt.Sprintf("export size (%d bytes) exceeds max size (%d bytes)", e.reached, e.maxSize)
 }
+
+// compactionConcurrencyOverride allows overriding the max concurrent
+// compactions.
+type compactionConcurrencyOverride struct {
+	override atomic.Uint32
+}
+
+// Set the override value. If the value is 0, the override is cancelled.
+func (cco *compactionConcurrencyOverride) Set(value uint32) {
+	cco.override.Store(value)
+}
+
+// Wrap a MaxConcurrentCompactions function to take into account the override.
+func (cco *compactionConcurrencyOverride) Wrap(maxConcurrentCompactions func() int) func() int {
+	return func() int {
+		if o := cco.override.Load(); o > 0 {
+			return int(o)
+		}
+		return maxConcurrentCompactions()
+	}
+}