main.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	targetAddr       string        // 目标程序的地址
	pprofDir         string        // pprof 文件保存目录
	periodicInterval time.Duration //时间间隔
	memoryThreshold  uint64        // 内存阈值，单位为GB
)

// MemStats via expvar
type ExpvarMemStats struct {
	Alloc uint64 `json:"Alloc"`
	// 可以根据需要添加其他 memstats 字段
}

type ExpvarResponse struct {
	Memstats ExpvarMemStats `json:"memstats"`
	// 可以根据需要添加其他 expvar 变量
}
type pprofGet struct {
	Name string `json:"name"`
	Url  string `json:"url"`
}

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile) //启用日志时间戳和文件名
	flag.StringVar(&targetAddr, "target", "127.0.0.1:6060", "目标Go程序的HTTP地址")
	flag.StringVar(&pprofDir, "save-dir", "./pprof_dumps", "指定proof文件保存目录")
	flag.Uint64Var(&memoryThreshold, "memory-threshold", 15, "指定生成pprof文件的内存阈值(单位GB)")
	flag.DurationVar(&periodicInterval, "periodic-interval", 10*time.Minute, "指定生成的pprof文件的间隔时间(单位：分钟)")
	flag.Parse()
	//尝试创建保存目录
	_, err := os.Stat(pprofDir)
	if err != nil {
		if os.IsNotExist(err) {
			if errMkdir := os.MkdirAll(pprofDir, 0755); errMkdir != nil {
				log.Fatalf("无法创建 pprof 目录或上级目录 %s: %v", pprofDir, errMkdir)
			}
		} else {
			log.Fatalf("无法访问 pprof 保存目录: %v", err)
		}
	}
	log.Println("生成pprof文件的脚本已启动...")
	log.Printf("监控目标程序的端口: %s，pprof文件保存目录: %s,间隔时间：%v \n", targetAddr, pprofDir, periodicInterval)
	log.Printf("每隔%v生成一次%s的pprof文件\n", periodicInterval, targetAddr)
	log.Printf("当%s占用内存达到%vGB时生成pprof文件\n", targetAddr, memoryThreshold)
	go periodicPprofGenerator()
	go memoryBasedPprofGenerator()
	select {} // 阻塞主 goroutine，保持程序运行
}
func periodicPprofGenerator() {
	ticker := time.NewTicker(periodicInterval)
	defer ticker.Stop()
	for range ticker.C {
		generateRemoteProfile("periodic")
	}
}
func memoryBasedPprofGenerator() {
	ticker := time.NewTicker(15 * time.Second) // 每15秒检查一次
	defer ticker.Stop()
	client := http.Client{Timeout: 10 * time.Second}
	for range ticker.C {
		expvarURL := fmt.Sprintf("http://%s/debug/vars", targetAddr)
		resp, err := client.Get(expvarURL)
		if err != nil {
			log.Printf("错误: 无法从 %s 获取 /debug/vars: %v \n", targetAddr, err)
			continue
		}
		defer func() {
			errClose := resp.Body.Close()
			if errClose != nil {
				log.Printf("错误: 关闭响应体时出错: %v \n", errClose)
			}
		}()
		if resp.StatusCode != http.StatusOK {
			log.Printf("错误: 从 %s 获取 /debug/vars 失败，状态码: %d \n", targetAddr, resp.StatusCode)
			bodyBytes, _ := io.ReadAll(resp.Body)
			log.Printf("响应体: %s \n", string(bodyBytes))
			continue
		}
		bodyBytes, err := io.ReadAll(resp.Body)
		if err != nil {
			log.Printf("错误: 无法读取 %s 的 /debug/vars 响应体: %v \n", targetAddr, err)
			continue
		}
		var data ExpvarResponse
		if err = json.NewDecoder(bytes.NewReader(bodyBytes)).Decode(&data); err != nil {
			log.Printf("错误: 无法解析 %s 的 /debug/vars 响应: %v \n", targetAddr, err)
			log.Printf("响应体: %s \n", string(bodyBytes))
			continue
		}
		thresholdBytes := memoryThreshold * 1024 * 1024 * 1024
		if data.Memstats.Alloc >= thresholdBytes {
			generateRemoteProfile(fmt.Sprintf("memory_threshold_%vGB", data.Memstats.Alloc))
			break
		}
	}
}
func generateRemoteProfile(reason string) {
	timestamp := time.Now().Format("20060102T150405")
	targetName := SanitizeFilename(targetAddr)
	// 生成 pprof 文件的逻辑
	var pprofUrl, fileName string

	pprofFileNames := make(map[string]pprofGet)

	pprofUrl = fmt.Sprintf("http://%s/debug/pprof/profile?seconds=%d", targetAddr, 5)
	fileName = fmt.Sprintf("%s/cpu_%s_%s_%s.pprof", pprofDir, targetName, reason, timestamp)
	pprofFileNames["cpu"] = pprofGet{fileName, pprofUrl}

	pprofUrl = fmt.Sprintf("http://%s/debug/pprof/trace?seconds=%d", targetAddr, 5)
	fileName = fmt.Sprintf("%s/trace_%s_%s_%s.pprof", pprofDir, targetName, reason, timestamp)
	pprofFileNames["trace"] = pprofGet{fileName, pprofUrl}

	pprofUrl = fmt.Sprintf("http://%s/debug/pprof/heap", targetAddr)
	fileName = fmt.Sprintf("%s/heap_%s_%s_%s.pprof", pprofDir, targetName, reason, timestamp)
	pprofFileNames["heap"] = pprofGet{fileName, pprofUrl}
	var wg sync.WaitGroup
	for _, pprof := range pprofFileNames {
		wg.Add(1)
		go func(url, name string) {
			defer wg.Done()
			log.Printf("正在为 %s 生成 pprof 文件%s (原因: %s)\n", targetAddr, name, reason)
			resp, err := http.Get(url)
			if err != nil {
				log.Printf("错误: 无法从 %s 获取 pprof 数据: %v \n", url, err)
				return
			}
			defer func() {
				err = resp.Body.Close()
				if err != nil {
					log.Printf("错误: 关闭响应体时出错: %v (%s) \n", err, url)
				}
			}()
			if resp.StatusCode != http.StatusOK {
				log.Printf("错误: 从 %s 获取 pprof 数据失败，状态码: %d \n", url, resp.StatusCode)
				return
			}
			file, err := os.Create(name)
			if err != nil {
				log.Printf("错误: 无法创建 pprof 文件 %s: %v \n", name, err)
				return
			}
			defer func() {
				err = file.Close()
				if err != nil {
					log.Printf("错误: 关闭 pprof 文件 %s 时出错: %v \n", name, err)
				}
			}()
			_, err = io.Copy(file, resp.Body)
			if err != nil {
				log.Printf("错误: 无法将 pprof 数据写入文件 %s: %v \n", name, err)
				return
			}
			log.Printf("成功为 %s 生成 pprof 文件: %s (原因: %s)\n", targetAddr, name, reason)
		}(pprof.Url, pprof.Name)
	}
	wg.Wait()
	log.Print("所有 pprof 文件已生成!")
}

// SanitizeFilename 用于整理文件名中的特殊字符
func SanitizeFilename(name string) string {
	r := []rune(name)
	for i, c := range r {
		if c < 32 || c > 126 || c == '/' || c == '\\' || c == ':' || c == '*' || c == '?' || c == '"' || c == '<' || c == '>' || c == '|' {
			r[i] = '_'
		}
	}
	return string(r)
}