Test connectivity of memcached clusters backed by Mcrouter

McrouterTest.java

/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2021 Chun-Kwong Wong
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 *
 * Test connectivity of memcached clusters backed by Mcrouter
 *
 * (see: https://github.com/facebook/mcrouter)
 *
 * How to use:
 *
 * $ javac ./McrouterTest.java
 * $ java [ -Dsocket.timeout=<NO_OF_MILLIS> ] \
 *	   McrouterTest INPUT [ INPUT [ INPUT ... ] ] ]
 *
 * NOTE:
 *
 * - NO_OF_MILLIS is the number of milliseconds of socket timeout.
 *
 * - INPUT is in the form of FQDN[:PORT] representing a memcached/mcrouter
 *	 node.	(e.g. "mcr.myrack.net:21211", "10.30.1.9", or "localhost:11211")
 */
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.Arrays;
import java.util.List;
import java.util.function.Predicate;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class McrouterTest {
	private static final Predicate<String> MATCHER_HOST
		= Pattern.compile("[^:\\s]+(:\\d+)?").asMatchPredicate();
	private static final Predicate<String> MATCHER_MCROUTER_VERSION
		= Pattern.compile("VERSION\\s+\\S+\\s+mcrouter").asMatchPredicate();
	private static final Predicate<String> MATCHER_CONNECT_ERROR
		= Pattern.compile
		(
			".*\\bconnect_error:[1-9][0-9]*\\b.*"
			, Pattern.DOTALL
		).asMatchPredicate();

	private static final String CMD_VERSION = "version";
	private static final String CMD_STATS_SERVERS = "stats servers";

	private List<InetSocketAddress> hostAddrs;
	private int soTimeout;

	public static void main(String[] args)
	{
		McrouterTest mcrouterTest = new McrouterTest
		(
			Integer.parseInt(System.getProperty("socket.timeout", "200"))
		);

		for(String someHostAddrs: args)
			mcrouterTest.identifyHosts(someHostAddrs);
	}

	public McrouterTest(int someSoTimeout)
	{
		soTimeout = someSoTimeout;
	}

	private void identifyHosts(String someHostAddrs)
	{
		hostAddrs = Arrays.asList(someHostAddrs.split("\\s*,\\s*"))
			.stream()
			.filter(MATCHER_HOST)
			.map
			(hostAddr -> {
				String[] hostAddrParts = hostAddr.split("\\s*:\\s*");
				int port;

				if (1 < hostAddrParts.length && !hostAddrParts[1].isEmpty())
					port = Integer.parseInt(hostAddrParts[1]);
				else
					port = 11211;

				return new InetSocketAddress(hostAddrParts[0], port);
			})
			.collect(Collectors.toList());

		hostAddrs.stream()
			.forEach
			(
				haddr -> System.out.println
				(
					String.format
					(
						"%s -> %s"
						, hostAddrs.toString()
						, checkHost(haddr) ? "OK" : "NG"
					)
				)
			);
	}

	private boolean checkHost(InetSocketAddress socketAddr)
	{
		try
		(
			Socket socket = new Socket
			(
				socketAddr.getAddress()
				, socketAddr.getPort()
			)
		)
		{
			socket.setReuseAddress(true);
			socket.setSoTimeout(soTimeout);

			System.out.println
			(
				String.format
				(
					"Sending VERSION to host %s"
					, socketAddr.toString()
				)
			);

			if (MATCHER_MCROUTER_VERSION.test(request(socket, CMD_VERSION)))
			{
				System.out.println
				(
					String.format
					(
						"Sending STATS SERVERS to host %s"
						, socketAddr.toString()
					)
				);

				return !MATCHER_CONNECT_ERROR.test
				(
					request(socket, CMD_STATS_SERVERS)
				);
			}
			else
			{
				return true;
			}
		}
		catch(Throwable t)
		{
			System.err.println(t.getMessage());
			t.printStackTrace(System.err);
			return false;
		}
	}

	private String request(Socket socket, String req) throws Exception
	{
		StringBuilder sb = new StringBuilder();
		InputStream is = socket.getInputStream();
		OutputStream os = socket.getOutputStream();

		System.out.println(String.format("request: [%s]", req));

		req += "\r\n";
		byte[] buf = req.getBytes();
		os.write(buf);
		os.flush();

		int n;
		do
		{
			n = is.read();
			if (n >= 0) sb.append((char) n);
		} while (is.available() > 0);

		List<String> response = Arrays.asList(sb.toString().split("\\r\\n"))
			.stream()
			.filter(s -> !s.equals("END"))
			.collect(Collectors.toList());

		return response.get(0);
	}
}