001/*
002 * Copyright 2012 Apache Software Foundation.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      http://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016package org.apache.oozie.util;
017
018import edu.uci.ics.jung.algorithms.layout.StaticLayout;
019import edu.uci.ics.jung.graph.DirectedSparseGraph;
020import edu.uci.ics.jung.graph.Graph;
021import edu.uci.ics.jung.graph.util.Context;
022import edu.uci.ics.jung.visualization.VisualizationImageServer;
023import edu.uci.ics.jung.visualization.renderers.Renderer;
024import edu.uci.ics.jung.visualization.util.ArrowFactory;
025import java.awt.*;
026import java.awt.geom.Ellipse2D;
027import java.awt.geom.Point2D;
028import java.awt.image.BufferedImage;
029import java.io.IOException;
030import java.io.OutputStream;
031import java.io.StringReader;
032import java.util.HashMap;
033import java.util.Iterator;
034import java.util.LinkedHashMap;
035import java.util.Map;
036import javax.imageio.ImageIO;
037import javax.xml.parsers.SAXParser;
038import javax.xml.parsers.SAXParserFactory;
039import org.apache.commons.collections15.Transformer;
040import org.apache.oozie.client.WorkflowAction;
041import org.apache.oozie.client.WorkflowAction.Status;
042import org.apache.oozie.client.WorkflowJob;
043import org.apache.oozie.client.rest.JsonWorkflowJob;
044import org.xml.sax.Attributes;
045import org.xml.sax.InputSource;
046import org.xml.sax.SAXException;
047import org.xml.sax.XMLReader;
048import org.xml.sax.helpers.DefaultHandler;
049
050/**
051 * Class to generate and plot runtime workflow DAG
052 */
053public class GraphGenerator {
054
055    private String xml;
056    private JsonWorkflowJob job;
057    private boolean showKill = false;
058
059    /**
060     * C'tor
061     * @param xml The workflow definition XML
062     * @param job Current status of the job
063     * @param showKill Flag to whether show 'kill' node
064     */
065    public GraphGenerator(String xml, JsonWorkflowJob job, boolean showKill) {
066        if(job == null) {
067            throw new IllegalArgumentException("JsonWorkflowJob can't be null");
068        }
069        this.xml = xml;
070        this.job = job;
071        this.showKill = showKill;
072    }
073
074    /**
075     * C'tor
076     * @param xml
077     * @param job
078     */
079    public GraphGenerator(String xml, JsonWorkflowJob job) {
080        this(xml, job, false);
081    }
082
083    /**
084     * Overridden to thwart finalizer attack
085     */
086    @Override
087    public final void finalize() {
088        // No-op; just to avoid finalizer attack
089        // as the constructor is throwing an exception
090    }
091
092    /**
093     * Stream the PNG file to client
094     * @param out
095     * @throws Exception
096     */
097    public void write(OutputStream out) throws Exception {
098        SAXParserFactory spf = SAXParserFactory.newInstance();
099        spf.setNamespaceAware(true);
100        SAXParser saxParser = spf.newSAXParser();
101        XMLReader xmlReader = saxParser.getXMLReader();
102        xmlReader.setContentHandler(new XMLParser(out));
103        xmlReader.parse(new InputSource(new StringReader(xml)));
104    }
105
106    private class XMLParser extends DefaultHandler {
107
108        private OutputStream out;
109        private LinkedHashMap<String, OozieWFNode> tags;
110
111        private String action = null;
112        private String actionOK = null;
113        private String actionErr = null;
114        private String actionType = null;
115        private String fork;
116        private String decision;
117
118        public XMLParser(OutputStream out) {
119            this.out = out;
120        }
121
122        @Override
123        public void startDocument() throws SAXException {
124            tags = new LinkedHashMap();
125        }
126
127        @Override
128        public void endDocument() throws SAXException {
129
130            if(tags.isEmpty()) {
131                // Nothing to do here!
132                return;
133            }
134
135            int maxX = Integer.MIN_VALUE;
136            int maxY = Integer.MIN_VALUE;
137            int minX = Integer.MAX_VALUE;
138            int currX = 45;
139            int currY = 45;
140            final int xMargin = 205;
141            final int yMargin = 50;
142            final int xIncr = 215; // The widest element is 200 pixels (Rectangle)
143            final int yIncr = 255; // The tallest element is 150 pixels; (Diamond)
144            HashMap<String, WorkflowAction> actionMap = new HashMap<String, WorkflowAction>();
145
146            // Create a hashmap for faster lookups
147            // Also override showKill if there's any failed action
148            boolean found = false;
149            for(WorkflowAction wfAction : job.getActions()) {
150                actionMap.put(wfAction.getName(), wfAction);
151                if(!found) {
152                    switch(wfAction.getStatus()) {
153                        case KILLED:
154                        case ERROR:
155                        case FAILED:
156                            showKill = true; // Assuming on error the workflow eventually ends with kill node
157                            found = true;
158                    }
159                }
160            }
161
162            // Start building the graph
163            DirectedSparseGraph<OozieWFNode, String> dg = new DirectedSparseGraph<OozieWFNode, String>();
164            for(Map.Entry<String, OozieWFNode> entry : tags.entrySet()) {
165                String name = entry.getKey();
166                OozieWFNode node = entry.getValue();
167                if(actionMap.containsKey(name)) {
168                    node.setStatus(actionMap.get(name).getStatus());
169                }
170
171                // Set (x,y) coords of the vertices if not already set
172                if(node.getLocation().equals(new Point(0, 0))) {
173                    node.setLocation(currX, currY);
174                }
175
176                float childStep = showKill ? -(((float)node.getArcs().size() - 1 ) / 2)
177                        : -((float)node.getArcs().size() / 2 - 1);
178                int nodeX = node.getLocation().x;
179                int nodeY = node.getLocation().y;
180                for(Map.Entry<String, Boolean> arc : node.getArcs().entrySet()) {
181                    if(!showKill && arc.getValue() && tags.get(arc.getKey()).getType().equals("kill")) {
182                        // Don't show kill node (assumption: only error goes to kill node;
183                        // No ok goes to kill node)
184                        continue;
185                    }
186                    OozieWFNode child = tags.get(arc.getKey());
187                    if(child == null) {
188                        continue; // or throw error?
189                    }
190                    dg.addEdge(name + "-->" + arc.getKey(), node, child);
191                    // TODO: Experimental -- should we set coords even if they're already set?
192                    //if(child.getLocation().equals(new Point(0, 0))) {
193                        int childX = (int)(nodeX + childStep * xIncr);
194                        int childY = nodeY + yIncr;
195                        child.setLocation(childX, childY);
196
197                        if(minX > childX) {
198                            minX = childX;
199                        }
200                        if(maxX < childX) {
201                            maxX = childX;
202                        }
203                        if(maxY < childY) {
204                            maxY = childY;
205                        }
206                    //}
207                    childStep += 1;
208                }
209
210                currY += yIncr;
211                currX = nodeX;
212                if(minX > nodeX) {
213                    minX = nodeX;
214                }
215                if(maxX < nodeX) {
216                    maxX = nodeX;
217                }
218                if(maxY < nodeY) {
219                    maxY = nodeY;
220                }
221            } // Done building graph
222
223            final int padX = minX < 0 ? -minX: 0;
224
225            Transformer<OozieWFNode, Point2D> locationInit = new Transformer<OozieWFNode, Point2D>() {
226
227                @Override
228                public Point2D transform(OozieWFNode node) {
229                    if(padX == 0) {
230                        return node.getLocation();
231                    } else {
232                        return new Point(node.getLocation().x + padX + xMargin, node.getLocation().y);
233                    }
234                }
235
236            };
237
238            StaticLayout<OozieWFNode, String> layout = new StaticLayout<OozieWFNode, String>(dg, locationInit, new Dimension(maxX + padX + xMargin, maxY));
239            layout.lock(true);
240            VisualizationImageServer<OozieWFNode, String> vis = new VisualizationImageServer<OozieWFNode, String>(layout, new Dimension(maxX + padX + 2 * xMargin, maxY + yMargin));
241
242            vis.getRenderContext().setEdgeArrowTransformer(new ArrowShapeTransformer());
243            vis.getRenderContext().setArrowDrawPaintTransformer(new ArcPaintTransformer());
244            vis.getRenderContext().setEdgeDrawPaintTransformer(new ArcPaintTransformer());
245            vis.getRenderContext().setEdgeStrokeTransformer(new ArcStrokeTransformer());
246            vis.getRenderContext().setVertexShapeTransformer(new NodeShapeTransformer());
247            vis.getRenderContext().setVertexFillPaintTransformer(new NodePaintTransformer());
248            vis.getRenderContext().setVertexStrokeTransformer(new NodeStrokeTransformer());
249            vis.getRenderContext().setVertexLabelTransformer(new NodeLabelTransformer());
250            vis.getRenderContext().setVertexFontTransformer(new NodeFontTransformer());
251            vis.getRenderer().getVertexLabelRenderer().setPosition(Renderer.VertexLabel.Position.CNTR);
252            vis.setBackground(Color.WHITE);
253
254            Dimension d = vis.getSize();
255            BufferedImage img = new BufferedImage(d.width, d.height, BufferedImage.TYPE_INT_RGB);
256            Graphics2D g = img.createGraphics();
257            vis.paintAll(g);
258
259            try {
260                ImageIO.write(img, "png", out);
261            }
262            catch (IOException ioe) {
263                throw new SAXException(ioe);
264            }
265            finally {
266                try {
267                    out.close(); //closing connection is imperative
268                                 //regardless of ImageIO.write throwing exception or not
269                                 //hence in finally block
270                }
271                catch (IOException e) {
272                    XLog.getLog(getClass()).trace("Exception while closing OutputStream");
273                }
274                out = null;
275                img.flush();
276                g.dispose();
277                vis.removeAll();
278            }
279        }
280
281        @Override
282        public void startElement(String namespaceURI,
283                                String localName,
284                                String qName,
285                                Attributes atts)
286            throws SAXException {
287
288            if(localName.equalsIgnoreCase("start")) {
289                String start = localName.toLowerCase();
290                if(!tags.containsKey(start)) {
291                    OozieWFNode v = new OozieWFNode(start, start);
292                    v.addArc(atts.getValue("to"));
293                    tags.put(start, v);
294                }
295            } else if(localName.equalsIgnoreCase("action")) {
296                action = atts.getValue("name");
297            } else if(action != null && actionType == null) {
298                actionType = localName.toLowerCase();
299            } else if(localName.equalsIgnoreCase("ok") && action != null && actionOK == null) {
300                    actionOK = atts.getValue("to");
301            } else if(localName.equalsIgnoreCase("error") && action != null && actionErr == null) {
302                    actionErr = atts.getValue("to");
303            } else if(localName.equalsIgnoreCase("fork")) {
304                fork = atts.getValue("name");
305                if(!tags.containsKey(fork)) {
306                    tags.put(fork, new OozieWFNode(fork, localName.toLowerCase()));
307                }
308            } else if(localName.equalsIgnoreCase("path")) {
309                tags.get(fork).addArc(atts.getValue("start"));
310            } else if(localName.equalsIgnoreCase("join")) {
311                String join = atts.getValue("name");
312                if(!tags.containsKey(join)) {
313                    OozieWFNode v = new OozieWFNode(join, localName.toLowerCase());
314                    v.addArc(atts.getValue("to"));
315                    tags.put(join, v);
316                }
317            } else if(localName.equalsIgnoreCase("decision")) {
318                decision = atts.getValue("name");
319                if(!tags.containsKey(decision)) {
320                    tags.put(decision, new OozieWFNode(decision, localName.toLowerCase()));
321                }
322            } else if(localName.equalsIgnoreCase("case")
323                    || localName.equalsIgnoreCase("default")) {
324                tags.get(decision).addArc(atts.getValue("to"));
325            } else if(localName.equalsIgnoreCase("kill")
326                    || localName.equalsIgnoreCase("end")) {
327                String name = atts.getValue("name");
328                if(!tags.containsKey(name)) {
329                    tags.put(name, new OozieWFNode(name, localName.toLowerCase()));
330                }
331            }
332        }
333
334        @Override
335        public void endElement(String namespaceURI,
336                                String localName,
337                                String qName)
338                throws SAXException {
339            if(localName.equalsIgnoreCase("action")) {
340                tags.put(action, new OozieWFNode(action, actionType));
341                tags.get(action).addArc(this.actionOK);
342                tags.get(action).addArc(this.actionErr, true);
343                action = null;
344                actionOK = null;
345                actionErr = null;
346                actionType = null;
347            }
348        }
349
350        private class OozieWFNode {
351            private String name;
352            private String type;
353            private Point loc;
354            private HashMap<String, Boolean> arcs;
355            private Status status = null;
356
357            public OozieWFNode(String name,
358                    String type,
359                    HashMap<String, Boolean> arcs,
360                    Point loc,
361                    Status status) {
362                this.name = name;
363                this.type = type;
364                this.arcs = arcs;
365                this.loc = loc;
366                this.status = status;
367            }
368
369            public OozieWFNode(String name, String type, HashMap<String, Boolean> arcs) {
370                this(name, type, arcs, new Point(0, 0), null);
371            }
372
373            public OozieWFNode(String name, String type) {
374                this(name, type, new HashMap<String, Boolean>(), new Point(0, 0), null);
375            }
376
377            public OozieWFNode(String name, String type, WorkflowAction.Status status) {
378                this(name, type, new HashMap<String, Boolean>(), new Point(0, 0), status);
379            }
380
381            public void addArc(String arc, boolean isError) {
382                arcs.put(arc, isError);
383            }
384
385            public void addArc(String arc) {
386                addArc(arc, false);
387            }
388
389            public void setName(String name) {
390                this.name = name;
391            }
392
393            public void setType(String type) {
394                this.type = type;
395            }
396
397            public void setLocation(Point loc) {
398                this.loc = loc;
399            }
400
401            public void setLocation(double x, double y) {
402                loc.setLocation(x, y);
403            }
404
405            public void setStatus(WorkflowAction.Status status) {
406                this.status = status;
407            }
408
409            public String getName() {
410                return name;
411            }
412
413            public String getType() {
414                return type;
415            }
416
417            public HashMap<String, Boolean> getArcs() {
418                return arcs;
419            }
420
421            public Point getLocation() {
422                return loc;
423            }
424
425            public WorkflowAction.Status getStatus() {
426                return status;
427            }
428
429            @Override
430            public String toString() {
431                StringBuilder s = new StringBuilder();
432
433                s.append("Node: ").append(name).append("\t");
434                s.append("Type: ").append(type).append("\t");
435                s.append("Location: (").append(loc.getX()).append(", ").append(loc.getY()).append(")\t");
436                s.append("Status: ").append(status).append("\n");
437                Iterator<Map.Entry<String, Boolean>> it = arcs.entrySet().iterator();
438                while(it.hasNext()) {
439                    Map.Entry<String, Boolean> entry = it.next();
440
441                    s.append("\t").append(entry.getKey());
442                    if(entry.getValue().booleanValue()) {
443                        s.append(" on error\n");
444                    } else {
445                        s.append("\n");
446                    }
447                }
448
449                return s.toString();
450            }
451        }
452
453        private class NodeFontTransformer implements Transformer<OozieWFNode, Font> {
454            private final Font font = new Font("Default", Font.BOLD, 15);
455
456            @Override
457            public Font transform(OozieWFNode node) {
458                return font;
459            }
460        }
461
462        private class ArrowShapeTransformer implements Transformer<Context<Graph<OozieWFNode, String>, String>,  Shape> {
463            private final Shape arrow = ArrowFactory.getWedgeArrow(10.0f, 20.0f);
464
465            @Override
466            public Shape transform(Context<Graph<OozieWFNode, String>, String> i) {
467                return arrow;
468            }
469        }
470
471        private class ArcPaintTransformer implements Transformer<String, Paint> {
472            // Paint based on transition
473            @Override
474            public Paint transform(String arc) {
475                int sep = arc.indexOf("-->");
476                String source = arc.substring(0, sep);
477                String target = arc.substring(sep + 3);
478                OozieWFNode src = tags.get(source);
479                OozieWFNode tgt = tags.get(target);
480
481                if(src.getType().equals("start")) {
482                    if(tgt.getStatus() == null) {
483                        return Color.LIGHT_GRAY;
484                    } else {
485                        return Color.GREEN;
486                    }
487                }
488
489                if(src.getArcs().get(target)) {
490                    // Dealing with error transition (i.e. target is error)
491                    if(src.getStatus() == null) {
492                        return Color.LIGHT_GRAY;
493                    }
494                    switch(src.getStatus()) {
495                        case KILLED:
496                        case ERROR:
497                        case FAILED:
498                            return Color.RED;
499                        default:
500                            return Color.LIGHT_GRAY;
501                    }
502                } else {
503                    // Non-error
504                    if(src.getType().equals("decision")) {
505                        // Check for target too
506                        if(tgt.getStatus() != null) {
507                            return Color.GREEN;
508                        } else {
509                            return Color.LIGHT_GRAY;
510                        }
511                    } else {
512                        if(src.getStatus() == null) {
513                            return Color.LIGHT_GRAY;
514                        }
515                        switch(src.getStatus()) {
516                            case OK:
517                            case DONE:
518                            case END_RETRY:
519                            case END_MANUAL:
520                                return Color.GREEN;
521                            default:
522                                return Color.LIGHT_GRAY;
523                        }
524                    }
525                }
526            }
527        }
528
529        private class NodeStrokeTransformer implements Transformer<OozieWFNode, Stroke> {
530            private final Stroke stroke1 = new BasicStroke(2.0f);
531            private final Stroke stroke2 = new BasicStroke(4.0f);
532
533            @Override
534            public Stroke transform(OozieWFNode node) {
535                if(node.getType().equals("start")
536                        || node.getType().equals("end")
537                        || node.getType().equals("kill")) {
538                    return stroke2;
539                }
540                return stroke1;
541            }
542        }
543
544        private class NodeLabelTransformer implements Transformer<OozieWFNode, String> {
545            /*
546            * 20 chars in rectangle in 2 rows max
547            * 14 chars in diamond in 2 rows max
548            * 9 in triangle in 2 rows max
549            * 8 in invtriangle in 2 rows max
550            * 8 in circle in 2 rows max
551            */
552            @Override
553            public String transform(OozieWFNode node) {
554                //return node.getType();
555                String name = node.getName();
556                String type = node.getType();
557                StringBuilder s = new StringBuilder();
558                if(type.equals("decision")) {
559                    if(name.length() <= 14) {
560                        return name;
561                    } else {
562                        s.append("<html>").append(name.substring(0, 12)).append("-<br />");
563                        if(name.substring(13).length() > 14) {
564                            s.append(name.substring(12, 25)).append("...");
565                        } else {
566                            s.append(name.substring(12));
567                        }
568                        s.append("</html>");
569                        return s.toString();
570                    }
571                } else if(type.equals("fork")) {
572                    if(name.length() <= 9) {
573                        return "<html><br />" + name + "</html>";
574                    } else {
575                        s.append("<html><br />").append(name.substring(0, 7)).append("-<br />");
576                        if(name.substring(8).length() > 9) {
577                            s.append(name.substring(7, 15)).append("...");
578                        } else {
579                            s.append(name.substring(7));
580                        }
581                        s.append("</html>");
582                        return s.toString();
583                    }
584                } else if(type.equals("join")) {
585                    if(name.length() <= 8) {
586                        return "<html>" + name + "</html>";
587                    } else {
588                        s.append("<html>").append(name.substring(0, 6)).append("-<br />");
589                        if(name.substring(7).length() > 8) {
590                            s.append(name.substring(6, 13)).append("...");
591                        } else {
592                            s.append(name.substring(6));
593                        }
594                        s.append("</html>");
595                        return s.toString();
596                    }
597                } else if(type.equals("start")
598                        || type.equals("end")
599                        || type.equals("kill")) {
600                    if(name.length() <= 8) {
601                        return "<html>" + name + "</html>";
602                    } else {
603                        s.append("<html>").append(name.substring(0, 6)).append("-<br />");
604                        if(name.substring(7).length() > 8) {
605                            s.append(name.substring(6, 13)).append("...");
606                        } else {
607                            s.append(name.substring(6));
608                        }
609                        s.append("</html>");
610                        return s.toString();
611                    }
612                }else {
613                    if(name.length() <= 20) {
614                        return name;
615                    } else {
616                        s.append("<html>").append(name.substring(0, 18)).append("-<br />");
617                        if(name.substring(19).length() > 20) {
618                            s.append(name.substring(18, 37)).append("...");
619                        } else {
620                            s.append(name.substring(18));
621                        }
622                        s.append("</html>");
623                        return s.toString();
624                    }
625                }
626            }
627        }
628
629        private class NodePaintTransformer implements Transformer<OozieWFNode, Paint> {
630            @Override
631            public Paint transform(OozieWFNode node) {
632                WorkflowJob.Status jobStatus = job.getStatus();
633                if(node.getType().equals("start")) {
634                    return Color.WHITE;
635                } else if(node.getType().equals("end")) {
636                    if(jobStatus == WorkflowJob.Status.SUCCEEDED) {
637                        return Color.GREEN;
638                    }
639                    return Color.BLACK;
640                } else if(node.getType().equals("kill")) {
641                    if(jobStatus == WorkflowJob.Status.FAILED
642                            || jobStatus == WorkflowJob.Status.KILLED) {
643                        return Color.RED;
644                    }
645                    return Color.WHITE;
646                }
647
648                // Paint based on status for rest
649                WorkflowAction.Status status = node.getStatus();
650                if(status == null) {
651                    return Color.LIGHT_GRAY;
652                }
653                switch(status) {
654                    case OK:
655                    case DONE:
656                    case END_RETRY:
657                    case END_MANUAL:
658                        return Color.GREEN;
659                    case PREP:
660                    case RUNNING:
661                    case USER_RETRY:
662                    case START_RETRY:
663                    case START_MANUAL:
664                        return Color.YELLOW;
665                    case KILLED:
666                    case ERROR:
667                    case FAILED:
668                        return Color.RED;
669                    default:
670                        return Color.LIGHT_GRAY;
671                }
672            }
673        }
674
675        private class NodeShapeTransformer implements Transformer<OozieWFNode, Shape> {
676            private final Ellipse2D.Double circle = new Ellipse2D.Double(-40, -40, 80, 80);
677            private final Rectangle rect = new Rectangle(-100, -30, 200, 60);
678            private final Polygon diamond = new Polygon(new int[]{-75, 0, 75, 0}, new int[]{0, 75, 0, -75}, 4);
679            private final Polygon triangle = new Polygon(new int[]{-85, 85, 0}, new int[]{0, 0, -148}, 3);
680            private final Polygon invtriangle = new Polygon(new int[]{-85, 85, 0}, new int[]{0, 0, 148}, 3);
681
682            @Override
683            public Shape transform(OozieWFNode node) {
684                if("start".equals(node.getType())
685                    || "end".equals(node.getType())
686                    || "kill".equals(node.getType())) {
687                    return circle;
688                }
689                if("fork".equals(node.getType())) {
690                    return triangle;
691                }
692                if("join".equals(node.getType())) {
693                    return invtriangle;
694                }
695                if("decision".equals(node.getType())) {
696                    return diamond;
697                }
698                return rect; // All action nodes
699            }
700        }
701
702        private class ArcStrokeTransformer implements Transformer<String, Stroke> {
703            private final Stroke stroke1 = new BasicStroke(2.0f);
704            private final Stroke dashed = new BasicStroke(1.0f, BasicStroke.CAP_BUTT, BasicStroke.JOIN_MITER, 10.0f, new float[] {10.0f}, 0.0f);
705
706            // Draw based on transition
707            @Override
708            public Stroke transform(String arc) {
709                int sep = arc.indexOf("-->");
710                String source = arc.substring(0, sep);
711                String target = arc.substring(sep + 3);
712                OozieWFNode src = tags.get(source);
713                if(src.getArcs().get(target)) {
714                        if(src.getStatus() == null) {
715                            return dashed;
716                        }
717                        switch(src.getStatus()) {
718                            case KILLED:
719                            case ERROR:
720                            case FAILED:
721                                return stroke1;
722                            default:
723                                return dashed;
724                        }
725                } else {
726                    return stroke1;
727                }
728            }
729        }
730    }
731}