Merge pull request #15288 from andyxiexu/java-sdk

Add google cloud heap profiling support to beam java sdk container

sdks/java/container/boot.go

@@ -48,6 +48,13 @@ var (
 	semiPersistDir    = flag.String("semi_persist_dir", "/tmp", "Local semi-persistent directory (optional).")
 )
 
+const (
+	enableGoogleCloudProfilerOption     = "enable_google_cloud_profiler"
+	enableGoogleCloudHeapSamplingOption = "enable_google_cloud_heap_sampling"
+	googleCloudProfilerAgentBaseArgs    = "-agentpath:/opt/google_cloud_profiler/profiler_java_agent.so=-logtostderr,-cprof_service=%s,-cprof_service_version=%s"
+	googleCloudProfilerAgentHeapArgs    = googleCloudProfilerAgentBaseArgs + ",-cprof_enable_heap_sampling,-cprof_heap_sampling_interval=2097152"
+)
+
 func main() {
 	flag.Parse()
 	if *id == "" {
@@ -151,12 +158,18 @@ func main() {
 		"-cp", strings.Join(cp, ":"),
 	}
 
-	enableGoogleCloudProfiler := strings.Contains(options, "enable_google_cloud_profiler")
+	enableGoogleCloudProfiler := strings.Contains(options, enableGoogleCloudProfilerOption)
+	enableGoogleCloudHeapSampling := strings.Contains(options, enableGoogleCloudHeapSamplingOption)
 	if enableGoogleCloudProfiler {
 		if metadata := info.GetMetadata(); metadata != nil {
 			if jobName, nameExists := metadata["job_name"]; nameExists {
 				if jobId, idExists := metadata["job_id"]; idExists {
-					args = append(args, fmt.Sprintf("-agentpath:/opt/google_cloud_profiler/profiler_java_agent.so=-cprof_service=%s,-cprof_service_version=%s", jobName, jobId))
+					if enableGoogleCloudHeapSampling {
+						args = append(args, fmt.Sprintf(googleCloudProfilerAgentHeapArgs, jobName, jobId))
+					} else {
+						args = append(args, fmt.Sprintf(googleCloudProfilerAgentBaseArgs, jobName, jobId))
+					}
+					log.Printf("Turning on Cloud Profiling. Profile heap: %t", enableGoogleCloudHeapSampling)
 				} else {
 					log.Println("Required job_id missing from metadata, profiling will not be enabled without it.")
 				}