com.aerospike.client.cluster.ConnectionRecover Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of aerospike-client-bc-jdk8 Show documentation
Aerospike Java client interface to Aerospike database server. Uses Bouncy Castle crypto library for RIPEMD-160 hashing.
There is a newer version: 9.0.0
Show newest version
/*
 * 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.cluster;

import java.io.EOFException;
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.util.concurrent.TimeUnit;

import com.aerospike.client.AerospikeException;
import com.aerospike.client.command.Buffer;
import com.aerospike.client.command.Command;

public final class ConnectionRecover {
	//private static final AtomicInteger Counter = new AtomicInteger();

	private final Connection conn;
	private final Node node;
	private byte[] headerBuf;
	//private int tranId;
	private int timeoutDelay;
	private int offset;
	private int length;
	private boolean isSingle;
	private boolean checkReturnCode;
	private boolean lastGroup;
	private byte state;

	public ConnectionRecover(Connection conn, Node node, int timeoutDelay, Connection.ReadTimeout crt, boolean isSingle) {
		//tranId = Counter.getAndIncrement();
		//System.out.println("" + tranId + " timeout:" + crt.state + ',' + crt.offset + ',' + crt.length);
		this.conn = conn;
		this.node = node;
		this.timeoutDelay = timeoutDelay;
		this.offset = crt.offset;

		try {
			switch (crt.state) {
			case Command.STATE_READ_AUTH_HEADER:
				this.length = 10;
				this.isSingle = true;
				this.checkReturnCode = true;
				this.state = Command.STATE_READ_HEADER;

				if (offset >= length) {
					if (crt.buffer[length - 1] != 0) {
						// Authentication failed.
						//System.out.println("" + tranId + " invalid user/password:");
						abort();
						return;
					}
					length = getSize(crt.buffer) - (offset - 8);
					offset = 0;
					state = Command.STATE_READ_DETAIL;
				}
				else if (offset > 0) {
					copyHeaderBuf(crt.buffer);
				}
				break;

			case Command.STATE_READ_HEADER:
				// Extend header length to 12 for multi-record responses to include
				// last group info3 bit at offset 11.
				this.length = isSingle ? 8 : 12;
				this.isSingle = isSingle;
				this.checkReturnCode = false;
				this.state = crt.state;

				if (offset >= length) {
					parseProto(crt.buffer, offset);
				}
				else if (offset > 0) {
					copyHeaderBuf(crt.buffer);
				}
				break;

			case Command.STATE_READ_DETAIL:
			default:
				this.length = crt.length;
				this.isSingle = isSingle;
				this.checkReturnCode = false;
				this.state = crt.state;
				break;
			}

			conn.updateLastUsed();
			conn.setTimeout(1);
		}
		catch (Throwable e) {
			//System.out.println("" + tranId + " init failed: " + e.getMessage());
			abort();
		}
	}

	/**
	 * Drain connection.
	 * @return true if draining is complete.
	 */
	public boolean drain(byte[] buf) {
		try {
			if (isSingle) {
				if (state == Command.STATE_READ_HEADER) {
					drainHeader(buf);
				}
				drainDetail(buf);
				recover();
				return true;
			}
			else {
				while (true) {
					if (state == Command.STATE_READ_HEADER) {
						drainHeader(buf);
					}
					drainDetail(buf);

					if (lastGroup) {
						break;
					}
					length = 12;
					offset = 0;
					state = Command.STATE_READ_HEADER;
				}
				recover();
				return true;
			}
		}
		catch (SocketTimeoutException ste) {
			if (System.nanoTime() - conn.getLastUsed() >= TimeUnit.MILLISECONDS.toNanos(timeoutDelay)) {
				// Forcibly close connection.
				//System.out.println("" + tranId + " timeout expired. close connection");
				abort();
				return true;
			}
			// Put back on queue for later draining.
			return false;
		}
		catch (Throwable e) {
			// Forcibly close connection.
			//System.out.println("" + tranId + " socket error:");
			//e.printStackTrace();
			abort();
			return true;
		}
	}

	/**
	 * Has connection been recovered or closed.
	 */
	public boolean isComplete() {
		return state == Command.STATE_COMPLETE;
	}

	/**
	 * Close connection.
	 */
	public void abort() {
		node.closeConnection(conn);
		state = Command.STATE_COMPLETE;
	}

	private void recover() {
		//System.out.println("" + tranId + " connection drained");
		conn.updateLastUsed();
		node.putConnection(conn);
		state = Command.STATE_COMPLETE;
	}

	private void drainHeader(byte[] buf) throws IOException {
		byte[] b = (offset == 0)? buf : headerBuf;

		while (true) {
			int count = conn.read(b, offset, length - offset);

			if (count < 0) {
				// Connection closed by server.
				throw new EOFException();
			}
			offset += count;

			if (offset >= length) {
				break;
			}

			// Partial read.
			if (b == buf) {
				// Convert to header buf.
				copyHeaderBuf(b);
				b = headerBuf;
			}
		}

		if (checkReturnCode) {
			if (b[length - 1] != 0) {
				// Authentication failed.
				//System.out.println("" + tranId + " invalid user/password:");
				abort();
				return;
			}
		}
		parseProto(b, length);
	}

	private void drainDetail(byte[] buf) throws IOException {
		while (offset < length) {
			int rem = length - offset;
			int len = (rem <= buf.length)? rem : buf.length;
			int count = conn.read(buf, 0, len);

			if (count < 0) {
				// Connection closed by server.
				throw new EOFException();
			}
			offset += count;
		}
	}

	private void copyHeaderBuf(byte[] buf) {
		if (headerBuf == null) {
			headerBuf = new byte[length];
		}

		for (int i = 0; i < offset; i++) {
			headerBuf[i] = buf[i];
		}
	}

	private int getSize(byte[] buf) {
		long proto = Buffer.bytesToLong(buf, 0);
		return (int)(proto & 0xFFFFFFFFFFFFL);
	}

	private void parseProto(byte[] buf, int bytesRead) {
		long proto = Buffer.bytesToLong(buf, 0);

		if (! isSingle) {
			// The last group trailer will never be compressed.
			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");
			}

			// 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 = buf[length - 1];

			if ((info3 & Command.INFO3_LAST) == Command.INFO3_LAST) {
				lastGroup = true;
			}
		}
		int size = (int)(proto & 0xFFFFFFFFFFFFL);
		length = size - (bytesRead - 8);
		offset = 0;
		state = Command.STATE_READ_DETAIL;
	}
}