gcp/observability: fix End() to cleanup global state correctly

gcp/observability/observability.go

@@ -29,12 +29,7 @@ import (
 	"context"
 	"fmt"
 
-	"contrib.go.opencensus.io/exporter/stackdriver"
-
-	"go.opencensus.io/stats/view"
-	"go.opencensus.io/trace"
 	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal"
 )
 
 var logger = grpclog.Component("observability")
@@ -85,18 +80,5 @@ func Start(ctx context.Context) error {
 // Note: this method should only be invoked once.
 func End() {
 	defaultLogger.Close()
-	if exporter != nil {
-		if sdExporter, ok := exporter.(*stackdriver.Exporter); ok {
-			sdExporter.Flush()
-			sdExporter.Close()
-		}
-
-		// Call these unconditionally, doesn't matter if not registered, will be
-		// a noop if not registered.
-		trace.UnregisterExporter(exporter)
-		view.UnregisterExporter(exporter)
-
-		internal.ClearExtraDialOptions()
-		internal.ClearExtraServerOptions()
-	}
+	stopOpenCensus()
 }

gcp/observability/opencensus.go

@@ -121,3 +121,22 @@ func startOpenCensus(config *config) error {
 
 	return nil
 }
+
+// stopOpenCensus flushes the exporter's and cleans up globals across all
+// packages if exporter was created.
+func stopOpenCensus() {
+	if exporter != nil {
+		if sdExporter, ok := exporter.(*stackdriver.Exporter); ok {
+			sdExporter.Flush()
+			sdExporter.Close()
+		}
+
+		// Call these unconditionally, doesn't matter if not registered, will be
+		// a noop if not registered.
+		trace.UnregisterExporter(exporter)
+		view.UnregisterExporter(exporter)
+
+		internal.ClearExtraDialOptions()
+		internal.ClearExtraServerOptions()
+	}
+}