forked from KNEMC/KLALB
445 lines
14 KiB
Java
445 lines
14 KiB
Java
package org.kne.cloud.network;
|
|
|
|
import java.io.IOException;
|
|
import java.nio.ByteBuffer;
|
|
import java.nio.channels.*;
|
|
import java.util.concurrent.atomic.AtomicBoolean;
|
|
import java.util.concurrent.locks.ReentrantLock;
|
|
|
|
/**
|
|
* 修复的BufferedChannel，解决Socket关闭时的数据一致性问题
|
|
*/
|
|
public class BufferedChannel implements ReadableByteChannel, WritableByteChannel, ScatteringByteChannel, GatheringByteChannel {
|
|
private final ReadableByteChannel inputChannel;
|
|
private final WritableByteChannel outputChannel;
|
|
private final ByteBuffer readBuffer;
|
|
private final ByteBuffer writeBuffer;
|
|
|
|
private final ReentrantLock readLock = new ReentrantLock();
|
|
private final ReentrantLock writeLock = new ReentrantLock();
|
|
|
|
private final AtomicBoolean closed = new AtomicBoolean(false);
|
|
private final AtomicBoolean inputShutdown = new AtomicBoolean(false);
|
|
private final AtomicBoolean outputShutdown = new AtomicBoolean(false);
|
|
|
|
public BufferedChannel(ReadableByteChannel channel, int bufferSize) {
|
|
this(channel, null, bufferSize);
|
|
}
|
|
|
|
public BufferedChannel(WritableByteChannel channel, int bufferSize) {
|
|
this(null, channel, bufferSize);
|
|
}
|
|
|
|
public BufferedChannel(ReadableByteChannel inputChannel, WritableByteChannel outputChannel, int bufferSize) {
|
|
this.inputChannel = inputChannel;
|
|
this.outputChannel = outputChannel;
|
|
this.readBuffer = ByteBuffer.allocateDirect(bufferSize);
|
|
this.writeBuffer = ByteBuffer.allocateDirect(bufferSize);
|
|
|
|
// 更安全的初始化
|
|
this.readBuffer.limit(0); // 明确设置为空
|
|
this.writeBuffer.clear();
|
|
}
|
|
|
|
@Override
|
|
public int read(ByteBuffer dst) throws IOException {
|
|
if (closed.get() || inputShutdown.get()) {
|
|
throw new ClosedChannelException();
|
|
}
|
|
|
|
readLock.lock();
|
|
try {
|
|
// 检查底层通道是否仍然打开
|
|
if (inputChannel != null && !inputChannel.isOpen()) {
|
|
inputShutdown.set(true);
|
|
return handleInputShutdown();
|
|
}
|
|
|
|
int totalRead = 0;
|
|
boolean eofEncountered = false;
|
|
|
|
while (dst.hasRemaining() && !eofEncountered) {
|
|
// 如果读缓冲区有数据，先从中读取
|
|
if (readBuffer.hasRemaining()) {
|
|
int bytesToCopy = Math.min(readBuffer.remaining(), dst.remaining());
|
|
|
|
// 使用绝对位置操作，避免修改缓冲区状态
|
|
int oldlimit=readBuffer.limit();
|
|
readBuffer.limit(readBuffer.position()+bytesToCopy);
|
|
dst.put(readBuffer);
|
|
readBuffer.limit(oldlimit);
|
|
|
|
totalRead += bytesToCopy;
|
|
continue;
|
|
}
|
|
|
|
// 读缓冲区空了，需要重新填充
|
|
readBuffer.clear();
|
|
int bytesRead;
|
|
try {
|
|
bytesRead = inputChannel.read(readBuffer);
|
|
} catch (IOException e) {
|
|
// 读取时发生IO异常，标记为关闭
|
|
inputShutdown.set(true);
|
|
throw e;
|
|
}
|
|
|
|
if (bytesRead == -1) {
|
|
// 到达EOF
|
|
inputShutdown.set(true);
|
|
eofEncountered = true;
|
|
} else if (bytesRead == 0) {
|
|
// 没有数据可用，可能是非阻塞模式
|
|
break;
|
|
} else {
|
|
readBuffer.flip();
|
|
}
|
|
}
|
|
|
|
// 如果遇到EOF且没有读取到任何数据，返回-1
|
|
if (eofEncountered && totalRead == 0) {
|
|
return -1;
|
|
}
|
|
|
|
return totalRead;
|
|
|
|
} finally {
|
|
readLock.unlock();
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 处理输入关闭的情况
|
|
*/
|
|
private int handleInputShutdown() throws IOException {
|
|
// 如果读缓冲区还有剩余数据，先返回这些数据
|
|
if (readBuffer.hasRemaining()) {
|
|
return readBuffer.remaining();
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
@Override
|
|
public long read(ByteBuffer[] dsts) throws IOException {
|
|
return read(dsts, 0, dsts.length);
|
|
}
|
|
|
|
@Override
|
|
public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
|
|
if (closed.get() || inputShutdown.get()) {
|
|
throw new ClosedChannelException();
|
|
}
|
|
|
|
long totalRead = 0;
|
|
for (int i = offset; i < offset + length; i++) {
|
|
ByteBuffer dst = dsts[i];
|
|
if (dst == null) {
|
|
throw new NullPointerException("Destination buffer is null");
|
|
}
|
|
|
|
int bytesRead = read(dst);
|
|
if (bytesRead == -1) {
|
|
// 只有在没有读取任何数据时才返回-1
|
|
return totalRead > 0 ? totalRead : -1;
|
|
}
|
|
totalRead += bytesRead;
|
|
}
|
|
|
|
return totalRead;
|
|
}
|
|
|
|
@Override
|
|
public int write(ByteBuffer src) throws IOException {
|
|
if (closed.get() || outputShutdown.get()) {
|
|
throw new ClosedChannelException();
|
|
}
|
|
|
|
writeLock.lock();
|
|
try {
|
|
// 检查底层通道是否仍然打开
|
|
if (outputChannel != null && !outputChannel.isOpen()) {
|
|
outputShutdown.set(true);
|
|
throw new ClosedChannelException();
|
|
}
|
|
|
|
int totalWritten = 0;
|
|
|
|
while (src.hasRemaining()) {
|
|
// 如果写缓冲区有空间，先填充
|
|
if (writeBuffer.hasRemaining()) {
|
|
int bytesToCopy = Math.min(writeBuffer.remaining(), src.remaining());
|
|
int oldlimit= src.limit();
|
|
src.limit(src.position()+bytesToCopy);
|
|
writeBuffer.put(src);
|
|
src.limit(oldlimit);
|
|
|
|
totalWritten += bytesToCopy;
|
|
}
|
|
|
|
// 如果写缓冲区满了或者源数据还很多，刷新缓冲区
|
|
if (!writeBuffer.hasRemaining() ) {
|
|
try {
|
|
flushInternal();
|
|
} catch (IOException e) {
|
|
outputShutdown.set(true);
|
|
throw e;
|
|
}
|
|
}
|
|
}
|
|
return totalWritten;
|
|
|
|
} finally {
|
|
writeLock.unlock();
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public long write(ByteBuffer[] srcs) throws IOException {
|
|
return write(srcs, 0, srcs.length);
|
|
}
|
|
|
|
@Override
|
|
public long write(ByteBuffer[] srcs, int offset, int length) throws IOException {
|
|
if (closed.get() || outputShutdown.get()) {
|
|
throw new ClosedChannelException();
|
|
}
|
|
|
|
long totalWritten = 0;
|
|
for (int i = offset; i < offset + length; i++) {
|
|
ByteBuffer src = srcs[i];
|
|
if (src == null) {
|
|
throw new NullPointerException("Source buffer is null");
|
|
}
|
|
|
|
int bytesWritten = write(src);
|
|
totalWritten += bytesWritten;
|
|
}
|
|
|
|
return totalWritten;
|
|
}
|
|
|
|
/**
|
|
* 安全的刷新方法
|
|
*/
|
|
public void flush() throws IOException {
|
|
if (closed.get()) {
|
|
throw new ClosedChannelException();
|
|
}
|
|
|
|
if (outputChannel == null || outputShutdown.get()) {
|
|
return;
|
|
}
|
|
|
|
writeLock.lock();
|
|
try {
|
|
flushInternal();
|
|
} catch (IOException e) {
|
|
outputShutdown.set(true);
|
|
throw e;
|
|
} finally {
|
|
writeLock.unlock();
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 内部刷新方法，假设已经持有writeLock
|
|
*/
|
|
private void flushInternal() throws IOException {
|
|
if (writeBuffer.position() > 0) {
|
|
writeBuffer.flip();
|
|
|
|
try {
|
|
while (writeBuffer.hasRemaining()) {
|
|
int written = outputChannel.write(writeBuffer);
|
|
// System.out.println("发送"+written);
|
|
if (written == 0) {
|
|
// 可能遇到阻塞或关闭
|
|
if (!outputChannel.isOpen()) {
|
|
outputShutdown.set(true);
|
|
throw new ClosedChannelException();
|
|
}
|
|
// 给其他操作机会
|
|
Thread.yield();
|
|
}
|
|
}
|
|
} finally {
|
|
// 无论发生什么，确保写缓冲区处于可写状态
|
|
writeBuffer.clear();
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 安全的关闭方法
|
|
*/
|
|
@Override
|
|
public void close() throws IOException {
|
|
if (!closed.compareAndSet(false, true)) {
|
|
return; // 已经关闭
|
|
}
|
|
|
|
IOException exception = null;
|
|
|
|
// 先刷新输出缓冲区
|
|
if (outputChannel != null && !outputShutdown.get()) {
|
|
writeLock.lock();
|
|
try {
|
|
if (writeBuffer.position() > 0) {
|
|
try {
|
|
flushInternal();
|
|
} catch (IOException e) {
|
|
exception = e;
|
|
}
|
|
}
|
|
} finally {
|
|
writeLock.unlock();
|
|
}
|
|
}
|
|
|
|
// 关闭底层通道
|
|
try {
|
|
if (inputChannel != null) {
|
|
inputChannel.close();
|
|
}
|
|
} catch (IOException e) {
|
|
if (exception == null) {
|
|
exception = e;
|
|
} else {
|
|
exception.addSuppressed(e);
|
|
}
|
|
}
|
|
|
|
try {
|
|
if (outputChannel != null) {
|
|
outputChannel.close();
|
|
}
|
|
} catch (IOException e) {
|
|
if (exception == null) {
|
|
exception = e;
|
|
} else {
|
|
exception.addSuppressed(e);
|
|
}
|
|
}
|
|
|
|
// 清理缓冲区状态
|
|
readLock.lock();
|
|
try {
|
|
readBuffer.clear();
|
|
readBuffer.limit(0); // 标记为已清空
|
|
} finally {
|
|
readLock.unlock();
|
|
}
|
|
|
|
writeLock.lock();
|
|
try {
|
|
writeBuffer.clear();
|
|
} finally {
|
|
writeLock.unlock();
|
|
}
|
|
|
|
// 设置关闭状态
|
|
inputShutdown.set(true);
|
|
outputShutdown.set(true);
|
|
|
|
if (exception != null) {
|
|
throw exception;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 优雅关闭 - 只关闭输入或输出
|
|
*/
|
|
public void shutdownInput() throws IOException {
|
|
inputShutdown.set(true);
|
|
readLock.lock();
|
|
try {
|
|
readBuffer.clear();
|
|
readBuffer.limit(0);
|
|
} finally {
|
|
readLock.unlock();
|
|
}
|
|
|
|
if (inputChannel instanceof SocketChannel) {
|
|
((SocketChannel) inputChannel).shutdownInput();
|
|
}
|
|
}
|
|
|
|
public void shutdownOutput() throws IOException {
|
|
outputShutdown.set(true);
|
|
writeLock.lock();
|
|
try {
|
|
if (writeBuffer.position() > 0) {
|
|
flushInternal();
|
|
}
|
|
} finally {
|
|
writeLock.unlock();
|
|
}
|
|
|
|
if (outputChannel instanceof SocketChannel) {
|
|
((SocketChannel) outputChannel).shutdownOutput();
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public boolean isOpen() {
|
|
return !closed.get() &&
|
|
(inputChannel == null || inputChannel.isOpen()) &&
|
|
(outputChannel == null || outputChannel.isOpen());
|
|
}
|
|
|
|
/**
|
|
* 检查是否还有可读数据（包括缓冲区中的）
|
|
*/
|
|
public boolean hasRemaining() throws IOException {
|
|
if (closed.get() || inputShutdown.get()) {
|
|
return false;
|
|
}
|
|
|
|
readLock.lock();
|
|
try {
|
|
return readBuffer.hasRemaining() ||
|
|
(inputChannel != null && inputChannel.isOpen());
|
|
} finally {
|
|
readLock.unlock();
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 获取缓冲区状态信息（用于调试）
|
|
*/
|
|
public String getBufferState() {
|
|
readLock.lock();
|
|
writeLock.lock();
|
|
try {
|
|
return String.format(
|
|
"ReadBuffer[pos=%d, lim=%d, cap=%d], WriteBuffer[pos=%d, lim=%d, cap=%d], " +
|
|
"closed=%b, inputShutdown=%b, outputShutdown=%b",
|
|
readBuffer.position(), readBuffer.limit(), readBuffer.capacity(),
|
|
writeBuffer.position(), writeBuffer.limit(), writeBuffer.capacity(),
|
|
closed.get(), inputShutdown.get(), outputShutdown.get()
|
|
);
|
|
} finally {
|
|
writeLock.unlock();
|
|
readLock.unlock();
|
|
}
|
|
}
|
|
|
|
// 其他方法保持不变...
|
|
public int available() throws IOException {
|
|
if (closed.get() || inputShutdown.get()) {
|
|
return 0;
|
|
}
|
|
|
|
readLock.lock();
|
|
try {
|
|
return readBuffer.remaining();
|
|
} finally {
|
|
readLock.unlock();
|
|
}
|
|
}
|
|
|
|
public long skip(long n) throws IOException {
|
|
// 实现保持不变...
|
|
return 0;
|
|
}
|
|
} |