rewrite using notifications on one connection

2 files changed, 217 insertions, 176 deletions

diff --git a/joblib.py b/joblib.py
index 8faf71f..ab107f8 100644
--- a/joblib.py
+++ b/joblib.py
@@ -8,15 +8,28 @@ import logging
 
 class JobQueue(object):
     def __init__(self):
-        self._pool = None
+        self._con = None
+        self._listens = {}
+
+    def _listener(self, con, pid,
+                        channel, payload):
+        """Fire the registered callback. Must not be async."""
+        self._listens[channel](payload)
+
+    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._pool = await asyncpg.create_pool(
-                    min_size=1, **args)
+                self._con = await asyncpg.connect(**args)
+                await self._hook_listener()
                 break
             except asyncpg.exceptions.CannotConnectNowError:
                 logging.warning(
@@ -29,166 +42,162 @@ class JobQueue(object):
 
     async def setup(self):
         """Initialize the database."""
-        async with self._pool.acquire() as con:
-            await 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 con.execute("""
-                CREATE UNIQUE INDEX ON jobs(priority, id)
-            """)
+        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 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):
         """Create a new job and return its id."""
-        async with self._pool.acquire() as con:
-            insert = await con.prepare("""
-                    INSERT INTO jobs(type, payload, priority)
-                    VALUES($1, $2, $3)
-                    RETURNING id
-            """)
-            if isinstance(payload, list):
-                payloads = payload
-            else:
-                payloads = [payload]
-            if isinstance(priority, list):
-                priorities = priority
-            else:
-                priorities = [priority] * len(payloads)
-            ids = []
-            async with con.transaction():
-                for pl, pr in zip(payloads, priorities):
-                    ids.append(await insert.fetchval(jobtype,
-                                                     json.dumps(pl),
-                                                     pr))
-            if isinstance(payload, list):
-                return ids
-            else:
-                return ids[0]
+        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.fetch(jobtype, pl, pr))
+            await self._con.execute("SELECT pg_notify($1 || '_open', '')",
+                                    jobtype)
+
+        if isinstance(payload, list):
+            return ids
+        else:
+            return ids[0]
 
     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."""
-        async with self._pool.acquire() as con:
-            update = await con.prepare("""
-                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
-            """)
-            if length is None:
-                limit = 1
-            else:
-                limit = length
-            while True:
-                try:
-                    # do not allow async access
-                    async with con.transaction(isolation="serializable"):
-                        result = await update.fetch(jobtype, limit)
-                        if len(result) == 0 and length is None:
-                            # no jobs available
-                            # backwards compatibility
-                            return None, None, None
+        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]
 
-                        jobs = []
-                        for record in result:
-                            jobs.append((record[0],
-                                         json.loads(record[1]),
-                                         record[2]))
+                    await self._con.execute("SELECT pg_notify($1 || '_running', '')",
+                                            jobtype)
 
-                        if length is None:
-                            return jobs[0]
-                        else:
-                            return jobs
-                except asyncpg.exceptions.SerializationError:
-                    # job is being picked up by another worker, try again
-                    logging.debug("serialization error, retrying")
-                    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
 
     async def status(self, jobid):
         """Return the status of a job."""
-        async with self._pool.acquire() as con:
-            return await con.fetchval("""
-                SELECT status
-                FROM jobs WHERE
-                id=$1
-            """, jobid)
+        return await self._con.fetchval("""
+            SELECT status
+            FROM jobs WHERE
+            id=$1
+        """, jobid)
 
-    async def finish(self, jobid):
+    async def finish(self, jobid, jobtype):
         """Mark jobs as completed."""
-        async with self._pool.acquire() as con:
-            update = await con.prepare("""
-                UPDATE jobs
-                SET status='finished'
+        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='finished'
                 WHERE id=$1
-            """)
-            if not isinstance(jobid, list):
-                jobids = [jobid]
-            else:
-                jobids = jobid
-            async with con.transaction():
-                for jid in jobids:
-                    await update.fetch(jid)
+            """, jobids)
+            await self._con.execute("""
+                SELECT pg_notify($1 || '_finished', '')
+            """, jobtype)
 
-    async def fail(self, jobid, reason):
+    async def fail(self, jobid, jobtype, reason):
         """Mark a job as failed."""
-        async with self._pool.acquire() as con:
-            update = await con.prepare("""
-                UPDATE jobs
-                SET status='failed', payload=payload||$2::jsonb
+        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
-            """)
-            if not isinstance(jobid, list):
-                jobids = [jobid]
-            else:
-                jobids = jobid
-            if not isinstance(reason, list):
-                reasons = [reason]
-            else:
-                reasons = reason
-            assert len(jobids) == len(reasons)
-            async with con.transaction():
-                for jid, rsn in zip(jobids, reasons):
-                    await update.fetch(jid,
-                                       json.dumps({"error": rsn}))
+            """, zip(jobids, reasons))
+            await self._con.execute("""
+                SELECT pg_notify($1 || '_failed', '')
+            """, jobtype)
 
-    async def reset(self, jobid):
+    async def reset(self, jobid, jobtype):
         """Mark a job as open."""
-        async with self._pool.acquire() as con:
-            update = await con.prepare("""
-                UPDATE jobs
-                SET status='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
-            """)
-            if not isinstance(jobid, list):
-                jobids = [jobid]
-            else:
-                jobids = jobid
-            async with con.transaction():
-                for jid in jobids:
-                    await update.fetch(jid)
+            """, jobids)
+            await self._con.execute("""
+                SELECT pg_notify($1 || '_open', '')
+            """, jobtype)
 
     async def cleanup(self):
         """Reopen all unfinished jobs."""
-        async with self._pool.acquire() as con:
-            while True:
-                try:
-                    async with con.transaction(isolation="serializable"):
-                        await con.execute("""
-                            UPDATE jobs
-                            SET status='open'
-                            WHERE status='running'
-                        """)
-                        return
-                except asyncpg.exceptions.SerializationError:
-                    pass
+        while True:
+            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
diff --git a/worker.py b/worker.py
index e5f0b25..d67223a 100644
--- a/worker.py
+++ b/worker.py
@@ -14,9 +14,15 @@ class Worker(object):
     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()
 
@@ -36,44 +42,70 @@ class Worker(object):
         # override
         pass
 
-    async def run(self, batchlimit=1):
-        """Start jobs forever."""
-        while True:
-            jobs = await self._queue.acquire(jobtype=self._jobtype,
-                                       length=batchlimit)
+    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")
+            return
 
-            if len(jobs) == 0:
-                await asyncio.sleep(1)
-                # nothing to do
-                continue
+        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)
 
-            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)
-                logging.warning("batch failed, reset")
-            else:
-                await self._queue.finish(succeeded)
+        for jobid, err in failed:
+            await self._queue.fail(jobid,
+                                   self._jobtype,
+                                   err)
 
-            for jobid, err in failed:
-                await self._queue.fail(jobid, err)
+        await self._teardown(failed=critical_error)
 
-            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)
+            logging.debug("polling")
+            self._need_update = False
+            await self.poll(batchlimit)
 
     async def start(self, batchlimit=1):
         """Start in background."""