Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
com.clickzetta.platform.client.api.multi.MultiSession Maven / Gradle / Ivy
package com.clickzetta.platform.client.api.multi;
import com.clickzetta.platform.arrow.ArrowTable;
import com.clickzetta.platform.client.*;
import com.clickzetta.platform.client.api.*;
import com.clickzetta.platform.client.message.RequestMessage;
import com.clickzetta.platform.client.message.ResponseMessage;
import com.clickzetta.platform.client.message.multi.MultiRequestMessage;
import com.clickzetta.platform.client.message.multi.MultiResponseMessage;
import com.clickzetta.platform.client.proxy.RpcProxy;
import com.clickzetta.platform.common.CZException;
import com.clickzetta.platform.common.Constant;
import com.clickzetta.platform.common.NotifyScheduledExecutorService;
import com.clickzetta.platform.connection.ChannelManager;
import com.clickzetta.platform.connection.ReconnectSupport;
import com.clickzetta.platform.flusher.*;
import com.clickzetta.platform.operator.WriteOperation;
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 org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Supplier;
import static com.clickzetta.platform.connection.ChannelManager.RPC_CLIENT_STREAM_ID;
/**
* only support protocol v2.
*/
public final class MultiSession
extends AbstractSession {
// not-thread-safe
private final Set pendingFlushTables;
private MultiSession(RpcProxy client, MultiStream stream, Options options) {
super(client, stream, options,
new RpcCallbackFactory<
RequestMessage,
ResponseMessage>() {
@Override
public RpcCallback<
RequestMessage,
ResponseMessage>
buildRpcResponseCallback(Session session, Options options) {
return new MultiRpcResponseRetryCallback(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.pendingFlushTables = new HashSet<>();
}
public static MultiSession build(CZClient czClient, MultiStream stream, Options options) {
return new MultiSession(czClient, stream, options);
}
@Override
public synchronized void init() {
super.init();
try {
this.flusher.close(true);
} catch (Throwable t) {
throw new RuntimeException(t);
}
this.flusher = Flusher.Builder.build(options, Buffer.Type.MULTI);
// shared exception from flusher.
this.rootCause = this.flusher.getException();
}
@Override
public void initRpcConnection(List> hostPorts) throws IOException {
validInitialize();
Object obj = options.getProperties().getOrDefault(Constant.SESSION_MULTI_RPC_CONNECTION, false);
boolean multiEnable = obj instanceof String ? Boolean.parseBoolean((String) obj) : (boolean) obj;
for (Tuple2 hostPort : hostPorts) {
channelManager.buildChannelInternal(multiEnable, hostPort._1, hostPort._2,
channel -> {
ResponseProxyStreamObserver proxyStreamObserver =
new ResponseProxyStreamObserver<>(responseStreamObserver);
StreamObserver reqStreamObserver = initGrpcStub(channel, callOptions, proxyStreamObserver);
return new ReferenceCountedStreamObserver<>(channel, reqStreamObserver, proxyStreamObserver);
});
}
// init register multi reconnect task.
channelManager.registerReconnectTask(new MultiReconnectTask(new HashSet<>()));
}
@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.multiMutate(streamObserver);
}
@Override
public RequestMessage buildRequestMessage(IngestionV2.MultiMutateRequest request) {
return new MultiRequestMessage(request);
}
@Override
public ResponseMessage buildResponseMessage(IngestionV2.MultiMutateResponse response) {
return new MultiResponseMessage(response);
}
@Override
public IngestionV2.MultiMutateResponse buildRetryResponseMessage(long requestId, IngestionV2.MultiMutateRequest request) {
IngestionV2.MultiMutateResponse.Builder builder = IngestionV2.MultiMutateResponse.newBuilder()
.setBatchId(requestId)
.setStatus(IngestionV2.ResponseStatus.newBuilder()
.setCode(IngestionV2.Code.STREAM_UNAVAILABLE).build());
long numRows = request.getMutateRequestsList().stream()
.map(mutateRequest -> mutateRequest.getDataBlock().getNumRows())
.reduce(Integer::sum).orElse(0).longValue();
builder.setNumRows(numRows);
for (int i = 0; i < request.getMutateRequestsCount(); i++) {
IngestionV2.MutateRequest mutateRequest = request.getMutateRequests(i);
IngestionV2.MutateResponse.Builder mutateResponseBuilder = IngestionV2.MutateResponse.newBuilder()
.setBatchId(requestId)
.setStatus(IngestionV2.ResponseStatus.newBuilder()
.setCode(IngestionV2.Code.STREAM_UNAVAILABLE).build())
.setNumRows(mutateRequest.getDataBlock().getNumRows());
for (int j = 0; j < mutateRequest.getDataBlock().getNumRows(); j++) {
IngestionV2.MutateRowStatus mutateRowStatus = IngestionV2.MutateRowStatus.newBuilder()
.setRowIndex(j)
.setCode(IngestionV2.Code.FAILED)
.setErrorMessage("hit need retry exception. build retry response message").build();
mutateResponseBuilder.addRowStatusList(mutateRowStatus);
}
builder.addMutateResponses(mutateResponseBuilder.build());
}
return builder.build();
}
@Override
public AbstractTask buildStreamTask(Session session, NotifyScheduledExecutorService retryThreadPool,
AtomicInteger retryRequestCnt, Table table, ClientContext context,
long internalMs, RetryMode retryMode,
IngestionV2.MultiMutateRequest request,
IngestionV2.MultiMutateResponse response,
Supplier> channelDataSupplier,
RpcRequestCallback rpcRequestCallback,
Listener listener) {
MultiRpcResponseRetryCallback multiRpcResponseRetryCallback = (MultiRpcResponseRetryCallback) rpcResponseCallback;
Set retryStatusCode = multiRpcResponseRetryCallback.getRetryStatus();
return new MultiStreamTask(session, retryThreadPool, retryRequestCnt, context,
internalMs, retryMode, retryStatusCode, request, response, channelDataSupplier, rpcRequestCallback, listener);
}
@Override
public AbstractTask buildFlushTask(long requestId, ClientContext clientContext, Buffer buffer,
Supplier> channelDataSupplier,
RpcRequestCallback requestCallback,
Listener listener) {
return new MultiFlushTask(stream.getTable(), requestId, clientContext, buffer,
channelDataSupplier, requestCallback, listener);
}
private class MultiReconnectTask implements ReconnectSupport.ReconnectTask {
private final Logger LOG = LoggerFactory.getLogger(MultiReconnectTask.class);
private final Lock tableLock = new ReentrantLock();
Set tables;
public MultiReconnectTask(Set tables) {
this.tables = tables;
}
@Override
public void mergeTask(ReconnectSupport.ReconnectTask task) {
MultiReconnectTask other = (MultiReconnectTask) task;
tableLock.lock();
try {
this.tables.addAll(other.tables);
} finally {
tableLock.unlock();
}
}
@Override
public void run() throws IOException {
if (this.tables == null || this.tables.isEmpty()) {
LOG.info("MultiReconnectTask tables is null or empty. skip run.");
return;
}
List rebuildTable = new ArrayList<>();
tableLock.lock();
try {
rebuildTable.addAll(this.tables);
this.tables.clear();
} finally {
tableLock.unlock();
}
Object obj = options.getProperties().getOrDefault(Constant.SESSION_MULTI_RPC_CONNECTION, false);
boolean multiEnable = obj instanceof String ? Boolean.parseBoolean((String) obj) : (boolean) obj;
List> reconnectHostPorts = client.rebuildIdleTablet(options, rebuildTable);
channelManager.rebuildChannels(multiEnable, reconnectHostPorts,
channel -> {
ResponseProxyStreamObserver proxyStreamObserver =
new ResponseProxyStreamObserver<>(responseStreamObserver);
StreamObserver reqStreamObserver = initGrpcStub(channel, callOptions, proxyStreamObserver);
return new ReferenceCountedStreamObserver<>(channel, reqStreamObserver, proxyStreamObserver);
});
}
}
@Override
public Supplier reportLastRpcStatus(long batchId, int statusCode,
IngestionV2.MultiMutateRequest request,
IngestionV2.MultiMutateResponse response,
Supplier> callback) throws IOException {
// collect which mutate request return stream unavailable.
Set uniqueTables = new HashSet<>();
MultiTable multiTable = stream.getTable();
// update server id for each table
if (request.getMutateRequestsList().size() == response.getMutateResponsesList().size()) {
for (int index = 0; index < response.getMutateResponsesList().size(); index++) {
IngestionV2.MutateRequest mutateRequest = request.getMutateRequests(index);
IngestionV2.MutateResponse mutateResponse = response.getMutateResponses(index);
String serverToken = mutateResponse.getServerToken();
if (serverToken.length() != 0) {
multiTable.safeAddServerToken(mutateRequest.getTableIdent().getSchemaName(),
mutateRequest.getTableIdent().getTableName(), mutateResponse.getServerToken());
}
}
}
int finalStatusCode = statusCode;
if (response.getStatus().getCode() != IngestionV2.Code.SUCCESS) {
// get all multi table.
Set reconnectStatusSet = ((ReconnectSupport) channelManager).getReconnectStatus();
for (int index = 0; index < response.getMutateResponsesList().size(); index++) {
if (reconnectStatusSet.contains(response.getMutateResponses(index).getStatus().getCode().getNumber())) {
IngestionV2.MutateRequest mutateRequest = request.getMutateRequests(index);
uniqueTables.add(multiTable.getTableSchema(mutateRequest.getTableIdent().getSchemaName(),
mutateRequest.getTableIdent().getTableName()));
finalStatusCode = response.getMutateResponses(index).getStatus().getCode().getNumber();
}
}
}
boolean inReconnect = this.channelManager.reportLastRpcStatus(finalStatusCode);
/**
* if rpc success. callback & supplier is null.
*/
Supplier supplier = null;
if (callback != null) {
supplier = callback.get();
}
if (!inReconnect) {
/**
* condition:
* 1、not support reconnect.
* 2、support reconnect and rpc status (success or failed) but no idle status receive before.
*/
this.channelManager.addReconnectFinishTask(null, condition -> condition && responseHandler.getRequests().size() == 0);
return supplier;
} else {
/**
* condition:
* 1、support reconnect and receive other rpc status (success or failed) after last rpc failed with idle status.
*/
if (!uniqueTables.isEmpty()) {
this.channelManager.registerReconnectTask(new MultiReconnectTask(uniqueTables));
}
this.channelManager.addReconnectFinishTask(supplier, condition -> condition && responseHandler.getRequests().size() == 0);
return null;
}
}
@Override
protected void timerTriggerAction() throws IOException {
runOrWaitInCommit(() -> {
Set waitingFlushTables;
synchronized (MultiSession.this) {
waitingFlushTables = new HashSet<>(pendingFlushTables);
pendingFlushTables.clear();
}
// buffer send in sendOneRpcMessage. and call waitOnNoBufferInFight.
super.flush();
commitIfNeeded(waitingFlushTables);
});
}
@Override
public void commitIfNeeded() throws IOException {
Set waitingFlushTables;
synchronized (MultiSession.this) {
waitingFlushTables = new HashSet<>(pendingFlushTables);
pendingFlushTables.clear();
}
commitIfNeeded(waitingFlushTables);
}
public void commitIfNeeded(Set waitingCommitTables) throws IOException {
boolean requireCommit = waitingCommitTables.stream().anyMatch(table -> ((ArrowTable) table).isRequireCommit());
// TODO remove commit lock.
commitWithLock(() -> {
if (requireCommit && !waitingCommitTables.isEmpty()) {
ClientContext ctx = client.getClientContext();
List tableList = new ArrayList<>(waitingCommitTables.size());
List serverTokenList = new ArrayList<>(waitingCommitTables.size());
boolean hasValidServerToken = false;
for (Table table : waitingCommitTables) {
tableList.add(new TableIdentifier(table.getSchemaName(), table.getTableName()));
MultiTable multiTable = MultiSession.this.stream.getTable();
String serverToken = multiTable.getServerToken(table.getSchemaName(), table.getTableName());
if (serverToken != null) {
hasValidServerToken = true;
serverTokenList.add(serverToken);
} else {
serverTokenList.add("");
}
}
long commitId = client.asyncCommit(ctx.instanceId(), ctx.workspace(), tableList,
hasValidServerToken ? serverTokenList : null);
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 /------/
*/
MultiTable multiTable = getStream().getTable();
waitingCommitTables.forEach(table -> {
if (((ArrowTable) table).isRequireCommit()) {
multiTable.safeAddServerToken(table.getSchemaName(), table.getTableName(), null, true);
}
});
}
});
}
@Override
public void apply(Row... rows) throws IOException {
validInitialize();
synchronized (this) {
for (Row row : rows) {
while (inCommit) {
try {
validInitialize();
wait(commitWaitTimeoutMs);
} catch (Throwable t) {
throw new IOException(t);
}
}
validInitialize();
if (currentBuffer == null) {
currentBuffer = flusher.acquireBuffer();
rowPoolSupportIfNeed(currentBuffer);
}
// add target table to pending flush.
pendingFlushTables.add(row.getWriteOperation().getTable());
WriteOperation operation = row.getWriteOperation();
int byteSize = operation.getRowAllocSize();
if (currentBuffer.isFull(byteSize)) {
currentBuffer.addOperation(operation);
sendOneRpcMessage();
} else {
currentBuffer.addOperation(operation);
}
}
}
}
@Override
public void flush() throws IOException {
runOrWaitInCommit(() -> {
Set waitingFlushTables;
synchronized (MultiSession.this) {
waitingFlushTables = new HashSet<>(pendingFlushTables);
pendingFlushTables.clear();
}
// buffer send in sendOneRpcMessage. and call waitOnNoBufferInFight.
super.flush();
commitIfNeeded(waitingFlushTables);
});
}
}
