async_dofn_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.
#

import logging
import random
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
from threading import Lock

import apache_beam as beam
import apache_beam.transforms.async_dofn as async_lib


class BasicDofn(beam.DoFn):
  def __init__(self, sleep_time=0):
    self.processed = 0
    self.sleep_time = sleep_time
    self.lock = Lock()

  def process(self, element):
    logging.info('start processing element %s', element)
    time.sleep(self.sleep_time)
    self.lock.acquire()
    self.processed += 1
    self.lock.release()
    logging.info('finish processing element %s', element)
    yield element

  def getProcessed(self):
    self.lock.acquire()
    result = self.processed
    self.lock.release()
    return result


# Used for testing multiple elements produced by a single input element.
class MultiElementDoFn(beam.DoFn):
  def finish_bundle(self):
    yield ('key', 'bundle end')

  def process(self, element):
    yield element
    yield element


class FakeBagState:
  def __init__(self, items):
    self.items = items
    # Normally SE would have a lock on the BT row protecting this from multiple
    # updates. Here without SE we must lock ourselvs.
    self.lock = Lock()

  def add(self, item):
    with self.lock:
      self.items.append(item)

  def clear(self):
    with self.lock:
      self.items = []

  def read(self):
    with self.lock:
      return self.items.copy()


class FakeTimer:
  def __init__(self, time):
    self.time = time

  def set(self, time):
    self.time = time


