Try to resolve state sessions race condition (#977)

* Try to resolve state sessions race condition

I was not able to reproduce the bug via the integration tests, but this
has definitely happened once for me.

This change tries to at lease the reduce the chance of the race
condition where we loop through the cache and one of the items has been
removed.

Changes:

1. Reduce frequency of trying to clear state cache
2. Copy keys from loop first
3. Add some additional checks in case cache key does not exist.

mesop/server/state_session.py

@@ -72,6 +72,7 @@ class MemoryStateSessionBackend(StateSessionBackend):
 
   def __init__(self):
     self.cache = {}
+    self.last_checked_sessions = datetime(1900, 1, 1)
 
   def restore(self, token: str, states: States):
     if token not in self.cache:
@@ -89,15 +90,36 @@ def save(self, token: str, states: States):
   def clear_stale_sessions(self):
     stale_keys = set()
 
+    # Avoid trying to clear sessions with every request to avoid unnecessary processing.
     current_time = _current_datetime()
-    for key, (timestamp, _) in self.cache.items():
+    if (
+      self.last_checked_sessions + timedelta(minutes=self._SESSION_TTL_MINUTES)
+      > current_time
+    ):
+      return
+
+    self.last_checked_sessions = current_time
+
+    # Copy cache keys first to avoid possibility of the keys being removed
+    # while processing stale keys.
+    cache_keys = list(self.cache)
+
+    for key in cache_keys:
+      timestamp, _ = self.cache.get(key, (None, None))
       if (
-        timestamp + timedelta(minutes=self._SESSION_TTL_MINUTES) < current_time
+        timestamp
+        and timestamp + timedelta(minutes=self._SESSION_TTL_MINUTES)
+        < current_time
       ):
         stale_keys.add(key)
 
     for key in stale_keys:
-      del self.cache[key]
+      try:
+        del self.cache[key]
+      except KeyError:
+        logging.warning(
+          f"Tried to delete non-existent memory cache entry {key}."
+        )
 
 
 class FileStateSessionBackend(StateSessionBackend):