├── .bumpversion.cfg ├── .github ├── FUNDING.yml ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── config.yml └── workflows │ ├── push.yml │ └── release.yml ├── .gitignore ├── CONTRIBUTING.md ├── CONTRIBUTORS.md ├── COPYING ├── COPYING.LESSER ├── MANIFEST.in ├── README.md ├── artwork ├── LICENSE ├── logo.ai └── logo.png ├── benchmarks ├── README.md ├── __init__.py ├── bench.py └── requirements.txt ├── bin └── dramatiq-gevent ├── docs ├── .gitignore ├── Makefile └── source │ ├── _static │ ├── franz-logo.png │ ├── logo.png │ ├── podcatcher-logo.png │ ├── sendcloud-logo.png │ ├── worker_architecture.mmd │ └── worker_architecture.svg │ ├── _templates │ ├── analytics.html │ ├── sidebarlogo.html │ ├── sponsors.html │ └── versions.html │ ├── advanced.rst │ ├── best_practices.rst │ ├── changelog.rst │ ├── conf.py │ ├── cookbook.rst │ ├── global.rst │ ├── guide.rst │ ├── index.rst │ ├── installation.rst │ ├── license.rst │ ├── motivation.rst │ ├── reference.rst │ ├── sitemap.py │ └── troubleshooting.rst ├── dramatiq ├── __init__.py ├── __main__.py ├── actor.py ├── asyncio.py ├── broker.py ├── brokers │ ├── __init__.py │ ├── rabbitmq.py │ ├── redis.py │ ├── redis │ │ ├── dispatch.lua │ │ └── maxstack.lua │ └── stub.py ├── canteen.py ├── cli.py ├── common.py ├── compat.py ├── composition.py ├── encoder.py ├── errors.py ├── generic.py ├── logging.py ├── message.py ├── middleware │ ├── __init__.py │ ├── age_limit.py │ ├── asyncio.py │ ├── callbacks.py │ ├── current_message.py │ ├── group_callbacks.py │ ├── middleware.py │ ├── pipelines.py │ ├── prometheus.py │ ├── retries.py │ ├── shutdown.py │ ├── threading.py │ └── time_limit.py ├── py.typed ├── rate_limits │ ├── __init__.py │ ├── backend.py │ ├── backends │ │ ├── __init__.py │ │ ├── memcached.py │ │ ├── redis.py │ │ └── stub.py │ ├── barrier.py │ ├── bucket.py │ ├── concurrent.py │ ├── rate_limiter.py │ └── window.py ├── results │ ├── __init__.py │ ├── backend.py │ ├── backends │ │ ├── __init__.py │ │ ├── memcached.py │ │ ├── redis.py │ │ └── stub.py │ ├── errors.py │ ├── middleware.py │ └── result.py ├── threading.py ├── watcher.py └── worker.py ├── examples ├── asyncio │ └── example.py ├── basic │ ├── README.md │ ├── __init__.py │ └── example.py ├── callable_broker │ └── app.py ├── composition │ ├── README.md │ ├── __init__.py │ └── example.py ├── crawler │ ├── README.md │ ├── __init__.py │ ├── example.py │ └── requirements.txt ├── issue_351 │ └── app.py ├── long_running │ ├── README.md │ ├── __init__.py │ └── example.py ├── max_tasks_per_child │ └── app.py ├── persistent │ ├── .gitignore │ ├── README.md │ ├── __init__.py │ └── example.py ├── results │ ├── README.md │ ├── __init__.py │ └── example.py ├── scheduling │ ├── README.md │ ├── __init__.py │ └── example.py └── time_limit │ ├── README.md │ ├── __init__.py │ └── example.py ├── mypy.ini ├── pytest-gevent.py ├── setup.cfg ├── setup.py ├── tests ├── __init__.py ├── benchmarks │ ├── __init__.py │ ├── test_rabbitmq_cli.py │ └── test_redis_cli.py ├── common.py ├── conftest.py ├── middleware │ ├── __init__.py │ ├── test_asyncio.py │ ├── test_prometheus.py │ ├── test_retries.py │ ├── test_shutdown.py │ ├── test_threading.py │ └── test_time_limit.py ├── test_actors.py ├── test_barrier.py ├── test_broker.py ├── test_bucket_rate_limiter.py ├── test_callbacks.py ├── test_canteen.py ├── test_cli.py ├── test_common.py ├── test_composition.py ├── test_concurrent_rate_limiter.py ├── test_encoders.py ├── test_generic_actors.py ├── test_messages.py ├── test_pidfile.py ├── test_rabbitmq.py ├── test_redis.py ├── test_results.py ├── test_stub_broker.py ├── test_watch.py ├── test_window_rate_limiter.py └── test_worker.py └── tox.ini /.bumpversion.cfg: -------------------------------------------------------------------------------- 1 | [bumpversion] 2 | current_version = 1.18.0 3 | message = chore: bump version {current_version} → {new_version} 4 | commit = True 5 | tag = True 6 | 7 | [bumpversion:file:dramatiq/__init__.py] 8 | search = __version__ = "{current_version}" 9 | replace = __version__ = "{new_version}" 10 | -------------------------------------------------------------------------------- /.github/FUNDING.yml: -------------------------------------------------------------------------------- 1 | github: Bogdanp 2 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Github issues are primarily for bugs 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## Checklist 11 | 12 | * [ ] Does your title concisely summarize the problem? 13 | * [ ] Did you include a minimal, reproducible example? 14 | * [ ] What OS are you using? 15 | * [ ] What version of Dramatiq are you using? 16 | * [ ] What version of Python are you using? 17 | * [ ] What did you do? 18 | * [ ] What did you expect would happen? 19 | * [ ] What happened? 20 | 21 | 22 | ## What OS are you using? 23 | 24 | 25 | 26 | 27 | ## What version of Dramatiq are you using? 28 | 29 | 30 | 31 | 32 | ## What version of Python are you using? 33 | 34 | 35 | 36 | 37 | ## What did you do? 38 | 39 | 40 | 41 | 42 | ## What did you expect would happen? 43 | 44 | 45 | 46 | 47 | ## What happened? 48 | 49 | 50 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/config.yml: -------------------------------------------------------------------------------- 1 | blank_issues_enabled: true 2 | contact_links: 3 | - name: Questions 4 | url: https://groups.io/g/dramatiq-users 5 | about: If you have questions, please ask them on the discussion Board. 6 | - name: Feature Requests 7 | url: https://groups.io/g/dramatiq-users 8 | about: If you want to request or suggest a feature, please start a discussion on the discussion Board. 9 | -------------------------------------------------------------------------------- /.github/workflows/push.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | on: [push, pull_request] 3 | jobs: 4 | lint: 5 | runs-on: ubuntu-24.04 6 | name: lint 7 | steps: 8 | - uses: actions/checkout@master 9 | - uses: actions/setup-python@v4 10 | with: 11 | python-version: "3.13" 12 | - run: | 13 | sudo apt-get update 14 | sudo apt-get remove libhashkit2 libmemcached11 || true 15 | sudo apt-get install -y libmemcached-dev 16 | - run: pip install tox 17 | - run: tox -e lint 18 | - run: tox -e docs 19 | 20 | build-unix: 21 | timeout-minutes: 30 22 | strategy: 23 | fail-fast: false 24 | matrix: 25 | os: ["ubuntu-24.04"] 26 | python: ["3.9", "3.10", "3.11", "3.12", "3.13"] 27 | concurrency: ["cpython", "gevent"] 28 | 29 | runs-on: ${{ matrix.os }} 30 | name: test on ${{ matrix.python }} (${{ matrix.os }}, ${{ matrix.concurrency }}) 31 | 32 | services: 33 | memcached: 34 | image: memcached:latest 35 | ports: 36 | - 11211:11211 37 | rabbitmq: 38 | image: rabbitmq:3 39 | env: 40 | RABBITMQ_DEFAULT_USER: "dramatiq" 41 | RABBITMQ_DEFAULT_PASS: "dramatiq" 42 | ports: 43 | - 5672:5672 44 | options: '--hostname "rmq" --health-cmd "rabbitmqctl status" --health-interval 10s --health-timeout 10s --health-retries 3 --health-start-period 60s' 45 | redis: 46 | image: redis:latest 47 | ports: 48 | - 6379:6379 49 | 50 | steps: 51 | - uses: actions/checkout@master 52 | - uses: actions/setup-python@v4 53 | with: 54 | python-version: ${{ matrix.python }} 55 | - run: | 56 | sudo apt-get update 57 | sudo apt-get remove libhashkit2 libmemcached11 || true 58 | sudo apt-get install -y libmemcached-dev 59 | - run: pip install -e '.[dev]' 60 | - run: pytest --benchmark-skip 61 | if: ${{ matrix.concurrency == 'cpython' }} 62 | env: 63 | RABBITMQ_USERNAME: "dramatiq" 64 | RABBITMQ_PASSWORD: "dramatiq" 65 | - run: python pytest-gevent.py --benchmark-skip 66 | if: ${{ matrix.concurrency == 'gevent' }} 67 | env: 68 | RABBITMQ_USERNAME: "dramatiq" 69 | RABBITMQ_PASSWORD: "dramatiq" 70 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: PyPI release 2 | 3 | on: 4 | release: 5 | types: [ published ] 6 | 7 | jobs: 8 | package: 9 | runs-on: ubuntu-latest 10 | 11 | steps: 12 | - uses: actions/checkout@v3 13 | - uses: actions/setup-python@v4 14 | with: 15 | python-version: "3.13" 16 | 17 | - name: Install dependencies 18 | run: python -m pip install build 19 | 20 | - name: Build dist packages 21 | run: python -m build . 22 | 23 | - name: Upload packages to PyPI 24 | uses: pypa/gh-action-pypi-publish@release/v1 25 | with: 26 | password: ${{ secrets.PYPI_API_TOKEN }} 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /.benchmarks 2 | /.cache 3 | /.coverage 4 | /.pytest_cache 5 | /.tox 6 | /benchmark*.svg 7 | /build 8 | /dist 9 | /*.egg-info 10 | /htmlcov 11 | __pycache__ 12 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing 2 | 3 | ## Code 4 | 5 | [Start a discussion] before attempting to make a contribution. Any 6 | contribution that doesn't fit my design goals for the project will be 7 | rejected so it's always better to start a discussion first! 8 | 9 | By submitting contributions, you disavow any rights or claims to any 10 | changes submitted to the Dramatiq project and assign the copyright of 11 | those changes to CLEARTYPE SRL. If you cannot or do not want to 12 | reassign those rights, you shouldn't submit a PR. Instead, you should 13 | open an issue and let someone else do that work. 14 | 15 | ### Pull Requests 16 | 17 | * Make sure any code changes are covered by tests. 18 | * Run [isort] on any modified files. 19 | * If this is your first contribution, add yourself to the [CONTRIBUTORS] file. 20 | * If your branch is behind master, [rebase] on top of it. 21 | 22 | Run the test suite with `tox`. The tests require running [RabbitMQ], 23 | [Redis] and [Memcached] servers. 24 | 25 | [CONTRIBUTORS]: https://github.com/Bogdanp/dramatiq/blob/master/CONTRIBUTORS.md 26 | [RabbitMQ]: https://www.rabbitmq.com/ 27 | [Redis]: https://redis.io 28 | [Memcached]: https://memcached.org/ 29 | [isort]: https://github.com/timothycrosley/isort 30 | [rebase]: https://github.com/edx/edx-platform/wiki/How-to-Rebase-a-Pull-Request 31 | 32 | 33 | ## Issues 34 | 35 | When you open an issue make sure you include the full stack trace and 36 | that you list all pertinent information (operating system, message 37 | broker, Python implementation) as part of the issue description. 38 | 39 | Please include a minimal, reproducible test case with every bug 40 | report. If the issue is actually a question, consider asking it on 41 | the [discussion board] or Stack Overflow first. 42 | 43 | [Start a discussion]: https://groups.io/g/dramatiq-users 44 | [discussion board]: https://groups.io/g/dramatiq-users 45 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include COPYING 2 | include COPYING.LESSER 3 | include README.md 4 | include dramatiq/py.typed 5 | include setup.cfg 6 | include setup.py 7 | 8 | recursive-include bin * 9 | recursive-include dramatiq/brokers/redis *.lua 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | # dramatiq 4 | 5 | [![Build Status](https://github.com/Bogdanp/dramatiq/workflows/CI/badge.svg)](https://github.com/Bogdanp/dramatiq/actions?query=workflow%3A%22CI%22) 6 | [![PyPI version](https://badge.fury.io/py/dramatiq.svg)](https://badge.fury.io/py/dramatiq) 7 | [![Documentation](https://img.shields.io/badge/doc-latest-brightgreen.svg)](http://dramatiq.io) 8 | [![Discuss](https://img.shields.io/badge/discuss-online-orange.svg)](https://groups.io/g/dramatiq-users) 9 | 10 | *A fast and reliable distributed task processing library for Python 3.* 11 | 12 |
13 | 14 | **Changelog**: https://dramatiq.io/changelog.html
15 | **Community**: https://groups.io/g/dramatiq-users
16 | **Documentation**: https://dramatiq.io
17 | 18 |
19 | 20 |

Sponsors

21 | 22 |

23 | 24 | 25 | 26 | 27 | 28 | 29 |

30 | 31 | 32 | ## Installation 33 | 34 | If you want to use it with [RabbitMQ] 35 | 36 | pip install 'dramatiq[rabbitmq, watch]' 37 | 38 | or if you want to use it with [Redis] 39 | 40 | pip install 'dramatiq[redis, watch]' 41 | 42 | 43 | ## Quickstart 44 | 45 | Make sure you've got [RabbitMQ] running, then create a new file called 46 | `example.py`: 47 | 48 | ``` python 49 | import dramatiq 50 | import requests 51 | import sys 52 | 53 | 54 | @dramatiq.actor 55 | def count_words(url): 56 | response = requests.get(url) 57 | count = len(response.text.split(" ")) 58 | print(f"There are {count} words at {url!r}.") 59 | 60 | 61 | if __name__ == "__main__": 62 | count_words.send(sys.argv[1]) 63 | ``` 64 | 65 | In one terminal, run your workers: 66 | 67 | dramatiq example 68 | 69 | In another, start enqueueing messages: 70 | 71 | python example.py http://example.com 72 | python example.py https://github.com 73 | python example.py https://news.ycombinator.com 74 | 75 | Check out the [user guide] to learn more! 76 | 77 | 78 | ## License 79 | 80 | dramatiq is licensed under the LGPL. Please see [COPYING] and 81 | [COPYING.LESSER] for licensing details. 82 | 83 | 84 | [COPYING.LESSER]: https://github.com/Bogdanp/dramatiq/blob/master/COPYING.LESSER 85 | [COPYING]: https://github.com/Bogdanp/dramatiq/blob/master/COPYING 86 | [RabbitMQ]: https://www.rabbitmq.com/ 87 | [Redis]: https://redis.io 88 | [user guide]: https://dramatiq.io/guide.html 89 | -------------------------------------------------------------------------------- /artwork/LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017 CLEARTYPE SRL 2 | 3 | All rights reserved. 4 | 5 | This logo or a modified version may be used by anyone to refer to the 6 | Dramatiq project, but does not indicate endorsement by the project. 7 | 8 | Redistribution and use in source (the AI file) and binary forms 9 | (rendered PNG files, etc.) of the image, with or without modification, 10 | are permitted provided that the following conditions are met: 11 | 12 | * Redistributions of source code must retain the above copyright 13 | notice and this list of conditions. 14 | 15 | * The names of the contributors to the Dramatiq software (see 16 | CONTRIBUTORS.md) may not be used to endorse or promote products 17 | derived from this software without specific prior written 18 | permission. 19 | 20 | We would appreciate it if you make the image a link to 21 | https://dramatiq.io/ if you use it on a web page. -------------------------------------------------------------------------------- /artwork/logo.ai: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/artwork/logo.ai -------------------------------------------------------------------------------- /artwork/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/artwork/logo.png -------------------------------------------------------------------------------- /benchmarks/README.md: -------------------------------------------------------------------------------- 1 | # Celery vs Dramatiq benchmarks 2 | 3 | ## Setup 4 | 5 | 1. Install [memcached][memcached] and [RabbitMQ][rabbitmq] or [Redis][redis] 6 | 1. Install Dramatiq: `pip install dramatiq[rabbitmq]` or `pip install dramatiq[redis]` 7 | 1. Install the dependencies: `pip install -r requirements.txt` 8 | 1. Run `redis-server` or `rabbitmq-server` in a terminal window. 9 | 1. In a separate terminal window, run `memcached`. 10 | 11 | ## Running the benchmarks 12 | 13 | Run `python bench.py` to run the Dramatiq benchmark, and `python 14 | bench.py --use-celery` to run the Celery benchmark. Prepend `env 15 | REDIS=1` to each command to run the benchmarks against Redis. 16 | 17 | Run `python bench.py --help` to see all the available options. 18 | 19 | ## Caveats 20 | 21 | As with any benchmark, take it with a grain of salt. Dramatiq has an 22 | advantage over Celery in most of these benchmarks by design. 23 | 24 | 25 | [memcached]: https://memcached.org 26 | [rabbitmq]: https://www.rabbitmq.com 27 | [redis]: https://redis.io 28 | -------------------------------------------------------------------------------- /benchmarks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/benchmarks/__init__.py -------------------------------------------------------------------------------- /benchmarks/bench.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import logging 3 | import os 4 | import random 5 | import subprocess 6 | import sys 7 | import time 8 | 9 | import celery 10 | import pylibmc 11 | 12 | import dramatiq 13 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 14 | from dramatiq.brokers.redis import RedisBroker 15 | 16 | logger = logging.getLogger("example") 17 | counter_key = "latench-bench-counter" 18 | memcache_client = pylibmc.Client(["localhost"], binary=True) 19 | memcache_pool = pylibmc.ClientPool(memcache_client, 8) 20 | random.seed(1337) 21 | 22 | if os.getenv("REDIS") == "1": 23 | broker = RedisBroker() 24 | dramatiq.set_broker(broker) 25 | celery_app = celery.Celery(broker="redis:///") 26 | 27 | else: 28 | broker = RabbitmqBroker(host="127.0.0.1") 29 | dramatiq.set_broker(broker) 30 | celery_app = celery.Celery(broker="amqp:///") 31 | 32 | 33 | def fib_bench(n): 34 | p, q = 0, 1 35 | while n > 0: 36 | p, q = q, p + q 37 | n -= 1 38 | 39 | with memcache_pool.reserve() as client: 40 | client.incr(counter_key) 41 | 42 | return p 43 | 44 | 45 | def latency_bench(): 46 | p = random.randint(1, 100) 47 | if p == 1: 48 | duration = 10 49 | elif p <= 30: 50 | duration = 5 51 | elif p <= 50: 52 | duration = 3 53 | else: 54 | duration = 1 55 | 56 | time.sleep(duration) 57 | with memcache_pool.reserve() as client: 58 | client.incr(counter_key) 59 | 60 | 61 | dramatiq_fib_bench = dramatiq.actor(fib_bench) 62 | dramatiq_latency_bench = dramatiq.actor(latency_bench) 63 | 64 | celery_fib_bench = celery_app.task(name="fib-bench", acks_late=True)(fib_bench) 65 | celery_latency_bench = celery_app.task(name="latency-bench", acks_late=True)(latency_bench) 66 | 67 | 68 | def benchmark_arg(value): 69 | benchmarks = ("fib", "latency") 70 | if value not in benchmarks: 71 | raise argparse.ArgumentTypeError(f"benchmark must be one of {benchmarks!r}") 72 | return value 73 | 74 | 75 | def parse_args(): 76 | parser = argparse.ArgumentParser() 77 | parser.add_argument( 78 | "--benchmark", help="the benchmark to run", 79 | type=benchmark_arg, default="latency", 80 | ) 81 | parser.add_argument( 82 | "--count", help="the number of messages to benchmark with", 83 | type=int, default=10000, 84 | ) 85 | parser.add_argument( 86 | "--use-green-threads", help="run workers with green threads rather than system threads", 87 | action="store_true", default=False, 88 | ) 89 | parser.add_argument( 90 | "--use-celery", help="run the benchmark under Celery", 91 | action="store_true", default=False, 92 | ) 93 | return parser.parse_args() 94 | 95 | 96 | def main(args): 97 | args = parse_args() 98 | for _ in range(args.count): 99 | if args.use_celery: 100 | if args.benchmark == "latency": 101 | celery_latency_bench.delay() 102 | 103 | elif args.benchmark == "fib": 104 | celery_fib_bench.delay(random.randint(1, 200)) 105 | 106 | else: 107 | if args.benchmark == "latency": 108 | dramatiq_latency_bench.send() 109 | 110 | elif args.benchmark == "fib": 111 | dramatiq_fib_bench.send(random.randint(1, 200)) 112 | 113 | print("Done enqueing messages. Booting workers...") 114 | with memcache_pool.reserve() as client: 115 | client.set(counter_key, 0) 116 | 117 | start_time = time.time() 118 | if args.use_celery: 119 | subprocess_args = ["celery", "worker", "-A", "bench.celery_app"] 120 | if args.use_green_threads: 121 | subprocess_args.extend(["-P", "eventlet", "-c", "2000"]) 122 | else: 123 | subprocess_args.extend(["-c", "8"]) 124 | 125 | else: 126 | if args.use_green_threads: 127 | subprocess_args = ["dramatiq-gevent", "bench", "-p", "8", "-t", "250"] 128 | 129 | else: 130 | subprocess_args = ["dramatiq", "bench", "-p", "8"] 131 | 132 | proc = subprocess.Popen(subprocess_args) 133 | processed = 0 134 | while processed < args.count: 135 | processed = client.get(counter_key) 136 | print(f"{processed}/{args.count} messages processed\r", end="") 137 | time.sleep(0.1) 138 | 139 | duration = time.time() - start_time 140 | proc.terminate() 141 | proc.wait() 142 | 143 | print(f"Took {duration} seconds to process {args.count} messages.") 144 | return 0 145 | 146 | 147 | if __name__ == "__main__": 148 | sys.exit(main(sys.argv)) 149 | -------------------------------------------------------------------------------- /benchmarks/requirements.txt: -------------------------------------------------------------------------------- 1 | celery 2 | pylibmc 3 | -------------------------------------------------------------------------------- /bin/dramatiq-gevent: -------------------------------------------------------------------------------- 1 | #!python 2 | try: 3 | from gevent import monkey; monkey.patch_all() # noqa 4 | except ImportError: 5 | import sys 6 | sys.stderr.write("error: gevent is missing. Run `pip install gevent`.") 7 | sys.exit(1) 8 | 9 | import sys 10 | 11 | from dramatiq.cli import main 12 | 13 | if __name__ == "__main__": 14 | sys.exit(main()) 15 | -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- 1 | .coverage 2 | build -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | SPHINXPROJ = dramatiq 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | 12 | .PHONY: all 13 | all: 14 | @$(MAKE) html 15 | 16 | 17 | .PHONY: deploy 18 | deploy: all 19 | rsync -avz --delete build/html/* dramatiq:~/www 20 | 21 | 22 | .PHONY: help 23 | help: 24 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 25 | 26 | 27 | .PHONY: Makefile 28 | # Catch-all target: route all unknown targets to Sphinx using the new 29 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 30 | %: Makefile 31 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 32 | -------------------------------------------------------------------------------- /docs/source/_static/franz-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/docs/source/_static/franz-logo.png -------------------------------------------------------------------------------- /docs/source/_static/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/docs/source/_static/logo.png -------------------------------------------------------------------------------- /docs/source/_static/podcatcher-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/docs/source/_static/podcatcher-logo.png -------------------------------------------------------------------------------- /docs/source/_static/sendcloud-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/docs/source/_static/sendcloud-logo.png -------------------------------------------------------------------------------- /docs/source/_static/worker_architecture.mmd: -------------------------------------------------------------------------------- 1 | sequenceDiagram 2 | participant CT as ConsumerThread("default") 3 | participant RMQ as RabbitMQ 4 | participant WorkQueue as WorkQueue in memory 5 | participant AcksQueue as AcksQueue in memory 6 | participant WT1 as WorkerThread1 7 | participant WT2 as WorkerThread2 8 | 9 | loop while running, read messages off the broker and put them on the work queue 10 | CT->>+RMQ: get_message(queue_name="default", block=True) 11 | RMQ-->>CT: Message(id=1, {...}) 12 | CT->>WorkQueue: put_work(Message(id=1, {...})) 13 | 14 | loop while pending acks 15 | CT->>AcksQueue: get_pending_ack() 16 | AcksQueue-->>CT: Message(id=1, {...}) 17 | CT->>RMQ: ack_message(id=1) 18 | end 19 | end 20 | 21 | alt concurrently 22 | loop while running, process work 23 | WT1->>WorkQueue: get_work(block=True) 24 | WorkQueue-->>WT1: Message(id=2, { ... }) 25 | WT1->>WT1: process_message(Message(id=2, {...})) 26 | WT1->>AcksQueue: ack_message(Message(id=2, { ... })) 27 | end 28 | else 29 | loop while running, process work 30 | WT2->>WorkQueue: get_work(block=True) 31 | WorkQueue-->>WT2: Message(id=1, { ... }) 32 | WT2->>WT2: process_message(Message(id=1, {...})) 33 | WT2->>AcksQueue: ack_message(Message(id=1, { ... })) 34 | end 35 | end -------------------------------------------------------------------------------- /docs/source/_templates/analytics.html: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/docs/source/_templates/analytics.html -------------------------------------------------------------------------------- /docs/source/_templates/sidebarlogo.html: -------------------------------------------------------------------------------- 1 | 7 | 8 |

