#include "Python.h"

#include "pycore_ceval.h" // _PyPerf_Callbacks

#include "pycore_frame.h"

#include "pycore_interp.h"

#ifdef PY_HAVE_PERF_TRAMPOLINE

#include <fcntl.h>

#include <stdio.h>

#include <stdlib.h>

#include <sys/mman.h> // mmap()

#include <sys/types.h>

#include <unistd.h> // sysconf()

#include <sys/time.h> // gettimeofday()

#if defined(__arm__) || defined(__arm64__) || defined(__aarch64__)

#define PY_HAVE_INVALIDATE_ICACHE

#if defined(__clang__) || defined(__GNUC__)

extern void __clear_cache(void *, void*);

static void invalidate_icache(char* begin, char*end) {

#if defined(__clang__) || defined(__GNUC__)

return __clear_cache(begin, end);

return;

}

typedef PyObject *(*py_evaluator)(PyThreadState *, _PyInterpreterFrame *,

int throwflag);

typedef PyObject *(*py_trampoline)(PyThreadState *, _PyInterpreterFrame *, int,

py_evaluator);

extern void *_Py_trampoline_func_start; // Start of the template of the

// assembly trampoline

extern void *

_Py_trampoline_func_end; // End of the template of the assembly trampoline

struct code_arena_st {

};

#define CODE_ALIGNMENT 32

typedef struct code_arena_st code_arena_t;

typedef struct trampoline_api_st trampoline_api_t;

enum perf_trampoline_type {

};

#define perf_status _PyRuntime.ceval.perf.status

#define extra_code_index _PyRuntime.ceval.perf.extra_code_index

#define perf_code_arena _PyRuntime.ceval.perf.code_arena

#define trampoline_api _PyRuntime.ceval.perf.trampoline_api

#define perf_map_file _PyRuntime.ceval.perf.map_file

#define persist_after_fork _PyRuntime.ceval.perf.persist_after_fork

#define perf_trampoline_type _PyRuntime.ceval.perf.perf_trampoline_type

#define prev_eval_frame _PyRuntime.ceval.perf.prev_eval_frame

#define trampoline_refcount _PyRuntime.ceval.perf.trampoline_refcount

#define code_watcher_id _PyRuntime.ceval.perf.code_watcher_id

static void free_code_arenas(void);

static void

perf_trampoline_clear_code_watcher(void)

{

if (code_watcher_id >= 0) {

PyCode_ClearWatcher(code_watcher_id);

code_watcher_id = -1;

}

extra_code_index = -1;

}

static void

perf_trampoline_reset_state(void)

{

free_code_arenas();

perf_trampoline_clear_code_watcher();

}

static int

perf_trampoline_code_watcher(PyCodeEvent event, PyCodeObject *co)

{

if (event != PY_CODE_EVENT_DESTROY) {

return 0;

}

if (extra_code_index == -1) {

return 0;

}

py_trampoline f = NULL;

int ret = _PyCode_GetExtra((PyObject *)co, extra_code_index, (void **)&f);

if (ret != 0 || f == NULL) {

return 0;

}

trampoline_refcount--;

if (trampoline_refcount == 0) {

perf_trampoline_reset_state();

}

return 0;

}

static void

perf_map_write_entry(void *state, const void *code_addr,

unsigned int code_size, PyCodeObject *co)

{

const char *entry = "";

if (co->co_qualname != NULL) {

entry = PyUnicode_AsUTF8(co->co_qualname);

if (entry == NULL) {

PyErr_Clear();

entry = "";

}

}

const char *filename = "";

if (co->co_filename != NULL) {

filename = PyUnicode_AsUTF8(co->co_filename);

if (filename == NULL) {

PyErr_Clear();

filename = "";

}

}

size_t perf_map_entry_size = snprintf(NULL, 0, "py::%s:%s", entry, filename) + 1;

char* perf_map_entry = (char*) PyMem_RawMalloc(perf_map_entry_size);

if (perf_map_entry == NULL) {

return;

}

snprintf(perf_map_entry, perf_map_entry_size, "py::%s:%s", entry, filename);

PyUnstable_WritePerfMapEntry(code_addr, code_size, perf_map_entry);

PyMem_RawFree(perf_map_entry);

}

