dxf: separate metered traffic into object storage and cluster and fix missing cluster traffic for import-into (#64523)

close #61702

Author: D3Hunter

br/pkg/storage/BUILD.bazel

@@ -111,6 +111,7 @@ go_test(
     shard_count = 50,
     deps = [
         "//br/pkg/mock",
+        "//br/pkg/storage/recording",
         "//pkg/util/intest",
         "@com_github_aws_aws_sdk_go//aws",
         "@com_github_aws_aws_sdk_go//aws/awserr",

br/pkg/storage/azblob.go

@@ -714,7 +714,7 @@ func (s *AzureBlobStorage) Create(_ context.Context, name string, _ *WriterOptio
 		cpkInfo:  s.cpkInfo,
 	}
 
-	uploaderWriter := newBufferedWriter(uploader, azblobChunkSize, NoCompression)
+	uploaderWriter := newBufferedWriter(uploader, azblobChunkSize, NoCompression, nil)
 	return uploaderWriter, nil
 }
 

br/pkg/storage/compress.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/pingcap/errors"
 	berrors "github.com/pingcap/tidb/br/pkg/errors"
+	"github.com/pingcap/tidb/br/pkg/storage/recording"
 )
 
 type withCompression struct {
@@ -38,7 +39,8 @@ func (w *withCompression) Create(ctx context.Context, name string, o *WriterOpti
 	if bw, ok := writer.(*bufferedWriter); ok {
 		writer = bw.writer
 	}
-	compressedWriter := newBufferedWriter(writer, hardcodedS3ChunkSize, w.compressType)
+	// the external storage will do access recording, so no need to pass it again.
+	compressedWriter := newBufferedWriter(writer, hardcodedS3ChunkSize, w.compressType, nil)
 	return compressedWriter, nil
 }
 
@@ -148,13 +150,15 @@ func (c *compressReader) GetFileSize() (int64, error) {
 }
 
 type flushStorageWriter struct {
-	writer  io.Writer
-	flusher flusher
-	closer  io.Closer
+	writer    io.Writer
+	flusher   flusher
+	closer    io.Closer
+	accessRec *recording.AccessStats
 }
 
 func (w *flushStorageWriter) Write(_ context.Context, data []byte) (int, error) {
 	n, err := w.writer.Write(data)
+	w.accessRec.RecWrite(n)
 	return n, errors.Trace(err)
 }
 
@@ -166,10 +170,11 @@ func (w *flushStorageWriter) Close(_ context.Context) error {
 	return w.closer.Close()
 }
 
-func newFlushStorageWriter(writer io.Writer, flusher2 flusher, closer io.Closer) *flushStorageWriter {
+func newFlushStorageWriter(writer io.Writer, flusher2 flusher, closer io.Closer, accessRec *recording.AccessStats) *flushStorageWriter {
 	return &flushStorageWriter{
-		writer:  writer,
-		flusher: flusher2,
-		closer:  closer,
+		writer:    writer,
+		flusher:   flusher2,
+		closer:    closer,
+		accessRec: accessRec,
 	}
 }

br/pkg/storage/gcs.go

@@ -117,6 +117,7 @@ type GCSStorage struct {
 	handles      []*storage.BucketHandle
 	clients      []*storage.Client
 	clientCancel context.CancelFunc
+	accessRec    *recording.AccessStats
 }
 
 // CopyFrom implements Copier.
@@ -205,6 +206,7 @@ func (s *GCSStorage) WriteFile(ctx context.Context, name string, data []byte) er
 	if err != nil {
 		return errors.Trace(err)
 	}
+	s.accessRec.RecWrite(len(data))
 	return wc.Close()
 }
 
@@ -228,6 +230,7 @@ func (s *GCSStorage) ReadFile(ctx context.Context, name string) ([]byte, error)
 		b = make([]byte, size)
 		_, err = io.ReadFull(rc, b)
 	}
+	s.accessRec.RecRead(len(b))
 	return b, errors.Trace(err)
 }
 
@@ -345,7 +348,7 @@ func (s *GCSStorage) Create(ctx context.Context, name string, wo *WriterOption)
 		wc := s.GetBucketHandle().Object(object).NewWriter(ctx)
 		wc.StorageClass = s.gcs.StorageClass
 		wc.PredefinedACL = s.gcs.PredefinedAcl
-		return newFlushStorageWriter(wc, &emptyFlusher{}, wc), nil
+		return newFlushStorageWriter(wc, &emptyFlusher{}, wc, s.accessRec), nil
 	}
 	uri := s.objectName(name)
 	// 5MB is the minimum part size for GCS.
@@ -354,8 +357,9 @@ func (s *GCSStorage) Create(ctx context.Context, name string, wo *WriterOption)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	fw := newFlushStorageWriter(w, &emptyFlusher{}, w)
-	bw := newBufferedWriter(fw, int(partSize), NoCompression)
+	fw := newFlushStorageWriter(w, &emptyFlusher{}, w, s.accessRec)
+	// we already pass the accessRec to flushStorageWriter.
+	bw := newBufferedWriter(fw, int(partSize), NoCompression, nil)
 	return bw, nil
 }
 
@@ -423,14 +427,14 @@ skipHandleCred:
 	}
 
 	httpClient := opts.HTTPClient
-	if reqRec := recording.GetRequests(ctx); reqRec != nil {
+	if opts.AccessRecording != nil {
 		if httpClient == nil {
 			transport, _ := http.DefaultTransport.(*http.Transport)
 			httpClient = &http.Client{Transport: transport.Clone()}
 		}
 		httpClient.Transport = &roundTripperWrapper{
 			RoundTripper: httpClient.Transport,
-			requestsRec:  reqRec,
+			accessRec:    opts.AccessRecording,
 		}
 	}
 	if httpClient != nil {
@@ -460,6 +464,7 @@ skipHandleCred:
 		idx:       atomic.NewInt64(0),
 		clientCnt: gcsClientCnt,
 		clientOps: clientOps,
+		accessRec: opts.AccessRecording,
 	}
 	if err := ret.Reset(ctx); err != nil {
 		return nil, errors.Trace(err)
@@ -632,6 +637,7 @@ func (r *gcsObjectReader) Read(p []byte) (n int, err error) {
 		}
 	}
 	n, err = r.reader.Read(p)
+	r.storage.accessRec.RecRead(n)
 	r.pos += int64(n)
 	return n, err
 }
@@ -700,10 +706,10 @@ func (r *gcsObjectReader) GetFileSize() (int64, error) {
 
 type roundTripperWrapper struct {
 	http.RoundTripper
-	requestsRec *recording.Requests
+	accessRec *recording.AccessStats
 }
 
 func (rt *roundTripperWrapper) RoundTrip(req *http.Request) (*http.Response, error) {
-	rt.requestsRec.Rec(req)
+	rt.accessRec.RecRequest(req)
 	return rt.RoundTripper.RoundTrip(req)
 }

br/pkg/storage/gcs_test.go

@@ -19,12 +19,22 @@ import (
 
 	"github.com/fsouza/fake-gcs-server/fakestorage"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
+	"github.com/pingcap/tidb/br/pkg/storage/recording"
 	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"
 )
 
-func TestGCS(t *testing.T) {
+func CheckAccessStats(t *testing.T, rec *recording.AccessStats, expectedGets, expectedPuts, expectedRead, expectWrite int) {
+	t.Helper()
+	require.EqualValues(t, expectedGets, rec.Requests.Get.Load())
+	require.EqualValues(t, expectedPuts, rec.Requests.Put.Load())
+	require.EqualValues(t, expectedRead, rec.Traffic.Read.Load())
+	require.EqualValues(t, expectWrite, rec.Traffic.Write.Load())
+}
+
+func prepareGCSStore(t *testing.T, bucketName string, accessRec *recording.AccessStats) (*fakestorage.Server, *GCSStorage) {
+	t.Helper()
 	require.True(t, intest.InTest)
 	ctx := context.Background()
 
@@ -33,7 +43,6 @@ func TestGCS(t *testing.T) {
 	}
 	server, err := fakestorage.NewServerWithOptions(opts)
 	require.NoError(t, err)
-	bucketName := "testbucket"
 	server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: bucketName})
 
 	gcs := &backuppb.GCS{
@@ -47,10 +56,19 @@ func TestGCS(t *testing.T) {
 		SendCredentials:  false,
 		CheckPermissions: []Permission{AccessBuckets},
 		HTTPClient:       server.HTTPClient(),
+		AccessRecording:  accessRec,
 	})
 	require.NoError(t, err)
+	return server, stg
+}
+
+func TestGCS(t *testing.T) {
+	require.True(t, intest.InTest)
+	ctx := context.Background()
+	bucketName := "testbucket"
+	server, stg := prepareGCSStore(t, bucketName, nil)
 
-	err = stg.WriteFile(ctx, "key", []byte("data"))
+	err := stg.WriteFile(ctx, "key", []byte("data"))
 	require.NoError(t, err)
 
 	err = stg.WriteFile(ctx, "key1", []byte("data1"))
@@ -629,3 +647,29 @@ func TestDeleteFiles(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, stg.DeleteFiles(ctx, []string{filename, "not-exist-file"}))
 }
+
+func TestGCSAccessRecording(t *testing.T) {
+	ctx := context.Background()
+	accessRec := &recording.AccessStats{}
+	_, store := prepareGCSStore(t, "testbucket", accessRec)
+	require.NoError(t, store.WriteFile(ctx, "a.txt", []byte("hello")))
+	CheckAccessStats(t, accessRec, 0, 1, 0, 5)
+	_, err := store.ReadFile(ctx, "a.txt")
+	require.NoError(t, err)
+	CheckAccessStats(t, accessRec, 1, 1, 5, 5)
+	writer, err := store.Create(ctx, "b.txt", nil)
+	require.NoError(t, err)
+	_, err = writer.Write(ctx, []byte(" world!"))
+	require.NoError(t, err)
+	require.NoError(t, writer.Close(ctx))
+	CheckAccessStats(t, accessRec, 1, 2, 5, 12)
+	reader, err := store.Open(ctx, "b.txt", nil)
+	require.NoError(t, err)
+	buf := make([]byte, 20)
+	n, err := reader.Read(buf)
+	require.NoError(t, err)
+	require.Equal(t, 7, n)
+	require.NoError(t, reader.Close())
+	// Open will use 2 get, one for get file size.
+	CheckAccessStats(t, accessRec, 3, 2, 12, 12)
+}

br/pkg/storage/ks3.go

@@ -52,8 +52,9 @@ const (
 
 // KS3Storage acts almost same as S3Storage except it's used for kingsoft s3.
 type KS3Storage struct {
-	svc     *s3.S3 // https://github.com/ks3sdklib/aws-sdk-go/issues/28
-	options *backuppb.S3
+	svc       *s3.S3 // https://github.com/ks3sdklib/aws-sdk-go/issues/28
+	options   *backuppb.S3
+	accessRec *recording.AccessStats
 }
 
 // NewKS3Storage initialize a new s3 storage for metadata.
@@ -106,11 +107,11 @@ func NewKS3Storage(
 	}
 	c := s3.New(awsConfig)
 
-	if reqRec := recording.GetRequests(ctx); reqRec != nil {
+	if opts.AccessRecording != nil {
 		// unlike AWS SDK, ks3 only support change handlers after we initialize
 		// the client, so no need to call defaults.Handlers().
 		c.Handlers.Send.PushBack(func(r *aws.Request) {
-			reqRec.Rec(r.HTTPRequest)
+			opts.AccessRecording.RecRequest(r.HTTPRequest)
 		})
 	}
 
@@ -126,8 +127,9 @@ func NewKS3Storage(
 	}
 
 	return &KS3Storage{
-		svc:     c,
-		options: &qs,
+		svc:       c,
+		options:   &qs,
+		accessRec: opts.AccessRecording,
 	}, nil
 }
 
@@ -281,6 +283,7 @@ func (rs *KS3Storage) WriteFile(ctx context.Context, file string, data []byte) e
 	// since aws-go-sdk already did it in #computeBodyHashes
 	// https://github.com/aws/aws-sdk-go/blob/bcb2cf3fc2263c8c28b3119b07d2dbb44d7c93a0/service/s3/body_hash.go#L30
 	_, err := rs.svc.PutObjectWithContext(ctx, input)
+	rs.accessRec.RecWrite(len(data))
 	return errors.Trace(err)
 }
 
@@ -301,6 +304,7 @@ func (rs *KS3Storage) ReadFile(ctx context.Context, file string) ([]byte, error)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
+	rs.accessRec.RecRead(len(data))
 	return data, nil
 }
 
@@ -590,6 +594,7 @@ func (r *ks3ObjectReader) Read(p []byte) (n int, err error) {
 		n, err = r.reader.Read(p[:maxCnt])
 	}
 
+	r.storage.accessRec.RecRead(n)
 	r.pos += int64(n)
 	return
 }
@@ -737,7 +742,7 @@ func (rs *KS3Storage) Create(ctx context.Context, name string, option *WriterOpt
 	if option != nil && option.PartSize > 0 {
 		bufSize = int(option.PartSize)
 	}
-	uploaderWriter := newBufferedWriter(uploader, bufSize, NoCompression)
+	uploaderWriter := newBufferedWriter(uploader, bufSize, NoCompression, rs.accessRec)
 	return uploaderWriter, nil
 }
 

br/pkg/storage/local.go

@@ -273,7 +273,7 @@ func (l *LocalStorage) Create(_ context.Context, name string, _ *WriterOption) (
 		return nil, errors.Trace(err)
 	}
 	buf := bufio.NewWriter(file)
-	return newFlushStorageWriter(buf, buf, file), nil
+	return newFlushStorageWriter(buf, buf, file, nil), nil
 }
 
 // Rename implements ExternalStorage interface.

br/pkg/storage/recording/BUILD.bazel

@@ -13,6 +13,5 @@ go_test(
     srcs = ["recording_test.go"],
     embed = [":recording"],
     flaky = True,
-    shard_count = 2,
     deps = ["@com_github_stretchr_testify//require"],
 )