ratomic 0.2.1 → 0.3.1

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 {
@@ -46,12 +46,14 @@ impl Counter {
         make_shareable(ruby, value)
     }
 
-    fn increment(&self, amt: u64) {
+    fn increment(&self, amt: u64) -> u64 {
         self.0.inc(amt);
+        self.0.read()
     }
 
-    fn decrement(&self, amt: u64) {
+    fn decrement(&self, amt: u64) -> u64 {
         self.0.dec(amt);
+        self.0.read()
     }
 
     fn read(&self) -> u64 {
@@ -64,7 +66,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 +78,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,12 +98,24 @@ impl HashMap {
         unsafe { value_from_raw(raw) }.into_value_with(ruby)
     }
 
-    fn set(&self, key: Value, value: Value) {
+    fn contains_key(&self, key: Value) -> bool {
+        self.0.contains_key(value_to_raw(key))
+    }
+
+    fn set(&self, key: Value, value: Value) -> Value {
         self.0.set(value_to_raw(key), value_to_raw(value));
+        value
     }
 
-    fn clear(&self) {
-        self.0.clear();
+    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(_ruby: &Ruby, value: Value) -> Result<Value, Error> {
+        let rb_self: &Self = TryConvert::try_convert(value)?;
+        rb_self.0.clear();
+        Ok(value)
     }
 
     fn size(&self) -> usize {
@@ -132,20 +148,113 @@ impl HashMap {
             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
@@ -203,9 +312,10 @@ impl Queue {
         make_shareable(ruby, value)
     }
 
-    fn push(&self, item: Value) {
+    fn push(ruby: &Ruby, rb_self: Value, item: Value) -> Result<Value, Error> {
+        let queue: &Self = TryConvert::try_convert(rb_self)?;
         let mut payload = PushPayload {
-            queue: &self.0,
+            queue: &queue.0,
             item: value_to_raw(item),
         };
         unsafe {
@@ -216,6 +326,7 @@ impl Queue {
                 std::ptr::null_mut(),
             );
         }
+        Ok(rb_self.into_value_with(ruby))
     }
 
     fn pop(&self) -> Value {
@@ -245,14 +356,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();
@@ -278,7 +391,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
@@ -296,7 +412,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(
@@ -351,7 +470,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();
@@ -383,14 +503,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/lib/ratomic/counter.rb

@@ -1,51 +1,57 @@
 # 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) # => 1
+  #   counter.read # => 1
+  #
+  # @!method self.new
+  #   Create a counter initialized to zero.
+  #
+  #   @return [Ratomic::Counter] a new shareable counter
+  #
+  # @!method read
+  #   Read the current counter value.
+  #
+  #   @return [Integer] the current counter value
+  #
+  # @!method increment(amt)
+  #   Increment the counter by +amt+.
+  #
+  #   @param amt [Integer] amount to add to the counter
+  #   @return [Integer] the updated counter value
+  #
+  # @!method decrement(amt)
+  #   Decrement the counter by +amt+.
+  #
+  #   @param amt [Integer] amount to subtract from the counter
+  #   @return [Integer] the updated counter value
+  class Counter
     # Read the current counter value.
     #
-    # @return [Integer]
+    # @return [Integer] the current counter value
     def value
       read
     end
 
     # Coerce the counter to an Integer snapshot.
     #
-    # @return [Integer]
+    # @return [Integer] the current counter value
     def to_i
       read
     end
 
     # Check whether the current counter value is zero.
     #
-    # @return [Boolean]
+    # @return [Boolean] true when the counter currently reads zero
     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