Repository: knative/pkg
Pull request #2572 (2022-08-15 20:58:08)

allow drainer to be reset
Replaces https://github.com/knative/pkg/pull/2491

We need a reset so that downstream for serving to work around containers that are failing their liveness probe see: https://github.com/knative/serving/issues/12571

knative-prow[bot]: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/knative/pkg/pull/2572#" title="Author self-approved">dprotaso</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=knative%2Fpkg).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[network/OWNERS](https://github.com/knative/pkg/blob/main/network/OWNERS)~~ [dprotaso]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
knative-prow[bot]: @dprotaso: The following test **failed**, say `/retest` to rerun all failed tests or `/retest-required` to rerun all mandatory failed tests:

Test name | Commit | Details | Required | Rerun command
--- | --- | --- | --- | ---
unit-tests_pkg_main | 5f85ca6bb7c81f61859c129d61b2d002e59bec54 | [link](https://prow.knative.dev/view/gs/knative-prow/pr-logs/pull/knative_pkg/2572/unit-tests_pkg_main/1559283327599382528) | true | `/test unit-tests_pkg_main`

[Your PR dashboard](https://prow.knative.dev/pr).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
dprotaso: /retest 
codecov[bot]: # [Codecov](https://codecov.io/gh/knative/pkg/pull/2572?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative) Report
> Merging [#2572](https://codecov.io/gh/knative/pkg/pull/2572?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative) (5f85ca6) into [main](https://codecov.io/gh/knative/pkg/commit/9a1245b7ec88474b0a1adf9fd223aecacfdcdb4f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative) (9a1245b) will **increase** coverage by `0.04%`.
> The diff coverage is `82.50%`.

```diff
@@            Coverage Diff             @@
##             main    #2572      +/-   ##
==========================================
+ Coverage   81.40%   81.44%   +0.04%     
==========================================
  Files         163      163              
  Lines        9740     9768      +28     
==========================================
+ Hits         7929     7956      +27     
  Misses       1574     1574              
- Partials      237      238       +1     
```


| [Impacted Files](https://codecov.io/gh/knative/pkg/pull/2572?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative) | Coverage Δ | |
|---|---|---|
| [network/handlers/drain.go](https://codecov.io/gh/knative/pkg/pull/2572/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative#diff-bmV0d29yay9oYW5kbGVycy9kcmFpbi5nbw==) | `90.47% <82.50%> (+2.16%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=knative)


Patch:
diff --git a/network/handlers/drain.go b/network/handlers/drain.go
index d6ef19f4e..5252016f0 100644
--- a/network/handlers/drain.go
+++ b/network/handlers/drain.go
@@ -70,12 +70,12 @@ type Drainer struct {
 	// after Drain is called before it may return.
 	QuietPeriod time.Duration
 
-	// once is used to initialize timer
-	once sync.Once
-
 	// timer is used to orchestrate the drain.
 	timer timer
 
+	// used to synchronize callers of Drain and Reset
+	ch chan struct{}
+
 	// HealthCheckUAPrefixes are the additional user agent prefixes that trigger the
 	// drainer's health check
 	HealthCheckUAPrefixes []string
@@ -106,7 +106,7 @@ func (d *Drainer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	d.reset()
+	d.resetTimer()
 	d.Inner.ServeHTTP(w, r)
 }
 
@@ -115,19 +115,42 @@ func (d *Drainer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 func (d *Drainer) Drain() {
 	// Note: until the first caller exits, the others
 	// will wait blocked as well.
-	d.once.Do(func() {
-		t := func() timer {
-			d.Lock()
-			defer d.Unlock()
-			if d.QuietPeriod <= 0 {
-				d.QuietPeriod = network.DefaultDrainTimeout
+	ch := func() chan struct{} {
+		d.Lock()
+		defer d.Unlock()
+		if d.ch != nil {
+			return d.ch
+		}
+
+		if d.QuietPeriod <= 0 {
+			d.QuietPeriod = network.DefaultDrainTimeout
+		}
+
+		timer := newTimer(d.QuietPeriod)
+		ch := make(chan struct{})
+
+		go func() {
+			select {
+			case <-ch:
+				// closed by reset
+			case <-timer.tickChan():
+				close(ch)
 			}
-			d.timer = newTimer(d.QuietPeriod)
-			return d.timer
 		}()
 
-		<-t.tickChan()
-	})
+		d.ch = ch
+		d.timer = timer
+		return ch
+	}()
+
+	<-ch
+}
+
+func drainTimer(tc <-chan time.Time) {
+	select {
+	case <-tc:
+	default:
+	}
 }
 
 // isHealthcheckRequest validates if the request has a user agent that is for healthcheck
@@ -145,8 +168,27 @@ func (d *Drainer) isHealthCheckRequest(r *http.Request) bool {
 	return false
 }
 
-// reset resets the drain timer to the full amount of time.
-func (d *Drainer) reset() {
+// Reset interrupts Drain and clears the drainers internal state
+// Thus further calls to Drain will block and wait for the entire QuietPeriod
+func (d *Drainer) Reset() {
+	d.Lock()
+	defer d.Unlock()
+
+	if d.timer != nil {
+		if d.timer.Stop() {
+			d.timer = nil
+		} else {
+			drainTimer(d.timer.tickChan())
+		}
+	}
+
+	if d.ch != nil {
+		close(d.ch)
+		d.ch = nil
+	}
+}
+
+func (d *Drainer) resetTimer() {
 	if func() bool {
 		d.RLock()
 		defer d.RUnlock()
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/network/handlers/drain_test.go b/network/handlers/drain_test.go
index 3970df925..22c55706b 100644
--- a/network/handlers/drain_test.go
+++ b/network/handlers/drain_test.go
@@ -490,3 +490,56 @@ func TestServeKProbe(t *testing.T) {
 		t.Errorf("Probe status = %d, wanted %d", got, want)
 	}
 }
+
+func TestReset(t *testing.T) {
+	d := Drainer{
+		QuietPeriod: 5 * time.Second,
+	}
+
+	drain1 := make(chan struct{})
+	drain2 := make(chan struct{})
+
+	go func() {
+		defer close(drain1)
+		d.Drain()
+	}()
+
+	go func() {
+		defer close(drain2)
+		d.Drain()
+	}()
+
+	// wait for draining to be active
+	time.Sleep(50 * time.Millisecond)
+
+	d.Reset()
+
+	select {
+	case <-drain1:
+	case <-time.After(time.Second):
+		t.Fatal("Reset didn't unblock first Drain")
+	}
+
+	select {
+	case <-drain2:
+	case <-time.After(time.Second):
+		t.Fatal("Reset didn't unblock second Drain")
+	}
+
+	// Calling reset again should be a noop
+	d.Reset()
+
+	d.QuietPeriod = time.Second / 2
+
+	start := time.Now()
+	d.Drain()
+	duration := time.Since(start)
+	diff := d.QuietPeriod - duration
+	if diff < 0 {
+		diff = -diff
+	}
+
+	if diff > 50*time.Millisecond {
+		t.Error("expected to drain to wait QuietPeriod time after reset")
+	}
+}
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())
 		})
 	}
 }