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

* 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

@@ -80,4 +80,5 @@ func Start(ctx context.Context) error {
 // Note: this method should only be invoked once.
 func End() {
 	defaultLogger.Close()
+	stopOpenCensus()
 }

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"
@@ -387,6 +386,12 @@ func (fe *fakeOpenCensusExporter) ExportSpan(vd *trace.SpanData) {
 	fe.t.Logf("Span[%v]", vd.Name)
 }
 
+func (fe *fakeOpenCensusExporter) Flush() {}
+
+func (fe *fakeOpenCensusExporter) Close() error {
+	return nil
+}
+
 func (s) TestLoggingForOkCall(t *testing.T) {
 	te := newTest(t)
 	defer te.tearDown()
@@ -875,12 +880,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

@@ -56,8 +56,12 @@ func tagsToTraceAttributes(tags map[string]string) map[string]interface{} {
 type tracingMetricsExporter interface {
 	trace.Exporter
 	view.Exporter
+	Flush()
+	Close() error
 }
 
+var exporter tracingMetricsExporter
+
 // global to stub out in tests
 var newExporter = newStackdriverExporter
 
@@ -87,7 +91,8 @@ func startOpenCensus(config *config) error {
 		return nil
 	}
 
-	exporter, err := newExporter(config)
+	var err error
+	exporter, err = newExporter(config)
 	if err != nil {
 		return err
 	}
@@ -118,3 +123,20 @@ 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 {
+		internal.ClearExtraDialOptions()
+		internal.ClearExtraServerOptions()
+
+		// Call these unconditionally, doesn't matter if not registered, will be
+		// a noop if not registered.
+		trace.UnregisterExporter(exporter)
+		view.UnregisterExporter(exporter)
+
+		exporter.Flush()
+		exporter.Close()
+	}
+}