com.clickzetta.platform.arrow.CZArrowSession 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.arrow;

import com.clickzetta.platform.client.*;
import com.clickzetta.platform.client.api.*;
import com.clickzetta.platform.client.api.multi.TableIdentifier;
import com.clickzetta.platform.client.message.RequestMessage;
import com.clickzetta.platform.client.message.ResponseMessage;
import com.clickzetta.platform.client.message.arrow.ArrowRequestMessage;
import com.clickzetta.platform.client.message.arrow.ArrowResponseMessage;
import com.clickzetta.platform.client.proxy.RpcProxy;
import com.clickzetta.platform.common.CZException;
import com.clickzetta.platform.common.NotifyScheduledExecutorService;
import com.clickzetta.platform.connection.ChannelManager;
import com.clickzetta.platform.flusher.AbstractTask;
import com.clickzetta.platform.flusher.ArrowFlushTask;
import com.clickzetta.platform.flusher.ArrowStreamTask;
import com.clickzetta.platform.flusher.Buffer;
import com.clickzetta.platform.util.Util;
import cz.proto.ingestion.v2.IngestionV2;
import cz.proto.ingestion.v2.IngestionWorkerServiceGrpc;
import io.grpc.CallCredentials;
import io.grpc.CallOptions;
import io.grpc.ManagedChannel;
import io.grpc.Metadata;
import io.grpc.stub.StreamObserver;

import java.io.IOException;
import java.util.Arrays;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

import static com.clickzetta.platform.connection.ChannelManager.RPC_CLIENT_STREAM_ID;

public final class CZArrowSession extends AbstractSession {

  // thread-safe
  private volatile boolean shouldFlush;

  private CZArrowSession(RpcProxy client, ArrowStream stream, Options options) {
    super(client, stream, options,
        new RpcCallbackFactory<
            RequestMessage,
            ResponseMessage>() {
          @Override
          public RpcCallback<
              RequestMessage,
              ResponseMessage>
          buildRpcResponseCallback(Session session, Options options) {
            return new ArrowRpcResponseRetryCallback(session, options);
          }

          @Override
          public RpcResponseHandler<
              RequestMessage,
              ResponseMessage> buildResponseHandler(Session session) {
            return new RpcResponseHandler,
                ResponseMessage>(session) {
              @Override
              public void handleCallback(
                  RpcCallback,
                      ResponseMessage> responseCallback,
                  RequestMessage req,
                  ResponseMessage resp) {
                if (resp.getStatusCode() == IngestionV2.Code.SUCCESS.getNumber()) {
                  responseCallback.onSuccess(req, resp);
                } else {
                  // TODO get error row handler.
                  responseCallback.onFailure(req, resp, new CZException(resp.getOriginal().getStatus().toString()));
                }
              }
            };
          }
        });
    this.shouldFlush = false;
  }

  public static CZArrowSession build(CZClient czClient, ArrowStream stream, Options options) {
    return new CZArrowSession(czClient, stream, options);
  }

  @Override
  public Supplier reportLastRpcStatus(long batchId, int statusCode,
                                            IngestionV2.MutateRequest request,
                                            IngestionV2.MutateResponse response,
                                            Supplier> callback) throws IOException {
    if (response.getStatus().getCode() == IngestionV2.Code.SUCCESS) {
      getStream().safeSetServerToken(response.getServerToken());
    }
    return super.reportLastRpcStatus(batchId, statusCode, request, response, callback);
  }

  @Override
  public StreamObserver buildStreamObserver(
      ManagedChannel channel, CallOptions callOptions,
      StreamObserver streamObserver) {
    IngestionWorkerServiceGrpc.IngestionWorkerServiceStub stub = IngestionWorkerServiceGrpc.newStub(channel);
    if (callOptions.getDeadline() != null) {
      stub = stub.withDeadline(callOptions.getDeadline());
    }
    if (callOptions.getCompressor() != null) {
      stub = stub.withCompression(callOptions.getCompressor());
    }
    if (callOptions.getMaxInboundMessageSize() != null) {
      stub = stub.withMaxInboundMessageSize(callOptions.getMaxInboundMessageSize());
    }
    if (callOptions.getMaxOutboundMessageSize() != null) {
      stub = stub.withMaxOutboundMessageSize(callOptions.getMaxOutboundMessageSize());
    }
    stub = stub.withCallCredentials(new CallCredentials() {
      @Override
      public void applyRequestMetadata(RequestInfo requestInfo, Executor executor, MetadataApplier metadataApplier) {
        Metadata metadata = new Metadata();
        metadata.put(Metadata.Key.of(RPC_CLIENT_STREAM_ID, Metadata.ASCII_STRING_MARSHALLER), Util.generatorStreamId(sessionId));
        metadataApplier.apply(metadata);
        LOG.info("session {} apply request metadata with streamId {}", sessionId, metadata);
      }

      @Override
      public void thisUsesUnstableApi() {
      }
    });
    return stub.mutate(streamObserver);
  }

  @Override
  public RequestMessage buildRequestMessage(IngestionV2.MutateRequest request) {
    return new ArrowRequestMessage(request);
  }

