From a9fce43b8d0c6f49989d37b2e5fd88881f869911 Mon Sep 17 00:00:00 2001
From: Jonatan Ivanov <jonatan.ivanov@gmail.com>
Date: Thu, 27 Jun 2024 08:16:49 -0700
Subject: [PATCH] Add MemoryInfo to ProcessInfo

---
 .../boot/info/ProcessInfo.java                | 71 +++++++++++++++++++
 .../boot/info/ProcessInfoTests.java           | 13 ++++
 2 files changed, 84 insertions(+)

diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/info/ProcessInfo.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/info/ProcessInfo.java
index 96099d479e4..01ed9649752 100644
--- a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/info/ProcessInfo.java
+++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/info/ProcessInfo.java
@@ -16,6 +16,10 @@
 
 package org.springframework.boot.info;
 
+import java.lang.management.ManagementFactory;
+import java.lang.management.MemoryMXBean;
+import java.lang.management.MemoryUsage;
+
 /**
  * Information about the process of the application.
  *
@@ -50,6 +54,24 @@ public class ProcessInfo {
 		return runtime.availableProcessors();
 	}
 
+	/**
+	 * Memory information for the process. These values can provide details about the
+	 * current memory usage and limits selected by the user or JVM ergonomics (init, max,
+	 * committed, used for heap and non-heap). If limits not set explicitly, it might not
+	 * be trivial to know what these values are runtime; especially in (containerized)
+	 * environments where resource usage can be isolated (for example using control
+	 * groups) or not necessarily trivial to discover. Other than that, these values can
+	 * indicate if the JVM can resize the heap (stop-the-world).
+	 * @return heap and non-heap memory information
+	 * @since 3.4.0
+	 * @see MemoryMXBean#getHeapMemoryUsage()
+	 * @see MemoryMXBean#getNonHeapMemoryUsage()
+	 * @see MemoryUsage
+	 */
+	public MemoryInfo getMemory() {
+		return new MemoryInfo();
+	}
+
 	public long getPid() {
 		return this.pid;
 	}
@@ -62,4 +84,53 @@ public class ProcessInfo {
 		return this.owner;
 	}
 
+	public static class MemoryInfo {
+
+		private static final MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean();
+
+		private final MemoryUsageInfo heap;
+
+		private final MemoryUsageInfo nonHeap;
+
+		MemoryInfo() {
+			this.heap = new MemoryUsageInfo(memoryMXBean.getHeapMemoryUsage());
+			this.nonHeap = new MemoryUsageInfo(memoryMXBean.getNonHeapMemoryUsage());
+		}
+
+		public MemoryUsageInfo getHeap() {
+			return this.heap;
+		}
+
+		public MemoryUsageInfo getNonHeap() {
+			return this.nonHeap;
+		}
+
+		public static class MemoryUsageInfo {
+
+			private final MemoryUsage memoryUsage;
+
+			MemoryUsageInfo(MemoryUsage memoryUsage) {
+				this.memoryUsage = memoryUsage;
+			}
+
+			public long getInit() {
+				return this.memoryUsage.getInit();
+			}
+
+			public long getUsed() {
+				return this.memoryUsage.getUsed();
+			}
+
+			public long getCommited() {
+				return this.memoryUsage.getCommitted();
+			}
+
+			public long getMax() {
+				return this.memoryUsage.getMax();
+			}
+
+		}
+
+	}
+
 }
diff --git a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/info/ProcessInfoTests.java b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/info/ProcessInfoTests.java
index 289581f5377..de198bc3e7c 100644
--- a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/info/ProcessInfoTests.java
+++ b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/info/ProcessInfoTests.java
@@ -18,6 +18,8 @@ package org.springframework.boot.info;
 
 import org.junit.jupiter.api.Test;
 
+import org.springframework.boot.info.ProcessInfo.MemoryInfo.MemoryUsageInfo;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
@@ -35,6 +37,17 @@ class ProcessInfoTests {
 		assertThat(processInfo.getPid()).isEqualTo(ProcessHandle.current().pid());
 		assertThat(processInfo.getParentPid())
 			.isEqualTo(ProcessHandle.current().parent().map(ProcessHandle::pid).orElse(null));
+
+		MemoryUsageInfo heapUsageInfo = processInfo.getMemory().getHeap();
+		MemoryUsageInfo nonHeapUsageInfo = processInfo.getMemory().getNonHeap();
+		assertThat(heapUsageInfo.getInit()).isPositive().isLessThanOrEqualTo(heapUsageInfo.getMax());
+		assertThat(heapUsageInfo.getUsed()).isPositive().isLessThanOrEqualTo(heapUsageInfo.getCommited());
+		assertThat(heapUsageInfo.getCommited()).isPositive().isLessThanOrEqualTo(heapUsageInfo.getMax());
+		assertThat(heapUsageInfo.getMax()).isPositive();
+		assertThat(nonHeapUsageInfo.getInit()).isPositive();
+		assertThat(nonHeapUsageInfo.getUsed()).isPositive().isLessThanOrEqualTo(heapUsageInfo.getCommited());
+		assertThat(nonHeapUsageInfo.getCommited()).isPositive();
+		assertThat(nonHeapUsageInfo.getMax()).isEqualTo(-1);
 	}
 
 }