// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Copyright 2023 TiKV Project 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,g
// 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 controller

import (
	"context"
	"fmt"
	"runtime"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	promtestutil "github.com/prometheus/client_golang/prometheus/testutil"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	"github.com/pingcap/errors"
	"github.com/pingcap/failpoint"
	"github.com/pingcap/kvproto/pkg/meta_storagepb"
	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
	pd "github.com/tikv/pd/client"
	"github.com/tikv/pd/client/errs"
)

func createTestGroupCostController(re *require.Assertions) *groupCostController {
	return createTestGroupCostControllerWithName(re, "test")
}

func createTestGroupCostControllerWithName(re *require.Assertions, name string) *groupCostController {
	return createTestGroupCostControllerWithConfig(re, name, DefaultRUConfig())
}

func createTestGroupCostControllerWithConfig(re *require.Assertions, name string, cfg *RUConfig) *groupCostController {
	return createTestGroupCostControllerWithClientID(re, name, cfg, testMetricClientUniqueID)
}

func createTestGroupCostControllerWithClientID(re *require.Assertions, name string, cfg *RUConfig, clientUniqueID uint64) *groupCostController {
	group := &rmpb.ResourceGroup{
		Name:     name,
		Mode:     rmpb.GroupMode_RUMode,
		Priority: 1,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate: 1000,
				},
			},
		},
		BackgroundSettings: &rmpb.BackgroundSettings{
			JobTypes: []string{"lightning", "br"},
		},
	}
	ch1 := make(chan notifyMsg, 1)
	ch2 := make(chan *groupCostController)
	gc, err := newGroupCostController(group, cfg, ch1, ch2, newGroupMetricIdentity(name, name, clientUniqueID))
	re.NoError(err)
	gc.bindMetrics()
	return gc
}

const testMetricClientUniqueID uint64 = 1

func testTokenLabelValues(gc *groupCostController, tokenType string) []string {
	return gc.metrics.tokenLabelValues(tokenType)
}

func testDirectionLabelValues(gc *groupCostController, direction string) []string {
	return gc.metrics.directionLabelValues(direction)
}

func testBaseLabelValues(gc *groupCostController) []string {
	return gc.metrics.baseLabelValues()
}

func waitForNonZeroProcessCPUTime(re *require.Assertions) {
	deadline := time.Now().Add(2 * time.Second)
	var sink uint64
	for getSQLProcessCPUTime(true) <= 0 && time.Now().Before(deadline) {
		for i := range 100000 {
			sink += uint64(i)
		}
	}
	runtime.KeepAlive(sink)
	re.Greater(getSQLProcessCPUTime(true), 0.0)
}

func hasMetricWithLabels(collector prometheus.Collector, expected prometheus.Labels) bool {
	ch := make(chan prometheus.Metric, 16)
	collector.Collect(ch)
	close(ch)
	for metric := range ch {
		var m dto.Metric
		if err := metric.Write(&m); err != nil {
			continue
		}
		matched := 0
		for _, label := range m.GetLabel() {
			expectedValue, ok := expected[label.GetName()]
			if ok && expectedValue == label.GetValue() {
				matched++
			}
		}
		if matched == len(expected) {
			return true
		}
	}
	return false
}

func hasClientInfoMetric(clientUniqueID, resourceGroupName string) bool {
	return hasMetricWithLabels(clientInfoGauge, prometheus.Labels{
		clientUniqueIDLabel:       clientUniqueID,
		newResourceGroupNameLabel: resourceGroupName,
	})
}

func hasClientTokenDemandMetric(gc *groupCostController, tokenType string) bool {
	return hasMetricWithLabels(clientTokenDemandGauge, prometheus.Labels{
		resourceGroupNameLabel:    gc.metricIdentity.oldName,
		newResourceGroupNameLabel: gc.metricIdentity.resourceGroup,
		clientUniqueIDLabel:       gc.metricIdentity.clientUniqueID,
		errType:                   tokenType,
	})
}

func hasTokenRequestDurationMetric(clientUniqueID, tokenRequestType string) bool {
	return hasMetricWithLabels(tokenRequestDuration, prometheus.Labels{
		clientUniqueIDLabel: clientUniqueID,
		errType:             tokenRequestType,
	})
}

func isClosed(ch <-chan struct{}) bool {
	select {
	case <-ch:
		return true
	default:
		return false
	}
}

func TestClientTokenMetricsIncludeClientUniqueID(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)

	gc.metrics.recordTokenDemand(tokenType, 11)

	re.Equal(11.0, promtestutil.ToFloat64(clientTokenDemandGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))
}

func TestLowTokenNotificationMetricIncludesClientUniqueIDAndType(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	before := promtestutil.ToFloat64(lowTokenRequestNotifyCounter.WithLabelValues(testTokenLabelValues(gc, tokenType)...))

	counter.limiter.mu.Lock()
	counter.limiter.notify()
	counter.limiter.mu.Unlock()

	re.Equal(before+1, promtestutil.ToFloat64(lowTokenRequestNotifyCounter.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))
}