  @Override
  public ResponseMessage buildResponseMessage(IngestionV2.MutateResponse response) {
    return new ArrowResponseMessage(response);
  }

  @Override
  public IngestionV2.MutateResponse buildRetryResponseMessage(long requestId, IngestionV2.MutateRequest request) {
    IngestionV2.MutateResponse.Builder builder = IngestionV2.MutateResponse.newBuilder()
        .setBatchId(requestId)
        .setStatus(IngestionV2.ResponseStatus.newBuilder()
            .setCode(IngestionV2.Code.STREAM_UNAVAILABLE).build())
        .setNumRows(request.getDataBlock().getNumRows());
    for (int i = 0; i < request.getDataBlock().getNumRows(); i++) {
      IngestionV2.MutateRowStatus mutateRowStatus = IngestionV2.MutateRowStatus.newBuilder()
          .setRowIndex(i)
          .setCode(IngestionV2.Code.FAILED)
          .setErrorMessage("hit need retry exception. build retry response message").build();
      builder.addRowStatusList(mutateRowStatus);
    }
    return builder.build();
  }

  @Override
  public AbstractTask buildStreamTask(Session session, NotifyScheduledExecutorService retryThreadPool,
                                      AtomicInteger retryRequestCnt, Table table, ClientContext context,
                                      long internalMs, RetryMode retryMode,
                                      IngestionV2.MutateRequest request,
                                      IngestionV2.MutateResponse response,
                                      Supplier> channelDataSupplier,
                                      RpcRequestCallback rpcRequestCallback,
                                      Listener listener) {
    return new ArrowStreamTask(session, retryThreadPool, retryRequestCnt, context,
        internalMs, retryMode, request, response, channelDataSupplier, rpcRequestCallback, listener);
  }

  @Override
  public AbstractTask buildFlushTask(long requestId, ClientContext clientContext, Buffer buffer,
                                     Supplier> channelDataSupplier,
                                     RpcRequestCallback requestCallback,
                                     Listener listener) {
    return new ArrowFlushTask(requestId, clientContext, buffer, stream.getServerToken(),
        channelDataSupplier, requestCallback, listener);
  }

  @Override
  protected void timerTriggerAction() throws IOException {
    runOrWaitInCommit(() -> {
      boolean needFlush;
      synchronized (CZArrowSession.this) {
        needFlush = shouldFlush;
        shouldFlush = false;
      }
      super.flush();
      commitIfNeeded(needFlush);
    });
  }

  @Override
  public void commitIfNeeded() throws IOException {
    boolean needFlush;
    synchronized (CZArrowSession.this) {
      needFlush = shouldFlush;
      shouldFlush = false;
    }
    commitIfNeeded(needFlush);
  }

  public void commitIfNeeded(boolean needFlush) throws IOException {
    // TODO remove commit lock.
    commitWithLock(() -> {
      if (client != null && stream != null) {
        ArrowTable table = stream.getTable();
        if (table.isRequireCommit() && needFlush) {
          ClientContext context = client.getClientContext();
          long commitId = client.asyncCommit(context.instanceId(), context.workspace(),
              Arrays.asList(new TableIdentifier(table.getSchemaName(), table.getTableName())),
              Arrays.asList(getStream().getServerToken() == null ? "" : getStream().getServerToken()));
          LOG.info("Async commit id: {}.", commitId);
          client.checkCommitResult(commitId);
          LOG.info("Commit id {} success.", commitId);

          /**
           * for each upsert mutate. token reset like:
           *
           * /clear(t0)/(null)(t1)(t1)(t1)/clear(t1)/(null)(t2)(t2)/clear(t2)/...../
           * / commit /------------------/ commit /---------------/ commit /------/
           */
          getStream().safeSetServerToken(null, true);
        }
      }
    });
  }

  @Override
  public void apply(Row... rows) throws IOException {
    validInitialize();
    for (Row row : rows) {
      synchronized (this) {
        while (inCommit) {
          try {
            validInitialize();
            wait(commitWaitTimeoutMs);
          } catch (Throwable t) {
            throw new IOException(t);
          }
        }
        validInitialize();
        ArrowRow arrowRow = (ArrowRow) row;
        arrowRow.validate();
        if (currentBuffer == null) {
          currentBuffer = flusher.acquireBuffer();
          rowPoolSupportIfNeed(currentBuffer);
        }
        // mark target table need to flush.
        shouldFlush = true;

        int byteSize = arrowRow.getRowAllocSize();
        if (currentBuffer.isFull(byteSize)) {
          currentBuffer.addOperation(arrowRow);
          sendOneRpcMessage();
        } else {
          currentBuffer.addOperation(arrowRow);
        }
      }
    }
  }

  @Override
  public void flush() throws IOException {
    runOrWaitInCommit(() -> {
      boolean needFlush;
      synchronized (CZArrowSession.this) {
        needFlush = shouldFlush;
        shouldFlush = false;
      }
      super.flush();
      commitIfNeeded(needFlush);
    });
  }
}