Skip to content

Commit e9318ac

Browse files
beikovbeikov
committed
HHH-19240 Improve memory consumption testing
(cherry picked from commit 411bae9)
1 parent 0b7c64b commit e9318ac

File tree

7 files changed

+456
-54
lines changed

7 files changed

+456
-54
lines changed

‎hibernate-core/src/test/java/org/hibernate/orm/test/hql/HqlParserMemoryUsageTest.java‎

Lines changed: 4 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import jakarta.persistence.Table;
1313
import org.hibernate.cfg.QuerySettings;
1414
import org.hibernate.query.hql.HqlTranslator;
15-
import org.hibernate.query.sqm.tree.SqmStatement;
15+
import org.hibernate.testing.memory.MemoryUsageUtil;
1616
import org.hibernate.testing.orm.junit.DomainModel;
1717
import org.hibernate.testing.orm.junit.Jira;
1818
import org.hibernate.testing.orm.junit.ServiceRegistry;
@@ -80,65 +80,15 @@ SELECT AVG(p3.price) FROM Product p3
8080
@Test
8181
public void testParserMemoryUsage(SessionFactoryScope scope) {
8282
final HqlTranslator hqlTranslator = scope.getSessionFactory().getQueryEngine().getHqlTranslator();
83-
final Runtime runtime = Runtime.getRuntime();
8483

8584
// Ensure classes and basic stuff is initialized in case this is the first test run
8685
hqlTranslator.translate( "from AppUser", AppUser.class );
87-
runtime.gc();
88-
runtime.gc();
89-
90-
// Track memory usage before execution
91-
long totalMemoryBefore = runtime.totalMemory();
92-
long usedMemoryBefore = totalMemoryBefore - runtime.freeMemory();
93-
94-
System.out.println("Memory Usage Before Create Query:");
95-
System.out.println("----------------------------");
96-
System.out.println("Total Memory: " + (totalMemoryBefore / 1024) + " KB");
97-
System.out.println("Used Memory : " + (usedMemoryBefore / 1024) + " KB");
98-
System.out.println();
99-
100-
// Create query
101-
SqmStatement<Long> statement = hqlTranslator.translate( HQL, Long.class );
102-
103-
// Track memory usage after execution
104-
long totalMemoryAfter = runtime.totalMemory();
105-
long usedMemoryAfter = totalMemoryAfter - runtime.freeMemory();
106-
107-
System.out.println("Memory Usage After Create Query:");
108-
System.out.println("----------------------------");
109-
System.out.println("Total Memory: " + (totalMemoryAfter / 1024) + " KB");
110-
System.out.println("Used Memory : " + (usedMemoryAfter / 1024) + " KB");
111-
System.out.println();
112-
113-
System.out.println("Memory increase After Parsing:");
114-
System.out.println("----------------------------");
115-
System.out.println("Total Memory increase: " + ((totalMemoryAfter - totalMemoryBefore) / 1024) + " KB");
116-
System.out.println("Used Memory increase : " + ((usedMemoryAfter - usedMemoryBefore) / 1024) + " KB");
117-
System.out.println();
118-
119-
runtime.gc();
120-
runtime.gc();
121-
122-
// Track memory usage after execution
123-
long totalMemoryAfterGc = runtime.totalMemory();
124-
long usedMemoryAfterGc = totalMemoryAfterGc - runtime.freeMemory();
125-
126-
System.out.println("Memory Usage After Create Query and GC:");
127-
System.out.println("----------------------------");
128-
System.out.println("Total Memory: " + (totalMemoryAfterGc / 1024) + " KB");
129-
System.out.println("Used Memory : " + (usedMemoryAfterGc / 1024) + " KB");
130-
System.out.println();
131-
132-
System.out.println("Memory overhead of Parsing:");
133-
System.out.println("----------------------------");
134-
System.out.println("Total Memory increase: " + ((totalMemoryAfter - totalMemoryAfterGc) / 1024) + " KB");
135-
System.out.println("Used Memory increase : " + ((usedMemoryAfter - usedMemoryAfterGc) / 1024) + " KB");
136-
System.out.println();
13786

13887
// During testing, before the fix for HHH-19240, the allocation was around 500+ MB,
13988
// and after the fix it dropped to 170 - 250 MB
140-
final long memoryConsumption = usedMemoryAfter - usedMemoryAfterGc;
141-
assertTrue( usedMemoryAfter - usedMemoryAfterGc < 256_000_000, "Parsing of queries consumes too much memory (" + ( memoryConsumption / 1024 ) + " KB), when at most 256 MB are expected" );
89+
final long memoryUsage = MemoryUsageUtil.estimateMemoryUsage( () -> hqlTranslator.translate( HQL, Long.class ) );
90+
System.out.println( "Memory Consumption: " + (memoryUsage / 1024) + " KB" );
91+
assertTrue( memoryUsage < 256_000_000, "Parsing of queries consumes too much memory (" + ( memoryUsage / 1024 ) + " KB), when at most 256 MB are expected" );
14292
}
14393

