72 lines
1.0 KiB
Go
72 lines
1.0 KiB
Go
package zw
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"io/fs"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type downloadReader struct {
|
|
recv func() ([]byte, error)
|
|
cancel context.CancelFunc
|
|
|
|
pending []byte
|
|
terminal error
|
|
|
|
cancelOnce sync.Once
|
|
closed atomic.Bool
|
|
}
|
|
|
|
func (r *downloadReader) Read(destination []byte) (int, error) {
|
|
if len(destination) == 0 {
|
|
return 0, nil
|
|
}
|
|
|
|
if r.closed.Load() {
|
|
return 0, fs.ErrClosed
|
|
}
|
|
|
|
if len(r.pending) > 0 {
|
|
n := copy(destination, r.pending)
|
|
r.pending = r.pending[n:]
|
|
return n, nil
|
|
}
|
|
|
|
if r.terminal != nil {
|
|
return 0, r.terminal
|
|
}
|
|
|
|
for {
|
|
chunk, err := r.recv()
|
|
if err != nil {
|
|
if r.closed.Load() {
|
|
return 0, fs.ErrClosed
|
|
}
|
|
|
|
if err == io.EOF {
|
|
r.terminal = io.EOF
|
|
} else {
|
|
r.terminal = convertRPCError(err)
|
|
}
|
|
r.cancelOnce.Do(r.cancel)
|
|
return 0, r.terminal
|
|
}
|
|
|
|
if len(chunk) == 0 {
|
|
continue
|
|
}
|
|
|
|
n := copy(destination, chunk)
|
|
r.pending = chunk[n:]
|
|
return n, nil
|
|
}
|
|
}
|
|
|
|
func (r *downloadReader) Close() error {
|
|
r.closed.Store(true)
|
|
r.cancelOnce.Do(r.cancel)
|
|
return nil
|
|
}
|