From ebb738668324eed3ec837c2476317553e7e5b79c Mon Sep 17 00:00:00 2001
From: Andrew Haberlandt <ahaberla@andrew.cmu.edu>
Date: Tue, 17 Dec 2024 23:54:27 +0000
Subject: [PATCH] fixes: dynamorio suspend flag, TERM in tests

---
 pyda_core/pyda_core.c    | 13 +++++++------
 pyda_core/pyda_threads.c |  1 +
 pyda_core/tool.c         |  2 ++
 tests/run_tests.py       | 22 ++++++++++++----------
 4 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/pyda_core/pyda_core.c b/pyda_core/pyda_core.c
index 59e5ac4..10415f2 100644
--- a/pyda_core/pyda_core.c
+++ b/pyda_core/pyda_core.c
@@ -323,7 +323,7 @@ void pyda_break(pyda_thread *t) {
     // Hack to tell dynamorio that dr_flush_region on another thread is OK
     // here -- this is not REALLY safe per the docs but we use
     // dr_redirect_execution so we *should* always return to a valid fragment...
-    dr_mark_safe_to_suspend(dr_get_current_drcontext(), true);
+    // dr_mark_safe_to_suspend(dr_get_current_drcontext(), true);
 
     // here we wait for the python to signal
     pthread_mutex_lock(&t->mutex);
@@ -332,7 +332,7 @@ void pyda_break(pyda_thread *t) {
     while (!t->python_yielded)
         pthread_cond_wait(&t->resume_cond, &t->mutex);
 
-    dr_mark_safe_to_suspend(dr_get_current_drcontext(), false);
+    // dr_mark_safe_to_suspend(dr_get_current_drcontext(), false);
 
     t->python_yielded = 0;
     pthread_mutex_unlock(&t->mutex);
@@ -507,13 +507,14 @@ int pyda_check_run_until(pyda_process *proc, void *test_pc) {
 
 static void thread_prepare_for_python_entry(PyGILState_STATE *gstate, pyda_thread *t, void* pc) {
     if (gstate) {
-        // HACK: This is not allowed per the docs. We get away with this
-        // because we check later to see if any flushes occurred during this period
+        // HACK: dr_mark_safe_to_suspend is not allowed in a cleancall, per the docs.
+        // We sortof get away with this because we check later to see if any flushes
+        // occurred during this period, and force a dr_redirect_execution if they did.
         t->flush_ts = t->proc->flush_count;
 
-        dr_mark_safe_to_suspend(dr_get_current_drcontext(), true);
+        // dr_mark_safe_to_suspend(dr_get_current_drcontext(), true);
         *gstate = PyGILState_Ensure();
-        dr_mark_safe_to_suspend(dr_get_current_drcontext(), false);
+        // dr_mark_safe_to_suspend(dr_get_current_drcontext(), false);
     }
 
     void *drcontext = dr_get_current_drcontext();
diff --git a/pyda_core/pyda_threads.c b/pyda_core/pyda_threads.c
index 9e03953..6527542 100644
--- a/pyda_core/pyda_threads.c
+++ b/pyda_core/pyda_threads.c
@@ -141,6 +141,7 @@ static void client_thread_init(void *arg) {
     void *tls = python_thread_init(NULL);
     ts->start_routine(ts->arg);
     DEBUG_PRINTF("start_routine returned\n");
+    dr_client_thread_set_suspendable(true);
     dr_thread_free(dr_get_current_drcontext(), tls, sizeof(void*) * 130);
     dr_global_free(ts, sizeof(struct thread_start));
 }
diff --git a/pyda_core/tool.c b/pyda_core/tool.c
index 219f5ee..cefb7fd 100644
--- a/pyda_core/tool.c
+++ b/pyda_core/tool.c
@@ -530,6 +530,8 @@ void python_aux_thread(void *arg) {
 
     PyGILState_Release(gstate);
 
+    dr_client_thread_set_suspendable(true);
+
     DEBUG_PRINTF("python_aux_thread 4\n");
 
     t->python_exited = 1;
diff --git a/tests/run_tests.py b/tests/run_tests.py
index db39289..e181981 100644
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -50,16 +50,16 @@ def no_warnings_or_errors(stdout: bytes, stderr: bytes) -> bool:
     )),
 
     # tests whether we can handle a large number of threads that do not get waited on
-    # ("threads_nojoin", "thread_nojoin.c", "../examples/ltrace_multithreaded.py", RunOpts(), ExpectedResult(
-    #     retcode=0,
-    #     checkers=[
-    #         output_checker,
-    #         no_warnings_or_errors,
-    #         lambda o, e: o.count(b"malloc") > 15000,
-    #         lambda o, e: o.count(b"free") > 15000,
-    #         lambda o, e: all((o.count(f"[thread {i}]".encode('utf-8')) == 40 for i in range(2, 100))),
-    #     ]
-    # )),
+    ("threads_nojoin", "thread_nojoin.c", "../examples/ltrace_multithreaded.py", RunOpts(), ExpectedResult(
+        retcode=0,
+        checkers=[
+            output_checker,
+            no_warnings_or_errors,
+            lambda o, e: o.count(b"malloc") > 15000,
+            lambda o, e: o.count(b"free") > 15000,
+            lambda o, e: all((o.count(f"[thread {i}]".encode('utf-8')) == 40 for i in range(2, 100))),
+        ]
+    )),
 
     # hook throws an exception
     ("err_hook_throw", "thread_1000.c", "err_hook.py", RunOpts(), ExpectedResult(
@@ -318,6 +318,8 @@ def run_test(c_file, python_file, run_opts, expected_result, test_name, debug, n
         if run_opts.no_pty:
             env["PYDA_NO_PTY"] = "1"
 
+        env["TERM"] = "linux"
+
         for trial in range(ntrials):
             result_str = ""
             try: