27052412546

Committed 06 Jun 2026 04:20AM UTC coverage: 50.862% (+1.0%) from 49.905%

Build # 27052412546

Build Type

Pull #73

github

Committed by

web-flow

Commit Message

Merge 734572730 into 88a3c8457

Pull Request Pull Request #73: add portable dependencies wheel support

Coverage Stats

31801 of 79859 branches covered (39.82%)

Branch coverage included in aggregate %.

32491 of 46545 relevant lines covered (69.81%)

9947.11 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

63.71

/src/dftracer/utils/python/indexer.cpp

#include <dftracer/utils/core/runtime.h>
#include <dftracer/utils/core/tasks/coro_scope.h>
#include <dftracer/utils/python/indexer.h>
#include <dftracer/utils/python/indexer_checkpoint.h>
#include <dftracer/utils/python/runtime.h>
#include <dftracer/utils/utilities/composites/dft/internal/utils.h>
#include <dftracer/utils/utilities/indexer/index_builder_utility.h>
#include <dftracer/utils/utilities/indexer/index_database.h>
#include <dftracer/utils/utilities/indexer/internal/helpers.h>
#include <structmember.h>

#include <cstring>
#include <memory>

static void CheckpointIndexer_dealloc(CheckpointIndexerObject *self) {
    if (self->handle) {
        // The Python wrapper owns only the native indexer handle. The
        // underlying RocksDB instance remains manager-owned and may continue to
        // live process-wide for the same .dftindex path.
        dft_indexer_destroy(self->handle);
        self->handle = NULL;
    }
    Py_XDECREF(self->gz_path);
    Py_XDECREF(self->index_path);
    Py_XDECREF(self->runtime_obj);
    Py_TYPE(self)->tp_free((PyObject *)self);
}

static void CheckpointIndexer_release_handle(CheckpointIndexerObject *self) {
    if (self->handle) {
        // Releasing the handle drops this wrapper's native indexer state only.
        // Shared RocksDB lifetime is managed separately by RocksDBManager.
        dft_indexer_destroy(self->handle);
        self->handle = NULL;
    }
}

static PyObject *CheckpointIndexer_new(PyTypeObject *type, PyObject *args,
                                       PyObject *kwds) {
    CheckpointIndexerObject *self;
    self = (CheckpointIndexerObject *)type->tp_alloc(type, 0);
    if (self != NULL) {
        self->handle = NULL;
        self->gz_path = NULL;
        self->index_path = NULL;
        self->checkpoint_size = 0;
        self->build_bloom = 0;
        self->build_manifest = 0;
        self->runtime_obj = NULL;
    }
    return (PyObject *)self;
}

static int CheckpointIndexer_init(CheckpointIndexerObject *self, PyObject *args,
                                  PyObject *kwds) {
    static const char *kwlist[] = {
        "gz_path",     "index_path",     "checkpoint_size", "force_rebuild",
        "build_bloom", "build_manifest", "runtime",         NULL};
    const char *gz_path;
    const char *index_path = NULL;
    std::uint64_t checkpoint_size =
        dftracer::utils::constants::indexer::DEFAULT_CHECKPOINT_SIZE;
    int force_rebuild = 0;
    int build_bloom = 0;
    int build_manifest = 0;
    PyObject *runtime_arg = NULL;

    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|snpppO", (char **)kwlist,
                                     &gz_path, &index_path, &checkpoint_size,
                                     &force_rebuild, &build_bloom,
                                     &build_manifest, &runtime_arg)) {
        return -1;
    }

    if (runtime_arg && runtime_arg != Py_None) {
        if (PyObject_TypeCheck(runtime_arg, &RuntimeType)) {
            Py_INCREF(runtime_arg);
            self->runtime_obj = runtime_arg;
        } else {
            PyObject *native = PyObject_GetAttrString(runtime_arg, "_native");
            if (native && PyObject_TypeCheck(native, &RuntimeType)) {
                self->runtime_obj = native;
            } else {
                Py_XDECREF(native);
                PyErr_SetString(PyExc_TypeError,
                                "runtime must be a Runtime instance or None");
                return -1;
            }
        }
    }

    self->gz_path = PyUnicode_FromString(gz_path);
    if (!self->gz_path) {
        return -1;
    }

    if (index_path) {
        self->index_path = PyUnicode_FromString(index_path);
    } else {
        const std::string index_path = dftracer::utils::utilities::composites::
            dft::internal::determine_index_path(gz_path, "");
        self->index_path = PyUnicode_FromString(index_path.c_str());
    }

    if (!self->index_path) {
        Py_DECREF(self->gz_path);
        return -1;
    }

    self->checkpoint_size = checkpoint_size;
    self->build_bloom = build_bloom;
    self->build_manifest = build_manifest;

    const char *index_path_str = PyUnicode_AsUTF8(self->index_path);
    if (!index_path_str) {
        return -1;
    }

    self->handle = dft_indexer_create(gz_path, index_path_str, checkpoint_size,
                                      force_rebuild);
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to create indexer");
        return -1;
    }

    return 0;
}

static dftracer::utils::Runtime *get_indexer_runtime(
    CheckpointIndexerObject *self) {
    if (self->runtime_obj) {
        return ((RuntimeObject *)self->runtime_obj)->runtime.get();
    }
    return get_default_runtime();
}

