elastic_cpu_work_queue.go

// Copyright 2021 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package admission

import (
	"context"
	"fmt"
	"time"

	"github.com/cockroachdb/cockroach/pkg/settings"
	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
)

const (
	// MinElasticCPUDuration is the minimum on-CPU time elastic requests can ask
	// when seeking admission.
	MinElasticCPUDuration = 10 * time.Millisecond

	// MaxElasticCPUDuration is the maximum on-CPU time elastic requests can ask
	// when seeking admission.
	MaxElasticCPUDuration = 100 * time.Millisecond
)

// ValidateElasticCPUDuration is a validation function for the {minimum,maximum}
// on-CPU time elastic requests can ask for when seeking admission.
func ValidateElasticCPUDuration(duration time.Duration) error {
	if duration < MinElasticCPUDuration {
		return fmt.Errorf("minimum CPU duration allowed is %s, got %s",
			MinElasticCPUDuration, duration)
	}
	if duration > MaxElasticCPUDuration {
		return fmt.Errorf("maximum CPU duration allowed is %s, got %s",
			MaxElasticCPUDuration, duration)
	}
	return nil
}

var (
	elasticCPUControlEnabled = settings.RegisterBoolSetting(
		settings.SystemOnly,
		"admission.elastic_cpu.enabled",
		"when true, elastic work (like backups) performed by the KV layer is subject to admission control",
		true,
	)
)

// ElasticCPUWorkQueue maintains a queue of elastic work waiting to be admitted.
type ElasticCPUWorkQueue struct {
	settings  *cluster.Settings
	workQueue elasticCPUInternalWorkQueue
	granter   granter
	metrics   *elasticCPUGranterMetrics

	testingEnabled bool
}

// elasticCPUInternalWorkQueue abstracts *WorkQueue for testing.
type elasticCPUInternalWorkQueue interface {
	requester
	Admit(ctx context.Context, info WorkInfo) (enabled bool, err error)
	SetTenantWeights(tenantWeights map[uint64]uint32)
}

func makeElasticCPUWorkQueue(
	settings *cluster.Settings,
	workQueue elasticCPUInternalWorkQueue,
	granter granter,
	metrics *elasticCPUGranterMetrics,
) *ElasticCPUWorkQueue {
	return &ElasticCPUWorkQueue{
		settings:  settings,
		workQueue: workQueue,
		granter:   granter,
		metrics:   metrics,
	}
}

// Admit is called when requesting admission for elastic CPU work.
func (e *ElasticCPUWorkQueue) Admit(
	ctx context.Context, duration time.Duration, info WorkInfo,
) (*ElasticCPUWorkHandle, error) {
	if !e.enabled() {
		return nil, nil
	}
	if duration < MinElasticCPUDuration {
		duration = MinElasticCPUDuration
	}
	if duration > MaxElasticCPUDuration {
		duration = MaxElasticCPUDuration
	}
	info.RequestedCount = duration.Nanoseconds()
	enabled, err := e.workQueue.Admit(ctx, info)
	if err != nil {
		return nil, err
	}
	if !enabled {
		return nil, nil
	}
	e.metrics.AcquiredNanos.Inc(duration.Nanoseconds())
	return newElasticCPUWorkHandle(duration), nil
}

// AdmittedWorkDone indicates to the queue that the admitted work has
// completed.
func (e *ElasticCPUWorkQueue) AdmittedWorkDone(h *ElasticCPUWorkHandle) {
	if h == nil {
		return // nothing to do
	}

	e.metrics.PreWorkNanos.Inc(h.preWork.Nanoseconds())
	_, difference := h.OverLimit()
	if difference > 0 {
		// We've used up our allotted slice, which we've already deducted tokens
		// for. But we've gone over by difference, which we now need to deduct
		// tokens for.
		e.granter.tookWithoutPermission(difference.Nanoseconds())
		e.metrics.AcquiredNanos.Inc(difference.Nanoseconds())
		e.metrics.OverLimitDuration.RecordValue(difference.Nanoseconds())
		return
	}

	e.granter.returnGrant(-difference.Nanoseconds())
	e.metrics.ReturnedNanos.Inc(-difference.Nanoseconds())
}

// SetTenantWeights passes through to WorkQueue.SetTenantWeights.
func (e *ElasticCPUWorkQueue) SetTenantWeights(tenantWeights map[uint64]uint32) {
	e.workQueue.SetTenantWeights(tenantWeights)
}

func (e *ElasticCPUWorkQueue) enabled() bool {
	if e.testingEnabled {
		return true
	}

	return elasticCPUControlEnabled.Get(&e.settings.SV)
}

func (e *ElasticCPUWorkQueue) close() {
	e.workQueue.close()
}