28521653886

Committed 01 Jul 2026 01:36PM UTC coverage: 50.92% (-1.4%) from 52.278%

Build # 28521653886

Build Type

Pull #83

github

Committed by

web-flow

Commit Message

Merge 9bdedb1e9 into 2efed6649

Pull Request Pull Request #83: refactor and improve code QoL

Coverage Stats

31893 of 80049 branches covered (39.84%)

Branch coverage included in aggregate %.

789 of 1613 new or added lines in 87 files covered. (48.92%)

5007 existing lines in 181 files now uncovered.

32812 of 47024 relevant lines covered (69.78%)

9905.42 hits per line

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

63.94

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

llnl / dftracer-utils / 28521653886

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