mirror of
https://github.com/versity/versitygw.git
synced 2026-08-16 20:26:16 +00:00
99 lines
2.6 KiB
Go
99 lines
2.6 KiB
Go
// Copyright 2026 Versity Software
|
|
// This file is licensed under the Apache License, Version 2.0
|
|
// (the "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
// Package bufferpool provides a pool of pre-registered RDMA buffers
|
|
// to avoid per-request allocation and registration overhead.
|
|
package bufferpool
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
|
|
"github.com/versity/versitygw/rdma"
|
|
)
|
|
|
|
// Pool manages a fixed set of RDMA-registered host memory buffers.
|
|
type Pool struct {
|
|
server *rdma.Server
|
|
bufSize int
|
|
ch chan *rdma.Buffer
|
|
all []*rdma.Buffer
|
|
}
|
|
|
|
// NewPool pre-allocates and registers count buffers of bufSize bytes each.
|
|
func NewPool(server *rdma.Server, bufSize, count int) (*Pool, error) {
|
|
if count <= 0 {
|
|
return nil, errors.New("bufferpool: count must be > 0")
|
|
}
|
|
if bufSize <= 0 {
|
|
return nil, errors.New("bufferpool: bufSize must be > 0")
|
|
}
|
|
|
|
p := &Pool{
|
|
server: server,
|
|
bufSize: bufSize,
|
|
ch: make(chan *rdma.Buffer, count),
|
|
all: make([]*rdma.Buffer, 0, count),
|
|
}
|
|
|
|
for i := range count {
|
|
ptr, err := server.AllocHostBuffer(bufSize)
|
|
if err != nil {
|
|
p.Close()
|
|
return nil, fmt.Errorf("bufferpool: alloc buffer %d: %w", i, err)
|
|
}
|
|
|
|
buf, err := server.RegisterBuffer(ptr, bufSize)
|
|
if err != nil {
|
|
server.FreeHostBuffer(ptr)
|
|
p.Close()
|
|
return nil, fmt.Errorf("bufferpool: register buffer %d: %w", i, err)
|
|
}
|
|
|
|
p.all = append(p.all, buf)
|
|
p.ch <- buf
|
|
}
|
|
|
|
return p, nil
|
|
}
|
|
|
|
// Acquire blocks until a buffer is available or ctx is cancelled.
|
|
// Returns the buffer and its backing byte slice.
|
|
func (p *Pool) Acquire(ctx context.Context) (*rdma.Buffer, []byte, error) {
|
|
select {
|
|
case buf := <-p.ch:
|
|
return buf, buf.Slice(), nil
|
|
case <-ctx.Done():
|
|
return nil, nil, ctx.Err()
|
|
}
|
|
}
|
|
|
|
// Release returns a buffer to the pool.
|
|
func (p *Pool) Release(buf *rdma.Buffer) {
|
|
p.ch <- buf
|
|
}
|
|
|
|
// Close deregisters and frees all buffers. The pool must not be used afterward.
|
|
func (p *Pool) Close() {
|
|
for _, buf := range p.all {
|
|
// Capture the host pointer before DeregisterBuffer clears it, so the
|
|
// underlying allocation from AllocHostBuffer can still be freed.
|
|
ptr := buf.HostPtr()
|
|
p.server.DeregisterBuffer(buf)
|
|
p.server.FreeHostBuffer(ptr)
|
|
}
|
|
p.all = nil
|
|
}
|