func TestHandleTokenBucketUpdateEventCanceledByInitCounterNotify(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]

	counter.notify.mu.Lock()
	counter.notify.setupNotificationTimer = time.NewTimer(time.Hour)
	counter.notify.setupNotificationCh = counter.notify.setupNotificationTimer.C
	counter.notify.setupNotificationThreshold = 1
	counter.notify.cancelCh = make(chan struct{})
	counter.notify.mu.Unlock()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan struct{})
	go func() {
		gc.handleTokenBucketUpdateEvent(ctx)
		close(done)
	}()

	re.Never(func() bool {
		return isClosed(done)
	}, 50*time.Millisecond, 5*time.Millisecond)

	initCounterNotify(counter)
	re.Eventually(func() bool {
		return isClosed(done)
	}, time.Second, 10*time.Millisecond)
}

func TestHandleTokenCounterUpdateEventCanceledCounterAllowsCallerToContinue(t *testing.T) {
	re := require.New(t)
	counter := createTestGroupCostControllerWithName(re, t.Name()).run.requestUnitTokens[rmpb.RequestUnitType_RU]
	cancelCh := make(chan struct{})

	counter.notify.mu.Lock()
	counter.notify.setupNotificationTimer = time.NewTimer(time.Hour)
	counter.notify.setupNotificationCh = counter.notify.setupNotificationTimer.C
	counter.notify.setupNotificationThreshold = 1
	counter.notify.cancelCh = cancelCh
	counter.notify.mu.Unlock()
	close(cancelCh)

	re.True(handleTokenCounterUpdateEvent(context.Background(), counter))
}

func TestHandleTokenBucketUpdateEventCleansNotifyOnTimer(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	threshold := 42.0

	counter.notify.mu.Lock()
	counter.notify.setupNotificationTimer = time.NewTimer(10 * time.Millisecond)
	counter.notify.setupNotificationCh = counter.notify.setupNotificationTimer.C
	counter.notify.setupNotificationThreshold = threshold
	counter.notify.cancelCh = make(chan struct{})
	counter.notify.mu.Unlock()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan struct{})
	go func() {
		gc.handleTokenBucketUpdateEvent(ctx)
		close(done)
	}()

	re.Eventually(func() bool {
		return isClosed(done)
	}, time.Second, 10*time.Millisecond)

	counter.notify.mu.Lock()
	re.Nil(counter.notify.setupNotificationTimer)
	re.Nil(counter.notify.setupNotificationCh)
	re.Nil(counter.notify.cancelCh)
	counter.notify.mu.Unlock()

	counter.limiter.mu.Lock()
	re.Equal(threshold, counter.limiter.notifyThreshold)
	counter.limiter.mu.Unlock()
}

func TestClientTokenDemandGrantAndStateMetrics(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	counter.avgRUPerSec = 1000
	gc.run.requestInProgress = true

	req := gc.collectRequestAndConsumption(periodicReport)
	re.NotNil(req)
	re.Equal(promtestutil.ToFloat64(clientTokenAvgRUPerSecGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...)), counter.avgRUPerSec)

	beforeGrant := promtestutil.ToFloat64(clientTokenGrantedCounter.WithLabelValues(testTokenLabelValues(gc, tokenType)...))
	gc.modifyTokenCounter(counter, tokenType, &rmpb.TokenBucket{
		Tokens: 250,
		Settings: &rmpb.TokenLimitSettings{
			FillRate:   100,
			BurstLimit: 1000,
		},
	}, 0)

	re.Equal(beforeGrant+250, promtestutil.ToFloat64(clientTokenGrantedCounter.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))
	re.Equal(100.0, promtestutil.ToFloat64(clientTokenFillRateGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))
	re.Equal(1000.0, promtestutil.ToFloat64(clientTokenBurstLimitGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))
	re.Equal(0.0, promtestutil.ToFloat64(clientTokenThrottledGauge.WithLabelValues(testBaseLabelValues(gc)...)))
}

func TestClientTokenGrantMetricRecordsRawGrant(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	gc.run.now = time.Now()
	counter.lastDeadline = gc.run.now.Add(10 * time.Second)
	counter.lastRate = 100
	beforeGrant := promtestutil.ToFloat64(clientTokenGrantedCounter.WithLabelValues(testTokenLabelValues(gc, tokenType)...))

	gc.modifyTokenCounter(counter, tokenType, &rmpb.TokenBucket{
		Tokens: 250,
		Settings: &rmpb.TokenLimitSettings{
			FillRate:   100,
			BurstLimit: 1000,
		},
	}, 0)

	re.Equal(beforeGrant+250, promtestutil.ToFloat64(clientTokenGrantedCounter.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))
}