static PyObject *CheckpointIndexer_build(CheckpointIndexerObject *self,
                                         PyObject *Py_UNUSED(ignored)) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    // Use IndexBatchBuilderUtility when bloom or manifest is requested.
    // Otherwise, use the simpler dft_indexer_build which only creates
    // checkpoints.
    if (self->build_bloom || self->build_manifest) {
        using namespace dftracer::utils;
        using namespace dftracer::utils::utilities::indexer;

        const char *gz = PyUnicode_AsUTF8(self->gz_path);
        const char *idx = PyUnicode_AsUTF8(self->index_path);
        if (!gz || !idx) {
            return NULL;
        }

        auto batch_config = std::make_shared<IndexBuildBatchConfig>();
        batch_config->file_paths.emplace_back(gz);
        batch_config->checkpoint_size =
            static_cast<std::size_t>(self->checkpoint_size);
        batch_config->build_manifest = self->build_manifest != 0;
        batch_config->parallelism = 1;
        batch_config->use_batch_write = true;
        batch_config->rebuild_root_summaries = true;

        std::string idx_str(idx);
        auto pos = idx_str.find_last_of('/');
        if (pos != std::string::npos) {
            batch_config->index_dir = idx_str.substr(0, pos);
        }

        Runtime *rt = get_indexer_runtime(self);
        IndexBuildBatchResult batch_result;

        try {
            Py_BEGIN_ALLOW_THREADS rt
                ->submit(
                    run_coro_scope(
                        rt->executor(),
                        [](CoroScope &scope,
                           std::shared_ptr<IndexBuildBatchConfig> cfg,
                           IndexBuildBatchResult *out) -> coro::CoroTask<void> {
                            *out = co_await IndexBatchBuilderUtility::process(
                                &scope, std::move(cfg));
                        },
                        batch_config, &batch_result),
                    "indexer-build")
                .get();
            Py_END_ALLOW_THREADS
        } catch (const std::exception &e) {
            PyErr_SetString(PyExc_RuntimeError, e.what());
            return NULL;
        }

        if (batch_result.failed > 0 && !batch_result.results.empty()) {
            const auto &result = batch_result.results[0];
            if (!result.success) {
                PyErr_SetString(PyExc_RuntimeError,
                                result.error_message.c_str());
                return NULL;
            }
        }
    } else {
        // Simple checkpoint-only build
        int result;
        Py_BEGIN_ALLOW_THREADS result = dft_indexer_build(self->handle);
        Py_END_ALLOW_THREADS

            if (result < 0) {
            PyErr_SetString(PyExc_RuntimeError, "Failed to build index");
            return NULL;
        }
    }

    Py_RETURN_NONE;
}

static PyObject *CheckpointIndexer_need_rebuild(CheckpointIndexerObject *self,
                                                PyObject *Py_UNUSED(ignored)) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    int result = dft_indexer_need_rebuild(self->handle);
    return PyBool_FromLong(result);
}

static PyObject *CheckpointIndexer_exists(CheckpointIndexerObject *self,
                                          PyObject *Py_UNUSED(ignored)) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    int result = dft_indexer_exists(self->handle);
    return PyBool_FromLong(result);
}

static PyObject *CheckpointIndexer_get_max_bytes(CheckpointIndexerObject *self,
                                                 PyObject *Py_UNUSED(ignored)) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    uint64_t result = dft_indexer_get_max_bytes(self->handle);
    return PyLong_FromUnsignedLongLong(result);
}

static PyObject *CheckpointIndexer_get_num_lines(CheckpointIndexerObject *self,
                                                 PyObject *Py_UNUSED(ignored)) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    uint64_t result = dft_indexer_get_num_lines(self->handle);
    return PyLong_FromUnsignedLongLong(result);
}

static PyObject *CheckpointIndexer_find_checkpoint(
    CheckpointIndexerObject *self, PyObject *args) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    std::size_t target_offset;
    if (!PyArg_ParseTuple(args, "n", &target_offset)) {
        return NULL;
    }

    dft_indexer_checkpoint_t checkpoint;
    int found =
        dft_indexer_find_checkpoint(self->handle, target_offset, &checkpoint);

    if (!found) {
        Py_RETURN_NONE;
    }

    // Create IndexerCheckpoint object
    IndexerCheckpointObject *cp_obj =
        (IndexerCheckpointObject *)IndexerCheckpoint_new(&IndexerCheckpointType,
                                                         NULL, NULL);
    if (!cp_obj) {
        return NULL;
    }

    cp_obj->checkpoint = checkpoint;
    return (PyObject *)cp_obj;
}

static PyObject *CheckpointIndexer_get_checkpoints(
    CheckpointIndexerObject *self, PyObject *Py_UNUSED(ignored)) {
    if (!self->handle) {
        PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");
        return NULL;
    }

    dft_indexer_checkpoint_t *checkpoints = NULL;
    std::size_t count = 0;

    int result =
        dft_indexer_get_checkpoints(self->handle, &checkpoints, &count);
    if (result != 0 || !checkpoints) {
        dft_indexer_free_checkpoints(checkpoints, count);
        PyObject *list = PyList_New(0);
        return list;
    }

    PyObject *list = PyList_New(count);
    if (!list) {
        dft_indexer_free_checkpoints(checkpoints, count);
        return NULL;
    }

    for (std::size_t i = 0; i < count; i++) {
        IndexerCheckpointObject *cp_obj =
            (IndexerCheckpointObject *)IndexerCheckpoint_new(
                &IndexerCheckpointType, NULL, NULL);
        if (!cp_obj) {
            Py_DECREF(list);
            dft_indexer_free_checkpoints(checkpoints, count);
            return NULL;
        }
        cp_obj->checkpoint = checkpoints[i];
        PyList_SetItem(list, i, (PyObject *)cp_obj);
    }

    dft_indexer_free_checkpoints(checkpoints, count);
    return list;
}

