Intermediate changes

2 files changed, 44 insertions, 4 deletions

diff --git a/library/python/func/__init__.py b/library/python/func/__init__.py
index 5eda75267a..1573b99c1a 100644
--- a/library/python/func/__init__.py
+++ b/library/python/func/__init__.py
@@ -150,12 +150,25 @@ def compose(*functions):
 
 
 class Singleton(type):
-    _instances = {}
+    __instances = {}
+    __lock = threading.Lock()
+
+    class _LockedObj:
+        def __init__(self, obj, lock):
+            self.obj = obj
+            self.lock = lock
 
     def __call__(cls, *args, **kwargs):
-        if cls not in cls._instances:
-            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
-        return cls._instances[cls]
+        if cls not in cls.__instances:
+            with cls.__lock:
+                if cls not in cls.__instances:
+                    cls.__instances[cls] = cls._LockedObj(None, threading.Lock())
+
+        if not cls.__instances[cls].obj:
+            with cls.__instances[cls].lock:
+                if not cls.__instances[cls].obj:
+                    cls.__instances[cls].obj = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls.__instances[cls].obj
 
 
 def stable_uniq(it):
diff --git a/library/python/func/ut/test_func.py b/library/python/func/ut/test_func.py
index 70a10d62cb..90fe125a11 100644
--- a/library/python/func/ut/test_func.py
+++ b/library/python/func/ut/test_func.py
@@ -3,6 +3,7 @@ import multiprocessing
 import random
 import threading
 import time
+import six
 
 import library.python.func as func
 
@@ -135,6 +136,32 @@ def test_flatten_dict():
     assert func.flatten_dict({"a": 1, "b": {"c": {"d": 2}}}, separator="/") == {"a": 1, "b/c/d": 2}
 
 
+def test_threadsafe_singleton():
+    class ShouldBeSingle(six.with_metaclass(func.Singleton, object)):
+        def __new__(cls, *args, **kwargs):
+            time.sleep(0.1)
+            return super(ShouldBeSingle, cls).__new__(cls, *args, **kwargs)
+
+    threads_count = 100
+    threads = [None] * threads_count
+    results = [None] * threads_count
+
+    def class_factory(results, i):
+        time.sleep(0.1)
+        results[i] = ShouldBeSingle()
+
+    for i in range(threads_count):
+        threads[i] = threading.Thread(target=class_factory, args=(results, i))
+
+    for i in range(threads_count):
+        threads[i].start()
+
+    for i in range(threads_count):
+        threads[i].join()
+
+    assert len(set(results)) == 1
+
+
 def test_memoize_thread_local():
     class Counter(object):
         def __init__(self, s):