feat(storage): trace span covers life of a Reader (#8390)

Make the tracing span created in NewRangeReader cover the entire lifecycle of the Reader. NewRangeReader returns before the data transfer is complete, so it's more useful to track the entire period until the Reader is closed.

NewRangeReader is still tracked with a span at the level of individual transports (HTTP and gRPC).

Author: tritone

storage/reader.go

@@ -87,8 +87,9 @@ func (o *ObjectHandle) NewReader(ctx context.Context) (*Reader, error) {
 // that file will be served back whole, regardless of the requested range as
 // Google Cloud Storage dictates.
 func (o *ObjectHandle) NewRangeReader(ctx context.Context, offset, length int64) (r *Reader, err error) {
-	ctx = trace.StartSpan(ctx, "cloud.google.com/go/storage.Object.NewRangeReader")
-	defer func() { trace.EndSpan(ctx, err) }()
+	// This span covers the life of the reader. It is closed via the context
+	// in Reader.Close.
+	ctx = trace.StartSpan(ctx, "cloud.google.com/go/storage.Object.Reader")
 
 	if err := o.validate(); err != nil {
 		return nil, err
@@ -117,6 +118,14 @@ func (o *ObjectHandle) NewRangeReader(ctx context.Context, offset, length int64)
 
 	r, err = o.c.tc.NewRangeReader(ctx, params, opts...)
 
+	// Pass the context so that the span can be closed in Reader.Close, or close the
+	// span now if there is an error.
+	if err == nil {
+		r.ctx = ctx
+	} else {
+		trace.EndSpan(ctx, err)
+	}
+
 	return r, err
 }
 
@@ -204,11 +213,14 @@ type Reader struct {
 	gotCRC             uint32 // running crc
 
 	reader io.ReadCloser
+	ctx    context.Context
 }
 
 // Close closes the Reader. It must be called when done reading.
 func (r *Reader) Close() error {
-	return r.reader.Close()
+	err := r.reader.Close()
+	trace.EndSpan(r.ctx, err)
+	return err
 }
 
 func (r *Reader) Read(p []byte) (int, error) {