Skip to content

Commit 306eaef

Browse files
hangc0276zymap
hangc0276
authored and
zymap
committed
Unify ByteBufAllocator for the DirectIO component (#3985)
### Motivation Some classes in the DirectIO component use `PooledByteBufAllocator` to create ByteBuf instead of using the BookKeeper-initiated allocator. When we configured the BookKeeper allocator OOM policy to shut down, the OOM policy can't apply to those classes which use `PooledByteBufAllocator` to create ByteBuf. ### Modifications - Unify ByteBufAllocator for the DirectIO component. (cherry picked from commit bf06642)
1 parent cec8551 commit 306eaef

File tree

10 files changed

+45
-39
lines changed

10 files changed

+45
-39
lines changed

‎bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/directentrylogger/Buffer.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import static org.apache.bookkeeper.common.util.ExceptionMessageHelper.exMsg;
2525

2626
import io.netty.buffer.ByteBuf;
27-
import io.netty.buffer.PooledByteBufAllocator;
27+
import io.netty.buffer.ByteBufAllocator;
2828
import io.netty.util.ReferenceCountUtil;
2929
import java.io.IOException;
3030
import java.nio.ByteBuffer;
@@ -60,12 +60,14 @@ class Buffer {
6060
final int bufferSize;
6161
ByteBuf buffer;
6262
ByteBuffer byteBuffer;
63+
ByteBufAllocator allocator;
6364
long pointer = 0;
6465

65-
Buffer(NativeIO nativeIO, int bufferSize) throws IOException {
66+
Buffer(NativeIO nativeIO, ByteBufAllocator allocator, int bufferSize) throws IOException {
6667
checkArgument(isAligned(bufferSize),
6768
"Buffer size not aligned %d", bufferSize);
6869

70+
this.allocator = allocator;
6971
this.buffer = allocateAligned(ALIGNMENT, bufferSize);
7072
this.nativeIO = nativeIO;
7173
this.bufferSize = bufferSize;
@@ -74,7 +76,7 @@ class Buffer {
7476
}
7577

7678
private ByteBuf allocateAligned(int alignment, int bufferSize) {
77-
ByteBuf buf = PooledByteBufAllocator.DEFAULT.directBuffer(bufferSize + alignment);
79+
ByteBuf buf = allocator.directBuffer(bufferSize + alignment);
7880
long addr = buf.memoryAddress();
7981
if ((addr & (alignment - 1)) == 0) {
8082
// The address is already aligned

‎bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/directentrylogger/BufferPool.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
*/
2121
package org.apache.bookkeeper.bookie.storage.directentrylogger;
2222

23+
import io.netty.buffer.ByteBufAllocator;
2324
import java.io.IOException;
2425
import java.util.concurrent.ArrayBlockingQueue;
2526
import org.apache.bookkeeper.common.util.nativeio.NativeIO;
@@ -30,10 +31,10 @@
3031
public class BufferPool implements AutoCloseable {
3132
private final ArrayBlockingQueue<Buffer> pool;
3233

33-
BufferPool(NativeIO nativeIO, int bufferSize, int maxPoolSize) throws IOException {
34+
BufferPool(NativeIO nativeIO, ByteBufAllocator allocator, int bufferSize, int maxPoolSize) throws IOException {
3435
pool = new ArrayBlockingQueue<>(maxPoolSize);
3536
for (int i = 0; i < maxPoolSize; i++) {
36-
pool.add(new Buffer(nativeIO, bufferSize));
37+
pool.add(new Buffer(nativeIO, allocator, bufferSize));
3738
}
3839
}
3940

‎bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/directentrylogger/DirectCompactionEntryLog.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ public void markCompacted() throws IOException {
207207
public void scan(EntryLogScanner scanner) throws IOException {
208208
try (LogReader reader = new DirectReader(dstLogId, compactedFile.toString(), allocator, nativeIO,
209209
readBufferSize, maxSaneEntrySize, readBlockStats)) {
210-
LogReaderScan.scan(reader, scanner);
210+
LogReaderScan.scan(allocator, reader, scanner);
211211
}
212212
}
213213
}

‎bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/directentrylogger/DirectEntryLogger.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ public DirectEntryLogger(File ledgerDir,
119119
this.allocator = allocator;
120120

121121
int singleWriteBufferSize = Buffer.nextAlignment((int) (totalWriteBufferSize / NUMBER_OF_WRITE_BUFFERS));
122-
this.writeBuffers = new BufferPool(nativeIO, singleWriteBufferSize, NUMBER_OF_WRITE_BUFFERS);
122+
this.writeBuffers = new BufferPool(nativeIO, allocator, singleWriteBufferSize, NUMBER_OF_WRITE_BUFFERS);
123123

124124
// The total read buffer memory needs to get split across all the read threads, since the caches
125125
// are thread-specific and we want to ensure we don't pass the total memory limit.
@@ -385,7 +385,7 @@ public boolean removeEntryLog(long entryLogId) {
385385
public void scanEntryLog(long entryLogId, EntryLogScanner scanner) throws IOException {
386386
checkArgument(entryLogId < Integer.MAX_VALUE, "Entry log id must be an int [%d]", entryLogId);
387387
try (LogReader reader = newDirectReader((int) entryLogId)) {
388-
LogReaderScan.scan(reader, scanner);
388+
LogReaderScan.scan(allocator, reader, scanner);
389389
}
390390
}
391391

‎bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/directentrylogger/DirectReader.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ class DirectReader implements LogReader {
6969
.kv("errno", ne.getErrno()).toString());
7070
}
7171
refreshMaxOffset();
72-
nativeBuffer = new Buffer(nativeIO, bufferSize);
72+
nativeBuffer = new Buffer(nativeIO, allocator, bufferSize);
7373
}
7474

7575
@Override

‎bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/storage/directentrylogger/LogReaderScan.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,16 +21,16 @@
2121
package org.apache.bookkeeper.bookie.storage.directentrylogger;
2222

2323
import io.netty.buffer.ByteBuf;
24-
import io.netty.buffer.PooledByteBufAllocator;
24+
import io.netty.buffer.ByteBufAllocator;
2525
import io.netty.util.ReferenceCountUtil;
2626
import java.io.IOException;
2727
import org.apache.bookkeeper.bookie.storage.EntryLogScanner;
2828

2929
class LogReaderScan {
30-
static void scan(LogReader reader, EntryLogScanner scanner) throws IOException {
30+
static void scan(ByteBufAllocator allocator, LogReader reader, EntryLogScanner scanner) throws IOException {
3131
int offset = Header.LOGFILE_LEGACY_HEADER_SIZE;
3232

33-
ByteBuf entry = PooledByteBufAllocator.DEFAULT.directBuffer(16 * 1024 * 1024);
33+
ByteBuf entry = allocator.directBuffer(16 * 1024 * 1024);
3434

3535
try {
3636
while (offset < reader.maxOffset()) {

‎bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/storage/directentrylogger/TestBuffer.java

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
// CHECKSTYLE.OFF: IllegalImport
2828
import io.netty.buffer.ByteBuf;
29+
import io.netty.buffer.ByteBufAllocator;
2930
import io.netty.buffer.Unpooled;
3031
import io.netty.util.internal.PlatformDependent;
3132
import java.io.IOException;
@@ -70,13 +71,13 @@ public void testMaxAlignment() throws Exception {
7071

7172
@Test(expected = IllegalArgumentException.class)
7273
public void testCreateUnaligned() throws Exception {
73-
new Buffer(new NativeIOImpl(), 1234);
74+
new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, 1234);
7475
}
7576

7677
@Test
7778
public void testWriteInt() throws Exception {
7879
int bufferSize = 1 << 20;
79-
Buffer b = new Buffer(new NativeIOImpl(), bufferSize);
80+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, bufferSize);
8081
assertTrue(b.hasSpace(bufferSize));
8182
assertEquals(0, b.position());
8283
b.writeInt(0xdeadbeef);
@@ -111,7 +112,7 @@ public void testWriteBuffer() throws Exception {
111112
ByteBuf bb = Unpooled.buffer(1021);
112113
fillByteBuf(bb, 0xdeadbeef);
113114
int bufferSize = 1 << 20;
114-
Buffer b = new Buffer(new NativeIOImpl(), bufferSize);
115+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, bufferSize);
115116
assertEquals(0, b.position());
116117
b.writeByteBuf(bb);
117118
assertEquals(1021, b.position());
@@ -138,7 +139,7 @@ public void testWriteBuffer() throws Exception {
138139
public void testPartialRead() throws Exception {
139140
ByteBuf bb = Unpooled.buffer(5000);
140141

141-
Buffer b = new Buffer(new NativeIOImpl(), 4096);
142+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, 4096);
142143
for (int i = 0; i < 4096 / Integer.BYTES; i++) {
143144
b.writeInt(0xdeadbeef);
144145
}
@@ -149,7 +150,7 @@ public void testPartialRead() throws Exception {
149150

150151
@Test(expected = IOException.class)
151152
public void testReadIntAtBoundary() throws Exception {
152-
Buffer b = new Buffer(new NativeIOImpl(), 4096);
153+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, 4096);
153154

154155
for (int i = 0; i < 4096 / Integer.BYTES; i++) {
155156
b.writeInt(0xdeadbeef);
@@ -163,7 +164,7 @@ public void testReadIntAtBoundary() throws Exception {
163164

164165
@Test(expected = IOException.class)
165166
public void testReadLongAtBoundary() throws Exception {
166-
Buffer b = new Buffer(new NativeIOImpl(), 4096);
167+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, 4096);
167168

168169
for (int i = 0; i < 4096 / Integer.BYTES; i++) {
169170
b.writeInt(0xdeadbeef);
@@ -177,7 +178,7 @@ public void testReadLongAtBoundary() throws Exception {
177178

178179
@Test
179180
public void testPadToAlignment() throws Exception {
180-
Buffer b = new Buffer(new NativeIOImpl(), 1 << 23);
181+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, 1 << 23);
181182

182183
for (int i = 0; i < 1025; i++) {
183184
b.writeInt(0xdededede);
@@ -194,7 +195,7 @@ public void testPadToAlignment() throws Exception {
194195

195196
@Test
196197
public void testFree() throws Exception {
197-
Buffer b = new Buffer(new NativeIOImpl(), 1 << 23);
198+
Buffer b = new Buffer(new NativeIOImpl(), ByteBufAllocator.DEFAULT, 1 << 23);
198199
b.free(); // success if process doesn't explode
199200
b.free();
200201
}

‎bookkeeper-server/src/test/java/org/apache/bookkeeper/bookie/storage/directentrylogger/TestDirectReader.java

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,7 @@ public void testReadBufferAcrossBoundary() throws Exception {
189189
File ledgerDir = tmpDirs.createNew("readBuffer", "logs");
190190

191191
writeFileWithPattern(ledgerDir, 1234, 0xbeefcafe, 1, 1 << 20);
192-
BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT * 4, 8);
192+
BufferPool buffers = new BufferPool(new NativeIOImpl(), ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT * 4, 8);
193193

194194
try (LogReader reader = new DirectReader(1234, logFilename(ledgerDir, 1234),
195195
ByteBufAllocator.DEFAULT,
@@ -268,7 +268,7 @@ public void testReadEntries() throws Exception {
268268

269269
int entrySize = Buffer.ALIGNMENT / 4 + 100;
270270
Map<Integer, Integer> offset2Pattern = new HashMap<>();
271-
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT, 8);
271+
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT, 8);
272272
LogWriter writer = new DirectWriter(1234, logFilename(ledgerDir, 1234),
273273
1 << 20, MoreExecutors.newDirectExecutorService(),
274274
buffers, new NativeIOImpl(), Slogger.CONSOLE)) {
@@ -315,7 +315,7 @@ public int fallocate(int fd, int mode, long offset, long len)
315315
return 0;
316316
}
317317
};
318-
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT, 8);
318+
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT, 8);
319319
LogWriter writer = new DirectWriter(1234, logFilename(ledgerDir, 1234),
320320
1 << 20, MoreExecutors.newDirectExecutorService(),
321321
buffers, new NativeIOImpl(), Slogger.CONSOLE);
@@ -353,7 +353,7 @@ public void testReadFromFileBeingWrittenReadInPreallocated() throws Exception {
353353

354354
int entrySize = Buffer.ALIGNMENT / 2 + 8;
355355

356-
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT, 8);
356+
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT, 8);
357357
LogWriter writer = new DirectWriter(1234, logFilename(ledgerDir, 1234),
358358
1 << 20, MoreExecutors.newDirectExecutorService(),
359359
buffers, new NativeIOImpl(), Slogger.CONSOLE);
@@ -403,7 +403,8 @@ public int fallocate(int fd, int mode, long offset, long len)
403403
return 0; // don't preallocate
404404
}
405405
};
406-
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT * 10, 8);
406+
try (BufferPool buffers = new BufferPool(new NativeIOImpl(),
407+
ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT * 10, 8);
407408
LogWriter writer = new DirectWriter(1234, logFilename(ledgerDir, 1234), 1 << 20,
408409
MoreExecutors.newDirectExecutorService(),
409410
buffers, new NativeIOImpl(), Slogger.CONSOLE)) {
@@ -452,7 +453,7 @@ public void testLargeEntry() throws Exception {
452453
int entrySize = Buffer.ALIGNMENT * 4;
453454

454455
int offset1, offset2;
455-
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT * 8, 8);
456+
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT * 8, 8);
456457
LogWriter writer = new DirectWriter(1234, logFilename(ledgerDir, 1234), 1 << 20,
457458
MoreExecutors.newDirectExecutorService(), buffers, new NativeIOImpl(),
458459
Slogger.CONSOLE)) {
@@ -496,7 +497,7 @@ public void testLargeEntry() throws Exception {
496497

497498
private static void writeFileWithPattern(File directory, int logId,
498499
int pattern, int blockIncrement, int fileSize) throws Exception {
499-
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), Buffer.ALIGNMENT, 8);
500+
try (BufferPool buffers = new BufferPool(new NativeIOImpl(), ByteBufAllocator.DEFAULT, Buffer.ALIGNMENT, 8);
500501
LogWriter writer = new DirectWriter(logId, logFilename(directory, logId),
501502
fileSize, MoreExecutors.newDirectExecutorService(),
502503
buffers, new NativeIOImpl(), Slogger.CONSOLE)) {

0 commit comments

Comments
 (0)