Support reader.get_working_set_size

Author: haochengxia

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: ...

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._obj_ids is not None:
+            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"""

libcachesim/trace_reader.py

@@ -1,12 +1,12 @@
 """Wrapper of Reader with S3 support."""
 
 import logging
-from typing import overload, Union, Optional
+from typing import Tuple, overload, Union, Optional
 from collections.abc import Iterator
 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

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_<sampler_type>(m, "SamplerType")
       .value("SPATIAL_SAMPLER", sampler_type::SPATIAL_SAMPLER)