github.com/qiniu/go-sdk/v7@v7.27.0/storagev2/uploader/schedulers.go (about)

     1  package uploader
     2  
     3  import (
     4  	"context"
     5  	"sort"
     6  	"sync"
     7  
     8  	"github.com/qiniu/go-sdk/v7/storagev2/uploader/source"
     9  	"golang.org/x/sync/errgroup"
    10  )
    11  
    12  type (
    13  	serialMultiPartsUploaderScheduler struct {
    14  		uploader MultiPartsUploader
    15  		partSize uint64
    16  	}
    17  
    18  	// 串行分片上传调度器选项
    19  	serialMultiPartsUploaderSchedulerOptions struct {
    20  		PartSize uint64 // 分片大小
    21  	}
    22  
    23  	concurrentMultiPartsUploaderScheduler struct {
    24  		uploader    MultiPartsUploader
    25  		partSize    uint64
    26  		concurrency int
    27  	}
    28  
    29  	// 并行分片上传调度器选项
    30  	concurrentMultiPartsUploaderSchedulerOptions struct {
    31  		PartSize    uint64 // 分片大小
    32  		Concurrency int    // 并发度
    33  	}
    34  )
    35  
    36  // 创建串行分片上传调度器
    37  func newSerialMultiPartsUploaderScheduler(uploader MultiPartsUploader, options *serialMultiPartsUploaderSchedulerOptions) multiPartsUploaderScheduler {
    38  	if options == nil {
    39  		options = &serialMultiPartsUploaderSchedulerOptions{}
    40  	}
    41  	partSize := options.PartSize
    42  	if partSize == 0 {
    43  		partSize = 1 << 22
    44  	} else if partSize < (1 << 20) {
    45  		partSize = 1 << 20
    46  	} else if partSize > (1 << 30) {
    47  		partSize = 1 << 30
    48  	}
    49  	return serialMultiPartsUploaderScheduler{uploader, partSize}
    50  }
    51  
    52  // 创建并行分片上传调度器
    53  func newConcurrentMultiPartsUploaderScheduler(uploader MultiPartsUploader, options *concurrentMultiPartsUploaderSchedulerOptions) multiPartsUploaderScheduler {
    54  	if options == nil {
    55  		options = &concurrentMultiPartsUploaderSchedulerOptions{}
    56  	}
    57  	partSize := options.PartSize
    58  	if partSize == 0 {
    59  		partSize = 1 << 22
    60  	} else if partSize < (1 << 20) {
    61  		partSize = 1 << 20
    62  	} else if partSize > (1 << 30) {
    63  		partSize = 1 << 30
    64  	}
    65  	concurrency := options.Concurrency
    66  	if concurrency <= 0 {
    67  		concurrency = 4
    68  	}
    69  
    70  	return concurrentMultiPartsUploaderScheduler{uploader, partSize, concurrency}
    71  }
    72  
    73  func (scheduler serialMultiPartsUploaderScheduler) UploadParts(ctx context.Context, initialized InitializedParts, src source.Source, options *UploadPartsOptions) ([]UploadedPart, error) {
    74  	parts := make([]UploadedPart, 0)
    75  	for {
    76  		part, err := src.Slice(scheduler.partSize)
    77  		if err != nil {
    78  			return nil, err
    79  		}
    80  		if part == nil {
    81  			break
    82  		}
    83  		var uploadPartParam UploadPartOptions
    84  		if options != nil && options.OnUploadingProgress != nil {
    85  			uploadPartParam.OnUploadingProgress = func(progress *UploadingPartProgress) {
    86  				options.OnUploadingProgress(part.PartNumber(), &UploadingPartProgress{Uploaded: progress.Uploaded, PartSize: part.Size()})
    87  			}
    88  		}
    89  		uploadedPart, err := scheduler.uploader.UploadPart(ctx, initialized, part, &uploadPartParam)
    90  		if err != nil {
    91  			return nil, err
    92  		}
    93  		if options != nil && options.OnPartUploaded != nil {
    94  			if err = options.OnPartUploaded(uploadedPart); err != nil {
    95  				return nil, err
    96  			}
    97  		}
    98  		parts = append(parts, uploadedPart)
    99  	}
   100  	return parts, nil
   101  }
   102  
   103  func (scheduler serialMultiPartsUploaderScheduler) MultiPartsUploader() MultiPartsUploader {
   104  	return scheduler.uploader
   105  }
   106  
   107  func (scheduler serialMultiPartsUploaderScheduler) PartSize() uint64 {
   108  	return scheduler.partSize
   109  }
   110  
   111  func (scheduler concurrentMultiPartsUploaderScheduler) UploadParts(ctx context.Context, initialized InitializedParts, src source.Source, options *UploadPartsOptions) ([]UploadedPart, error) {
   112  	var (
   113  		parts     []UploadedPart
   114  		partsLock sync.Mutex
   115  	)
   116  	if ss, ok := src.(source.SizedSource); ok {
   117  		totalSize, err := ss.TotalSize()
   118  		if err != nil {
   119  			return nil, err
   120  		}
   121  		partsCount := (totalSize + scheduler.partSize - 1) / scheduler.partSize
   122  		parts = make([]UploadedPart, 0, partsCount)
   123  	}
   124  	g, ctx := errgroup.WithContext(ctx)
   125  	g.SetLimit(scheduler.concurrency)
   126  	var onUploadingProgressMutex sync.Mutex
   127  	for {
   128  		part, err := src.Slice(scheduler.partSize)
   129  		if err != nil {
   130  			return nil, err
   131  		}
   132  		if part == nil {
   133  			break
   134  		}
   135  		g.Go(func() error {
   136  			var uploadPartParam UploadPartOptions
   137  			if options != nil && options.OnUploadingProgress != nil {
   138  				uploadPartParam.OnUploadingProgress = func(progress *UploadingPartProgress) {
   139  					onUploadingProgressMutex.Lock()
   140  					defer onUploadingProgressMutex.Unlock()
   141  					options.OnUploadingProgress(part.PartNumber(), progress)
   142  				}
   143  			}
   144  			uploadedPart, err := scheduler.uploader.UploadPart(ctx, initialized, part, &uploadPartParam)
   145  			if err != nil {
   146  				return err
   147  			}
   148  			if options != nil && options.OnPartUploaded != nil {
   149  				if err = options.OnPartUploaded(uploadedPart); err != nil {
   150  					return err
   151  				}
   152  			}
   153  
   154  			partsLock.Lock()
   155  			defer partsLock.Unlock()
   156  			parts = append(parts, uploadedPart)
   157  			return nil
   158  		})
   159  	}
   160  	if err := g.Wait(); err != nil {
   161  		return nil, err
   162  	}
   163  	sort.Slice(parts, func(i, j int) bool {
   164  		return parts[i].Offset() < parts[j].Offset()
   165  	})
   166  	return parts, nil
   167  }
   168  
   169  func (scheduler concurrentMultiPartsUploaderScheduler) MultiPartsUploader() MultiPartsUploader {
   170  	return scheduler.uploader
   171  }
   172  
   173  func (scheduler concurrentMultiPartsUploaderScheduler) PartSize() uint64 {
   174  	return scheduler.partSize
   175  }