/*
* 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.net;
import java.io.EOFException;
import java.io.IOException;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import com.alibaba.cobar.config.ErrorCode;
import com.alibaba.cobar.net.buffer.BufferPool;
import com.alibaba.cobar.net.buffer.BufferQueue;
import com.alibaba.cobar.util.TimeUtil;
/**
* @author xianmao.hexm
*/
public abstract class AbstractConnection implements NIOConnection {
private static final int OP_NOT_READ = ~SelectionKey.OP_READ;
private static final int OP_NOT_WRITE = ~SelectionKey.OP_WRITE;
protected final SocketChannel channel;
protected NIOProcessor processor;
protected SelectionKey processKey;
protected final ReentrantLock keyLock;
protected int packetHeaderSize;
protected int maxPacketSize;
protected int readBufferOffset;
protected ByteBuffer readBuffer;
protected BufferQueue writeQueue;
protected final ReentrantLock writeLock;
protected boolean isRegistered;
protected final AtomicBoolean isClosed;
protected boolean isSocketClosed;
protected long startupTime;
protected long lastReadTime;
protected long lastWriteTime;
protected long netInBytes;
protected long netOutBytes;
protected int writeAttempts;
public AbstractConnection(SocketChannel channel) {
this.channel = channel;
this.keyLock = new ReentrantLock();
this.writeLock = new ReentrantLock();
this.isClosed = new AtomicBoolean(false);
this.startupTime = TimeUtil.currentTimeMillis();
this.lastReadTime = startupTime;
this.lastWriteTime = startupTime;
}
public SocketChannel getChannel() {
return channel;
}
public int getPacketHeaderSize() {
return packetHeaderSize;
}
public void setPacketHeaderSize(int packetHeaderSize) {
this.packetHeaderSize = packetHeaderSize;
}
public int getMaxPacketSize() {
return maxPacketSize;
}
public void setMaxPacketSize(int maxPacketSize) {
this.maxPacketSize = maxPacketSize;
}
public long getStartupTime() {
return startupTime;
}
public long getLastReadTime() {
return lastReadTime;
}
public long getLastWriteTime() {
return lastWriteTime;
}
public long getNetInBytes() {
return netInBytes;
}
public long getNetOutBytes() {
return netOutBytes;
}
public int getWriteAttempts() {
return writeAttempts;
}
public NIOProcessor getProcessor() {
return processor;
}
public ByteBuffer getReadBuffer() {
return readBuffer;
}
public BufferQueue getWriteQueue() {
return writeQueue;
}
public void setWriteQueue(BufferQueue writeQueue) {
this.writeQueue = writeQueue;
}
/**
* 分配缓存
*/
public ByteBuffer allocate() {
return processor.getBufferPool().allocate();
}
/**
* 回收缓存
*/
public void recycle(ByteBuffer buffer) {
processor.getBufferPool().recycle(buffer);
}
@Override
public void register(Selector selector) throws IOException {
try {
processKey = channel.register(selector, SelectionKey.OP_READ, this);
isRegistered = true;
} finally {
if (isClosed.get()) {
clearSelectionKey();
}
}
}
@Override
public void read() throws IOException {
ByteBuffer buffer = this.readBuffer;
int got = channel.read(buffer);
lastReadTime = TimeUtil.currentTimeMillis();
if (got < 0) {
throw new EOFException();
}
netInBytes += got;
processor.addNetInBytes(got);
// 处理数据
int offset = readBufferOffset, length = 0, position = buffer.position();
for (;;) {
length = getPacketLength(buffer, offset);
if (length == -1) {// 未达到可计算数据包长度的数据
if (!buffer.hasRemaining()) {
checkReadBuffer(buffer, offset, position);
}
break;
}
if (position >= offset + length) {
// 提取一个数据包的数据进行处理
buffer.position(offset);
byte[] data = new byte[length];
buffer.get(data, 0, length);
handle(data);
// 设置偏移量
offset += length;
if (position == offset) {// 数据正好全部处理完毕
if (readBufferOffset != 0) {
readBufferOffset = 0;
}
buffer.clear();
break;
} else {// 还有剩余数据未处理
readBufferOffset = offset;
buffer.position(position);
continue;
}
} else {// 未到达一个数据包的数据
if (!buffer.hasRemaining()) {
checkReadBuffer(buffer, offset, position);
}
break;
}
}
}
public void write(byte[] data) {
ByteBuffer buffer = allocate();
buffer = writeToBuffer(data, buffer);
write(buffer);
}
@Override
public void write(ByteBuffer buffer) {
if (isClosed.get()) {
processor.getBufferPool().recycle(buffer);
return;
}
if (isRegistered) {
try {
writeQueue.put(buffer);
} catch (InterruptedException e) {
error(ErrorCode.ERR_PUT_WRITE_QUEUE, e);
return;
}
processor.postWrite(this);
} else {
processor.getBufferPool().recycle(buffer);
close();
}
}
@Override
public void writeByQueue() throws IOException {
if (isClosed.get()) {
return;
}
final ReentrantLock lock = this.writeLock;
lock.lock();
try {
// 满足以下两个条件时，切换到基于事件的写操作。
// 1.当前key对写事件不该兴趣。
// 2.write0()返回false。
if ((processKey.interestOps() & SelectionKey.OP_WRITE) == 0 && !write0()) {
enableWrite();
}
} finally {
lock.unlock();
}
}
@Override
public void writeByEvent() throws IOException {
if (isClosed.get()) {
return;
}
final ReentrantLock lock = this.writeLock;
lock.lock();
try {
// 满足以下两个条件时，切换到基于队列的写操作。
// 1.write0()返回true。
// 2.发送队列的buffer为空。
if (write0() && writeQueue.size() == 0) {
disableWrite();
}
} finally {
lock.unlock();
}
}
/**
* 打开读事件
*/
public void enableRead() {
final Lock lock = this.keyLock;
lock.lock();
try {
SelectionKey key = this.processKey;
key.interestOps(key.interestOps() | SelectionKey.OP_READ);
} finally {
lock.unlock();
}
processKey.selector().wakeup();
}
/**
* 关闭读事件
*/
public void disableRead() {
final Lock lock = this.keyLock;
lock.lock();
try {
SelectionKey key = this.processKey;
key.interestOps(key.interestOps() & OP_NOT_READ);
} finally {
lock.unlock();
}
}
/**
* 检查WriteBuffer容量，不够则写出当前缓存块并申请新的缓存块。
*/
public ByteBuffer checkWriteBuffer(ByteBuffer buffer, int capacity) {
if (capacity > buffer.remaining()) {
write(buffer);
return processor.getBufferPool().allocate();
} else {
return buffer;
}
}
/**
* 把数据写到给定的缓存中，如果满了则提交当前缓存并申请新的缓存。
*/
public ByteBuffer writeToBuffer(byte[] src, ByteBuffer buffer) {
int offset = 0;
int length = src.length;
int remaining = buffer.remaining();
while (length > 0) {
if (remaining >= length) {
buffer.put(src, offset, length);
break;
} else {
buffer.put(src, offset, remaining);
write(buffer);
buffer = processor.getBufferPool().allocate();
offset += remaining;
length -= remaining;
remaining = buffer.remaining();
continue;
}
}
return buffer;
}
@Override
public boolean close() {
if (isClosed.get()) {
return false;
} else {
if (closeSocket()) {
return isClosed.compareAndSet(false, true);
} else {
return false;
}
}
}
public boolean isClosed() {
return isClosed.get();
}
/**
* 由Processor调用的空闲检查
*/
protected abstract void idleCheck();
/**
* 清理遗留资源
*/
protected void cleanup() {
BufferPool pool = processor.getBufferPool();
ByteBuffer buffer = null;
// 回收接收缓存
buffer = this.readBuffer;
if (buffer != null) {
this.readBuffer = null;
pool.recycle(buffer);
}
// 回收发送缓存
while ((buffer = writeQueue.poll()) != null) {
pool.recycle(buffer);
}
}
/**
* 获取数据包长度，默认是MySQL数据包，其他数据包重载此方法。
*/
protected int getPacketLength(ByteBuffer buffer, int offset) {
if (buffer.position() < offset + packetHeaderSize) {
return -1;
} else {
int length = buffer.get(offset) & 0xff;
length |= (buffer.get(++offset) & 0xff) << 8;
length |= (buffer.get(++offset) & 0xff) << 16;
return length + packetHeaderSize;
}
}
/**
* 检查ReadBuffer容量，不够则扩展当前缓存，直到最大值。
*/
private ByteBuffer checkReadBuffer(ByteBuffer buffer, int offset, int position) {
// 当偏移量为0时需要扩容，否则移动数据至偏移量为0的位置。
if (offset == 0) {
if (buffer.capacity() >= maxPacketSize) {
throw new IllegalArgumentException("Packet size over the limit.");
}
int size = buffer.capacity() << 1;
size = (size > maxPacketSize) ? maxPacketSize : size;
ByteBuffer newBuffer = ByteBuffer.allocate(size);
buffer.position(offset);
newBuffer.put(buffer);
readBuffer = newBuffer;
// 回收扩容前的缓存块
processor.getBufferPool().recycle(buffer);
return newBuffer;
} else {
buffer.position(offset);
buffer.compact();
readBufferOffset = 0;
return buffer;
}
}
private boolean write0() throws IOException {
// 检查是否有遗留数据未写出
ByteBuffer buffer = writeQueue.attachment();
if (buffer != null) {
int written = channel.write(buffer);
if (written > 0) {
netOutBytes += written;
processor.addNetOutBytes(written);
}
lastWriteTime = TimeUtil.currentTimeMillis();
if (buffer.hasRemaining()) {
writeAttempts++;
return false;
} else {
writeQueue.attach(null);
processor.getBufferPool().recycle(buffer);
}
}
// 写出发送队列中的数据块
if ((buffer = writeQueue.poll()) != null) {
// 如果是一块未使用过的buffer，则执行关闭连接。
if (buffer.position() == 0) {
processor.getBufferPool().recycle(buffer);
close();
return true;
}
buffer.flip();
int written = channel.write(buffer);
if (written > 0) {
netOutBytes += written;
processor.addNetOutBytes(written);
}
lastWriteTime = TimeUtil.currentTimeMillis();
if (buffer.hasRemaining()) {
writeQueue.attach(buffer);
writeAttempts++;
return false;
} else {
processor.getBufferPool().recycle(buffer);
}
}
return true;
}
/**
* 打开写事件
*/
private void enableWrite() {
final Lock lock = this.keyLock;
lock.lock();
try {
SelectionKey key = this.processKey;
key.interestOps(key.interestOps() | SelectionKey.OP_WRITE);
} finally {
lock.unlock();
}
processKey.selector().wakeup();
}
/**
* 关闭写事件
*/
private void disableWrite() {
final Lock lock = this.keyLock;
lock.lock();
try {
SelectionKey key = this.processKey;
key.interestOps(key.interestOps() & OP_NOT_WRITE);
} finally {
lock.unlock();
}
}
private void clearSelectionKey() {
final Lock lock = this.keyLock;
lock.lock();
try {
SelectionKey key = this.processKey;
if (key != null && key.isValid()) {
key.attach(null);
key.cancel();
}
} finally {
lock.unlock();
}
}
private boolean closeSocket() {
clearSelectionKey();
SocketChannel channel = this.channel;
if (channel != null) {
boolean isSocketClosed = true;
Socket socket = channel.socket();
if (socket != null) {
try {
socket.close();
} catch (Throwable e) {
}
isSocketClosed = socket.isClosed();
}
try {
channel.close();
} catch (Throwable e) {
}
return isSocketClosed && (!channel.isOpen());
} else {
return true;
}
}
}