Fix flaky tests

pinot-common/pom.xml

@@ -46,6 +46,8 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-surefire-plugin</artifactId>
         <configuration combine.children="override">
+          <forkCount>1</forkCount>
+          <reuseForks>true</reuseForks>
           <properties>
             <property>
               <name>usedefaultlisteners</name>

pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardBrokerPrometheusMetricsTest.java

@@ -19,6 +19,9 @@
 
 package org.apache.pinot.common.metrics.prometheus.dropwizard;
 
+import org.apache.pinot.common.metrics.BrokerGauge;
+import org.apache.pinot.common.metrics.BrokerMeter;
+import org.apache.pinot.common.metrics.BrokerTimer;
 import org.apache.pinot.common.metrics.prometheus.BrokerPrometheusMetricsTest;
 import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory;
 import org.apache.pinot.spi.annotations.metrics.PinotMetricsFactory;
@@ -28,7 +31,7 @@
 /**
  * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory}
  */
-@Test(enabled = false)
+@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij
 public class DropwizardBrokerPrometheusMetricsTest extends BrokerPrometheusMetricsTest {
   @Override
   protected PinotMetricsFactory getPinotMetricsFactory() {
@@ -40,4 +43,19 @@ protected String getConfigFile() {
     //todo: return the correct dir once this test is enabled
     return null;
   }
+
+  @Test(dataProvider = "brokerGauges", enabled = false)
+  public void timerTest(BrokerTimer timer) {
+    super.timerTest(timer);
+  }
+
+  @Test(dataProvider = "brokerMeters", enabled = false)
+  public void meterTest(BrokerMeter meter) {
+    super.meterTest(meter);
+  }
+
+  @Test(dataProvider = "brokerGauges", enabled = false)
+  public void gaugeTest(BrokerGauge gauge) {
+    super.gaugeTest(gauge);
+  }
 }

pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardControllerPrometheusMetricsTest.java

@@ -19,6 +19,9 @@
 
 package org.apache.pinot.common.metrics.prometheus.dropwizard;
 
+import org.apache.pinot.common.metrics.ControllerGauge;
+import org.apache.pinot.common.metrics.ControllerMeter;
+import org.apache.pinot.common.metrics.ControllerTimer;
 import org.apache.pinot.common.metrics.prometheus.ControllerPrometheusMetricsTest;
 import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory;
 import org.apache.pinot.spi.annotations.metrics.PinotMetricsFactory;
@@ -28,7 +31,7 @@
 /**
  * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory}
  */
-@Test(enabled = false)
+@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij
 public class DropwizardControllerPrometheusMetricsTest extends ControllerPrometheusMetricsTest {
 
   @Override
@@ -41,4 +44,19 @@ protected String getConfigFile() {
     //todo: return the correct dir once this test is enabled
     return null;
   }
+
+  @Test(dataProvider = "controllerTimers", enabled = false)
+  public void timerTest(ControllerTimer controllerTimer) {
+    super.timerTest(controllerTimer);
+  }
+
+  @Test(dataProvider = "controllerMeters", enabled = false)
+  public void meterTest(ControllerMeter meter) {
+    super.meterTest(meter);
+  }
+
+  @Test(dataProvider = "controllerGauges", enabled = false)
+  public void gaugeTest(ControllerGauge controllerGauge) {
+    super.gaugeTest(controllerGauge);
+  }
 }

pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardMinionPrometheusMetricsTest.java

@@ -19,6 +19,9 @@
 
 package org.apache.pinot.common.metrics.prometheus.dropwizard;
 
+import org.apache.pinot.common.metrics.MinionGauge;
+import org.apache.pinot.common.metrics.MinionMeter;
+import org.apache.pinot.common.metrics.MinionTimer;
 import org.apache.pinot.common.metrics.prometheus.MinionPrometheusMetricsTest;
 import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory;
 import org.apache.pinot.spi.annotations.metrics.PinotMetricsFactory;
@@ -28,7 +31,7 @@
 /**
  * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory}
  */
-@Test(enabled = false)
+@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij
 public class DropwizardMinionPrometheusMetricsTest extends MinionPrometheusMetricsTest {
 
   @Override
@@ -41,4 +44,19 @@ protected String getConfigFile() {
     //todo: return the correct dir once this test is enabled
     return null;
   }
+
+  @Test(dataProvider = "minionTimers", enabled = false)
+  public void timerTest(MinionTimer timer) {
+    super.timerTest(timer);
+  }
+
+  @Test(dataProvider = "minionMeters", enabled = false)
+  public void meterTest(MinionMeter meter) {
+    super.meterTest(meter);
+  }
+
+  @Test(dataProvider = "minionGauges", enabled = false)
+  public void gaugeTest(MinionGauge gauge) {
+    super.gaugeTest(gauge);
+  }
 }

pinot-common/src/test/java/org/apache/pinot/common/metrics/prometheus/dropwizard/DropwizardServerPrometheusMetricsTest.java

@@ -19,6 +19,9 @@
 
 package org.apache.pinot.common.metrics.prometheus.dropwizard;
 
+import org.apache.pinot.common.metrics.ServerGauge;
+import org.apache.pinot.common.metrics.ServerMeter;
+import org.apache.pinot.common.metrics.ServerTimer;
 import org.apache.pinot.common.metrics.prometheus.ServerPrometheusMetricsTest;
 import org.apache.pinot.plugin.metrics.dropwizard.DropwizardMetricsFactory;
 import org.apache.pinot.plugin.metrics.yammer.YammerMetricsFactory;
@@ -29,7 +32,7 @@
 /**
  * Disabling tests as Pinot currently uses Yammer and these tests fail for for {@link DropwizardMetricsFactory}
  */
-@Test(enabled = false)
+@Test(enabled = false) // enabled=false on class level doesn't seem to work in intellij
 public class DropwizardServerPrometheusMetricsTest extends ServerPrometheusMetricsTest {
 
   @Override
@@ -42,4 +45,19 @@ protected String getConfigFile() {
     //todo: return the correct dir once this test is enabled
     return null;
   }
+
+  @Test(dataProvider = "serverTimers", enabled = false)
+  public void timerTest(ServerTimer serverTimer) {
+    super.timerTest(serverTimer);
+  }
+
+  @Test(dataProvider = "serverMeters", enabled = false)
+  public void meterTest(ServerMeter serverMeter) {
+    super.meterTest(serverMeter);
+  }
+
+  @Test(dataProvider = "serverGauges", enabled = false)
+  public void gaugeTest(ServerGauge serverGauge) {
+    super.gaugeTest(serverGauge);
+  }
 }

pinot-common/src/test/java/org/apache/pinot/common/utils/RoundRobinURIProviderTest.java

@@ -48,7 +48,7 @@ public void testHostAddressRoundRobin()
         InetAddress.getByAddress("localhost", InetAddresses.forString("0:0:0:0:0:0:0:1").getAddress())
     };
 