static PyObject *CheckpointIndexer_has_bloom(CheckpointIndexerObject *self,
                                             void *closure) {
    const char *idx = PyUnicode_AsUTF8(self->index_path);
    const char *gz = PyUnicode_AsUTF8(self->gz_path);
    if (!idx || !gz) {
        Py_RETURN_FALSE;
    }
    try {
        using namespace dftracer::utils::utilities::indexer;
        using namespace dftracer::utils::utilities::indexer::internal;
        IndexDatabase db(
            idx, dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);
        std::string logical = get_logical_path(gz);
        int fid = db.get_file_info_id(logical);
        if (fid >= 0 && db.has_bloom_data(fid)) {
            Py_RETURN_TRUE;
        }
    } catch (...) {
    }
    Py_RETURN_FALSE;
}

static PyObject *CheckpointIndexer_has_manifest(CheckpointIndexerObject *self,
                                                void *closure) {
    const char *idx = PyUnicode_AsUTF8(self->index_path);
    const char *gz = PyUnicode_AsUTF8(self->gz_path);
    if (!idx || !gz) {
        Py_RETURN_FALSE;
    }
    try {
        using namespace dftracer::utils::utilities::indexer;
        using namespace dftracer::utils::utilities::indexer::internal;
        IndexDatabase db(
            idx, dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);
        std::string logical = get_logical_path(gz);
        int fid = db.get_file_info_id(logical);
        if (fid >= 0 && db.has_manifest_data(fid)) {
            Py_RETURN_TRUE;
        }
    } catch (...) {
    }
    Py_RETURN_FALSE;
}

static PyObject *CheckpointIndexer_gz_path(CheckpointIndexerObject *self,
                                           void *closure) {
    Py_INCREF(self->gz_path);
    return self->gz_path;
}

static PyObject *CheckpointIndexer_index_path(CheckpointIndexerObject *self,
                                              void *closure) {
    Py_INCREF(self->index_path);
    return self->index_path;
}

static PyObject *CheckpointIndexer_checkpoint_size(
    CheckpointIndexerObject *self, void *closure) {
    return PyLong_FromUnsignedLongLong(self->checkpoint_size);
}

static PyObject *CheckpointIndexer_enter(CheckpointIndexerObject *self,
                                         PyObject *Py_UNUSED(ignored)) {
    Py_INCREF(self);
    return (PyObject *)self;
}

static PyObject *CheckpointIndexer_close(CheckpointIndexerObject *self,
                                         PyObject *Py_UNUSED(ignored)) {
    CheckpointIndexer_release_handle(self);
    Py_RETURN_NONE;
}

static PyObject *CheckpointIndexer_exit(CheckpointIndexerObject *self,
                                        PyObject *args) {
    CheckpointIndexer_release_handle(self);
    Py_RETURN_NONE;
}

static PyMethodDef CheckpointIndexer_methods[] = {
    {"build", (PyCFunction)CheckpointIndexer_build, METH_NOARGS,
     "build()\n"
     "--\n"
     "\n"
     "Build or rebuild the index.\n"},
    {"need_rebuild", (PyCFunction)CheckpointIndexer_need_rebuild, METH_NOARGS,
     "Check if a rebuild is needed."},
    {"exists", (PyCFunction)CheckpointIndexer_exists, METH_NOARGS,
     "Check if the .dftindex store exists."},
    {"get_max_bytes", (PyCFunction)CheckpointIndexer_get_max_bytes, METH_NOARGS,
     "Get the maximum uncompressed bytes in the indexed file."},
    {"get_num_lines", (PyCFunction)CheckpointIndexer_get_num_lines, METH_NOARGS,
     "Get the total number of lines in the indexed file."},
    {"find_checkpoint", (PyCFunction)CheckpointIndexer_find_checkpoint,
     METH_VARARGS,
     "Find the best checkpoint for a given uncompressed offset.\n"
     "\n"
     "Args:\n"
     "    offset (int): Uncompressed byte offset.\n"},
    {"get_checkpoints", (PyCFunction)CheckpointIndexer_get_checkpoints,
     METH_NOARGS, "Get all checkpoints for this file as a list."},
    {"close", (PyCFunction)CheckpointIndexer_close, METH_NOARGS,
     "Release this Python wrapper's native indexer handle.\n"
     "\n"
     "The shared RocksDB instance for the same .dftindex path remains managed\n"
     "by the native RocksDBManager cache."},
    {"__enter__", (PyCFunction)CheckpointIndexer_enter, METH_NOARGS,
     "Enter the runtime context for the with statement."},
    {"__exit__", (PyCFunction)CheckpointIndexer_exit, METH_VARARGS,
     "Release this Python wrapper on context exit.\n"
     "\n"
     "This does not force-close the shared RocksDB instance for the same\n"
     ".dftindex path."},
    {NULL} /* Sentinel */
};

static PyGetSetDef CheckpointIndexer_getsetters[] = {
    {"gz_path", (getter)CheckpointIndexer_gz_path, NULL,
     "Path to the gzip file", NULL},
    {"index_path", (getter)CheckpointIndexer_index_path, NULL,
     "Path to the .dftindex store", NULL},
    {"checkpoint_size", (getter)CheckpointIndexer_checkpoint_size, NULL,
     "Checkpoint size in bytes", NULL},
    {"has_bloom", (getter)CheckpointIndexer_has_bloom, NULL,
     "Whether bloom data exists in index", NULL},
    {"has_manifest", (getter)CheckpointIndexer_has_manifest, NULL,
     "Whether manifest data exists in index", NULL},
    {NULL} /* Sentinel */
};

