/*
* 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.server.session;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.log4j.Logger;
import com.alibaba.cobar.CobarConfig;
import com.alibaba.cobar.CobarServer;
import com.alibaba.cobar.config.ErrorCode;
import com.alibaba.cobar.mysql.MySQLDataNode;
import com.alibaba.cobar.mysql.nio.MySQLConnection;
import com.alibaba.cobar.mysql.nio.handler.CommitNodeHandler;
import com.alibaba.cobar.mysql.nio.handler.KillConnectionHandler;
import com.alibaba.cobar.mysql.nio.handler.MultiNodeQueryHandler;
import com.alibaba.cobar.mysql.nio.handler.RollbackNodeHandler;
import com.alibaba.cobar.mysql.nio.handler.RollbackReleaseHandler;
import com.alibaba.cobar.mysql.nio.handler.SingleNodeHandler;
import com.alibaba.cobar.mysql.nio.handler.Terminatable;
import com.alibaba.cobar.net.FrontendConnection;
import com.alibaba.cobar.net.mysql.OkPacket;
import com.alibaba.cobar.route.RouteResultset;
import com.alibaba.cobar.route.RouteResultsetNode;
import com.alibaba.cobar.server.ServerConnection;
import com.alibaba.cobar.server.parser.ServerParse;
/**
* @author xianmao.hexm 2012-4-12
* @author <a href="mailto:shuo.qius@alibaba-inc.com">QIU Shuo</a>
*/
public class NonBlockingSession implements Session {
private static final Logger LOGGER = Logger.getLogger(NonBlockingSession.class);
private final ServerConnection source;
private final ConcurrentHashMap<RouteResultsetNode, MySQLConnection> target;
private final AtomicBoolean terminating;
// life-cycle: each sql execution
private volatile SingleNodeHandler singleNodeHandler;
private volatile MultiNodeQueryHandler multiNodeHandler;
private volatile CommitNodeHandler commitHandler;
private volatile RollbackNodeHandler rollbackHandler;
public NonBlockingSession(ServerConnection source) {
this.source = source;
this.target = new ConcurrentHashMap<RouteResultsetNode, MySQLConnection>(2, 1);
this.terminating = new AtomicBoolean(false);
}
@Override
public ServerConnection getSource() {
return source;
}
@Override
public int getTargetCount() {
return target.size();
}
public Set<RouteResultsetNode> getTargetKeys() {
return target.keySet();
}
public MySQLConnection getTarget(RouteResultsetNode key) {
return target.get(key);
}
public MySQLConnection removeTarget(RouteResultsetNode key) {
return target.remove(key);
}
@Override
public void execute(RouteResultset rrs, int type) {
if (LOGGER.isDebugEnabled()) {
StringBuilder s = new StringBuilder();
LOGGER.debug(s.append(source).append(rrs).toString());
}
// 检查路由结果是否为空
RouteResultsetNode[] nodes = rrs.getNodes();
if (nodes == null || nodes.length == 0) {
source.writeErrMessage(ErrorCode.ER_NO_DB_ERROR, "No dataNode selected");
return;
}
if (nodes.length == 1) {
singleNodeHandler = new SingleNodeHandler(nodes[0], this);
// singleNodeHandler.execute();
} else {
boolean autocommit = source.isAutocommit();
if (autocommit && isModifySQL(type)) {
autocommit = false;
}
multiNodeHandler = new MultiNodeQueryHandler(nodes, autocommit, this);
// multiNodeHandler.execute();
}
}
public void commit() {
final int initCount = target.size();
if (initCount <= 0) {
ByteBuffer buffer = source.allocate();
buffer = source.writeToBuffer(OkPacket.OK, buffer);
source.write(buffer);
return;
}
commitHandler = new CommitNodeHandler(this);
commitHandler.commit();
}
public void rollback() {
final int initCount = target.size();
if (initCount <= 0) {
ByteBuffer buffer = source.allocate();
buffer = source.writeToBuffer(OkPacket.OK, buffer);
source.write(buffer);
return;
}
rollbackHandler = new RollbackNodeHandler(this);
rollbackHandler.rollback();
}
@Override
public void cancel(FrontendConnection sponsor) {
// TODO Auto-generated method stub
}
/**
* {@link ServerConnection#isClosed()} must be true before invoking this
*/
public void terminate() {
if (!terminating.compareAndSet(false, true)) {
return;
}
kill(new Runnable() {
@Override
public void run() {
new Terminator().nextInvocation(singleNodeHandler)
.nextInvocation(multiNodeHandler)
.nextInvocation(commitHandler)
.nextInvocation(rollbackHandler)
.nextInvocation(new Terminatable() {
@Override
public void terminate(Runnable runnable) {
clearConnections(false);
}
})
.nextInvocation(new Terminatable() {
@Override
public void terminate(Runnable runnable) {
terminating.set(false);
}
})
.invoke();
}
});
}
public boolean closeConnection(RouteResultsetNode key) {
MySQLConnection conn = target.remove(key);
if (conn != null) {
conn.close();
return true;
}
return false;
}
public void setConnectionRunning(RouteResultsetNode[] route) {
for (RouteResultsetNode rrn : route) {
MySQLConnection c = target.get(rrn);
if (c != null) {
c.setRunning(true);
}
}
}
public void clearConnections() {
clearConnections(true);
}
public void releaseConnections() {
for (RouteResultsetNode rrn : target.keySet()) {
MySQLConnection c = target.remove(rrn);
if (c != null) {
if (c.isRunning()) {
c.close();
try {
throw new IllegalStateException("running connection is found: " + c);
} catch (Exception e) {
LOGGER.error(e.getMessage(), e);
}
} else if (!c.isClosedOrQuit()) {
if (source.isClosed()) {
c.quit();
} else {
c.release();
}
}
}
}
}
/**
* @return previous bound connection
*/
public MySQLConnection bindConnection(RouteResultsetNode key, MySQLConnection conn) {
return target.put(key, conn);
}
private static class Terminator {
private LinkedList<Terminatable> list = new LinkedList<Terminatable>();
private Iterator<Terminatable> iter;
public Terminator nextInvocation(Terminatable term) {
list.add(term);
return this;
}
public void invoke() {
iter = list.iterator();
terminate();
}
private void terminate() {
if (iter.hasNext()) {
Terminatable term = iter.next();
if (term != null) {
term.terminate(new Runnable() {
@Override
public void run() {
terminate();
}
});
} else {
terminate();
}
}
}
}
private void kill(Runnable run) {
boolean hooked = false;
AtomicInteger count = null;
Map<RouteResultsetNode, MySQLConnection> killees = null;
for (RouteResultsetNode node : target.keySet()) {
MySQLConnection c = target.get(node);
if (c != null && c.isRunning()) {
if (!hooked) {
hooked = true;
killees = new HashMap<RouteResultsetNode, MySQLConnection>();
count = new AtomicInteger(0);
}
killees.put(node, c);
count.incrementAndGet();
}
}
if (hooked) {
for (Entry<RouteResultsetNode, MySQLConnection> en : killees.entrySet()) {
KillConnectionHandler kill = new KillConnectionHandler(en.getValue(), this, run, count);
CobarConfig conf = CobarServer.getInstance().getConfig();
MySQLDataNode dn = conf.getDataNodes().get(en.getKey().getName());
try {
dn.getConnection(kill, en.getKey());
} catch (Exception e) {
LOGGER.error("get killer connection failed for " + en.getKey(), e);
kill.connectionError(e, null);
}
}
} else {
run.run();
}
}
private void clearConnections(boolean pessimisticRelease) {
for (RouteResultsetNode node : target.keySet()) {
MySQLConnection c = target.remove(node);
if (c == null || c.isClosedOrQuit()) {
continue;
}
// 如果通道正在运行中，则关闭当前通道。
if (c.isRunning() || (pessimisticRelease && source.isClosed())) {
c.close();
continue;
}
// 非事务中的通道，直接释放资源。
if (c.isAutocommit()) {
c.release();
continue;
}
c.setResponseHandler(new RollbackReleaseHandler());
c.rollback();
}
}
public boolean closed() {
return source.isClosed();
}
private static boolean isModifySQL(int type) {
switch (type) {
case ServerParse.INSERT:
case ServerParse.DELETE:
case ServerParse.UPDATE:
case ServerParse.REPLACE:
return true;
default:
return false;
}
}
}