(wip) use rabbitmqdevelop2.0.0

3 files changed, 62 insertions, 412 deletions

diff --git a/joblib.py b/joblib.py
index 04a37d0..4cc3b94 100644
--- a/joblib.py
+++ b/joblib.py
@@ -1,206 +1,84 @@
 #!/usr/bin/python3
 
-import asyncio
-import asyncpg
 import json
 import logging
+import pika
 
 
 class JobQueue(object):
     def __init__(self):
-        self._con = None
-        self._listens = {}
+        self._qcon = None
+        self._qchan = None
 
-    def _listener(self, con, pid,
-                        channel, payload):
-        """Fire the registered callback. Must not be async."""
-        self._listens[channel](payload)
+    def connect(self, **args):
+        """Connect to RabbitMQ."""
+        logging.info("connecting to broker")
+        self._qcon = pika.BlockingConnection(pika.ConnectionParameters(
+            **args))
+        self._qchan = self._qcon.channel()
 
-    async def _hook_listener(self):
-        """Register callbacks."""
-        if len(self._listens.keys()) > 0:
-            for channel in self._listens.keys():
-                await self._con.add_listener(channel,
-                                             self._listener)
-
-    async def connect(self, **args):
-        """Connect the database."""
-        logging.info("connecting to queue database")
-        while True:
-            try:
-                self._con = await asyncpg.connect(**args)
-                await self._hook_listener()
-                break
-            except asyncpg.exceptions.CannotConnectNowError:
-                logging.warning(
-                    "queue database is not ready yet, retrying")
-                await asyncio.sleep(1)
-            except asyncpg.exceptions.TooManyConnectionsError:
-                logging.warning(
-                    "queue database has too many clients, retrying")
-                await asyncio.sleep(1)
-
-    async def setup(self):
-        """Initialize the database."""
-        await self._con.execute("""
-            CREATE TABLE IF NOT EXISTS
-            jobs (
-                id SERIAL PRIMARY KEY,
-                priority INT DEFAULT 0,
-                status TEXT DEFAULT 'open',
-                type TEXT,
-                payload JSONB
-            )
-        """)
-        await self._con.execute("""
-            CREATE UNIQUE INDEX
-            IF NOT EXISTS
-            jobs_priority_id_idx
-            ON jobs(priority, id)
-        """)
-
-    async def listen(self, channel, callback):
-        """Hook a notification listener."""
-        self._listens[channel] = callback
-
-    async def request(self, jobtype, payload, priority=0):
+    def request(self, queue, payload):
         """Create a new job and return its id."""
-        if isinstance(payload, list):
-            payloads = [json.dumps(p)
-                        for p in payload]
-        else:
-            payloads = [json.dumps(payload)]
-        if isinstance(priority, list):
-            priorities = priority
-        else:
-            priorities = [priority] * len(payloads)
-        insert = await self._con.prepare("""
-            INSERT INTO jobs(type, payload, priority)
-            VALUES($1, $2, $3)
-            RETURNING id
-        """)
-        ids = []
-        async with self._con.transaction():
-            for pl, pr in zip(payloads, priorities):
-                ids.append(await insert.fetchval(jobtype, pl, pr))
-            await self._con.execute("SELECT pg_notify($1 || '_open', '')",
-                                    jobtype)
+        body = json.dumps(payload)
+        self._qchan.queue_declare(queue=queue, durable=True)
+        self._qchan.basic_publish(exchange="",
+                                  routing_key=queue,
+                                  body=body,
+                                  properties=pika.BasicProperties(
+                                      delivery_mode = 2
+                                  ))
 
-        if isinstance(payload, list):
-            return ids
-        else:
-            return ids[0]
+class JobFailed(Exception):
+    pass
 
