storm-client/src/jvm/org/apache/storm/pacemaker/PacemakerClient.java - storm - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.  The ASF licenses this file to you 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 org.apache.storm.pacemaker;

 import java.net.InetSocketAddress;
 import java.util.Map;
 import java.util.Timer;
 import java.util.TimerTask;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.security.auth.login.Configuration;
 import org.apache.storm.Config;
 import org.apache.storm.generated.HBMessage;
 import org.apache.storm.messaging.netty.ISaslClient;
 import org.apache.storm.messaging.netty.NettyRenameThreadFactory;
 import org.apache.storm.pacemaker.codec.ThriftNettyClientCodec;
 import org.apache.storm.security.auth.ClientAuthUtils;
 import org.apache.storm.shade.io.netty.bootstrap.Bootstrap;
 import org.apache.storm.shade.io.netty.buffer.PooledByteBufAllocator;
 import org.apache.storm.shade.io.netty.channel.Channel;
 import org.apache.storm.shade.io.netty.channel.ChannelOption;
 import org.apache.storm.shade.io.netty.channel.EventLoopGroup;
 import org.apache.storm.shade.io.netty.channel.WriteBufferWaterMark;
 import org.apache.storm.shade.io.netty.channel.nio.NioEventLoopGroup;
 import org.apache.storm.shade.io.netty.channel.socket.nio.NioSocketChannel;
 import org.apache.storm.utils.StormBoundedExponentialBackoffRetry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public class PacemakerClient implements ISaslClient {

     private static final Logger LOG = LoggerFactory.getLogger(PacemakerClient.class);
     private static Timer timer = new Timer(true);
     private final Bootstrap bootstrap;
     private final EventLoopGroup workerEventLoopGroup;
     private String clientName;
     private String secret;
     private AtomicBoolean ready;
     private AtomicBoolean shutdown;
     private AtomicReference<Channel> channelRef;
     private InetSocketAddress remoteAddr;
     private int maxPending = 100;
     private HBMessage[] messages;
     private LinkedBlockingQueue<Integer> availableMessageSlots;
     private ThriftNettyClientCodec.AuthMethod authMethod;
     private static final int maxRetries = 10;
     private String host;

     private StormBoundedExponentialBackoffRetry backoff = new StormBoundedExponentialBackoffRetry(100, 5000, 20);
     private int retryTimes = 0;

     public PacemakerClient(Map<String, Object> config, String host) {
         this.host = host;
         clientName = (String) config.get(Config.TOPOLOGY_NAME);
         if (clientName == null) {
             clientName = "pacemaker-client";
         }

         String auth = (String) config.get(Config.PACEMAKER_AUTH_METHOD);

         switch (auth) {

             case "DIGEST":
                 Configuration loginConf = ClientAuthUtils.getConfiguration(config);
                 authMethod = ThriftNettyClientCodec.AuthMethod.DIGEST;
                 secret = ClientAuthUtils.makeDigestPayload(loginConf, ClientAuthUtils.LOGIN_CONTEXT_PACEMAKER_DIGEST);
                 if (secret == null) {
                     LOG.error("Can't start pacemaker server without digest secret.");
                     throw new RuntimeException("Can't start pacemaker server without digest secret.");
                 }
                 break;

             case "KERBEROS":
                 authMethod = ThriftNettyClientCodec.AuthMethod.KERBEROS;
                 break;

             case "NONE":
                 authMethod = ThriftNettyClientCodec.AuthMethod.NONE;
                 break;

             default:
                 authMethod = ThriftNettyClientCodec.AuthMethod.NONE;
                 LOG.warn("Invalid auth scheme: '{}'. Falling back to 'NONE'", auth);
                 break;
         }

         ready = new AtomicBoolean(false);
         shutdown = new AtomicBoolean(false);
         channelRef = new AtomicReference<>(null);
         setupMessaging();

         ThreadFactory workerFactory = new NettyRenameThreadFactory(this.host + "-pm");
         // 0 means DEFAULT_EVENT_LOOP_THREADS
         // https://github.com/netty/netty/blob/netty-4.1.24.Final/transport/src/main/java/io/netty/channel/MultithreadEventLoopGroup.java#L40
         int maxWorkers = (int)config.get(Config.PACEMAKER_CLIENT_MAX_THREADS);
         this.workerEventLoopGroup = new NioEventLoopGroup(maxWorkers > 0 ? maxWorkers : 0, workerFactory);
         int thriftMessageMaxSize = (Integer) config.get(Config.PACEMAKER_THRIFT_MESSAGE_SIZE_MAX);
         bootstrap = new Bootstrap()
             .group(workerEventLoopGroup)
             .channel(NioSocketChannel.class)
             .option(ChannelOption.TCP_NODELAY, true)
             .option(ChannelOption.SO_SNDBUF, 5242880)
             .option(ChannelOption.SO_KEEPALIVE, true)
             .option(ChannelOption.WRITE_BUFFER_WATER_MARK, new WriteBufferWaterMark(8 * 1024, 32 * 1024))
             .option(ChannelOption.ALLOCATOR, PooledByteBufAllocator.DEFAULT)
             .handler(new ThriftNettyClientCodec(this, config, authMethod, host, thriftMessageMaxSize));

         int port = (int) config.get(Config.PACEMAKER_PORT);
         remoteAddr = new InetSocketAddress(host, port);
         bootstrap.connect(remoteAddr);
     }

     private void setupMessaging() {
         messages = new HBMessage[maxPending];
         availableMessageSlots = new LinkedBlockingQueue<Integer>();
         for (int i = 0; i < maxPending; i++) {
             availableMessageSlots.add(i);
         }
     }

     @Override
     public synchronized void channelReady(Channel channel) {
         Channel oldChannel = channelRef.get();
         if (oldChannel != null) {
             LOG.debug("Closing oldChannel is connected: {}", oldChannel.toString());
             close_channel();
         }

         channelRef.set(channel);
         retryTimes = 0;
         LOG.debug("Channel is ready: {}", channel.toString());
         ready.set(true);
         this.notifyAll();
     }

     @Override
     public String name() {
         return clientName;
     }

     @Override
     public String secretKey() {
         return secret;
     }

     public HBMessage send(HBMessage m) throws PacemakerConnectionException, InterruptedException {
         LOG.debug("Sending pacemaker message to {}: {}", host, m);

         int next = availableMessageSlots.take();
         synchronized (m) {
             m.set_message_id(next);
             messages[next] = m;
             LOG.debug("Put message in slot: {} for {}", Integer.toString(next), host);
             int retry = maxRetries;
             while (true) {
                 try {
                     waitUntilReady();
                     Channel channel = channelRef.get();
                     if (channel != null) {
                         channel.writeAndFlush(m, channel.voidPromise());
                         m.wait(1000);
                     }
                     if (messages[next] != m && messages[next] != null) {
                         // messages[next] == null can happen if we lost the connection and subsequently reconnected or timed out.
                         HBMessage ret = messages[next];
                         messages[next] = null;
                         LOG.debug("Got Response: {}", ret);
                         return ret;
                     }
                 } catch (PacemakerConnectionException e) {
                     if (retry <= 0) {
                         throw e;
                     }
                     LOG.error("error attempting to write to a channel {}.", e.getMessage());
                 }
                 if (retry <= 0) {
                     throw new PacemakerConnectionException("couldn't get response after " + maxRetries + " attempts.");
                 }
                 retry--;
                 LOG.warn("Not getting response or getting null response. Making {} more attempts for {}.", retry, host);
             }
         }
     }

     private void waitUntilReady() throws PacemakerConnectionException, InterruptedException {
         // Wait for 'ready' (channel connected and maybe authentication)
         if (!ready.get() || channelRef.get() == null) {
             synchronized (this) {
                 if (!ready.get()) {
                     LOG.debug("Waiting for netty channel to be ready.");
                     this.wait(1000);
                     if (!ready.get() || channelRef.get() == null) {
                         throw new PacemakerConnectionException("Timed out waiting for channel ready.");
                     }
                 }
             }
         }
     }

     public void gotMessage(HBMessage m) {
         int messageId = m.get_message_id();
         if (messageId >= 0 && messageId < maxPending) {

             LOG.debug("Pacemaker client got message: {}", m.toString());
             HBMessage request = messages[messageId];

             if (request == null) {
                 LOG.debug("No message for slot: {}", Integer.toString(messageId));
             } else {
                 synchronized (request) {
                     messages[messageId] = m;
                     request.notifyAll();
                     availableMessageSlots.add(messageId);
                 }
             }
         } else {
             LOG.error("Got Message with bad id: {}", m.toString());
         }
     }

     public void reconnect() {
         final PacemakerClient client = this;
         timer.schedule(new TimerTask() {
             @Override
             public void run() {
                 client.doReconnect();
             }
         }, backoff.getSleepTimeMs(retryTimes++, 0));
         ready.set(false);
         setupMessaging();
     }

     public synchronized void doReconnect() {
         LOG.info("reconnecting to {}", host);
         close_channel();
         if (!shutdown.get()) {
             bootstrap.connect(remoteAddr);
         }
     }

     public void shutdown() {
         shutdown.set(true);
         workerEventLoopGroup.shutdownGracefully().awaitUninterruptibly();
     }

     private synchronized void close_channel() {
         if (channelRef.get() != null) {
             channelRef.get().close();
             LOG.debug("channel {} closed", remoteAddr);
             channelRef.set(null);
         }
     }

     public void close() {
         close_channel();
     }
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership. The ASF licenses this file to you 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 org.apache.storm.pacemaker;

	import java.net.InetSocketAddress;
	import java.util.Map;
	import java.util.Timer;
	import java.util.TimerTask;
	import java.util.concurrent.LinkedBlockingQueue;
	import java.util.concurrent.ThreadFactory;
	import java.util.concurrent.atomic.AtomicBoolean;
	import java.util.concurrent.atomic.AtomicReference;
	import javax.security.auth.login.Configuration;
	import org.apache.storm.Config;
	import org.apache.storm.generated.HBMessage;
	import org.apache.storm.messaging.netty.ISaslClient;
	import org.apache.storm.messaging.netty.NettyRenameThreadFactory;
	import org.apache.storm.pacemaker.codec.ThriftNettyClientCodec;
	import org.apache.storm.security.auth.ClientAuthUtils;
	import org.apache.storm.shade.io.netty.bootstrap.Bootstrap;
	import org.apache.storm.shade.io.netty.buffer.PooledByteBufAllocator;
	import org.apache.storm.shade.io.netty.channel.Channel;
	import org.apache.storm.shade.io.netty.channel.ChannelOption;
	import org.apache.storm.shade.io.netty.channel.EventLoopGroup;
	import org.apache.storm.shade.io.netty.channel.WriteBufferWaterMark;
	import org.apache.storm.shade.io.netty.channel.nio.NioEventLoopGroup;
	import org.apache.storm.shade.io.netty.channel.socket.nio.NioSocketChannel;
	import org.apache.storm.utils.StormBoundedExponentialBackoffRetry;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public class PacemakerClient implements ISaslClient {

	private static final Logger LOG = LoggerFactory.getLogger(PacemakerClient.class);
	private static Timer timer = new Timer(true);
	private final Bootstrap bootstrap;
	private final EventLoopGroup workerEventLoopGroup;
	private String clientName;
	private String secret;
	private AtomicBoolean ready;
	private AtomicBoolean shutdown;
	private AtomicReference<Channel> channelRef;
	private InetSocketAddress remoteAddr;
	private int maxPending = 100;
	private HBMessage[] messages;
	private LinkedBlockingQueue<Integer> availableMessageSlots;
	private ThriftNettyClientCodec.AuthMethod authMethod;
	private static final int maxRetries = 10;
	private String host;

	private StormBoundedExponentialBackoffRetry backoff = new StormBoundedExponentialBackoffRetry(100, 5000, 20);
	private int retryTimes = 0;

	public PacemakerClient(Map<String, Object> config, String host) {
	this.host = host;
	clientName = (String) config.get(Config.TOPOLOGY_NAME);
	if (clientName == null) {
	clientName = "pacemaker-client";
	}

	String auth = (String) config.get(Config.PACEMAKER_AUTH_METHOD);

	switch (auth) {

	case "DIGEST":
	Configuration loginConf = ClientAuthUtils.getConfiguration(config);
	authMethod = ThriftNettyClientCodec.AuthMethod.DIGEST;
	secret = ClientAuthUtils.makeDigestPayload(loginConf, ClientAuthUtils.LOGIN_CONTEXT_PACEMAKER_DIGEST);
	if (secret == null) {
	LOG.error("Can't start pacemaker server without digest secret.");
	throw new RuntimeException("Can't start pacemaker server without digest secret.");
	}
	break;

	case "KERBEROS":
	authMethod = ThriftNettyClientCodec.AuthMethod.KERBEROS;
	break;

	case "NONE":
	authMethod = ThriftNettyClientCodec.AuthMethod.NONE;
	break;

	default:
	authMethod = ThriftNettyClientCodec.AuthMethod.NONE;
	LOG.warn("Invalid auth scheme: '{}'. Falling back to 'NONE'", auth);
	break;
	}

	ready = new AtomicBoolean(false);
	shutdown = new AtomicBoolean(false);
	channelRef = new AtomicReference<>(null);
	setupMessaging();

	ThreadFactory workerFactory = new NettyRenameThreadFactory(this.host + "-pm");
	// 0 means DEFAULT_EVENT_LOOP_THREADS
	// https://github.com/netty/netty/blob/netty-4.1.24.Final/transport/src/main/java/io/netty/channel/MultithreadEventLoopGroup.java#L40
	int maxWorkers = (int)config.get(Config.PACEMAKER_CLIENT_MAX_THREADS);
	this.workerEventLoopGroup = new NioEventLoopGroup(maxWorkers > 0 ? maxWorkers : 0, workerFactory);
	int thriftMessageMaxSize = (Integer) config.get(Config.PACEMAKER_THRIFT_MESSAGE_SIZE_MAX);
	bootstrap = new Bootstrap()
	.group(workerEventLoopGroup)
	.channel(NioSocketChannel.class)
	.option(ChannelOption.TCP_NODELAY, true)
	.option(ChannelOption.SO_SNDBUF, 5242880)
	.option(ChannelOption.SO_KEEPALIVE, true)
	.option(ChannelOption.WRITE_BUFFER_WATER_MARK, new WriteBufferWaterMark(8 * 1024, 32 * 1024))
	.option(ChannelOption.ALLOCATOR, PooledByteBufAllocator.DEFAULT)
	.handler(new ThriftNettyClientCodec(this, config, authMethod, host, thriftMessageMaxSize));

	int port = (int) config.get(Config.PACEMAKER_PORT);
	remoteAddr = new InetSocketAddress(host, port);
	bootstrap.connect(remoteAddr);
	}

	private void setupMessaging() {
	messages = new HBMessage[maxPending];
	availableMessageSlots = new LinkedBlockingQueue<Integer>();
	for (int i = 0; i < maxPending; i++) {
	availableMessageSlots.add(i);
	}
	}

	@Override
	public synchronized void channelReady(Channel channel) {
	Channel oldChannel = channelRef.get();
	if (oldChannel != null) {
	LOG.debug("Closing oldChannel is connected: {}", oldChannel.toString());
	close_channel();
	}

	channelRef.set(channel);
	retryTimes = 0;
	LOG.debug("Channel is ready: {}", channel.toString());
	ready.set(true);
	this.notifyAll();
	}

	@Override
	public String name() {
	return clientName;
	}

	@Override
	public String secretKey() {
	return secret;
	}

	public HBMessage send(HBMessage m) throws PacemakerConnectionException, InterruptedException {
	LOG.debug("Sending pacemaker message to {}: {}", host, m);

	int next = availableMessageSlots.take();
	synchronized (m) {
	m.set_message_id(next);
	messages[next] = m;
	LOG.debug("Put message in slot: {} for {}", Integer.toString(next), host);
	int retry = maxRetries;
	while (true) {
	try {
	waitUntilReady();
	Channel channel = channelRef.get();
	if (channel != null) {
	channel.writeAndFlush(m, channel.voidPromise());
	m.wait(1000);
	}
	if (messages[next] != m && messages[next] != null) {
	// messages[next] == null can happen if we lost the connection and subsequently reconnected or timed out.
	HBMessage ret = messages[next];
	messages[next] = null;
	LOG.debug("Got Response: {}", ret);
	return ret;
	}
	} catch (PacemakerConnectionException e) {
	if (retry <= 0) {
	throw e;
	}
	LOG.error("error attempting to write to a channel {}.", e.getMessage());
	}
	if (retry <= 0) {
	throw new PacemakerConnectionException("couldn't get response after " + maxRetries + " attempts.");
	}
	retry--;
	LOG.warn("Not getting response or getting null response. Making {} more attempts for {}.", retry, host);
	}
	}
	}

	private void waitUntilReady() throws PacemakerConnectionException, InterruptedException {
	// Wait for 'ready' (channel connected and maybe authentication)
	if (!ready.get() \|\| channelRef.get() == null) {
	synchronized (this) {
	if (!ready.get()) {
	LOG.debug("Waiting for netty channel to be ready.");
	this.wait(1000);
	if (!ready.get() \|\| channelRef.get() == null) {
	throw new PacemakerConnectionException("Timed out waiting for channel ready.");
	}
	}
	}
	}
	}

	public void gotMessage(HBMessage m) {
	int messageId = m.get_message_id();
	if (messageId >= 0 && messageId < maxPending) {

	LOG.debug("Pacemaker client got message: {}", m.toString());
	HBMessage request = messages[messageId];

	if (request == null) {
	LOG.debug("No message for slot: {}", Integer.toString(messageId));
	} else {
	synchronized (request) {
	messages[messageId] = m;
	request.notifyAll();
	availableMessageSlots.add(messageId);
	}
	}
	} else {
	LOG.error("Got Message with bad id: {}", m.toString());
	}
	}

	public void reconnect() {
	final PacemakerClient client = this;
	timer.schedule(new TimerTask() {
	@Override
	public void run() {
	client.doReconnect();
	}
	}, backoff.getSleepTimeMs(retryTimes++, 0));
	ready.set(false);
	setupMessaging();
	}

	public synchronized void doReconnect() {
	LOG.info("reconnecting to {}", host);
	close_channel();
	if (!shutdown.get()) {
	bootstrap.connect(remoteAddr);
	}
	}

	public void shutdown() {
	shutdown.set(true);
	workerEventLoopGroup.shutdownGracefully().awaitUninterruptibly();
	}

	private synchronized void close_channel() {
	if (channelRef.get() != null) {
	channelRef.get().close();
	LOG.debug("channel {} closed", remoteAddr);
	channelRef.set(null);
	}
	}

	public void close() {
	close_channel();
	}
	}