/*
* 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.base.Charsets;
import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.AbstractService;
import com.google.common.util.concurrent.AsyncFunction;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.Service;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import org.apache.twill.api.RunId;
import org.apache.twill.api.ServiceController;
import org.apache.twill.common.ServiceListenerAdapter;
import org.apache.twill.common.Threads;
import org.apache.twill.internal.json.StackTraceElementCodec;
import org.apache.twill.internal.json.StateNodeCodec;
import org.apache.twill.internal.state.Message;
import org.apache.twill.internal.state.MessageCallback;
import org.apache.twill.internal.state.MessageCodec;
import org.apache.twill.internal.state.StateNode;
import org.apache.twill.internal.state.SystemMessages;
import org.apache.twill.zookeeper.NodeChildren;
import org.apache.twill.zookeeper.NodeData;
import org.apache.twill.zookeeper.OperationFuture;
import org.apache.twill.zookeeper.ZKClient;
import org.apache.twill.zookeeper.ZKOperations;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import javax.annotation.Nullable;
/**
* A {@link Service} decorator that wrap another {@link Service} with the service states reflected
* to ZooKeeper.
*/
public final class ZKServiceDecorator extends AbstractService {
private static final Logger LOG = LoggerFactory.getLogger(ZKServiceDecorator.class);
private final ZKClient zkClient;
private final RunId id;
private final Supplier<? extends JsonElement> liveNodeData;
private final Service decoratedService;
private final MessageCallbackCaller messageCallback;
private ExecutorService callbackExecutor;
public ZKServiceDecorator(ZKClient zkClient, RunId id, Supplier<? extends JsonElement> liveNodeData,
Service decoratedService) {
this(zkClient, id, liveNodeData, decoratedService, null);
}
/**
* Creates a ZKServiceDecorator.
* @param zkClient ZooKeeper client
* @param id The run id of the service
* @param liveNodeData A supplier for providing information writing to live node.
* @param decoratedService The Service for monitoring state changes
* @param finalizer An optional Runnable to run when this decorator terminated.
*/
public ZKServiceDecorator(ZKClient zkClient, RunId id, Supplier <? extends JsonElement> liveNodeData,
Service decoratedService, @Nullable Runnable finalizer) {
this.zkClient = zkClient;
this.id = id;
this.liveNodeData = liveNodeData;
this.decoratedService = decoratedService;
if (decoratedService instanceof MessageCallback) {
this.messageCallback = new MessageCallbackCaller((MessageCallback) decoratedService, zkClient);
} else {
this.messageCallback = new MessageCallbackCaller(zkClient);
}
if (finalizer != null) {
addFinalizer(finalizer);
}
}
@Override
protected void doStart() {
callbackExecutor = Executors.newSingleThreadExecutor(Threads.createDaemonThreadFactory("message-callback"));
// Create the live node, if succeeded, start the decorated service, otherwise fail out.
Futures.addCallback(createLiveNode(), new FutureCallback<String>() {
@Override
public void onSuccess(String result) {
// Create nodes for states and messaging
StateNode stateNode = new StateNode(ServiceController.State.STARTING);
final ListenableFuture<List<String>> createFuture = Futures.allAsList(
ZKOperations.ignoreError(zkClient.create(getZKPath("messages"), null, CreateMode.PERSISTENT),
KeeperException.NodeExistsException.class, null),
zkClient.create(getZKPath("state"), encodeStateNode(stateNode), CreateMode.PERSISTENT)
);
createFuture.addListener(new Runnable() {
@Override
public void run() {
try {
createFuture.get();
// Starts the decorated service
decoratedService.addListener(createListener(), Threads.SAME_THREAD_EXECUTOR);
decoratedService.start();
} catch (Exception e) {
notifyFailed(e);
}
}
}, Threads.SAME_THREAD_EXECUTOR);
}
@Override
public void onFailure(Throwable t) {
notifyFailed(t);
}
});
// Watch for session expiration, recreate the live node if reconnected after expiration.
zkClient.addConnectionWatcher(new Watcher() {
private boolean expired = false;
@Override
public void process(WatchedEvent event) {
if (event.getState() == Event.KeeperState.Expired) {
LOG.warn("ZK Session expired for service {} with runId {}.", decoratedService, id.getId());
expired = true;
} else if (event.getState() == Event.KeeperState.SyncConnected && expired) {
LOG.info("Reconnected after expiration for service {} with runId {}", decoratedService, id.getId());
expired = false;
Futures.addCallback(createLiveNode(), new FutureCallback<String>() {
@Override
public void onSuccess(String result) {
// All good, no-op
}
@Override
public void onFailure(Throwable t) {
notifyFailed(t);
}
}, Threads.SAME_THREAD_EXECUTOR);
}
}
});
}
@Override
protected void doStop() {
// Stops the decorated service
decoratedService.stop();
callbackExecutor.shutdownNow();
}
private void addFinalizer(final Runnable finalizer) {
addListener(new ServiceListenerAdapter() {
@Override
public void terminated(State from) {
try {
finalizer.run();
} catch (Throwable t) {
LOG.warn("Exception when running finalizer.", t);
}
}
@Override
public void failed(State from, Throwable failure) {
try {
finalizer.run();
} catch (Throwable t) {
LOG.warn("Exception when running finalizer.", t);
}
}
}, Threads.SAME_THREAD_EXECUTOR);
}
private OperationFuture<String> createLiveNode() {
String liveNode = getLiveNodePath();
LOG.info("Create live node {}{}", zkClient.getConnectString(), liveNode);
JsonObject content = new JsonObject();
content.add("data", liveNodeData.get());
return ZKOperations.ignoreError(zkClient.create(liveNode, encodeJson(content), CreateMode.EPHEMERAL),
KeeperException.NodeExistsException.class, liveNode);
}
private OperationFuture<String> removeLiveNode() {
String liveNode = getLiveNodePath();
LOG.info("Remove live node {}{}", zkClient.getConnectString(), liveNode);
return ZKOperations.ignoreError(zkClient.delete(liveNode), KeeperException.NoNodeException.class, liveNode);
}
private OperationFuture<String> removeServiceNode() {
String serviceNode = String.format("/%s", id.getId());
LOG.info("Remove service node {}{}", zkClient.getConnectString(), serviceNode);
return ZKOperations.recursiveDelete(zkClient, serviceNode);
}
private void watchMessages() {
final String messagesPath = getZKPath("messages");
Futures.addCallback(zkClient.getChildren(messagesPath, new Watcher() {
@Override
public void process(WatchedEvent event) {
// TODO: Do we need to deal with other type of events?
if (event.getType() == Event.EventType.NodeChildrenChanged && decoratedService.isRunning()) {
watchMessages();
}
}
}), new FutureCallback<NodeChildren>() {
@Override
public void onSuccess(NodeChildren result) {
// Sort by the name, which is the messageId. Assumption is that message ids is ordered by time.
List<String> messages = Lists.newArrayList(result.getChildren());
Collections.sort(messages);
for (String messageId : messages) {
processMessage(messagesPath + "/" + messageId, messageId);
}
}
@Override
public void onFailure(Throwable t) {
// TODO: what could be done besides just logging?
LOG.error("Failed to watch messages.", t);
}
});
}
private void processMessage(final String path, final String messageId) {
Futures.addCallback(zkClient.getData(path), new FutureCallback<NodeData>() {
@Override
public void onSuccess(NodeData result) {
Message message = MessageCodec.decode(result.getData());
if (message == null) {
LOG.error("Failed to decode message for " + messageId + " in " + path);
listenFailure(zkClient.delete(path, result.getStat().getVersion()));
return;
}
if (LOG.isDebugEnabled()) {
LOG.debug("Message received from " + path + ": " + new String(MessageCodec.encode(message), Charsets.UTF_8));
}
if (handleStopMessage(message, getDeleteSupplier(path, result.getStat().getVersion()))) {
return;
}
messageCallback.onReceived(callbackExecutor, path, result.getStat().getVersion(), messageId, message);
}
@Override
public void onFailure(Throwable t) {
LOG.error("Failed to fetch message content.", t);
}
});
}
private <V> boolean handleStopMessage(Message message, final Supplier<OperationFuture<V>> postHandleSupplier) {
if (message.getType() == Message.Type.SYSTEM && SystemMessages.STOP_COMMAND.equals(message.getCommand())) {
callbackExecutor.execute(new Runnable() {
@Override
public void run() {
decoratedService.stop().addListener(new Runnable() {
@Override
public void run() {
stopServiceOnComplete(postHandleSupplier.get(), ZKServiceDecorator.this);
}
}, MoreExecutors.sameThreadExecutor());
}
});
return true;
}
return false;
}
private Supplier<OperationFuture<String>> getDeleteSupplier(final String path, final int version) {
return new Supplier<OperationFuture<String>>() {
@Override
public OperationFuture<String> get() {
return zkClient.delete(path, version);
}
};
}
private Listener createListener() {
return new DecoratedServiceListener();
}
private <V> byte[] encode(V data, Class<? extends V> clz) {
return new GsonBuilder().registerTypeAdapter(StateNode.class, new StateNodeCodec())
.registerTypeAdapter(StackTraceElement.class, new StackTraceElementCodec())
.create()
.toJson(data, clz).getBytes(Charsets.UTF_8);
}
private byte[] encodeStateNode(StateNode stateNode) {
return encode(stateNode, StateNode.class);
}
private <V extends JsonElement> byte[] encodeJson(V json) {
return new Gson().toJson(json).getBytes(Charsets.UTF_8);
}
private String getZKPath(String path) {
return String.format("/%s/%s", id, path);
}
private String getLiveNodePath() {
return "/instances/" + id;
}
private static <V> OperationFuture<V> listenFailure(final OperationFuture<V> operationFuture) {
operationFuture.addListener(new Runnable() {
@Override
public void run() {
try {
if (!operationFuture.isCancelled()) {
operationFuture.get();
}
} catch (Exception e) {
// TODO: what could be done besides just logging?
LOG.error("Operation execution failed for " + operationFuture.getRequestPath(), e);
}
}
}, Threads.SAME_THREAD_EXECUTOR);
return operationFuture;
}
private static final class MessageCallbackCaller {
private final MessageCallback callback;
private final ZKClient zkClient;
private MessageCallbackCaller(ZKClient zkClient) {
this(null, zkClient);
}
private MessageCallbackCaller(MessageCallback callback, ZKClient zkClient) {
this.callback = callback;
this.zkClient = zkClient;
}
public void onReceived(Executor executor, final String path,
final int version, final String id, final Message message) {
if (callback == null) {
// Simply delete the message
if (LOG.isDebugEnabled()) {
LOG.debug("Ignoring incoming message from " + path + ": " + message);
}
listenFailure(zkClient.delete(path, version));
return;
}
executor.execute(new Runnable() {
@Override
public void run() {
try {
// Message process is synchronous for now. Making it async needs more thoughts about race conditions.
// The executor is the callbackExecutor which is a single thread executor.
callback.onReceived(id, message).get();
} catch (Throwable t) {
LOG.error("Exception when processing message: {}, {}, {}", id, message, path, t);
} finally {
listenFailure(zkClient.delete(path, version));
}
}
});
}
}
private final class DecoratedServiceListener implements Listener {
private volatile boolean zkFailure = false;
@Override
public void starting() {
LOG.info("Starting: " + id);
saveState(ServiceController.State.STARTING);
}
@Override
public void running() {
LOG.info("Running: " + id);
notifyStarted();
watchMessages();
saveState(ServiceController.State.RUNNING);
}
@Override
public void stopping(State from) {
LOG.info("Stopping: " + id);
saveState(ServiceController.State.STOPPING);
}
@Override
public void terminated(State from) {
LOG.info("Terminated: " + from + " " + id);
if (zkFailure) {
return;
}
ImmutableList<OperationFuture<String>> futures = ImmutableList.of(removeLiveNode(), removeServiceNode());
final ListenableFuture<List<String>> future = Futures.allAsList(futures);
Futures.successfulAsList(futures).addListener(new Runnable() {
@Override
public void run() {
try {
future.get();
LOG.info("Service and state node removed");
notifyStopped();
} catch (Exception e) {
LOG.warn("Failed to remove ZK nodes.", e);
notifyFailed(e);
}
}
}, Threads.SAME_THREAD_EXECUTOR);
}
@Override
public void failed(State from, final Throwable failure) {
LOG.info("Failed: {} {}.", from, id, failure);
if (zkFailure) {
return;
}
ImmutableList<OperationFuture<String>> futures = ImmutableList.of(removeLiveNode(), removeServiceNode());
Futures.successfulAsList(futures).addListener(new Runnable() {
@Override
public void run() {
LOG.info("Service and state node removed");
notifyFailed(failure);
}
}, Threads.SAME_THREAD_EXECUTOR);
}
private void saveState(ServiceController.State state) {
if (zkFailure) {
return;
}
StateNode stateNode = new StateNode(state);
stopOnFailure(zkClient.setData(getZKPath("state"), encodeStateNode(stateNode)));
}
private <V> void stopOnFailure(final OperationFuture<V> future) {
future.addListener(new Runnable() {
@Override
public void run() {
try {
future.get();
} catch (final Exception e) {
LOG.error("ZK operation failed", e);
zkFailure = true;
decoratedService.stop().addListener(new Runnable() {
@Override
public void run() {
notifyFailed(e);
}
}, Threads.SAME_THREAD_EXECUTOR);
}
}
}, Threads.SAME_THREAD_EXECUTOR);
}
}
private <V> ListenableFuture<State> stopServiceOnComplete(ListenableFuture <V> future, final Service service) {
return Futures.transform(future, new AsyncFunction<V, State>() {
@Override
public ListenableFuture<State> apply(V input) throws Exception {
return service.stop();
}
});
}
}