PyTypeObject CheckpointIndexerType = {
    PyVarObject_HEAD_INIT(
        NULL, 0) "dftracer_utils_ext.CheckpointIndexer", /* tp_name */
    sizeof(CheckpointIndexerObject),                     /* tp_basicsize */
    0,                                                   /* tp_itemsize */
    (destructor)CheckpointIndexer_dealloc,               /* tp_dealloc */
    0,                                        /* tp_vectorcall_offset */
    0,                                        /* tp_getattr */
    0,                                        /* tp_setattr */
    0,                                        /* tp_as_async */
    0,                                        /* tp_repr */
    0,                                        /* tp_as_number */
    0,                                        /* tp_as_sequence */
    0,                                        /* tp_as_mapping */
    0,                                        /* tp_hash */
    0,                                        /* tp_call */
    0,                                        /* tp_str */
    0,                                        /* tp_getattro */
    0,                                        /* tp_setattro */
    0,                                        /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
    "CheckpointIndexer(gz_path, index_path=None, checkpoint_size=1048576, "
    "force_rebuild=False, build_bloom=False, build_manifest=False, "
    "runtime=None)\n"
    "--\n"
    "\n"
    "Checkpoint indexer for single-file checkpoint-level operations on a "
    "gzip trace.\n"
    "\n"
    "Args:\n"
    "    gz_path (str): Path to the gzip trace file.\n"
    "    index_path (str or None): Path to the .dftindex store. If None,\n"
    "        uses the root-local \".dftindex\" next to gz_path.\n"
    "    checkpoint_size (int): Checkpoint size in bytes for index\n"
    "        building (default 1 MB).\n"
    "    force_rebuild (bool): If True, rebuild the index even if it\n"
    "        exists.\n"
    "    build_bloom (bool): If True, build bloom filter data in the\n"
    "        index.\n"
    "    build_manifest (bool): If True, build manifest data in the\n"
    "        store.\n"
    "    runtime (Runtime or None): Runtime instance for thread pool\n"
    "        control. If None, uses the default global Runtime.\n", /* tp_doc */
    0,                                /* tp_traverse */
    0,                                /* tp_clear */
    0,                                /* tp_richcompare */
    0,                                /* tp_weaklistoffset */
    0,                                /* tp_iter */
    0,                                /* tp_iternext */
    CheckpointIndexer_methods,        /* tp_methods */
    0,                                /* tp_members */
    CheckpointIndexer_getsetters,     /* tp_getset */
    0,                                /* tp_base */
    0,                                /* tp_dict */
    0,                                /* tp_descr_get */
    0,                                /* tp_descr_set */
    0,                                /* tp_dictoffset */
    (initproc)CheckpointIndexer_init, /* tp_init */
    0,                                /* tp_alloc */
    CheckpointIndexer_new,            /* tp_new */
};

int init_checkpoint_indexer(PyObject *m) {
    if (PyType_Ready(&CheckpointIndexerType) < 0) return -1;

    Py_INCREF(&CheckpointIndexerType);
    if (PyModule_AddObject(m, "CheckpointIndexer",
                           (PyObject *)&CheckpointIndexerType) < 0) {
        Py_DECREF(&CheckpointIndexerType);
        Py_DECREF(m);
        return -1;
    }

    return 0;
}

