From 0bdb38ba6be208064a514c12a9b80328645689f8 Mon Sep 17 00:00:00 2001
From: Koichi Sasada <ko1@atdot.net>
Date: Fri, 13 Dec 2024 04:52:34 +0900
Subject: [PATCH] `Ractor.set_if_absent(key)`

to initialize ractor local storage in thread-safety.
[Feature #20875]
---
 bootstraptest/test_ractor.rb | 15 +++++++++++
 ractor.c                     | 50 ++++++++++++++++++++++++++++++++++++
 ractor.rb                    |  4 +++
 ractor_core.h                |  1 +
 4 files changed, 70 insertions(+)

diff --git a/bootstraptest/test_ractor.rb b/bootstraptest/test_ractor.rb
index 22e1533364..08f269dda9 100644
--- a/bootstraptest/test_ractor.rb
+++ b/bootstraptest/test_ractor.rb
@@ -1503,6 +1503,21 @@ assert_equal '[nil, "b", "a"]', %q{
   ans << Ractor.current[:key]
 }
 
+assert_equal '1', %q{
+  N = 1_000
+  Ractor.new{
+    a = []
+    1_000.times.map{|i|
+      Thread.new(i){|i|
+        Thread.pass if i < N
+        a << Ractor.store_if_absent(:i){ i }
+        a << Ractor.current[:i]
+      }
+    }.each(&:join)
+    a.uniq.size
+  }.take
+}
+
 ###
 ### Synchronization tests
 ###
diff --git a/ractor.c b/ractor.c
index 76d697a96f..9db904dc03 100644
--- a/ractor.c
+++ b/ractor.c
@@ -3882,6 +3882,56 @@ ractor_local_value_set(rb_execution_context_t *ec, VALUE self, VALUE sym, VALUE
     return val;
 }
 
+struct ractor_local_storage_store_data {
+    rb_execution_context_t *ec;
+    struct rb_id_table *tbl;
+    ID id;
+    VALUE sym;
+};
+
+static VALUE
+ractor_local_value_store_i(VALUE ptr)
+{
+    VALUE val;
+    struct ractor_local_storage_store_data *data = (struct ractor_local_storage_store_data *)ptr;
+
+    if (rb_id_table_lookup(data->tbl, data->id, &val)) {
+        // after synchronization, we found already registerred entry
+    }
+    else {
+        val = rb_yield(Qnil);
+        ractor_local_value_set(data->ec, Qnil, data->sym, val);
+    }
+    return val;
+}
+
+static VALUE
+ractor_local_value_store_if_absent(rb_execution_context_t *ec, VALUE self, VALUE sym)
+{
+    rb_ractor_t *cr = rb_ec_ractor_ptr(ec);
+    struct ractor_local_storage_store_data data = {
+        .ec = ec,
+        .sym = sym,
+        .id = SYM2ID(rb_to_symbol(sym)),
+        .tbl = cr->idkey_local_storage,
+    };
+    VALUE val;
+
+    if (data.tbl == NULL) {
+        data.tbl = cr->idkey_local_storage = rb_id_table_create(2);
+    }
+    else if (rb_id_table_lookup(data.tbl, data.id, &val)) {
+        // already set
+        return val;
+    }
+
+    if (!cr->local_storage_store_lock) {
+        cr->local_storage_store_lock = rb_mutex_new();
+    }
+
+    return rb_mutex_synchronize(cr->local_storage_store_lock, ractor_local_value_store_i, (VALUE)&data);
+}
+
 // Ractor::Channel (emulate with Ractor)
 
 typedef rb_ractor_t rb_ractor_channel_t;
diff --git a/ractor.rb b/ractor.rb
index 663f93132a..da0ca1d4c2 100644
--- a/ractor.rb
+++ b/ractor.rb
@@ -856,6 +856,10 @@ class Ractor
     Primitive.ractor_local_value_set(sym, val)
   end
 
+  def self.store_if_absent(sym)
+    Primitive.ractor_local_value_store_if_absent(sym)
+  end
+
   # returns main ractor
   def self.main
     __builtin_cexpr! %q{
diff --git a/ractor_core.h b/ractor_core.h
index 82fc7047e5..1e860edb4b 100644
--- a/ractor_core.h
+++ b/ractor_core.h
@@ -180,6 +180,7 @@ struct rb_ractor_struct {
 
     st_table *local_storage;
     struct rb_id_table *idkey_local_storage;
+    VALUE local_storage_store_lock;
 
     VALUE r_stdin;
     VALUE r_stdout;