diff --git a/libcachesim/protocols.py b/libcachesim/protocols.py index 74a45f8..9741b78 100644 --- a/libcachesim/protocols.py +++ b/libcachesim/protocols.py @@ -28,6 +28,7 @@ def skip_n_req(self, n: int) -> int: ... def reset(self) -> None: ... def close(self) -> None: ... def clone(self) -> "ReaderProtocol": ... + def get_working_set_size(self) -> tuple[int, int]: ... def __iter__(self) -> Iterator[Request]: ... def __next__(self) -> Request: ... def __len__(self) -> int: ... diff --git a/libcachesim/synthetic_reader.py b/libcachesim/synthetic_reader.py index 936f29d..c5e7df6 100644 --- a/libcachesim/synthetic_reader.py +++ b/libcachesim/synthetic_reader.py @@ -181,6 +181,15 @@ def set_read_pos(self, pos: float) -> None: def get_read_pos(self) -> float: """Get current read position""" return float(self.current_pos) + + def get_working_set_size(self) -> Tuple[int, int]: + """Calculate working set size""" + wss_obj, wss_byte = 0, 0 + if self.current_pos > 0: + unique_ids = np.unique(self.obj_ids[:self.current_pos]) + wss_obj = len(unique_ids) + wss_byte = wss_obj * self.obj_size + return wss_obj, wss_byte def __iter__(self) -> Iterator[Request]: """Iterator implementation""" diff --git a/libcachesim/trace_reader.py b/libcachesim/trace_reader.py index e593dbb..deb0312 100644 --- a/libcachesim/trace_reader.py +++ b/libcachesim/trace_reader.py @@ -6,7 +6,7 @@ from urllib.parse import urlparse from .protocols import ReaderProtocol -from .libcachesim_python import TraceType, SamplerType, Request, ReaderInitParam, Reader, Sampler, ReadDirection +from .libcachesim_python import TraceType, SamplerType, Request, ReaderInitParam, Reader, Sampler, ReadDirection, cal_working_set_size from ._s3_cache import get_data_loader logger = logging.getLogger(__name__) @@ -276,6 +276,9 @@ def go_back_one_req(self) -> None: def set_read_pos(self, pos: float) -> None: self._reader.set_read_pos(pos) + def get_working_set_size(self) -> tuple[int, int]: + return cal_working_set_size(self._reader) + def __iter__(self) -> Iterator[Request]: self._reader.reset() return self diff --git a/src/export_reader.cpp b/src/export_reader.cpp index 8f286f3..9f7df86 100644 --- a/src/export_reader.cpp +++ b/src/export_reader.cpp @@ -61,7 +61,15 @@ struct SamplerDeleter { } }; + void export_reader(py::module& m) { + /* Helper function(s) */ + m.def("cal_working_set_size", [](reader_t& reader) { + int64_t wss_obj = 0, wss_byte = 0; + cal_working_set_size(&reader, &wss_obj, &wss_byte); + return std::make_tuple(wss_obj, wss_byte); + }, "reader"_a); + // Sampler type enumeration py::enum_(m, "SamplerType") .value("SPATIAL_SAMPLER", sampler_type::SPATIAL_SAMPLER)