Stowage/cpython
2
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-12-31 12:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity
cpython/Modules/_remote_debugging/binary_io_writer.c
Pablo Galindo Salgado 7c44f37170
gh-138122: Extend binary profiling format with full source location and opcode (#143088) 
Co-authored-by: Stan Ulbrych <stan@ulbrych.org>
2025-12-24 16:15:11 +00:00

1242 lines
41 KiB
C
Raw Blame History

/******************************************************************************
* Python Remote Debugging Module - Binary Writer Implementation
*
* High-performance binary file writer for profiling data with optional zstd
* streaming compression.
******************************************************************************/
#ifndef Py_BUILD_CORE_MODULE
# define Py_BUILD_CORE_MODULE
#endif
#include "binary_io.h"
#include "_remote_debugging.h"
#include "pycore_opcode_utils.h" // MAX_REAL_OPCODE
#include <string.h>
#ifdef HAVE_ZSTD
#include <zstd.h>
#endif
/* ============================================================================
* CONSTANTS FOR BINARY FORMAT SIZES
* ============================================================================ */
/* Sample header sizes */
#define SAMPLE_HEADER_FIXED_SIZE 13 /* thread_id(8) + interpreter_id(4) + encoding(1) */
#define SAMPLE_HEADER_MAX_SIZE 26 /* fixed + max_varint(10) + status(1) + margin */
#define MAX_VARINT_SIZE 10 /* Maximum bytes for a varint64 */
#define MAX_VARINT_SIZE_U32 5 /* Maximum bytes for a varint32 */
/* Frame buffer: depth varint (max 2 bytes for 256) + 256 frames * 5 bytes/varint + margin */
#define MAX_FRAME_BUFFER_SIZE ((MAX_STACK_DEPTH * MAX_VARINT_SIZE_U32) + MAX_VARINT_SIZE_U32 + 16)
/* File structure sizes */
#define FILE_FOOTER_SIZE 32
/* Helper macro: convert PyLong to int32, using default_val if conversion fails */
#define PYLONG_TO_INT32_OR_DEFAULT(obj, var, default_val) \
do { \
(var) = (int32_t)PyLong_AsLong(obj); \
if (UNLIKELY(PyErr_Occurred() != NULL)) { \
PyErr_Clear(); \
(var) = (default_val); \
} \
} while (0)
/* ============================================================================
* WRITER-SPECIFIC UTILITY HELPERS
* ============================================================================ */
/* Grow two parallel arrays together (e.g., strings and string_lengths).
* Returns 0 on success, -1 on error (sets PyErr).
* On error, original arrays are preserved (truly atomic update). */
static inline int
grow_parallel_arrays(void **array1, void **array2, size_t *capacity,
size_t elem_size1, size_t elem_size2)
{
size_t old_cap = *capacity;
if (old_cap > SIZE_MAX / 2) {
PyErr_SetString(PyExc_OverflowError, "Array capacity overflow");
return -1;
}
size_t new_cap = old_cap * 2;
if (new_cap > SIZE_MAX / elem_size1 || new_cap > SIZE_MAX / elem_size2) {
PyErr_SetString(PyExc_OverflowError, "Array allocation size overflow");
return -1;
}
size_t new_size1 = new_cap * elem_size1;
size_t new_size2 = new_cap * elem_size2;
size_t old_size1 = old_cap * elem_size1;
size_t old_size2 = old_cap * elem_size2;
/* Allocate fresh memory blocks (not realloc) to ensure atomicity.
* If either allocation fails, original arrays are completely unchanged. */
void *new_array1 = PyMem_Malloc(new_size1);
if (!new_array1) {
PyErr_NoMemory();
return -1;
}
void *new_array2 = PyMem_Malloc(new_size2);
if (!new_array2) {
/* Second allocation failed - free first and return with no state change */
PyMem_Free(new_array1);
PyErr_NoMemory();
return -1;
}
/* Both allocations succeeded - copy data and update pointers atomically */
memcpy(new_array1, *array1, old_size1);
memcpy(new_array2, *array2, old_size2);
PyMem_Free(*array1);
PyMem_Free(*array2);
*array1 = new_array1;
*array2 = new_array2;
*capacity = new_cap;
return 0;
}
/* Checked fwrite with GIL release - returns 0 on success, -1 on error (sets PyErr).
* This version releases the GIL during the write operation to allow other Python
* threads to run during potentially blocking I/O. */
static inline int
fwrite_checked_allow_threads(const void *data, size_t size, FILE *fp)
{
size_t written;
Py_BEGIN_ALLOW_THREADS
written = fwrite(data, 1, size, fp);
Py_END_ALLOW_THREADS
if (written != size) {
PyErr_SetFromErrno(PyExc_IOError);
return -1;
}
return 0;
}
/* Forward declaration for writer_write_bytes */
static inline int writer_write_bytes(BinaryWriter *writer, const void *data, size_t size);
/* Encode and write a varint u32 - returns 0 on success, -1 on error */
static inline int
writer_write_varint_u32(BinaryWriter *writer, uint32_t value)
{
uint8_t buf[MAX_VARINT_SIZE];
size_t len = encode_varint_u32(buf, value);
return writer_write_bytes(writer, buf, len);
}
/* Encode and write a varint u64 - returns 0 on success, -1 on error */
static inline int
writer_write_varint_u64(BinaryWriter *writer, uint64_t value)
{
uint8_t buf[MAX_VARINT_SIZE];
size_t len = encode_varint_u64(buf, value);
return writer_write_bytes(writer, buf, len);
}
/* ============================================================================
* UTILITY FUNCTIONS
* ============================================================================ */
int
binary_io_zstd_available(void)
{
#ifdef HAVE_ZSTD
return 1;
#else
return 0;
#endif
}
int
binary_io_get_best_compression(void)
{
#ifdef HAVE_ZSTD
return COMPRESSION_ZSTD;
#else
return COMPRESSION_NONE;
#endif
}
/* ============================================================================
* BINARY WRITER IMPLEMENTATION
* ============================================================================ */
static int
writer_init_zstd(BinaryWriter *writer)
{
#ifdef HAVE_ZSTD
writer->zstd.cctx = ZSTD_createCCtx();
if (!writer->zstd.cctx) {
PyErr_SetString(PyExc_MemoryError, "Failed to create zstd compression context");
return -1;
}
/* Compression level 5: better ratio for repetitive profiling data */
size_t result = ZSTD_CCtx_setParameter(writer->zstd.cctx,
ZSTD_c_compressionLevel, 5);
if (ZSTD_isError(result)) {
PyErr_Format(PyExc_RuntimeError, "Failed to set zstd compression level: %s",
ZSTD_getErrorName(result));
ZSTD_freeCCtx(writer->zstd.cctx);
writer->zstd.cctx = NULL;
return -1;
}
/* Use large buffer (512KB) for fewer I/O syscalls */
writer->zstd.compressed_buffer = PyMem_Malloc(COMPRESSED_BUFFER_SIZE);
if (!writer->zstd.compressed_buffer) {
ZSTD_freeCCtx(writer->zstd.cctx);
writer->zstd.cctx = NULL;
PyErr_NoMemory();
return -1;
}
writer->zstd.compressed_buffer_size = COMPRESSED_BUFFER_SIZE;
return 0;
#else
PyErr_SetString(PyExc_RuntimeError,
"zstd compression requested but not available (HAVE_ZSTD not defined)");
return -1;
#endif
}
static int
writer_flush_buffer(BinaryWriter *writer)
{
if (writer->buffer_pos == 0) {
return 0;
}
#ifdef HAVE_ZSTD
if (writer->compression_type == COMPRESSION_ZSTD) {
ZSTD_inBuffer input = { writer->write_buffer, writer->buffer_pos, 0 };
while (input.pos < input.size) {
ZSTD_outBuffer output = {
writer->zstd.compressed_buffer,
writer->zstd.compressed_buffer_size,
0
};
size_t result = ZSTD_compressStream2(
writer->zstd.cctx, &output, &input, ZSTD_e_continue
);
if (ZSTD_isError(result)) {
PyErr_Format(PyExc_IOError, "zstd compression error: %s",
ZSTD_getErrorName(result));
return -1;
}
if (output.pos > 0) {
if (fwrite_checked_allow_threads(writer->zstd.compressed_buffer, output.pos, writer->fp) < 0) {
return -1;
}
}
}
} else
#endif
{
if (fwrite_checked_allow_threads(writer->write_buffer, writer->buffer_pos, writer->fp) < 0) {
return -1;
}
}
writer->buffer_pos = 0;
return 0;
}
static inline int
writer_write_bytes(BinaryWriter *writer, const void *data, size_t size)
{
const uint8_t *src = (const uint8_t *)data;
size_t original_size = size;
while (size > 0) {
size_t space = writer->buffer_size - writer->buffer_pos;
size_t to_copy = (size < space) ? size : space;
memcpy(writer->write_buffer + writer->buffer_pos, src, to_copy);
writer->buffer_pos += to_copy;
src += to_copy;
size -= to_copy;
if (writer->buffer_pos == writer->buffer_size) {
if (writer_flush_buffer(writer) < 0) {
return -1;
}
}
}
writer->stats.bytes_written += original_size;
return 0;
}
/* ============================================================================
* HASH TABLE SUPPORT FUNCTIONS (using _Py_hashtable)
* ============================================================================ */
/* Hash function for Python strings - uses Python's cached hash */
static Py_uhash_t
string_hash_func(const void *key)
{
PyObject *str = (PyObject *)key;
Py_hash_t hash = PyObject_Hash(str);
if (hash == -1) {
PyErr_Clear();
return 0;
}
return (Py_uhash_t)hash;
}
static int
string_compare_func(const void *key1, const void *key2)
{
PyObject *str1 = (PyObject *)key1;
PyObject *str2 = (PyObject *)key2;
if (str1 == str2) {
return 1;
}
int result = PyObject_RichCompareBool(str1, str2, Py_EQ);
if (result == -1) {
PyErr_Clear();
return 0;
}
return result;
}
static void
string_key_destroy(void *key)
{
Py_XDECREF((PyObject *)key);
}
static Py_uhash_t
frame_key_hash_func(const void *key)
{
const FrameKey *fk = (const FrameKey *)key;
/* FNV-1a style hash combining all fields */
Py_uhash_t hash = 2166136261u;
hash ^= fk->filename_idx;
hash *= 16777619u;
hash ^= fk->funcname_idx;
hash *= 16777619u;
hash ^= (uint32_t)fk->lineno;
hash *= 16777619u;
hash ^= (uint32_t)fk->end_lineno;
hash *= 16777619u;
hash ^= (uint32_t)fk->column;
hash *= 16777619u;
hash ^= (uint32_t)fk->end_column;
hash *= 16777619u;
hash ^= fk->opcode;
hash *= 16777619u;
return hash;
}
static int
frame_key_compare_func(const void *key1, const void *key2)
{
const FrameKey *fk1 = (const FrameKey *)key1;
const FrameKey *fk2 = (const FrameKey *)key2;
return (fk1->filename_idx == fk2->filename_idx &&
fk1->funcname_idx == fk2->funcname_idx &&
fk1->lineno == fk2->lineno &&
fk1->end_lineno == fk2->end_lineno &&
fk1->column == fk2->column &&
fk1->end_column == fk2->end_column &&
fk1->opcode == fk2->opcode);
}
static void
frame_key_destroy(void *key)
{
PyMem_Free(key);
}
static inline int
writer_intern_string(BinaryWriter *writer, PyObject *string, uint32_t *index)
{
void *existing = _Py_hashtable_get(writer->string_hash, string);
if (existing != NULL) {
*index = (uint32_t)(uintptr_t)existing - 1; /* index+1 stored to distinguish from NULL */
return 0;
}
if (writer->string_count >= writer->string_capacity) {
if (grow_parallel_arrays((void **)&writer->strings,
(void **)&writer->string_lengths,
&writer->string_capacity,
sizeof(char *), sizeof(size_t)) < 0) {
return -1;
}
}
Py_ssize_t str_len;
const char *str_data = PyUnicode_AsUTF8AndSize(string, &str_len);
if (!str_data) {
return -1;
}
char *str_copy = PyMem_Malloc(str_len + 1);
if (!str_copy) {
PyErr_NoMemory();
return -1;
}
memcpy(str_copy, str_data, str_len + 1);
*index = (uint32_t)writer->string_count;
/* Add to hash table FIRST to ensure atomic rollback on failure */
Py_INCREF(string);
if (_Py_hashtable_set(writer->string_hash, string, (void *)(uintptr_t)(*index + 1)) < 0) {
Py_DECREF(string);
PyMem_Free(str_copy);
PyErr_NoMemory();
return -1;
}
writer->strings[writer->string_count] = str_copy;
writer->string_lengths[writer->string_count] = str_len;
writer->string_count++;
return 0;
}
static inline int
writer_intern_frame(BinaryWriter *writer, const FrameEntry *entry, uint32_t *index)
{
FrameKey lookup_key = {
entry->filename_idx, entry->funcname_idx,
entry->lineno, entry->end_lineno,
entry->column, entry->end_column,
entry->opcode
};
void *existing = _Py_hashtable_get(writer->frame_hash, &lookup_key);
if (existing != NULL) {
*index = (uint32_t)(uintptr_t)existing - 1; /* index+1 stored to distinguish from NULL */
return 0;
}
if (GROW_ARRAY(writer->frame_entries, writer->frame_count,
writer->frame_capacity, FrameEntry) < 0) {
return -1;
}
FrameKey *key = PyMem_Malloc(sizeof(FrameKey));
if (!key) {
PyErr_NoMemory();
return -1;
}
*key = lookup_key;
*index = (uint32_t)writer->frame_count;
writer->frame_entries[writer->frame_count] = *entry;
if (_Py_hashtable_set(writer->frame_hash, key, (void *)(uintptr_t)(*index + 1)) < 0) {
PyMem_Free(key);
PyErr_NoMemory();
return -1;
}
writer->frame_count++;
return 0;
}
/* Get or create a thread entry for the given thread_id.
* Returns pointer to ThreadEntry, or NULL on allocation failure.
* If is_new is non-NULL, sets it to 1 if this is a new thread, 0 otherwise. */
static ThreadEntry *
writer_get_or_create_thread_entry(BinaryWriter *writer, uint64_t thread_id,
uint32_t interpreter_id, int *is_new)
{
/* Linear search is OK for small number of threads.
* Key is (thread_id, interpreter_id) since same thread_id can exist in different interpreters. */
for (size_t i = 0; i < writer->thread_count; i++) {
if (writer->thread_entries[i].thread_id == thread_id &&
writer->thread_entries[i].interpreter_id == interpreter_id) {
if (is_new) {
*is_new = 0;
}
return &writer->thread_entries[i];
}
}
if (writer->thread_count >= writer->thread_capacity) {
ThreadEntry *new_entries = grow_array(writer->thread_entries,
&writer->thread_capacity,
sizeof(ThreadEntry));
if (!new_entries) {
return NULL;
}
writer->thread_entries = new_entries;
}
ThreadEntry *entry = &writer->thread_entries[writer->thread_count];
memset(entry, 0, sizeof(ThreadEntry));
entry->thread_id = thread_id;
entry->interpreter_id = interpreter_id;
entry->prev_timestamp = writer->start_time_us;
entry->prev_stack_capacity = MAX_STACK_DEPTH;
entry->pending_rle_capacity = INITIAL_RLE_CAPACITY;
entry->prev_stack = PyMem_Malloc(entry->prev_stack_capacity * sizeof(uint32_t));
if (!entry->prev_stack) {
PyErr_NoMemory();
return NULL;
}
entry->pending_rle = PyMem_Malloc(entry->pending_rle_capacity * sizeof(PendingRLESample));
if (!entry->pending_rle) {
PyMem_Free(entry->prev_stack);
PyErr_NoMemory();
return NULL;
}
writer->thread_count++;
if (is_new) {
*is_new = 1;
}
return entry;
}
/* Compare two stacks and return the encoding type and parameters.
* Sets:
* - shared_count: number of frames matching from bottom of stack
* - pop_count: frames to remove from prev stack
* - push_count: new frames to add
*
* Returns the best encoding type to use. */
static int
compare_stacks(const uint32_t *prev_stack, size_t prev_depth,
const uint32_t *curr_stack, size_t curr_depth,
size_t *shared_count, size_t *pop_count, size_t *push_count)
{
if (prev_depth == curr_depth) {
int identical = 1;
for (size_t i = 0; i < prev_depth; i++) {
if (prev_stack[i] != curr_stack[i]) {
identical = 0;
break;
}
}
if (identical) {
*shared_count = prev_depth;
*pop_count = 0;
*push_count = 0;
return STACK_REPEAT;
}
}
/* Find longest common suffix (frames at the bottom/outer part of stack).
* Stacks are stored innermost-first, so suffix is at the end. */
size_t suffix_len = 0;
size_t min_depth = (prev_depth < curr_depth) ? prev_depth : curr_depth;
for (size_t i = 0; i < min_depth; i++) {
size_t prev_idx = prev_depth - 1 - i;
size_t curr_idx = curr_depth - 1 - i;
if (prev_stack[prev_idx] == curr_stack[curr_idx]) {
suffix_len++;
} else {
break;
}
}
*shared_count = suffix_len;
*pop_count = prev_depth - suffix_len;
*push_count = curr_depth - suffix_len;
/* Choose best encoding based on byte cost */
/* STACK_FULL: 1 (type) + 1-2 (depth) + sum(frame varints) */
/* STACK_SUFFIX: 1 (type) + 1-2 (shared) + 1-2 (new_count) + sum(new frame varints) */
/* STACK_POP_PUSH: 1 (type) + 1-2 (pop) + 1-2 (push) + sum(new frame varints) */
/* If no common suffix, use full stack */
if (suffix_len == 0) {
return STACK_FULL;
}
/* If only adding frames (suffix == prev_depth), use SUFFIX */
if (*pop_count == 0 && *push_count > 0) {
return STACK_SUFFIX;
}
/* If popping and/or pushing, use POP_PUSH if it saves bytes */
/* Heuristic: POP_PUSH is better when we're modifying top frames */
if (*pop_count > 0 || *push_count > 0) {
/* Use full stack if sharing less than half the frames */
if (suffix_len < curr_depth / 2) {
return STACK_FULL;
}
return STACK_POP_PUSH;
}
return STACK_FULL;
}
/* Write common sample header: thread_id(8) + interpreter_id(4) + encoding(1).
* Returns 0 on success, -1 on failure. */
static inline int
write_sample_header(BinaryWriter *writer, ThreadEntry *entry, uint8_t encoding)
{
uint8_t header[SAMPLE_HEADER_FIXED_SIZE];
memcpy(header, &entry->thread_id, 8);
memcpy(header + 8, &entry->interpreter_id, 4);
header[12] = encoding;
return writer_write_bytes(writer, header, SAMPLE_HEADER_FIXED_SIZE);
}
/* Flush pending RLE samples for a thread.
* Writes the RLE record to the output buffer.
* Returns 0 on success, -1 on failure. */
static int
flush_pending_rle(BinaryWriter *writer, ThreadEntry *entry)
{
if (!entry->has_pending_rle || entry->pending_rle_count == 0) {
return 0;
}
/* Write RLE record:
* [thread_id: 8] [interpreter_id: 4] [STACK_REPEAT: 1] [count: varint]
* [timestamp_delta_1: varint] [status_1: 1] ... [timestamp_delta_N: varint] [status_N: 1]
*/
if (write_sample_header(writer, entry, STACK_REPEAT) < 0) {
return -1;
}
if (writer_write_varint_u32(writer, (uint32_t)entry->pending_rle_count) < 0) {
return -1;
}
for (size_t i = 0; i < entry->pending_rle_count; i++) {
if (writer_write_varint_u64(writer, entry->pending_rle[i].timestamp_delta) < 0) {
return -1;
}
if (writer_write_bytes(writer, &entry->pending_rle[i].status, 1) < 0) {
return -1;
}
writer->total_samples++;
}
writer->stats.repeat_records++;
writer->stats.repeat_samples += entry->pending_rle_count;
/* Each RLE sample saves writing the entire stack */
writer->stats.frames_saved += entry->pending_rle_count * entry->prev_stack_depth;
entry->pending_rle_count = 0;
entry->has_pending_rle = 0;
return 0;
}
/* Write a single sample with the specified encoding.
* Returns 0 on success, -1 on failure. */
static int
write_sample_with_encoding(BinaryWriter *writer, ThreadEntry *entry,
uint64_t timestamp_delta, uint8_t status,
int encoding_type,
const uint32_t *frame_indices, size_t stack_depth,
size_t shared_count, size_t pop_count, size_t push_count)
{
/* Header: thread_id(8) + interpreter_id(4) + encoding(1) + delta(varint) + status(1) */
uint8_t header_buf[SAMPLE_HEADER_MAX_SIZE];
memcpy(header_buf, &entry->thread_id, 8);
memcpy(header_buf + 8, &entry->interpreter_id, 4);
header_buf[12] = (uint8_t)encoding_type;
size_t varint_len = encode_varint_u64(header_buf + 13, timestamp_delta);
header_buf[13 + varint_len] = status;
if (writer_write_bytes(writer, header_buf, 14 + varint_len) < 0) {
return -1;
}
uint8_t frame_buf[MAX_FRAME_BUFFER_SIZE];
size_t frame_buf_pos = 0;
size_t frames_written = 0;
switch (encoding_type) {
case STACK_FULL:
/* [depth: varint] [frame_idx: varint]... */
frame_buf_pos += encode_varint_u32(frame_buf, (uint32_t)stack_depth);
for (size_t i = 0; i < stack_depth; i++) {
frame_buf_pos += encode_varint_u32(frame_buf + frame_buf_pos, frame_indices[i]);
}
frames_written = stack_depth;
writer->stats.full_records++;
break;
case STACK_SUFFIX:
/* [shared_count: varint] [new_count: varint] [new_frame_idx: varint]... */
frame_buf_pos += encode_varint_u32(frame_buf, (uint32_t)shared_count);
frame_buf_pos += encode_varint_u32(frame_buf + frame_buf_pos, (uint32_t)push_count);
/* New frames are at the top (beginning) of current stack */
for (size_t i = 0; i < push_count; i++) {
frame_buf_pos += encode_varint_u32(frame_buf + frame_buf_pos, frame_indices[i]);
}
frames_written = push_count;
writer->stats.suffix_records++;
/* Saved writing shared_count frames */
writer->stats.frames_saved += shared_count;
break;
case STACK_POP_PUSH:
/* [pop_count: varint] [push_count: varint] [new_frame_idx: varint]... */
frame_buf_pos += encode_varint_u32(frame_buf, (uint32_t)pop_count);
frame_buf_pos += encode_varint_u32(frame_buf + frame_buf_pos, (uint32_t)push_count);
/* New frames are at the top (beginning) of current stack */
for (size_t i = 0; i < push_count; i++) {
frame_buf_pos += encode_varint_u32(frame_buf + frame_buf_pos, frame_indices[i]);
}
frames_written = push_count;
writer->stats.pop_push_records++;
/* Saved writing shared_count frames (stack_depth - push_count if we had written full) */
writer->stats.frames_saved += shared_count;
break;
default:
PyErr_SetString(PyExc_RuntimeError, "Invalid stack encoding type");
return -1;
}
if (writer_write_bytes(writer, frame_buf, frame_buf_pos) < 0) {
return -1;
}
writer->stats.total_frames_written += frames_written;
writer->total_samples++;
return 0;
}
BinaryWriter *
binary_writer_create(const char *filename, uint64_t sample_interval_us, int compression_type,
uint64_t start_time_us)
{
BinaryWriter *writer = PyMem_Calloc(1, sizeof(BinaryWriter));
if (!writer) {
PyErr_NoMemory();
return NULL;
}
writer->filename = PyMem_Malloc(strlen(filename) + 1);
if (!writer->filename) {
PyMem_Free(writer);
PyErr_NoMemory();
return NULL;
}
strcpy(writer->filename, filename);
writer->start_time_us = start_time_us;
writer->sample_interval_us = sample_interval_us;
writer->compression_type = compression_type;
writer->write_buffer = PyMem_Malloc(WRITE_BUFFER_SIZE);
if (!writer->write_buffer) {
goto error;
}
writer->buffer_size = WRITE_BUFFER_SIZE;
writer->string_hash = _Py_hashtable_new_full(
string_hash_func,
string_compare_func,
string_key_destroy, /* Key destroy: decref the Python string */
NULL, /* Value destroy: values are just indices, not pointers */
NULL /* Use default allocator */
);
if (!writer->string_hash) {
goto error;
}
writer->strings = PyMem_Malloc(INITIAL_STRING_CAPACITY * sizeof(char *));
if (!writer->strings) {
goto error;
}
writer->string_lengths = PyMem_Malloc(INITIAL_STRING_CAPACITY * sizeof(size_t));
if (!writer->string_lengths) {
goto error;
}
writer->string_capacity = INITIAL_STRING_CAPACITY;
writer->frame_hash = _Py_hashtable_new_full(
frame_key_hash_func,
frame_key_compare_func,
frame_key_destroy, /* Key destroy: free the FrameKey */
NULL, /* Value destroy: values are just indices, not pointers */
NULL /* Use default allocator */
);
if (!writer->frame_hash) {
goto error;
}
writer->frame_entries = PyMem_Malloc(INITIAL_FRAME_CAPACITY * sizeof(FrameEntry));
if (!writer->frame_entries) {
goto error;
}
writer->frame_capacity = INITIAL_FRAME_CAPACITY;
writer->thread_entries = PyMem_Malloc(INITIAL_THREAD_CAPACITY * sizeof(ThreadEntry));
if (!writer->thread_entries) {
goto error;
}
writer->thread_capacity = INITIAL_THREAD_CAPACITY;
if (compression_type == COMPRESSION_ZSTD) {
if (writer_init_zstd(writer) < 0) {
goto error;
}
}
writer->fp = fopen(filename, "wb");
if (!writer->fp) {
PyErr_SetFromErrnoWithFilename(PyExc_IOError, filename);
goto error;
}
/* Hint sequential write pattern to kernel for better I/O scheduling */
#if defined(__linux__) && defined(POSIX_FADV_SEQUENTIAL)
{
int fd = fileno(writer->fp);
if (fd >= 0) {
(void)posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL);
}
}
#endif
uint8_t header[FILE_HEADER_PLACEHOLDER_SIZE] = {0};
if (fwrite_checked_allow_threads(header, FILE_HEADER_PLACEHOLDER_SIZE, writer->fp) < 0) {
goto error;
}
return writer;
error:
binary_writer_destroy(writer);
return NULL;
}
/* Build a frame stack from Python frame list by interning all strings and frames.
* Returns 0 on success, -1 on error. */
static int
build_frame_stack(BinaryWriter *writer, PyObject *frame_list,
uint32_t *curr_stack, size_t *curr_depth)
{
Py_ssize_t stack_depth = PyList_Size(frame_list);
*curr_depth = (stack_depth < MAX_STACK_DEPTH) ? stack_depth : MAX_STACK_DEPTH;
for (Py_ssize_t k = 0; k < (Py_ssize_t)*curr_depth; k++) {
/* Use unchecked accessors since we control the data structures */
PyObject *frame_info = PyList_GET_ITEM(frame_list, k);
/* Get filename, location, funcname, opcode from FrameInfo using unchecked access */
PyObject *filename = PyStructSequence_GET_ITEM(frame_info, 0);
PyObject *location = PyStructSequence_GET_ITEM(frame_info, 1);
PyObject *funcname = PyStructSequence_GET_ITEM(frame_info, 2);
PyObject *opcode_obj = PyStructSequence_GET_ITEM(frame_info, 3);
/* Extract location fields (can be None for synthetic frames) */
int32_t lineno = LOCATION_NOT_AVAILABLE;
int32_t end_lineno = LOCATION_NOT_AVAILABLE;
int32_t column = LOCATION_NOT_AVAILABLE;
int32_t end_column = LOCATION_NOT_AVAILABLE;
if (location != Py_None) {
/* LocationInfo is a struct sequence or tuple with:
* (lineno, end_lineno, column, end_column) */
PyObject *lineno_obj = PyTuple_Check(location) ?
PyTuple_GET_ITEM(location, 0) :
PyStructSequence_GET_ITEM(location, 0);
PyObject *end_lineno_obj = PyTuple_Check(location) ?
PyTuple_GET_ITEM(location, 1) :
PyStructSequence_GET_ITEM(location, 1);
PyObject *column_obj = PyTuple_Check(location) ?
PyTuple_GET_ITEM(location, 2) :
PyStructSequence_GET_ITEM(location, 2);
PyObject *end_column_obj = PyTuple_Check(location) ?
PyTuple_GET_ITEM(location, 3) :
PyStructSequence_GET_ITEM(location, 3);
PYLONG_TO_INT32_OR_DEFAULT(lineno_obj, lineno, LOCATION_NOT_AVAILABLE);
PYLONG_TO_INT32_OR_DEFAULT(end_lineno_obj, end_lineno, LOCATION_NOT_AVAILABLE);
PYLONG_TO_INT32_OR_DEFAULT(column_obj, column, LOCATION_NOT_AVAILABLE);
PYLONG_TO_INT32_OR_DEFAULT(end_column_obj, end_column, LOCATION_NOT_AVAILABLE);
}
/* Extract opcode (can be None) */
uint8_t opcode = OPCODE_NONE;
if (opcode_obj != Py_None) {
long opcode_long = PyLong_AsLong(opcode_obj);
if (UNLIKELY(PyErr_Occurred() != NULL)) {
PyErr_Clear();
opcode = OPCODE_NONE;
} else if (opcode_long >= 0 && opcode_long <= MAX_REAL_OPCODE) {
opcode = (uint8_t)opcode_long;
}
}
/* Intern filename */
uint32_t filename_idx;
if (writer_intern_string(writer, filename, &filename_idx) < 0) {
return -1;
}
/* Intern funcname */
uint32_t funcname_idx;
if (writer_intern_string(writer, funcname, &funcname_idx) < 0) {
return -1;
}
/* Intern frame with full location info */
FrameEntry frame_entry = {
.filename_idx = filename_idx,
.funcname_idx = funcname_idx,
.lineno = lineno,
.end_lineno = end_lineno,
.column = column,
.end_column = end_column,
.opcode = opcode
};
uint32_t frame_idx;
if (writer_intern_frame(writer, &frame_entry, &frame_idx) < 0) {
return -1;
}
curr_stack[k] = frame_idx;
}
return 0;
}
/* Process a single thread's sample.
* Returns 0 on success, -1 on error. */
static int
process_thread_sample(BinaryWriter *writer, PyObject *thread_info,
uint32_t interpreter_id, uint64_t timestamp_us)
{
PyObject *thread_id_obj = PyStructSequence_GET_ITEM(thread_info, 0);
PyObject *status_obj = PyStructSequence_GET_ITEM(thread_info, 1);
PyObject *frame_list = PyStructSequence_GET_ITEM(thread_info, 2);
uint64_t thread_id = PyLong_AsUnsignedLongLong(thread_id_obj);
if (thread_id == (uint64_t)-1 && PyErr_Occurred()) {
return -1;
}
long status_long = PyLong_AsLong(status_obj);
if (status_long == -1 && PyErr_Occurred()) {
return -1;
}
uint8_t status = (uint8_t)status_long;
int is_new_thread = 0;
ThreadEntry *entry = writer_get_or_create_thread_entry(
writer, thread_id, interpreter_id, &is_new_thread);
if (!entry) {
return -1;
}
/* Calculate timestamp delta */
uint64_t delta = timestamp_us - entry->prev_timestamp;
entry->prev_timestamp = timestamp_us;
/* Process frames and build current stack */
uint32_t curr_stack[MAX_STACK_DEPTH];
size_t curr_depth;
if (build_frame_stack(writer, frame_list, curr_stack, &curr_depth) < 0) {
return -1;
}
/* Compare with previous stack to determine encoding */
size_t shared_count, pop_count, push_count;
int encoding = compare_stacks(
entry->prev_stack, entry->prev_stack_depth,
curr_stack, curr_depth,
&shared_count, &pop_count, &push_count);
if (encoding == STACK_REPEAT && !is_new_thread) {
/* Buffer this sample for RLE */
if (GROW_ARRAY(entry->pending_rle, entry->pending_rle_count,
entry->pending_rle_capacity, PendingRLESample) < 0) {
return -1;
}
entry->pending_rle[entry->pending_rle_count].timestamp_delta = delta;
entry->pending_rle[entry->pending_rle_count].status = status;
entry->pending_rle_count++;
entry->has_pending_rle = 1;
} else {
/* Stack changed - flush any pending RLE first */
if (entry->has_pending_rle) {
if (flush_pending_rle(writer, entry) < 0) {
return -1;
}
}
if (write_sample_with_encoding(writer, entry, delta, status, encoding,
curr_stack, curr_depth,
shared_count, pop_count, push_count) < 0) {
return -1;
}
memcpy(entry->prev_stack, curr_stack, curr_depth * sizeof(uint32_t));
entry->prev_stack_depth = curr_depth;
}
return 0;
}
int
binary_writer_write_sample(BinaryWriter *writer, PyObject *stack_frames, uint64_t timestamp_us)
{
if (!PyList_Check(stack_frames)) {
PyErr_SetString(PyExc_TypeError, "stack_frames must be a list");
return -1;
}
Py_ssize_t num_interpreters = PyList_GET_SIZE(stack_frames);
for (Py_ssize_t i = 0; i < num_interpreters; i++) {
PyObject *interp_info = PyList_GET_ITEM(stack_frames, i);
PyObject *interp_id_obj = PyStructSequence_GET_ITEM(interp_info, 0);
PyObject *threads = PyStructSequence_GET_ITEM(interp_info, 1);
unsigned long interp_id_long = PyLong_AsUnsignedLong(interp_id_obj);
if (interp_id_long == (unsigned long)-1 && PyErr_Occurred()) {
return -1;
}
/* Bounds check: interpreter_id is stored as uint32_t in binary format */
if (interp_id_long > UINT32_MAX) {
PyErr_Format(PyExc_OverflowError,
"interpreter_id %lu exceeds maximum value %lu",
interp_id_long, (unsigned long)UINT32_MAX);
return -1;
}
uint32_t interpreter_id = (uint32_t)interp_id_long;
Py_ssize_t num_threads = PyList_GET_SIZE(threads);
for (Py_ssize_t j = 0; j < num_threads; j++) {
PyObject *thread_info = PyList_GET_ITEM(threads, j);
if (process_thread_sample(writer, thread_info, interpreter_id, timestamp_us) < 0) {
return -1;
}
}
}
return 0;
}
int
binary_writer_finalize(BinaryWriter *writer)
{
for (size_t i = 0; i < writer->thread_count; i++) {
if (writer->thread_entries[i].has_pending_rle) {
if (flush_pending_rle(writer, &writer->thread_entries[i]) < 0) {
return -1;
}
}
}
if (writer_flush_buffer(writer) < 0) {
return -1;
}
#ifdef HAVE_ZSTD
/* Finalize compression stream */
if (writer->compression_type == COMPRESSION_ZSTD && writer->zstd.cctx) {
ZSTD_inBuffer input = { NULL, 0, 0 };
size_t remaining;
do {
ZSTD_outBuffer output = {
writer->zstd.compressed_buffer,
writer->zstd.compressed_buffer_size,
0
};
remaining = ZSTD_compressStream2(writer->zstd.cctx, &output, &input, ZSTD_e_end);
if (ZSTD_isError(remaining)) {
PyErr_Format(PyExc_IOError, "zstd finalization error: %s",
ZSTD_getErrorName(remaining));
return -1;
}
if (output.pos > 0) {
if (fwrite_checked_allow_threads(writer->zstd.compressed_buffer, output.pos, writer->fp) < 0) {
return -1;
}
}
} while (remaining > 0);
}
#endif
/* Use 64-bit file position for >2GB files */
file_offset_t string_table_offset = FTELL64(writer->fp);
if (string_table_offset < 0) {
PyErr_SetFromErrno(PyExc_IOError);
return -1;
}
/* Release GIL during potentially large writes */
for (size_t i = 0; i < writer->string_count; i++) {
uint8_t len_buf[10];
size_t len_size = encode_varint_u32(len_buf, (uint32_t)writer->string_lengths[i]);
if (fwrite_checked_allow_threads(len_buf, len_size, writer->fp) < 0 ||
fwrite_checked_allow_threads(writer->strings[i], writer->string_lengths[i], writer->fp) < 0) {
return -1;
}
}
file_offset_t frame_table_offset = FTELL64(writer->fp);
if (frame_table_offset < 0) {
PyErr_SetFromErrno(PyExc_IOError);
return -1;
}
for (size_t i = 0; i < writer->frame_count; i++) {
FrameEntry *entry = &writer->frame_entries[i];
uint8_t buf[64]; /* Increased buffer for additional fields */
size_t pos = encode_varint_u32(buf, entry->filename_idx);
pos += encode_varint_u32(buf + pos, entry->funcname_idx);
pos += encode_varint_i32(buf + pos, entry->lineno);
/* Delta encode end_lineno: store (end_lineno - lineno) as zigzag.
* When lineno is -1, store delta as 0 (result will be -1). */
int32_t end_lineno_delta = 0;
if (entry->lineno != LOCATION_NOT_AVAILABLE &&
entry->end_lineno != LOCATION_NOT_AVAILABLE) {
end_lineno_delta = entry->end_lineno - entry->lineno;
}
pos += encode_varint_i32(buf + pos, end_lineno_delta);
pos += encode_varint_i32(buf + pos, entry->column);
/* Delta encode end_column: store (end_column - column) as zigzag.
* When column is -1, store delta as 0 (result will be -1). */
int32_t end_column_delta = 0;
if (entry->column != LOCATION_NOT_AVAILABLE &&
entry->end_column != LOCATION_NOT_AVAILABLE) {
end_column_delta = entry->end_column - entry->column;
}
pos += encode_varint_i32(buf + pos, end_column_delta);
buf[pos++] = entry->opcode;
if (fwrite_checked_allow_threads(buf, pos, writer->fp) < 0) {
return -1;
}
}
/* Footer: string_count(4) + frame_count(4) + file_size(8) + checksum(16) */
file_offset_t footer_offset = FTELL64(writer->fp);
if (footer_offset < 0) {
PyErr_SetFromErrno(PyExc_IOError);
return -1;
}
uint64_t file_size = (uint64_t)footer_offset + 32;
uint8_t footer[32] = {0};
/* Cast size_t to uint32_t before memcpy to ensure correct bytes are copied
* on both little-endian and big-endian systems (size_t is 8 bytes on 64-bit) */
uint32_t string_count_u32 = (uint32_t)writer->string_count;
uint32_t frame_count_u32 = (uint32_t)writer->frame_count;
memcpy(footer + 0, &string_count_u32, 4);
memcpy(footer + 4, &frame_count_u32, 4);
memcpy(footer + 8, &file_size, 8);
/* bytes 16-31: checksum placeholder (zeros) */
if (fwrite_checked_allow_threads(footer, 32, writer->fp) < 0) {
return -1;
}
if (FSEEK64(writer->fp, 0, SEEK_SET) < 0) {
PyErr_SetFromErrno(PyExc_IOError);
return -1;
}
/* Convert file offsets and counts to fixed-width types for portable header format.
* This ensures correct behavior on both little-endian and big-endian systems. */
uint64_t string_table_offset_u64 = (uint64_t)string_table_offset;
uint64_t frame_table_offset_u64 = (uint64_t)frame_table_offset;
uint32_t thread_count_u32 = (uint32_t)writer->thread_count;
uint32_t compression_type_u32 = (uint32_t)writer->compression_type;
uint8_t header[FILE_HEADER_SIZE] = {0};
uint32_t magic = BINARY_FORMAT_MAGIC;
uint32_t version = BINARY_FORMAT_VERSION;
memcpy(header + HDR_OFF_MAGIC, &magic, HDR_SIZE_MAGIC);
memcpy(header + HDR_OFF_VERSION, &version, HDR_SIZE_VERSION);
header[HDR_OFF_PY_MAJOR] = PY_MAJOR_VERSION;
header[HDR_OFF_PY_MINOR] = PY_MINOR_VERSION;
header[HDR_OFF_PY_MICRO] = PY_MICRO_VERSION;
memcpy(header + HDR_OFF_START_TIME, &writer->start_time_us, HDR_SIZE_START_TIME);
memcpy(header + HDR_OFF_INTERVAL, &writer->sample_interval_us, HDR_SIZE_INTERVAL);
memcpy(header + HDR_OFF_SAMPLES, &writer->total_samples, HDR_SIZE_SAMPLES);
memcpy(header + HDR_OFF_THREADS, &thread_count_u32, HDR_SIZE_THREADS);
memcpy(header + HDR_OFF_STR_TABLE, &string_table_offset_u64, HDR_SIZE_STR_TABLE);
memcpy(header + HDR_OFF_FRAME_TABLE, &frame_table_offset_u64, HDR_SIZE_FRAME_TABLE);
memcpy(header + HDR_OFF_COMPRESSION, &compression_type_u32, HDR_SIZE_COMPRESSION);
if (fwrite_checked_allow_threads(header, FILE_HEADER_SIZE, writer->fp) < 0) {
return -1;
}
if (fclose(writer->fp) != 0) {
writer->fp = NULL;
PyErr_SetFromErrno(PyExc_IOError);
return -1;
}
writer->fp = NULL;
return 0;
}
void
binary_writer_destroy(BinaryWriter *writer)
{
if (!writer) {
return;
}
if (writer->fp) {
fclose(writer->fp);
}
PyMem_Free(writer->filename);
PyMem_Free(writer->write_buffer);
#ifdef HAVE_ZSTD
if (writer->zstd.cctx) {
ZSTD_freeCCtx(writer->zstd.cctx);
}
PyMem_Free(writer->zstd.compressed_buffer);
#endif
if (writer->string_hash) {
_Py_hashtable_destroy(writer->string_hash);
}
if (writer->strings) {
for (size_t i = 0; i < writer->string_count; i++) {
PyMem_Free(writer->strings[i]);
}
PyMem_Free(writer->strings);
}
PyMem_Free(writer->string_lengths);
if (writer->frame_hash) {
_Py_hashtable_destroy(writer->frame_hash);
}
PyMem_Free(writer->frame_entries);
if (writer->thread_entries) {
for (size_t i = 0; i < writer->thread_count; i++) {
PyMem_Free(writer->thread_entries[i].prev_stack);
PyMem_Free(writer->thread_entries[i].pending_rle);
}
PyMem_Free(writer->thread_entries);
}
PyMem_Free(writer);
}
#undef PYLONG_TO_INT32_OR_DEFAULT
Reference in a new issue View git blame Copy permalink