Fix s3fifo plugin cache

Author: haochengxia

README.md

@@ -76,11 +76,7 @@ Implement custom cache replacement algorithms using pure Python functions - **no
 
 ### Python Hook Cache Overview
 
-The `PluginCache` allows you to define custom caching behavior through Python callback functions without without any C/C++ compilation. 
-
-### Hook Functions
-
-You need to implement these callback functions:
+The `PluginCache` allows you to define custom caching behavior through Python callback functions. You need to implement these callback functions:
 
 | Function | Signature | Description |
 |----------|-----------|-------------|

docs/src/en/index.md

@@ -65,4 +65,4 @@ We welcome contributions! Please see our [GitHub repository](https://github.com/
 
 ## License
 
-This project is licensed under the Apache License 2.0.
+This project is licensed under the GPL-3.0 License.

docs/src/zh/index.md

@@ -65,4 +65,4 @@ pip install -e .
 
 ## 许可证
 
-本项目采用 Apache License 2.0 许可证。
+本项目采用 GPL-3.0 许可证。

examples/plugin_cache/lru.py

@@ -62,7 +62,7 @@ def cache_free_hook(cache):
 
 reader = SyntheticReader(
     num_of_req=100000,
-    num_objects=100,
+    num_objects=10000,
     obj_size=100,
     seed=42,
     alpha=0.8,

examples/plugin_cache/s3fifo.py

@@ -1,6 +1,12 @@
 # An example of plugin for s3fifo
+
+# NOTE(haocheng): the one shows that with plugin system, we can make cache as lego blocks
+# Happy caching!
+
+import libcachesim as lcs
 from collections import OrderedDict
-from libcachesim import PluginCache, CommonCacheParams, Request, S3FIFO, SyntheticReader
+from collections import deque
+from libcachesim import PluginCache, CommonCacheParams, Request, S3FIFO, FIFO, SyntheticReader
 
 # NOTE(haocheng): we only support ignore object size for now
 class StandaloneS3FIFO:
@@ -9,160 +15,157 @@ def __init__(self,
                  ghost_size_ratio: float = 0.9,
                  move_to_main_threshold: int = 2,
                  cache_size: int = 1024):
-        # S3-FIFO uses three queues with OrderedDict for O(1) operations
-        self.small_fifo = OrderedDict()
-        self.main_fifo = OrderedDict()
-        self.ghost_fifo = OrderedDict()
-        
-        # Size limits
-        self.small_max_size = int(small_size_ratio * cache_size)
-        self.main_max_size = int(cache_size - small_size_ratio * cache_size)
-        self.ghost_max_size = int(ghost_size_ratio * cache_size)
+        self.cache_size = cache_size
+        small_fifo_size = int(small_size_ratio * cache_size)
+        main_fifo_size = cache_size - small_fifo_size
+        ghost_fifo_size = int(ghost_size_ratio * cache_size)
+
+        self.small_set = set()
+        self.main_set = set()
+        self.ghost_set = deque(maxlen=ghost_fifo_size)
+
+        self.small_fifo = FIFO(small_fifo_size)
+        self.main_fifo = FIFO(main_fifo_size)
+        self.ghost_fifo = FIFO(ghost_fifo_size)
 
         # Frequency tracking
-        self.small_freq = {}
-        self.main_freq = {}
-        self.ghost_freq = {}
+        self.freq = {}
 
         # Other parameters
         self.max_freq = 3
         self.move_to_main_threshold = move_to_main_threshold
 
-    def cache_hit(self, obj_id):
-        """
-        Cache hit can happen in two cases:
-        1. Small FIFO cache hit (small_fifo)
-        2. Main FIFO cache hit (main_fifo)
-        """
-        if obj_id in self.main_fifo:
-            self.main_freq[obj_id] += 1
-        elif obj_id in self.small_fifo:
-            self.small_freq[obj_id] += 1
-        else:
-            print(f"Cache hit for obj_id {obj_id} but not found in any queue")
-            print(f"small_fifo: {list(self.small_fifo.keys())}")
-            print(f"main_fifo: {list(self.main_fifo.keys())}")
-            print(f"ghost_fifo: {list(self.ghost_fifo.keys())}")
-            assert False, "Cache hit should happen in small_fifo or main_fifo"
+        self.has_evicted = False # Mark if we start to evict, only after full we will start eviction
+        self.hit_on_ghost = False
+
+    def cache_hit(self, req: Request):
+        hit_small = False
+        hit_main = False
+        if self.small_fifo.find(req, update_cache=False):
+            self.freq[req.obj_id] += 1
+
+        if self.main_fifo.find(req, update_cache=False):
+            self.freq[req.obj_id] += 1
 
-    def cache_miss(self, obj_id, obj_size=1):
-        """
-        Cache miss can happen in three cases:
-        1. Miss in small and main but hit in ghost
-        2. Miss all three queues
-        """
-        if obj_id in self.ghost_fifo:
-            del self.ghost_fifo[obj_id]
-            del self.ghost_freq[obj_id]
-            self.insert_to_main(obj_id)
+    def cache_miss(self, req: Request):
+        if not self.hit_on_ghost:
+            obj = self.ghost_fifo.find(req, update_cache=False)
+            if obj is not None:
+                self.hit_on_ghost = True
+                # remove from ghost set
+                self.ghost_fifo.remove(req.obj_id)
+                self.ghost_set.remove(req.obj_id)
+
+
+        # NOTE(haocheng): first we need to know this miss object has record in ghost or not
+        if not self.hit_on_ghost:
+            if req.obj_size >= self.small_fifo.cache_size:
+                # If object is too large, we do not process it
+                return
+
+            # If is initialization state, we need to insert to small fifo, 
+            # then we can insert to main fifo
+            if not self.has_evicted and self.small_fifo.get_occupied_byte() >= self.small_fifo.cache_size:
+                obj = self.main_fifo.insert(req)
+                self.main_set.add(obj.obj_id)
+            else:
+                obj = self.small_fifo.insert(req)
+                self.small_set.add(obj.obj_id)
         else:
-            # Miss all three queues
-            cond = (obj_id not in self.small_fifo) and (obj_id not in self.main_fifo)
-            assert cond, "Should not be in small_fifo or main_fifo"
-
-            # Then we need to insert to small fifo queue
-            self.insert_to_small(obj_id)
-
-    def insert_to_small(self, obj_id):
-        if len(self.small_fifo) >= self.small_max_size:
-            self.cache_evict_small()
-        self.small_fifo[obj_id] = None  # OrderedDict value doesn't matter
-        self.small_freq[obj_id] = 0
-
-    def insert_to_main(self, obj_id):
-        if len(self.main_fifo) >= self.main_max_size:
-            self.cache_evict_main()
-        self.main_fifo[obj_id] = None
-        self.main_freq[obj_id] = 0
-    
-    def insert_to_ghost(self, obj_id, original_freq=0):
-        if len(self.ghost_fifo) >= self.ghost_max_size:
-            # Remove oldest item
-            oldest_id = next(iter(self.ghost_fifo))
-            del self.ghost_fifo[oldest_id]
-            del self.ghost_freq[oldest_id]
-        self.ghost_fifo[obj_id] = None
-        self.ghost_freq[obj_id] = original_freq
+            obj = self.main_fifo.insert(req)
+            self.main_set.add(req.obj_id)
+            self.hit_on_ghost = False
+        self.freq[obj.obj_id] = 0
 
-    def cache_evict_small(self):
+    def cache_evict_small(self, req: Request):
         has_evicted = False
         evicted_id = None
-        while not has_evicted and len(self.small_fifo) > 0:
-            obj_to_evict = next(iter(self.small_fifo))  # Get first item
-            if self.small_freq[obj_to_evict] >= self.move_to_main_threshold:
-                # Move to main fifo cache (not real evict, just move)
-                del self.small_fifo[obj_to_evict]
-                del self.small_freq[obj_to_evict]
-                self.insert_to_main(obj_to_evict)
+        real_evicted_id = None
+        while not has_evicted and self.small_fifo.get_occupied_byte() > 0:
+            obj_to_evict = self.small_fifo.to_evict(req)
+            evicted_id = obj_to_evict.obj_id  # Store the ID before any operations
+            if self.freq[obj_to_evict.obj_id] >= self.move_to_main_threshold:
+                new_req = Request(obj_id=evicted_id, obj_size=1)
+                self.main_fifo.insert(new_req)
+                self.main_set.add(evicted_id)
+                # Reset frequency
+                self.freq[evicted_id] = 0
             else:
-                evicted_id = obj_to_evict
-                # Insert to ghost fifo cache (real evict)
-                del self.small_fifo[obj_to_evict]
-                del self.small_freq[obj_to_evict]
-                self.insert_to_ghost(obj_to_evict)
+                new_req = Request(obj_id=evicted_id, obj_size=1)
+                self.ghost_fifo.get(new_req)
+                self.ghost_set.append(evicted_id)
                 has_evicted = True
-        return evicted_id
+                real_evicted_id = evicted_id
+            flag = self.small_fifo.remove(evicted_id)
+            self.small_set.remove(evicted_id)
+            assert flag, "Should be able to remove"
+        return real_evicted_id
 
-    def cache_evict_main(self):
+    def cache_evict_main(self, req: Request):
         has_evicted = False
         evicted_id = None
-        while not has_evicted and len(self.main_fifo) > 0:
-            obj_to_evict = next(iter(self.main_fifo))  # Get first item
-            freq = self.main_freq[obj_to_evict]
+        while not has_evicted and self.main_fifo.get_occupied_byte() > 0:
+            obj_to_evict = self.main_fifo.to_evict(req)
+            assert obj_to_evict is not None
+            evicted_id = obj_to_evict.obj_id  # Store the ID before any operations
+            freq = self.freq[evicted_id]
             if freq >= 1:
                 # Reinsert with decremented frequency
-                del self.main_fifo[obj_to_evict]
-                del self.main_freq[obj_to_evict]
-                self.insert_to_main(obj_to_evict)
-                self.main_freq[obj_to_evict] = min(freq, self.max_freq) - 1
+                self.main_fifo.remove(evicted_id)
+                self.main_set.remove(evicted_id)
+                new_req = Request(obj_id=evicted_id, obj_size=1)
+                self.main_fifo.insert(new_req)
+                self.main_set.add(evicted_id)
+                self.freq[evicted_id] = min(freq, self.max_freq) - 1
             else:
-                evicted_id = obj_to_evict
-                # Real eviction
-                del self.main_fifo[obj_to_evict]
-                del self.main_freq[obj_to_evict]
+                flag = self.main_fifo.remove(evicted_id)
+                self.main_set.remove(evicted_id)
                 has_evicted = True
+            # print(f"Evicted {evicted_id}")
         return evicted_id
 
-    def cache_evict(self):
-        evicted_id = None
-        # if main is full or small is empty, evict main
-        if len(self.main_fifo) >= self.main_max_size or len(self.small_fifo) == 0:
-            evicted_id = self.cache_evict_main()
-        # if small is not empty, evict small
+    def cache_evict(self, req: Request):
+        if not self.hit_on_ghost:
+            obj = self.ghost_fifo.find(req, update_cache=False)
+            if obj is not None:
+                self.hit_on_ghost = True
+                # remove from ghost set
+                self.ghost_fifo.remove(req.obj_id)
+                self.ghost_set.remove(req.obj_id)
+
+        self.has_evicted = True
+        cond = (self.main_fifo.get_occupied_byte() > self.main_fifo.cache_size)
+        if (cond or (self.small_fifo.get_occupied_byte() == 0)):
+            obj_id = self.cache_evict_main(req)
         else:
-            evicted_id = self.cache_evict_small()
-        if evicted_id is None:
-            assert False, "Should not be None"
-        return evicted_id
+            obj_id = self.cache_evict_small(req)
+
+        if obj_id is not None:
+            del self.freq[obj_id]
+        
+        return obj_id
 
     def cache_remove(self, obj_id):
         removed = False
-        if obj_id in self.small_fifo:
-            del self.small_fifo[obj_id]
-            del self.small_freq[obj_id]
-            removed = True
-        elif obj_id in self.ghost_fifo:
-            del self.ghost_fifo[obj_id]
-            del self.ghost_freq[obj_id]
-            removed = True
-        elif obj_id in self.main_fifo:
-            del self.main_fifo[obj_id]
-            del self.main_freq[obj_id]
-            removed = True
+        removed |= self.small_fifo.remove(obj_id)
+        removed |= self.ghost_fifo.remove(obj_id)
+        removed |= self.main_fifo.remove(obj_id)
         return removed
 
 def cache_init_hook(common_cache_params: CommonCacheParams):
     return StandaloneS3FIFO(cache_size=common_cache_params.cache_size)
 
 def cache_hit_hook(cache, request: Request):
-    cache.cache_hit(request.obj_id)
+    cache.cache_hit(request)
 
 def cache_miss_hook(cache, request: Request):
-    cache.cache_miss(request.obj_id, request.obj_size)
+    cache.cache_miss(request)
 
 def cache_eviction_hook(cache, request: Request):
-    # NOTE(haocheng): never called
-    pass
+    evicted_id = None
+    while evicted_id is None:
+        evicted_id = cache.cache_evict(request)
+    return evicted_id
 
 def cache_remove_hook(cache, obj_id):
     cache.cache_remove(obj_id)
@@ -176,7 +179,7 @@ def cache_free_hook(cache):
     cache.main_freq.clear()
 
 cache = PluginCache(
-    cache_size=1024*1024,
+    cache_size=1024,
     cache_init_hook=cache_init_hook,
     cache_hit_hook=cache_hit_hook,
     cache_miss_hook=cache_miss_hook,
@@ -185,20 +188,29 @@ def cache_free_hook(cache):
     cache_free_hook=cache_free_hook,
     cache_name="S3FIFO")
 
-ref_s3fifo = S3FIFO(cache_size=1024)
+URI = "cache_dataset_oracleGeneral/2007_msr/msr_hm_0.oracleGeneral.zst"
+dl = lcs.DataLoader()
+dl.load(URI)
 
-reader = SyntheticReader(
-    num_of_req=1000000,
-    num_objects=100,
-    obj_size=1,
-    seed=42,
-    alpha=0.8,
-    dist="zipf",
+# Step 2: Open trace and process efficiently
+reader = lcs.TraceReader(
+    trace = dl.get_cache_path(URI),
+    trace_type = lcs.TraceType.ORACLE_GENERAL_TRACE,
+    reader_init_params = lcs.ReaderInitParam(ignore_obj_size=True)
 )
 
-for req in reader:
-    plugin_hit = cache.get(req)
-    ref_hit = ref_s3fifo.get(req)
-    assert plugin_hit == ref_hit, f"Cache hit mismatch: {plugin_hit} != {ref_hit}"
+ref_s3fifo = S3FIFO(cache_size=1024, small_size_ratio=0.1, ghost_size_ratio=0.9, move_to_main_threshold=2)
+
+# for req in reader:
+#     hit = cache.get(req)
+#     ref_hit = ref_s3fifo.get(req)
+#     assert hit == ref_hit, f"Cache hit mismatch: {hit} != {ref_hit}"
+
+req_miss_ratio, byte_miss_ratio = cache.process_trace(reader)
+ref_req_miss_ratio, ref_byte_miss_ratio = ref_s3fifo.process_trace(reader)
+print(f"Plugin req miss ratio: {req_miss_ratio}, ref req miss ratio: {ref_req_miss_ratio}")
+print(f"Plugin byte miss ratio: {byte_miss_ratio}, ref byte miss ratio: {ref_byte_miss_ratio}")
 
+assert req_miss_ratio == ref_req_miss_ratio
+assert byte_miss_ratio == ref_byte_miss_ratio
 print("All requests processed successfully. Plugin cache matches reference S3FIFO cache.")

libcachesim/cache.py

@@ -54,7 +54,7 @@ def __init__(self, _cache: Cache):
     def get(self, req: Request) -> bool:
         return self._cache.get(req)
 
-    def find(self, req: Request, update_cache: bool = True) -> CacheObject:
+    def find(self, req: Request, update_cache: bool = True) -> Optional[CacheObject]:
         return self._cache.find(req, update_cache)
 
     def can_insert(self, req: Request) -> bool:

src/export_cache.cpp

@@ -281,7 +281,7 @@ void export_cache(py::module& m) {
           "find",
           [](cache_t& self, const request_t& req, const bool update_cache) {
             cache_obj_t* obj = self.find(&self, &req, update_cache);
-            return py::cast(obj, py::return_value_policy::reference);
+            return obj ? py::cast(obj, py::return_value_policy::reference) : py::none();
           },
           "req"_a, "update_cache"_a = true)
       .def(