static void*

perf_map_init_state(void)

{

PyUnstable_PerfMapState_Init();

trampoline_api.code_padding = 0;

perf_trampoline_type = PERF_TRAMPOLINE_TYPE_MAP;

return NULL;

}

static int

perf_map_free_state(void *state)

{

PyUnstable_PerfMapState_Fini();

return 0;

}

_PyPerf_Callbacks _Py_perfmap_callbacks = {

&perf_map_init_state,

&perf_map_write_entry,

&perf_map_free_state,

};

static size_t round_up(int64_t value, int64_t multiple) {

if (multiple == 0) {

// Avoid division by zero

return value;

}

int64_t remainder = value % multiple;

if (remainder == 0) {

// Value is already a multiple of 'multiple'

return value;

}

// Calculate the difference to the next multiple

int64_t difference = multiple - remainder;

// Add the difference to the value

int64_t rounded_up_value = value + difference;

return rounded_up_value;

}

static int

new_code_arena(void)

{

// non-trivial programs typically need 64 to 256 kiB.

size_t mem_size = 4096 * 16;

assert(mem_size % sysconf(_SC_PAGESIZE) == 0);

char *memory =

mmap(NULL, // address

mem_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS,

-1, // fd (not used here)

0); // offset (not used here)

if (memory == MAP_FAILED) {

PyErr_SetFromErrno(PyExc_OSError);

PyErr_FormatUnraisable("Failed to create new mmap for perf trampoline");

perf_status = PERF_STATUS_FAILED;

return -1;

}

void *start = &_Py_trampoline_func_start;

void *end = &_Py_trampoline_func_end;

size_t code_size = end - start;

size_t unaligned_size = code_size + trampoline_api.code_padding;

size_t chunk_size = round_up(unaligned_size, CODE_ALIGNMENT);

assert(chunk_size % CODE_ALIGNMENT == 0);

// TODO: Check the effect of alignment of the code chunks. Initial investigation

// showed that this has no effect on performance in x86-64 or aarch64 and the current

// version has the advantage that the unwinder in GDB can unwind across JIT-ed code.

//

// We should check the values in the future and see if there is a

// measurable performance improvement by rounding trampolines up to 32-bit

// or 64-bit alignment.

size_t n_copies = mem_size / chunk_size;

for (size_t i = 0; i < n_copies; i++) {

memcpy(memory + i * chunk_size, start, code_size * sizeof(char));

}

// Some systems may prevent us from creating executable code on the fly.

int res = mprotect(memory, mem_size, PROT_READ | PROT_EXEC);

if (res == -1) {

PyErr_SetFromErrno(PyExc_OSError);

munmap(memory, mem_size);

PyErr_FormatUnraisable("Failed to set mmap for perf trampoline to "

"PROT_READ | PROT_EXEC");

return -1;

}

#ifdef PY_HAVE_INVALIDATE_ICACHE

// Before the JIT can run a block of code that has been emitted it must invalidate

// the instruction cache on some platforms like arm and aarch64.

invalidate_icache(memory, memory + mem_size);

code_arena_t *new_arena = PyMem_RawCalloc(1, sizeof(code_arena_t));

if (new_arena == NULL) {

PyErr_NoMemory();

munmap(memory, mem_size);

PyErr_FormatUnraisable("Failed to allocate new code arena struct for perf trampoline");

return -1;

}

new_arena->start_addr = memory;

new_arena->current_addr = memory;

new_arena->size = mem_size;

new_arena->size_left = mem_size;

new_arena->code_size = code_size;

new_arena->prev = perf_code_arena;

perf_code_arena = new_arena;

return 0;

}

static void

free_code_arenas(void)

{

code_arena_t *cur = perf_code_arena;

code_arena_t *prev;

perf_code_arena = NULL; // invalid static pointer

while (cur) {

munmap(cur->start_addr, cur->size);

prev = cur->prev;

PyMem_RawFree(cur);

cur = prev;

}

}

