多线程编程之两阶段终止模式

public interface Terminatable {
  void terminate();
}

public abstract class TerminatableSupport extends Thread implements Terminatable {
  public final TerminationToken terminationToken;  // 记录当前的标志位

  public TerminatableSupport() {
    this(new TerminationToken());  // 初始化当前标志位
  }

  public TerminatableSupport(TerminationToken terminationToken) {
    super();
    this.terminationToken = terminationToken;  // 初始化标志位
    terminationToken.register(this);  // 注册当前对象的标志位
  }

  protected abstract void doRun() throws Exception;  // 供子类实现具体任务的方法

  // 钩子方法，用于子类进行一些清理工作
  protected void doCleanup(Exception cause) {}

  // 钩子方法，用于子类进行终止时的一些定制化操作
  protected void doTerminate() {}

  @Override
  public void run() {
    Exception ex = null;
    try {
      // 在当前线程中执行任务时，会判断是否标识为终止，并且剩余任务数小于等于0，是才会真正终止当前线程
      while (!terminationToken.isToShutdown() || terminationToken.reservations.get() > 0) {
        doRun();
      }
    } catch (Exception e) {
      ex = e;
    } finally {
      try {
        doCleanup(ex);  // 当前线程终止后需要执行的操作
      } finally {
        terminationToken.notifyThreadTermination(this);
      }
    }
  }

  @Override
  public void interrupt() {
    terminate();
  }

  @Override
  public void terminate() {
    terminationToken.setToShutdown(true);  // 设置终止状态
    try {
      doTerminate();  // 执行客户端定制的终止操作
    } finally {
      if (terminationToken.reservations.get() <= 0) {
        super.interrupt();  // 如果当前线程处于终止状态，则强制终止当前线程
      }
    }
  }

  // 提供给客户端调用的，即客户端线程必须等待终止完成之后才会继续往下执行
  public void terminate(boolean waitUntilThreadTerminated) {
    terminate();
    if (waitUntilThreadTerminated) {
      try {
        this.join();
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
      }
    }
  }
}

public class TerminationToken {
  protected volatile boolean toShutdown = false;  // 终止状态的标志位
  public final AtomicInteger reservations = new AtomicInteger(0);  // 记录当前剩余任务数

  // 记录了所有注册了TerminationToken的实例，这里使用Queue是因为可能会有多个
  // Terminatable实例共享同一个TeraminationToken，如果是共享的，那么reservations
  // 实例就保存了所有共享当前TerminationToken实例的线程所需要执行的任务总数
  private final Queue<WeakReference<Terminatable>> coordinatedThreads;

  public TerminationToken() {
    coordinatedThreads = new ConcurrentLinkedQueue<>();
  }

  public boolean isToShutdown() {
    return toShutdown;
  }

  public void setToShutdown(boolean toShutdown) {
    this.toShutdown = toShutdown;
  }

  // 将当前Terminatable实例注册到当前TerminationToken中
  protected void register(Terminatable thread) {
    coordinatedThreads.add(new WeakReference<>(thread));
  }

  // 如果是多个Terminatable实例注册到当前TerminationToken中，
  // 则广播当前的终止状态，使得这些实例都会终止
  protected void notifyThreadTermination(Terminatable thread) {
    WeakReference<Terminatable> wrThread;
    Terminatable otherThread;
    while (null != (wrThread = coordinatedThreads.poll())) {
      otherThread = wrThread.get();
      if (null != otherThread && otherThread != thread) {
        otherThread.terminate();
      }
    }
  }
}

public class SomeService {
  private final BlockingQueue<String> queue = new ArrayBlockingQueue<>(100);

  private final Producer producer = new Producer();
  private final Consumer consumer = new Consumer();

  public static void main(String[] args) throws InterruptedException {
    SomeService ss = new SomeService();
    ss.init();
    TimeUnit.SECONDS.sleep(500);
    ss.shutdown();
  }

  // 停止生产者和消费者的执行
  public void shutdown() {
    producer.terminate(true);  // 先停止生产者，只有在生产者完全停止之后才会停止消费者
    consumer.terminate();  // 停止消费者
  }

  // 启动生产者和消费者
  public void init() {
    producer.start();
    consumer.start();
  }

  // 生产者
  private class Producer extends TerminatableSupport {
    private int i = 0;

    @Override
    protected void doRun() throws Exception {
      queue.put(String.valueOf(i++));  // 将任务添加到任务队列中
      consumer.terminationToken.reservations.incrementAndGet();  // 更新需要执行的任务数量
    }
  }

  // 消费者
  private class Consumer extends TerminatableSupport {
    @Override
    protected void doRun() throws Exception {
      String product = queue.take();  // 获取任务
      System.out.println("Processing product: " + product);
      try {
        TimeUnit.SECONDS.sleep(new Random().nextInt(100));  // 模拟消费者对任务的执行
      } catch (InterruptedException e) {
        // ignore
      } finally {
        terminationToken.reservations.decrementAndGet();  // 更新需要执行的任务数量
      }
    }
  }
}