14494
@Entity(name = "Address")

‎hibernate-testing/src/main/java/org/hibernate/testing/memory/GlobalMemoryUsageSnapshotter.java‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright Red Hat Inc. and Hibernate Authors
4+
*/
5+
package org.hibernate.testing.memory;
6+
7+
import java.lang.management.ManagementFactory;
8+
import java.lang.management.MemoryPoolMXBean;
9+
import java.util.List;
10+
import java.util.Objects;
11+
12+
final class GlobalMemoryUsageSnapshotter implements MemoryAllocationSnapshotter {
13+
14+
private static final GlobalMemoryUsageSnapshotter INSTANCE = new GlobalMemoryUsageSnapshotter(
15+
ManagementFactory.getMemoryPoolMXBeans()
16+
);
17+
18+
private final List<MemoryPoolMXBean> heapPoolBeans;
19+
private final Runnable gcAndWait;
20+
21+
private GlobalMemoryUsageSnapshotter(List<MemoryPoolMXBean> heapPoolBeans) {
22+
this.heapPoolBeans = heapPoolBeans;
23+
this.gcAndWait = () -> {
24+
for (int i = 0; i < 3; i++) {
25+
System.gc();
26+
try { Thread.sleep(50); } catch (InterruptedException ignored) {}
27+
}
28+
};
29+
}
30+
31+
public static GlobalMemoryUsageSnapshotter getInstance() {
32+
return INSTANCE;
33+
}
34+
35+
@Override
36+
public MemoryAllocationSnapshot snapshot() {
37+
final long peakUsage = heapPoolBeans.stream().mapToLong(p -> p.getPeakUsage().getUsed()).sum();
38+
gcAndWait.run();
39+
final long retainedUsage = heapPoolBeans.stream().mapToLong(p -> p.getUsage().getUsed()).sum();
40+
heapPoolBeans.forEach(MemoryPoolMXBean::resetPeakUsage);
41+
return new GlobalMemoryAllocationSnapshot( peakUsage, retainedUsage );
42+
}
43+
44+
final static class GlobalMemoryAllocationSnapshot implements MemoryAllocationSnapshot {
45+
private final long peakUsage;
46+
private final long retainedUsage;
47+
48+
GlobalMemoryAllocationSnapshot(long peakUsage, long retainedUsage) {
49+
this.peakUsage = peakUsage;
50+
this.retainedUsage = retainedUsage;
51+
}
52+
53+
public long peakUsage() {
54+
return peakUsage;
55+
}
56+
57+
public long retainedUsage() {
58+
return retainedUsage;
59+
}
60+
61+
@Override
62+
public long difference(MemoryAllocationSnapshot before) {
63+
// When doing the "before" snapshot, the peak usage is reset.
64+
// Since this object is the "after" snapshot, we can simply estimate the memory usage of an operation
65+
// to be the peak usage of that operation minus the usage after GC
66+
return peakUsage - retainedUsage;
67+
}
68+
69+
@Override
70+
public boolean equals(Object obj) {
71+
if ( obj == this ) {
72+
return true;
73+
}
74+
if ( obj == null || obj.getClass() != this.getClass() ) {
75+
return false;
76+
}
77+
var that = (GlobalMemoryAllocationSnapshot) obj;
78+
return this.peakUsage == that.peakUsage &&
79+
this.retainedUsage == that.retainedUsage;
80+
}
81+
82+
@Override
83+
public int hashCode() {
84+
return Objects.hash( peakUsage, retainedUsage );
85+
}
86+
87+
@Override
88+
public String toString() {
89+
return "GlobalMemoryAllocationSnapshot[" +
90+
"peakUsage=" + peakUsage + ", " +
91+
"retainedUsage=" + retainedUsage + ']';
92+
}
93+
}
94+
}

