Adjust JVM heap size for extremely large memory machine

sdks/java/container/boot.go

@@ -31,7 +31,6 @@ import (
 
 	"github.com/apache/beam/sdks/v2/go/container/tools"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/artifact"
-	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
 	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/execx"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/grpcx"
@@ -159,9 +158,19 @@ func main() {
 		cp = append(cp, filepath.Join(dir, filepath.FromSlash(name)))
 	}
 
-	var setRecommendedMaxXmx = strings.Contains(options, "set_recommended_max_xmx")
+	var lim uint64
+	if strings.Contains(options, "set_recommended_max_xmx") {
+    lim = 32 << 30
+  } else {
+    size, err := syscallx.PhysicalMemorySize()
+    if err != nil {
+      size = 0
+    }
+    lim = HeapSizeLimit(size)
+  }
+
 	args := []string{
-		"-Xmx" + strconv.FormatUint(heapSizeLimit(info, setRecommendedMaxXmx), 10),
+		"-Xmx" + strconv.FormatUint(lim, 10),
 		// ParallelGC the most adequate for high throughput and lower CPU utilization
 		// It is the default GC in Java 8, but not on newer versions
 		"-XX:+UseParallelGC",
@@ -248,20 +257,22 @@ func main() {
 }
 
 // heapSizeLimit returns 80% of the runner limit, if provided. If not provided,
-// it returns 70% of the physical memory on the machine. If it cannot determine
-// that value, it returns 1GB. This is an imperfect heuristic. It aims to
-// ensure there is memory for non-heap use and other overhead, while also not
-// underutilizing the machine. if set_recommended_max_xmx experiment is enabled,
-// sets xmx to 32G. Under 32G JVM enables CompressedOops. CompressedOops
-// utilizes memory more efficiently, and has positive impact on GC performance
-// and cache hit rate.
-func heapSizeLimit(info *fnpb.ProvisionInfo, setRecommendedMaxXmx bool) uint64 {
-	if setRecommendedMaxXmx {
-		return 32 << 30
-	} else if size, err := syscallx.PhysicalMemorySize(); err == nil {
-		return (size * 70) / 100
+// it returns max(70% size, size - 32GB). Set size=0 if the physical memory on
+// the machine was undetermined, then it returns 1GB. This is an imperfect
+// heuristic. It aims to ensure there is memory for non-heap use and other
+// overhead, while also not underutilizing the machine.
+// if set_recommended_max_xmx experiment is enabled, sets xmx to 32G. Under 32G
+// JVM enables CompressedOops. CompressedOops utilizes memory more efficiently,
+// and has positive impact on GC performance and cache hit rate.
+func HeapSizeLimit(size uint64) uint64 {
+	if size == 0 {
+		return 1 << 30
+	}
+	lim := (size * 70) / 100
+	if size-lim < 32<<30 {
+		return lim
 	}
-	return 1 << 30
+	return size - (32 << 30)
 }
 
 // Options represents java VM invocation options in a simple,

sdks/java/container/boot_test.go

@@ -78,3 +78,15 @@ func TestBuildOptions(t *testing.T) {
 		t.Errorf("BuildOptions(%v).JavaProperties = %v, want %v", metaOptions, javaOptions.Properties, wantProperties)
 	}
 }
+
+func TestHeapSizeLimit(t *testing.T) {
+	if lim := HeapSizeLimit(0); lim != 1 << 30 {
+		t.Errorf("HeapSizeLimit(0). Actual (%d). want 1 GB", lim)
+	}
+	if lim := HeapSizeLimit(2 << 30); lim != (2 << 30) * 7 / 10 {
+		t.Errorf("HeapSizeLimit(0). Actual (%d). want 700 MB", lim)
+	}
+	if lim := HeapSizeLimit(200 << 30); lim != (200 - 32) << 30 {
+		t.Errorf("HeapSizeLimit(0). Actual (%d). want 168 GB", lim)
+	}
+}