Fixing ThreadLocalVar to use 'ref' gem

ThreadLocalVar needs to be required manually.

Author: pitr-ch

lib/concurrent/atomic/thread_local_var.rb

@@ -2,71 +2,60 @@
 
 module Concurrent
 
-  module ThreadLocalRubyStorage
+  class AbstractThreadLocalVar
 
-    protected
+    module ThreadLocalRubyStorage
 
-    def allocate_storage
-      @storage = Atomic.new Hash.new
-    end
+      protected
 
-    def get
-      @storage.get[Thread.current.object_id]
-    end
+      begin
+        require 'ref'
+      rescue LoadError
+        raise LoadError,
+              'ThreadLocalVar requires ref gem installed on MRI. It\'s not concurrent-ruby dependency.'
+      end
 
-    def set(value, &block)
-      key = Thread.current.object_id
+      def allocate_storage
+        @storage = Ref::WeakKeyMap.new
+      end
 
-      @storage.update do |s|
-        s.merge(key => value)
+      def get
+        @storage[Thread.current]
       end
 
-      if block_given?
-        begin
-          block.call
-        ensure
-          @storage.update do |s|
-            s.clone.tap { |h| h.delete key }
-          end
-        end
+      def set(value, &block)
+        key = Thread.current
 
-      else
-        unless ThreadLocalRubyStorage.i_know_it_may_leak_values?
-          warn "it may leak values if used without block\n#{caller[0]}"
+        @storage[key] = value
+
+        if block_given?
+          begin
+            block.call
+          ensure
+            @storage.delete key
+          end
         end
       end
     end
 
-    def self.i_know_it_may_leak_values!
-      @leak_acknowledged = true
-    end
-
-    def self.i_know_it_may_leak_values?
-      @leak_acknowledged
-    end
-
-  end
+    module ThreadLocalJavaStorage
 
-  module ThreadLocalJavaStorage
+      protected
 
-    protected
+      def allocate_storage
+        @var = java.lang.ThreadLocal.new
+      end
 
-    def allocate_storage
-      @var = java.lang.ThreadLocal.new
-    end
+      def get
+        @var.get
+      end
 
-    def get
-      @var.get
-    end
+      def set(value)
+        @var.set(value)
+      end
 
-    def set(value)
-      @var.set(value)
     end
 
-  end
-
-  class AbstractThreadLocalVar
-
     NIL_SENTINEL = Object.new
 
     def initialize(default = nil)

lib/concurrent/atomics.rb

@@ -7,5 +7,4 @@
 require 'concurrent/atomic/cyclic_barrier'
 require 'concurrent/atomic/count_down_latch'
 require 'concurrent/atomic/event'
-require 'concurrent/atomic/thread_local_var'
 require 'concurrent/atomic/synchronization'

spec/concurrent/atomic/thread_local_var_spec.rb

@@ -3,7 +3,7 @@
 
 module Concurrent
 
-  ThreadLocalRubyStorage.i_know_it_may_leak_values!
+  require 'concurrent/atomic/thread_local_var'
 
   describe ThreadLocalVar do
 
@@ -31,11 +31,11 @@ module Concurrent
 
       if jruby?
         it 'uses ThreadLocalJavaStorage' do
-          expect(subject.class.ancestors).to include(Concurrent::ThreadLocalJavaStorage)
+          expect(subject.class.ancestors).to include(Concurrent::AbstractThreadLocalVar::ThreadLocalJavaStorage)
         end
       else
         it 'uses ThreadLocalNewStorage' do
-          expect(subject.class.ancestors).to include(Concurrent::ThreadLocalRubyStorage)
+          expect(subject.class.ancestors).to include(Concurrent::AbstractThreadLocalVar::ThreadLocalRubyStorage)
         end
       end
     end
@@ -48,7 +48,20 @@ module Concurrent
             Thread.new { var.bind(i) { var.value } }
           end.each(&:join)
           var.value = 0
-          expect(var.instance_variable_get(:@storage).get.size).to be == 1
+          expect(var.instance_variable_get(:@storage).keys.size).to be == 1
+        end
+
+        it 'does not leave values behind when bind is not used' do
+          var = ThreadLocalVar.new(0)
+          100.times.map do |i|
+            Thread.new { var.value = i; var.value }
+          end.each(&:join)
+          var.value = 0
+          sleep 0.1
+          GC.start
+          sleep 0.1
+          GC.start
+          expect(var.instance_variable_get(:@storage).keys.size).to be == 1
         end
       end
     end