content: /**
   * 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.activemq.thread;
  
  import java.lang.reflect.Constructor;
  import java.lang.reflect.InvocationTargetException;
  import java.lang.reflect.Method;
  import java.util.concurrent.Executor;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.RejectedExecutionHandler;
  import java.util.concurrent.SynchronousQueue;
  import java.util.concurrent.ThreadFactory;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.concurrent.atomic.AtomicReference;
  
  import org.apache.activemq.util.ThreadPoolUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * Manages the thread pool for long running tasks. Long running tasks are not
   * always active but when they are active, they may need a few iterations of
   * processing for them to become idle. The manager ensures that each task is
   * processes but that no one task overtakes the system. This is somewhat like
   * cooperative multitasking.
   *
   * @org.apache.xbean.XBean
   */
  public class TaskRunnerFactory implements Executor {
  
      private static final Logger LOG = LoggerFactory.getLogger(TaskRunnerFactory.class);
      private final AtomicReference<ExecutorService> executorRef = new AtomicReference<>();
      private int maxIterationsPerRun;
      private String name;
      private int priority;
      private boolean daemon;
      private final AtomicLong id = new AtomicLong(0);
      private boolean dedicatedTaskRunner;
      private long shutdownAwaitTermination = 30000;
      private final AtomicBoolean initDone = new AtomicBoolean(false);
      private int maxThreadPoolSize = getDefaultMaximumPoolSize();
      private RejectedExecutionHandler rejectedTaskHandler = null;
      private ClassLoader threadClassLoader;
  
      // [AMQ-9394] Virtual Thread support
      private boolean virtualThreadTaskRunner = false;
  
      public TaskRunnerFactory() {
          this("ActiveMQ Task");
      }
  
      public TaskRunnerFactory(String name) {
          this(name, Thread.NORM_PRIORITY, true, 1000);
      }
  
      private TaskRunnerFactory(String name, int priority, boolean daemon, int maxIterationsPerRun) {
          this(name, priority, daemon, maxIterationsPerRun, false);
      }
  
      public TaskRunnerFactory(String name, int priority, boolean daemon, int maxIterationsPerRun, boolean dedicatedTaskRunner) {
          this(name, priority, daemon, maxIterationsPerRun, dedicatedTaskRunner, getDefaultMaximumPoolSize());
      }
  
      public TaskRunnerFactory(String name, int priority, boolean daemon, int maxIterationsPerRun, boolean dedicatedTaskRunner, boolean virtualThreadTaskRunner) {
          this(name, priority, daemon, maxIterationsPerRun, dedicatedTaskRunner, getDefaultMaximumPoolSize(), virtualThreadTaskRunner);
      }
  
      public TaskRunnerFactory(String name, int priority, boolean daemon, int maxIterationsPerRun, boolean dedicatedTaskRunner, int maxThreadPoolSize) {
          this(name, priority, daemon, maxIterationsPerRun, dedicatedTaskRunner, maxThreadPoolSize, false);
      }
  
      public TaskRunnerFactory(String name, int priority, boolean daemon, int maxIterationsPerRun, boolean dedicatedTaskRunner, int maxThreadPoolSize, boolean virtualThreadTaskRunner) {
          this.name = name;
          this.priority = priority;
          this.daemon = daemon;
          this.maxIterationsPerRun = maxIterationsPerRun;
          this.dedicatedTaskRunner = dedicatedTaskRunner;
          this.maxThreadPoolSize = maxThreadPoolSize;
          this.virtualThreadTaskRunner = virtualThreadTaskRunner;
      }
  
      public void init() {
          if (!initDone.get()) {
              // If your OS/JVM combination has a good thread model, you may want to
              // avoid using a thread pool to run tasks and use a DedicatedTaskRunner instead.
              //AMQ-6602 - lock instead of using compareAndSet to prevent threads from seeing a null value
              //for executorRef inside createTaskRunner() on contention and creating a DedicatedTaskRunner
              synchronized(this) {
                  //need to recheck if initDone is true under the lock
                  if (!initDone.get()) {
                      if (virtualThreadTaskRunner || "true".equalsIgnoreCase(System.getProperty("org.apache.activemq.UseVirtualThreadTaskRunner")) ) {
                          executorRef.compareAndSet(null, createVirtualThreadExecutor());
                      } else if (dedicatedTaskRunner || "true".equalsIgnoreCase(System.getProperty("org.apache.activemq.UseDedicatedTaskRunner"))) {
                          executorRef.set(null);
                      } else {
                          executorRef.compareAndSet(null, createDefaultExecutor());
                      }
                      LOG.debug("Initialized TaskRunnerFactory[{}] using ExecutorService: {}", name, executorRef.get());
                      initDone.set(true);
                  }
              }
          }
      }
  
      /**
       * Performs a shutdown only, by which the thread pool is shutdown by not graceful nor aggressively.
       *
       * @see ThreadPoolUtils#shutdown(java.util.concurrent.ExecutorService)
       */
      public void shutdown() {
          ExecutorService executor = executorRef.get();
          if (executor != null) {
              ThreadPoolUtils.shutdown(executor);
          }
          clearExecutor();
      }
  
      /**
       * Performs a shutdown now (aggressively) on the thread pool.
       *
       * @see ThreadPoolUtils#shutdownNow(java.util.concurrent.ExecutorService)
       */
      public void shutdownNow() {
          ExecutorService executor = executorRef.get();
          if (executor != null) {
              ThreadPoolUtils.shutdownNow(executor);
          }
          clearExecutor();
      }
  
      /**
       * Performs a graceful shutdown.
       *
       * @see ThreadPoolUtils#shutdownGraceful(java.util.concurrent.ExecutorService)
       */
      public void shutdownGraceful() {
          ExecutorService executor = executorRef.get();
          if (executor != null) {
              ThreadPoolUtils.shutdownGraceful(executor, shutdownAwaitTermination);
          }
          clearExecutor();
      }
  
      private void clearExecutor() {
          //clear under a lock to prevent threads from seeing initDone == true
          //but then getting null from executorRef
          synchronized(this) {
              executorRef.set(null);
              initDone.set(false);
          }
      }
  
      public TaskRunner createTaskRunner(Task task, String name) {
          init();
          ExecutorService executor = executorRef.get();
          if (executor != null) {
              if(isVirtualThreadTaskRunner()) {
                  return createVirtualThreadTaskRunner(executor, task, maxIterationsPerRun);
              } else {
                  return new PooledTaskRunner(executor, task, maxIterationsPerRun);
              }
          } else {
              return new DedicatedTaskRunner(task, name, priority, daemon);
          }
      }
  
      @Override
      public void execute(Runnable runnable) {
          execute(runnable, name);
      }
  
      public void execute(Runnable runnable, String name) {
          init();
          LOG.trace("Execute[{}] runnable: {}", name, runnable);
          ExecutorService executor = executorRef.get();
          if (executor != null) {
              executor.execute(runnable);
          } else {
              doExecuteNewThread(runnable, name);
          }
      }
  
      private void doExecuteNewThread(Runnable runnable, String name) {
          String threadName = name + "-" + id.incrementAndGet();
          Thread thread = new Thread(runnable, threadName);
          thread.setDaemon(daemon);
  
          LOG.trace("Created and running thread[{}]: {}", threadName, thread);
          thread.start();
      }
  
      protected ExecutorService createDefaultExecutor() {
          ThreadPoolExecutor rc = new ThreadPoolExecutor(getDefaultCorePoolSize(), getMaxThreadPoolSize(), getDefaultKeepAliveTime(), TimeUnit.SECONDS, new SynchronousQueue<Runnable>(), new ThreadFactory() {
              @Override
              public Thread newThread(Runnable runnable) {
                  String threadName = name + "-" + id.incrementAndGet();
                  Thread thread = new Thread(runnable, threadName);
                  thread.setDaemon(daemon);
                  thread.setPriority(priority);
                  if (threadClassLoader != null) {
                      thread.setContextClassLoader(threadClassLoader);
                  }
                  thread.setUncaughtExceptionHandler(new Thread.UncaughtExceptionHandler() {
                      @Override
                      public void uncaughtException(final Thread t, final Throwable e) {
                          LOG.error("Error in thread '{}'", t.getName(), e);
                      }
                  });
  
                  LOG.trace("Created thread[{}]: {}", threadName, thread);
                  return thread;
              }
          });
  
          if (rejectedTaskHandler != null) {
              rc.setRejectedExecutionHandler(rejectedTaskHandler);
          } else {
              rc.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
          }
  
          return rc;
      }
  
      protected ExecutorService createVirtualThreadExecutor() {
          if(!(Runtime.version().feature() >= 21)) {
              LOG.error("Virtual Thread support requires JDK 21 or higher");
              throw new IllegalStateException("Virtual Thread support requires JDK 21 or higher");
          }
  
          try {
              Class<?> virtualThreadExecutorClass = Class.forName("org.apache.activemq.thread.VirtualThreadExecutor", false, threadClassLoader);
              Method method = virtualThreadExecutorClass.getMethod("createVirtualThreadExecutorService", String.class, AtomicLong.class, Logger.class);
              Object result = method.invoke(null, name, id, LOG);
              if(!ExecutorService.class.isAssignableFrom(result.getClass())) {
                  throw new IllegalStateException("VirtualThreadExecutor not returned");
              }
              LOG.info("VirtualThreadExecutor initialized name:{}", name);
              return ExecutorService.class.cast(result);
          } catch (ClassNotFoundException | NoSuchMethodException | SecurityException | IllegalAccessException | InvocationTargetException e) {
              LOG.error("VirtualThreadExecutor class failed to load", e);
              throw new IllegalStateException(e);
          }
      }
  
      protected TaskRunner createVirtualThreadTaskRunner(Executor executor, Task task, int maxIterations) {
          if(!(Runtime.version().feature() >= 21)) {
              LOG.error("Virtual Thread support requires JDK 21 or higher");
              throw new IllegalStateException("Virtual Thread support requires JDK 21 or higher");
          }
  
          try {
              Class<?> virtualThreadTaskRunnerClass = Class.forName("org.apache.activemq.thread.VirtualThreadTaskRunner", false, threadClassLoader);
              Constructor<?> constructor = virtualThreadTaskRunnerClass.getConstructor(Executor.class, Task.class, Integer.TYPE);
              Object result = constructor.newInstance(executor, task, maxIterations);
              if(!TaskRunner.class.isAssignableFrom(result.getClass())) {
                  throw new IllegalStateException("VirtualThreadTaskRunner not returned");
              }
              return TaskRunner.class.cast(result);
          } catch (ClassNotFoundException | NoSuchMethodException | SecurityException | IllegalAccessException | InvocationTargetException | InstantiationException | IllegalArgumentException e) {
              LOG.error("VirtualThreadTaskRunner class failed to load", e);
              throw new IllegalStateException(e);
          }
      }
  
  
      public ExecutorService getExecutor() {
          return executorRef.get();
      }
  
      public void setExecutor(ExecutorService executor) {
          this.executorRef.set(executor);
      }
  
      public int getMaxIterationsPerRun() {
          return maxIterationsPerRun;
      }
  
      public void setMaxIterationsPerRun(int maxIterationsPerRun) {
          this.maxIterationsPerRun = maxIterationsPerRun;
      }
  
      public String getName() {
          return name;
      }
  
      public void setName(String name) {
          this.name = name;
      }
  
      public int getPriority() {
          return priority;
      }
  
      public void setPriority(int priority) {
          this.priority = priority;
      }
  
      public boolean isDaemon() {
          return daemon;
      }
  
      public void setDaemon(boolean daemon) {
          this.daemon = daemon;
      }
  
      public boolean isDedicatedTaskRunner() {
          return dedicatedTaskRunner;
      }
  
      public void setDedicatedTaskRunner(boolean dedicatedTaskRunner) {
          this.dedicatedTaskRunner = dedicatedTaskRunner;
      }
  
      public boolean isVirtualThreadTaskRunner() {
          return virtualThreadTaskRunner;
      }
  
      public void setVirtualThreadTaskRunner(boolean virtualThreadTaskRunner) {
          this.virtualThreadTaskRunner = virtualThreadTaskRunner;
      }
  
      public int getMaxThreadPoolSize() {
          return maxThreadPoolSize;
      }
  
      public void setMaxThreadPoolSize(int maxThreadPoolSize) {
          this.maxThreadPoolSize = maxThreadPoolSize;
      }
  
      public void setThreadClassLoader(ClassLoader threadClassLoader) {
          this.threadClassLoader = threadClassLoader;
      }
  
      public RejectedExecutionHandler getRejectedTaskHandler() {
          return rejectedTaskHandler;
      }
  
      public void setRejectedTaskHandler(RejectedExecutionHandler rejectedTaskHandler) {
          this.rejectedTaskHandler = rejectedTaskHandler;
      }
  
      public long getShutdownAwaitTermination() {
          return shutdownAwaitTermination;
      }
  
      public void setShutdownAwaitTermination(long shutdownAwaitTermination) {
          this.shutdownAwaitTermination = shutdownAwaitTermination;
      }
  
      private static int getDefaultCorePoolSize() {
          return Integer.getInteger("org.apache.activemq.thread.TaskRunnerFactory.corePoolSize", 0);
      }
  
      private static int getDefaultMaximumPoolSize() {
          return Integer.getInteger("org.apache.activemq.thread.TaskRunnerFactory.maximumPoolSize", Integer.MAX_VALUE);
      }
  
      private static int getDefaultKeepAliveTime() {
          return Integer.getInteger("org.apache.activemq.thread.TaskRunnerFactory.keepAliveTime", 30);
      }
  }
  