| // Licensed to the Apache Software Foundation (ASF) under one or more |
| // contributor license agreements. See the NOTICE file distributed with |
| // this work for additional information regarding copyright ownership. |
| // The ASF licenses this file to You under the Apache License, Version 2.0 |
| // (the "License"); you may not use this file except in compliance with |
| // the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| package appflag |
| |
| import ( |
| "context" |
| "fmt" |
| "os" |
| "time" |
| ) |
| |
| import ( |
| "github.com/pkg/profile" |
| |
| "github.com/spf13/pflag" |
| |
| "go.opentelemetry.io/otel" |
| |
| "go.uber.org/multierr" |
| |
| "go.uber.org/zap" |
| ) |
| |
| import ( |
| "github.com/apache/dubbo-kubernetes/pkg/bufman/pkg/app" |
| "github.com/apache/dubbo-kubernetes/pkg/bufman/pkg/app/applog" |
| "github.com/apache/dubbo-kubernetes/pkg/bufman/pkg/app/appverbose" |
| "github.com/apache/dubbo-kubernetes/pkg/bufman/pkg/observabilityzap" |
| ) |
| |
| type builder struct { |
| appName string |
| |
| verbose bool |
| debug bool |
| noWarn bool |
| logFormat string |
| |
| profile bool |
| profilePath string |
| profileLoops int |
| profileType string |
| profileAllowError bool |
| |
| timeout time.Duration |
| |
| defaultTimeout time.Duration |
| |
| tracing bool |
| } |
| |
| func newBuilder(appName string, options ...BuilderOption) *builder { |
| builder := &builder{ |
| appName: appName, |
| } |
| for _, option := range options { |
| option(builder) |
| } |
| return builder |
| } |
| |
| func (b *builder) BindRoot(flagSet *pflag.FlagSet) { |
| flagSet.BoolVarP(&b.verbose, "verbose", "v", false, "Turn on verbose mode") |
| flagSet.BoolVar(&b.debug, "debug", false, "Turn on debug logging") |
| flagSet.StringVar(&b.logFormat, "log-format", "color", "The log format [text,color,json]") |
| if b.defaultTimeout > 0 { |
| flagSet.DurationVar(&b.timeout, "timeout", b.defaultTimeout, `The duration until timing out, setting it to zero means no timeout`) |
| } |
| |
| flagSet.BoolVar(&b.profile, "profile", false, "Run profiling") |
| _ = flagSet.MarkHidden("profile") |
| flagSet.StringVar(&b.profilePath, "profile-path", "", "The profile base directory path") |
| _ = flagSet.MarkHidden("profile-path") |
| flagSet.IntVar(&b.profileLoops, "profile-loops", 1, "The number of loops to run") |
| _ = flagSet.MarkHidden("profile-loops") |
| flagSet.StringVar(&b.profileType, "profile-type", "cpu", "The profile type [cpu,mem,block,mutex]") |
| _ = flagSet.MarkHidden("profile-type") |
| flagSet.BoolVar(&b.profileAllowError, "profile-allow-error", false, "Allow errors for profiled commands") |
| _ = flagSet.MarkHidden("profile-allow-error") |
| |
| // We do not officially support this flag, this is for testing, where we need warnings turned off. |
| flagSet.BoolVar(&b.noWarn, "no-warn", false, "Turn off warn logging") |
| _ = flagSet.MarkHidden("no-warn") |
| } |
| |
| func (b *builder) NewRunFunc( |
| f func(context.Context, Container) error, |
| interceptors ...Interceptor, |
| ) func(context.Context, app.Container) error { |
| interceptor := chainInterceptors(interceptors...) |
| return func(ctx context.Context, appContainer app.Container) error { |
| if interceptor != nil { |
| return b.run(ctx, appContainer, interceptor(f)) |
| } |
| return b.run(ctx, appContainer, f) |
| } |
| } |
| |
| func (b *builder) run( |
| ctx context.Context, |
| appContainer app.Container, |
| f func(context.Context, Container) error, |
| ) (retErr error) { |
| logLevel, err := getLogLevel(b.debug, b.noWarn) |
| if err != nil { |
| return err |
| } |
| logger, err := applog.NewLogger(appContainer.Stderr(), logLevel, b.logFormat) |
| if err != nil { |
| return err |
| } |
| defer func() { |
| retErr = multierr.Append(retErr, logger.Sugar().Sync()) |
| }() |
| verbosePrinter := appverbose.NewVerbosePrinter(appContainer.Stderr(), b.appName, b.verbose) |
| container, err := newContainer(appContainer, b.appName, logger, verbosePrinter) |
| if err != nil { |
| return err |
| } |
| |
| var cancel context.CancelFunc |
| if !b.profile && b.timeout != 0 { |
| ctx, cancel = context.WithTimeout(ctx, b.timeout) |
| defer cancel() |
| } |
| |
| if b.tracing { |
| closer := observabilityzap.Start(logger) |
| defer func() { |
| retErr = multierr.Append(retErr, closer.Close()) |
| }() |
| _, span := otel.GetTracerProvider().Tracer("bufbuild/buf").Start(ctx, "command") |
| defer span.End() |
| } |
| if !b.profile { |
| return f(ctx, container) |
| } |
| return runProfile( |
| logger, |
| b.profilePath, |
| b.profileType, |
| b.profileLoops, |
| b.profileAllowError, |
| func() error { |
| return f(ctx, container) |
| }, |
| ) |
| } |
| |
| // runProfile profiles the function. |
| func runProfile( |
| logger *zap.Logger, |
| profilePath string, |
| profileType string, |
| profileLoops int, |
| profileAllowError bool, |
| f func() error, |
| ) error { |
| var err error |
| if profilePath == "" { |
| profilePath, err = os.MkdirTemp("", "") |
| if err != nil { |
| return err |
| } |
| } |
| logger.Sugar().Debug("profile", zap.String("path", profilePath)) |
| if profileType == "" { |
| profileType = "cpu" |
| } |
| if profileLoops == 0 { |
| profileLoops = 10 |
| } |
| var profileFunc func(*profile.Profile) |
| switch profileType { |
| case "cpu": |
| profileFunc = profile.CPUProfile |
| case "mem": |
| profileFunc = profile.MemProfile |
| case "block": |
| profileFunc = profile.BlockProfile |
| case "mutex": |
| profileFunc = profile.MutexProfile |
| default: |
| return fmt.Errorf("unknown profile type: %q", profileType) |
| } |
| stop := profile.Start( |
| profile.Quiet, |
| profile.ProfilePath(profilePath), |
| profileFunc, |
| ) |
| for i := 0; i < profileLoops; i++ { |
| if err := f(); err != nil { |
| if !profileAllowError { |
| return err |
| } |
| } |
| } |
| stop.Stop() |
| return nil |
| } |
| |
| func getLogLevel(debugFlag bool, noWarnFlag bool) (string, error) { |
| if debugFlag && noWarnFlag { |
| return "", fmt.Errorf("cannot set both --debug and --no-warn") |
| } |
| if noWarnFlag { |
| return "error", nil |
| } |
| if debugFlag { |
| return "debug", nil |
| } |
| return "info", nil |
| } |
| |
| // chainInterceptors consolidates the given interceptors into one. |
| // The interceptors are applied in the order they are declared. |
| func chainInterceptors(interceptors ...Interceptor) Interceptor { |
| filtered := make([]Interceptor, 0, len(interceptors)) |
| for _, interceptor := range interceptors { |
| if interceptor != nil { |
| filtered = append(filtered, interceptor) |
| } |
| } |
| switch n := len(filtered); n { |
| case 0: |
| return nil |
| case 1: |
| return filtered[0] |
| default: |
| first := filtered[0] |
| return func(next func(context.Context, Container) error) func(context.Context, Container) error { |
| for i := len(filtered) - 1; i > 0; i-- { |
| next = filtered[i](next) |
| } |
| return first(next) |
| } |
| } |
| } |