ratomic 0.2.0 → 0.3.0

data/ext/ratomic/src/lib.rs

@@ -7,17 +7,17 @@ mod sem;
 
 use counter::AtomicCounter;
 use fixed_size_object_pool::FixedSizeObjectPool;
-use hashmap::ConcurrentHashMap;
+use hashmap::MapStore;
 use magnus::{
-    data_type_builder, method, IntoValue,
+    data_type_builder, method,
     prelude::*,
     typed_data::{DataType, DataTypeFunctions},
     value::Lazy,
-    Error, RClass, Ruby, TryConvert, TypedData, Value,
+    Error, IntoValue, RClass, Ruby, TryConvert, TypedData, Value,
 };
 use mpmc_queue::MpmcQueue;
-use rb_sys::{rb_ext_ractor_safe, rb_thread_call_without_gvl, ruby_special_consts, VALUE};
 use parking_lot::Mutex;
+use rb_sys::{rb_ext_ractor_safe, rb_thread_call_without_gvl, ruby_special_consts, VALUE};
 use std::{ffi::c_void, mem::transmute};
 
 fn value_to_raw(value: Value) -> VALUE {
@@ -64,7 +64,8 @@ impl DataTypeFunctions for Counter {}
 unsafe impl TypedData for Counter {
     fn class(ruby: &Ruby) -> RClass {
         static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class = ruby.define_module("Ratomic")
+            let class = ruby
+                .define_module("Ratomic")
                 .unwrap()
                 .define_class("Counter", ruby.class_object())
                 .unwrap();
@@ -75,17 +76,18 @@ unsafe impl TypedData for Counter {
     }
 
     fn data_type() -> &'static DataType {
-        static DATA_TYPE: DataType =
-            data_type_builder!(Counter, "ratomic/counter").frozen_shareable().build();
+        static DATA_TYPE: DataType = data_type_builder!(Counter, "ratomic/counter")
+            .frozen_shareable()
+            .build();
         &DATA_TYPE
     }
 }
 
-struct HashMap(ConcurrentHashMap);
+struct HashMap(MapStore);
 
 impl HashMap {
     fn new(ruby: &Ruby, class: RClass) -> Result<Value, Error> {
-        let value = ruby.wrap_as(Self(ConcurrentHashMap::new()), class).as_value();
+        let value = ruby.wrap_as(Self(MapStore::new()), class).as_value();
         make_shareable(ruby, value)
     }
 
@@ -94,10 +96,19 @@ impl HashMap {
         unsafe { value_from_raw(raw) }.into_value_with(ruby)
     }
 
+    fn contains_key(&self, key: Value) -> bool {
+        self.0.contains_key(value_to_raw(key))
+    }
+
     fn set(&self, key: Value, value: Value) {
         self.0.set(value_to_raw(key), value_to_raw(value));
     }
 
+    fn delete(ruby: &Ruby, rb_self: &Self, key: Value) -> Value {
+        let raw = rb_self.0.delete(value_to_raw(key)).unwrap_or_else(qnil_raw);
+        unsafe { value_from_raw(raw) }.into_value_with(ruby)
+    }
+
     fn clear(&self) {
         self.0.clear();
     }
@@ -115,26 +126,130 @@ impl HashMap {
         }
 
         let proc = ruby.block_proc()?;
+        let mut error = None;
         rb_self.0.fetch_and_modify(value_to_raw(key), |value| {
-            let result: Value = proc.call((unsafe { value_from_raw(value) },)).unwrap();
-            value_to_raw(result)
+            match proc.call::<_, Value>((unsafe { value_from_raw(value) },)) {
+                Ok(result) => value_to_raw(result),
+                Err(err) => {
+                    error = Some(err);
+                    value
+                }
+            }
         });
-        Ok(())
+
+        if let Some(err) = error {
+            Err(err)
+        } else {
+            Ok(())
+        }
+    }
+
+    fn compute(ruby: &Ruby, rb_self: &Self, key: Value) -> Result<Value, Error> {
+        if !ruby.block_given() {
+            return Err(Error::new(
+                ruby.exception_local_jump_error(),
+                "no block given",
+            ));
+        }
+
+        let proc = ruby.block_proc()?;
+        let raw = rb_self.0.compute(value_to_raw(key), qnil_raw(), |value| {
+            proc.call::<_, Value>((unsafe { value_from_raw(value) },))
+                .map(value_to_raw)
+        })?;
+
+        Ok(unsafe { value_from_raw(raw) }.into_value_with(ruby))
+    }
+
+    fn fetch_or_store(ruby: &Ruby, rb_self: &Self, key: Value) -> Result<Value, Error> {
+        if !ruby.block_given() {
+            return Err(Error::new(
+                ruby.exception_local_jump_error(),
+                "no block given",
+            ));
+        }
+
+        let proc = ruby.block_proc()?;
+        let raw = rb_self.0.fetch_or_store(value_to_raw(key), || {
+            proc.call::<_, Value>(()).map(value_to_raw)
+        })?;
+
+        Ok(unsafe { value_from_raw(raw) }.into_value_with(ruby))
+    }
+
+    fn upsert(ruby: &Ruby, rb_self: &Self, key: Value, initial: Value) -> Result<Value, Error> {
+        if !ruby.block_given() {
+            return Err(Error::new(
+                ruby.exception_local_jump_error(),
+                "no block given",
+            ));
+        }
+
+        let proc = ruby.block_proc()?;
+        let raw = rb_self
+            .0
+            .upsert(value_to_raw(key), value_to_raw(initial), |value| {
+                proc.call::<_, Value>((unsafe { value_from_raw(value) },))
+                    .map(value_to_raw)
+            })?;
+
+        Ok(unsafe { value_from_raw(raw) }.into_value_with(ruby))
+    }
+
+    fn increment_numeric(
+        ruby: &Ruby,
+        rb_self: &Self,
+        key: Value,
+        by: Value,
+    ) -> Result<Value, Error> {
+        let key_inspect = key.funcall::<_, _, String>("inspect", ())?;
+        let numeric_class: RClass = ruby.class_object().const_get("Numeric")?;
+        let raw = rb_self.0.update(value_to_raw(key), |current| {
+            let next = match current {
+                Some(value) if value == qnil_raw() => {
+                    return Err(Error::new(
+                        ruby.exception_type_error(),
+                        format!("existing value for {key_inspect} must be numeric: nil"),
+                    ));
+                }
+                Some(value) => {
+                    let old_value = unsafe { value_from_raw(value) };
+                    if !old_value.funcall::<_, _, bool>("is_a?", (numeric_class,))? {
+                        let old_value_inspect = old_value.funcall::<_, _, String>("inspect", ())?;
+                        return Err(Error::new(
+                            ruby.exception_type_error(),
+                            format!(
+                                "existing value for {key_inspect} must be numeric: {old_value_inspect}"
+                            ),
+                        ));
+                    }
+
+                    old_value.funcall::<_, _, Value>("+", (by,))?
+                }
+                None => by,
+            };
+
+            Ok(value_to_raw(next))
+        })?;
+
+        Ok(unsafe { value_from_raw(raw) }.into_value_with(ruby))
     }
 }
 
 impl DataTypeFunctions for HashMap {
     fn mark(&self, marker: &magnus::gc::Marker) {
-        self.0.mark(|value| marker.mark(unsafe { value_from_raw(value) }));
+        self.0
+            .mark(|value| marker.mark(unsafe { value_from_raw(value) }));
     }
 }
 
 unsafe impl TypedData for HashMap {
     fn class(ruby: &Ruby) -> RClass {
         static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class = ruby.define_module("Ratomic")
+            let class = ruby
+                .define_module("Ratomic")
                 .unwrap()
-                .define_class("ConcurrentHashMap", ruby.class_object())
+                .define_class("Map", ruby.class_object())
                 .unwrap();
             class.undef_default_alloc_func();
             class
@@ -234,14 +349,16 @@ impl Queue {
 
 impl DataTypeFunctions for Queue {
     fn mark(&self, marker: &magnus::gc::Marker) {
-        self.0.mark(|value| marker.mark(unsafe { value_from_raw(value) }));
+        self.0
+            .mark(|value| marker.mark(unsafe { value_from_raw(value) }));
     }
 }
 
 unsafe impl TypedData for Queue {
     fn class(ruby: &Ruby) -> RClass {
         static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class = ruby.define_module("Ratomic")
+            let class = ruby
+                .define_module("Ratomic")
                 .unwrap()
                 .define_class("Queue", ruby.class_object())
                 .unwrap();
@@ -267,7 +384,10 @@ impl Pool {
         if args.len() > 2 {
             return Err(Error::new(
                 ruby.exception_arg_error(),
-                format!("wrong number of arguments (given {}, expected 0..2)", args.len()),
+                format!(
+                    "wrong number of arguments (given {}, expected 0..2)",
+                    args.len()
+                ),
             ));
         }
         let size = args
@@ -285,7 +405,10 @@ impl Pool {
             .unwrap_or(1000);
 
         if size == 0 {
-            return Err(Error::new(ruby.exception_arg_error(), "pool size must be positive"));
+            return Err(Error::new(
+                ruby.exception_arg_error(),
+                "pool size must be positive",
+            ));
         }
         if !ruby.block_given() {
             return Err(Error::new(
@@ -340,7 +463,8 @@ impl DataTypeFunctions for Pool {
 unsafe impl TypedData for Pool {
     fn class(ruby: &Ruby) -> RClass {
         static CLASS: Lazy<RClass> = Lazy::new(|ruby| {
-            let class = ruby.define_module("Ratomic")
+            let class = ruby
+                .define_module("Ratomic")
                 .unwrap()
                 .define_class("FixedSizeObjectPool", ruby.class_object())
                 .unwrap();
@@ -372,14 +496,23 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     counter.define_method("decrement", method!(Counter::decrement, 1))?;
     counter.define_method("read", method!(Counter::read, 0))?;
 
-    let hashmap = root.define_class("ConcurrentHashMap", ruby.class_object())?;
+    let hashmap = root.define_class("Map", ruby.class_object())?;
     hashmap.undef_default_alloc_func();
     hashmap.define_singleton_method("new", method!(HashMap::new, 0))?;
     hashmap.define_method("get", method!(HashMap::get, 1))?;
+    hashmap.define_method("key?", method!(HashMap::contains_key, 1))?;
     hashmap.define_method("set", method!(HashMap::set, 2))?;
+    hashmap.define_method("delete", method!(HashMap::delete, 1))?;
     hashmap.define_method("clear", method!(HashMap::clear, 0))?;
     hashmap.define_method("size", method!(HashMap::size, 0))?;
     hashmap.define_method("fetch_and_modify", method!(HashMap::fetch_and_modify, 1))?;
+    hashmap.define_method("compute", method!(HashMap::compute, 1))?;
+    hashmap.define_method("fetch_or_store", method!(HashMap::fetch_or_store, 1))?;
+    hashmap.define_method("upsert", method!(HashMap::upsert, 2))?;
+    hashmap.define_private_method(
+        "__increment_numeric",
+        method!(HashMap::increment_numeric, 2),
+    )?;
 
     let queue = root.define_class("Queue", ruby.class_object())?;
     queue.undef_default_alloc_func();

data/ext/ratomic/src/mpmc_queue.rs

@@ -15,7 +15,7 @@ unsafe impl Sync for QueueElement {}
 
 pub struct MpmcQueue {
     buffer: Vec<QueueElement>,
-    buffer_mask: usize,
+    buffer_size: usize,
     enqueue_pos: AtomicUsize,
     dequeue_pos: AtomicUsize,
     gc_guard: GcGuard,
@@ -25,7 +25,7 @@ pub struct MpmcQueue {
 
 impl MpmcQueue {
     pub fn new(buffer_size: usize, default: VALUE) -> Self {
-        let mut buffer = Vec::with_capacity(buffer_size.next_power_of_two());
+        let mut buffer = Vec::with_capacity(buffer_size);
         for i in 0..buffer_size {
             buffer.push(QueueElement {
                 sequence: AtomicUsize::new(i),
@@ -42,7 +42,7 @@ impl MpmcQueue {
 
         Self {
             buffer,
-            buffer_mask: buffer_size - 1,
+            buffer_size,
             enqueue_pos: AtomicUsize::new(0),
             dequeue_pos: AtomicUsize::new(0),
             gc_guard,
@@ -55,7 +55,7 @@ impl MpmcQueue {
         let mut cell;
         let mut pos = self.enqueue_pos.load(Ordering::Relaxed);
         loop {
-            cell = &self.buffer[pos & self.buffer_mask];
+            cell = &self.buffer[pos % self.buffer_size];
             let seq = cell.sequence.load(Ordering::Acquire);
             let diff = seq as isize - pos as isize;
             if diff == 0 {
@@ -82,7 +82,7 @@ impl MpmcQueue {
         let mut cell;
         let mut pos = self.dequeue_pos.load(Ordering::Relaxed);
         loop {
-            cell = &self.buffer[pos & self.buffer_mask];
+            cell = &self.buffer[pos % self.buffer_size];
             let seq = cell.sequence.load(Ordering::Acquire);
             let diff = seq as isize - (pos + 1) as isize;
             if diff == 0 {
@@ -102,7 +102,7 @@ impl MpmcQueue {
 
         let data = cell.data.get();
         cell.sequence
-            .store(pos + self.buffer_mask + 1, Ordering::Release);
+            .store(pos + self.buffer_size, Ordering::Release);
         self.write_sem.post();
 
         #[cfg(feature = "simulation")]
@@ -131,7 +131,7 @@ impl MpmcQueue {
 
     pub fn peek(&self) -> Option<VALUE> {
         let pos = self.dequeue_pos.load(Ordering::Relaxed);
-        let cell = &self.buffer[pos & self.buffer_mask];
+        let cell = &self.buffer[pos % self.buffer_size];
         let seq = cell.sequence.load(Ordering::Acquire);
         let diff = seq as isize - (pos + 1) as isize;
         if diff == 0 {

data/lib/ratomic/counter.rb

@@ -1,8 +1,38 @@
 # frozen_string_literal: true
 
 module Ratomic
-  # Ruby convenience methods for {Counter}.
-  module CounterMethods
+  # A Ractor-shareable atomic counter.
+  #
+  # Counter stores an unsigned integer in native Rust atomics and can be shared
+  # safely across Ractors.
+  #
+  # @example Count work across Ractors
+  #   counter = Ratomic::Counter.new
+  #   counter.increment(1)
+  #   counter.read # => 1
+  #
+  # @!method self.new
+  #   Create a counter initialized to zero.
+  #
+  #   @return [Ratomic::Counter]
+  #
+  # @!method read
+  #   Read the current counter value.
+  #
+  #   @return [Integer]
+  #
+  # @!method increment(amt)
+  #   Increment the counter by +amt+.
+  #
+  #   @param amt [Integer]
+  #   @return [void]
+  #
+  # @!method decrement(amt)
+  #   Decrement the counter by +amt+.
+  #
+  #   @param amt [Integer]
+  #   @return [void]
+  class Counter
     # Read the current counter value.
     #
     # @return [Integer]
@@ -23,29 +53,5 @@ module Ratomic
     def zero?
       read.zero?
     end
-
-    # Increment the counter.
-    #
-    # @param amt [Integer] amount to add
-    # @raise [ArgumentError] if +amt+ is negative
-    # @return [void]
-    def inc(amt = 1)
-      raise ArgumentError, "amount must be positive: #{amt}" if amt.negative?
-
-      increment(amt)
-    end
-
-    # Decrement the counter.
-    #
-    # @param amt [Integer] amount to subtract
-    # @raise [ArgumentError] if +amt+ is negative
-    # @return [void]
-    def dec(amt = 1)
-      raise ArgumentError, "amount must be positive: #{amt}" if amt.negative?
-
-      decrement(amt)
-    end
   end
-
-  Counter.prepend(CounterMethods)
 end