com.aerospike.client.async.NioRecover Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2012-2023 Aerospike, Inc.
 *
 * Portions may be licensed to Aerospike, Inc. under one or more contributor
 * license agreements WHICH ARE COMPATIBLE WITH THE APACHE LICENSE, VERSION 2.0.
 *
 * 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.aerospike.client.async;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.util.concurrent.TimeUnit;

import com.aerospike.client.AerospikeException;
import com.aerospike.client.Log;
import com.aerospike.client.async.HashedWheelTimer.HashedWheelTimeout;
import com.aerospike.client.cluster.Node;
import com.aerospike.client.command.Command;
import com.aerospike.client.util.Util;

public final class NioRecover implements INioCommand, TimerTask {
	//private static final AtomicInteger Counter = new AtomicInteger();

	private final NioEventLoop eventLoop;
	private final EventState eventState;
	private final Node node;
	private final NioConnection conn;
	private final HashedWheelTimeout timeoutTask;
	private final ByteBuffer byteBuffer;
	//private final int tranId;
	private int offset;
	private int length;
	private int state;
	private final boolean isSingle;
	private final boolean checkReturnCode;
	private boolean isLastGroup;

	public NioRecover(NioCommand cmd) {
		AsyncCommand a = cmd.command;
		this.eventLoop = cmd.eventLoop;
		this.eventState = cmd.eventState;
		this.node = cmd.node;
		this.conn = cmd.conn;
		this.byteBuffer = cmd.byteBuffer;

		switch (cmd.state) {
		case AsyncCommand.AUTH_READ_HEADER:
			this.offset = byteBuffer.position();
			this.length = byteBuffer.limit();
			this.state = AsyncCommand.COMMAND_READ_HEADER;
			this.isSingle = true;
			this.checkReturnCode = true;
			break;

		case AsyncCommand.AUTH_READ_BODY:
			this.offset = byteBuffer.position();
			this.length = byteBuffer.limit();
			this.state = AsyncCommand.COMMAND_READ_BODY;
			this.isSingle = true;
			this.checkReturnCode = true;
			break;

		case AsyncCommand.COMMAND_READ_BODY:
			if (a.isSingle) {
				this.offset = byteBuffer.position();
				this.length = byteBuffer.limit();
			}
			else {
				// Multi-record detail handles offsets differently.
				this.offset = a.dataOffset;
				this.length = a.receiveSize;

				if (a.dataOffset >= 4) {
					// Warning: The following code assumes multi-record responses always end with a separate proto
					// that only contains one header with the info3 last group bit.  This is always true for batch
					// and scan, but query does not conform.  Therefore, connection recovery for queries will
					// likely fail.
					byte info3 = a.dataBuffer[3];

					if ((info3 & Command.INFO3_LAST) != 0) {
						isLastGroup = true;
					}
				}
			}
			this.state = cmd.state;
			this.isSingle = a.isSingle;
			this.checkReturnCode = false;
			break;

		default:
			this.offset = byteBuffer.position();
			this.length = byteBuffer.limit();
			this.state = cmd.state;
			this.isSingle = a.isSingle;
			this.checkReturnCode = false;
			break;
		}

		//tranId = Counter.getAndIncrement();
		//System.out.println("" + tranId + " timeout:" + a.isSingle + ',' + cmd.state + ',' + offset + ',' + length);

		conn.attach(this);
		timeoutTask = new HashedWheelTimeout(this);
		eventLoop.timer.addTimeout(timeoutTask, System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(a.policy.timeoutDelay));
	}

	@Override
	public final void timeout() {
		//System.out.println("" + tranId + " timeout expired. close connection");

		// Transaction has been delayed long enough.
		// User has already been notified.
		// timeoutTask has already been removed, so avoid cancel.
		abort(false);
	}

	@Override
	public final void processEvent(SelectionKey key) {
		try {
			int ops = key.readyOps();

			if ((ops & SelectionKey.OP_READ) != 0) {
				switch (state) {
				case AsyncCommand.COMMAND_READ_HEADER:
					if (isSingle) {
						drainSingleHeader();
					}
					else {
						if (drainMultiHeader()) {
							drainMultiBody();
						}
					}
					break;

				case AsyncCommand.COMMAND_READ_BODY:
					if (isSingle) {
						drainSingleBody();
					}
					else {
						drainMultiBody();
					}
					break;
				}
			}
		}
		catch (Throwable e) {
			//System.out.println("" + tranId + " socket error:");
			//e.printStackTrace();
			abort(true);
		}
	}

