github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/shared_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Test for Shared class."""

import gc
import threading
import time
import unittest

from apache_beam.utils import shared


class Count(object):
  def __init__(self):
    self._lock = threading.Lock()
    self._total = 0
    self._active = 0

  def add_ref(self):
    with self._lock:
      self._total += 1
      self._active += 1

  def release_ref(self):
    with self._lock:
      self._active -= 1

  def get_active(self):
    with self._lock:
      return self._active

  def get_total(self):
    with self._lock:
      return self._total


class Marker(object):
  def __init__(self, count):
    self._count = count
    self._count.add_ref()

  def __del__(self):
    self._count.release_ref()


class NamedObject(object):
  def __init__(self, name):
    self._name = name

  def get_name(self):
    return self._name


class Sequence(object):
  def __init__(self):
    self._sequence = 0

  def make_acquire_fn(self):
    # Every time acquire_fn is called, increases the sequence number and returns
    # a NamedObject with that sequenece number.
    def acquire_fn():
      self._sequence += 1
      return NamedObject('sequence%d' % self._sequence)

    return acquire_fn


class SharedTest(unittest.TestCase):
  def testKeepalive(self):
    count = Count()
    shared_handle = shared.Shared()
    other_shared_handle = shared.Shared()

    def dummy_acquire_fn():
      return None

    def acquire_fn():
      return Marker(count)

    p1 = shared_handle.acquire(acquire_fn)
    self.assertEqual(1, count.get_total())
    self.assertEqual(1, count.get_active())
    del p1
    gc.collect()
    # Won't be garbage collected, because of the keep-alive
    self.assertEqual(1, count.get_active())

    # Reacquire.
    p2 = shared_handle.acquire(acquire_fn)
    self.assertEqual(1, count.get_total())  # No reinitialisation.
    self.assertEqual(1, count.get_active())

    # Get rid of the keepalive
    other_shared_handle.acquire(dummy_acquire_fn)
    del p2
    gc.collect()
    self.assertEqual(0, count.get_active())

  def testMultiple(self):
    count = Count()
    shared_handle = shared.Shared()
    other_shared_handle = shared.Shared()

    def dummy_acquire_fn():
      return None

    def acquire_fn():
      return Marker(count)

    p = shared_handle.acquire(acquire_fn)
    other_shared_handle.acquire(dummy_acquire_fn)  # Get rid of the keepalive
    self.assertEqual(1, count.get_total())
    self.assertEqual(1, count.get_active())
    del p
    gc.collect()
    self.assertEqual(0, count.get_active())
    # Shared value should be garbage collected.

    # Acquiring multiple times only results in one initialisation
    p1 = shared_handle.acquire(acquire_fn)
    # Since shared value was released, expect a reinitialisation.
    self.assertEqual(2, count.get_total())
    self.assertEqual(1, count.get_active())
    p2 = shared_handle.acquire(acquire_fn)
    self.assertEqual(2, count.get_total())
    self.assertEqual(1, count.get_active())

    other_shared_handle.acquire(dummy_acquire_fn)  # Get rid of the keepalive

    # Check that shared object isn't destroyed if there's still a reference to
    # it.
    del p2
    gc.collect()
    self.assertEqual(1, count.get_active())

    del p1
    gc.collect()
    self.assertEqual(0, count.get_active())

  def testConcurrentCallsDeduped(self):
    # Test that only one among many calls to acquire will actually run the
    # initialisation function.

    count = Count()
    shared_handle = shared.Shared()
    other_shared_handle = shared.Shared()

    refs = []
    ref_lock = threading.Lock()

    def dummy_acquire_fn():
      return None

    def acquire_fn():
      time.sleep(1)
      return Marker(count)

    def thread_fn():
      p = shared_handle.acquire(acquire_fn)
      with ref_lock:
        refs.append(p)

    threads = []
    for _ in range(100):
      t = threading.Thread(target=thread_fn)
      threads.append(t)
      t.start()

    for t in threads:
      t.join()

    self.assertEqual(1, count.get_total())
    self.assertEqual(1, count.get_active())

    other_shared_handle.acquire(dummy_acquire_fn)  # Get rid of the keepalive

    with ref_lock:
      del refs[:]
    gc.collect()

    self.assertEqual(0, count.get_active())

  def testDifferentObjects(self):
    sequence = Sequence()

    def dummy_acquire_fn():
      return None

    first_handle = shared.Shared()
    second_handle = shared.Shared()
    dummy_handle = shared.Shared()

    f1 = first_handle.acquire(sequence.make_acquire_fn())
    s1 = second_handle.acquire(sequence.make_acquire_fn())

    self.assertEqual('sequence1', f1.get_name())
    self.assertEqual('sequence2', s1.get_name())

    f2 = first_handle.acquire(sequence.make_acquire_fn())
    s2 = second_handle.acquire(sequence.make_acquire_fn())

    # Check that the repeated acquisitions return the earlier objects
    self.assertEqual('sequence1', f2.get_name())
    self.assertEqual('sequence2', s2.get_name())

    # Release all references and force garbage-collection
    del f1
    del f2
    del s1
    del s2
    dummy_handle.acquire(dummy_acquire_fn)  # Get rid of the keepalive
    gc.collect()

    # Check that acquiring again after they're released gives new objects
    f3 = first_handle.acquire(sequence.make_acquire_fn())
    s3 = second_handle.acquire(sequence.make_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'

  def testTagReturnsCached(self):
    sequence = Sequence()
    handle = shared.Shared()

    f1 = handle.acquire(sequence.make_acquire_fn(), tag='1')
    self.assertEqual('sequence1', f1.get_name())

    # should return cached
    f1 = handle.acquire(sequence.make_acquire_fn(), tag='1')
    self.assertEqual('sequence1', f1.get_name())


if __name__ == '__main__':
  unittest.main()