WIP

open-policy-agent · May 7, 2024 · 8a11868 · 8a11868
1 parent 493ba40
commit 8a11868
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 6 deletions.
diff --git a/runtime/runtime.go b/runtime/runtime.go
@@ -542,6 +542,13 @@ func (rt *Runtime) Serve(ctx context.Context) error {
 		}
 	}
 
+	defer func() {
+		err := rt.traceExporter.Shutdown(ctx)
+		if err != nil {
+			rt.logger.WithFields(map[string]interface{}{"err": err}).Error("Failed to shutdown OpenTelemetry trace exporter gracefully.")
+		}
+	}()
+
 	rt.server = server.New().
 		WithRouter(rt.Params.Router).
 		WithStore(rt.Store).
@@ -857,12 +864,12 @@ func (rt *Runtime) gracefulServerShutdown(s *server.Server) error {
 	}
 	rt.logger.Info("Server shutdown.")
 
-	if rt.traceExporter != nil {
-		err = rt.traceExporter.Shutdown(ctx)
-		if err != nil {
-			rt.logger.WithFields(map[string]interface{}{"err": err}).Error("Failed to shutdown OpenTelemetry trace exporter gracefully.")
-		}
-	}
+	//if rt.traceExporter != nil {
+	//	err = rt.traceExporter.Shutdown(ctx)
+	//	if err != nil {
+	//		rt.logger.WithFields(map[string]interface{}{"err": err}).Error("Failed to shutdown OpenTelemetry trace exporter gracefully.")
+	//	}
+	//}
 	return nil
 }
 

diff --git a/runtime/runtime_test.go b/runtime/runtime_test.go
@@ -9,6 +9,10 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"github.com/open-policy-agent/opa/tracing"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -1235,6 +1239,43 @@ func TestServerInitializedWithBundleRegoVersion(t *testing.T) {
 	}
 }
 
+func TestGracefulTracerShutdown(t *testing.T) {
+	ctx := context.Background()
+
+	logger := logging.New()
+	stdout := bytes.NewBuffer(nil)
+	logger.SetOutput(stdout)
+	logger.SetLevel(logging.Error)
+
+	spanExp := tracetest.NewInMemoryExporter()
+	options := tracing.NewOptions(
+		otelhttp.WithTracerProvider(trace.NewTracerProvider(trace.WithSpanProcessor(trace.NewSimpleSpanProcessor(spanExp)))),
+	)
+
+	params := NewParams()
+	params.Addrs = &[]string{"localhost:0"}
+	params.GracefulShutdownPeriod = 1
+	params.Logger = logger
+
+	params.DistributedTracingOpts = options
+
+	rt, err := NewRuntime(ctx, params)
+	if err != nil {
+		t.Fatalf("Unexpected error %v", err)
+	}
+	err = rt.Serve(ctx)
+	if err != nil {
+		t.Fatalf("Unexpected error %v", err)
+	}
+
+	ctx.Done()
+
+	expected := "Failed to shutdown OpenTelemetry trace exporter gracefully."
+	if strings.Contains(stdout.String(), expected) {
+		t.Fatalf("Expected no output containing: \"%v\"", expected)
+	}
+}
+
 func TestUrlPathToConfigOverride(t *testing.T) {
 	params := NewParams()
 	params.Paths = []string{"https://www.example.com/bundles/bundle.tar.gz"}