Skip to content

Commit

Permalink
Add timeout mechaism for cache.
Browse files Browse the repository at this point in the history
  • Loading branch information
ralgond committed Dec 5, 2024
1 parent 5f38c65 commit 18bddc1
Show file tree
Hide file tree
Showing 4 changed files with 34 additions and 14 deletions.
21 changes: 12 additions & 9 deletions mybatis/cache.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import json
import pickle
import time
from typing import Dict, Any, Optional

from pympler import asizeof
Expand All @@ -9,18 +10,11 @@ def __init__(self, sql, param_list):
self.sql = sql
self.param_list = param_list

def __hash__(self):
return hash((self.sql, self.param_list))

def __eq__(self, other):
if not isinstance(other, CacheKey):
return False
return self.sql == other.sql and self.param_list == other.param_list

class Cache(object):
def __init__(self, memory_limit:int):
def __init__(self, memory_limit:int, max_live_ms:int):
self.memory_limit = memory_limit
self.memory_used = 0
self.max_live_ms = max_live_ms
self.table : Dict[str, CacheNode] = {}
self.list = CacheList()

Expand Down Expand Up @@ -69,7 +63,15 @@ def get(self, raw_key: CacheKey) -> Optional[Any]:
key = json.dumps(raw_key.__dict__)
if key not in self.table:
return None

node = self.table[key]
current_time_ms = int(time.time() * 1000)
if current_time_ms - node.timestamp > self.max_live_ms:
del self.table[node.key]
self.list.remove(node)
self.memory_used -= node.memory_usage
return None

self.list.move_to_head(node)
return json.loads(node.value)

Expand All @@ -89,6 +91,7 @@ def __init__(self, key : Any, value : Any):
self.memory_usage = asizeof.asizeof(key) + asizeof.asizeof(value)
self.key = key
self.value = value
self.timestamp = int(time.time() * 1000)

class CacheList:
def __init__(self):
Expand Down
6 changes: 3 additions & 3 deletions mybatis/mybatis.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,14 @@
import os

class Mybatis(object):
def __init__(self, conn, mapper_path:str, cache_memory_limit:Optional[int]=None):
def __init__(self, conn, mapper_path:str, cache_memory_limit:Optional[int]=None, cache_max_live_ms:int=5*1000):
self.conn = conn
self.mapper_manager = MapperManager()

if cache_memory_limit is not None:
self.cache = Cache(cache_memory_limit)
self.cache = Cache(cache_memory_limit, cache_max_live_ms)
else:
self.cache = Cache(0)
self.cache = Cache(0, cache_max_live_ms)

mapper_file_name_l = [name for name in os.listdir(mapper_path) if name.endswith(".xml")]
for file_name in mapper_file_name_l:
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

setup(
name='mybatis',
version='0.0.6',
version='0.0.7',
description='A python ORM like mybatis.',
long_description=open('README.md').read(),
long_description_content_type='text/markdown', # 如果你使用的是Markdown格式的README
Expand Down
19 changes: 18 additions & 1 deletion test/test_cache.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import time

import pytest
from mybatis import Cache, CacheKey

def test_basic():
cache = Cache(memory_limit=555) # 50MB
cache = Cache(memory_limit=555, max_live_ms=10*1000) # 50MB, 10sec
cache.put(CacheKey("a", [1, 'a', None]), [{"a1": 1}, {"a2": 2}])
cache.put(CacheKey("b", [2, 'b', None]), "2")
cache.put(CacheKey("c", [3, 'c', None]), "3")
Expand Down Expand Up @@ -42,3 +44,18 @@ def test_basic():

assert l[2][0] == '{"sql": "d", "param_list": [4, "d", null]}'
assert l[2][1] == None

def test_timeout():
cache = Cache(memory_limit=555, max_live_ms=1 * 1000) # 50MB, 10sec
cache.put(CacheKey("a", [1, 'a', None]), [{"a1": 1}, {"a2": 2}])
cache.put(CacheKey("b", [2, 'b', None]), "2")
cache.put(CacheKey("c", [3, 'c', None]), "3")
cache.put(CacheKey("d", [4, 'd', None]), None)

time.sleep(2)
assert cache.get(CacheKey("a", [1, 'a', None])) == None
assert cache.get(CacheKey("b", [2, 'b', None])) == None
assert cache.get(CacheKey("c", [3, 'c', None])) == None
assert cache.get(CacheKey("d", [4, 'd', None])) == None

assert cache.memory_used == 0

0 comments on commit 18bddc1

Please sign in to comment.