resource_control.go

// Copyright 2023 TiKV Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package resourcecontrol

import (
	"strings"
	"time"

	"github.com/pingcap/kvproto/pkg/coprocessor"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
	"github.com/tikv/client-go/v2/config"
	"github.com/tikv/client-go/v2/kv"
	"github.com/tikv/client-go/v2/tikvrpc"
	"github.com/tikv/client-go/v2/util"
	"github.com/tikv/pd/client/resource_group/controller"
)

// RequestInfo contains information about a request that is able to calculate the RU cost
// before the request is sent. Specifically, the write bytes RU cost of a write request
// could be calculated by its key size to write.
type RequestInfo struct {
	// writeBytes is the actual write size if the request is a write request,
	// or -1 if it's a read request.
	writeBytes    int64
	storeID       uint64
	replicaNumber int64
	requestSize   uint64
	accessType    controller.AccessLocationType
	// predictedReadBytes is an optional caller-supplied read-bytes
	// estimate. The PD controller decides whether to use it for paging
	// accounting based on request type; non-cop hints may be ignored.
	predictedReadBytes uint64
	// isCop is true when the underlying tikvrpc.Request targets the
	// coprocessor endpoint (CmdCop / CmdCopStream). PD uses this to
	// scope paging_* metrics to coprocessor reads only.
	isCop bool
	// bypass indicates whether the request should be bypassed.
	// some internal request should be bypassed, such as Privilege request.
	bypass bool
}

func toPDAccessLocationType(accessType kv.AccessLocationType) controller.AccessLocationType {
	switch accessType {
	case kv.AccessLocalZone:
		return controller.AccessLocalZone
	case kv.AccessCrossZone:
		return controller.AccessCrossZone
	default:
		return controller.AccessUnknown
	}
}

// reqTypeAnalyze is the type of analyze coprocessor request.
// ref: https://github.com/pingcap/tidb/blob/ee4eac2ccb83e1ea653b8131d9a43495019cb5ac/pkg/kv/kv.go#L340
const reqTypeAnalyze = 104

func shouldBypass(req *tikvrpc.Request) bool {
	requestSource := req.GetRequestSource()
	// Check both coprocessor request type and the request source to ensure the request is an internal analyze request.
	// Internal analyze request may consume a lot of resources, bypass it to avoid affecting the user experience.
	// This bypass currently only works with NextGen.
	if config.NextGen && strings.Contains(requestSource, util.InternalTxnStats) {
		var tp int64
		switch req.Type {
		case tikvrpc.CmdBatchCop:
			tp = req.BatchCop().GetTp()
		case tikvrpc.CmdCop, tikvrpc.CmdCopStream:
			tp = req.Cop().GetTp()
		}
		if tp == reqTypeAnalyze {
			return true
		}
	}
	// Some internal requests should be bypassed, which may affect the user experience.
	// For example, the `alter user password` request completely bypasses resource control.
	// Although it does not consume many resources, it can still impact the user experience.
	return strings.Contains(requestSource, util.InternalRequestPrefix+util.InternalTxnOthers)
}

// MakeRequestInfo extracts the relevant information from a BatchRequest.
func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo {
	bypass := shouldBypass(req)
	storeID := req.Context.GetPeer().GetStoreId()
	if !req.IsTxnWriteRequest() && !req.IsRawWriteRequest() {
		return &RequestInfo{
			writeBytes:         -1,
			storeID:            storeID,
			bypass:             bypass,
			requestSize:        uint64(req.GetSize()),
			accessType:         toPDAccessLocationType(req.AccessLocation),
			predictedReadBytes: req.PredictedReadBytes,
			isCop:              isCopRequest(req),
		}
	}

	var writeBytes int64
	switch r := req.Req.(type) {
	case *kvrpcpb.PrewriteRequest:
		for _, m := range r.Mutations {
			writeBytes += int64(len(m.Key)) + int64(len(m.Value))
		}
		writeBytes += int64(len(r.PrimaryLock))
		for _, l := range r.Secondaries {
			writeBytes += int64(len(l))
		}
	case *kvrpcpb.CommitRequest:
		for _, k := range r.Keys {
			writeBytes += int64(len(k))
		}
	}
	return &RequestInfo{
		writeBytes:    writeBytes,
		storeID:       storeID,
		replicaNumber: req.ReplicaNumber,
		bypass:        bypass,
		requestSize:   uint64(req.GetSize()),
		accessType:    toPDAccessLocationType(req.AccessLocation),
	}
}

// IsWrite returns whether the request is a write request.
func (req *RequestInfo) IsWrite() bool {
	return req.writeBytes > -1
}

// WriteBytes returns the actual write size of the request,
// -1 will be returned if it's not a write request.
func (req *RequestInfo) WriteBytes() uint64 {
	if req.writeBytes > 0 {
		return uint64(req.writeBytes)
	}
	return 0
}

func (req *RequestInfo) ReplicaNumber() int64 {
	return req.replicaNumber
}

// Bypass returns whether the request should be bypassed.
func (req *RequestInfo) Bypass() bool {
	return req.bypass
}

func (req *RequestInfo) StoreID() uint64 {
	return req.storeID
}

func (req *RequestInfo) RequestSize() uint64 {
	return req.requestSize
}

func (req *RequestInfo) AccessLocationType() controller.AccessLocationType {
	return req.accessType
}

// PredictedReadBytes implements PD's controller.RequestInfo interface,
// supplying the read-bytes hint. The PD controller decides whether the hint is
// eligible for paging accounting based on the request type.
func (req *RequestInfo) PredictedReadBytes() uint64 {
	return req.predictedReadBytes
}

