Fix uploads larger than 256MB (Azure#19658)

* Fix uploads larger than 256MB

Uploads larger than 256MB could fail depending on their size.
Refactored some internal code to follow convention.

* add another test

* fix comment

* refine changelog entry

Author: jhendrixMSFT

sdk/storage/azblob/CHANGELOG.md

@@ -19,6 +19,7 @@
 
 * Corrected signing of User Delegation SAS. Fixes [#19372](https://github.com/Azure/azure-sdk-for-go/issues/19372) and [#19454](https://github.com/Azure/azure-sdk-for-go/issues/19454)
 * Added formatting of start and expiry time in [SetAccessPolicy](https://learn.microsoft.com/rest/api/storageservices/set-container-acl#request-body). Fixes [#18712](https://github.com/Azure/azure-sdk-for-go/issues/18712)
+* Uploading block blobs larger than 256MB can fail in some cases with error `net/http: HTTP/1.x transport connection broken`.
 
 ### Other Changes
 

sdk/storage/azblob/blob/client.go

@@ -336,8 +336,7 @@ func (b *Client) download(ctx context.Context, writer io.WriterAt, o downloadOpt
 		TransferSize:  count,
 		ChunkSize:     o.BlockSize,
 		Concurrency:   o.Concurrency,
-		Operation: func(chunkStart int64, count int64, ctx context.Context) error {
-
+		Operation: func(ctx context.Context, chunkStart int64, count int64) error {
 			downloadBlobOptions := o.getDownloadBlobOptions(HTTPRange{
 				Offset: chunkStart + o.Range.Offset,
 				Count:  count,

sdk/storage/azblob/blockblob/client.go

@@ -362,7 +362,8 @@ func (bb *Client) CopyFromURL(ctx context.Context, copySource string, o *blob.Co
 // Concurrent Upload Functions -----------------------------------------------------------------------------------------
 
 // uploadFromReader uploads a buffer in blocks to a block blob.
-func (bb *Client) uploadFromReader(ctx context.Context, reader io.ReaderAt, readerSize int64, o *uploadFromReaderOptions) (uploadFromReaderResponse, error) {
+func (bb *Client) uploadFromReader(ctx context.Context, reader io.ReaderAt, actualSize int64, o *uploadFromReaderOptions) (uploadFromReaderResponse, error) {
+	readerSize := actualSize
 	if o.BlockSize == 0 {
 		// If bufferSize > (MaxStageBlockBytes * MaxBlocks), then error
 		if readerSize > MaxStageBlockBytes*MaxBlocks {
@@ -412,11 +413,17 @@ func (bb *Client) uploadFromReader(ctx context.Context, reader io.ReaderAt, read
 		TransferSize:  readerSize,
 		ChunkSize:     o.BlockSize,
 		Concurrency:   o.Concurrency,
-		Operation: func(offset int64, count int64, ctx context.Context) error {
+		Operation: func(ctx context.Context, offset int64, chunkSize int64) error {
 			// This function is called once per block.
 			// It is passed this block's offset within the buffer and its count of bytes
 			// Prepare to read the proper block/section of the buffer
-			var body io.ReadSeeker = io.NewSectionReader(reader, offset, count)
+			if chunkSize < o.BlockSize {
+				// this is the last block.  its actual size might be less
+				// than the calculated size due to rounding up of the payload
+				// size to fit in a whole number of blocks.
+				chunkSize = (actualSize - offset)
+			}
+			var body io.ReadSeeker = io.NewSectionReader(reader, offset, chunkSize)
 			blockNum := offset / o.BlockSize
 			if o.Progress != nil {
 				blockProgress := int64(0)

sdk/storage/azblob/blockblob/client_test.go

@@ -11,6 +11,7 @@ import (
 	"context"
 	"crypto/md5"
 	"encoding/base64"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -20,6 +21,7 @@ import (
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/recording"
@@ -3843,3 +3845,80 @@ func (s *BlockBlobUnrecordedTestsSuite) TestLargeBlockBufferedUploadInParallel()
 	_require.Equal(*(committed[0].Size), largeBlockSize)
 	_require.Equal(*(committed[1].Size), largeBlockSize)
 }
+
+type fakeBlockBlob struct {
+	totalStaged int64
+}
+
+func (f *fakeBlockBlob) Do(req *http.Request) (*http.Response, error) {
+	// verify that the number of bytes read matches what was specified
+	data := make([]byte, req.ContentLength)
+	read, err := req.Body.Read(data)
+	if err != nil && !errors.Is(err, io.EOF) {
+		return nil, err
+	} else if int64(read) < req.ContentLength {
+		return nil, fmt.Errorf("expected %d bytes, read %d", req.ContentLength, read)
+	}
+	qp := req.URL.Query()
+	if comp := qp.Get("comp"); comp == "block" {
+		// staging a block, record its size
+		f.totalStaged += int64(read)
+	}
+	return &http.Response{
+		Request:    req,
+		Status:     "Created",
+		StatusCode: http.StatusCreated,
+		Header:     http.Header{},
+		Body:       http.NoBody,
+	}, nil
+}
+
+func TestUploadBufferUnevenBlockSize(t *testing.T) {
+	fbb := &fakeBlockBlob{}
+	client, err := blockblob.NewClientWithNoCredential("https://fake/blob/path", &blockblob.ClientOptions{
+		ClientOptions: policy.ClientOptions{
+			Transport: fbb,
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, client)
+
+	// create fake source that's not evenly divisible by 50000 (max number of blocks)
+	// and greater than MaxUploadBlobBytes (256MB) so that it doesn't fit into a single upload.
+
+	buffer := make([]byte, 263*1024*1024)
+	for i := 0; i < len(buffer); i++ {
+		buffer[i] = 1
+	}
+
+	_, err = client.UploadBuffer(context.Background(), buffer, &blockblob.UploadBufferOptions{
+		Concurrency: 1,
+	})
+	require.NoError(t, err)
+	require.Equal(t, int64(len(buffer)), fbb.totalStaged)
+}
+
+func TestUploadBufferEvenBlockSize(t *testing.T) {
+	fbb := &fakeBlockBlob{}
+	client, err := blockblob.NewClientWithNoCredential("https://fake/blob/path", &blockblob.ClientOptions{
+		ClientOptions: policy.ClientOptions{
+			Transport: fbb,
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, client)
+
+	// create fake source that's evenly divisible by 50000 (max number of blocks)
+	// and greater than MaxUploadBlobBytes (256MB) so that it doesn't fit into a single upload.
+
+	buffer := make([]byte, 270000000)
+	for i := 0; i < len(buffer); i++ {
+		buffer[i] = 1
+	}
+
+	_, err = client.UploadBuffer(context.Background(), buffer, &blockblob.UploadBufferOptions{
+		Concurrency: 1,
+	})
+	require.NoError(t, err)
+	require.Equal(t, int64(len(buffer)), fbb.totalStaged)
+}

sdk/storage/azblob/client_test.go

@@ -512,7 +512,7 @@ func (s *AZBlobUnrecordedTestsSuite) TestBasicDoBatchTransfer() {
 			TransferSize: test.transferSize,
 			ChunkSize:    test.chunkSize,
 			Concurrency:  test.concurrency,
-			Operation: func(offset int64, chunkSize int64, ctx context.Context) error {
+			Operation: func(ctx context.Context, offset int64, chunkSize int64) error {
 				atomic.AddInt64(&totalSizeCount, chunkSize)
 				atomic.AddInt64(&runCount, 1)
 				return nil
@@ -554,7 +554,7 @@ func (s *AZBlobUnrecordedTestsSuite) TestDoBatchTransferWithError() {
 		TransferSize: 5,
 		ChunkSize:    1,
 		Concurrency:  5,
-		Operation: func(offset int64, chunkSize int64, ctx context.Context) error {
+		Operation: func(ctx context.Context, offset int64, chunkSize int64) error {
 			// simulate doing some work (HTTP call in real scenarios)
 			// later chunks later longer to finish
 			time.Sleep(time.Second * time.Duration(offset))

sdk/storage/azblob/internal/shared/batch_transfer.go

@@ -16,7 +16,7 @@ type BatchTransferOptions struct {
 	TransferSize  int64
 	ChunkSize     int64
 	Concurrency   uint16
-	Operation     func(offset int64, chunkSize int64, ctx context.Context) error
+	Operation     func(ctx context.Context, offset int64, chunkSize int64) error
 	OperationName string
 }
 
@@ -57,9 +57,8 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {
 			curChunkSize = o.TransferSize - (int64(chunkNum) * o.ChunkSize) // Remove size of all transferred chunks from total
 		}
 		offset := int64(chunkNum) * o.ChunkSize
-
 		operationChannel <- func() error {
-			return o.Operation(offset, curChunkSize, ctx)
+			return o.Operation(ctx, offset, curChunkSize)
 		}
 	}
 	close(operationChannel)