package org.apache.maven.surefire.junitcore.pc;
   
   /*
    * 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.
   */
  
  import org.apache.maven.surefire.api.report.ConsoleStream;
  import org.junit.runner.Description;
  import org.junit.runners.model.RunnerScheduler;
  
  import java.io.ByteArrayOutputStream;
  import java.io.PrintStream;
  import java.util.Collection;
  import java.util.Set;
  import java.util.concurrent.ConcurrentLinkedQueue;
  import java.util.concurrent.CopyOnWriteArraySet;
  import java.util.concurrent.RejectedExecutionException;
  import java.util.concurrent.RejectedExecutionHandler;
  import java.util.concurrent.ThreadPoolExecutor;
  
  /**
   * Schedules tests, controls thread resources, awaiting tests and other schedulers finished, and
   * a master scheduler can shutdown slaves.
   * <br>
   * The scheduler objects should be first created (and wired) and set in runners
   * {@link org.junit.runners.ParentRunner#setScheduler(org.junit.runners.model.RunnerScheduler)}.
   * <br>
   * A new instance of scheduling strategy should be passed to the constructor of this scheduler.
   *
   * @author Tibor Digana (tibor17)
   * @since 2.16
   */
  public class Scheduler
      implements RunnerScheduler
  {
      private final Balancer balancer;
  
      private final SchedulingStrategy strategy;
  
      private final Set<Controller> slaves = new CopyOnWriteArraySet<>();
  
      private final Description description;
  
      private final ConsoleStream logger;
  
      private volatile boolean shutdown = false;
  
      private volatile boolean started = false;
  
      private volatile boolean finished = false;
  
      private volatile Controller masterController;
  
      /**
       * Use e.g. parallel classes have own non-shared thread pool, and methods another pool.
       * <br>
       * You can use it with one infinite thread pool shared in strategies across all
       * suites, class runners, etc.
       *
       * @param logger          console logger
       * @param description     JUnit description of class
       * @param strategy        scheduling strategy
       */
      public Scheduler( ConsoleStream logger, Description description, SchedulingStrategy strategy )
      {
          this( logger, description, strategy, -1 );
      }
  
      /**
       * Should be used if schedulers in parallel children and parent use one instance of bounded thread pool.
       * <br>
       * Set this scheduler in a e.g. one suite of classes, then every individual class runner should reference
       * {@link #Scheduler(ConsoleStream, org.junit.runner.Description, Scheduler, SchedulingStrategy)}
       * or {@link #Scheduler(ConsoleStream, org.junit.runner.Description, Scheduler, SchedulingStrategy, int)}.
       *
       * @param logger current logger implementation
       * @param description description of current runner
       * @param strategy    scheduling strategy with a shared thread pool
       * @param concurrency determines maximum concurrent children scheduled a time via {@link #schedule(Runnable)}
       * @throws NullPointerException if null <tt>strategy</tt>
       */
      public Scheduler( ConsoleStream logger, Description description, SchedulingStrategy strategy, int concurrency )
      {
          this( logger, description, strategy, BalancerFactory.createBalancer( concurrency ) );
     }
 
     /**
      * New instances should be used by schedulers with limited concurrency by <tt>balancer</tt>
      * against other groups of schedulers. The schedulers share one pool.
      * <br>
      * Unlike in {@link #Scheduler(ConsoleStream, org.junit.runner.Description, SchedulingStrategy, int)} which was
      * limiting the <tt>concurrency</tt> of children of a runner where this scheduler was set, {@code this}
      * <tt>balancer</tt> is limiting the concurrency of all children in runners having schedulers created by this
      * constructor.
      *
      * @param logger current logger implementation
      * @param description description of current runner
      * @param strategy    scheduling strategy which may share threads with other strategy
      * @param balancer    determines maximum concurrent children scheduled a time via {@link #schedule(Runnable)}
      * @throws NullPointerException if null <tt>strategy</tt> or <tt>balancer</tt>
      */
     public Scheduler( ConsoleStream logger, Description description, SchedulingStrategy strategy, Balancer balancer )
     {
         strategy.setDefaultShutdownHandler( newShutdownHandler() );
         this.logger = logger;
         this.description = description;
         this.strategy = strategy;
         this.balancer = balancer;
         masterController = null;
     }
 
     /**
      * Can be used by e.g. a runner having parallel classes in use case with parallel
      * suites, classes and methods sharing the same thread pool.
      *
      * @param logger current logger implementation
      * @param description     description of current runner
      * @param masterScheduler scheduler sharing own threads with this slave
      * @param strategy        scheduling strategy for this scheduler
      * @param balancer        determines maximum concurrent children scheduled a time via {@link #schedule(Runnable)}
      * @throws NullPointerException if null <tt>masterScheduler</tt>, <tt>strategy</tt> or <tt>balancer</tt>
      */
     public Scheduler( ConsoleStream logger, Description description, Scheduler masterScheduler,
                       SchedulingStrategy strategy, Balancer balancer )
     {
         this( logger, description, strategy, balancer );
         strategy.setDefaultShutdownHandler( newShutdownHandler() );
         masterScheduler.register( this );
     }
 
     /**
      * @param logger          console logger
      * @param description     JUnit description of class
      * @param masterScheduler a reference to
      * {@link #Scheduler(ConsoleStream, org.junit.runner.Description, SchedulingStrategy, int)}
      *                        or {@link #Scheduler(ConsoleStream, org.junit.runner.Description, SchedulingStrategy)}
      * @param strategy        scheduling strategy
      * @param concurrency     determines maximum concurrent children scheduled a time via {@link #schedule(Runnable)}
      *
      * @see #Scheduler(ConsoleStream, org.junit.runner.Description, SchedulingStrategy)
      * @see #Scheduler(ConsoleStream, org.junit.runner.Description, SchedulingStrategy, int)
      */
     public Scheduler( ConsoleStream logger, Description description, Scheduler masterScheduler,
                       SchedulingStrategy strategy, int concurrency )
     {
         this( logger, description, strategy, concurrency );
         strategy.setDefaultShutdownHandler( newShutdownHandler() );
         masterScheduler.register( this );
     }
 
     /**
      * Should be used with individual pools on suites, classes and methods, see
      * {@link org.apache.maven.surefire.junitcore.pc.ParallelComputerBuilder#useSeparatePools()}.
      * <br>
      * Cached thread pool is infinite and can be always shared.
      *
      * @param logger          console logger
      * @param description     JUnit description of class
      * @param masterScheduler parent scheduler
      * @param strategy        scheduling strategy
      */
     public Scheduler( ConsoleStream logger, Description description, Scheduler masterScheduler,
                       SchedulingStrategy strategy )
     {
         this( logger, description, masterScheduler, strategy, 0 );
     }
 
     private void setController( Controller masterController )
     {
         if ( masterController == null )
         {
             throw new NullPointerException( "null ExecutionController" );
         }
         this.masterController = masterController;
     }
 
     /**
      * @param slave a slave scheduler to register
      * @return {@code true} if successfully registered the <tt>slave</tt>.
      */
     private boolean register( Scheduler slave )
     {
         boolean canRegister = slave != null && slave != this;
         if ( canRegister )
         {
             Controller controller = new Controller( slave );
             canRegister = !slaves.contains( controller );
             if ( canRegister )
             {
                 slaves.add( controller );
                 slave.setController( controller );
             }
         }
         return canRegister;
     }
 
     /**
      * @return {@code true} if new tasks can be scheduled.
      */
     private boolean canSchedule()
     {
         return !shutdown && ( masterController == null || masterController.canSchedule() );
     }
 
     protected void logQuietly( Throwable t )
     {
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         try ( PrintStream stream = new PrintStream( out ) )
         {
             t.printStackTrace( stream );
         }
         logger.println( out.toString() );
     }
 
     protected void logQuietly( String msg )
     {
         logger.println( msg );
     }
 
     /**
      * Attempts to stop all actively executing tasks and immediately returns a collection
      * of descriptions of those tasks which have started prior to this call.
      * <br>
      * This scheduler and other registered schedulers will stop, see {@link #register(Scheduler)}.
      * If <tt>shutdownNow</tt> is set, waiting methods will be interrupted via {@link Thread#interrupt}.
      *
      * @param stopNow if {@code true} interrupts waiting test methods
      * @return collection of descriptions started before shutting down
      */
     protected ShutdownResult describeStopped( boolean stopNow )
     {
         Collection<Description> executedTests = new ConcurrentLinkedQueue<>();
         Collection<Description> incompleteTests = new ConcurrentLinkedQueue<>();
         stop( executedTests, incompleteTests, false, stopNow );
         return new ShutdownResult( executedTests, incompleteTests );
     }
 
     /**
      * Stop/Shutdown/Interrupt scheduler and its children (if any).
      *
      * @param executedTests       Started tests which have finished normally or abruptly till called this method.
      * @param incompleteTests     Started tests which have finished incomplete due to shutdown.
      * @param tryCancelFutures    Useful to set to {@code false} if a timeout is specified in plugin config.
      *                            When the runner of
      *                            {@link ParallelComputer#getSuite(org.junit.runners.model.RunnerBuilder, Class[])}
      *                            is finished in
      *                            {@link org.junit.runners.Suite#run(org.junit.runner.notification.RunNotifier)}
      *                            all the thread-pools created by {@link ParallelComputerBuilder.PC} are already dead.
      *                            See the unit test {@code ParallelComputerBuilder#timeoutAndForcedShutdown()}.
      * @param stopNow             Interrupting tests by {@link java.util.concurrent.ExecutorService#shutdownNow()} or
      *                            {@link java.util.concurrent.Future#cancel(boolean) Future#cancel(true)} or
      *                            {@link Thread#interrupt()}.
      */
     private void stop( Collection<Description> executedTests, Collection<Description> incompleteTests,
                        boolean tryCancelFutures, boolean stopNow )
     {
         shutdown = true;
         try
         {
             if ( started && !ParallelComputerUtil.isUnusedDescription( description ) )
             {
                 if ( executedTests != null )
                 {
                     executedTests.add( description );
                 }
 
                 if ( incompleteTests != null && !finished )
                 {
                     incompleteTests.add( description );
                 }
             }
 
             for ( Controller slave : slaves )
             {
                 slave.stop( executedTests, incompleteTests, tryCancelFutures, stopNow );
             }
         }
         finally
         {
             try
             {
                 balancer.releaseAllPermits();
             }
             finally
             {
                 if ( stopNow )
                 {
                     strategy.stopNow();
                 }
                 else if ( tryCancelFutures )
                 {
                     strategy.stop();
                 }
                 else
                 {
                     strategy.disable();
                 }
             }
         }
     }
 
     protected boolean shutdownThreadPoolsAwaitingKilled()
     {
         if ( masterController == null )
         {
             stop( null, null, true, false );
             boolean isNotInterrupted = true;
             if ( strategy != null )
             {
                 isNotInterrupted = strategy.destroy();
             }
             for ( Controller slave : slaves )
             {
                 isNotInterrupted &= slave.destroy();
             }
             return isNotInterrupted;
         }
         else
         {
             throw new UnsupportedOperationException( "cannot call this method if this is not a master scheduler" );
         }
     }
 
     protected void beforeExecute()
     {
     }
 
     protected void afterExecute()
     {
     }
 
     @Override
     public void schedule( Runnable childStatement )
     {
         if ( childStatement == null )
         {
             logQuietly( "cannot schedule null" );
         }
         else if ( canSchedule() && strategy.canSchedule() )
         {
             try
             {
                 boolean isNotInterrupted = balancer.acquirePermit();
                 if ( isNotInterrupted && !shutdown )
                 {
                     Runnable task = wrapTask( childStatement );
                     strategy.schedule( task );
                     started = true;
                 }
             }
             catch ( RejectedExecutionException e )
             {
                 stop( null, null, true, false );
             }
             catch ( Throwable t )
             {
                 balancer.releasePermit();
                 logQuietly( t );
             }
         }
     }
 
     @Override
     public void finished()
     {
         try
         {
             strategy.finished();
         }
         catch ( InterruptedException e )
         {
             logQuietly( e );
         }
         finally
         {
             finished = true;
         }
     }
 
     private Runnable wrapTask( final Runnable task )
     {
         return new Runnable()
         {
             @Override
             public void run()
             {
                 try
                 {
                     beforeExecute();
                     task.run();
                 }
                 finally
                 {
                     try
                     {
                         afterExecute();
                     }
                     finally
                     {
                         balancer.releasePermit();
                     }
                 }
             }
         };
     }
 
     protected ShutdownHandler newShutdownHandler()
     {
         return new ShutdownHandler();
     }
 
     /**
      * If this is a master scheduler, the slaves can stop scheduling by the master through the controller.
      */
     private final class Controller
     {
         private final Scheduler slave;
 
         private Controller( Scheduler slave )
         {
             this.slave = slave;
         }
 
         /**
          * @return {@code true} if new children can be scheduled.
          */
         boolean canSchedule()
         {
             return Scheduler.this.canSchedule();
         }
 
         void stop( Collection<Description> executedTests, Collection<Description> incompleteTests,
                    boolean tryCancelFutures, boolean shutdownNow )
         {
             slave.stop( executedTests, incompleteTests, tryCancelFutures, shutdownNow );
         }
 
         /**
          * @see org.apache.maven.surefire.junitcore.pc.Destroyable#destroy()
          */
         boolean destroy()
         {
             return slave.strategy.destroy();
         }
 
         @Override
         public int hashCode()
         {
             return slave.hashCode();
         }
 
         @Override
         public boolean equals( Object o )
         {
             return o == this || ( o instanceof Controller ) && slave.equals( ( (Controller) o ).slave );
         }
     }
 
     /**
      * There is a way to shutdown the hierarchy of schedulers. You can do it in master scheduler via
      * {@link #shutdownThreadPoolsAwaitingKilled()} which kills the current master and children recursively.
      * If alternatively a shared {@link java.util.concurrent.ExecutorService} used by the master and children
      * schedulers is shutdown from outside, then the {@link ShutdownHandler} is a hook calling current
      * {@link #describeStopped(boolean)}. The method {@link #describeStopped(boolean)} is again shutting down children
      * schedulers recursively as well.
      */
     public class ShutdownHandler
         implements RejectedExecutionHandler
     {
         private volatile RejectedExecutionHandler poolHandler;
 
         protected ShutdownHandler()
         {
             poolHandler = null;
         }
 
         public void setRejectedExecutionHandler( RejectedExecutionHandler poolHandler )
         {
             this.poolHandler = poolHandler;
         }
 
         @Override
         public void rejectedExecution( Runnable r, ThreadPoolExecutor executor )
         {
             if ( executor.isShutdown() )
             {
                 Scheduler.this.stop( null, null, true, false );
             }
             final RejectedExecutionHandler poolHandler = this.poolHandler;
             if ( poolHandler != null )
             {
                 poolHandler.rejectedExecution( r, executor );
             }
         }
     }
 }