1
1
mirror of https://github.com/wader/fq.git synced 2024-12-24 22:05:31 +03:00
fq/internal/progressreadseeker/progressreaderseeker.go
2021-11-21 21:55:53 +01:00

67 lines
1.4 KiB
Go

package progressreadseeker
import (
"io"
)
type ProgressFn func(approxReadBytes int64, totalSize int64)
type Reader struct {
rs io.ReadSeeker
pos int64
totalSize int64
partitionSize int64
partitions []bool
partitionsReadCount int64
progressFn ProgressFn
}
func New(rs io.ReadSeeker, precision int64, totalSize int64, fn ProgressFn) *Reader {
partitionSize := totalSize / precision
if totalSize%precision != 0 {
partitionSize++
}
return &Reader{
rs: rs,
totalSize: totalSize,
partitionSize: partitionSize,
partitions: make([]bool, precision),
progressFn: fn,
}
}
func (prs *Reader) Read(p []byte) (n int, err error) {
n, err = prs.rs.Read(p)
newPos := prs.pos + int64(n)
lastPartitionsReadCount := prs.partitionsReadCount
partStart := prs.pos / prs.partitionSize
partEnd := newPos / prs.partitionSize
for i := partStart; i < partEnd; i++ {
if prs.partitions[i] {
continue
}
prs.partitions[i] = true
prs.partitionsReadCount++
}
if lastPartitionsReadCount != prs.partitionsReadCount {
readBytes := prs.partitionSize * prs.partitionsReadCount
if readBytes > prs.totalSize {
readBytes = prs.totalSize
}
prs.progressFn(readBytes, prs.totalSize)
}
prs.pos = newPos
return n, err
}
func (prs *Reader) Seek(offset int64, whence int) (int64, error) {
pos, err := prs.rs.Seek(offset, whence)
prs.pos = pos
return pos, err
}