client-go/internal/resourcecontrol/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
	// 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.Context.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 && (req.BatchCop().GetTp() == reqTypeAnalyze || req.Cop().GetTp() == reqTypeAnalyze) &&
		strings.Contains(requestSource, util.InternalTxnStats) {
		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),
		}
	}

	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
}

// 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
	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
		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.Response.GetExecDetailsV2()
			details = r.Response.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 {
		readBytes = scanDetail.GetProcessedVersionsSize()
	}
	// Get the KV CPU time in milliseconds from the execution time details.
	kvCPU := getKVCPU(detailsV2, details)
	return &ResponseInfo{
		readBytes: readBytes,
		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
}

// 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
}