com.clickzetta.platform.client.proxy.router.AbstractRouterRpcProxy 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.proxy.router;

import com.clickzetta.platform.client.api.ClientContext;
import com.clickzetta.platform.client.proxy.AbstractRpcProxy;
import com.clickzetta.platform.client.proxy.RequestConstructor;
import com.clickzetta.platform.client.proxy.ResponseConstructor;
import com.clickzetta.platform.client.proxy.RpcProxy;
import com.clickzetta.platform.common.Constant;
import com.google.protobuf.AbstractMessage;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionV2;

import java.io.IOException;
import java.util.concurrent.CompletableFuture;

public abstract class AbstractRouterRpcProxy extends AbstractRpcProxy {

  // real rpc proxy.
  volatile private RpcProxy rpcProxy;

  // check if need redirect.
  volatile boolean needRedirect = false;

  public AbstractRouterRpcProxy(ClientContext context) {
    super(context);
  }

  private synchronized void validRpcProxyInit(long instanceId) throws IOException {
    validRpcProxyInit(instanceId, "", "", "");
  }

  private synchronized void validRpcProxyInit(long instanceId, String workspace) throws IOException {
    validRpcProxyInit(instanceId, workspace, "", "");
  }

  private synchronized void validRpcProxyInit(long instanceId, String workspace, String schemaName) throws IOException {
    validRpcProxyInit(instanceId, workspace, schemaName, "");
  }

  private synchronized void validRpcProxyInit(long instanceId, String workspace,
                                              String schemaName, String tableName) throws IOException {
    if (needRedirect && rpcProxy != null) {
      // close current rpcProxy & require a new one.
      LOG.info("router hit need redirect. try to redirect with instanceId [{}] workspace [{}] schema [{}.{}]",
          instanceId, workspace, schemaName, tableName);
      try {
        rpcProxy.close();
      } finally {
        needRedirect = false;
        rpcProxy = null;
      }
    }
    if (rpcProxy == null) {
      String address = getRouterServiceAddress(instanceId, workspace, schemaName, tableName);

      // open internal rpc proxy to call rpc message.
      ClientContext contextProxy = getClientContext().mirror();
      contextProxy.getProperties().put("crlAddrs", address);
      contextProxy.refreshAll(contextProxy);
      rpcProxy = RpcProxy.Builder.build(contextProxy);
      try {
        rpcProxy.open();
      } catch (Exception e) {
        if (rpcProxy != null) {
          try {
            rpcProxy.close();
          } catch (Exception ce) {
            LOG.warn("fallback rpc proxy init to close rpc proxy failed.", ce);
          }
          rpcProxy = null;
        }
        throw e;
      }
    }
  }

  protected abstract 
  CompletableFuture routerRpcProxyCallAsync(Ingestion.MethodEnum method, T request) throws IOException;

  @Override
  public 
  CompletableFuture rpcProxyCallAsync(Ingestion.MethodEnum method, T request) throws IOException {
    if (Ingestion.MethodEnum.GET_ROUTER_CONTROLLER_ADDRESS == method) {
      return routerRpcProxyCallAsync(method, request);
    }
    // init rpc proxy first.
    Ingestion.GatewayRequest wrapGatewayRequest = null;
    if (method == Ingestion.MethodEnum.GATEWAY_RPC_CALL) {
      wrapGatewayRequest = (Ingestion.GatewayRequest) request;
    } else {
      // wrap with gateway rpc request with direct rpc call.
      Ingestion.GatewayRequest.Builder builder = Ingestion.GatewayRequest.newBuilder();
      builder.setMethodEnumValue(method.getNumber())
          .setMessage(RequestConstructor.toReqString(request))
          .setVersionInfo(Ingestion.VersionInfo.newBuilder()
              .setName(Constant.VERSION_NAME)
              .setVersion(Constant.VERSION_NUMBER).build())
          .setInstanceId(getClientContext().instanceId());
      if (context.authentication()) {
        if (getClientContext().userId() != null) {
          builder.setUserId(getClientContext().userId());
        }
      }
      wrapGatewayRequest = builder.build();
    }
    long instanceId = wrapGatewayRequest.getInstanceId();
    validRpcProxyInit(instanceId);

    CompletableFuture wrapCompletableFuture =
        rpcProxy.rpcProxyCallAsync(Ingestion.MethodEnum.GATEWAY_RPC_CALL, wrapGatewayRequest);

    Ingestion.MethodEnum gatewayInternalMethod = Ingestion.MethodEnum.forNumber(wrapGatewayRequest.getMethodEnumValue());
    CompletableFuture future = new CompletableFuture<>();
    wrapCompletableFuture.whenCompleteAsync((gatewayResponse, throwable) -> {
      if (throwable != null) {
        future.completeExceptionally(throwable);
      } else {
        try {
          if (gatewayResponse.hasStatus() && gatewayResponse.getStatus().getCode() != Ingestion.Code.SUCCESS) {
            throw new IOException(gatewayResponse.getStatus().getMessage());
          }
          R finalResult;
          if (method != Ingestion.MethodEnum.GATEWAY_RPC_CALL) {
            R response = ResponseConstructor.getResponse(method, gatewayResponse.getMessage());
            finalResult = response;

            // check if need needRedirect.
            if (method == Ingestion.MethodEnum.CREATE_OR_GET_STREAM_V2) {
              IngestionV2.CreateOrGetStreamResponse createOrGetStreamResponse = (IngestionV2.CreateOrGetStreamResponse) response;
              boolean predicate = createOrGetStreamResponse.hasStatus() &&
                  createOrGetStreamResponse.getStatus().getCode() == IngestionV2.Code.NEED_REDIRECT;
              if (!needRedirect && predicate) {
                needRedirect = true;
              }
            }
          } else {
            finalResult = (R) gatewayResponse;

            // check if need needRedirect.
            R response = ResponseConstructor.getResponse(gatewayInternalMethod, gatewayResponse.getMessage());
            if (gatewayInternalMethod == Ingestion.MethodEnum.CREATE_OR_GET_STREAM_V2) {
              IngestionV2.CreateOrGetStreamResponse createOrGetStreamResponse = (IngestionV2.CreateOrGetStreamResponse) response;
              boolean predicate = createOrGetStreamResponse.hasStatus() &&
                  createOrGetStreamResponse.getStatus().getCode() == IngestionV2.Code.NEED_REDIRECT;
              if (!needRedirect && predicate) {
                needRedirect = true;
              }
            }
          }
          future.complete(finalResult);
        } catch (Throwable t) {
          future.completeExceptionally(t);
        }
      }
    });
    return future;
  }

  @Override
  public void close(long wait_time_ms) throws IOException {
    if (rpcProxy != null) {
      try {
        rpcProxy.close(wait_time_ms);
      } catch (Throwable t) {
        // ignore.
      }
    }
    super.close(wait_time_ms);
  }
}