Flush OpenTelemetry data with a deadline

Issue: PGO-1954

Author: cbandy

cmd/postgres-operator/main.go

@@ -9,8 +9,10 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 	"unicode"
 
@@ -58,8 +60,8 @@ func initLogging() {
 
 //+kubebuilder:rbac:groups="coordination.k8s.io",resources="leases",verbs={get,create,update,watch}
 
-func initManager() (runtime.Options, error) {
-	log := logging.FromContext(context.Background())
+func initManager(ctx context.Context) (runtime.Options, error) {
+	log := logging.FromContext(ctx)
 
 	options := runtime.Options{}
 	options.Cache.SyncPeriod = initialize.Pointer(time.Hour)
@@ -120,45 +122,63 @@ func initManager() (runtime.Options, error) {
 }
 
 func main() {
-	// This context is canceled by SIGINT, SIGTERM, or by calling shutdown.
-	ctx, shutdown := context.WithCancel(runtime.SignalHandler())
-
-	otelFlush, err := initOpenTelemetry()
-	assertNoError(err)
-	defer otelFlush()
+	running, stopRunning := context.WithCancel(context.Background())
+	defer stopRunning()
 
 	initLogging()
-
-	log := logging.FromContext(ctx)
+	log := logging.FromContext(running)
 	log.V(1).Info("debug flag set to true")
 
+	// Start a goroutine that waits for SIGINT or SIGTERM.
+	{
+		signals := []os.Signal{os.Interrupt, syscall.SIGTERM}
+		receive := make(chan os.Signal, len(signals))
+		signal.Notify(receive, signals...)
+		go func() {
+			// Wait for a signal then immediately restore the default signal handlers.
+			// After this, a SIGHUP, SIGINT, or SIGTERM causes the program to exit.
+			// - https://pkg.go.dev/os/signal#hdr-Default_behavior_of_signals_in_Go_programs
+			s := <-receive
+			signal.Stop(receive)
+
+			log.Info("received signal from OS", "signal", s.String())
+			stopRunning()
+		}()
+	}
+
 	features := feature.NewGate()
 	assertNoError(features.Set(os.Getenv("PGO_FEATURE_GATES")))
 
-	ctx = feature.NewContext(ctx, features)
+	running = feature.NewContext(running, features)
 	log.Info("feature gates",
 		// These are set by the user
-		"PGO_FEATURE_GATES", feature.ShowAssigned(ctx),
+		"PGO_FEATURE_GATES", feature.ShowAssigned(running),
 		// These are enabled, including features that are on by default
-		"enabled", feature.ShowEnabled(ctx))
+		"enabled", feature.ShowEnabled(running))
+
+	// Initialize OpenTelemetry and flush data when there is a panic.
+	otelFinish, err := initOpenTelemetry(running)
+	assertNoError(err)
+	defer otelFinish(running)
 
 	cfg, err := runtime.GetConfig()
 	assertNoError(err)
 
 	cfg.Wrap(otelTransportWrapper())
 
+	// TODO(controller-runtime): Set config.WarningHandler instead after v0.19.0.
 	// Configure client-go to suppress warnings when warning headers are encountered. This prevents
 	// warnings from being logged over and over again during reconciliation (e.g. this will suppress
 	// deprecation warnings when using an older version of a resource for backwards compatibility).
 	rest.SetDefaultWarningHandler(rest.NoWarnings{})
 
 	k8s, err := kubernetes.NewDiscoveryRunner(cfg)
 	assertNoError(err)
-	assertNoError(k8s.Read(ctx))
+	assertNoError(k8s.Read(running))
 
-	log.Info("Connected to Kubernetes", "api", k8s.Version().String(), "openshift", k8s.IsOpenShift())
+	log.Info("connected to Kubernetes", "api", k8s.Version().String(), "openshift", k8s.IsOpenShift())
 
-	options, err := initManager()
+	options, err := initManager(running)
 	assertNoError(err)
 
 	// Add to the Context that Manager passes to Reconciler.Start, Runnable.Start,
@@ -174,7 +194,7 @@ func main() {
 	assertNoError(err)
 	assertNoError(mgr.Add(k8s))
 
-	registrar, err := registration.NewRunner(os.Getenv("RSA_KEY"), os.Getenv("TOKEN_PATH"), shutdown)
+	registrar, err := registration.NewRunner(os.Getenv("RSA_KEY"), os.Getenv("TOKEN_PATH"), stopRunning)
 	assertNoError(err)
 	assertNoError(mgr.Add(registrar))
 	token, _ := registrar.CheckToken()
@@ -212,10 +232,30 @@ func main() {
 	assertNoError(mgr.AddHealthzCheck("health", healthz.Ping))
 	assertNoError(mgr.AddReadyzCheck("check", healthz.Ping))
 
-	log.Info("starting controller runtime manager and will wait for signal to exit")
+	// Start the manager and wait for its context to be canceled.
+	stopped := make(chan error, 1)
+	go func() { stopped <- mgr.Start(running) }()
+	<-running.Done()
+
+	// Set a deadline for graceful termination.
+	log.Info("shutting down")
+	stopping, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+
+	// Wait for the manager to return or the deadline to pass.
+	select {
+	case err = <-stopped:
+	case <-stopping.Done():
+		err = stopping.Err()
+	}
 
-	assertNoError(mgr.Start(ctx))
-	log.Info("signal received, exiting")
+	// Flush any telemetry with the remaining time we have.
+	otelFinish(stopping)
+	if err != nil {
+		log.Error(err, "shutdown failed")
+	} else {
+		log.Info("shutdown complete")
+	}
 }
 
 // addControllersToManager adds all PostgreSQL Operator controllers to the provided controller

cmd/postgres-operator/main_test.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"reflect"
 	"testing"
 	"time"
@@ -14,8 +15,10 @@ import (
 )
 
 func TestInitManager(t *testing.T) {
+	ctx := context.Background()
+
 	t.Run("Defaults", func(t *testing.T) {
-		options, err := initManager()
+		options, err := initManager(ctx)
 		assert.NilError(t, err)
 
 		if assert.Check(t, options.Cache.SyncPeriod != nil) {
@@ -48,7 +51,7 @@ func TestInitManager(t *testing.T) {
 		t.Run("Invalid", func(t *testing.T) {
 			t.Setenv("PGO_CONTROLLER_LEASE_NAME", "INVALID_NAME")
 
-			options, err := initManager()
+			options, err := initManager(ctx)
 			assert.ErrorContains(t, err, "PGO_CONTROLLER_LEASE_NAME")
 			assert.ErrorContains(t, err, "invalid")
 
@@ -59,7 +62,7 @@ func TestInitManager(t *testing.T) {
 		t.Run("Valid", func(t *testing.T) {
 			t.Setenv("PGO_CONTROLLER_LEASE_NAME", "valid-name")
 
-			options, err := initManager()
+			options, err := initManager(ctx)
 			assert.NilError(t, err)
 			assert.Assert(t, options.LeaderElection == true)
 			assert.Equal(t, options.LeaderElectionNamespace, "test-namespace")
@@ -70,7 +73,7 @@ func TestInitManager(t *testing.T) {
 	t.Run("PGO_TARGET_NAMESPACE", func(t *testing.T) {
 		t.Setenv("PGO_TARGET_NAMESPACE", "some-such")
 
-		options, err := initManager()
+		options, err := initManager(ctx)
 		assert.NilError(t, err)
 		assert.Assert(t, cmp.Len(options.Cache.DefaultNamespaces, 1),
 			"expected only one configured namespace")
@@ -81,7 +84,7 @@ func TestInitManager(t *testing.T) {
 	t.Run("PGO_TARGET_NAMESPACES", func(t *testing.T) {
 		t.Setenv("PGO_TARGET_NAMESPACES", "some-such,another-one")
 
-		options, err := initManager()
+		options, err := initManager(ctx)
 		assert.NilError(t, err)
 		assert.Assert(t, cmp.Len(options.Cache.DefaultNamespaces, 2),
 			"expect two configured namespaces")
@@ -95,7 +98,7 @@ func TestInitManager(t *testing.T) {
 			for _, v := range []string{"-3", "0", "3.14"} {
 				t.Setenv("PGO_WORKERS", v)
 
-				options, err := initManager()
+				options, err := initManager(ctx)
 				assert.NilError(t, err)
 				assert.DeepEqual(t, options.Controller.GroupKindConcurrency,
 					map[string]int{
@@ -107,7 +110,7 @@ func TestInitManager(t *testing.T) {
 		t.Run("Valid", func(t *testing.T) {
 			t.Setenv("PGO_WORKERS", "19")
 
-			options, err := initManager()
+			options, err := initManager(ctx)
 			assert.NilError(t, err)
 			assert.DeepEqual(t, options.Controller.GroupKindConcurrency,
 				map[string]int{

cmd/postgres-operator/open_telemetry.go

@@ -19,7 +19,7 @@ import (
 	"go.opentelemetry.io/otel/sdk/trace"
 )
 
-func initOpenTelemetry() (func(), error) {
+func initOpenTelemetry(ctx context.Context) (func(context.Context), error) {
 	// At the time of this writing, the SDK (go.opentelemetry.io/[email protected])
 	// does not automatically initialize any exporter. We import the OTLP and
 	// stdout exporters and configure them below. Much of the OTLP exporter can
@@ -49,8 +49,8 @@ func initOpenTelemetry() (func(), error) {
 		}
 
 		provider := trace.NewTracerProvider(trace.WithBatcher(exporter))
-		flush := func() {
-			_ = provider.Shutdown(context.TODO())
+		flush := func(ctx context.Context) {
+			_ = provider.Shutdown(ctx)
 			if closer != nil {
 				_ = closer.Close()
 			}
@@ -67,8 +67,8 @@ func initOpenTelemetry() (func(), error) {
 		}
 
 		provider := trace.NewTracerProvider(trace.WithBatcher(exporter))
-		flush := func() {
-			_ = provider.Shutdown(context.TODO())
+		flush := func(ctx context.Context) {
+			_ = provider.Shutdown(ctx)
 		}
 
 		otel.SetTracerProvider(provider)
@@ -78,7 +78,7 @@ func initOpenTelemetry() (func(), error) {
 	// $OTEL_TRACES_EXPORTER is unset or unknown, so no TracerProvider has been assigned.
 	// The default at this time is a single "no-op" tracer.
 
-	return func() {}, nil
+	return func(context.Context) {}, nil
 }
 
 // otelTransportWrapper creates a function that wraps the provided net/http.RoundTripper

internal/controller/runtime/runtime.go

@@ -5,16 +5,13 @@
 package runtime
 
 import (
-	"context"
-
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/kubernetes/scheme"
 	"k8s.io/client-go/rest"
 	"sigs.k8s.io/controller-runtime/pkg/cache"
 	"sigs.k8s.io/controller-runtime/pkg/client/config"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
-	"sigs.k8s.io/controller-runtime/pkg/manager/signals"
 
 	"github.com/crunchydata/postgres-operator/internal/logging"
 	"github.com/crunchydata/postgres-operator/pkg/apis/postgres-operator.crunchydata.com/v1beta1"
@@ -71,6 +68,3 @@ func NewManager(config *rest.Config, options manager.Options) (manager.Manager,
 
 // SetLogger assigns the default Logger used by [sigs.k8s.io/controller-runtime].
 func SetLogger(logger logging.Logger) { log.SetLogger(logger) }
-
-// SignalHandler returns a Context that is canceled on SIGINT or SIGTERM.
-func SignalHandler() context.Context { return signals.SetupSignalHandler() }