lightning: record network usage

br/pkg/lightning/BUILD.bazel

@@ -28,7 +28,10 @@ go_library(
         "//br/pkg/version/build",
         "//expression",
         "//planner/core",
+        "//util",
         "//util/promutil",
+        "@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_kvproto//pkg/import_sstpb",

br/pkg/lightning/common/util.go

@@ -54,6 +54,7 @@ type MySQLConnectParam struct {
 	SQLMode          string
 	MaxAllowedPacket uint64
 	TLS              string
+	Net              string
 	Vars             map[string]string
 }
 
@@ -64,6 +65,9 @@ func (param *MySQLConnectParam) ToDriverConfig() *mysql.Config {
 	cfg.User = param.User
 	cfg.Passwd = param.Password
 	cfg.Net = "tcp"
+	if param.Net != "" {
+		cfg.Net = param.Net
+	}
 	cfg.Addr = net.JoinHostPort(param.Host, strconv.Itoa(param.Port))
 	cfg.Params["charset"] = "utf8mb4"
 	cfg.Params["sql_mode"] = fmt.Sprintf("'%s'", param.SQLMode)

br/pkg/lightning/config/config.go

@@ -135,6 +135,9 @@ type DBStore struct {
 	IndexSerialScanConcurrency int               `toml:"index-serial-scan-concurrency" json:"index-serial-scan-concurrency"`
 	ChecksumTableConcurrency   int               `toml:"checksum-table-concurrency" json:"checksum-table-concurrency"`
 	Vars                       map[string]string `toml:"-" json:"vars"`
+
+	IOTotalBytes *atomic.Uint64 `toml:"-" json:"-"`
+	UUID         string         `toml:"-" json:"-"`
 }
 
 type Config struct {

br/pkg/lightning/lightning.go

@@ -33,6 +33,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/go-sql-driver/mysql"
+	"github.com/google/uuid"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/import_sstpb"
@@ -53,11 +55,13 @@ import (
 	"github.com/pingcap/tidb/br/pkg/version/build"
 	_ "github.com/pingcap/tidb/expression" // get rid of `import cycle`: just init expression.RewriteAstExpr,and called at package `backend.kv`.
 	_ "github.com/pingcap/tidb/planner/core"
+	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/promutil"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/shurcooL/httpgzip"
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 	"golang.org/x/exp/slices"
@@ -370,6 +374,36 @@ func (l *Lightning) RunOnceWithOptions(taskCtx context.Context, taskCfg *config.
 		taskCfg.TaskID = int64(val.(int))
 	})
 
+	failpoint.Inject("SetIOTotalBytes", func(_ failpoint.Value) {
+		o.logger.Info("set io total bytes")
+		taskCfg.TiDB.IOTotalBytes = atomic.NewUint64(0)
+		taskCfg.TiDB.UUID = uuid.New().String()
+		go func() {
+			for {
+				time.Sleep(time.Millisecond * 10)
+				log.L().Info("IOTotalBytes", zap.Uint64("IOTotalBytes", taskCfg.TiDB.IOTotalBytes.Load()))
+			}
+		}()
+	})
+	if taskCfg.TiDB.IOTotalBytes != nil {
+		o.logger.Info("found IO total bytes counter")
+		mysql.RegisterDialContext(taskCfg.TiDB.UUID, func(ctx context.Context, addr string) (net.Conn, error) {
+			o.logger.Debug("connection with IO bytes counter")
+			d := &net.Dialer{}
+			conn, err := d.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 {
+				o.logger.Warn("set TCP keep alive failed", zap.Error(err))
+			}
+			return util.NewTCPConnWithIOCounter(tcpConn, taskCfg.TiDB.IOTotalBytes), nil
+		})
+	}
+
 	return l.run(taskCtx, taskCfg, o)
 }
 

br/pkg/lightning/restore/tidb.go

@@ -73,6 +73,7 @@ func DBFromConfig(ctx context.Context, dsn config.DBStore) (*sql.DB, error) {
 		SQLMode:          dsn.StrSQLMode,
 		MaxAllowedPacket: dsn.MaxAllowedPacket,
 		TLS:              dsn.TLS,
+		Net:              dsn.UUID,
 	}
 
 	db, err := param.Connect()

br/tests/lightning_record_network/config.toml

@@ -0,0 +1,2 @@
+[tikv-importer]
+backend = 'tidb'

br/tests/lightning_record_network/data/db-schema-create.sql

@@ -0,0 +1 @@
+create database db;

br/tests/lightning_record_network/data/db.test-schema.sql

@@ -0,0 +1 @@
+create table test ( id int primary key, a int, b int );

br/tests/lightning_record_network/data/db.test.1.sql

@@ -0,0 +1,21 @@
+insert into db.test values
+(1,1,1),
+(2,1,1),
+(3,1,1),
+(4,1,1),
+(5,1,1),
+(6,1,1),
+(7,1,1),
+(8,1,1),
+(9,1,1),
+(10,1,1),
+(11,1,1),
+(12,1,1),
+(13,1,1),
+(14,1,1),
+(15,1,1),
+(16,1,1),
+(17,1,1),
+(18,1,1),
+(19,1,1),
+(20,1,1);

br/tests/lightning_record_network/run.sh

@@ -0,0 +1,22 @@
+#!/bin/sh
+#
+# Copyright 2022 PingCAP, Inc.
+#
+# 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.
+
+set -euE
+
+export GO_FAILPOINTS="github.com/pingcap/tidb/br/pkg/lightning/SetIOTotalBytes=return(1)"
+run_lightning
+
+grep 'IOTotal' "$TEST_DIR/lightning.log" | grep -v 'IOTotalBytes=0'

util/BUILD.bazel

@@ -46,6 +46,7 @@ go_library(
         "@io_etcd_go_etcd_client_v3//:client",
         "@io_etcd_go_etcd_client_v3//concurrency",
         "@org_golang_google_grpc//:grpc",
+        "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],
 )

util/util.go

@@ -19,13 +19,15 @@ import (
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
+	"net"
 	"net/http"
 	"strconv"
 	"strings"
 	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/parser"
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 )
 
@@ -202,3 +204,30 @@ func FmtNonASCIIPrintableCharToHex(str string) string {
 	}
 	return b.String()
 }
+
+// TCPConnWithIOCounter is a wrapper of net.TCPConn with counter that accumulates
+// the bytes this connection reads/writes.
+type TCPConnWithIOCounter struct {
+	*net.TCPConn
+	c *atomic.Uint64
+}
+
+// NewTCPConnWithIOCounter creates a new TCPConnWithIOCounter.
+func NewTCPConnWithIOCounter(conn *net.TCPConn, c *atomic.Uint64) net.Conn {
+	return &TCPConnWithIOCounter{
+		TCPConn: conn,
+		c:       c,
+	}
+}
+
+func (t *TCPConnWithIOCounter) Read(b []byte) (n int, err error) {
+	n, err = t.TCPConn.Read(b)
+	t.c.Add(uint64(n))
+	return n, err
+}
+
+func (t *TCPConnWithIOCounter) Write(b []byte) (n int, err error) {
+	n, err = t.TCPConn.Write(b)
+	t.c.Add(uint64(n))
+	return n, err
+}