public void _testOperationEncodeTimeout() throws Exception {
memcachedClient.set("name", 0, "dennis");
assertEquals("dennis", memcachedClient.get("name"));
long writeMessageCount = memcachedClient.getConnector().getStatistics().getWriteMessageCount();
CountDownLatch latch = new CountDownLatch(1);
Command errorCommand = null;
if (memcachedClient.getProtocol() == Protocol.Text) {
errorCommand =
new MockEncodeTimeoutTextGetOneCommand(
"name", "name".getBytes(), CommandType.GET_ONE, latch, 1000);
} else {
errorCommand =
new MockEncodeTimeoutBinaryGetCommand(
"name", "name".getBytes(), CommandType.GET_ONE, latch, OpCode.GET, false, 1000);
}
memcachedClient.getConnector().send(errorCommand);
// Force write thread to encode command
errorCommand.setIoBuffer(null);
// wait 100 milliseconds,the operation will be timeout
if (!latch.await(100, TimeUnit.MILLISECONDS)) {
errorCommand.cancel();
}
Thread.sleep(1000);
// It is not written to channel,because it is canceled.
assertEquals(
writeMessageCount, memcachedClient.getConnector().getStatistics().getWriteMessageCount());
// It works
assertEquals("dennis", memcachedClient.get("name"));
}
public void visit(Command command) {
if (wasFirst) {
append(command.getKey());
wasFirst = false;
} else {
append(" ");
append(command.getKey());
}
}
/**
* Merge get operation to multi-get operation
*
* @param currentCmd
* @param mergeCommands
* @return
* @throws InterruptedException
*/
@SuppressWarnings("unchecked")
public final Command optimiezeGet(
final Queue writeQueue, final Queue<Command> executingCmds, Command optimiezeCommand) {
if (optimiezeCommand.getCommandType() == CommandType.GET_ONE
|| optimiezeCommand.getCommandType() == CommandType.GETS_ONE) {
// 浼��get���
if (optimiezeGet) {
optimiezeCommand =
mergeGetCommands(
optimiezeCommand, writeQueue, executingCmds, optimiezeCommand.getCommandType());
}
}
return optimiezeCommand;
}
/**
* merge buffers to fit socket's send buffer size
*
* @param currentCommand
* @return
* @throws InterruptedException
*/
@SuppressWarnings("unchecked")
public final Command optimiezeMergeBuffer(
Command optimiezeCommand,
final Queue writeQueue,
final Queue<Command> executingCmds,
int sendBufferSize) {
if (log.isDebugEnabled()) {
log.debug("Optimieze merge buffer:" + optimiezeCommand.toString());
}
if (optimiezeMergeBuffer && optimiezeCommand.getIoBuffer().remaining() < sendBufferSize) {
optimiezeCommand = mergeBuffer(optimiezeCommand, writeQueue, executingCmds, sendBufferSize);
}
return optimiezeCommand;
}
public void _testErrorCommand() throws Exception {
Command nonexisCmd =
new Command() {
@Override
public boolean decode(MemcachedTCPSession session, ByteBuffer buffer) {
return decodeError(ByteUtils.nextLine(buffer));
}
@Override
public void encode() {
ioBuffer = IoBuffer.wrap(ByteBuffer.wrap("test\r\n".getBytes()));
}
};
nonexisCmd.setKey("test");
nonexisCmd.setLatch(new CountDownLatch(1));
memcachedClient.getConnector().send(nonexisCmd);
// this.memcachedClient.flushAll();
nonexisCmd.getLatch().await();
assertNotNull(nonexisCmd.getException());
assertEquals(
"Nonexist command,check your memcached version please.",
nonexisCmd.getException().getMessage());
assertTrue(nonexisCmd.getException() instanceof UnknownCommandException);
memcachedClient.set("name", 0, "dennis");
assertEquals("dennis", memcachedClient.get("name"));
}
public void testAddEncodeAndDecode() {
Command command =
this.commandFactory.createCASCommand(
this.key, this.keyBytes, 0, this.value, 9L, this.noreply, this.transcoder);
command.encode();
ByteBuffer encodeBuffer = command.getIoBuffer().buf();
assertNotNull(encodeBuffer);
assertEquals(42, encodeBuffer.capacity());
byte opCode = encodeBuffer.get(1);
// cas use set command
assertEquals(OpCode.SET.fieldValue(), opCode);
ByteBuffer buffer =
constructResponse(
OpCode.SET.fieldValue(),
(short) 0,
(byte) 0,
(byte) 0,
(short) 0,
0,
0,
10L,
null,
null,
null);
assertTrue(command.decode(null, buffer));
assertTrue((Boolean) command.getResult());
assertEquals(0, buffer.remaining());
buffer =
constructResponse(
OpCode.SET.fieldValue(),
(short) 0,
(byte) 0,
(byte) 0,
(short) 0x0005,
0,
0,
0L,
null,
null,
null);
command =
this.commandFactory.createCASCommand(
this.key, this.keyBytes, 0, this.value, 9L, this.noreply, this.transcoder);
assertTrue(command.decode(null, buffer));
assertFalse((Boolean) command.getResult());
assertEquals(0, buffer.remaining());
}
public void testOperationDecodeTimeOut() throws Exception {
memcachedClient.set("name", 0, "dennis");
assertEquals("dennis", memcachedClient.get("name"));
CountDownLatch latch = new CountDownLatch(1);
Command errorCommand = null;
if (memcachedClient.getProtocol() == Protocol.Text) {
errorCommand =
new MockDecodeTimeoutTextGetOneCommand(
"name", "name".getBytes(), CommandType.GET_ONE, latch, 1000);
} else {
errorCommand =
new MockDecodeTimeoutBinaryGetOneCommand(
"name", "name".getBytes(), CommandType.GET_ONE, latch, OpCode.GET, false, 1000);
}
memcachedClient.getConnector().send(errorCommand);
// wait 100 milliseconds,the operation will be timeout
latch.await(100, TimeUnit.MILLISECONDS);
assertNull(errorCommand.getResult());
Thread.sleep(1000);
// It works.
assertNotNull(errorCommand.getResult());
assertEquals("dennis", memcachedClient.get("name"));
}
public void finish() {
// prev command is the last command,last command must be getk,ensure
// getq commands send response back
Command lastGetKCommand =
new BinaryGetCommand(
prevCommand.getKey(),
prevCommand.getKeyBytes(),
CommandType.GET_ONE,
new CountDownLatch(1),
OpCode.GET_KEY,
false);
lastGetKCommand.encode();
bufferList.add(lastGetKCommand.getIoBuffer());
totalBytes += lastGetKCommand.getIoBuffer().remaining();
}
public void send(final Command msg) throws MemcachedException {
MemcachedSession session = (MemcachedSession) this.findSessionByKey(msg.getKey());
if (session == null) {
throw new MemcachedException("There is no available connection at this moment");
}
// If session was closed,try to use standby memcached node
if (session.isClosed()) {
session = this.findStandbySession(session);
}
if (session.isClosed()) {
throw new MemcachedException(
"Session("
+ SystemUtils.getRawAddress(session.getRemoteSocketAddress())
+ ":"
+ session.getRemoteSocketAddress().getPort()
+ ") has been closed");
}
if (session.isAuthFailed()) {
throw new MemcachedException("Auth failed to connection " + session.getRemoteSocketAddress());
}
session.write(msg);
}
public void visit(Command command) {
// Encode prev command
if (prevCommand != null) {
// first n-1 send getq command
Command getqCommand =
new BinaryGetCommand(
prevCommand.getKey(),
prevCommand.getKeyBytes(),
null,
null,
OpCode.GET_KEY_QUIETLY,
true);
getqCommand.encode();
totalBytes += getqCommand.getIoBuffer().remaining();
bufferList.add(getqCommand.getIoBuffer());
}
prevCommand = command;
}
@SuppressWarnings("unchecked")
private final Command mergeGetCommands(
final Command currentCmd,
final Queue writeQueue,
final Queue<Command> executingCmds,
CommandType expectedCommandType) {
Map<Object, Command> mergeCommands = null;
int mergeCount = 1;
final CommandCollector commandCollector = creatCommandCollector();
currentCmd.setStatus(OperationStatus.WRITING);
commandCollector.visit(currentCmd);
while (mergeCount < mergeFactor) {
Command nextCmd = (Command) writeQueue.peek();
if (nextCmd == null) {
break;
}
if (nextCmd.isCancel()) {
writeQueue.remove();
continue;
}
if (nextCmd.getCommandType() == expectedCommandType) {
if (mergeCommands == null) { // lazy initialize
mergeCommands = new HashMap<Object, Command>(mergeFactor / 2);
mergeCommands.put(currentCmd.getKey(), currentCmd);
}
if (log.isDebugEnabled()) {
log.debug("Merge get command:" + nextCmd.toString());
}
nextCmd.setStatus(OperationStatus.WRITING);
Command removedCommand = (Command) writeQueue.remove();
// If the key is exists,add the command to associated list.
if (mergeCommands.containsKey(removedCommand.getKey())) {
final AssocCommandAware mergedGetCommand =
(AssocCommandAware) mergeCommands.get(removedCommand.getKey());
if (mergedGetCommand.getAssocCommands() == null) {
mergedGetCommand.setAssocCommands(new ArrayList<Command>(5));
}
mergedGetCommand.getAssocCommands().add(removedCommand);
} else {
commandCollector.visit(nextCmd);
mergeCommands.put(removedCommand.getKey(), removedCommand);
}
mergeCount++;
} else {
break;
}
}
commandCollector.finish();
if (mergeCount == 1) {
return currentCmd;
} else {
if (log.isDebugEnabled()) {
log.debug("Merge optimieze:merge " + mergeCount + " get commands");
}
return newMergedCommand(mergeCommands, mergeCount, commandCollector, expectedCommandType);
}
}
@SuppressWarnings("unchecked")
private final Command mergeBuffer(
final Command firstCommand,
final Queue writeQueue,
final Queue<Command> executingCmds,
final int sendBufferSize) {
Command lastCommand = firstCommand; // ��苟������涓�ommand
Command nextCmd = (Command) writeQueue.peek();
if (nextCmd == null) {
return lastCommand;
}
final List<Command> commands = getLocalList();
final ByteBuffer firstBuffer = firstCommand.getIoBuffer().buf();
int totalBytes = firstBuffer.remaining();
commands.add(firstCommand);
boolean wasFirst = true;
while (totalBytes + nextCmd.getIoBuffer().remaining() <= sendBufferSize
&& (nextCmd = (Command) writeQueue.peek()) != null) {
if (nextCmd.getStatus() == OperationStatus.WRITING) {
break;
}
if (nextCmd.isCancel()) {
writeQueue.remove();
continue;
}
nextCmd.setStatus(OperationStatus.WRITING);
writeQueue.remove();
if (wasFirst) {
wasFirst = false;
}
// if it is get_one command,try to merge get commands
if ((nextCmd.getCommandType() == CommandType.GET_ONE
|| nextCmd.getCommandType() == CommandType.GETS_ONE)
&& optimiezeGet) {
nextCmd = mergeGetCommands(nextCmd, writeQueue, executingCmds, nextCmd.getCommandType());
}
commands.add(nextCmd);
lastCommand = nextCmd;
totalBytes += nextCmd.getIoBuffer().remaining();
if (totalBytes > sendBufferSize) {
break;
}
}
if (commands.size() > 1) {
byte[] buf = new byte[totalBytes];
int offset = 0;
for (Command command : commands) {
byte[] ba = command.getIoBuffer().array();
System.arraycopy(ba, 0, buf, offset, ba.length);
offset += ba.length;
if (command != lastCommand
&& (!command.isNoreply() || command instanceof BaseBinaryCommand)) {
executingCmds.add(command);
}
}
lastCommand.setIoBuffer(IoBuffer.wrap(buf));
}
return lastCommand;
}
private void doAuth() {
SaslClient saslClient = null;
try {
saslClient =
Sasl.createSaslClient(
authInfo.getMechanisms(),
null,
"memcached",
memcachedTCPSession.getRemoteSocketAddress().toString(),
null,
this.authInfo.getCallbackHandler());
final AtomicBoolean done = new AtomicBoolean(false);
byte[] response =
saslClient.hasInitialResponse() ? saslClient.evaluateChallenge(EMPTY_BYTES) : EMPTY_BYTES;
CountDownLatch latch = new CountDownLatch(1);
Command command =
this.commandFactory.createAuthStartCommand(
saslClient.getMechanismName(), latch, response);
if (!this.memcachedTCPSession.isClosed()) this.memcachedTCPSession.write(command);
else {
log.error("Authentication fail,because the connection has been closed");
throw new RuntimeException("Authentication fai,connection has been close");
}
while (!done.get()) {
try {
latch.await();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
done.set(true);
}
ResponseStatus responseStatus = ((BaseBinaryCommand) command).getResponseStatus();
switch (responseStatus) {
case NO_ERROR:
done.set(true);
log.info(
"Authentication to "
+ this.memcachedTCPSession.getRemoteSocketAddress()
+ " successfully");
break;
case AUTH_REQUIRED:
log.error(
"Authentication failed to " + this.memcachedTCPSession.getRemoteSocketAddress());
log.warn(
"Reopen connection to "
+ this.memcachedTCPSession.getRemoteSocketAddress()
+ ",beacause auth fail");
this.memcachedTCPSession.setAuthFailed(true);
// It it is not first time,try to sleep 1 second
if (!this.authInfo.isFirstTime()) {
Thread.sleep(1000);
}
this.memcachedTCPSession.close();
done.set(true);
break;
case FUTHER_AUTH_REQUIRED:
String result = (String) command.getResult();
response = saslClient.evaluateChallenge(ByteUtils.getBytes(result));
latch = new CountDownLatch(1);
command =
commandFactory.createAuthStepCommand(
saslClient.getMechanismName(), latch, response);
if (!this.memcachedTCPSession.isClosed()) this.memcachedTCPSession.write(command);
else {
log.error("Authentication fail,because the connection has been closed");
throw new RuntimeException("Authentication fai,connection has been close");
}
break;
default:
done.set(true);
log.error(
"Authentication failed to "
+ this.memcachedTCPSession.getRemoteSocketAddress()
+ ",response status="
+ responseStatus);
break;
}
}
} catch (Exception e) {
log.error("Create saslClient error", e);
} finally {
if (saslClient != null) {
try {
saslClient.dispose();
} catch (SaslException e) {
log.error("Dispose saslClient error", e);
}
}
}
}