9 | 11 |
12 |

13 | -------------------------------------------------------------------------------- /docs/source/_templates/sponsors.html: -------------------------------------------------------------------------------- 1 |
2 |

Sponsors

3 | 4 | 5 | 6 | Franz — Desktop Client for Apache Kafka 7 | 8 | 9 |
10 | 11 | 12 | 13 | Podcatcher — iOS Podcast Player 14 | 15 |
16 | -------------------------------------------------------------------------------- /docs/source/_templates/versions.html: -------------------------------------------------------------------------------- 1 | {% if versions %} 2 |
3 |

Versions

4 | 5 | 22 | 23 | 29 |
30 | {% endif %} 31 | -------------------------------------------------------------------------------- /docs/source/best_practices.rst: -------------------------------------------------------------------------------- 1 | .. include:: global.rst 2 | 3 | Best Practices 4 | ============== 5 | 6 | Concurrent Actors 7 | ----------------- 8 | 9 | Your actor will run concurrently with other actors in the system. You 10 | need to be mindful of the impact this has on your database, any third 11 | party services you might be calling and the resources available on the 12 | systems running your workers. Additionally, you need to be mindful of 13 | data races between actors that manipulate the same objects in your 14 | database. 15 | 16 | 17 | Retriable Actors 18 | ---------------- 19 | 20 | Dramatiq actors may receive the same message multiple times in the 21 | event of a worker failure (hardware, network or power failure). This 22 | means that, for any given message, running your actor multiple times 23 | must be safe. This is also known as being "idempotent". 24 | 25 | 26 | Simple Messages 27 | --------------- 28 | 29 | Attempting to send an actor any object that can't be encoded to JSON 30 | by the standard ``json`` package will fail immediately so you'll want 31 | to limit your actor parameters to the following object types: `bool`, 32 | `int`, `float`, `bytes`, `string`, `list` and `dict`. 33 | 34 | Additionally, since messages are sent over the wire you'll want to 35 | keep them as short as possible. For example, if you've got an actor 36 | that operates over ``User`` objects in your system, send that actor 37 | the user's id rather than the serialized user. 38 | 39 | 40 | Error Reporting 41 | --------------- 42 | 43 | Invariably, you're probably going to introduce issues in production 44 | every now and then and some of those issues are going to affect your 45 | tasks. You should use an error reporting service such as Sentry_ so 46 | you get notified of these errors as soon as they occur. 47 | 48 | 49 | .. _Sentry: https://sentry.io/welcome/ 50 | -------------------------------------------------------------------------------- /docs/source/global.rst: -------------------------------------------------------------------------------- 1 | .. References 2 | 3 | .. |AgeLimit| replace:: :class:`AgeLimit` 4 | .. |Barriers| replace:: :class:`Barriers` 5 | .. |Brokers| replace:: :class:`Brokers` 6 | .. |Broker| replace:: :class:`Broker` 7 | .. |Callbacks| replace:: :class:`Callbacks` 8 | .. |CurrentMessage| replace:: :class:`CurrentMessage` 9 | .. |DramatiqError| replace:: :class:`DramatiqError` 10 | .. |Encoders| replace:: :class:`Encoders` 11 | .. |GenericActors| replace:: :class:`class-based actors` 12 | .. |Groups| replace:: :func:`Groups` 13 | .. |Interrupt| replace:: :class:`Interrupt` 14 | .. |MemcachedRLBackend| replace:: :class:`Memcached` 15 | .. |Messages| replace:: :class:`Messages` 16 | .. |MiddlewareError| replace:: :class:`MiddlewareError` 17 | .. |Prometheus| replace:: :class:`Prometheus` 18 | .. |RabbitmqBroker_join| replace:: :meth:`join` 19 | .. |RabbitmqBroker| replace:: :class:`RabbitmqBroker` 20 | .. |RateLimitExceeded| replace:: :class:`RateLimitExceeded` 21 | .. |RateLimiters| replace:: :class:`RateLimiters` 22 | .. |RedisBroker| replace:: :class:`RedisBroker` 23 | .. |RedisRLBackend| replace:: :class:`Redis` 24 | .. |RedisResBackend| replace:: :class:`Redis` 25 | .. |ResultBackends| replace:: :class:`ResultBackends` 26 | .. |ResultBackend| replace:: :class:`ResultBackend` 27 | .. |ResultFailure| replace:: :class:`ResultFailure` 28 | .. |ResultMissing| replace:: :class:`ResultMissing` 29 | .. |ResultTimeout| replace:: :class:`ResultTimeout` 30 | .. |Results| replace:: :class:`Results` 31 | .. |Retries| replace:: :class:`Retries` 32 | .. |ShutdownNotifications| replace:: :class:`ShutdownNotifications` 33 | .. |Shutdown| replace:: :class:`Shutdown` 34 | .. |SkipMessage| replace:: :class:`SkipMessage` 35 | .. |StubBroker_flush_all| replace:: :meth:`StubBroker.flush_all` 36 | .. |StubBroker_flush| replace:: :meth:`StubBroker.flush` 37 | .. |StubBroker_join| replace:: :meth:`StubBroker.join` 38 | .. |StubBroker| replace:: :class:`StubBroker` 39 | .. |TimeLimitExceeded| replace:: :class:`TimeLimitExceeded` 40 | .. |TimeLimit| replace:: :class:`TimeLimit` 41 | .. |URLRabbitmqBroker| replace:: :class:`URLRabbitmqBroker` 42 | .. |WindowRateLimiter| replace:: :class:`WindowRateLimiter` 43 | .. |Worker_join| replace:: :meth:`Worker.join` 44 | .. |Worker_pause| replace:: :meth:`Worker.pause` 45 | .. |Worker_resume| replace:: :meth:`Worker.resume` 46 | .. |Worker| replace:: :meth:`Worker` 47 | .. |actor| replace:: :func:`actor` 48 | .. |add_middleware| replace:: :meth:`add_middleware` 49 | .. |after_skip_message| replace:: :meth:`after_skip_message` 50 | .. |before_consumer_thread_shutdown| replace:: :meth:`before_consumer_thread_shutdown` 51 | .. |before_worker_thread_shutdown| replace:: :meth:`before_worker_thread_shutdown` 52 | .. |dramatiq| replace:: :mod:`dramatiq` 53 | .. |get_current_message| replace:: :meth:`get_current_message` 54 | .. |group| replace:: :func:`group` 55 | .. |pipeline_get_results| replace:: :meth:`get_results` 56 | .. |pipeline_get_result| replace:: :meth:`get_result` 57 | .. |pipeline| replace:: :func:`pipeline` 58 | .. |rate_limits| replace:: :mod:`dramatiq.rate_limits` 59 | .. |send_with_options| replace:: :meth:`send_with_options` 60 | .. |send| replace:: :meth:`send` 61 | 62 | .. _gevent: http://www.gevent.org/ 63 | .. _Memcached: http://memcached.org 64 | .. _RabbitMQ: https://www.rabbitmq.com 65 | .. _Redis: https://redis.io 66 | -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | .. include:: global.rst 2 | 3 | Dramatiq: background tasks 4 | ========================== 5 | 6 | Release v\ |release|. (:doc:`installation`, :doc:`changelog`, `Discuss`_, `Source Code`_) 7 | 8 | .. _Discuss: https://groups.io/g/dramatiq-users 9 | .. _Source Code: https://github.com/Bogdanp/dramatiq 10 | 11 | .. image:: https://img.shields.io/badge/license-LGPL-blue.svg 12 | :target: license.html 13 | .. image:: https://github.com/Bogdanp/dramatiq/workflows/CI/badge.svg 14 | :target: https://github.com/Bogdanp/dramatiq/actions?query=workflow%3A%22CI%22 15 | .. image:: https://badge.fury.io/py/dramatiq.svg 16 | :target: https://badge.fury.io/py/dramatiq 17 | 18 | **Dramatiq** is a background task processing library for Python with a 19 | focus on simplicity, reliability and performance. 20 | 21 | .. raw:: html 22 | 23 | 24 | 25 | Here's what it looks like: 26 | 27 | :: 28 | 29 | import dramatiq 30 | import requests 31 | 32 | 33 | @dramatiq.actor 34 | def count_words(url): 35 | response = requests.get(url) 36 | count = len(response.text.split(" ")) 37 | print(f"There are {count} words at {url!r}.") 38 | 39 | 40 | # Synchronously count the words on example.com in the current process 41 | count_words("http://example.com") 42 | 43 | # or send the actor a message so that it may perform the count 44 | # later, in a separate process. 45 | count_words.send("http://example.com") 46 | 47 | **Dramatiq** is :doc:`licensed` under the LGPL and it 48 | officially supports Python 3.9 and later. 49 | 50 | 51 | Get It Now 52 | ---------- 53 | 54 | If you want to use it with RabbitMQ_:: 55 | 56 | $ pip install -U 'dramatiq[rabbitmq, watch]' 57 | 58 | Or if you want to use it with Redis_:: 59 | 60 | $ pip install -U 'dramatiq[redis, watch]' 61 | 62 | Read the :doc:`motivation` behind it or the :doc:`guide` if you're 63 | ready to get started. 64 | 65 | 66 | User Guide 67 | ---------- 68 | 69 | This part of the documentation is focused primarily on teaching you 70 | how to use Dramatiq. 71 | 72 | .. toctree:: 73 | :maxdepth: 2 74 | 75 | installation 76 | motivation 77 | guide 78 | best_practices 79 | troubleshooting 80 | advanced 81 | cookbook 82 | 83 | 84 | API Reference 85 | ------------- 86 | 87 | This part of the documentation is focused on detailing the various 88 | bits and pieces of the Dramatiq developer interface. 89 | 90 | .. toctree:: 91 | :maxdepth: 2 92 | 93 | reference 94 | 95 | 96 | Project Info 97 | ------------ 98 | 99 | .. toctree:: 100 | :maxdepth: 1 101 | 102 | Source Code 103 | changelog 104 | Contributing 105 | Discussion Board 106 | license 107 | -------------------------------------------------------------------------------- /docs/source/installation.rst: -------------------------------------------------------------------------------- 1 | .. include:: global.rst 2 | 3 | Installation 4 | ============ 5 | 6 | Dramatiq supports Python versions 3.9 and up and is installable via 7 | `pip`_ or from source. 8 | 9 | 10 | Via pip 11 | ------- 12 | 13 | To install dramatiq, simply run the following command in a terminal:: 14 | 15 | $ pip install -U 'dramatiq[rabbitmq, watch]' 16 | 17 | RabbitMQ_ is the recommended message broker, but Dramatiq also 18 | supports Redis_. 19 | 20 | If you would like to use it with Redis_ then run:: 21 | 22 | $ pip install -U 'dramatiq[redis, watch]' 23 | 24 | If you don't have `pip`_ installed, check out `this guide`_. 25 | 26 | Extra Requirements 27 | ^^^^^^^^^^^^^^^^^^ 28 | 29 | When installing the package via pip you can specify the following 30 | extra requirements: 31 | 32 | ============= ======================================================================================= 33 | Name Description 34 | ============= ======================================================================================= 35 | ``memcached`` Installs the required dependencies for the Memcached rate limiter backend. 36 | ``rabbitmq`` Installs the required dependencies for using Dramatiq with RabbitMQ. 37 | ``redis`` Installs the required dependencies for using Dramatiq with Redis. 38 | ``watch`` Installs the required dependencies for the ``--watch`` flag. (Supported only on UNIX) 39 | ============= ======================================================================================= 40 | 41 | If you want to install Dramatiq with all available features, run:: 42 | 43 | $ pip install -U 'dramatiq[all]' 44 | 45 | Optional Requirements 46 | ^^^^^^^^^^^^^^^^^^^^^ 47 | 48 | If you're using Redis as your broker and aren't planning on using PyPy 49 | then you should additionally install the ``hiredis`` package to get an 50 | increase in throughput. 51 | 52 | 53 | From Source 54 | ----------- 55 | 56 | To install the latest development version of dramatiq from source, 57 | clone the repo from `GitHub`_ 58 | 59 | :: 60 | 61 | $ git clone https://github.com/Bogdanp/dramatiq 62 | 63 | then install it to your local site-packages by running 64 | 65 | :: 66 | 67 | $ python setup.py install 68 | 69 | in the cloned directory. 70 | 71 | 72 | .. _GitHub: https://github.com/Bogdanp/dramatiq 73 | .. _pip: https://pip.pypa.io/en/stable/ 74 | .. _this guide: http://docs.python-guide.org/en/latest/starting/installation/ 75 | -------------------------------------------------------------------------------- /docs/source/license.rst: -------------------------------------------------------------------------------- 1 | Project License 2 | =============== 3 | 4 | Copyright (C) 2017,2018,2019 CLEARTYPE SRL 5 | 6 | Dramatiq is free software; you can redistribute it and/or modify it 7 | under the terms of the GNU Lesser General Public License as published by 8 | the Free Software Foundation, either version 3 of the License, or (at 9 | your option) any later version. 10 | 11 | Dramatiq is distributed in the hope that it will be useful, but WITHOUT 12 | ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 13 | FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 14 | License for more details. 15 | 16 | You should have received a copy of the GNU Lesser General Public License 17 | along with this program. If not, see http://www.gnu.org/licenses/. 18 | -------------------------------------------------------------------------------- /docs/source/sitemap.py: -------------------------------------------------------------------------------- 1 | import os 2 | import xml.etree.ElementTree as ET 3 | from datetime import datetime 4 | 5 | 6 | def setup(app): 7 | app.connect("build-finished", build_sitemap) 8 | return {"version": "1.0"} 9 | 10 | 11 | def collect_pages(basedir): 12 | for root, _dirs, files in os.walk(basedir): 13 | for filename in files: 14 | if not filename.endswith(".html"): 15 | continue 16 | 17 | if filename == "index.html": 18 | filename = "" 19 | 20 | priority = 0.8 21 | if "genindex" in filename or "modindex" in filename: 22 | priority = 0.4 23 | elif "_modules" in root: 24 | priority = 0.6 25 | elif filename == "": 26 | priority = 1 27 | 28 | path = os.path.join(root, filename) 29 | stat = os.stat(path) 30 | last_mod = datetime.fromtimestamp(stat.st_mtime).date().isoformat() 31 | 32 | page_path = path[len(basedir):].lstrip("/") 33 | yield page_path, last_mod, priority 34 | 35 | 36 | def build_sitemap(app, exception): 37 | pages = collect_pages(app.outdir) 38 | root = ET.Element("urlset") 39 | root.set("xmlns", "http://www.sitemaps.org/schemas/sitemap/0.9") 40 | for page, mod, prio in sorted(pages, key=lambda t: -t[2]): 41 | url = ET.SubElement(root, "url") 42 | loc = ET.SubElement(url, "loc") 43 | lastmod = ET.SubElement(url, "lastmod") 44 | priority = ET.SubElement(url, "priority") 45 | loc.text = f"https://dramatiq.io/{page}" 46 | lastmod.text = mod 47 | priority.text = f"{prio:.02f}" 48 | 49 | filename = f"{app.outdir}/sitemap.xml" 50 | tree = ET.ElementTree(root) 51 | tree.write(filename, xml_declaration=True, encoding="utf-8", method="xml") 52 | -------------------------------------------------------------------------------- /docs/source/troubleshooting.rst: -------------------------------------------------------------------------------- 1 | .. include:: global.rst 2 | 3 | Troubleshooting 4 | =============== 5 | 6 | This part of the documentation contains solutions to common problems 7 | you may encounter in the real world. 8 | 9 | 10 | ``FileNotFoundError`` when Enqueueing 11 | ------------------------------------- 12 | 13 | Dramatiq operations on builtin brokers are thread-safe, however they 14 | are not *process* safe so, if you use a pre-forking web server that 15 | forks *after* loading all of your code, then it's likely you'll run 16 | into issues enqueueing messages. That is because fork has 17 | copy-on-write semantics on most systems so any file descriptors open 18 | before forking will be shared between all of the processes. 19 | 20 | ``gunicorn`` Workaround 21 | ^^^^^^^^^^^^^^^^^^^^^^^ 22 | 23 | This problem should not occur under gunicorn_ since it loads the 24 | application after forking by default. 25 | 26 | .. _gunicorn: https://gunicorn.org/ 27 | 28 | ``uwsgi`` Workaround 29 | ^^^^^^^^^^^^^^^^^^^^ 30 | 31 | To work around this problem in uwsgi_, you have to turn on `lazy apps 32 | mode`_. This will ensure that all your app code is loaded after each 33 | worker process is forked. The tradeoff you make by turning on this 34 | option is your application will use slightly more memory. 35 | 36 | .. _uwsgi: https://uwsgi-docs.readthedocs.io/en/latest 37 | .. _lazy apps mode: https://uwsgi-docs.readthedocs.io/en/latest/Options.html#lazy-apps 38 | 39 | 40 | Integration Tests Hang 41 | ---------------------- 42 | 43 | During integration tests, actors are executed in a separate thread 44 | from the main thread that is running your code, just like they would 45 | be in the real world. In that sense, the |StubBroker| is great 46 | because it helps you simulate real-world execution conditions when 47 | you're testing your controller code. 48 | 49 | The main drawback to this approach is that -- because the actors are 50 | run in a separate thread -- your testing code has no way of knowing 51 | when an actor fails so, often, your tests may hang waiting for a 52 | message to be processed. An easy way to notice when these types of 53 | issues occur, is to turn on logging for your tests. If you use 54 | pytest_, then you can easily do this from the command line using the 55 | ``--log-cli-level`` flag:: 56 | 57 | $ py.test --log-cli-level=warning 58 | 59 | You can also pass ``fail_fast=True`` as a parameter to |StubBroker_join| 60 | in order to make it reraise whatever exception caused the actor to 61 | fail in the main thread. Note, however, that the actor is only 62 | considered to fail once all of its retries have been used up; meaning 63 | that unless you specify custom retry limits for the actors or for your 64 | tests as a whole (by configuring the |Retries| middleware), then each 65 | actor will retry for up to about 30 days before exhausting its 66 | available retries! 67 | 68 | .. _pytest: https://docs.pytest.org/en/latest/ 69 | -------------------------------------------------------------------------------- /dramatiq/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .actor import Actor, actor 19 | from .broker import Broker, Consumer, MessageProxy, get_broker, set_broker 20 | from .composition import group, pipeline 21 | from .encoder import Encoder, JSONEncoder, PickleEncoder 22 | from .errors import ( 23 | ActorNotFound, BrokerError, ConnectionClosed, ConnectionError, ConnectionFailed, DecodeError, DramatiqError, 24 | QueueJoinTimeout, QueueNotFound, RateLimitExceeded, Retry 25 | ) 26 | from .generic import GenericActor 27 | from .logging import get_logger 28 | from .message import Message, get_encoder, set_encoder 29 | from .middleware import Middleware 30 | from .worker import Worker 31 | 32 | __all__ = [ 33 | # Actors 34 | "Actor", "GenericActor", "actor", 35 | 36 | # Brokers 37 | "Broker", "Consumer", "MessageProxy", "get_broker", "set_broker", 38 | 39 | # Composition 40 | "group", "pipeline", 41 | 42 | # Encoding 43 | "Encoder", "JSONEncoder", "PickleEncoder", 44 | 45 | # Errors 46 | "DramatiqError", 47 | "BrokerError", "DecodeError", 48 | "ActorNotFound", "QueueNotFound", "QueueJoinTimeout", 49 | "ConnectionError", "ConnectionClosed", "ConnectionFailed", 50 | "RateLimitExceeded", "Retry", 51 | 52 | # Logging 53 | "get_logger", 54 | 55 | # Messages 56 | "Message", "get_encoder", "set_encoder", 57 | 58 | # Middlware 59 | "Middleware", 60 | 61 | # Workers 62 | "Worker", 63 | ] 64 | 65 | __version__ = "1.18.0" 66 | -------------------------------------------------------------------------------- /dramatiq/__main__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import sys 19 | 20 | from dramatiq.cli import main 21 | 22 | if __name__ == "__main__": 23 | sys.exit(main()) 24 | -------------------------------------------------------------------------------- /dramatiq/brokers/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | -------------------------------------------------------------------------------- /dramatiq/brokers/redis/maxstack.lua: -------------------------------------------------------------------------------- 1 | -- This file is a part of Dramatiq. 2 | -- 3 | -- Copyright (C) 2020 CLEARTYPE SRL 4 | -- 5 | -- Dramatiq is free software; you can redistribute it and/or modify it 6 | -- under the terms of the GNU Lesser General Public License as published by 7 | -- the Free Software Foundation, either version 3 of the License, or (at 8 | -- your option) any later version. 9 | -- 10 | -- Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | -- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | -- FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | -- License for more details. 14 | -- 15 | -- You should have received a copy of the GNU Lesser General Public License 16 | -- along with this program. If not, see . 17 | 18 | local function unpack_with_size(n) 19 | local items = {} 20 | for i = 0, n do 21 | items[i] = 0 22 | end 23 | 24 | (table.unpack or unpack)(items) 25 | end 26 | 27 | -- Programmers tend to pick multiples of 2 for their limits. Assuming 28 | -- that whomever sets LUAI_MAXCSTACK picks a value over 1024, starting 29 | -- from 999 will let us find a close-enough value to the limit in very 30 | -- few steps. 31 | local function find_max_unpack_size() 32 | local size = 999 33 | while true do 34 | if pcall(function() unpack_with_size(size * 2) end) then 35 | size = size * 2 36 | else 37 | return size 38 | end 39 | end 40 | end 41 | 42 | return find_max_unpack_size() 43 | -------------------------------------------------------------------------------- /dramatiq/canteen.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2019 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | # Don't depend on *anything* in this module. The contents of this 19 | # module can and *will* change without notice. 20 | 21 | import time 22 | from contextlib import contextmanager 23 | from ctypes import Array, Structure, c_bool, c_byte, c_int 24 | 25 | 26 | class Buffer(Array): 27 | _length_ = 1024 * 1024 28 | _type_ = c_byte 29 | 30 | 31 | # Canteen is the collective noun for a set of cutlery. 32 | # It's OK to be cute every once in a while. 33 | class Canteen(Structure): 34 | _fields_ = [ 35 | ("initialized", c_bool), 36 | ("last_position", c_int), 37 | ("paths", Buffer) 38 | ] 39 | 40 | 41 | def canteen_add(canteen, path): 42 | lo = canteen.last_position 43 | hi = canteen.last_position + len(path) + 1 44 | if hi > len(canteen.paths): 45 | raise RuntimeError("canteen is full") 46 | 47 | canteen.paths[lo:hi] = path.encode("utf-8") + b";" 48 | canteen.last_position = hi 49 | 50 | 51 | def canteen_get(canteen, timeout=1): 52 | if not wait(canteen, timeout): 53 | return [] 54 | 55 | data = bytes(canteen.paths[:canteen.last_position]) 56 | return data.decode("utf-8").split(";")[:-1] 57 | 58 | 59 | @contextmanager 60 | def canteen_try_init(cv): 61 | if cv.initialized: 62 | yield False 63 | return 64 | 65 | with cv.get_lock(): 66 | if cv.initialized: 67 | yield False 68 | return 69 | 70 | yield True 71 | cv.initialized = True 72 | 73 | 74 | def wait(canteen, timeout): 75 | deadline = time.monotonic() + timeout 76 | while not canteen.initialized: 77 | if time.monotonic() > deadline: 78 | return False 79 | 80 | time.sleep(0) 81 | 82 | return True 83 | -------------------------------------------------------------------------------- /dramatiq/common.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | import threading 18 | from os import getenv 19 | from queue import Empty 20 | from random import uniform 21 | from time import time 22 | 23 | from .errors import QueueJoinTimeout 24 | 25 | 26 | def getenv_int(name): 27 | """Parse an optional environment variable as an integer. 28 | """ 29 | v = getenv(name, None) 30 | if v is None: 31 | return None 32 | try: 33 | return int(v) 34 | except ValueError: 35 | raise ValueError("invalid integer value for env var %r: %r" % (name, v)) from None 36 | 37 | 38 | def compute_backoff(attempts, *, factor=5, jitter=True, max_backoff=2000, max_exponent=32): 39 | """Compute an exponential backoff value based on some number of attempts. 40 | 41 | Parameters: 42 | attempts(int): The number of attempts there have been so far. 43 | factor(int): The number of milliseconds to multiply each backoff by. 44 | max_backoff(int): The max number of milliseconds to backoff by. 45 | max_exponent(int): The maximum backoff exponent. 46 | 47 | Returns: 48 | tuple: The new number of attempts and the backoff in milliseconds. 49 | """ 50 | exponent = min(attempts, max_exponent) 51 | backoff = min(factor * 2 ** exponent, max_backoff) 52 | if jitter: 53 | backoff /= 2 54 | backoff = int(backoff + uniform(0, backoff)) 55 | return attempts + 1, backoff 56 | 57 | 58 | def current_millis(): 59 | """Returns the current UNIX time in milliseconds. 60 | """ 61 | return int(time() * 1000) 62 | 63 | 64 | def iter_queue(queue): 65 | """Iterate over the messages on a queue until it's empty. Does 66 | not block while waiting for messages. 67 | 68 | Parameters: 69 | queue(Queue): The queue to iterate over. 70 | 71 | Returns: 72 | generator[object]: The iterator. 73 | """ 74 | while True: 75 | try: 76 | yield queue.get_nowait() 77 | except Empty: 78 | break 79 | 80 | 81 | def join_queue(queue, timeout=None): 82 | """The join() method of standard queues in Python doesn't support 83 | timeouts. This implements the same functionality as that method, 84 | with optional timeout support, using only exposed Queue interfaces. 85 | 86 | Raises: 87 | QueueJoinTimeout: When the timeout is reached. 88 | 89 | Parameters: 90 | queue(Queue) 91 | timeout(Optional[float]) 92 | """ 93 | if timeout is None: 94 | queue.join() 95 | return 96 | 97 | join_complete = threading.Event() 98 | 99 | def join_and_signal(): 100 | queue.join() 101 | join_complete.set() 102 | 103 | join_thread = threading.Thread(target=join_and_signal, name="join_and_signal_thread") 104 | join_thread.daemon = True 105 | join_thread.start() 106 | 107 | # Wait for the join to complete or timeout 108 | if not join_complete.wait(timeout): 109 | raise QueueJoinTimeout("timed out after %.02f seconds" % timeout) 110 | 111 | 112 | def join_all(joinables, timeout): 113 | """Wait on a list of objects that can be joined with a total 114 | timeout represented by ``timeout``. 115 | 116 | Parameters: 117 | joinables(object): Objects with a join method. 118 | timeout(int): The total timeout in milliseconds. 119 | """ 120 | started, elapsed = current_millis(), 0 121 | for ob in joinables: 122 | ob.join(timeout=timeout / 1000) 123 | elapsed = current_millis() - started 124 | timeout = max(0, timeout - elapsed) 125 | 126 | 127 | def q_name(queue_name): 128 | """Returns the canonical queue name for a given queue. 129 | """ 130 | if queue_name.endswith(".DQ") or queue_name.endswith(".XQ"): 131 | return queue_name[:-3] 132 | return queue_name 133 | 134 | 135 | def dq_name(queue_name): 136 | """Returns the delayed queue name for a given queue. If the given 137 | queue name already belongs to a delayed queue, then it is returned 138 | unchanged. 139 | """ 140 | if queue_name.endswith(".DQ"): 141 | return queue_name 142 | 143 | if queue_name.endswith(".XQ"): 144 | queue_name = queue_name[:-3] 145 | return queue_name + ".DQ" 146 | 147 | 148 | def xq_name(queue_name): 149 | """Returns the dead letter queue name for a given queue. If the 150 | given queue name belongs to a delayed queue, the dead letter queue 151 | name for the original queue is generated. 152 | """ 153 | if queue_name.endswith(".XQ"): 154 | return queue_name 155 | 156 | if queue_name.endswith(".DQ"): 157 | queue_name = queue_name[:-3] 158 | return queue_name + ".XQ" 159 | -------------------------------------------------------------------------------- /dramatiq/compat.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | # Don't depend on *anything* in this module. The contents of this 19 | # module can and *will* change without notice. 20 | 21 | import sys 22 | from contextlib import contextmanager 23 | 24 | 25 | class StreamablePipe: 26 | """Wrap a multiprocessing.connection.Connection so it can be used 27 | with logging's StreamHandler. 28 | 29 | Parameters: 30 | pipe(multiprocessing.connection.Connection): writable end of the 31 | pipe to be used for transmitting child worker logging data to 32 | parent. 33 | """ 34 | 35 | def __init__(self, pipe, *, encoding="utf-8"): 36 | self.encoding = encoding 37 | self.pipe = pipe 38 | 39 | def fileno(self): 40 | return self.pipe.fileno() 41 | 42 | def isatty(self): 43 | return False 44 | 45 | def flush(self): 46 | pass 47 | 48 | def close(self): 49 | self.pipe.close() 50 | 51 | def read(self): # pragma: no cover 52 | raise NotImplementedError("StreamablePipes cannot be read from!") 53 | 54 | def write(self, s): 55 | self.pipe.send_bytes(s.encode(self.encoding, errors="replace")) 56 | 57 | @property 58 | def closed(self): 59 | return self.pipe.closed 60 | 61 | 62 | def file_or_stderr(filename, *, mode="a", encoding="utf-8"): 63 | """Returns a context object wrapping either the given file or 64 | stderr (if filename is None). This makes dealing with log files 65 | more convenient. 66 | """ 67 | if filename is not None: 68 | return open(filename, mode, encoding=encoding) 69 | 70 | @contextmanager 71 | def stderr_wrapper(): 72 | yield sys.stderr 73 | 74 | return stderr_wrapper() 75 | -------------------------------------------------------------------------------- /dramatiq/encoder.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import abc 19 | import json 20 | import pickle 21 | import typing 22 | 23 | from .errors import DecodeError 24 | 25 | #: Represents the contents of a Message object as a dict. 26 | MessageData = typing.Dict[str, typing.Any] 27 | 28 | 29 | class Encoder(abc.ABC): 30 | """Base class for message encoders. 31 | """ 32 | 33 | @abc.abstractmethod 34 | def encode(self, data: MessageData) -> bytes: # pragma: no cover 35 | """Convert message metadata into a bytestring. 36 | """ 37 | raise NotImplementedError 38 | 39 | @abc.abstractmethod 40 | def decode(self, data: bytes) -> MessageData: # pragma: no cover 41 | """Convert a bytestring into message metadata. 42 | """ 43 | raise NotImplementedError 44 | 45 | 46 | class JSONEncoder(Encoder): 47 | """Encodes messages as JSON. This is the default encoder. 48 | """ 49 | 50 | def encode(self, data: MessageData) -> bytes: 51 | return json.dumps(data, separators=(",", ":")).encode("utf-8") 52 | 53 | def decode(self, data: bytes) -> MessageData: 54 | try: 55 | data_str = data.decode("utf-8") 56 | except UnicodeDecodeError as e: 57 | raise DecodeError("failed to decode data %r" % (data,), data, e) from None 58 | 59 | try: 60 | return json.loads(data_str) 61 | except json.decoder.JSONDecodeError as e: 62 | raise DecodeError("failed to decode message %r" % (data_str,), data_str, e) from None 63 | 64 | 65 | class PickleEncoder(Encoder): 66 | """Pickles messages. 67 | 68 | Warning: 69 | This encoder is not secure against maliciously-constructed data. 70 | Use it at your own risk. 71 | """ 72 | 73 | def encode(self, data: MessageData) -> bytes: 74 | return pickle.dumps(data) 75 | 76 | def decode(self, data: bytes) -> MessageData: 77 | return pickle.loads(data) 78 | -------------------------------------------------------------------------------- /dramatiq/errors.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | 19 | class DramatiqError(Exception): # pragma: no cover 20 | """Base class for all dramatiq errors. 21 | """ 22 | 23 | def __init__(self, message): 24 | self.message = message 25 | 26 | def __str__(self): 27 | return str(self.message) or repr(self.message) 28 | 29 | 30 | class DecodeError(DramatiqError): 31 | """Raised when a message fails to decode. 32 | """ 33 | 34 | def __init__(self, message, data, error): 35 | super().__init__(message) 36 | self.data = data 37 | self.error = error 38 | 39 | 40 | class BrokerError(DramatiqError): 41 | """Base class for broker-related errors. 42 | """ 43 | 44 | 45 | class ActorNotFound(BrokerError): 46 | """Raised when a message is sent to an actor that hasn't been declared. 47 | """ 48 | 49 | 50 | class QueueNotFound(BrokerError): 51 | """Raised when a message is sent to an queue that hasn't been declared. 52 | """ 53 | 54 | 55 | class QueueJoinTimeout(DramatiqError): 56 | """Raised by brokers that support joining on queues when the join 57 | operation times out. 58 | """ 59 | 60 | 61 | class ConnectionError(BrokerError): 62 | """Base class for broker connection-related errors. 63 | """ 64 | 65 | 66 | class ConnectionFailed(ConnectionError): 67 | """Raised when a broker connection could not be opened. 68 | """ 69 | 70 | 71 | class ConnectionClosed(ConnectionError): 72 | """Raised when a broker connection is suddenly closed. 73 | """ 74 | 75 | 76 | class RateLimitExceeded(DramatiqError): 77 | """Raised when a rate limit has been exceeded. 78 | """ 79 | 80 | 81 | class Retry(DramatiqError): 82 | """Actors may raise this error when they should be retried. This 83 | behaves just like any other exception from the perspective of the 84 | :class:`Retries` middleware, the only 85 | difference is it doesn't get logged as an error. 86 | 87 | If the ``delay`` argument is provided, then the message will be 88 | retried after at least that amount of time (in milliseconds). 89 | """ 90 | 91 | def __init__(self, message="", delay=None): 92 | super().__init__(message) 93 | self.delay = delay 94 | -------------------------------------------------------------------------------- /dramatiq/generic.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from . import actor 19 | 20 | 21 | class generic_actor(type): 22 | """Meta for class-based actors. 23 | """ 24 | 25 | def __new__(metacls, name, bases, attrs): 26 | clazz = super().__new__(metacls, name, bases, attrs) 27 | meta = getattr(clazz, "Meta", object()) 28 | if not getattr(meta, "abstract", False): 29 | options = {name: getattr(meta, name) for name in vars(meta) if not name.startswith("_")} 30 | options.pop("abstract", False) 31 | 32 | clazz_instance = clazz() 33 | actor_registry = options.pop("actor", actor) 34 | actor_instance = actor_registry(clazz_instance, **options) 35 | setattr(clazz, "__getattr__", generic_actor.__getattr__) 36 | setattr(clazz_instance, "__actor__", actor_instance) 37 | return clazz_instance 38 | 39 | setattr(meta, "abstract", False) 40 | return clazz 41 | 42 | def __getattr__(cls, name): 43 | if "__actor__" not in cls.__dict__: 44 | # avoid infinite recursion on GenericActor 45 | raise AttributeError(f"type object {cls.__name__!r} has no attribute {name!r}") 46 | return getattr(cls.__actor__, name) 47 | 48 | 49 | class GenericActor(metaclass=generic_actor): 50 | """Base-class for class-based actors. 51 | 52 | Each subclass may define an inner class named ``Meta``. You can 53 | use the meta class to provide broker options for the actor. 54 | 55 | Classes that have ``abstract = True`` in their meta class are 56 | considered abstract base classes and are not converted into 57 | actors. You can't send these classes messages, you can only 58 | inherit from them. Actors that subclass abstract base classes 59 | inherit their parents' meta classes. 60 | 61 | Example: 62 | 63 | >>> class BaseTask(GenericActor): 64 | ... class Meta: 65 | ... abstract = True 66 | ... queue_name = "tasks" 67 | ... max_retries = 20 68 | ... 69 | ... def get_task_name(self): 70 | ... raise NotImplementedError 71 | ... 72 | ... def perform(self): 73 | ... print(f"Hello from {self.get_task_name()}!") 74 | 75 | >>> class FooTask(BaseTask): 76 | ... def get_task_name(self): 77 | ... return "Foo" 78 | 79 | >>> class BarTask(BaseTask): 80 | ... def get_task_name(self): 81 | ... return "Bar" 82 | 83 | >>> FooTask.send() 84 | >>> BarTask.send() 85 | 86 | Attributes: 87 | logger(Logger): The actor's logger. 88 | broker(Broker): The broker this actor is bound to. 89 | actor_name(str): The actor's name. 90 | queue_name(str): The actor's queue. 91 | priority(int): The actor's priority. 92 | options(dict): Arbitrary options that are passed to the broker 93 | and middleware. 94 | """ 95 | 96 | class Meta: 97 | abstract = True 98 | 99 | @property 100 | def __name__(self): 101 | """The default name of this actor. 102 | """ 103 | return type(self).__name__ 104 | 105 | def __call__(self, *args, **kwargs): 106 | return self.perform(*args, **kwargs) 107 | 108 | def perform(self): 109 | """This is the method that gets called when the actor receives 110 | a message. All non-abstract subclasses must implement this 111 | method. 112 | """ 113 | raise NotImplementedError("%s does not implement perform()" % self.__name__) 114 | -------------------------------------------------------------------------------- /dramatiq/logging.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import inspect 19 | import logging 20 | 21 | 22 | def get_logger(module, name=None): 23 | logger_fqn = module 24 | if name is not None: 25 | if inspect.isclass(name): 26 | name = name.__name__ 27 | logger_fqn += "." + name 28 | 29 | return logging.getLogger(logger_fqn) 30 | -------------------------------------------------------------------------------- /dramatiq/middleware/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018,2019 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .age_limit import AgeLimit 19 | from .asyncio import AsyncIO 20 | from .callbacks import Callbacks 21 | from .current_message import CurrentMessage 22 | from .group_callbacks import GroupCallbacks 23 | from .middleware import Middleware, MiddlewareError, SkipMessage 24 | from .pipelines import Pipelines 25 | from .prometheus import Prometheus 26 | from .retries import Retries 27 | from .shutdown import Shutdown, ShutdownNotifications 28 | from .threading import Interrupt, raise_thread_exception 29 | from .time_limit import TimeLimit, TimeLimitExceeded 30 | 31 | __all__ = [ 32 | # Basics 33 | "Middleware", "MiddlewareError", "SkipMessage", 34 | 35 | # Threading 36 | "Interrupt", "raise_thread_exception", 37 | 38 | # Middlewares 39 | "AgeLimit", "AsyncIO", "Callbacks", "CurrentMessage", "GroupCallbacks", 40 | "Pipelines", "Retries", "Shutdown", "ShutdownNotifications", "TimeLimit", 41 | "TimeLimitExceeded", "Prometheus", 42 | ] 43 | 44 | 45 | #: The list of middleware that are enabled by default. 46 | default_middleware = [ 47 | Prometheus, AgeLimit, TimeLimit, ShutdownNotifications, Callbacks, Pipelines, Retries 48 | ] 49 | -------------------------------------------------------------------------------- /dramatiq/middleware/age_limit.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from ..common import current_millis 19 | from ..logging import get_logger 20 | from .middleware import Middleware, SkipMessage 21 | 22 | 23 | class AgeLimit(Middleware): 24 | """Middleware that drops messages that have been in the queue for 25 | too long. 26 | 27 | Parameters: 28 | max_age(int): The default message age limit in milliseconds. 29 | Defaults to ``None``, meaning that messages can exist 30 | indefinitely. 31 | """ 32 | 33 | def __init__(self, *, max_age=None): 34 | self.logger = get_logger(__name__, type(self)) 35 | self.max_age = max_age 36 | 37 | @property 38 | def actor_options(self): 39 | return {"max_age"} 40 | 41 | def before_process_message(self, broker, message): 42 | actor = broker.get_actor(message.actor_name) 43 | max_age = message.options.get("max_age") or actor.options.get("max_age", self.max_age) 44 | if not max_age: 45 | return 46 | 47 | if current_millis() - message.message_timestamp >= max_age: 48 | self.logger.warning("Message %r has exceeded its age limit.", message.message_id) 49 | message.fail() 50 | raise SkipMessage("Message age limit exceeded") 51 | -------------------------------------------------------------------------------- /dramatiq/middleware/asyncio.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2023 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from __future__ import annotations 19 | 20 | from ..asyncio import EventLoopThread, get_event_loop_thread, set_event_loop_thread 21 | from ..logging import get_logger 22 | from .middleware import Middleware 23 | 24 | 25 | class AsyncIO(Middleware): 26 | """This middleware manages the event loop thread for async actors. 27 | """ 28 | 29 | def __init__(self): 30 | self.logger = get_logger(__name__, type(self)) 31 | 32 | def before_worker_boot(self, broker, worker): 33 | event_loop_thread = EventLoopThread(self.logger) 34 | event_loop_thread.start(timeout=1.0) 35 | set_event_loop_thread(event_loop_thread) 36 | 37 | def after_worker_shutdown(self, broker, worker): 38 | event_loop_thread = get_event_loop_thread() 39 | event_loop_thread.stop() 40 | event_loop_thread.join() 41 | set_event_loop_thread(None) 42 | -------------------------------------------------------------------------------- /dramatiq/middleware/callbacks.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .middleware import Middleware 19 | 20 | 21 | class Callbacks(Middleware): 22 | """Middleware that lets you chain success and failure callbacks 23 | onto Actors. 24 | 25 | Parameters: 26 | on_failure(str): The name of an actor to send a message to on 27 | failure. 28 | on_success(str): The name of an actor to send a message to on 29 | success. 30 | """ 31 | 32 | @property 33 | def actor_options(self): 34 | return { 35 | "on_failure", 36 | "on_success", 37 | } 38 | 39 | def after_process_message(self, broker, message, *, result=None, exception=None): 40 | actor = broker.get_actor(message.actor_name) 41 | if exception is None: 42 | target_actor_name = message.options.get("on_success") or actor.options.get("on_success") 43 | if target_actor_name: 44 | target_actor = broker.get_actor(target_actor_name) 45 | target_actor.send(message.asdict(), result) 46 | 47 | else: 48 | target_actor_name = message.options.get("on_failure") or actor.options.get("on_failure") 49 | if target_actor_name: 50 | target_actor = broker.get_actor(target_actor_name) 51 | target_actor.send(message.asdict(), { 52 | "type": type(exception).__name__, 53 | "message": str(exception), 54 | }) 55 | -------------------------------------------------------------------------------- /dramatiq/middleware/current_message.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2019 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import contextvars 19 | from typing import TYPE_CHECKING, Any, Optional 20 | 21 | from .middleware import Middleware 22 | 23 | if TYPE_CHECKING: 24 | from ..message import Message 25 | 26 | 27 | class CurrentMessage(Middleware): 28 | """Middleware that exposes the current message via a thread-local 29 | variable. 30 | 31 | Example: 32 | >>> import dramatiq 33 | >>> from dramatiq.middleware import CurrentMessage 34 | 35 | >>> @dramatiq.actor 36 | ... def example(x): 37 | ... print(CurrentMessage.get_current_message()) 38 | ... 39 | >>> example.send(1) 40 | 41 | """ 42 | 43 | _MESSAGE: contextvars.ContextVar[ 44 | "Optional[Message[Any]]" 45 | ] = contextvars.ContextVar("_MESSAGE", default=None) 46 | 47 | @classmethod 48 | def get_current_message(cls) -> "Optional[Message[Any]]": 49 | """Get the message that triggered the current actor. Messages 50 | are thread local so this returns ``None`` when called outside 51 | of actor code. 52 | """ 53 | return cls._MESSAGE.get() 54 | 55 | def before_process_message(self, broker, message): 56 | self._MESSAGE.set(message) 57 | 58 | def after_process_message(self, broker, message, *, result=None, exception=None): 59 | self._MESSAGE.set(None) 60 | -------------------------------------------------------------------------------- /dramatiq/middleware/group_callbacks.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import os 19 | 20 | from ..rate_limits import Barrier 21 | from .middleware import Middleware 22 | 23 | GROUP_CALLBACK_BARRIER_TTL = int(os.getenv("dramatiq_group_callback_barrier_ttl", "86400000")) 24 | 25 | 26 | class GroupCallbacks(Middleware): 27 | def __init__(self, rate_limiter_backend): 28 | self.rate_limiter_backend = rate_limiter_backend 29 | 30 | def after_process_message(self, broker, message, *, result=None, exception=None): 31 | from ..message import Message 32 | 33 | if exception is None: 34 | group_completion_uuid = message.options.get("group_completion_uuid") 35 | group_completion_callbacks = message.options.get("group_completion_callbacks") 36 | if group_completion_uuid and group_completion_callbacks: 37 | barrier = Barrier(self.rate_limiter_backend, group_completion_uuid, ttl=GROUP_CALLBACK_BARRIER_TTL) 38 | if barrier.wait(block=False): 39 | for message in group_completion_callbacks: 40 | broker.enqueue(Message(**message)) 41 | -------------------------------------------------------------------------------- /dramatiq/middleware/pipelines.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .middleware import Middleware 19 | 20 | 21 | class Pipelines(Middleware): 22 | """Middleware that lets you pipe actors together so that the 23 | output of one actor feeds into the input of another. 24 | 25 | Parameters: 26 | pipe_ignore(bool): When True, ignores the result of the previous 27 | actor in the pipeline. 28 | pipe_target(dict): A message representing the actor the current 29 | result should be fed into. 30 | """ 31 | 32 | @property 33 | def actor_options(self): 34 | return { 35 | "pipe_ignore", 36 | "pipe_target", 37 | } 38 | 39 | def after_process_message(self, broker, message, *, result=None, exception=None): 40 | # Since Pipelines is a default middleware, this import has to 41 | # happen at runtime in order to avoid a cyclic dependency 42 | # from broker -> pipelines -> messages -> broker. 43 | from ..message import Message 44 | 45 | if exception is not None or message.failed: 46 | return 47 | 48 | actor = broker.get_actor(message.actor_name) 49 | message_data = message.options.get("pipe_target") 50 | if message_data is not None: 51 | next_message = Message(**message_data) 52 | pipe_ignore = next_message.options.get("pipe_ignore") or actor.options.get("pipe_ignore") 53 | if not pipe_ignore: 54 | next_message = next_message.copy(args=next_message.args + (result,)) 55 | 56 | broker.enqueue(next_message, delay=next_message.options.get("delay")) 57 | -------------------------------------------------------------------------------- /dramatiq/middleware/threading.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2023 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from ..threading import Interrupt, current_platform, is_gevent_active, raise_thread_exception, supported_platforms 19 | 20 | __all__ = [ 21 | "Interrupt", 22 | "current_platform", 23 | "is_gevent_active", 24 | "raise_thread_exception", 25 | "supported_platforms", 26 | ] 27 | -------------------------------------------------------------------------------- /dramatiq/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/dramatiq/py.typed -------------------------------------------------------------------------------- /dramatiq/rate_limits/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .backend import RateLimiterBackend 19 | from .barrier import Barrier 20 | from .bucket import BucketRateLimiter 21 | from .concurrent import ConcurrentRateLimiter 22 | from .rate_limiter import RateLimiter, RateLimitExceeded 23 | from .window import WindowRateLimiter 24 | 25 | __all__ = [ 26 | "RateLimiterBackend", "RateLimiter", "RateLimitExceeded", "Barrier", 27 | "BucketRateLimiter", "ConcurrentRateLimiter", "WindowRateLimiter", 28 | ] 29 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/backend.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | 19 | class RateLimiterBackend: 20 | """ABC for rate limiter backends. 21 | """ 22 | 23 | def add(self, key, value, ttl): # pragma: no cover 24 | """Add a key to the backend iff it doesn't exist. 25 | 26 | Parameters: 27 | key(str): The key to add. 28 | value(int): The value to add. 29 | ttl(int): The max amount of time in milliseconds the key can 30 | live in the backend for. 31 | """ 32 | raise NotImplementedError 33 | 34 | def incr(self, key, amount, maximum, ttl): # pragma: no cover 35 | """Atomically increment a key in the backend up to the given 36 | maximum. 37 | 38 | Parameters: 39 | key(str): The key to increment. 40 | amount(int): The amount to increment the value by. 41 | maximum(int): The maximum amount the value can have. 42 | ttl(int): The max amount of time in milliseconds the key can 43 | live in the backend for. 44 | 45 | Returns: 46 | bool: True if the key was successfully incremented. 47 | """ 48 | raise NotImplementedError 49 | 50 | def decr(self, key, amount, minimum, ttl): # pragma: no cover 51 | """Atomically decrement a key in the backend up to the given 52 | maximum. 53 | 54 | Parameters: 55 | key(str): The key to decrement. 56 | amount(int): The amount to decrement the value by. 57 | minimum(int): The minimum amount the value can have. 58 | ttl(int): The max amount of time in milliseconds the key can 59 | live in the backend for. 60 | 61 | Returns: 62 | bool: True if the key was successfully decremented. 63 | """ 64 | raise NotImplementedError 65 | 66 | def incr_and_sum(self, key, keys, amount, maximum, ttl): # pragma: no cover 67 | """Atomically increment a key unless the sum of keys is greater 68 | than the given maximum. 69 | 70 | Parameters: 71 | key(str): The key to increment. 72 | keys(callable): A callable to return the list of keys to be 73 | summed over. 74 | amount(int): The amount to decrement the value by. 75 | maximum(int): The maximum sum of the keys. 76 | ttl(int): The max amount of time in milliseconds the key can 77 | live in the backend for. 78 | 79 | Returns: 80 | bool: True if the key was successfully incremented. 81 | """ 82 | raise NotImplementedError 83 | 84 | def wait(self, key, timeout): # pragma: no cover 85 | """Wait until an event is published to the given key or the 86 | timeout expires. This is used to implement efficient blocking 87 | against a synchronized resource. 88 | 89 | Parameters: 90 | key(str): The key to wait on. 91 | timeout(int): The timeout in milliseconds. 92 | 93 | Returns: 94 | bool: True if en event was published before the timeout. 95 | """ 96 | raise NotImplementedError 97 | 98 | def wait_notify(self, key, ttl): # pragma: no cover 99 | """Notify parties wait()ing on a key that an event has 100 | occurred. The default implementation is a no-op. 101 | 102 | Parameters: 103 | key(str): The key to notify on. 104 | ttl(int): The max amount of time in milliseconds that the 105 | notification should exist for. 106 | 107 | Returns: 108 | None 109 | """ 110 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/backends/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import warnings 19 | 20 | from .stub import StubBackend 21 | 22 | try: 23 | from .memcached import MemcachedBackend 24 | except ImportError: # pragma: no cover 25 | warnings.warn( 26 | "MemcachedBackend is not available. Run `pip install dramatiq[memcached]` " 27 | "to add support for that backend.", 28 | category=ImportWarning, 29 | stacklevel=2, 30 | ) 31 | 32 | try: 33 | from .redis import RedisBackend 34 | except ImportError: # pragma: no cover 35 | warnings.warn( 36 | "RedisBackend is not available. Run `pip install dramatiq[redis]` " 37 | "to add support for that backend.", 38 | category=ImportWarning, 39 | stacklevel=2, 40 | ) 41 | 42 | 43 | __all__ = ["StubBackend", "MemcachedBackend", "RedisBackend"] 44 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/backends/memcached.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from pylibmc import Client, ClientPool, NotFound 19 | 20 | from ..backend import RateLimiterBackend 21 | 22 | 23 | class MemcachedBackend(RateLimiterBackend): 24 | """A rate limiter backend for Memcached_. 25 | 26 | Examples: 27 | 28 | >>> from dramatiq.rate_limits.backends import MemcachedBackend 29 | >>> backend = MemcachedBackend(servers=["127.0.0.1"], binary=True) 30 | 31 | Parameters: 32 | pool(ClientPool): An optional pylibmc client pool to use. If 33 | this is passed, all other connection params are ignored. 34 | pool_size(int): The size of the connection pool to use. 35 | **parameters: Connection parameters are passed directly 36 | to :class:`pylibmc.Client`. 37 | 38 | .. _memcached: https://memcached.org 39 | """ 40 | 41 | def __init__(self, *, pool=None, pool_size=8, **parameters): 42 | behaviors = parameters.setdefault("behaviors", {}) 43 | behaviors["cas"] = True 44 | 45 | self.pool = pool or ClientPool(Client(**parameters), pool_size) 46 | 47 | def add(self, key, value, ttl): 48 | with self.pool.reserve(block=True) as client: 49 | return client.add(key, value, time=int(ttl / 1000)) 50 | 51 | def incr(self, key, amount, maximum, ttl): 52 | with self.pool.reserve(block=True) as client: 53 | return client.incr(key, amount) <= maximum 54 | 55 | def decr(self, key, amount, minimum, ttl): 56 | with self.pool.reserve(block=True) as client: 57 | return client.decr(key, amount) >= minimum 58 | 59 | def incr_and_sum(self, key, keys, amount, maximum, ttl): 60 | ttl = int(ttl / 1000) 61 | with self.pool.reserve(block=True) as client: 62 | client.add(key, 0, time=ttl) 63 | 64 | while True: 65 | value, cid = client.gets(key) 66 | if cid is None: 67 | return False 68 | 69 | value += amount 70 | if value > maximum: 71 | return False 72 | 73 | # TODO: Drop non-callable keys in Dramatiq v2. 74 | key_list = keys() if callable(keys) else keys 75 | mapping = client.get_multi(key_list) 76 | total = amount + sum(mapping.values()) 77 | if total > maximum: 78 | return False 79 | 80 | try: 81 | swapped = client.cas(key, value, cid, ttl) 82 | if swapped: 83 | return True 84 | except NotFound: # pragma: no cover 85 | continue 86 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/backends/redis.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import redis 19 | 20 | from ..backend import RateLimiterBackend 21 | 22 | 23 | class RedisBackend(RateLimiterBackend): 24 | """A rate limiter backend for Redis_. 25 | 26 | Parameters: 27 | client(Redis): An optional client. If this is passed, 28 | then all other parameters are ignored. 29 | url(str): An optional connection URL. If both a URL and 30 | connection parameters are provided, the URL is used. 31 | **parameters: Connection parameters are passed directly 32 | to :class:`redis.Redis`. 33 | 34 | .. _redis: https://redis.io 35 | """ 36 | 37 | def __init__(self, *, client=None, url=None, **parameters): 38 | if url is not None: 39 | parameters["connection_pool"] = redis.ConnectionPool.from_url(url) 40 | 41 | # TODO: Replace usages of StrictRedis (redis-py 2.x) with Redis in Dramatiq 2.0. 42 | self.client = client or redis.StrictRedis(**parameters) 43 | 44 | def add(self, key, value, ttl): 45 | return bool(self.client.set(key, value, px=ttl, nx=True)) 46 | 47 | def incr(self, key, amount, maximum, ttl): 48 | with self.client.pipeline() as pipe: 49 | while True: 50 | try: 51 | pipe.watch(key) 52 | value = int(pipe.get(key) or b"0") 53 | value += amount 54 | if value > maximum: 55 | return False 56 | 57 | pipe.multi() 58 | pipe.set(key, value, px=ttl) 59 | pipe.execute() 60 | return True 61 | except redis.WatchError: 62 | continue 63 | 64 | def decr(self, key, amount, minimum, ttl): 65 | with self.client.pipeline() as pipe: 66 | while True: 67 | try: 68 | pipe.watch(key) 69 | value = int(pipe.get(key) or b"0") 70 | value -= amount 71 | if value < minimum: 72 | return False 73 | 74 | pipe.multi() 75 | pipe.set(key, value, px=ttl) 76 | pipe.execute() 77 | return True 78 | except redis.WatchError: 79 | continue 80 | 81 | def incr_and_sum(self, key, keys, amount, maximum, ttl): 82 | with self.client.pipeline() as pipe: 83 | while True: 84 | try: 85 | # TODO: Drop non-callable keys in Dramatiq v2. 86 | key_list = keys() if callable(keys) else keys 87 | pipe.watch(key, *key_list) 88 | value = int(pipe.get(key) or b"0") 89 | value += amount 90 | if value > maximum: 91 | return False 92 | 93 | # Fetch keys again to account for net/server latency. 94 | values = pipe.mget(keys() if callable(keys) else keys) 95 | total = amount + sum(int(n) for n in values if n) 96 | if total > maximum: 97 | return False 98 | 99 | pipe.multi() 100 | pipe.set(key, value, px=ttl) 101 | pipe.execute() 102 | return True 103 | except redis.WatchError: 104 | continue 105 | 106 | def wait(self, key, timeout): 107 | assert timeout is None or timeout >= 1000, "wait timeouts must be >= 1000" 108 | event = self.client.brpoplpush(key, key, (timeout or 0) // 1000) 109 | return event == b"x" 110 | 111 | def wait_notify(self, key, ttl): 112 | with self.client.pipeline() as pipe: 113 | pipe.rpush(key, b"x") 114 | pipe.pexpire(key, ttl) 115 | pipe.execute() 116 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/backends/stub.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import time 19 | from collections import defaultdict 20 | from threading import Condition, Lock 21 | 22 | from ..backend import RateLimiterBackend 23 | 24 | 25 | class StubBackend(RateLimiterBackend): 26 | """An in-memory rate limiter backend. For use in unit tests. 27 | """ 28 | 29 | def __init__(self): 30 | self.conditions = defaultdict(lambda: Condition(self.mutex)) 31 | self.mutex = Lock() 32 | self.db = {} 33 | 34 | def add(self, key, value, ttl): 35 | with self.mutex: 36 | res = self._get(key) 37 | if res is not None: 38 | return False 39 | 40 | return self._put(key, value, ttl) 41 | 42 | def incr(self, key, amount, maximum, ttl): 43 | with self.mutex: 44 | value = self._get(key, default=0) + amount 45 | if value > maximum: 46 | return False 47 | 48 | return self._put(key, value, ttl) 49 | 50 | def decr(self, key, amount, minimum, ttl): 51 | with self.mutex: 52 | value = self._get(key, default=0) - amount 53 | if value < minimum: 54 | return False 55 | 56 | return self._put(key, value, ttl) 57 | 58 | def incr_and_sum(self, key, keys, amount, maximum, ttl): 59 | self.add(key, 0, ttl) 60 | with self.mutex: 61 | value = self._get(key, default=0) + amount 62 | if value > maximum: 63 | return False 64 | 65 | # TODO: Drop non-callable keys in Dramatiq v2. 66 | key_list = keys() if callable(keys) else keys 67 | values = sum(self._get(k, default=0) for k in key_list) 68 | total = amount + values 69 | if total > maximum: 70 | return False 71 | 72 | return self._put(key, value, ttl) 73 | 74 | def wait(self, key, timeout): 75 | cond = self.conditions[key] 76 | with cond: 77 | return cond.wait(timeout=timeout / 1000) 78 | 79 | def wait_notify(self, key, ttl): 80 | cond = self.conditions[key] 81 | with cond: 82 | cond.notify_all() 83 | 84 | def _get(self, key, *, default=None): 85 | value, expiration = self.db.get(key, (None, None)) 86 | if expiration and time.monotonic() < expiration: 87 | return value 88 | return default 89 | 90 | def _put(self, key, value, ttl): 91 | self.db[key] = (value, time.monotonic() + ttl / 1000) 92 | return True 93 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/barrier.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | 19 | class Barrier: 20 | """A distributed barrier. 21 | 22 | Examples: 23 | 24 | >>> from dramatiq.rate_limits import Barrier 25 | >>> from dramatiq.rate_limits.backends import RedisBackend 26 | 27 | >>> backend = RedisBackend() 28 | >>> barrier = Barrier(backend, "some-barrier", ttl=30_000) 29 | 30 | >>> created = barrier.create(parties=3) 31 | >>> barrier.wait(block=False) 32 | False 33 | >>> barrier.wait(block=False) 34 | False 35 | >>> barrier.wait(block=False) 36 | True 37 | 38 | Parameters: 39 | backend(BarrierBackend): The barrier backend to use. 40 | key(str): The key for the barrier. 41 | ttl(int): The TTL for the barrier key, in milliseconds. 42 | """ 43 | 44 | def __init__(self, backend, key, *, ttl=900000): 45 | self.backend = backend 46 | self.key = key 47 | self.key_events = key + "@events" 48 | self.ttl = ttl 49 | 50 | def create(self, parties): 51 | """Create the barrier for the given number of parties. 52 | 53 | Parameters: 54 | parties(int): The number of parties to wait for. 55 | 56 | Returns: 57 | bool: Whether or not the new barrier was successfully created. 58 | """ 59 | assert parties > 0, "parties must be a positive integer." 60 | return self.backend.add(self.key, parties, self.ttl) 61 | 62 | def wait(self, *, block=True, timeout=None): 63 | """Signal that a party has reached the barrier. 64 | 65 | Warning: 66 | Barrier blocking is currently only supported by the stub and 67 | Redis backends. 68 | 69 | Warning: 70 | Re-using keys between blocking calls may lead to undefined 71 | behaviour. Make sure your barrier keys are always unique 72 | (use a UUID). 73 | 74 | Parameters: 75 | block(bool): Whether or not to block while waiting for the 76 | other parties. 77 | timeout(int): The maximum number of milliseconds to wait for 78 | the barrier to be cleared. 79 | 80 | Returns: 81 | bool: Whether or not the barrier has been reached by all parties. 82 | """ 83 | cleared = not self.backend.decr(self.key, 1, 1, self.ttl) 84 | if cleared: 85 | self.backend.wait_notify(self.key_events, self.ttl) 86 | return True 87 | 88 | if block: 89 | return self.backend.wait(self.key_events, timeout) 90 | 91 | return False 92 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/bucket.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import time 19 | 20 | from .rate_limiter import RateLimiter 21 | 22 | 23 | class BucketRateLimiter(RateLimiter): 24 | """A rate limiter that ensures that only up to `limit` operations 25 | may happen over some time interval. 26 | 27 | Examples: 28 | 29 | Up to 10 operations every second: 30 | 31 | >>> BucketRateLimiter(backend, "some-key", limit=10, bucket=1_000) 32 | 33 | Up to 1 operation every minute: 34 | 35 | >>> BucketRateLimiter(backend, "some-key", limit=1, bucket=60_000) 36 | 37 | Warning: 38 | 39 | Bucket rate limits are cheap to maintain but are susceptible to 40 | burst "attacks". Given a bucket rate limit of 100 per minute, 41 | an attacker could make a burst of 100 calls in the last second 42 | of a minute and then another 100 calls in the first second of 43 | the subsequent minute. 44 | 45 | For a rate limiter that doesn't have this problem (but is more 46 | expensive to maintain), see |WindowRateLimiter|. 47 | 48 | Parameters: 49 | backend(RateLimiterBackend): The backend to use. 50 | key(str): The key to rate limit on. 51 | limit(int): The maximum number of operations per bucket per key. 52 | bucket(int): The bucket interval in milliseconds. 53 | 54 | .. |WindowRateLimiter| replace:: :class:`WindowRateLimiter` 55 | """ 56 | 57 | def __init__(self, backend, key, *, limit=1, bucket=1000): 58 | assert limit >= 1, "limit must be positive" 59 | 60 | super().__init__(backend, key) 61 | self.limit = limit 62 | self.bucket = bucket 63 | 64 | def _acquire(self): 65 | timestamp = int(time.time() * 1000) 66 | current_timestamp = timestamp - (timestamp % self.bucket) 67 | current_key = "%s@%d" % (self.key, current_timestamp) 68 | added = self.backend.add(current_key, 1, ttl=self.bucket) 69 | if added: 70 | return True 71 | 72 | return self.backend.incr(current_key, 1, maximum=self.limit, ttl=self.bucket) 73 | 74 | def _release(self): 75 | pass 76 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/concurrent.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .rate_limiter import RateLimiter 19 | 20 | 21 | class ConcurrentRateLimiter(RateLimiter): 22 | """A rate limiter that ensures that only `limit` concurrent 23 | operations may happen at the same time. 24 | 25 | Note: 26 | 27 | You can use a concurrent rate limiter of size 1 to get a 28 | distributed mutex. 29 | 30 | Parameters: 31 | backend(RateLimiterBackend): The backend to use. 32 | key(str): The key to rate limit on. 33 | limit(int): The maximum number of concurrent operations per key. 34 | ttl(int): The time in milliseconds that keys may live for. 35 | """ 36 | 37 | def __init__(self, backend, key, *, limit=1, ttl=900000): 38 | assert limit >= 1, "limit must be positive" 39 | 40 | super().__init__(backend, key) 41 | self.limit = limit 42 | self.ttl = ttl 43 | 44 | def _acquire(self): 45 | added = self.backend.add(self.key, 1, ttl=self.ttl) 46 | if added: 47 | return True 48 | 49 | return self.backend.incr(self.key, 1, maximum=self.limit, ttl=self.ttl) 50 | 51 | def _release(self): 52 | return self.backend.decr(self.key, 1, minimum=0, ttl=self.ttl) 53 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/rate_limiter.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from contextlib import contextmanager 19 | 20 | from ..errors import RateLimitExceeded 21 | 22 | 23 | class RateLimiter: 24 | """ABC for rate limiters. 25 | 26 | Examples: 27 | 28 | >>> from dramatiq.rate_limits.backends import RedisBackend 29 | 30 | >>> backend = RedisBackend() 31 | >>> limiter = ConcurrentRateLimiter(backend, "distributed-mutex", limit=1) 32 | 33 | >>> with limiter.acquire(raise_on_failure=False) as acquired: 34 | ... if not acquired: 35 | ... print("Mutex not acquired.") 36 | ... return 37 | ... 38 | ... print("Mutex acquired.") 39 | 40 | Parameters: 41 | backend(RateLimiterBackend): The rate limiting backend to use. 42 | key(str): The key to rate limit on. 43 | """ 44 | 45 | def __init__(self, backend, key): 46 | self.backend = backend 47 | self.key = key 48 | 49 | def _acquire(self): # pragma: no cover 50 | raise NotImplementedError 51 | 52 | def _release(self): # pragma: no cover 53 | raise NotImplementedError 54 | 55 | @contextmanager 56 | def acquire(self, *, raise_on_failure=True): 57 | """Attempt to acquire a slot under this rate limiter. 58 | 59 | Parameters: 60 | raise_on_failure(bool): Whether or not failures should raise an 61 | exception. If this is false, the context manager will instead 62 | return a boolean value representing whether or not the rate 63 | limit slot was acquired. 64 | 65 | Returns: 66 | bool: Whether or not the slot could be acquired. 67 | """ 68 | acquired = False 69 | 70 | try: 71 | acquired = self._acquire() 72 | if raise_on_failure and not acquired: 73 | raise RateLimitExceeded("rate limit exceeded for key %(key)r" % vars(self)) 74 | 75 | yield acquired 76 | finally: 77 | if acquired: 78 | self._release() 79 | -------------------------------------------------------------------------------- /dramatiq/rate_limits/window.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import time 19 | 20 | from .rate_limiter import RateLimiter 21 | 22 | 23 | class WindowRateLimiter(RateLimiter): 24 | """A rate limiter that ensures that only `limit` operations may 25 | happen over some sliding window. 26 | 27 | Note: 28 | 29 | Windows are in seconds rather that milliseconds. This is 30 | different from most durations and intervals used in Dramatiq, 31 | because keeping metadata at the millisecond level is far too 32 | expensive for most use cases. 33 | 34 | Parameters: 35 | backend(RateLimiterBackend): The backend to use. 36 | key(str): The key to rate limit on. 37 | limit(int): The maximum number of operations per window per key. 38 | window(int): The window size in *seconds*. The wider the 39 | window, the more expensive it is to maintain. 40 | """ 41 | 42 | def __init__(self, backend, key, *, limit=1, window=1): 43 | assert limit >= 1, "limit must be positive" 44 | assert window >= 1, "window must be positive" 45 | 46 | super().__init__(backend, key) 47 | self.limit = limit 48 | self.window = window 49 | self.window_millis = window * 1000 50 | 51 | def _get_keys(self): 52 | timestamp = int(time.time()) 53 | return ["%s@%s" % (self.key, timestamp - i) for i in range(self.window)] 54 | 55 | def _acquire(self): 56 | keys = self._get_keys() 57 | return self.backend.incr_and_sum( 58 | keys[0], self._get_keys, 1, 59 | maximum=self.limit, 60 | ttl=self.window_millis, 61 | ) 62 | 63 | def _release(self): 64 | pass 65 | -------------------------------------------------------------------------------- /dramatiq/results/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .backend import Missing, ResultBackend 19 | from .errors import ResultError, ResultFailure, ResultMissing, ResultTimeout 20 | from .middleware import Results 21 | 22 | __all__ = ["Missing", "ResultBackend", "ResultError", "ResultFailure", "ResultTimeout", "ResultMissing", "Results"] 23 | -------------------------------------------------------------------------------- /dramatiq/results/backends/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import warnings 19 | 20 | from .stub import StubBackend 21 | 22 | __all__ = ["StubBackend", "MemcachedBackend", "RedisBackend"] 23 | 24 | 25 | def __getattr__(name): 26 | module_importer = _module_importers.get(name) 27 | if module_importer is None: 28 | raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 29 | 30 | return module_importer() 31 | 32 | 33 | def import_memcached(): 34 | try: 35 | from .memcached import MemcachedBackend 36 | 37 | return MemcachedBackend 38 | except ModuleNotFoundError: 39 | warnings.warn( 40 | "MemcachedBackend is not available. Run `pip install dramatiq[memcached]` " 41 | "to add support for that backend.", 42 | category=ImportWarning, 43 | stacklevel=2, 44 | ) 45 | raise 46 | 47 | 48 | def import_redis(): 49 | try: 50 | from .redis import RedisBackend 51 | 52 | return RedisBackend 53 | except ModuleNotFoundError: 54 | warnings.warn( 55 | "RedisBackend is not available. Run `pip install dramatiq[redis]` " 56 | "to add support for that backend.", 57 | category=ImportWarning, 58 | stacklevel=2, 59 | ) 60 | raise 61 | 62 | 63 | _module_importers = { 64 | "MemcachedBackend": import_memcached, 65 | "RedisBackend": import_redis, 66 | } 67 | -------------------------------------------------------------------------------- /dramatiq/results/backends/memcached.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from pylibmc import Client, ClientPool 19 | 20 | from ..backend import Missing, ResultBackend 21 | 22 | 23 | class MemcachedBackend(ResultBackend): 24 | """A result backend for Memcached_. This backend uses long 25 | polling to retrieve results. 26 | 27 | Parameters: 28 | namespace(str): A string with which to prefix result keys. 29 | encoder(Encoder): The encoder to use when storing and retrieving 30 | result data. Defaults to :class:`.JSONEncoder`. 31 | pool(ClientPool): An optional pylibmc client pool to use. If 32 | this is passed, all other connection params are ignored. 33 | pool_size(int): The size of the connection pool to use. 34 | **parameters: Connection parameters are passed directly 35 | to :class:`pylibmc.Client`. 36 | 37 | .. _memcached: https://memcached.org 38 | """ 39 | 40 | def __init__(self, *, namespace="dramatiq-results", encoder=None, pool=None, pool_size=8, **parameters): 41 | super().__init__(namespace=namespace, encoder=encoder) 42 | self.pool = pool or ClientPool(Client(**parameters), pool_size) 43 | 44 | def _get(self, message_key): 45 | with self.pool.reserve(block=True) as client: 46 | data = client.get(message_key) 47 | if data is not None: 48 | return self.encoder.decode(data) 49 | return Missing 50 | 51 | def _store(self, message_key, result, ttl): 52 | result_data = self.encoder.encode(result) 53 | with self.pool.reserve(block=True) as client: 54 | client.set(message_key, result_data, time=int(ttl / 1000)) 55 | -------------------------------------------------------------------------------- /dramatiq/results/backends/redis.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018,2019,2020 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import redis 19 | 20 | from ..backend import DEFAULT_TIMEOUT, ResultBackend, ResultMissing, ResultTimeout 21 | 22 | 23 | class RedisBackend(ResultBackend): 24 | """A result backend for Redis_. This is the recommended result 25 | backend as waiting for a result is resource efficient. 26 | 27 | Parameters: 28 | namespace(str): A string with which to prefix result keys. 29 | encoder(Encoder): The encoder to use when storing and retrieving 30 | result data. Defaults to :class:`.JSONEncoder`. 31 | client(Redis): An optional client. If this is passed, 32 | then all other parameters are ignored. 33 | url(str): An optional connection URL. If both a URL and 34 | connection parameters are provided, the URL is used. 35 | **parameters: Connection parameters are passed directly 36 | to :class:`redis.Redis`. 37 | 38 | .. _redis: https://redis.io 39 | """ 40 | 41 | def __init__(self, *, namespace="dramatiq-results", encoder=None, client=None, url=None, **parameters): 42 | super().__init__(namespace=namespace, encoder=encoder) 43 | 44 | if url: 45 | parameters["connection_pool"] = redis.ConnectionPool.from_url(url) 46 | 47 | # TODO: Replace usages of StrictRedis (redis-py 2.x) with Redis in Dramatiq 2.0. 48 | self.client = client or redis.StrictRedis(**parameters) 49 | 50 | def get_result(self, message, *, block=False, timeout=None): 51 | """Get a result from the backend. 52 | 53 | Warning: 54 | Sub-second timeouts are not respected by this backend. 55 | 56 | Parameters: 57 | message(Message) 58 | block(bool): Whether or not to block until a result is set. 59 | timeout(int): The maximum amount of time, in ms, to wait for 60 | a result when block is True. Defaults to 10 seconds. 61 | 62 | Raises: 63 | ResultMissing: When block is False and the result isn't set. 64 | ResultTimeout: When waiting for a result times out. 65 | 66 | Returns: 67 | object: The result. 68 | """ 69 | if timeout is None: 70 | timeout = DEFAULT_TIMEOUT 71 | 72 | message_key = self.build_message_key(message) 73 | if block: 74 | timeout = int(timeout / 1000) 75 | if timeout == 0: 76 | data = self.client.rpoplpush(message_key, message_key) 77 | else: 78 | data = self.client.brpoplpush(message_key, message_key, timeout) 79 | 80 | if data is None: 81 | raise ResultTimeout(message) 82 | 83 | else: 84 | data = self.client.lindex(message_key, 0) 85 | if data is None: 86 | raise ResultMissing(message) 87 | 88 | return self.unwrap_result(self.encoder.decode(data)) 89 | 90 | def _store(self, message_key, result, ttl): 91 | with self.client.pipeline() as pipe: 92 | pipe.delete(message_key) 93 | pipe.lpush(message_key, self.encoder.encode(result)) 94 | pipe.pexpire(message_key, ttl) 95 | pipe.execute() 96 | -------------------------------------------------------------------------------- /dramatiq/results/backends/stub.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import time 19 | from typing import Dict, Optional, Tuple 20 | 21 | from ..backend import Missing, ResultBackend 22 | 23 | 24 | class StubBackend(ResultBackend): 25 | """An in-memory result backend. For use in unit tests. 26 | 27 | Parameters: 28 | namespace(str): A string with which to prefix result keys. 29 | encoder(Encoder): The encoder to use when storing and retrieving 30 | result data. Defaults to :class:`.JSONEncoder`. 31 | """ 32 | 33 | results: Dict[str, Tuple[Optional[str], Optional[float]]] = {} 34 | 35 | def _get(self, message_key): 36 | data, expiration = self.results.get(message_key, (None, None)) 37 | if data is not None and time.monotonic() < expiration: 38 | return self.encoder.decode(data) 39 | return Missing 40 | 41 | def _store(self, message_key, result, ttl): 42 | result_data = self.encoder.encode(result) 43 | expiration = time.monotonic() + int(ttl / 1000) 44 | self.results[message_key] = (result_data, expiration) 45 | -------------------------------------------------------------------------------- /dramatiq/results/errors.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from ..errors import DramatiqError 19 | 20 | 21 | class ResultError(DramatiqError): 22 | """Base class for result errors. 23 | """ 24 | 25 | 26 | class ResultTimeout(ResultError): 27 | """Raised when waiting for a result times out. 28 | """ 29 | 30 | 31 | class ResultMissing(ResultError): 32 | """Raised when a result can't be found. 33 | """ 34 | 35 | 36 | class ResultFailure(ResultError): 37 | """Raised when getting a result from an actor that failed. 38 | """ 39 | 40 | def __init__(self, message="", orig_exc_type="", orig_exc_msg=""): 41 | super().__init__(message) 42 | 43 | self.orig_exc_type = orig_exc_type 44 | self.orig_exc_msg = orig_exc_msg 45 | -------------------------------------------------------------------------------- /dramatiq/results/middleware.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | from ..errors import ActorNotFound 18 | from ..logging import get_logger 19 | from ..middleware import Middleware 20 | 21 | #: The maximum amount of milliseconds results are allowed to exist in 22 | #: the backend. 23 | DEFAULT_RESULT_TTL = 600000 24 | 25 | 26 | class Results(Middleware): 27 | """Middleware that automatically stores actor results. 28 | 29 | Example: 30 | 31 | >>> from dramatiq.results import Results 32 | >>> from dramatiq.results.backends import RedisBackend 33 | >>> backend = RedisBackend() 34 | >>> broker.add_middleware(Results(backend=backend)) 35 | 36 | >>> @dramatiq.actor(store_results=True) 37 | ... def add(x, y): 38 | ... return x + y 39 | 40 | >>> message = add.send(1, 2) 41 | >>> message.get_result(backend=backend) 42 | 3 43 | 44 | >>> @dramatiq.actor(store_results=True) 45 | ... def fail(): 46 | ... raise Exception("failed") 47 | 48 | >>> message = fail.send() 49 | >>> message.get_result(backend=backend) 50 | Traceback (most recent call last): 51 | ... 52 | ResultFailure: actor raised Exception: failed 53 | 54 | Parameters: 55 | backend(ResultBackend): The result backend to use when storing 56 | results. 57 | store_results(bool): Whether or not actor results should be 58 | stored. Defaults to False and can be set on a per-actor 59 | basis. 60 | result_ttl(int): The maximum number of milliseconds results are 61 | allowed to exist in the backend. Defaults to 10 minutes and 62 | can be set on a per-actor basis. 63 | 64 | Warning: 65 | If you have retries turned on for an actor that stores results, 66 | then the result of a message may be delayed until its retries 67 | run out! 68 | """ 69 | 70 | def __init__(self, *, backend=None, store_results=False, result_ttl=None): 71 | self.logger = get_logger(__name__, type(self)) 72 | self.backend = backend 73 | self.store_results = store_results 74 | self.result_ttl = result_ttl or DEFAULT_RESULT_TTL 75 | 76 | @property 77 | def actor_options(self): 78 | return { 79 | "store_results", 80 | "result_ttl", 81 | } 82 | 83 | def _lookup_options(self, broker, message): 84 | try: 85 | actor = broker.get_actor(message.actor_name) 86 | store_results = message.options.get("store_results", actor.options.get("store_results", self.store_results)) 87 | result_ttl = message.options.get("result_ttl", actor.options.get("result_ttl", self.result_ttl)) 88 | return store_results, result_ttl 89 | except ActorNotFound: 90 | return False, 0 91 | 92 | def after_process_message(self, broker, message, *, result=None, exception=None): 93 | store_results, result_ttl = self._lookup_options(broker, message) 94 | if store_results and exception is None: 95 | self.backend.store_result(message, result, result_ttl) 96 | if not store_results \ 97 | and result is not None \ 98 | and message.options.get("pipe_target") is None: 99 | self.logger.warning( 100 | "Actor '%s' returned a value that is not None, but you " 101 | "haven't set its `store_results' option to `True' so " 102 | "the value has been discarded." % message.actor_name 103 | ) 104 | 105 | def after_skip_message(self, broker, message): 106 | """If the message was skipped but not failed, then store None. 107 | Let after_nack handle the case where the message was skipped and failed. 108 | """ 109 | store_results, result_ttl = self._lookup_options(broker, message) 110 | if store_results and not message.failed: 111 | self.backend.store_result(message, None, result_ttl) 112 | 113 | def after_nack(self, broker, message): 114 | store_results, result_ttl = self._lookup_options(broker, message) 115 | if store_results and message.failed: 116 | exception = message._exception or Exception("unknown") 117 | self.backend.store_exception(message, exception, result_ttl) 118 | -------------------------------------------------------------------------------- /dramatiq/results/result.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2020 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | from .errors import ResultFailure 19 | 20 | # We need to deal with backwards-compatibility here in case the user 21 | # is migrating between dramatiq versions so the canary is used to 22 | # detect exception results. 23 | _CANARY = "dramatiq.results.Result" 24 | 25 | 26 | def wrap_result(res): 27 | # This is a no-op for forwards-compatibility. If the user deploys 28 | # a new version of dramatiq and then is forced to roll back, then 29 | # this will minimize the fallout. 30 | return res 31 | 32 | 33 | def wrap_exception(e): 34 | return { 35 | "__t": _CANARY, 36 | "exn": { 37 | "type": type(e).__name__, 38 | "msg": str(e), 39 | } 40 | } 41 | 42 | 43 | def unwrap_result(res): 44 | if isinstance(res, dict) and res.get("__t") == _CANARY: 45 | message = "actor raised %s: %s" % (res["exn"]["type"], res["exn"]["msg"]) 46 | raise ResultFailure(message, res["exn"]["type"], res["exn"]["msg"]) 47 | return res 48 | -------------------------------------------------------------------------------- /dramatiq/threading.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import ctypes 19 | import inspect 20 | import platform 21 | 22 | from .logging import get_logger 23 | 24 | __all__ = [ 25 | "Interrupt", 26 | "current_platform", 27 | "is_gevent_active", 28 | "raise_thread_exception", 29 | "supported_platforms", 30 | ] 31 | 32 | 33 | logger = get_logger(__name__) 34 | 35 | current_platform = platform.python_implementation() 36 | python_version = platform.python_version_tuple() 37 | thread_id_ctype = ctypes.c_long if python_version < ("3", "7") else ctypes.c_ulong 38 | supported_platforms = {"CPython"} 39 | 40 | 41 | def is_gevent_active(): 42 | """Detect if gevent monkey patching is active.""" 43 | try: 44 | from gevent import monkey 45 | except ImportError: # pragma: no cover 46 | return False 47 | return bool(monkey.saved) 48 | 49 | 50 | class Interrupt(BaseException): 51 | """Base class for exceptions used to asynchronously interrupt a 52 | thread's execution. An actor may catch these exceptions in order 53 | to respond gracefully, such as performing any necessary cleanup. 54 | 55 | This is *not* a subclass of ``DramatiqError`` to avoid it being 56 | caught unintentionally. 57 | """ 58 | 59 | 60 | def raise_thread_exception(thread_id, exception): 61 | """Raise an exception in a thread. 62 | 63 | Currently, this is only available on CPython. 64 | 65 | Note: 66 | This works by setting an async exception in the thread. This means 67 | that the exception will only get called the next time that thread 68 | acquires the GIL. Concretely, this means that this middleware can't 69 | cancel system calls. 70 | """ 71 | if current_platform == "CPython": 72 | _raise_thread_exception_cpython(thread_id, exception) 73 | else: 74 | message = "Setting thread exceptions (%s) is not supported for your current platform (%r)." 75 | exctype = (exception if inspect.isclass(exception) else type(exception)).__name__ 76 | logger.critical(message, exctype, current_platform) 77 | 78 | 79 | def _raise_thread_exception_cpython(thread_id, exception): 80 | exctype = (exception if inspect.isclass(exception) else type(exception)).__name__ 81 | thread_id = thread_id_ctype(thread_id) 82 | exception = ctypes.py_object(exception) 83 | count = ctypes.pythonapi.PyThreadState_SetAsyncExc(thread_id, exception) 84 | if count == 0: 85 | logger.critical("Failed to set exception (%s) in thread %r.", exctype, thread_id.value) 86 | elif count > 1: # pragma: no cover 87 | logger.critical("Exception (%s) was set in multiple threads. Undoing...", exctype) 88 | ctypes.pythonapi.PyThreadState_SetAsyncExc(thread_id, ctypes.c_long(0)) 89 | -------------------------------------------------------------------------------- /dramatiq/watcher.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import signal 4 | 5 | import watchdog.events 6 | import watchdog.observers.polling 7 | 8 | try: 9 | import watchdog_gevent 10 | 11 | EVENTED_OBSERVER = watchdog_gevent.Observer 12 | except ImportError: 13 | EVENTED_OBSERVER = watchdog.observers.Observer 14 | 15 | 16 | def setup_file_watcher(path, use_polling=False, include_patterns=None, exclude_patterns=None): 17 | """Sets up a background thread that watches for source changes and 18 | automatically sends SIGHUP to the current process whenever a file 19 | changes. 20 | """ 21 | if use_polling: 22 | observer_class = watchdog.observers.polling.PollingObserver 23 | else: 24 | observer_class = EVENTED_OBSERVER 25 | 26 | if include_patterns is None: 27 | include_patterns = ["*.py"] 28 | 29 | file_event_handler = _SourceChangesHandler( 30 | patterns=include_patterns, ignore_patterns=exclude_patterns 31 | ) 32 | file_watcher = observer_class() 33 | file_watcher.schedule(file_event_handler, path, recursive=True) 34 | file_watcher.start() 35 | return file_watcher 36 | 37 | 38 | class _SourceChangesHandler(watchdog.events.PatternMatchingEventHandler): # pragma: no cover 39 | """Handles source code change events by sending a HUP signal to 40 | the current process. 41 | """ 42 | 43 | def on_any_event(self, event: watchdog.events.FileSystemEvent): 44 | # watchdog >= 2.3 emits an extra event when a file is opened that will cause unnecessary 45 | # restarts. This affects any usage of watchdog where a custom event handler is used, including 46 | # this _SourceChangesHandler. 47 | # 48 | # For more context, see: 49 | # https://github.com/gorakhargosh/watchdog/issues/949 50 | # https://github.com/pallets/werkzeug/pull/2604/files 51 | if event.event_type == "opened": 52 | return 53 | 54 | # watchdog >= 5.0.0 added events for when a file is closed without being written. 55 | # Also ignore these so they don't cause unnecessay restarts. 56 | # See https://github.com/gorakhargosh/watchdog/pull/1059 57 | if event.event_type == "closed_no_write": 58 | return 59 | 60 | logger = logging.getLogger("SourceChangesHandler") 61 | logger.info("Detected changes to %r.", event.src_path) 62 | os.kill(os.getpid(), signal.SIGHUP) 63 | -------------------------------------------------------------------------------- /examples/asyncio/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import asyncio 3 | import os 4 | import random 5 | import sys 6 | 7 | import dramatiq 8 | from dramatiq.middleware.asyncio import AsyncIO 9 | 10 | if os.getenv("REDIS") == "1": 11 | from dramatiq.brokers.redis import RedisBroker 12 | broker = RedisBroker() 13 | dramatiq.set_broker(broker) 14 | 15 | dramatiq.get_broker().add_middleware(AsyncIO()) 16 | 17 | 18 | @dramatiq.actor 19 | async def add(x, y): 20 | add.logger.info("About to compute the sum of %d and %d.", x, y) 21 | await asyncio.sleep(random.uniform(0.1, 0.3)) 22 | add.logger.info("The sum of %d and %d is %d.", x, y, x + y) 23 | 24 | 25 | def main(args): 26 | parser = argparse.ArgumentParser() 27 | parser.add_argument("count", type=int, help="the number of messages to enqueue") 28 | args = parser.parse_args() 29 | for _ in range(args.count): 30 | add.send(random.randint(0, 1000), random.randint(0, 1000)) 31 | 32 | 33 | if __name__ == "__main__": 34 | sys.exit(main(sys.argv)) 35 | -------------------------------------------------------------------------------- /examples/basic/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Basic Example 2 | 3 | This example demonstrates how easy it is to get started with Dramatiq. 4 | 5 | ## Running the Example 6 | 7 | 1. Install [RabbitMQ][rabbitmq] or [Redis][redis] 8 | 1. Install dramatiq: `pip install dramatiq[rabbitmq]` or `pip install dramatiq[redis]` 9 | 1. Run RabbitMQ: `rabbitmq-server` or Redis: `redis-server` 10 | 1. In a separate terminal window, run the workers: `dramatiq example`. 11 | Add `REDIS=1` before `dramatiq` to use the Redis broker. 12 | 1. In another terminal, run `python -m example 100` to enqueue 100 13 | `add` tasks. 14 | 15 | 16 | [rabbitmq]: https://www.rabbitmq.com 17 | [redis]: https://redis.io 18 | -------------------------------------------------------------------------------- /examples/basic/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/basic/__init__.py -------------------------------------------------------------------------------- /examples/basic/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import random 4 | import sys 5 | 6 | import dramatiq 7 | 8 | if os.getenv("REDIS") == "1": 9 | from dramatiq.brokers.redis import RedisBroker 10 | broker = RedisBroker() 11 | dramatiq.set_broker(broker) 12 | 13 | 14 | @dramatiq.actor 15 | def add(x, y): 16 | add.logger.info("The sum of %d and %d is %d.", x, y, x + y) 17 | 18 | 19 | def main(args): 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("count", type=int, help="the number of messages to enqueue") 22 | args = parser.parse_args() 23 | for _ in range(args.count): 24 | add.send(random.randint(0, 1000), random.randint(0, 1000)) 25 | 26 | 27 | if __name__ == "__main__": 28 | sys.exit(main(sys.argv)) 29 | -------------------------------------------------------------------------------- /examples/callable_broker/app.py: -------------------------------------------------------------------------------- 1 | import dramatiq 2 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 3 | 4 | 5 | def setup_broker(): 6 | print("Setting up broker...") 7 | broker = RabbitmqBroker() 8 | dramatiq.set_broker(broker) 9 | -------------------------------------------------------------------------------- /examples/composition/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Composition Example 2 | 3 | This example demonstrates how to use Dramatiq's high level composition 4 | abstractions. 5 | 6 | ## Running the Example 7 | 8 | 1. Install [RabbitMQ][rabbitmq] and [Redis][redis] 9 | 1. Install dramatiq and requests: `pip install dramatiq[rabbitmq,redis] requests` 10 | 1. Run RabbitMQ: `rabbitmq-server` 11 | 1. Run Redis: `redis-server` 12 | 1. In a separate terminal window, run the workers: `dramatiq example`. 13 | 1. In another terminal, run `python -m example`. 14 | 15 | 16 | [rabbitmq]: https://www.rabbitmq.com 17 | [redis]: https://redis.io 18 | -------------------------------------------------------------------------------- /examples/composition/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/composition/__init__.py -------------------------------------------------------------------------------- /examples/composition/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import sys 3 | 4 | import requests 5 | 6 | import dramatiq 7 | from dramatiq import group 8 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 9 | from dramatiq.encoder import PickleEncoder 10 | from dramatiq.results import Results 11 | from dramatiq.results.backends import RedisBackend 12 | 13 | encoder = PickleEncoder() 14 | backend = RedisBackend(encoder=encoder) 15 | broker = RabbitmqBroker(host="127.0.0.1") 16 | broker.add_middleware(Results(backend=backend)) 17 | dramatiq.set_broker(broker) 18 | dramatiq.set_encoder(encoder) 19 | 20 | 21 | @dramatiq.actor 22 | def request(uri): 23 | return requests.get(uri) 24 | 25 | 26 | @dramatiq.actor(store_results=True) 27 | def count_words(response): 28 | return len(response.text.split(" ")) 29 | 30 | 31 | def main(): 32 | parser = argparse.ArgumentParser() 33 | parser.add_argument("uri", nargs="+", help="A website URI.") 34 | 35 | arguments = parser.parse_args() 36 | jobs = group(request.message(uri) | count_words.message() for uri in arguments.uri).run() 37 | for uri, count in zip(arguments.uri, jobs.get_results(block=True)): 38 | print(f" * {uri} has {count} words") 39 | 40 | return 0 41 | 42 | 43 | if __name__ == "__main__": 44 | sys.exit(main()) 45 | -------------------------------------------------------------------------------- /examples/crawler/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Web Crawler Example 2 | 3 | This example implements a very simple distributed web crawler using 4 | Dramatiq. 5 | 6 | ## Running the Example 7 | 8 | 1. Install [RabbitMQ][rabbitmq] or [Redis][redis] 9 | 1. Install dramatiq: `pip install dramatiq[rabbitmq]` or `pip install dramatiq[redis]` 10 | 1. Install the example's dependencies: `pip install -r requirements.txt` 11 | 1. Run `redis-server` or `rabbitmq-server` in a terminal window. 12 | 1. In a separate terminal, run `dramatiq example` to run the workers. 13 | If you want to use the Redis broker, add `REDIS=1` before 14 | `dramatiq`. 15 | 1. Finally, run `python example.py https://example.com` to begin 16 | crawling http://example.com. 17 | 18 | 19 | [rabbitmq]: https://www.rabbitmq.com 20 | [redis]: https://redis.io 21 | -------------------------------------------------------------------------------- /examples/crawler/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/crawler/__init__.py -------------------------------------------------------------------------------- /examples/crawler/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import hashlib 3 | import logging 4 | import os 5 | import re 6 | import sys 7 | from contextlib import closing 8 | from threading import local 9 | 10 | import pylibmc 11 | import requests 12 | 13 | import dramatiq 14 | 15 | logger = logging.getLogger("example") 16 | memcache_client = pylibmc.Client(["localhost"], binary=True) 17 | memcache_pool = pylibmc.ThreadMappedPool(memcache_client) 18 | anchor_re = re.compile(rb'') 19 | state = local() 20 | 21 | if os.getenv("REDIS") == "1": 22 | from dramatiq.brokers.redis import RedisBroker 23 | broker = RedisBroker() 24 | dramatiq.set_broker(broker) 25 | 26 | 27 | def get_session(): 28 | session = getattr(state, "session", None) 29 | if session is None: 30 | session = state.session = requests.Session() 31 | return session 32 | 33 | 34 | @dramatiq.actor(max_retries=3, time_limit=10000) 35 | def crawl(url): 36 | url_hash = hashlib.md5(url.encode("utf-8")).hexdigest() 37 | with memcache_pool.reserve() as client: 38 | added = client.add(url_hash, b"", time=3600) 39 | if not added: 40 | logger.warning("URL %r has already been visited. Skipping...", url) 41 | return 42 | 43 | logger.info("Crawling %r...", url) 44 | matches = 0 45 | session = get_session() 46 | with closing(session.get(url, timeout=(3.05, 5), stream=True)) as response: 47 | if not response.headers.get("content-type", "").startswith("text/html"): 48 | logger.warning("Skipping URL %r since it's not HTML.", url) 49 | return 50 | 51 | for match in anchor_re.finditer(response.content): 52 | anchor = match.group(1).decode("utf-8") 53 | if anchor.startswith("http://") or anchor.startswith("https://"): 54 | crawl.send(anchor) 55 | matches += 1 56 | 57 | logger.info("Done crawling %r. Found %d anchors.", url, matches) 58 | 59 | 60 | def main(args): 61 | parser = argparse.ArgumentParser() 62 | parser.add_argument("url", type=str, help="a URL to crawl") 63 | args = parser.parse_args() 64 | crawl.send(args.url) 65 | return 0 66 | 67 | 68 | if __name__ == "__main__": 69 | sys.exit(main(sys.argv)) 70 | -------------------------------------------------------------------------------- /examples/crawler/requirements.txt: -------------------------------------------------------------------------------- 1 | pylibmc 2 | requests[security] 3 | -------------------------------------------------------------------------------- /examples/issue_351/app.py: -------------------------------------------------------------------------------- 1 | import dramatiq 2 | 3 | 4 | @dramatiq.actor 5 | def foo(): 6 | a = tuple(range(5000000)) # noqa 7 | raise Exception("bar") 8 | 9 | 10 | if __name__ == "__main__": 11 | for _ in range(10): 12 | foo.send() 13 | -------------------------------------------------------------------------------- /examples/long_running/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Long Running Example 2 | 3 | This example demonstrates extremely long-running tasks in Dramatiq. 4 | 5 | ## Running the Example 6 | 7 | 1. Install [RabbitMQ][rabbitmq] or [Redis][redis] 8 | 1. Install dramatiq: `pip install dramatiq[rabbitmq]` or `pip install dramatiq[redis]` 9 | 1. Run RabbitMQ: `rabbitmq-server` or Redis: `redis-server` 10 | 1. In a separate terminal window, run the workers: `dramatiq example`. 11 | Add `REDIS=1` before `dramatiq` to use the Redis broker. 12 | 1. In another terminal, run `python -m example` to enqueue a task. 13 | 14 | 15 | [rabbitmq]: https://www.rabbitmq.com 16 | [redis]: https://redis.io 17 | -------------------------------------------------------------------------------- /examples/long_running/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/long_running/__init__.py -------------------------------------------------------------------------------- /examples/long_running/example.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import sys 4 | import time 5 | 6 | import dramatiq 7 | 8 | if os.getenv("REDIS") == "1": 9 | from dramatiq.brokers.redis import RedisBroker 10 | broker = RedisBroker() 11 | dramatiq.set_broker(broker) 12 | 13 | 14 | def fib(n): 15 | x, y = 1, 1 16 | while n > 2: 17 | x, y = x + y, x 18 | n -= 1 19 | return x 20 | 21 | 22 | @dramatiq.actor(time_limit=86_400_000, max_retries=0) 23 | def long_running(duration): 24 | deadline = time.monotonic() + duration 25 | while time.monotonic() < deadline: 26 | long_running.logger.info("%d seconds remaining.", deadline - time.monotonic()) 27 | 28 | n = random.randint(1_000, 1_000_000) 29 | long_running.logger.debug("Computing fib(%d).", n) 30 | 31 | fib(n) 32 | long_running.logger.debug("Computed fib(%d).", n) 33 | 34 | sleep = random.randint(1, 30) 35 | long_running.logger.debug("Sleeping for %d seconds...", sleep) 36 | time.sleep(sleep) 37 | 38 | 39 | def main(args): 40 | for _ in range(1_000): 41 | long_running.send(random.randint(3_600, 14_400)) 42 | time.sleep(random.randint(60, 3600)) 43 | 44 | 45 | if __name__ == "__main__": 46 | sys.exit(main(sys.argv)) 47 | -------------------------------------------------------------------------------- /examples/max_tasks_per_child/app.py: -------------------------------------------------------------------------------- 1 | import os 2 | import signal 3 | from threading import Lock 4 | 5 | import dramatiq 6 | from dramatiq import Middleware, get_logger 7 | 8 | 9 | class MaxTasksPerChild(Middleware): 10 | def __init__(self, max_tasks=100): 11 | self.counter_mu = Lock() 12 | self.counter = max_tasks 13 | self.signaled = False 14 | self.logger = get_logger("max_tasks_per_child.app", MaxTasksPerChild) 15 | 16 | def after_process_message(self, broker, message, *, result=None, exception=None): 17 | with self.counter_mu: 18 | self.counter -= 1 19 | self.logger.debug("Remaining tasks: %d.", self.counter) 20 | if self.counter <= 0 and not self.signaled: 21 | self.logger.warning("Counter reached zero. Signaling current process.") 22 | os.kill(os.getppid(), getattr(signal, "SIGHUP", signal.SIGTERM)) 23 | self.signaled = True 24 | 25 | 26 | broker = dramatiq.get_broker() 27 | broker.add_middleware(MaxTasksPerChild()) 28 | 29 | 30 | @dramatiq.actor 31 | def example(): 32 | pass 33 | 34 | 35 | if __name__ == "__main__": 36 | for _ in range(105): 37 | example.send() 38 | -------------------------------------------------------------------------------- /examples/persistent/.gitignore: -------------------------------------------------------------------------------- 1 | states -------------------------------------------------------------------------------- /examples/persistent/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Persistent Example 2 | 3 | This example demonstrates long-running tasks in Dramatiq that are durable to 4 | worker shutdowns by using simple state persistence. 5 | 6 | ## Running the Example 7 | 8 | 1. Install [RabbitMQ][rabbitmq] or [Redis][redis] 9 | 1. Install dramatiq: `pip install dramatiq[rabbitmq]` or `pip install dramatiq[redis]` 10 | 1. Run RabbitMQ: `rabbitmq-server` or Redis: `redis-server` 11 | 1. In a separate terminal window, run the workers: `dramatiq example`. 12 | Add `REDIS=1` before `dramatiq` to use the Redis broker. 13 | 1. In another terminal, run `python -m example ` to enqueue a task. 14 | 1. To test the persistence, terminate the worker process with `crtl-c`, then 15 | restart with the same `n` value. 16 | 17 | 18 | [rabbitmq]: https://www.rabbitmq.com 19 | [redis]: https://redis.io 20 | -------------------------------------------------------------------------------- /examples/persistent/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/persistent/__init__.py -------------------------------------------------------------------------------- /examples/persistent/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import sys 5 | import time 6 | 7 | import dramatiq 8 | from dramatiq.middleware import Shutdown 9 | 10 | if os.getenv("REDIS") == "1": 11 | from dramatiq.brokers.redis import RedisBroker 12 | broker = RedisBroker() 13 | dramatiq.set_broker(broker) 14 | 15 | 16 | def path_to(*xs): 17 | return os.path.abspath(os.path.join(os.path.dirname(__file__), *(str(x) for x in xs))) 18 | 19 | 20 | def load_state(n): 21 | try: 22 | with open(path_to("states", n), "r") as f: 23 | data = json.load(f) 24 | except Exception: 25 | fib.logger.info("Could not read state file, using defaults.") 26 | return 1, 1, 0 27 | 28 | i, x2, x1 = data["i"], data["x2"], data["x1"] 29 | fib.logger.info("Resuming fib(%d) from iteration %d.", n, i) 30 | return i, x2, x1 31 | 32 | 33 | def dump_state(n, state): 34 | os.makedirs("states", exist_ok=True) 35 | with open(path_to("states", n), "w") as f: 36 | json.dump(state, f) 37 | 38 | fib.logger.info("Dumped fib(%d) state for iteration %d.", n, state["i"]) 39 | 40 | 41 | def remove_state(n): 42 | try: 43 | os.remove(path_to("states", n)) 44 | except OSError: 45 | pass 46 | 47 | fib.logger.info("Deleted state for fib(%d).", n) 48 | 49 | 50 | @dramatiq.actor(time_limit=float("inf"), notify_shutdown=True, max_retries=0) 51 | def fib(n): 52 | i, x2, x1 = load_state(n) 53 | 54 | try: 55 | for i in range(i, n + 1): 56 | state = { 57 | "i": i, 58 | "x2": x2, 59 | "x1": x1, 60 | } 61 | 62 | x2, x1 = x1, x2 + x1 63 | fib.logger.info("fib(%d): %d", i, x1) 64 | time.sleep(.1) 65 | 66 | remove_state(n) 67 | fib.logger.info("Done!") 68 | except Shutdown: 69 | dump_state(n, state) 70 | 71 | 72 | def main(args): 73 | parser = argparse.ArgumentParser(description="Calculates fib(n)") 74 | parser.add_argument("n", type=int, help="must be a positive number") 75 | args = parser.parse_args() 76 | fib.send(args.n) 77 | return 0 78 | 79 | 80 | if __name__ == "__main__": 81 | sys.exit(main(sys.argv)) 82 | -------------------------------------------------------------------------------- /examples/results/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Results Example 2 | 3 | This example demonstrates how to use Dramatiq actors to store and 4 | retrieve task results. 5 | 6 | ## Running the Example 7 | 8 | 1. Install [RabbitMQ][rabbitmq] and [Redis][redis] 9 | 1. Install dramatiq: `pip install dramatiq[rabbitmq,redis]` 10 | 1. Run RabbitMQ: `rabbitmq-server` 11 | 1. Run Redis: `redis-server` 12 | 1. In a separate terminal window, run the workers: `dramatiq example`. 13 | 1. In another terminal, run `python -m example`. 14 | 15 | 16 | [rabbitmq]: https://www.rabbitmq.com 17 | [redis]: https://redis.io 18 | -------------------------------------------------------------------------------- /examples/results/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/results/__init__.py -------------------------------------------------------------------------------- /examples/results/example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import random 3 | import sys 4 | import time 5 | 6 | import dramatiq 7 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 8 | from dramatiq.encoder import PickleEncoder 9 | from dramatiq.results import Results 10 | from dramatiq.results.backends import RedisBackend 11 | 12 | result_backend = RedisBackend(encoder=PickleEncoder()) 13 | broker = RabbitmqBroker() 14 | broker.add_middleware(Results(backend=result_backend)) 15 | dramatiq.set_broker(broker) 16 | 17 | 18 | @dramatiq.actor(store_results=True) 19 | def sleep_then_add(t, x, y): 20 | time.sleep(t) 21 | return x + y 22 | 23 | 24 | def main(args): 25 | parser = argparse.ArgumentParser() 26 | parser.add_argument("count", type=int, help="the number of messages to enqueue") 27 | args = parser.parse_args() 28 | 29 | messages = [] 30 | for _ in range(args.count): 31 | messages.append(sleep_then_add.send( 32 | random.randint(1, 5), 33 | random.randint(0, 1000), 34 | random.randint(0, 1000) 35 | )) 36 | 37 | for message in messages: 38 | print(message.get_result(block=True)) 39 | 40 | 41 | if __name__ == "__main__": 42 | sys.exit(main(sys.argv)) 43 | -------------------------------------------------------------------------------- /examples/scheduling/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Scheduling Example 2 | 3 | This example demonstrates how to use Dramatiq in conjunction with 4 | [APScheduler] in order to schedule messages. 5 | 6 | ## Running the Example 7 | 8 | 1. Install [RabbitMQ][rabbitmq] and [Redis][redis] 9 | 1. Install dramatiq: `pip install dramatiq[rabbitmq,redis]` 10 | 1. Install apscheduler: `pip install apscheduler` 11 | 1. Run RabbitMQ: `rabbitmq-server` 12 | 1. Run Redis: `redis-server` 13 | 1. In a separate terminal window, run the workers: `dramatiq example`. 14 | 1. In another terminal, run `python -m example`. 15 | 16 | 17 | [APScheduler]: https://apscheduler.readthedocs.io/en/latest/ 18 | [rabbitmq]: https://www.rabbitmq.com 19 | [redis]: https://redis.io 20 | -------------------------------------------------------------------------------- /examples/scheduling/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/scheduling/__init__.py -------------------------------------------------------------------------------- /examples/scheduling/example.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from datetime import datetime 3 | 4 | from apscheduler.schedulers.blocking import BlockingScheduler 5 | from apscheduler.triggers.cron import CronTrigger 6 | 7 | import dramatiq 8 | 9 | 10 | @dramatiq.actor 11 | def print_current_date(): 12 | print(datetime.now()) 13 | 14 | 15 | def main(args): 16 | scheduler = BlockingScheduler() 17 | scheduler.add_job( 18 | print_current_date.send, 19 | CronTrigger.from_crontab("* * * * *"), 20 | ) 21 | try: 22 | scheduler.start() 23 | except KeyboardInterrupt: 24 | scheduler.shutdown() 25 | 26 | return 0 27 | 28 | 29 | if __name__ == "__main__": 30 | sys.exit(main(sys.argv)) 31 | -------------------------------------------------------------------------------- /examples/time_limit/README.md: -------------------------------------------------------------------------------- 1 | # Dramatiq Time Limit Example 2 | 3 | This example demonstrates the time limit feature of Dramatiq. 4 | 5 | ## Running the Example 6 | 7 | 1. Install [RabbitMQ][rabbitmq] or [Redis][redis] 8 | 1. Install dramatiq: `pip install dramatiq[rabbitmq]` or `pip install dramatiq[redis]` 9 | 1. Run RabbitMQ: `rabbitmq-server` or Redis: `redis-server` 10 | 1. In a separate terminal window, run the workers: `dramatiq example`. 11 | Add `REDIS=1` before `dramatiq` to use the Redis broker. 12 | 1. In another terminal, run `python -m example` to enqueue a task. 13 | 14 | 15 | [rabbitmq]: https://www.rabbitmq.com 16 | [redis]: https://redis.io 17 | -------------------------------------------------------------------------------- /examples/time_limit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/examples/time_limit/__init__.py -------------------------------------------------------------------------------- /examples/time_limit/example.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import time 4 | 5 | import dramatiq 6 | 7 | if os.getenv("REDIS") == "1": 8 | from dramatiq.brokers.redis import RedisBroker 9 | broker = RedisBroker() 10 | dramatiq.set_broker(broker) 11 | 12 | 13 | @dramatiq.actor(time_limit=5000, max_retries=3) 14 | def long_running(): 15 | logger = long_running.logger 16 | 17 | while True: 18 | logger.info("Sleeping...") 19 | time.sleep(1) 20 | 21 | 22 | @dramatiq.actor(time_limit=5000, max_retries=0) 23 | def long_running_with_catch(): 24 | logger = long_running_with_catch.logger 25 | 26 | try: 27 | while True: 28 | logger.info("Sleeping...") 29 | time.sleep(1) 30 | except dramatiq.middleware.time_limit.TimeLimitExceeded: 31 | logger.warning("Time limit exceeded. Aborting...") 32 | 33 | 34 | def main(args): 35 | long_running.send() 36 | long_running_with_catch.send() 37 | 38 | 39 | if __name__ == "__main__": 40 | sys.exit(main(sys.argv)) 41 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | pretty = True 3 | show_column_numbers = True 4 | show_error_codes = True 5 | show_error_context = True 6 | 7 | [mypy-gevent] 8 | ignore_missing_imports = True 9 | 10 | [mypy-greenlet] 11 | ignore_missing_imports = True 12 | 13 | [mypy-pika.*] 14 | ignore_missing_imports = True 15 | 16 | [mypy-prometheus_client] 17 | ignore_missing_imports = True 18 | 19 | [mypy-pylibmc] 20 | ignore_missing_imports = True 21 | 22 | [mypy-redis] 23 | ignore_missing_imports = True 24 | 25 | [mypy-watchdog.*] 26 | ignore_missing_imports = True 27 | 28 | [mypy-watchdog_gevent] 29 | ignore_missing_imports = True 30 | -------------------------------------------------------------------------------- /pytest-gevent.py: -------------------------------------------------------------------------------- 1 | #!python 2 | try: 3 | from gevent import monkey; monkey.patch_all() # noqa 4 | except ImportError: 5 | import sys 6 | 7 | sys.stderr.write("error: gevent is missing. Run `pip install gevent`.") 8 | sys.exit(1) 9 | 10 | import sys 11 | 12 | import pytest 13 | 14 | if __name__ == "__main__": 15 | sys.exit(pytest.main(args=sys.argv[1:])) 16 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [tool:pytest] 2 | testpaths = tests 3 | addopts = --cov dramatiq --cov-report html --benchmark-autosave --benchmark-compare 4 | 5 | [pep8] 6 | max-line-length = 120 7 | 8 | [flake8] 9 | max-complexity = 20 10 | max-line-length = 120 11 | select = C,E,F,W,B,B9,Q0 12 | ignore = E127,E402,E501,F403,F811,W504,B010,B020,B036,B905,B908 13 | 14 | inline-quotes = double 15 | multiline-quotes = double 16 | 17 | [isort] 18 | line_length = 120 19 | known_first_party = dramatiq 20 | multi_line_output = 5 21 | order_by_type = true 22 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # This file is a part of Dramatiq. 2 | # 3 | # Copyright (C) 2017,2018,2019 CLEARTYPE SRL 4 | # 5 | # Dramatiq is free software; you can redistribute it and/or modify it 6 | # under the terms of the GNU Lesser General Public License as published by 7 | # the Free Software Foundation, either version 3 of the License, or (at 8 | # your option) any later version. 9 | # 10 | # Dramatiq is distributed in the hope that it will be useful, but WITHOUT 11 | # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 12 | # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public 13 | # License for more details. 14 | # 15 | # You should have received a copy of the GNU Lesser General Public License 16 | # along with this program. If not, see . 17 | 18 | import os 19 | 20 | from setuptools import setup 21 | 22 | 23 | def rel(*xs): 24 | return os.path.join(os.path.abspath(os.path.dirname(__file__)), *xs) 25 | 26 | 27 | with open(rel("README.md")) as f: 28 | long_description = f.read() 29 | 30 | 31 | with open(rel("dramatiq", "__init__.py"), "r") as f: 32 | version_marker = "__version__ = " 33 | for line in f: 34 | if line.startswith(version_marker): 35 | _, version = line.split(version_marker) 36 | version = version.strip().strip('"') 37 | break 38 | else: 39 | raise RuntimeError("Version marker not found.") 40 | 41 | 42 | dependencies = [ 43 | "prometheus-client>=0.2", 44 | ] 45 | 46 | extra_dependencies = { 47 | "gevent": [ 48 | "gevent>=1.1", 49 | ], 50 | 51 | "memcached": [ 52 | "pylibmc>=1.5,<2.0", 53 | ], 54 | 55 | "rabbitmq": [ 56 | "pika>=1.0,<2.0", 57 | ], 58 | 59 | "redis": [ 60 | "redis>=2.0,<7.0", 61 | ], 62 | 63 | "watch": [ 64 | "watchdog>=4.0", 65 | "watchdog_gevent>=0.2", 66 | ], 67 | } 68 | 69 | extra_dependencies["all"] = list(set(sum(extra_dependencies.values(), []))) 70 | extra_dependencies["dev"] = extra_dependencies["all"] + [ 71 | # Docs 72 | "alabaster", 73 | "sphinx", 74 | "sphinxcontrib-napoleon", 75 | 76 | # Linting 77 | "flake8", 78 | "flake8-bugbear", 79 | "flake8-quotes", 80 | "isort", 81 | "mypy", 82 | 83 | # Misc 84 | "bumpversion", 85 | "hiredis", 86 | "twine", 87 | "wheel", 88 | 89 | # Testing 90 | "pytest", 91 | "pytest-benchmark[histogram]", 92 | "pytest-cov", 93 | "tox", 94 | ] 95 | 96 | setup( 97 | name="dramatiq", 98 | version=version, 99 | author="Bogdan Popa", 100 | author_email="bogdan@cleartype.io", 101 | project_urls={ 102 | "Documentation": "https://dramatiq.io", 103 | "Source": "https://github.com/Bogdanp/dramatiq", 104 | }, 105 | description="Background Processing for Python 3.", 106 | long_description=long_description, 107 | long_description_content_type="text/markdown", 108 | packages=[ 109 | "dramatiq", 110 | "dramatiq.brokers", 111 | "dramatiq.middleware", 112 | "dramatiq.rate_limits", 113 | "dramatiq.rate_limits.backends", 114 | "dramatiq.results", 115 | "dramatiq.results.backends", 116 | ], 117 | include_package_data=True, 118 | install_requires=dependencies, 119 | python_requires=">=3.9", 120 | extras_require=extra_dependencies, 121 | entry_points={"console_scripts": ["dramatiq = dramatiq.__main__:main"]}, 122 | scripts=["bin/dramatiq-gevent"], 123 | classifiers=[ 124 | "Programming Language :: Python :: 3.9", 125 | "Programming Language :: Python :: 3.10", 126 | "Programming Language :: Python :: 3.11", 127 | "Programming Language :: Python :: 3.12", 128 | "Programming Language :: Python :: 3.13", 129 | "Programming Language :: Python :: 3 :: Only", 130 | "Topic :: System :: Distributed Computing", 131 | "License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)", 132 | ], 133 | ) 134 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/tests/__init__.py -------------------------------------------------------------------------------- /tests/benchmarks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/tests/benchmarks/__init__.py -------------------------------------------------------------------------------- /tests/benchmarks/test_rabbitmq_cli.py: -------------------------------------------------------------------------------- 1 | import random 2 | import time 3 | 4 | import pytest 5 | 6 | import dramatiq 7 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 8 | 9 | from ..common import RABBITMQ_CREDENTIALS 10 | 11 | broker = RabbitmqBroker( 12 | host="127.0.0.1", 13 | credentials=RABBITMQ_CREDENTIALS, 14 | ) 15 | 16 | 17 | @dramatiq.actor(queue_name="benchmark-throughput", broker=broker) 18 | def throughput(): 19 | pass 20 | 21 | 22 | @dramatiq.actor(queue_name="benchmark-fib", broker=broker) 23 | def fib(n): 24 | x, y = 1, 1 25 | while n > 2: 26 | x, y = x + y, x 27 | n -= 1 28 | return x 29 | 30 | 31 | @dramatiq.actor(queue_name="benchmark-latency", broker=broker) 32 | def latency(): 33 | p = random.randint(1, 100) 34 | if p == 1: 35 | durations = [3, 3, 3, 1] 36 | elif p <= 10: 37 | durations = [2, 3] 38 | elif p <= 40: 39 | durations = [1, 2] 40 | else: 41 | durations = [1] 42 | 43 | for duration in durations: 44 | time.sleep(duration) 45 | 46 | 47 | @pytest.mark.benchmark(group="rabbitmq-100k-throughput") 48 | def test_rabbitmq_process_100k_messages_with_cli(benchmark, info_logging, start_cli): 49 | # Given that I've loaded 100k messages into RabbitMQ 50 | def setup(): 51 | # The connection error tests have the side-effect of 52 | # disconnecting this broker so we need to force it to 53 | # reconnect. 54 | del broker.channel 55 | del broker.connection 56 | 57 | for _ in range(100000): 58 | throughput.send() 59 | 60 | start_cli("tests.benchmarks.test_rabbitmq_cli:broker") 61 | 62 | # I expect processing those messages with the CLI to be consistently fast 63 | benchmark.pedantic(broker.join, args=(throughput.queue_name,), setup=setup) 64 | 65 | 66 | @pytest.mark.benchmark(group="rabbitmq-10k-fib") 67 | def test_rabbitmq_process_10k_fib_with_cli(benchmark, info_logging, start_cli): 68 | # Given that I've loaded 10k messages into RabbitMQ 69 | def setup(): 70 | # The connection error tests have the side-effect of 71 | # disconnecting this broker so we need to force it to 72 | # reconnect. 73 | del broker.channel 74 | del broker.connection 75 | 76 | for _ in range(10000): 77 | fib.send(random.choice([1, 512, 1024, 2048, 4096, 8192])) 78 | 79 | start_cli("tests.benchmarks.test_rabbitmq_cli:broker") 80 | 81 | # I expect processing those messages with the CLI to be consistently fast 82 | benchmark.pedantic(broker.join, args=(fib.queue_name,), setup=setup) 83 | 84 | 85 | @pytest.mark.benchmark(group="rabbitmq-1k-latency") 86 | def test_rabbitmq_process_1k_latency_with_cli(benchmark, info_logging, start_cli): 87 | # Given that I've loaded 1k messages into RabbitMQ 88 | def setup(): 89 | # The connection error tests have the side-effect of 90 | # disconnecting this broker so we need to force it to 91 | # reconnect. 92 | del broker.channel 93 | del broker.connection 94 | 95 | for _ in range(1000): 96 | latency.send() 97 | 98 | start_cli("tests.benchmarks.test_rabbitmq_cli:broker") 99 | 100 | # I expect processing those messages with the CLI to be consistently fast 101 | benchmark.pedantic(broker.join, args=(latency.queue_name,), setup=setup) 102 | -------------------------------------------------------------------------------- /tests/benchmarks/test_redis_cli.py: -------------------------------------------------------------------------------- 1 | import random 2 | import time 3 | 4 | import pytest 5 | 6 | import dramatiq 7 | from dramatiq.brokers.redis import RedisBroker 8 | 9 | broker = RedisBroker() 10 | 11 | 12 | @dramatiq.actor(queue_name="benchmark-throughput", broker=broker) 13 | def throughput(): 14 | pass 15 | 16 | 17 | @dramatiq.actor(queue_name="benchmark-fib", broker=broker) 18 | def fib(n): 19 | x, y = 1, 1 20 | while n > 2: 21 | x, y = x + y, x 22 | n -= 1 23 | return x 24 | 25 | 26 | @dramatiq.actor(queue_name="benchmark-latency", broker=broker) 27 | def latency(): 28 | p = random.randint(1, 100) 29 | if p == 1: 30 | durations = [3, 3, 3, 1] 31 | elif p <= 10: 32 | durations = [2, 3] 33 | elif p <= 40: 34 | durations = [1, 2] 35 | else: 36 | durations = [1] 37 | 38 | for duration in durations: 39 | time.sleep(duration) 40 | 41 | 42 | @pytest.mark.benchmark(group="redis-100k-throughput") 43 | def test_redis_process_100k_messages_with_cli(benchmark, info_logging, start_cli): 44 | # Given that I've loaded 100k messages into Redis 45 | def setup(): 46 | for _ in range(100000): 47 | throughput.send() 48 | 49 | start_cli("tests.benchmarks.test_redis_cli:broker") 50 | 51 | # I expect processing those messages with the CLI to be consistently fast 52 | benchmark.pedantic(broker.join, args=(throughput.queue_name,), setup=setup) 53 | 54 | 55 | @pytest.mark.benchmark(group="redis-10k-fib") 56 | def test_redis_process_10k_fib_with_cli(benchmark, info_logging, start_cli): 57 | # Given that I've loaded 1k messages into Redis 58 | def setup(): 59 | for _ in range(10000): 60 | fib.send(random.choice([1, 512, 1024, 2048, 4096, 8192])) 61 | 62 | start_cli("tests.benchmarks.test_redis_cli:broker") 63 | 64 | # I expect processing those messages with the CLI to be consistently fast 65 | benchmark.pedantic(broker.join, args=(fib.queue_name,), setup=setup) 66 | 67 | 68 | @pytest.mark.benchmark(group="redis-1k-latency") 69 | def test_redis_process_1k_latency_with_cli(benchmark, info_logging, start_cli): 70 | # Given that I've loaded 1k messages into Redis 71 | def setup(): 72 | for _ in range(1000): 73 | latency.send() 74 | 75 | start_cli("tests.benchmarks.test_redis_cli:broker") 76 | 77 | # I expect processing those messages with the CLI to be consistently fast 78 | benchmark.pedantic(broker.join, args=(latency.queue_name,), setup=setup) 79 | -------------------------------------------------------------------------------- /tests/common.py: -------------------------------------------------------------------------------- 1 | import os 2 | import platform 3 | from contextlib import contextmanager 4 | 5 | import pika 6 | import pika.exceptions 7 | import pytest 8 | 9 | from dramatiq import Worker 10 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 11 | from dramatiq.threading import is_gevent_active 12 | 13 | 14 | @contextmanager 15 | def worker(*args, **kwargs): 16 | try: 17 | worker = Worker(*args, **kwargs) 18 | worker.start() 19 | yield worker 20 | finally: 21 | worker.stop() 22 | 23 | 24 | skip_in_ci = pytest.mark.skipif( 25 | os.getenv("APPVEYOR") is not None or 26 | os.getenv("GITHUB_ACTION") is not None, 27 | reason="test skipped in CI" 28 | ) 29 | 30 | skip_on_windows = pytest.mark.skipif(platform.system() == "Windows", reason="test skipped on Windows") 31 | skip_on_pypy = pytest.mark.skipif(platform.python_implementation() == "PyPy", reason="Time limits are not supported under PyPy.") 32 | skip_with_gevent = pytest.mark.skipif(is_gevent_active(), reason="Behaviour with gevent is different.") 33 | skip_without_gevent = pytest.mark.skipif(not is_gevent_active(), reason="Behaviour without gevent is different.") 34 | 35 | RABBITMQ_USERNAME = os.getenv("RABBITMQ_USERNAME", "guest") 36 | RABBITMQ_PASSWORD = os.getenv("RABBITMQ_PASSWORD", "guest") 37 | RABBITMQ_CREDENTIALS = pika.credentials.PlainCredentials(RABBITMQ_USERNAME, RABBITMQ_PASSWORD) 38 | 39 | 40 | def rabbit_mq_is_unreachable(): 41 | broker = RabbitmqBroker( 42 | host="127.0.0.1", 43 | max_priority=10, 44 | credentials=RABBITMQ_CREDENTIALS, 45 | ) 46 | try: 47 | broker.connection 48 | except pika.exceptions.AMQPConnectionError: 49 | # RabbitMQ is unreachable 50 | return True 51 | return False 52 | 53 | 54 | skip_unless_rabbit_mq = pytest.mark.skipif(rabbit_mq_is_unreachable(), reason="RabbitMQ is unreachable") 55 | -------------------------------------------------------------------------------- /tests/middleware/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Bogdanp/dramatiq/0ebf911d3fc24ff1284871692589ec4bbc307031/tests/middleware/__init__.py -------------------------------------------------------------------------------- /tests/middleware/test_prometheus.py: -------------------------------------------------------------------------------- 1 | import time 2 | import urllib.request as request 3 | from threading import Thread 4 | 5 | from dramatiq.middleware.prometheus import _run_exposition_server 6 | 7 | 8 | def test_prometheus_middleware_exposes_metrics(): 9 | # Given an instance of the exposition server 10 | thread = Thread(target=_run_exposition_server, daemon=True) 11 | thread.start() 12 | 13 | # When I give it time to boot up 14 | time.sleep(1) 15 | 16 | # And I request metrics via HTTP 17 | with request.urlopen("http://127.0.0.1:9191") as resp: 18 | # Then the response should be successful 19 | assert resp.getcode() == 200 20 | -------------------------------------------------------------------------------- /tests/middleware/test_threading.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import time 3 | from threading import Thread 4 | 5 | import pytest 6 | 7 | from dramatiq import threading 8 | 9 | not_supported = threading.current_platform not in threading.supported_platforms 10 | 11 | 12 | @pytest.mark.skipif(not_supported, reason="Threading not supported on this platform.") 13 | @pytest.mark.skipif(threading.is_gevent_active(), reason="Thread exceptions not supported with gevent.") 14 | def test_raise_thread_exception(): 15 | # Given that I have a database 16 | caught = [] 17 | 18 | # And a function that waits for an interrupt 19 | def work(): 20 | try: 21 | for _ in range(10): 22 | time.sleep(.1) 23 | except threading.Interrupt: 24 | caught.append(1) 25 | 26 | # When I start the thread 27 | t = Thread(target=work) 28 | t.start() 29 | time.sleep(.1) 30 | 31 | # And raise the interrupt and join on the thread 32 | threading.raise_thread_exception(t.ident, threading.Interrupt) 33 | t.join() 34 | 35 | # I expect the interrupt to have been caught 36 | assert sum(caught) == 1 37 | 38 | 39 | @pytest.mark.skipif(not_supported, reason="Threading not supported on this platform.") 40 | @pytest.mark.skipif(threading.is_gevent_active(), reason="Thread exceptions not supported with gevent.") 41 | def test_raise_thread_exception_on_nonexistent_thread(caplog): 42 | # When an interrupt is raised on a nonexistent thread 43 | thread_id = 2 ** 31 - 1 44 | threading.raise_thread_exception(thread_id, threading.Interrupt) 45 | 46 | # I expect a 'failed to set exception' critical message to be logged 47 | expected_message = "Failed to set exception (Interrupt) in thread %d." % thread_id 48 | assert caplog.record_tuples == [ 49 | ("dramatiq.threading", logging.CRITICAL, expected_message), 50 | ] 51 | 52 | 53 | def test_raise_thread_exception_unsupported_platform(caplog, monkeypatch): 54 | # monkeypatch fake platform to test logging. 55 | monkeypatch.setattr(threading, "current_platform", "not supported") 56 | 57 | # When raising a thread exception on an unsupported platform 58 | threading.raise_thread_exception(1, threading.Interrupt) 59 | 60 | # I expect a 'platform not supported' critical message to be logged 61 | assert caplog.record_tuples == [ 62 | ("dramatiq.threading", logging.CRITICAL, ( 63 | "Setting thread exceptions (Interrupt) is not supported " 64 | "for your current platform ('not supported')." 65 | )), 66 | ] 67 | -------------------------------------------------------------------------------- /tests/test_barrier.py: -------------------------------------------------------------------------------- 1 | import time 2 | from concurrent.futures import ThreadPoolExecutor 3 | 4 | import pytest 5 | 6 | from dramatiq.rate_limits import Barrier 7 | 8 | 9 | def test_barrier(rate_limiter_backend): 10 | # Given that I have a barrier of two parties 11 | barrier = Barrier(rate_limiter_backend, "sequential-barrier", ttl=30000) 12 | assert barrier.create(parties=2) 13 | 14 | # When I try to recreate it 15 | # Then that should return False 16 | assert not barrier.create(parties=10) 17 | 18 | # The first party to wait should get back False 19 | assert not barrier.wait(block=False) 20 | 21 | # The second party to wait should get back True 22 | assert barrier.wait(block=False) 23 | 24 | 25 | def test_barriers_can_block(rate_limiter_backend): 26 | # Given that I have a barrier of two parties 27 | barrier = Barrier(rate_limiter_backend, "sequential-barrier", ttl=30000) 28 | assert barrier.create(parties=2) 29 | 30 | # And I have a worker function that waits on the barrier and writes its timestamp 31 | times = [] 32 | 33 | def worker(): 34 | time.sleep(0.1) 35 | assert barrier.wait(timeout=1000) 36 | times.append(time.monotonic()) 37 | 38 | try: 39 | # When I run those workers 40 | with ThreadPoolExecutor(max_workers=8) as e: 41 | for future in [e.submit(worker), e.submit(worker)]: 42 | future.result() 43 | except NotImplementedError: 44 | pytest.skip("Waiting is not supported under this backend.") 45 | 46 | # Then their execution times should be really close to one another 47 | assert abs(times[0] - times[1]) <= 0.01 48 | 49 | 50 | def test_barriers_can_timeout(rate_limiter_backend): 51 | # Given that I have a barrier of two parties 52 | barrier = Barrier(rate_limiter_backend, "sequential-barrier", ttl=30000) 53 | assert barrier.create(parties=2) 54 | 55 | try: 56 | # When I wait on the barrier with a timeout 57 | # Then I should get False back 58 | assert not barrier.wait(timeout=1000) 59 | except NotImplementedError: 60 | pytest.skip("Waiting is not supported under this backend.") 61 | -------------------------------------------------------------------------------- /tests/test_broker.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import pytest 3 | 4 | import dramatiq 5 | import dramatiq.broker 6 | from dramatiq.brokers.rabbitmq import RabbitmqBroker 7 | from dramatiq.middleware import Middleware 8 | 9 | from .common import RABBITMQ_CREDENTIALS, skip_on_windows 10 | 11 | 12 | class EmptyMiddleware(Middleware): 13 | pass 14 | 15 | 16 | def test_broker_uses_rabbitmq_if_not_set(): 17 | # Given that no global broker is set 18 | dramatiq.broker.global_broker = None 19 | 20 | # If I try to get the global broker 21 | broker = dramatiq.get_broker() 22 | 23 | # I expect it to be a RabbitmqBroker instance 24 | assert isinstance(broker, RabbitmqBroker) 25 | 26 | 27 | @skip_on_windows 28 | def test_broker_middleware_can_be_added_before_other_middleware(stub_broker): 29 | from dramatiq.middleware import Prometheus 30 | 31 | # Given that I have a custom middleware 32 | empty_middleware = EmptyMiddleware() 33 | 34 | # If I add it before the Prometheus middleware 35 | stub_broker.add_middleware(empty_middleware, before=Prometheus) 36 | 37 | # I expect it to be the first middleware 38 | assert stub_broker.middleware[0] == empty_middleware 39 | 40 | 41 | @skip_on_windows 42 | def test_broker_middleware_can_be_added_after_other_middleware(stub_broker): 43 | from dramatiq.middleware import Prometheus 44 | 45 | # Given that I have a custom middleware 46 | empty_middleware = EmptyMiddleware() 47 | 48 | # If I add it after the Prometheus middleware 49 | stub_broker.add_middleware(empty_middleware, after=Prometheus) 50 | 51 | # I expect it to be the second middleware 52 | assert stub_broker.middleware[1] == empty_middleware 53 | 54 | 55 | def test_broker_middleware_can_fail_to_be_added_before_or_after_missing_middleware(stub_broker): 56 | # Given that I have a custom middleware 57 | empty_middleware = EmptyMiddleware() 58 | 59 | # If I add it after a middleware that isn't registered 60 | # I expect a ValueError to be raised 61 | with pytest.raises(ValueError): 62 | stub_broker.add_middleware(empty_middleware, after=EmptyMiddleware) 63 | 64 | 65 | @skip_on_windows 66 | def test_broker_middleware_cannot_be_addwed_both_before_and_after(stub_broker): 67 | from dramatiq.middleware import Prometheus 68 | 69 | # Given that I have a custom middleware 70 | empty_middleware = EmptyMiddleware() 71 | 72 | # If I add it with both before and after parameters 73 | # I expect an AssertionError to be raised 74 | with pytest.raises(AssertionError): 75 | stub_broker.add_middleware(empty_middleware, before=Prometheus, after=Prometheus) 76 | 77 | 78 | def test_can_instantiate_brokers_without_middleware(): 79 | # Given that I have an empty list of middleware 80 | # When I pass that to the RMQ Broker 81 | broker = RabbitmqBroker(middleware=[], credentials=RABBITMQ_CREDENTIALS) 82 | 83 | # Then I should get back a broker with no middleware 84 | assert not broker.middleware 85 | 86 | 87 | def test_broker_middleware_logs_warning_when_added_twice(stub_broker, caplog): 88 | # Set the log level to capture warnings 89 | caplog.set_level(logging.WARNING) 90 | 91 | # Given that I have a custom middleware 92 | empty_middleware1 = EmptyMiddleware() 93 | empty_middleware2 = EmptyMiddleware() 94 | 95 | # When I add the first middleware 96 | stub_broker.add_middleware(empty_middleware1) 97 | 98 | # And I add another middleware of the same type 99 | stub_broker.add_middleware(empty_middleware2) 100 | 101 | # Then I expect a warning to be logged 102 | assert any("You're adding a middleware of the same type twice" in record.message 103 | for record in caplog.records 104 | if record.levelname == "WARNING") 105 | 106 | # And I expect both middlewares to be added 107 | assert empty_middleware1 in stub_broker.middleware 108 | assert empty_middleware2 in stub_broker.middleware 109 | -------------------------------------------------------------------------------- /tests/test_bucket_rate_limiter.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dramatiq.rate_limits import BucketRateLimiter 4 | 5 | from .common import skip_in_ci 6 | 7 | 8 | @skip_in_ci 9 | def test_bucket_rate_limiter_limits_per_bucket(rate_limiter_backend): 10 | # Given that I have a bucket rate limiter and a call database 11 | limiter = BucketRateLimiter(rate_limiter_backend, "sequential-test", limit=2) 12 | calls = 0 13 | 14 | for _ in range(2): 15 | # And I wait until the next second starts 16 | now = time.time() 17 | time.sleep(1 - (now - int(now))) 18 | 19 | # And I acquire it multiple times sequentially 20 | for _ in range(8): 21 | with limiter.acquire(raise_on_failure=False) as acquired: 22 | if not acquired: 23 | continue 24 | 25 | calls += 1 26 | 27 | # I expect it to have succeeded four times 28 | assert calls == 4 29 | -------------------------------------------------------------------------------- /tests/test_callbacks.py: -------------------------------------------------------------------------------- 1 | from collections import Counter 2 | 3 | import pytest 4 | 5 | import dramatiq 6 | 7 | 8 | def test_actors_can_define_success_callbacks(stub_broker, stub_worker): 9 | # Given an actor that returns the sum of two numbers 10 | @dramatiq.actor 11 | def add(x, y): 12 | return x + y 13 | 14 | # And an actor that takes in a number and stores it in a db 15 | db = [] 16 | 17 | @dramatiq.actor 18 | def save(message_data, result): 19 | db.append(result) 20 | 21 | # When I send the first actor a message and tell it to call the 22 | # second actor on success 23 | add.send_with_options(args=(1, 2), on_success=save) 24 | 25 | # And join on the broker and worker 26 | stub_broker.join(add.queue_name) 27 | stub_worker.join() 28 | 29 | # Then my db should contain the result 30 | assert db == [3] 31 | 32 | 33 | def test_actors_can_define_failure_callbacks(stub_broker, stub_worker): 34 | # Given an actor that fails with an exception 35 | @dramatiq.actor(max_retries=0) 36 | def do_work(): 37 | raise Exception() 38 | 39 | # And an actor that reports on exceptions 40 | exceptions = Counter() 41 | 42 | @dramatiq.actor 43 | def report_exceptions(message_data, exception_data): 44 | exceptions.update({message_data["actor_name"]}) 45 | 46 | # When I send the first actor a message and tell it to call the 47 | # second actor on failure 48 | do_work.send_with_options(on_failure="report_exceptions") 49 | 50 | # And join on the broker and worker 51 | stub_broker.join(do_work.queue_name) 52 | stub_worker.join() 53 | 54 | # Then my db should contain the result 55 | assert exceptions[do_work.actor_name] == 1 56 | 57 | 58 | def test_actor_callbacks_raise_type_error_when_given_a_normal_callable(stub_broker): 59 | # Given an actor that does nothing 60 | @dramatiq.actor 61 | def do_work(): 62 | pass 63 | 64 | # And a non-actor callable 65 | def callback(message, res): 66 | pass 67 | 68 | # When I try to set that callable as an on success callback 69 | # Then I should get back a TypeError 70 | with pytest.raises(TypeError): 71 | do_work.send_with_options(on_success=callback) 72 | 73 | 74 | def test_actor_callback_knows_correct_number_of_retries(stub_broker, stub_worker): 75 | MAX_RETRIES = 3 76 | attempts, retries = [], [] 77 | 78 | # Given a callback that handles failure only after the last retry 79 | @dramatiq.actor 80 | def my_callback(message_data, exception_data): 81 | global handled 82 | handled = False 83 | retry = message_data["options"]["retries"] 84 | retries.append(retry) 85 | # Handle failure after last retry 86 | if retry > MAX_RETRIES: 87 | handled = True 88 | 89 | # And an actor that fails every time 90 | @dramatiq.actor(max_retries=MAX_RETRIES, max_backoff=100, on_failure=my_callback.actor_name) 91 | def do_work(): 92 | attempts.append(1) 93 | raise RuntimeError("failure") 94 | 95 | # When I send it a message 96 | do_work.send() 97 | 98 | # And join on the queue 99 | stub_broker.join(do_work.queue_name) 100 | stub_worker.join() 101 | 102 | # Then I expect 4 attempts to have occurred 103 | assert len(attempts) == 4 104 | assert len(retries) == len(attempts) 105 | 106 | # And I expect the retry number to increase every time 107 | assert retries == [1, 2, 3, 4] 108 | # And I expect the callback to have handled the failure 109 | assert handled 110 | -------------------------------------------------------------------------------- /tests/test_canteen.py: -------------------------------------------------------------------------------- 1 | import multiprocessing 2 | 3 | import pytest 4 | 5 | from dramatiq.canteen import Canteen, canteen_add, canteen_get, canteen_try_init 6 | 7 | 8 | def test_canteen_add_adds_paths(): 9 | # Given that I have a Canteen 10 | c = multiprocessing.Value(Canteen) 11 | 12 | # When I append a couple of paths and mark it ready 13 | with canteen_try_init(c): 14 | canteen_add(c, "hello") 15 | canteen_add(c, "there") 16 | 17 | # Then those paths should be stored in the canteen 18 | assert canteen_get(c) == ["hello", "there"] 19 | 20 | 21 | def test_canteen_add_fails_when_adding_too_many_paths(): 22 | # Given that I have a Canteen 23 | c = Canteen() 24 | 25 | # When I append too many paths 26 | # Then a RuntimeError should be raised 27 | with pytest.raises(RuntimeError): 28 | for _ in range(1024): 29 | canteen_add(c, "0" * 1024) 30 | 31 | 32 | def test_canteen_try_init_runs_at_most_once(): 33 | # Given that I have a Canteen 34 | c = multiprocessing.Value(Canteen) 35 | 36 | # When I run two canteen_try_init blocks 37 | with canteen_try_init(c) as acquired: 38 | if acquired: 39 | canteen_add(c, "hello") 40 | 41 | with canteen_try_init(c) as acquired: 42 | if acquired: 43 | canteen_add(c, "goodbye") 44 | 45 | # Then only the first one should run 46 | assert canteen_get(c) == ["hello"] 47 | -------------------------------------------------------------------------------- /tests/test_common.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from dramatiq.common import dq_name, q_name, xq_name 4 | 5 | 6 | @pytest.mark.parametrize("given,expected", [ 7 | ("default", "default"), 8 | ("default.DQ", "default"), 9 | ("default.XQ", "default"), 10 | ]) 11 | def test_q_name_returns_canonical_names(given, expected): 12 | assert q_name(given) == expected 13 | 14 | 15 | @pytest.mark.parametrize("given,expected", [ 16 | ("default", "default.DQ"), 17 | ("default.DQ", "default.DQ"), 18 | ("default.XQ", "default.DQ"), 19 | ]) 20 | def test_dq_name_returns_canonical_delay_names(given, expected): 21 | assert dq_name(given) == expected 22 | 23 | 24 | @pytest.mark.parametrize("given,expected", [ 25 | ("default", "default.XQ"), 26 | ("default.DQ", "default.XQ"), 27 | ("default.XQ", "default.XQ"), 28 | ]) 29 | def test_xq_name_returns_delay_names(given, expected): 30 | assert xq_name(given) == expected 31 | -------------------------------------------------------------------------------- /tests/test_concurrent_rate_limiter.py: -------------------------------------------------------------------------------- 1 | import time 2 | from concurrent.futures import ThreadPoolExecutor 3 | 4 | from dramatiq.rate_limits import ConcurrentRateLimiter, RateLimitExceeded 5 | 6 | 7 | def test_concurrent_rate_limiter_releases_the_lock_after_each_call(rate_limiter_backend): 8 | # Given that I have a distributed mutex and a call database 9 | mutex = ConcurrentRateLimiter(rate_limiter_backend, "sequential-test", limit=1) 10 | calls = 0 11 | 12 | # And I acquire it multiple times sequentially 13 | for _ in range(8): 14 | with mutex.acquire(raise_on_failure=False) as acquired: 15 | if not acquired: 16 | continue 17 | 18 | calls += 1 19 | 20 | # I expect it to have succeeded that number of times 21 | assert calls == 8 22 | 23 | 24 | def test_concurrent_rate_limiter_can_act_as_a_mutex(rate_limiter_backend): 25 | # Given that I have a distributed mutex and a call database 26 | mutex = ConcurrentRateLimiter(rate_limiter_backend, "concurrent-test", limit=1) 27 | calls = [] 28 | 29 | # And a function that adds calls to the database after acquiring the mutex 30 | def work(): 31 | with mutex.acquire(raise_on_failure=False) as acquired: 32 | if not acquired: 33 | return 34 | 35 | calls.append(1) 36 | time.sleep(0.3) 37 | 38 | # If I execute multiple workers concurrently 39 | with ThreadPoolExecutor(max_workers=8) as e: 40 | futures = [] 41 | for _ in range(8): 42 | futures.append(e.submit(work)) 43 | 44 | for future in futures: 45 | future.result() 46 | 47 | # I expect only one call to have succeeded 48 | assert sum(calls) == 1 49 | 50 | 51 | def test_concurrent_rate_limiter_limits_concurrency(rate_limiter_backend): 52 | # Given that I have a distributed rate limiter and a call database 53 | mutex = ConcurrentRateLimiter(rate_limiter_backend, "concurrent-test", limit=4) 54 | calls = [] 55 | 56 | # And a function that adds calls to the database after acquiring the rate limit 57 | def work(): 58 | try: 59 | with mutex.acquire(): 60 | calls.append(1) 61 | time.sleep(0.3) 62 | except RateLimitExceeded: 63 | pass 64 | 65 | # If I execute multiple workers concurrently 66 | with ThreadPoolExecutor(max_workers=32) as e: 67 | futures = [] 68 | for _ in range(32): 69 | futures.append(e.submit(work)) 70 | 71 | for future in futures: 72 | future.result() 73 | 74 | # I expect at most 4 calls to have succeeded 75 | assert 3 <= sum(calls) <= 4 76 | -------------------------------------------------------------------------------- /tests/test_encoders.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | import dramatiq 4 | 5 | 6 | @pytest.fixture 7 | def pickle_encoder(): 8 | old_encoder = dramatiq.get_encoder() 9 | new_encoder = dramatiq.PickleEncoder() 10 | dramatiq.set_encoder(new_encoder) 11 | yield new_encoder 12 | dramatiq.set_encoder(old_encoder) 13 | 14 | 15 | def test_set_encoder_sets_the_global_encoder(pickle_encoder): 16 | # Given that I've set a Pickle encoder as the global encoder 17 | # When I get the global encoder 18 | encoder = dramatiq.get_encoder() 19 | 20 | # Then it should be the same as the encoder that was set 21 | assert encoder == pickle_encoder 22 | 23 | 24 | def test_pickle_encoder(pickle_encoder, stub_broker, stub_worker): 25 | # Given that I've set a Pickle encoder as the global encoder 26 | # And I have an actor that adds a value to a db 27 | db = [] 28 | 29 | @dramatiq.actor 30 | def add_value(x): 31 | db.append(x) 32 | 33 | # When I send that actor a message 34 | add_value.send(1) 35 | 36 | # And wait on the broker and worker 37 | stub_broker.join(add_value.queue_name) 38 | stub_worker.join() 39 | 40 | # Then I expect the message to have been processed 41 | assert db == [1] 42 | -------------------------------------------------------------------------------- /tests/test_generic_actors.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import Mock 2 | 3 | import pytest 4 | 5 | import dramatiq 6 | 7 | 8 | def test_generic_actors_can_be_defined(stub_broker): 9 | # Given that I've subclassed GenericActor 10 | class Add(dramatiq.GenericActor): 11 | def perform(self, x, y): 12 | return x + y 13 | 14 | # Then Add.__actor__ should be an instance of Actor 15 | assert isinstance(Add.__actor__, dramatiq.Actor) 16 | 17 | # And it should be callable 18 | assert Add(1, 2) == 3 19 | 20 | 21 | def test_generic_actors_can_be_assigned_options(stub_broker): 22 | # Given that I've subclassed GenericActor 23 | class Add(dramatiq.GenericActor): 24 | # When I set its max_retries value to 32 25 | class Meta: 26 | max_retries = 32 27 | 28 | def perform(self, x, y): 29 | return x + y 30 | 31 | # Then the resulting actor should have that option set 32 | assert Add.options["max_retries"] == 32 33 | 34 | 35 | def test_generic_actors_raise_not_implemented_if_perform_is_missing(stub_broker): 36 | # Given that I've subclassed GenericActor without implementing perform 37 | class Foo(dramatiq.GenericActor): 38 | pass 39 | 40 | # When I call that actor 41 | # Then a NotImplementedError should be raised 42 | with pytest.raises(NotImplementedError): 43 | Foo() 44 | 45 | 46 | def test_generic_actors_can_be_abstract(stub_broker, stub_worker): 47 | # Given that I have a calls database 48 | calls = set() 49 | 50 | # And I've subclassed GenericActor 51 | class BaseTask(dramatiq.GenericActor): 52 | # When I set abstract to True 53 | class Meta: 54 | abstract = True 55 | queue_name = "tasks" 56 | 57 | def get_task_name(self): 58 | raise NotImplementedError 59 | 60 | def perform(self): 61 | calls.add(self.get_task_name()) 62 | 63 | # Then BaseTask should not be an Actor 64 | assert not isinstance(BaseTask, dramatiq.Actor) 65 | 66 | # When I subclass BaseTask 67 | class FooTask(BaseTask): 68 | def get_task_name(self): 69 | return "Foo" 70 | 71 | class BarTask(BaseTask): 72 | def get_task_name(self): 73 | return "Bar" 74 | 75 | # Then both subclasses should be actors 76 | # And they should inherit the parent's meta 77 | assert isinstance(FooTask.__actor__, dramatiq.Actor) 78 | assert isinstance(BarTask.__actor__, dramatiq.Actor) 79 | assert FooTask.queue_name == BarTask.queue_name == "tasks" 80 | 81 | # When I send both actors a message 82 | # And wait for them to get processed 83 | FooTask.send() 84 | BarTask.send() 85 | stub_broker.join(queue_name=BaseTask.Meta.queue_name) 86 | stub_worker.join() 87 | 88 | # Then my calls database should contain both task names 89 | assert calls == {"Foo", "Bar"} 90 | 91 | 92 | def test_generic_actors_can_have_class_attributes(stub_broker): 93 | # Given a generic actor with class attributes 94 | class DoSomething(dramatiq.GenericActor): 95 | STATUS_RUNNING = "running" 96 | STATUS_DONE = "done" 97 | 98 | # When I access one of it class attributes 99 | # Then I should get back that attribute's value 100 | assert DoSomething.STATUS_DONE == "done" 101 | 102 | 103 | def test_generic_actors_can_accept_custom_actor_registry(stub_broker): 104 | # Given a generic actor with a custom actor registry 105 | actor_instance = Mock() 106 | actor_registry = Mock(return_value=actor_instance) 107 | 108 | class CustomActor(dramatiq.GenericActor): 109 | class Meta: 110 | actor = actor_registry 111 | 112 | def perform(self): 113 | pass 114 | 115 | # Then CustomActor.__actor__ should be the actor instance 116 | assert CustomActor.__actor__ is actor_instance 117 | 118 | # And the actor registry should be called with CustomActor 119 | actor_registry.assert_called_once_with(CustomActor) 120 | 121 | 122 | def test_getattr_generic_actor(): 123 | with pytest.raises(AttributeError): 124 | dramatiq.GenericActor.missing_property 125 | -------------------------------------------------------------------------------- /tests/test_messages.py: -------------------------------------------------------------------------------- 1 | import dramatiq 2 | 3 | 4 | def test_messages_have_namedtuple_methods(stub_broker): 5 | @dramatiq.actor 6 | def add(x, y): 7 | return x + y 8 | 9 | msg1 = add.message(1, 2) 10 | assert msg1.asdict() == msg1._asdict() 11 | 12 | assert msg1._field_defaults == {} 13 | assert msg1._fields == ( 14 | "queue_name", 15 | "actor_name", 16 | "args", 17 | "kwargs", 18 | "options", 19 | "message_id", 20 | "message_timestamp", 21 | ) 22 | 23 | msg2 = msg1._replace(queue_name="example") 24 | assert msg2._asdict()["queue_name"] == "example" 25 | 26 | 27 | def test_messageproxy_representation(stub_broker): 28 | """Ensure that MessageProxy defines ``__repr__``. 29 | 30 | While Dramatiq logs messages and MessageProxies formatted using the "%s" placeholder, other tools like Sentry-SDK 31 | force ``repr`` of any log param that isn't a primitive. This means that MessageProxy has to implement ``__repr__`` 32 | in order for the captured messages to be more helpful than: 33 | 34 | > Failed to process message with unhandled exception. 35 | """ 36 | @dramatiq.actor 37 | def actor(arg): 38 | return arg 39 | 40 | message = actor.send("input") 41 | 42 | consumer = stub_broker.consume("default") 43 | message_proxy = next(consumer) 44 | 45 | assert repr(message) in repr(message_proxy), "Expecting MessageProxy repr to contain Message repr" 46 | assert str(message) == str(message_proxy), "Expecting identical __str__ of MessageProxy and Message" 47 | -------------------------------------------------------------------------------- /tests/test_pidfile.py: -------------------------------------------------------------------------------- 1 | import os 2 | import signal 3 | import time 4 | 5 | from dramatiq.brokers.stub import StubBroker 6 | 7 | from .common import skip_on_windows 8 | 9 | broker = StubBroker() 10 | 11 | 12 | def remove(filename): 13 | try: 14 | os.remove(filename) 15 | except OSError: 16 | pass 17 | 18 | 19 | @skip_on_windows 20 | def test_cli_scrubs_stale_pid_files(start_cli): 21 | try: 22 | # Given that I have an existing file containing an old pid 23 | filename = "test_scrub.pid" 24 | with open(filename, "w") as f: 25 | f.write("999999") 26 | 27 | # When I try to start the cli and pass that file as a PID file 28 | proc = start_cli("tests.test_pidfile:broker", extra_args=["--pid-file", filename]) 29 | 30 | # And I wait for it to write the pid file 31 | time.sleep(1) 32 | 33 | # Then the process should write its pid to the file 34 | with open(filename, "r") as f: 35 | pid = int(f.read()) 36 | 37 | assert pid == proc.pid 38 | 39 | # When I stop the process 40 | proc.terminate() 41 | proc.wait() 42 | 43 | # Then the process should exit with return code 0 44 | assert proc.returncode == 0 45 | 46 | # And the file should be removed 47 | assert not os.path.exists(filename) 48 | finally: 49 | remove(filename) 50 | 51 | 52 | def test_cli_aborts_when_pidfile_contains_garbage(start_cli): 53 | try: 54 | # Given that I have an existing file containing important information 55 | filename = "test_garbage.pid" 56 | with open(filename, "w") as f: 57 | f.write("important!") 58 | 59 | # When I try to start the cli and pass that file as a PID file 60 | proc = start_cli("tests.test_pidfile:broker", extra_args=["--pid-file", filename]) 61 | proc.wait() 62 | 63 | # Then the process should exit with return code 4 64 | assert proc.returncode == 4 65 | finally: 66 | remove(filename) 67 | 68 | 69 | @skip_on_windows 70 | def test_cli_with_pidfile_can_be_reloaded(start_cli): 71 | try: 72 | # Given that I have a PID file 73 | filename = "test_reload.pid" 74 | 75 | # When I try to start the cli and pass that file as a PID file 76 | proc = start_cli("tests.test_pidfile:broker", extra_args=["--pid-file", filename]) 77 | time.sleep(1) 78 | 79 | # And send the proc a HUP signal 80 | proc.send_signal(signal.SIGHUP) 81 | time.sleep(5) 82 | 83 | # And then terminate the process 84 | proc.terminate() 85 | proc.wait() 86 | 87 | # Then the process should exit with return code 0 88 | assert proc.returncode == 0 89 | finally: 90 | remove(filename) 91 | -------------------------------------------------------------------------------- /tests/test_stub_broker.py: -------------------------------------------------------------------------------- 1 | import time 2 | from unittest.mock import Mock 3 | 4 | import pytest 5 | 6 | import dramatiq 7 | from dramatiq import QueueJoinTimeout, QueueNotFound 8 | 9 | 10 | def test_stub_broker_raises_queue_error_when_consuming_undeclared_queues(stub_broker): 11 | # Given that I have a stub broker 12 | # If I attempt to consume a queue that wasn't declared 13 | # I expect a QueueNotFound error to be raised 14 | with pytest.raises(QueueNotFound): 15 | stub_broker.consume("idontexist") 16 | 17 | 18 | def test_stub_broker_raises_queue_error_when_enqueueing_messages_on_undeclared_queues(stub_broker): 19 | # Given that I have a stub broker 20 | # If I attempt to enqueue a message on a queue that wasn't declared 21 | # I expect a QueueNotFound error to be raised 22 | with pytest.raises(QueueNotFound): 23 | stub_broker.enqueue(Mock(queue_name="idontexist")) 24 | 25 | 26 | def test_stub_broker_raises_queue_error_when_joining_on_undeclared_queues(stub_broker): 27 | # Given that I have a stub broker 28 | # If I attempt to join on a queue that wasn't declared 29 | # I expect a QueueNotFound error to be raised 30 | with pytest.raises(QueueNotFound): 31 | stub_broker.join("idontexist") 32 | 33 | 34 | def test_stub_broker_can_be_flushed(stub_broker): 35 | # Given that I have an actor 36 | @dramatiq.actor 37 | def do_work(): 38 | pass 39 | 40 | # When I send that actor a message 41 | do_work.send() 42 | 43 | # And when there is already a message on that actor's dead-letter queue 44 | stub_broker.dead_letters_by_queue[do_work.queue_name].append("dead letter") 45 | 46 | # Then its queue should contain the right number of messages 47 | assert stub_broker.queues[do_work.queue_name].qsize() == 1 48 | assert len(stub_broker.dead_letters) == 1 49 | 50 | # When I flush all of the queues 51 | stub_broker.flush_all() 52 | 53 | # Then the queue should be empty 54 | assert stub_broker.queues[do_work.queue_name].qsize() == 0 55 | # and it should contain no in-progress tasks 56 | assert stub_broker.queues[do_work.queue_name].unfinished_tasks == 0 57 | # and the dead-letter queue should be empty 58 | assert len(stub_broker.dead_letters) == 0 59 | 60 | 61 | def test_stub_broker_can_join_with_timeout(stub_broker, stub_worker): 62 | # Given that I have an actor that takes a long time to run 63 | @dramatiq.actor 64 | def do_work(): 65 | time.sleep(1) 66 | 67 | # When I send that actor a message 68 | do_work.send() 69 | 70 | # And join on its queue with a timeout 71 | # Then I expect a QueueJoinTimeout to be raised 72 | with pytest.raises(QueueJoinTimeout): 73 | stub_broker.join(do_work.queue_name, timeout=500) 74 | 75 | 76 | def test_stub_broker_join_reraises_actor_exceptions_in_the_joining_current_thread(stub_broker, stub_worker): 77 | # Given that I have an actor that always fails with a custom exception 78 | class CustomError(Exception): 79 | pass 80 | 81 | @dramatiq.actor(max_retries=0) 82 | def do_work(): 83 | raise CustomError("well, shit") 84 | 85 | # When I send that actor a message 86 | do_work.send() 87 | 88 | # And join on its queue 89 | # Then that exception should be raised in my thread 90 | with pytest.raises(CustomError): 91 | stub_broker.join(do_work.queue_name, fail_fast=True) 92 | -------------------------------------------------------------------------------- /tests/test_watch.py: -------------------------------------------------------------------------------- 1 | import time 2 | from pathlib import Path 3 | 4 | import pytest 5 | 6 | import dramatiq 7 | from dramatiq.brokers.redis import RedisBroker 8 | from dramatiq.common import current_millis 9 | 10 | from .common import skip_in_ci, skip_on_pypy, skip_on_windows 11 | 12 | broker = RedisBroker() 13 | loaded_at = current_millis() 14 | 15 | 16 | @dramatiq.actor(broker=broker) 17 | def write_loaded_at(filename): 18 | with open(filename, "w") as f: 19 | f.write(str(loaded_at)) 20 | 21 | 22 | @skip_in_ci 23 | @skip_on_windows 24 | @skip_on_pypy 25 | @pytest.mark.parametrize("extra_args", [ 26 | (), 27 | ("--watch-use-polling",), 28 | ]) 29 | def test_cli_can_watch_for_source_code_changes(start_cli, extra_args): 30 | # Given that I have a shared file the processes can use to communicate with 31 | filename = "/tmp/dramatiq-loaded-at" 32 | 33 | # When I start my workers 34 | start_cli("tests.test_watch:broker", extra_args=[ 35 | "--processes", "1", 36 | "--threads", "1", 37 | "--watch", "tests", 38 | *extra_args, 39 | ]) 40 | 41 | # And enqueue a task to write the loaded timestamp 42 | write_loaded_at.send(filename) 43 | broker.join(write_loaded_at.queue_name) 44 | 45 | # Then I expect a timestamp to have been written to the file 46 | with open(filename, "r") as f: 47 | timestamp_1 = int(f.read()) 48 | 49 | # When I then update a watched file's mtime 50 | (Path("tests") / "test_watch.py").touch() 51 | 52 | # And wait for the workers to reload 53 | time.sleep(5) 54 | 55 | # And write another timestamp 56 | write_loaded_at.send(filename) 57 | broker.join(write_loaded_at.queue_name) 58 | 59 | # Then I expect another timestamp to have been written to the file 60 | with open(filename, "r") as f: 61 | timestamp_2 = int(f.read()) 62 | 63 | # And the second time to be at least a second apart from the first 64 | assert timestamp_2 - timestamp_1 >= 1000 65 | 66 | # When I open a watched file, this should not trigger a reload 67 | last_loaded_at = timestamp_2 68 | with (Path("tests") / "test_watch.py").open("r"): 69 | time.sleep(5) 70 | write_loaded_at.send(filename) 71 | broker.join(write_loaded_at.queue_name) 72 | 73 | # Then I expect another timestamp to have been written to the file 74 | with open(filename, "r") as f: 75 | timestamp_3 = int(f.read()) 76 | 77 | assert last_loaded_at == timestamp_3 78 | -------------------------------------------------------------------------------- /tests/test_window_rate_limiter.py: -------------------------------------------------------------------------------- 1 | import time 2 | from collections import defaultdict 3 | from concurrent.futures import ThreadPoolExecutor 4 | 5 | from dramatiq.rate_limits import WindowRateLimiter 6 | 7 | from .common import skip_in_ci 8 | 9 | 10 | @skip_in_ci 11 | def test_window_rate_limiter_limits_per_window(rate_limiter_backend): 12 | # Given that I have a bucket rate limiter and a call database 13 | limiter = WindowRateLimiter(rate_limiter_backend, "window-test", limit=2, window=5) 14 | calls = defaultdict(lambda: 0) 15 | 16 | # And a function that increments keys over the span of 20 seconds 17 | def work(): 18 | for _ in range(20): 19 | for _ in range(8): 20 | with limiter.acquire(raise_on_failure=False) as acquired: 21 | if not acquired: 22 | continue 23 | 24 | calls[int(time.time())] += 1 25 | 26 | time.sleep(1) 27 | 28 | # If I run that function multiple times concurrently 29 | with ThreadPoolExecutor(max_workers=8) as e: 30 | futures = [] 31 | for _ in range(8): 32 | futures.append(e.submit(work)) 33 | 34 | for future in futures: 35 | future.result() 36 | 37 | # I expect between 8 and 10 calls to have been made in total 38 | assert 8 <= sum(calls.values()) <= 10 39 | -------------------------------------------------------------------------------- /tests/test_worker.py: -------------------------------------------------------------------------------- 1 | from .common import worker 2 | 3 | 4 | def test_workers_dont_register_queues_that_arent_whitelisted(stub_broker): 5 | # Given that I have a worker object with a restricted set of queues 6 | with worker(stub_broker, queues={"a", "b"}) as stub_worker: 7 | # When I try to register a consumer for a queue that hasn't been whitelisted 8 | stub_broker.declare_queue("c") 9 | stub_broker.declare_queue("c.DQ") 10 | 11 | # Then a consumer should not get spun up for that queue 12 | assert "c" not in stub_worker.consumers 13 | assert "c.DQ" not in stub_worker.consumers 14 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist= 3 | py{39,310,311,312,313}{,-gevent} 4 | docs 5 | lint 6 | 7 | [testenv] 8 | setenv = 9 | PYTHONTRACEMALLOC=1 10 | extras= 11 | dev 12 | commands= 13 | python -Wall -m pytest --benchmark-skip {posargs} 14 | passenv= 15 | TRAVIS 16 | 17 | [testenv:py{39,310,311,312,313}-gevent] 18 | setenv = 19 | PYTHONTRACEMALLOC=1 20 | extras= 21 | dev 22 | commands= 23 | python -Wall {toxinidir}/pytest-gevent.py --benchmark-skip {posargs} 24 | passenv= 25 | TRAVIS 26 | 27 | [testenv:docs] 28 | allowlist_externals=make 29 | changedir=docs 30 | commands= 31 | make html 32 | 33 | [testenv:lint] 34 | extras = 35 | dev 36 | commands= 37 | flake8 {toxinidir}/dramatiq {toxinidir}/examples {toxinidir}/tests 38 | isort -c {toxinidir}/dramatiq 39 | mypy {toxinidir}/dramatiq {toxinidir}/tests 40 | --------------------------------------------------------------------------------