-    MockedStatic<InetAddress> mock = Mockito.mockStatic(InetAddress.class);
+    try (MockedStatic<InetAddress> mock = Mockito.mockStatic(InetAddress.class)) {
     mock.when(() -> InetAddress.getAllByName("localhost")).thenReturn(localHostAddresses);
     mock.when(() -> InetAddress.getAllByName("testweb.com")).thenReturn(testWebAddresses);
 
@@ -176,6 +176,7 @@ public void testHostAddressRoundRobin()
         previousIndex = currentIndex;
       }
     }
+    }
   }
 
   static class TestCase {

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/forward/BaseChunkForwardIndexReader.java

@@ -447,7 +447,7 @@ public void readValuesSV(int[] docIds, int length, double[] values, ChunkReaderC
   public void close()
       throws IOException {
     // NOTE: DO NOT close the PinotDataBuffer here because it is tracked by the caller and might be reused later. The
-    // caller is responsible of closing the PinotDataBuffer.
+    // caller is responsible for closing the PinotDataBuffer.
     _chunkDecompressor.close();
   }
 

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/text/NativeTextIndexReader.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.segment.local.segment.index.readers.text;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import java.io.File;
 import java.io.IOException;
@@ -119,6 +120,15 @@ public MutableRoaringBitmap getDocIds(String searchQuery) {
   @Override
   public void close()
       throws IOException {
+    // TODO: this method doesn't release native buffers held by _fst
     _buffer.close();
   }
+
+  @VisibleForTesting
+  public void closeInTest()
+      throws IOException {
+    if (_fst instanceof ImmutableFST) {
+      ((ImmutableFST) _fst)._mutableBytesStore.close();
+    }
+  }
 }

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/store/FilePerIndexDirectory.java

@@ -99,6 +99,7 @@ public void close()
 
   @Override
   public void removeIndex(String columnName, IndexType<?, ?, ?> indexType) {
+    // TODO: this leaks the removed data buffer (it's not going to be freed in close() method)
     _indexBuffers.remove(new IndexKey(columnName, indexType));
     if (indexType == StandardIndexes.text()) {
       TextIndexUtils.cleanupTextIndex(_segmentDirectory, columnName);

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/nativefst/NativeFSTIndexReader.java

@@ -18,6 +18,7 @@
  */
 package org.apache.pinot.segment.local.utils.nativefst;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Collections;
@@ -71,5 +72,15 @@ public ImmutableRoaringBitmap getDictIds(String searchQuery) {
   @Override
   public void close()
       throws IOException {
+    //TODO: why does this class not close FST ?
+  }
+
+  @VisibleForTesting
+  public void closeInTest()
+      throws IOException {
+    // immutable fst contains native data buffers that need to be closed
+    if (_fst instanceof ImmutableFST) {
+      ((ImmutableFST) _fst)._mutableBytesStore.close();
+    }
   }
 }

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterClassCheckRule.java

@@ -0,0 +1,33 @@
+/**
+ * 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.pinot.segment.local;
+
+import org.testng.ITestContext;
+import org.testng.annotations.AfterClass;
+
+
+// Checks that tests don't leak buffers after executing all test methods
+// Meant for tests that share fixture that contains buffers and can't be verified with an 'after' method
+public interface PinotBuffersAfterClassCheckRule {
+
+  @AfterClass
+  default void checkPinotBuffersAfterClass(ITestContext context) {
+    PinotBuffersAfterMethodCheckRule.assertPinotBuffers(context.getAllTestMethods()[0].getRealClass().getName());
+  }
+}

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/PinotBuffersAfterMethodCheckRule.java

@@ -0,0 +1,70 @@
+/**
+ * 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.pinot.segment.local;
+
+import java.util.List;
+import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.testng.ITestResult;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+
+// Check for pinot buffer leaks after each test method
+public interface PinotBuffersAfterMethodCheckRule {
+
+  @AfterMethod
+  default void checkPinotBuffers(ITestResult result) {
+    assertPinotBuffers(result);
+  }
+
+  @Test(enabled = false)
+  static void assertPinotBuffers(Object target) {
+    List<String> bufferInfos = PinotDataBuffer.getBufferInfo();
+    if (bufferInfos.size() > 0) {
+
+      StringBuilder builder;
+      if (target instanceof ITestResult) {
+        ITestResult result = (ITestResult) target;
+        if (result.getStatus() != ITestResult.SUCCESS) {
+          // don't override result of failed test
+          PinotDataBuffer.closeOpenBuffers();
+          return;
+        }
+
+        builder = new StringBuilder("Test method: ").append(result.getTestClass().getRealClass().getSimpleName())
+            .append('.').append(result.getMethod().getMethodName());
+      } else {
+        builder = new StringBuilder("Test class: ").append(target);
+      }
+
+      StringBuilder message =
+          builder.append(target)
+              .append(" did not release ")
+              .append(bufferInfos.size())
+              .append(" pinot buffer(s): \n");
+      for (String bufferInfo : bufferInfos) {
+        message.append(bufferInfo).append('\n');
+      }
+      //close all open buffers otherwise one test will fail all following tests
+      PinotDataBuffer.closeOpenBuffers();
+
+      throw new AssertionError(message.toString());
+    }
+  }
+}

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java

@@ -24,12 +24,14 @@
 import java.util.HashSet;
 import org.apache.pinot.common.metrics.ServerMeter;
 import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.segment.local.PinotBuffersAfterMethodCheckRule;
 import org.apache.pinot.spi.config.table.JsonIndexConfig;
 import org.apache.pinot.spi.data.FieldSpec;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.GenericRow;
 import org.apache.pinot.spi.stream.StreamMessageMetadata;
 import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -40,7 +42,7 @@
 import static org.testng.Assert.assertTrue;
 
 
-public class IndexingFailureTest {
+public class IndexingFailureTest implements PinotBuffersAfterMethodCheckRule {
   private static final String TABLE_NAME = "testTable";
   private static final String INT_COL = "int_col";
   private static final String STRING_COL = "string_col";
@@ -61,6 +63,11 @@ public void setup() {
             Collections.singletonMap(JSON_COL, new JsonIndexConfig()), _serverMetrics);
   }
 
+  @AfterMethod
+  public void tearDown() {
+    _mutableSegment.destroy();
+  }
+
   @Test
   public void testIndexingFailures()
       throws IOException {