[pulsar] 04/06: [Tests] Fix OutOfMemoryError and NoClassDefFoundErrors in tests (#15911)

[lh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lhotari pushed a commit to branch branch-2.10
in repository https://gitbox.apache.org/repos/asf/pulsar.git

commit 3f3d4525c495e3da3b98b6929613761f30871a52
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Fri Jun 3 03:52:40 2022 +0300

    [Tests] Fix OutOfMemoryError and NoClassDefFoundErrors in tests (#15911)
    
    Fixes #15689
    Fixes #15793
    Fixes #15897
    
    ### Motivation
    
    `PulsarByteBufAllocator.DEFAULT` instance was replaced in BatchMessageContainerImplTest, but never restored. This caused strange issues in many unit tests. OutOfMemoryError and NoClassDefFoundErrors  were most likely caused by the BatchMessageContainerImplTest.
    
    ### Modifications
    
    Fix the problems in BatchMessageContainerImplTest and reset the state after the test completes.
    
    ### Additional context
    
    The flaky test was introduced by #12170 which has been cherry-picked to branch-2.8, branch-2.9 and included also in branch-2.10 .
    
    (cherry picked from commit ef1e0aa73592b4a81bd89a43628096a839b0bc26)
---
 .../client/impl/BatchMessageContainerImplTest.java | 33 ++++++++++++++++------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java
index 54af93c87e3..8fc018b3199 100644
--- a/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java
+++ b/pulsar-client/src/test/java/org/apache/pulsar/client/impl/BatchMessageContainerImplTest.java
@@ -18,31 +18,38 @@
  */
 package org.apache.pulsar.client.impl;
 
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.when;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.bookkeeper.common.allocator.impl.ByteBufAllocatorImpl;
 import org.apache.pulsar.client.api.CompressionType;
 import org.apache.pulsar.client.api.Schema;
 import org.apache.pulsar.client.impl.conf.ProducerConfigurationData;
+import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
 import org.apache.pulsar.common.api.proto.MessageMetadata;
 import org.mockito.MockedConstruction;
 import org.mockito.Mockito;
+import org.powermock.reflect.Whitebox;
 import org.testng.annotations.Test;
 
-import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
-
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.when;
-
 public class BatchMessageContainerImplTest {
 
     @Test
     public void recoveryAfterOom() throws Exception {
+        AtomicBoolean called = new AtomicBoolean();
         try (MockedConstruction<ByteBufAllocatorImpl> mocked = Mockito.mockConstruction(ByteBufAllocatorImpl.class,
                 (mockAllocator, context) -> {
+                    called.set(true);
                     doThrow(new OutOfMemoryError("test")).when(mockAllocator).buffer(anyInt(), anyInt());
                 })) {
-
+            if (PulsarByteBufAllocator.DEFAULT != null && !called.get()) {
+                replaceByteBufAllocator();
+            }
             final ProducerImpl producer = Mockito.mock(ProducerImpl.class);
             final ProducerConfigurationData producerConfigurationData = new ProducerConfigurationData();
             producerConfigurationData.setCompressionType(CompressionType.NONE);
@@ -64,7 +71,17 @@ public class BatchMessageContainerImplTest {
             final MessageImpl<byte[]> message2 = MessageImpl.create(messageMetadata2, payload2, Schema.BYTES, null);
             // after oom, our add can self-healing, won't throw exception
             batchMessageContainer.add(message2, null);
+        } finally {
+            replaceByteBufAllocator();
         }
+
+    }
+
+    private void replaceByteBufAllocator() throws NoSuchMethodException, IllegalAccessException, InvocationTargetException {
+        Method createByteBufAllocatorMethod = PulsarByteBufAllocator.class.getDeclaredMethod("createByteBufAllocator");
+        createByteBufAllocatorMethod.setAccessible(true);
+        Whitebox.setInternalState(PulsarByteBufAllocator.class, "DEFAULT",
+                createByteBufAllocatorMethod.invoke(null));
     }
 
 }