com.clickzetta.platform.client.RpcResponseRetryCallback 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.client;

import com.clickzetta.platform.client.api.ErrorTypeHandler;
import com.clickzetta.platform.client.api.Message;
import com.clickzetta.platform.client.api.Options;
import com.clickzetta.platform.client.api.RetryMode;
import com.clickzetta.platform.client.message.RequestMessage;
import com.clickzetta.platform.client.message.ResponseMessage;
import com.google.protobuf.AbstractMessage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

public abstract class RpcResponseRetryCallback<
    Req extends RequestMessage,
    Resp extends ResponseMessage> implements RpcCallback {

  protected static final Logger LOG = LoggerFactory.getLogger(RpcResponseRetryCallback.class);

  protected final Session session;
  protected final Options options;
  protected final boolean retryEnable;
  protected final RetryMode retryMode;
  protected final int maxRetry;
  protected final long retryInternalMs;
  protected final Map requestRetryTimesMap;
  protected final Set retryStatus;
  protected ErrorTypeHandler errorTypeHandler;

  public RpcResponseRetryCallback(Session session, Options options, List retryStatus) {
    this.session = session;
    this.options = options;
    this.errorTypeHandler = options.getErrorTypeHandler();
    this.retryEnable = options.getRequestFailedRetryEnable();
    this.retryMode = options.getRequestFailedRetryMode();
    this.maxRetry = options.getRequestFailedRetryTimes();
    this.retryInternalMs = options.getRequestFailedRetryInternalMs();
    this.retryStatus = new HashSet<>(retryStatus);
    this.requestRetryTimesMap = new ConcurrentHashMap<>();
  }

  public abstract Message buildSuccessMessage(Req request, Resp response);

  public abstract Message buildFailureMessage(Table table, Req request, Resp response);

  public abstract Message buildOnErrorMessage(Table table, Req request);

  public boolean retryCondition(Set retryStatus, Resp response) {
    return retryStatus.contains(response.getStatusCode());
  }

  @Override
  public void onSuccess(Req request, Resp response) {
    // add all metrics if defined.
    session.getMetrics().getPushDataQPS().mark();
    session.getMetrics().getPushDataRecordCount().mark(request.getBatchCount());
    session.getMetrics().getPushDataE2ELatency().update(
        System.currentTimeMillis() - request.getTimestamp(), TimeUnit.MILLISECONDS);
    session.getMetrics().getPushDataThroughput().mark(request.messageSize());
    requestRetryTimesMap.remove(request.getRequestId());
    errorTypeHandler.onSuccess(buildSuccessMessage(request, response));

    try {
      this.session.reportLastRpcStatus(response.getRequestId(), response.getStatusCode(),
          request.getOriginal(), response.getOriginal(), null);
    } catch (Throwable t) {
      Message logMsg = buildFailureMessage(session.getStream().getTable(), request, response);
      String errMsg = String.format("mutate data with sessionId %s table [%s.%s] batch id %s failed in onSuccess. " +
              "rows %s error row nums %s", logMsg.getSessionId(), logMsg.getSchemaName(), logMsg.getTableName(),
          logMsg.getBatchId(), logMsg.getTotalRowsCounts(), logMsg.getErrorRowsCounts());
      LOG.error(errMsg, t);
    }
  }

  @Override
  public void onFailure(Req request, Resp response, Throwable e) {
    // TODO add more response status check and do retry.
    Message failureMsg = buildFailureMessage(session.getStream().getTable(), request, response);
    if (retryEnable && retryCondition(this.retryStatus, response)) {
      int retryCnt = requestRetryTimesMap.computeIfAbsent(request.getRequestId(), aLong -> 0);
      if (retryCnt < maxRetry) {
        if (options.getRequestFailedRetryLogDebugEnable()) {
          String errMsg = String.format("mutate data with sessionId %s table [%s.%s] batch id %s failed. " +
                  "rows %s error row nums %s. will retry in %s times.",
              failureMsg.getSessionId(), failureMsg.getSchemaName(), failureMsg.getTableName(),
              failureMsg.getBatchId(), failureMsg.getTotalRowsCounts(), failureMsg.getErrorRowsCounts(), retryCnt + 1);
          LOG.info(errMsg, e);
        }
        requestRetryTimesMap.put(request.getRequestId(), retryCnt + 1);
        try {
          Supplier> callback = () -> {
            try {
              return RpcResponseRetryCallback.this.session.sendStreamRequest(
                  requestRetryTimesMap.get(request.getRequestId()) * retryInternalMs,
                  retryMode, request.getOriginal(), response.getOriginal());
            } catch (IOException ex) {
              throw new RuntimeException(ex);
            }
          };
          Supplier supplier = this.session.reportLastRpcStatus(response.getRequestId(), response.getStatusCode(),
              request.getOriginal(), response.getOriginal(), callback);
          if (supplier != null) {
            supplier.get();
          }
          return;
        } catch (Throwable t) {
          String errMsg = String.format("mutate data with sessionId %s table [%s.%s] batch id %s failed in onFailure. " +
                  "rows %s error row nums %s.", failureMsg.getSessionId(), failureMsg.getSchemaName(), failureMsg.getTableName(),
              failureMsg.getBatchId(), failureMsg.getTotalRowsCounts(), failureMsg.getErrorRowsCounts());
          LOG.error(errMsg, t);
        }
      }
    }
    session.getMetrics().getPushDataFailedQPS().mark();
    requestRetryTimesMap.remove(request.getRequestId());
    if (errorTypeHandler.getTerminateIfFailure()) {
      // CZException.
      this.session.setRootCause(new IOException(e));
    }
    errorTypeHandler.onFailure(failureMsg, e);
  }

  @Override
  public void onError(Req request, Resp response, Throwable e) {
    session.getMetrics().getPushDataFailedQPS().mark(1L);
    requestRetryTimesMap.remove(request.getRequestId());
    errorTypeHandler.onFailure(buildOnErrorMessage(session.getStream().getTable(), request), e);
  }
}