func TestClientTokenDemandMetricSamplesPreThrottlingConsumption(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	gc.mainCfg.WaitRetryTimes = 1
	gc.mainCfg.WaitRetryInterval = 0
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	now := time.Now()
	counter.limiter.Reconfigure(now, tokenBucketReconfigureArgs{
		NewTokens: 0,
		NewRate:   1,
		NewBurst:  1,
	})
	counter.avgLastTime = now.Add(-time.Second)
	counter.demandAvgLastTime = now.Add(-time.Second)
	gc.run.now = now
	before := promtestutil.ToFloat64(clientTokenDemandGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...))

	_, _, _, _, err := gc.onRequestWaitImpl(context.Background(), NewTestRequestInfo(true, 64*1024, 1))
	re.Error(err)
	re.True(errs.ErrClientResourceGroupThrottled.Equal(err))
	re.Equal(before, promtestutil.ToFloat64(clientTokenDemandGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...)))

	gc.updateAvgRequestResourcePerSec()
	re.Greater(promtestutil.ToFloat64(clientTokenDemandGauge.WithLabelValues(testTokenLabelValues(gc, tokenType)...)), before)
}

func TestInitializesSQLCPUConsumptionBaseline(t *testing.T) {
	re := require.New(t)
	waitForNonZeroProcessCPUTime(re)
	cfg := DefaultRUConfig()
	cfg.isSingleGroupByKeyspace = true
	gc := createTestGroupCostControllerWithConfig(re, t.Name(), cfg)

	gc.mu.Lock()
	consumption := *gc.mu.consumption
	gc.mu.Unlock()

	re.Greater(consumption.SqlLayerCpuTimeMs, 0.0)
	re.Equal(consumption.SqlLayerCpuTimeMs, consumption.TotalCpuTimeMs)
	re.Equal(consumption, *gc.run.consumption)
	re.Equal(consumption, *gc.run.lastRequestConsumption)
}

func TestCleanUpResourceGroupDeletesTombstoneMetricsByResourceGroupName(t *testing.T) {
	re := require.New(t)
	const clientUniqueID = 1042
	customGroupName := t.Name() + "-custom"
	defaultGC := createTestGroupCostControllerWithClientID(re, defaultResourceGroupName, DefaultRUConfig(), clientUniqueID)
	customGC := createTestGroupCostControllerWithClientID(re, customGroupName, DefaultRUConfig(), clientUniqueID)
	controller := &ResourceGroupsController{
		clientUniqueID:        clientUniqueID,
		ruConfig:              DefaultRUConfig(),
		lowTokenNotifyChan:    make(chan notifyMsg, 1),
		tokenBucketUpdateChan: make(chan *groupCostController, maxNotificationChanLen),
	}
	controller.groupsController.Store(defaultResourceGroupName, defaultGC)
	controller.groupsController.Store(customGroupName, customGC)
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)
	defaultGC.metrics.recordTokenDemand(tokenType, 1)
	customGC.metrics.recordTokenDemand(tokenType, 1)
	clientInfoGauge.WithLabelValues(formatClientUniqueID(clientUniqueID), defaultResourceGroupName).Set(1)
	clientInfoGauge.WithLabelValues(formatClientUniqueID(clientUniqueID), customGroupName).Set(1)
	t.Cleanup(func() {
		defaultGC.metrics.deleteLabelValues()
		customGC.metrics.deleteLabelValues()
		clientInfoGauge.DeleteLabelValues(formatClientUniqueID(clientUniqueID), defaultResourceGroupName)
		clientInfoGauge.DeleteLabelValues(formatClientUniqueID(clientUniqueID), customGroupName)
		resourceGroupStatusGauge.DeleteLabelValues(customGroupName, customGroupName, formatClientUniqueID(clientUniqueID))
	})

	controller.tombstoneGroupCostController(customGroupName)
	controller.cleanUpResourceGroup()

	re.True(clientTokenDemandGauge.DeleteLabelValues(defaultResourceGroupName, defaultResourceGroupName, formatClientUniqueID(clientUniqueID), tokenType))
	re.False(clientTokenDemandGauge.DeleteLabelValues(customGroupName, customGroupName, formatClientUniqueID(clientUniqueID), tokenType))
	re.True(clientInfoGauge.DeleteLabelValues(formatClientUniqueID(clientUniqueID), defaultResourceGroupName))
	re.False(clientInfoGauge.DeleteLabelValues(formatClientUniqueID(clientUniqueID), customGroupName))
}

func TestStopDeletesClientInfoMetrics(t *testing.T) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	const clientUniqueID = 2042
	mockProvider := newMockResourceGroupProvider()
	controller, err := NewResourceGroupController(ctx, clientUniqueID, mockProvider, nil, 0)
	re.NoError(err)
	controller.Start(ctx)
	defer func() {
		_ = controller.Stop()
	}()

	resourceGroup := &rmpb.ResourceGroup{
		Name: defaultResourceGroupName,
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate: 1000000,
				},
			},
		},
	}
	mockProvider.On("GetResourceGroup", mock.Anything, defaultResourceGroupName, mock.Anything).Return(resourceGroup, nil)
	gc, err := controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, false)
	re.NoError(err)
	tokenType := requestUnitMetricLabel(rmpb.RequestUnitType_RU)
	gc.metrics.recordTokenDemand(tokenType, 9)
	re.True(hasClientInfoMetric(formatClientUniqueID(clientUniqueID), defaultResourceGroupName))
	re.True(hasClientTokenDemandMetric(gc, tokenType))

	re.NoError(controller.Stop())
	re.Eventually(func() bool {
		return !hasClientInfoMetric(formatClientUniqueID(clientUniqueID), defaultResourceGroupName) &&
			!hasClientTokenDemandMetric(gc, tokenType)
	}, time.Second, 10*time.Millisecond)
}

