Fix Thread leaking in ThreadLocalVariables

- introduce #bind method to avoid value leaks
- warn on using just #value=

Author: pitr-ch

lib/concurrent/atomic/thread_local_var.rb

@@ -4,16 +4,45 @@ module Concurrent
 
   module ThreadLocalRubyStorage
 
+    protected
+
     def allocate_storage
       @storage = Atomic.new Hash.new
     end
 
     def get
-      @storage.get[Thread.current]
+      @storage.get[Thread.current.object_id]
     end
 
-    def set(value)
-      @storage.update { |s| s.merge Thread.current => value }
+    def set(value, &block)
+      key = Thread.current.object_id
+
+      @storage.update do |s|
+        s.merge(key => value)
+      end
+
+      if block_given?
+        begin
+          block.call
+        ensure
+          @storage.update do |s|
+            s.clone.tap { |h| h.delete key }
+          end
+        end
+
+      else
+        unless ThreadLocalRubyStorage.i_know_it_may_leak_values?
+          warn "it may leak values if used without block\n#{caller[0]}"
+        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
@@ -57,14 +86,19 @@ def value
       end
     end
 
+    # may leak the value, #bind is preferred
     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

spec/concurrent/atomic/thread_local_var_spec.rb

@@ -3,9 +3,11 @@
 
 module Concurrent
 
+  ThreadLocalRubyStorage.i_know_it_may_leak_values!
+
   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
@@ -38,6 +40,17 @@ module Concurrent
       end
     end
 
+    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).get.size).to be == 1
+      end
+    end
+
     context '#value' do
 
       it 'returns the current value' do
@@ -46,7 +59,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 +69,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 +80,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)