mirror of
https://github.com/yusing/godoxy.git
synced 2025-06-09 04:52:35 +02:00
small memory usage optimization
This commit is contained in:
parent
2fe8531e51
commit
dc88a037eb
1 changed files with 42 additions and 38 deletions
|
@ -132,6 +132,17 @@ const (
|
|||
// This is a copy of io.Copy with context and HTTP flusher handling
|
||||
// Author: yusing <yusing@6uo.me>.
|
||||
func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
|
||||
// If the reader has a WriteTo method, use it to do the copy.
|
||||
// Avoids an allocation and a copy.
|
||||
if wt, ok := src.Reader.(io.WriterTo); ok {
|
||||
_, err = wt.WriteTo(dst)
|
||||
return
|
||||
}
|
||||
// Similarly, if the writer has a ReadFrom method, use it to do the copy.
|
||||
if rf, ok := dst.Writer.(io.ReaderFrom); ok {
|
||||
_, err = rf.ReadFrom(src)
|
||||
return
|
||||
}
|
||||
var buf []byte
|
||||
if l, ok := src.Reader.(*io.LimitedReader); ok {
|
||||
size := copyBufSize
|
||||
|
@ -142,7 +153,7 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
|
|||
size = int(l.N)
|
||||
}
|
||||
}
|
||||
buf = make([]byte, size)
|
||||
buf = make([]byte, 0, size)
|
||||
} else {
|
||||
buf = copyBufPool.Get().([]byte)
|
||||
defer copyBufPool.Put(buf)
|
||||
|
@ -179,13 +190,7 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
|
|||
flusher := getHttpFlusher(dst.Writer)
|
||||
canFlush := flusher != nil
|
||||
for {
|
||||
select {
|
||||
case <-src.ctx.Done():
|
||||
return src.ctx.Err()
|
||||
case <-dst.ctx.Done():
|
||||
return dst.ctx.Err()
|
||||
default:
|
||||
nr, er := src.Reader.Read(buf)
|
||||
nr, er := src.Reader.Read(buf[:copyBufSize])
|
||||
if nr > 0 {
|
||||
nw, ew := dst.Writer.Write(buf[0:nr])
|
||||
if nw < 0 || nr < nw {
|
||||
|
@ -222,7 +227,6 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
|
|||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func CopyCloseWithContext(ctx context.Context, dst io.Writer, src io.Reader) (err error) {
|
||||
return CopyClose(NewContextWriter(ctx, dst), NewContextReader(ctx, src))
|
||||
|
|
Loading…
Add table
Reference in a new issue