/**
 * 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.core.query.scheduler.resources;

import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.pinot.core.query.scheduler.SchedulerGroupAccountant;
import org.testng.annotations.Test;

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.testng.Assert.assertEquals;


public class BoundedAccountingExecutorTest {
  private AtomicInteger running = new AtomicInteger(0);

  private class Syncer {
    CyclicBarrier validationBarrier;
    CyclicBarrier startupBarrier;
  }

  // Created bounded executor with limit 3 and add 5 jobs. Only 3 can be running at a time
  @Test
  public void testBoundsWithinThreadCount()
      throws BrokenBarrierException, InterruptedException {
    SchedulerGroupAccountant accountant = mock(SchedulerGroupAccountant.class);
    // Test below relies on jobs > limit
    final int limit = 3;
    final int jobs = 5;
    // we want total threads > limit
    Executor es = Executors.newFixedThreadPool(2 * limit);
    final BoundedAccountingExecutor bes = new BoundedAccountingExecutor(es, limit, accountant);
    final Syncer syncer = new Syncer();
    // barrier parties: all the executables plus 1 for main testing thread
    // startup barrier pauses main thread till all the threads have started
    // validation barrier allows for validation to complete before proceeding further
    syncer.startupBarrier = new CyclicBarrier(limit + 1);
    syncer.validationBarrier = new CyclicBarrier(limit + 1);

    // start adding jobs in new thread. We need to add jobs in new thread
    // because the thread adding jobs is expected to block at limit
    new Thread(new Runnable() {
      @Override
      public void run() {
        for (int i = 0; i < jobs; i++) {
          bes.execute(new Runnable() {
            @Override
            public void run() {
              try {
                running.incrementAndGet();
                syncer.startupBarrier.await();
                syncer.validationBarrier.await();
                running.decrementAndGet();
              } catch (Exception e) {
                throw new RuntimeException(e);
              }
            }
          });
        }
      }
    }).start();

    syncer.startupBarrier.await();
    assertEquals(running.get(), limit);
    verify(accountant, times(limit)).incrementThreads();
    // reset will clear the counts on incrementThreads
    reset(accountant);
    int pendingJobs = jobs - limit;
    syncer.startupBarrier = new CyclicBarrier(pendingJobs + 1);
    syncer.validationBarrier.await();
    // verify additional jobs are run as soon as current job finishes
    syncer.validationBarrier = new CyclicBarrier(pendingJobs + 1);
    syncer.startupBarrier.await();
    assertEquals(running.get(), pendingJobs);
    verify(accountant, times(pendingJobs)).incrementThreads();
    syncer.validationBarrier.await();
  }
}


@@ -1168,155 +1168,73 @@
 org.
-testng.annotations.Test;%0A%0Aimport static org.mockito.Mockito.mock;%0Aimport static org.mockito.Mockito.reset;%0Aimport static org.mockito.Mockito.times;
+apache.pinot.util.TestUtils;%0Aimport org.testng.annotations.Test;%0A
 %0Aimp
@@ -1264,22 +1264,17 @@
 Mockito.
-verify
+*
 ;%0Aimport
@@ -1299,28 +1299,17 @@
 .Assert.
-assertEquals
+*
 ;%0A%0A%0Apubl
@@ -3213,32 +3213,299 @@
 arrier.await();%0A
+    // At this point, 'limit' jobs should have executed the startupBarrier.await() call.%0A    // The other jobs should be waiting on a semaphore for permits inside the BoundedAccountingExecutor.execute()%0A    // so they have not executed running.incrementAndGet() yet.%0A
     assertEquals
@@ -3667,16 +3667,22 @@
 t);%0A    
+final 
 int pend
@@ -3713,317 +3713,1250 @@
     
-syncer.startupBarrier = new CyclicBarrier(pendingJobs + 1);%0A    syncer.validationBarrier.await();%0A    // verify additional jobs are run as soon as current job finishes%0A    syncer.validationBarrier = new CyclicBarrier(pendingJobs + 1);%0A    syncer.startupBarrier.await();%0A    assertEquals(running.get(),
+// Before the pendingJobs get to startupBarrier, reset it to a different value%0A    // since we cannot change the limit of the CyclicBarrier once created.%0A    // The new limit will be pending jobs plus the await we will call in this thread.%0A    syncer.startupBarrier = new CyclicBarrier(pendingJobs + 1);%0A%0A    // Now let the running threads complete and call running.decrementAndGet. As%0A    // they exit, the pending jobs will acquire permits and start to increment%0A    // the running counter and wait on startupBarrier.await().%0A    syncer.validationBarrier.await();%0A    // verify additional jobs are run as soon as current job finishes%0A    syncer.validationBarrier = new CyclicBarrier(pendingJobs + 1);%0A    // When we run the test in a small number of cores, it is possible that the running jobs%0A    // have not yet gotten to execute running.decrementAndGet(), but the pending jobs have already%0A    // done the increment. So, we need to wait until we check the running counter to equal the%0A    // pending jobs.%0A    TestUtils.waitForCondition(aVoid -%3E running.get() == pendingJobs,%0A        10_000, %22Invalid number of running jobs%22 + running.get());%0A%0A    // Now that there are no jobs running, we can let the new ones in.%0A    // All the
  pending
 Jobs
@@ -4951,20 +4951,145 @@
  pending
-J
+ j
 obs
+ will wait on the validationBarrier after we let them pass%0A    // the startupbarrier below.%0A    syncer.startupBarrier.await(
 );%0A    v

Fix flakiness in BoundedAccountingExecutorTest (#4674)