Skip to content

Commit ef0d629

Browse files
PengznaPengzna
committed
complete allocate memory test and fix bug
1 parent aaeacb5 commit ef0d629

17 files changed

+162
-42
lines changed

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/MemoryManager.java

+6-6
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,8 @@ public class MemoryManager {
4646
public static final long MAX_MEMORY_CAPACITY_IN_BYTES = Bytes.GB;
4747
private final AtomicLong currentAvailableMemoryInBytes =
4848
new AtomicLong(MAX_MEMORY_CAPACITY_IN_BYTES);
49-
private final AtomicLong currentOffHeapAllocatedMemory = new AtomicLong(0);
50-
private final AtomicLong currentOnHeapAllocatedMemory = new AtomicLong(0);
49+
private final AtomicLong currentOffHeapAllocatedMemoryInBytes = new AtomicLong(0);
50+
private final AtomicLong currentOnHeapAllocatedMemoryInBytes = new AtomicLong(0);
5151
private final Queue<MemoryPool> queryMemoryPools =
5252
new PriorityQueue<>((o1, o2) -> (int) (o2.getFreeBytes() - o1.getFreeBytes()));
5353
private final MemoryArbitrator memoryArbitrator;
@@ -136,12 +136,12 @@ public void consumeAvailableMemory(long size) {
136136
currentAvailableMemoryInBytes.addAndGet(-size);
137137
}
138138

139-
public AtomicLong getCurrentOnHeapAllocatedMemory() {
140-
return currentOnHeapAllocatedMemory;
139+
public AtomicLong getCurrentOnHeapAllocatedMemoryInBytes() {
140+
return currentOnHeapAllocatedMemoryInBytes;
141141
}
142142

143-
public AtomicLong getCurrentOffHeapAllocatedMemory() {
144-
return currentOffHeapAllocatedMemory;
143+
public AtomicLong getCurrentOffHeapAllocatedMemoryInBytes() {
144+
return currentOffHeapAllocatedMemoryInBytes;
145145
}
146146

147147
private static class MemoryManagerHolder {

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/allocator/NettyMemoryAllocator.java

+6-12
Original file line numberDiff line numberDiff line change
@@ -37,22 +37,24 @@ public NettyMemoryAllocator(MemoryManager memoryManager) {
3737

3838
@Override
3939
public ByteBuf forceAllocate(long size) {
40+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().addAndGet(size);
4041
return offHeapAllocator.directBuffer((int) size);
4142
}
4243

4344
@Override
4445
public ByteBuf tryToAllocate(long size) {
45-
if (memoryManager.getCurrentOnHeapAllocatedMemory().get() +
46-
memoryManager.getCurrentOffHeapAllocatedMemory().get() + size <
46+
if (memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().get() +
47+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().get() + size <
4748
MemoryManager.MAX_MEMORY_CAPACITY_IN_BYTES) {
49+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().addAndGet(size);
4850
return offHeapAllocator.directBuffer((int) size);
4951
}
5052
return null;
5153
}
5254

5355
@Override
5456
public void returnMemoryToManager(long size) {
55-
memoryManager.getCurrentOffHeapAllocatedMemory().addAndGet(-size);
57+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().addAndGet(-size);
5658
}
5759

5860
@Override
@@ -61,14 +63,6 @@ public void releaseMemoryBlock(Object memoryBlock) {
6163
throw new IllegalArgumentException("memoryBlock must be ByteBuf");
6264
}
6365
ByteBuf buf = (ByteBuf) memoryBlock;
64-
ReferenceCountUtil.safeRelease(buf, ReferenceCountUtil.refCnt(buf));
65-
}
66-
67-
public static void main(String[] args) {
68-
MemoryAllocator netty = new NettyMemoryAllocator(null);
69-
ByteBuf buf = (ByteBuf) netty.forceAllocate(1024);
70-
System.out.println(ReferenceCountUtil.refCnt(buf));
71-
netty.releaseMemoryBlock(buf);
72-
System.out.println(ReferenceCountUtil.refCnt(buf));
66+
ReferenceCountUtil.safeRelease(buf);
7367
}
7468
}

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/allocator/OnHeapMemoryStrategy.java

+14-13
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@
1717

1818
package org.apache.hugegraph.memory.allocator;
1919

20+
import java.util.concurrent.atomic.AtomicReference;
21+
2022
import org.apache.hugegraph.memory.MemoryManager;
21-
import org.apache.hugegraph.util.Bytes;
2223

2324
public class OnHeapMemoryStrategy implements MemoryAllocator {
2425

@@ -29,31 +30,31 @@ public OnHeapMemoryStrategy(MemoryManager memoryManager) {
2930
}
3031

3132
@Override
32-
public byte[] tryToAllocate(long size) {
33-
if (memoryManager.getCurrentOnHeapAllocatedMemory().get() +
34-
memoryManager.getCurrentOffHeapAllocatedMemory().get() + size <
33+
public AtomicReference<byte[]> tryToAllocate(long size) {
34+
if (memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().get() +
35+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().get() + size <
3536
MemoryManager.MAX_MEMORY_CAPACITY_IN_BYTES) {
36-
int sizeOfByte = (int) (size / Bytes.BASE);
37-
memoryManager.getCurrentOnHeapAllocatedMemory().addAndGet(sizeOfByte);
38-
return new byte[sizeOfByte];
37+
memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().addAndGet(size);
38+
byte[] memoryBlock = new byte[(int) size];
39+
return new AtomicReference<>(memoryBlock);
3940
}
4041
return null;
4142
}
4243

4344
@Override
44-
public byte[] forceAllocate(long size) {
45-
int sizeOfByte = (int) (size / Bytes.BASE);
46-
memoryManager.getCurrentOnHeapAllocatedMemory().addAndGet(sizeOfByte);
47-
return new byte[sizeOfByte];
45+
public AtomicReference<byte[]> forceAllocate(long size) {
46+
memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().addAndGet(size);
47+
byte[] memoryBlock = new byte[(int) size];
48+
return new AtomicReference<>(memoryBlock);
4849
}
4950

5051
@Override
5152
public void returnMemoryToManager(long size) {
52-
memoryManager.getCurrentOnHeapAllocatedMemory().addAndGet(-size);
53+
memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().addAndGet(-size);
5354
}
5455

5556
@Override
5657
public void releaseMemoryBlock(Object memoryBlock) {
57-
memoryBlock = null;
58+
((AtomicReference<byte[]>) memoryBlock).set(null);
5859
}
5960
}

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/MemoryConsumer.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -49,5 +49,5 @@ public interface MemoryConsumer {
4949

5050
MemoryPool getOperatorMemoryPool();
5151

52-
List<ByteBuf> getAllOffHeapByteBuf();
52+
List<ByteBuf> getAllMemoryBlock();
5353
}

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/BinaryIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ public MemoryPool getOperatorMemoryPool() {
6363
}
6464

6565
@Override
66-
public List<ByteBuf> getAllOffHeapByteBuf() {
66+
public List<ByteBuf> getAllMemoryBlock() {
6767
return Collections.singletonList(bytesOffHeap);
6868
}
6969

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/EdgeIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,7 @@ public MemoryPool getOperatorMemoryPool() {
144144
}
145145

146146
@Override
147-
public List<ByteBuf> getAllOffHeapByteBuf() {
147+
public List<ByteBuf> getAllMemoryBlock() {
148148
return Lists.newArrayList(this.sortValuesOffHeap, this.cacheOffHeap);
149149
}
150150

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/LongIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public MemoryPool getOperatorMemoryPool() {
7878
}
7979

8080
@Override
81-
public List<ByteBuf> getAllOffHeapByteBuf() {
81+
public List<ByteBuf> getAllMemoryBlock() {
8282
return Collections.singletonList(idOffHeap);
8383
}
8484

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/ObjectIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ public MemoryPool getOperatorMemoryPool() {
6666
}
6767

6868
@Override
69-
public List<ByteBuf> getAllOffHeapByteBuf() {
69+
public List<ByteBuf> getAllMemoryBlock() {
7070
return Collections.singletonList(objectOffHeap);
7171
}
7272

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/QueryIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ public MemoryPool getOperatorMemoryPool() {
7272
}
7373

7474
@Override
75-
public List<ByteBuf> getAllOffHeapByteBuf() {
75+
public List<ByteBuf> getAllMemoryBlock() {
7676
return Collections.singletonList(queryOffHeap);
7777
}
7878

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/StringIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ public MemoryPool getOperatorMemoryPool() {
7878
}
7979

8080
@Override
81-
public List<ByteBuf> getAllOffHeapByteBuf() {
81+
public List<ByteBuf> getAllMemoryBlock() {
8282
return Collections.singletonList(idOffHeap);
8383
}
8484

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/UuidIdOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ public MemoryPool getOperatorMemoryPool() {
9090
}
9191

9292
@Override
93-
public List<ByteBuf> getAllOffHeapByteBuf() {
93+
public List<ByteBuf> getAllMemoryBlock() {
9494
return Collections.singletonList(idOffHeap);
9595
}
9696

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/property/HugeEdgePropertyOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ public MemoryPool getOperatorMemoryPool() {
7171
}
7272

7373
@Override
74-
public List<ByteBuf> getAllOffHeapByteBuf() {
74+
public List<ByteBuf> getAllMemoryBlock() {
7575
return Collections.singletonList(valueOffHeap);
7676
}
7777

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/property/HugeVertexPropertyOffHeap.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public MemoryPool getOperatorMemoryPool() {
7070
}
7171

7272
@Override
73-
public List<ByteBuf> getAllOffHeapByteBuf() {
73+
public List<ByteBuf> getAllMemoryBlock() {
7474
return Collections.singletonList(valueOffHeap);
7575
}
7676

‎hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/pool/impl/OperatorMemoryPool.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public void bindMemoryConsumer(MemoryConsumer memoryConsumer) {
5252
public synchronized void releaseSelf(String reason) {
5353
this.memoryAllocator.returnMemoryToManager(getAllocatedBytes());
5454
this.memoryConsumers.forEach(memoryConsumer -> {
55-
memoryConsumer.getAllOffHeapByteBuf().forEach(memoryAllocator::releaseMemoryBlock);
55+
memoryConsumer.getAllMemoryBlock().forEach(memoryAllocator::releaseMemoryBlock);
5656
});
5757
this.memoryConsumers.clear();
5858
super.releaseSelf(reason);

‎hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/memory/MemoryAllocateTest.java

+121
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.hugegraph.core.memory;
19+
20+
import java.util.concurrent.atomic.AtomicReference;
21+
22+
import org.apache.hugegraph.memory.MemoryManager;
23+
import org.apache.hugegraph.memory.allocator.MemoryAllocator;
24+
import org.apache.hugegraph.memory.allocator.NettyMemoryAllocator;
25+
import org.apache.hugegraph.memory.allocator.OnHeapMemoryStrategy;
26+
import org.apache.hugegraph.util.Bytes;
27+
import org.junit.Assert;
28+
import org.junit.Before;
29+
import org.junit.BeforeClass;
30+
import org.junit.Test;
31+
32+
import io.netty.buffer.ByteBuf;
33+
import io.netty.util.IllegalReferenceCountException;
34+
import io.netty.util.ReferenceCountUtil;
35+
36+
public class MemoryAllocateTest {
37+
38+
private static MemoryAllocator nettyAllocator;
39+
private static MemoryAllocator onHeapAllocator;
40+
private static MemoryManager memoryManager;
41+
42+
@Before
43+
public void setup() {
44+
MemoryManager.getInstance().getCurrentOffHeapAllocatedMemoryInBytes().set(0);
45+
MemoryManager.getInstance().getCurrentOnHeapAllocatedMemoryInBytes().set(0);
46+
}
47+
48+
@BeforeClass
49+
public static void beforeClass() {
50+
nettyAllocator = new NettyMemoryAllocator(MemoryManager.getInstance());
51+
onHeapAllocator = new OnHeapMemoryStrategy(MemoryManager.getInstance());
52+
memoryManager = MemoryManager.getInstance();
53+
}
54+
55+
@Test
56+
public void testNettyAllocate() {
57+
ByteBuf memoryBlock1 = (ByteBuf) nettyAllocator.tryToAllocate(Bytes.KB);
58+
ByteBuf memoryBlock2 = (ByteBuf) nettyAllocator.tryToAllocate(Bytes.MB);
59+
Assert.assertNotNull(memoryBlock1);
60+
Assert.assertEquals(Bytes.KB, memoryBlock1.capacity());
61+
Assert.assertNotNull(memoryBlock2);
62+
Assert.assertEquals(Bytes.MB, memoryBlock2.capacity());
63+
Assert.assertEquals(Bytes.KB + Bytes.MB,
64+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().get());
65+
ByteBuf memoryBlock3 = (ByteBuf) nettyAllocator.tryToAllocate(Bytes.GB);
66+
Assert.assertNull(memoryBlock3);
67+
memoryBlock3 = (ByteBuf) nettyAllocator.forceAllocate(Bytes.GB);
68+
Assert.assertNotNull(memoryBlock3);
69+
Assert.assertEquals(Bytes.GB, memoryBlock3.capacity());
70+
}
71+
72+
@Test
73+
public void testNettyDeallocate() {
74+
ByteBuf buf = (ByteBuf) nettyAllocator.tryToAllocate(Bytes.KB);
75+
Assert.assertNotNull(buf);
76+
Assert.assertTrue(buf.isWritable());
77+
Assert.assertEquals(buf.capacity(),
78+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().get());
79+
Assert.assertEquals(1, ReferenceCountUtil.refCnt(buf));
80+
nettyAllocator.releaseMemoryBlock(buf);
81+
Assert.assertThrows(IllegalReferenceCountException.class, buf::memoryAddress);
82+
Assert.assertEquals(0, ReferenceCountUtil.refCnt(buf));
83+
nettyAllocator.returnMemoryToManager(buf.capacity());
84+
Assert.assertEquals(0,
85+
memoryManager.getCurrentOffHeapAllocatedMemoryInBytes().get());
86+
}
87+
88+
@Test
89+
public void testOnHeapAllocate() {
90+
AtomicReference<byte[]> memoryBlock1 =
91+
(AtomicReference<byte[]>) onHeapAllocator.tryToAllocate(Bytes.KB);
92+
AtomicReference<byte[]> memoryBlock2 =
93+
(AtomicReference<byte[]>) onHeapAllocator.tryToAllocate(Bytes.MB);
94+
Assert.assertNotNull(memoryBlock1);
95+
Assert.assertEquals(Bytes.KB, memoryBlock1.get().length);
96+
Assert.assertNotNull(memoryBlock2);
97+
Assert.assertEquals(Bytes.MB, memoryBlock2.get().length);
98+
Assert.assertEquals(Bytes.KB + Bytes.MB,
99+
memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().get());
100+
AtomicReference<byte[]> memoryBlock3 =
101+
(AtomicReference<byte[]>) onHeapAllocator.tryToAllocate(Bytes.GB);
102+
Assert.assertNull(memoryBlock3);
103+
memoryBlock3 = (AtomicReference<byte[]>) onHeapAllocator.forceAllocate(Bytes.GB);
104+
Assert.assertNotNull(memoryBlock3);
105+
Assert.assertEquals(Bytes.GB, memoryBlock3.get().length);
106+
}
107+
108+
@Test
109+
public void testOnHeapDeallocate() {
110+
AtomicReference<byte[]> buf =
111+
(AtomicReference<byte[]>) onHeapAllocator.tryToAllocate(Bytes.KB);
112+
Assert.assertNotNull(buf);
113+
Assert.assertEquals(buf.get().length,
114+
memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().get());
115+
onHeapAllocator.returnMemoryToManager(buf.get().length);
116+
Assert.assertEquals(0,
117+
memoryManager.getCurrentOnHeapAllocatedMemoryInBytes().get());
118+
onHeapAllocator.releaseMemoryBlock(buf);
119+
Assert.assertNull(buf.get());
120+
}
121+
}

‎hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/memory/MemoryConsumerTest.java

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
package org.apache.hugegraph.core.memory;public class MemoryConsumerTest {
2+
}

‎hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/memory/MemoryManagerTest.java

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
package org.apache.hugegraph.core.memory;public class MemoryManagerTest {
2+
}

0 commit comments

Comments
 (0)