-
Notifications
You must be signed in to change notification settings - Fork 2
/
io_utils.go
47 lines (41 loc) · 857 Bytes
/
io_utils.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
package utils
import (
"context"
"io"
"time"
"github.com/pkg/errors"
)
type readerFunc func(p []byte) (n int, err error)
func (f readerFunc) Read(p []byte) (n int, err error) {
return f(p)
}
func CopyClose(ctx context.Context, dst io.Writer, src io.ReadCloser) (int64, error) {
defer src.Close()
return io.Copy(dst, readerFunc(func(p []byte) (int, error) {
select {
case <-ctx.Done():
return 0, ctx.Err()
default:
return src.Read(p)
}
}))
}
func WaitForNBytes(r io.ReadSeeker, N int, interval time.Duration, retryTime int) (err error) {
var (
t []byte
i int
)
for i = 0; i < retryTime; i++ {
t = make([]byte, N)
_, err = io.ReadFull(r, t)
_, _ = r.Seek(0, 0)
if err != io.EOF {
break
}
<-time.After(interval)
}
if i >= retryTime {
err = errors.Errorf("Reader is less than %d bytes", N)
}
return err
}