add unittest for partial order

Author: lzle

jobq/jobq.py

@@ -22,6 +22,10 @@ class Finish(object):
     pass
 
 
+class NotMatch(object):
+    pass
+
+
 class JobWorkerError(Exception):
     pass
 
@@ -176,18 +180,18 @@ def _exec(self, input_q, output_q, thread_index):
                 with self.probe['probe_lock']:
                     self.probe['out'] += 1
 
-            if self.partial_order:
+            # If rst is an iterator, it procures more than one args to next job.
+            # In order to be accurate, we only count an iterator as one.
+
+            _put_rst(output_q, rst)
+
+            if self.partial_order and args is not None:
                 k = str(args[0])
                 with self.buffer_lock:
                     if k in self.in_working:
                         del self.in_working[k]
                         self.buffer_queue.not_match.set()
 
-            # If rst is an iterator, it procures more than one args to next job.
-            # In order to be accurate, we only count an iterator as one.
-
-            _put_rst(output_q, rst)
-
     def _exec_in_order(self, input_q, output_q, thread_index):
 
         while self.running:
@@ -270,16 +274,22 @@ def check_expired():
             args = self.buffer_queue.get(exclude=self.in_working)
             if args is Finish:
                 return
-            elif args is None:
+            elif args is NotMatch:
                 self.buffer_queue.not_match.wait()
+                self.buffer_queue.not_match.clear()
             else:
-                with self.buffer_lock:
-                    self.in_working[str(args[0])] = time.time()
+                if args is not None:
+                    with self.buffer_lock:
+                        self.in_working[str(args[0])] = time.time()
+
                 _put_rst(self.input_queue, args)
 
 class JobManager(object):
 
-    def __init__(self, workers, queue_size=1024, expire=3000, probe=None, keep_order=False, partial_order=False):
+    def __init__(self, workers, queue_size=1024, expire=None, probe=None, keep_order=False, partial_order=False):
+
+        if expire is None:
+            expire = 60 * 5
 
         if probe is None:
             probe = {}
@@ -489,10 +499,8 @@ def get(self, block=True, timeout=None, exclude=None):
                         raise Queue.Empty
                     self.not_empty.wait(remaining)
             item = self._get(exclude)
-            if item is not None:
+            if item is not NotMatch:
                 self.not_full.notify()
-            else:
-                self.not_match.clear()
             return item
 
     def _qsize(self):
@@ -502,8 +510,10 @@ def _put(self, item):
         self.queue.append(item)
 
     def _get(self, exclude):
-        item = None
+        item = NotMatch
         index = None
+        selected = {}
+
         for i, v in enumerate(self.queue):
             if v is Finish:
                 if i == 0:
@@ -514,9 +524,13 @@ def _get(self, exclude):
                 break
             else:
                 key = str(v[0])
+                if key in selected:
+                    continue
                 if not exclude or key not in exclude:
                     index, item = i, v
                     break
+                selected[key] = True
+
         if index is not None:
             del self.queue[index]
         return item

jobq/t.sh

@@ -1,6 +1,6 @@
 #!/bin/sh
 
-python2 -m unittest -v discover
+python2 -m unittest discover -v
 
 # python2 -m unittest -v test.test_jobq
 # python2 -m unittest -v test.test_jobq.TestProbe

jobq/test/test_jobq.py

@@ -1,9 +1,13 @@
+#!/usr/bin/env python2
+# coding: utf-8
+
 import logging
 import random
 import threading
 import time
 import unittest
 
+
 from pykit import jobq
 from pykit import threadutil
 from pykit import ututil
@@ -36,9 +40,25 @@ def sleep_5(args):
     return args
 
 
+def is_partial_order(rst):
+    cache = {}
+
+    for item in rst:
+        k = item[0]
+        if k not in cache:
+            cache[k] = item
+
+        last = cache[k]
+        if last[1] > item[1]:
+            return False
+        cache[k] = item
+
+    return True
+
+
 class TestProbe(unittest.TestCase):
 
-    def _start_jobq_in_thread(self, n_items, n_worker, keep_order=False):
+    def _start_jobq_in_thread(self, items, n_worker, keep_order=False, partial_order=False):
 
         def _sleep_1(args):
             time.sleep(0.1)
@@ -48,11 +68,12 @@ def _nothing(args):
             return args
 
         probe = {}
-        th = threading.Thread(target=lambda: jobq.run(range(n_items),
+        th = threading.Thread(target=lambda: jobq.run(items,
                                                       [(_sleep_1, n_worker),
                                                        _nothing],
                                                       probe=probe,
                                                       keep_order=keep_order,
+                                                      partial_order=partial_order,
                                                       ))
         th.daemon = True
         th.start()
@@ -67,13 +88,12 @@ def test_probe_single_thread(self):
                 (0.2,  0, 'all done'),
         )
 
-        th, probe = self._start_jobq_in_thread(3, 1)
+        th, probe = self._start_jobq_in_thread(range(3), 1)
 
         for sleep_time, doing, case_mes in cases:
 
             time.sleep(sleep_time)
             stat = jobq.stat(probe)
-
             self.assertEqual(doing, stat['doing'], case_mes)
 
             # qsize() is not reliable. do not test the value of it.
@@ -100,7 +120,7 @@ def test_probe_3_thread(self):
                 (0.4,  0, 'all done'),
         )
 
