def _create_closure(self, cancellation_mgr): @def_function.function() def some_function(): return 1.0 return coordinator_lib.Closure(some_function, cancellation_mgr)
def testProcessAtLeaseOnce(self): closure_queue = coordinator_lib._CoordinatedClosureQueue() labels = ['A', 'B', 'C', 'D', 'E'] processed_count = collections.defaultdict(int) coord = coordinator.Coordinator(clean_stop_exception_types=[]) def process_queue(): with coord.stop_on_exception(): has_been_put_back = False while True: closure = closure_queue.get(timeout=30) if closure is None: break if not has_been_put_back: has_been_put_back = True closure_queue.put_back(closure) continue closure._function() closure_queue.mark_finished() def get_func(label): def func(): time.sleep(3) processed_count[label] += 1 return func cm = cancellation.CancellationManager() for label in labels: closure_queue.put(coordinator_lib.Closure(get_func(label), cm)) t1 = threading.Thread(target=process_queue, daemon=True) t1.start() t2 = threading.Thread(target=process_queue, daemon=True) t2.start() # Make sure multiple wait() calls are fine. closure_queue.wait() closure_queue.wait() closure_queue.wait() closure_queue.wait() self.assertEqual(processed_count, collections.Counter(labels)) coord.join([t1, t2])
def testNotifyBeforeWait(self): closure_queue = coordinator_lib._CoordinatedClosureQueue() def func(): logging.info('func running') coord = coordinator.Coordinator(clean_stop_exception_types=[]) def process_queue(): with coord.stop_on_exception(): closure_queue.get() closure_queue.mark_finished() closure_queue.put( coordinator_lib.Closure(func, closure_queue._cancellation_mgr)) t = threading.Thread(target=process_queue) t.start() coord.join([t]) # This test asserts that waiting at the time the function has been processed # doesn't time out. closure_queue.wait()