Fix: memoize multithreding optimization

Fix: memoize multithreding optimization

1 files changed, 21 insertions, 7 deletions

diff --git a/library/python/func/__init__.py b/library/python/func/__init__.py
index 12a280bddc..15f2137f1d 100644
--- a/library/python/func/__init__.py
+++ b/library/python/func/__init__.py
@@ -1,6 +1,8 @@
 import functools
 import threading
 import collections
+import contextlib
+import six
 
 
 def map0(func, value):
@@ -76,20 +78,32 @@ class lazy_classproperty(object):
         return getattr(owner, attr_name)
 
 
-def memoize(limit=0, thread_local=False):
+class nullcontext(object):
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+
+def memoize(limit=0, thread_local=False, thread_safe=True):
     assert limit >= 0
+    assert limit <= 0 or thread_safe, 'memoize() it not thread safe enough to work in limiting and non-thread safe mode'
 
     def decorator(func):
         memory = {}
-        lock = threading.Lock()
+
+        if six.PY3:
+            lock = contextlib.nullcontext()
+        else:
+            lock = nullcontext()
+        lock = threading.Lock() if thread_safe else lock
 
         if limit:
             keys = collections.deque()
 
             def get(args):
-                try:
-                    return memory[args]
-                except KeyError:
+                if args not in memory:
                     with lock:
                         if args not in memory:
                             fargs = args[-1]
@@ -97,7 +111,7 @@ def memoize(limit=0, thread_local=False):
                             keys.append(args)
                             if len(keys) > limit:
                                 del memory[keys.popleft()]
-                        return memory[args]
+                return memory[args]
 
         else:
 
@@ -106,7 +120,7 @@ def memoize(limit=0, thread_local=False):
                     with lock:
                         if args not in memory:
                             fargs = args[-1]
-                            memory[args] = func(*fargs)
+                            memory.setdefault(args, func(*fargs))
                 return memory[args]
 
         if thread_local: