Allocator support exitOnOutOfMemory config. (#3984)

* Allocator support exitOnOutOfMemory config.

bookkeeper-common-allocator/pom.xml

@@ -29,6 +29,11 @@
  <groupId>io.netty</groupId>
  <artifactId>netty-buffer</artifactId>
  </dependency>
+ <dependency>
+ <groupId>org.mockito</groupId>
+ <artifactId>mockito-core</artifactId>
+ <scope>test</scope>
+ </dependency>
  </dependencies>
  <build>
  <plugins>

bookkeeper-common-allocator/src/main/java/org/apache/bookkeeper/common/allocator/ByteBufAllocatorBuilder.java

@@ -92,4 +92,6 @@ static ByteBufAllocatorBuilder create() {
  * <p>Default is {@link LeakDetectionPolicy#Disabled}
  */
  ByteBufAllocatorBuilder leakDetectionPolicy(LeakDetectionPolicy leakDetectionPolicy);
+
+ ByteBufAllocatorBuilder exitOnOutOfMemory(boolean exitOnOutOfMemory);
 }

bookkeeper-common-allocator/src/main/java/org/apache/bookkeeper/common/allocator/impl/ByteBufAllocatorBuilderImpl.java

@@ -37,11 +37,12 @@ public class ByteBufAllocatorBuilderImpl implements ByteBufAllocatorBuilder {
  OutOfMemoryPolicy outOfMemoryPolicy = OutOfMemoryPolicy.FallbackToHeap;
  Consumer<OutOfMemoryError> outOfMemoryListener = null;
  LeakDetectionPolicy leakDetectionPolicy = LeakDetectionPolicy.Disabled;
+ boolean exitOnOutOfMemory = false;
 
  @Override
  public ByteBufAllocatorWithOomHandler build() {
  return new ByteBufAllocatorImpl(pooledAllocator, unpooledAllocator, poolingPolicy, poolingConcurrency,
- outOfMemoryPolicy, outOfMemoryListener, leakDetectionPolicy);
+ outOfMemoryPolicy, outOfMemoryListener, leakDetectionPolicy, exitOnOutOfMemory);
  }
 
  @Override
@@ -86,4 +87,10 @@ public ByteBufAllocatorBuilder leakDetectionPolicy(LeakDetectionPolicy leakDetec
  return this;
  }
 
+ @Override
+ public ByteBufAllocatorBuilder exitOnOutOfMemory(boolean exitOnOutOfMemory) {
+ this.exitOnOutOfMemory = exitOnOutOfMemory;
+ return this;
+ }
+
 }

bookkeeper-common-allocator/src/main/java/org/apache/bookkeeper/common/allocator/impl/ByteBufAllocatorImpl.java

@@ -29,6 +29,7 @@
 import org.apache.bookkeeper.common.allocator.LeakDetectionPolicy;
 import org.apache.bookkeeper.common.allocator.OutOfMemoryPolicy;
 import org.apache.bookkeeper.common.allocator.PoolingPolicy;
+import org.apache.bookkeeper.common.util.ShutdownUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -48,15 +49,24 @@ public class ByteBufAllocatorImpl extends AbstractByteBufAllocator implements By
  private final PoolingPolicy poolingPolicy;
  private final OutOfMemoryPolicy outOfMemoryPolicy;
  private Consumer<OutOfMemoryError> outOfMemoryListener;
+ private final boolean exitOnOutOfMemory;
 
  ByteBufAllocatorImpl(ByteBufAllocator pooledAllocator, ByteBufAllocator unpooledAllocator,
  PoolingPolicy poolingPolicy, int poolingConcurrency, OutOfMemoryPolicy outOfMemoryPolicy,
  Consumer<OutOfMemoryError> outOfMemoryListener,
  LeakDetectionPolicy leakDetectionPolicy) {
- super(poolingPolicy == PoolingPolicy.PooledDirect /* preferDirect */);
+ this(pooledAllocator, unpooledAllocator, poolingPolicy, poolingConcurrency, outOfMemoryPolicy,
+ outOfMemoryListener, leakDetectionPolicy, false);
+ }
 
+ ByteBufAllocatorImpl(ByteBufAllocator pooledAllocator, ByteBufAllocator unpooledAllocator,
+ PoolingPolicy poolingPolicy, int poolingConcurrency, OutOfMemoryPolicy outOfMemoryPolicy,
+ Consumer<OutOfMemoryError> outOfMemoryListener,
+ LeakDetectionPolicy leakDetectionPolicy, boolean exitOnOutOfMemory) {
+ super(poolingPolicy == PoolingPolicy.PooledDirect /* preferDirect */);
  this.poolingPolicy = poolingPolicy;
  this.outOfMemoryPolicy = outOfMemoryPolicy;
+ this.exitOnOutOfMemory = exitOnOutOfMemory;
  if (outOfMemoryListener == null) {
  this.outOfMemoryListener = (v) -> {
  log.error("Unable to allocate memory", v);
@@ -146,7 +156,7 @@ protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) {
  : unpooledAllocator;
  return alloc.heapBuffer(initialCapacity, maxCapacity);
  } catch (OutOfMemoryError e) {
- outOfMemoryListener.accept(e);
+ consumeOOMError(e);
  throw e;
  }
  }
@@ -166,12 +176,12 @@ private ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity, boolean ca
  try {
  return unpooledAllocator.heapBuffer(initialCapacity, maxCapacity);
  } catch (OutOfMemoryError e2) {
- outOfMemoryListener.accept(e2);
+ consumeOOMError(e2);
  throw e2;
  }
  } else {
  // ThrowException
- outOfMemoryListener.accept(e);
+ consumeOOMError(e);
  throw e;
  }
  }
