com.clickzetta.platform.flusher.AsyncFlusher Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of clickzetta-java Show documentation
The java SDK for clickzetta's Lakehouse
There is a newer version: 2.0.0
package com.clickzetta.platform.flusher;

import com.clickzetta.platform.client.api.Options;
import com.clickzetta.platform.common.Constant;
import com.clickzetta.platform.util.Util;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class AsyncFlusher extends BaseFlusher {

  private ExecutorService flushThreadPool;
  private int nThreads;

  public AsyncFlusher(Options options, Buffer.Type bufferType) {
    super(options, bufferType);
  }

  @Override
  public void init() {
    Object obj = options.getProperties().getOrDefault(Constant.ASYNC_FLUSH_THREAD, 8);
    this.nThreads = obj instanceof String ? Integer.parseInt((String) obj) : (int) obj;
    LOG.info("AsyncFlusher init with nThreads {}", this.nThreads);
    this.flushThreadPool = Executors.newFixedThreadPool(this.nThreads, new ThreadFactoryBuilder()
        .setNameFormat("Session-Flush-Thread-%d")
        .setDaemon(true).build());
    this.flushThreadPool.execute(() -> {
      while (!Thread.currentThread().isInterrupted()) {
        try {
          Task task = flushQueue.poll(200, TimeUnit.MILLISECONDS);
          if (task != null) {
            CompletableFuture future;
            // no need to do task call when hit exception in preTask.call.
            if (getException().get() != null) {
              future = task.skipCall(getException().get());
            } else {
              future = task.call();
            }
            future.whenComplete((o, t) -> {
              if (t != null) {
                setException(t);
              }
              returnBuffer(task.getBuffer());
            });
          }
        } catch (InterruptedException ite) {
          Thread.currentThread().interrupt();
        }
      }
    });
    super.init();
  }

  @Override
  public void executorTask(Task task) throws IOException {
    try {
      validCallException();
      flushQueue.put(task);
    } catch (InterruptedException ite) {
      throw new IOException(ite);
    }
  }

  @Override
  public void close(boolean wait) throws IOException {
    Util.tryWithFinally(() -> {
      try {
        if (wait) {
          synchronized (AsyncFlusher.this) {
            while (flushQueue.size() != 0) {
              validCallException();
              wait(200);
            }
          }
        }
      } catch (InterruptedException ite) {
        throw new IOException(ite);
      }
      validCallException();
      super.close(wait);
    }, () -> {
      try {
        if (this.flushThreadPool != null) {
          flushThreadPool.shutdown();
          boolean success = flushThreadPool.awaitTermination(5 * 1000, TimeUnit.MILLISECONDS);
          if (!success) {
            flushThreadPool.shutdownNow();
          }
        }
      } catch (Throwable t) {
        throw new IOException(t);
      }
    });
  }
}