https://git.ruby-lang.org/ruby.git/commit/?id=35471a9487

From 35471a948739ca13b85fe900871e081d553f68e6 Mon Sep 17 00:00:00 2001
From: Koichi Sasada <ko1@a...>
Date: Tue, 22 Dec 2020 01:55:15 +0900
Subject: add Ractor#[]/#[]= for ractor local storage

This API is similar to plain old Thread#[]/Fiber#[] interface
with symbol key.

diff --git a/bootstraptest/test_ractor.rb b/bootstraptest/test_ractor.rb
index a853e91..06bd739 100644
--- a/bootstraptest/test_ractor.rb
+++ b/bootstraptest/test_ractor.rb
@@ -1245,6 +1245,20 @@ assert_equal '[:ok, :ok]', %q{ https://github.com/ruby/ruby/blob/trunk/bootstraptest/test_ractor.rb#L1245
   end
 }
 
+# Ractor-local storage
+assert_equal '[nil, "b", "a"]', %q{
+  ans = []
+  Ractor.current[:key] = 'a'
+  r = Ractor.new{
+    Ractor.yield self[:key]
+    self[:key] = 'b'
+    self[:key]
+  }
+  ans << r.take
+  ans << r.take
+  ans << Ractor.current[:key]
+}
+
 ###
 ### Synchronization tests
 ###
diff --git a/ractor.c b/ractor.c
index fab4646..5937ac6 100644
--- a/ractor.c
+++ b/ractor.c
@@ -2928,6 +2928,13 @@ ractor_local_storage_mark_i(st_data_t key, st_data_t val, st_data_t dmy) https://github.com/ruby/ruby/blob/trunk/ractor.c#L2928
     return ST_CONTINUE;
 }
 
+static enum rb_id_table_iterator_result
+idkey_local_storage_mark_i(ID id, VALUE val, void *dmy)
+{
+    rb_gc_mark(val);
+    return ID_TABLE_CONTINUE;
+}
+
 static void
 ractor_local_storage_mark(rb_ractor_t *r)
 {
@@ -2943,6 +2950,10 @@ ractor_local_storage_mark(rb_ractor_t *r) https://github.com/ruby/ruby/blob/trunk/ractor.c#L2950
             }
         }
     }
+
+    if (r->idkey_local_storage) {
+        rb_id_table_foreach(r->idkey_local_storage, idkey_local_storage_mark_i, NULL);
+    }
 }
 
 static int
@@ -2960,6 +2971,10 @@ ractor_local_storage_free(rb_ractor_t *r) https://github.com/ruby/ruby/blob/trunk/ractor.c#L2971
         st_foreach(r->local_storage, ractor_local_storage_free_i, 0);
         st_free_table(r->local_storage);
     }
+
+    if (r->idkey_local_storage) {
+        rb_id_table_free(r->idkey_local_storage);
+    }
 }
 
 static void
@@ -3103,4 +3118,35 @@ rb_ractor_finish_marking(void) https://github.com/ruby/ruby/blob/trunk/ractor.c#L3118
     }
 }
 
+static VALUE
+ractor_local_value(rb_execution_context_t *ec, VALUE self, VALUE sym)
+{
+    rb_ractor_t *cr = rb_ec_ractor_ptr(ec);
+    ID id = rb_check_id(&sym);
+    struct rb_id_table *tbl = cr->idkey_local_storage;
+    VALUE val;
+
+    if (id && tbl && rb_id_table_lookup(tbl, id, &val)) {
+        rp(val);
+        return val;
+    }
+    else {
+        return Qnil;
+    }
+}
+
+static VALUE
+ractor_local_value_set(rb_execution_context_t *ec, VALUE self, VALUE sym, VALUE val)
+{
+    rb_ractor_t *cr = rb_ec_ractor_ptr(ec);
+    ID id = SYM2ID(rb_to_symbol(sym));
+    struct rb_id_table *tbl = cr->idkey_local_storage;
+
+    if (tbl == NULL) {
+        tbl = cr->idkey_local_storage = rb_id_table_create(2);
+    }
+    rb_id_table_insert(tbl, id, val);
+    return val;
+}
+
 #include "ractor.rbinc"
diff --git a/ractor.rb b/ractor.rb
index c73c574..d72e02b 100644
--- a/ractor.rb
+++ b/ractor.rb
@@ -758,4 +758,14 @@ class Ractor https://github.com/ruby/ruby/blob/trunk/ractor.rb#L758
       }
     end
   end
+
+  # get a value from ractor-local storage
+  def [](sym)
+    Primitive.ractor_local_value(sym)
+  end
+
+  # set a value in ractor-local storage
+  def []=(sym, val)
+    Primitive.ractor_local_value_set(sym, val)
+  end
 end
diff --git a/ractor_core.h b/ractor_core.h
index 1f1e7aa..4451e06 100644
--- a/ractor_core.h
+++ b/ractor_core.h
@@ -130,6 +130,7 @@ struct rb_ractor_struct { https://github.com/ruby/ruby/blob/trunk/ractor_core.h#L130
     // ractor local data
 
     st_table *local_storage;
+    struct rb_id_table *idkey_local_storage;
 
     VALUE r_stdin;
     VALUE r_stdout;
-- 
cgit v0.10.2


--
ML: ruby-changes@q...
Info: http://www.atdot.net/~ko1/quickml/