Calculate the padding automatically

Author: pablogsal

Python/perf_jit_trampoline.c

@@ -101,10 +101,9 @@
  * /tmp/jitted-PID-0.so: [headers][.text][unwind_info][padding]
  * /tmp/jitted-PID-1.so:                                       [headers][.text][unwind_info][padding]
  *
- * The padding size (0x100) is chosen to accommodate typical unwind info sizes
- * while maintaining 16-byte alignment requirements.
+ * The padding size is now calculated automatically during initialization
+ * based on the actual unwind information requirements.
  */
-#define PERF_JIT_CODE_PADDING 0x100
 
 
 /* These constants are defined inside <elf.h>, which we can't use outside of linux. */
@@ -666,6 +665,8 @@ static void elfctx_append_uleb128(ELFObjectContext* ctx, uint32_t v) {
 //                              DWARF EH FRAME GENERATION
 // =============================================================================
 
+static void elf_init_ehframe(ELFObjectContext* ctx);
+
 /*
  * Initialize DWARF .eh_frame section for a code region
  *
@@ -680,6 +681,23 @@ static void elfctx_append_uleb128(ELFObjectContext* ctx, uint32_t v) {
  * Args:
  *   ctx: ELF object context containing code size and buffer pointers
  */
+static size_t calculate_eh_frame_size(void) {
+    /* Calculate the EH frame size for the trampoline function */
+    extern void *_Py_trampoline_func_start;
+    extern void *_Py_trampoline_func_end;
+
+    size_t code_size = (char*)&_Py_trampoline_func_end - (char*)&_Py_trampoline_func_start;
+
+    ELFObjectContext ctx;
+    char buffer[1024];  // Buffer for DWARF data (1KB should be sufficient)
+    ctx.code_size = code_size;
+    ctx.startp = ctx.p = (uint8_t*)buffer;
+    ctx.fde_p = NULL;
+
+    elf_init_ehframe(&ctx);
+    return ctx.p - ctx.startp;
+}
+
 static void elf_init_ehframe(ELFObjectContext* ctx) {
     uint8_t* p = ctx->p;
     uint8_t* framep = p;  // Remember start of frame data
@@ -876,7 +894,7 @@ static void elf_init_ehframe(ELFObjectContext* ctx) {
      *
      * The FDE describes unwinding information specific to this function.
      * It references the CIE and provides function-specific CFI instructions.
-     * 
+     *
      * The PC-relative offset is calculated after the entire EH frame is built
      * to ensure accurate positioning relative to the synthesized DSO layout.
      */
@@ -901,16 +919,16 @@ static void elf_init_ehframe(ELFObjectContext* ctx) {
 #  endif
         DWRF_U8(DWRF_CFA_advance_loc | 1);    // Advance past push %rbp (1 byte)
         DWRF_U8(DWRF_CFA_def_cfa_offset);     // def_cfa_offset 16
-        DWRF_UV(16);
+        DWRF_UV(16);                          // New offset: SP + 16
         DWRF_U8(DWRF_CFA_offset | DWRF_REG_BP); // offset r6 at cfa-16
-        DWRF_UV(2);
+        DWRF_UV(2);                           // Offset factor: 2 * 8 = 16 bytes
         DWRF_U8(DWRF_CFA_advance_loc | 3);    // Advance past mov %rsp,%rbp (3 bytes)
         DWRF_U8(DWRF_CFA_def_cfa_register);   // def_cfa_register r6
-        DWRF_UV(DWRF_REG_BP);
+        DWRF_UV(DWRF_REG_BP);                 // Use base pointer register
         DWRF_U8(DWRF_CFA_advance_loc | 3);    // Advance past call *%rcx (2 bytes) + pop %rbp (1 byte) = 3
         DWRF_U8(DWRF_CFA_def_cfa);            // def_cfa r7 ofs 8
-        DWRF_UV(DWRF_REG_SP);
-        DWRF_UV(8);
+        DWRF_UV(DWRF_REG_SP);                 // Use stack pointer register
+        DWRF_UV(8);                           // New offset: SP + 8
 #elif defined(__aarch64__) && defined(__AARCH64EL__) && !defined(__ILP32__)
         /* AArch64 calling convention unwinding rules */
         DWRF_U8(DWRF_CFA_advance_loc | 1);        // Advance location by 1 instruction (stp x29, x30)
@@ -934,11 +952,11 @@ static void elf_init_ehframe(ELFObjectContext* ctx) {
     )
 
     ctx->p = p;  // Update context pointer to end of generated data
-    
+
     /* Calculate and update the PC-relative offset in the FDE
-     * 
+     *
      * When perf processes the jitdump, it creates a synthesized DSO with this layout:
-     * 
+     *
      *     Synthesized DSO Memory Layout:
      *     ┌─────────────────────────────────────────────────────────────┐ < code_start
      *     │                        Code Section                         │
@@ -956,33 +974,33 @@ static void elf_init_ehframe(ELFObjectContext* ctx) {
      *     │  │ CFI Instructions...                                 │    │
      *     │  └─────────────────────────────────────────────────────┘    │
      *     ├─────────────────────────────────────────────────────────────┤ < reference_point
-     *     │                    EhFrameHeader                            │ 
+     *     │                    EhFrameHeader                            │
      *     │                 (navigation metadata)                       │
      *     └─────────────────────────────────────────────────────────────┘
-     * 
+     *
      * The PC offset field in the FDE must contain the distance from itself to code_start:
-     * 
+     *
      *   distance = code_start - fde_pc_field
-     *   
+     *
      * Where:
      *   fde_pc_field_location = reference_point - eh_frame_size + fde_offset_in_frame
      *   code_start_location = reference_point - eh_frame_size - round_up(code_size, 8)
-     *   
+     *
      * Therefore:
      *   distance = code_start_location - fde_pc_field_location
      *            = (ref - eh_frame_size - rounded_code_size) - (ref - eh_frame_size + fde_offset_in_frame)
      *            = -rounded_code_size - fde_offset_in_frame
      *            = -(round_up(code_size, 8) + fde_offset_in_frame)
      *
      * Note: fde_offset_in_frame is the offset from EH frame start to the PC offset field,
-     * 
+     *
      */
     if (ctx->fde_p != NULL) {
         int32_t fde_offset_in_frame = (ctx->fde_p - ctx->startp);
         int32_t rounded_code_size = round_up(ctx->code_size, 8);
         int32_t pc_relative_offset = -(rounded_code_size + fde_offset_in_frame);
-        
-        
+
+
         // Update the PC-relative offset in the FDE
         *(int32_t*)ctx->fde_p = pc_relative_offset;
     }
@@ -1091,8 +1109,10 @@ static void* perf_map_jit_init(void) {
     /* Initialize code ID counter */
     perf_jit_map_state.code_id = 0;
 
-    /* Configure trampoline API with padding information */
-    trampoline_api.code_padding = PERF_JIT_CODE_PADDING;
+    /* Calculate padding size based on actual unwind info requirements */
+    size_t eh_frame_size = calculate_eh_frame_size();
+    size_t unwind_data_size = sizeof(EhFrameHeader) + eh_frame_size;
+    trampoline_api.code_padding = round_up(unwind_data_size, 16);
 
     return &perf_jit_map_state;
 }
@@ -1200,7 +1220,7 @@ static void perf_map_jit_write_entry(void *state, const void *code_addr,
     ev2.unwind_data_size = sizeof(EhFrameHeader) + eh_frame_size;
 
     /* Verify we don't exceed our padding budget */
-    assert(ev2.unwind_data_size <= PERF_JIT_CODE_PADDING);
+    assert(ev2.unwind_data_size <= (uint64_t)trampoline_api.code_padding);
 
     ev2.eh_frame_hdr_size = sizeof(EhFrameHeader);
     ev2.mapped_size = round_up(ev2.unwind_data_size, 16);  // 16-byte alignment