twill-core/src/main/java/org/apache/twill/internal/AbstractZKServiceController.java - twill - 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.twill.internal;

 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import org.apache.twill.api.Command;
 import org.apache.twill.api.RunId;
 import org.apache.twill.api.ServiceController;
 import org.apache.twill.common.Threads;
 import org.apache.twill.internal.state.Message;
 import org.apache.twill.internal.state.Messages;
 import org.apache.twill.internal.state.SystemMessages;
 import org.apache.twill.zookeeper.NodeData;
 import org.apache.twill.zookeeper.ZKClient;
 import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.WatchedEvent;
 import org.apache.zookeeper.Watcher;
 import org.apache.zookeeper.data.Stat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;

 /**
  * An abstract base class for implementing a {@link ServiceController} using ZooKeeper as a means for
  * communicating with the remote service. This is designed to work in pair with the {@link AbstractTwillService}.
  */
 public abstract class AbstractZKServiceController extends AbstractExecutionServiceController {

   private static final Logger LOG = LoggerFactory.getLogger(AbstractZKServiceController.class);

   protected final ZKClient zkClient;
   private final InstanceNodeDataCallback instanceNodeDataCallback;
   private final List<ListenableFuture<?>> messageFutures;
   private ListenableFuture<State> stopMessageFuture;

   protected AbstractZKServiceController(RunId runId, ZKClient zkClient) {
     super(runId);
     this.zkClient = zkClient;
     this.instanceNodeDataCallback = new InstanceNodeDataCallback();
     this.messageFutures = Lists.newLinkedList();
   }

   @Override
   public final ListenableFuture<Command> sendCommand(Command command) {
     return sendMessage(Messages.createForAll(command), command);
   }

   @Override
   public final ListenableFuture<Command> sendCommand(String runnableName, Command command) {
     return sendMessage(Messages.createForRunnable(runnableName, command), command);
   }

   @Override
   protected final void startUp() {
     doStartUp();

     // Watch for instance node existence.
     actOnExists(getInstancePath(), new Runnable() {
       @Override
       public void run() {
         watchInstanceNode();
       }
     });
   }

   @Override
   protected final synchronized void shutDown() {
     if (stopMessageFuture == null) {
       stopMessageFuture = ZKMessages.sendMessage(zkClient, getMessagePrefix(),
                                                  SystemMessages.stopApplication(), State.TERMINATED);
     }

     // Cancel all pending message futures.
     for (ListenableFuture<?> future : messageFutures) {
       future.cancel(true);
     }

     doShutDown();
   }

   /**
    * Sends a {@link Message} to the remote service. Returns a future that will be completed when the message
    * has been processed.
    * @param message The message to send.
    * @param result Object to set into the future when message is being processed.
    * @param <V> Type of the result.
    * @return A {@link ListenableFuture} that will be completed when the message has been processed.
    */
   protected final synchronized <V> ListenableFuture<V> sendMessage(Message message, V result) {
     if (!isRunning()) {
       return Futures.immediateFailedFuture(new IllegalStateException("Cannot send message to non-running application"));
     }
     final ListenableFuture<V> messageFuture = ZKMessages.sendMessage(zkClient, getMessagePrefix(), message, result);
     messageFutures.add(messageFuture);
     messageFuture.addListener(new Runnable() {
       @Override
       public void run() {
         // If the completion is triggered when stopping, do nothing.
         if (state() == State.STOPPING) {
           return;
         }
         synchronized (AbstractZKServiceController.this) {
           messageFutures.remove(messageFuture);
         }
       }
     }, Threads.SAME_THREAD_EXECUTOR);

     return messageFuture;
   }

   protected final ListenableFuture<State> getStopMessageFuture() {
     return stopMessageFuture;
   }

   /**
    * Called during startup. Executed in the startup thread.
    */
   protected abstract void doStartUp();

   /**
    * Called during shutdown. Executed in the shutdown thread.
    */
   protected abstract void doShutDown();

   /**
    * Called when an update on the live instance node is detected.
    * @param nodeData The updated live instance node data or {@code null} if there is an error when fetching
    *                 the node data.
    */
   protected abstract void instanceNodeUpdated(NodeData nodeData);

   /**
    * Called when failed to fetch from live instance node.
    * @param cause The cause of failure of {@code null} if cause is unknown.
    */
   protected abstract void instanceNodeFailed(Throwable cause);

   protected synchronized void forceShutDown() {
     if (stopMessageFuture == null) {
       // In force shutdown, don't send message.
       stopMessageFuture = Futures.immediateFuture(State.TERMINATED);
     }
     stop();
   }


   private void actOnExists(final String path, final Runnable action) {
     // Watch for node existence.
     final AtomicBoolean nodeExists = new AtomicBoolean(false);
     Futures.addCallback(zkClient.exists(path, new Watcher() {
       @Override
       public void process(WatchedEvent event) {
         if (!shouldProcessZKEvent()) {
           return;
         }
         // When node is created, call the action.
         // Other event type would be handled by the action.
         if (event.getType() == Event.EventType.NodeCreated && nodeExists.compareAndSet(false, true)) {
           action.run();
         }
       }
     }), new FutureCallback<Stat>() {
       @Override
       public void onSuccess(Stat result) {
         if (result != null && nodeExists.compareAndSet(false, true)) {
           action.run();
         }
       }

       @Override
       public void onFailure(Throwable t) {
         LOG.error("Failed in exists call to {}. Shutting down service.", path, t);
         forceShutDown();
       }
     }, Threads.SAME_THREAD_EXECUTOR);
   }

   protected final void watchInstanceNode() {
     if (!shouldProcessZKEvent()) {
       return;
     }
     Futures.addCallback(zkClient.getData(getInstancePath(), new Watcher() {
       @Override
       public void process(WatchedEvent event) {
         if (!shouldProcessZKEvent()) {
           return;
         }
         switch (event.getType()) {
           case NodeDataChanged:
             watchInstanceNode();
             break;
           case NodeDeleted:
             instanceNodeFailed(KeeperException.create(KeeperException.Code.NONODE, getInstancePath()));
             break;
           default:
             LOG.info("Ignore ZK event for instance node: {}", event);
         }
       }
     }), instanceNodeDataCallback, Threads.SAME_THREAD_EXECUTOR);
   }

   /**
    * Returns true if ZK events needs to be processed.
    */
   private boolean shouldProcessZKEvent() {
     State state = state();
     return (state == State.NEW || state == State.STARTING || state == State.RUNNING);
   }

   /**
    * Returns the path prefix for creating sequential message node for the remote service.
    */
   private String getMessagePrefix() {
     return getZKPath("messages/msg");
   }

   /**
    * Returns the zookeeper node path for the ephemeral instance node for this runId.
    */
   protected final String getInstancePath() {
     return String.format("/instances/%s", getRunId().getId());
   }

   private String getZKPath(String path) {
     return String.format("/%s/%s", getRunId().getId(), path);
   }

   private final class InstanceNodeDataCallback implements FutureCallback<NodeData> {

     @Override
     public void onSuccess(NodeData result) {
       if (shouldProcessZKEvent()) {
         instanceNodeUpdated(result);
       }
     }

     @Override
     public void onFailure(Throwable t) {
       LOG.error("Failed in fetching instance node data.", t);
       if (shouldProcessZKEvent()) {
         instanceNodeFailed(t);
       }
     }
   }
 }
	/*
	* 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.twill.internal;

	import com.google.common.collect.Lists;
	import com.google.common.util.concurrent.FutureCallback;
	import com.google.common.util.concurrent.Futures;
	import com.google.common.util.concurrent.ListenableFuture;
	import org.apache.twill.api.Command;
	import org.apache.twill.api.RunId;
	import org.apache.twill.api.ServiceController;
	import org.apache.twill.common.Threads;
	import org.apache.twill.internal.state.Message;
	import org.apache.twill.internal.state.Messages;
	import org.apache.twill.internal.state.SystemMessages;
	import org.apache.twill.zookeeper.NodeData;
	import org.apache.twill.zookeeper.ZKClient;
	import org.apache.zookeeper.KeeperException;
	import org.apache.zookeeper.WatchedEvent;
	import org.apache.zookeeper.Watcher;
	import org.apache.zookeeper.data.Stat;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.util.List;
	import java.util.concurrent.atomic.AtomicBoolean;

	/**
	* An abstract base class for implementing a {@link ServiceController} using ZooKeeper as a means for
	* communicating with the remote service. This is designed to work in pair with the {@link AbstractTwillService}.
	*/
	public abstract class AbstractZKServiceController extends AbstractExecutionServiceController {

	private static final Logger LOG = LoggerFactory.getLogger(AbstractZKServiceController.class);

	protected final ZKClient zkClient;
	private final InstanceNodeDataCallback instanceNodeDataCallback;
	private final List<ListenableFuture<?>> messageFutures;
	private ListenableFuture<State> stopMessageFuture;

	protected AbstractZKServiceController(RunId runId, ZKClient zkClient) {
	super(runId);
	this.zkClient = zkClient;
	this.instanceNodeDataCallback = new InstanceNodeDataCallback();
	this.messageFutures = Lists.newLinkedList();
	}

	@Override
	public final ListenableFuture<Command> sendCommand(Command command) {
	return sendMessage(Messages.createForAll(command), command);
	}

	@Override
	public final ListenableFuture<Command> sendCommand(String runnableName, Command command) {
	return sendMessage(Messages.createForRunnable(runnableName, command), command);
	}

	@Override
	protected final void startUp() {
	doStartUp();

	// Watch for instance node existence.
	actOnExists(getInstancePath(), new Runnable() {
	@Override
	public void run() {
	watchInstanceNode();
	}
	});
	}

	@Override
	protected final synchronized void shutDown() {
	if (stopMessageFuture == null) {
	stopMessageFuture = ZKMessages.sendMessage(zkClient, getMessagePrefix(),
	SystemMessages.stopApplication(), State.TERMINATED);
	}

	// Cancel all pending message futures.
	for (ListenableFuture<?> future : messageFutures) {
	future.cancel(true);
	}

	doShutDown();
	}

	/**
	* Sends a {@link Message} to the remote service. Returns a future that will be completed when the message
	* has been processed.
	* @param message The message to send.
	* @param result Object to set into the future when message is being processed.
	* @param <V> Type of the result.
	* @return A {@link ListenableFuture} that will be completed when the message has been processed.
	*/
	protected final synchronized <V> ListenableFuture<V> sendMessage(Message message, V result) {
	if (!isRunning()) {
	return Futures.immediateFailedFuture(new IllegalStateException("Cannot send message to non-running application"));
	}
	final ListenableFuture<V> messageFuture = ZKMessages.sendMessage(zkClient, getMessagePrefix(), message, result);
	messageFutures.add(messageFuture);
	messageFuture.addListener(new Runnable() {
	@Override
	public void run() {
	// If the completion is triggered when stopping, do nothing.
	if (state() == State.STOPPING) {
	return;
	}
	synchronized (AbstractZKServiceController.this) {
	messageFutures.remove(messageFuture);
	}
	}
	}, Threads.SAME_THREAD_EXECUTOR);

	return messageFuture;
	}

	protected final ListenableFuture<State> getStopMessageFuture() {
	return stopMessageFuture;
	}

	/**
	* Called during startup. Executed in the startup thread.
	*/
	protected abstract void doStartUp();

	/**
	* Called during shutdown. Executed in the shutdown thread.
	*/
	protected abstract void doShutDown();

	/**
	* Called when an update on the live instance node is detected.
	* @param nodeData The updated live instance node data or {@code null} if there is an error when fetching
	* the node data.
	*/
	protected abstract void instanceNodeUpdated(NodeData nodeData);

	/**
	* Called when failed to fetch from live instance node.
	* @param cause The cause of failure of {@code null} if cause is unknown.
	*/
	protected abstract void instanceNodeFailed(Throwable cause);

	protected synchronized void forceShutDown() {
	if (stopMessageFuture == null) {
	// In force shutdown, don't send message.
	stopMessageFuture = Futures.immediateFuture(State.TERMINATED);
	}
	stop();
	}


	private void actOnExists(final String path, final Runnable action) {
	// Watch for node existence.
	final AtomicBoolean nodeExists = new AtomicBoolean(false);
	Futures.addCallback(zkClient.exists(path, new Watcher() {
	@Override
	public void process(WatchedEvent event) {
	if (!shouldProcessZKEvent()) {
	return;
	}
	// When node is created, call the action.
	// Other event type would be handled by the action.
	if (event.getType() == Event.EventType.NodeCreated && nodeExists.compareAndSet(false, true)) {
	action.run();
	}
	}
	}), new FutureCallback<Stat>() {
	@Override
	public void onSuccess(Stat result) {
	if (result != null && nodeExists.compareAndSet(false, true)) {
	action.run();
	}
	}

	@Override
	public void onFailure(Throwable t) {
	LOG.error("Failed in exists call to {}. Shutting down service.", path, t);
	forceShutDown();
	}
	}, Threads.SAME_THREAD_EXECUTOR);
	}

	protected final void watchInstanceNode() {
	if (!shouldProcessZKEvent()) {
	return;
	}
	Futures.addCallback(zkClient.getData(getInstancePath(), new Watcher() {
	@Override
	public void process(WatchedEvent event) {
	if (!shouldProcessZKEvent()) {
	return;
	}
	switch (event.getType()) {
	case NodeDataChanged:
	watchInstanceNode();
	break;
	case NodeDeleted:
	instanceNodeFailed(KeeperException.create(KeeperException.Code.NONODE, getInstancePath()));
	break;
	default:
	LOG.info("Ignore ZK event for instance node: {}", event);
	}
	}
	}), instanceNodeDataCallback, Threads.SAME_THREAD_EXECUTOR);
	}

	/**
	* Returns true if ZK events needs to be processed.
	*/
	private boolean shouldProcessZKEvent() {
	State state = state();
	return (state == State.NEW \|\| state == State.STARTING \|\| state == State.RUNNING);
	}

	/**
	* Returns the path prefix for creating sequential message node for the remote service.
	*/
	private String getMessagePrefix() {
	return getZKPath("messages/msg");
	}

	/**
	* Returns the zookeeper node path for the ephemeral instance node for this runId.
	*/
	protected final String getInstancePath() {
	return String.format("/instances/%s", getRunId().getId());
	}

	private String getZKPath(String path) {
	return String.format("/%s/%s", getRunId().getId(), path);
	}

	private final class InstanceNodeDataCallback implements FutureCallback<NodeData> {

	@Override
	public void onSuccess(NodeData result) {
	if (shouldProcessZKEvent()) {
	instanceNodeUpdated(result);
	}
	}

	@Override
	public void onFailure(Throwable t) {
	LOG.error("Failed in fetching instance node data.", t);
	if (shouldProcessZKEvent()) {
	instanceNodeFailed(t);
	}
	}
	}
	}