1package profiler
  2
  3import (
  4    "context"
  5    "fmt"
  6    "os"
  7    "runtime"
  8    "runtime/pprof"
  9    "runtime/trace"
 10    "sync"
 11    "time"
 12)
 13
 14type ProfileType int
 15
 16const (
 17    CPUProfile ProfileType = iota
 18    MemProfile
 19    HeapProfile
 20    GoroutineProfile
 21    BlockProfile
 22    MutexProfile
 23    ThreadProfile
 24    AllocProfile
 25)
 26
 27func String() string {
 28    switch pt {
 29    case CPUProfile:
 30        return "cpu"
 31    case MemProfile:
 32        return "mem"
 33    case HeapProfile:
 34        return "heap"
 35    case GoroutineProfile:
 36        return "goroutine"
 37    case BlockProfile:
 38        return "block"
 39    case MutexProfile:
 40        return "mutex"
 41    case ThreadProfile:
 42        return "threadcreate"
 43    case AllocProfile:
 44        return "allocs"
 45    default:
 46        return "unknown"
 47    }
 48}
 49
 50type ProfileStats struct {
 51    NumGoroutines  int
 52    MemAllocated   uint64
 53    MemTotal       uint64
 54    NumGC          uint32
 55    CPUSamples     int
 56    BlockingSamples int
 57}
 58
 59type Profiler struct {
 60    types      []ProfileType
 61    interval   time.Duration
 62    outputDir  string
 63    cpuFile    *os.File
 64    traceFile  *os.File
 65    mu         sync.Mutex
 66    collecting bool
 67    stats      ProfileStats
 68}
 69
 70func New(types []ProfileType, interval time.Duration, outputDir string) *Profiler {
 71    return &Profiler{
 72        types:     types,
 73        interval:  interval,
 74        outputDir: outputDir,
 75    }
 76}
 77
 78func Start(ctx context.Context) error {
 79    p.mu.Lock()
 80    if p.collecting {
 81        p.mu.Unlock()
 82        return fmt.Errorf("profiler already running")
 83    }
 84    p.collecting = true
 85    p.mu.Unlock()
 86
 87    // Enable profiling for specific types
 88    for _, typ := range p.types {
 89        switch typ {
 90        case BlockProfile:
 91            runtime.SetBlockProfileRate(1)
 92        case MutexProfile:
 93            runtime.SetMutexProfileFraction(1)
 94        }
 95    }
 96
 97    // Start CPU profiling if requested
 98    if p.hasProfileType(CPUProfile) {
 99        if err := p.startCPUProfile(); err != nil {
100            return err
101        }
102    }
103
104    // Start periodic collection
105    go p.periodicCollection(ctx)
106
107    return nil
108}
109
110func Stop() error {
111    p.mu.Lock()
112    if !p.collecting {
113        p.mu.Unlock()
114        return fmt.Errorf("profiler not running")
115    }
116    p.collecting = false
117    p.mu.Unlock()
118
119    // Stop CPU profiling
120    if p.cpuFile != nil {
121        pprof.StopCPUProfile()
122        p.cpuFile.Close()
123        p.cpuFile = nil
124    }
125
126    // Stop trace
127    if p.traceFile != nil {
128        trace.Stop()
129        p.traceFile.Close()
130        p.traceFile = nil
131    }
132
133    return nil
134}
135
136func startCPUProfile() error {
137    filename := fmt.Sprintf("%s/cpu_%d.prof", p.outputDir, time.Now().Unix())
138    f, err := os.Create(filename)
139    if err != nil {
140        return err
141    }
142
143    if err := pprof.StartCPUProfile(f); err != nil {
144        f.Close()
145        return err
146    }
147
148    p.cpuFile = f
149    return nil
150}
151
152func periodicCollection(ctx context.Context) {
153    ticker := time.NewTicker(p.interval)
154    defer ticker.Stop()
155
156    for {
157        select {
158        case <-ctx.Done():
159            return
160        case <-ticker.C:
161            p.collectStats()
162            p.writeSnapshots()
163        }
164    }
165}
166
167func collectStats() {
168    var m runtime.MemStats
169    runtime.ReadMemStats(&m)
170
171    p.mu.Lock()
172    p.stats.NumGoroutines = runtime.NumGoroutine()
173    p.stats.MemAllocated = m.Alloc
174    p.stats.MemTotal = m.TotalAlloc
175    p.stats.NumGC = m.NumGC
176    p.mu.Unlock()
177}
178
179func writeSnapshots() {
180    timestamp := time.Now().Unix()
181
182    for _, typ := range p.types {
183        if typ == CPUProfile {
184            continue // CPU profile is written separately
185        }
186
187        filename := fmt.Sprintf("%s/%s_%d.prof",
188            p.outputDir, typ.String(), timestamp)
189        if err := p.WriteProfile(typ, filename); err != nil {
190            // Log error
191        }
192    }
193}
194
195func WriteProfile(typ ProfileType, filename string) error {
196    f, err := os.Create(filename)
197    if err != nil {
198        return err
199    }
200    defer f.Close()
201
202    var profile *pprof.Profile
203
204    switch typ {
205    case HeapProfile, MemProfile:
206        profile = pprof.Lookup("heap")
207    case GoroutineProfile:
208        profile = pprof.Lookup("goroutine")
209    case BlockProfile:
210        profile = pprof.Lookup("block")
211    case MutexProfile:
212        profile = pprof.Lookup("mutex")
213    case ThreadProfile:
214        profile = pprof.Lookup("threadcreate")
215    case AllocProfile:
216        profile = pprof.Lookup("allocs")
217    default:
218        return fmt.Errorf("unknown profile type: %v", typ)
219    }
220
221    if profile == nil {
222        return fmt.Errorf("profile not found: %v", typ)
223    }
224
225    return profile.WriteTo(f, 0)
226}
227
228func Stats() ProfileStats {
229    p.mu.Lock()
230    defer p.mu.Unlock()
231    return p.stats
232}
233
234func hasProfileType(typ ProfileType) bool {
235    for _, t := range p.types {
236        if t == typ {
237            return true
238        }
239    }
240    return false
241}
242
243// StartTrace starts execution tracing
244func StartTrace(filename string) error {
245    f, err := os.Create(filename)
246    if err != nil {
247        return err
248    }
249
250    if err := trace.Start(f); err != nil {
251        f.Close()
252        return err
253    }
254
255    p.traceFile = f
256    return nil
257}
258
259// Snapshot takes an immediate snapshot of all profiles
260func Snapshot() error {
261    p.writeSnapshots()
262    return nil
263}
264
265// GetGoroutineStacks returns stack traces of all goroutines
266func GetGoroutineStacks() []byte {
267    buf := make([]byte, 1<<20) // 1MB buffer
268    n := runtime.Stack(buf, true)
269    return buf[:n]
270}

