support configurable concurrency

3 files changed, 29 insertions, 4 deletions

diff --git a/lib/chef/cookbook_uploader.rb b/lib/chef/cookbook_uploader.rb
index acd913e72c..2ae4703f01 100644
--- a/lib/chef/cookbook_uploader.rb
+++ b/lib/chef/cookbook_uploader.rb
@@ -17,10 +17,10 @@ class Chef
       @work_queue ||= Queue.new
     end
 
-    def self.setup_worker_threads
+    def self.setup_worker_threads(concurrency=10)
       @worker_threads ||= begin
         work_queue
-        (1...10).map do
+        (1...concurrency).map do
           Thread.new do
             loop do
               work_queue.pop.call
@@ -34,6 +34,7 @@ class Chef
     attr_reader :path
     attr_reader :opts
     attr_reader :rest
+    attr_reader :concurrency
 
     # Creates a new CookbookUploader.
     # ===Arguments:
@@ -50,10 +51,13 @@ class Chef
     # * :rest   A Chef::REST object that you have configured the way you like it.
     #           If you don't provide this, one will be created using the values
     #           in Chef::Config.
+    # * :concurrency   An integer that decided how many threads will be used to
+    #           perform concurrent uploads
     def initialize(cookbooks, path, opts={})
       @path, @opts = path, opts
       @cookbooks = Array(cookbooks)
       @rest = opts[:rest] || Chef::REST.new(Chef::Config[:chef_server_url])
+      @concurrency = opts[:concurrency] || 10
     end
 
     def upload_cookbooks
@@ -73,7 +77,7 @@ class Chef
 
       Chef::Log.info("Uploading files")
 
-      self.class.setup_worker_threads
+      self.class.setup_worker_threads(concurrency)
 
       checksums_to_upload = Set.new
 
diff --git a/lib/chef/knife/cookbook_upload.rb b/lib/chef/knife/cookbook_upload.rb
index d4a66e6925..a882cd7109 100644
--- a/lib/chef/knife/cookbook_upload.rb
+++ b/lib/chef/knife/cookbook_upload.rb
@@ -57,6 +57,12 @@ class Chef
         :boolean => true,
         :description => "Update cookbook versions even if they have been frozen"
 
+      option :concurrency,
+        :long => '--concurrency NUMBER_OF_THREADS',
+        :description => "How many concurrent threads will be used",
+        :default => 10,
+        :proc => lambda { |o| o.to_i }
+
       option :environment,
         :short => '-E',
         :long  => '--environment ENVIRONMENT',
@@ -228,7 +234,7 @@ WARNING
           check_for_broken_links!(cb)
           check_for_dependencies!(cb)
         end
-        Chef::CookbookUploader.new(cookbooks, config[:cookbook_path], :force => config[:force]).upload_cookbooks
+        Chef::CookbookUploader.new(cookbooks, config[:cookbook_path], :force => config[:force], :concurrency => config[:concurrency]).upload_cookbooks
       rescue Chef::Exceptions::CookbookFrozen => e
         ui.error e
         raise
diff --git a/spec/unit/knife/cookbook_upload_spec.rb b/spec/unit/knife/cookbook_upload_spec.rb
index d9f1b8f1ce..302670278d 100644
--- a/spec/unit/knife/cookbook_upload_spec.rb
+++ b/spec/unit/knife/cookbook_upload_spec.rb
@@ -21,6 +21,7 @@ require File.expand_path(File.join(File.dirname(__FILE__), "..", "..", "spec_hel
 
 require 'chef/cookbook_uploader'
 require 'timeout'
+require 'ostruct'
 
 describe Chef::Knife::CookbookUpload do
   before(:each) do
@@ -40,6 +41,20 @@ describe Chef::Knife::CookbookUpload do
     @knife.ui.stub!(:stderr).and_return(@output)
   end
 
+  describe 'with --concurrency' do
+    it 'should upload cookbooks with predefined concurrency' do
+      @cookbook_uploader = stub(:upload_cookbooks => nil)
+      Chef::CookbookVersion.stub(:list_all_versions).and_return({})
+      @knife.config[:concurrency] = 3
+      @test_cookbook = Chef::CookbookVersion.new('test_cookbook')
+      @cookbook_loader.stub!(:each).and_yield("test_cookbook", @test_cookbook)
+      @cookbook_loader.stub!(:cookbook_names).and_return(["test_cookbook"])
+      Chef::CookbookUploader.should_receive(:new).with( kind_of(Array),  kind_of(Array),
+        {:force=>nil, :concurrency => 3}).and_return(OpenStruct.new({:upload_cookbooks=> true}))
+      @knife.run
+    end
+  end
+
   describe 'run' do
     before(:each) do
       @cookbook_uploader = stub(:upload_cookbooks => nil)