// IsCop implements PD's controller.RequestInfo interface. It reports whether
// the underlying tikvrpc.Request targets the coprocessor endpoint, so that
// PD can scope paging_* metrics to coprocessor reads and ignore point gets,
// batch gets, scans, and other bounded-size reads that share the same RC
// interceptor path.
func (req *RequestInfo) IsCop() bool {
	return req.isCop
}

// isCopRequest reports whether req is a coprocessor read RPC.
func isCopRequest(req *tikvrpc.Request) bool {
	return req.Type == tikvrpc.CmdCop || req.Type == tikvrpc.CmdCopStream
}

// ResponseInfo contains information about a response that is able to calculate the RU cost
// after the response is received. Specifically, the read bytes RU cost of a read request
// could be calculated by its response size, and the KV CPU time RU cost of a request could
// be calculated by its execution details info.
type ResponseInfo struct {
	readBytes       uint64
	remoteReadBytes uint64
	kvCPU           time.Duration
	respSize        uint64
}

// MakeResponseInfo extracts the relevant information from a BatchResponse.
func MakeResponseInfo(resp *tikvrpc.Response) *ResponseInfo {
	if resp.Resp == nil {
		return &ResponseInfo{}
	}
	// Parse the response to extract the info.
	var (
		readBytes       uint64
		remoteReadBytes uint64
		detailsV2       *kvrpcpb.ExecDetailsV2
		details         *kvrpcpb.ExecDetails
	)
	switch r := resp.Resp.(type) {
	case *coprocessor.Response:
		detailsV2 = r.GetExecDetailsV2()
		details = r.GetExecDetails()
		readBytes = uint64(r.Data.Size())
	case *tikvrpc.CopStreamResponse:
		// Streaming request returns `io.EOF``, so the first `CopStreamResponse.Response`` may be nil.
		if r.Response != nil {
			detailsV2 = r.GetExecDetailsV2()
			details = r.GetExecDetails()
		}
		readBytes = uint64(r.Data.Size())
	case *kvrpcpb.GetResponse:
		detailsV2 = r.GetExecDetailsV2()
	case *kvrpcpb.BatchGetResponse:
		detailsV2 = r.GetExecDetailsV2()
	case *kvrpcpb.ScanResponse:
		// TODO: using a more accurate size rather than using the whole response size as the read bytes.
		readBytes = uint64(r.Size())
	default:
		return &ResponseInfo{}
	}
	// Try to get read bytes from the `detailsV2`.
	// TODO: clarify whether we should count the underlying storage engine read bytes or not.
	if scanDetail := detailsV2.GetScanDetailV2(); scanDetail != nil {
		if config.NextGen {
			// Using the total versions size as the read bytes, which includes not
			// only processed versions size, but also skipped MVCC versions size.
			// It can reflect the actual read bytes more accurately, especially for
			// the request with a large number of MVCC versions.
			//
			// For compatibility with older versions of TiKV, if the
			// processed versions size is greater than the total versions size,
			// we use the processed versions size as the read bytes.
			readBytes = max(scanDetail.GetTotalVersionsSize(), scanDetail.GetProcessedVersionsSize())
			remoteReadBytes = max(
				scanDetail.GetRemoteTotalVersionsSize(),
				scanDetail.GetRemoteProcessedVersionsSize(),
			)
		} else {
			// NOTE: The original design intended to account for all MVCC read
			// overhead, but TotalVersionsSize did not exist at the time, so
			// ProcessedVersionsSize was used instead, which only counts the
			// versions that are actually processed and excludes skipped MVCC
			// versions. Since this behavior has already been released, switching
			// to TotalVersionsSize would change the RU calculation. To avoid
			// unexpected impact on existing users, we only fix this in NextGen
			// and keep the legacy behavior here for now.
			readBytes = scanDetail.GetProcessedVersionsSize()
		}
	}
	// Get the KV CPU time in milliseconds from the execution time details.
	kvCPU := getKVCPU(detailsV2, details)
	return &ResponseInfo{
		readBytes:       readBytes,
		remoteReadBytes: remoteReadBytes,
		kvCPU:           kvCPU,
		respSize:        uint64(resp.GetSize()),
	}
}

// TODO: find out a more accurate way to get the actual KV CPU time.
func getKVCPU(detailsV2 *kvrpcpb.ExecDetailsV2, details *kvrpcpb.ExecDetails) time.Duration {
	if timeDetail := detailsV2.GetTimeDetailV2(); timeDetail != nil {
		return time.Duration(timeDetail.GetProcessWallTimeNs())
	}
	if timeDetail := detailsV2.GetTimeDetail(); timeDetail != nil {
		return time.Duration(timeDetail.GetProcessWallTimeMs()) * time.Millisecond
	}
	if timeDetail := details.GetTimeDetail(); timeDetail != nil {
		return time.Duration(timeDetail.GetProcessWallTimeMs()) * time.Millisecond
	}
	return time.Duration(0)
}

// ReadBytes returns the read bytes of the response.
func (res *ResponseInfo) ReadBytes() uint64 {
	return res.readBytes
}

// RemoteReadBytes returns the factual subset of ReadBytes processed by a
// remote coprocessor. It is zero outside NextGen so legacy RU pricing remains
// unchanged.
func (res *ResponseInfo) RemoteReadBytes() uint64 {
	return res.remoteReadBytes
}

// KVCPU returns the KV CPU time of the response.
func (res *ResponseInfo) KVCPU() time.Duration {
	return res.kvCPU
}

// Succeed returns whether the KV request is successful.
// Todo: to fit https://github.com/tikv/pd/pull/5941
func (res *ResponseInfo) Succeed() bool {
	return true
}

func (res *ResponseInfo) ResponseSize() uint64 {
	return res.respSize
}