github.com/qjfoidnh/BaiduPCS-Go@v0.0.0-20231011165705-caa18a3765f3/requester/uploader/multiworker.go (about)

     1  package uploader
     2  
     3  import (
     4  	"context"
     5  	"github.com/oleiade/lane"
     6  	"github.com/qjfoidnh/BaiduPCS-Go/pcsutil/waitgroup"
     7  	"os"
     8  )
     9  
    10  type (
    11  	worker struct {
    12  		id         int
    13  		partOffset int64
    14  		splitUnit  SplitUnit
    15  		checksum   string
    16  	}
    17  
    18  	workerList []*worker
    19  )
    20  
    21  // CheckSumList 返回所以worker的checksum
    22  // TODO: 实现sort
    23  func (werl *workerList) CheckSumList() []string {
    24  	checksumList := make([]string, 0, len(*werl))
    25  	for _, wer := range *werl {
    26  		checksumList = append(checksumList, wer.checksum)
    27  	}
    28  	return checksumList
    29  }
    30  
    31  func (werl *workerList) Readed() int64 {
    32  	var readed int64
    33  	for _, wer := range *werl {
    34  		readed += wer.splitUnit.Readed()
    35  	}
    36  	return readed
    37  }
    38  
    39  func (muer *MultiUploader) upload() (uperr error) {
    40  	err := muer.multiUpload.Precreate(muer.file.Len(), muer.config.Policy)
    41  	if err != nil {
    42  		return err
    43  	}
    44  	var (
    45  		uploadDeque = lane.NewDeque()
    46  	)
    47  
    48  	// 加入队列
    49  	for _, wer := range muer.workers {
    50  		if wer.checksum == "" {
    51  			uploadDeque.Append(wer)
    52  		}
    53  	}
    54  
    55  	for {
    56  		wg := waitgroup.NewWaitGroup(muer.config.Parallel)
    57  		for {
    58  			e := uploadDeque.Shift()
    59  			if e == nil { // 任务为空
    60  				break
    61  			}
    62  
    63  			wer := e.(*worker)
    64  			wg.AddDelta()
    65  			go func() {
    66  				defer wg.Done()
    67  
    68  				var (
    69  					ctx, cancel = context.WithCancel(context.Background())
    70  					doneChan    = make(chan struct{})
    71  					checksum    string
    72  					terr        error
    73  				)
    74  				go func() {
    75  					checksum, terr = muer.multiUpload.TmpFile(ctx, int(wer.id), wer.partOffset, wer.splitUnit)
    76  					close(doneChan)
    77  				}()
    78  				select {
    79  				case <-muer.canceled:
    80  					cancel()
    81  					return
    82  				case <-doneChan:
    83  					// continue
    84  				}
    85  				cancel()
    86  				if terr != nil {
    87  					if me, ok := terr.(*MultiError); ok {
    88  						if me.Terminated { // 终止
    89  							muer.closeCanceledOnce.Do(func() { // 只关闭一次
    90  								close(muer.canceled)
    91  							})
    92  							uperr = me.Err
    93  							return
    94  						}
    95  					}
    96  
    97  					uploaderVerbose.Warnf("upload err: %s, id: %d\n", terr, wer.id)
    98  					wer.splitUnit.Seek(0, os.SEEK_SET)
    99  					uploadDeque.Append(wer)
   100  					return
   101  				}
   102  				wer.checksum = checksum
   103  
   104  				// 通知更新
   105  				if muer.updateInstanceStateChan != nil && len(muer.updateInstanceStateChan) < cap(muer.updateInstanceStateChan) {
   106  					muer.updateInstanceStateChan <- struct{}{}
   107  				}
   108  			}()
   109  		}
   110  		wg.Wait()
   111  
   112  		// 没有任务了
   113  		if uploadDeque.Size() == 0 {
   114  			break
   115  		}
   116  	}
   117  
   118  	select {
   119  	case <-muer.canceled:
   120  		if uperr != nil {
   121  			return uperr
   122  		}
   123  		return context.Canceled
   124  	default:
   125  	}
   126  
   127  	cerr := muer.multiUpload.CreateSuperFile(muer.config.Policy, muer.workers.CheckSumList()...)
   128  	if cerr != nil {
   129  		return cerr
   130  	}
   131  
   132  	return
   133  }