另一个Java基于阻塞的定时消费内存队列(依赖guava)

本文的代码是对一个Java基于阻塞的定时消费内存队列 - Jackie_JK - 博客园 (cnblogs.com)

方法的改进，完善了包装以及部分细节，非jdk21可能需要更换线程池实例。

消费类型：

@Getter
@AllArgsConstructor
public enum PushType {
  ELASTIC,
  SQL,
  ;
}

队列参数枚举：

@Getter
@AllArgsConstructor
public enum QueueEnum {
  A(PushType.SQL, 3000, 2, 200, Duration.of(20, ChronoUnit.SECONDS)),
  B(PushType.SQL, 3000, 2, 200, Duration.of(20, ChronoUnit.SECONDS)),
  C(PushType.SQL, 3000, 2, 200, Duration.of(20, ChronoUnit.SECONDS)),
  D(PushType.SQL, 3000, 2, 200, Duration.of(20, ChronoUnit.SECONDS)),
  ;

  private final PushType pushType;

  //队列长度
  private final int capacity;

  //消费线程数
  private final int threads;

  //单次消费数量
  private final int batchNum;

  //最长阻塞时间
  private final Duration timeout;


}

初始化以及消费方法：

@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class InsertQueue {

  private static final Map<QueueEnum, BlockingQueue<Object>> TYPE_QUEUE_MAP = new EnumMap<>(QueueEnum.class);

  private static final ExecutorService EXECUTOR_SERVICE = Executors.newVirtualThreadPerTaskExecutor();

  private static boolean watching;

  @FunctionalInterface
  public interface Consumer<T> {
    void accept(T t) throws Exception;
  }

  private static void watch() {
    Thread.startVirtualThread(() -> {
      while (true) {
        watchQueueSize();
        try {
          TimeUnit.MINUTES.sleep(1);
        } catch (InterruptedException e) {
          Thread.currentThread().interrupt();
        }
      }
    }).setName(InsertQueue.class.getSimpleName() + "Monitor");
  }

  public static void watchQueueSize() {
    for (Map.Entry<QueueEnum, BlockingQueue<Object>> entry : TYPE_QUEUE_MAP.entrySet()) {
      BlockingQueue<Object> queue = entry.getValue();
      QueueEnum type = entry.getKey();
      int capacity = type.getCapacity();
      if (capacity <= 0) {
        continue;
      }
      if (queue.size() == capacity) {
        log.warn("{}队列满", type);
      } else if (queue.size() >= capacity * 0.8) {
        log.warn("{}队列占用高", type);
      } else if (queue.size() >= capacity * 0.5) {
        log.warn("{}队列占用超过一半", type);
      }
    }
  }

  public static <T> void init(QueueEnum type, Consumer<List<T>> task) {
    init(type, type.getCapacity(), type.getThreads(), type.getBatchNum(), type.getTimeout(), task);
  }

  public static <T> void init(QueueEnum type,
                              int capacity,
                              int threads,
                              int batchNum,
                              Duration timeout,
                              Consumer<List<T>> task) {
    if (TYPE_QUEUE_MAP.containsKey(type)) {
      throw new IllegalStateException("重复初始化");
    }
    if (!watching) {
      watching = true;
      watch();
    }
    log.info("[{}]{}队列长度{}", InsertQueue.class.getSimpleName(), type, capacity);
    if (capacity != Integer.MAX_VALUE) {
      TYPE_QUEUE_MAP.put(type, new ArrayBlockingQueue<>(capacity));
    } else {
      TYPE_QUEUE_MAP.put(type, new LinkedBlockingQueue<>(capacity));
    }
    for (int i = 0; i < threads; i++) {
      EXECUTOR_SERVICE.submit(() -> {
        while (true) {
          try {
            List<T> pop = pop(type, batchNum, timeout);
            task.accept(pop);
          } catch (Exception e) {
            log.error(type.toString(), e);
          }
        }
      });
    }
  }

  /*
   * 非阻塞
   * 务必处理返回值
   * */
  public static boolean offer(QueueEnum type, Object obj) {
    BlockingQueue<Object> blockingQueue = TYPE_QUEUE_MAP.get(type);
    if (blockingQueue == null) {
      throw new IllegalStateException("未初始化");
    }
    return blockingQueue.offer(obj);
  }


  /*
   * 阻塞
   * */
  public static boolean put(QueueEnum type, Object obj) {
    BlockingQueue<Object> blockingQueue = TYPE_QUEUE_MAP.get(type);
    if (blockingQueue == null) {
      throw new IllegalStateException("未初始化");
    }
    try {
      blockingQueue.put(obj);
    } catch (Exception e) {
      return false;
    }
    return true;
  }

  private static <T> List<T> pop(
      QueueEnum type,
      int batchNum,
      Duration timeout
  ) throws InterruptedException {
    //拿到batchSize个元素或超过maxWaitTime才返回，返回容器不会为空
    BlockingQueue<Object> blockingQueue = TYPE_QUEUE_MAP.get(type);
    BlockingQueue<T> queue = (BlockingQueue<T>) blockingQueue;
    List<T> buffer = new ArrayList<>(batchNum);
    try {
      int drain = 0;
      do {
        drain = Queues.drain(queue, buffer, batchNum, timeout);
      } while (drain == 0);
      return buffer;
    } catch (InterruptedException ex) {
      log.error("[{}]{}队列获取元素中断", InsertQueue.class.getSimpleName(), type);
      throw ex;
    }
  }

  public static int getQueueSize(QueueEnum type) {
    BlockingQueue<Object> queue = TYPE_QUEUE_MAP.get(type);
    if (queue == null) {
      return 0;
    }
    return queue.size();
  }

  /*
   * 优雅停机
   * */
  public static void closeWait() {
    Timer timer = new Timer();
    Thread thread = Thread.currentThread();
    timer.schedule(new TimerTask() {
      @Override
      public void run() {
        for (QueueEnum value : QueueEnum.values()) {
          if (InsertQueue.getQueueSize(value) > 0) {
            log.info("正在等待消费完成");
            return;
          }
        }
        LockSupport.unpark(thread);
      }
    }, 0, 1000);
    LockSupport.park();
  }
}

使用前需要调用#init(QueueEnum type, Consumer<List<T>> task)

往后只需要调用#push(QueueEnum type, Object obj) 方法即可