feat: support to get pprof of an extension (#307)

Co-authored-by: rick <[email protected]>

Author: LinuxSuRen

cmd/server.go

@@ -58,6 +58,7 @@ import (
 	"golang.org/x/oauth2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/reflection"
 )
 
@@ -275,7 +276,7 @@ func (o *serverOption) runE(cmd *cobra.Command, args []string) (err error) {
 			promhttp.HandlerFor(reg, promhttp.HandlerOpts{Registry: reg}).ServeHTTP(w, r)
 		})
 
-		debugHandler(mux)
+		debugHandler(mux, remoteServer)
 		o.httpServer.WithHandler(mux)
 		log.Printf("HTTP server listening at %v", httplis.Addr())
 		log.Printf("Server is running.")
@@ -346,9 +347,32 @@ func frontEndHandlerWithLocation(consolePath string) func(w http.ResponseWriter,
 	}
 }
 
-func debugHandler(mux *runtime.ServeMux) {
+func debugHandler(mux *runtime.ServeMux, remoteServer server.RunnerServer) {
 	mux.HandlePath(http.MethodGet, "/debug/pprof/{sub}", func(w http.ResponseWriter, r *http.Request, pathParams map[string]string) {
-		switch sub := pathParams["sub"]; sub {
+		sub := pathParams["sub"]
+		extName := r.URL.Query().Get("name")
+		if extName != "" && remoteServer != nil {
+			log.Println("get pprof of extension:", extName)
+
+			ctx := metadata.NewIncomingContext(r.Context(), metadata.New(map[string]string{
+				server.HeaderKeyStoreName: extName,
+			}))
+
+			data, err := remoteServer.PProf(ctx, &server.PProfRequest{
+				Name: sub,
+			})
+			if err == nil {
+				w.Header().Set("Content-Type", "application/octet-stream")
+				w.Write(data.Data)
+			} else {
+				w.WriteHeader(http.StatusBadRequest)
+				w.Write([]byte(err.Error()))
+			}
+
+			return
+		}
+
+		switch sub {
 		case "cmdline":
 			pprof.Cmdline(w, r)
 		case "profile":

cmd/server_test.go

@@ -132,7 +132,7 @@ func TestFrontEndHandlerWithLocation(t *testing.T) {
 			"allocs", "block", "goroutine", "heap", "mutex", "threadcreate"}
 
 		mu := runtime.NewServeMux()
-		debugHandler(mu)
+		debugHandler(mu, nil)
 
 		ready := make(chan struct{})
 		var err error

extensions/store-etcd/pkg/server.go

@@ -30,6 +30,7 @@ import (
 	"log"
 	"time"
 
+	"github.com/linuxsuren/api-testing/pkg/extension"
 	"github.com/linuxsuren/api-testing/pkg/server"
 	"github.com/linuxsuren/api-testing/pkg/testing"
 	"github.com/linuxsuren/api-testing/pkg/testing/remote"
@@ -263,6 +264,14 @@ func (s *remoteserver) Verify(ctx context.Context, in *server.Empty) (reply *ser
 	}
 	return
 }
+func (s *remoteserver) PProf(ctx context.Context, in *server.PProfRequest) (data *server.PProfData, err error) {
+	log.Println("pprof", in.Name)
+
+	data = &server.PProfData{
+		Data: extension.LoadPProf(in.Name),
+	}
+	return
+}
 
 func connectTestOption() []clientv3.OpOption {
 	return []clientv3.OpOption{clientv3.WithLimit(1), clientv3.WithPrefix(),

extensions/store-git/pkg/server.go

@@ -39,6 +39,7 @@ import (
 	"github.com/go-git/go-git/v5/plumbing"
 	"github.com/go-git/go-git/v5/plumbing/object"
 	"github.com/go-git/go-git/v5/plumbing/transport/http"
+	"github.com/linuxsuren/api-testing/pkg/extension"
 	"github.com/linuxsuren/api-testing/pkg/server"
 	"github.com/linuxsuren/api-testing/pkg/testing"
 	"github.com/linuxsuren/api-testing/pkg/testing/remote"
@@ -268,6 +269,14 @@ func (s *gitClient) Verify(ctx context.Context, in *server.Empty) (reply *server
 	}
 	return
 }
+func (s *gitClient) PProf(ctx context.Context, in *server.PProfRequest) (data *server.PProfData, err error) {
+	log.Println("pprof", in.Name)
+
+	data = &server.PProfData{
+		Data: extension.LoadPProf(in.Name),
+	}
+	return
+}
 func (s *gitClient) getClient(ctx context.Context) (opt *gitOptions, err error) {
 	store := remote.GetStoreFromContext(ctx)
 	if store == nil {

extensions/store-mongodb/pkg/server.go

@@ -26,7 +26,9 @@ package pkg
 
 import (
 	"context"
+	"log"
 
+	"github.com/linuxsuren/api-testing/pkg/extension"
 	"github.com/linuxsuren/api-testing/pkg/server"
 	"github.com/linuxsuren/api-testing/pkg/testing"
 	"github.com/linuxsuren/api-testing/pkg/testing/remote"
@@ -265,6 +267,14 @@ func (s *dbserver) Verify(ctx context.Context, in *server.Empty) (reply *server.
 	}
 	return
 }
+func (s *dbserver) PProf(ctx context.Context, in *server.PProfRequest) (data *server.PProfData, err error) {
+	log.Println("pprof", in.Name)
+
+	data = &server.PProfData{
+		Data: extension.LoadPProf(in.Name),
+	}
+	return
+}
 
 type SuiteNameGetter interface {
 	GetSuiteName() string

extensions/store-orm/pkg/server.go

@@ -31,6 +31,7 @@ import (
 	"log"
 	"strings"
 
+	"github.com/linuxsuren/api-testing/pkg/extension"
 	"github.com/linuxsuren/api-testing/pkg/server"
 	"github.com/linuxsuren/api-testing/pkg/testing/remote"
 	"github.com/linuxsuren/api-testing/pkg/util"
@@ -268,6 +269,14 @@ func (s *dbserver) Verify(ctx context.Context, in *server.Empty) (reply *server.
 	}
 	return
 }
+func (s *dbserver) PProf(ctx context.Context, in *server.PProfRequest) (data *server.PProfData, err error) {
+	log.Println("pprof", in.Name)
+
+	data = &server.PProfData{
+		Data: extension.LoadPProf(in.Name),
+	}
+	return
+}
 
 func testCaseIdentiy(db *gorm.DB, testcase *TestCase) *gorm.DB {
 	return db.Model(testcase).Where(fmt.Sprintf("suite_name = '%s' AND name = '%s'", testcase.SuiteName, testcase.Name))

extensions/store-s3/pkg/s3_server.go

@@ -35,6 +35,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/linuxsuren/api-testing/pkg/extension"
 	"github.com/linuxsuren/api-testing/pkg/server"
 	"github.com/linuxsuren/api-testing/pkg/testing"
 	"github.com/linuxsuren/api-testing/pkg/testing/remote"
@@ -205,6 +206,14 @@ func (s *s3Client) Verify(ctx context.Context, in *server.Empty) (reply *server.
 	}
 	return
 }
+func (s *s3Client) PProf(ctx context.Context, in *server.PProfRequest) (data *server.PProfData, err error) {
+	log.Println("pprof", in.Name)
+
+	data = &server.PProfData{
+		Data: extension.LoadPProf(in.Name),
+	}
+	return
+}
 func (s *s3Client) getClient(ctx context.Context) (db *s3WithBucket, err error) {
 	store := remote.GetStoreFromContext(ctx)
 	if store == nil {

pkg/extension/pprof.go

@@ -0,0 +1,38 @@
+/**
+MIT License
+
+Copyright (c) 2023 API Testing Authors.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+package extension
+
+import (
+	"bytes"
+	"runtime/pprof"
+)
+
+func LoadPProf(name string) []byte {
+	buf := new(bytes.Buffer)
+	if p := pprof.Lookup(name); p != nil {
+		p.WriteTo(buf, 0)
+	}
+	return buf.Bytes()
+}

pkg/server/remote_server.go

@@ -180,7 +180,7 @@ func (s *server) getLoader(ctx context.Context) (loader testing.Writer) {
 	if mdd, ok = metadata.FromIncomingContext(ctx); ok {
 		storeNameMeta := mdd.Get(HeaderKeyStoreName)
 		if len(storeNameMeta) > 0 {
-			storeName := storeNameMeta[0]
+			storeName := strings.TrimSpace(storeNameMeta[0])
 			if storeName == "local" || storeName == "" {
 				return
 			}
@@ -813,6 +813,13 @@ func (s *server) DeleteSecret(ctx context.Context, in *Secret) (reply *CommonRes
 func (s *server) UpdateSecret(ctx context.Context, in *Secret) (reply *CommonResult, err error) {
 	return s.secretServer.UpdateSecret(ctx, in)
 }
+func (s *server) PProf(ctx context.Context, in *PProfRequest) (reply *PProfData, err error) {
+	loader := s.getLoader(ctx)
+	reply = &PProfData{
+		Data: loader.PProf(in.Name),
+	}
+	return
+}
 
 func (s *server) getLoaderByStoreName(storeName string) (loader testing.Writer, err error) {
 	var store *testing.Store