Add profiling, integrate into Butler/Annotee.

logdog/client/cmd/logdog_butler/main.go

 // Copyright 2015 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package main
 
 import (
         "bytes"
         "flag"
         "fmt"
         "net/http"
         "os"
         "os/signal"
-        "runtime/pprof"
         "sort"
         "strings"
         "time"
 
         "github.com/maruel/subcommands"
         "golang.org/x/net/context"
 
         "github.com/luci/luci-go/client/authcli"
         "github.com/luci/luci-go/common/auth"
         "github.com/luci/luci-go/common/cli"
         "github.com/luci/luci-go/common/clock/clockflag"
         "github.com/luci/luci-go/common/config"
         "github.com/luci/luci-go/common/errors"
         "github.com/luci/luci-go/common/flag/multiflag"
         log "github.com/luci/luci-go/common/logging"
         "github.com/luci/luci-go/common/logging/gologger"
         "github.com/luci/luci-go/common/runtime/paniccatcher"
+        "github.com/luci/luci-go/common/runtime/profiling"
         grpcLogging "github.com/luci/luci-go/grpc/logging"
         "github.com/luci/luci-go/logdog/client/butler"
         "github.com/luci/luci-go/logdog/client/butler/output"
         "github.com/luci/luci-go/logdog/common/types"
 )
 
 const (
         // flagErrorReturnCode is returned when there is an error with the Butler's
         // command-line configuration.
         configErrorReturnCode = 2
@@ skipping 17 matching lines @@
         project       config.ProjectName
         prefix        types.StreamName
         outputWorkers int
         outputConfig  outputConfigFlag
 
         authFlags authcli.Flags
 
         maxBufferAge clockflag.Duration
         noBufferLogs bool
 
-        cpuProfile string
+        prof profiling.Profiler
 
         client *http.Client
 
         // ncCtx is a context that will not be cancelled when cancelFunc is called.
         ncCtx      context.Context
         cancelFunc func()
 }
 
 func (a *application) addFlags(fs *flag.FlagSet) {
         a.outputConfig.Output = os.Stdout
         a.outputConfig.Description = "Select and configure message output adapter."
         a.outputConfig.Options = []multiflag.Option{
                 multiflag.HelpOption(&a.outputConfig.MultiFlag),
         }
 
         // Add registered conditional (build tag) options.
         for _, f := range getOutputFactories() {
                 a.outputConfig.AddFactory(f)
         }
 
         a.maxBufferAge = clockflag.Duration(butler.DefaultMaxBufferAge)
 
         fs.Var(&a.project, "project",
                 "The log prefix's project name (required).")
         fs.Var(&a.prefix, "prefix",
                 "Prefix to apply to all stream names.")
         fs.Var(&a.outputConfig, "output",
                 "The output name and configuration. Specify 'help' for more information.")
-        fs.StringVar(&a.cpuProfile,
-                "cpuprofile", "", "If specified, enables CPU profiling and profiles to the specified path.")
         fs.IntVar(&a.outputWorkers, "output-workers", butler.DefaultOutputWorkers,
                 "The maximum number of parallel output dispatches.")
         fs.Var(&a.maxBufferAge, "output-max-buffer-age",
                 "Send buffered messages if they've been held for longer than this period.")
         fs.BoolVar(&a.noBufferLogs, "output-no-buffer", false,
                 "If true, dispatch logs immediately. Setting this flag simplifies output at the expense "+
                         "of wire-format efficiency.")
 }
 
 func (a *application) authenticator(ctx context.Context) (*auth.Authenticator, error) {
@@ skipping 10 matching lines @@
         if factory == nil {
                 return nil, errors.New("main: No output is configured")
         }
         return factory, nil
 }
 
 // runWithButler is an execution harness that adds application-level management
 // to a Butler run.
 func (a *application) runWithButler(out output.Output, runFunc func(*butler.Butler) error) error {
 
-        // Enable CPU profiling if specified
-        if a.cpuProfile != "" {
-                f, err := os.Create(a.cpuProfile)
-                if err != nil {
-                        return fmt.Errorf("failed to create CPU profile output: %v", err)
-                }
-                pprof.StartCPUProfile(f)
-                defer pprof.StopCPUProfile()
+        // Start our Profiler.
+        a.prof.Logger = log.Get(a)
+        if err := a.prof.Start(); err != nil {
+                return fmt.Errorf("failed to start Profiler: %v", err)
         }
+        defer a.prof.Stop()
 
         // Instantiate our Butler.
         butlerOpts := butler.Config{
                 Project:       a.project,
                 Prefix:        a.prefix,
                 MaxBufferAge:  time.Duration(a.maxBufferAge),
                 BufferLogs:    !a.noBufferLogs,
                 Output:        out,
                 OutputWorkers: a.outputWorkers,
                 TeeStdout:     os.Stdout,
@@ skipping 96 matching lines @@
                 },
         }
         // Install logging configuration flags.
         flags := flag.NewFlagSet("flags", flag.ExitOnError)
         logConfig := log.Config{
                 Level: log.Warning,
         }
         logConfig.AddFlags(flags)
         a.addFlags(flags)
         a.authFlags.Register(flags, authOptions)
+        a.prof.AddFlags(flags)
 
         // Parse the top-level flag set.
         if err := flags.Parse(argv); err != nil {
                 log.WithError(err).Errorf(a, "Failed to parse command-line.")
                 return configErrorReturnCode
         }
 
         a.Context = logConfig.Set(a.Context)
 
         // Install a global gRPC logger adapter. This routes gRPC log messages that
@@ skipping 74 matching lines @@
 
         paniccatcher.Do(func() {
                 rc = mainImpl(ctx, os.Args[1:])
         }, func(p *paniccatcher.Panic) {
                 log.Fields{
                         "panic.error": p.Reason,
                 }.Errorf(ctx, "Panic caught in main:\n%s", p.Stack)
                 rc = runtimeErrorReturnCode
         })
 }