56 lines
1.1 KiB
Go
56 lines
1.1 KiB
Go
package fetch
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
bufferSize = 1 << 20
|
|
stallTimeout = 2 * time.Minute
|
|
)
|
|
|
|
func copyBody(destination io.Writer, body io.ReadCloser, remaining int64, record func(int64) error) error {
|
|
var stalled atomic.Bool
|
|
stallTimer := time.AfterFunc(stallTimeout, func() {
|
|
stalled.Store(true)
|
|
body.Close()
|
|
})
|
|
defer stallTimer.Stop()
|
|
|
|
buffer := make([]byte, bufferSize)
|
|
for remaining != 0 {
|
|
readSize := len(buffer)
|
|
if remaining > 0 && remaining < int64(readSize) {
|
|
readSize = int(remaining)
|
|
}
|
|
|
|
count, readErr := body.Read(buffer[:readSize])
|
|
if count > 0 {
|
|
stallTimer.Reset(stallTimeout)
|
|
if _, err := destination.Write(buffer[:count]); err != nil {
|
|
return err
|
|
}
|
|
if err := record(int64(count)); err != nil {
|
|
return err
|
|
}
|
|
if remaining > 0 {
|
|
remaining -= int64(count)
|
|
}
|
|
}
|
|
if errors.Is(readErr, io.EOF) {
|
|
return nil
|
|
}
|
|
if readErr != nil {
|
|
if stalled.Load() {
|
|
return ErrStalled
|
|
}
|
|
return fmt.Errorf("%w: %w", errConnectionFailed, readErr)
|
|
}
|
|
}
|
|
return nil
|
|
}
|