Skip to content

Commit

Permalink
azfile, azdatalake: Updating the type of number of chunks to uint64 (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
souravgupta-msft authored Feb 29, 2024
1 parent 2b23e14 commit da33ad0
Show file tree
Hide file tree
Showing 10 changed files with 171 additions and 29 deletions.
1 change: 1 addition & 0 deletions sdk/storage/azdatalake/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
* Re-enabled `SharedKeyCredential` authentication mode for non TLS protected endpoints.

### Other Changes
* Updated version of azblob to `1.3.1`

## 1.1.0 (2024-02-14)

Expand Down
77 changes: 77 additions & 0 deletions sdk/storage/azdatalake/file/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,15 @@ import (
"crypto/md5"
"encoding/binary"
"fmt"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
"hash/crc64"
"io"
"math/rand"
"net/http"
"os"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"

Expand Down Expand Up @@ -4986,3 +4988,78 @@ func (s *UnrecordedTestSuite) TestFileCreateDeleteUsingOAuth() {
_, err = fClient.GetProperties(context.Background(), nil)
_require.NoError(err)
}

type fakeUploadBlob struct {
numChunks uint64
}

// nolint
func (f *fakeUploadBlob) Do(req *http.Request) (*http.Response, error) {
statusCode := http.StatusOK

// check how many times append data is called
if val, ok := req.URL.Query()["action"]; ok {
if val[0] == "append" {
statusCode = http.StatusAccepted
atomic.AddUint64(&f.numChunks, 1)
}
}

return &http.Response{
Request: req,
Status: "Created",
StatusCode: statusCode,
Header: http.Header{},
Body: http.NoBody,
}, nil
}

func TestUploadSmallChunkSize(t *testing.T) {
_require := require.New(t)

fileSize := int64(100 * 1024 * 1024)
chunkSize := int64(1024)
numChunks := uint64(((fileSize - 1) / chunkSize) + 1)
fbb := &fakeUploadBlob{}

log.SetListener(nil) // no logging

fClient, err := file.NewClientWithNoCredential("https://fake/blob/path", &file.ClientOptions{
ClientOptions: policy.ClientOptions{
Transport: fbb,
},
})
_require.NoError(err)
_require.NotNil(fClient)

// create local file
_, content := generateData(int(fileSize))
err = os.WriteFile("testFile", content, 0644)
_require.NoError(err)

defer func() {
err = os.Remove("testFile")
_require.NoError(err)
}()

fh, err := os.Open("testFile")
_require.NoError(err)

defer func(fh *os.File) {
err := fh.Close()
_require.NoError(err)
}(fh)

err = fClient.UploadFile(context.Background(), fh, &file.UploadFileOptions{ChunkSize: chunkSize})
_require.NoError(err)

_require.Equal(atomic.LoadUint64(&fbb.numChunks), numChunks)

// reset counter
atomic.StoreUint64(&fbb.numChunks, 0)

err = fClient.UploadBuffer(context.Background(), content, &file.UploadBufferOptions{ChunkSize: chunkSize})
_require.NoError(err)

_require.Equal(atomic.LoadUint64(&fbb.numChunks), numChunks)
}
2 changes: 1 addition & 1 deletion sdk/storage/azdatalake/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ require (
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.9.2
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.5.1
github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.0
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.1
github.com/stretchr/testify v1.8.4
)

Expand Down
4 changes: 2 additions & 2 deletions sdk/storage/azdatalake/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.5.1/go.mod h1:h8hyGFDsU5HMiv
github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2 h1:LqbJ/WzJUwBf8UiaSzgX7aMclParm9/5Vgp+TY51uBQ=
github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2/go.mod h1:yInRyqWXAuaPrgI7p70+lDDgh3mlBohis29jGMISnmc=
github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage v1.5.0 h1:AifHbc4mg0x9zW52WOpKbsHaDKuRhlI7TVl47thgQ70=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.0 h1:IfFdxTUDiV58iZqPKgyWiz4X4fCxZeQ1pTQPImLYXpY=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.0/go.mod h1:SUZc9YRRHfx2+FAQKNDGrssXehqLpxmwRv2mC/5ntj4=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.1 h1:fXPMAmuh0gDuRDey0atC8cXBuKIlqCzCkL8sm1n9Ov0=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.1/go.mod h1:SUZc9YRRHfx2+FAQKNDGrssXehqLpxmwRv2mC/5ntj4=
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1 h1:DzHpqpoJVaCgOUdVHxE8QB52S6NiVdDQvGlny1qvPqA=
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
Expand Down
6 changes: 3 additions & 3 deletions sdk/storage/azdatalake/internal/shared/batch_transfer.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {
}

// Prepare and do parallel operations.
numChunks := uint16(((o.TransferSize - 1) / o.ChunkSize) + 1)
numChunks := uint64(((o.TransferSize - 1) / o.ChunkSize) + 1)
operationChannel := make(chan func() error, o.Concurrency) // Create the channel that release 'concurrency' goroutines concurrently
operationResponseChannel := make(chan error, numChunks) // Holds each response
ctx, cancel := context.WithCancel(ctx)
Expand All @@ -50,7 +50,7 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {
}

// Add each chunk's operation to the channel.
for chunkNum := uint16(0); chunkNum < numChunks; chunkNum++ {
for chunkNum := uint64(0); chunkNum < numChunks; chunkNum++ {
curChunkSize := o.ChunkSize

if chunkNum == numChunks-1 { // Last chunk
Expand All @@ -65,7 +65,7 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {

// Wait for the operations to complete.
var firstErr error = nil
for chunkNum := uint16(0); chunkNum < numChunks; chunkNum++ {
for chunkNum := uint64(0); chunkNum < numChunks; chunkNum++ {
responseError := <-operationResponseChannel
// record the first error (the original error which should cause the other chunks to fail with canceled context)
if responseError != nil && firstErr == nil {
Expand Down
2 changes: 2 additions & 0 deletions sdk/storage/azfile/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@

### Other Changes

* Updated `azidentity` version to `1.5.1`.

## 1.2.0 (2024-02-12)

### Other Changes
Expand Down
62 changes: 62 additions & 0 deletions sdk/storage/azfile/file/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"crypto/rand"
"encoding/binary"
"fmt"
"github.com/Azure/azure-sdk-for-go/sdk/azcore/log"
"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"
Expand All @@ -34,6 +35,7 @@ import (
"net/http"
"os"
"strings"
"sync/atomic"
"testing"
"time"
)
Expand Down Expand Up @@ -4624,4 +4626,64 @@ func (f *FileRecordedTestsSuite) TestFileClientAudienceNegative() {
testcommon.ValidateFileErrorCode(_require, err, fileerror.AuthenticationFailed)
}

type fakeDownloadFile struct {
contentSize int64
numChunks uint64
}

// nolint
func (f *fakeDownloadFile) Do(req *http.Request) (*http.Response, error) {
// check how many times range based get file is called
if _, ok := req.Header["x-ms-range"]; ok {
atomic.AddUint64(&f.numChunks, 1)
}
return &http.Response{
Request: req,
Status: "Created",
StatusCode: http.StatusOK,
Header: http.Header{"Content-Length": []string{fmt.Sprintf("%v", f.contentSize)}},
Body: http.NoBody,
}, nil
}

func TestDownloadSmallChunkSize(t *testing.T) {
_require := require.New(t)

fileSize := int64(100 * 1024 * 1024)
chunkSize := int64(1024)
numChunks := uint64(((fileSize - 1) / chunkSize) + 1)
fbb := &fakeDownloadFile{
contentSize: fileSize,
}

log.SetListener(nil) // no logging

fileClient, err := file.NewClientWithNoCredential("https://fake/file/path", &file.ClientOptions{
ClientOptions: policy.ClientOptions{
Transport: fbb,
},
})
_require.NoError(err)
_require.NotNil(fileClient)

// download to a temp file and verify contents
tmp, err := os.CreateTemp("", "")
_require.NoError(err)
defer tmp.Close()

_, err = fileClient.DownloadFile(context.Background(), tmp, &file.DownloadFileOptions{ChunkSize: chunkSize})
_require.NoError(err)

_require.Equal(atomic.LoadUint64(&fbb.numChunks), numChunks)

// reset counter
atomic.StoreUint64(&fbb.numChunks, 0)

buff := make([]byte, fileSize)
_, err = fileClient.DownloadBuffer(context.Background(), buff, &file.DownloadBufferOptions{ChunkSize: chunkSize})
_require.NoError(err)

_require.Equal(atomic.LoadUint64(&fbb.numChunks), numChunks)
}

// TODO: Add tests for retry header options
12 changes: 6 additions & 6 deletions sdk/storage/azfile/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,20 +4,20 @@ go 1.18

require (
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.9.2
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.4.0
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.5.1
github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.2.0
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.1
github.com/stretchr/testify v1.8.4
)

require (
github.com/AzureAD/microsoft-authentication-library-for-go v1.1.1 // indirect
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/dnaeon/go-vcr v1.2.0 // indirect
github.com/golang-jwt/jwt/v5 v5.0.0 // indirect
github.com/google/uuid v1.3.1 // indirect
github.com/golang-jwt/jwt/v5 v5.2.0 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/kylelemons/godebug v1.1.0 // indirect
github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 // indirect
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
golang.org/x/crypto v0.18.0 // indirect
golang.org/x/net v0.20.0 // indirect
Expand Down
28 changes: 14 additions & 14 deletions sdk/storage/azfile/go.sum
Original file line number Diff line number Diff line change
@@ -1,27 +1,27 @@
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.9.2 h1:c4k2FIYIh4xtwqrQwV0Ct1v5+ehlNXj5NI/MWVsiTkQ=
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.9.2/go.mod h1:5FDJtLEO/GxwNgUxbwrY3LP0pEoThTQJtk2oysdXHxM=
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.4.0 h1:BMAjVKJM0U/CYF27gA0ZMmXGkOcvfFtD0oHVZ1TIPRI=
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.4.0/go.mod h1:1fXstnBMas5kzG+S3q8UoJcmyU6nUeunJcMDHcRYHhs=
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.5.1 h1:sO0/P7g68FrryJzljemN+6GTssUXdANk6aJ7T1ZxnsQ=
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.5.1/go.mod h1:h8hyGFDsU5HMivxiS2iYFZsgDbU9OnnJ163x5UGVKYo=
github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2 h1:LqbJ/WzJUwBf8UiaSzgX7aMclParm9/5Vgp+TY51uBQ=
github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2/go.mod h1:yInRyqWXAuaPrgI7p70+lDDgh3mlBohis29jGMISnmc=
github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage v1.2.0 h1:Ma67P/GGprNwsslzEH6+Kb8nybI8jpDTm4Wmzu2ReK8=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.2.0 h1:gggzg0SUMs6SQbEw+3LoSsYf9YMjkupeAnHMX8O9mmY=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.2.0/go.mod h1:+6KLcKIVgxoBDMqMO/Nvy7bZ9a0nbU3I1DtFQK3YvB4=
github.com/AzureAD/microsoft-authentication-library-for-go v1.1.1 h1:WpB/QDNLpMw72xHJc34BNNykqSOeEJDAWkhf0u12/Jk=
github.com/AzureAD/microsoft-authentication-library-for-go v1.1.1/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage v1.5.0 h1:AifHbc4mg0x9zW52WOpKbsHaDKuRhlI7TVl47thgQ70=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.1 h1:fXPMAmuh0gDuRDey0atC8cXBuKIlqCzCkL8sm1n9Ov0=
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.3.1/go.mod h1:SUZc9YRRHfx2+FAQKNDGrssXehqLpxmwRv2mC/5ntj4=
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1 h1:DzHpqpoJVaCgOUdVHxE8QB52S6NiVdDQvGlny1qvPqA=
github.com/AzureAD/microsoft-authentication-library-for-go v1.2.1/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI=
github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
github.com/golang-jwt/jwt/v5 v5.0.0 h1:1n1XNM9hk7O9mnQoNBGolZvzebBQ7p93ULHRc28XJUE=
github.com/golang-jwt/jwt/v5 v5.0.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
github.com/google/uuid v1.3.1 h1:KjJaJ9iWZ3jOFZIf1Lqf4laDRCasjl0BCmnEGxkdLb4=
github.com/google/uuid v1.3.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/golang-jwt/jwt/v5 v5.2.0 h1:d/ix8ftRUorsN+5eMIlF4T6J8CAt9rch3My2winC1Jw=
github.com/golang-jwt/jwt/v5 v5.2.0/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
github.com/modocache/gover v0.0.0-20171022184752-b58185e213c5/go.mod h1:caMODM3PzxT8aQXRPkAt8xlV/e7d7w8GM5g0fa5F0D8=
github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 h1:KoWmjvw+nsYOo29YJK9vDA65RGE3NrOnUtO7a+RF9HU=
github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8/go.mod h1:HKlIX3XHQyzLZPlr7++PzdhaXEj94dEiJgZDTsxEqUI=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
Expand All @@ -30,7 +30,7 @@ golang.org/x/crypto v0.18.0 h1:PGVlW0xEltQnzFZ55hkuX5+KLyrMYhHld1YHO4AKcdc=
golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1mg=
golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo=
golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
golang.org/x/sys v0.0.0-20210616045830-e2b7044e8c71/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
Expand Down
6 changes: 3 additions & 3 deletions sdk/storage/azfile/internal/shared/batch_transfer.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {
}

// Prepare and do parallel operations.
numChunks := uint16(((o.TransferSize - 1) / o.ChunkSize) + 1)
numChunks := uint64(((o.TransferSize - 1) / o.ChunkSize) + 1)
operationChannel := make(chan func() error, o.Concurrency) // Create the channel that release 'concurrency' goroutines concurrently
operationResponseChannel := make(chan error, numChunks) // Holds each response
ctx, cancel := context.WithCancel(ctx)
Expand All @@ -50,7 +50,7 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {
}

// Add each chunk's operation to the channel.
for chunkNum := uint16(0); chunkNum < numChunks; chunkNum++ {
for chunkNum := uint64(0); chunkNum < numChunks; chunkNum++ {
curChunkSize := o.ChunkSize

if chunkNum == numChunks-1 { // Last chunk
Expand All @@ -65,7 +65,7 @@ func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {

// Wait for the operations to complete.
var firstErr error = nil
for chunkNum := uint16(0); chunkNum < numChunks; chunkNum++ {
for chunkNum := uint64(0); chunkNum < numChunks; chunkNum++ {
responseError := <-operationResponseChannel
// record the first error (the original error which should cause the other chunks to fail with canceled context)
if responseError != nil && firstErr == nil {
Expand Down

0 comments on commit da33ad0

Please sign in to comment.