Usage Example

 1package main
 2
 3import (
 4    "context"
 5    "fmt"
 6    "log"
 7    "os"
 8    "os/signal"
 9    "syscall"
10    "time"
11)
12
13func main() {
14    // Create output directory
15    os.MkdirAll("profiles", 0755)
16
17    // Create profiler
18    profiler := profiler.New(
19        []profiler.ProfileType{
20            profiler.CPUProfile,
21            profiler.HeapProfile,
22            profiler.GoroutineProfile,
23            profiler.BlockProfile,
24        },
25        30*time.Second, // Snapshot every 30 seconds
26        "profiles",
27    )
28
29    // Start profiling
30    ctx, cancel := context.WithCancel(context.Background())
31    defer cancel()
32
33    if err := profiler.Start(ctx); err != nil {
34        log.Fatal(err)
35    }
36    defer profiler.Stop()
37
38    // Optional: Start execution trace
39    if err := profiler.StartTrace("profiles/trace.out"); err != nil {
40        log.Printf("Failed to start trace: %v", err)
41    }
42
43    // Run application
44    go runApplication()
45
46    // Monitor stats
47    go func() {
48        ticker := time.NewTicker(10 * time.Second)
49        defer ticker.Stop()
50
51        for {
52            select {
53            case <-ctx.Done():
54                return
55            case <-ticker.C:
56                stats := profiler.Stats()
57                fmt.Printf("Stats: Goroutines=%d MemAlloc=%d MB NumGC=%d\n",
58                    stats.NumGoroutines,
59                    stats.MemAllocated/(1024*1024),
60                    stats.NumGC)
61            }
62        }
63    }()
64
65    // Wait for interrupt
66    sigCh := make(chan os.Signal, 1)
67    signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
68    <-sigCh
69
70    fmt.Println("Shutting down...")
71}
72
73func runApplication() {
74    // Your application logic here
75    for {
76        time.Sleep(time.Second)
77    }
78}

Analyze profiles:

 1# CPU profile
 2go tool pprof profiles/cpu_*.prof
 3
 4# Memory profile
 5go tool pprof profiles/heap_*.prof
 6
 7# Goroutine profile
 8go tool pprof profiles/goroutine_*.prof
 9
10# View in web browser
11go tool pprof -http=:8080 profiles/cpu_*.prof
12
13# Execution trace
14go tool trace profiles/trace.out