From d852640a1e1067b34cf8312c0107592814601f6a Mon Sep 17 00:00:00 2001
From: Gautham <41098605+ahgamut@users.noreply.github.com>
Date: Wed, 13 Oct 2021 23:30:25 +0530
Subject: [PATCH] Add Python ftrace contextmanager (#285)

---
 libc/runtime/ftracer.c                  |  4 +-
 libc/runtime/runtime.h                  |  1 +
 third_party/python/Python/cosmomodule.c | 97 ++++++++++++++++++++-----
 3 files changed, 82 insertions(+), 20 deletions(-)

diff --git a/libc/runtime/ftracer.c b/libc/runtime/ftracer.c
index 790c95e99..b8c15658f 100644
--- a/libc/runtime/ftracer.c
+++ b/libc/runtime/ftracer.c
@@ -43,6 +43,7 @@
 
 void ftrace_hook(void);
 
+bool ftrace_enabled;
 static int g_skew;
 static int g_lastsymbol;
 static uint64_t laststamp;
@@ -80,7 +81,7 @@ privileged noinstrument noasan void ftracer(void) {
   static bool noreentry;
   struct StackFrame *frame;
   if (!cmpxchg(&noreentry, 0, 1)) return;
-  if (g_symbols) {
+  if (ftrace_enabled && g_symbols) {
     stamp = rdtsc();
     frame = __builtin_frame_address(0);
     frame = frame->next;
@@ -103,6 +104,7 @@ textstartup void ftrace_install(void) {
       laststamp = kStartTsc;
       g_lastsymbol = -1;
       g_skew = GetNestingLevelImpl(__builtin_frame_address(0));
+      ftrace_enabled = 1;
       __hook(ftrace_hook, g_symbols);
     } else {
       __printf("error: --ftrace failed to open symbol table\r\n");
diff --git a/libc/runtime/runtime.h b/libc/runtime/runtime.h
index 9875b2caa..c9c5344b4 100644
--- a/libc/runtime/runtime.h
+++ b/libc/runtime/runtime.h
@@ -34,6 +34,7 @@ extern unsigned char *__relo_start[];               /* αpε */
 extern unsigned char *__relo_end[];                 /* αpε */
 extern uint8_t __zip_start[];                       /* αpε */
 extern uint8_t __zip_end[];                         /* αpε */
+extern bool ftrace_enabled;
 
 void mcount(void);
 unsigned long getauxval(unsigned long);
diff --git a/third_party/python/Python/cosmomodule.c b/third_party/python/Python/cosmomodule.c
index 73b8ff439..6b21a2633 100644
--- a/third_party/python/Python/cosmomodule.c
+++ b/third_party/python/Python/cosmomodule.c
@@ -33,9 +33,11 @@
 #include "third_party/python/Include/methodobject.h"
 #include "third_party/python/Include/modsupport.h"
 #include "third_party/python/Include/moduleobject.h"
+#include "third_party/python/Include/object.h"
 #include "third_party/python/Include/pyerrors.h"
 #include "third_party/python/Include/pymacro.h"
 #include "third_party/python/Include/pyport.h"
+#include "third_party/python/Include/structmember.h"
 #include "third_party/python/Include/yoink.h"
 #include "third_party/xed/x86.h"
 /* clang-format off */
@@ -114,25 +116,6 @@ cosmo_getcpunode(PyObject *self, PyObject *noargs)
     return PyLong_FromUnsignedLong(TSC_AUX_NODE(rdpid()));
 }
 
-PyDoc_STRVAR(ftrace_doc,
-"ftrace($module)\n\
---\n\n\
-Enables logging of C function calls to stderr, e.g.\n\
-\n\
-    cosmo.ftrace()\n\
-    WeirdFunction()\n\
-    cosmo.exit1()\n\
-\n\
-Please be warned this prints massive amount of text. In order for it\n\
-to work, the concomitant .com.dbg binary needs to be present.");
-
-static PyObject *
-cosmo_ftrace(PyObject *self, PyObject *noargs)
-{
-    ftrace_install();
-    Py_RETURN_NONE;
-}
-
 PyDoc_STRVAR(crc32c_doc,
 "crc32c($module, bytes, init=0)\n\
 --\n\n\
@@ -218,6 +201,78 @@ cosmo_exit1(PyObject *self, PyObject *args)
     _Exit(1);
 }
 
+static bool ftrace_installed = 0;
+
+typedef struct {
+    PyObject_HEAD
+} TracerObject;
+
+static int TracerObject_init(PyObject* self, PyObject *args, PyObject *kwargs)
+{
+    if (!ftrace_installed) {
+        ftrace_install();
+        ftrace_installed = 1;
+        ftrace_enabled = 0;
+    }
+    return 0;
+}
+
+static PyObject* TracerObject_enter(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    ftrace_enabled = 1;
+    return self;
+}
+
+static PyObject* TracerObject_exit(PyObject *self, PyObject *args)
+{
+    ftrace_enabled = 0;
+    return self;
+}
+
+static PyMethodDef TracerObject_methods[] = {
+    {"__enter__", (PyCFunction) TracerObject_enter, METH_NOARGS,
+     "enable ftrace to start logging"
+    },
+    {"__exit__", (PyCFunction) TracerObject_exit, METH_VARARGS,
+     "disable ftrace to stop logging"
+    },
+    {0}
+};
+
+static PyTypeObject TracerType = {
+    PyVarObject_HEAD_INIT(NULL, 0)
+    .tp_name = "cosmo.Tracer",
+    .tp_doc = "wrapping ftrace with a context manager",
+    .tp_basicsize = sizeof(TracerObject),
+    .tp_itemsize = 0,
+    .tp_flags = Py_TPFLAGS_DEFAULT,
+    .tp_new = PyType_GenericNew,
+    .tp_init = (initproc) TracerObject_init,
+    .tp_methods = TracerObject_methods,
+};
+
+PyDoc_STRVAR(ftrace_doc,
+"ftrace($module)\n\
+--\n\n\
+Enables logging of C function calls to stderr, e.g.\n\
+\n\
+    with cosmo.ftrace() as F:\n\
+        WeirdFunction()\n\
+\n\
+Please be warned this prints massive amount of text. In order for it\n\
+to work, the concomitant .com.dbg binary needs to be present.");
+
+static PyObject *
+cosmo_ftrace(PyObject *self, PyObject *noargs)
+{
+    PyObject *tracer = PyType_GenericNew(&TracerType, NULL, NULL);
+    if (tracer == NULL) Py_RETURN_NONE;
+    TracerObject_init(tracer, NULL, NULL);
+    return tracer;
+}
+
+
+
 static PyMethodDef cosmo_methods[] = {
     {"exit1", cosmo_exit1, METH_NOARGS, exit1_doc},
     {"rdtsc", cosmo_rdtsc, METH_NOARGS, rdtsc_doc},
@@ -266,11 +321,15 @@ PyMODINIT_FUNC
 PyInit_cosmo(void)
 {
     PyObject *m;
+    if (PyType_Ready(&TracerType) < 0) return 0;
     if (!(m = PyModule_Create(&cosmomodule))) return 0;
     PyModule_AddStringConstant(m, "MODE", MODE);
     PyModule_AddIntConstant(m, "IMAGE_BASE_VIRTUAL", IMAGE_BASE_VIRTUAL);
     PyModule_AddStringConstant(m, "kernel", GetKernelName());
     PyModule_AddIntConstant(m, "kStartTsc", kStartTsc);
+
+    Py_INCREF(&TracerType);
+    // PyModule_AddObject(m, "Tracer", (PyObject *) &TracerType);
     return !PyErr_Occurred() ? m : 0;
 }