static inline py_trampoline

code_arena_new_code(code_arena_t *code_arena)

{

py_trampoline trampoline = (py_trampoline)code_arena->current_addr;

size_t total_code_size = round_up(code_arena->code_size + trampoline_api.code_padding,

CODE_ALIGNMENT);

assert(total_code_size % CODE_ALIGNMENT == 0);

code_arena->size_left -= total_code_size;

code_arena->current_addr += total_code_size;

return trampoline;

}

static inline py_trampoline

compile_trampoline(void)

{

size_t total_code_size = round_up(perf_code_arena->code_size + trampoline_api.code_padding, 16);

if ((perf_code_arena == NULL) ||

(perf_code_arena->size_left <= total_code_size)) {

if (new_code_arena() < 0) {

return NULL;

}

}

assert(perf_code_arena->size_left <= perf_code_arena->size);

return code_arena_new_code(perf_code_arena);

}

static PyObject *

py_trampoline_evaluator(PyThreadState *ts, _PyInterpreterFrame *frame,

int throw)

{

if (perf_status == PERF_STATUS_FAILED ||

perf_status == PERF_STATUS_NO_INIT) {

goto default_eval;

}

PyCodeObject *co = _PyFrame_GetCode(frame);

py_trampoline f = NULL;

assert(extra_code_index != -1);

int ret = _PyCode_GetExtra((PyObject *)co, extra_code_index, (void **)&f);

if (ret != 0 || f == NULL) {

// This is the first time we see this code object so we need

// to compile a trampoline for it.

py_trampoline new_trampoline = compile_trampoline();

if (new_trampoline == NULL) {

goto default_eval;

}

trampoline_api.write_state(trampoline_api.state, new_trampoline,

perf_code_arena->code_size, co);

_PyCode_SetExtra((PyObject *)co, extra_code_index,

(void *)new_trampoline);

trampoline_refcount++;

f = new_trampoline;

}

assert(f != NULL);

return f(ts, frame, throw, _PyEval_EvalFrameDefault);

default_eval:

// Something failed, fall back to the default evaluator.

return _PyEval_EvalFrameDefault(ts, frame, throw);

}

#endif // PY_HAVE_PERF_TRAMPOLINE

int PyUnstable_PerfTrampoline_CompileCode(PyCodeObject *co)

{

#ifdef PY_HAVE_PERF_TRAMPOLINE

py_trampoline f = NULL;

assert(extra_code_index != -1);

int ret = _PyCode_GetExtra((PyObject *)co, extra_code_index, (void **)&f);

if (ret != 0 || f == NULL) {

py_trampoline new_trampoline = compile_trampoline();

if (new_trampoline == NULL) {

return 0;

}

trampoline_api.write_state(trampoline_api.state, new_trampoline,

perf_code_arena->code_size, co);

trampoline_refcount++;

return _PyCode_SetExtra((PyObject *)co, extra_code_index,

(void *)new_trampoline);

}

#endif // PY_HAVE_PERF_TRAMPOLINE

return 0;

}

_PyIsPerfTrampolineActive(void)

{

#ifdef PY_HAVE_PERF_TRAMPOLINE

PyThreadState *tstate = _PyThreadState_GET();

return tstate->interp->eval_frame == py_trampoline_evaluator;

return 0;

}

_PyPerfTrampoline_GetCallbacks(_PyPerf_Callbacks *callbacks)

{

if (callbacks == NULL) {

return;

}

#ifdef PY_HAVE_PERF_TRAMPOLINE

callbacks->init_state = trampoline_api.init_state;

callbacks->write_state = trampoline_api.write_state;

callbacks->free_state = trampoline_api.free_state;

return;

}

_PyPerfTrampoline_SetCallbacks(_PyPerf_Callbacks *callbacks)

{

if (callbacks == NULL) {

return -1;

}

#ifdef PY_HAVE_PERF_TRAMPOLINE

if (trampoline_api.state) {

_PyPerfTrampoline_Fini();

}

trampoline_api.init_state = callbacks->init_state;

trampoline_api.write_state = callbacks->write_state;

trampoline_api.free_state = callbacks->free_state;

trampoline_api.state = NULL;

return 0;

}

