Use internal mutability for SharedMemoryLimiter, removing one `RefC…

…ell`. (#221)

src/lib.rs

@@ -168,7 +168,7 @@ cfg_if! {
             EndTag, Serialize, StartTag, Token, TokenCaptureFlags, Mutations
         };
 
-        pub use self::memory::MemoryLimiter;
+        pub use self::memory::SharedMemoryLimiter;
         pub use self::html::{LocalName, LocalNameHash, Tag, Namespace};
     } else {
         mod selectors_vm;

src/memory/arena.rs

@@ -10,7 +10,7 @@ pub struct Arena {
 
 impl Arena {
     pub fn new(limiter: SharedMemoryLimiter, preallocated_size: usize) -> Self {
-        limiter.borrow_mut().preallocate(preallocated_size);
+        limiter.preallocate(preallocated_size);
 
         Arena {
             limiter,
@@ -27,9 +27,9 @@ impl Arena {
 
             // NOTE: approximate usage, as `Vec::reserve_exact` doesn't
             // give guarantees about exact capacity value :).
-            self.limiter.borrow_mut().increase_usage(additional)?;
+            self.limiter.increase_usage(additional)?;
 
-            // NOTE: with wicely chosen preallocated size this branch should be
+            // NOTE: with wisely chosen preallocated size this branch should be
             // executed quite rarely. We can't afford to use double capacity
             // strategy used by default (see: https://github.com/rust-lang/rust/blob/bdfd698f37184da42254a03ed466ab1f90e6fb6c/src/liballoc/raw_vec.rs#L424)
             // as we'll run out of the space allowance quite quickly.
@@ -58,26 +58,25 @@ impl Arena {
 
 #[cfg(test)]
 mod tests {
-    use super::super::limiter::MemoryLimiter;
+    use super::super::limiter::SharedMemoryLimiter;
     use super::*;
-    use std::rc::Rc;
 
     #[test]
     fn append() {
-        let limiter = MemoryLimiter::new_shared(10);
-        let mut arena = Arena::new(Rc::clone(&limiter), 2);
+        let limiter = SharedMemoryLimiter::new(10);
+        let mut arena = Arena::new(limiter.clone(), 2);
 
         arena.append(&[1, 2]).unwrap();
         assert_eq!(arena.bytes(), &[1, 2]);
-        assert_eq!(limiter.borrow().current_usage(), 2);
+        assert_eq!(limiter.current_usage(), 2);
 
         arena.append(&[3, 4]).unwrap();
         assert_eq!(arena.bytes(), &[1, 2, 3, 4]);
-        assert_eq!(limiter.borrow().current_usage(), 4);
+        assert_eq!(limiter.current_usage(), 4);
 
         arena.append(&[5, 6, 7, 8, 9, 10]).unwrap();
         assert_eq!(arena.bytes(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
-        assert_eq!(limiter.borrow().current_usage(), 10);
+        assert_eq!(limiter.current_usage(), 10);
 
         let err = arena.append(&[11]).unwrap_err();
 
@@ -86,24 +85,24 @@ mod tests {
 
     #[test]
     fn init_with() {
-        let limiter = MemoryLimiter::new_shared(5);
-        let mut arena = Arena::new(Rc::clone(&limiter), 0);
+        let limiter = SharedMemoryLimiter::new(5);
+        let mut arena = Arena::new(limiter.clone(), 0);
 
         arena.init_with(&[1]).unwrap();
         assert_eq!(arena.bytes(), &[1]);
-        assert_eq!(limiter.borrow().current_usage(), 1);
+        assert_eq!(limiter.current_usage(), 1);
 
         arena.append(&[1, 2]).unwrap();
         assert_eq!(arena.bytes(), &[1, 1, 2]);
-        assert_eq!(limiter.borrow().current_usage(), 3);
+        assert_eq!(limiter.current_usage(), 3);
 
         arena.init_with(&[1, 2, 3]).unwrap();
         assert_eq!(arena.bytes(), &[1, 2, 3]);
-        assert_eq!(limiter.borrow().current_usage(), 3);
+        assert_eq!(limiter.current_usage(), 3);
 
         arena.init_with(&[]).unwrap();
         assert_eq!(arena.bytes(), &[]);
-        assert_eq!(limiter.borrow().current_usage(), 3);
+        assert_eq!(limiter.current_usage(), 3);
 
         let err = arena.init_with(&[1, 2, 3, 4, 5, 6, 7]).unwrap_err();
 
@@ -112,25 +111,25 @@ mod tests {
 
     #[test]
     fn shift() {
-        let limiter = MemoryLimiter::new_shared(10);
-        let mut arena = Arena::new(Rc::clone(&limiter), 0);
+        let limiter = SharedMemoryLimiter::new(10);
+        let mut arena = Arena::new(limiter.clone(), 0);
 
         arena.append(&[0, 1, 2, 3]).unwrap();
         arena.shift(2);
         assert_eq!(arena.bytes(), &[2, 3]);
-        assert_eq!(limiter.borrow().current_usage(), 4);
+        assert_eq!(limiter.current_usage(), 4);
 
         arena.append(&[0, 1]).unwrap();
         assert_eq!(arena.bytes(), &[2, 3, 0, 1]);
-        assert_eq!(limiter.borrow().current_usage(), 4);
+        assert_eq!(limiter.current_usage(), 4);
 
         arena.shift(3);
         assert_eq!(arena.bytes(), &[1]);
-        assert_eq!(limiter.borrow().current_usage(), 4);
+        assert_eq!(limiter.current_usage(), 4);
 
         arena.append(&[2, 3, 4, 5]).unwrap();
         arena.shift(1);
         assert_eq!(arena.bytes(), &[2, 3, 4, 5]);
-        assert_eq!(limiter.borrow().current_usage(), 5);
+        assert_eq!(limiter.current_usage(), 5);
     }
 }

src/memory/limited_vec.rs

@@ -21,7 +21,7 @@ impl<T> LimitedVec<T> {
     }
 
     pub fn push(&mut self, element: T) -> Result<(), MemoryLimitExceededError> {
-        self.limiter.borrow_mut().increase_usage(size_of::<T>())?;
+        self.limiter.increase_usage(size_of::<T>())?;
         self.vec.push(element);
         Ok(())
     }
@@ -64,9 +64,7 @@ impl<T> LimitedVec<T> {
             Unbounded => self.len(),
         };
 
-        self.limiter
-            .borrow_mut()
-            .decrease_usage(size_of::<T>() * (end - start));
+        self.limiter.decrease_usage(size_of::<T>() * (end - start));
 
         self.vec.drain(range)
     }
@@ -91,43 +89,40 @@ impl<T> Index<usize> for LimitedVec<T> {
 
 impl<T> Drop for LimitedVec<T> {
     fn drop(&mut self) {
-        self.limiter
-            .borrow_mut()
-            .decrease_usage(size_of::<T>() * self.vec.len());
+        self.limiter.decrease_usage(size_of::<T>() * self.vec.len());
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::super::MemoryLimiter;
+    use super::super::SharedMemoryLimiter;
     use super::*;
-    use std::rc::Rc;
 
     #[test]
     fn current_usage() {
         {
-            let limiter = MemoryLimiter::new_shared(10);
-            let mut vec_u8: LimitedVec<u8> = LimitedVec::new(Rc::clone(&limiter));
+            let limiter = SharedMemoryLimiter::new(10);
+            let mut vec_u8: LimitedVec<u8> = LimitedVec::new(limiter.clone());
 
             vec_u8.push(1).unwrap();
             vec_u8.push(2).unwrap();
-            assert_eq!(limiter.borrow().current_usage(), 2);
+            assert_eq!(limiter.current_usage(), 2);
         }
 
         {
-            let limiter = MemoryLimiter::new_shared(10);
-            let mut vec_u32: LimitedVec<u32> = LimitedVec::new(Rc::clone(&limiter));
+            let limiter = SharedMemoryLimiter::new(10);
+            let mut vec_u32: LimitedVec<u32> = LimitedVec::new(limiter.clone());
 
             vec_u32.push(1).unwrap();
             vec_u32.push(2).unwrap();
-            assert_eq!(limiter.borrow().current_usage(), 8);
+            assert_eq!(limiter.current_usage(), 8);
         }
     }
 
     #[test]
     fn max_limit() {
-        let limiter = MemoryLimiter::new_shared(2);
-        let mut vector: LimitedVec<u8> = LimitedVec::new(Rc::clone(&limiter));
+        let limiter = SharedMemoryLimiter::new(2);
+        let mut vector: LimitedVec<u8> = LimitedVec::new(limiter.clone());
 
         vector.push(1).unwrap();
         vector.push(2).unwrap();
@@ -139,38 +134,38 @@ mod tests {
 
     #[test]
     fn drop() {
-        let limiter = MemoryLimiter::new_shared(1);
+        let limiter = SharedMemoryLimiter::new(1);
 
         {
-            let mut vector: LimitedVec<u8> = LimitedVec::new(Rc::clone(&limiter));
+            let mut vector: LimitedVec<u8> = LimitedVec::new(limiter.clone());
 
             vector.push(1).unwrap();
-            assert_eq!(limiter.borrow().current_usage(), 1);
+            assert_eq!(limiter.current_usage(), 1);
         }
 
-        assert_eq!(limiter.borrow().current_usage(), 0);
+        assert_eq!(limiter.current_usage(), 0);
     }
 
     #[test]
     fn drain() {
-        let limiter = MemoryLimiter::new_shared(10);
-        let mut vector: LimitedVec<u8> = LimitedVec::new(Rc::clone(&limiter));
+        let limiter = SharedMemoryLimiter::new(10);
+        let mut vector: LimitedVec<u8> = LimitedVec::new(limiter.clone());
 
         vector.push(1).unwrap();
         vector.push(2).unwrap();
         vector.push(3).unwrap();
-        assert_eq!(limiter.borrow().current_usage(), 3);
+        assert_eq!(limiter.current_usage(), 3);
 
         vector.drain(0..3);
-        assert_eq!(limiter.borrow().current_usage(), 0);
+        assert_eq!(limiter.current_usage(), 0);
 
         vector.push(1).unwrap();
         vector.push(2).unwrap();
         vector.push(3).unwrap();
         vector.push(4).unwrap();
-        assert_eq!(limiter.borrow().current_usage(), 4);
+        assert_eq!(limiter.current_usage(), 4);
 
         vector.drain(1..=2);
-        assert_eq!(limiter.borrow().current_usage(), 2);
+        assert_eq!(limiter.current_usage(), 2);
     }
 }

src/memory/limiter.rs

@@ -1,9 +1,7 @@
-use std::cell::RefCell;
+use std::cell::Cell;
 use std::rc::Rc;
 use thiserror::Error;
 
-pub type SharedMemoryLimiter = Rc<RefCell<MemoryLimiter>>;
-
 /// An error that occures when rewriter exceedes the memory limit specified in the
 /// [`MemorySettings`].
 ///
@@ -12,46 +10,52 @@ pub type SharedMemoryLimiter = Rc<RefCell<MemoryLimiter>>;
 #[error("The memory limit has been exceeded.")]
 pub struct MemoryLimitExceededError;
 
-#[derive(Debug)]
-pub struct MemoryLimiter {
-    current_usage: usize,
+#[derive(Debug, Clone)]
+pub struct SharedMemoryLimiter {
+    current_usage: Rc<Cell<usize>>,
     max: usize,
 }
 
-impl MemoryLimiter {
-    pub fn new_shared(max: usize) -> SharedMemoryLimiter {
-        Rc::new(RefCell::new(MemoryLimiter {
+impl SharedMemoryLimiter {
+    pub fn new(max: usize) -> SharedMemoryLimiter {
+        SharedMemoryLimiter {
+            current_usage: Rc::new(Cell::new(0)),
             max,
-            current_usage: 0,
-        }))
+        }
     }
 
     #[cfg(test)]
     pub fn current_usage(&self) -> usize {
-        self.current_usage
+        self.current_usage.get()
     }
 
     #[inline]
-    pub fn increase_usage(&mut self, byte_count: usize) -> Result<(), MemoryLimitExceededError> {
-        self.current_usage += byte_count;
+    pub fn increase_usage(&self, byte_count: usize) -> Result<(), MemoryLimitExceededError> {
+        let previous_usage = self.current_usage.get();
+        let current_usage = previous_usage + byte_count;
 
-        if self.current_usage > self.max {
+        self.current_usage.set(current_usage);
+
+        if current_usage > self.max {
             Err(MemoryLimitExceededError)
         } else {
             Ok(())
         }
     }
 
     #[inline]
-    pub fn preallocate(&mut self, byte_count: usize) {
+    pub fn preallocate(&self, byte_count: usize) {
         self.increase_usage(byte_count).expect(
             "Total preallocated memory size should be less than `MemorySettings::max_allowed_memory_usage`.",
         );
     }
 
     #[inline]
-    pub fn decrease_usage(&mut self, byte_count: usize) {
-        self.current_usage -= byte_count;
+    pub fn decrease_usage(&self, byte_count: usize) {
+        let previous_usage = self.current_usage.get();
+        let current_usage = previous_usage - byte_count;
+
+        self.current_usage.set(current_usage);
     }
 }
 
@@ -61,8 +65,7 @@ mod tests {
 
     #[test]
     fn current_usage() {
-        let limiter = MemoryLimiter::new_shared(10);
-        let mut limiter = limiter.borrow_mut();
+        let limiter = SharedMemoryLimiter::new(10);
 
         assert_eq!(limiter.current_usage(), 0);
 
@@ -85,8 +88,7 @@ mod tests {
         expected = "Total preallocated memory size should be less than `MemorySettings::max_allowed_memory_usage`."
     )]
     fn preallocate() {
-        let limiter = MemoryLimiter::new_shared(10);
-        let mut limiter = limiter.borrow_mut();
+        let limiter = SharedMemoryLimiter::new(10);
 
         limiter.preallocate(8);
         assert_eq!(limiter.current_usage(), 8);

src/memory/mod.rs

@@ -4,4 +4,4 @@ mod limiter;
 
 pub use arena::Arena;
 pub use limited_vec::LimitedVec;
-pub use limiter::{MemoryLimitExceededError, MemoryLimiter, SharedMemoryLimiter};
+pub use limiter::{MemoryLimitExceededError, SharedMemoryLimiter};

src/rewriter/mod.rs

@@ -7,8 +7,7 @@ mod settings;
 use self::handlers_dispatcher::ContentHandlersDispatcher;
 use self::rewrite_controller::*;
 use crate::base::SharedEncoding;
-use crate::memory::MemoryLimitExceededError;
-use crate::memory::MemoryLimiter;
+use crate::memory::{MemoryLimitExceededError, SharedMemoryLimiter};
 use crate::parser::ParsingAmbiguityError;
 use crate::selectors_vm::{self, SelectorMatchingVm};
 use crate::transform_stream::*;
@@ -17,7 +16,6 @@ use mime::Mime;
 use std::borrow::Cow;
 use std::error::Error as StdError;
 use std::fmt::{self, Debug};
-use std::rc::Rc;
 use thiserror::Error;
 
 pub use self::settings::*;
@@ -192,13 +190,13 @@ impl<'h, O: OutputSink> HtmlRewriter<'h, O> {
         }
 
         let memory_limiter =
-            MemoryLimiter::new_shared(settings.memory_settings.max_allowed_memory_usage);
+            SharedMemoryLimiter::new(settings.memory_settings.max_allowed_memory_usage);
 
         let selector_matching_vm = if has_selectors {
             Some(SelectorMatchingVm::new(
                 selectors_ast,
                 settings.encoding.into(),
-                Rc::clone(&memory_limiter),
+                memory_limiter.clone(),
                 settings.enable_esi_tags,
             ))
         } else {
@@ -578,7 +576,7 @@ mod tests {
         macro_rules! create_handlers {
             ($sel:expr, $idx:expr) => {
                 element!($sel, {
-                    let handlers_executed = Rc::clone(&handlers_executed);
+                    let handlers_executed = ::std::rc::Rc::clone(&handlers_executed);
 
                     move |_| {
                         handlers_executed.borrow_mut().push($idx);