Skip to content

Commit e407327

Browse files
awolverpawolverp
committed
Support pickle for TTLCache/VTTLCache
1 parent 008577a commit e407327

15 files changed

Lines changed: 244 additions & 22 deletions

File tree

‎src/internal/onceinit.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,11 @@ impl<T> OnceInit<T> {
6363
.into()
6464
}
6565

66+
#[inline]
67+
pub fn is_initialized(&self) -> bool {
68+
self.0.state.load(atomic::Ordering::Acquire) == INIT
69+
}
70+
6671
/// Initializes the container with `val`, transitioning state from `UNINIT` to `INIT`.
6772
///
6873
/// Intended to be called from the PyO3 `__init__` handler once the Python-side

‎src/internal/utils.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,13 @@ impl From<chrono::NaiveDateTime> for ExpiresAt {
296296
}
297297
}
298298

299+
impl From<std::time::SystemTime> for ExpiresAt {
300+
#[inline]
301+
fn from(value: std::time::SystemTime) -> Self {
302+
Self::Instant(value.into())
303+
}
304+
}
305+
299306
impl From<ExpiresAt> for std::time::SystemTime {
300307
#[inline]
301308
fn from(value: ExpiresAt) -> Self {

‎src/policies/common.rs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -114,15 +114,12 @@ impl Shared {
114114
/// Creates a new [`Shared`].
115115
#[inline]
116116
pub fn new(maxsize: usize, getsizeof: Option<alias::PyObject>) -> Self {
117-
unsafe { Self::with_ttl(maxsize, getsizeof, None) }
117+
Self::with_ttl(maxsize, getsizeof, None)
118118
}
119119

120120
/// Creates a new [`Shared`] with configured TTL.
121-
///
122-
/// # Safety
123-
/// `ttl` should not be negative or zero.
124121
#[inline]
125-
pub unsafe fn with_ttl(
122+
pub fn with_ttl(
126123
maxsize: usize,
127124
getsizeof: Option<alias::PyObject>,
128125
ttl: Option<std::time::Duration>,

‎src/policies/traits.rs‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -144,17 +144,13 @@ pub trait PolicyExt: Sized {
144144
fn build_pickle(
145145
&self,
146146
tuple: &mut pickle::TupleBuilder<'_, pickle::PickleBuilder>,
147-
) -> pyo3::PyResult<()> {
148-
todo!()
149-
}
147+
) -> pyo3::PyResult<()>;
150148

151149
/// Loads the builded pickle.
152150
fn from_pickle(
153151
maxsize: usize,
154152
getsizeof: Option<alias::PyObject>,
155153
global_ttl: Option<std::time::Duration>,
156154
builded: pyo3::Bound<'_, pyo3::types::PyTuple>,
157-
) -> pyo3::PyResult<(Self::Shared, Self)> {
158-
todo!()
159-
}
155+
) -> pyo3::PyResult<(Self::Shared, Self)>;
160156
}

‎src/policies/ttlpolicy.rs‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::collections::VecDeque;
22

33
use crate::hashbrown;
44
use crate::internal::alias;
5+
use crate::internal::pickle::Builder;
56
use crate::internal::utils;
67
use crate::policies::traits;
78
use crate::policies::traits::HandleExt;
@@ -524,4 +525,83 @@ impl PolicyExt for TTLPolicy {
524525
front_offset: self.front_offset,
525526
}
526527
}
528+
529+
fn build_pickle(
530+
&self,
531+
tuple: &mut crate::internal::pickle::TupleBuilder<
532+
'_,
533+
crate::internal::pickle::PickleBuilder,
534+
>,
535+
) -> pyo3::PyResult<()> {
536+
let mut list = tuple.begin_list()?;
537+
538+
for handle in self.entries.iter() {
539+
let mut tuple = list.begin_tuple(3)?;
540+
tuple.push(handle.key().as_ref())?;
541+
tuple.push(handle.value())?;
542+
tuple.push(
543+
handle
544+
.expires_at
545+
.duration_since(std::time::UNIX_EPOCH)
546+
.unwrap(),
547+
)?;
548+
tuple.end()?;
549+
}
550+
551+
list.end()
552+
}
553+
554+
fn from_pickle(
555+
maxsize: usize,
556+
getsizeof: Option<crate::internal::alias::PyObject>,
557+
global_ttl: Option<std::time::Duration>,
558+
builded: pyo3::Bound<'_, pyo3::types::PyTuple>,
559+
) -> pyo3::PyResult<(Self::Shared, Self)> {
560+
use pyo3::types::PyAnyMethods;
561+
use pyo3::types::PyListMethods;
562+
use pyo3::types::PyTupleMethods;
563+
564+
if global_ttl.is_none_or(|x| x.is_zero()) {
565+
return Err(new_py_error!(PyValueError, "global_ttl is zero"));
566+
}
567+
568+
let list = builded.get_item(0)?.cast_into::<pyo3::types::PyList>()?;
569+
let list_length = list.len();
570+
571+
if list_length > maxsize {
572+
return Err(new_py_error!(
573+
PyValueError,
574+
"list size is incompatible with maxsize"
575+
));
576+
}
577+
578+
let shared = Shared::with_ttl(maxsize, getsizeof, global_ttl);
579+
let mut slf = Self::new(list.len());
580+
581+
for bound in list.iter() {
582+
let (key, value, timestamp) =
583+
bound.extract::<(alias::PyObject, alias::PyObject, f64)>()?;
584+
585+
let handle = ExpiringHandle::new(
586+
bound.py(),
587+
shared.getsizeof(),
588+
(std::time::UNIX_EPOCH + std::time::Duration::from_secs_f64(timestamp)).into(),
589+
key,
590+
value,
591+
)?;
592+
593+
slf.currsize = slf.currsize.saturating_add(handle.size());
594+
595+
unsafe {
596+
slf.table.insert_no_grow(
597+
handle.key().hash(),
598+
// Adding `slf.front_offset` is unnecessary here
599+
slf.entries.len(),
600+
);
601+
}
602+
slf.entries.push_back(handle);
603+
}
604+
605+
Ok((shared, slf))
606+
}
527607
}

‎src/policies/vttlpolicy.rs‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crate::hashbrown;
22
use crate::internal::alias;
33
use crate::internal::lazyheap;
4+
use crate::internal::pickle::Builder;
45
use crate::internal::utils;
56
use crate::policies::traits;
67
use crate::policies::traits::HandleExt;
@@ -475,4 +476,79 @@ impl PolicyExt for VTTLPolicy {
475476
currsize: self.currsize,
476477
}
477478
}
479+
fn build_pickle(
480+
&self,
481+
tuple: &mut crate::internal::pickle::TupleBuilder<
482+
'_,
483+
crate::internal::pickle::PickleBuilder,
484+
>,
485+
) -> pyo3::PyResult<()> {
486+
let mut list = tuple.begin_list()?;
487+
488+
for cursor in unsafe { self.table.iter() } {
489+
let handle = unsafe { cursor.as_ref().element() };
490+
491+
let mut tuple = list.begin_tuple(3)?;
492+
tuple.push(handle.key.as_ref())?;
493+
tuple.push(handle.value())?;
494+
tuple.push(
495+
handle
496+
.expires_at
497+
.map(|x| x.duration_since(std::time::UNIX_EPOCH).unwrap()),
498+
)?;
499+
tuple.end()?;
500+
}
501+
502+
list.end()
503+
}
504+
505+
fn from_pickle(
506+
maxsize: usize,
507+
getsizeof: Option<crate::internal::alias::PyObject>,
508+
_global_ttl: Option<std::time::Duration>,
509+
builded: pyo3::Bound<'_, pyo3::types::PyTuple>,
510+
) -> pyo3::PyResult<(Self::Shared, Self)> {
511+
use pyo3::types::PyAnyMethods;
512+
use pyo3::types::PyListMethods;
513+
use pyo3::types::PyTupleMethods;
514+
515+
let list = builded.get_item(0)?.cast_into::<pyo3::types::PyList>()?;
516+
let list_length = list.len();
517+
518+
if list_length > maxsize {
519+
return Err(new_py_error!(
520+
PyValueError,
521+
"list size is incompatible with maxsize"
522+
));
523+
}
524+
525+
let shared = Shared::new(maxsize, getsizeof);
526+
let mut slf = Self::new(list.len());
527+
528+
for bound in list.iter() {
529+
let (key, value, timestamp) =
530+
bound.extract::<(alias::PyObject, alias::PyObject, Option<f64>)>()?;
531+
532+
let handle = ExpiringHandle::new(
533+
bound.py(),
534+
shared.getsizeof(),
535+
timestamp
536+
.map(|x| std::time::UNIX_EPOCH + std::time::Duration::from_secs_f64(x))
537+
.map(Into::into),
538+
key,
539+
value,
540+
)?;
541+
542+
slf.currsize = slf.currsize.saturating_add(handle.size());
543+
544+
let hash = handle.key().hash();
545+
let cursor = slf.heap.push(handle);
546+
unsafe {
547+
slf.table.insert_no_grow(hash, cursor);
548+
}
549+
}
550+
551+
slf.heap.sort_by(compare_fn!());
552+
Ok((shared, slf))
553+
}
478554
}

