Merge pull request #164 from ruby-concurrency/thread-local-var

Fix Thread leaking in ThreadLocalVariables

Tested the exclusion of the `Ref` gem in the pre-compiled Java build (using our automated build process) and it worked as intended. The `Ref` gem will not be installed under JRuby when using the precompiled Java build.

Author: jdantonio

concurrent-ruby.gemspec

@@ -31,4 +31,8 @@ Gem::Specification.new do |s|
   end
 
   s.required_ruby_version = '>= 1.9.3'
+
+  unless defined?(JRUBY_VERSION)
+    s.add_dependency 'ref', '~> 1.0.5'
+  end
 end

lib/concurrent/atomic/thread_local_var.rb

@@ -2,41 +2,56 @@
 
 module Concurrent
 
-  module ThreadLocalRubyStorage
+  class AbstractThreadLocalVar
 
-    def allocate_storage
-      @storage = Atomic.new Hash.new
-    end
+    module ThreadLocalRubyStorage
 
-    def get
-      @storage.get[Thread.current]
-    end
+      protected
 
-    def set(value)
-      @storage.update { |s| s.merge Thread.current => value }
-    end
+      unless RUBY_PLATFORM == 'java'
+        require 'ref'
+      end
 
-  end
+      def allocate_storage
+        @storage = Ref::WeakKeyMap.new
+      end
 
-  module ThreadLocalJavaStorage
+      def get
+        @storage[Thread.current]
+      end
 
-    protected
+      def set(value, &block)
+        key = Thread.current
 
-    def allocate_storage
-      @var = java.lang.ThreadLocal.new
-    end
+        @storage[key] = value
 
-    def get
-      @var.get
+        if block_given?
+          begin
+            block.call
+          ensure
+            @storage.delete key
+          end
+        end
+      end
     end
 
-    def set(value)
-      @var.set(value)
-    end
+    module ThreadLocalJavaStorage
 
-  end
+      protected
 
-  class AbstractThreadLocalVar
+      def allocate_storage
+        @var = java.lang.ThreadLocal.new
+      end
+
+      def get
+        @var.get
+      end
+
+      def set(value)
+        @var.set(value)
+      end
+
+    end
 
     NIL_SENTINEL = Object.new
 
@@ -58,13 +73,17 @@ def value
     end
 
     def value=(value)
+      bind value
+    end
+
+    def bind(value, &block)
       if value.nil?
         stored_value = NIL_SENTINEL
       else
         stored_value = value
       end
 
-      set stored_value
+      set stored_value, &block
 
       value
     end

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,9 +3,11 @@
 
 module Concurrent
 
+  require 'concurrent/atomic/thread_local_var'
+
   describe ThreadLocalVar do
 
-    subject{ ThreadLocalVar.new }
+    subject { ThreadLocalVar.new }
 
     context '#initialize' do
 
@@ -20,7 +22,7 @@ module Concurrent
       end
 
       it 'sets the same initial value for all threads' do
-        v = ThreadLocalVar.new(14)
+        v  = ThreadLocalVar.new(14)
         t1 = Thread.new { v.value }
         t2 = Thread.new { v.value }
         expect(t1.value).to eq 14
@@ -29,11 +31,37 @@ 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
+
+    unless jruby?
+      context 'GC' do
+        it 'does not leave values behind when bind is used' do
+          var = ThreadLocalVar.new(0)
+          100.times.map do |i|
+            Thread.new { var.bind(i) { var.value } }
+          end.each(&:join)
+          var.value = 0
+          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
@@ -46,7 +74,7 @@ module Concurrent
       end
 
       it 'returns the value after modification' do
-        v = ThreadLocalVar.new(14)
+        v       = ThreadLocalVar.new(14)
         v.value = 2
         expect(v.value).to eq 2
       end
@@ -56,7 +84,7 @@ module Concurrent
     context '#value=' do
 
       it 'sets a new value' do
-        v = ThreadLocalVar.new(14)
+        v       = ThreadLocalVar.new(14)
         v.value = 2
         expect(v.value).to eq 2
       end
@@ -67,14 +95,14 @@ module Concurrent
       end
 
       it 'does not modify the initial value for other threads' do
-        v = ThreadLocalVar.new(14)
+        v       = ThreadLocalVar.new(14)
         v.value = 2
-        t = Thread.new { v.value }
+        t       = Thread.new { v.value }
         expect(t.value).to eq 14
       end
 
       it 'does not modify the value for other threads' do
-        v = ThreadLocalVar.new(14)
+        v       = ThreadLocalVar.new(14)
         v.value = 2
 
         b1 = CountDownLatch.new(2)