-    async def acquire(self, jobtype, length=None):
-        """Mark a job as running, return id, payload and priority.
-        Return (None, None, None) if no job is available."""
-        if length is None:
-            limit = 1
-        else:
-            limit = length
-        while True:
-            try:
-                # do not allow async access
-                async with self._con.transaction(isolation="serializable"):
-                    result = await self._con.fetch("""
-                        UPDATE jobs SET STATUS='running'
-                        FROM (
-                            SELECT id FROM jobs
-                            WHERE status='open' AND type=$1
-                            ORDER BY priority, id
-                            LIMIT $2
-                        ) AS open_jobs
-                        WHERE jobs.id=open_jobs.id
-                        RETURNING jobs.id, jobs.payload, jobs.priority
-                    """, jobtype, limit)
-                    if len(result) == 0 and length is None:
-                        # no jobs available
-                        # backwards compatibility
-                        return None, None, None
 
-                    jobs = [(r[0], json.loads(r[1]), r[2]) for r in result]
+class Worker(JobQueue):
+    """Abstract service worker class."""
+    def __init__(self, jobtype):
+        super().__init__()
+        self._queue = jobtype
+        self._up = False
+        self._notifq = []
 
-                    await self._con.execute("SELECT pg_notify($1 || '_running', '')",
-                                            jobtype)
+    def setup(self):
+        # override
+        pass
 
-                    if length is None:
-                        return jobs[0]
-                    else:
-                        return jobs
-            except asyncpg.exceptions.SerializationError:
-                # job is being picked up by another worker, try again
-                pass
+    def work(self, payload):
+        # override
+        pass
 
-    async def status(self, jobid):
-        """Return the status of a job."""
-        return await self._con.fetchval("""
-            SELECT status
-            FROM jobs WHERE
-            id=$1
-        """, jobid)
+    def commit(self, failed):
+        # override
+        pass
 
-    async def finish(self, jobid, jobtype):
-        """Mark jobs as completed."""
-        if not isinstance(jobid, list):
-            jobids = [(jobid,)]
+    def request(self, queue, payload, now=True):
+        if now:
+            super().request(queue, payload)
         else:
-            jobids = [(jid,) for jid in jobid]
-        async with self._con.transaction():
-            await self._con.executemany("""
-                UPDATE jobs SET status='finished'
-                WHERE id=$1
-            """, jobids)
-            await self._con.execute("""
-                SELECT pg_notify($1 || '_finished', '')
-            """, jobtype)
+            # send after COMMIT
+            self._notifq.append((queue, payload))
 
-    async def fail(self, jobid, jobtype, reason):
-        """Mark a job as failed."""
-        if not isinstance(jobid, list):
-            jobids = [jobid]
-        else:
-            jobids = jobid
-        if not isinstance(reason, list):
-            reasons = [json.dumps({"error": reason})]
-        else:
-            reasons = [json.dumps({"error": r})
-                       for r in reason]
-        assert len(jobids) == len(reasons)
-        async with self._con.transaction():
-            await self._con.executemany("""
-                UPDATE jobs SET status='failed',
-                    payload=payload||$2::jsonb
-                WHERE id=$1
-            """, zip(jobids, reasons))
-            await self._con.execute("""
-                SELECT pg_notify($1 || '_failed', '')
-            """, jobtype)
-
-    async def reset(self, jobid, jobtype):
-        """Mark a job as open."""
-        if not isinstance(jobid, list):
-            jobids = [(jobid,)]
-        else:
-            jobids = [(jid,) for jid in jobid]
-        async with self._con.transaction():
-            await self._con.executemany("""
-                UPDATE jobs SET status='open'
-                WHERE id=$1
-            """, jobids)
-            await self._con.execute("""
-                SELECT pg_notify($1 || '_open', '')
-            """, jobtype)
+    def run(self):
+        self._qchan.queue_declare(queue=self._queue, durable=True)
+        self._qchan.basic_qos(prefetch_count=1)
+        for msg in self._qchan.consume(queue=self._queue,
+                                       inactivity_timeout=1):
+            # TODO force commit after n, c+=1
+            if msg is None:
+                # idling
+                self.commit(False)
+                # send notifs dependant on COMMIT
+                for notif in self._notifq:
+                    self.request(*notif)
+                self._notifq = []
+                continue
 
-    async def cleanup(self):
-        """Reopen all unfinished jobs."""
-        while True:
+            method, properties, body = msg
+            # run a job
+            payload = json.loads(body)
             try:
-                async with self._con.transaction(isolation="serializable"):
-                    await self._con.execute("""
-                        UPDATE jobs
-                        SET status='open'
-                        WHERE status='running'
-                    """)
-                    return
-            except asyncpg.exceptions.SerializationError:
-                pass
+                self.work(payload)
+            except JobFailed as err:
+                pass  # TODO
+            self._qchan.basic_ack(delivery_tag=method.delivery_tag)
diff --git a/test_joblib.py b/test_joblib.py
deleted file mode 100644
index fc9640b..0000000
--- a/test_joblib.py
+++ /dev/null
@@ -1,116 +0,0 @@
-#!/usr/bin/python3
-
-import os
-import asyncio
-import json
-import asyncpg
-import pytest
-import joblib
-
-class TestJoblib:
-    # async fixtures are not available yet, so we use a workaround
-    async def queue_helper(self, q):
-        await q.connect(
-            host=os.environ["POSTGRESQL_HOST"],
-            port=os.environ["POSTGRESQL_PORT"],
-            user=os.environ["POSTGRESQL_USER"],
-            password=os.environ["POSTGRESQL_PASSWORD"],
-            database=os.environ["POSTGRESQL_DB"]
-        )
-        async with q._pool.acquire() as con:
-            await con.execute("DROP TABLE IF EXISTS jobs")
-
-        # clean db table
-        await q.setup()
-
-    @pytest.fixture
-    def queue(self, event_loop):
-        queue = joblib.JobQueue()
-        event_loop.run_until_complete(self.queue_helper(queue))
-        return queue
-
-    @pytest.fixture
-    def payload(self):
-        return {
-            "key": "value",
-            "dict": {
-                "foo": 1,
-                "bar": "baz"
-            }
-        }
-
-    @pytest.mark.asyncio
-    async def test_request_and_acquire(self, queue, payload):
-        # request should succeed
-        await queue.request(jobtype="testing", payload=payload)
-        # acquire should return same payload
-        assert payload == (await queue.acquire(jobtype="testing"))[1]
-        # there should not be another job
-        assert None == (await queue.acquire(jobtype="testing"))[1]
-
-    @pytest.mark.asyncio
-    async def test_request_and_acquire_batch(self, queue, payload):
-        # request should succeed
-        await queue.request(jobtype="testing", payload=[payload]*5)
-        # acquire should return same payload
-        for _ in range(5):
-            assert payload == (await queue.acquire(jobtype="testing"))[1]
-        # there should not be another job
-        assert None == (await queue.acquire(jobtype="testing"))[1]
-
-    @pytest.mark.asyncio
-    async def test_request_and_acquire_batch_batched(self, queue, payload):
-        # request should succeed
-        await queue.request(jobtype="testing", payload=[payload]*5)
-        # acquire should return same payload
-        assert payload == (await queue.acquire(jobtype="testing", length=5))[0][1]
-        # there should not be another job
-        assert None == (await queue.acquire(jobtype="testing"))[1]
-
-    @pytest.mark.asyncio
-    async def test_cleanup(self, queue, payload):
-        await queue.request(jobtype="testing", payload=payload)
-        # mark job as processing
-        jobid_1, payload_1, _ = await queue.acquire(jobtype="testing")
-        assert payload_1 == payload
-        await queue.cleanup()
-        # same job should be available again
-        jobid_2, payload_2, _ = await queue.acquire(jobtype="testing")
-        assert jobid_1 == jobid_2 and payload_1 == payload_2
-
-    @pytest.mark.asyncio
-    async def test_priority(self, queue, payload):
-        await queue.request(jobtype="testing", payload=payload, priority=9)
-        assert 9 == (await queue.acquire(jobtype="testing"))[2]
-
-    @pytest.mark.asyncio
-    async def test_fail(self, queue, payload):
-        err = "testing errors"
-        await queue.request(jobtype="testing", payload=payload)
-        jobid, _, _ = await queue.acquire(jobtype="testing")
-        await queue.fail(jobid, err)
-        async with queue._pool.acquire() as con:
-            jid, pl = await con.fetchrow(
-                "SELECT id, payload FROM jobs WHERE status='failed'")
-            assert jid == jobid
-            payload["error"] = err
-            assert json.loads(pl) == payload
-
-    @pytest.mark.asyncio
-    async def test_status(self, queue, payload):
-        assert await queue.status(-1) == None
-        jobid = await queue.request(jobtype="testing",
-                                    payload=payload)
-        assert await queue.status(jobid) == "open"
-        await queue.acquire(jobtype="testing")
-        assert await queue.status(jobid) == "running"
-        await queue.finish(jobid)
-        assert await queue.status(jobid) == "finished"
-
-    @pytest.mark.asyncio
-    async def test_finish(self, queue, payload):
-        await queue.request(jobtype="testing", payload=payload)
-        jobid, _, _ = await queue.acquire(jobtype="testing")
-        await queue.finish(jobid)
-        # job should not be available again
-        assert None == (await queue.acquire(jobtype="testing"))[1]
diff --git a/worker.py b/worker.py
deleted file mode 100644
index c27d95e..0000000
--- a/worker.py
+++ /dev/null
@@ -1,112 +0,0 @@
-#!/usr/bin/python
-
-import asyncio
-import logging
-import joblib.joblib
-
-
-class JobFailed(Exception):
-    pass
-
-
-class Worker(object):
-    """Abstract service worker class."""
-    def __init__(self, jobtype):
-        self._queue = None
-        self._jobtype = jobtype
-        self._need_update = True
-        self._children = []
-
-    async def connect(self, **queuedb):
-        self._queue = joblib.joblib.JobQueue()
-        # register event handler on new open job
-        await self._queue.listen(
-            self._jobtype + "_open",
-            self._pushed_new)
-        await self._queue.connect(**queuedb)
-        await self._queue.setup()
-
-    async def setup(self):
-        # override
-        pass
-
-    async def _windup(self):
-        # override
-        pass
-
-    async def _execute_job(self, jobid, payload, priority):
-        # override
-        pass
-
-    async def _teardown(self, failed):
-        # override
-        pass
-
-    async def request(self, jobtype, payload, priority=None):
-        """Queue child jobs to be executed after teardown."""
-        self._children.append((jobtype, payload, priority))
-
-    async def poll(self, batchlimit=1):
-        """Start a batch of jobs."""
-        jobs = await self._queue.acquire(jobtype=self._jobtype,
-                                         length=batchlimit)
-
-        if len(jobs) == 0:
-            logging.debug("no jobs")
-            self._need_update = False
-            return
-
-        await self._windup()
-        critical_error = False
-        failed = []
-        succeeded = []
-        for jobid, payload, priority in jobs:
-            try:
-                await self._execute_job(jobid, payload, priority)
-                succeeded.append(jobid)
-            except JobFailed as err:
-                failed.append((jobid, err.args[0]))
-                if len(err.args) > 1:
-                    # rollback
-                    critical_error = err.args[1]
-                    if critical_error:
-                        break
-                else:
-                    critical_error = True  # default
-        if critical_error:
-            await self._queue.reset(succeeded,
-                                    self._jobtype)
-            logging.warning("batch failed, reset")
-        else:
-            await self._queue.finish(succeeded,
-                                     self._jobtype)
-
-        for jobid, err in failed:
-            await self._queue.fail(jobid,
-                                   self._jobtype,
-                                   err)
-
-        await self._teardown(failed=critical_error)
-
-        # spawn child jobs
-        for jobt, payloads, priorities in self._children:
-            await self._queue.request(
-                jobt, payloads, priorities)
-        self._children = []
-
-    def _pushed_new(self, _):
-        """New job event handler."""
-        logging.debug("received a notification")
-        self._need_update = True
-
-    async def run(self, batchlimit):
-        """Request a poll after a push or a timer."""
-        while True:
-            while not self._need_update:
-                await asyncio.sleep(0.01)
-            logging.debug("polling")
-            await self.poll(batchlimit)
-
-    async def start(self, batchlimit=1):
-        """Start in background."""
-        asyncio.ensure_future(self.run(batchlimit))