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

.github/workflows/testing.yml

@@ -109,7 +109,7 @@ jobs:
           cd "${GITHUB_WORKSPACE}"
           for MOD_FILE in $(find . -name 'go.mod' | grep -Ev '^\./go\.mod'); do
             pushd "$(dirname ${MOD_FILE})"
-            go test ${{ matrix.testflags }} -timeout 2m ./...
+            go test ${{ matrix.testflags }} -cpu 1,4 -timeout 2m ./...
             popd
           done
 

gcp/observability/observability.go

@@ -29,7 +29,12 @@ 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")
@@ -80,4 +85,18 @@ 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()
+	}
 }

gcp/observability/observability_test.go

@@ -35,7 +35,6 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	grpclogrecordpb "google.golang.org/grpc/gcp/observability/internal/logging"
-	"google.golang.org/grpc/internal"
 	iblog "google.golang.org/grpc/internal/binarylog"
 	"google.golang.org/grpc/internal/grpctest"
 	"google.golang.org/grpc/internal/leakcheck"
@@ -875,12 +874,6 @@ func (s) TestCustomTagsTracingMetrics(t *testing.T) {
 	cleanup, err := createTmpConfigInFileSystem(configJSON)
 	defer cleanup()
 
-	// To clear globally registered tracing and metrics exporters.
-	defer func() {
-		internal.ClearExtraDialOptions()
-		internal.ClearExtraServerOptions()
-	}()
-
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 	err = Start(ctx)

gcp/observability/opencensus.go

@@ -58,6 +58,8 @@ type tracingMetricsExporter interface {
 	view.Exporter
 }
 
+var exporter tracingMetricsExporter
+
 // global to stub out in tests
 var newExporter = newStackdriverExporter
 
@@ -87,7 +89,8 @@ func startOpenCensus(config *config) error {
 		return nil
 	}
 
-	exporter, err := newExporter(config)
+	var err error
+	exporter, err = newExporter(config)
 	if err != nil {
 		return err
 	}