001/**
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *      http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018
019package org.apache.oozie.util.graph;
020
021import com.google.common.collect.ArrayListMultimap;
022import com.google.common.collect.Multimap;
023import guru.nidi.graphviz.attribute.Color;
024import guru.nidi.graphviz.attribute.RankDir;
025import guru.nidi.graphviz.attribute.Shape;
026import guru.nidi.graphviz.engine.Engine;
027import guru.nidi.graphviz.engine.Format;
028import guru.nidi.graphviz.engine.Graphviz;
029import guru.nidi.graphviz.engine.Rasterizer;
030import guru.nidi.graphviz.model.Factory;
031import guru.nidi.graphviz.model.Graph;
032import guru.nidi.graphviz.model.Node;
033import org.apache.oozie.client.WorkflowAction;
034import org.apache.oozie.service.ConfigurationService;
035
036import java.awt.image.BufferedImage;
037import java.util.LinkedHashMap;
038import java.util.Map;
039import java.util.concurrent.Callable;
040import java.util.concurrent.ExecutionException;
041import java.util.concurrent.ExecutorService;
042import java.util.concurrent.Executors;
043import java.util.concurrent.Future;
044import java.util.concurrent.TimeUnit;
045import java.util.concurrent.TimeoutException;
046
047public class GraphvizRenderer implements GraphRenderer {
048
049    /**
050     * We need this single-thread executor because we have to make sure:
051     * <ul>
052     * <li>all GraphViz rendering operations happen on the same thread.
053     * This is because of {@link com.eclipsesource.v8.V8} thread handling</li>
054     * <li>GraphViz rendering operations don't timeout</li>
055     * <li>GraphViz rendering operations don't overlap</li>
056     * </ul>
057     */
058    private static final ExecutorService EXECUTOR_SERVICE = Executors.newSingleThreadExecutor();
059    private static final long GRAPHVIZ_TIMEOUT_SECONDS = ConfigurationService.getLong("oozie.graphviz.timeout.seconds");
060
061    private Graph graphvizGraph = Factory.graph().generalAttr().with(RankDir.TOP_TO_BOTTOM).directed();
062    private final Map<String, Node> graphvizNodes = new LinkedHashMap<>();
063    private final Multimap<String, String> edges = ArrayListMultimap.create();
064    ;
065    private int arcCount = 0;
066
067    @Override
068    public void addNode(final WorkflowActionNode node) {
069        final Shape shape = getShape(node.getType());
070        final Color color = getColor(node.getStatus());
071
072        final Node graphvizNode = Factory.node(node.getName()).with(shape).with(color);
073
074        graphvizNodes.put(node.getName(), graphvizNode);
075    }
076
077    private Shape getShape(final String type) {
078        final Shape shape;
079
080        switch (type) {
081            case "start":
082                shape = Shape.CIRCLE;
083                break;
084            case "end":
085                shape = Shape.DOUBLE_CIRCLE;
086                break;
087            case "kill":
088                shape = Shape.OCTAGON;
089                break;
090            case "decision":
091                shape = Shape.DIAMOND;
092                break;
093            case "fork":
094                shape = Shape.TRIANGLE;
095                break;
096            case "join":
097                shape = Shape.INV_TRIANGLE;
098                break;
099            default:
100                shape = Shape.RECTANGLE;
101                break;
102        }
103
104        return shape;
105    }
106
107    private Color getColor(final WorkflowAction.Status status) {
108        if (status == null) {
109            return Color.BLACK;
110        }
111
112        final Color color;
113
114        switch (status) {
115            case PREP:
116            case USER_RETRY:
117            case START_RETRY:
118            case START_MANUAL:
119                color = Color.GREY;
120                break;
121            case RUNNING:
122            case END_RETRY:
123            case END_MANUAL:
124                color = Color.YELLOW;
125                break;
126            case OK:
127            case DONE:
128                color = Color.GREEN;
129                break;
130            case ERROR:
131            case FAILED:
132            case KILLED:
133                color = Color.RED;
134                break;
135            default:
136                color = Color.BLACK;
137        }
138
139        return color;
140    }
141
142    private Node createOrGetGraphvizNode(final WorkflowActionNode node) {
143        if (graphvizNodes.containsKey(node.getName())) {
144            return graphvizNodes.get(node.getName());
145        }
146
147        addNode(node);
148
149        return graphvizNodes.get(node.getName());
150    }
151
152    @Override
153    public void addEdge(final WorkflowActionNode parent, final WorkflowActionNode child) {
154        if (edges.containsEntry(parent.getName(), child.getName())) {
155            return;
156        }
157
158        Node graphvizParent = createOrGetGraphvizNode(parent);
159
160        graphvizParent = graphvizParent.link(
161                Factory.to(createOrGetGraphvizNode(child)).with(calculateEdgeColor(child.getStatus())));
162        graphvizNodes.put(parent.getName(), graphvizParent);
163
164        edges.put(parent.getName(), child.getName());
165        arcCount++;
166    }
167
168    private Color calculateEdgeColor(final WorkflowAction.Status childStatus) {
169        if (childStatus == null) {
170            return Color.BLACK;
171        }
172
173        if (childStatus.equals(WorkflowAction.Status.RUNNING)) {
174            return Color.GREEN;
175        }
176
177        return getColor(childStatus);
178    }
179
180    @Override
181    public void persist(final WorkflowActionNode node) {
182        final Node graphvizNode = graphvizNodes.get(node.getName());
183        graphvizGraph = graphvizGraph.with(graphvizNode);
184    }
185
186    @Override
187    public BufferedImage renderPng() {
188        final Future<BufferedImage> pngFuture = EXECUTOR_SERVICE.submit(new PngRenderer());
189
190        try {
191            return pngFuture.get(GRAPHVIZ_TIMEOUT_SECONDS, TimeUnit.SECONDS);
192        } catch (final InterruptedException | ExecutionException | TimeoutException e) {
193            throw new RuntimeException(e);
194        }
195    }
196
197    private int calculateHeight(final int arcCount) {
198        return Math.min(arcCount * 100, 2000);
199    }
200
201    @Override
202    public String renderDot() {
203        return graphvizGraph.toString();
204    }
205
206    @Override
207    public String renderSvg() {
208        final Future<String> svgFuture = EXECUTOR_SERVICE.submit(new SvgRenderer());
209
210        try {
211            return svgFuture.get(GRAPHVIZ_TIMEOUT_SECONDS, TimeUnit.SECONDS);
212        } catch (final InterruptedException | ExecutionException | TimeoutException e) {
213            throw new RuntimeException(e);
214        }
215    }
216
217    private class PngRenderer implements Callable<BufferedImage> {
218        @Override
219        public BufferedImage call() throws Exception {
220            final Graphviz graphviz = newGraphviz();
221
222            return graphviz.render(Format.PNG).toImage();
223        }
224    }
225
226    private class SvgRenderer implements Callable<String> {
227
228        @Override
229        public String call() throws Exception {
230            final Graphviz graphviz = newGraphviz();
231
232            return graphviz.render(Format.SVG).toString();
233        }
234    }
235
236    private Graphviz newGraphviz() {
237        return Graphviz.fromGraph(graphvizGraph)
238                .rasterizer(Rasterizer.BATIK)
239                .engine(Engine.DOT)
240                .height(calculateHeight(arcCount));
241    }
242}