2
\$\begingroup\$

I've been using Beanstaldk (not Amazon's one) for years and loved the possibility to handle delays and priorities in the tasks, but I recently had to filters the tasks for specific reasons, and Beanstalk doesn't allow that easily.

So, after searching a bit, I decided to use Mongo as it seems to be pretty fast, and I wrote a (basic) Task Queue system for this in Python.

# -*- coding:utf-8 -*-


from pymongo import MongoClient, ASCENDING
from pymongo.collection import ReturnDocument
from bson.objectid import ObjectId
import time


class MongoQ:
    """
    MongoQ is a queue system based on MongoDB.
    It offers the possibility to:
        * Delay
        * Filters
        * Priorize (from 1 (highest) to 1024 (lowest))
    """

    def __init__(self, host, queue, port=27017, database='_queue', precision=1000 * 1000):
        # connect
        self.client = MongoClient(host, port)
        self.db = self.client[database]
        self.queue = queue
        self._precision = precision

    def _parse_filters(self, filters):
        if filters is None:
            return {}

        assert isinstance(filters, dict)
        results = {}
        for k in filters:
            if k[0:1] == '$':
                sub = filters[k]
                if isinstance(sub, dict):
                    results[k] = self._parse_filters(sub)
                elif isinstance(sub, (tuple, list)):
                    results[k] = []
                    for value in sub:
                        results[k].append(self._parse_filters(value))
            else:
                results['task.{0}'.format(k)] = filters[k]

        return results

    def use(self, queue):
        self.queue = queue

    def using(self):
        return self.queue

    def stats(self):
        ts = int(time.time() * self._precision)
        return {
            'jobs': self.db[self.queue].estimated_document_count(),
            'available': self.db[self.queue].count_documents({
                '_available': {'$lte': ts},
                '$or': [
                    {'_reserved': None},
                    {'_reserved': {'$lte': ts}}
                ]
            }),
            'reserved': self.db[self.queue].count_documents({'_reserved': {'$gt': ts}})
        }

    def peek(self, _id):
        if isinstance(_id, str):
            _id = ObjectId(_id)

        return self.db[self.queue].find_one({'_id': _id})

    def reserve(self, filters=None, ttr=5):
        """
        Returns a tuple (ObjectId, Task)
        ttr is in seconds
        """
        assert ttr >= 0

        filters = self._parse_filters(filters)

        ts = int(time.time() * self._precision)
        filtering = {
            '_available': {'$lte': ts}
        }

        _or_reserved_filter = {'$or': [
            {'_reserved': None},
            {'_reserved': {'$lte': ts}}
        ]}

        if filters.get('$or') is not None:
            if filters.get('$and') is not None:
                filtering['$and'] = filters.get('$and')
            else:
                filtering['$and'] = []

            filtering['$and'].append({'$or': filters['$or']})
            filtering['$and'].append(_or_reserved_filter)
        else:
            filtering['$or'] = _or_reserved_filter['$or']

        obj = self.db[self.queue].find_one_and_update(
            filter=filtering,
            update={'$set': {'_reserved': ts + int(ttr * self._precision)}},
            sort=[('_priority', ASCENDING), ('_available', ASCENDING)],
            return_document=ReturnDocument.AFTER
        )

        if not obj:
            return None

        return Job(self.db[self.queue], obj)

    def reserve_many(self, filters=None, ttr=5, limit=1):
        """
        ttr is in seconds but float is allowed
        """
        jobs = []
        while True:
            job = self.reserve(filters, ttr)
            if not job:
                break

            jobs.append(job)
            if len(jobs) >= limit:
                break

        return jobs

    def put(self, task, priority=1, delay=0):
        assert isinstance(task, dict)
        assert isinstance(priority, int)
        assert priority >= 0
        assert priority <= 1024
        assert isinstance(delay, int)
        assert delay >= 0

        ts = int(time.time() * self._precision)

        return self.db[self.queue].insert_one({
            '_available': ts + (delay * self._precision),
            '_reserved': None,
            '_priority': priority,
            'task': task
        }).inserted_id

    def put_many(self, tasks, priority=1, delay=0):
        assert isinstance(tasks, (list, tuple))
        assert isinstance(priority, int)
        assert priority >= 0
        assert priority <= 1024
        assert isinstance(delay, int)
        assert delay >= 0

        ts = int(time.time() * self._precision)

        documents = []
        for d in tasks:
            assert isinstance(d, dict)
            documents.append({
                '_available': ts + (delay * self._precision),
                '_reserved': None,
                '_priority': priority,
                'task': d
            })

        return self.db[self.queue].insert_many(documents).inserted_ids

    def remove(self, _id):
        assert isinstance(_id, ObjectId)
        self.db[self.queue].delete_one({'_id': _id})

    def remove_many(self, ids):
        assert isinstance(ids, (list, tuple))
        for i in ids:
            assert isinstance(i, ObjectId)

        self.db[self.queue].delete_many({'_id': {'$in': ids}})


class Job:
    """
    Represents one Job.
    Allows to remove or release it depending on the need.
    """
    def __init__(self, queue, job):
        self.queue = queue
        self.job = job

    @property
    def priority(self):
        return self.job['_priority']

    @property
    def id(self):
        return self.job['_id']

    @property
    def task(self):
        return self.job['task']

    def remove(self):
        self.queue.delete_one({'_id': self.job['_id']})

    def release(self):
        self.queue.update_one({'_id': self.job['_id']}, {'$set': {'_reserved': None}})

It works as follow:

# Server side:
q = MongoQ('127.0.0.1', 'my-queue')
q.put({'some': 'task', 'to': 'run'}, delay=5, priority=10)

# Client side:
q = MongoQ('127.0.0.1', 'test-a')
job = q.reserve('my-queue')
print(job.task['some'])
job.remove()  # you can also do q.remove(job.id)

What do you think?

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.