server: support download optimize trace file

domain/optimize_trace.go

@@ -0,0 +1,27 @@
+// Copyright 2021 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.
+
+package domain
+
+import (
+	"os"
+	"path/filepath"
+	"strconv"
+)
+
+// GetOptimizerTraceDirName returns optimizer trace directory path.
+// The path is related to the process id.
+func GetOptimizerTraceDirName() string {
+	return filepath.Join(os.TempDir(), "optimizer_trace", strconv.Itoa(os.Getpid()))
+}

executor/trace.go

@@ -24,12 +24,12 @@ import (
 	"os"
 	"path/filepath"
 	"sort"
-	"strconv"
 	"time"
 
 	"github.com/opentracing/basictracer-go"
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/parser/terror"
@@ -303,7 +303,7 @@ func generateLogResult(allSpans []basictracer.RawSpan, chk *chunk.Chunk) {
 }
 
 func generateOptimizerTraceFile() (*os.File, string, error) {
-	dirPath := getOptimizerTraceDirName()
+	dirPath := domain.GetOptimizerTraceDirName()
 	// Create path
 	err := os.MkdirAll(dirPath, os.ModePerm)
 	if err != nil {
@@ -324,9 +324,3 @@ func generateOptimizerTraceFile() (*os.File, string, error) {
 	}
 	return zf, fileName, nil
 }
-
-// getOptimizerTraceDirName returns optimizer trace directory path.
-// The path is related to the process id.
-func getOptimizerTraceDirName() string {
-	return filepath.Join(os.TempDir(), "optimizer_trace", strconv.Itoa(os.Getpid()))
-}

server/http_status.go

@@ -203,6 +203,8 @@ func (s *Server) startHTTPServer() {
 
 	router.Handle("/plan_replayer/dump/{filename}", s.newPlanReplayerHandler()).Name("PlanReplayerDump")
 
+	router.Handle("/optimize_trace/dump/{filename}", s.newOptimizeTraceHandler()).Name("OptimizeTraceDump")
+
 	tikvHandlerTool := s.newTikvHandlerTool()
 	router.Handle("/settings", settingsHandler{tikvHandlerTool}).Name("Settings")
 	router.Handle("/binlog/recover", binlogRecover{}).Name("BinlogRecover")

server/optimize_trace.go

@@ -0,0 +1,60 @@
+// Copyright 2021 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.
+
+package server
+
+import (
+	"fmt"
+	"net/http"
+	"path/filepath"
+
+	"github.com/gorilla/mux"
+	"github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/domain"
+	"github.com/pingcap/tidb/domain/infosync"
+)
+
+// OptimizeTraceHandler serve http
+type OptimizeTraceHandler struct {
+	infoGetter *infosync.InfoSyncer
+	address    string
+	statusPort uint
+}
+
+func (s *Server) newOptimizeTraceHandler() *OptimizeTraceHandler {
+	cfg := config.GetGlobalConfig()
+	oth := &OptimizeTraceHandler{
+		address:    cfg.AdvertiseAddress,
+		statusPort: cfg.Status.StatusPort,
+	}
+	if s.dom != nil && s.dom.InfoSyncer() != nil {
+		oth.infoGetter = s.dom.InfoSyncer()
+	}
+	return oth
+}
+
+func (oth OptimizeTraceHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	params := mux.Vars(req)
+	name := params[pFileName]
+	handler := downloadFileHandler{
+		filePath:           filepath.Join(domain.GetOptimizerTraceDirName(), name),
+		fileName:           name,
+		infoGetter:         oth.infoGetter,
+		address:            oth.address,
+		statusPort:         oth.statusPort,
+		urlPath:            fmt.Sprintf("optimize_trace/dump/%s", name),
+		downloadedFilename: "optimize_trace",
+	}
+	handleDownloadFile(handler, w, req)
+}

server/optimize_trace_test.go

@@ -0,0 +1,91 @@
+// Copyright 2021 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.
+
+package server
+
+import (
+	"database/sql"
+	"net/http"
+	"path/filepath"
+	"testing"
+
+	"github.com/gorilla/mux"
+	"github.com/pingcap/tidb/session"
+	"github.com/pingcap/tidb/testkit"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDumpOptimizeTraceAPI(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+
+	driver := NewTiDBDriver(store)
+	client := newTestServerClient()
+	cfg := newTestConfig()
+	cfg.Port = client.port
+	cfg.Status.StatusPort = client.statusPort
+	cfg.Status.ReportStatus = true
+
+	server, err := NewServer(cfg, driver)
+	require.NoError(t, err)
+	defer server.Close()
+
+	client.port = getPortFromTCPAddr(server.listener.Addr())
+	client.statusPort = getPortFromTCPAddr(server.statusListener.Addr())
+	go func() {
+		err := server.Run()
+		require.NoError(t, err)
+	}()
+	client.waitUntilServerOnline()
+
+	dom, err := session.GetDomain(store)
+	require.NoError(t, err)
+	statsHandler := &StatsHandler{dom}
+
+	otHandler := &OptimizeTraceHandler{}
+	filename := prepareData4OptimizeTrace(t, client, statsHandler)
+
+	router := mux.NewRouter()
+	router.Handle("/optimize_trace/dump/{filename}", otHandler)
+
+	resp0, err := client.fetchStatus(filepath.Join("/optimize_trace/dump/", filename))
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, resp0.Body.Close())
+	}()
+	require.Equal(t, http.StatusOK, resp0.StatusCode)
+}
+
+func prepareData4OptimizeTrace(t *testing.T, client *testServerClient, statHandle *StatsHandler) string {
+	db, err := sql.Open("mysql", client.getDSN())
+	require.NoError(t, err, "Error connecting")
+	defer func() {
+		err := db.Close()
+		require.NoError(t, err)
+	}()
+	tk := testkit.NewDBTestKit(t, db)
+
+	h := statHandle.do.StatsHandle()
+	tk.MustExec("create database optimizeTrace")
+	tk.MustExec("use optimizeTrace")
+	tk.MustExec("create table t(a int)")
+	err = h.HandleDDLEvent(<-h.DDLEventCh())
+	require.NoError(t, err)
+	rows := tk.MustQuery("trace plan select * from t")
+	require.True(t, rows.Next(), "unexpected data")
+	var filename string
+	err = rows.Scan(&filename)
+	require.NoError(t, err)
+	return filename
+}

server/plan_replayer.go

@@ -53,10 +53,25 @@ func (s *Server) newPlanReplayerHandler() *PlanReplayerHandler {
 func (prh PlanReplayerHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	params := mux.Vars(req)
 	name := params[pFileName]
-	path := filepath.Join(domain.GetPlanReplayerDirName(), name)
+	handler := downloadFileHandler{
+		filePath:           filepath.Join(domain.GetPlanReplayerDirName(), name),
+		fileName:           name,
+		infoGetter:         prh.infoGetter,
+		address:            prh.address,
+		statusPort:         prh.statusPort,
+		urlPath:            fmt.Sprintf("plan_replyaer/dump/%s", name),
+		downloadedFilename: "plan_replayer",
+	}
+	handleDownloadFile(handler, w, req)
+}
+
+func handleDownloadFile(handler downloadFileHandler, w http.ResponseWriter, req *http.Request) {
+	params := mux.Vars(req)
+	name := params[pFileName]
+	path := handler.filePath
 	if isExists(path) {
 		w.Header().Set("Content-Type", "application/zip")
-		w.Header().Set("Content-Disposition", "attachment; filename=\"plan_replayer.zip\"")
+		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s.zip\"", handler.downloadedFilename))
 		file, err := os.Open(path)
 		if err != nil {
 			writeError(w, err)
@@ -80,7 +95,7 @@ func (prh PlanReplayerHandler) ServeHTTP(w http.ResponseWriter, req *http.Reques
 		w.WriteHeader(http.StatusOK)
 		return
 	}
-	if prh.infoGetter == nil {
+	if handler.infoGetter == nil {
 		w.WriteHeader(http.StatusNotFound)
 		return
 	}
@@ -91,17 +106,17 @@ func (prh PlanReplayerHandler) ServeHTTP(w http.ResponseWriter, req *http.Reques
 		return
 	}
 	// If we didn't find file in origin request, try to broadcast the request to all remote tidb-servers
-	topos, err := prh.infoGetter.GetAllTiDBTopology(req.Context())
+	topos, err := handler.infoGetter.GetAllTiDBTopology(req.Context())
 	if err != nil {
 		writeError(w, err)
 		return
 	}
 	// transfer each remote tidb-server and try to find dump file
 	for _, topo := range topos {
-		if topo.IP == prh.address && topo.StatusPort == prh.statusPort {
+		if topo.IP == handler.address && topo.StatusPort == handler.statusPort {
 			continue
 		}
-		url := fmt.Sprintf("http://%s:%v/plan_replayer/dump/%s?forward=true", topo.IP, topo.StatusPort, name)
+		url := fmt.Sprintf("http://%s:%v/%s?forward=true", topo.IP, topo.StatusPort, handler.urlPath)
 		resp, err := http.Get(url) // #nosec G107
 		if err != nil {
 			terror.Log(errors.Trace(err))
@@ -113,7 +128,7 @@ func (prh PlanReplayerHandler) ServeHTTP(w http.ResponseWriter, req *http.Reques
 		}
 		// find dump file in one remote tidb-server, return file directly
 		w.Header().Set("Content-Type", "application/zip")
-		w.Header().Set("Content-Disposition", "attachment; filename=\"plan_replayer.zip\"")
+		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s.zip\"", handler.downloadedFilename))
 		_, err = io.Copy(w, resp.Body)
 		if err != nil {
 			writeError(w, err)
@@ -132,6 +147,16 @@ func (prh PlanReplayerHandler) ServeHTTP(w http.ResponseWriter, req *http.Reques
 	w.WriteHeader(http.StatusNotFound)
 }
 
+type downloadFileHandler struct {
+	filePath           string
+	fileName           string
+	infoGetter         *infosync.InfoSyncer
+	address            string
+	statusPort         uint
+	urlPath            string
+	downloadedFilename string
+}
+
 func isExists(path string) bool {
 	_, err := os.Stat(path)
 	if err != nil && !os.IsExist(err) {