github.com/upyun/upx@v0.4.7-0.20240419023638-b184a7cb7c10/partial/downloader.go (about)

     1  package partial
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  	"os"
     8  	"sync"
     9  )
    10  
    11  const DefaultChunkSize = 1024 * 1024 * 10
    12  
    13  type ChunkDownFunc func(start, end int64) ([]byte, error)
    14  
    15  type MultiPartialDownloader struct {
    16  
    17  	// 文件路径
    18  	filePath string
    19  
    20  	// 最终文件大小
    21  	finalSize int64
    22  
    23  	// 本地文件大小
    24  	localSize int64
    25  
    26  	//分片大小
    27  	chunkSize int64
    28  
    29  	writer   io.Writer
    30  	works    int
    31  	downFunc ChunkDownFunc
    32  }
    33  
    34  func NewMultiPartialDownloader(filePath string, finalSize, chunkSize int64, writer io.Writer, works int, fn ChunkDownFunc) *MultiPartialDownloader {
    35  	return &MultiPartialDownloader{
    36  		filePath:  filePath,
    37  		finalSize: finalSize,
    38  		works:     works,
    39  		writer:    writer,
    40  		chunkSize: chunkSize,
    41  		downFunc:  fn,
    42  	}
    43  }
    44  
    45  func (p *MultiPartialDownloader) Download() error {
    46  	fileinfo, err := os.Stat(p.filePath)
    47  
    48  	// 如果异常
    49  	// - 文件不存在异常: localSize 默认值 0
    50  	// - 不是文件不存在异常: 报错
    51  	if err != nil && !os.IsNotExist(err) {
    52  		return err
    53  	}
    54  	if err == nil {
    55  		p.localSize = fileinfo.Size()
    56  	}
    57  
    58  	// 计算需要下载的块数
    59  	needDownSize := p.finalSize - p.localSize
    60  	chunkCount := needDownSize / p.chunkSize
    61  	if needDownSize%p.chunkSize != 0 {
    62  		chunkCount++
    63  	}
    64  
    65  	chunksSorter := NewChunksSorter(
    66  		chunkCount,
    67  		p.works,
    68  	)
    69  
    70  	// 下载切片任务
    71  	var wg sync.WaitGroup
    72  	ctx, cancel := context.WithCancel(context.Background())
    73  	defer func() {
    74  		// 取消切片下载任务，并等待
    75  		cancel()
    76  		wg.Wait()
    77  	}()
    78  
    79  	for i := 0; i < p.works; i++ {
    80  		wg.Add(1)
    81  		go func(ctx context.Context, workId int) {
    82  			defer func() {
    83  				// 关闭 workId 下的接收通道
    84  				chunksSorter.Close(workId)
    85  				wg.Done()
    86  			}()
    87  
    88  			// 每个 work 取自己倍数的 chunk
    89  			for j := workId; j < int(chunkCount); j += p.works {
    90  				select {
    91  				case <-ctx.Done():
    92  					return
    93  				default:
    94  					var (
    95  						err    error
    96  						buffer []byte
    97  					)
    98  					start := p.localSize + int64(j)*p.chunkSize
    99  					end := p.localSize + int64(j+1)*p.chunkSize
   100  					if end > p.finalSize {
   101  						end = p.finalSize
   102  					}
   103  					chunk := NewChunk(int64(j), start, end)
   104  
   105  					// 重试三次
   106  					for t := 0; t < 3; t++ {
   107  						// ? 由于长度是从1开始，而数据是从0地址开始
   108  						// ? 计算字节时容量会多出开头的一位，所以末尾需要减少一位
   109  						buffer, err = p.downFunc(chunk.start, chunk.end-1)
   110  						if err == nil {
   111  							break
   112  						}
   113  					}
   114  					chunk.SetData(buffer)
   115  					chunk.SetError(err)
   116  					chunksSorter.Write(chunk)
   117  
   118  					if err != nil {
   119  						return
   120  					}
   121  				}
   122  			}
   123  		}(ctx, i)
   124  	}
   125  
   126  	// 将分片顺序写入到文件
   127  	for {
   128  		chunk := chunksSorter.Read()
   129  		if chunk == nil {
   130  			break
   131  		}
   132  		if chunk.Error() != nil {
   133  			return chunk.Error()
   134  		}
   135  		if len(chunk.Data()) == 0 {
   136  			return errors.New("chunk buffer download but size is 0")
   137  		}
   138  		p.writer.Write(chunk.Data())
   139  	}
   140  	return nil
   141  }