Repository: knative/pkg (go)

Issue:
Make sure OpenCensusTracer from tracing package can be correctly shut down (flushed)
The problem is that currently the components that use [SetupStaticPublishing](https://github.com/knative/pkg/blob/main/tracing/setup.go#L32) or [SetupDynamicPublishing](https://github.com/knative/pkg/blob/main/tracing/setup.go#L45) can drop some Spans/Traces when being shut down or evicted (e.g. during upgrades).
The goal is to prevent dropping buffered traces before they're sent as indicated by [upstream issue](https://github.com/census-instrumentation/opencensus-go/issues/862) in OpenCensus.

The OpenCensusTracer has a [Finish](https://github.com/knative/pkg/blob/main/tracing/opencensus.go#L84) function that can be called to properly [Close](https://github.com/knative/pkg/blob/main/vendor/github.com/openzipkin/zipkin-go/reporter/http/http.go#L71) the OpenZipkin http reporter and flush buffered traces. However, the SetupStaticPublishing/SetupDynamicPublishing functions don't make the OpenCensusTracer available so that anyone can call the Finish function. Another problem is that the Finish function [passes nil](https://github.com/knative/pkg/blob/main/tracing/opencensus.go#L92) which triggers nil pointer exception when calling the Finish function.

I've verified that calling the Finish function on the returned OpenCensusTracer sends all remaining traces and makes them available in Zipkin backend. It is only required to pass `&config.Config{}` on [this line](https://github.com/knative/pkg/blob/main/tracing/opencensus.go#L92) instead of nil.

It would be good to enforce calling Finish before components are shutdown as long as there's no sleep time in that component by default before it is shut down or evicted (I think some components have the sleep time such as Knative Serving Activator).

Related discussion [here](https://github.com/knative/eventing/pull/6219#discussion_r836446571) and [here](https://github.com/knative/eventing/pull/6219#discussion_r835613964)

<!--
Pro-tip: You can leave this block commented, and it still works!

Select the appropriate areas for your issue:

/area API
/area monitoring
/area networking
/area test-and-release

Classify what kind of issue this is:

/kind bug

-->

## Expected Behavior

It is possible to flush traces on demand (possibly before shutting down a component).

## Actual Behavior

Traces that were buffered before the component is shutdown are not properly sent to tracing backend.

## Additional Info


Pull request:
SetupStaticPublishing and SetupDynamicPublishing returns Tracer with Shutdown function
…acer

* The caller can then call Shutdown() to properly flush buffered Spans
before shutdown

This is incompatible change that will require updates in many repositories. I can either prepare those PRs so that they can be merged when this change propagates to those repositories or I can keep the old implementation and also provide a new impl that returns the OpenCensusTracer. Whatever works best.
Initially, I will provide an example PR for knative/eventing that will show the required updates.

<!-- Thanks for sending a pull request! -->

# Changes

<!-- 
Describe your changes here- ideally you can get that description straight from
your descriptive commit message(s)! 

- :gift: Add new feature
- :bug: Fix bug
- :broom: Update or clean up current behavior
- :wastebasket: Remove feature or internal logic
-->


<!--
In addition, categorize the changes you're making using the "/kind" Prow command, example:

/kind enhancement

Supported kinds are: api-change, bug, cleanup, deprecation, removal, documentation, enhancement, performance

-->
/kind enhancement

<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #2475

<!-- Please include the 'why' behind your changes if no issue exists -->

**Release Note**

<!--
:page_facing_up: If this change has user-visible impact, write a release note in the block
below. Include the string "action required" if additional action is required of
users switching to the new release, for example in case of a breaking change.

Write as if you are speaking to users, not other Knative contributors. If this
change has no user-visible impact, no release-note is needed.
-->
```release-note

```

**Docs**

<!--
:book: If this change has user-visible impact, link to an issue or PR in
https://github.com/knative/docs.

See here for guidance on getting permanent links to files: https://help.github.com/en/articles/getting-permanent-links-to-files

Please use the following format for linking documentation:
- [knative/docs]: <issue or pr link>
- [Feature Track]: <link>
- [Usage]: <link>
- [Other doc]: <link>
-->
```docs

```


Patch:
diff --git a/.github/workflows/knative-releasability.yaml b/.github/workflows/knative-releasability.yaml
index 815a31eb3..930ad8e62 100644
--- a/.github/workflows/knative-releasability.yaml
+++ b/.github/workflows/knative-releasability.yaml
@@ -25,8 +25,8 @@ jobs:
   releasability:
     uses: knative/actions/.github/workflows/releasability.yaml@main
     with:
-      releaseFamily: ${{ github.event.inputs.releaseFamily || 'v1.6' }}
-      moduleReleaseFamily: ${{ github.event.inputs.moduleReleaseFamily || 'v0.33' }}
-      slackChannel: ${{ github.event.inputs.slackChannel || 'release-1dot6' }}
+      releaseFamily: ${{ github.event.inputs.releaseFamily || 'v1.7' }}
+      moduleReleaseFamily: ${{ github.event.inputs.moduleReleaseFamily || 'v0.34' }}
+      slackChannel: ${{ github.event.inputs.slackChannel || 'release' }}
     secrets:
       SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }}
diff --git a/OWNERS_ALIASES b/OWNERS_ALIASES
index 529327dc1..196409e33 100644
--- a/OWNERS_ALIASES
+++ b/OWNERS_ALIASES
@@ -67,9 +67,9 @@ aliases:
   - knative-test-reporter-robot
   - kvmware
   - lance
+  - matzew
   - n3wscott
   - pmorie
-  - psschwei
   - smoser-ibm
   - spencerdillard
   - thisisnotapril
@@ -77,8 +77,8 @@ aliases:
   - vaikas
   - zroubalik
   knative-release-leads:
-  - evankanderson
-  - psschwei
+  - dprotaso
+  - matzew
   knative-robots:
   - knative-automation
   - knative-prow-releaser-robot
@@ -179,6 +179,7 @@ aliases:
   - whaught
   serving-wg-leads:
   - dprotaso
+  - psschwei
   serving-writers:
   - dprotaso
   - julz
diff --git a/pool/OWNERS b/pool/OWNERS
deleted file mode 100644
index 758a612a0..000000000
--- a/pool/OWNERS
+++ /dev/null
@@ -1,10 +0,0 @@
-# The OWNERS file is used by prow to automatically merge approved PRs.
-
-approvers:
-- pkg-controller-writers
-
-reviewers:
-- pkg-controller-reviewers
-
-labels:
-- area/API
diff --git a/pool/doc.go b/pool/doc.go
deleted file mode 100644
index 052b9ee93..000000000
--- a/pool/doc.go
+++ /dev/null
@@ -1,21 +0,0 @@
-/*
-Copyright 2020 The Knative Authors
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    https://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-// Package pool contains a simple threadpool implementation that accepts
-// work in the form of `func() error` function.  The intent is for it to
-// support similar workloads to errgroup, but with a maximum number of
-// concurrent worker threads.
-package pool
diff --git a/pool/interface.go b/pool/interface.go
deleted file mode 100644
index 86cc8e48c..000000000
--- a/pool/interface.go
+++ /dev/null
@@ -1,36 +0,0 @@
-/*
-Copyright 2020 The Knative Authors
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    https://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package pool
-
-import (
-	"golang.org/x/sync/errgroup"
-)
-
-// Interface defines an errgroup-compatible interface for interacting with
-// our threadpool.
-type Interface interface {
-	// Go queues a single unit of work for execution on this pool. All calls
-	// to Go must be finished before Wait is called.
-	Go(func() error)
-
-	// Wait blocks until all work is complete, returning the first
-	// error returned by any of the work.
-	Wait() error
-}
-
-// errgroup.Group implements Interface
-var _ Interface = (*errgroup.Group)(nil)
diff --git a/pool/pool.go b/pool/pool.go
deleted file mode 100644
index 9f298d41b..000000000
--- a/pool/pool.go
+++ /dev/null
@@ -1,117 +0,0 @@
-/*
-Copyright 2020 The Knative Authors
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    https://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package pool
-
-import (
-	"context"
-	"sync"
-)
-
-type impl struct {
-	wg     sync.WaitGroup
-	workCh chan func() error
-
-	// Ensure that we Wait exactly once and memoize
-	// the result.
-	waitOnce sync.Once
-
-	cancel context.CancelFunc
-
-	// We're only interested in the first result so
-	// only set it once.
-	resultOnce sync.Once
-	result     error
-}
-
-// impl implements Interface
-var _ Interface = (*impl)(nil)
-
-// defaultCapacity is the number of work items or errors that we
-// can queue up before calls to Go will block, or work will
-// block until Wait is called.
-const defaultCapacity = 50
-
-// New creates a fresh worker pool with the specified size.
-func New(workers int) Interface {
-	return NewWithCapacity(workers, defaultCapacity)
-}
-
-// NewWithCapacity creates a fresh worker pool with the specified size.
-func NewWithCapacity(workers, capacity int) Interface {
-	i, _ := NewWithContext(context.Background(), workers, capacity)
-	return i
-}
-
-// NewWithContext creates a pool that is driven by a cancelable context.
-// Just like errgroup.Group on first error the context will be canceled as well.
-func NewWithContext(ctx context.Context, workers, capacity int) (Interface, context.Context) {
-	ctx, cancel := context.WithCancel(ctx)
-	i := &impl{
-		cancel: cancel,
-		workCh: make(chan func() error, capacity),
-	}
-
-	// Start a go routine for each worker, which:
-	// 1. reads off of the work channel,
-	// 2. (optionally) sets the error as the result,
-	// 3. marks work as done in our sync.WaitGroup.
-	for idx := 0; idx < workers; idx++ {
-		go func() {
-			for work := range i.workCh {
-				i.exec(work)
-			}
-		}()
-	}
-	return i, ctx
-}
-
-func (i *impl) exec(w func() error) {
-	defer i.wg.Done()
-	if err := w(); err != nil {
-		i.resultOnce.Do(func() {
-			if i.cancel != nil {
-				i.cancel()
-			}
-			i.result = err
-		})
-	}
-}
-
-// Go implements Interface.
-func (i *impl) Go(w func() error) {
-	// Increment the amount of outstanding work we're waiting on.
-	i.wg.Add(1)
-	// Send the work along the queue.
-	i.workCh <- w
-}
-
-// Wait implements Interface.
-func (i *impl) Wait() error {
-	i.waitOnce.Do(func() {
-		// Wait for queued work to complete.
-		i.wg.Wait()
-		// Notify the context, that it's done now.
-		if i.cancel != nil {
-			i.cancel()
-		}
-
-		// Now we know there are definitely no new items arriving.
-		close(i.workCh)
-	})
-
-	return i.result
-}
diff --git a/tracing/opencensus.go b/tracing/opencensus.go
index 9793e5c0f..7471c1336 100644
--- a/tracing/opencensus.go
+++ b/tracing/opencensus.go
@@ -17,6 +17,7 @@ limitations under the License.
 package tracing
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -81,7 +82,12 @@ func (oct *OpenCensusTracer) ApplyConfig(cfg *config.Config) error {
 	return nil
 }
 
+// Deprecated: Use Shutdown.
 func (oct *OpenCensusTracer) Finish() error {
+	return oct.Shutdown(context.Background())
+}
+
+func (oct *OpenCensusTracer) Shutdown(ctx context.Context) error {
 	err := oct.acquireGlobal()
 	defer octMutex.Unlock()
 	if err != nil {
@@ -89,8 +95,13 @@ func (oct *OpenCensusTracer) Finish() error {
 	}
 
 	for _, configOpt := range oct.configOptions {
-		if err = configOpt(nil); err != nil {
-			return err
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			if err = configOpt(nil); err != nil {
+				return err
+			}
 		}
 	}
 	globalOct = nil
@@ -144,33 +155,36 @@ func WithExporterFull(name, host string, logger *zap.SugaredLogger) ConfigOption
 			exporter trace.Exporter
 			closer   io.Closer
 		)
-		switch cfg.Backend {
-		case config.Zipkin:
-			// If host isn't specified, then zipkin.NewEndpoint will return an error saying that it
-			// can't find the host named ''. So, if not specified, default it to this machine's
-			// hostname.
-			if host == "" {
-				n, err := os.Hostname()
+		if cfg != nil {
+			switch cfg.Backend {
+			case config.Zipkin:
+				// If host isn't specified, then zipkin.NewEndpoint will return an error saying that it
+				// can't find the host named ''. So, if not specified, default it to this machine's
+				// hostname.
+				if host == "" {
+					n, err := os.Hostname()
+					if err != nil {
+						return fmt.Errorf("unable to get hostname: %w", err)
+					}
+					host = n
+				}
+				if name == "" {
+					name = host
+				}
+				zipEP, err := zipkin.NewEndpoint(name, host)
 				if err != nil {
-					return fmt.Errorf("unable to get hostname: %w", err)
+					logger.Errorw("error building zipkin endpoint", zap.Error(err))
+					return err
 				}
-				host = n
-			}
-			if name == "" {
-				name = host
-			}
-			zipEP, err := zipkin.NewEndpoint(name, host)
-			if err != nil {
-				logger.Errorw("error building zipkin endpoint", zap.Error(err))
-				return err
-			}
-			reporter := httpreporter.NewReporter(cfg.ZipkinEndpoint)
-			exporter = oczipkin.NewExporter(reporter, zipEP)
-			closer = reporter
+				reporter := httpreporter.NewReporter(cfg.ZipkinEndpoint)
+				exporter = oczipkin.NewExporter(reporter, zipEP)
+				closer = reporter
 
-		default:
-			// Disables tracing.
+			default:
+				// Disables tracing.
+			}
 		}
+
 		if exporter != nil {
 			trace.RegisterExporter(exporter)
 		}
diff --git a/tracing/setup.go b/tracing/setup.go
index 1940c40d7..0c3f23503 100644
--- a/tracing/setup.go
+++ b/tracing/setup.go
@@ -17,6 +17,7 @@ limitations under the License.
 package tracing
 
 import (
+	"context"
 	"fmt"
 
 	"go.uber.org/zap"
@@ -26,24 +27,38 @@ import (
 	"knative.dev/pkg/tracing/config"
 )
 
-// SetupStaticPublishing sets up trace publishing for the process. Note that other
-// pieces still need to generate the traces, this just ensures that if generated, they are collected
+type Tracer interface {
+	// Shutdown allows for proper cleanup/flush.
+	Shutdown(ctx context.Context) error
+}
+
+// SetupPublishingWithStaticConfig sets up trace publishing for the process. The caller should call .Shutdown() on
+// the returned Tracer before shutdown to make sure all traces are properly flushed.
+// Note that other pieces still need to generate the traces, this just ensures that if generated, they are collected
 // appropriately. This is normally done by using tracing.HTTPSpanMiddleware as a middleware HTTP
 // handler. The configuration will not be dynamically updated.
-func SetupStaticPublishing(logger *zap.SugaredLogger, serviceName string, cfg *config.Config) error {
+func SetupPublishingWithStaticConfig(logger *zap.SugaredLogger, serviceName string, cfg *config.Config) (Tracer, error) {
 	oct := NewOpenCensusTracer(WithExporter(serviceName, logger))
 	if err := oct.ApplyConfig(cfg); err != nil {
-		return fmt.Errorf("unable to set OpenCensusTracing config: %w", err)
+		return oct, fmt.Errorf("unable to set OpenCensusTracing config: %w", err)
 	}
-	return nil
+	return oct, nil
+}
+
+// Deprecated: Use SetupPublishingWithStaticConfig.
+func SetupStaticPublishing(logger *zap.SugaredLogger, serviceName string, cfg *config.Config) error {
+	_, err := SetupPublishingWithStaticConfig(logger, serviceName, cfg)
+	return err
 }
 
-// SetupDynamicPublishing sets up trace publishing for the process, by watching a
-// ConfigMap for the configuration. Note that other pieces still need to generate the traces, this
+// SetupPublishingWithDynamicConfig sets up trace publishing for the process, by watching a
+// ConfigMap for the configuration. The caller should call .Shutdown() on
+// the returned OpenCensusTracer before shutdown to make sure all traces are properly flushed.
+// Note that other pieces still need to generate the traces, this
 // just ensures that if generated, they are collected appropriately. This is normally done by using
 // tracing.HTTPSpanMiddleware as a middleware HTTP handler. The configuration will be dynamically
 // updated when the ConfigMap is updated.
-func SetupDynamicPublishing(logger *zap.SugaredLogger, configMapWatcher configmap.Watcher, serviceName, tracingConfigName string) error {
+func SetupPublishingWithDynamicConfig(logger *zap.SugaredLogger, configMapWatcher configmap.Watcher, serviceName, tracingConfigName string) (Tracer, error) {
 	oct := NewOpenCensusTracer(WithExporter(serviceName, logger))
 
 	tracerUpdater := func(name string, value interface{}) {
@@ -65,22 +80,28 @@ func SetupDynamicPublishing(logger *zap.SugaredLogger, configMapWatcher configma
 		},
 		tracerUpdater)
 	configStore.WatchConfigs(configMapWatcher)
-	return nil
+	return oct, nil
+}
+
+// Deprecated: Use SetupPublishingWithDynamicConfig.
+func SetupDynamicPublishing(logger *zap.SugaredLogger, configMapWatcher configmap.Watcher, serviceName, tracingConfigName string) error {
+	_, err := SetupPublishingWithDynamicConfig(logger, configMapWatcher, serviceName, tracingConfigName)
+	return err
 }
 
-// SetupDynamicPublishingWithInitialValue sets up the trace publishing for the process with an
+// SetupPublishingWithDynamicConfigAndInitialValue sets up the trace publishing for the process with an
 // initial value, by watching a ConfigMap for the configuration. Note that other pieces still
 // need to generate the traces, this just ensures that if generated, they are collected
 // appropriately. This is normally done by using tracing.HTTPSpanMiddleware as a middleware
 // HTTP handler. The configuration will be dynamically updated when the ConfigMap is updated.
-func SetupDynamicPublishingWithInitialValue(logger *zap.SugaredLogger, configMapWatcher configmap.Watcher, serviceName string, configm *corev1.ConfigMap) error {
+func SetupPublishingWithDynamicConfigAndInitialValue(logger *zap.SugaredLogger, configMapWatcher configmap.Watcher, serviceName string, configm *corev1.ConfigMap) (Tracer, error) {
 	oct := NewOpenCensusTracer(WithExporter(serviceName, logger))
 	cfg, err := config.NewTracingConfigFromConfigMap(configm)
 	if err != nil {
-		return fmt.Errorf("error while parsing %s config map: %w", configm.Name, err)
+		return oct, fmt.Errorf("error while parsing %s config map: %w", configm.Name, err)
 	}
 	if err = oct.ApplyConfig(cfg); err != nil {
-		return fmt.Errorf("unable to set OpenCensusTracing config: %w", err)
+		return oct, fmt.Errorf("unable to set OpenCensusTracing config: %w", err)
 	}
 
 	tracerUpdater := func(name string, value interface{}) {
@@ -102,5 +123,5 @@ func SetupDynamicPublishingWithInitialValue(logger *zap.SugaredLogger, configMap
 		},
 		tracerUpdater)
 	configStore.WatchConfigs(configMapWatcher)
-	return nil
+	return oct, nil
 }


Test patch:
diff --git a/pool/pool_test.go b/pool/pool_test.go
deleted file mode 100644
index 600a2c7e0..000000000
--- a/pool/pool_test.go
+++ /dev/null
@@ -1,247 +0,0 @@
-/*
-Copyright 2020 The Knative Authors
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    https://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-*/
-
-package pool
-
-import (
-	"context"
-	"errors"
-	"sync"
-	"testing"
-	"time"
-
-	"go.uber.org/atomic"
-)
-
-func TestParallelismNoErrors(t *testing.T) {
-	tests := []struct {
-		name string
-		size int
-		work int
-	}{{
-		name: "single threaded",
-		size: 1,
-		work: 3,
-	}, {
-		name: "three workers",
-		size: 3,
-		work: 10,
-	}, {
-		name: "ten workers",
-		size: 10,
-		work: 100,
-	}}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			var (
-				// m guards max.
-				m      sync.Mutex
-				max    int32
-				active atomic.Int32
-			)
-
-			// Use our own waitgroup to ensure that the work
-			// can all complete before we block on the error
-			// result.
-			wg := &sync.WaitGroup{}
-
-			worker := func() error {
-				defer wg.Done()
-				active.Inc()
-				defer active.Dec()
-
-				func() {
-					m.Lock()
-					defer m.Unlock()
-					if v := active.Load(); max < v {
-						max = v
-					}
-				}()
-
-				// Sleep a small amount to simulate work. This should be
-				// sufficient to saturate the threadpool before the first
-				// one wakes up.
-				time.Sleep(10 * time.Millisecond)
-				return nil
-			}
-
-			p := New(tc.size)
-			for idx := 0; idx < tc.work; idx++ {
-				wg.Add(1)
-				p.Go(worker)
-			}
-
-			// First wait for the waitgroup to finish, so that
-			// we are sure it isn't the Wait call that flushes
-			// remaining work.
-			wg.Wait()
-
-			if err := p.Wait(); err != nil {
-				t.Error("Wait() =", err)
-			}
-
-			if err := p.Wait(); err != nil {
-				t.Error("Wait() =", err)
-			}
-
-			if got, want := max, int32(tc.size); got != want {
-				t.Errorf("max active = %v, wanted %v", got, want)
-			}
-		})
-	}
-}
-
-func TestParallelismWithErrors(t *testing.T) {
-	tests := []struct {
-		name string
-		size int
-		work int
-	}{{
-		name: "single threaded",
-		size: 1,
-		work: 3,
-	}, {
-		name: "three workers",
-		size: 3,
-		work: 10,
-	}, {
-		name: "ten workers",
-		size: 10,
-		work: defaultCapacity + 10,
-	}}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			var (
-				// m guards max.
-				m      sync.Mutex
-				max    int32
-				active atomic.Int32
-			)
-
-			// Use our own waitgroup to ensure that the work
-			// can all complete before we block on the error
-			// result.
-			wg := &sync.WaitGroup{}
-
-			// The barrier holds all requests we don't expect to
-			// finish until we validate the condition.
-			barrier := make(chan struct{})
-
-			errExpected := errors.New("this is what I expect")
-			workerFactory := func(err error) func() error {
-				return func() error {
-					defer wg.Done()
-					active.Inc()
-					defer active.Dec()
-
-					func() {
-						m.Lock()
-						defer m.Unlock()
-						if v := active.Load(); max < v {
-							max = v
-						}
-					}()
-
-					// Sleep a small amount to simulate work. This should be
-					// sufficient to saturate the threadpool before the first
-					// one wakes up.
-					time.Sleep(10 * time.Millisecond)
-
-					// Make all unexpected errors wait.
-					if !errors.Is(err, errExpected) {
-						<-barrier
-					}
-					return err
-				}
-			}
-
-			p, ctx := NewWithContext(context.Background(), tc.size, defaultCapacity)
-
-			for i := 0; i < tc.work; i++ {
-				err := errors.New("this is not what I expect")
-				if i == tc.size-1 {
-					// Just fit in a request with the error we expect.
-					err = errExpected
-				}
-
-				wg.Add(1)
-				p.Go(workerFactory(err))
-			}
-
-			// Wait for the error to propagate.
-			<-ctx.Done()
-			// Allow all remaining requests to finish.
-			close(barrier)
-
-			// First wait for the waitgroup to finish, so that
-			// we are sure it isn't the Wait call that flushes
-			// remaining work.
-			wg.Wait()
-
-			if err := p.Wait(); !errors.Is(err, errExpected) {
-				t.Errorf("Wait() = %v, wanted %v", err, errExpected)
-			}
-
-			if got, want := max, int32(tc.size); got != want {
-				t.Errorf("max active = %v, wanted %v", got, want)
-			}
-		})
-	}
-}
-
-func TestWithContextWaitCancels(t *testing.T) {
-	pool, ctx := NewWithContext(context.Background(), 1 /*1 thread*/, 10 /*capacity*/)
-	for i := 0; i < 10; i++ {
-		pool.Go(func() error {
-			time.Sleep(10 * time.Millisecond)
-			return nil
-		})
-	}
-	if err := pool.Wait(); err != nil {
-		t.Fatal("pool.Wait =", err)
-	}
-	select {
-	case <-ctx.Done():
-	default:
-		t.Error("ctx is not canceled")
-	}
-}
-
-func TestErrorCancelsContext(t *testing.T) {
-	want := errors.New("i failed, sorry")
-	pool, ctx := NewWithContext(context.Background(), 1 /*1 thread*/, 10 /*capacity*/)
-	pool.Go(func() error {
-		return want
-	})
-	// Those don't matter, but generate load.
-	for i := 0; i < 10; i++ {
-		pool.Go(func() error {
-			time.Sleep(100 * time.Millisecond)
-			return nil
-		})
-	}
-	// This should be triggered basically immediately.
-	select {
-	case <-ctx.Done():
-	case <-time.After(100 * time.Millisecond):
-		t.Error("ctx is not canceled due to the first error")
-	}
-	if err := pool.Wait(); !errors.Is(err, want) {
-		t.Fatalf("pool.Wait() = %v, want: %v", err, want)
-	}
-}
diff --git a/tracing/http_test.go b/tracing/http_test.go
index f3b46166c..0b31005ac 100644
--- a/tracing/http_test.go
+++ b/tracing/http_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package tracing_test
 
 import (
+	"context"
 	"net/http"
 	"net/url"
 	"testing"
@@ -70,7 +71,7 @@ func TestHTTPSpanMiddleware(t *testing.T) {
 	oct := NewOpenCensusTracer(co)
 	t.Cleanup(func() {
 		reporter.Close()
-		oct.Finish()
+		oct.Shutdown(context.Background())
 	})
 
 	if err := oct.ApplyConfig(&cfg); err != nil {
@@ -117,7 +118,7 @@ func TestHTTPSpanIgnoringPaths(t *testing.T) {
 	oct := NewOpenCensusTracer(co)
 	t.Cleanup(func() {
 		reporter.Close()
-		oct.Finish()
+		oct.Shutdown(context.Background())
 	})
 
 	if err := oct.ApplyConfig(&cfg); err != nil {
@@ -190,7 +191,7 @@ func BenchmarkSpanMiddleware(b *testing.B) {
 	oct := NewOpenCensusTracer(co)
 	b.Cleanup(func() {
 		reporter.Close()
-		oct.Finish()
+		oct.Shutdown(context.Background())
 	})
 
 	if err := oct.ApplyConfig(&cfg); err != nil {
diff --git a/tracing/opencensus_test.go b/tracing/opencensus_test.go
index 40233890f..17fa7ec7c 100644
--- a/tracing/opencensus_test.go
+++ b/tracing/opencensus_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package tracing_test
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -39,14 +40,14 @@ func TestOpenCensusTracerGlobalLifecycle(t *testing.T) {
 		t.Fatalf("Expected error when applying config to second OCT.")
 	}
 
-	if err := oct.Finish(); err != nil {
+	if err := oct.Shutdown(context.Background()); err != nil {
 		t.Fatal("Failed to finish OCT:", err)
 	}
 
 	if err := otherOCT.ApplyConfig(&config.Config{}); err != nil {
 		t.Fatal("Failed to ApplyConfig on OtherOCT after finishing OCT:", err)
 	}
-	otherOCT.Finish()
+	otherOCT.Shutdown(context.Background())
 }
 
 func TestOpenCensusTraceApplyConfigFailingConfigOption(t *testing.T) {
@@ -60,7 +61,7 @@ func TestOpenCensusTraceApplyConfigFailingConfigOption(t *testing.T) {
 	if err := oct.ApplyConfig(&config.Config{}); !errors.Is(err, coErr) {
 		t.Errorf("Expected error not seen. Got %q. Want %q", err, coErr)
 	}
-	if err := oct.Finish(); err != nil {
+	if err := oct.Shutdown(context.Background()); err != nil {
 		t.Errorf("Unexpected error Finishing: %q", err)
 	}
 }
@@ -81,10 +82,10 @@ func TestOpenCensusTraceFinishFailingConfigOption(t *testing.T) {
 	if err := oct.ApplyConfig(&config.Config{}); err != nil {
 		t.Errorf("Unexpected error Applying Config: %q", err)
 	}
-	if err := oct.Finish(); !errors.Is(err, coErr) {
+	if err := oct.Shutdown(context.Background()); !errors.Is(err, coErr) {
 		t.Errorf("Expected error not seen. Got %q. Want %q", err, coErr)
 	}
-	if err := oct.Finish(); err != nil {
-		t.Errorf("Unexpected error on second Finish (global state mutated, other tests may fail oddly): %q", err)
+	if err := oct.Shutdown(context.Background()); err != nil {
+		t.Errorf("Unexpected error on second Shutdown (global state mutated, other tests may fail oddly): %q", err)
 	}
 }
diff --git a/tracing/zipkin_test.go b/tracing/zipkin_test.go
index 1dfec3e1d..a491c5ec6 100644
--- a/tracing/zipkin_test.go
+++ b/tracing/zipkin_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package tracing_test
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -72,7 +73,7 @@ func TestOpenCensusTracerApplyConfig(t *testing.T) {
 				t.Error("Got tracer config (-want, +got) =", diff)
 			}
 
-			oct.Finish()
+			oct.Shutdown(context.Background())
 		})
 	}
 }