From 7dd16b618a5e07445539fbd42bc347b952de2819 Mon Sep 17 00:00:00 2001
From: Timothy Palpant <palpant@dropbox.com>
Date: Wed, 7 Aug 2019 20:15:58 -0400
Subject: [PATCH] Flush buffers before reading in addresses

---
 src/_vmprof.c      | 21 ++++++++++++++++++++-
 src/vmprof_mt.c    | 32 +++++++++++++++++++++++++-------
 src/vmprof_mt.h    |  1 +
 vmprof/__init__.py | 17 ++++++++---------
 4 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/src/_vmprof.c b/src/_vmprof.c
index bdefdc28..cdadbfae 100644
--- a/src/_vmprof.c
+++ b/src/_vmprof.c
@@ -363,6 +363,23 @@ start_sampling(PyObject *module, PyObject *noargs)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+vmp_flush_buffers(PyObject *module, PyObject *noargs)
+{
+    int ret = 0;
+
+#ifdef VMPROF_UNIX
+    ret = flush_buffers(vmp_profile_fileno());
+    if (ret != 0) {
+        return PyLong_NEW(ret);
+    }
+
+    ret = fsync(fd);
+#endif
+
+    return PyLong_NEW(ret);
+}
+
 #ifdef VMPROF_UNIX
 static PyObject * vmp_get_profile_path(PyObject *module, PyObject *noargs) {
     PyObject * o;
@@ -437,7 +454,9 @@ static PyMethodDef VMProfMethods[] = {
     {"stop_sampling", stop_sampling, METH_NOARGS,
         "Blocks signals to occur and returns the file descriptor"},
     {"start_sampling", start_sampling, METH_NOARGS,
-        "Unblocks vmprof signals. After compeltion vmprof will sample again"},
+        "Unblocks vmprof signals. After completion vmprof will sample again"},
+    {"flush_buffers", vmp_flush_buffers, METH_NOARGS,
+        "Flushes all pending buffers to the file descriptor"},
 #ifdef VMP_SUPPORTS_NATIVE_PROFILING
     {"resolve_addr", resolve_addr, METH_VARARGS,
         "Returns the name of the given address"},
diff --git a/src/vmprof_mt.c b/src/vmprof_mt.c
index f43276dc..906c4235 100644
--- a/src/vmprof_mt.c
+++ b/src/vmprof_mt.c
@@ -161,21 +161,39 @@ void cancel_buffer(struct profbuf_s *buf)
     profbuf_state[i] = PROFBUF_UNUSED;
 }
 
-int shutdown_concurrent_bufs(int fd)
+int flush_buffers(int fd)
 {
-    /* no signal handler can be running concurrently here, because we
-       already did vmprof_ignore_signals(1) */
-    assert(profbuf_write_lock == 0);
-    profbuf_write_lock = 2;
+    if (!__sync_bool_compare_and_swap(&profbuf_write_lock, 0, 1)) {
+        return -1;
+    }
 
-    /* last attempt to flush buffers */
     int i;
     for (i = 0; i < MAX_NUM_BUFFERS; i++) {
         while (profbuf_state[i] == PROFBUF_READY) {
-            if (_write_single_ready_buffer(fd, i) < 0)
+            if (_write_single_ready_buffer(fd, i) < 0) {
+                profbuf_write_lock = 0;
                 return -1;
+            }
         }
     }
+
+    profbuf_write_lock = 0;
+    return 0;
+}
+
+int shutdown_concurrent_bufs(int fd)
+{
+    /* no signal handler can be running concurrently here, because we
+       already did vmprof_ignore_signals(1) */
+    assert(profbuf_write_lock == 0);
+
+    /* last attempt to flush buffers */
+    int flushed = flush_buffers(fd);
+    if (flushed != 0) {
+      return flushed;
+    }
+
+    profbuf_write_lock = 2;
     unprepare_concurrent_bufs();
     return 0;
 }
diff --git a/src/vmprof_mt.h b/src/vmprof_mt.h
index 7113a0d7..726a9a0e 100644
--- a/src/vmprof_mt.h
+++ b/src/vmprof_mt.h
@@ -48,4 +48,5 @@ int prepare_concurrent_bufs(void);
 struct profbuf_s *reserve_buffer(int fd);
 void commit_buffer(int fd, struct profbuf_s *buf);
 void cancel_buffer(struct profbuf_s *buf);
+int flush_buffers(int fd);
 int shutdown_concurrent_bufs(int fd);
diff --git a/vmprof/__init__.py b/vmprof/__init__.py
index 1a9c7a58..f1210122 100644
--- a/vmprof/__init__.py
+++ b/vmprof/__init__.py
@@ -29,15 +29,14 @@
 def disable():
     try:
         # fish the file descriptor that is still open!
-        if hasattr(_vmprof, 'stop_sampling'):
-            fileno = _vmprof.stop_sampling()
-            if fileno >= 0:
-                # TODO does fileobj leak the fd? I dont think so, but need to check 
-                fileobj = FdWrapper(fileno)
-                l = LogReaderDumpNative(fileobj, LogReaderState())
-                l.read_all()
-                if hasattr(_vmprof, 'write_all_code_objects'):
-                    _vmprof.write_all_code_objects(l.dedup)
+        fileno = _vmprof.stop_sampling()
+        if fileno >= 0:
+            _vmprof.flush_buffers()
+            # TODO does fileobj leak the fd? I dont think so, but need to check 
+            fileobj = FdWrapper(fileno)
+            l = LogReaderDumpNative(fileobj, LogReaderState())
+            l.read_all()
+            _vmprof.write_all_code_objects(l.dedup)
         _vmprof.disable()
     except IOError as e:
         raise Exception("Error while writing profile: " + str(e))