Skip to main content

zenoh_sync/
object_pool.rs

1//
2// Copyright (c) 2023 ZettaScale Technology
3//
4// This program and the accompanying materials are made available under the
5// terms of the Eclipse Public License 2.0 which is available at
6// http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
7// which is available at https://www.apache.org/licenses/LICENSE-2.0.
8//
9// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
10//
11// Contributors:
12//   ZettaScale Zenoh Team, <zenoh@zettascale.tech>
13//
14use std::{
15    any::Any,
16    fmt,
17    ops::{Deref, DerefMut, Drop},
18    sync::{Arc, Weak},
19};
20
21use zenoh_buffers::ZSliceBuffer;
22
23use super::LifoQueue;
24
25/// Provides a pool of pre-allocated objects that are automaticlaly reinserted into
26/// the pool when dropped.
27#[derive(Debug)]
28pub struct RecyclingObjectPool<T, F>
29where
30    F: Fn() -> T,
31{
32    inner: Arc<LifoQueue<T>>,
33    f: F,
34}
35
36impl<T, F: Fn() -> T + Clone> Clone for RecyclingObjectPool<T, F> {
37    fn clone(&self) -> Self {
38        Self {
39            inner: self.inner.clone(),
40            f: self.f.clone(),
41        }
42    }
43}
44
45impl<T, F: Fn() -> T> RecyclingObjectPool<T, F> {
46    pub fn new(num: usize, f: F) -> RecyclingObjectPool<T, F> {
47        let inner: Arc<LifoQueue<T>> = Arc::new(LifoQueue::new(num));
48        for _ in 0..num {
49            let obj = (f)();
50            inner.try_push(obj);
51        }
52        RecyclingObjectPool { inner, f }
53    }
54
55    pub fn alloc(&self) -> RecyclingObject<T> {
56        RecyclingObject::new((self.f)(), Weak::new())
57    }
58
59    pub fn try_take(&self) -> Option<RecyclingObject<T>> {
60        self.inner
61            .try_pull()
62            .map(|obj| RecyclingObject::new(obj, Arc::downgrade(&self.inner)))
63    }
64
65    pub fn take(&self) -> RecyclingObject<T> {
66        let obj = self.inner.pull();
67        RecyclingObject::new(obj, Arc::downgrade(&self.inner))
68    }
69}
70
71#[derive(Clone)]
72pub struct RecyclingObject<T> {
73    pool: Weak<LifoQueue<T>>,
74    object: Option<T>,
75}
76
77impl<T> RecyclingObject<T> {
78    pub fn new(obj: T, pool: Weak<LifoQueue<T>>) -> RecyclingObject<T> {
79        RecyclingObject {
80            pool,
81            object: Some(obj),
82        }
83    }
84
85    pub fn recycle(mut self) {
86        if let Some(pool) = self.pool.upgrade() {
87            if let Some(obj) = self.object.take() {
88                pool.push(obj);
89            }
90        }
91    }
92}
93
94impl<T: PartialEq> Eq for RecyclingObject<T> {}
95
96impl<T: PartialEq> PartialEq for RecyclingObject<T> {
97    fn eq(&self, other: &Self) -> bool {
98        self.object == other.object
99    }
100}
101
102impl<T> Deref for RecyclingObject<T> {
103    type Target = T;
104    #[inline]
105    fn deref(&self) -> &Self::Target {
106        self.object.as_ref().unwrap()
107    }
108}
109
110impl<T> DerefMut for RecyclingObject<T> {
111    #[inline]
112    fn deref_mut(&mut self) -> &mut Self::Target {
113        self.object.as_mut().unwrap()
114    }
115}
116
117impl<T> From<T> for RecyclingObject<T> {
118    fn from(obj: T) -> RecyclingObject<T> {
119        RecyclingObject::new(obj, Weak::new())
120    }
121}
122
123impl<T> Drop for RecyclingObject<T> {
124    fn drop(&mut self) {
125        if let Some(pool) = self.pool.upgrade() {
126            if let Some(obj) = self.object.take() {
127                pool.push(obj);
128            }
129        }
130    }
131}
132
133impl<T: fmt::Debug> fmt::Debug for RecyclingObject<T> {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        f.debug_struct("").field("inner", &self.object).finish()
136    }
137}
138
139// Buffer impl
140impl AsRef<[u8]> for RecyclingObject<Box<[u8]>> {
141    fn as_ref(&self) -> &[u8] {
142        self.deref()
143    }
144}
145
146impl AsMut<[u8]> for RecyclingObject<Box<[u8]>> {
147    fn as_mut(&mut self) -> &mut [u8] {
148        self.deref_mut()
149    }
150}
151
152impl ZSliceBuffer for RecyclingObject<Box<[u8]>> {
153    fn as_slice(&self) -> &[u8] {
154        self.as_ref()
155    }
156
157    fn as_any(&self) -> &dyn Any {
158        self
159    }
160
161    fn as_any_mut(&mut self) -> &mut dyn Any {
162        self
163    }
164}