-
Notifications
You must be signed in to change notification settings - Fork 798
/
batch_transfer.go
77 lines (66 loc) · 2.41 KB
/
batch_transfer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
//go:build go1.18
// +build go1.18
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License. See License.txt in the project root for license information.
package shared
import (
"context"
"errors"
)
// BatchTransferOptions identifies options used by doBatchTransfer.
type BatchTransferOptions struct {
TransferSize int64
ChunkSize int64
Concurrency uint16
Operation func(ctx context.Context, offset int64, chunkSize int64) error
OperationName string
}
// DoBatchTransfer helps to execute operations in a batch manner.
// Can be used by users to customize batch works (for other scenarios that the SDK does not provide)
func DoBatchTransfer(ctx context.Context, o *BatchTransferOptions) error {
if o.ChunkSize == 0 {
return errors.New("ChunkSize cannot be 0")
}
if o.Concurrency == 0 {
o.Concurrency = 5 // default concurrency
}
// Prepare and do parallel operations.
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)
defer cancel()
// Create the goroutines that process each operation (in parallel).
for g := uint16(0); g < o.Concurrency; g++ {
//grIndex := g
go func() {
for f := range operationChannel {
err := f()
operationResponseChannel <- err
}
}()
}
// Add each chunk's operation to the channel.
for chunkNum := uint64(0); chunkNum < numChunks; chunkNum++ {
curChunkSize := o.ChunkSize
if chunkNum == numChunks-1 { // Last chunk
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(ctx, offset, curChunkSize)
}
}
close(operationChannel)
// Wait for the operations to complete.
var firstErr error = nil
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 {
cancel() // As soon as any operation fails, cancel all remaining operation calls
firstErr = responseError
}
}
return firstErr
}