From 5cd01474750b6183a89df0731e601b450ee0de74 Mon Sep 17 00:00:00 2001
From: Arushi Gaur <arushigaur@google.com>
Date: Mon, 9 Dec 2024 13:23:48 -0800
Subject: [PATCH] Setting MemoryMax for postgres template service in debian
 systems

PiperOrigin-RevId: 704396036
---
 perfkitbenchmarker/linux_packages/postgresql16.py | 14 +++++++++++++-
 perfkitbenchmarker/os_types.py                    |  7 +++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/perfkitbenchmarker/linux_packages/postgresql16.py b/perfkitbenchmarker/linux_packages/postgresql16.py
index 3ae56edbe..b6a01aa29 100644
--- a/perfkitbenchmarker/linux_packages/postgresql16.py
+++ b/perfkitbenchmarker/linux_packages/postgresql16.py
@@ -49,6 +49,7 @@
         'conf_dir': '/etc/postgresql/16/main',
         'disk_mount_point': '/etc/postgresql/16/data',
         'postgres_service_name': 'postgresql',
+        'postgres_template_service_name': 'postgresql@16-main',
     },
     'amazonlinux': {
         'data_dir': '/var/lib/pgsql/data',
@@ -166,6 +167,11 @@ def GetOSDependentDefaults(os_type: str) -> dict[str, str]:
     return OS_DEPENDENT_DEFAULTS['debian']
 
 
+def IsUbuntu(vm):
+  """Returns whether the VM is Debian."""
+  return vm.OS_TYPE in os_types.UBUNTU_OS_TYPES
+
+
 def ConfigureAndRestart(vm, run_uri, buffer_size):
   """Configure and restart postgres.
 
@@ -235,7 +241,13 @@ def ConfigureAndRestart(vm, run_uri, buffer_size):
       f' {GetOSDependentDefaults(vm.OS_TYPE)["postgres_service_name"]}.service'
       f' MemoryMax={SHARED_BUFFERS_CONF[buffer_size_key]["max_memory"]}'
   )
-  vm.RemoteCommand('sudo sync; echo 3 | sudo tee -a /proc/sys/vm/drop_caches')
+  if IsUbuntu(vm):
+    vm.RemoteCommand(
+        'sudo systemctl set-property'
+        f' {GetOSDependentDefaults(vm.OS_TYPE)["postgres_template_service_name"]}.service'
+        f' MemoryMax={SHARED_BUFFERS_CONF[buffer_size_key]["max_memory"]}'
+    )
+  vm.RemoteCommand('sudo sync; echo 3 | sudo tee /proc/sys/vm/drop_caches')
   vm.RemoteCommand(
       f'cat /etc/systemd/system.control/{GetOSDependentDefaults(vm.OS_TYPE)["postgres_service_name"]}.service.d/50-MemoryMax.conf'
   )
diff --git a/perfkitbenchmarker/os_types.py b/perfkitbenchmarker/os_types.py
index e6ea6db7c..c07db4e1b 100644
--- a/perfkitbenchmarker/os_types.py
+++ b/perfkitbenchmarker/os_types.py
@@ -148,6 +148,13 @@
 WINDOWS_OS_TYPES = (
     WINDOWS_CORE_OS_TYPES + WINDOWS_DESKOP_OS_TYPES + WINDOWS_SQLSERVER_OS_TYPES
 )
+UBUNTU_OS_TYPES = [
+    UBUNTU2004,
+    UBUNTU2004_EFA,
+    UBUNTU2004_DL,
+    UBUNTU2204,
+    UBUNTU2404,
+]
 
 AMAZONLINUX_TYPES = [
     AMAZONLINUX2,