mirror of
https://github.com/yusing/godoxy.git
synced 2026-03-23 17:41:05 +01:00
fix: improved sync.Pool handling
This commit is contained in:
@@ -9,6 +9,7 @@ import (
|
||||
"syscall"
|
||||
|
||||
"github.com/yusing/go-proxy/internal/gperr"
|
||||
"github.com/yusing/go-proxy/internal/utils/synk"
|
||||
)
|
||||
|
||||
// TODO: move to "utils/io".
|
||||
@@ -117,24 +118,21 @@ func getHttpFlusher(dst io.Writer) httpFlusher {
|
||||
return nil
|
||||
}
|
||||
|
||||
const (
|
||||
copyBufSize = 32 * 1024
|
||||
)
|
||||
const copyBufSize = 32 * 1024
|
||||
|
||||
var copyBufPool = sync.Pool{
|
||||
New: func() any {
|
||||
return make([]byte, copyBufSize)
|
||||
},
|
||||
}
|
||||
var copyBufPool = synk.NewBytesPool(copyBufSize, synk.DefaultMaxBytes)
|
||||
|
||||
// Copyright 2009 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// 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) {
|
||||
var buf []byte
|
||||
buf := copyBufPool.Get()
|
||||
defer copyBufPool.Put(buf)
|
||||
|
||||
var size int
|
||||
if l, ok := src.Reader.(*io.LimitedReader); ok {
|
||||
size := copyBufSize
|
||||
size = copyBufSize
|
||||
if int64(size) > l.N {
|
||||
if l.N < 1 {
|
||||
size = 1
|
||||
@@ -142,10 +140,8 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
|
||||
size = int(l.N)
|
||||
}
|
||||
}
|
||||
buf = make([]byte, 0, size)
|
||||
} else {
|
||||
buf = copyBufPool.Get().([]byte)
|
||||
defer copyBufPool.Put(buf[:0])
|
||||
size = cap(buf)
|
||||
}
|
||||
// close both as soon as one of them is done
|
||||
wCloser, wCanClose := dst.Writer.(io.Closer)
|
||||
@@ -179,7 +175,7 @@ func CopyClose(dst *ContextWriter, src *ContextReader) (err error) {
|
||||
flusher := getHttpFlusher(dst.Writer)
|
||||
canFlush := flusher != nil
|
||||
for {
|
||||
nr, er := src.Reader.Read(buf[:copyBufSize])
|
||||
nr, er := src.Reader.Read(buf[:size])
|
||||
if nr > 0 {
|
||||
nw, ew := dst.Writer.Write(buf[0:nr])
|
||||
if nw < 0 || nr < nw {
|
||||
|
||||
42
internal/utils/synk/pool.go
Normal file
42
internal/utils/synk/pool.go
Normal file
@@ -0,0 +1,42 @@
|
||||
package synk
|
||||
|
||||
import "sync"
|
||||
|
||||
type (
|
||||
// Pool is a wrapper of sync.Pool that limits the size of the object.
|
||||
Pool[T any] struct {
|
||||
pool sync.Pool
|
||||
maxSize int
|
||||
}
|
||||
BytesPool = Pool[byte]
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultInitBytes = 1024
|
||||
DefaultMaxBytes = 1024 * 1024
|
||||
)
|
||||
|
||||
func NewPool[T any](initSize int, maxSize int) *Pool[T] {
|
||||
return &Pool[T]{
|
||||
pool: sync.Pool{
|
||||
New: func() any {
|
||||
return make([]T, 0, initSize)
|
||||
},
|
||||
},
|
||||
maxSize: maxSize,
|
||||
}
|
||||
}
|
||||
|
||||
func NewBytesPool(initSize int, maxSize int) *BytesPool {
|
||||
return NewPool[byte](initSize, maxSize)
|
||||
}
|
||||
|
||||
func (p *Pool[T]) Get() []T {
|
||||
return p.pool.Get().([]T)
|
||||
}
|
||||
|
||||
func (p *Pool[T]) Put(b []T) {
|
||||
if cap(b) <= p.maxSize {
|
||||
p.pool.Put(b[:0])
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user