diff --git a/cmd/webserver/main.go b/cmd/webserver/main.go index c46945f..d363a13 100644 --- a/cmd/webserver/main.go +++ b/cmd/webserver/main.go @@ -102,8 +102,8 @@ func main() { recvSignal := <-sig l.Logger.Info().Str("received_signal", recvSignal.String()).Msg("Performing a graceful shutdown") + pprofServer.Shutdown(ctx) httpServer.Shutdown(ctx) - pprofServer.Stop(ctx) l.Logger.Info().Msgf("Service %s exiting", serviceName) } diff --git a/pkg/profiling/http.go b/pkg/profiling/http.go index e493016..ed92cdd 100644 --- a/pkg/profiling/http.go +++ b/pkg/profiling/http.go @@ -27,6 +27,7 @@ func NewPprofServer() *PprofServer { listenAddress := net.JoinHostPort(hostname, port) return &PprofServer{ listenAddr: listenAddress, + once: &sync.Once{}, server: &http.Server{ Addr: listenAddress, Handler: http.DefaultServeMux, @@ -45,7 +46,7 @@ func (p *PprofServer) Start() { }() } -func (p *PprofServer) Stop(ctx context.Context) { +func (p *PprofServer) Shutdown(ctx context.Context) { p.once.Do(func() { cCtx, cancel := context.WithTimeoutCause( ctx, diff --git a/pkg/tracer/otel.go b/pkg/tracer/otel.go index 9b6cbb8..e40bb17 100644 --- a/pkg/tracer/otel.go +++ b/pkg/tracer/otel.go @@ -6,6 +6,7 @@ import ( "fmt" "net" "sync" + "time" l "weezel/example-gin/pkg/logger" @@ -59,6 +60,7 @@ func NewOtelTracerMetrics( serviceName: serviceName, res: res, modes: modes, + closeOnce: &sync.Once{}, } otelTracerMetrics.connection, err = grpc.NewClient( @@ -98,14 +100,19 @@ func (o *OtelTracerMetrics) Close(ctx context.Context) { var connErr error var tracerErr error var metricsErr error - if o.connection != nil { - connErr = o.connection.Close() - } + + timeout := 4 * time.Second + cCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + if o.tracer != nil { - tracerErr = o.tracer.Shutdown(ctx) + tracerErr = o.tracer.Shutdown(cCtx) } if o.metrics != nil { - metricsErr = o.metrics.Shutdown(ctx) + metricsErr = o.metrics.Shutdown(cCtx) + } + if o.connection != nil { + connErr = o.connection.Close() } errs := errors.Join(connErr, tracerErr, metricsErr) if errs != nil {