PL-433 retry on 500 concurrency fixes

2 files changed, 20 insertions, 16 deletions

diff --git a/chef/lib/chef/config.rb b/chef/lib/chef/config.rb
index 5365518713..9e72b7ba15 100644
--- a/chef/lib/chef/config.rb
+++ b/chef/lib/chef/config.rb
@@ -100,6 +100,7 @@ class Chef
 
     authorized_openid_identifiers nil
     authorized_openid_providers nil
+    client_registration_retries 5
     cookbook_path [ "/var/chef/cookbooks", "/var/chef/site-cookbooks" ]
     cookbook_tarball_path "/var/chef/cookbook-tarballs"
     couchdb_database "chef"
diff --git a/chef/lib/chef/rest.rb b/chef/lib/chef/rest.rb
index c491e7a3c5..50b3f10074 100644
--- a/chef/lib/chef/rest.rb
+++ b/chef/lib/chef/rest.rb
@@ -66,26 +66,29 @@ class Chef
     
     # Register the client 
     def register(name=Chef::Config[:node_name], destination=Chef::Config[:client_key])
-
-      if File.exists?(destination)
-        raise Chef::Exceptions::CannotWritePrivateKey, "I cannot write your private key to #{destination} - check permissions?" unless File.writable?(destination)
-      end
+      raise Chef::Exceptions::CannotWritePrivateKey, "I cannot write your private key to #{destination} - check permissions?" if (File.exists?(destination) &&  !File.writable?(destination))
 
       nc = Chef::ApiClient.new
       nc.name(name)
-      response = nc.save(true, true)
-
-      Chef::Log.debug("Registration response: #{response.inspect}")
-
-      raise Chef::Exceptions::CannotWritePrivateKey, "The response from the server did not include a private key!" unless response.has_key?("private_key")
 
-      begin
-        # Write out the private key
-        file = File.open(destination, File::WRONLY|File::EXCL|File::CREAT, 0600)
-        file.print(response["private_key"])
-        file.close
-      rescue 
-        raise Chef::Exceptions::CannotWritePrivateKey, "I cannot write your private key to #{destination}"
+      catch(:done) do
+        retries = Chef::Config[:client_registration_retries] || 5
+        retries.downto(0) do
+          begin
+            response = nc.save(true, true)
+            Chef::Log.debug("Registration response: #{response.inspect}")
+            raise Chef::Exceptions::CannotWritePrivateKey, "The response from the server did not include a private key!" unless response.has_key?("private_key")
+            # Write out the private key
+            file = File.open(destination, "w")
+            file.print(response["private_key"])
+            file.close
+            throw :done
+          rescue IOError
+            raise Chef::Exceptions::CannotWritePrivateKey, "I cannot write your private key to #{destination}"
+          rescue Net::HTTPFatalError => e
+            raise unless e.response.code == "500"
+          end
+        end
       end
 
       true