Skip to content

Commit 0c6f5b2

Browse files
jduojduo
committed
Refactor memory modules for JPMS support
Based on apache#13072 - Avoid having multiple memory modules contribute to the same package - Introduce memory-netty-buffer-patch module for patching classes into Netty modules - Avoid using BaseAllocator in tests and use RootAllocator instead - Move TestBaseAllocator#testMemoryUsage() to a new test in memory-netty TestNettyAllocator because TestBaseAllocator is now in memory-core, but that specific test has Netty dependencies.
1 parent 0a416d2 commit 0c6f5b2

30 files changed

+314
-130
lines changed

‎java/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java

+33-6
Original file line numberDiff line numberDiff line change
@@ -31,20 +31,35 @@
3131
*/
3232
final class CheckAllocator {
3333
private static final Logger logger = LoggerFactory.getLogger(CheckAllocator.class);
34-
private static final String ALLOCATOR_PATH = "org/apache/arrow/memory/DefaultAllocationManagerFactory.class";
34+
// unique package names needed by JPMS module naming
35+
private static final String ALLOCATOR_PATH_CORE =
36+
"org/apache/arrow/memory/DefaultAllocationManagerFactory.class";
37+
private static final String ALLOCATOR_PATH_UNSAFE =
38+
"org/apache/arrow/memory/unsafe/DefaultAllocationManagerFactory.class";
39+
private static final String ALLOCATOR_PATH_NETTY =
40+
"org/apache/arrow/memory/netty/DefaultAllocationManagerFactory.class";
3541

3642
private CheckAllocator() {
37-
3843
}
3944

4045
static String check() {
4146
Set<URL> urls = scanClasspath();
4247
URL rootAllocator = assertOnlyOne(urls);
4348
reportResult(rootAllocator);
44-
return "org.apache.arrow.memory.DefaultAllocationManagerFactory";
49+
if (rootAllocator.getPath().contains("memory-core") ||
50+
rootAllocator.getPath().contains("/org/apache/arrow/memory/core/")) {
51+
return "org.apache.arrow.memory.DefaultAllocationManagerFactory";
52+
} else if (rootAllocator.getPath().contains("memory-unsafe") ||
53+
rootAllocator.getPath().contains("/org/apache/arrow/memory/unsafe/")) {
54+
return "org.apache.arrow.memory.unsafe.DefaultAllocationManagerFactory";
55+
} else if (rootAllocator.getPath().contains("memory-netty") ||
56+
rootAllocator.getPath().contains("/org/apache/arrow/memory/netty/")) {
57+
return "org.apache.arrow.memory.netty.DefaultAllocationManagerFactory";
58+
} else {
59+
throw new IllegalStateException("Unknown allocation manager type to infer. Current: " + rootAllocator.getPath());
60+
}
4561
}
4662

47-
4863
private static Set<URL> scanClasspath() {
4964
// LinkedHashSet appropriate here because it preserves insertion order
5065
// during iteration
@@ -53,9 +68,21 @@ private static Set<URL> scanClasspath() {
5368
ClassLoader allocatorClassLoader = CheckAllocator.class.getClassLoader();
5469
Enumeration<URL> paths;
5570
if (allocatorClassLoader == null) {
56-
paths = ClassLoader.getSystemResources(ALLOCATOR_PATH);
71+
paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_CORE);
72+
if (!paths.hasMoreElements()) {
73+
paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_UNSAFE);
74+
}
75+
if (!paths.hasMoreElements()) {
76+
paths = ClassLoader.getSystemResources(ALLOCATOR_PATH_NETTY);
77+
}
5778
} else {
58-
paths = allocatorClassLoader.getResources(ALLOCATOR_PATH);
79+
paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_CORE);
80+
if (!paths.hasMoreElements()) {
81+
paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_UNSAFE);
82+
}
83+
if (!paths.hasMoreElements()) {
84+
paths = allocatorClassLoader.getResources(ALLOCATOR_PATH_NETTY);
85+
}
5986
}
6087
while (paths.hasMoreElements()) {
6188
URL path = paths.nextElement();

‎java/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java

+10-3
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919

2020
import java.lang.reflect.Field;
2121

22+
import org.apache.arrow.util.VisibleForTesting;
23+
2224
/**
2325
* A class for choosing the default allocation manager.
2426
*/
@@ -61,7 +63,12 @@ public enum AllocationManagerType {
6163
Unknown,
6264
}
6365

64-
static AllocationManagerType getDefaultAllocationManagerType() {
66+
/**
67+
* Returns the default allocation manager type.
68+
* @return the default allocation manager type.
69+
*/
70+
@VisibleForTesting
71+
public static AllocationManagerType getDefaultAllocationManagerType() {
6572
AllocationManagerType ret = AllocationManagerType.Unknown;
6673

6774
try {
@@ -115,7 +122,7 @@ private static AllocationManager.Factory getFactory(String clazzName) {
115122

116123
private static AllocationManager.Factory getUnsafeFactory() {
117124
try {
118-
return getFactory("org.apache.arrow.memory.UnsafeAllocationManager");
125+
return getFactory("org.apache.arrow.memory.unsafe.UnsafeAllocationManager");
119126
} catch (RuntimeException e) {
120127
throw new RuntimeException("Please add arrow-memory-unsafe to your classpath," +
121128
" No DefaultAllocationManager found to instantiate an UnsafeAllocationManager", e);
@@ -124,7 +131,7 @@ private static AllocationManager.Factory getUnsafeFactory() {
124131

125132
private static AllocationManager.Factory getNettyFactory() {
126133
try {
127-
return getFactory("org.apache.arrow.memory.NettyAllocationManager");
134+
return getFactory("org.apache.arrow.memory.netty.NettyAllocationManager");
128135
} catch (RuntimeException e) {
129136
throw new RuntimeException("Please add arrow-memory-netty to your classpath," +
130137
" No DefaultAllocationManager found to instantiate an NettyAllocationManager", e);

‎java/memory/memory-netty/src/test/java/org/apache/arrow/memory/CountingAllocationListener.java ‎java/memory/memory-core/src/test/java/org/apache/arrow/memory/CountingAllocationListener.java

+4
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,10 @@
1717

1818
package org.apache.arrow.memory;
1919

20+
import org.apache.arrow.memory.AllocationListener;
21+
import org.apache.arrow.memory.AllocationOutcome;
22+
import org.apache.arrow.memory.BufferAllocator;
23+
2024
/**
2125
* Counting allocation listener.
2226
* It counts the number of times it has been invoked, and how much memory allocation it has seen

‎java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java ‎java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java

+46-88
Original file line numberDiff line numberDiff line change
@@ -31,22 +31,16 @@
3131
import java.util.Collection;
3232
import java.util.Collections;
3333
import java.util.Iterator;
34-
import java.util.stream.Collectors;
3534

3635
import org.apache.arrow.memory.AllocationOutcomeDetails.Entry;
3736
import org.apache.arrow.memory.rounding.RoundingPolicy;
3837
import org.apache.arrow.memory.rounding.SegmentRoundingPolicy;
3938
import org.apache.arrow.memory.util.AssertionUtil;
39+
import org.junit.Assert;
4040
import org.junit.Ignore;
4141
import org.junit.Test;
4242
import org.junit.jupiter.api.Assertions;
43-
import org.slf4j.LoggerFactory;
4443

45-
import ch.qos.logback.classic.Level;
46-
import ch.qos.logback.classic.Logger;
47-
import ch.qos.logback.classic.spi.ILoggingEvent;
48-
import ch.qos.logback.core.read.ListAppender;
49-
import io.netty.buffer.PooledByteBufAllocatorL;
5044
import sun.misc.Unsafe;
5145

5246
public class TestBaseAllocator {
@@ -448,73 +442,73 @@ public ArrowBuf empty() {
448442
@Test
449443
public void testRootAllocator_listeners() throws Exception {
450444
CountingAllocationListener l1 = new CountingAllocationListener();
451-
assertEquals(0, l1.getNumPreCalls());
452-
assertEquals(0, l1.getNumCalls());
453-
assertEquals(0, l1.getNumReleaseCalls());
454-
assertEquals(0, l1.getNumChildren());
455-
assertEquals(0, l1.getTotalMem());
445+
Assert.assertEquals(0, l1.getNumPreCalls());
446+
Assert.assertEquals(0, l1.getNumCalls());
447+
Assert.assertEquals(0, l1.getNumReleaseCalls());
448+
Assert.assertEquals(0, l1.getNumChildren());
449+
Assert.assertEquals(0, l1.getTotalMem());
456450
CountingAllocationListener l2 = new CountingAllocationListener();
457-
assertEquals(0, l2.getNumPreCalls());
458-
assertEquals(0, l2.getNumCalls());
459-
assertEquals(0, l2.getNumReleaseCalls());
460-
assertEquals(0, l2.getNumChildren());
461-
assertEquals(0, l2.getTotalMem());
451+
Assert.assertEquals(0, l2.getNumPreCalls());
452+
Assert.assertEquals(0, l2.getNumCalls());
453+
Assert.assertEquals(0, l2.getNumReleaseCalls());
454+
Assert.assertEquals(0, l2.getNumChildren());
455+
Assert.assertEquals(0, l2.getTotalMem());
462456
// root and first-level child share the first listener
463457
// second-level and third-level child share the second listener
464458
try (final RootAllocator rootAllocator = new RootAllocator(l1, MAX_ALLOCATION)) {
465459
try (final BufferAllocator c1 = rootAllocator.newChildAllocator("c1", 0, MAX_ALLOCATION)) {
466-
assertEquals(1, l1.getNumChildren());
460+
Assert.assertEquals(1, l1.getNumChildren());
467461
final ArrowBuf buf1 = c1.buffer(16);
468462
assertNotNull("allocation failed", buf1);
469-
assertEquals(1, l1.getNumPreCalls());
470-
assertEquals(1, l1.getNumCalls());
471-
assertEquals(0, l1.getNumReleaseCalls());
472-
assertEquals(16, l1.getTotalMem());
463+
Assert.assertEquals(1, l1.getNumPreCalls());
464+
Assert.assertEquals(1, l1.getNumCalls());
465+
Assert.assertEquals(0, l1.getNumReleaseCalls());
466+
Assert.assertEquals(16, l1.getTotalMem());
473467
buf1.getReferenceManager().release();
474468
try (final BufferAllocator c2 = c1.newChildAllocator("c2", l2, 0, MAX_ALLOCATION)) {
475-
assertEquals(2, l1.getNumChildren()); // c1 got a new child, so c1's listener (l1) is notified
476-
assertEquals(0, l2.getNumChildren());
469+
Assert.assertEquals(2, l1.getNumChildren()); // c1 got a new child, so c1's listener (l1) is notified
470+
Assert.assertEquals(0, l2.getNumChildren());
477471
final ArrowBuf buf2 = c2.buffer(32);
478472
assertNotNull("allocation failed", buf2);
479-
assertEquals(1, l1.getNumCalls());
480-
assertEquals(16, l1.getTotalMem());
481-
assertEquals(1, l2.getNumPreCalls());
482-
assertEquals(1, l2.getNumCalls());
483-
assertEquals(0, l2.getNumReleaseCalls());
484-
assertEquals(32, l2.getTotalMem());
473+
Assert.assertEquals(1, l1.getNumCalls());
474+
Assert.assertEquals(16, l1.getTotalMem());
475+
Assert.assertEquals(1, l2.getNumPreCalls());
476+
Assert.assertEquals(1, l2.getNumCalls());
477+
Assert.assertEquals(0, l2.getNumReleaseCalls());
478+
Assert.assertEquals(32, l2.getTotalMem());
485479
buf2.getReferenceManager().release();
486480
try (final BufferAllocator c3 = c2.newChildAllocator("c3", 0, MAX_ALLOCATION)) {
487-
assertEquals(2, l1.getNumChildren());
488-
assertEquals(1, l2.getNumChildren());
481+
Assert.assertEquals(2, l1.getNumChildren());
482+
Assert.assertEquals(1, l2.getNumChildren());
489483
final ArrowBuf buf3 = c3.buffer(64);
490484
assertNotNull("allocation failed", buf3);
491-
assertEquals(1, l1.getNumPreCalls());
492-
assertEquals(1, l1.getNumCalls());
493-
assertEquals(1, l1.getNumReleaseCalls());
494-
assertEquals(16, l1.getTotalMem());
495-
assertEquals(2, l2.getNumPreCalls());
496-
assertEquals(2, l2.getNumCalls());
497-
assertEquals(1, l2.getNumReleaseCalls());
498-
assertEquals(32 + 64, l2.getTotalMem());
485+
Assert.assertEquals(1, l1.getNumPreCalls());
486+
Assert.assertEquals(1, l1.getNumCalls());
487+
Assert.assertEquals(1, l1.getNumReleaseCalls());
488+
Assert.assertEquals(16, l1.getTotalMem());
489+
Assert.assertEquals(2, l2.getNumPreCalls());
490+
Assert.assertEquals(2, l2.getNumCalls());
491+
Assert.assertEquals(1, l2.getNumReleaseCalls());
492+
Assert.assertEquals(32 + 64, l2.getTotalMem());
499493
buf3.getReferenceManager().release();
500494
}
501-
assertEquals(2, l1.getNumChildren());
502-
assertEquals(0, l2.getNumChildren()); // third-level child removed
495+
Assert.assertEquals(2, l1.getNumChildren());
496+
Assert.assertEquals(0, l2.getNumChildren()); // third-level child removed
503497
}
504-
assertEquals(1, l1.getNumChildren()); // second-level child removed
505-
assertEquals(0, l2.getNumChildren());
498+
Assert.assertEquals(1, l1.getNumChildren()); // second-level child removed
499+
Assert.assertEquals(0, l2.getNumChildren());
506500
}
507-
assertEquals(0, l1.getNumChildren()); // first-level child removed
501+
Assert.assertEquals(0, l1.getNumChildren()); // first-level child removed
508502

509-
assertEquals(2, l2.getNumReleaseCalls());
503+
Assert.assertEquals(2, l2.getNumReleaseCalls());
510504
}
511505
}
512506

513507
@Test
514508
public void testRootAllocator_listenerAllocationFail() throws Exception {
515509
CountingAllocationListener l1 = new CountingAllocationListener();
516-
assertEquals(0, l1.getNumCalls());
517-
assertEquals(0, l1.getTotalMem());
510+
Assert.assertEquals(0, l1.getNumCalls());
511+
Assert.assertEquals(0, l1.getTotalMem());
518512
// Test attempts to allocate too much from a child whose limit is set to half of the max
519513
// allocation. The listener's callback triggers, expanding the child allocator's limit, so then
520514
// the allocation succeeds.
@@ -527,14 +521,14 @@ public void testRootAllocator_listenerAllocationFail() throws Exception {
527521
} catch (OutOfMemoryException e) {
528522
// expected
529523
}
530-
assertEquals(0, l1.getNumCalls());
531-
assertEquals(0, l1.getTotalMem());
524+
Assert.assertEquals(0, l1.getNumCalls());
525+
Assert.assertEquals(0, l1.getTotalMem());
532526

533527
l1.setExpandOnFail(c1, MAX_ALLOCATION);
534528
ArrowBuf arrowBuf = c1.buffer(MAX_ALLOCATION);
535529
assertNotNull("allocation failed", arrowBuf);
536-
assertEquals(1, l1.getNumCalls());
537-
assertEquals(MAX_ALLOCATION, l1.getTotalMem());
530+
Assert.assertEquals(1, l1.getNumCalls());
531+
Assert.assertEquals(MAX_ALLOCATION, l1.getTotalMem());
538532
arrowBuf.getReferenceManager().release();
539533
}
540534
}
@@ -1098,42 +1092,6 @@ public void testMemoryLeakWithReservation() throws Exception {
10981092
}
10991093
}
11001094

1101-
@Test
1102-
public void testMemoryUsage() {
1103-
ListAppender<ILoggingEvent> memoryLogsAppender = new ListAppender<>();
1104-
Logger logger = (Logger) LoggerFactory.getLogger("arrow.allocator");
1105-
try {
1106-
logger.setLevel(Level.TRACE);
1107-
logger.addAppender(memoryLogsAppender);
1108-
memoryLogsAppender.start();
1109-
try (ArrowBuf buf = new ArrowBuf(ReferenceManager.NO_OP, null,
1110-
1024, new PooledByteBufAllocatorL().empty.memoryAddress())) {
1111-
buf.memoryAddress();
1112-
}
1113-
boolean result = false;
1114-
long startTime = System.currentTimeMillis();
1115-
while ((System.currentTimeMillis() - startTime) < 10000) { // 10 seconds maximum for time to read logs
1116-
result = memoryLogsAppender.list.stream()
1117-
.anyMatch(
1118-
log -> log.toString().contains("Memory Usage: \n") &&
1119-
log.toString().contains("Large buffers outstanding: ") &&
1120-
log.toString().contains("Normal buffers outstanding: ") &&
1121-
log.getLevel().equals(Level.TRACE)
1122-
);
1123-
if (result) {
1124-
break;
1125-
}
1126-
}
1127-
assertTrue("Log messages are:\n" +
1128-
memoryLogsAppender.list.stream().map(ILoggingEvent::toString).collect(Collectors.joining("\n")),
1129-
result);
1130-
} finally {
1131-
memoryLogsAppender.stop();
1132-
logger.detachAppender(memoryLogsAppender);
1133-
logger.setLevel(null);
1134-
}
1135-
}
1136-
11371095
@Test
11381096
public void testOverlimit() {
11391097
try (BufferAllocator allocator = new RootAllocator(1024)) {

‎java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestForeignAllocation.java ‎java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestForeignAllocation.java

+7
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,13 @@
2323
import java.util.ArrayList;
2424
import java.util.List;
2525

26+
import org.apache.arrow.memory.AllocationListener;
27+
import org.apache.arrow.memory.AllocationOutcome;
28+
import org.apache.arrow.memory.ArrowBuf;
29+
import org.apache.arrow.memory.BufferAllocator;
30+
import org.apache.arrow.memory.ForeignAllocation;
31+
import org.apache.arrow.memory.OutOfMemoryException;
32+
import org.apache.arrow.memory.RootAllocator;
2633
import org.apache.arrow.memory.util.MemoryUtil;
2734
import org.junit.After;
2835
import org.junit.Before;

‎java/memory/memory-netty-buffer-patch/pom.xml

+44
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!-- Licensed to the Apache Software Foundation (ASF) under one or more contributor
3+
license agreements. See the NOTICE file distributed with this work for additional
4+
information regarding copyright ownership. The ASF licenses this file to
5+
You under the Apache License, Version 2.0 (the "License"); you may not use
6+
this file except in compliance with the License. You may obtain a copy of
7+
the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required
8+
by applicable law or agreed to in writing, software distributed under the
9+
License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS
10+
OF ANY KIND, either express or implied. See the License for the specific
11+
language governing permissions and limitations under the License. -->
12+
<project xmlns="http://maven.apache.org/POM/4.0.0"
13+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
14+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
15+
<parent>
16+
<artifactId>arrow-memory</artifactId>
17+
<groupId>org.apache.arrow</groupId>
18+
<version>15.0.0-SNAPSHOT</version>
19+
</parent>
20+
<modelVersion>4.0.0</modelVersion>
21+
22+
<artifactId>arrow-memory-netty-buffer-patch</artifactId>
23+
<name>Arrow Memory - Netty Buffer</name>
24+
<description>Netty Buffer needed to patch that is consumed by Arrow Memory Netty</description>
25+
26+
<dependencies>
27+
<dependency>
28+
<groupId>org.apache.arrow</groupId>
29+
<artifactId>arrow-memory-core</artifactId>
30+
</dependency>
31+
<dependency>
32+
<groupId>io.netty</groupId>
33+
<artifactId>netty-buffer</artifactId>
34+
</dependency>
35+
<dependency>
36+
<groupId>io.netty</groupId>
37+
<artifactId>netty-common</artifactId>
38+
</dependency>
39+
<dependency>
40+
<groupId>org.slf4j</groupId>
41+
<artifactId>slf4j-api</artifactId>
42+
</dependency>
43+
</dependencies>
44+
</project>

‎java/memory/memory-netty/src/main/java/io/netty/buffer/ExpandableByteBuf.java ‎java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/ExpandableByteBuf.java

File renamed without changes.

‎java/memory/memory-netty/src/main/java/io/netty/buffer/LargeBuffer.java ‎java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/LargeBuffer.java

File renamed without changes.

‎java/memory/memory-netty/src/main/java/io/netty/buffer/MutableWrappedByteBuf.java ‎java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/MutableWrappedByteBuf.java

File renamed without changes.

0 commit comments

Comments
 (0)