/*
* Copyright © 2014 Cask Data, Inc.
*
* 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 co.cask.tigon.internal.app.runtime.flow;
import co.cask.tigon.api.flow.FlowSpecification;
import co.cask.tigon.api.flow.FlowletDefinition;
import co.cask.tigon.api.flow.flowlet.FlowletSpecification;
import co.cask.tigon.app.program.Program;
import co.cask.tigon.app.program.ProgramType;
import co.cask.tigon.data.queue.QueueName;
import co.cask.tigon.data.transaction.queue.QueueAdmin;
import co.cask.tigon.internal.app.runtime.AbstractProgramController;
import co.cask.tigon.internal.app.runtime.Arguments;
import co.cask.tigon.internal.app.runtime.BasicArguments;
import co.cask.tigon.internal.app.runtime.ProgramController;
import co.cask.tigon.internal.app.runtime.ProgramOptionConstants;
import co.cask.tigon.internal.app.runtime.ProgramOptions;
import co.cask.tigon.internal.app.runtime.ProgramRunner;
import co.cask.tigon.internal.app.runtime.ProgramRunnerFactory;
import co.cask.tigon.internal.app.runtime.SimpleProgramOptions;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.HashBasedTable;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Table;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.inject.Inject;
import org.apache.twill.api.RunId;
import org.apache.twill.discovery.DiscoveryServiceClient;
import org.apache.twill.discovery.ServiceDiscovered;
import org.apache.twill.internal.RunIds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.List;
import java.util.Map;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
/**
*
*/
public final class FlowProgramRunner implements ProgramRunner {
private static final Logger LOG = LoggerFactory.getLogger(FlowProgramRunner.class);
private final ProgramRunnerFactory programRunnerFactory;
private final Map<RunId, ProgramOptions> programOptions = Maps.newHashMap();
private final QueueAdmin queueAdmin;
private final DiscoveryServiceClient discoveryServiceClient;
@Inject
public FlowProgramRunner(ProgramRunnerFactory programRunnerFactory, QueueAdmin queueAdmin,
DiscoveryServiceClient discoveryServiceClient) {
this.programRunnerFactory = programRunnerFactory;
this.queueAdmin = queueAdmin;
this.discoveryServiceClient = discoveryServiceClient;
}
@Override
public ProgramController run(Program program, ProgramOptions options) {
// Extract and verify parameters
FlowSpecification flowSpec = program.getSpecification();
ProgramType processorType = program.getType();
Preconditions.checkNotNull(processorType, "Missing processor type.");
Preconditions.checkArgument(processorType == ProgramType.FLOW, "Only FLOW process type is supported.");
Preconditions.checkNotNull(flowSpec, "Missing FlowSpecification for %s", program.getName());
for (FlowletDefinition flowletDefinition : flowSpec.getFlowlets().values()) {
int maxInstances = flowletDefinition.getFlowletSpec().getMaxInstances();
Preconditions.checkArgument(flowletDefinition.getInstances() <= maxInstances,
"Flowlet %s can have a maximum of %s instances",
flowletDefinition.getFlowletSpec().getName(), maxInstances);
}
try {
// Launch flowlet program runners
RunId runId = RunIds.generate();
programOptions.put(runId, options);
Multimap<String, QueueName> consumerQueues = FlowUtils.configureQueue(program, flowSpec, queueAdmin);
final Table<String, Integer, ProgramController> flowlets = createFlowlets(program, runId, flowSpec);
return new FlowProgramController(flowlets, runId, program, flowSpec, consumerQueues, discoveryServiceClient);
} catch (Exception e) {
throw Throwables.propagate(e);
}
}
/**
* Starts all flowlets in the flow program.
* @param program Program to run
* @param flowSpec The {@link FlowSpecification}.
* @return A {@link com.google.common.collect.Table} with row as flowlet id, column as instance id,
* cell as the {@link ProgramController} for the flowlet.
*/
private Table<String, Integer, ProgramController> createFlowlets(Program program, RunId runId,
FlowSpecification flowSpec) {
Table<String, Integer, ProgramController> flowlets = HashBasedTable.create();
try {
for (Map.Entry<String, FlowletDefinition> entry : flowSpec.getFlowlets().entrySet()) {
int instanceCount = entry.getValue().getInstances();
for (int instanceId = 0; instanceId < instanceCount; instanceId++) {
flowlets.put(entry.getKey(), instanceId,
startFlowlet(program, createFlowletOptions(entry.getKey(), instanceId, instanceCount, runId)));
}
}
} catch (Throwable t) {
try {
// Need to stop all started flowlets
Futures.successfulAsList(Iterables.transform(flowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.stop();
}
})).get();
} catch (Exception e) {
LOG.error("Fail to stop all flowlets on failure.");
}
throw Throwables.propagate(t);
}
return flowlets;
}
private ProgramController startFlowlet(Program program, ProgramOptions options) {
return programRunnerFactory.create(ProgramRunnerFactory.Type.FLOWLET)
.run(program, options);
}
private ProgramOptions createFlowletOptions(String name, int instanceId, int instances, RunId runId) {
// Get the right user arguments.
Arguments userArguments = new BasicArguments();
if (programOptions.containsKey(runId)) {
userArguments = programOptions.get(runId).getUserArguments();
}
return new SimpleProgramOptions(name, new BasicArguments(
ImmutableMap.of(
ProgramOptionConstants.INSTANCE_ID, Integer.toString(instanceId),
ProgramOptionConstants.INSTANCES, Integer.toString(instances),
ProgramOptionConstants.RUN_ID, runId.getId()
)), userArguments
);
}
private final class FlowProgramController extends AbstractProgramController {
private final Table<String, Integer, ProgramController> flowlets;
private final Program program;
private final FlowSpecification flowSpec;
private final Lock lock = new ReentrantLock();
private final Multimap<String, QueueName> consumerQueues;
private final DiscoveryServiceClient discoveryServiceClient;
FlowProgramController(Table<String, Integer, ProgramController> flowlets, RunId runId,
Program program, FlowSpecification flowSpec, Multimap<String, QueueName> consumerQueues,
DiscoveryServiceClient discoveryServiceClient) {
super(program.getName(), runId);
this.flowlets = flowlets;
this.program = program;
this.flowSpec = flowSpec;
this.consumerQueues = consumerQueues;
this.discoveryServiceClient = discoveryServiceClient;
started();
}
@Override
protected void doSuspend() throws Exception {
LOG.info("Suspending flow: " + flowSpec.getName());
lock.lock();
try {
Futures.successfulAsList(
Iterables.transform(flowlets.values(),
new Function<ProgramController, ListenableFuture<ProgramController>>() {
@Override
public ListenableFuture<ProgramController> apply(ProgramController input) {
return input.suspend();
}
})).get();
} finally {
lock.unlock();
}
LOG.info("Flow suspended: " + flowSpec.getName());
}
@Override
protected void doResume() throws Exception {
LOG.info("Resuming flow: " + flowSpec.getName());
lock.lock();
try {
Futures.successfulAsList(
Iterables.transform(flowlets.values(),
new Function<ProgramController, ListenableFuture<ProgramController>>() {
@Override
public ListenableFuture<ProgramController> apply(ProgramController input) {
return input.resume();
}
})).get();
} finally {
lock.unlock();
}
LOG.info("Flow resumed: " + flowSpec.getName());
}
@Override
protected void doStop() throws Exception {
LOG.info("Stopping flow: " + flowSpec.getName());
lock.lock();
try {
Futures.successfulAsList(
Iterables.transform(flowlets.values(),
new Function<ProgramController, ListenableFuture<ProgramController>>() {
@Override
public ListenableFuture<ProgramController> apply(ProgramController input) {
return input.stop();
}
})).get();
} finally {
lock.unlock();
}
LOG.info("Flow stopped: " + flowSpec.getName());
}
@Override
@SuppressWarnings("unchecked")
protected void doCommand(String name, Object value) throws Exception {
if (!ProgramOptionConstants.FLOWLET_INSTANCES.equals(name) || !(value instanceof Map)) {
return;
}
Map<String, String> command = (Map<String, String>) value;
lock.lock();
try {
changeInstances(command.get("flowlet"), Integer.valueOf(command.get("newInstances")));
} catch (Throwable t) {
LOG.error(String.format("Fail to change instances: %s", command), t);
} finally {
lock.unlock();
}
}
/**
* Change the number of instances of the running flowlet. Notice that this method needs to be
* synchronized as change of instances involves multiple steps that need to be completed all at once.
* @param flowletName Name of the flowlet
* @param newInstanceCount New instance count
* @throws java.util.concurrent.ExecutionException
* @throws InterruptedException
*/
private synchronized void changeInstances(String flowletName, final int newInstanceCount) throws Exception {
Map<Integer, ProgramController> liveFlowlets = flowlets.row(flowletName);
int liveCount = liveFlowlets.size();
if (liveCount == newInstanceCount) {
return;
}
if (liveCount < newInstanceCount) {
increaseInstances(flowletName, newInstanceCount, liveFlowlets, liveCount);
return;
}
decreaseInstances(flowletName, newInstanceCount, liveFlowlets, liveCount);
}
private synchronized void increaseInstances(String flowletName, final int newInstanceCount,
Map<Integer, ProgramController> liveFlowlets,
int liveCount) throws Exception {
FlowletProgramController flowletProgramController =
(FlowletProgramController) Iterables.getFirst(liveFlowlets.values(), null);
FlowletSpecification flowletSpecification = flowletProgramController.getFlowletContext().getSpecification();
int flowletMaxInstances = flowletSpecification.getMaxInstances();
Preconditions.checkArgument(newInstanceCount <= flowletMaxInstances,
"Flowlet %s can have a maximum of %s instances",
flowletSpecification.getName(), flowletMaxInstances);
// First pause all flowlets
Futures.successfulAsList(Iterables.transform(
liveFlowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.suspend();
}
})).get();
// Then reconfigure stream/queue consumers
FlowUtils.reconfigure(consumerQueues.get(flowletName),
FlowUtils.generateConsumerGroupId(program, flowletName), newInstanceCount, queueAdmin);
// Then change instance count of current flowlets
Futures.successfulAsList(Iterables.transform(
liveFlowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.command(ProgramOptionConstants.INSTANCES, newInstanceCount);
}
})).get();
// Next resume all current flowlets
Futures.successfulAsList(Iterables.transform(
liveFlowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.resume();
}
})).get();
// Last create more instances
for (int instanceId = liveCount; instanceId < newInstanceCount; instanceId++) {
flowlets.put(flowletName, instanceId,
startFlowlet(program,
createFlowletOptions(flowletName, instanceId, newInstanceCount, getRunId())));
}
}
private synchronized void decreaseInstances(String flowletName, final int newInstanceCount,
Map<Integer, ProgramController> liveFlowlets,
int liveCount) throws Exception {
// Shrink number of flowlets
// First stop the extra flowlets
List<ListenableFuture<?>> futures = Lists.newArrayListWithCapacity(liveCount - newInstanceCount);
for (int instanceId = liveCount - 1; instanceId >= newInstanceCount; instanceId--) {
futures.add(flowlets.remove(flowletName, instanceId).stop());
}
Futures.successfulAsList(futures).get();
// Then pause all flowlets
Futures.successfulAsList(Iterables.transform(
liveFlowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.suspend();
}
})).get();
// Then reconfigure stream/queue consumers
FlowUtils.reconfigure(consumerQueues.get(flowletName),
FlowUtils.generateConsumerGroupId(program, flowletName), newInstanceCount, queueAdmin);
// Next updates instance count for each flowlets
Futures.successfulAsList(Iterables.transform(
liveFlowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.command(ProgramOptionConstants.INSTANCES, newInstanceCount);
}
})).get();
// Last resume all remaing flowlets
Futures.successfulAsList(Iterables.transform(
liveFlowlets.values(),
new Function<ProgramController, ListenableFuture<?>>() {
@Override
public ListenableFuture<?> apply(ProgramController controller) {
return controller.resume();
}
})).get();
}
@Override
public ServiceDiscovered discover(String service) {
return discoveryServiceClient.discover(service);
}
}
}