Repository: pingcap/tidb
File: pkg/disttask/framework/metering/metering_test.go (Go)

Code before:
		t.Skip("metering is a feature of nextgen")
	}
	meter, reader := createLocalMeter(t, t.TempDir())
	setupMeterForTest(t, meter)
	curTime := time.Now()

	writeTime := curTime.Truncate(time.Minute).Add(time.Minute)
	data := readMeteringData(t, reader, writeTime.Add(-time.Minute).Unix())
	require.Len(t, data, 0)
	recorder := RegisterRecorder(&proto.TaskBase{ID: 1, Keyspace: "ks1"})
	recorder.objStoreAccess.Requests.Get.Add(10)
	recorder.objStoreAccess.Requests.Put.Add(20)
	recorder.objStoreAccess.RecRead(11)
	recorder.objStoreAccess.RecWrite(22)
	recorder.IncClusterReadBytes(300)
	recorder.IncClusterWriteBytes(400)
	meter.flush(context.Background(), writeTime.Unix())
	data = readMeteringData(t, reader, writeTime.Add(-time.Minute).Unix())
	require.Len(t, data, 0)
	data = readMeteringData(t, reader, writeTime.Unix())
	require.Len(t, data, 1)
	require.Equal(t, "1", data[0]["version"])
	require.Equal(t, "ks1", data[0]["cluster_id"])
	require.Equal(t, "dxf", data[0]["source_name"])
	require.Equal(t, float64(10), data[0][getRequestsField].(float64))
	require.Equal(t, float64(20), data[0][putRequestsField].(float64))
	require.Equal(t, float64(11), data[0][objStoreReadBytesField].(float64))
	require.Equal(t, float64(22), data[0][objStoreWriteBytesField].(float64))
	require.Equal(t, float64(300), data[0][clusterReadBytesField].(float64))
	require.Equal(t, float64(400), data[0][clusterWriteBytesField].(float64))
}

func createLocalMeter(t *testing.T, dir string) (*Meter, *meteringreader.MeteringReader) {
	t.Helper()
	meterConfig := &mconfig.MeteringConfig{
		Type:   "s3",
		Bucket: "bucket",
		AWS: &mconfig.MeteringAWSConfig{
			AssumeRoleARN: "test-role-arn",
		},
	}
	m, err := NewMeter(meterConfig)
	require.NoError(t, err)
	t.Cleanup(func() {
		require.NoError(t, m.Close())
	})

	// Replace the S3 writer with the local writer.
	localConfig := &storage.ProviderConfig{
		Type: storage.ProviderTypeLocalFS,
		LocalFS: &storage.LocalFSConfig{

Review comment (bug):
Direct access to `meter.pendingRetryData.data[1000000].retryCnt` without holding the lock creates a data race. The test should either lock `pendingRetryData.mu` before accessing the retry count or use a getter method that provides thread-safe access to these values.

Code after:
		meter.pendingRetryData.addFailedData(2000000, []map[string]any{{"get_requests": 23}})
		meter.pendingRetryData.addFailedData(3000000, []map[string]any{{"get_requests": 34}})
		require.Len(t, meter.pendingRetryData.data, 3)
		thirdRetryCnt := 0
		mockWriter.EXPECT().Write(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, data any) error {
			md := data.(*common.MeteringData)
			require.Len(t, md.Data, 1)
			switch md.Timestamp {
			case 1000000:
				return goerrors.New("some err")
			case 2000000:
			case 3000000:
				thirdRetryCnt++
				if thirdRetryCnt < 5 {
					return goerrors.New("some err")
				}
			}
			return nil
		}).Times(16)
		for r := range 10 {
			meter.retryWrite(ctx)
			retryCnt := r + 1
			if retryCnt < 5 {
				require.Len(t, meter.pendingRetryData.data, 2)
				require.EqualValues(t, retryCnt, meter.pendingRetryData.data[1000000].retryCnt)
				require.EqualValues(t, retryCnt, meter.pendingRetryData.data[3000000].retryCnt)
			} else if retryCnt == 10 {
				// after 10 retries, the first item should be dropped
				require.Empty(t, meter.pendingRetryData.data)
			} else {
				require.Len(t, meter.pendingRetryData.data, 1)
				require.EqualValues(t, retryCnt, meter.pendingRetryData.data[1000000].retryCnt)
			}
		}
		require.True(t, ctrl.Satisfied())
	})
}

func TestMeterStartFlushLoop(t *testing.T) {
	if kerneltype.IsClassic() {
		t.Skip("metering is a feature of nextgen")
	}
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	logger := zap.Must(zap.NewDevelopment())
	runOneTest := func(t *testing.T, runRecordersFn func() map[int64]uint64, writeFn func(ctx context.Context, data any) error) map[int64]uint64 {
		ctx, cancel := context.WithCancel(context.Background())
		mockWriter := writermock.NewMockMeteringWriter(ctrl)
		mockWriter.EXPECT().Write(gomock.Any(), gomock.Any()).DoAndReturn(writeFn).AnyTimes()
		mockWriter.EXPECT().Close().Return(nil)
		meter := newMeterWithWriter(logger, mockWriter)