From 92a04ef0240e84d5db5a6e00a5b7fb8b10d4b2a1 Mon Sep 17 00:00:00 2001
From: Zejun Li <lizejun@pingcap.com>
Date: Wed, 23 Sep 2020 18:45:38 +0800
Subject: [PATCH] trace: tee logs to runtime/trace output (#19994)

---
 executor/adapter.go |  4 ++++
 server/conn.go      |  9 ++++++++-
 util/logutil/log.go | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/executor/adapter.go b/executor/adapter.go
index be9da134f5f6c..823e196e7db44 100644
--- a/executor/adapter.go
+++ b/executor/adapter.go
@@ -17,6 +17,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"runtime/trace"
 	"strconv"
 	"strings"
 	"sync/atomic"
@@ -913,6 +914,9 @@ func (a *ExecStmt) LogSlowQuery(txnTS uint64, succ bool, hasMoreResults bool) {
 	if _, ok := a.StmtNode.(*ast.CommitStmt); ok {
 		slowItems.PrevStmt = sessVars.PrevStmt.String()
 	}
+	if trace.IsEnabled() {
+		trace.Log(a.GoCtx, "details", sessVars.SlowLogFormat(slowItems))
+	}
 	if costTime < threshold {
 		logutil.SlowQueryLogger.Debug(sessVars.SlowLogFormat(slowItems))
 	} else {
diff --git a/server/conn.go b/server/conn.go
index d47286dbe0981..fe6d8abaaa52f 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -50,6 +50,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
@@ -928,8 +929,14 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error {
 		if len(sqlType) > 0 {
 			var task *trace.Task
 			ctx, task = trace.NewTask(ctx, sqlType)
-			trace.Log(ctx, "sql", lc.String())
 			defer task.End()
+
+			trace.Log(ctx, "sql", lc.String())
+			ctx = logutil.WithTraceLogger(ctx, cc.connectionID)
+
+			taskID := *(*uint64)(unsafe.Pointer(task))
+			ctx = pprof.WithLabels(ctx, pprof.Labels("trace", strconv.FormatUint(taskID, 10)))
+			pprof.SetGoroutineLabels(ctx)
 		}
 	}
 	token := cc.server.getToken()
diff --git a/util/logutil/log.go b/util/logutil/log.go
index 49429c70a965f..55bbc950c3ebd 100644
--- a/util/logutil/log.go
+++ b/util/logutil/log.go
@@ -20,6 +20,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"runtime/trace"
 	"sort"
 	"strings"
 	"time"
@@ -355,6 +356,43 @@ func WithConnID(ctx context.Context, connID uint32) context.Context {
 	return context.WithValue(ctx, ctxLogKey, logger.With(zap.Uint32("conn", connID)))
 }
 
+// WithTraceLogger attaches trace identifier to context
+func WithTraceLogger(ctx context.Context, connID uint32) context.Context {
+	var logger *zap.Logger
+	if ctxLogger, ok := ctx.Value(ctxLogKey).(*zap.Logger); ok {
+		logger = ctxLogger
+	} else {
+		logger = zaplog.L()
+	}
+	return context.WithValue(ctx, ctxLogKey, wrapTraceLogger(ctx, connID, logger))
+}
+
+func wrapTraceLogger(ctx context.Context, connID uint32, logger *zap.Logger) *zap.Logger {
+	return logger.WithOptions(zap.WrapCore(func(core zapcore.Core) zapcore.Core {
+		tl := &traceLog{ctx: ctx}
+		traceCore := zaplog.NewTextCore(zaplog.NewTextEncoder(&zaplog.Config{}), tl, tl).
+			With([]zapcore.Field{zap.Uint32("conn", connID)})
+		return zapcore.NewTee(traceCore, core)
+	}))
+}
+
+type traceLog struct {
+	ctx context.Context
+}
+
+func (t *traceLog) Enabled(_ zapcore.Level) bool {
+	return true
+}
+
+func (t *traceLog) Write(p []byte) (n int, err error) {
+	trace.Log(t.ctx, "log", string(p))
+	return len(p), nil
+}
+
+func (t *traceLog) Sync() error {
+	return nil
+}
+
 // WithKeyValue attaches key/value to context.
 func WithKeyValue(ctx context.Context, key, value string) context.Context {
 	var logger *zap.Logger