From a61087c955a49916c5ccfb7ee73be6446be562a3 Mon Sep 17 00:00:00 2001
From: Pritham Marupaka <mpritham@users.noreply.github.com>
Date: Thu, 16 Nov 2023 17:53:00 -0500
Subject: [PATCH] Add test to check that heap size is not computed when cgroups
 memory limits are over 1TB (#363)

Add test to check that heap size is not computed when cgroups memory limits are over 1TB
---
 changelog/@unreleased/pr-363.v2.yml           |  6 ++++++
 .../go_java_launcher_integration_test.go      | 20 +++++++++++++++++--
 launchlib/launcher.go                         | 11 ++++++----
 3 files changed, 31 insertions(+), 6 deletions(-)
 create mode 100644 changelog/@unreleased/pr-363.v2.yml

diff --git a/changelog/@unreleased/pr-363.v2.yml b/changelog/@unreleased/pr-363.v2.yml
new file mode 100644
index 00000000..ac46500b
--- /dev/null
+++ b/changelog/@unreleased/pr-363.v2.yml
@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Add test to check that heap size is not computed when cgroups memory
+    limits are over 1TB
+  links:
+  - https://github.com/palantir/go-java-launcher/pull/363
diff --git a/integration_test/go_java_launcher_integration_test.go b/integration_test/go_java_launcher_integration_test.go
index 9bee79de..1839dadd 100644
--- a/integration_test/go_java_launcher_integration_test.go
+++ b/integration_test/go_java_launcher_integration_test.go
@@ -209,6 +209,7 @@ func TestComputeJVMHeapSize(t *testing.T) {
 		numHostProcessors   int
 		memoryLimit         uint64
 		expectedMaxHeapSize uint64
+		expectError         bool
 	}{
 		{
 			name:              "at least 50% of heap",
@@ -216,6 +217,7 @@ func TestComputeJVMHeapSize(t *testing.T) {
 			memoryLimit:       10 * launchlib.BytesInMebibyte,
 			// 75% of heap - 3mb*processors = 4.5mb
 			expectedMaxHeapSize: 5 * launchlib.BytesInMebibyte,
+			expectError:         false,
 		},
 		{
 			name:              "computes 75% of heap minus 3mb per processor",
@@ -223,6 +225,7 @@ func TestComputeJVMHeapSize(t *testing.T) {
 			memoryLimit:       16 * launchlib.BytesInMebibyte,
 			// 75% of heap - 3mb*processors = 9mb
 			expectedMaxHeapSize: 9 * launchlib.BytesInMebibyte,
+			expectError:         false,
 		},
 		{
 			name:              "multiple processors",
@@ -230,11 +233,24 @@ func TestComputeJVMHeapSize(t *testing.T) {
 			memoryLimit:       120 * launchlib.BytesInMebibyte,
 			// 75% of heap - 3mb*processors = 81mb
 			expectedMaxHeapSize: 81 * launchlib.BytesInMebibyte,
+			expectError:         false,
+		},
+		{
+			name:                "memory limit too large",
+			numHostProcessors:   1,
+			memoryLimit:         1_000_001 * launchlib.BytesInMebibyte,
+			expectedMaxHeapSize: 0,
+			expectError:         true,
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			heapSizeInBytes := launchlib.ComputeJVMHeapSizeInBytes(tc.numHostProcessors, tc.memoryLimit)
-			assert.Equal(t, heapSizeInBytes, tc.expectedMaxHeapSize)
+			heapSizeInBytes, err := launchlib.ComputeJVMHeapSizeInBytes(tc.numHostProcessors, tc.memoryLimit)
+			if tc.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, heapSizeInBytes, tc.expectedMaxHeapSize)
+			}
 		})
 	}
 }
diff --git a/launchlib/launcher.go b/launchlib/launcher.go
index 40fe105c..ebfc39f0 100644
--- a/launchlib/launcher.go
+++ b/launchlib/launcher.go
@@ -355,10 +355,10 @@ func filterHeapSizeArgsV2(args []string) ([]string, error) {
 		if err != nil {
 			return filtered, errors.Wrap(err, "failed to get cgroup memory limit")
 		}
-		if cgroupMemoryLimitInBytes > 1_000_000*BytesInMebibyte {
+		jvmHeapSizeInBytes, err := ComputeJVMHeapSizeInBytes(runtime.NumCPU(), cgroupMemoryLimitInBytes)
+		if err != nil {
 			return filtered, errors.New("cgroups memory limit is unusually high. Not setting JVM heap size options")
 		}
-		jvmHeapSizeInBytes := ComputeJVMHeapSizeInBytes(runtime.NumCPU(), cgroupMemoryLimitInBytes)
 		filtered = append(filtered, fmt.Sprintf("-Xms%d", jvmHeapSizeInBytes))
 		filtered = append(filtered, fmt.Sprintf("-Xmx%d", jvmHeapSizeInBytes))
 	}
@@ -427,9 +427,12 @@ func isInitialRAMPercentage(arg string) bool {
 
 // ComputeJVMHeapSizeInBytes If the experimental `ContainerV2` is set, compute the heap size to be 75% of
 // the heap minus 3mb per processor, with a minimum value of 50% of the heap.
-func ComputeJVMHeapSizeInBytes(hostProcessorCount int, cgroupMemoryLimitInBytes uint64) uint64 {
+func ComputeJVMHeapSizeInBytes(hostProcessorCount int, cgroupMemoryLimitInBytes uint64) (uint64, error) {
+	if cgroupMemoryLimitInBytes > 1_000_000*BytesInMebibyte {
+		return 0, errors.New("cgroups memory limit is unusually high. Not computing JVM heap size")
+	}
 	var memoryLimit = float64(cgroupMemoryLimitInBytes)
 	var processorAdjustment = 3 * BytesInMebibyte * float64(hostProcessorCount)
 	var computedHeapSize = max(0.5*memoryLimit, 0.75*memoryLimit-processorAdjustment)
-	return uint64(computedHeapSize)
+	return uint64(computedHeapSize), nil
 }