Redmage/api/reddit/image_download_reader.go

114 lines
2.2 KiB
Go

package reddit
import (
"errors"
"io"
"net/http"
"sync"
"sync/atomic"
"time"
"github.com/alecthomas/units"
)
var ErrIdleTimeoutReached = errors.New("download idle timeout reached")
type ImageDownloadReader struct {
OnProgress func(downloaded int64, contentLength int64, err error)
OnClose func(downloaded int64, contentLength int64, closeErr error)
IdleTimeout time.Duration
IdleSpeedThreshold units.MetricBytes
errCancel error
cancelDebounce *time.Timer
reader io.ReadCloser
contentLength int64
downloaded atomic.Int64
deltastart time.Time
deltavalue atomic.Int64
end time.Time
exit chan struct{}
mu sync.Mutex
}
func (idr *ImageDownloadReader) WrapHTTPResponse(resp *http.Response) *http.Response {
idr.reader = resp.Body
idr.contentLength = resp.ContentLength
idr.exit = make(chan struct{}, 1)
resp.Body = idr
go func() {
ticker := time.NewTicker(time.Second)
defer ticker.Stop()
for {
select {
case <-idr.exit:
return
case <-ticker.C:
idr.checkSpeed()
}
}
}()
return resp
}
func (idr *ImageDownloadReader) checkSpeed() {
now := time.Now()
if idr.deltastart.IsZero() {
idr.deltastart = now
}
if idr.cancelDebounce == nil {
idr.cancelDebounce = time.AfterFunc(idr.IdleTimeout, func() {
idr.mu.Lock()
defer idr.mu.Unlock()
idr.errCancel = ErrIdleTimeoutReached
})
}
if now.Sub(idr.deltastart) < time.Second {
return
}
idr.deltastart = now
delta := idr.deltavalue.Load()
if delta >= int64(idr.IdleSpeedThreshold) {
idr.deltavalue.Store(0)
idr.cancelDebounce.Stop()
idr.cancelDebounce = nil
}
}
func (idr *ImageDownloadReader) Read(p []byte) (n int, err error) {
n, err = idr.reader.Read(p)
idr.deltavalue.Add(int64(n))
newd := idr.downloaded.Add(int64(n))
if idr.OnProgress != nil {
idr.OnProgress(newd, idr.contentLength, err)
}
idr.mu.Lock()
if idr.errCancel != nil {
idr.mu.Unlock()
idr.OnProgress(newd, idr.contentLength, idr.errCancel)
return n, idr.errCancel
}
idr.mu.Unlock()
return n, err
}
func (idr *ImageDownloadReader) Close() error {
idr.exit <- struct{}{}
err := idr.reader.Close()
if idr.OnClose != nil {
idr.OnClose(idr.downloaded.Load(), idr.contentLength, err)
}
return err
}