Refactor BufferPool #62

Merged
fyrchik merged 1 commit from achuprov/frostfs-api-go:buffer into master 2024-01-12 15:40:41 +00:00
4 changed files with 73 additions and 58 deletions

54
util/pool/buffer.go Normal file
View file

@ -0,0 +1,54 @@
package pool
import (
"sync"
)
// Buffer contains a byte slice.
type Buffer struct {
Data []byte
}
// BufferPool manages a pool of Buffers.
type BufferPool struct {
poolSliceSize uint32 // Size for the buffer slices in the pool.
buffersPool *sync.Pool
}
// NewBufferPool creates a BufferPool with a specified size.
func NewBufferPool(poolSliceSize uint32) BufferPool {
pool := sync.Pool{
New: func() any {
return new(Buffer)
},
}
return BufferPool{poolSliceSize: poolSliceSize, buffersPool: &pool}
}
// Get retrieves a Buffer from the pool or creates a new one if necessary.
// It ensures the buffer's capacity is at least the specified size.
func (pool BufferPool) Get(size uint32) *Buffer {
result := pool.buffersPool.Get().(*Buffer)
if cap(result.Data) < int(size) {
result.Data = make([]byte, size)
} else {
result.Data = result.Data[:size]
}
return result
}
// Put returns a Buffer to the pool if its capacity does not exceed poolSliceSize.
func (pool BufferPool) Put(buf *Buffer) {
if cap(buf.Data) > int(pool.poolSliceSize) {
return
}
buf.Data = buf.Data[:0]
pool.buffersPool.Put(buf)
}
// PoolSliceSize returns the size for buffer slices in the pool.
func (pool BufferPool) PoolSliceSize() uint32 {
return uint32(pool.poolSliceSize)
}

View file

@ -1,40 +0,0 @@
package signature
import "sync"
const poolSliceMaxSize = 128 * 1024
type buffer struct {
data []byte
}
var buffersPool = sync.Pool{
New: func() any {
return new(buffer)
},
}
func tryGetNewBufferFromPool(size int) (*buffer, bool) {
if size > poolSliceMaxSize {
return &buffer{}, false
}
return newBufferFromPool(size), true
}
func newBufferFromPool(size int) *buffer {
result := buffersPool.Get().(*buffer)
if cap(result.data) < size {
result.data = make([]byte, size)
} else {
result.data = result.data[:size]
}
return result
}
func returnBufferToPool(buf *buffer) {
if cap(buf.data) > poolSliceMaxSize {
return
}
buf.data = buf.data[:0]
buffersPool.Put(buf)
}

View file

@ -4,9 +4,14 @@ import (
"crypto/ecdsa"
"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/refs"
"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/pool"
crypto "git.frostfs.info/TrueCloudLab/frostfs-crypto"
)
const poolSliceMaxSize = 128 * 1024
var buffersPool = pool.NewBufferPool(poolSliceMaxSize)
type DataSource interface {
ReadSignedData([]byte) ([]byte, error)
SignedDataSize() int
@ -35,12 +40,10 @@ func SignDataWithHandler(key *ecdsa.PrivateKey, src DataSource, handler KeySigna
opts[i](cfg)
}
buffer, ok := tryGetNewBufferFromPool(src.SignedDataSize())
if ok {
defer returnBufferToPool(buffer)
}
buffer := buffersPool.Get(uint32(src.SignedDataSize()))
defer buffersPool.Put(buffer)
data, err := src.ReadSignedData(buffer.data)
data, err := src.ReadSignedData(buffer.Data)
if err != nil {
return err
}
@ -66,12 +69,10 @@ func VerifyDataWithSource(dataSrc DataSource, sigSrc KeySignatureSource, opts ..
opts[i](cfg)
}
buffer, ok := tryGetNewBufferFromPool(dataSrc.SignedDataSize())
if ok {
defer returnBufferToPool(buffer)
}
buffer := buffersPool.Get(uint32(dataSrc.SignedDataSize()))
defer buffersPool.Put(buffer)
data, err := dataSrc.ReadSignedData(buffer.data)
data, err := dataSrc.ReadSignedData(buffer.Data)
if err != nil {
return err
}

View file

@ -35,10 +35,10 @@ func verify(cfg *cfg, data []byte, sig *refs.Signature) error {
case refs.ECDSA_RFC6979_SHA256:
return crypto.VerifyRFC6979(pub, data, sig.GetSign())
case refs.ECDSA_RFC6979_SHA256_WALLET_CONNECT:
buffer := newBufferFromPool(base64.StdEncoding.EncodedLen(len(data)))
defer returnBufferToPool(buffer)
base64.StdEncoding.Encode(buffer.data, data)
if !walletconnect.Verify(pub, buffer.data, sig.GetSign()) {
buffer := buffersPool.Get(uint32(base64.StdEncoding.EncodedLen(len(data))))
defer buffersPool.Put(buffer)
base64.StdEncoding.Encode(buffer.Data, data)
if !walletconnect.Verify(pub, buffer.Data, sig.GetSign()) {
return crypto.ErrInvalidSignature
}
return nil
@ -54,10 +54,10 @@ func sign(cfg *cfg, key *ecdsa.PrivateKey, data []byte) ([]byte, error) {
case refs.ECDSA_RFC6979_SHA256:
return crypto.SignRFC6979(key, data)
case refs.ECDSA_RFC6979_SHA256_WALLET_CONNECT:
buffer := newBufferFromPool(base64.StdEncoding.EncodedLen(len(data)))
defer returnBufferToPool(buffer)
base64.StdEncoding.Encode(buffer.data, data)
return walletconnect.Sign(key, buffer.data)
buffer := buffersPool.Get(uint32(base64.StdEncoding.EncodedLen(len(data))))
defer buffersPool.Put(buffer)
base64.StdEncoding.Encode(buffer.Data, data)
return walletconnect.Sign(key, buffer.Data)
default:
panic(fmt.Sprintf("unsupported scheme %s", cfg.scheme))
}