‎src/policies/wrapped.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -281,10 +281,20 @@ impl<P: PolicyExt> Wrapped<P> {
281281

282282
let maxsize: usize = tuple.get_item(0)?.extract()?;
283283
let getsizeof: Option<alias::PyObject> = tuple.get_item(1)?.extract()?;
284-
let global_ttl: Option<std::time::Duration> = tuple.get_item(2)?.extract()?;
284+
let global_ttl: Option<f64> = tuple.get_item(2)?.extract()?;
285+
286+
if global_ttl.is_some_and(|x| x < 0.0) {
287+
return Err(new_py_error!(PyValueError, "global_ttl is negative"));
288+
}
289+
285290
let builded = tuple.get_item(3)?.cast_into::<pyo3::types::PyTuple>()?;
286291

287-
let (shared, inner) = P::from_pickle(maxsize, getsizeof, global_ttl, builded)?;
292+
let (shared, inner) = P::from_pickle(
293+
maxsize,
294+
getsizeof,
295+
global_ttl.map(|x| std::time::Duration::from_secs_f64(x)),
296+
builded,
297+
)?;
288298

289299
Ok(Self {
290300
shared,

‎src/pyclasses/cache.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -589,6 +589,10 @@ impl PyCache {
589589
}
590590

591591
fn __traverse__(&self, visit: pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
592+
if self.0.is_initialized() {
593+
return Ok(());
594+
}
595+
592596
let inner = self.0.get();
593597
let policy = inner.policy();
594598

@@ -602,6 +606,10 @@ impl PyCache {
602606
}
603607

604608
fn __clear__(&self) {
609+
if self.0.is_initialized() {
610+
return;
611+
}
612+
605613
let inner = self.0.get();
606614
let mut policy = inner.policy();
607615
policy.clear(inner.shared());

‎src/pyclasses/fifocache.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -619,6 +619,10 @@ impl PyFIFOCache {
619619
}
620620

621621
fn __traverse__(&self, visit: pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
622+
if self.0.is_initialized() {
623+
return Ok(());
624+
}
625+
622626
let inner = self.0.get();
623627
let policy = inner.policy();
624628

@@ -630,6 +634,10 @@ impl PyFIFOCache {
630634
}
631635

632636
fn __clear__(&self) {
637+
if self.0.is_initialized() {
638+
return;
639+
}
640+
633641
let inner = self.0.get();
634642
let mut policy = inner.policy();
635643
policy.clear(inner.shared());

‎src/pyclasses/lfucache.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,10 @@ impl PyLFUCache {
683683
}
684684

685685
fn __traverse__(&self, visit: pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {
686+
if self.0.is_initialized() {
687+
return Ok(());
688+
}
689+
686690
let inner = self.0.get();
687691
let policy = inner.policy();
688692

@@ -696,6 +700,10 @@ impl PyLFUCache {
696700
}
697701

698702
fn __clear__(&self) {
703+
if self.0.is_initialized() {
704+
return;
705+
}
706+
699707
let inner = self.0.get();
700708
let mut policy = inner.policy();
701709
policy.clear(inner.shared());

0 commit comments

Comments
 (0)