lukecwik commented on code in PR #23046:
URL: https://github.com/apache/beam/pull/23046#discussion_r966417187
##########
sdks/python/apache_beam/runners/worker/statecache.py:
##########
@@ -111,28 +140,86 @@ def __init__(self, max_weight):
self._max_weight = max_weight
self._current_weight = 0
self._cache = collections.OrderedDict(
- ) # type: collections.OrderedDict[Tuple[bytes, Optional[bytes]],
WeightedValue]
+ ) # type: collections.OrderedDict[Any, WeightedValue]
self._hit_count = 0
self._miss_count = 0
self._evict_count = 0
+ self._load_time_ns = 0
+ self._load_count = 0
self._lock = threading.RLock()
- def get(self, state_key, cache_token):
- # type: (bytes, Optional[bytes]) -> Any
- assert cache_token and self.is_cache_enabled()
- key = (state_key, cache_token)
+ def peek(self, key):
+ # type: (Any) -> Any
+ assert self.is_cache_enabled()
with self._lock:
value = self._cache.get(key, None)
- if value is None:
+ if value is None or isinstance(value, _LoadingValue):
self._miss_count += 1
return None
+
self._cache.move_to_end(key)
self._hit_count += 1
- return value.value()
+ return value.value()
+
+ def get(self, key, loading_fn):
+ # type: (Any, Callable[[Any], Any]) -> Any
+ assert self.is_cache_enabled() and callable(loading_fn)
+
+ self._lock.acquire()
+ value = self._cache.get(key, None)
+ if value is None:
+ self._miss_count += 1
+ loading_value = _LoadingValue()
Review Comment:
I updated the test to ensure that the value after the load isn't a
`_LoadingValue`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