_PyPerfTrampoline_Init(int activate)

{

#ifdef PY_HAVE_PERF_TRAMPOLINE

PyThreadState *tstate = _PyThreadState_GET();

if (code_watcher_id == 0) {

// Initialize to -1 since 0 is a valid watcher ID

code_watcher_id = -1;

}

if (tstate->interp->eval_frame &&

tstate->interp->eval_frame != py_trampoline_evaluator) {

PyErr_SetString(PyExc_RuntimeError,

"Trampoline cannot be initialized as a custom eval "

"frame is already present");

return -1;

}

if (!activate) {

tstate->interp->eval_frame = NULL;

perf_status = PERF_STATUS_NO_INIT;

}

else {

tstate->interp->eval_frame = py_trampoline_evaluator;

extra_code_index = _PyEval_RequestCodeExtraIndex(NULL);

if (extra_code_index == -1) {

return -1;

}

if (trampoline_api.state == NULL && trampoline_api.init_state != NULL) {

trampoline_api.state = trampoline_api.init_state();

}

if (new_code_arena() < 0) {

return -1;

}

code_watcher_id = PyCode_AddWatcher(perf_trampoline_code_watcher);

if (code_watcher_id < 0) {

PyErr_FormatUnraisable("Failed to register code watcher for perf trampoline");

free_code_arenas();

return -1;

}

trampoline_refcount = 1; // Base refcount held by the system

perf_status = PERF_STATUS_OK;

}

return 0;

}

_PyPerfTrampoline_Fini(void)

{

#ifdef PY_HAVE_PERF_TRAMPOLINE

if (perf_status != PERF_STATUS_OK) {

return 0;

}

PyThreadState *tstate = _PyThreadState_GET();

if (tstate->interp->eval_frame == py_trampoline_evaluator) {

tstate->interp->eval_frame = NULL;

}

if (perf_status == PERF_STATUS_OK) {

trampoline_api.free_state(trampoline_api.state);

perf_trampoline_type = PERF_TRAMPOLINE_UNSET;

}

// Prevent new trampolines from being created

perf_status = PERF_STATUS_NO_INIT;

// Decrement base refcount. If refcount reaches 0, all code objects are already

// dead so clean up now. Otherwise, watcher remains active to clean up when last

// code object dies; extra_code_index stays valid so watcher can identify them.

trampoline_refcount--;

if (trampoline_refcount == 0) {

perf_trampoline_reset_state();

}

return 0;

}

PyUnstable_PerfTrampoline_SetPersistAfterFork(int enable){

#ifdef PY_HAVE_PERF_TRAMPOLINE

persist_after_fork = enable;

return persist_after_fork;

return 0;

}

_PyPerfTrampoline_AfterFork_Child(void)

{

#ifdef PY_HAVE_PERF_TRAMPOLINE

if (persist_after_fork) {

if (perf_trampoline_type != PERF_TRAMPOLINE_TYPE_MAP) {

return PyStatus_Error("Failed to copy perf map file as perf trampoline type is not type map.");

}

_PyPerfTrampoline_Fini();

char filename[256];

pid_t parent_pid = getppid();

snprintf(filename, sizeof(filename), "/tmp/perf-%d.map", parent_pid);

if (PyUnstable_CopyPerfMapFile(filename) != 0) {

return PyStatus_Error("Failed to copy perf map file.");

}

} else {

// Restart trampoline in file in child.

int was_active = _PyIsPerfTrampolineActive();

_PyPerfTrampoline_Fini();

if (was_active) {

// After fork, Fini may leave the old code watcher registered

// if trampolined code objects from the parent still exist

// (trampoline_refcount > 0). Clear it unconditionally before

// Init registers a new one, but keep the old arenas mapped: the

// child may still need to return through trampoline frames that

// were on the C stack at fork().

perf_trampoline_clear_code_watcher();

_PyPerfTrampoline_Init(1);

}

}

return PyStatus_Ok();

}