fix io.NopCloser issue
io.NopCloser in go 1.18 and below doesn't check for io.WriterTo
Vladimir Bauer
3 years ago
| 5 | 5 | ) |
| 6 | 6 | |
| 7 | 7 | type proxyReader struct { |
| 8 | io.ReadCloser | |
| 8 | io.Reader | |
| 9 | 9 | bar *Bar |
| 10 | 10 | } |
| 11 | 11 | |
| 12 | 12 | func (x proxyReader) Read(p []byte) (int, error) { |
| 13 | n, err := x.ReadCloser.Read(p) | |
| 13 | n, err := x.Reader.Read(p) | |
| 14 | 14 | x.bar.IncrBy(n) |
| 15 | 15 | return n, err |
| 16 | 16 | } |
| 20 | 20 | } |
| 21 | 21 | |
| 22 | 22 | func (x proxyWriterTo) WriteTo(w io.Writer) (int64, error) { |
| 23 | n, err := x.ReadCloser.(io.WriterTo).WriteTo(w) | |
| 23 | n, err := x.Reader.(io.WriterTo).WriteTo(w) | |
| 24 | 24 | x.bar.IncrInt64(n) |
| 25 | 25 | return n, err |
| 26 | 26 | } |
| 44 | 44 | |
| 45 | 45 | func (x ewmaProxyWriterTo) WriteTo(w io.Writer) (int64, error) { |
| 46 | 46 | start := time.Now() |
| 47 | n, err := x.ReadCloser.(io.WriterTo).WriteTo(w) | |
| 47 | n, err := x.Reader.(io.WriterTo).WriteTo(w) | |
| 48 | 48 | if n > 0 { |
| 49 | 49 | x.bar.DecoratorEwmaUpdate(time.Since(start)) |
| 50 | 50 | } |
| 52 | 52 | } |
| 53 | 53 | |
| 54 | 54 | func (b *Bar) newProxyReader(r io.Reader, hasEwma bool) io.ReadCloser { |
| 55 | pr := proxyReader{toReadCloser(r), b} | |
| 55 | return toReadCloser(toReader(r, b, hasEwma)) | |
| 56 | } | |
| 57 | ||
| 58 | func toReader(r io.Reader, b *Bar, hasEwma bool) io.Reader { | |
| 59 | pr := proxyReader{r, b} | |
| 56 | 60 | if hasEwma { |
| 57 | 61 | epr := ewmaProxyReader{pr} |
| 58 | 62 | if _, ok := r.(io.WriterTo); ok { |
| 70 | 74 | if rc, ok := r.(io.ReadCloser); ok { |
| 71 | 75 | return rc |
| 72 | 76 | } |
| 73 | return io.NopCloser(r) | |
| 77 | return toNopCloser(r) | |
| 74 | 78 | } |
| 79 | ||
| 80 | func toNopCloser(r io.Reader) io.ReadCloser { | |
| 81 | if _, ok := r.(io.WriterTo); ok { | |
| 82 | return nopCloserWriterTo{r} | |
| 83 | } | |
| 84 | return nopCloser{r} | |
| 85 | } | |
| 86 | ||
| 87 | type nopCloser struct { | |
| 88 | io.Reader | |
| 89 | } | |
| 90 | ||
| 91 | func (nopCloser) Close() error { return nil } | |
| 92 | ||
| 93 | type nopCloserWriterTo struct { | |
| 94 | io.Reader | |
| 95 | } | |
| 96 | ||
| 97 | func (nopCloserWriterTo) Close() error { return nil } | |
| 98 | ||
| 99 | func (c nopCloserWriterTo) WriteTo(w io.Writer) (n int64, err error) { | |
| 100 | return c.Reader.(io.WriterTo).WriteTo(w) | |
| 101 | } | |