-        th, probe = self._start_jobq_in_thread(10, 3)
+        th, probe = self._start_jobq_in_thread(range(10), 3)
 
         for sleep_time, doing, case_mes in cases:
 
@@ -124,7 +144,7 @@ def test_probe_3_thread_keep_order(self):
                 (0.4,  0, 'all done'),
         )
 
-        th, probe = self._start_jobq_in_thread(10, 3, keep_order=True)
+        th, probe = self._start_jobq_in_thread(range(10), 3, keep_order=True)
 
         for sleep_time, doing, case_mes in cases:
 
@@ -140,6 +160,28 @@ def test_probe_3_thread_keep_order(self):
 
         th.join()
 
+    def test_probe_3_thread_partial_order(self):
+        cases = (
+            (0.05, 3, '_sleep_1 is working on 1st 3 items'),
+            (0.1, 3, '_sleep_1 is working on 2nd 3 items'),
+            (0.4, 0, 'all done'),
+        )
+
+        th, probe = self._start_jobq_in_thread(([i, 0] for i in range(10)), 3, partial_order=True)
+
+        for sleep_time, doing, case_mes in cases:
+            time.sleep(sleep_time)
+            stat = jobq.stat(probe)
+
+            self.assertEqual(doing, stat['doing'], case_mes)
+
+        # use the last stat
+
+        workers = stat['workers']
+        self.assertEqual(2, len(workers))
+
+        th.join()
+
 
 class TestDispatcher(unittest.TestCase):
 
@@ -264,6 +306,37 @@ def collect(args):
         self.assertEqual([0], rst)
 
 
+class TestExpire(unittest.TestCase):
+
+    def test_expire(self):
+
+        def _sleep_1(args):
+            sleep_got.append(args)
+            time.sleep(0.5)
+            return args
+
+        def collect(args):
+            rst.append(args)
+
+        rst = []
+        sleep_got = []
+
+        jm = jobq.JobManager([(_sleep_1, 3), collect], expire=0.2, partial_order=True)
+
+        n = 6
+        for i in range(n):
+            jm.put((i % 2, i))
+
+        time.sleep(0.1)
+        self.assertEqual(set([(0, 0), (1, 1)]), set(sleep_got))
+
+        time.sleep(0.3)
+        self.assertEqual(3, len(sleep_got))
+
+        jm.join()
+        self.assertEqual(set([(i % 2, i) for i in range(n)]), set(rst))
+
+
 class TestJobManager(unittest.TestCase):
 
     def test_manager(self):
@@ -434,6 +507,42 @@ def _change_thread_nr():
         for th in ths:
             th.join()
 
+    def test_set_thread_num_partial_order(self):
+
+        def _pass(args):
+            return args
+
+        rst = []
+
+        jm = jobq.JobManager([_pass, rst.append], partial_order=True)
+
+        setter = {'running': True}
+
+        def _change_thread_nr():
+            while setter['running']:
+                jm.set_thread_num(_pass, random.randint(1, 4))
+                time.sleep(0.5)
+
+        ths = []
+        for ii in range(3):
+            th = threadutil.start_daemon_thread(_change_thread_nr)
+            ths.append(th)
+
+        n = 10240
+        for i in range(n):
+            jm.put([i % 2, i])
+
+        jm.join()
+
+        self.assertEqual(n, len(rst))
+
+        self.assertTrue(is_partial_order(rst))
+
+        setter['running'] = False
+
+        for th in ths:
+            th.join()
+
 
 class TestJobQ(unittest.TestCase):
 
@@ -506,6 +615,36 @@ def collect(args):
             jobq.run(inp, workers + [collect], keep_order=True)
             self.assertEqual(out, rst)
 
+        def add_list(args):
+            args[1] += 1
+            return args
+
+        def multi2_list(args):
+            args[1] *= 2
+            return args
+
+        def multi2_list_sleep(args):
+            time.sleep(0.02)
+            args[1] *= 2
+            return args
+
+        cases = (
+                (list([i % 2, i] for i in range(100)), [add_list, (multi2_list_sleep, 10)],
+                 list([i % 2, (i + 1) * 2] for i in range(100))
+                 ),
+                (list([i % 50, i] for i in range(1024 * 10)), [add_list, (multi2_list, 10)],
+                 list([i % 50, (i + 1) * 2] for i in range(1024 * 10))
+                 ),
+        )
+        for inp, workers, out in cases:
+            rst = []
+            jobq.run(inp, workers + [collect], partial_order=True)
+            self.assertTrue(is_partial_order(rst))
+
+            out.sort()
+            rst.sort()
+            self.assertEqual(out, rst)
+
     def test_generator(self):
 
         def gen(args):
@@ -514,6 +653,7 @@ def gen(args):
                 time.sleep(0.1)
 
         def collect(args):
+            time.sleep(random.uniform(0.005, 0.02))
             rst.append(args)
 
         rst = []
@@ -528,6 +668,17 @@ def collect(args):
 
         self.assertEqual(9, len(rst), 'nr of elts')
 
+        def _gen(args):
+            k = args[0]
+            for i in range(3):
+                yield (k, i)
+
+        rst = []
+        jobq.run([(k, 0) for k in range(3)], [(_gen, 3), (collect, 3)], partial_order=True)
+        self.assertEqual(set([(k, v) for k in range(3) for v in range(3)]), set(rst),
+                         "generator should get all")
+        self.assertTrue(is_partial_order(rst))
+
 
 class TestDefaultTimeout(unittest.TestCase):