Merge pull request pingcap#9 from ti2sky/trace

add trace

executor/compiler.go

@@ -18,6 +18,8 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/SkyAPM/go2sky"
+	language_agent "github.com/SkyAPM/go2sky/reporter/grpc/language-agent"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/config"
@@ -54,6 +56,17 @@ func (c *Compiler) Compile(ctx context.Context, stmtNode ast.StmtNode) (*ExecStm
 		defer span1.Finish()
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
+	if span := go2sky.SpanFromContext(ctx); span != nil {
+		tracer := (*span).Tracer()
+		s, c, e := tracer.CreateLocalSpan(ctx)
+		if e != nil {
+			return nil, e
+		}
+		s.SetOperationName("executor.Compile")
+		s.SetSpanLayer(language_agent.SpanLayer_Database)
+		defer s.End()
+		ctx = c
+	}
 
 	ret := &plannercore.PreprocessorReturn{}
 	pe := &plannercore.PreprocessExecuteISUpdate{ExecuteInfoSchemaUpdate: planner.GetExecuteForUpdateReadIS, Node: stmtNode}

executor/executor.go

@@ -27,6 +27,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/SkyAPM/go2sky"
+	language_agent "github.com/SkyAPM/go2sky/reporter/grpc/language-agent"
 	"github.com/cznic/mathutil"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
@@ -280,6 +282,17 @@ func Next(ctx context.Context, e Executor, req *chunk.Chunk) error {
 		defer span1.Finish()
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
+	if span := go2sky.SpanFromContext(ctx); span != nil {
+		tracer := (*span).Tracer()
+		s, c, err := tracer.CreateLocalSpan(ctx)
+		if err != nil {
+			return err
+		}
+		s.SetOperationName(fmt.Sprintf("%T.Next", e))
+		s.SetSpanLayer(language_agent.SpanLayer_Database)
+		defer s.End()
+		ctx = c
+	}
 	if trace.IsEnabled() {
 		defer trace.StartRegion(ctx, fmt.Sprintf("%T.Next", e)).End()
 	}

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/DATA-DOG/go-sqlmock v1.5.0
 	github.com/HdrHistogram/hdrhistogram-go v1.1.0 // indirect
 	github.com/Jeffail/gabs/v2 v2.5.1
+	github.com/SkyAPM/go2sky v0.5.0
 	github.com/aws/aws-sdk-go v1.35.3
 	github.com/blacktear23/go-proxyprotocol v0.0.0-20180807104634-af7a81e8dd0d
 	github.com/carlmjohnson/flagext v0.21.0
@@ -31,7 +32,7 @@ require (
 	github.com/golang/snappy v0.0.3
 	github.com/google/btree v1.0.0
 	github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1
-	github.com/google/uuid v1.1.2
+	github.com/google/uuid v1.3.0
 	github.com/gorilla/handlers v1.5.1 // indirect
 	github.com/gorilla/mux v1.8.0
 	github.com/grpc-ecosystem/go-grpc-middleware v1.1.0
@@ -100,5 +101,7 @@ replace google.golang.org/grpc => google.golang.org/grpc v1.29.1
 
 replace github.com/pingcap/tidb/parser => ./parser
 
+replace github.com/SkyAPM/go2sky => github.com/ti2sky/go2sky v0.5.1-0.20211225145356-52b22ddc18cc
+
 // fix potential security issue(CVE-2020-26160) introduced by indirect dependency.
 replace github.com/dgrijalva/jwt-go => github.com/form3tech-oss/jwt-go v3.2.6-0.20210809144907-32ab6a8243d7+incompatible

go.sum

@@ -357,8 +357,8 @@ github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm4
 github.com/google/shlex v0.0.0-20181106134648-c34317bd91bf/go.mod h1:RpwtwJQFrIEPstU94h88MWPXP2ektJZ8cZ0YntAmXiE=
 github.com/google/uuid v1.0.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y=
-github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5 h1:sjZBwGj9Jlw33ImPtvFviGYvseOtDM7hkSKB7+Tv3SM=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
@@ -707,6 +707,8 @@ github.com/swaggo/swag v1.6.6-0.20200529100950-7c765ddd0476/go.mod h1:xDhTyuFIuj
 github.com/syndtr/goleveldb v1.0.1-0.20190318030020-c3a204f8e965 h1:1oFLiOyVl+W7bnBzGhf7BbIv9loSFQcieWWYIjLqcAw=
 github.com/syndtr/goleveldb v1.0.1-0.20190318030020-c3a204f8e965/go.mod h1:9OrXJhf154huy1nPWmuSrkgjPUtUNhA+Zmy+6AESzuA=
 github.com/thoas/go-funk v0.8.0/go.mod h1:+IWnUfUmFO1+WVYQWQtIJHeRRdaIyyYglZN7xzUPe4Q=
+github.com/ti2sky/go2sky v0.5.1-0.20211225145356-52b22ddc18cc h1:0mmHuDlJcEWj42oH9oEkd2RlfAK+lpO/RSg8Bdu9twI=
+github.com/ti2sky/go2sky v0.5.1-0.20211225145356-52b22ddc18cc/go.mod h1:TANzYw5EvIlTidGWvQxtvO87rM6C746HkM0xkWqnPQw=
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJfDRtkanvQPiooDH8HvJ2FBh+iKT/OmiQQ=
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
 github.com/tidwall/gjson v1.3.5/go.mod h1:P256ACg0Mn+j1RXIDXoss50DeIABTYK1PULOJHhxOls=
@@ -1028,6 +1030,7 @@ golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9sn
 golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.1-0.20181010134911-4d1c5fb19474/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=

server/conn.go

@@ -55,6 +55,9 @@ import (
 	"time"
 	"unsafe"
 
+	"github.com/SkyAPM/go2sky"
+	"github.com/SkyAPM/go2sky/reporter"
+	language_agent "github.com/SkyAPM/go2sky/reporter/grpc/language-agent"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
@@ -1005,6 +1008,8 @@ func (cc *clientConn) initConnect(ctx context.Context) error {
 	return nil
 }
 
+var tracerKey = struct{}{}
+
 // Run reads client query and writes query result to client in for loop, if there is a panic during query handling,
 // it will be recovered and log the panic error.
 // This function returns and the connection is closed if there is an IO error or there is a panic.
@@ -1030,6 +1035,21 @@ func (cc *clientConn) Run(ctx context.Context) {
 			terror.Log(err)
 		}
 	}()
+
+	swReporter, err := reporter.NewGRPCReporter("127.0.0.1:11800")
+	if err != nil {
+		terror.Log(err)
+		return
+	}
+	defer swReporter.Close()
+
+	swTracer, err := go2sky.NewTracer("tidb", go2sky.WithReporter(swReporter))
+	if err != nil {
+		terror.Log(err)
+		return
+	}
+	ctx = context.WithValue(ctx, tracerKey, swTracer)
+
 	// Usually, client connection status changes between [dispatching] <=> [reading].
 	// When some event happens, server may notify this client connection by setting
 	// the status to special values, for example: kill or graceful shutdown.
@@ -1233,6 +1253,10 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error {
 	if cfg.OpenTracing.Enable {
 		ctx = opentracing.ContextWithSpan(ctx, span)
 	}
+	var tracer *go2sky.Tracer
+	if swTracer := ctx.Value(tracerKey); swTracer != nil {
+		tracer = swTracer.(*go2sky.Tracer)
+	}
 
 	var cancelFunc context.CancelFunc
 	ctx, cancelFunc = context.WithCancel(ctx)
@@ -1294,8 +1318,33 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error {
 	case mysql.ComPing, mysql.ComStmtClose, mysql.ComStmtSendLongData, mysql.ComStmtReset,
 		mysql.ComSetOption, mysql.ComChangeUser:
 		cc.ctx.SetProcessInfo("", t, cmd, 0)
+		if tracer != nil {
+			span, nCtx, err := tracer.CreateEntrySpan(ctx, strconv.Itoa(int(cmd)), func() (string, error) {
+				return "", nil
+			})
+			if err != nil {
+				terror.Log(err)
+			} else {
+				defer span.End()
+				span.SetSpanLayer(language_agent.SpanLayer_Database)
+				ctx = nCtx
+			}
+		}
 	case mysql.ComInitDB:
 		cc.ctx.SetProcessInfo("use "+dataStr, t, cmd, 0)
+		if tracer != nil {
+			span, nCtx, err := tracer.CreateEntrySpan(ctx, "TiDB/Command/Query", func() (string, error) {
+				return "", nil
+			})
+			if err != nil {
+				terror.Log(err)
+			} else {
+				defer span.End()
+				span.SetSpanLayer(language_agent.SpanLayer_Database)
+				span.Tag("db.statement", dataStr)
+				ctx = nCtx
+			}
+		}
 	}
 
 	switch cmd {

session/session.go

@@ -32,8 +32,9 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/SkyAPM/go2sky"
+	language_agent "github.com/SkyAPM/go2sky/reporter/grpc/language-agent"
 	"github.com/google/uuid"
-
 	"github.com/ngaut/pools"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
@@ -1523,6 +1524,17 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 		defer span1.Finish()
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
+	if span := go2sky.SpanFromContext(ctx); span != nil {
+		tracer := (*span).Tracer()
+		s, c, e := tracer.CreateLocalSpan(ctx)
+		if e != nil {
+			return nil, e
+		}
+		s.SetOperationName("session.ExecuteStmt")
+		s.SetSpanLayer(language_agent.SpanLayer_Database)
+		defer s.End()
+		ctx = c
+	}
 	s.PrepareTxnCtx(ctx)
 	if err := s.loadCommonGlobalVariablesIfNeeded(); err != nil {
 		return nil, err
@@ -1666,6 +1678,17 @@ func runStmt(ctx context.Context, se *session, s sqlexec.Statement) (rs sqlexec.
 		defer span1.Finish()
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
+	if span := go2sky.SpanFromContext(ctx); span != nil {
+		tracer := (*span).Tracer()
+		s, c, e := tracer.CreateLocalSpan(ctx)
+		if e != nil {
+			return nil, e
+		}
+		s.SetOperationName("session.runStmt")
+		s.SetSpanLayer(language_agent.SpanLayer_Database)
+		defer s.End()
+		ctx = c
+	}
 	se.SetValue(sessionctx.QueryString, s.OriginText())
 	if _, ok := s.(*executor.ExecStmt).StmtNode.(ast.DDLNode); ok {
 		se.SetValue(sessionctx.LastExecuteDDL, true)