github.com/fzfile/BaiduPCS-Go@v0.0.0-20200606205115-4408961cf336/requester/uploader/block.go (about)

     1  package uploader
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"github.com/fzfile/BaiduPCS-Go/requester/rio/speeds"
     7  	"github.com/fzfile/BaiduPCS-Go/requester/transfer"
     8  	"io"
     9  	"os"
    10  	"sync"
    11  )
    12  
    13  type (
    14  	// SplitUnit 将 io.ReaderAt 分割单元
    15  	SplitUnit interface {
    16  		Readed64
    17  		io.Seeker
    18  		Range() transfer.Range
    19  		Left() int64
    20  	}
    21  
    22  	fileBlock struct {
    23  		readRange     transfer.Range
    24  		readed        int64
    25  		readerAt      io.ReaderAt
    26  		speedsStatRef *speeds.Speeds
    27  		rateLimit     *speeds.RateLimit
    28  		mu            sync.Mutex
    29  	}
    30  
    31  	bufioFileBlock struct {
    32  		*fileBlock
    33  		bufio *bufio.Reader
    34  	}
    35  )
    36  
    37  // SplitBlock 文件分块
    38  func SplitBlock(fileSize, blockSize int64) (blockList []*BlockState) {
    39  	gen := transfer.NewRangeListGenBlockSize(fileSize, 0, blockSize)
    40  	rangeCount := gen.RangeCount()
    41  	blockList = make([]*BlockState, 0, rangeCount)
    42  	for i := 0; i < rangeCount; i++ {
    43  		id, r := gen.GenRange()
    44  		blockList = append(blockList, &BlockState{
    45  			ID:    id,
    46  			Range: *r,
    47  		})
    48  	}
    49  	return
    50  }
    51  
    52  // NewBufioSplitUnit io.ReaderAt实现SplitUnit接口, 有Buffer支持
    53  func NewBufioSplitUnit(readerAt io.ReaderAt, readRange transfer.Range, speedsStat *speeds.Speeds, rateLimit *speeds.RateLimit) SplitUnit {
    54  	su := &fileBlock{
    55  		readerAt:      readerAt,
    56  		readRange:     readRange,
    57  		speedsStatRef: speedsStat,
    58  		rateLimit:     rateLimit,
    59  	}
    60  	return &bufioFileBlock{
    61  		fileBlock: su,
    62  		bufio:     bufio.NewReaderSize(su, BufioReadSize),
    63  	}
    64  }
    65  
    66  func (bfb *bufioFileBlock) Read(b []byte) (n int, err error) {
    67  	return bfb.bufio.Read(b) // 间接调用fileBlock 的Read
    68  }
    69  
    70  // Read 只允许一个线程读同一个文件
    71  func (fb *fileBlock) Read(b []byte) (n int, err error) {
    72  	fb.mu.Lock()
    73  	defer fb.mu.Unlock()
    74  
    75  	left := int(fb.Left())
    76  	if left <= 0 {
    77  		return 0, io.EOF
    78  	}
    79  
    80  	if len(b) > left {
    81  		n, err = fb.readerAt.ReadAt(b[:left], fb.readed+fb.readRange.Begin)
    82  	} else {
    83  		n, err = fb.readerAt.ReadAt(b, fb.readed+fb.readRange.Begin)
    84  	}
    85  
    86  	n64 := int64(n)
    87  	fb.readed += n64
    88  	if fb.rateLimit != nil {
    89  		fb.rateLimit.Add(n64) // 限速阻塞
    90  	}
    91  	if fb.speedsStatRef != nil {
    92  		fb.speedsStatRef.Add(n64)
    93  	}
    94  	return
    95  }
    96  
    97  func (fb *fileBlock) Seek(offset int64, whence int) (int64, error) {
    98  	fb.mu.Lock()
    99  	defer fb.mu.Unlock()
   100  
   101  	switch whence {
   102  	case os.SEEK_SET:
   103  		fb.readed = offset
   104  	case os.SEEK_CUR:
   105  		fb.readed += offset
   106  	case os.SEEK_END:
   107  		fb.readed = fb.readRange.End - fb.readRange.Begin + offset
   108  	default:
   109  		return 0, fmt.Errorf("unsupport whence: %d", whence)
   110  	}
   111  	if fb.readed < 0 {
   112  		fb.readed = 0
   113  	}
   114  	return fb.readed, nil
   115  }
   116  
   117  func (fb *fileBlock) Len() int64 {
   118  	return fb.readRange.End - fb.readRange.Begin
   119  }
   120  
   121  func (fb *fileBlock) Left() int64 {
   122  	return fb.readRange.End - fb.readRange.Begin - fb.readed
   123  }
   124  
   125  func (fb *fileBlock) Range() transfer.Range {
   126  	return fb.readRange
   127  }
   128  
   129  func (fb *fileBlock) Readed() int64 {
   130  	return fb.readed
   131  }