func TestCanceledTokenRequestDoesNotRecreateMetricsAfterCleanup(t *testing.T) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	const clientUniqueID = 3042
	mockProvider := newMockResourceGroupProvider()
	controller, err := NewResourceGroupController(ctx, clientUniqueID, mockProvider, nil, 0)
	re.NoError(err)

	requestStarted := make(chan struct{})
	requestReturned := make(chan struct{})
	mockProvider.On("AcquireTokenBuckets", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
		close(requestStarted)
		<-args.Get(0).(context.Context).Done()
		close(requestReturned)
	}).Return([]*rmpb.TokenBucketResponse(nil), context.Canceled).Once()

	controller.sendTokenBucketRequests(ctx, []*rmpb.TokenBucketRequest{{ResourceGroupName: defaultResourceGroupName}}, FromPeriodReport, notifyMsg{})
	<-requestStarted
	cancel()
	controller.deleteClientMetrics()
	<-requestReturned

	clientUniqueIDLabelValue := formatClientUniqueID(clientUniqueID)
	re.Never(func() bool {
		return hasTokenRequestDurationMetric(clientUniqueIDLabelValue, "fail")
	}, 100*time.Millisecond, 10*time.Millisecond)
}

func TestClientConsumptionMetricsRecordChargeAndRefund(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())
	gc.burstable.Store(true)

	req := NewTestRequestInfo(true, 64*1024, 1)
	resp := NewTestResponseInfo(16*1024, 4*time.Millisecond, false)

	writeRUChargeBefore := promtestutil.ToFloat64(clientWriteRequestUnitCost.WithLabelValues(testDirectionLabelValues(gc, chargeDirection)...))
	writeRURefundBefore := promtestutil.ToFloat64(clientWriteRequestUnitCost.WithLabelValues(testDirectionLabelValues(gc, refundDirection)...))
	writeByteChargeBefore := promtestutil.ToFloat64(clientWriteByteCost.WithLabelValues(testDirectionLabelValues(gc, chargeDirection)...))
	writeByteRefundBefore := promtestutil.ToFloat64(clientWriteByteCost.WithLabelValues(testDirectionLabelValues(gc, refundDirection)...))
	readRUBefore := promtestutil.ToFloat64(clientReadRequestUnitCost.WithLabelValues(testBaseLabelValues(gc)...))
	readByteBefore := promtestutil.ToFloat64(clientReadByteCost.WithLabelValues(testBaseLabelValues(gc)...))

	requestDelta, _, _, _, err := gc.onRequestWaitImpl(context.Background(), req)
	re.NoError(err)
	re.Positive(requestDelta.WRU)
	re.Positive(requestDelta.WriteBytes)

	responseDelta, err := gc.onResponseImpl(req, resp)
	re.NoError(err)
	re.Negative(responseDelta.WRU)
	re.Negative(responseDelta.WriteBytes)
	re.Positive(responseDelta.RRU)
	re.Positive(responseDelta.ReadBytes)

	re.Equal(writeRUChargeBefore+requestDelta.WRU, promtestutil.ToFloat64(clientWriteRequestUnitCost.WithLabelValues(testDirectionLabelValues(gc, chargeDirection)...)))
	re.Equal(writeRURefundBefore-responseDelta.WRU, promtestutil.ToFloat64(clientWriteRequestUnitCost.WithLabelValues(testDirectionLabelValues(gc, refundDirection)...)))
	re.Equal(writeByteChargeBefore+requestDelta.WriteBytes, promtestutil.ToFloat64(clientWriteByteCost.WithLabelValues(testDirectionLabelValues(gc, chargeDirection)...)))
	re.Equal(writeByteRefundBefore-responseDelta.WriteBytes, promtestutil.ToFloat64(clientWriteByteCost.WithLabelValues(testDirectionLabelValues(gc, refundDirection)...)))
	re.Equal(readRUBefore+responseDelta.RRU, promtestutil.ToFloat64(clientReadRequestUnitCost.WithLabelValues(testBaseLabelValues(gc)...)))
	re.Equal(readByteBefore+responseDelta.ReadBytes, promtestutil.ToFloat64(clientReadByteCost.WithLabelValues(testBaseLabelValues(gc)...)))
}

func TestAddRUConsumptionIgnoresNilConsumption(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostControllerWithName(re, t.Name())

	writeRUChargeBefore := promtestutil.ToFloat64(clientWriteRequestUnitCost.WithLabelValues(testDirectionLabelValues(gc, chargeDirection)...))
	re.NotPanics(func() {
		gc.addRUConsumption(nil)
	})
	re.Equal(writeRUChargeBefore, promtestutil.ToFloat64(clientWriteRequestUnitCost.WithLabelValues(testDirectionLabelValues(gc, chargeDirection)...)))
}

