com.clickzetta.platform.flusher.StreamTask 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.Table;
import com.clickzetta.platform.client.api.*;
import com.clickzetta.platform.common.NotifyScheduledExecutorService;
import com.clickzetta.platform.connection.ChannelManager;
import com.clickzetta.platform.operator.OperationsEncoder;
import com.clickzetta.platform.operator.WriteOperation;
import cz.proto.ingestion.Ingestion;
import org.apache.kudu.RowOperations;
import scala.Tuple2;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

public class StreamTask extends AbstractTask {

  private NotifyScheduledExecutorService retryThreadPool;
  private AtomicInteger retryRequestCnt;
  private Session session;
  private long batchId;
  private Table table;
  private ClientContext context;
  private long internalMs;
  private RetryMode retryMode;
  private Ingestion.DataMutateRequest request;
  private Ingestion.DataMutateResponse response;
  private Supplier> channelDataSupplier;
  private RpcRequestCallback rpcRequestCallback;
  private Listener listener;

  public StreamTask(
      Session session, NotifyScheduledExecutorService retryThreadPool,
      AtomicInteger retryRequestCnt, Table table, ClientContext context,
      long internalMs, RetryMode retryMode,
      Ingestion.DataMutateRequest request,
      Ingestion.DataMutateResponse response,
      Supplier> channelDataSupplier,
      RpcRequestCallback rpcRequestCallback,
      Listener listener) {
    this.batchId = request.getBatchId();
    this.table = table;
    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) {
          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 {
          Message message = new FailureMessage(session.getSessionId(), table, request, response);
          List rows = message.getErrorRows();
          if (rows != null && !rows.isEmpty()) {
            List operations = rows.stream()
                .map((Function) Row::getWriteOperation).collect(Collectors.toList());
            Tuple2 pbData =
                OperationsEncoder.encodeOperations(operations);
            Ingestion.DataMutateRequest.Builder builder = Ingestion.DataMutateRequest.newBuilder()
                .setBatchId(batchId)
                .setInstanceId(context.instanceId())
                .setWorkspace(context.workspace())
                .setWriteTimestamp(System.currentTimeMillis())
                .setSchemaName(table.getSchemaName())
                .setTableName(table.getTableName())
                .setTableType(table.getIgsTableType())
                .setSchema(table.getMeta())
                .setSchemaPb(table.getSchemaPB())
                .setRowOperations(pbData._1)
                .setKeySchemaPb(table.getKeySchemaPB())
                .setKeyRowOperations(pbData._2)
                .setBucketsCount(table.getBucketsNum())
                .addAllKeyNames(table.getSchema().getKeyIndexMap().keySet())
                .addAllSortNames(table.getSchema().getSortIndexMap().keySet())
                .setBatchCount(rows.size())
                .setIsDispatch(false);
            if (context.authentication()) {
              Ingestion.Account.Builder accountBuilder = Ingestion.Account.newBuilder();
              accountBuilder.setUserName(context.userName());
              if (context.userId() != null) {
                accountBuilder.setUserId(context.userId());
              }
              builder.setToken(context.token()).setAccount(accountBuilder.build());
            }
            Ingestion.DataMutateRequest request = builder.build();
            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 {
            future.complete(true);
          }
        }
      } catch (Throwable t) {
        future.completeExceptionally(t);
      } finally {
        retryRequestCnt.decrementAndGet();
        listener.onRetryTrigger();
      }
    }, internalMs, TimeUnit.MILLISECONDS);
  }
}