IMPALA-13033: Parse WebUI thrift profile downloads

WebUI thrift profile downloads contain the archived profile string
without the timestamp and query id prefix used by profile log lines.
impala-profile-tool currently treats these files as malformed input.

Retry failed log-line parses as bare archived profile input while
keeping query id and timestamp filters limited to log entries that
include metadata. Add regression coverage for bare profile input using
the existing public profile fixtures.

Generated-by: OpenAI Codex
Change-Id: I2baed77cce3de43ef153b853329ad899f999f473
Reviewed-on: http://gerrit.cloudera.org:8080/24357
Reviewed-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>
Tested-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>

Author: alexandrefimov

be/src/util/impala-profile-tool.cc

@@ -25,19 +25,22 @@
 #include <gflags/gflags.h>
 
 #include "common/object-pool.h"
+#include "gutil/strings/strip.h"
 #include "util/os-info.h"
 #include "util/runtime-profile.h"
 
 #include "common/names.h"
 
 static const char* USAGE =
-    "Utility to decode an Impala profile log from standard input.\n"
+    "Utility to decode an Impala profile log or WebUI thrift profile download"
+    " from standard input.\n"
     "\n"
-    "The profile log is consumed from standard input and each successfully parsed entry"
+    "The input is consumed from standard input and each successfully parsed profile"
     " is pretty-printed to standard output.\n"
     "\n"
-    "Usage:"
+    "Usage:\n"
     "  impala-profile-tool < impala_profile_log_1.1-1607057366897\n"
+    "  impala-profile-tool < thrift_profile_<query id>.txt\n"
     "\n"
     "The following options are supported:\n"
     "Output options:\n"
@@ -54,7 +57,9 @@ static const char* USAGE =
     "--min_timestamp=<integer timestamp>: only process profiles at or after this"
     " timestamp\n"
     "--max_timestamp=<integer timestamp>: only process profiles at or before this"
-    " timestamp\n";
+    " timestamp\n"
+    "Filtering options only apply to profile log entries that include timestamp and"
+    " query id metadata.\n";
 
 DEFINE_string(
     profile_format, "text", "Profile format to output: either text, json or prettyjson");
@@ -104,23 +109,25 @@ int main(int argc, char** argv) {
   int profiles_emitted = 0;
   string line;
   int lineno = 1;
-  // Read profile log lines from stdin.
+  // Read profile log or WebUI thrift profile lines from stdin.
   for (; getline(cin, line); ++lineno) {
-    // Parse out fields from the line.
+    // Profile logs prefix each encoded profile with timestamp and query id. WebUI
+    // thrift profile downloads contain only the encoded profile.
     istringstream liness(line);
-    int64_t timestamp;
+    int64_t timestamp = -1;
     string query_id, encoded_profile;
     liness >> timestamp >> query_id >> encoded_profile;
-    if (liness.fail()) {
-      cerr << "Error parsing line " << lineno << ": '" << line << "'\n";
-      ++errors;
-      continue;
+    const bool has_log_metadata = !liness.fail();
+    if (!has_log_metadata) {
+      encoded_profile = line;
+      StripWhiteSpace(&encoded_profile);
     }
 
     // Skip decoding entries that don't match our parameters.
-    if ((FLAGS_query_id != "" && FLAGS_query_id != query_id) ||
-        (FLAGS_min_timestamp != -1 && timestamp < FLAGS_min_timestamp) ||
-        (FLAGS_max_timestamp != -1 && timestamp > FLAGS_max_timestamp)) {
+    if (has_log_metadata
+        && ((FLAGS_query_id != "" && FLAGS_query_id != query_id) ||
+            (FLAGS_min_timestamp != -1 && timestamp < FLAGS_min_timestamp) ||
+            (FLAGS_max_timestamp != -1 && timestamp > FLAGS_max_timestamp))) {
       continue;
     }
 

tests/observability/test_profile_tool.py

@@ -80,12 +80,53 @@ def test_json_output_profile_v2(self):
         get_profile_path(
             'impala_profile_log_tpcds_compute_stats_v2_extended.expected.pretty.json'))
 
+  def test_webui_thrift_profile_text_output(self):
+    # WebUI thrift profile downloads contain only the archived profile string, without
+    # the timestamp and query id prefix found in profile log lines.
+    self._compare_webui_thrift_profile_output([],
+        get_profile_path('impala_profile_log_tpcds_compute_stats'))
+    self._compare_webui_thrift_profile_output([],
+        get_profile_path('impala_profile_log_tpcds_compute_stats_v2'))
+
+  def test_webui_thrift_profile_prettyjson_output(self):
+    self._compare_webui_thrift_profile_output(['--profile_format=prettyjson'],
+        get_profile_path('impala_profile_log_tpcds_compute_stats'))
+
+  def test_webui_thrift_profile_ignores_surrounding_whitespace(self):
+    self._compare_webui_thrift_profile_output([],
+        get_profile_path('impala_profile_log_tpcds_compute_stats'),
+        profile_prefix=' \t', profile_suffix=' \r')
+
   def _compare_profile_tool_output(self, args, input_log, expected_output):
     """Run impala-profile-tool on input_log and compare it to the contents of the
     file at 'expected_output'."""
     with tempfile.NamedTemporaryFile() as tmp:
-      with open(input_log, 'r') as f:
-        check_call([os.path.join(IMPALA_HOME, "bin/run-binary.sh"),
-                    os.path.join(impalad_basedir, 'util/impala-profile-tool')] + args,
-                    stdin=f, stdout=tmp)
-        check_call(['diff', expected_output, tmp.name])
+      self._run_profile_tool(args, input_log, tmp)
+      check_call(['diff', expected_output, tmp.name])
+
+  def _compare_webui_thrift_profile_output(
+      self, args, input_log, profile_prefix='', profile_suffix=''):
+    """Compare a bare WebUI thrift profile to the same profile in a profile log."""
+    query_id, encoded_profile = self._get_first_profile_log_entry(input_log)
+    with tempfile.NamedTemporaryFile() as thrift_profile:
+      with tempfile.NamedTemporaryFile() as expected_output:
+        with tempfile.NamedTemporaryFile() as actual_output:
+          profile_input = profile_prefix + encoded_profile + profile_suffix + '\n'
+          thrift_profile.write(profile_input.encode('utf-8'))
+          thrift_profile.flush()
+          self._run_profile_tool(
+              args + ['--query_id=%s' % query_id], input_log, expected_output)
+          self._run_profile_tool(args, thrift_profile.name, actual_output)
+          check_call(['diff', expected_output.name, actual_output.name])
+
+  def _get_first_profile_log_entry(self, input_log):
+    with open(input_log, 'r') as f:
+      _timestamp, query_id, encoded_profile = f.readline().split(None, 2)
+    return query_id, encoded_profile.rstrip()
+
+  def _run_profile_tool(self, args, input_log, output=None):
+    with open(input_log, 'r') as f:
+      command = [os.path.join(IMPALA_HOME, "bin/run-binary.sh"),
+                 os.path.join(impalad_basedir, 'util/impala-profile-tool')] + args
+      check_call(command, stdin=f, stdout=output)
+      output.flush()