bigquery/storage/managedwriter: EnableWriteRetries seems to lead to panics after ManagedStream context cancellation

[zapo]

Client

Bigquery

Environment

GKE (also reproducible locally)

Go Environment

$ go version
go version go1.19 linux/amd64

Code

This simulates how we use the managedwriter with concurrent streams for the same client.

package main

import (
	"context"
	"log"
	"sync"
	"test/pb"

	"cloud.google.com/go/bigquery/storage/managedwriter"
	"google.golang.org/protobuf/proto"
)

func main() {
	ctx := context.Background()
	client, err := managedwriter.NewClient(ctx, "<redacted>")
	if err != nil {
		log.Fatal(err)
	}

	var wg sync.WaitGroup

	for i := 0; i < 2; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()

			ctx, cancel := context.WithCancel(ctx)
			defer cancel()

			stream, err := client.NewManagedStream(ctx,
				managedwriter.WithType(managedwriter.DefaultStream),
				managedwriter.WithDestinationTable("<redacted>"),
				managedwriter.EnableWriteRetries(true),
			)
			if err != nil {
				log.Fatal(err)
			}

			for i := 0; i < 10000; i++ {
				rec := &pb.TestRecord{}
				buf, _ := proto.Marshal(rec)

				if _, err := stream.AppendRows(ctx, [][]byte{buf}); err != nil {
					log.Fatalf("append row: %v", err)
				}
			}

			if err := stream.Close(); err != nil {
				log.Fatalf("stream close: %v", err)
			}
		}()
	}
	wg.Wait()

	if err := client.Close(); err != nil {
		log.Fatalf("client close: %v", err)
	}
}

Expected behavior

Does not panic

Actual behavior

Panics with: panic: close of closed channel when using EnableWriteRetries(true).
However this seems related to cancellation of context passed to NewManagedStream and concurrency.

Additional context

panic: close of closed channel
goroutine 10580 [running]:
[cloud.google.com/go/bigquery/storage/managedwriter.(*pendingWrite).markDone(0xc047e496c0](http://cloud.google.com/go/bigquery/storage/managedwriter.(*pendingWrite).markDone(0xc047e496c0), 0x2fd4a60?, {0x2fd4a60?, 0x45024e0?}, 0xc049f6af30)
	/home/circleci/go/pkg/mod/cloud.google.com/go/bigquery@v1.43.0/storage/managedwriter/appendresult.go:218 +0x9a
[cloud.google.com/go/bigquery/storage/managedwriter.(*ManagedStream).processRetry(0xc048c13ef0](http://cloud.google.com/go/bigquery/storage/managedwriter.(*ManagedStream).processRetry(0xc048c13ef0), 0xc047e496c0, 0x0?, {0x2fcd9e0?, 0xc0698ef300?})
	/home/circleci/go/pkg/mod/cloud.google.com/go/bigquery@v1.43.0/storage/managedwriter/managed_stream.go:530 +0x19e
[cloud.google.com/go/bigquery/storage/managedwriter.recvProcessor(0xc048c13ef0](http://cloud.google.com/go/bigquery/storage/managedwriter.recvProcessor(0xc048c13ef0), {0x2ff9f38, 0xc049f53760}, 0xc049f47c20)
	/home/circleci/go/pkg/mod/cloud.google.com/go/bigquery@v1.43.0/storage/managedwriter/managed_stream.go:493 +0x128
created by [cloud.google.com/go/bigquery/storage/managedwriter.(*ManagedStream).openWithRetry](http://cloud.google.com/go/bigquery/storage/managedwriter.(*ManagedStream).openWithRetry)
	/home/circleci/go/pkg/mod/cloud.google.com/go/bigquery@v1.43.0/storage/managedwriter/managed_stream.go:252 +0x36a

[fsaintjacques]

Note that this also applies to 1.46.0

[shollyman]

Thanks for the report, will take a look!