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

     1  package partial
     2  
     3  import (
     4  	"sync/atomic"
     5  )
     6  
     7  type Chunk struct {
     8  	// 切片的顺序
     9  	index int64
    10  
    11  	// 切片内容的在源文件的开始地址
    12  	start int64
    13  
    14  	// 切片内容在源文件的结束地址
    15  	end int64
    16  
    17  	// 切片任务的下载错误
    18  	err error
    19  
    20  	// 下载完的切片的具体内容
    21  	buffer []byte
    22  }
    23  
    24  func NewChunk(index, start, end int64) *Chunk {
    25  	chunk := &Chunk{
    26  		start: start,
    27  		end:   end,
    28  		index: index,
    29  	}
    30  	return chunk
    31  }
    32  
    33  func (p *Chunk) SetData(bytes []byte) {
    34  	p.buffer = bytes
    35  }
    36  
    37  func (p *Chunk) SetError(err error) {
    38  	p.err = err
    39  }
    40  
    41  func (p *Chunk) Error() error {
    42  	return p.err
    43  }
    44  
    45  func (p *Chunk) Data() []byte {
    46  	return p.buffer
    47  }
    48  
    49  // 切片乱序写入后，将切片顺序读取
    50  type ChunksSorter struct {
    51  	// 已经读取的切片数量
    52  	readCount int64
    53  
    54  	// 切片的所有总数
    55  	chunkCount int64
    56  
    57  	// 线程数，用于阻塞写入
    58  	works int64
    59  
    60  	// 存储切片的缓存区
    61  	chunks []chan *Chunk
    62  }
    63  
    64  func NewChunksSorter(chunkCount int64, works int) *ChunksSorter {
    65  	chunks := make([]chan *Chunk, works)
    66  	for i := 0; i < len(chunks); i++ {
    67  		chunks[i] = make(chan *Chunk)
    68  	}
    69  
    70  	return &ChunksSorter{
    71  		chunkCount: chunkCount,
    72  		works:      int64(works),
    73  		chunks:     chunks,
    74  	}
    75  }
    76  
    77  // 将数据写入到缓存区，如果该缓存已满，则会被阻塞
    78  func (p *ChunksSorter) Write(chunk *Chunk) {
    79  	p.chunks[chunk.index%p.works] <- chunk
    80  }
    81  
    82  // 关闭 workId 下的通道
    83  func (p *ChunksSorter) Close(workId int) {
    84  	if (len(p.chunks) - 1) >= workId {
    85  		close(p.chunks[workId])
    86  	}
    87  }
    88  
    89  // 顺序读取切片，如果下一个切片没有下载完，则会被阻塞
    90  func (p *ChunksSorter) Read() *Chunk {
    91  	if p.chunkCount == 0 {
    92  		return nil
    93  	}
    94  	i := atomic.AddInt64(&p.readCount, 1)
    95  	chunk := <-p.chunks[(i-1)%p.works]
    96  	return chunk
    97  }