Disable overcommit accounting on E2 VMs (#865)

* Disable overcommit accounting on E2 VMs

2 files changed, 16 insertions, 7 deletions

diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
index cb1a2a6..dbd3d8f 100755
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/instance_setup.py
@@ -28,7 +28,7 @@ from google_compute_engine import file_utils
 from google_compute_engine import logger
 from google_compute_engine import metadata_watcher
 from google_compute_engine.boto import boto_config
-from google_compute_engine.compat import urlerror
+from google_compute_engine.compat import distro_name, urlerror
 from google_compute_engine.compat import urlrequest
 from google_compute_engine.instance_setup import instance_config
 
@@ -65,17 +65,27 @@ class InstanceSetup(object):
       instance_config_metadata = self._GetInstanceConfig()
       self.instance_config = instance_config.InstanceConfig(
           logger=self.logger, instance_config_metadata=instance_config_metadata)
+
       if self.instance_config.GetOptionBool('InstanceSetup', 'set_host_keys'):
         host_key_types = self.instance_config.GetOptionString(
             'InstanceSetup', 'host_key_types')
         self._SetSshHostKeys(host_key_types=host_key_types)
+
       if self.instance_config.GetOptionBool('InstanceSetup', 'set_boto_config'):
         self._SetupBotoConfig()
+
+      # machineType is e.g. u'projects/00000000000000/machineTypes/n1-standard-1'
+      machineType = self.metadata_dict['instance']['machineType'].split('/')[-1]
+      if machineType.startswith("e2-") and 'bsd' not in distro_name:  # Not yet supported on BSD.
+        subprocess.call(["sysctl", "vm.overcommit_memory=1"])
+
     if self.instance_config.GetOptionBool(
         'InstanceSetup', 'optimize_local_ssd'):
       self._RunScript('google_optimize_local_ssd')
+
     if self.instance_config.GetOptionBool('InstanceSetup', 'set_multiqueue'):
       self._RunScript('google_set_multiqueue')
+
     try:
       self.instance_config.WriteConfig()
     except (IOError, OSError) as e:
@@ -205,10 +215,9 @@ class InstanceSetup(object):
     Args:
       host_key_types: string, a comma separated list of host key types.
     """
-    section = 'Instance'
     instance_id = self._GetInstanceId()
     if instance_id != self.instance_config.GetOptionString(
-        section, 'instance_id'):
+        'instance', 'instance_id'):
       self.logger.info('Generating SSH host keys for instance %s.', instance_id)
       file_regex = re.compile(r'ssh_host_(?P<type>[a-z0-9]*)_key\Z')
       key_dir = '/etc/ssh'
@@ -222,7 +231,7 @@ class InstanceSetup(object):
         if key_data:
           self._WriteHostKeyToGuestAttributes(key_data[0], key_data[1])
       self._StartSshd()
-      self.instance_config.SetOption(section, 'instance_id', str(instance_id))
+      self.instance_config.SetOption('Instance', 'instance_id', str(instance_id))
 
   def _GetNumericProjectId(self):
     """Get the numeric project ID.
diff --git a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
index a418219..e255f86 100644
--- a/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
+++ b/packages/python-google-compute-engine/google_compute_engine/instance_setup/tests/instance_setup_test.py
@@ -46,7 +46,7 @@ class InstanceSetupTest(unittest.TestCase):
     mock_logger_instance = mock.Mock()
     mock_logger.Logger.return_value = mock_logger_instance
     mock_watcher_instance = mock.Mock()
-    mock_watcher_instance.GetMetadata.return_value = {'hello': 'world'}
+    mock_watcher_instance.GetMetadata.return_value = {'hello': 'world', 'instance': {'machineType': 'fake'}}
     mock_watcher.MetadataWatcher.return_value = mock_watcher_instance
     mock_config_instance = mock.Mock()
     mock_config_instance.GetOptionBool.return_value = True
@@ -91,7 +91,7 @@ class InstanceSetupTest(unittest.TestCase):
         mock.call.config.InstanceConfig().WriteConfig(),
     ]
     self.assertEqual(mocks.mock_calls, expected_calls)
-    self.assertEqual(mock_setup.metadata_dict, {'hello': 'world'})
+    self.assertEqual(mock_setup.metadata_dict, {'hello': 'world', 'instance': {'machineType': 'fake'}})
 
   @mock.patch('google_compute_engine.instance_setup.instance_setup.instance_config')
   @mock.patch('google_compute_engine.instance_setup.instance_setup.metadata_watcher')
@@ -386,7 +386,7 @@ class InstanceSetupTest(unittest.TestCase):
     ]
 
     self.assertEqual(sorted(mock_generate_key.mock_calls), expected_calls)
-    self.mock_instance_config.SetOption.assert_called_once_with(
+    self.mock_instance_config.SetOption.assert_called_with(
         'Instance', 'instance_id', '123')
 
   def testGetNumericProjectId(self):