@@ -181,12 +191,24 @@ private ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity, boolean ca
  try {
  return unpooledAllocator.directBuffer(initialCapacity, maxCapacity);
  } catch (OutOfMemoryError e) {
- outOfMemoryListener.accept(e);
-  throw e;
+ consumeOOMError(e);
+ throw e;
  }
  }
  }
 
+ private void consumeOOMError(OutOfMemoryError outOfMemoryError) {
+ try {
+ outOfMemoryListener.accept(outOfMemoryError);
+ } catch (Throwable e) {
+ log.warn("Consume outOfMemory error failed.", e);
+ }
+ if (exitOnOutOfMemory) {
+ log.info("Exiting JVM process for OOM error: {}", outOfMemoryError.getMessage(), outOfMemoryError);
+ ShutdownUtil.triggerImmediateForcefulShutdown();
+ }
+ }
+
  @Override
  public boolean isDirectBufferPooled() {
  return pooledAllocator != null && pooledAllocator.isDirectBufferPooled();

bookkeeper-common-allocator/src/main/java/org/apache/bookkeeper/common/util/ShutdownUtil.java

@@ -0,0 +1,86 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.bookkeeper.common.util;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Forked from <a href="https://github.com/apache/pulsar">Pulsar</a>.
+ */
+@Slf4j
+public class ShutdownUtil {
+ private static final Method log4j2ShutdownMethod;
+
+ static {
+ // use reflection to find org.apache.logging.log4j.LogManager.shutdown method
+ Method shutdownMethod = null;
+ try {
+ shutdownMethod = Class.forName("org.apache.logging.log4j.LogManager")
+ .getMethod("shutdown");
+ } catch (ClassNotFoundException | NoSuchMethodException e) {
+ // ignore when Log4j2 isn't found, log at debug level
+ log.debug("Cannot find org.apache.logging.log4j.LogManager.shutdown method", e);
+ }
+ log4j2ShutdownMethod = shutdownMethod;
+ }
+
+ /**
+ * Triggers an immediate forceful shutdown of the current process.
+ *
+ * @param status Termination status. By convention, a nonzero status code indicates abnormal termination.
+ * @see Runtime#halt(int)
+ */
+ public static void triggerImmediateForcefulShutdown(int status) {
+ triggerImmediateForcefulShutdown(status, true);
+ }
+ public static void triggerImmediateForcefulShutdown(int status, boolean logging) {
+ try {
+ if (status != 0 && logging) {
+ log.warn("Triggering immediate shutdown of current process with status {}", status,
+ new Exception("Stacktrace for immediate shutdown"));
+ }
+ shutdownLogging();
+ } finally {
+ Runtime.getRuntime().halt(status);
+ }
+ }
+
+ private static void shutdownLogging() {
+ // flush log buffers and shutdown log4j2 logging to prevent log truncation
+ if (log4j2ShutdownMethod != null) {
+ try {
+ // use reflection to call org.apache.logging.log4j.LogManager.shutdown()
+ log4j2ShutdownMethod.invoke(null);
+ } catch (IllegalAccessException | InvocationTargetException e) {
+ log.error("Unable to call org.apache.logging.log4j.LogManager.shutdown using reflection.", e);
+ }
+ }
+ }
+
+ /**
+ * Triggers an immediate forceful shutdown of the current process using 1 as the status code.
+ *
+ * @see Runtime#halt(int)
+ */
+ public static void triggerImmediateForcefulShutdown() {
+ triggerImmediateForcefulShutdown(1);
+ }
+}

bookkeeper-common-allocator/src/main/java/org/apache/bookkeeper/common/util/package-info.java

@@ -0,0 +1,21 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/**
+ * defines the utilities for allocator used across the project.
+ */
+package org.apache.bookkeeper.common.util;

bookkeeper-common-allocator/src/test/java/org/apache/bookkeeper/common/allocator/impl/ByteBufAllocatorBuilderTest.java

@@ -23,6 +23,7 @@
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.when;
 
 import io.netty.buffer.ByteBuf;
@@ -35,7 +36,10 @@
 import org.apache.bookkeeper.common.allocator.ByteBufAllocatorBuilder;
 import org.apache.bookkeeper.common.allocator.OutOfMemoryPolicy;
 import org.apache.bookkeeper.common.allocator.PoolingPolicy;
+import org.apache.bookkeeper.common.util.ShutdownUtil;
 import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 /**
  * Tests for {@link ByteBufAllocatorBuilderImpl}.
@@ -87,6 +91,30 @@ public void testOomWithException() {
  assertEquals(outOfDirectMemException, receivedException.get());
  }
 
+ @Test()
+ public void testOomExit() {
+ ByteBufAllocator baseAlloc = mock(ByteBufAllocator.class);
+ when(baseAlloc.directBuffer(anyInt(), anyInt())).thenThrow(outOfDirectMemException);
+
+ ByteBufAllocator alloc = ByteBufAllocatorBuilder.create()
+ .pooledAllocator(baseAlloc)
+ .outOfMemoryPolicy(OutOfMemoryPolicy.ThrowException)
+ .exitOnOutOfMemory(true)
+ .build();
+
+ MockedStatic<ShutdownUtil> mockedStatic = mockStatic(ShutdownUtil.class);
+
+ try {
+ alloc.buffer();
+ fail("Should have thrown exception");
+ } catch (OutOfMemoryError e) {
+ // Expected
+ assertEquals(outOfDirectMemException, e);
+ }
+
+ mockedStatic.verify(() -> ShutdownUtil.triggerImmediateForcefulShutdown(), Mockito.times(1));
+ }
+
  @Test
  public void testOomWithFallback() {
  ByteBufAllocator baseAlloc = mock(ByteBufAllocator.class);

bookkeeper-server/src/main/java/org/apache/bookkeeper/bookie/BookieResources.java

@@ -70,6 +70,7 @@ public static ByteBufAllocatorWithOomHandler createAllocator(ServerConfiguration
  .poolingConcurrency(conf.getAllocatorPoolingConcurrency())
  .outOfMemoryPolicy(conf.getAllocatorOutOfMemoryPolicy())
  .leakDetectionPolicy(conf.getAllocatorLeakDetectionPolicy())
+ .exitOnOutOfMemory(conf.exitOnOutOfMemory())
  .build();
  }
 

bookkeeper-server/src/main/java/org/apache/bookkeeper/client/BookKeeper.java

@@ -478,6 +478,7 @@ public BookKeeper(ClientConfiguration conf, ZooKeeper zk, EventLoopGroup eventLo
  .poolingConcurrency(conf.getAllocatorPoolingConcurrency())
  .outOfMemoryPolicy(conf.getAllocatorOutOfMemoryPolicy())
  .leakDetectionPolicy(conf.getAllocatorLeakDetectionPolicy())
+ .exitOnOutOfMemory(conf.exitOnOutOfMemory())
  .build();
  }
 

bookkeeper-server/src/main/java/org/apache/bookkeeper/conf/AbstractConfiguration.java

@@ -184,6 +184,7 @@ public abstract class AbstractConfiguration<T extends AbstractConfiguration>
  protected static final String ALLOCATOR_POOLING_CONCURRENCY = "allocatorPoolingConcurrency";
  protected static final String ALLOCATOR_OOM_POLICY = "allocatorOutOfMemoryPolicy";
  protected static final String ALLOCATOR_LEAK_DETECTION_POLICY = "allocatorLeakDetectionPolicy";
+ protected static final String ALLOCATOR_EXIT_ON_OUT_OF_MEMORY = "allocatorExitOnOutOfMemory";
 
  // option to limit stats logging
  public static final String LIMIT_STATS_LOGGING = "limitStatsLogging";
@@ -1157,6 +1158,15 @@ public T setAllocatorLeakDetectionPolicy(LeakDetectionPolicy leakDetectionPolicy
  return getThis();
  }
 
+ public T setExitOnOutOfMemory(boolean exitOnOutOfMemory) {
+ this.setProperty(ALLOCATOR_EXIT_ON_OUT_OF_MEMORY, exitOnOutOfMemory);
+ return getThis();
+ }
+
+ public boolean exitOnOutOfMemory() {
+ return getBoolean(ALLOCATOR_EXIT_ON_OUT_OF_MEMORY, false);
+ }
+
  /**
  * Return whether the busy-wait is enabled for BookKeeper and Netty IO threads.
  *

bookkeeper-server/src/test/java/org/apache/bookkeeper/conf/AbstractConfigurationTest.java

@@ -19,6 +19,8 @@
 package org.apache.bookkeeper.conf;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.CALLS_REAL_METHODS;
 import static org.mockito.Mockito.mock;
 
@@ -179,4 +181,12 @@ public void testAllocatorLeakDetectionPolicy() {
  System.getProperties().put(nettyLevelKey, nettyLevelStr);
  }
  }
+
+ @Test
+ public void testExitOnOutOfMemory() {
+ assertFalse(conf.exitOnOutOfMemory());
+ conf.setExitOnOutOfMemory(true);
+ assertTrue(conf.exitOnOutOfMemory());
+ }
+
 }

tools/perf/src/main/java/org/apache/bookkeeper/tools/perf/journal/JournalWriter.java

@@ -495,6 +495,7 @@ private static ByteBufAllocator getAllocator(ServerConfiguration conf) {
  log.error("Unable to allocate memory, exiting bookie", ex);
  })
  .leakDetectionPolicy(conf.getAllocatorLeakDetectionPolicy())
+ .exitOnOutOfMemory(conf.exitOnOutOfMemory())
  .build();
  }