func TestGroupControlBurstable(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostController(re)
	args := tokenBucketReconfigureArgs{
		NewRate:  1000,
		NewBurst: -1,
	}
	for _, counter := range gc.run.requestUnitTokens {
		counter.limiter.Reconfigure(time.Now(), args)
	}
	gc.updateAvgRequestResourcePerSec()
	re.True(gc.burstable.Load())
}

func TestRequestAndResponseConsumption(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostController(re)
	testCases := []struct {
		req  *TestRequestInfo
		resp *TestResponseInfo
	}{
		// Write request
		{
			req: &TestRequestInfo{
				isWrite:    true,
				writeBytes: 100,
			},
			resp: &TestResponseInfo{
				readBytes: 100,
				succeed:   true,
			},
		},
		// Read request
		{
			req: &TestRequestInfo{
				isWrite:    false,
				writeBytes: 0,
			},
			resp: &TestResponseInfo{
				readBytes: 100,
				kvCPU:     100 * time.Millisecond,
				succeed:   true,
			},
		},
	}
	kvCalculator := gc.getKVCalculator()
	for idx, testCase := range testCases {
		caseNum := fmt.Sprintf("case %d", idx)
		consumption, _, _, priority, err := gc.onRequestWaitImpl(context.TODO(), testCase.req)
		re.NoError(err, caseNum)
		re.Equal(priority, gc.meta.Priority)
		expectedConsumption := &rmpb.Consumption{}
		if testCase.req.IsWrite() {
			kvCalculator.calculateWriteCost(expectedConsumption, testCase.req)
			re.Equal(expectedConsumption.WRU, consumption.WRU)
		}
		consumption, err = gc.onResponseImpl(testCase.req, testCase.resp)
		re.NoError(err, caseNum)
		kvCalculator.calculateReadCost(expectedConsumption, testCase.resp)
		kvCalculator.calculateCPUCost(expectedConsumption, testCase.resp)
		re.Equal(expectedConsumption.RRU, consumption.RRU, caseNum)
		re.Equal(expectedConsumption.TotalCpuTimeMs, consumption.TotalCpuTimeMs, caseNum)
	}
}

func TestOnResponseWaitConsumption(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostController(re)

	req := &TestRequestInfo{
		isWrite: false,
	}
	resp := &TestResponseInfo{
		readBytes: 2000 * 64 * 1024, // 2000RU
		succeed:   true,
	}

	consumption, waitTIme, err := gc.onResponseWaitImpl(context.TODO(), req, resp)
	re.NoError(err)
	re.Zero(waitTIme)
	verify := func() {
		expectedConsumption := &rmpb.Consumption{}
		kvCalculator := gc.getKVCalculator()
		kvCalculator.calculateReadCost(expectedConsumption, resp)
		re.Equal(expectedConsumption.RRU, consumption.RRU)
	}
	verify()

	// modify the counter, then on response should has wait time.
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	gc.modifyTokenCounter(counter, requestUnitMetricLabel(rmpb.RequestUnitType_RU), &rmpb.TokenBucket{
		Settings: &rmpb.TokenLimitSettings{
			FillRate:   1000,
			BurstLimit: 1000,
		},
	},
		int64(5*time.Second/time.Millisecond),
	)

	consumption, waitTIme, err = gc.onResponseWaitImpl(context.TODO(), req, resp)
	re.NoError(err)
	re.NotZero(waitTIme)
	verify()
}

func TestResourceGroupThrottledError(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostController(re)
	req := &TestRequestInfo{
		isWrite:    true,
		writeBytes: 10000000,
	}
	// The group is throttled
	_, _, _, _, err := gc.onRequestWaitImpl(context.TODO(), req)
	re.Error(err)
	re.True(errs.ErrClientResourceGroupThrottled.Equal(err))
}

// MockResourceGroupProvider is a mock implementation of the ResourceGroupProvider interface.
type MockResourceGroupProvider struct {
	mock.Mock
}

func newMockResourceGroupProvider() *MockResourceGroupProvider {
	mockProvider := &MockResourceGroupProvider{}
	mockProvider.On("Get", mock.Anything, mock.Anything, mock.Anything).Return(&meta_storagepb.GetResponse{}, nil)
	mockProvider.On("LoadResourceGroups", mock.Anything).Return([]*rmpb.ResourceGroup{}, int64(0), nil)
	mockProvider.On("Watch", mock.Anything, mock.Anything, mock.Anything).Return(make(chan *pd.WatchResponse), nil)
	return mockProvider
}

func (m *MockResourceGroupProvider) GetResourceGroup(ctx context.Context, resourceGroupName string, opts ...pd.GetResourceGroupOption) (*rmpb.ResourceGroup, error) {
	var err error
	failpoint.Inject("gerResourceGroupError", func() {
		err = errors.New("fake get resource group error")
	})
	if err != nil {
		return nil, &errs.ErrClientGetResourceGroup{ResourceGroupName: resourceGroupName, Cause: err.Error()}
	}

	args := m.Called(ctx, resourceGroupName, opts)
	return args.Get(0).(*rmpb.ResourceGroup), args.Error(1)
}

