Skip to content

Instantly share code, notes, and snippets.

@Petelin
Created November 18, 2019 11:51
Show Gist options
  • Save Petelin/0aabf793a15fa86cfa44cbe74cfbc417 to your computer and use it in GitHub Desktop.
Save Petelin/0aabf793a15fa86cfa44cbe74cfbc417 to your computer and use it in GitHub Desktop.
上面是一个middleware实现, 最终cpu信息会记录到profMap中. (要多次调用接口才能看到效果).
func Handler(meta endpoint.Meta, next endpoint.Endpoint) endpoint.Endpoint {
profMapLock.Lock()
defer profMapLock.Unlock()
traceBytes := []byte{}
tracePtr := &traceBytes
var lock sync.Mutex
var buf bytes.Buffer
var traceBuf bytes.Buffer
var hasMerged bool
setEnabled()
// Prepare buffers for this endpoint
merged := new(profile.Profile)
profKey := meta.Service + "." + meta.Name
profMap[profKey] = &profileDump{
profile: merged,
trace: tracePtr,
}
profWriter := bufio.NewWriter(&buf)
traceWriter := bufio.NewWriter(&traceBuf)
return func(ctx context.Context, request interface{}) (metadata metadata.Metadata, response interface{}, err error) {
// Lock to ensure only one endpoint is profiled at a time
lock.Lock()
defer lock.Unlock()
// Profile just this endpoint
if err := pprof.StartCPUProfile(profWriter); err != nil {
// If already enabled, don't profile
return next(ctx, request)
}
if err := trace.Start(traceWriter); err != nil {
// If already enabled, don't trace
return next(ctx, request)
}
metadata, response, err = next(ctx, request)
trace.Stop()
_ = traceWriter.Flush()
pprof.StopCPUProfile()
_ = profWriter.Flush()
*tracePtr = traceBuf.Bytes()
// Parse and merge profiles ready for converting to flamegraph
// This is done eagerly but could be deferred until the flamegraph is called
thisProf, parseErr := profile.Parse(bufio.NewReader(&buf))
buf.Reset()
traceBuf.Reset()
if parseErr == nil {
if hasMerged {
m, merr := profile.Merge([]*profile.Profile{merged, thisProf})
if merr == nil {
*merged = *m
}
return
}
*merged = *thisProf
hasMerged = true
}
return
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment