cloud: have resumable reader manage span lifetime

This change adds a span to the cloud.ResumableReader that is tracked for
the lifetime of the reader. This ensures that if the reader
implementation needs to issue a new HTTP request, the span is still
valid and can be used to annotate logs.

Fixes: #153347
Release note: none

Author: jeffswenson

pkg/cloud/azure/azure_storage.go

@@ -366,6 +366,7 @@ func (s *azureStorage) ReadFile(
 			}
 		}
 	}
+	// BUG: we should follow the azure retry setting here.
 	reader := resp.NewRetryReader(ctx, &azblob.RetryReaderOptions{MaxRetries: 3})
 	return ioctx.ReadCloserAdapter(reader), fileSize, nil
 }

pkg/cloud/cloud_io.go

@@ -242,6 +242,7 @@ type ReaderOpenerAt func(ctx context.Context, pos int64) (io.ReadCloser, int64,
 type ResumingReader struct {
 	Opener       ReaderOpenerAt   // Get additional content
 	Reader       io.ReadCloser    // Currently opened reader
+	ReaderSpan   *tracing.Span    // Span for the current reader, if Reader is non-nil
 	Filename     string           // Used for logging
 	Pos          int64            // How much data was received so far
 	Size         int64            // Total size of the file
@@ -284,6 +285,10 @@ func NewResumingReader(
 
 // Open opens the reader at its current offset.
 func (r *ResumingReader) Open(ctx context.Context) error {
+	if r.Reader != nil {
+		return errors.AssertionFailedf("reader already open")
+	}
+
 	if r.Size > 0 && r.Pos >= r.Size {
 		// Don't try to open a file if the size has been set and the position is
 		// at size. This generally results in an invalid range error for the
@@ -294,10 +299,18 @@ func (r *ResumingReader) Open(ctx context.Context) error {
 
 	return DelayedRetry(ctx, "Open", r.ErrFn, func() error {
 		var readErr error
+
+		ctx, span := tracing.ForkSpan(ctx, "resuming-reader")
 		r.Reader, r.Size, readErr = r.Opener(ctx, r.Pos)
 		if readErr != nil {
+			span.Finish()
 			return errors.Wrapf(readErr, "open %s", r.Filename)
 		}
+
+		// We hold onto the span for the lifetime of the reader because the reader
+		// may issue new HTTP requests after Open returns.
+		r.ReaderSpan = span
+
 		return nil
 	})
 }
@@ -340,10 +353,9 @@ func (r *ResumingReader) Read(ctx context.Context, p []byte) (int, error) {
 			}
 			log.Dev.Errorf(ctx, "Retry IO error: %s", lastErr)
 			lastErr = nil
-			if r.Reader != nil {
-				r.Reader.Close()
-			}
-			r.Reader = nil
+			// Ignore the error from Close(). We are already handling a read error
+			// so we know the handle is in a bad state.
+			_ = r.Close(ctx)
 		}
 	}
 
@@ -356,10 +368,14 @@ func (r *ResumingReader) Read(ctx context.Context, p []byte) (int, error) {
 
 // Close implements io.Closer.
 func (r *ResumingReader) Close(ctx context.Context) error {
-	if r.Reader != nil {
-		return r.Reader.Close()
+	if r.Reader == nil {
+		return nil
 	}
-	return nil
+
+	err := r.Reader.Close()
+	r.ReaderSpan.Finish()
+	r.Reader = nil
+	return err
 }
 
 // CheckHTTPContentRangeHeader parses Content-Range header and ensures that

pkg/cloud/cloudtestutils/BUILD.bazel

@@ -23,6 +23,8 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/sysutil",
         "//pkg/util/timeutil",
+        "//pkg/util/tracing",
+        "//pkg/util/tracing/tracingpb",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
         "@org_golang_x_sync//errgroup",

pkg/cloud/cloudtestutils/cloud_nemesis.go

@@ -19,6 +19,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cloud"
 	"github.com/cockroachdb/cockroach/pkg/util/ioctx"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing/tracingpb"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"
@@ -39,9 +41,15 @@ func RunCloudNemesisTest(t *testing.T, storage cloud.ExternalStorage) {
 		listConcurrency:  1,
 	}
 
-	// We create a context here because we don't want to support a caller supplied
-	// cancelation signal.
-	ctx := context.Background()
+	// Create a root tracing span for the test to ensure ops create spans
+	// which will help flush out span lifetime bugs.
+	tracer := tracing.NewTracerWithOpt(
+		context.Background(),
+		tracing.WithUseAfterFinishOpt(true, true),
+		tracing.WithSpanReusePercent(0))
+	ctx, rootSpan := tracer.StartSpanCtx(context.Background(), "cloud-nemesis-test", tracing.WithRecording(tracingpb.RecordingStructured))
+	defer rootSpan.Finish()
+
 	if err := nemesis.run(ctx, 2*time.Minute); err != nil {
 		t.Fatalf("%+v", err)
 	}
@@ -85,9 +93,13 @@ func (c *cloudNemesis) run(ctx context.Context, duration time.Duration) error {
 	g, ctx := errgroup.WithContext(ctx)
 
 	g.Go(func() error {
-		time.Sleep(duration)
+		wait := time.NewTimer(duration)
+		select {
+		case <-ctx.Done():
+		case <-wait.C:
+		}
 		close(done)
-		return nil
+		return ctx.Err()
 	})
 
 	for i := 0; i < c.writeConcurrency; i++ {