go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/cli/profile.go (about)

     1  // Copyright 2017 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  //go:build include_profiler
    16  // +build include_profiler
    17  
    18  package cli
    19  
    20  import (
    21  	"context"
    22  
    23  	"github.com/maruel/subcommands"
    24  
    25  	"go.chromium.org/luci/common/clock"
    26  	"go.chromium.org/luci/common/logging"
    27  	"go.chromium.org/luci/common/runtime/profiling"
    28  )
    29  
    30  type profilingExt struct {
    31  	added bool
    32  }
    33  
    34  func (p *profilingExt) addProfiling(cmds []*subcommands.Command) {
    35  	if !p.added {
    36  		for _, cmd := range cmds {
    37  			cmd.CommandRun = wrapCmdRun(cmd.CommandRun)
    38  		}
    39  		p.added = true
    40  	}
    41  }
    42  
    43  func wrapCmdRun(orig func() subcommands.CommandRun) func() subcommands.CommandRun {
    44  	return func() subcommands.CommandRun {
    45  		r := &wrappedCmdRun{CommandRun: orig()}
    46  		r.prof.AddFlags(r.GetFlags())
    47  		return r
    48  	}
    49  }
    50  
    51  type wrappedCmdRun struct {
    52  	subcommands.CommandRun
    53  	prof profiling.Profiler
    54  }
    55  
    56  // Run is part of CommandRun interface.
    57  func (r *wrappedCmdRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
    58  	ctx := GetContext(a, r, env)
    59  	r.prof.Logger = logging.Get(ctx)
    60  	r.prof.Clock = clock.Get(ctx)
    61  
    62  	if err := r.prof.Start(); err != nil {
    63  		logging.WithError(err).Errorf(ctx, "Failed to start profiling")
    64  		return 1
    65  	}
    66  	defer r.prof.Stop()
    67  
    68  	return r.CommandRun.Run(a, args, env)
    69  }
    70  
    71  // ModifyContext is part of ContextModificator interface.
    72  //
    73  // Need to explicitly define it, since embedding original CommandRun in
    74  // wrappedCmdRun "disables" the sniffing of ContextModificator in GetContext.
    75  func (r *wrappedCmdRun) ModifyContext(ctx context.Context) context.Context {
    76  	if m, _ := r.CommandRun.(ContextModificator); m != nil {
    77  		return m.ModifyContext(ctx)
    78  	}
    79  	return ctx
    80  }