All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.taosdata.jdbc.ws.Transport Maven / Gradle / Ivy

There is a newer version: 3.4.0
Show newest version
package com.taosdata.jdbc.ws;

import com.taosdata.jdbc.TSDBError;
import com.taosdata.jdbc.TSDBErrorNumbers;
import com.taosdata.jdbc.common.SerializeBlock;
import com.taosdata.jdbc.enums.WSFunction;
import com.taosdata.jdbc.rs.ConnectionParam;
import com.taosdata.jdbc.utils.CompletableFutureTimeout;
import com.taosdata.jdbc.utils.ReqId;
import com.taosdata.jdbc.utils.StringUtils;
import com.taosdata.jdbc.ws.entity.*;
import org.java_websocket.exceptions.WebsocketNotConnectedException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;

import javax.net.SocketFactory;
import javax.net.ssl.*;
import java.net.URI;
import java.security.cert.X509Certificate;

import static com.taosdata.jdbc.TSDBErrorNumbers.ERROR_CONNECTION_TIMEOUT;

/**
 * send message
 */
public class Transport implements AutoCloseable {
    private final Logger log = LoggerFactory.getLogger(Transport.class);

    public static final int DEFAULT_MESSAGE_WAIT_TIMEOUT = 60_000;

    public static final int TSDB_CODE_RPC_NETWORK_UNAVAIL = 0x0B;
    public static final int TSDB_CODE_RPC_SOMENODE_NOT_CONNECTED = 0x20;



    private final ArrayList clientArr = new ArrayList<>();;
    private final InFlightRequest inFlightRequest;
    private long timeout;
    private volatile boolean  closed = false;

    private final ConnectionParam connectionParam;
    private final WSFunction wsFunction;
    public static final byte[] EMPTY_BYTE_ARRAY = new byte[0];

    private int currentNodeIndex = 0;
    public Transport(WSFunction function, ConnectionParam param, InFlightRequest inFlightRequest) throws SQLException {
        WSClient master = WSClient.getInstance(param, function, this);
        WSClient slave = WSClient.getSlaveInstance(param, function, this);

        if (param.isDisableSslCertValidation()){
            // creat a TrustManager that trusts all certificates
            TrustManager[] trustAllCerts = new TrustManager[]{
                    new X509TrustManager() {
                        public X509Certificate[] getAcceptedIssuers() {
                            return new X509Certificate[0];
                        }
                        @SuppressWarnings("unused")
                        public void checkClientTrusted(X509Certificate[] certs, String authType) {
                            // Intentionally left blank to accept all certificates
                        }
                        @SuppressWarnings("unused")
                        public void checkServerTrusted(X509Certificate[] certs, String authType) {
                            // Intentionally left blank to accept all certificates
                        }
                    }
            };

            try{
                // create a custom SSLContext
                SSLContext sslContext = SSLContext.getInstance("TLS");
                sslContext.init(null, trustAllCerts, new java.security.SecureRandom());
                // get SSLContext SocketFactory
                final SSLSocketFactory sslSocketFactory = sslContext.getSocketFactory();
                master.setSocketFactory(sslSocketFactory);

                if (slave != null){
                    slave.setSocketFactory(sslSocketFactory);
                }
            } catch (Exception e){
                throw new SQLException("setSocketFactory failed ", e);
            }
        }

        this.clientArr.add(master);
        if (slave != null){
            this.clientArr.add(slave);
        }

        this.inFlightRequest = inFlightRequest;
        this.connectionParam = param;
        this.wsFunction = function;

        this.timeout = param.getRequestTimeout();
    }

    public void setTextMessageHandler(Consumer textMessageHandler) {
        for (WSClient wsClient : clientArr){
            wsClient.setTextMessageHandler(textMessageHandler);
        }
    }

    public void setBinaryMessageHandler(Consumer binaryMessageHandler) {
        for (WSClient wsClient : clientArr) {
            wsClient.setBinaryMessageHandler(binaryMessageHandler);
        }
    }

    public void setTimeout(long timeout) {
        this.timeout = timeout;
    }