	private final void drainSingleHeader() throws IOException {
		if (! conn.read(byteBuffer)) {
			return;
		}

		byteBuffer.position(0);
		length = ((int) (byteBuffer.getLong() & 0xFFFFFFFFFFFFL));
		byteBuffer.clear();

		if (length < byteBuffer.capacity()) {
			byteBuffer.limit(length);
		}

		offset = 0;
		state = AsyncCommand.COMMAND_READ_BODY;
		drainSingleBody();
	}

	private final void drainSingleBody() throws IOException {
		while (true) {
			if (! conn.read(byteBuffer)) {
				return;
			}

			if (checkReturnCode) {
				int resultCode = byteBuffer.get(1) & 0xFF;

				if (resultCode != 0) {
					// Authentication failed.
					//System.out.println("" + tranId + " invalid user/password:");
					abort(false);
					return;
				}
			}

			if (resetBuffer()) {
				continue;
			}

			recover();
			return;
		}
	}

	private final boolean drainMultiHeader() throws IOException {
		while (true) {
			if (! conn.read(byteBuffer)) {
				return false;
			}

			byteBuffer.position(0);
			long proto = byteBuffer.getLong();
			length = ((int)(proto & 0xFFFFFFFFFFFFL));

			if (length <= 0) {
				// Received zero length block. Read next header.
				byteBuffer.clear();
				byteBuffer.limit(8);
				continue;
			}

			boolean compressed = ((proto >> 48) & 0xff) == Command.MSG_TYPE_COMPRESSED;

			if (compressed) {
				// Do not recover connections with compressed data because that would
				// require saving large buffers with associated state and performing decompression
				// just to drain the connection.
				throw new AerospikeException("Recovering connections with compressed multi-record data is not supported");
			}
			break;
		}

		byteBuffer.clear();

		if (length < byteBuffer.capacity()) {
			byteBuffer.limit(length);
		}

		offset = 0;
		state = AsyncCommand.COMMAND_READ_BODY;
		return true;
	}

	private final void drainMultiBody() throws IOException {
		while (true) {
			if (! conn.read(byteBuffer)) {
				return;
			}

			// Only check for last group on the header of the
			// first record in the group.
			if (offset < 4) {
				byte info3 = byteBuffer.get(3);

				if ((info3 & Command.INFO3_LAST) != 0) {
					isLastGroup = true;
				}
			}

			if (resetBuffer()) {
				continue;
			}

			if (isLastGroup) {
				recover();
				return;
			}

			byteBuffer.clear();
			byteBuffer.limit(8);
			state = AsyncCommand.COMMAND_READ_HEADER;

			if (! drainMultiHeader()) {
				return;
			}
		}
	}

	private final boolean resetBuffer() {
		offset += byteBuffer.limit();

		if (offset < length) {
			byteBuffer.clear();

			int remaining = length - offset;

			if (remaining < byteBuffer.capacity()) {
				byteBuffer.limit(remaining);
			}
			return true;
		}
		return false;
	}

	private final void recover() {
		//System.out.println("" + tranId + " connection drained");
		if (state == AsyncCommand.COMPLETE) {
			return;
		}
		state = AsyncCommand.COMPLETE;

		try {
			conn.unregister();
			conn.updateLastUsed();
			node.putAsyncConnection(conn, eventLoop.index);
			close(true);
		}
		catch (Throwable e) {
			if (! eventState.closed) {
				Log.error("NioRecover recover failed: " + Util.getStackTrace(e));
			}
		}
	}

	private final void abort(boolean cancelTimeout) {
		if (state == AsyncCommand.COMPLETE) {
			return;
		}
		state = AsyncCommand.COMPLETE;

		try {
			node.closeAsyncConnection(conn, eventLoop.index);
			close(cancelTimeout);
		}
		catch (Throwable e) {
			if (! eventState.closed) {
				Log.error("NioRecover abort failed: " + Util.getStackTrace(e));
			}
		}
	}

	private final void close(boolean cancelTimeout) {
		if (cancelTimeout) {
			timeoutTask.cancel();
		}

		if (byteBuffer != null) {
			eventLoop.putByteBuffer(byteBuffer);
		}
	}
}