allow the process to exit early, fixes #25

src/exc.h

@@ -23,4 +23,10 @@ class FatalException : public std::runtime_error {
   explicit FatalException(const std::string &what_arg)
       : std::runtime_error(what_arg) {}
 };
+
+class PtraceException : public std::runtime_error {
+ public:
+  explicit PtraceException(const std::string &what_arg)
+      : std::runtime_error(what_arg) {}
+};
 }  // namespace pyflame

src/ptrace.cc

@@ -29,20 +29,20 @@ void PtraceAttach(pid_t pid) {
   if (ptrace(PTRACE_ATTACH, pid, 0, 0)) {
     std::ostringstream ss;
     ss << "Failed to attach to PID " << pid << ": " << strerror(errno);
-    throw FatalException(ss.str());
+    throw PtraceException(ss.str());
   }
   if (wait(nullptr) == -1) {
     std::ostringstream ss;
     ss << "Failed to wait on PID " << pid << ": " << strerror(errno);
-    throw FatalException(ss.str());
+    throw PtraceException(ss.str());
   }
 }
 
 void PtraceDetach(pid_t pid) {
   if (ptrace(PTRACE_DETACH, pid, 0, 0)) {
     std::ostringstream ss;
     ss << "Failed to detach PID " << pid << ": " << strerror(errno);
-    throw FatalException(ss.str());
+    throw PtraceException(ss.str());
   }
 }
 
@@ -53,7 +53,7 @@ long PtracePeek(pid_t pid, unsigned long addr) {
     std::ostringstream ss;
     ss << "Failed to PTRACE_PEEKDATA at " << reinterpret_cast<void *>(addr)
        << ": " << strerror(errno);
-    throw FatalException(ss.str());
+    throw PtraceException(ss.str());
   }
   return data;
 }

src/pyflame.cc

@@ -43,6 +43,23 @@ const char usage_str[] =
      "(default 0.001)\n"
      "  -v, --version        Show the version\n"
      "  -x, --exclude-idle   Exclude idle time from statistics\n");
+
+typedef std::unordered_map<frames_t, size_t, FrameHash> buckets_t;
+
+void PrintBuckets(const buckets_t &buckets) {
+  for (const auto &kv : buckets) {
+    if (kv.first.empty()) {
+      std::cerr << "fatal error\n";
+      return;
+    }
+    auto last = kv.first.rend();
+    last--;
+    for (auto it = kv.first.rbegin(); it != last; ++it) {
+      std::cout << *it << ";";
+    }
+    std::cout << *last << " " << kv.second << "\n";
+  }
+}
 }  // namespace
 
 int main(int argc, char **argv) {
@@ -104,14 +121,14 @@ int main(int argc, char **argv) {
     std::cerr << "PID " << pid << " is out of valid PID range.\n";
     return 1;
   }
+  buckets_t buckets;
   try {
     PtraceAttach(pid);
     Namespace ns(pid);
     const unsigned long tstate_addr = ThreadStateAddr(pid, &ns);
     if (seconds) {
       const std::chrono::microseconds interval{
           static_cast<long>(sample_rate * 1000000)};
-      std::unordered_map<frames_t, size_t, FrameHash> buckets;
       size_t idle = 0;
       auto end =
           std::chrono::system_clock::now() +
@@ -142,19 +159,7 @@ int main(int argc, char **argv) {
       if (idle) {
         std::cout << "(idle) " << idle << "\n";
       }
-      // process the frames
-      for (const auto &kv : buckets) {
-        if (kv.first.empty()) {
-          std::cerr << "uh oh\n";
-          return 1;
-        }
-        auto last = kv.first.rend();
-        last--;
-        for (auto it = kv.first.rbegin(); it != last; ++it) {
-          std::cout << *it << ";";
-        }
-        std::cout << *last << " " << kv.second << "\n";
-      }
+      PrintBuckets(buckets);
     } else {
       const unsigned long frame_addr = FirstFrameAddr(pid, tstate_addr);
       if (frame_addr) {
@@ -166,6 +171,15 @@ int main(int argc, char **argv) {
         std::cout << "(idle)\n";
       }
     }
+  } catch (const PtraceException &exc) {
+    // If the process terminates early then we just print the buckets up until
+    // that point in time.
+    if (!buckets.empty()) {
+      PrintBuckets(buckets);
+    } else {
+      std::cerr << exc.what() << std::endl;
+      return 1;
+    }
   } catch (const std::exception &exc) {
     std::cerr << exc.what() << std::endl;
     return 1;

tests/exit_early.py

@@ -0,0 +1,36 @@
+# Copyright 2016 Uber Technologies, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import time
+
+
+def main():
+    sys.stdout.write('%d\n' % (os.getpid(),))
+    sys.stdout.flush()
+    max_time = time.time() + 5
+    while True:
+        time.sleep(0.1)
+        target = time.time() + 0.1
+        while True:
+            now = time.time()
+            if now >= max_time:
+                return
+            if now >= target:
+                break
+
+
+if __name__ == '__main__':
+    main()

tests/test_end_to_end.py

@@ -49,6 +49,12 @@ def sleeper():
         yield p
 
 
+@pytest.yield_fixture
+def exit_early():
+    with proc('exit_early.py') as p:
+        yield p
+
+
 def test_monitor(dijkstra):
     """Basic test for the monitor mode."""
     proc = subprocess.Popen(['./src/pyflame', str(dijkstra.pid)],
@@ -94,3 +100,15 @@ def test_exclude_idle(sleeper):
     for line in lines:
         assert FLAMEGRAPH_RE.match(line) is not None
         assert not IDLE_RE.match(line)
+
+def test_exit_early(exit_early):
+    proc = subprocess.Popen(['./src/pyflame', '-s', '10', str(exit_early.pid)],
+                            stdout=subprocess.PIPE,
+                            stderr=subprocess.PIPE)
+    out, err = proc.communicate()
+    assert not err
+    assert proc.returncode == 0
+    lines = out.split('\n')
+    assert lines.pop(-1) == ''  # output should end in a newline
+    for line in lines:
+        assert FLAMEGRAPH_RE.match(line) or IDLE_RE.match(line)