    private void reconnect() throws SQLException {
        for (int i = 0; i < clientArr.size() && this.connectionParam.isEnableAutoConnect(); i++){
            boolean reconnected = reconnectCurNode();
            if (reconnected){
                log.debug("reconnect success to {}", StringUtils.getBasicUrl(clientArr.get(currentNodeIndex).serverUri));
                return;
            }

            log.debug("reconnect failed to {}", StringUtils.getBasicUrl(clientArr.get(currentNodeIndex).serverUri));

            currentNodeIndex =  (currentNodeIndex + 1) % clientArr.size();
        }

        close();
        throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_CONNECTION_CLOSED, "Websocket Not Connected Exception");
    }

    private void tmqRethrowConnectionCloseException() throws SQLException {
        // TMQ reconnect will be handled in poll
        if (WSFunction.TMQ.equals(this.wsFunction)){
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_CONNECTION_CLOSED, "Websocket Not Connected Exception");
        }
    }
    @SuppressWarnings("all")
    public Response send(Request request) throws SQLException {
        if (isClosed()){
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_CONNECTION_CLOSED, "Websocket Not Connected Exception");
        }

        Response response = null;
        CompletableFuture completableFuture = new CompletableFuture<>();
        String reqString = request.toString();

        try {
            inFlightRequest.put(new FutureResponse(request.getAction(), request.id(), completableFuture));
        } catch (InterruptedException | TimeoutException e) {
            throw new SQLException(e);
        }

        try {
            clientArr.get(currentNodeIndex).send(reqString);
        } catch (WebsocketNotConnectedException e) {
            tmqRethrowConnectionCloseException();
            reconnect();
            try {
                clientArr.get(currentNodeIndex).send(reqString);
            }catch (Exception ex){
                inFlightRequest.remove(request.getAction(), request.id());
                throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_RESTFul_Client_IOException, e.getMessage());
            }
        }

        CompletableFuture responseFuture = CompletableFutureTimeout.orTimeout(
                completableFuture, timeout, TimeUnit.MILLISECONDS, reqString);
        try {
            response = responseFuture.get();
            handleErrInMasterSlaveMode(response);
        } catch (InterruptedException | ExecutionException e) {
            inFlightRequest.remove(request.getAction(), request.id());
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_QUERY_TIMEOUT, e.getMessage());
        }
        return response;
    }
    public Response send(String action, long reqId, long resultId, long type, byte[] rawData) throws SQLException {
        return send(action, reqId, resultId, type, rawData, EMPTY_BYTE_ARRAY);
    }

    public Response send(String action, long reqId, long resultId, long type, byte[] rawData, byte[] rawData2) throws SQLException {
        if (isClosed()){
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_CONNECTION_CLOSED, "Websocket Not Connected Exception");
        }

        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        try {
            buffer.write(SerializeBlock.longToBytes(reqId));
            buffer.write(SerializeBlock.longToBytes(resultId));
            buffer.write(SerializeBlock.longToBytes(type));
            buffer.write(rawData);
            buffer.write(rawData2);
        } catch (IOException e) {
            throw new SQLException("data serialize error!", e);
        }

        Response response;
        CompletableFuture completableFuture = new CompletableFuture<>();
        try {
            inFlightRequest.put(new FutureResponse(action, reqId, completableFuture));
        } catch (InterruptedException | TimeoutException e) {
            throw new SQLException(e);
        }

        try {
            clientArr.get(currentNodeIndex).send(buffer.toByteArray());
        } catch (WebsocketNotConnectedException e) {
            tmqRethrowConnectionCloseException();
            reconnect();
            try {
                clientArr.get(currentNodeIndex).send(buffer.toByteArray());
            }catch (Exception ex){
                inFlightRequest.remove(action, reqId);
                throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_RESTFul_Client_IOException, e.getMessage());
            }
        }

        String reqString = "action:" + action + ", reqId:" + reqId + ", resultId:" + resultId + ", actionType" + type;
        CompletableFuture responseFuture = CompletableFutureTimeout.orTimeout(completableFuture, timeout, TimeUnit.MILLISECONDS, reqString);
        try {
            response = responseFuture.get();
            handleErrInMasterSlaveMode(response);
        } catch (InterruptedException | ExecutionException e) {
            inFlightRequest.remove(action, reqId);
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_QUERY_TIMEOUT, e.getMessage());
        }
        return response;
    }
    private void handleErrInMasterSlaveMode(Response response) throws InterruptedException{
        if (clientArr.size() > 1 && response instanceof CommonResp){
            CommonResp commonResp = (CommonResp) response;
            if (TSDB_CODE_RPC_NETWORK_UNAVAIL == commonResp.getCode() || TSDB_CODE_RPC_SOMENODE_NOT_CONNECTED == commonResp.getCode()) {
                clientArr.get(currentNodeIndex).closeBlocking();
            }
        }
    }

    public Response sendWithoutRetry(Request request) throws SQLException {
        if (isClosed()){
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_CONNECTION_CLOSED, "Websocket Not Connected Exception");
        }

        Response response;
        CompletableFuture completableFuture = new CompletableFuture<>();
        String reqString = request.toString();

        try {
            inFlightRequest.put(new FutureResponse(request.getAction(), request.id(), completableFuture));
        } catch (InterruptedException | TimeoutException e) {
            throw new SQLException(e);
        }

        try {
            clientArr.get(currentNodeIndex).send(reqString);
        } catch (Exception e) {
            inFlightRequest.remove(request.getAction(), request.id());
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_RESTFul_Client_IOException, e.getMessage());
        }

        CompletableFuture responseFuture = CompletableFutureTimeout.orTimeout(
                completableFuture, timeout, TimeUnit.MILLISECONDS, reqString);
        try {
            response = responseFuture.get();
        } catch (InterruptedException | ExecutionException e) {
            inFlightRequest.remove(request.getAction(), request.id());
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_QUERY_TIMEOUT, e.getMessage());
        }
        return response;
    }

    public void sendWithoutResponse(Request request) throws SQLException  {
        if (isClosed()){
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_CONNECTION_CLOSED, "Websocket Not Connected Exception");
        }

        try {
            clientArr.get(currentNodeIndex).send(request.toString());
        } catch (WebsocketNotConnectedException e) {
            tmqRethrowConnectionCloseException();
            reconnect();
            try {
                clientArr.get(currentNodeIndex).send(request.toString());
            }catch (Exception ex){
                throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_RESTFul_Client_IOException, e.getMessage());
            }
        }
    }

    public boolean isClosed() {
        return closed;
    }

    public boolean isConnectionLost() {
        return clientArr.get(currentNodeIndex).isClosed();
    }

    public void disconnectAndReconnect() throws SQLException {
        try {
            clientArr.get(currentNodeIndex).closeBlocking();
            if (!clientArr.get(currentNodeIndex).reconnectBlockingWithoutRetry()){
                throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_RESTFul_Client_IOException, "websocket reconnect failed!");
            }
        } catch (Exception e) {
            throw TSDBError.createSQLException(TSDBErrorNumbers.ERROR_RESTFul_Client_IOException, e.getMessage());
        }
    }
    @Override
    public synchronized void close() {
        if (isClosed()){
            return;
        }
        closed = true;
        inFlightRequest.close();
        for (WSClient wsClient : clientArr){
            wsClient.shutdown();
        }
    }

    public void checkConnection(int connectTimeout) throws SQLException {
        try {
            if (WSConnection.g_FirstConnection && clientArr.size() > 1) {
                // 测试所有节点,如果连接失败,直接异常
                for (WSClient wsClient : clientArr){
                    if (!wsClient.connectBlocking(connectTimeout, TimeUnit.MILLISECONDS)) {
                        close();
                        throw TSDBError.createSQLException(ERROR_CONNECTION_TIMEOUT,
                                "can't create connection with server " + wsClient.serverUri.toString() + " within: " + connectTimeout + " milliseconds");
                    }
                    log.debug("connect success to {}", StringUtils.getBasicUrl(wsClient.serverUri));
                }

                // 断开其他节点
                for (int i = 0; i < clientArr.size(); i++){
                    if (i != currentNodeIndex) {
                        clientArr.get(i).closeBlocking();
                        log.debug("disconnect success to {}", StringUtils.getBasicUrl(clientArr.get(i).serverUri));
                    }
                }
            } else {
                if (!clientArr.get(currentNodeIndex).connectBlocking(connectTimeout, TimeUnit.MILLISECONDS)) {
                    close();
                    throw TSDBError.createSQLException(ERROR_CONNECTION_TIMEOUT,
                            "can't create connection with server within: " + connectTimeout + " milliseconds");
                }
                log.debug("connect success to {}", StringUtils.getBasicUrl(clientArr.get(currentNodeIndex).serverUri));
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            close();
            throw new SQLException("create websocket connection has been Interrupted ", e);
        }
    }

    public void shutdown() {
        closed = true;
        if (inFlightRequest.hasInFlightRequest()) {
            CompletableFuture future = CompletableFuture.runAsync(() -> {
                try {
                    TimeUnit.MILLISECONDS.sleep(timeout);
                } catch (InterruptedException e) {
                    // ignore
                }
            });
            future.thenRun(this::close);
        } else {
            close();
        }
    }

    public boolean doReconnectCurNode() throws SQLException {
        boolean reconnected = false;
        for (int retryTimes = 0; retryTimes < connectionParam.getReconnectRetryCount(); retryTimes++) {
            try {
                reconnected = clientArr.get(currentNodeIndex).reconnectBlocking();
                if (reconnected) {
                    break;
                }
                Thread.sleep(connectionParam.getReconnectIntervalMs());
            } catch (Exception e) {
                log.error("try connect remote server failed!", e);
            }
        }
        return reconnected;
    }

    public boolean reconnectCurNode() throws SQLException {
        for (int retryTimes = 0; retryTimes < connectionParam.getReconnectRetryCount(); retryTimes++) {
            try {
                boolean reconnected = clientArr.get(currentNodeIndex).reconnectBlocking();
                if (reconnected) {
                    // send con msgs
                    ConnectReq connectReq = new ConnectReq();
                    connectReq.setReqId(ReqId.getReqID());
                    connectReq.setUser(connectionParam.getUser());
                    connectReq.setPassword(connectionParam.getPassword());
                    connectReq.setDb(connectionParam.getDatabase());

                    if (connectionParam.getConnectMode() != 0) {
                        connectReq.setMode(connectionParam.getConnectMode());
                    }

                    ConnectResp auth;
                    auth = (ConnectResp) sendWithoutRetry(new Request(Action.CONN.getAction(), connectReq));

                    if (Code.SUCCESS.getCode() == auth.getCode()) {
                        return true;
                    } else {
                        clientArr.get(currentNodeIndex).closeBlocking();
                        log.error("reconnect failed, code: {}, msg: {}", auth.getCode(), auth.getMessage());
                    }
                }
                Thread.sleep(connectionParam.getReconnectIntervalMs());
            } catch (Exception e) {
                log.error("try connect remote server failed!", e);
            }
        }
        return false;
    }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy