com.clickzetta.platform.connection.AbstractReconnectSupport 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.connection;

import com.clickzetta.platform.client.api.Options;
import com.clickzetta.platform.common.Constant;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import cz.proto.ingestion.v2.IngestionV2;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.regex.Pattern;

public abstract class AbstractReconnectSupport implements ReconnectSupport {

  private static final Logger LOG = LoggerFactory.getLogger(ReconnectSupport.class);

  private final Set EMPTY_SET = new HashSet<>(0);

  /**
   * reconnect config.
   */
  private Options options;
  private volatile boolean reconnectSupport;
  private Set reconnectStatus;

  /**
   * reconnect status.
   */
  private volatile boolean inReconnect;

  private ReentrantLock taskLock;
  private Condition waitCondition;

  private AtomicReference atomicReference;

  /**
   * reconnect async task.
   */
  private CompletableExecutorService executorService;
  private volatile ReconnectTask reconnectTask;
  private volatile CompletableFuture taskFuture;

  @Override
  public void init(Options options) {
    this.options = options;
    {
      Object obj = options.getProperties().getOrDefault(Constant.TABLET_IDLE_RECREATE_SUPPORT, true);
      this.reconnectSupport = obj instanceof String ? Boolean.parseBoolean((String) obj) : (boolean) obj;
    }
    if (this.reconnectSupport) {
      // register reconnect status first.
      this.reconnectStatus = new HashSet<>();
      this.reconnectStatus.add(IngestionV2.Code.STREAM_UNAVAILABLE.getNumber());
      {
        String obj = (String) options.getProperties().getOrDefault(Constant.TABLET_IDLE_STATUS, "");
        if (!StringUtils.isEmpty(obj)) {
          for (String status : obj.split(Pattern.quote(","))) {
            try {
              IngestionV2.Code code = IngestionV2.Code.valueOf(status);
              this.reconnectStatus.add(code.getNumber());
            } catch (Throwable t) {
              LOG.warn("ignore unknown code with: {}. Error details: {}", status, t);
            }
          }
        }
      }

      this.taskLock = new ReentrantLock();
      this.waitCondition = this.taskLock.newCondition();
      this.atomicReference = new AtomicReference<>();
      this.executorService = new CompletableExecutorService(
          Executors.newSingleThreadExecutor(
              new ThreadFactoryBuilder()
                  .setNameFormat("idle-reconnect-thread-%d")
                  .setDaemon(true).build()));
    }
  }

  @Override
  public boolean supportReconnect() {
    return reconnectSupport;
  }

  @Override
  public void registerReconnectTask(ReconnectTask task) {
    if (supportReconnect()) {
      this.taskLock.lock();
      try {
        if (this.reconnectTask == null) {
          this.reconnectTask = task;
        } else {
          this.reconnectTask.mergeTask(task);
        }
      } finally {
        this.taskLock.unlock();
      }
    }
  }

  @Override
  public boolean reportLastRpcStatus(int statusCode) {
    if (supportReconnect() && this.reconnectStatus.contains(statusCode)) {
      if (!this.inReconnect) {
        this.taskLock.lock();
        try {
          if (!this.inReconnect) {
            this.inReconnect = true;
          }
        } finally {
          this.taskLock.unlock();
        }
      }
    }
    return this.inReconnect;
  }

  @Override
  public void resetReconnectStatus(boolean status) {
    this.inReconnect = status;
  }

  @Override
  public Set getReconnectStatus() {
    if (supportReconnect()) {
      return this.reconnectStatus;
    } else {
      return EMPTY_SET;
    }
  }

  @Override
  public void addReconnectFinishTask(Supplier supplier, Predicate triggerCondition) {
    if (!supportReconnect()) {
      return;
    }
    if (supplier != null) {
      // first recall task callback like resend stream task.
      // no exception throws.
      supplier.get();
    }
    // concurrent control. sync with lock.
    this.taskLock.lock();
    try {
      if (this.inReconnect && triggerCondition.test(this.taskFuture == null)) {
        LOG.info("triggerCondition start.");
        long startTimeMs = System.currentTimeMillis();
        this.taskFuture = this.executorService.submit(() -> reconnectTask.run());
        this.taskFuture.whenComplete((aBoolean, throwable) -> {
          if (throwable != null) {
            atomicReference.compareAndSet(null, new IOException(throwable));
          }
          this.taskLock.lock();
          try {
            this.taskFuture = null;
            this.inReconnect = false;
            this.waitCondition.signalAll();
          } finally {
            this.taskLock.unlock();
          }
          LOG.info("triggerCondition end with cost {} ms", System.currentTimeMillis() - startTimeMs);
        });
      }
    } finally {
      this.taskLock.unlock();
    }
  }

  private void validReconnectTaskException() throws IOException {
    if (atomicReference.get() != null) {
      throw atomicReference.get();
    }
  }

  @Override
  public void waitOnNoInReconnect() throws IOException {
    if (supportReconnect()) {
      if (inReconnect) {
        this.taskLock.lock();
        try {
          while (inReconnect) {
            validReconnectTaskException();
            this.waitCondition.await(200, TimeUnit.MILLISECONDS);
          }
        } catch (Throwable t) {
          throw new IOException(t);
        } finally {
          this.taskLock.unlock();
        }
      }
      validReconnectTaskException();
    }
  }

  protected void close(long wait_time_ms) {
    if (executorService != null) {
      executorService.close(wait_time_ms);
    }
  }

  private static class CompletableExecutorService {
    private final ExecutorService executorService;

    public CompletableExecutorService(ExecutorService executorService) {
      this.executorService = executorService;
    }

    public CompletableFuture submit(ReconnectTask task) {
      CompletableFuture completableFuture = new CompletableFuture<>();
      this.executorService.submit(new Runnable() {
        @Override
        public void run() {
          try {
            task.run();
            completableFuture.complete(true);
          } catch (Throwable t) {
            completableFuture.completeExceptionally(t);
          }
        }
      });
      return completableFuture;
    }

    public void close(long wait_time_ms) {
      if (executorService != null) {
        executorService.shutdown();
        try {
          boolean closed = executorService.awaitTermination(wait_time_ms, TimeUnit.MILLISECONDS);
          if (!closed) {
            executorService.shutdownNow();
          }
        } catch (InterruptedException ite) {
          // ignore.
        }
      }
    }
  }
}