MySQLConnection.java

/*
 * Copyright 1999-2012 Alibaba Group.
 *  
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.cobar.mysql.nio;

import java.io.UnsupportedEncodingException;
import java.nio.channels.SocketChannel;
import java.security.NoSuchAlgorithmException;
import java.util.concurrent.atomic.AtomicBoolean;

import org.apache.log4j.Logger;

import com.alibaba.cobar.config.Capabilities;
import com.alibaba.cobar.config.ErrorCode;
import com.alibaba.cobar.config.Isolations;
import com.alibaba.cobar.exception.UnknownTxIsolationException;
import com.alibaba.cobar.mysql.CharsetUtil;
import com.alibaba.cobar.mysql.SecurityUtil;
import com.alibaba.cobar.mysql.nio.handler.ResponseHandler;
import com.alibaba.cobar.net.BackendConnection;
import com.alibaba.cobar.net.mysql.AuthPacket;
import com.alibaba.cobar.net.mysql.CommandPacket;
import com.alibaba.cobar.net.mysql.HandshakePacket;
import com.alibaba.cobar.net.mysql.MySQLPacket;
import com.alibaba.cobar.net.mysql.QuitPacket;
import com.alibaba.cobar.route.RouteResultsetNode;
import com.alibaba.cobar.server.ServerConnection;
import com.alibaba.cobar.statistic.SQLRecord;
import com.alibaba.cobar.statistic.SQLRecorder;
import com.alibaba.cobar.util.TimeUtil;

/**
 * @author xianmao.hexm
 */
public class MySQLConnection extends BackendConnection {
    private static final Logger LOGGER = Logger.getLogger(MySQLConnection.class);
    private static final long CLIENT_FLAGS = initClientFlags();

    private static long initClientFlags() {
        int flag = 0;
        flag |= Capabilities.CLIENT_LONG_PASSWORD;
        flag |= Capabilities.CLIENT_FOUND_ROWS;
        flag |= Capabilities.CLIENT_LONG_FLAG;
        flag |= Capabilities.CLIENT_CONNECT_WITH_DB;
        // flag |= Capabilities.CLIENT_NO_SCHEMA;
        // flag |= Capabilities.CLIENT_COMPRESS;
        flag |= Capabilities.CLIENT_ODBC;
        // flag |= Capabilities.CLIENT_LOCAL_FILES;
        flag |= Capabilities.CLIENT_IGNORE_SPACE;
        flag |= Capabilities.CLIENT_PROTOCOL_41;
        flag |= Capabilities.CLIENT_INTERACTIVE;
        // flag |= Capabilities.CLIENT_SSL;
        flag |= Capabilities.CLIENT_IGNORE_SIGPIPE;
        flag |= Capabilities.CLIENT_TRANSACTIONS;
        // flag |= Capabilities.CLIENT_RESERVED;
        flag |= Capabilities.CLIENT_SECURE_CONNECTION;
        // client extension
        // flag |= Capabilities.CLIENT_MULTI_STATEMENTS;
        // flag |= Capabilities.CLIENT_MULTI_RESULTS;
        return flag;
    }

