org.testng.internal.thread.graph.GraphThreadPoolExecutor Maven / Gradle / Ivy

package org.testng.internal.thread.graph;

import org.testng.TestNGException;
import org.testng.collections.Lists;
import org.testng.internal.DynamicGraph;
import org.testng.internal.DynamicGraph.Status;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * An Executor that launches tasks per batches. It takes a {@code DynamicGraph}
 * of tasks to be run and a {@code IThreadWorkerFactory} to initialize/create
 * {@code Runnable} wrappers around those tasks
 */
public class GraphThreadPoolExecutor extends ThreadPoolExecutor {
  private static final boolean DEBUG = false;
  /** Set to true if you want to generate GraphViz graphs */
  private static final boolean DOT_FILES = false;

  private DynamicGraph m_graph;
  private List m_activeRunnables = Lists.newArrayList();
  private IThreadWorkerFactory m_factory;
  private List m_dotFiles = Lists.newArrayList();
  private int m_threadCount;

  public GraphThreadPoolExecutor(DynamicGraph graph, IThreadWorkerFactory factory, int corePoolSize,
      int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue workQueue) {
    super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue
        /* , new TestNGThreadPoolFactory() */);
    ppp("Initializing executor with " + corePoolSize + " threads and following graph " + graph);
    m_threadCount = maximumPoolSize;
    m_graph = graph;
    m_factory = factory;

    if (m_graph.getFreeNodes().isEmpty()) {
      throw new TestNGException("The graph of methods contains a cycle:" + graph.getEdges());
    }
  }

  public void run() {
    synchronized(m_graph) {
      if (DOT_FILES) {
        m_dotFiles.add(m_graph.toDot());
      }
      List freeNodes = m_graph.getFreeNodes();
      runNodes(freeNodes);
    }
  }

  /**
   * Create one worker per node and execute them.
   */
  private void runNodes(List freeNodes) {
    List> runnables = m_factory.createWorkers(freeNodes);
    for (IWorker r : runnables) {
      m_activeRunnables.add(r);
      ppp("Added to active runnable");
      setStatus(r, Status.RUNNING);
      ppp("Executing: " + r);
      try {
        execute(r);
//        if (m_threadCount > 1) execute(r);
//        else r.run();
      }
      catch(Exception ex) {
        ex.printStackTrace();
      }
    }
  }

  private void setStatus(IWorker worker, Status status) {
    ppp("Set status:" + worker + " status:" + status);
    if (status == Status.FINISHED) {
      m_activeRunnables.remove(worker);
    }
    synchronized(m_graph) {
      for (T m : worker.getTasks()) {
        m_graph.setStatus(m, status);
      }
    }
  }

  @Override
  public void afterExecute(Runnable r, Throwable t) {
    ppp("Finished runnable:" + r);
    setStatus((IWorker) r, Status.FINISHED);
    synchronized(m_graph) {
      ppp("Node count:" + m_graph.getNodeCount() + " and "
          + m_graph.getNodeCountWithStatus(Status.FINISHED) + " finished");
      if (m_graph.getNodeCount() == m_graph.getNodeCountWithStatus(Status.FINISHED)) {
        ppp("Shutting down executor " + this);
        if (DOT_FILES) {
          generateFiles(m_dotFiles);
        }
        shutdown();
      } else {
        if (DOT_FILES) {
          m_dotFiles.add(m_graph.toDot());
        }
        List freeNodes = m_graph.getFreeNodes();
        runNodes(freeNodes);
      }
    }
//    if (m_activeRunnables.isEmpty() && m_index < m_runnables.getSize()) {
//      runNodes(m_index++);
//    }
  }

  private void generateFiles(List files) {
    try {
      File dir = File.createTempFile("TestNG-", "");
      dir.delete();
      dir.mkdir();
      for (int i = 0; i < files.size(); i++) {
        File f = new File(dir, "" + (i < 10 ? "0" : "") + i + ".dot");
        BufferedWriter bw = new BufferedWriter(new FileWriter(f));
        bw.append(files.get(i));
        bw.close();
      }
      if (DOT_FILES) {
        System.out.println("Created graph files in " + dir);
      }
    } catch(IOException ex) {
      ex.printStackTrace();
    }
  }

  private void ppp(String string) {
    if (DEBUG) {
      System.out.println("============ [GraphThreadPoolExecutor] " + Thread.currentThread().getId() + " "
          + string);
    }
  }

}

class TestNGThreadPoolFactory implements ThreadFactory {
  private int m_count = 0;

  @Override
  public Thread newThread(Runnable r) {
    Thread result = new Thread(r);
    result.setName("TestNG-" + m_count++);
    return result;
  }
}