28356348514

Committed 29 Jun 2026 07:40AM UTC coverage: 52.174% (-0.1%) from 52.278%

Build # 28356348514

Build Type

Pull #83

github

Committed by

web-flow

Commit Message

Merge 278203630 into 2efed6649

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

Coverage Stats

37276 of 92891 branches covered (40.13%)

Branch coverage included in aggregate %.

671 of 1173 new or added lines in 58 files covered. (57.2%)

66 existing lines in 30 files now uncovered.

33619 of 42991 relevant lines covered (78.2%)

20387.45 hits per line

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

63.32

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

llnl / dftracer-utils / 28356348514

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