    private static final CommandPacket _READ_UNCOMMITTED = new CommandPacket();
    private static final CommandPacket _READ_COMMITTED = new CommandPacket();
    private static final CommandPacket _REPEATED_READ = new CommandPacket();
    private static final CommandPacket _SERIALIZABLE = new CommandPacket();
    private static final CommandPacket _AUTOCOMMIT_ON = new CommandPacket();
    private static final CommandPacket _AUTOCOMMIT_OFF = new CommandPacket();
    private static final CommandPacket _COMMIT = new CommandPacket();
    private static final CommandPacket _ROLLBACK = new CommandPacket();
    static {
        _READ_UNCOMMITTED.packetId = 0;
        _READ_UNCOMMITTED.command = MySQLPacket.COM_QUERY;
        _READ_UNCOMMITTED.arg = "SET SESSION TRANSACTION ISOLATION LEVEL READ UNCOMMITTED".getBytes();
        _READ_COMMITTED.packetId = 0;
        _READ_COMMITTED.command = MySQLPacket.COM_QUERY;
        _READ_COMMITTED.arg = "SET SESSION TRANSACTION ISOLATION LEVEL READ COMMITTED".getBytes();
        _REPEATED_READ.packetId = 0;
        _REPEATED_READ.command = MySQLPacket.COM_QUERY;
        _REPEATED_READ.arg = "SET SESSION TRANSACTION ISOLATION LEVEL REPEATABLE READ".getBytes();
        _SERIALIZABLE.packetId = 0;
        _SERIALIZABLE.command = MySQLPacket.COM_QUERY;
        _SERIALIZABLE.arg = "SET SESSION TRANSACTION ISOLATION LEVEL SERIALIZABLE".getBytes();
        _AUTOCOMMIT_ON.packetId = 0;
        _AUTOCOMMIT_ON.command = MySQLPacket.COM_QUERY;
        _AUTOCOMMIT_ON.arg = "SET autocommit=1".getBytes();
        _AUTOCOMMIT_OFF.packetId = 0;
        _AUTOCOMMIT_OFF.command = MySQLPacket.COM_QUERY;
        _AUTOCOMMIT_OFF.arg = "SET autocommit=0".getBytes();
        _COMMIT.packetId = 0;
        _COMMIT.command = MySQLPacket.COM_QUERY;
        _COMMIT.arg = "commit".getBytes();
        _ROLLBACK.packetId = 0;
        _ROLLBACK.command = MySQLPacket.COM_QUERY;
        _ROLLBACK.arg = "rollback".getBytes();
    }

    private MySQLConnectionPool pool;
    private long threadId;
    private HandshakePacket handshake;

    private int charsetIndex;
    private String charset;
    private String dbCharset;

    private volatile int txIsolation;
    private volatile boolean autocommit;
    private long clientFlags;
    private boolean isAuthenticated;
    private String user;
    private String password;
    private String schema;
    private Object attachment;

    private final AtomicBoolean isRunning;
    private long lastTime; // QS_TODO
    private final AtomicBoolean isQuit;
    private volatile StatusSync statusSync;

    public MySQLConnection(SocketChannel channel) {
        super(channel);
        this.clientFlags = CLIENT_FLAGS;
        this.lastTime = TimeUtil.currentTimeMillis();
        this.isRunning = new AtomicBoolean(false);
        this.isQuit = new AtomicBoolean(false);
        this.autocommit = true;
    }

    public MySQLConnectionPool getPool() {
        return pool;
    }

    public void setPool(MySQLConnectionPool pool) {
        this.pool = pool;
    }

    public String getUser() {
        return user;
    }

    public void setUser(String user) {
        this.user = user;
    }

    public String getSchema() {
        return schema;
    }

    public void setSchema(String schema) {
        this.schema = schema;
    }

    public void setPassword(String password) {
        this.password = password;
    }

    public HandshakePacket getHandshake() {
        return handshake;
    }

    public void setHandshake(HandshakePacket handshake) {
        this.handshake = handshake;
    }

    public void setCharsetIndex(int charsetIndex) {
        this.charsetIndex = charsetIndex;
    }

    public long getThreadId() {
        return threadId;
    }

    public void setThreadId(long threadId) {
        this.threadId = threadId;
    }

    public String getCharset() {
        return charset;
    }

    public void setCharset(String charset) {
        this.charset = charset;
    }

    public boolean isAuthenticated() {
        return isAuthenticated;
    }

    public void setAuthenticated(boolean isAuthenticated) {
        this.isAuthenticated = isAuthenticated;
    }

    public String getPassword() {
        return password;
    }

