From 3884b28d378ea0d8c36496258dbc7ce06d2a7234 Mon Sep 17 00:00:00 2001
From: Obliviate <756541536@qq.com>
Date: Sat, 10 Dec 2022 20:03:41 +0800
Subject: [PATCH] dumpling: meter network usage (#39774)

close pingcap/tidb#39674
---
 dumpling/export/BUILD.bazel |  1 +
 dumpling/export/config.go   |  7 +++++++
 dumpling/export/dump.go     | 31 +++++++++++++++++++++++++++++++
 dumpling/tests/basic/run.sh | 11 +++++++++++
 4 files changed, 50 insertions(+)

diff --git a/dumpling/export/BUILD.bazel b/dumpling/export/BUILD.bazel
index c9aebfb577c8a..08cc7fe2e664e 100644
--- a/dumpling/export/BUILD.bazel
+++ b/dumpling/export/BUILD.bazel
@@ -49,6 +49,7 @@ go_library(
         "@com_github_coreos_go_semver//semver",
         "@com_github_docker_go_units//:go-units",
         "@com_github_go_sql_driver_mysql//:mysql",
+        "@com_github_google_uuid//:uuid",
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_log//:log",
diff --git a/dumpling/export/config.go b/dumpling/export/config.go
index d724d4fb4354e..36af37b30e924 100644
--- a/dumpling/export/config.go
+++ b/dumpling/export/config.go
@@ -25,6 +25,7 @@ import (
 	filter "github.com/pingcap/tidb/util/table-filter"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/pflag"
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 )
 
@@ -144,6 +145,9 @@ type Config struct {
 	PromFactory  promutil.Factory        `json:"-"`
 	PromRegistry promutil.Registry       `json:"-"`
 	ExtStorage   storage.ExternalStorage `json:"-"`
+
+	IOTotalBytes *atomic.Uint64
+	Net          string
 }
 
 // ServerInfoUnknown is the unknown database type to dumpling
@@ -212,6 +216,9 @@ func (conf *Config) GetDriverConfig(db string) *mysql.Config {
 	driverCfg.User = conf.User
 	driverCfg.Passwd = conf.Password
 	driverCfg.Net = "tcp"
+	if conf.Net != "" {
+		driverCfg.Net = conf.Net
+	}
 	driverCfg.Addr = hostPort
 	driverCfg.DBName = db
 	driverCfg.Collation = "utf8mb4_general_ci"
diff --git a/dumpling/export/dump.go b/dumpling/export/dump.go
index 291b02b30bff9..818f91f937231 100644
--- a/dumpling/export/dump.go
+++ b/dumpling/export/dump.go
@@ -10,6 +10,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"math/big"
+	"net"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -17,6 +18,7 @@ import (
 
 	// import mysql driver
 	"github.com/go-sql-driver/mysql"
+	"github.com/google/uuid"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
 	pclog "github.com/pingcap/log"
@@ -31,8 +33,10 @@ import (
 	"github.com/pingcap/tidb/parser/format"
 	"github.com/pingcap/tidb/store/helper"
 	"github.com/pingcap/tidb/tablecodec"
+	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/codec"
 	pd "github.com/tikv/pd/client"
+	gatomic "go.uber.org/atomic"
 	"go.uber.org/zap"
 	"golang.org/x/exp/slices"
 	"golang.org/x/sync/errgroup"
@@ -101,6 +105,17 @@ func NewDumper(ctx context.Context, conf *Config) (*Dumper, error) {
 	if err != nil {
 		return nil, err
 	}
+	failpoint.Inject("SetIOTotalBytes", func(_ failpoint.Value) {
+		d.conf.IOTotalBytes = gatomic.NewUint64(0)
+		d.conf.Net = uuid.New().String()
+		go func() {
+			for {
+				time.Sleep(10 * time.Millisecond)
+				d.tctx.L().Logger.Info("IOTotalBytes", zap.Uint64("IOTotalBytes", d.conf.IOTotalBytes.Load()))
+			}
+		}()
+	})
+
 	err = runSteps(d,
 		initLogger,
 		createExternalStore,
@@ -1330,6 +1345,22 @@ func startHTTPService(d *Dumper) error {
 
 // openSQLDB is an initialization step of Dumper.
 func openSQLDB(d *Dumper) error {
+	if d.conf.IOTotalBytes != nil {
+		mysql.RegisterDialContext(d.conf.Net, func(ctx context.Context, addr string) (net.Conn, error) {
+			dial := &net.Dialer{}
+			conn, err := dial.DialContext(ctx, "tcp", addr)
+			if err != nil {
+				return nil, err
+			}
+			tcpConn := conn.(*net.TCPConn)
+			// try https://github.com/go-sql-driver/mysql/blob/bcc459a906419e2890a50fc2c99ea6dd927a88f2/connector.go#L56-L64
+			err = tcpConn.SetKeepAlive(true)
+			if err != nil {
+				d.tctx.L().Logger.Warn("fail to keep alive", zap.Error(err))
+			}
+			return util.NewTCPConnWithIOCounter(tcpConn, d.conf.IOTotalBytes), nil
+		})
+	}
 	conf := d.conf
 	c, err := mysql.NewConnector(conf.GetDriverConfig(""))
 	if err != nil {
diff --git a/dumpling/tests/basic/run.sh b/dumpling/tests/basic/run.sh
index 51e33de43fbcb..5b549cf857e79 100644
--- a/dumpling/tests/basic/run.sh
+++ b/dumpling/tests/basic/run.sh
@@ -128,3 +128,14 @@ run_dumpling --consistency lock -B "$DB_NAME" -L ${DUMPLING_OUTPUT_DIR}/dumpling
 cnt=$(grep -w "$DB_NAME" ${DUMPLING_OUTPUT_DIR}/${DB_NAME}-schema-create.sql|wc -l)
 echo "records count is ${cnt}"
 [ "$cnt" = 1 ]
+
+# Test for recording network usage
+run_sql "drop database if exists test_db;"
+run_sql "create database test_db;"
+run_sql "create table test_db.test_table (a int primary key);"
+run_sql "insert into test_db.test_table values (1),(2),(3),(4),(5),(6),(7),(8);"
+
+export GO_FAILPOINTS="github.com/pingcap/tidb/dumpling/export/SetIOTotalBytes=return(1)"
+run_dumpling -B "test_db" -L ${DUMPLING_OUTPUT_DIR}/dumpling.log
+cnt=$(grep "IOTotalBytes=" ${DUMPLING_OUTPUT_DIR}/dumpling.log | grep -v "IOTotalBytes=0" | wc -l)
+[ "$cnt" -ge 1 ]
\ No newline at end of file