func (m *MockResourceGroupProvider) ListResourceGroups(ctx context.Context, opts ...pd.GetResourceGroupOption) ([]*rmpb.ResourceGroup, error) {
	args := m.Called(ctx, opts)
	return args.Get(0).([]*rmpb.ResourceGroup), args.Error(1)
}

func (m *MockResourceGroupProvider) AddResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
	args := m.Called(ctx, metaGroup)
	return args.String(0), args.Error(1)
}

func (m *MockResourceGroupProvider) ModifyResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
	args := m.Called(ctx, metaGroup)
	return args.String(0), args.Error(1)
}

func (m *MockResourceGroupProvider) DeleteResourceGroup(ctx context.Context, resourceGroupName string) (string, error) {
	args := m.Called(ctx, resourceGroupName)
	return args.String(0), args.Error(1)
}

func (m *MockResourceGroupProvider) AcquireTokenBuckets(ctx context.Context, request *rmpb.TokenBucketsRequest) ([]*rmpb.TokenBucketResponse, error) {
	args := m.Called(ctx, request)
	return args.Get(0).([]*rmpb.TokenBucketResponse), args.Error(1)
}

func (m *MockResourceGroupProvider) LoadResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, int64, error) {
	args := m.Called(ctx)
	return args.Get(0).([]*rmpb.ResourceGroup), args.Get(1).(int64), args.Error(2)
}

func (m *MockResourceGroupProvider) Watch(ctx context.Context, key []byte, opts ...pd.OpOption) (chan *pd.WatchResponse, error) {
	args := m.Called(ctx, key, opts)
	return args.Get(0).(chan *pd.WatchResponse), args.Error(1)
}

func (m *MockResourceGroupProvider) Get(ctx context.Context, key []byte, opts ...pd.OpOption) (*meta_storagepb.GetResponse, error) {
	args := m.Called(ctx, key, opts)
	return args.Get(0).(*meta_storagepb.GetResponse), args.Error(1)
}

func TestControllerReportConsumption(t *testing.T) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	mockProvider := newMockResourceGroupProvider()
	controller, _ := NewResourceGroupController(ctx, 1, mockProvider, nil, 0)
	controller.Start(ctx)

	testResourceGroup := &rmpb.ResourceGroup{
		Name: "test-group",
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate: 1000000,
				},
			},
		},
	}
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).Return(testResourceGroup, nil)

	gc, err := controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.Equal(testResourceGroup, gc.meta)

	var totalConsumption rmpb.Consumption
	gc.mu.Lock()
	totalConsumption = *gc.mu.consumption
	gc.mu.Unlock()

	delta := &rmpb.Consumption{
		RRU:               1.0,
		WRU:               2.0,
		ReadBytes:         10,
		WriteBytes:        20,
		TotalCpuTimeMs:    30.0,
		SqlLayerCpuTimeMs: 40.0,
		KvReadRpcCount:    50,
		KvWriteRpcCount:   60,
	}
	controller.ReportConsumption("test-group", delta)

	gc.mu.Lock()
	add(&totalConsumption, delta)
	require.Equal(t, gc.mu.consumption, &totalConsumption)
	gc.mu.Unlock()

	controller.ReportRUV2Consumption("test-group", 3.0, 4.0, 5.0)
	gc.mu.Lock()
	totalConsumption.TikvRUV2 += 3.0
	totalConsumption.TidbRUV2 += 4.0
	totalConsumption.TiflashRUV2 += 5.0
	require.Equal(t, gc.mu.consumption, &totalConsumption)
	gc.mu.Unlock()

	controller.ReportConsumption("unknown-name", delta)
	controller.ReportRUV2Consumption("unknown-name", 1.0, 1.0, 1.0)
}

func TestControllerWithTwoGroupRequestConcurrency(t *testing.T) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	re.NoError(failpoint.Enable("github.com/tikv/pd/client/resource_group/controller/triggerLowRUReport", fmt.Sprintf("return(\"%s\")", "test-group")))
	defer failpoint.Disable("github.com/tikv/pd/client/resource_group/controller/triggerLowRUReport")

	mockProvider := newMockResourceGroupProvider()
	controller, _ := NewResourceGroupController(ctx, 1, mockProvider, nil, 0)
	controller.Start(ctx)

	defaultResourceGroup := &rmpb.ResourceGroup{Name: defaultResourceGroupName, Mode: rmpb.GroupMode_RUMode, RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}}}
	testResourceGroup := &rmpb.ResourceGroup{Name: "test-group", Mode: rmpb.GroupMode_RUMode, RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}}}
	mockProvider.On("GetResourceGroup", mock.Anything, defaultResourceGroupName, mock.Anything).Return(defaultResourceGroup, nil)
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).Return(testResourceGroup, nil)

	c1, err := controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, false)
	re.NoError(err)
	re.Equal(defaultResourceGroup, c1.meta)

	c2, err := controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.Equal(testResourceGroup, c2.meta)

	var expectResp []*rmpb.TokenBucketResponse
	recTestGroupAcquireTokenRequest := make(chan bool)
	mockProvider.On("AcquireTokenBuckets", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
		request := args.Get(1).(*rmpb.TokenBucketsRequest)
		var responses []*rmpb.TokenBucketResponse
		for _, req := range request.Requests {
			responses = append(responses, &rmpb.TokenBucketResponse{
				ResourceGroupName: req.ResourceGroupName,
				GrantedRUTokens: []*rmpb.GrantedRUTokenBucket{
					{
						GrantedTokens: &rmpb.TokenBucket{
							Tokens: 100000,
						},
					},
				},
			})
		}
		// receive test-group request
		if len(request.Requests) == 1 && request.Requests[0].ResourceGroupName == "test-group" {
			recTestGroupAcquireTokenRequest <- true
		}
		expectResp = responses
	}).Return(expectResp, nil)
	// Simulate an unresolved token request for another resource group.
	controller.run.currentRequests = []*rmpb.TokenBucketRequest{{ResourceGroupName: defaultResourceGroupName}}
	counter := c2.run.requestUnitTokens[0]
	counter.limiter.mu.Lock()
	counter.limiter.notify()
	counter.limiter.mu.Unlock()
	select {
	case res := <-recTestGroupAcquireTokenRequest:
		re.True(res)
	case <-time.After(5 * time.Second):
		re.Fail("timeout")
	}
}