‎hibernate-testing/src/main/java/org/hibernate/testing/memory/HotspotPerThreadAllocationSnapshotter.java‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
* Copyright Red Hat Inc. and Hibernate Authors
4+
*/
5+
package org.hibernate.testing.memory;
6+
7+
import org.checkerframework.checker.nullness.qual.Nullable;
8+
9+
import java.lang.management.ManagementFactory;
10+
import java.lang.management.ThreadMXBean;
11+
import java.lang.reflect.Method;
12+
import java.util.HashMap;
13+
import java.util.Objects;
14+
15+
final class HotspotPerThreadAllocationSnapshotter implements MemoryAllocationSnapshotter {
16+
17+
private static final @Nullable HotspotPerThreadAllocationSnapshotter INSTANCE;
18+
private static final Method GET_THREAD_ALLOCATED_BYTES;
19+
20+
static {
21+
ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
22+
Method method = null;
23+
try {
24+
@SuppressWarnings("unchecked")
25+
Class<? extends ThreadMXBean> hotspotInterface =
26+
(Class<? extends ThreadMXBean>) Class.forName( "com.sun.management.ThreadMXBean" );
27+
try {
28+
method = hotspotInterface.getMethod( "getThreadAllocatedBytes", long[].class );
29+
}
30+
catch (Exception e) {
31+
// Ignore
32+
}
33+
34+
if ( !hotspotInterface.isInstance( threadMXBean ) ) {
35+
threadMXBean = ManagementFactory.getPlatformMXBean( hotspotInterface );
36+
}
37+
}
38+
catch (Throwable e) {
39+
// Ignore
40+
}
41+
42+
GET_THREAD_ALLOCATED_BYTES = method;
43+
44+
HotspotPerThreadAllocationSnapshotter instance = null;
45+
if ( method != null && threadMXBean != null ) {
46+
try {
47+
instance = new HotspotPerThreadAllocationSnapshotter( threadMXBean );
48+
instance.snapshot();
49+
}
50+
catch (Exception e) {
51+
instance = null;
52+
}
53+
}
54+
INSTANCE = instance;
55+
}
56+
57+
public static @Nullable HotspotPerThreadAllocationSnapshotter getInstance() {
58+
return INSTANCE;
59+
}
60+
61+
@Override
62+
public MemoryAllocationSnapshot snapshot() {
63+
long[] threadIds = threadMXBean.getAllThreadIds();
64+
try {
65+
return new PerThreadMemoryAllocationSnapshot(
66+
threadIds,
67+
(long[]) GET_THREAD_ALLOCATED_BYTES.invoke( threadMXBean, (Object) threadIds )
68+
);
69+
}
70+
catch (Exception e) {
71+
throw new RuntimeException( e );
72+
}
73+
}
74+
75+
final static class PerThreadMemoryAllocationSnapshot implements MemoryAllocationSnapshot {
76+
private final long[] threadIds;
77+
private final long[] threadAllocatedBytes;
78+
79+
PerThreadMemoryAllocationSnapshot(long[] threadIds, long[] threadAllocatedBytes) {
80+
this.threadIds = threadIds;
81+
this.threadAllocatedBytes = threadAllocatedBytes;
82+
}
83+
84+
public long[] threadIds() {
85+
return threadIds;
86+
}
87+
88+
public long[] threadAllocatedBytes() {
89+
return threadAllocatedBytes;
90+
}
91+
92+
@Override
93+
public long difference(MemoryAllocationSnapshot before) {
94+
final PerThreadMemoryAllocationSnapshot other = (PerThreadMemoryAllocationSnapshot) before;
95+
final HashMap<Long, Integer> previousThreadIdToIndexMap = new HashMap<>();
96+
for ( int i = 0; i < other.threadIds.length; i++ ) {
97+
previousThreadIdToIndexMap.put( other.threadIds[i], i );
98+
}
99+
long allocatedBytes = 0;
100+
for ( int i = 0; i < threadIds.length; i++ ) {
101+
allocatedBytes += threadAllocatedBytes[i];
102+
final Integer previousThreadIndex = previousThreadIdToIndexMap.get( threadIds[i] );
103+
if ( previousThreadIndex != null ) {
104+
allocatedBytes -= other.threadAllocatedBytes[previousThreadIndex];
105+
}
106+
}
107+
return allocatedBytes;
108+
}
109+
110+
@Override
111+
public boolean equals(Object obj) {
112+
if ( obj == this ) {
113+
return true;
114+
}
115+
if ( obj == null || obj.getClass() != this.getClass() ) {
116+
return false;
117+
}
118+
var that = (PerThreadMemoryAllocationSnapshot) obj;
119+
return Objects.equals( this.threadIds, that.threadIds ) &&
120+
Objects.equals( this.threadAllocatedBytes, that.threadAllocatedBytes );
121+
}
122+
123+
@Override
124+
public int hashCode() {
125+
return Objects.hash( threadIds, threadAllocatedBytes );
126+
}
127+
128+
@Override
129+
public String toString() {
130+
return "PerThreadMemoryAllocationSnapshot[" +
131+
"threadIds=" + threadIds + ", " +
132+
"threadAllocatedBytes=" + threadAllocatedBytes + ']';
133+
}
134+
}
135+
private final ThreadMXBean threadMXBean;
136+
137+
HotspotPerThreadAllocationSnapshotter(ThreadMXBean threadMXBean) {
138+
this.threadMXBean = threadMXBean;
139+
}
140+
141+
public ThreadMXBean threadMXBean() {
142+
return threadMXBean;
143+
}
144+
145+
@Override
146+
public boolean equals(Object obj) {
147+
if ( obj == this ) {
148+
return true;
149+
}
150+
if ( obj == null || obj.getClass() != this.getClass() ) {
151+
return false;
152+
}
153+
var that = (HotspotPerThreadAllocationSnapshotter) obj;
154+
return Objects.equals( this.threadMXBean, that.threadMXBean );
155+
}
156+
157+
@Override
158+
public int hashCode() {
159+
return Objects.hash( threadMXBean );
160+
}
161+
162+
@Override
163+
public String toString() {
164+
return "HotspotPerThreadAllocationSnapshotter[" +
165+
"threadMXBean=" + threadMXBean + ']';
166+
}
167+
}

0 commit comments

Comments
 (0)