diff --git a/uptracehelper/helper.go b/uptracehelper/helper.go index 31706f3..1e7d02e 100644 --- a/uptracehelper/helper.go +++ b/uptracehelper/helper.go @@ -2,6 +2,7 @@ package uptracehelper import ( "context" + "time" "github.com/uptrace/uptrace-go/uptrace" ) @@ -10,6 +11,7 @@ type Config struct { ServiceName string ServiceVersion string DeploymentEnvironment string + ShutdownTimeout time.Duration } func Setup(cfg Config) { @@ -23,19 +25,21 @@ func Setup(cfg Config) { uptrace.ConfigureOpentelemetry(opts...) } -// GoStop 初始化并等待 uptrace +// Bootstrap 初始化并等待 uptrace // 配合 [git.blauwelle.com/go/crate/exegroup] 使用 // env // - UPTRACE_DISABLED 存在就不再初始化 // - UPTRACE_DSN 服务端地址 -func GoStop(cfg Config) (func(ctx context.Context) error, func(ctx context.Context)) { +func Bootstrap(cfg Config) func(ctx context.Context) error { Setup(cfg) - shutdownErr := make(chan error, 1) - goFunc := func(context.Context) error { - return <-shutdownErr + return func(ctx context.Context) error { + <-ctx.Done() + shutdownCtx := context.Background() + if cfg.ShutdownTimeout >= 0 { + var cancel func() + shutdownCtx, cancel = context.WithTimeout(shutdownCtx, cfg.ShutdownTimeout) + defer cancel() + } + return uptrace.Shutdown(shutdownCtx) } - stopFunc := func(ctx context.Context) { - shutdownErr <- uptrace.Shutdown(ctx) - } - return goFunc, stopFunc }