1	#include <dftracer/utils/core/runtime.h>
2	#include <dftracer/utils/core/tasks/coro_scope.h>
3	#include <dftracer/utils/python/indexer.h>
4	#include <dftracer/utils/python/indexer_checkpoint.h>
5	#include <dftracer/utils/python/runtime.h>
6	#include <dftracer/utils/utilities/composites/dft/internal/utils.h>
7	#include <dftracer/utils/utilities/indexer/index_builder_utility.h>
8	#include <dftracer/utils/utilities/indexer/index_database.h>
9	#include <dftracer/utils/utilities/indexer/internal/helpers.h>
10	#include <structmember.h>
11
12	#include <cstring>
13	#include <memory>
14
15	static void CheckpointIndexer_dealloc(CheckpointIndexerObject *self) {	71✔
16	if (self->handle) {	71✔
17	// The Python wrapper owns only the native indexer handle. The
18	// underlying RocksDB instance remains manager-owned and may continue to
19	// live process-wide for the same .dftindex path.
20	dft_indexer_destroy(self->handle);	6✔
21	self->handle = NULL;	6✔
22	}	6✔
23	Py_XDECREF(self->gz_path);	71✔
24	Py_XDECREF(self->index_path);	71✔
25	Py_XDECREF(self->runtime_obj);	71✔
26	Py_TYPE(self)->tp_free((PyObject *)self);	71✔
27	}	71✔
28
29	static void CheckpointIndexer_release_handle(CheckpointIndexerObject *self) {	64✔
30	if (self->handle) {	64!
31	// Releasing the handle drops this wrapper's native indexer state only.
32	// Shared RocksDB lifetime is managed separately by RocksDBManager.
33	dft_indexer_destroy(self->handle);	64✔
34	self->handle = NULL;	64✔
35	}	64✔
36	}	64✔
37
38	static PyObject CheckpointIndexer_new(PyTypeObject type, PyObject *args,	65✔
39	PyObject *kwds) {
40	CheckpointIndexerObject *self;
41	self = (CheckpointIndexerObject *)type->tp_alloc(type, 0);	65✔
42	if (self != NULL) {	65!
43	self->handle = NULL;	65✔
44	self->gz_path = NULL;	65✔
45	self->index_path = NULL;	65✔
46	self->checkpoint_size = 0;	65✔
47	self->build_bloom = 0;	65✔
48	self->build_manifest = 0;	65✔
49	self->runtime_obj = NULL;	65✔
50	}	65✔
51	return (PyObject *)self;	65✔
52	}
53
54	static int CheckpointIndexer_init(CheckpointIndexerObject self, PyObject args,	65✔
55	PyObject *kwds) {
56	static const char *kwlist[] = {
57	"gz_path", "index_path", "checkpoint_size", "force_rebuild",
58	"build_bloom", "build_manifest", "runtime", NULL};
59	const char *gz_path;
60	const char *index_path = NULL;	65✔
61	std::uint64_t checkpoint_size =	65✔
62	dftracer::utils::constants::indexer::DEFAULT_CHECKPOINT_SIZE;
63	int force_rebuild = 0;	65✔
64	int build_bloom = 0;	65✔
65	int build_manifest = 0;	65✔
66	PyObject *runtime_arg = NULL;	65✔
67
68	if (!PyArg_ParseTupleAndKeywords(args, kwds, "s\|snpppO", (char **)kwlist,	65!
69	&gz_path, &index_path, &checkpoint_size,
70	&force_rebuild, &build_bloom,
71	&build_manifest, &runtime_arg)) {
72	return -1;	×
73	}
74
75	if (runtime_arg && runtime_arg != Py_None) {	65!
76	if (PyObject_TypeCheck(runtime_arg, &RuntimeType)) {	×
77	Py_INCREF(runtime_arg);	×
78	self->runtime_obj = runtime_arg;	×
79	} else {	×
80	PyObject *native = PyObject_GetAttrString(runtime_arg, "_native");	×
81	if (native && PyObject_TypeCheck(native, &RuntimeType)) {	×
82	self->runtime_obj = native;	×
83	} else {	×
84	Py_XDECREF(native);	×
85	PyErr_SetString(PyExc_TypeError,	×
86	"runtime must be a Runtime instance or None");
87	return -1;	×
88	}
89	}
90	}	×
91
92	self->gz_path = PyUnicode_FromString(gz_path);	65✔
93	if (!self->gz_path) {	65!
94	return -1;	×
95	}
96
97	if (index_path) {	65✔
98	self->index_path = PyUnicode_FromString(index_path);	64✔
99	} else {	64✔
100	const std::string index_path = dftracer::utils::utilities::composites::	1!
101	dft::internal::determine_index_path(gz_path, "");	1!
102	self->index_path = PyUnicode_FromString(index_path.c_str());	1!
103	}	1✔
104
105	if (!self->index_path) {	65!
106	Py_DECREF(self->gz_path);	×
107	return -1;	×
108	}
109
110	self->checkpoint_size = checkpoint_size;	65✔
111	self->build_bloom = build_bloom;	65✔
112	self->build_manifest = build_manifest;	65✔
113
114	const char *index_path_str = PyUnicode_AsUTF8(self->index_path);	65✔
115	if (!index_path_str) {	65!
116	return -1;	×
117	}
118
119	self->handle = dft_indexer_create(gz_path, index_path_str, checkpoint_size,	130✔
120	force_rebuild);	65✔
121	if (!self->handle) {	65✔
122	PyErr_SetString(PyExc_RuntimeError, "Failed to create indexer");	1✔
123	return -1;	1✔
124	}
125
126	return 0;	64✔
127	}	65✔
128
129	static dftracer::utils::Runtime *get_indexer_runtime(	16✔
130	CheckpointIndexerObject *self) {
131	if (self->runtime_obj) {	16!
132	return ((RuntimeObject *)self->runtime_obj)->runtime.get();	×
133	}
134	return get_default_runtime();	16✔
135	}	16✔
136
137	static PyObject CheckpointIndexer_build(CheckpointIndexerObject self,	62✔
138	PyObject *Py_UNUSED(ignored)) {
139	if (!self->handle) {	62!
140	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
141	return NULL;	×
142	}
143
144	// Use IndexBatchBuilderUtility when bloom or manifest is requested.
145	// Otherwise, use the simpler dft_indexer_build which only creates
146	// checkpoints.
147	if (self->build_bloom \|\| self->build_manifest) {	62✔
148	using namespace dftracer::utils;
149	using namespace dftracer::utils::utilities::indexer;
150
151	const char *gz = PyUnicode_AsUTF8(self->gz_path);	16✔
152	const char *idx = PyUnicode_AsUTF8(self->index_path);	16✔
153	if (!gz \|\| !idx) {	16!
154	return NULL;	×
155	}
156
157	auto batch_config = std::make_shared<IndexBuildBatchConfig>();	16✔
158	batch_config->file_paths.emplace_back(gz);	16!
159	batch_config->checkpoint_size =	16✔
160	static_cast<std::size_t>(self->checkpoint_size);	16✔
161	batch_config->build_manifest = self->build_manifest != 0;	16✔
162	batch_config->parallelism = 1;	16✔
163	batch_config->use_batch_write = true;	16✔
164	batch_config->rebuild_root_summaries = true;	16✔
165
166	std::string idx_str(idx);	16!
167	auto pos = idx_str.find_last_of('/');	16✔
168	if (pos != std::string::npos) {	16!
169	batch_config->index_dir = idx_str.substr(0, pos);	16!
170	}	16✔
171
172	Runtime *rt = get_indexer_runtime(self);	16!
173	IndexBuildBatchResult batch_result;	16✔
174
175	try {
176	Py_BEGIN_ALLOW_THREADS rt	32!
177	->submit(	16!
178	run_coro_scope(	16!
179	rt->executor(),	16!
180	[](CoroScope &scope,	112!
181	std::shared_ptr<IndexBuildBatchConfig> cfg,
182	IndexBuildBatchResult *out) -> coro::CoroTask<void> {	16!
183	*out = co_await IndexBatchBuilderUtility::process(	128!
184	&scope, std::move(cfg));	48✔
185	},	32✔
186	batch_config, &batch_result),	16✔
187	"indexer-build")	16!
188	.get();	16!
189	Py_END_ALLOW_THREADS	16!
190	} catch (const std::exception &e) {	16!
191	PyErr_SetString(PyExc_RuntimeError, e.what());	×
192	return NULL;	×
193	}	×
194
195	if (batch_result.failed > 0 && !batch_result.results.empty()) {	16!
196	const auto &result = batch_result.results[0];	×
197	if (!result.success) {	×
198	PyErr_SetString(PyExc_RuntimeError,	×
199	result.error_message.c_str());	×
200	return NULL;	×
201	}
202	}	×
203	} else {	16!
204	// Simple checkpoint-only build
205	int result;
206	Py_BEGIN_ALLOW_THREADS result = dft_indexer_build(self->handle);	46✔
207	Py_END_ALLOW_THREADS	46✔
208
209	if (result < 0) {	46!
210	PyErr_SetString(PyExc_RuntimeError, "Failed to build index");	×
211	return NULL;	×
212	}
213	}
214
215	Py_RETURN_NONE;	62✔
216	}	62✔
217
218	static PyObject CheckpointIndexer_need_rebuild(CheckpointIndexerObject self,	30✔
219	PyObject *Py_UNUSED(ignored)) {
220	if (!self->handle) {	30!
221	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
222	return NULL;	×
223	}
224
225	int result = dft_indexer_need_rebuild(self->handle);	30✔
226	return PyBool_FromLong(result);	30✔
227	}	30✔
228
229	static PyObject CheckpointIndexer_exists(CheckpointIndexerObject self,	×
230	PyObject *Py_UNUSED(ignored)) {
231	if (!self->handle) {	×
232	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
233	return NULL;	×
234	}
235
236	int result = dft_indexer_exists(self->handle);	×
237	return PyBool_FromLong(result);	×
238	}	×
239
240	static PyObject CheckpointIndexer_get_max_bytes(CheckpointIndexerObject self,	4✔
241	PyObject *Py_UNUSED(ignored)) {
242	if (!self->handle) {	4!
243	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
244	return NULL;	×
245	}
246
247	uint64_t result = dft_indexer_get_max_bytes(self->handle);	4✔
248	return PyLong_FromUnsignedLongLong(result);	4✔
249	}	4✔
250
251	static PyObject CheckpointIndexer_get_num_lines(CheckpointIndexerObject self,	5✔
252	PyObject *Py_UNUSED(ignored)) {
253	if (!self->handle) {	5!
254	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
255	return NULL;	×
256	}
257
258	uint64_t result = dft_indexer_get_num_lines(self->handle);	5✔
259	return PyLong_FromUnsignedLongLong(result);	5✔
260	}	5✔
261
262	static PyObject *CheckpointIndexer_find_checkpoint(	3✔
263	CheckpointIndexerObject self, PyObject args) {
264	if (!self->handle) {	3!
265	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
266	return NULL;	×
267	}
268
269	std::size_t target_offset;
270	if (!PyArg_ParseTuple(args, "n", &target_offset)) {	3!
271	return NULL;	×
272	}
273
274	dft_indexer_checkpoint_t checkpoint;
275	int found =	3✔
276	dft_indexer_find_checkpoint(self->handle, target_offset, &checkpoint);	3✔
277
278	if (!found) {	3✔
279	Py_RETURN_NONE;	1✔
280	}
281
282	// Create IndexerCheckpoint object
283	IndexerCheckpointObject *cp_obj =	2✔
284	(IndexerCheckpointObject *)IndexerCheckpoint_new(&IndexerCheckpointType,	2✔
285	NULL, NULL);
286	if (!cp_obj) {	2!
287	return NULL;	×
288	}
289
290	cp_obj->checkpoint = checkpoint;	2✔
291	return (PyObject *)cp_obj;	2✔
292	}	3✔
293
294	static PyObject *CheckpointIndexer_get_checkpoints(	3✔
295	CheckpointIndexerObject self, PyObject Py_UNUSED(ignored)) {
296	if (!self->handle) {	3!
297	PyErr_SetString(PyExc_RuntimeError, "Indexer not initialized");	×
298	return NULL;	×
299	}
300
301	dft_indexer_checkpoint_t *checkpoints = NULL;	3✔
302	std::size_t count = 0;	3✔
303
304	int result =	3✔
305	dft_indexer_get_checkpoints(self->handle, &checkpoints, &count);	3✔
306	if (result != 0 \|\| !checkpoints) {	3!
307	dft_indexer_free_checkpoints(checkpoints, count);	1✔
308	PyObject *list = PyList_New(0);	1✔
309	return list;	1✔
310	}
311
312	PyObject *list = PyList_New(count);	2✔
313	if (!list) {	2!
314	dft_indexer_free_checkpoints(checkpoints, count);	×
315	return NULL;	×
316	}
317
318	for (std::size_t i = 0; i < count; i++) {	85✔
319	IndexerCheckpointObject *cp_obj =	83✔
320	(IndexerCheckpointObject *)IndexerCheckpoint_new(	83✔
321	&IndexerCheckpointType, NULL, NULL);
322	if (!cp_obj) {	83!
323	Py_DECREF(list);	×
324	dft_indexer_free_checkpoints(checkpoints, count);	×
325	return NULL;	×
326	}
327	cp_obj->checkpoint = checkpoints[i];	83✔
328	PyList_SetItem(list, i, (PyObject *)cp_obj);	83✔
329	}	83✔
330
331	dft_indexer_free_checkpoints(checkpoints, count);	2✔
332	return list;	2✔
333	}	3✔
334
335	static PyObject CheckpointIndexer_has_bloom(CheckpointIndexerObject self,	6✔
336	void *closure) {
337	const char *idx = PyUnicode_AsUTF8(self->index_path);	6✔
338	const char *gz = PyUnicode_AsUTF8(self->gz_path);	6✔
339	if (!idx \|\| !gz) {	6!
340	Py_RETURN_FALSE;	×
341	}
342	try {
343	using namespace dftracer::utils::utilities::indexer;
344	using namespace dftracer::utils::utilities::indexer::internal;
345	IndexDatabase db(	6!
346	idx, dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);	6!
347	std::string logical = get_logical_path(gz);	6!
348	int fid = db.get_file_info_id(logical);	6!
349	if (fid >= 0 && db.has_bloom_data(fid)) {	6!
350	Py_RETURN_TRUE;	1✔
351	}
352	} catch (...) {	6!
353	}	×
354	Py_RETURN_FALSE;	5✔
355	}	6✔
356
357	static PyObject CheckpointIndexer_has_manifest(CheckpointIndexerObject self,	3✔
358	void *closure) {
359	const char *idx = PyUnicode_AsUTF8(self->index_path);	3✔
360	const char *gz = PyUnicode_AsUTF8(self->gz_path);	3✔
361	if (!idx \|\| !gz) {	3!
362	Py_RETURN_FALSE;	×
363	}
364	try {
365	using namespace dftracer::utils::utilities::indexer;
366	using namespace dftracer::utils::utilities::indexer::internal;
367	IndexDatabase db(	3!
368	idx, dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);	3!
369	std::string logical = get_logical_path(gz);	3!
370	int fid = db.get_file_info_id(logical);	3!
371	if (fid >= 0 && db.has_manifest_data(fid)) {	3!
372	Py_RETURN_TRUE;	1✔
373	}
374	} catch (...) {	3!
375	}	×
376	Py_RETURN_FALSE;	2✔
377	}	3✔
378
379	static PyObject CheckpointIndexer_gz_path(CheckpointIndexerObject self,	3✔
380	void *closure) {
381	Py_INCREF(self->gz_path);	3✔
382	return self->gz_path;	3✔
383	}
384
385	static PyObject CheckpointIndexer_index_path(CheckpointIndexerObject self,	2✔
386	void *closure) {
387	Py_INCREF(self->index_path);	2✔
388	return self->index_path;	2✔
389	}
390
391	static PyObject *CheckpointIndexer_checkpoint_size(	2✔
392	CheckpointIndexerObject self, void closure) {
393	return PyLong_FromUnsignedLongLong(self->checkpoint_size);	2✔
394	}
395
396	static PyObject CheckpointIndexer_enter(CheckpointIndexerObject self,	63✔
397	PyObject *Py_UNUSED(ignored)) {
398	Py_INCREF(self);	63✔
399	return (PyObject *)self;	63✔
400	}
401
402	static PyObject CheckpointIndexer_close(CheckpointIndexerObject self,	1✔
403	PyObject *Py_UNUSED(ignored)) {
404	CheckpointIndexer_release_handle(self);	1✔
405	Py_RETURN_NONE;	1✔
406	}
407
408	static PyObject CheckpointIndexer_exit(CheckpointIndexerObject self,	63✔
409	PyObject *args) {
410	CheckpointIndexer_release_handle(self);	63✔
411	Py_RETURN_NONE;	63✔
412	}
413
414	static PyMethodDef CheckpointIndexer_methods[] = {
415	{"build", (PyCFunction)CheckpointIndexer_build, METH_NOARGS,
416	"build()\n"
417	"--\n"
418	"\n"
419	"Build or rebuild the index.\n"},
420	{"need_rebuild", (PyCFunction)CheckpointIndexer_need_rebuild, METH_NOARGS,
421	"Check if a rebuild is needed."},
422	{"exists", (PyCFunction)CheckpointIndexer_exists, METH_NOARGS,
423	"Check if the .dftindex store exists."},
424	{"get_max_bytes", (PyCFunction)CheckpointIndexer_get_max_bytes, METH_NOARGS,
425	"Get the maximum uncompressed bytes in the indexed file."},
426	{"get_num_lines", (PyCFunction)CheckpointIndexer_get_num_lines, METH_NOARGS,
427	"Get the total number of lines in the indexed file."},
428	{"find_checkpoint", (PyCFunction)CheckpointIndexer_find_checkpoint,
429	METH_VARARGS,
430	"Find the best checkpoint for a given uncompressed offset.\n"
431	"\n"
432	"Args:\n"
433	" offset (int): Uncompressed byte offset.\n"},
434	{"get_checkpoints", (PyCFunction)CheckpointIndexer_get_checkpoints,
435	METH_NOARGS, "Get all checkpoints for this file as a list."},
436	{"close", (PyCFunction)CheckpointIndexer_close, METH_NOARGS,
437	"Release this Python wrapper's native indexer handle.\n"
438	"\n"
439	"The shared RocksDB instance for the same .dftindex path remains managed\n"
440	"by the native RocksDBManager cache."},
441	{"__enter__", (PyCFunction)CheckpointIndexer_enter, METH_NOARGS,
442	"Enter the runtime context for the with statement."},
443	{"__exit__", (PyCFunction)CheckpointIndexer_exit, METH_VARARGS,
444	"Release this Python wrapper on context exit.\n"
445	"\n"
446	"This does not force-close the shared RocksDB instance for the same\n"
447	".dftindex path."},
448	{NULL} /* Sentinel */
449	};
450
451	static PyGetSetDef CheckpointIndexer_getsetters[] = {
452	{"gz_path", (getter)CheckpointIndexer_gz_path, NULL,
453	"Path to the gzip file", NULL},
454	{"index_path", (getter)CheckpointIndexer_index_path, NULL,
455	"Path to the .dftindex store", NULL},
456	{"checkpoint_size", (getter)CheckpointIndexer_checkpoint_size, NULL,
457	"Checkpoint size in bytes", NULL},
458	{"has_bloom", (getter)CheckpointIndexer_has_bloom, NULL,
459	"Whether bloom data exists in index", NULL},
460	{"has_manifest", (getter)CheckpointIndexer_has_manifest, NULL,
461	"Whether manifest data exists in index", NULL},
462	{NULL} /* Sentinel */
463	};
464
465	PyTypeObject CheckpointIndexerType = {
466	PyVarObject_HEAD_INIT(
467	NULL, 0) "dftracer_utils_ext.CheckpointIndexer", /* tp_name */
468	sizeof(CheckpointIndexerObject), /* tp_basicsize */
469	0, /* tp_itemsize */
470	(destructor)CheckpointIndexer_dealloc, /* tp_dealloc */
471	0, /* tp_vectorcall_offset */
472	0, /* tp_getattr */
473	0, /* tp_setattr */
474	0, /* tp_as_async */
475	0, /* tp_repr */
476	0, /* tp_as_number */
477	0, /* tp_as_sequence */
478	0, /* tp_as_mapping */
479	0, /* tp_hash */
480	0, /* tp_call */
481	0, /* tp_str */
482	0, /* tp_getattro */
483	0, /* tp_setattro */
484	0, /* tp_as_buffer */
485	Py_TPFLAGS_DEFAULT \| Py_TPFLAGS_BASETYPE, /* tp_flags */
486	"CheckpointIndexer(gz_path, index_path=None, checkpoint_size=1048576, "
487	"force_rebuild=False, build_bloom=False, build_manifest=False, "
488	"runtime=None)\n"
489	"--\n"
490	"\n"
491	"Checkpoint indexer for single-file checkpoint-level operations on a "
492	"gzip trace.\n"
493	"\n"
494	"Args:\n"
495	" gz_path (str): Path to the gzip trace file.\n"
496	" index_path (str or None): Path to the .dftindex store. If None,\n"
497	" uses the root-local \".dftindex\" next to gz_path.\n"
498	" checkpoint_size (int): Checkpoint size in bytes for index\n"
499	" building (default 1 MB).\n"
500	" force_rebuild (bool): If True, rebuild the index even if it\n"
501	" exists.\n"
502	" build_bloom (bool): If True, build bloom filter data in the\n"
503	" index.\n"
504	" build_manifest (bool): If True, build manifest data in the\n"
505	" store.\n"
506	" runtime (Runtime or None): Runtime instance for thread pool\n"
507	" control. If None, uses the default global Runtime.\n", /* tp_doc */
508	0, /* tp_traverse */
509	0, /* tp_clear */
510	0, /* tp_richcompare */
511	0, /* tp_weaklistoffset */
512	0, /* tp_iter */
513	0, /* tp_iternext */
514	CheckpointIndexer_methods, /* tp_methods */
515	0, /* tp_members */
516	CheckpointIndexer_getsetters, /* tp_getset */
517	0, /* tp_base */
518	0, /* tp_dict */
519	0, /* tp_descr_get */
520	0, /* tp_descr_set */
521	0, /* tp_dictoffset */
522	(initproc)CheckpointIndexer_init, /* tp_init */
523	0, /* tp_alloc */
524	CheckpointIndexer_new, /* tp_new */
525	};
526
527	int init_checkpoint_indexer(PyObject *m) {	1✔
528	if (PyType_Ready(&CheckpointIndexerType) < 0) return -1;	1!
529
530	Py_INCREF(&CheckpointIndexerType);	1✔
531	if (PyModule_AddObject(m, "CheckpointIndexer",	2!
532	(PyObject *)&CheckpointIndexerType) < 0) {	1✔
533	Py_DECREF(&CheckpointIndexerType);	×
534	Py_DECREF(m);	×
535	return -1;	×
536	}
537
538	return 0;	1✔
539	}	1✔

llnl / dftracer-utils / 27052412546

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous