Migrate shared tag from tfx-bsl

sdks/python/apache_beam/utils/shared.py

@@ -102,10 +102,12 @@ class _SharedControlBlock(object):
   def __init__(self):
     self._lock = threading.Lock()
     self._ref = None
+    self._tag = None
 
   def acquire(
       self,
-      constructor_fn  # type: Callable[[], Any]
+      constructor_fn,  # type: Callable[[], Any]
+      tag=None  # type: Any
   ):
     # type: (...) -> Any
 
@@ -116,6 +118,9 @@ def acquire(
         present in the cache. This function should take no arguments. It should
         return an initialised object, or None if the object could not be
         initialised / constructed.
+      tag: an optional indentifier to store with the cached object. If
+        subsequent calls to acquire use different tags, the object will be
+        reloaded rather than returned from cache.
 
     Returns:
       An initialised object, either from a previous initialisation, or
@@ -124,7 +129,8 @@ def acquire(
     with self._lock:
       # self._ref is None if this is a new control block.
       # self._ref() is None if the weak reference was GCed.
-      if self._ref is None or self._ref() is None:
+      # self._tag != tag if user specifies a new identifier
+      if self._ref is None or self._ref() is None or self._tag != tag:
         result = constructor_fn()
         if result is None:
           return None
@@ -205,6 +211,7 @@ def acquire(
       self,
       key,  # type: Text
       constructor_fn,  # type: Callable[[], Any]
+      tag=None  # type: Any
   ):
     # type: (...) -> Any
 
@@ -216,6 +223,9 @@ def acquire(
         present in the cache. This function should take no arguments. It should
         return an initialised object, or None if the object could not be
         initialised / constructed.
+      tag: an optional indentifier to store with the cached object. If
+        subsequent calls to acquire use different tags, the object will be
+        reloaded rather than returned from cache.
 
     Returns:
       A reference to the initialised object, either from the cache, or
@@ -227,7 +237,7 @@ def acquire(
         control_block = _SharedControlBlock()
         self._cache_map[key] = control_block
 
-    result = control_block.acquire(constructor_fn)
+    result = control_block.acquire(constructor_fn, tag)
 
     # Because we release the lock in between, if we acquire multiple Shareds
     # in a short time, there's no guarantee as to which one will be kept alive.
@@ -266,7 +276,8 @@ def __init__(self):
 
   def acquire(
       self,
-      constructor_fn  # type: Callable[[], Any]
+      constructor_fn,  # type: Callable[[], Any]
+      tag=None  # type: Any
   ):
     # type: (...) -> Any
 
@@ -277,9 +288,12 @@ def acquire(
         present in the cache. This function should take no arguments. It should
         return an initialised object, or None if the object could not be
         initialised / constructed.
+      tag: an optional indentifier to store with the cached object. If
+        subsequent calls to acquire use different tags, the object will be
+        reloaded rather than returned from cache.
 
     Returns:
       A reference to an initialised object, either from the cache, or
       newly-constructed.
     """
-    return _shared_map.acquire(self._key, constructor_fn)
+    return _shared_map.acquire(self._key, constructor_fn, tag)

sdks/python/apache_beam/utils/shared_test.py

@@ -236,6 +236,28 @@ def dummy_acquire_fn():
     self.assertEqual('sequence3', f3.get_name())
     self.assertEqual('sequence4', s3.get_name())
 
+  def testTagCacheEviction(self):
+    shared1 = shared.Shared()
+    shared2 = shared.Shared()
+
+    def acquire_fn_1():
+      return NamedObject('obj_1')
+
+    def acquire_fn_2():
+      return NamedObject('obj_2')
+
+    # with no tag, shared handle does not know when to evict objects
+    p1 = shared1.acquire(acquire_fn_1)
+    assert p1.get_name() == 'obj_1'
+    p2 = shared1.acquire(acquire_fn_2)
+    assert p2.get_name() == 'obj_1'
+
+    # cache eviction can be forced by specifying different tags
+    p1 = shared2.acquire(acquire_fn_1, tag='1')
+    assert p1.get_name() == 'obj_1'
+    p2 = shared2.acquire(acquire_fn_2, tag='2')
+    assert p2.get_name() == 'obj_2'
+
 
 if __name__ == '__main__':
   unittest.main()