Fix several live mode bugs

news/147.bugfix.1.rst

@@ -0,0 +1 @@
+Fix the handling of the thread switch commands in the live mode TUI before the first allocation has been seen.

news/147.bugfix.2.rst

@@ -0,0 +1 @@
+Fix a bug that made ``memray run --live -c`` fail if the command to run contained double quotes.

news/147.bugfix.3.rst

@@ -0,0 +1 @@
+Ensure our TUI isn't displaying stale data by periodically flushing the latest available data from the tracker (rather than only flushing when a buffer fills up).

src/memray/_memray/record_writer.h

@@ -162,7 +162,7 @@ bool inline RecordWriter::writeRecordUnsafe(const MemoryRecord& record)
 {
     RecordTypeAndFlags token{RecordType::MEMORY_RECORD, 0};
     return writeSimpleType(token) && writeVarint(record.rss)
-           && writeVarint(record.ms_since_epoch - d_stats.start_time);
+           && writeVarint(record.ms_since_epoch - d_stats.start_time) && d_sink->flush();
 }
 
 bool inline RecordWriter::writeRecordUnsafe(const ContextSwitch& record)

src/memray/_memray/sink.h

@@ -16,6 +16,10 @@ class Sink
     virtual bool writeAll(const char* data, size_t length) = 0;
     virtual bool seek(off_t offset, int whence) = 0;
     virtual std::unique_ptr<Sink> cloneInChildProcess() = 0;
+    virtual bool flush()
+    {
+        return true;
+    }
 };
 
 class FileSink : public memray::io::Sink
@@ -64,11 +68,11 @@ class SocketSink : public Sink
     bool writeAll(const char* data, size_t length) override;
     bool seek(off_t offset, int whence) override;
     std::unique_ptr<Sink> cloneInChildProcess() override;
+    bool flush() override;
 
   private:
     size_t freeSpaceInBuffer();
     void open();
-    bool flush();
 
     const std::string d_host;
     uint16_t d_port;

src/memray/commands/run.py

@@ -95,13 +95,12 @@ def _run_child_process_and_attach(args: argparse.Namespace) -> None:
 
     arguments = (
         f"{port},{args.native},{args.run_as_module},{args.run_as_cmd},{args.quiet},"
-        f'"{args.script}",{args.script_args}'
+        f"{args.script!r},{args.script_args}"
     )
     tracked_app_cmd = [
         sys.executable,
         "-c",
-        f"from memray.commands.run import _child_process;"
-        f"_child_process({arguments})",
+        f"from memray.commands.run import _child_process;_child_process({arguments})",
     ]
     with contextlib.suppress(KeyboardInterrupt):
         with subprocess.Popen(

src/memray/reporters/tui.py

@@ -197,6 +197,11 @@ class TUI:
         "a": 5,
     }
 
+    # Start with a non-empty list of threads so that we always have something
+    # to display. This avoids "Thread 1 of 0" and fixes a DivideByZeroError
+    # when switching threads before the first allocation is seen.
+    _DUMMY_THREAD_LIST = [0]
+
     def __init__(self, pid: Optional[int], cmd_line: Optional[str], native: bool):
         self.pid = pid or "???"
         if not cmd_line:
@@ -207,7 +212,7 @@ def __init__(self, pid: Optional[int], cmd_line: Optional[str], native: bool):
         self._native = native
         self._thread_idx = 0
         self._seen_threads: Set[int] = set()
-        self._threads: List[int] = []
+        self._threads: List[int] = self._DUMMY_THREAD_LIST
         self.n_samples = 0
         self.start = datetime.now()
         self._last_update = datetime.now()
@@ -360,8 +365,6 @@ def message(self, message: str) -> None:
 
     @property
     def current_thread(self) -> int:
-        if not self._threads:
-            return 0
         return self._threads[self._thread_idx]
 
     def next_thread(self) -> None:
@@ -382,6 +385,8 @@ def update_snapshot(self, snapshot: Iterable[AllocationRecord]) -> None:
         for record in self._snapshot:
             if record.tid in self._seen_threads:
                 continue
+            if self._threads is self._DUMMY_THREAD_LIST:
+                self._threads = []
             self._threads.append(record.tid)
             self._seen_threads.add(record.tid)
         self.n_samples += 1

tests/unit/test_cli.py

@@ -166,7 +166,7 @@ def test_run_with_live(
                 sys.executable,
                 "-c",
                 "from memray.commands.run import _child_process;"
-                '_child_process(1234,False,False,False,False,"./directory/foobar.py",'
+                "_child_process(1234,False,False,False,False,'./directory/foobar.py',"
                 "['arg1', 'arg2'])",
             ],
             stderr=-1,

tests/unit/test_tui_reporter.py

@@ -50,7 +50,7 @@ def test_pid(self, pid, out_str):
             "2021",
             "                                               ╭─ Memory ─╮",
             f"(∩｀-´)⊃━☆ﾟ.*…PID: {out_str}      CMD: ???           │          │",
-            "              TID: 0x0      Thread 1 of 0      │          │",
+            "              TID: 0x0      Thread 1 of 1      │          │",
             "              Samples: 1    Duration: 0.0      │          │",
             "                            seconds            │          │",
             "                                               ╰──────────╯",
@@ -75,7 +75,7 @@ def test_command_line(self):
             "                                               ╭─ Memory ─╮",
             "(∩｀-´)⊃━☆ﾟ.*…PID: 123      CMD: python3       │          │",
             "                            some_command_to_te…│          │",
-            "              TID: 0x0      Thread 1 of 0      │          │",
+            "              TID: 0x0      Thread 1 of 1      │          │",
             "              Samples: 1    Duration: 0.0      │          │",
             "                            seconds            ╰──────────╯",
         ]
@@ -99,7 +99,7 @@ def test_too_long_command_line_is_trimmed(self):
             "                                               ╭─ Memory ─╮",
             "(∩｀-´)⊃━☆ﾟ.*…PID: 123      CMD: python3       │          │",
             "                            aaaaaaaaaaaaaaaaaa…│          │",
-            "              TID: 0x0      Thread 1 of 0      │          │",
+            "              TID: 0x0      Thread 1 of 1      │          │",
             "              Samples: 1    Duration: 0.0      │          │",
             "                            seconds            ╰──────────╯",
         ]
@@ -124,7 +124,7 @@ def test_with_no_allocations(self):
             "                                               ╭─ Memory ─╮",
             "(∩｀-´)⊃━☆ﾟ.*…PID: 123      CMD: python3       │          │",
             "                            some_program.py    │          │",
-            "              TID: 0x0      Thread 1 of 0      │          │",
+            "              TID: 0x0      Thread 1 of 1      │          │",
             "              Samples: 1    Duration: 0.0      │          │",
             "                            seconds            ╰──────────╯",
         ]