Disable overcommit accounting on E2 VMs (again) (#874)

2 files changed, 57 insertions, 2 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..efe100b 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,6 +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 distro_name
 from google_compute_engine.compat import urlerror
 from google_compute_engine.compat import urlrequest
 from google_compute_engine.instance_setup import instance_config
@@ -65,22 +66,38 @@ 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()
+
+      self._DisableOvercommit()
+
     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:
       self.logger.warning(str(e))
 
+  def _DisableOvercommit(self, distro=distro_name):
+    """Disable overcommit accounting on E2 machine types."""
+
+    # Expected machine type format:
+    # 'projects/00000000000/machineTypes/n1-standard-1'
+    machine_type = self.metadata_dict['instance']['machineType'].split('/')[-1]
+    if machine_type.startswith('e2-') and 'bsd' not in distro:
+      subprocess.call(['sysctl', 'vm.overcommit_memory=1'])
+
   def _GetInstanceConfig(self):
     """Get the instance configuration specified in metadata.
 
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..5c6e513 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,10 @@ 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
@@ -79,6 +82,7 @@ class InstanceSetupTest(unittest.TestCase):
         mock.call.config.InstanceConfig().GetOptionBool(
             'InstanceSetup', 'set_boto_config'),
         mock.call.setup._SetupBotoConfig(),
+        mock.call.setup._DisableOvercommit(),
         # Setup for local SSD.
         mock.call.config.InstanceConfig().GetOptionBool(
             'InstanceSetup', 'optimize_local_ssd'),
@@ -91,7 +95,10 @@ 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')
@@ -421,6 +428,37 @@ class InstanceSetupTest(unittest.TestCase):
     instance_setup.InstanceSetup._SetupBotoConfig(self.mock_setup)
     self.mock_logger.warning.assert_called_once_with('Test Error')
 
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess')
+  def testDisableOvercommitNonE2(self, mock_subprocess):
+    self.mock_setup.metadata_dict = {
+        'instance': {
+            'machineType': 'projects/00000000000/machineTypes/n1-standard-1',
+        }
+    }
+    instance_setup.InstanceSetup._DisableOvercommit(self.mock_setup)
+    mock_subprocess.call.assert_not_called()
+
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess')
+  def testDisableOvercommitE2(self, mock_subprocess):
+    self.mock_setup.metadata_dict = {
+        'instance': {
+            'machineType': 'projects/00000000000/machineTypes/e2-standard-1',
+        }
+    }
+    instance_setup.InstanceSetup._DisableOvercommit(self.mock_setup)
+    mock_subprocess.call.assert_called_once_with(
+        ['sysctl', 'vm.overcommit_memory=1'])
+
+  @mock.patch('google_compute_engine.instance_setup.instance_setup.subprocess')
+  def testDisableOvercommitBSD(self, mock_subprocess):
+    self.mock_setup.metadata_dict = {
+        'instance': {
+            'machineType': 'projects/00000000000/machineTypes/e2-standard-1',
+        }
+    }
+    instance_setup.InstanceSetup._DisableOvercommit(self.mock_setup, 'bsd')
+    mock_subprocess.call.assert_not_called()
+
 
 if __name__ == '__main__':
   unittest.main()