    public void authenticate() {
        AuthPacket packet = new AuthPacket();
        packet.packetId = 1;
        packet.clientFlags = clientFlags;
        packet.maxPacketSize = maxPacketSize;
        packet.charsetIndex = charsetIndex;
        packet.user = user;
        try {
            packet.password = passwd(password, handshake);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e.getMessage());
        }
        packet.database = schema;
        packet.write(this);
    }

    public long getLastTime() {
        return lastTime;
    }

    public void setLastTime(long lastTime) {
        this.lastTime = lastTime;
    }

    public void setRunning(boolean running) {
        isRunning.set(running);
    }

    public boolean isRunning() {
        return isRunning.get();
    }

    public boolean isAutocommit() {
        return autocommit;
    }

    public Object getAttachment() {
        return attachment;
    }

    public void setAttachment(Object attachment) {
        this.attachment = attachment;
    }

    public boolean isClosedOrQuit() {
        return isClosed() || isQuit.get();
    }

    private static class StatusSync {
        private final RouteResultsetNode rrn;
        private final MySQLConnection conn;
        private CommandPacket charCmd;
        private CommandPacket isoCmd;
        private CommandPacket acCmd;
        private final int charIndex;
        private final int txIsolation;
        private final boolean autocommit;
        private volatile boolean executed;

        public StatusSync(MySQLConnection conn, RouteResultsetNode rrn, ServerConnection sc, boolean autocommit) {
            this.conn = conn;
            this.rrn = rrn;
            this.charIndex = sc.getCharsetIndex();
            this.charCmd = conn.charsetIndex != charIndex ? getCharsetCommand(charIndex) : null;
            this.txIsolation = sc.getTxIsolation();
            this.isoCmd = conn.txIsolation != txIsolation ? getTxIsolationCommand(txIsolation) : null;
            this.autocommit = autocommit;
            this.acCmd = conn.autocommit != autocommit ? (autocommit ? _AUTOCOMMIT_ON : _AUTOCOMMIT_OFF) : null;
        }

        private Runnable updater;

        public boolean isExecuted() {
            return executed;
        }

        public boolean isSync() {
            return charCmd == null && isoCmd == null && acCmd == null;
        }

        public void update() {
            Runnable updater = this.updater;
            if (updater != null) {
                updater.run();
            }
        }

        /**
         * @return false if sync complete
         */
        public boolean sync() {
            CommandPacket cmd;
            if (charCmd != null) {
                updater = new Runnable() {
                    @Override
                    public void run() {
                        int ci = StatusSync.this.charIndex;
                        conn.charsetIndex = ci;
                        conn.charset = CharsetUtil.getCharset(ci);
                    }
                };
                cmd = charCmd;
                charCmd = null;
                cmd.write(conn);
                return true;
            }
            if (isoCmd != null) {
                updater = new Runnable() {
                    @Override
                    public void run() {
                        conn.txIsolation = StatusSync.this.txIsolation;
                    }
                };
                cmd = isoCmd;
                isoCmd = null;
                cmd.write(conn);
                return true;
            }
            if (acCmd != null) {
                updater = new Runnable() {
                    @Override
                    public void run() {
                        conn.autocommit = StatusSync.this.autocommit;
                    }
                };
                cmd = acCmd;
                acCmd = null;
                cmd.write(conn);
                return true;
            }
            return false;
        }

        public void execute() throws UnsupportedEncodingException {
            executed = true;
            CommandPacket packet = new CommandPacket();
            packet.packetId = 0;
            packet.command = MySQLPacket.COM_QUERY;
            packet.arg = rrn.getStatement().getBytes(conn.getCharset());
            conn.lastTime = TimeUtil.currentTimeMillis();
            packet.write(conn);
        }

        private static CommandPacket getTxIsolationCommand(int txIsolation) {
            switch (txIsolation) {
            case Isolations.READ_UNCOMMITTED:
                return _READ_UNCOMMITTED;
            case Isolations.READ_COMMITTED:
                return _READ_COMMITTED;
            case Isolations.REPEATED_READ:
                return _REPEATED_READ;
            case Isolations.SERIALIZABLE:
                return _SERIALIZABLE;
            default:
                throw new UnknownTxIsolationException("txIsolation:" + txIsolation);
            }
        }

        private static CommandPacket getCharsetCommand(int ci) {
            String charset = CharsetUtil.getCharset(ci);
            StringBuilder s = new StringBuilder();
            s.append("SET names ").append(charset);
            CommandPacket cmd = new CommandPacket();
            cmd.packetId = 0;
            cmd.command = MySQLPacket.COM_QUERY;
            cmd.arg = s.toString().getBytes();
            return cmd;
        }
    }

    /**
     * @return if synchronization finished and execute-sql has already been sent
     *         before
     */
    public boolean syncAndExcute() throws UnsupportedEncodingException {
        StatusSync sync = statusSync;
        if (sync.isExecuted()) {
            return true;
        }
        if (sync.isSync()) {
            sync.update();
            sync.execute();
        } else {
            sync.update();
            sync.sync();
        }
        return false;
    }

    public void execute(RouteResultsetNode rrn, ServerConnection sc, boolean autocommit)
            throws UnsupportedEncodingException {
        StatusSync sync = new StatusSync(this, rrn, sc, autocommit);
        statusSync = sync;
        if (sync.isSync() || !sync.sync()) {
            sync.execute();
        }
    }

    public void quit() {
        if (isQuit.compareAndSet(false, true) && !isClosed()) {
            if (isAuthenticated) {
                // QS_TODO check
                write(writeToBuffer(QuitPacket.QUIT, allocate()));
                write(processor.getBufferPool().allocate());
            } else {
                close();
            }
        }
    }

    @Override
    public boolean close() {
        isQuit.set(true);
        boolean closed = super.close();
        if (closed) {
            pool.deActive();
        }
        return closed;
    }

    public void commit() {
        _COMMIT.write(this);
    }

    public void rollback() {
        _ROLLBACK.write(this);
    }

    public void release() {
        attachment = null;
        statusSync = null;
        setResponseHandler(null);
        pool.releaseChannel(this);
    }

    @Override
    public void error(int errCode, Throwable t) {
        LOGGER.warn(toString(), t);
        switch (errCode) {
        case ErrorCode.ERR_HANDLE_DATA:
            // handle error ..
            break;
        case ErrorCode.ERR_PUT_WRITE_QUEUE:
            // QS_TODO
            break;
        default:
            close();
            if (handler instanceof MySQLConnectionHandler) {
                ((MySQLConnectionHandler) handler).connectionError(t);
            }
        }
    }

    public boolean setResponseHandler(ResponseHandler queryHandler) {
        if (handler instanceof MySQLConnectionHandler) {
            ((MySQLConnectionHandler) handler).setResponseHandler(queryHandler);
            return true;
        }
        return false;
    }

    /**
     * 记录sql执行信息
     */
    public void recordSql(String host, String schema, String stmt) {
        final long now = TimeUtil.currentTimeMillis();
        if (now > this.lastTime) {
            long time = now - this.lastTime;
            SQLRecorder sqlRecorder = this.pool.getSqlRecorder();
            if (sqlRecorder.check(time)) {
                SQLRecord recorder = new SQLRecord();
                recorder.host = host;
                recorder.schema = schema;
                recorder.statement = stmt;
                recorder.startTime = lastTime;
                recorder.executeTime = time;
                recorder.dataNode = pool.getName();
                recorder.dataNodeIndex = pool.getIndex();
                sqlRecorder.add(recorder);
            }
        }
        this.lastTime = now;
    }

    private static byte[] passwd(String pass, HandshakePacket hs) throws NoSuchAlgorithmException {
        if (pass == null || pass.length() == 0) {
            return null;
        }
        byte[] passwd = pass.getBytes();
        int sl1 = hs.seed.length;
        int sl2 = hs.restOfScrambleBuff.length;
        byte[] seed = new byte[sl1 + sl2];
        System.arraycopy(hs.seed, 0, seed, 0, sl1);
        System.arraycopy(hs.restOfScrambleBuff, 0, seed, sl1, sl2);
        return SecurityUtil.scramble411(passwd, seed);
    }

    public String getDbCharset() {
        return dbCharset;
    }

    public void setDbCharset(String dbCharset) {
        this.dbCharset = dbCharset;
    }
}