func TestTryGetController(t *testing.T) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	mockProvider := newMockResourceGroupProvider()
	controller, _ := NewResourceGroupController(ctx, 1, mockProvider, nil, 0)
	controller.Start(ctx)

	defaultResourceGroup := &rmpb.ResourceGroup{Name: defaultResourceGroupName, Mode: rmpb.GroupMode_RUMode, RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}}}
	testResourceGroup := &rmpb.ResourceGroup{Name: "test-group", Mode: rmpb.GroupMode_RUMode, RUSettings: &rmpb.GroupRequestUnitSettings{RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}}}}
	mockProvider.On("GetResourceGroup", mock.Anything, defaultResourceGroupName, mock.Anything).Return(defaultResourceGroup, nil)
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).Return(testResourceGroup, nil)
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group-non-existent", mock.Anything).Return((*rmpb.ResourceGroup)(nil), nil)

	gc, err := controller.tryGetResourceGroupController(ctx, "test-group-non-existent", false)
	re.Error(err)
	re.Nil(gc)
	gc, err = controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, false)
	re.NoError(err)
	re.Equal(defaultResourceGroup, gc.getMeta())
	gc, err = controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.Equal(testResourceGroup, gc.getMeta())
	requestInfo, responseInfo := NewTestRequestInfo(true, 1, 1), NewTestResponseInfo(1, time.Millisecond, true)
	_, _, _, _, err = controller.OnRequestWait(ctx, "test-group", requestInfo)
	re.NoError(err)
	consumption, err := controller.OnResponse("test-group", requestInfo, responseInfo)
	re.NoError(err)
	re.NotEmpty(consumption)
	// Mark the tombstone manually to test the fallback case.
	gc, err = controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.NotNil(gc)
	controller.tombstoneGroupCostController("test-group")
	gc, err = controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.Error(err)
	re.Nil(gc)
	gc, err = controller.tryGetResourceGroupController(ctx, "test-group", true)
	re.NoError(err)
	re.Equal(defaultResourceGroup, gc.getMeta())
	_, _, _, _, err = controller.OnRequestWait(ctx, "test-group", requestInfo)
	re.NoError(err)
	consumption, err = controller.OnResponse("test-group", requestInfo, responseInfo)
	re.NoError(err)
	re.NotEmpty(consumption)
	// Test the default group protection.
	gc, err = controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, false)
	re.NoError(err)
	re.Equal(defaultResourceGroup, gc.getMeta())
	controller.tombstoneGroupCostController(defaultResourceGroupName)
	gc, err = controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, false)
	re.NoError(err)
	re.Equal(defaultResourceGroup, gc.getMeta())
	gc, err = controller.tryGetResourceGroupController(ctx, defaultResourceGroupName, true)
	re.NoError(err)
	re.Equal(defaultResourceGroup, gc.getMeta())
	_, _, _, _, err = controller.OnRequestWait(ctx, defaultResourceGroupName, requestInfo)
	re.NoError(err)
	consumption, err = controller.OnResponse(defaultResourceGroupName, requestInfo, responseInfo)
	re.NoError(err)
	re.NotEmpty(consumption)
}

