com.clickzetta.platform.flusher.ArrowStreamTask 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.RpcRequestCallback;
import com.clickzetta.platform.client.Session;
import com.clickzetta.platform.client.api.ArrowStream;
import com.clickzetta.platform.client.api.ClientContext;
import com.clickzetta.platform.client.api.Listener;
import com.clickzetta.platform.client.api.RetryMode;
import com.clickzetta.platform.common.NotifyScheduledExecutorService;
import com.clickzetta.platform.connection.ChannelManager;
import cz.proto.ingestion.v2.IngestionV2;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

public class ArrowStreamTask extends AbstractTask {

  private NotifyScheduledExecutorService retryThreadPool;
  private AtomicInteger retryRequestCnt;
  private Session session;
  private long batchId;
  private ClientContext context;
  private long internalMs;
  private RetryMode retryMode;
  private IngestionV2.MutateRequest request;
  private IngestionV2.MutateResponse response;
  private Supplier> channelDataSupplier;
  private RpcRequestCallback rpcRequestCallback;
  private Listener listener;

  public ArrowStreamTask(
      Session session, NotifyScheduledExecutorService retryThreadPool,
      AtomicInteger retryRequestCnt, ClientContext context,
      long internalMs, RetryMode retryMode,
      IngestionV2.MutateRequest request,
      IngestionV2.MutateResponse response,
      Supplier> channelDataSupplier,
      RpcRequestCallback rpcRequestCallback,
      Listener listener) {
    this.batchId = request.getBatchId();
    this.context = context;
    this.session = session;
    this.retryRequestCnt = retryRequestCnt;
    this.retryThreadPool = retryThreadPool;
    this.internalMs = internalMs;
    this.retryMode = retryMode;
    this.request = request;
    this.response = response;
    this.channelDataSupplier = channelDataSupplier;
    this.rpcRequestCallback = rpcRequestCallback;
    this.listener = listener;
  }

  @Override
  public long getId() {
    return batchId;
  }

  @Override
  public Buffer getBuffer() {
    return null;
  }

  @Override
  public CompletableFuture skipCall(Throwable t) {
    try {
      return super.skipCall(t);
    } finally {
      retryRequestCnt.decrementAndGet();
    }
  }

  @Override
  public void callInternal() throws Exception {
    this.retryThreadPool.schedule(() -> {
      try {
        if (retryMode == RetryMode.BATCH_REQUEST_MODE || retryMode == RetryMode.ROW_REQUEST_MODE) {
          ChannelManager.ChannelData channelData = channelDataSupplier.get();
          rpcRequestCallback.setTargetHost(channelData.hostPort.toString());

          // reset server token.
          String serverToken = ((ArrowStream) session.getStream()).getServerToken();
          if (serverToken == null) {
            request = request.toBuilder().clearServerToken().build();
          } else {
            request = request.toBuilder().setServerToken(serverToken).build();
          }

          rpcRequestCallback.onSuccess(request, future);
          try {
            synchronized (channelData.streamObserver) {
              channelData.streamObserver.onNext(request);
            }
          } catch (Throwable t) {
            rpcRequestCallback.onFailure(request, future, t);
          }
        }
      } catch (Throwable t) {
        future.completeExceptionally(t);
      } finally {
        retryRequestCnt.decrementAndGet();
        listener.onRetryTrigger();
      }
    }, internalMs, TimeUnit.MILLISECONDS);
  }
}