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

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

public class MultiStreamTask extends AbstractTask {

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

  public MultiStreamTask(
      Session session, NotifyScheduledExecutorService retryThreadPool,
      AtomicInteger retryRequestCnt, ClientContext context,
      long internalMs, RetryMode retryMode,
      Set retryStatusCode,
      IngestionV2.MultiMutateRequest request,
      IngestionV2.MultiMutateResponse 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.retryStatusCode = retryStatusCode;
    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) {
          // try decode request & re-encode all failed sub-request in multi mutate request.
          if (response.getMutateResponsesList().isEmpty()) {
            ChannelManager.ChannelData channelData = channelDataSupplier.get();
            rpcRequestCallback.setTargetHost(channelData.hostPort.toString());
            rpcRequestCallback.onSuccess(request, future);
            try {
              synchronized (channelData.streamObserver) {
                channelData.streamObserver.onNext(request);
              }
            } catch (Throwable t) {
              rpcRequestCallback.onFailure(request, future, t);
            }
          } else {
            IngestionV2.MultiMutateRequest.Builder builder = request.toBuilder().clearMutateRequests();
            builder.setWriteTimestamp(System.currentTimeMillis());

            ChannelManager.ChannelData channelData = channelDataSupplier.get();
            rpcRequestCallback.setTargetHost(channelData.hostPort.toString());

            // hit retry status & reset server token.
            // retryStatusCode size == 0 means not support retry.
            for (int index = 0; index < response.getMutateResponsesList().size(); index++) {
              IngestionV2.MutateResponse mutateResponse = response.getMutateResponses(index);
              if (mutateResponse.hasStatus() && mutateResponse.getStatus().getCode() == IngestionV2.Code.SUCCESS) {
                continue;
              }
              if (mutateResponse.hasStatus() && !this.retryStatusCode.isEmpty() &&
                  !this.retryStatusCode.contains(mutateResponse.getStatus().getCode().getNumber())) {
                throw new IOException(String.format("MutateResponse contains not retryable status. %s.", mutateResponse.getStatus()));
              } else {
                IngestionV2.MutateRequest mutateRequest = request.getMutateRequests(index);
                MultiTable multiTable = session.getStream().getTable();
                String schemaName = mutateRequest.getTableIdent().getSchemaName();
                String tableName = mutateRequest.getTableIdent().getTableName();
                String serverToken = multiTable.getServerToken(schemaName, tableName);
                if (serverToken == null) {
                  mutateRequest = mutateRequest.toBuilder().clearServerToken().build();
                } else {
                  mutateRequest = mutateRequest.toBuilder().setServerToken(serverToken).build();
                }
                builder.addMutateRequests(mutateRequest);
              }
            }
            IngestionV2.MultiMutateRequest retryRequest = builder.build();

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