From ca1ec49eebcc6b531340fcf2735acae196050e4f Mon Sep 17 00:00:00 2001
From: liruichen <anchor_lee@163.com>
Date: Fri, 28 Jul 2023 11:48:05 +0800
Subject: [PATCH] feat: add needed file filter

---
 cmd/profile.go                     |  4 +++-
 pkg/client/profie.go               |  5 +++--
 pkg/client/rest/profile/profile.go |  9 ++++++++
 pkg/server/api.go                  | 34 ++++++++++++++++++++++++------
 4 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/cmd/profile.go b/cmd/profile.go
index 6048908..aac0092 100644
--- a/cmd/profile.go
+++ b/cmd/profile.go
@@ -32,6 +32,7 @@ var (
 	profileIds         []string
 	profileSkipPattern []string
 	profileExtra       string
+	profileNeedPattern []string
 )
 
 func init() {
@@ -45,6 +46,7 @@ func init() {
 	add2Flags := func(f *pflag.FlagSet) {
 		f.StringVarP(&profileOutput, "output", "o", "", "download cover profile")
 		f.StringSliceVar(&profileSkipPattern, "skip", nil, "skip specific packages in the profile")
+		f.StringSliceVarP(&profileNeedPattern, "need", "n", nil, "find specific packages in the profile")
 	}
 
 	add1Flags(getProfileCmd.Flags())
@@ -63,7 +65,7 @@ var getProfileCmd = &cobra.Command{
 }
 
 func getProfile(cmd *cobra.Command, args []string) {
-	client.GetProfile(profileHost, profileIds, profileSkipPattern, profileExtra, profileOutput)
+	client.GetProfile(profileHost, profileIds, profileSkipPattern, profileExtra, profileOutput, profileNeedPattern)
 }
 
 var clearProfileCmd = &cobra.Command{
diff --git a/pkg/client/profie.go b/pkg/client/profie.go
index 4964380..3c9eccd 100644
--- a/pkg/client/profie.go
+++ b/pkg/client/profie.go
@@ -25,12 +25,13 @@ import (
 	"github.com/qiniu/goc/v2/pkg/log"
 )
 
-func GetProfile(host string, ids []string, skips []string, extra string, output string) {
+func GetProfile(host string, ids []string, skips []string, extra string, output string, need []string) {
 	gocClient := rest.NewV2Client(host)
 
 	profiles, err := gocClient.Profile().Get(ids,
 		profile.WithPackagePattern(skips),
-		profile.WithExtraPattern(extra))
+		profile.WithExtraPattern(extra),
+		profile.WithNeed(need))
 	if err != nil {
 		log.Fatalf("fail to get profile from the goc server: %v, response: %v", err, profiles)
 	}
diff --git a/pkg/client/rest/profile/profile.go b/pkg/client/rest/profile/profile.go
index 3744833..acf8f26 100644
--- a/pkg/client/rest/profile/profile.go
+++ b/pkg/client/rest/profile/profile.go
@@ -39,6 +39,7 @@ type profileClient struct {
 	c            *resty.Client
 	skipPatterns []string
 	extraPattern string
+	needPattern  []string
 }
 
 func NewProfileClient(c *resty.Client) *profileClient {
@@ -61,6 +62,12 @@ func WithExtraPattern(pattern string) profileOption {
 	}
 }
 
+func WithNeed(need []string) profileOption {
+	return func(pc *profileClient) {
+		pc.needPattern = need
+	}
+}
+
 func (p *profileClient) Get(ids []string, opts ...profileOption) (string, error) {
 	for _, opt := range opts {
 		opt(p)
@@ -70,10 +77,12 @@ func (p *profileClient) Get(ids []string, opts ...profileOption) (string, error)
 
 	idQuery := strings.Join(ids, ",")
 	skipQuery := strings.Join(p.skipPatterns, ",")
+	needPattern := strings.Join(p.needPattern, ",")
 
 	req.QueryParam.Add("id", idQuery)
 	req.QueryParam.Add("skippattern", skipQuery)
 	req.QueryParam.Add("extra", p.extraPattern)
+	req.QueryParam.Add("needpattern", needPattern)
 
 	res := struct {
 		Data string `json:"profile,omitempty"`
diff --git a/pkg/server/api.go b/pkg/server/api.go
index e68d7c8..18d92f2 100644
--- a/pkg/server/api.go
+++ b/pkg/server/api.go
@@ -106,6 +106,11 @@ func (gs *gocServer) getProfiles(c *gin.Context) {
 	if skippatternRaw != "" {
 		skippattern = strings.Split(skippatternRaw, ",")
 	}
+	neerpatternRaw := c.Query("needpattern")
+	var neerpattern []string
+	if neerpatternRaw != "" {
+		neerpattern = strings.Split(neerpatternRaw, ",")
+	}
 
 	extra := c.Query("extra")
 	isExtra := filterExtra(extra)
@@ -180,7 +185,7 @@ func (gs *gocServer) getProfiles(c *gin.Context) {
 			}
 
 			// check if skippattern matches
-			newProfile := filterProfileByPattern(skippattern, profile)
+			newProfile := filterProfileByPattern(skippattern, neerpattern, profile)
 
 			mu.Lock()
 			defer mu.Unlock()
@@ -315,13 +320,14 @@ func (gs *gocServer) watchProfileUpdate(c *gin.Context) {
 	<-gwc.exitCh
 }
 
-func filterProfileByPattern(skippattern []string, profiles []*cover.Profile) []*cover.Profile {
-
-	if len(skippattern) == 0 {
-		return profiles
-	}
+func filterProfileByPattern(skippattern []string, needpattern []string, profiles []*cover.Profile) []*cover.Profile {
 
 	var out = make([]*cover.Profile, 0)
+
+	if len(skippattern) == 0 {
+		goto need
+	}
+
 	for _, profile := range profiles {
 		skip := false
 		for _, pattern := range skippattern {
@@ -335,6 +341,22 @@ func filterProfileByPattern(skippattern []string, profiles []*cover.Profile) []*
 			out = append(out, profile)
 		}
 	}
+need:
+	if len(needpattern) == 0 {
+		return profiles
+	}
+	for _, profile := range profiles {
+		need := false
+		for _, pattern := range needpattern {
+			if strings.Contains(profile.FileName, pattern) {
+				need = true
+				break
+			}
+		}
+		if need {
+			out = append(out, profile)
+		}
+	}
 
 	return out
 }