class _AsyncTestBase:
  use_asyncio: bool

  def setUp(self):
    super().setUp()
    async_lib.AsyncWrapper.reset_state()

  def wait_for_empty(self, async_dofn, timeout=10):
    count = 0
    increment = 1
    while not async_dofn.is_empty():
      time.sleep(1)
      count += increment
      if count > timeout:
        raise TimeoutError('Timed out waiting for async dofn to be empty')
    # Usually we are waiting for an item to be ready to output not just no
    # longer in the queue. These are not the same atomic operation. Sleep a bit
    # to make sure the item is also ready to output.
    time.sleep(1)

  def check_output(self, result, expected_output):
    output = []
    for x in result:
      output.append(x)
    # Ordering is not guaranteed so sort both lists before comparing.
    output.sort()
    expected_output.sort()
    self.assertEqual(output, expected_output)

  def check_items_in_buffer(self, async_dofn, expected_count):
    self.assertEqual(
        async_lib.AsyncWrapper._items_in_buffer[async_dofn._uuid],
        expected_count,
    )

  def test_custom_id_fn(self):
    class CustomIdObject:
      def __init__(self, element_id, value):
        self.element_id = element_id
        self.value = value

      def __hash__(self):
        return hash(self.element_id)

      def __eq__(self, other):
        return self.element_id == other.element_id

    dofn = BasicDofn()
    async_dofn = async_lib.AsyncWrapper(
        dofn, id_fn=lambda x: x.element_id, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg1 = ('key1', CustomIdObject(1, 'a'))
    msg2 = ('key1', CustomIdObject(1, 'b'))

    result = async_dofn.process(
        msg1, to_process=fake_bag_state, timer=fake_timer)
    self.assertEqual(result, [])

    # The second message should be a no-op as it has the same id.
    result = async_dofn.process(
        msg2, to_process=fake_bag_state, timer=fake_timer)
    self.assertEqual(result, [])

    self.wait_for_empty(async_dofn)
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [('key1', msg1[1])])
    self.assertEqual(fake_bag_state.items, [])

  def test_basic(self):
    # Setup an async dofn and send a message in to process.
    dofn = BasicDofn()
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg = ('key1', 1)
    result = async_dofn.process(
        msg, to_process=fake_bag_state, timer=fake_timer)

    # nothing should be produced because this is async.
    self.assertEqual(result, [])
    # there should be an element to process in the bag state.
    self.assertEqual(fake_bag_state.items, [msg])
    # The timer should be set.
    self.assertNotEqual(fake_timer.time, 0)

    # Give time for the async element to process.
    self.wait_for_empty(async_dofn)
    # Fire the (fake) timer and check that the input message has produced output
    # and the message has been removed from bag state.
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg])
    self.assertTrue(dofn.getProcessed() == 1)
    self.assertEqual(fake_bag_state.items, [])

  def test_multi_key(self):
    # Send in two messages with different keys..
    dofn = BasicDofn()
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state_key1 = FakeBagState([])
    fake_bag_state_key2 = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg1 = ('key1', 1)
    msg2 = ('key2', 2)
    async_dofn.process(msg1, to_process=fake_bag_state_key1, timer=fake_timer)
    async_dofn.process(msg2, to_process=fake_bag_state_key2, timer=fake_timer)
    self.wait_for_empty(async_dofn)

    # The result should depend on which key the timer fires for. We are firing
    # 'key2' first so expect its message to be produced and msg1 to remain in
    # state.
    result = async_dofn.commit_finished_items(fake_bag_state_key2, fake_timer)
    self.check_output(result, [msg2])
    self.assertEqual(fake_bag_state_key1.items, [msg1])
    self.assertEqual(fake_bag_state_key2.items, [])

    # Now fire 'key1' and expect its message to be produced.
    result = async_dofn.commit_finished_items(fake_bag_state_key1, fake_timer)
    self.check_output(result, [msg1])
    self.assertEqual(fake_bag_state_key1.items, [])
    self.assertEqual(fake_bag_state_key2.items, [])

  def test_long_item(self):
    # Test that everything still works with a long running time for the dofn.
    dofn = BasicDofn(sleep_time=5)
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg = ('key1', 1)
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)

    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [])
    self.assertTrue(dofn.getProcessed() == 0)
    self.assertEqual(fake_bag_state.items, [msg])

    self.wait_for_empty(async_dofn, 20)

    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg])
    self.assertTrue(dofn.getProcessed() == 1)
    self.assertEqual(fake_bag_state.items, [])

  def test_lost_item(self):
    # Setup an element in the bag stat thats not in processing state.
    # The async dofn should reschedule this element.
    dofn = BasicDofn()
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_timer = FakeTimer(0)
    msg = ('key1', 1)
    fake_bag_state = FakeBagState([msg])

    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [])
    # The element should be rescheduled and still processed.
    self.wait_for_empty(async_dofn)
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg])

  def test_cancelled_item(self):
    # Test that an item gets removed for processing and does not get output when
    # it is not present in the bag state. Either this item moved or a commit
    # failed making the local state and bag stat inconsistent.
    dofn = BasicDofn()
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    msg = ('key1', 1)
    msg2 = ('key1', 2)
    fake_timer = FakeTimer(0)
    fake_bag_state = FakeBagState([])
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    async_dofn.process(msg2, to_process=fake_bag_state, timer=fake_timer)

    self.wait_for_empty(async_dofn)
    fake_bag_state.clear()
    fake_bag_state.add(msg2)
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg2])
    self.assertEqual(fake_bag_state.items, [])

  def test_multi_element_dofn(self):
    # Test that async works when a dofn produces multiple elements in process
    # and finish_bundle.
    dofn = MultiElementDoFn()
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg = ('key1', 1)
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)

    self.wait_for_empty(async_dofn)

    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg, msg, ('key', 'bundle end')])
    self.assertEqual(fake_bag_state.items, [])

  def test_duplicates(self):
    # Test that async will produce a single output when a given input is sent
    # multiple times.
    dofn = BasicDofn(5)
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg = ('key1', 1)
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    # SE will only deliver the message again if the commit failed and so bag
    # state should still be empty.
    fake_bag_state.clear()
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    self.assertEqual(fake_bag_state.items, [msg])

    self.wait_for_empty(async_dofn)
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg])
    self.assertEqual(fake_bag_state.items, [])

  def test_slow_duplicates(self):
    # Test that async will produce a single output when a given input is sent
    # multiple times.
    dofn = BasicDofn(5)
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_bag_state = FakeBagState([])
    fake_timer = FakeTimer(0)
    msg = ('key1', 1)
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    time.sleep(10)
    # SE will only deliver the message again if the commit failed and so bag
    # state should still be empty.
    fake_bag_state.clear()
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [])
    self.assertEqual(fake_bag_state.items, [])

    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    self.assertEqual(fake_bag_state.items, [msg])
    self.wait_for_empty(async_dofn)

    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_output(result, [msg])
    self.assertEqual(fake_bag_state.items, [])

  def test_buffer_count(self):
    # Test that the buffer count is correctly incremented when adding items.
    dofn = BasicDofn(5)
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    msg = ('key1', 1)
    fake_timer = FakeTimer(0)
    fake_bag_state = FakeBagState([])
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    self.check_items_in_buffer(async_dofn, 1)
    self.wait_for_empty(async_dofn)
    self.check_items_in_buffer(async_dofn, 0)

    # Commiting already finished items should not change the count.
    async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_items_in_buffer(async_dofn, 0)

  def test_buffer_stops_accepting_items(self):
    # Test that the buffer stops accepting items when it is full.
    dofn = BasicDofn(5)
    async_dofn = async_lib.AsyncWrapper(
        dofn,
        parallelism=1,
        max_items_to_buffer=5,
        use_asyncio=self.use_asyncio)
    async_dofn.setup()
    fake_timer = FakeTimer(0)
    fake_bag_state = FakeBagState([])

    # Send more elements than the buffer can hold.
    pool = ThreadPoolExecutor(max_workers=10)
    expected_output = []
    for i in range(0, 10):

      def add_item(i):
        item = ('key', i)
        expected_output.append(item)
        pool.submit(
            lambda: async_dofn.process(
                item, to_process=fake_bag_state, timer=fake_timer))

      add_item(i)

    # Assert that the buffer stops accepting items at its limit.
    time.sleep(1)
    self.assertEqual(async_dofn._max_items_to_buffer, 5)
    self.check_items_in_buffer(async_dofn, 5)

    # Wait for all buffered items to finish.
    self.wait_for_empty(async_dofn, 100)
    # This will commit buffered items and add new items which didn't fit in the
    # buffer.
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)

    # Wait for the new buffered items to finish.
    self.wait_for_empty(async_dofn, 100)
    result.extend(async_dofn.commit_finished_items(fake_bag_state, fake_timer))
    self.check_output(result, expected_output)
    self.check_items_in_buffer(async_dofn, 0)

  def test_buffer_with_cancellation(self):
    dofn = BasicDofn(3)
    async_dofn = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    msg = ('key1', 1)
    msg2 = ('key1', 2)
    fake_timer = FakeTimer(0)
    fake_bag_state = FakeBagState([])
    async_dofn.process(msg, to_process=fake_bag_state, timer=fake_timer)
    async_dofn.process(msg2, to_process=fake_bag_state, timer=fake_timer)
    self.check_items_in_buffer(async_dofn, 2)
    self.assertEqual(fake_bag_state.items, [msg, msg2])

    # Remove one item from the bag state to simulate the user worker being out
    # of sync with SE state. The item should be cancelled and also removed from
    # the buffer.
    fake_bag_state.clear()
    fake_bag_state.add(msg2)
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    # Note buffer may still be 2 as msg1 can start processing before we get to
    # cancel it but also not be finished yet.
    self.check_output(result, [])
    self.assertEqual(fake_bag_state.items, [msg2])

    self.wait_for_empty(async_dofn)
    result = async_dofn.commit_finished_items(fake_bag_state, fake_timer)
    self.check_items_in_buffer(async_dofn, 0)
    self.check_output(result, [msg2])
    self.assertEqual(fake_bag_state.items, [])

  def test_load_correctness(self):
    # Test AsyncDofn over heavy load.
    dofn = BasicDofn(1)
    max_sleep = 10
    async_dofn = async_lib.AsyncWrapper(
        dofn, max_wait_time=max_sleep, use_asyncio=self.use_asyncio)
    async_dofn.setup()
    bag_states = {}
    timers = {}
    expected_outputs = {}
    pool = ThreadPoolExecutor(max_workers=10)
    futures = []

    for i in range(0, 10):
      bag_states['key' + str(i)] = FakeBagState([])
      timers['key' + str(i)] = FakeTimer(0)
      expected_outputs['key' + str(i)] = []

    for i in range(0, 100):

      def add_item(i):
        item = ('key' + str(random.randint(0, 9)), i)
        expected_outputs[item[0]].append(item)
        futures.append(
            pool.submit(
                lambda: async_dofn.process(
                    item, to_process=bag_states[item[0]], timer=timers[item[0]])
            ))

      add_item(i)
      time.sleep(random.random())

    # Run for a while. Should be enough to start all items but not finish them
    # all.
    time.sleep(random.randint(30, 50))

    done = False
    results = [[] for _ in range(0, 10)]
    while not done:
      done = True
      for i in range(0, 10):
        results[i].extend(
            async_dofn.commit_finished_items(
                bag_states['key' + str(i)], timers['key' + str(i)]))
        if not bag_states['key' + str(i)].items:
          self.check_output(results[i], expected_outputs['key' + str(i)])
        else:
          done = False
      time.sleep(random.randint(10, 30))

    for i in range(0, 10):
      self.check_output(results[i], expected_outputs['key' + str(i)])
      self.assertEqual(bag_states['key' + str(i)].items, [])


class AsyncTestThreadPool(_AsyncTestBase, unittest.TestCase):
  use_asyncio = False


class AsyncTestAsyncio(_AsyncTestBase, unittest.TestCase):
  use_asyncio = True


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