Merge pull request #557 from ivoanjo/thread-local-block-initialization

Allow ThreadLocalVar to be initialized with a block

Author: jdantonio

lib/concurrent/atomic/abstract_thread_local_var.rb

@@ -8,8 +8,17 @@ module Concurrent
   class AbstractThreadLocalVar
 
     # @!macro thread_local_var_method_initialize
-    def initialize(default = nil)
-      @default = default
+    def initialize(default = nil, &default_block)
+      if default && block_given?
+        raise ArgumentError, "Cannot use both value and block as default value"
+      end
+
+      if block_given?
+        @default_block = default_block
+      else
+        @default = default
+      end
+
       allocate_storage
     end
 
@@ -25,7 +34,15 @@ def value=(value)
 
     # @!macro thread_local_var_method_bind
     def bind(value, &block)
-      raise NotImplementedError
+      if block_given?
+        old_value = self.value
+        begin
+          self.value = value
+          yield
+        ensure
+          self.value = old_value
+        end
+      end
     end
 
     protected
@@ -34,5 +51,14 @@ def bind(value, &block)
     def allocate_storage
       raise NotImplementedError
     end
+
+    # @!visibility private
+    def default
+      if @default_block
+        self.value = @default_block.call
+      else
+        @default
+      end
+    end
   end
 end

lib/concurrent/atomic/java_thread_local_var.rb

@@ -13,7 +13,7 @@ def value
         value = @var.get
 
         if value.nil?
-          @default
+          default
         elsif value == NULL
           nil
         else
@@ -26,19 +26,6 @@ def value=(value)
         @var.set(value)
       end
 
-      # @!macro thread_local_var_method_bind
-      def bind(value, &block)
-        if block_given?
-          old_value = @var.get
-          begin
-            @var.set(value)
-            yield
-          ensure
-            @var.set(old_value)
-          end
-        end
-      end
-
       protected
 
       # @!visibility private

lib/concurrent/atomic/ruby_thread_local_var.rb

@@ -35,29 +35,19 @@ class RubyThreadLocalVar < AbstractThreadLocalVar
     @@next = 0
     private_constant :FREE, :LOCK, :ARRAYS
 
-    # @!macro [attach] thread_local_var_method_initialize
-    #
-    #   Creates a thread local variable.
-    #
-    #   @param [Object] default the default value when otherwise unset
-    def initialize(default = nil)
-      @default = default
-      allocate_storage
-    end
-
     # @!macro thread_local_var_method_get
     def value
       if array = get_threadlocal_array
         value = array[@index]
         if value.nil?
-          @default
+          default
         elsif value.equal?(NULL)
           nil
         else
           value
         end
       else
-        @default
+        default
       end
     end
 
@@ -76,19 +66,6 @@ def value=(value)
       value
     end
 
-    # @!macro thread_local_var_method_bind
-    def bind(value, &block)
-      if block_given?
-        old_value = self.value
-        begin
-          self.value = value
-          yield
-        ensure
-          self.value = old_value
-        end
-      end
-    end
-
     protected
 
     # @!visibility private
@@ -158,12 +135,20 @@ def value_for(thread)
       if array = get_threadlocal_array(thread)
         value = array[@index]
         if value.nil?
-          @default
+          default_for(thread)
         elsif value.equal?(NULL)
           nil
         else
           value
         end
+      else
+        default_for(thread)
+      end
+    end
+
+    def default_for(thread)
+      if @default_block
+        raise "Cannot use default_for with default block"
       else
         @default
       end

lib/concurrent/atomic/thread_local_var.rb

@@ -11,6 +11,8 @@ module Concurrent
   #   Creates a thread local variable.
   #
   #   @param [Object] default the default value when otherwise unset
+  #   @param [Proc] block Optional block that gets called to obtain the
+  #     default value for each thread
 
   # @!macro [new] thread_local_var_method_get
   #

spec/concurrent/atomic/thread_local_var_spec.rb

@@ -6,22 +6,20 @@ module Concurrent
 
   describe ThreadLocalVar do
 
-    subject { ThreadLocalVar.new }
-
     context '#initialize' do
 
       it 'can set an initial value' do
-        v = ThreadLocalVar.new(14)
+        v = described_class.new(14)
         expect(v.value).to eq 14
       end
 
       it 'sets nil as a default initial value' do
-        v = ThreadLocalVar.new
+        v = described_class.new
         expect(v.value).to be_nil
       end
 
       it 'sets the same initial value for all threads' do
-        v  = ThreadLocalVar.new(14)
+        v  = described_class.new(14)
         t1 = Thread.new { v.value }
         t2 = Thread.new { v.value }
         expect(t1.value).to eq 14
@@ -30,51 +28,88 @@ module Concurrent
 
       if Concurrent.on_jruby?
         it 'extends JavaThreadLocalVar' do
-          expect(subject.class.ancestors).to include(Concurrent::JavaThreadLocalVar)
+          expect(described_class.ancestors).to include(Concurrent::JavaThreadLocalVar)
         end
       else
         it 'extends RubyThreadLocalVar' do
-          expect(subject.class.ancestors).to include(Concurrent::RubyThreadLocalVar)
+          expect(described_class.ancestors).to include(Concurrent::RubyThreadLocalVar)
+        end
+      end
+
+      it 'can set a block to be called to get the initial value' do
+        v = described_class.new { 14 }
+        expect(v.value).to eq 14
+      end
+
+      context 'when attempting to set both an initial value and a block' do
+        it do
+          expect { described_class.new(14) { 14 } }.to raise_error(ArgumentError)
         end
       end
     end
 
     context '#value' do
+      let(:v) { described_class.new(14) }
 
       it 'returns the current value' do
-        v = ThreadLocalVar.new(14)
         expect(v.value).to eq 14
       end
 
       it 'returns the value after modification' do
-        v = ThreadLocalVar.new(14)
         v.value = 2
         expect(v.value).to eq 2
       end
+
+      context 'when using a block to initialize the value' do
+        it 'calls the block to initialize the value' do
+          block = proc { }
+
+          expect(block).to receive(:call)
+
+          v = described_class.new(&block)
+          v.value
+        end
+
+        it 'sets the block return value as the current value' do
+          value = 13
+
+          v = described_class.new { value += 1 }
+
+          v.value
+          expect(v.value).to be 14
+        end
+
+        it 'calls the block to initialize the value for each thread' do
+          block = proc { }
+
+          expect(block).to receive(:call).twice
+
+          v = described_class.new(&block)
+          Thread.new { v.value }.join
+          Thread.new { v.value }.join
+        end
+      end
     end
 
     context '#value=' do
+      let(:v) { described_class.new(14) }
 
       it 'sets a new value' do
-        v = ThreadLocalVar.new(14)
         v.value = 2
         expect(v.value).to eq 2
       end
 
       it 'returns the new value' do
-        v = ThreadLocalVar.new(14)
         expect(v.value = 2).to eq 2
       end
 
       it 'does not modify the initial value for other threads' do
-        v = ThreadLocalVar.new(14)
         v.value = 2
         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.value = 2
 
         b1 = CountDownLatch.new(2)