SaveAny-Bot/common/utils/ioutil/progress_reader.go

package ioutil

import (
	"io"
	"sync/atomic"
)

var _ io.ReadSeeker = (*ProgressReadSeeker)(nil)

// ProgressReadSeeker wraps an io.ReadSeeker and tracks read progress
type ProgressReadSeeker struct {
	reader     io.ReadSeeker
	total      atomic.Int64
	read       atomic.Int64
	onProgress func(read int64, total int64)
}

// Seek implements io.ReadSeeker.
func (pr *ProgressReadSeeker) Seek(offset int64, whence int) (int64, error) {
	position, err := pr.reader.Seek(offset, whence)
	if err == nil {
		pr.read.Store(position)
	}
	return position, err
}

// NewProgressReader creates a new ProgressReader
func NewProgressReader(rs io.ReadSeeker, total int64, onProgress func(read int64, total int64)) *ProgressReadSeeker {
	prs := &ProgressReadSeeker{
		reader:     rs,
		total:      atomic.Int64{},
		read:       atomic.Int64{},
		onProgress: onProgress,
	}
	prs.total.Store(total)
	return prs
}

// Read implements io.Reader
func (pr *ProgressReadSeeker) Read(p []byte) (int, error) {
	n, err := pr.reader.Read(p)
	if n > 0 {
		pr.read.Add(int64(n))
		read := pr.read.Load()

		if pr.onProgress != nil {
			pr.onProgress(read, pr.total.Load())
		}
	}
	return n, err
}

// Progress returns the current progress as a float64 between 0 and 1
func (pr *ProgressReadSeeker) Progress() float64 {
	if pr.total.Load() <= 0 {
		return 0
	}
	return float64(pr.read.Load()) / float64(pr.total.Load())
}

// BytesRead returns the current tracked reader position.
func (pr *ProgressReadSeeker) BytesRead() int64 {
	return pr.read.Load()
}

// Total returns the total number of bytes
func (pr *ProgressReadSeeker) Total() int64 {
	return pr.total.Load()
}