From 0c34e6b35b61c6572ec29611cdeea90a075f3717 Mon Sep 17 00:00:00 2001
From: Kaxil Naik <kaxilnaik@gmail.com>
Date: Wed, 1 Apr 2020 22:40:51 +0100
Subject: [PATCH] [AIRFLOW-6871] optimize tree view for large DAGs (#7492)

This change reduces page size by more than 10X and
reduces page load time by 3-5X. As a result, the
tree view can now load large DAGs that were causing
5XX error without the patch.

List of optimizations applied to the view handler:
* only seralize used task instance attributes to json instead of the
  whole ORM object
* encode task instance attributes as array instead of dict
* encode datetime in unix timestamp instead of iso formmat string
* push task instance attribute construction into client side JS
* remove redundant task instance attributes
* simplify reduce_nodes() logic, remove unnecessary if statements
* seralize JSON as string to be used with JSON.parse on the client side
  to speed up browser JS parse time
* remove spaces in seralized JSON string to reduce payload size

Co-Authored-By: QP Hou <houqp@users.noreply.github.com>

(cherry-picked from c1c2d6a1a13350672de52b01a9456103a8dc9b2e)
---
 airflow/www_rbac/templates/airflow/tree.html | 100 ++++++++++++---
 airflow/www_rbac/views.py                    | 126 ++++++++++++-------
 2 files changed, 167 insertions(+), 59 deletions(-)

diff --git a/airflow/www_rbac/templates/airflow/tree.html b/airflow/www_rbac/templates/airflow/tree.html
index 9bf29c8643af32..7b6288fd72213a 100644
--- a/airflow/www_rbac/templates/airflow/tree.html
+++ b/airflow/www_rbac/templates/airflow/tree.html
@@ -84,8 +84,59 @@
 <script>
 $('span.status_square').tooltip({html: true});
 
+function ts_to_dtstr(ts) {
+  var dt = new Date(ts * 1000);
+  return dt.toISOString();
+}
+
+function is_dag_run(d) {
+  return d.run_id !== undefined;
+}
+
+var now_ts = Date.now()/1000;
+
+function populate_taskinstance_properties(node) {
+  // populate task instance properties for display purpose
+  var j;
+  for (j=0; j<node.instances.length; j++) {
+    var dr_instance = data.instances[j];
+    var row = node.instances[j];
+
+    if (row === null) {
+      node.instances[j] = {
+        task_id: node.name,
+        execution_date: dr_instance.execution_date,
+      };
+      continue;
+    }
+
+    var task_instance = {
+      state: row[0],
+      try_number: row[1],
+      start_ts: row[2],
+      duration: row[3],
+    };
+    node.instances[j] = task_instance;
+
+    task_instance.task_id = node.name;
+    task_instance.execution_date = dr_instance.execution_date;
+    task_instance.external_trigger = dr_instance.external_trigger;
+
+    // compute start_date and end_date if applicable
+    if (task_instance.start_ts !== null) {
+      task_instance.start_date = ts_to_dtstr(task_instance.start_ts);
+      if (task_instance.state === "running") {
+        task_instance.duration = now_ts - task_instance.start_ts;
+      } else if (task_instance.duration !== null) {
+        task_instance.end_date = ts_to_dtstr(task_instance.start_ts + task_instance.duration);
+      }
+    }
+  }
+}
+
 var devicePixelRatio = window.devicePixelRatio || 1;
-var data = {{ data|tojson|safe }};
+// JSON.parse is faster for large payloads than an object literal (because the JSON grammer is simpler!)
+var data = JSON.parse('{{ data }}');
 var barHeight = 20;
 var axisHeight = 40;
 var square_x = parseInt(500 * devicePixelRatio);
@@ -103,8 +154,25 @@
 var nodes = tree.nodes(data);
 var nodeobj = {};
 for (i=0; i<nodes.length; i++) {
-    node = nodes[i];
-    nodeobj[node.name] = node;
+  var node = nodes[i];
+  nodeobj[node.name] = node;
+
+  if (node.name === "[DAG]") {
+    // skip synthetic root node since it's doesn't contain actual task instances
+    continue;
+  }
+
+  if (node.start_ts !== undefined) {
+    node.start_date = ts_to_dtstr(node.start_ts);
+  }
+  if (node.end_ts !== undefined) {
+    node.end_date = ts_to_dtstr(node.end_ts);
+  }
+  if (node.depends_on_past === undefined) {
+    node.depends_on_past = false;
+  }
+
+  populate_taskinstance_properties(node);
 }
 
 var diagonal = d3.svg.diagonal()
@@ -213,6 +281,7 @@
           if (d.operator != undefined) {
             tt += "operator: " + escapeHtml(d.operator) + "<br/>";
           }
+
           tt += "depends_on_past: " + escapeHtml(d.depends_on_past) + "<br/>";
           tt += "upstream: " + escapeHtml(d.num_dep) + "<br/>";
           tt += "retries: " + escapeHtml(d.retries) + "<br/>";
@@ -228,7 +297,7 @@
       .on('mouseout', function(d, i) {
         taskTip.hide(d)
         d3.select(this).transition()
-         .style("stroke-width", function(d) {return (d.run_id != undefined)? "2": "1"})
+         .style("stroke-width", function(d) {return is_dag_run(d)? "2": "1"})
       })
       .attr("height", barHeight)
       .attr("width", function(d, i) {return barWidth - d.y;})
@@ -262,10 +331,10 @@
       })
       .attr("class", function(d) {return "state " + d.state})
       .attr("data-toggle", "tooltip")
