mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-11 07:36:53 +00:00
houseme
89 lines
2.6 KiB
Rust
89 lines
2.6 KiB
Rust
// Copyright 2024 RustFS Team
|
|
//
|
|
// 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.
|
|
|
|
#[derive(Debug, Default)]
|
|
pub(crate) struct ShardBufferPool {
|
|
buffers: Vec<Option<Vec<u8>>>,
|
|
}
|
|
|
|
impl ShardBufferPool {
|
|
#[inline]
|
|
pub(crate) fn new(slot_count: usize) -> Self {
|
|
let mut buffers = Vec::with_capacity(slot_count);
|
|
buffers.resize_with(slot_count, || None);
|
|
Self { buffers }
|
|
}
|
|
|
|
#[inline]
|
|
pub(crate) fn ensure_slots(&mut self, slot_count: usize) {
|
|
if self.buffers.len() < slot_count {
|
|
self.buffers.resize_with(slot_count, || None);
|
|
}
|
|
}
|
|
|
|
// Returned bytes may contain stale scratch data and must be overwritten before use.
|
|
#[inline]
|
|
pub(crate) fn take(&mut self, index: usize, len: usize) -> Vec<u8> {
|
|
self.ensure_slots(index + 1);
|
|
let mut buf = self.buffers[index].take().unwrap_or_else(|| Vec::with_capacity(len));
|
|
if buf.capacity() < len {
|
|
buf.reserve_exact(len - buf.capacity());
|
|
}
|
|
buf.resize(len, 0);
|
|
buf
|
|
}
|
|
|
|
#[inline]
|
|
pub(crate) fn put(&mut self, index: usize, buf: Vec<u8>) {
|
|
self.ensure_slots(index + 1);
|
|
self.buffers[index] = Some(buf);
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn stored_capacity(&self, index: usize) -> Option<usize> {
|
|
self.buffers.get(index).and_then(|buf| buf.as_ref().map(Vec::capacity))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn shard_buffer_pool_reuses_slot_without_clearing() {
|
|
let mut pool = ShardBufferPool::new(2);
|
|
let mut buf = pool.take(1, 16);
|
|
assert_eq!(buf.len(), 16);
|
|
buf[0] = 42;
|
|
let capacity = buf.capacity();
|
|
|
|
pool.put(1, buf);
|
|
assert_eq!(pool.stored_capacity(1), Some(capacity));
|
|
|
|
let reused = pool.take(1, 8);
|
|
assert_eq!(reused.len(), 8);
|
|
assert!(reused.capacity() >= capacity);
|
|
assert_eq!(reused[0], 42);
|
|
}
|
|
|
|
#[test]
|
|
fn shard_buffer_pool_grows_for_sparse_slot_indexes() {
|
|
let mut pool = ShardBufferPool::new(0);
|
|
let buf = pool.take(3, 4);
|
|
|
|
assert_eq!(buf.len(), 4);
|
|
assert_eq!(pool.buffers.len(), 4);
|
|
}
|
|
}
|