Support RR extension commands in LLDB

CMakeLists.txt

@@ -592,6 +592,7 @@ set(RR_SOURCES
   src/kernel_metadata.cc
   src/launch_debugger.cc
   src/log.cc
+  src/LldbInitCommand.cc
   src/LsCommand.cc
   src/main.cc
   src/MagicSaveDataMonitor.cc

src/DebuggerExtensionCommand.cc

@@ -183,8 +183,11 @@ static SimpleDebuggerExtensionCommand info_checkpoints(
   "list all checkpoints created with the 'checkpoint' command",
   invoke_info_checkpoints);
 
-/*static*/ void DebuggerExtensionCommand::init_auto_args() {
-  checkpoint.add_auto_arg("rr-where");
+void DebuggerExtensionCommand::init_auto_args() {
+  static __attribute__((unused)) int dummy = []() {
+    checkpoint.add_auto_arg("rr-where");
+    return 0;
+  }();
 }
 
 } // namespace rr

src/DebuggerExtensionCommandHandler.cc

@@ -15,32 +15,27 @@ namespace rr {
 // and linear search is fine.
 static vector<DebuggerExtensionCommand*>* debugger_command_list;
 
-static string gdb_macro_binding(const DebuggerExtensionCommand& cmd) {
-  string auto_args_str = "[";
-  for (size_t i = 0; i < cmd.auto_args().size(); i++) {
+static void print_python_string_array(ostream& out, const std::vector<std::string>& strs) {
+  out << "[";
+  for (size_t i = 0; i < strs.size(); i++) {
     if (i > 0) {
-      auto_args_str += ", ";
+      out << ", ";
     }
-    auto_args_str += "'" + cmd.auto_args()[i] + "'";
+    out << "'" << strs[i] << "'";
   }
-  auto_args_str += "]";
-  string ret = "python RRCmd('" + cmd.name() + "', " + auto_args_str + ")\n";
+  out << "]";
+}
+
+static void gdb_macro_binding(ostream& out, const DebuggerExtensionCommand& cmd) {
+  out << "python RRCmd('" << cmd.name() << "', ";
+  print_python_string_array(out, cmd.auto_args());
+  out << ")\n";
   if (!cmd.docs().empty()) {
-    ret += "document " + cmd.name() + "\n" + cmd.docs() + "\nend\n";
+    out << "document " << cmd.name() << "\n" << cmd.docs() << "\nend\n";
   }
-  return ret;
 }
 
-/* static */ string DebuggerExtensionCommandHandler::gdb_macros() {
-  DebuggerExtensionCommand::init_auto_args();
-  stringstream ss;
-  ss << string(R"Delimiter(
-
-set python print-stack full
-python
-
-import re
-
+static const char shared_python[] = R"Delimiter(
 def hex_unescape(string):
     str_len = len(string)
     if str_len % 2: # check for unexpected string length
@@ -64,6 +59,18 @@ def hex_escape(string):
         result += format(curr_char, '02x')
     return result
 
+)Delimiter";
+
+string DebuggerExtensionCommandHandler::gdb_macros() {
+  DebuggerExtensionCommand::init_auto_args();
+  stringstream ss;
+  ss << R"Delimiter(set python print-stack full
+python
+)Delimiter"
+    << shared_python
+    << R"Delimiter(
+import re
+
 class RRWhere(gdb.Command):
     """Helper to get the location for checkpoints/history. Used by auto-args"""
     def __init__(self):
@@ -163,11 +170,11 @@ RRSetSuppressRunHook()
 gdb.events.stop.connect(history_push)
 
 end
-)Delimiter");
+)Delimiter";
 
   if (debugger_command_list) {
     for (auto& it : *debugger_command_list) {
-      ss << gdb_macro_binding(*it);
+      gdb_macro_binding(ss, *it);
     }
   }
 
@@ -186,6 +193,77 @@ end
   return ss.str();
 }
 
+static void lldb_macro_binding(ostream& ss, const DebuggerExtensionCommand& cmd) {
+  string func_name = "rr_command_";
+  for (char ch : cmd.name()) {
+    if (ch == ' ') {
+      // We don't support commands with spaces in LLDB yet
+      return;
+    }
+    if (ch == '-') {
+      ch = '_';
+    }
+    func_name.push_back(ch);
+  }
+  ss << "def " << func_name << "(debugger, command, exe_ctx, result, internal_dict):\n";
+  if (!cmd.docs().empty()) {
+    ss << "    \"\"\"" << cmd.docs() << "\"\"\"\n";
+  }
+  ss << "    cmd_name = '" << cmd.name() << "'\n"
+     << "    auto_args = ";
+  print_python_string_array(ss, cmd.auto_args());
+  ss << "\n"
+     << "    command_impl(debugger, command, exe_ctx, result, cmd_name, auto_args)\n"
+     << "\n"
+     << "lldb.debugger.HandleCommand('command script add -f " << func_name
+     << " " << cmd.name() << "')\n\n";
+}
+
+string DebuggerExtensionCommandHandler::lldb_python_macros() {
+  DebuggerExtensionCommand::init_auto_args();
+  stringstream ss;
+  ss << shared_python
+     << R"Delimiter(
+import lldb
+import re
+import shlex
+
+def run_command_and_get_output(debugger, command):
+    result = lldb.SBCommandReturnObject()
+    debugger.GetCommandInterpreter().HandleCommand(command, result)
+    assert result.Succeeded()
+    return result.GetOutput()
+
+def command_impl(debugger, command, exe_ctx, result, cmd_name, auto_args):
+    interpreter = debugger.GetCommandInterpreter()
+    args = shlex.split(command)
+    # Ensure lldb tells rr its current thread
+    curr_thread = exe_ctx.thread
+    cmd_prefix = ("process plugin packet send qRRCmd:%s:%d"%
+        (cmd_name, -1 if curr_thread is None else curr_thread.GetThreadID()))
+    arg_strs = []
+    for auto_arg in auto_args:
+        arg_strs.append(":" + hex_escape(run_command_and_get_output(debugger, auto_arg)))
+    for arg in args:
+        arg_strs.append(":" + hex_escape(arg))
+    rv = run_command_and_get_output(debugger, cmd_prefix + ''.join(arg_strs));
+    rv_match = re.search('response: (.*)$', rv, re.MULTILINE);
+    if not rv_match:
+        result.SetError(None, "Invalid response: %s" % rv)
+        return
+    response = hex_unescape(rv_match.group(1))
+    result.Print(response.strip())
+
+)Delimiter";
+
+  if (debugger_command_list) {
+    for (auto& it : *debugger_command_list) {
+      lldb_macro_binding(ss, *it);
+    }
+  }
+  return ss.str();
+}
+
 /*static*/ DebuggerExtensionCommand* DebuggerExtensionCommandHandler::command_for_name(const string& name) {
   if (!debugger_command_list) {
     return nullptr;

src/DebuggerExtensionCommandHandler.h

@@ -14,19 +14,23 @@ class GdbServer;
 class Task;
 
 /**
- * rr extends debuggers (GDB) with custom commands such as `when`.
+ * rr extends debuggers (GDB, LLDB) with custom commands such as `when`.
  * This class manages those commands.
  */
 class DebuggerExtensionCommandHandler {
 public:
   // Declare any registered command with supporting
-  // wrapper code.
+  // wrapper code --- GDB script.
   static std::string gdb_macros();
 
+  // Declare any registered command with supporting
+  // wrapper code --- LLDB Python script.
+  static std::string lldb_python_macros();
+
   static void register_command(DebuggerExtensionCommand& cmd);
 
   /**
-   * Process an incoming GDB payload of the following form:
+   * Process an incoming debugger payload of the following form:
    *   <command name>:<arg1>:<arg2>:...
    *
    * NOTE: RR Command are typically sent with the qRRCmd: prefix which

src/LldbInitCommand.cc

@@ -0,0 +1,35 @@
+/* -*- Mode: C++; tab-width: 8; c-basic-offset: 2; indent-tabs-mode: nil; -*- */
+
+#include "Command.h"
+#include "GdbServer.h"
+#include "launch_debugger.h"
+#include "main.h"
+
+using namespace std;
+
+namespace rr {
+
+class LldbInitCommand : public Command {
+public:
+  virtual int run(vector<string>& args) override;
+
+protected:
+  LldbInitCommand(const char* name, const char* help) : Command(name, help) {}
+
+  static LldbInitCommand singleton;
+};
+
+LldbInitCommand LldbInitCommand::singleton("lldbinit", " rr lldbinit\n");
+
+int LldbInitCommand::run(vector<string>& args) {
+  while (parse_global_option(args)) {
+  }
+
+  puts("# This is a Python script. Save it to a file and run it from LLDB using");
+  puts("#     script exec(open('<filename>').read())");
+  puts("# or similar.");
+  fputs(lldb_init_script().c_str(), stdout);
+  return 0;
+}
+
+} // namespace rr

src/launch_debugger.cc

@@ -138,15 +138,13 @@ static const string& gdb_rr_macros() {
   return s;
 }
 
-// Special-sauce macros defined by rr when launching the gdb client,
-// which implement functionality outside of the gdb remote protocol.
-// (Don't stare at them too long or you'll go blind ;).)
-static const string& lldb_rr_macros() {
+static const string& lldb_python_rr_macros() {
   static string s;
 
   if (s.empty()) {
     stringstream ss;
-    ss << "set set prompt \"(rr) \"\n";
+    ss << DebuggerExtensionCommandHandler::lldb_python_macros()
+       << "lldb.debugger.HandleCommand('set set prompt \"(rr) \"')\n";
     s = ss.str();
   }
   return s;
@@ -332,9 +330,12 @@ void launch_debugger(ScopedFd& params_pipe_fd,
     }
     case DebuggerType::LLDB: {
       cmd.push_back("--source-quietly");
-      string lldb_command_file = create_command_file(lldb_rr_macros());
-      cmd.push_back("--source-before-file");
-      cmd.push_back(lldb_command_file);
+      // We have to load the commands as a Python script. If we
+      // use the "script" command to launch a nested Python interpreter,
+      // Python emits some annoying text that we dont want to see.
+      string lldb_command_file = create_command_file(lldb_python_rr_macros());
+      cmd.push_back("-o");
+      cmd.push_back("script exec(open('" + lldb_command_file + "').read())");
       cmd.insert(cmd.end(), options.begin(), options.end());
       push_lldb_target_remote_cmd(cmd, socket_domain, host, port);
       env.push_back("LLDB_UNDER_RR=1");
@@ -403,4 +404,6 @@ void emergency_debug(Task* t) {
 
 string gdb_init_script() { return gdb_rr_macros(); }
 
+string lldb_init_script() { return lldb_python_rr_macros(); }
+
 } // namespace rr

src/launch_debugger.h

@@ -55,6 +55,11 @@ void emergency_debug(Task* t);
  */
 std::string gdb_init_script();
 
+/**
+ * A string containing the default lldbinit script that we load into lldb.
+ */
+std::string lldb_init_script();
+
 } // namespace rr
 
 #endif /* RR_LAUNCH_DEBUGGER_H_ */

src/test/util.py

@@ -9,7 +9,8 @@
             'delete_watchpoint', 'expect_signal_stop',
             'set_breakpoint_commands', 'select_thread',
             'scheduler_locking_on', 'scheduler_locking_off',
-            'expect_expression', 'expect_threads',  ]
+            'expect_expression', 'expect_threads',
+            'send_custom_command' ]
 
 # Don't use python timeout. Use test-monitor timeout instead.
 TIMEOUT_SEC = 10000
@@ -165,6 +166,9 @@ def scheduler_locking_off():
     if debugger_type == 'GDB':
         send_gdb('set scheduler-locking off')
 
+def send_custom_command(cmd):
+    send_debugger(cmd, cmd)
+
 def send_debugger(gdb_cmd, lldb_cmd):
     if debugger_type == 'GDB':
         send_gdb(gdb_cmd)