-      .attr("rx", function(d) {return (d.run_id != undefined)? "5": "0"})
-      .attr("ry", function(d) {return (d.run_id != undefined)? "5": "0"})
-      .style("shape-rendering", function(d) {return (d.run_id != undefined)? "auto": "crispEdges"})
-      .style("stroke-width", function(d) {return (d.run_id != undefined)? "2": "1"})
+      .attr("rx", function(d) {return is_dag_run(d)? "5": "0"})
+      .attr("ry", function(d) {return is_dag_run(d)? "5": "0"})
+      .style("shape-rendering", function(d) {return is_dag_run(d)? "auto": "crispEdges"})
+      .style("stroke-width", function(d) {return is_dag_run(d)? "2": "1"})
       .style("stroke-opacity", function(d) {return d.external_trigger ? "0": "1"})
       .on("mouseover", function(d){
         var s = "";
@@ -273,15 +342,18 @@
           s += "Task_id: " + escapeHtml(d.task_id) + "<br>";
         }
         s += "Run: " + escapeHtml(d.execution_date) + "<br>";
-        if(d.run_id != undefined){
+        if(is_dag_run(d)){
           s += "run_id: <nobr>" + escapeHtml(d.run_id) + "</nobr><br>";
         }
-        if (d.operator != undefined) {
-          s += "Operator: " + escapeHtml(d.operator) + "<br>"
-        }
+
+        var task = nodeobj[d.task_id];
+
+        s += "Operator: " + escapeHtml(task.operator) + "<br>"
         if(d.start_date != undefined){
           s += "Started: " + escapeHtml(d.start_date) + "<br>";
-          s += "Ended: " + escapeHtml(d.end_date) + "<br>";
+          if(d.end_date != undefined){
+            s += "Ended: " + escapeHtml(d.end_date) + "<br>";
+          }
           if (d.duration != undefined) {
             s += "Duration: " + escapeHtml(convertSecsToHumanReadable(d.duration)) + "<br>";
           }
@@ -295,7 +367,7 @@
       .on('mouseout', function(d,i) {
         taskTip.hide(d)
         d3.select(this).transition()
-         .style("stroke-width", function(d) {return (d.run_id != undefined)? "2": "1"})
+         .style("stroke-width", function(d) {return is_dag_run(d)? "2": "1"})
       })
       .attr('x', function(d, i) {return (i*(square_size+square_spacing));})
       .attr('y', -square_size/2)
diff --git a/airflow/www_rbac/views.py b/airflow/www_rbac/views.py
index 52958028aa7e06..233ab262dfdbe8 100644
--- a/airflow/www_rbac/views.py
+++ b/airflow/www_rbac/views.py
@@ -30,6 +30,7 @@
 from datetime import timedelta
 from urllib.parse import unquote
 
+import six
 from six.moves.urllib.parse import quote
 
 import markdown
@@ -45,6 +46,7 @@
 from flask_appbuilder.models.sqla.filters import BaseFilter
 from flask_babel import lazy_gettext
 import lazy_object_proxy
+from jinja2.utils import htmlsafe_json_dumps  # type: ignore
 from pygments import highlight, lexers
 from pygments.formatters import HtmlFormatter
 from sqlalchemy import and_, desc, func, or_, union_all
@@ -140,7 +142,6 @@ def get_date_time_num_runs_dag_runs_form_data(request, session, dag):
 ######################################################################################
 #                                    BaseViews
 ######################################################################################
-
 @app.errorhandler(404)
 def circles(error):
     return render_template(
@@ -1361,7 +1362,6 @@ def success(self):
     @action_logging
     @provide_session
     def tree(self, session=None):
-        default_dag_run = conf.getint('webserver', 'default_dag_run_display_number')
         dag_id = request.args.get('dag_id')
         blur = conf.getboolean('webserver', 'demo_mode')
         dag = dagbag.get_dag(dag_id)
@@ -1378,7 +1378,10 @@ def tree(self, session=None):
 
         base_date = request.args.get('base_date')
         num_runs = request.args.get('num_runs')
-        num_runs = int(num_runs) if num_runs else default_dag_run
+        if num_runs:
+            num_runs = int(num_runs)
+        else:
+            num_runs = conf.getint('webserver', 'default_dag_run_display_number')
 
         if base_date:
             base_date = timezone.parse(base_date)
@@ -1396,7 +1399,8 @@ def tree(self, session=None):
             .all()
         )
         dag_runs = {
-            dr.execution_date: alchemy_to_dict(dr) for dr in dag_runs}
+            dr.execution_date: alchemy_to_dict(dr) for dr in dag_runs
+        }
 
         dates = sorted(list(dag_runs.keys()))
         max_date = max(dates) if dates else None
@@ -1406,82 +1410,114 @@ def tree(self, session=None):
             start_date=min_date, end_date=base_date, session=session)
         task_instances = {}
         for ti in tis:
-            tid = alchemy_to_dict(ti)
-            dr = dag_runs.get(ti.execution_date)
-            tid['external_trigger'] = dr['external_trigger'] if dr else False
-            task_instances[(ti.task_id, ti.execution_date)] = tid
+            task_instances[(ti.task_id, ti.execution_date)] = ti
 
-        expanded = []
+        expanded = set()
         # The default recursion traces every path so that tree view has full
         # expand/collapse functionality. After 5,000 nodes we stop and fall
         # back on a quick DFS search for performance. See PR #320.
         node_count = [0]
         node_limit = 5000 / max(1, len(dag.leaves))
 
+        def encode_ti(ti):
+            if not ti:
+                return None
+
+            # NOTE: order of entry is important here because client JS relies on it for
+            # tree node reconstruction. Remember to change JS code in tree.html
+            # whenever order is altered.
+            data = [
+                ti.state,
+                ti.try_number,
+                None,  # start_ts
+                None,  # duration
+            ]
+
+            if ti.start_date:
+                # round to seconds to reduce payload size
+                if six.PY2:
+                    data[2] = int(pendulum.instance(ti.start_date).timestamp())
+                else:
+                    data[2] = int(ti.start_date.timestamp())
+                if ti.duration is not None:
+                    data[3] = int(ti.duration)
+
+            return data
+
         def recurse_nodes(task, visited):
-            visited.add(task)
             node_count[0] += 1
+            visited.add(task)
+            task_id = task.task_id
 
-            children = [
-                recurse_nodes(t, visited) for t in task.downstream_list
-                if node_count[0] < node_limit or t not in visited]
-
-            # D3 tree uses children vs _children to define what is
-            # expanded or not. The following block makes it such that
-            # repeated nodes are collapsed by default.
-            children_key = 'children'
-            if task.task_id not in expanded:
-                expanded.append(task.task_id)
-            elif children:
-                children_key = "_children"
-
-            def set_duration(tid):
-                if (isinstance(tid, dict) and tid.get("state") == State.RUNNING and
-                        tid["start_date"] is not None):
-                    d = timezone.utcnow() - pendulum.parse(tid["start_date"])
-                    tid["duration"] = d.total_seconds()
-                return tid
-
-            return {
+            node = {
                 'name': task.task_id,
                 'instances': [
-                    set_duration(task_instances.get((task.task_id, d))) or {
-                        'execution_date': d.isoformat(),
-                        'task_id': task.task_id
-                    }
-                    for d in dates],
-                children_key: children,
+                    encode_ti(task_instances.get((task_id, d)))
+                    for d in dates
+                ],
                 'num_dep': len(task.downstream_list),
                 'operator': task.task_type,
                 'retries': task.retries,
                 'owner': task.owner,
-                'start_date': task.start_date,
-                'end_date': task.end_date,
-                'depends_on_past': task.depends_on_past,
                 'ui_color': task.ui_color,
-                'extra_links': task.extra_links,
             }
 
+            if task.downstream_list:
+                children = [
+                    recurse_nodes(t, visited) for t in task.downstream_list
+                    if node_count[0] < node_limit or t not in visited]
+
+                # D3 tree uses children vs _children to define what is
+                # expanded or not. The following block makes it such that
+                # repeated nodes are collapsed by default.
+                if task.task_id not in expanded:
+                    children_key = 'children'
+                    expanded.add(task.task_id)
+                else:
+                    children_key = "_children"
+                node[children_key] = children
+
+            if task.depends_on_past:
+                node['depends_on_past'] = task.depends_on_past
+            if task.start_date:
+                # round to seconds to reduce payload size
+                if six.PY2:
+                    node['start_ts'] = int(pendulum.instance(task.start_date).timestamp())
+                else:
+                    node['start_ts'] = int(task.start_date.timestamp())
+                if task.end_date:
+                    # round to seconds to reduce payload size
+                    if six.PY2:
+                        node['end_ts'] = int(pendulum.instance(task.end_date).timestamp())
+                    else:
+                        node['end_ts'] = int(task.end_date.timestamp())
+            if task.extra_links:
+                node['extra_links'] = task.extra_links
+            return node
+
         data = {
             'name': '[DAG]',
             'children': [recurse_nodes(t, set()) for t in dag.roots],
             'instances': [
                 dag_runs.get(d) or {'execution_date': d.isoformat()}
-                for d in dates],
+                for d in dates
+            ],
         }
 
-        session.commit()
-
         form = DateTimeWithNumRunsForm(data={'base_date': max_date,
                                              'num_runs': num_runs})
         external_logs = conf.get('elasticsearch', 'frontend')
+
         return self.render_template(
             'airflow/tree.html',
             operators=sorted({op.task_type: op for op in dag.tasks}.values(),
                              key=lambda x: x.task_type),
             root=root,
             form=form,
-            dag=dag, data=data, blur=blur, num_runs=num_runs,
+            dag=dag,
+            # avoid spaces to reduce payload size
+            data=htmlsafe_json_dumps(data, separators=(',', ':')),
+            blur=blur, num_runs=num_runs,
             show_external_logs=bool(external_logs))
 
     @expose('/graph')