将 watchagent 从全局注入改为随 main 包注入

pkg/cover/agentwatch.tpl

@@ -1,4 +1,4 @@
-package coverdef
+package cover
 
 import (
 	"fmt"
@@ -6,19 +6,11 @@ import (
 	"os"
 	"log"
 	"strconv"
-	"strings"
 	"net/url"
 
 	"{{.GlobalCoverVarImportPath}}/websocket"
-)
 
-var (
+	_cover "{{.GlobalCoverVarImportPath}}"
-	watchChannel = make(chan *blockInfo, 1024)
-
-	watchEnabled = false
-
-	waitDelay time.Duration = 10 * time.Second
-	host      string        = "{{.Host}}"
 )
 
 func init() {
@@ -55,7 +47,7 @@ func init() {
 			}
 
 			// 连接成功
-			watchEnabled = true
+			_cover.WatchEnabled_{{.Random}} = true
 			log.Printf("[goc][Info] watch connected to goc server")
 
 			ticker := time.NewTicker(time.Second)
@@ -74,18 +66,18 @@ func init() {
 			Loop:
 			for {
 				select {
-				case block := <-watchChannel:
+				case block := <-_cover.WatchChannel_{{.Random}}:
-					i := block.i
+					i := block.I
 
-					cov := fmt.Sprintf("%s:%d.%d,%d.%d %d %d", block.name,
+					cov := fmt.Sprintf("%s:%d.%d,%d.%d %d %d", block.Name,
-						block.pos[3*i+0], uint16(block.pos[3*i+2]),
+						block.Pos[3*i+0], uint16(block.Pos[3*i+2]),
-						block.pos[3*i+1], uint16(block.pos[3*i+2] >> 16),
+						block.Pos[3*i+1], uint16(block.Pos[3*i+2] >> 16),
-						block.stmts,
+						block.Stmts,
 						1)
 
 					err = ws.WriteMessage(websocket.TextMessage, []byte(cov))
 					if err != nil {
-						watchEnabled = false
+						_cover.WatchEnabled_{{.Random}} = false
 						log.Println("[goc][Error] push coverage failed: %v", err)
 						time.Sleep(waitDelay)
 						break Loop
@@ -99,56 +91,3 @@ func init() {
 		}
 	}()
 }
-
-// get process meta info for register
-type processInfo struct {
-	hostname string
-	pid      int
-	cmdline  string
-}
-
-func getRegisterInfo() (*processInfo, error) {
-	hostname, err := os.Hostname()
-	if err != nil {
-		log.Printf("[goc][Error] fail to get hostname: %v", hostname)
-		return nil, err
-	}
-
-	pid := os.Getpid()
-
-	cmdline := strings.Join(os.Args, " ")
-
-	return &processInfo{
-		hostname: hostname,
-		pid:      pid,
-		cmdline:  cmdline,
-	}, nil
-}
-
-//
-
-type blockInfo struct {
-	name  string
-	pos   []uint32
-	i     int
-	stmts int
-}
-
-// UploadCoverChangeEvent_{{.Random}} is non-blocking
-func UploadCoverChangeEvent_{{.Random}}(name string, pos []uint32, i int, stmts uint16) {
-
-	if watchEnabled == false {
-		return
-	}
-
-	// make sure send is non-blocking
-	select {
-	case watchChannel <- &blockInfo{
-		name:  name,
-		pos:   pos,
-		i:     i,
-		stmts: int(stmts),
-	}:
-	default:
-	}
-}

pkg/cover/inject.go

@@ -1,6 +1,7 @@
 package cover
 
 import (
+	"fmt"
 	"os"
 	"path"
 	"path/filepath"
@@ -56,12 +57,7 @@ func Inject() {
 	}
 	// 在工程根目录注入所有插桩变量的声明+定义
 	injectGlobalCoverVarFile(allDecl)
-	// 在工程根目录注入 watch agent 的定义
+
-	if config.GocConfig.Mode == "watch" {
-		log.Infof("watch mode is enabled")
-		injectWatchAgentFile()
-		log.Donef("watch handler injected")
-	}
 	// 添加自定义 websocket 依赖
 	// 用户代码可能有 gorrila/websocket 的依赖，为避免依赖冲突，以及可能的 replace/vendor，
 	// 这里直接注入一份完整的 gorrila/websocket 实现
@@ -120,7 +116,8 @@ func getPkgTmpDir(pkgDir string) string {
 // - goc-cover-agent-apis-auto-generated-11111-22222-bridge.go
 // - goc-cover-agent-apis-auto-generated-11111-22222-package
 //  |
-//  -- init.go
+//  -- rpcagent.go
+//  -- watchagent.go
 //
 // 11111_22222_bridge.go 仅仅用于引用 11111_22222_package, where package contains ws agent main logic.
 // 使用 bridge.go 文件是为了避免插桩逻辑中的变量名污染 main 包
@@ -153,7 +150,7 @@ func injectGocAgent(where string, covers []*config.PackageCover) {
 	}
 
 	// create ws agent files
-	dest := filepath.Join(wherePkg, "init.go")
+	dest := filepath.Join(wherePkg, "rpcagent.go")
 
 	f2, err := os.Create(dest)
 	if err != nil {
@@ -184,6 +181,29 @@ func injectGocAgent(where string, covers []*config.PackageCover) {
 	if err := coverMainTmpl.Execute(f2, tmplData); err != nil {
 		log.Fatalf("fail to generate cover agent handlers in temporary project: %v", err)
 	}
+
+	// 写入 watch
+	if config.GocConfig.Mode != "watch" {
+		return
+	}
+	f, err := os.Create(filepath.Join(wherePkg, "watchagent.go"))
+	if err != nil {
+		log.Fatalf("fail to create watchagent file: %v", err)
+	}
+
+	tmplwatchData := struct {
+		Random                   string
+		Host                     string
+		GlobalCoverVarImportPath string
+	}{
+		Random:                   filepath.Base(config.GocConfig.TmpModProjectDir),
+		Host:                     config.GocConfig.Host,
+		GlobalCoverVarImportPath: config.GocConfig.GlobalCoverVarImportPath,
+	}
+
+	if err := coverWatchTmpl.Execute(f, tmplwatchData); err != nil {
+		log.Fatalf("fail to generate watchagent in temporary project: %v", err)
+	}
 }
 
 // injectGlobalCoverVarFile 写入所有插桩变量的全局定义至一个单独的文件
@@ -205,32 +225,42 @@ func injectGlobalCoverVarFile(decl string) {
 
 	packageName := "package coverdef\n\n"
 
-	_, err = coverFile.WriteString(packageName + decl)
+	random := filepath.Base(config.GocConfig.TmpModProjectDir)
+	varWatchDef := fmt.Sprintf(`
+var WatchChannel_%v = make(chan *blockInfo, 1024)
+
+var WatchEnabled_%v = false
+
+type blockInfo struct {
+	Name  string
+	Pos   []uint32
+	I     int
+	Stmts int
+}
+
+// UploadCoverChangeEvent_%v is non-blocking
+func UploadCoverChangeEvent_%v(name string, pos []uint32, i int, stmts uint16) {
+
+	if WatchEnabled_%v == false {
+		return
+	}
+
+	// make sure send is non-blocking
+	select {
+	case WatchChannel_%v <- &blockInfo{
+		Name:  name,
+		Pos:   pos,
+		I:     i,
+		Stmts: int(stmts),
+	}:
+	default:
+	}
+}
+
+`, random, random, random, random, random, random)
+
+	_, err = coverFile.WriteString(packageName + varWatchDef + decl)
 	if err != nil {
 		log.Fatalf("fail to write to global cover definition file: %v", err)
 	}
 }
-
-func injectWatchAgentFile() {
-	globalCoverVarPackage := path.Base(config.GocConfig.GlobalCoverVarImportPath)
-	globalCoverDef := filepath.Join(config.GocConfig.TmpModProjectDir, globalCoverVarPackage)
-
-	f, err := os.Create(filepath.Join(globalCoverDef, "watchagent.go"))
-	if err != nil {
-		log.Fatalf("fail to create watchagent file: %v", err)
-	}
-
-	tmplData := struct {
-		Random                   string
-		Host                     string
-		GlobalCoverVarImportPath string
-	}{
-		Random:                   filepath.Base(config.GocConfig.TmpModProjectDir),
-		Host:                     config.GocConfig.Host,
-		GlobalCoverVarImportPath: config.GocConfig.GlobalCoverVarImportPath,
-	}
-
-	if err := coverWatchTmpl.Execute(f, tmplData); err != nil {
-		log.Fatalf("fail to generate watchagent in temporary project: %v", err)
-	}
-}