fix: fix main.go

cloudwego · May 20, 2024 · 3d04807 · 3d04807
1 parent 0125858
commit 3d04807
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 45 deletions.
diff --git a/generator/generator.go b/generator/generator.go
@@ -17,12 +17,13 @@ package generator
 import (
 	"errors"
 	"fmt"
-	"github.com/cloudwego/thriftgo/generator/backend"
-	"github.com/cloudwego/thriftgo/plugin"
-	"github.com/cloudwego/thriftgo/utils/dir_utils"
 	"io/ioutil"
 	"os"
 	"path/filepath"
+
+	"github.com/cloudwego/thriftgo/generator/backend"
+	"github.com/cloudwego/thriftgo/plugin"
+	"github.com/cloudwego/thriftgo/utils/dir_utils"
 )
 
 // LangSpec is the parameter to specify which language to generate codes for and
@@ -105,12 +106,6 @@ func (g *Generator) preparePlugins(be backend.Backend, pds []*plugin.Desc) error
 
 // Generate generates codes for the target language and executes plugins specified.
 func (g *Generator) Generate(args *Arguments) (res *plugin.Response) {
-
-	//startTime := time.Now()
-	//defer func() {
-	//	fmt.Printf("Generate Cost: %s\n", time.Since(startTime))
-	//}()
-
 	out, req, log := args.Out, args.Req, args.Log
 
 	log.Info(fmt.Sprintf(`Generating: "%s"`, out.Language))
@@ -147,10 +142,6 @@ func (g *Generator) Generate(args *Arguments) (res *plugin.Response) {
 	}
 
 	if len(out.SDKPlugins) > 0 {
-		//startTime := time.Now()
-		//defer func() {
-		//	fmt.Printf("SDK Plugin Cost: %s\n", time.Since(startTime))
-		//}()
 		for _, sdk := range out.SDKPlugins {
 			req.PluginParameters = sdk.GetPluginParameters()
 			extra := sdk.Invoke(req)
@@ -161,26 +152,25 @@ func (g *Generator) Generate(args *Arguments) (res *plugin.Response) {
 				return plugin.BuildErrorResponse(err.Error())
 			}
 		}
-	} else {
-		//startTime := time.Now()
-		//defer func() {
-		//	fmt.Printf("Sub-process Plugin Cost: %s\n", time.Since(startTime))
-		//}()
-		for i, p := range g.plugins {
-			log.Info(fmt.Sprintf(`Run plugin "%s"`, p.Name()))
-
-			req.PluginParameters = plugin.Pack(out.UsedPlugins[i].Options)
-			extra := p.Execute(req)
-			log.MultiWarn(extra.Warnings)
+		res = g.files.BuildResponse()
+		return res
+	}
 
-			if err := extra.GetError(); err != "" {
-				return plugin.BuildErrorResponse(err)
-			}
-			if err := g.files.Feed(p.Name(), extra.Contents); err != nil {
-				return plugin.BuildErrorResponse(err.Error())
-			}
+	for i, p := range g.plugins {
+		log.Info(fmt.Sprintf(`Run plugin "%s"`, p.Name()))
+
+		req.PluginParameters = plugin.Pack(out.UsedPlugins[i].Options)
+		extra := p.Execute(req)
+		log.MultiWarn(extra.Warnings)
+
+		if err := extra.GetError(); err != "" {
+			return plugin.BuildErrorResponse(err)
+		}
+		if err := g.files.Feed(p.Name(), extra.Contents); err != nil {
+			return plugin.BuildErrorResponse(err.Error())
 		}
 	}
+
 	res = g.files.BuildResponse()
 	return res
 }

diff --git a/main.go b/main.go
@@ -16,15 +16,19 @@ package main
 
 import (
 	"fmt"
+	"os"
+	"runtime/debug"
+	"runtime/pprof"
+
+	"time"
+
 	"github.com/cloudwego/thriftgo/args"
 	"github.com/cloudwego/thriftgo/generator"
 	"github.com/cloudwego/thriftgo/generator/golang"
 	"github.com/cloudwego/thriftgo/parser"
 	"github.com/cloudwego/thriftgo/plugin"
 	"github.com/cloudwego/thriftgo/semantic"
 	"github.com/cloudwego/thriftgo/version"
-	"os"
-	"runtime/debug"
 )
 
 var (
@@ -37,7 +41,7 @@ var debugMode bool
 func init() {
 	_ = g.RegisterBackend(new(golang.GoBackend))
 	// export THRIFTGO_DEBUG=1
-	debugMode = true //os.Getenv("THRIFTGO_DEBUG") == "1"
+	debugMode = os.Getenv("THRIFTGO_DEBUG") == "1"
 }
 
 func check(err error) {
@@ -48,17 +52,17 @@ func check(err error) {
 }
 
 func main() {
-	//if debugMode {
-	//	f, _ := os.Create("thriftgo-cpu.pprof")
-	//	defer f.Close()
-	//	_ = pprof.StartCPUProfile(f)
-	//	defer pprof.StopCPUProfile()
-	//
-	//	startTime := time.Now()
-	//	defer func() {
-	//		fmt.Printf("Cost: %s\n", time.Since(startTime))
-	//	}()
-	//}
+	if debugMode {
+		f, _ := os.Create("thriftgo-cpu.pprof")
+		defer f.Close()
+		_ = pprof.StartCPUProfile(f)
+		defer pprof.StopCPUProfile()
+
+		startTime := time.Now()
+		defer func() {
+			fmt.Printf("Cost: %s\n", time.Since(startTime))
+		}()
+	}
 
 	defer handlePanic()
 	check(a.Parse(os.Args))

diff --git a/sdk/run.go → sdk/sdk.go b/sdk/run.go → sdk/sdk.go
@@ -16,6 +16,8 @@ package sdk
 
 import (
 	"fmt"
+	"os"
+
 	"github.com/cloudwego/thriftgo/args"
 	"github.com/cloudwego/thriftgo/generator"
 	"github.com/cloudwego/thriftgo/generator/backend"
@@ -25,7 +27,6 @@ import (
 	"github.com/cloudwego/thriftgo/semantic"
 	"github.com/cloudwego/thriftgo/utils/dir_utils"
 	"github.com/cloudwego/thriftgo/version"
-	"os"
 )
 
 func init() {