Assigning `nil` to fiber storage deletes the association. (#7378)

Also avoid allocations when looking up `Fiber#storage` if not needed.

2 files changed, 31 insertions, 13 deletions

diff --git a/cont.c b/cont.c
index 3962572b4e..a932bebf41 100644
--- a/cont.c
+++ b/cont.c
@@ -2059,10 +2059,10 @@ fiber_storage_set(struct rb_fiber_struct *fiber, VALUE storage)
 }
 
 static inline VALUE
-fiber_storage_get(rb_fiber_t *fiber)
+fiber_storage_get(rb_fiber_t *fiber, int allocate)
 {
     VALUE storage = fiber->cont.saved_ec.storage;
-    if (storage == Qnil) {
+    if (storage == Qnil && allocate) {
         storage = rb_hash_new();
         fiber_storage_set(fiber, storage);
     }
@@ -2089,7 +2089,14 @@ static VALUE
 rb_fiber_storage_get(VALUE self)
 {
     storage_access_must_be_from_same_fiber(self);
-    return rb_obj_dup(fiber_storage_get(fiber_ptr(self)));
+
+    VALUE storage = fiber_storage_get(fiber_ptr(self), FALSE);
+
+    if (storage == Qnil) {
+        return Qnil;
+    } else {
+        return rb_obj_dup(storage);
+    }
 }
 
 static int
@@ -2170,8 +2177,7 @@ rb_fiber_storage_aref(VALUE class, VALUE key)
     ID id = rb_check_id(&key);
     if (!id) return Qnil;
 
-    VALUE storage = fiber_storage_get(fiber_current());
-
+    VALUE storage = fiber_storage_get(fiber_current(), FALSE);
     if (storage == Qnil) return Qnil;
 
     return rb_hash_aref(storage, key);
@@ -2193,9 +2199,14 @@ rb_fiber_storage_aset(VALUE class, VALUE key, VALUE value)
     ID id = rb_check_id(&key);
     if (!id) return Qnil;
 
-    VALUE storage = fiber_storage_get(fiber_current());
+    VALUE storage = fiber_storage_get(fiber_current(), value != Qnil);
+    if (storage == Qnil) return Qnil;
 
-    return rb_hash_aset(storage, key, value);
+    if (value == Qnil) {
+        return rb_hash_delete(storage, key);
+    } else {
+        return rb_hash_aset(storage, key, value);
+    }
 }
 
 static VALUE
diff --git a/spec/ruby/core/fiber/storage_spec.rb b/spec/ruby/core/fiber/storage_spec.rb
index 98215ebd59..e99fe6e4df 100644
--- a/spec/ruby/core/fiber/storage_spec.rb
+++ b/spec/ruby/core/fiber/storage_spec.rb
@@ -11,7 +11,7 @@ describe "Fiber.new(storage:)" do
     end
 
     it "creates a fiber with lazily initialized storage" do
-      Fiber.new(storage: nil) { Fiber.current.storage }.resume.should == {}
+      Fiber.new(storage: nil) { Fiber[:x] = 10; Fiber.current.storage }.resume.should == {x: 10}
     end
 
     it "creates a fiber by inheriting the storage of the parent fiber" do
@@ -30,18 +30,19 @@ end
 describe "Fiber#storage=" do
   ruby_version_is "3.2" do
     it "can clear the storage of the fiber" do
-      fiber = Fiber.new(storage: {life: 42}) {
+      fiber = Fiber.new(storage: {life: 42}) do
         Fiber.current.storage = nil
+        Fiber[:x] = 10
         Fiber.current.storage
-      }
-      fiber.resume.should == {}
+      end
+      fiber.resume.should == {x: 10}
     end
 
     it "can set the storage of the fiber" do
-      fiber = Fiber.new(storage: {life: 42}) {
+      fiber = Fiber.new(storage: {life: 42}) do
         Fiber.current.storage = {life: 43}
         Fiber.current.storage
-      }
+      end
       fiber.resume.should == {life: 43}
     end
 
@@ -89,6 +90,12 @@ describe "Fiber.[]=" do
       Fiber.new { Fiber[:life] = 43; Fiber[:life] }.resume.should == 43
     end
   end
+
+  ruby_version_is "3.3" do
+    it "deletes the fiber storage key when assigning nil" do
+      Fiber.new(storage: {life: 42}) { Fiber[:life] = nil; Fiber.current.storage }.resume.should == {}
+    end
+  end
 end
 
 describe "Thread.new" do