From 6890a1a7b9331aa82371fc22b9c3f2be019d7651 Mon Sep 17 00:00:00 2001
From: Pritham Marupaka <pmarupaka@palantir.com>
Date: Tue, 14 Nov 2023 13:59:16 -0500
Subject: [PATCH] add test for ComputeJVMHeapSizeInBytes

---
 .../go_java_launcher_integration_test.go      | 36 +++++++++++++++++++
 launchlib/launcher.go                         | 19 ++++------
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/integration_test/go_java_launcher_integration_test.go b/integration_test/go_java_launcher_integration_test.go
index ae5d92b3..c7b19e66 100644
--- a/integration_test/go_java_launcher_integration_test.go
+++ b/integration_test/go_java_launcher_integration_test.go
@@ -203,6 +203,42 @@ func TestSubProcessesParsedMonitorSignals(t *testing.T) {
 	assert.Len(t, trapped.FindAll(output.Bytes(), -1), 2, "expect two messages that SIGPOLL was caught")
 }
 
+func TestComputeJVMHeapSize(t *testing.T) {
+	for _, tc := range []struct {
+		name                string
+		numHostProcessors   int
+		memoryLimit         uint64
+		expectedMaxHeapSize uint64
+	}{
+		{
+			name:              "at least 50% of heap",
+			numHostProcessors: 1,
+			memoryLimit:       10 * launchlib.BytesInMebibyte,
+			// 75% of heap - 3mb*processors = 4.5mb
+			expectedMaxHeapSize: 5 * launchlib.BytesInMebibyte,
+		},
+		{
+			name:              "computes 75% of heap minus 3mb per processor",
+			numHostProcessors: 1,
+			memoryLimit:       12 * launchlib.BytesInMebibyte,
+			// 75% of heap - 3mb*processors = 6mb
+			expectedMaxHeapSize: 6 * launchlib.BytesInMebibyte,
+		},
+		{
+			name:              "multiple processors",
+			numHostProcessors: 3,
+			memoryLimit:       120 * launchlib.BytesInMebibyte,
+			// 75% of heap - 3mb*processors = 81mb
+			expectedMaxHeapSize: 81 * launchlib.BytesInMebibyte,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			heapSizeInBytes := launchlib.ComputeJVMHeapSizeInBytes(tc.numHostProcessors, tc.memoryLimit)
+			assert.Equal(t, heapSizeInBytes, tc.expectedMaxHeapSize)
+		})
+	}
+}
+
 func runMainWithArgs(t *testing.T, staticConfigFile, customConfigFile string, env ...string) (string, error) {
 	jdkDir := "jdk"
 	javaHome, err := filepath.Abs(jdkDir)
diff --git a/launchlib/launcher.go b/launchlib/launcher.go
index 21fb3b30..9b1037d0 100644
--- a/launchlib/launcher.go
+++ b/launchlib/launcher.go
@@ -347,11 +347,11 @@ func filterHeapSizeArgsV2(args []string, logger io.Writer) ([]string, error) {
 	}
 
 	if !hasInitialRAMPercentage && !hasMaxRAMPercentage {
-		jvmHeapSizeInBytes, err := computeJVMHeapSizeInBytes()
+		cgroupMemoryLimitInBytes, err := DefaultMemoryLimit.MemoryLimitInBytes()
 		if err != nil {
-			_, _ = fmt.Fprintln(logger, "Failed to compute JVM heap size", err.Error())
-			return filtered, err
+			return filtered, errors.Wrap(err, "failed to get cgroup memory limit")
 		}
+		jvmHeapSizeInBytes := ComputeJVMHeapSizeInBytes(runtime.NumCPU(), cgroupMemoryLimitInBytes)
 		filtered = append(filtered, fmt.Sprintf("-Xms%d", jvmHeapSizeInBytes))
 		filtered = append(filtered, fmt.Sprintf("-Xmx%d", jvmHeapSizeInBytes))
 	}
@@ -418,16 +418,11 @@ func isInitialRAMPercentage(arg string) bool {
 	return strings.HasPrefix(arg, "-XX:InitialRAMPercentage=")
 }
 
-// If the experimental `ExperimentalContainerV2` 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() (uint64, error) {
-	hostProcessorCount := runtime.NumCPU()
-	cgroupMemoryLimitInBytes, err := DefaultMemoryLimit.MemoryLimitInBytes()
-	if err != nil {
-		return 0, errors.Wrap(err, "failed to get cgroup memory limit")
-	}
+// ComputeJVMHeapSizeInBytes If the experimental `ExperimentalContainerV2` 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 {
 	var heapLimit = float64(cgroupMemoryLimitInBytes)
 	var processorAdjustment = 3 * BytesInMebibyte * float64(hostProcessorCount)
 	var computedHeapSize = max(0.5*heapLimit, 0.75*heapLimit-processorAdjustment)
-	return uint64(computedHeapSize), nil
+	return uint64(computedHeapSize)
 }