Improved error handling and support for COS. (#374)

6 files changed, 28 insertions, 14 deletions

diff --git a/README.md b/README.md
index 940e456..8e1cc3b 100644
--- a/README.md
+++ b/README.md
@@ -247,6 +247,7 @@ InstanceSetup     | set_boto_config      | `false` skips setting up a boto confi
 InstanceSetup     | set_host_keys        | `false` skips generating host keys on first boot.
 InstanceSetup     | set_multiqueue       | `false` skips multiqueue driver support.
 IpForwarding      | ethernet_proto_id    | Protocol ID string for daemon added routes.
+MetadataScripts   | run_dir              | String base directory where metadata scripts are executed.
 MetadataScripts   | startup              | `false` disables startup script execution.
 MetadataScripts   | shutdown             | `false` disables shutdown script execution.
 NetworkInterfaces | dhcp_command         | String to execute to enable network interfaces.
diff --git a/google_compute_engine/accounts/accounts_utils.py b/google_compute_engine/accounts/accounts_utils.py
index fbdd5ef..09b1761 100644
--- a/google_compute_engine/accounts/accounts_utils.py
+++ b/google_compute_engine/accounts/accounts_utils.py
@@ -167,6 +167,8 @@ class AccountsUtils(object):
 
     Raises:
       IOError, raised when there is an exception updating a file.
+      OSError, raised when setting permissions or writing to a read-only
+          file system.
     """
     pw_entry = self._GetUser(user)
     if not pw_entry:
@@ -302,7 +304,7 @@ class AccountsUtils(object):
 
     try:
       self._UpdateAuthorizedKeys(user, ssh_keys)
-    except IOError as e:
+    except (IOError, OSError) as e:
       message = 'Could not update the authorized keys file for user %s. %s.'
       self.logger.warning(message, user, str(e))
       return False
diff --git a/google_compute_engine/instance_setup/instance_config.py b/google_compute_engine/instance_setup/instance_config.py
index 675db3a..10e7c71 100644
--- a/google_compute_engine/instance_setup/instance_config.py
+++ b/google_compute_engine/instance_setup/instance_config.py
@@ -62,6 +62,7 @@ class InstanceConfig(config_manager.ConfigManager):
           'ethernet_proto_id': '66',
       },
       'MetadataScripts': {
+          'run_dir': '',
           'startup': 'true',
           'shutdown': 'true',
       },
diff --git a/google_compute_engine/metadata_scripts/script_manager.py b/google_compute_engine/metadata_scripts/script_manager.py
index 2993c36..5cd0ebe 100755
--- a/google_compute_engine/metadata_scripts/script_manager.py
+++ b/google_compute_engine/metadata_scripts/script_manager.py
@@ -29,16 +29,17 @@ from google_compute_engine.metadata_scripts import script_retriever
 
 
 @contextlib.contextmanager
-def _CreateTempDir(prefix):
+def _CreateTempDir(prefix, run_dir=None):
   """Context manager for creating a temporary directory.
 
   Args:
     prefix: string, the prefix for the temporary directory.
+    run_dir: string, the base directory location of the temporary directory.
 
   Yields:
     string, the temporary directory created.
   """
-  temp_dir = tempfile.mkdtemp(prefix=prefix + '-')
+  temp_dir = tempfile.mkdtemp(prefix=prefix + '-', dir=run_dir)
   try:
     yield temp_dir
   finally:
@@ -48,11 +49,12 @@ def _CreateTempDir(prefix):
 class ScriptManager(object):
   """A class for retrieving and executing metadata scripts."""
 
-  def __init__(self, script_type, debug=False):
+  def __init__(self, script_type, run_dir=None, debug=False):
     """Constructor.
 
     Args:
       script_type: string, the metadata script type to run.
+      run_dir: string, the base directory location of the temporary directory.
       debug: bool, True if debug output should write to the console.
     """
     self.script_type = script_type
@@ -61,11 +63,15 @@ class ScriptManager(object):
     self.logger = logger.Logger(name=name, debug=debug, facility=facility)
     self.retriever = script_retriever.ScriptRetriever(self.logger, script_type)
     self.executor = script_executor.ScriptExecutor(self.logger, script_type)
-    self._RunScripts()
+    self._RunScripts(run_dir=run_dir)
 
-  def _RunScripts(self):
-    """Retrieve metadata scripts and execute them."""
-    with _CreateTempDir(self.script_type) as dest_dir:
+  def _RunScripts(self, run_dir=None):
+    """Retrieve metadata scripts and execute them.
+
+    Args:
+      run_dir: string, the base directory location of the temporary directory.
+    """
+    with _CreateTempDir(self.script_type, run_dir=run_dir) as dest_dir:
       try:
         self.logger.info('Starting %s scripts.', self.script_type)
         script_dict = self.retriever.GetScripts(dest_dir)
@@ -92,7 +98,10 @@ def main():
 
   instance_config = config_manager.ConfigManager()
   if instance_config.GetOptionBool('MetadataScripts', script_type):
-    ScriptManager(script_type, debug=bool(options.debug))
+    ScriptManager(
+        script_type,
+        run_dir=instance_config.GetOptionString('MetadataScripts', 'run_dir'),
+        debug=bool(options.debug))
 
 
 if __name__ == '__main__':
diff --git a/google_compute_engine/metadata_scripts/tests/script_manager_test.py b/google_compute_engine/metadata_scripts/tests/script_manager_test.py
index 83fa879..ec7d61f 100644
--- a/google_compute_engine/metadata_scripts/tests/script_manager_test.py
+++ b/google_compute_engine/metadata_scripts/tests/script_manager_test.py
@@ -40,6 +40,7 @@ class ScriptManagerTest(unittest.TestCase):
     mocks.attach_mock(mock_executor, 'executor')
     mocks.attach_mock(mock_logger, 'logger')
     mocks.attach_mock(mock_retriever, 'retriever')
+    run_dir = '/var/run'
     script_type = 'test'
     script_name = '%s-script' % script_type
     script_prefix = '%s-' % script_type
@@ -48,13 +49,13 @@ class ScriptManagerTest(unittest.TestCase):
     mock_mkdir.return_value = test_dir
     mock_retriever_instance.GetScripts.return_value = test_dict
 
-    script_manager.ScriptManager(script_type)
+    script_manager.ScriptManager(script_type, run_dir=run_dir)
     expected_calls = [
         mock.call.logger.Logger(
             name=script_name, debug=False, facility=mock.ANY),
         mock.call.retriever.ScriptRetriever(mock_logger_instance, script_type),
         mock.call.executor.ScriptExecutor(mock_logger_instance, script_type),
-        mock.call.mkdir(prefix=script_prefix),
+        mock.call.mkdir(prefix=script_prefix, dir=run_dir),
         mock.call.logger.Logger().info(mock.ANY, script_type),
         mock.call.retriever.ScriptRetriever().GetScripts(test_dir),
         mock.call.executor.ScriptExecutor().RunScripts(test_dict),
diff --git a/scripts/set_multiqueue b/scripts/set_multiqueue
index 5c722e3..636ee90 100755
--- a/scripts/set_multiqueue
+++ b/scripts/set_multiqueue
@@ -39,13 +39,13 @@ echo "Running $(basename "$0")."
 NET_DEVS=/sys/bus/virtio/drivers/virtio_net/virtio*
 
 # Loop through all the virtionet devices and enable multi-queue
-if [ -x /sbin/ethtool ]; then
+if [ -x "$(command -v ethtool)" ]; then
   for dev in $NET_DEVS; do
     ETH_DEVS=${dev}/net/*
     for eth_dev in $ETH_DEVS; do
       eth_dev=$(basename "$eth_dev")
       if ! errormsg=$(ethtool -l "$eth_dev" 2>&1); then
-        echo "/sbin/ethtool says that $eth_dev does not support virtionet multiqueue: $errormsg."
+        echo "ethtool says that $eth_dev does not support virtionet multiqueue: $errormsg."
         continue
       fi
       num_max_channels=$(ethtool -l "$eth_dev" | grep -m 1 Combined | cut -f2)
@@ -59,7 +59,7 @@ if [ -x /sbin/ethtool ]; then
     done
   done
 else
-  echo "/sbin/ethtool not found: cannot configure virtionet multiqueue."
+  echo "ethtool not found: cannot configure virtionet multiqueue."
 fi
 
 for dev in $NET_DEVS