func TestGetResourceGroup(t *testing.T) {
	re := require.New(t)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	mockProvider := newMockResourceGroupProvider()

	expectResourceGroupName := "test-group"
	expectRUSettings := &rmpb.GroupRequestUnitSettings{
		RU: &rmpb.TokenBucket{
			Settings: &rmpb.TokenLimitSettings{
				FillRate:   uint64(50),
				BurstLimit: int64(100),
			},
		},
	}
	expectResourceGroup := &rmpb.ResourceGroup{
		Name:       expectResourceGroupName,
		Mode:       rmpb.GroupMode_RUMode,
		RUSettings: expectRUSettings,
	}

	opts := []ResourceControlCreateOption{WithDegradedRUSettings(expectRUSettings)}

	controller, _ := NewResourceGroupController(ctx, 1, mockProvider, nil, 0, opts...)
	controller.Start(ctx)

	testResourceGroup := &rmpb.ResourceGroup{
		Name: "test-group",
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{
					FillRate: 1000000,
				},
			},
		},
	}
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).
		Return((*rmpb.ResourceGroup)(nil), errors.New("fake get resource group error")).
		Twice()

	// case1: when GetResourceGroup return error, it should return the expectResourceGroup.
	gc, err := controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.Equal(expectResourceGroup, gc.getMeta())

	// case2: when GetResourceGroup return error again, It should still return the expectResourceGroup.
	gc, err = controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.Equal(expectResourceGroup, gc.getMeta())

	// case3: If `GetResourceGroup` returns no error (`nil`), it should return the testResourceGroup.
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).Return(testResourceGroup, nil).Once()
	gc, err = controller.tryGetResourceGroupController(ctx, "test-group", false)
	re.NoError(err)
	re.Equal(testResourceGroup, gc.getMeta())

	// case4: when we don't set degradedRUSettings and GetResourceGroup return error, tryGetResourceGroupController will return err.
	mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).
		Return((*rmpb.ResourceGroup)(nil), errors.New("fake get resource group error")).
		Once()

	controller02, _ := NewResourceGroupController(ctx, 1, mockProvider, nil, 0)
	controller02.Start(ctx)

	gc02, err := controller02.tryGetResourceGroupController(ctx, "test-group", false)
	re.Error(err)
	re.Nil(gc02)
}

func TestAcquireTokensSignalAwareWait(t *testing.T) {
	re := require.New(t)

	// Build controller with a buffered lowRUNotifyChan so the test can
	// observe the notify() call inside reserveN as a synchronization point.
	group := &rmpb.ResourceGroup{
		Name: "test",
		Mode: rmpb.GroupMode_RUMode,
		RUSettings: &rmpb.GroupRequestUnitSettings{
			RU: &rmpb.TokenBucket{
				Settings: &rmpb.TokenLimitSettings{FillRate: 1000},
			},
		},
	}
	notifyCh := make(chan notifyMsg, 1)
	cfg := DefaultRUConfig()
	cfg.WaitRetryInterval = 5 * time.Second
	cfg.WaitRetryTimes = 3
	gc, err := newGroupCostController(group, cfg, notifyCh, make(chan *groupCostController, 1), newGroupMetricIdentity(group.Name, group.Name, testMetricClientUniqueID))
	re.NoError(err)
	gc.bindMetrics()

	// Set fillRate=0 so reservation always fails with InfDuration,
	// which is the exact scenario described in issue #10251.
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	counter.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{
		NewTokens: 1000,
		NewRate:   0,
		NewBurst:  0,
	})

	delta := &rmpb.Consumption{RRU: 5000}
	type acquireResult struct {
		err          error
		waitDuration time.Duration
	}
	resultCh := make(chan acquireResult, 1)
	go func() {
		var waitDuration time.Duration
		_, err := gc.acquireTokens(context.Background(), delta, &waitDuration, false)
		resultCh <- acquireResult{err, waitDuration}
	}()

	// Wait for notify — Reserve has failed and the retry path is entered.
	select {
	case <-notifyCh:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for low-RU notification")
	}

	// Now reconfigure with enough tokens and a real fillRate.
	// This closes the reconfiguredCh (waking the select) and provides
	// tokens so the next Reserve() succeeds.
	counter.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{
		NewTokens: 100000,
		NewRate:   100000,
		NewBurst:  0,
	})

	select {
	case r := <-resultCh:
		re.NoError(r.err)
		re.Less(r.waitDuration, cfg.WaitRetryInterval)
	case <-time.After(cfg.WaitRetryInterval):
		t.Fatal("acquireTokens was not woken up promptly by Reconfigure signal")
	}
}

func TestAcquireTokensFallbackToTimer(t *testing.T) {
	re := require.New(t)
	gc := createTestGroupCostController(re)
	// Short retry interval so the test runs fast.
	gc.mainCfg.WaitRetryInterval = 50 * time.Millisecond
	gc.mainCfg.WaitRetryTimes = 3
	gc.mainCfg.LTBMaxWaitDuration = 100 * time.Millisecond

	// Set fillRate=0 and never reconfigure — no signal will arrive.
	counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
	counter.limiter.Reconfigure(time.Now(), tokenBucketReconfigureArgs{
		NewTokens: 1000,
		NewRate:   0,
		NewBurst:  0,
	})

	delta := &rmpb.Consumption{RRU: 5000}
	ctx := context.Background()
	var waitDuration time.Duration
	_, err := gc.acquireTokens(ctx, delta, &waitDuration, false)

	// Without a Reconfigure signal, all retries should exhaust and return an error.
	re.Error(err)
	re.True(errs.ErrClientResourceGroupThrottled.Equal(err))
	// waitDuration should be roughly retryTimes * retryInterval.
	re.GreaterOrEqual(waitDuration, gc.mainCfg.WaitRetryInterval*time.Duration(gc.mainCfg.WaitRetryTimes))
}