From 6000d3363615927ea193eadc21f71fdc2e777829 Mon Sep 17 00:00:00 2001
From: Jeff Swenson <jeffswenson@betterthannull.com>
Date: Fri, 10 Oct 2025 16:36:18 -0400
Subject: [PATCH] azure: use ResumingReader

Previously, the azure client did not use the ResumingReader. This was
causing tracing span use after free issues as the client would issue
http retries with a span that was already finished. The ResumingReader
manages a span for the lifetime of the reader.

There is also some evidence in #154483 that the RetryReader built into
the azure client did not properly retry errors that occur while reading
the body.

Fixes: 154849
Release note: none
---
 pkg/cloud/azure/azure_storage.go | 50 +++++++++++++++++++++-----------
 1 file changed, 33 insertions(+), 17 deletions(-)

diff --git a/pkg/cloud/azure/azure_storage.go b/pkg/cloud/azure/azure_storage.go
index afaf00ee7361..e53df16fe4dc 100644
--- a/pkg/cloud/azure/azure_storage.go
+++ b/pkg/cloud/azure/azure_storage.go
@@ -356,31 +356,47 @@ func isNotFoundErr(err error) bool {
 func (s *azureStorage) ReadFile(
 	ctx context.Context, basename string, opts cloud.ReadOptions,
 ) (_ ioctx.ReadCloserCtx, fileSize int64, _ error) {
+	object := path.Join(s.prefix, basename)
+
 	ctx, sp := tracing.ChildSpan(ctx, "azure.ReadFile")
 	defer sp.Finish()
-	sp.SetTag("path", attribute.StringValue(path.Join(s.prefix, basename)))
-	resp, err := s.getBlob(basename).DownloadStream(ctx, &azblob.DownloadStreamOptions{Range: azblob.
-		HTTPRange{Offset: opts.Offset}})
-	if err != nil {
-		if isNotFoundErr(err) {
-			return nil, 0, cloud.WrapErrFileDoesNotExist(err, "azure blob does not exist")
-		}
-		return nil, 0, errors.Wrapf(err, "failed to create azure reader")
-	}
+	sp.SetTag("path", attribute.StringValue(object))
 
-	if !opts.NoFileSize {
-		if opts.Offset == 0 {
-			fileSize = *resp.ContentLength
-		} else {
-			fileSize, err = cloud.CheckHTTPContentRangeHeader(*resp.ContentRange, opts.Offset)
+	r := cloud.NewResumingReader(ctx,
+		func(ctx context.Context, pos int64) (io.ReadCloser, int64, error) {
+			resp, err := s.getBlob(basename).DownloadStream(ctx, &azblob.DownloadStreamOptions{
+				Range: azblob.HTTPRange{Offset: pos},
+			})
 			if err != nil {
 				return nil, 0, err
 			}
+
+			length := *resp.ContentLength
+			if 0 < pos {
+				length, err = cloud.CheckHTTPContentRangeHeader(*resp.ContentRange, pos)
+				if err != nil {
+					return nil, 0, err
+				}
+			}
+
+			return resp.Body, length, nil
+		},
+		nil, // reader
+		opts.Offset,
+		0,
+		object,
+		cloud.ResumingReaderRetryOnErrFnForSettings(ctx, s.settings),
+		nil, // errFn
+	)
+
+	if err := r.Open(ctx); err != nil {
+		if isNotFoundErr(err) {
+			return nil, 0, cloud.WrapErrFileDoesNotExist(err, "azure blob does not exist")
 		}
+		return nil, 0, err
 	}
-	// BUG: we should follow the azure retry setting here.
-	reader := resp.NewRetryReader(ctx, &azblob.RetryReaderOptions{MaxRetries: 3})
-	return ioctx.ReadCloserAdapter(reader), fileSize, nil
+
+	return r, r.Size, nil
 }
 
 func (s *azureStorage) List(ctx context.Context, prefix, delim string, fn cloud.ListingFn) error {