├── .github ├── dependabot.yml └── workflows │ ├── release.yml │ └── test.yaml ├── .gitignore ├── CHANGELOG.md ├── Dockerfile ├── LICENSE ├── MANIFEST.in ├── README.rst ├── docker-compose.yml ├── pyproject.toml ├── requirements-lint.txt ├── requirements-test.txt ├── requirements.txt ├── scripts └── redis_scan.py ├── setup.cfg ├── setup.py ├── tasktiger ├── __init__.py ├── _internal.py ├── exceptions.py ├── executor.py ├── flask_script.py ├── logging.py ├── lua │ ├── move_task.lua │ └── semaphore.lua ├── migrations.py ├── redis_scripts.py ├── redis_semaphore.py ├── retry.py ├── rollbar.py ├── runner.py ├── schedule.py ├── stats.py ├── task.py ├── tasktiger.py ├── test_helpers.py ├── timeouts.py ├── types.py ├── utils.py └── worker.py ├── tests ├── __init__.py ├── config.py ├── conftest.py ├── tasks.py ├── tasks_periodic.py ├── test_base.py ├── test_context_manager.py ├── test_lazy_init.py ├── test_logging.py ├── test_migrations.py ├── test_periodic.py ├── test_queue_size.py ├── test_redis_scripts.py ├── test_semaphore.py ├── test_stats.py ├── test_task.py ├── test_workers.py └── utils.py └── tox.ini /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | # Docs: https://help.github.com/en/github/administering-a-repository/configuration-options-for-dependency-updates 2 | 3 | version: 2 4 | updates: 5 | # Enable version updates for python 6 | - package-ecosystem: "pip" 7 | directory: "/" 8 | open-pull-requests-limit: 5 9 | schedule: 10 | interval: "weekly" 11 | pull-request-branch-name: 12 | # so it's compatible with docker tags 13 | separator: "-" 14 | assignees: 15 | - AlecRosenbaum 16 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release To PyPI 2 | 3 | on: 4 | workflow_dispatch: 5 | inputs: 6 | requested_release_tag: 7 | description: 'The tag to use for this release (e.g., `v2.3.0`)' 8 | required: true 9 | 10 | jobs: 11 | build_and_upload: 12 | runs-on: 'ubuntu-20.04' 13 | environment: production 14 | permissions: 15 | # id-token for the trusted publisher setup 16 | id-token: write 17 | # for tagging the commit 18 | contents: write 19 | steps: 20 | - uses: actions/checkout@v2 21 | 22 | - uses: actions/setup-python@v2 23 | name: Install Python 24 | with: 25 | python-version: 3.8 26 | 27 | - run: | 28 | pip install packaging 29 | - name: Normalize the release version 30 | run: | 31 | echo "release_version=`echo '${{ github.event.inputs.requested_release_tag }}' | sed 's/^v//'`" >> $GITHUB_ENV 32 | - name: Normalize the release tag 33 | run: | 34 | echo "release_tag=v${release_version}" >> $GITHUB_ENV 35 | - name: Get the VERSION from setup.py 36 | run: | 37 | echo "package_version=`grep -Po '__version__ = "\K[^"]*' tasktiger/__init__.py`" >> $GITHUB_ENV 38 | - name: Get the latest version from PyPI 39 | run: | 40 | curl https://pypi.org/pypi/tasktiger/json | python -c 'import json, sys; contents=sys.stdin.read(); parsed = json.loads(contents); print("pypi_version=" + parsed["info"]["version"])' >> $GITHUB_ENV 41 | - name: Log all the things 42 | run: | 43 | echo 'Requested release tag `${{ github.event.inputs.requested_release_tag }}`' 44 | echo 'Release version `${{ env.release_version }}`' 45 | echo 'Release tag `${{ env.release_tag }}`' 46 | echo 'version in package `${{ env.package_version }}`' 47 | echo 'Version in PyPI `${{ env.pypi_version }}`' 48 | - name: Verify that the version string we produced looks like a version string 49 | run: | 50 | echo "${{ env.release_version }}" | sed '/^[0-9]\+\.[0-9]\+\.[0-9]\+$/!{q1}' 51 | - name: Verify that the version tag we produced looks like a version tag 52 | run: | 53 | echo "${{ env.release_tag }}" | sed '/^v[0-9]\+\.[0-9]\+\.[0-9]\+$/!{q1}' 54 | - name: Verify that the release version matches the VERSION in the package source 55 | run: | 56 | [[ ${{ env.release_version }} == ${{ env.package_version }} ]] 57 | - name: Verify that the `release_version` is larger/newer than the existing release in PyPI 58 | run: | 59 | python -c 'import sys; from packaging import version; code = 0 if version.parse("${{ env.package_version }}") > version.parse("${{ env.pypi_version }}") else 1; sys.exit(code)' 60 | - name: Verify that the `release_version` is present in the CHANGELOG 61 | run: | 62 | grep ${{ env.release_version }} CHANGELOG.md 63 | - name: Serialize normalized release values as outputs 64 | run: | 65 | echo "release_version=$release_version" 66 | echo "release_tag=$release_tag" 67 | echo "release_version=$release_version" >> $GITHUB_OUTPUT 68 | echo "release_tag=$release_tag" >> $GITHUB_OUTPUT 69 | - name: Tag commit 70 | uses: actions/github-script@v7.0.1 71 | with: 72 | script: | 73 | github.rest.git.createRef({ 74 | owner: context.repo.owner, 75 | repo: context.repo.repo, 76 | ref: 'refs/tags/${{ env.release_tag }}', 77 | sha: context.sha 78 | }) 79 | - name: Build Source Distribution 80 | run: | 81 | python setup.py sdist 82 | - name: Upload to PyPI 83 | uses: closeio/gh-action-pypi-publish@ec4db0b4ddc65acdf4bff5fa45ac92d78b56bdf0 # v1.9 84 | -------------------------------------------------------------------------------- /.github/workflows/test.yaml: -------------------------------------------------------------------------------- 1 | name: test-workflow 2 | on: 3 | # When any branch in the repository is pushed 4 | push: 5 | # When a pull request is created 6 | pull_request: 7 | # When manually triggered to run 8 | workflow_dispatch: 9 | 10 | jobs: 11 | lint: 12 | strategy: 13 | matrix: 14 | python-version: [ '3.8', '3.9', '3.10', '3.11' ] 15 | name: Lint ${{ matrix.python-version }} 16 | runs-on: 'ubuntu-20.04' 17 | container: python:${{ matrix.python-version }} 18 | steps: 19 | - name: Checkout code 20 | uses: actions/checkout@v2 21 | 22 | - name: Lint code 23 | run: | 24 | pip install -c requirements.txt -r requirements-lint.txt 25 | lintlizard --ci 26 | 27 | # Run tests 28 | test: 29 | strategy: 30 | matrix: 31 | python-version: ['3.8', '3.9', '3.10', '3.11'] 32 | os: ['ubuntu-20.04'] 33 | redis-version: [4, 5, "6.2.6", "7.0.9"] 34 | redis-py-version: [3.3.0, 4.6.0] 35 | # Do not cancel any jobs when a single job fails 36 | fail-fast: false 37 | name: Python ${{ matrix.python-version }} on ${{ matrix.os }} with Redis ${{ matrix.redis-version }} and redis-py==${{ matrix.redis-py-version }} 38 | runs-on: ${{ matrix.os }} 39 | container: python:${{ matrix.python-version }} 40 | services: 41 | redis: 42 | image: redis:${{ matrix.redis-version }} 43 | # Set health checks to wait until redis has started 44 | options: >- 45 | --health-cmd "redis-cli ping" 46 | --health-interval 10s 47 | --health-timeout 5s 48 | --health-retries 5 49 | steps: 50 | - name: Checkout code 51 | uses: actions/checkout@v2 52 | 53 | - name: Install dependencies 54 | run: | 55 | pip install -r requirements.txt -r requirements-test.txt 56 | pip install redis==${{ matrix.redis-py-version }} 57 | 58 | - name: Run tests 59 | run: pytest 60 | env: 61 | # The hostname used to communicate with the Redis service container 62 | REDIS_HOST: redis 63 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .dmypy.json 2 | *.pyc 3 | .DS_Store 4 | build/ 5 | dist/ 6 | *.egg-info/ 7 | .tox/ 8 | .idea/ 9 | venv/ 10 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | ## Version 0.21.0 4 | 5 | * When raising `RetryException` with no `method`, use task decorator retry method if set ([356](https://github.com/closeio/tasktiger/pull/356)) 6 | 7 | ## Version 0.20.0 8 | 9 | * Add `tiger.get_sizes_for_queues_and_states` ([352](https://github.com/closeio/tasktiger/pull/352)) 10 | 11 | ## Version 0.19.5 12 | 13 | * First version using the automated-release process 14 | 15 | ## Version 0.19.4 16 | 17 | * Log task processing in sync executor ([347](https://github.com/closeio/tasktiger/pull/347)) 18 | * Test and Docstring improvements ([338](https://github.com/closeio/tasktiger/pull/338), [343](https://github.com/closeio/tasktiger/pull/343)) 19 | 20 | ## Version 0.19.3 21 | 22 | * Stop heartbeat thread in case of unhandled exceptions ([335](https://github.com/closeio/tasktiger/pull/335)) 23 | 24 | ## Version 0.19.2 25 | 26 | * Heartbeat threading-related fixes with synchronous worker ([333](https://github.com/closeio/tasktiger/pull/333)) 27 | 28 | ## Version 0.19.1 29 | 30 | * Implement heartbeat with synchronous TaskTiger worker ([331](https://github.com/closeio/tasktiger/pull/331)) 31 | 32 | ## Version 0.19 33 | 34 | * Adding synchronous (non-forking) executor ([319](https://github.com/closeio/tasktiger/pull/319), [320](https://github.com/closeio/tasktiger/pull/320)) 35 | * If possible, retry tasks that fail with "execution not found" ([323](https://github.com/closeio/tasktiger/pull/323)) 36 | * Option to exit TaskTiger after a certain amount of time ([324](https://github.com/closeio/tasktiger/pull/324)) 37 | 38 | ## Version 0.18.2 39 | 40 | * Purge errored tasks even if task object is not found ([310](https://github.com/closeio/tasktiger/pull/310)) 41 | 42 | ## Version 0.18.1 43 | 44 | * Added `current_serialized_func` property to `TaskTiger` object ([296](https://github.com/closeio/tasktiger/pull/296)) 45 | 46 | ## Version 0.18.0 47 | 48 | * Added support for Redis >= 6.2.7 ([268](https://github.com/closeio/tasktiger/issues/268)) 49 | 50 | ### ⚠️ Breaking changes 51 | 52 | * Removed `execute_pipeline` script ([284](https://github.com/closeio/tasktiger/pull/284)) 53 | 54 | ### Other changes 55 | 56 | * Added typing to more parts of the codebase 57 | * Dropped Python 3.7 support, added Python 3.11 support 58 | * Added CI checks to ensure compatibility on redis-py versions (currently >=3.3.0,<5) 59 | 60 | ## Version 0.17.1 61 | 62 | ### Other changes 63 | 64 | * Add 'task_func' to logging processor ([265](https://github.com/closeio/tasktiger/pull/265)) 65 | * Deprecate Flask-Script integration ([260](https://github.com/closeio/tasktiger/pull/260)) 66 | 67 | ## Version 0.17.0 68 | 69 | ### ⚠️ Breaking changes 70 | 71 | #### Allow truncating task executions ([251](https://github.com/closeio/tasktiger/pull/251)) 72 | 73 | ##### Overview 74 | 75 | This version of TaskTiger switches to using the `t:task::executions_count` Redis key to determine the total number of task executions. In previous versions this was accomplished by obtaining the length of `t:task::executions`. This change was required for the introduction of a parameter to enable the truncation of task execution entries. This is useful for tasks with many retries, where execution entries consume a lot of memory. 76 | 77 | This behavior is incompatible with the previous mechanism and requires a migration to populate the task execution counters. 78 | Without the migration, the execution counters will behave as though they were reset, which may result in existing tasks retrying more times than they should. 79 | 80 | ##### Migration 81 | 82 | The migration can be executed fully live without concern for data integrity. 83 | 84 | 1. Upgrade TaskTiger to `0.16.2` if running a version lower than that. 85 | 2. Call `tasktiger.migrations.migrate_executions_count` with your `TaskTiger` instance, e.g.: 86 | ```py 87 | from tasktiger import TaskTiger 88 | from tasktiger.migrations import migrate_executions_count 89 | 90 | # Instantiate directly or import from your application module 91 | tiger = TaskTiger(...) 92 | 93 | # This could take a while depending on the 94 | # number of failed/retrying tasks you have 95 | migrate_executions_count(tiger) 96 | ``` 97 | 3. Upgrade TaskTiger to `0.17.0`. Done! 98 | 99 | #### Import cleanup ([258](https://github.com/closeio/tasktiger/pull/258)) 100 | 101 | Due to a cleanup of imports, some internal TaskTiger objects can no longer be imported from the public modules. This shouldn't cause problems for most users, but it's a good idea to double check that all imports from the TaskTiger package continue to function correctly in your application. 102 | 103 | ## Version 0.16.2 104 | 105 | ### Other changes 106 | 107 | * Prefilter polled queues ([242](https://github.com/closeio/tasktiger/pull/242)) 108 | * Use SSCAN to prefilter queues in scheduled state ([248](https://github.com/closeio/tasktiger/pull/248)) 109 | * Add task execution counter ([252](https://github.com/closeio/tasktiger/pull/252)) 110 | 111 | ## Version 0.16.1 112 | 113 | ### Other changes 114 | 115 | * Add function name to tasktiger done log messages ([203](https://github.com/closeio/tasktiger/pull/203)) 116 | * Add task args / kwargs to the task_error log statement ([215](https://github.com/closeio/tasktiger/pull/215)) 117 | * Fix `hard_timeout` in parent process when stored on task function ([235](https://github.com/closeio/tasktiger/pull/235)) 118 | 119 | ## Version 0.16 120 | 121 | ### Other changes 122 | 123 | * Handle hard timeout in parent process ([f3b3e24](https://github.com/closeio/tasktiger/commit/f3b3e24485497a2b87281a1b809966bcb525c5fc)) 124 | * Add queue name to logs ([a090d00](https://github.com/closeio/tasktiger/commit/a090d00bca496082f149f2187b026ff96a0d4fac)) 125 | 126 | ## Version 0.15 127 | 128 | ### Other changes 129 | 130 | * Populate `Task.ts` field in `Task.from_id` function ([019bf18](https://github.com/closeio/tasktiger/commit/019bf189c9622b299691dbe3b71cefa0bf2ee8dc)) 131 | * Add `TaskTiger.would_process_configured_queue()` function ([217152d](https://github.com/closeio/tasktiger/commit/217152d16ff21a87b70643a0a2571efc91c0aeb9)) 132 | 133 | ## Version 0.14 134 | 135 | ### Other changes 136 | 137 | * Add `Task.time_last_queued` property getter ([6d2285d](https://github.com/closeio/tasktiger/commit/6d2285da5bd5f82455765e6b132594d4ceab2d82)) 138 | 139 | ## Version 0.13 140 | 141 | ### ⚠️ Breaking changes 142 | 143 | #### Changing the locking mechanism 144 | 145 | This new version of TaskTiger uses a new locking mechanism: the `Lock` provided by redis-py. It is incompatible with the old locking mechanism we were using, and several core functions in TaskTiger depends on locking to work correctly, so this warrants a careful migration process. 146 | 147 | You can perform this migration in two ways: via a live migration, or via a downtime migration. After the migration, there's an optional cleanup step. 148 | 149 | ##### The live migration 150 | 151 | 1. Update your environment to TaskTiger 0.12 as usual. 152 | 1. Deploy TaskTiger as it is in the commit SHA `cf600449d594ac22e6d8393dc1009a84b52be0c1`. In `pip` parlance, it would be: 153 | 154 | -e git+ssh://git@github.com/closeio/tasktiger.git@cf600449d594ac22e6d8393dc1009a84b52be0c1#egg=tasktiger 155 | 156 | 1. Wait at least 2-3 minutes with it running in production in all your TaskTiger workers. This is to give time for the old locks to expire, and after that the new locks will be fully in effect. 157 | 1. Deploy TaskTiger 0.13. Your system is migrated. 158 | 159 | ##### The downtime migration 160 | 161 | 1. Update your environment to TaskTiger 0.12 as usual. 162 | 1. Scale your TaskTiger workers down to zero. 163 | 1. Deploy TaskTiger 0.13. Your system is migrated. 164 | 165 | ##### The cleanup step 166 | 167 | Run the script in `scripts/redis_scan.py` to delete the old lock keys from your Redis instance: 168 | 169 | ./scripts/redis_scan.py --host HOST --port PORT --db DB --print --match "t:lock:*" --ttl 300 170 | 171 | The flags: 172 | 173 | - `--host`: The Redis host. Required. 174 | - `--port`: The port the Redis instance is listening on. Defaults to `6379`. 175 | - `--db`: The Redis database. Defaults to `0`. 176 | - `--print`: If you want the script to print which keys it is modifying, use this. 177 | - `--match`: What pattern to look for. If you didn't change the default prefix TaskTiger uses for keys, this will be `t:lock:*`, otherwise it will be `PREFIX:lock:*`. By default, scans all keys. 178 | - `--ttl`: A TTL to set. A TTL of 300 will give you time to undo if you want to halt the migration for whatever reason. (Just call this command again with `--ttl -1`.) By default, does not change keys' TTLs. 179 | 180 | Plus, there is: 181 | 182 | - `--file`: A log file that will receive the changes made. Defaults to `redis-stats.log` in the current working directory. 183 | - `--delay`: How long, in seconds, to wait between `SCAN` iterations. Defaults to `0.1`. 184 | 185 | ## Version 0.12 186 | 187 | ### ⚠️ Breaking changes 188 | 189 | * Drop support for redis-py 2 ([#183](https://github.com/closeio/tasktiger/pull/183)) 190 | 191 | ### Other changes 192 | 193 | * Make the `TaskTiger` instance available for the task via global state ([#170](https://github.com/closeio/tasktiger/pull/170)) 194 | * Support for custom task runners ([#175](https://github.com/closeio/tasktiger/pull/175)) 195 | * Add ability to configure a poll- vs push-method for task runners to discover new tasks ([#176](https://github.com/closeio/tasktiger/pull/176)) 196 | * `unique_key` specifies the list of kwargs to use to construct the unique key ([#180](https://github.com/closeio/tasktiger/pull/180)) 197 | 198 | ### Bugfixes 199 | 200 | * Ensure task exists in the given queue when retrieving it ([#184](https://github.com/closeio/tasktiger/pull/184)) 201 | * Clear retried executions from successful periodic tasks ([#188](https://github.com/closeio/tasktiger/pull/188)) 202 | 203 | ## Version 0.11 204 | 205 | ### Breaking changes 206 | 207 | * Drop support for Python 3.4 and add testing for Python 3.7 ([#163](https://github.com/closeio/tasktiger/pull/163)) 208 | 209 | ### Other changes 210 | 211 | * Add support for redis-py 3 ([#163](https://github.com/closeio/tasktiger/pull/163)) 212 | * Fix test timings ([#164](https://github.com/closeio/tasktiger/pull/164)) 213 | * Allow custom context managers to see task errors ([#165](https://github.com/closeio/tasktiger/pull/165)). Thanks @igor47 214 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM circleci/python:3.8 2 | 3 | WORKDIR /src 4 | COPY requirements.txt . 5 | COPY requirements-test.txt . 6 | RUN pip install --user -r requirements.txt -r requirements-test.txt 7 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2020 Elastic Inc. (Close) 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include tox.ini LICENSE 2 | recursive-include tests *.py 3 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: "3.7" 2 | services: 3 | redis: 4 | image: redis:7.0.9 5 | expose: 6 | - 6379 7 | tasktiger: 8 | build: 9 | context: . 10 | dockerfile: Dockerfile 11 | environment: 12 | REDIS_HOST: redis 13 | volumes: 14 | - .:/src 15 | depends_on: 16 | - redis 17 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.black] 2 | line-length = 79 3 | exclude = ''' 4 | /( 5 | \.git 6 | )/ 7 | ''' 8 | 9 | [tool.isort] 10 | skip = ['.git', 'venv'] 11 | known_tests = 'tests' 12 | sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', 'TESTS', 'LOCALFOLDER'] 13 | default_section = 'THIRDPARTY' 14 | use_parentheses = true 15 | multi_line_output = 3 16 | include_trailing_comma = true 17 | force_grid_wrap = 0 18 | combine_as_imports = true 19 | line_length = 79 20 | float_to_top = true 21 | -------------------------------------------------------------------------------- /requirements-lint.txt: -------------------------------------------------------------------------------- 1 | lintlizard==0.26.0 2 | types-redis 3 | -------------------------------------------------------------------------------- /requirements-test.txt: -------------------------------------------------------------------------------- 1 | freezefrog==0.4.1 2 | psutil==5.9.8 3 | pytest==8.1.1 -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | click==8.1.7 2 | redis==4.5.2 3 | structlog==24.1.0 4 | croniter 5 | -------------------------------------------------------------------------------- /scripts/redis_scan.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | """ 3 | Safely scan Redis instance and report key stats. 4 | 5 | Can also set the TTL for keys to facilitate removing data. 6 | """ 7 | import signal 8 | import sys 9 | import time 10 | 11 | import click 12 | import redis 13 | 14 | loop = True 15 | 16 | 17 | def signal_handler(signum, frame): 18 | """Signal handler.""" 19 | 20 | global loop 21 | print("Caught ctrl-c, finishing up.") 22 | loop = False 23 | 24 | 25 | def get_size(client, key, key_type): 26 | """Get size of key.""" 27 | 28 | size = -1 29 | if key_type == "string": 30 | size = client.strlen(key) 31 | elif key_type == "zset": 32 | size = client.zcard(key) 33 | elif key_type == "set": 34 | size = client.scard(key) 35 | elif key_type == "list": 36 | size = client.llen(key) 37 | elif key_type == "hash": 38 | size = client.hlen(key) 39 | 40 | return size 41 | 42 | 43 | @click.command() 44 | @click.option("--file", "file_name", default="redis-stats.log") 45 | @click.option("--match", default=None) 46 | @click.option( 47 | "--ttl", 48 | "set_ttl", 49 | default=None, 50 | type=click.INT, 51 | help="Set TTL if one isn't already set (-1 will remove TTL)", 52 | ) 53 | @click.option("--host", required=True) 54 | @click.option("--port", type=click.INT, default=6379) 55 | @click.option("--db", type=click.INT, default=0) 56 | @click.option("--delay", type=click.FLOAT, default=0.1) 57 | @click.option("--print", "print_it", is_flag=True) 58 | def run(host, port, db, delay, file_name, print_it, match, set_ttl=None): 59 | """Run scan.""" 60 | 61 | if set_ttl is not None and match is None: 62 | print("You must specify match when setting TTLs!") 63 | sys.exit(1) 64 | 65 | client = redis.Redis(host=host, port=port, db=db) 66 | 67 | if match: 68 | print(f"Scanning redis keys with match: {match}\n") 69 | else: 70 | print("Scanning all redis keys\n") 71 | 72 | # This is a string because we want the `while` below to run, and the string 73 | # will be correctly interpreted by `client.scan` as a request for a new 74 | # cursor. 75 | cursor = "0" 76 | 77 | signal.signal(signal.SIGINT, signal_handler) 78 | 79 | with open(file_name, "w") as log_file: 80 | while cursor != 0 and loop: 81 | cursor, data = client.scan(cursor=cursor, match=match) 82 | 83 | for key in data: 84 | key_type = client.type(key) 85 | size = get_size(client, key, key_type) 86 | 87 | ttl = client.ttl(key) 88 | new_ttl = None 89 | if set_ttl == -1: 90 | if ttl != -1: 91 | client.persist(key) 92 | new_ttl = -1 93 | elif set_ttl is not None and ttl == -1: 94 | # Only change TTLs for keys with no TTL 95 | client.expire(key, set_ttl) 96 | new_ttl = set_ttl 97 | 98 | line = f"{key} {key_type} {ttl} {new_ttl} {size}" 99 | log_file.write(line + "\n") 100 | if print_it: 101 | print(line) 102 | 103 | log_file.flush() 104 | time.sleep(delay) 105 | 106 | 107 | if __name__ == "__main__": 108 | run() 109 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [flake8] 2 | ignore= 3 | # TODO: Go over these and fix the ones we want to 4 | B007, 5 | B011, 6 | D100, 7 | D101, 8 | D102, 9 | D103, 10 | D104, 11 | D105, 12 | D107, 13 | D200, 14 | D202, 15 | D205, 16 | D400, 17 | D401, 18 | D402, 19 | D403, 20 | D412, 21 | E501, 22 | M105, 23 | M114, 24 | M210, 25 | M300, 26 | M401, 27 | M908, 28 | N801, 29 | N804, 30 | N805, 31 | N806, 32 | N818, 33 | S101, 34 | S201, 35 | S301, 36 | W503, 37 | SFS, 38 | SIM, 39 | 40 | [tool:pytest] 41 | testpaths=tests 42 | 43 | [mypy] 44 | warn_unused_configs = True 45 | ignore_missing_imports = False 46 | disallow_untyped_defs = True 47 | disallow_incomplete_defs = True 48 | no_implicit_optional = True 49 | strict_equality = True 50 | warn_unreachable = True 51 | warn_unused_ignores = True 52 | show_error_context = True 53 | pretty = True 54 | check_untyped_defs = True 55 | python_version = 3.8 56 | files = tasktiger 57 | 58 | [mypy-flask_script.*] 59 | ignore_missing_imports = True 60 | 61 | [mypy-rollbar.*] 62 | ignore_missing_imports = True 63 | 64 | [mypy-structlog.*] 65 | ignore_missing_imports = True 66 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | from setuptools import setup 4 | 5 | VERSION_FILE = "tasktiger/__init__.py" 6 | with open(VERSION_FILE, encoding="utf8") as fd: 7 | version = re.search(r'__version__ = ([\'"])(.*?)\1', fd.read()).group(2) 8 | 9 | with open("README.rst", encoding="utf-8") as file: 10 | long_description = file.read() 11 | 12 | install_requires = ["click", "redis>=3.3.0,<5", "structlog"] 13 | 14 | tests_require = install_requires + ["freezefrog", "pytest", "psutil"] 15 | 16 | setup( 17 | name="tasktiger", 18 | version=version, 19 | url="http://github.com/closeio/tasktiger", 20 | license="MIT", 21 | description="Python task queue", 22 | long_description=long_description, 23 | test_suite="tests", 24 | tests_require=tests_require, 25 | install_requires=install_requires, 26 | classifiers=[ 27 | "Intended Audience :: Developers", 28 | "License :: OSI Approved :: MIT License", 29 | "Operating System :: MacOS :: MacOS X", 30 | "Operating System :: POSIX", 31 | "Operating System :: POSIX :: Linux", 32 | "Topic :: Software Development :: Libraries :: Python Modules", 33 | "Programming Language :: Python", 34 | "Programming Language :: Python :: 3", 35 | "Programming Language :: Python :: 3.8", 36 | "Programming Language :: Python :: 3.9", 37 | "Programming Language :: Python :: 3.10", 38 | "Programming Language :: Python :: 3.11", 39 | ], 40 | packages=["tasktiger"], 41 | package_data={"tasktiger": ["lua/*.lua"]}, 42 | entry_points={"console_scripts": ["tasktiger = tasktiger:run_worker"]}, 43 | ) 44 | -------------------------------------------------------------------------------- /tasktiger/__init__.py: -------------------------------------------------------------------------------- 1 | from .exceptions import ( 2 | JobTimeoutException, 3 | QueueFullException, 4 | RetryException, 5 | StopRetry, 6 | TaskImportError, 7 | TaskNotFound, 8 | ) 9 | from .retry import exponential, fixed, linear 10 | from .schedule import cron_expr, periodic 11 | from .task import Task 12 | from .tasktiger import TaskTiger, run_worker 13 | from .worker import Worker 14 | 15 | __version__ = "0.21.0" 16 | __all__ = [ 17 | "TaskTiger", 18 | "Worker", 19 | "Task", 20 | # Exceptions 21 | "JobTimeoutException", 22 | "RetryException", 23 | "StopRetry", 24 | "TaskImportError", 25 | "TaskNotFound", 26 | "QueueFullException", 27 | # Retry methods 28 | "fixed", 29 | "linear", 30 | "exponential", 31 | # Schedules 32 | "periodic", 33 | "cron_expr", 34 | ] 35 | 36 | 37 | if __name__ == "__main__": 38 | run_worker() 39 | -------------------------------------------------------------------------------- /tasktiger/_internal.py: -------------------------------------------------------------------------------- 1 | import binascii 2 | import calendar 3 | import datetime 4 | import hashlib 5 | import importlib 6 | import json 7 | import operator 8 | import os 9 | import threading 10 | from typing import ( 11 | TYPE_CHECKING, 12 | Any, 13 | Callable, 14 | Collection, 15 | Iterable, 16 | List, 17 | Optional, 18 | Tuple, 19 | Type, 20 | TypedDict, 21 | Union, 22 | ) 23 | 24 | from .exceptions import TaskImportError 25 | 26 | if TYPE_CHECKING: 27 | from .task import Task 28 | from .tasktiger import TaskTiger 29 | 30 | # Task states (represented by different queues) 31 | # Note some client code may rely on the string values (e.g. get_queue_stats). 32 | QUEUED = "queued" 33 | ACTIVE = "active" 34 | SCHEDULED = "scheduled" 35 | ERROR = "error" 36 | 37 | # This lock is acquired in the main process when forking, and must be acquired 38 | # in any thread of the main process when performing an operation that triggers a 39 | # lock that a child process might want to acquire. 40 | # 41 | # Specifically, we use this lock when logging in the StatsThread to prevent a 42 | # deadlock in the child process when the child is forked while the stats thread 43 | # is logging a message. This issue happens because Python acquires a lock while 44 | # logging, so a child process could be stuck forever trying to acquire that 45 | # lock. See http://bugs.python.org/issue6721 for more details. 46 | g_fork_lock = threading.Lock() 47 | 48 | # Global task context. We store this globally (and not on the TaskTiger 49 | # instance) for consistent results just in case the user has multiple TaskTiger 50 | # instances. 51 | 52 | 53 | class _G(TypedDict): 54 | tiger: Optional["TaskTiger"] 55 | current_task_is_batch: Optional[bool] 56 | current_tasks: Optional[List["Task"]] 57 | 58 | 59 | g: _G = {"tiger": None, "current_task_is_batch": None, "current_tasks": None} 60 | 61 | 62 | # from rq 63 | def import_attribute(name: str) -> Any: 64 | """Return an attribute from a dotted path name (e.g. "path.to.func").""" 65 | try: 66 | sep = ":" if ":" in name else "." # For backwards compatibility 67 | module_name, attribute = name.rsplit(sep, 1) 68 | module = importlib.import_module(module_name) 69 | return operator.attrgetter(attribute)(module) 70 | except (ValueError, ImportError, AttributeError) as e: 71 | raise TaskImportError(e) 72 | 73 | 74 | def gen_id() -> str: 75 | """ 76 | Generates and returns a random hex-encoded 256-bit unique ID. 77 | """ 78 | return binascii.b2a_hex(os.urandom(32)).decode("utf8") 79 | 80 | 81 | def gen_unique_id(serialized_name: str, args: Any, kwargs: Any) -> str: 82 | """ 83 | Generates and returns a hex-encoded 256-bit ID for the given task name and 84 | args. Used to generate IDs for unique tasks or for task locks. 85 | """ 86 | return hashlib.sha256( 87 | json.dumps( 88 | {"func": serialized_name, "args": args, "kwargs": kwargs}, 89 | sort_keys=True, 90 | ).encode("utf8") 91 | ).hexdigest() 92 | 93 | 94 | def serialize_func_name(func: Union[Callable, Type]) -> str: 95 | """ 96 | Returns the dotted serialized path to the passed function. 97 | """ 98 | if func.__module__ == "__main__": 99 | raise ValueError( 100 | "Functions from the __main__ module cannot be processed by " 101 | "workers." 102 | ) 103 | try: 104 | # This will only work on Python 3.3 or above, but it will allow us to use static/classmethods 105 | func_name = func.__qualname__ 106 | except AttributeError: 107 | func_name = func.__name__ 108 | return ":".join([func.__module__, func_name]) 109 | 110 | 111 | def dotted_parts(s: str) -> Iterable[str]: 112 | """ 113 | For a string "a.b.c", yields "a", "a.b", "a.b.c". 114 | """ 115 | idx = -1 116 | while s: 117 | idx = s.find(".", idx + 1) 118 | if idx == -1: 119 | yield s 120 | break 121 | yield s[:idx] 122 | 123 | 124 | def reversed_dotted_parts(s: str) -> Iterable[str]: 125 | """ 126 | For a string "a.b.c", yields "a.b.c", "a.b", "a". 127 | """ 128 | idx = -1 129 | if s: 130 | yield s 131 | while s: 132 | idx = s.rfind(".", 0, idx) 133 | if idx == -1: 134 | break 135 | yield s[:idx] 136 | 137 | 138 | def serialize_retry_method(retry_method: Any) -> Tuple[str, Tuple]: 139 | if callable(retry_method): 140 | return (serialize_func_name(retry_method), ()) 141 | else: 142 | return (serialize_func_name(retry_method[0]), retry_method[1]) 143 | 144 | 145 | def get_timestamp( 146 | when: Optional[Union[datetime.timedelta, datetime.datetime]] 147 | ) -> Optional[float]: 148 | # convert timedelta to datetime 149 | if isinstance(when, datetime.timedelta): 150 | when = datetime.datetime.utcnow() + when 151 | 152 | if when: 153 | # Convert to unixtime: utctimetuple drops microseconds so we add 154 | # them manually. 155 | return calendar.timegm(when.utctimetuple()) + when.microsecond / 1.0e6 156 | return None 157 | 158 | 159 | def queue_matches( 160 | queue: str, 161 | only_queues: Optional[Collection[str]] = None, 162 | exclude_queues: Optional[Collection[str]] = None, 163 | ) -> bool: 164 | """Checks if the given queue matches against only/exclude constraints 165 | 166 | Returns whether the given queue should be included by checking each part of 167 | the queue name. 168 | 169 | :param str queue: The queue name to check 170 | :param iterable(str) only_queues: Limit to only these queues 171 | :param iterable(str) exclude_queues: Specifically excluded queues 172 | 173 | :returns: A boolean indicating whether this queue matches against the given 174 | ``only`` and ``excludes`` constraints 175 | """ 176 | # Check arguments to prevent a common footgun of passing 'my_queue' instead 177 | # of ``['my_queue']`` 178 | error_template = ( 179 | "{kwarg} should be an iterable of strings, not a string directly. " 180 | "Did you mean `{kwarg}=['{val}']`?" 181 | ) 182 | assert not isinstance(only_queues, str), error_template.format( 183 | kwarg="queues", val=only_queues 184 | ) 185 | assert not isinstance(exclude_queues, str), error_template.format( 186 | kwarg="exclude_queues", val=exclude_queues 187 | ) 188 | 189 | only_queues = only_queues or [] 190 | exclude_queues = exclude_queues or [] 191 | for part in reversed_dotted_parts(queue): 192 | if part in exclude_queues: 193 | return False 194 | if part in only_queues: 195 | return True 196 | return not only_queues 197 | 198 | 199 | class classproperty(property): 200 | """ 201 | Simple class property implementation. 202 | 203 | Works like @property but on classes. 204 | """ 205 | 206 | def __get__(desc, self, cls): # type:ignore[no-untyped-def] 207 | return desc.fget(cls) # type:ignore[misc] 208 | -------------------------------------------------------------------------------- /tasktiger/exceptions.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from typing import Any 3 | 4 | 5 | class TaskImportError(ImportError): 6 | """ 7 | Raised when a task could not be imported. 8 | """ 9 | 10 | 11 | class JobTimeoutException(BaseException): 12 | """ 13 | Raised when a job takes longer to complete than the allowed maximum timeout 14 | value. 15 | """ 16 | 17 | 18 | class QueueFullException(BaseException): 19 | """ 20 | Raised when a task is attempted to be queued using max_queue_size and the 21 | total queue size (QUEUED + SCHEDULED + ACTIVE) is greater than or equal to 22 | max_queue_size. 23 | """ 24 | 25 | 26 | class StopRetry(Exception): 27 | """ 28 | Raised by a retry function to indicate that the task shouldn't be retried. 29 | """ 30 | 31 | 32 | class RetryException(BaseException): 33 | """ 34 | Alternative to the `retry_on` parameter for retrying a task. 35 | 36 | If this exception is raised within a task, the task will be retried as long 37 | as the retry method permits. 38 | 39 | The default retry method (specified in the task or in DEFAULT_RETRY_METHOD) 40 | may be overridden using the method argument. 41 | 42 | If original_traceback is True and RetryException is raised from within an 43 | `except` block, the original traceback will be logged. 44 | 45 | If `log_error` is set to False and the task fails permanently, a warning 46 | will be logged instead of an error, and the task will be removed from Redis 47 | when it completes. 48 | """ 49 | 50 | def __init__( 51 | self, 52 | method: Any = None, 53 | original_traceback: bool = False, 54 | log_error: bool = True, 55 | ): 56 | self.method = method 57 | self.exc_info = sys.exc_info() if original_traceback else None 58 | self.log_error = log_error 59 | 60 | 61 | class TaskNotFound(Exception): 62 | """ 63 | The task was not found or does not exist in the given queue/state. 64 | """ 65 | -------------------------------------------------------------------------------- /tasktiger/executor.py: -------------------------------------------------------------------------------- 1 | import errno 2 | import fcntl 3 | import os 4 | import random 5 | import select 6 | import signal 7 | import socket 8 | import sys 9 | import threading 10 | import time 11 | import traceback 12 | from contextlib import ExitStack 13 | from typing import ( 14 | TYPE_CHECKING, 15 | Any, 16 | Collection, 17 | ContextManager, 18 | Dict, 19 | List, 20 | Optional, 21 | ) 22 | 23 | from redis.exceptions import LockError 24 | from redis.lock import Lock 25 | from structlog.stdlib import BoundLogger 26 | 27 | from ._internal import ( 28 | g, 29 | g_fork_lock, 30 | serialize_func_name, 31 | serialize_retry_method, 32 | ) 33 | from .exceptions import RetryException 34 | from .redis_semaphore import Semaphore 35 | from .runner import get_runner_class 36 | from .task import Task 37 | from .timeouts import JobTimeoutException 38 | 39 | if TYPE_CHECKING: 40 | from .worker import Worker 41 | 42 | 43 | def sigchld_handler(*args: Any) -> None: 44 | # Nothing to do here. This is just a dummy handler that we set up to catch 45 | # the child process exiting. 46 | pass 47 | 48 | 49 | class WorkerContextManagerStack(ExitStack): 50 | def __init__(self, context_managers: List[ContextManager]) -> None: 51 | super(WorkerContextManagerStack, self).__init__() 52 | 53 | for mgr in context_managers: 54 | self.enter_context(mgr) 55 | 56 | 57 | class Executor: 58 | exit_worker_on_job_timeout = False 59 | 60 | def __init__(self, worker: "Worker"): 61 | self.tiger = worker.tiger 62 | self.worker = worker 63 | self.connection = worker.connection 64 | self.config = worker.config 65 | 66 | def heartbeat( 67 | self, 68 | queue: str, 69 | task_ids: Collection[str], 70 | log: BoundLogger, 71 | locks: Collection[Lock], 72 | queue_lock: Optional[Semaphore], 73 | ) -> None: 74 | self.worker.heartbeat(queue, task_ids) 75 | for lock in locks: 76 | try: 77 | lock.reacquire() 78 | except LockError: 79 | log.warning("could not reacquire lock", lock=lock.name) 80 | if queue_lock: 81 | acquired, current_locks = queue_lock.renew() 82 | if not acquired: 83 | log.debug("queue lock renew failure") 84 | 85 | def execute( 86 | self, 87 | queue: str, 88 | tasks: List[Task], 89 | log: BoundLogger, 90 | locks: Collection[Lock], 91 | queue_lock: Optional[Semaphore], 92 | ) -> bool: 93 | """ 94 | Executes the given tasks. Returns a boolean indicating whether 95 | the tasks were executed successfully. 96 | 97 | Args: 98 | queue: Name of the task queue. 99 | tasks: List of tasks to execute, 100 | log: Logger. 101 | locks: List of task locks to renew periodically. 102 | queue_lock: Optional queue lock to renew periodically for max 103 | workers per queue. 104 | 105 | Returns: 106 | Whether task execution was successful. 107 | """ 108 | raise NotImplementedError 109 | 110 | def execute_tasks(self, tasks: List[Task], log: BoundLogger) -> bool: 111 | """ 112 | Executes the tasks in the current process. Multiple tasks can be passed 113 | for batch processing. However, they must all use the same function and 114 | will share the execution entry. 115 | """ 116 | success = False 117 | 118 | execution: Dict[str, Any] = {} 119 | 120 | assert len(tasks) 121 | task_func = tasks[0].serialized_func 122 | assert all([task_func == task.serialized_func for task in tasks[1:]]) 123 | 124 | execution["time_started"] = time.time() 125 | 126 | try: 127 | func = tasks[0].func 128 | 129 | runner_class = get_runner_class(log, tasks) 130 | runner = runner_class(self.tiger) 131 | 132 | is_batch_func = getattr(func, "_task_batch", False) 133 | g["tiger"] = self.tiger 134 | g["current_task_is_batch"] = is_batch_func 135 | 136 | hard_timeouts = self.worker.get_hard_timeouts(func, tasks) 137 | 138 | with WorkerContextManagerStack( 139 | self.config["CHILD_CONTEXT_MANAGERS"] 140 | ): 141 | if is_batch_func: 142 | # Batch process if the task supports it. 143 | g["current_tasks"] = tasks 144 | runner.run_batch_tasks(tasks, hard_timeouts[0]) 145 | else: 146 | # Process sequentially. 147 | for task, hard_timeout in zip(tasks, hard_timeouts): 148 | g["current_tasks"] = [task] 149 | runner.run_single_task(task, hard_timeout) 150 | 151 | except RetryException as exc: 152 | execution["retry"] = True 153 | if exc.method: 154 | execution["retry_method"] = serialize_retry_method(exc.method) 155 | execution["log_error"] = exc.log_error 156 | execution["exception_name"] = serialize_func_name(exc.__class__) 157 | exc_info = exc.exc_info or sys.exc_info() 158 | except (JobTimeoutException, Exception) as exc: 159 | execution["exception_name"] = serialize_func_name(exc.__class__) 160 | exc_info = sys.exc_info() 161 | else: 162 | success = True 163 | 164 | if not success: 165 | execution["time_failed"] = time.time() 166 | if self.worker.store_tracebacks: 167 | # Currently we only log failed task executions to Redis. 168 | execution["traceback"] = "".join( 169 | traceback.format_exception(*exc_info) 170 | ) 171 | execution["success"] = success 172 | execution["host"] = socket.gethostname() 173 | 174 | self.worker.store_task_execution(tasks, execution) 175 | 176 | g["current_task_is_batch"] = None 177 | g["current_tasks"] = None 178 | g["tiger"] = None 179 | 180 | return success 181 | 182 | 183 | class ForkExecutor(Executor): 184 | """ 185 | Executor that runs tasks in a forked process. 186 | 187 | Child process is killed after a hard timeout + margin. 188 | """ 189 | 190 | def execute( 191 | self, 192 | queue: str, 193 | tasks: List[Task], 194 | log: BoundLogger, 195 | locks: Collection[Lock], 196 | queue_lock: Optional[Semaphore], 197 | ) -> bool: 198 | task_func = tasks[0].func 199 | serialized_task_func = tasks[0].serialized_func 200 | 201 | all_task_ids = {task.id for task in tasks} 202 | with g_fork_lock: 203 | child_pid = os.fork() 204 | 205 | if child_pid == 0: 206 | # Child process 207 | log = log.bind(child_pid=os.getpid()) 208 | assert isinstance(log, BoundLogger) 209 | 210 | # Disconnect the Redis connection inherited from the main process. 211 | # Note that this doesn't disconnect the socket in the main process. 212 | self.connection.connection_pool.disconnect() 213 | 214 | random.seed() 215 | 216 | # Ignore Ctrl+C in the child so we don't abort the job -- the main 217 | # process already takes care of a graceful shutdown. 218 | signal.signal(signal.SIGINT, signal.SIG_IGN) 219 | 220 | # Run the tasks. 221 | success = self.execute_tasks(tasks, log) 222 | 223 | # Wait for any threads that might be running in the child, just 224 | # like sys.exit() would. Note we don't call sys.exit() directly 225 | # because it would perform additional cleanup (e.g. calling atexit 226 | # handlers twice). See also: https://bugs.python.org/issue18966 227 | threading._shutdown() # type: ignore[attr-defined] 228 | 229 | os._exit(int(not success)) 230 | else: 231 | # Main process 232 | log = log.bind(child_pid=child_pid) 233 | assert isinstance(log, BoundLogger) 234 | for task in tasks: 235 | log.info( 236 | "processing", 237 | func=serialized_task_func, 238 | task_id=task.id, 239 | params={"args": task.args, "kwargs": task.kwargs}, 240 | ) 241 | 242 | # Attach a signal handler to SIGCHLD (sent when the child process 243 | # exits) so we can capture it. 244 | signal.signal(signal.SIGCHLD, sigchld_handler) 245 | 246 | # Since newer Python versions retry interrupted system calls we can't 247 | # rely on the fact that select() is interrupted with EINTR. Instead, 248 | # we'll set up a wake-up file descriptor below. 249 | 250 | # Create a new pipe and apply the non-blocking flag (required for 251 | # set_wakeup_fd). 252 | pipe_r, pipe_w = os.pipe() 253 | 254 | opened_fd = os.fdopen(pipe_r) 255 | flags = fcntl.fcntl(pipe_r, fcntl.F_GETFL, 0) 256 | flags = flags | os.O_NONBLOCK 257 | fcntl.fcntl(pipe_r, fcntl.F_SETFL, flags) 258 | 259 | flags = fcntl.fcntl(pipe_w, fcntl.F_GETFL, 0) 260 | flags = flags | os.O_NONBLOCK 261 | fcntl.fcntl(pipe_w, fcntl.F_SETFL, flags) 262 | 263 | # A byte will be written to pipe_w if a signal occurs (and can be 264 | # read from pipe_r). 265 | old_wakeup_fd = signal.set_wakeup_fd(pipe_w) 266 | 267 | def check_child_exit() -> Optional[int]: 268 | """ 269 | Do a non-blocking check to see if the child process exited. 270 | Returns None if the process is still running, or the exit code 271 | value of the child process. 272 | """ 273 | try: 274 | pid, return_code = os.waitpid(child_pid, os.WNOHANG) 275 | if pid != 0: # The child process is done. 276 | return return_code 277 | except OSError as e: 278 | # Of course EINTR can happen if the child process exits 279 | # while we're checking whether it exited. In this case it 280 | # should be safe to retry. 281 | if e.errno == errno.EINTR: 282 | return check_child_exit() 283 | else: 284 | raise 285 | return None 286 | 287 | hard_timeouts = self.worker.get_hard_timeouts(task_func, tasks) 288 | time_started = time.time() 289 | 290 | # Upper bound for when we expect the child processes to finish. 291 | # Since the hard timeout doesn't cover any processing overhead, 292 | # we're adding an extra buffer of ACTIVE_TASK_UPDATE_TIMEOUT 293 | # (which is the same time we use to determine if a task has 294 | # expired). 295 | timeout_at = ( 296 | time_started 297 | + sum(hard_timeouts) 298 | + self.config["ACTIVE_TASK_UPDATE_TIMEOUT"] 299 | ) 300 | 301 | # Wait for the child to exit and perform a periodic heartbeat. 302 | # We check for the child twice in this loop so that we avoid 303 | # unnecessary waiting if the child exited just before entering 304 | # the while loop or while renewing heartbeat/locks. 305 | while True: 306 | return_code = check_child_exit() 307 | if return_code is not None: 308 | break 309 | 310 | # Wait until the timeout or a signal / child exit occurs. 311 | try: 312 | # If observed the following behavior will be seen 313 | # in the pipe when the parent process receives a 314 | # SIGTERM while a task is running in a child process: 315 | # Linux: 316 | # - 0 when parent receives SIGTERM 317 | # - select() exits with EINTR when child exit 318 | # triggers signal, so the signal in the 319 | # pipe is never seen since check_child_exit() 320 | # will see the child is gone 321 | # 322 | # macOS: 323 | # - 15 (SIGTERM) when parent receives SIGTERM 324 | # - 20 (SIGCHLD) when child exits 325 | results = select.select( 326 | [pipe_r], 327 | [], 328 | [], 329 | self.config["ACTIVE_TASK_UPDATE_TIMER"], 330 | ) 331 | 332 | if results[0]: 333 | # Purge pipe so select will pause on next call 334 | try: 335 | # Behavior of a would be blocking read() 336 | # Linux: 337 | # Python 2.7 Raises OSError 338 | # Python 3.x returns empty string 339 | # 340 | # macOS: 341 | # Returns empty string 342 | opened_fd.read(1) 343 | except OSError: 344 | pass 345 | 346 | except OSError as e: 347 | if e.args[0] != errno.EINTR: 348 | raise 349 | 350 | return_code = check_child_exit() 351 | if return_code is not None: 352 | break 353 | 354 | now = time.time() 355 | if now > timeout_at: 356 | log.error("hard timeout elapsed in parent process") 357 | os.kill(child_pid, signal.SIGKILL) 358 | pid, return_code = os.waitpid(child_pid, 0) 359 | log.error("child killed", return_code=return_code) 360 | execution = { 361 | "time_started": time_started, 362 | "time_failed": now, 363 | "exception_name": serialize_func_name( 364 | JobTimeoutException 365 | ), 366 | "success": False, 367 | "host": socket.gethostname(), 368 | } 369 | self.worker.store_task_execution(tasks, execution) 370 | break 371 | 372 | try: 373 | self.heartbeat(queue, all_task_ids, log, locks, queue_lock) 374 | except OSError as e: 375 | # EINTR happens if the task completed. Since we're just 376 | # renewing locks/heartbeat it's okay if we get interrupted. 377 | if e.errno != errno.EINTR: 378 | raise 379 | 380 | # Restore signals / clean up 381 | signal.signal(signal.SIGCHLD, signal.SIG_DFL) 382 | signal.set_wakeup_fd(old_wakeup_fd) 383 | opened_fd.close() 384 | os.close(pipe_w) 385 | 386 | success = return_code == 0 387 | return success 388 | 389 | 390 | class SyncExecutor(Executor): 391 | """ 392 | Executor that runs tasks in the current thread/process. 393 | """ 394 | 395 | exit_worker_on_job_timeout = True 396 | 397 | def _periodic_heartbeat( 398 | self, 399 | queue: str, 400 | task_ids: Collection[str], 401 | log: BoundLogger, 402 | locks: Collection[Lock], 403 | queue_lock: Optional[Semaphore], 404 | stop_event: threading.Event, 405 | ) -> None: 406 | while not stop_event.wait(self.config["ACTIVE_TASK_UPDATE_TIMER"]): 407 | try: 408 | self.heartbeat(queue, task_ids, log, locks, queue_lock) 409 | except Exception: 410 | log.exception("task heartbeat failed") 411 | 412 | def execute( 413 | self, 414 | queue: str, 415 | tasks: List[Task], 416 | log: BoundLogger, 417 | locks: Collection[Lock], 418 | queue_lock: Optional[Semaphore], 419 | ) -> bool: 420 | assert tasks 421 | 422 | # Run heartbeat thread. 423 | all_task_ids = {task.id for task in tasks} 424 | stop_event = threading.Event() 425 | heartbeat_thread = threading.Thread( 426 | target=self._periodic_heartbeat, 427 | kwargs={ 428 | "queue": queue, 429 | "task_ids": all_task_ids, 430 | "log": log, 431 | "locks": locks, 432 | "queue_lock": queue_lock, 433 | "stop_event": stop_event, 434 | }, 435 | ) 436 | heartbeat_thread.start() 437 | 438 | serialized_task_func = tasks[0].serialized_func 439 | for task in tasks: 440 | log.info( 441 | "processing", 442 | func=serialized_task_func, 443 | task_id=task.id, 444 | params={"args": task.args, "kwargs": task.kwargs}, 445 | ) 446 | 447 | # Run the tasks. 448 | try: 449 | result = self.execute_tasks(tasks, log) 450 | # Always stop the heartbeat thread -- even in case of an unhandled 451 | # exception after running the task code, or when an unhandled 452 | # BaseException is raised from within the task. 453 | finally: 454 | stop_event.set() 455 | heartbeat_thread.join() 456 | 457 | return result 458 | -------------------------------------------------------------------------------- /tasktiger/flask_script.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from typing import TYPE_CHECKING, Any, List 3 | 4 | from flask_script import Command 5 | 6 | if TYPE_CHECKING: 7 | from tasktiger import TaskTiger 8 | 9 | 10 | class TaskTigerCommand(Command): 11 | """ 12 | This class is deprecated and may be removed in future versions. 13 | That is because Flask-Script is no longer supported (since 2017). 14 | """ 15 | 16 | capture_all_args = True 17 | help = "Run a TaskTiger worker" 18 | 19 | def __init__(self, tiger: "TaskTiger") -> None: 20 | super(TaskTigerCommand, self).__init__() 21 | self.tiger = tiger 22 | 23 | def create_parser( 24 | self, *args: Any, **kwargs: Any 25 | ) -> argparse.ArgumentParser: 26 | # Override the default parser so we can pass all arguments to the 27 | # TaskTiger parser. 28 | func_stack = kwargs.pop("func_stack", ()) 29 | parent = kwargs.pop("parent", None) 30 | parser = argparse.ArgumentParser(*args, add_help=False, **kwargs) # type: ignore[misc] 31 | parser.set_defaults(func_stack=func_stack + (self,)) 32 | self.parser = parser 33 | self.parent = parent 34 | return parser 35 | 36 | def setup(self) -> None: 37 | """ 38 | Override this method to implement custom setup (e.g. logging) before 39 | running the worker. 40 | """ 41 | 42 | def run(self, args: List[str]) -> None: 43 | # Allow passing a callable that returns the TaskTiger instance. 44 | if callable(self.tiger): 45 | self.tiger = self.tiger() 46 | self.setup() 47 | self.tiger.run_worker_with_args(args) 48 | -------------------------------------------------------------------------------- /tasktiger/logging.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict 2 | 3 | from ._internal import g 4 | 5 | 6 | def tasktiger_processor( 7 | logger: Any, method_name: Any, event_dict: Dict[str, Any] 8 | ) -> Dict[str, Any]: 9 | """ 10 | TaskTiger structlog processor. 11 | 12 | Inject the current task ID and queue for non-batch tasks. 13 | """ 14 | 15 | if g["current_tasks"] is not None and not g["current_task_is_batch"]: 16 | current_task = g["current_tasks"][0] 17 | event_dict["task_id"] = current_task.id 18 | event_dict["task_func"] = current_task.serialized_func 19 | event_dict["queue"] = current_task.queue 20 | 21 | return event_dict 22 | -------------------------------------------------------------------------------- /tasktiger/lua/move_task.lua: -------------------------------------------------------------------------------- 1 | local function zadd_w_mode(key, score, member, mode) 2 | if mode == "" then 3 | redis.call('zadd', key, score, member) 4 | elseif mode == "nx" then 5 | zadd_noupdate({ key }, { score, member }) 6 | elseif mode == "min" then 7 | zadd_update_min({ key }, { score, member }) 8 | else 9 | error("mode " .. mode .. " unsupported") 10 | end 11 | end 12 | 13 | 14 | local key_task_id = KEYS[1] 15 | local key_task_id_executions = KEYS[2] 16 | local key_task_id_executions_count = KEYS[3] 17 | local key_from_state = KEYS[4] 18 | local key_to_state = KEYS[5] 19 | local key_active_queue = KEYS[6] 20 | local key_queued_queue = KEYS[7] 21 | local key_error_queue = KEYS[8] 22 | local key_scheduled_queue = KEYS[9] 23 | local key_activity = KEYS[10] 24 | 25 | local id = ARGV[1] 26 | local queue = ARGV[2] 27 | local from_state = ARGV[3] 28 | local to_state = ARGV[4] 29 | local unique = ARGV[5] 30 | local when = ARGV[6] 31 | local mode = ARGV[7] 32 | local publish_queued_tasks = ARGV[8] 33 | 34 | local state_queues_keys_by_state = { 35 | active = key_active_queue, 36 | queued = key_queued_queue, 37 | error = key_error_queue, 38 | scheduled = key_scheduled_queue, 39 | } 40 | local key_from_state_queue = state_queues_keys_by_state[from_state] 41 | local key_to_state_queue = state_queues_keys_by_state[to_state] 42 | 43 | assert(redis.call('zscore', key_from_state_queue, id), '') 44 | 45 | if to_state ~= "" then 46 | zadd_w_mode(key_to_state_queue, when, id, mode) 47 | redis.call('sadd', key_to_state, queue) 48 | end 49 | redis.call('zrem', key_from_state_queue, id) 50 | 51 | if to_state == "" then -- Remove the task if necessary 52 | if unique == 'true' then 53 | -- Delete executions if there were no errors 54 | local to_delete = { 55 | key_task_id_executions, 56 | key_task_id_executions_count, 57 | } 58 | local keys = { unpack(to_delete) } 59 | if from_state ~= 'error' then 60 | table.insert(keys, key_error_queue) 61 | end 62 | -- keys=[to_delete + zsets], args=[len(to_delete), value] 63 | delete_if_not_in_zsets(keys, { #to_delete, id }) 64 | 65 | -- Only delete task if it's not in any other queue 66 | local to_delete = { key_task_id } 67 | local zsets = {} 68 | for i, v in pairs({ 'active', 'queued', 'error', 'scheduled' }) do 69 | if v ~= from_state then 70 | table.insert(zsets, state_queues_keys_by_state[v]) 71 | end 72 | end 73 | -- keys=[to_delete + zsets], args=[len(to_delete), value] 74 | delete_if_not_in_zsets({ unpack(to_delete), unpack(zsets) }, { #to_delete, id }) 75 | else 76 | -- Safe to remove 77 | redis.call( 78 | 'del', 79 | key_task_id, 80 | key_task_id_executions, 81 | key_task_id_executions_count 82 | ) 83 | end 84 | end 85 | 86 | -- keys=[key, other_key], args=[member] 87 | srem_if_not_exists({ key_from_state, key_from_state_queue }, { queue }) 88 | 89 | if to_state == 'queued' and publish_queued_tasks == 'true' then 90 | redis.call('publish', key_activity, queue) 91 | end 92 | -------------------------------------------------------------------------------- /tasktiger/lua/semaphore.lua: -------------------------------------------------------------------------------- 1 | -- Semaphore lock 2 | -- 3 | -- KEYS = { semaphore key } 4 | -- ARGV = { lock id (must be unique across callers), 5 | -- semaphore size, 6 | -- lock timeout in seconds, 7 | -- current time in seconds } 8 | -- 9 | -- Returns: { lock acquired (True or False), 10 | -- number of locks in semaphore } 11 | 12 | -- Using redis server time ensures consistent time across callers but it requires 13 | -- Redis v3.2+ because the replicate_commands command is needed to replicate 14 | -- the actions in this script to a slave instance. 15 | --redis.replicate_commands() 16 | 17 | local semaphore_key = KEYS[1] 18 | 19 | local lock_id = ARGV[1] 20 | local semaphore_size = tonumber(ARGV[2]) 21 | local timeout = tonumber(ARGV[3]) 22 | -- TODO: Dynamically enable Redis time usage in Redis 3.2+ 23 | --local time = redis.call("time") 24 | --local now = tonumber(time[1]) 25 | local now = tonumber(ARGV[4]) 26 | 27 | --Remove expired locks 28 | redis.call("ZREMRANGEBYSCORE", semaphore_key, 0, now) 29 | 30 | --Check if there is a system lock which will override all other locks 31 | if redis.call("ZSCORE", semaphore_key, "SYSTEM_LOCK") ~= false then 32 | return {false, -1} 33 | end 34 | 35 | --Update TTL for semaphore key. This is done after checking 36 | --for a system lock so we don't accidently make the TTL 37 | --less than the system lock timeout. 38 | redis.call("EXPIRE", semaphore_key, math.ceil(timeout * 2)) 39 | 40 | -- Get current count of active locks 41 | local lock_count = redis.call("ZCARD", semaphore_key) 42 | 43 | -- Check if this lock_id already has an active lock 44 | local updating_lock = redis.call("ZSCORE", semaphore_key, lock_id) 45 | 46 | -- The lock count will increase by 1 if we are getting a new lock 47 | local current_lock_count = lock_count 48 | if updating_lock == false then 49 | lock_count = lock_count + 1 50 | end 51 | 52 | -- Check if we should allow this lock 53 | if lock_count > semaphore_size then 54 | return {false, current_lock_count} 55 | else 56 | -- This also handles renewing an existing lock 57 | -- Score is set to the time this lock expires 58 | redis.call("ZADD", semaphore_key, now + timeout, lock_id) 59 | return {true, lock_count} 60 | end 61 | -------------------------------------------------------------------------------- /tasktiger/migrations.py: -------------------------------------------------------------------------------- 1 | from typing import TYPE_CHECKING 2 | 3 | from .utils import redis_glob_escape 4 | 5 | if TYPE_CHECKING: 6 | from . import TaskTiger 7 | 8 | 9 | def migrate_executions_count(tiger: "TaskTiger") -> None: 10 | """ 11 | Backfills ``t:task::executions_count`` by counting 12 | elements in ``t:task::executions``. 13 | """ 14 | 15 | migrate_task = tiger.connection.register_script( 16 | """ 17 | local count = redis.call('llen', KEYS[1]) 18 | if tonumber(redis.call('get', KEYS[2]) or 0) < count then 19 | redis.call('set', KEYS[2], count) 20 | end 21 | """ 22 | ) 23 | 24 | match = ( 25 | redis_glob_escape(tiger.config["REDIS_PREFIX"]) + ":task:*:executions" 26 | ) 27 | 28 | for key in tiger.connection.scan_iter(count=100, match=match): 29 | migrate_task(keys=[key, key + "_count"]) 30 | -------------------------------------------------------------------------------- /tasktiger/redis_scripts.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Any, Callable, List, Literal, Optional, Tuple, Union 3 | 4 | from redis import Redis 5 | 6 | from ._internal import ACTIVE, ERROR, QUEUED, SCHEDULED 7 | 8 | try: 9 | from redis.commands.core import Script 10 | except ImportError: 11 | # redis-py < 4.0 : https://github.com/redis/redis-py/pull/1534 12 | from redis.client import Script # type: ignore 13 | 14 | 15 | LOCAL_FUNC_TEMPLATE = """ 16 | local function {func_name}(KEYS, ARGV) 17 | {func_body} 18 | end 19 | 20 | """ 21 | 22 | # ARGV = { score, member } 23 | ZADD_NOUPDATE_TEMPLATE = """ 24 | if {condition} redis.call('zscore', {key}, {member}) then 25 | redis.call('zadd', {key}, {score}, {member}) 26 | end 27 | """ 28 | ZADD_NOUPDATE = ZADD_NOUPDATE_TEMPLATE.format( 29 | key="KEYS[1]", score="ARGV[1]", member="ARGV[2]", condition="not" 30 | ) 31 | ZADD_UPDATE_EXISTING = ZADD_NOUPDATE_TEMPLATE.format( 32 | key="KEYS[1]", score="ARGV[1]", member="ARGV[2]", condition="" 33 | ) 34 | ZADD_UPDATE_TEMPLATE = """ 35 | local score = redis.call('zscore', {key}, {member}) 36 | local new_score 37 | if score then 38 | new_score = math.{f}(score, {score}) 39 | else 40 | new_score = {score} 41 | end 42 | {ret} redis.call('zadd', {key}, new_score, {member}) 43 | """ 44 | ZADD_UPDATE_MIN = ZADD_UPDATE_TEMPLATE.format( 45 | f="min", key="KEYS[1]", score="ARGV[1]", member="ARGV[2]", ret="return" 46 | ) 47 | ZADD_UPDATE_MAX = ZADD_UPDATE_TEMPLATE.format( 48 | f="max", key="KEYS[1]", score="ARGV[1]", member="ARGV[2]", ret="return" 49 | ) 50 | 51 | _ZPOPPUSH_EXISTS_TEMPLATE = """ 52 | -- Load keys and arguments 53 | local source = KEYS[1] 54 | local destination = KEYS[2] 55 | local remove_from_set = KEYS[3] 56 | local add_to_set = KEYS[4] 57 | local add_to_set_if_exists = KEYS[5] 58 | local if_exists_key = KEYS[6] 59 | 60 | local score = ARGV[1] 61 | local count = ARGV[2] 62 | local new_score = ARGV[3] 63 | local set_value = ARGV[4] 64 | local if_exists_score = ARGV[5] 65 | 66 | -- Fetch affected members from the source set. 67 | local members = redis.call('zrangebyscore', source, '-inf', score, 'LIMIT', 0, count) 68 | 69 | -- Tables to keep track of the members that we're moving to the destination 70 | -- (moved_members), along with their new scores (new_scoremembers), and 71 | -- members that already exist at the destination (existing_members). 72 | local new_scoremembers = {{}} 73 | local moved_members = {{}} 74 | local existing_members = {{}} 75 | 76 | -- Counters so we can quickly append to the tables 77 | local existing_idx = 0 78 | local moved_idx = 0 79 | 80 | -- Populate the tables defined above. 81 | for i, member in ipairs(members) do 82 | if redis.call('zscore', destination, member) then 83 | existing_idx = existing_idx + 1 84 | existing_members[existing_idx] = member 85 | else 86 | moved_idx = moved_idx + 1 87 | new_scoremembers[2*moved_idx] = member 88 | new_scoremembers[2*moved_idx-1] = new_score 89 | moved_members[moved_idx] = member 90 | end 91 | end 92 | 93 | if #members > 0 then 94 | -- If we matched any members, remove them from the source. 95 | redis.call('zremrangebyrank', source, 0, #members-1) 96 | 97 | -- Add members to the destination. 98 | if #new_scoremembers > 0 then 99 | redis.call('zadd', destination, unpack(new_scoremembers)) 100 | end 101 | 102 | -- Perform the "if exists" action for members that exist at the 103 | -- destination. 104 | for i, member in ipairs(existing_members) do 105 | {if_exists_template} 106 | end 107 | 108 | -- Perform any "on success" action. 109 | {on_success} 110 | 111 | -- If we moved any members to the if_exists_key, add the set_value 112 | -- to the add_to_set_if_exists (see zpoppush docstring). 113 | if if_exists_key then 114 | local if_exists_key_exists = redis.call('exists', if_exists_key) 115 | if if_exists_key_exists == 1 then 116 | redis.call('sadd', add_to_set_if_exists, set_value) 117 | end 118 | end 119 | 120 | end 121 | 122 | -- Return just the moved members. 123 | return moved_members 124 | """ 125 | 126 | _ON_SUCCESS_UPDATE_SETS_TEMPLATE = """ 127 | local src_exists = redis.call('exists', source) 128 | if src_exists == 0 then 129 | redis.call('srem', {remove_from_set}, {set_value}) 130 | end 131 | redis.call('sadd', {add_to_set}, {set_value}) 132 | """ 133 | 134 | # KEYS = { source, destination, remove_from_set, add_to_set, 135 | # add_to_set_if_exists, if_exists_key } 136 | # ARGV = { score, count, new_score, set_value, if_exists_score } 137 | ZPOPPUSH_EXISTS_MIN_UPDATE_SETS = _ZPOPPUSH_EXISTS_TEMPLATE.format( 138 | if_exists_template=ZADD_UPDATE_TEMPLATE.format( 139 | f="min", 140 | key="if_exists_key", 141 | score="if_exists_score", 142 | member="member", 143 | ret="", 144 | ), 145 | on_success=_ON_SUCCESS_UPDATE_SETS_TEMPLATE.format( 146 | set_value="set_value", 147 | add_to_set="add_to_set", 148 | remove_from_set="remove_from_set", 149 | ), 150 | ) 151 | 152 | # KEYS = { source, destination, remove_from_set, add_to_set } 153 | # ARGV = { score, count, new_score, set_value } 154 | ZPOPPUSH_EXISTS_IGNORE_UPDATE_SETS = _ZPOPPUSH_EXISTS_TEMPLATE.format( 155 | if_exists_template="", 156 | on_success=_ON_SUCCESS_UPDATE_SETS_TEMPLATE.format( 157 | set_value="set_value", 158 | add_to_set="add_to_set", 159 | remove_from_set="remove_from_set", 160 | ), 161 | ) 162 | 163 | # KEYS = { source, destination, ... } 164 | # ARGV = { score, count, new_score, ... } 165 | _ZPOPPUSH_TEMPLATE = """ 166 | -- Load keys and arguments 167 | local source = KEYS[1] 168 | local destination = KEYS[2] 169 | 170 | local score = ARGV[1] 171 | local count = ARGV[2] 172 | local new_score = ARGV[3] 173 | 174 | -- Fetch affected members from the source set. 175 | local members = redis.call('zrangebyscore', source, '-inf', score, 'LIMIT', 0, count) 176 | 177 | -- Table to keep track of the members along with their new scores, which is 178 | -- passed to ZADD. 179 | local new_scoremembers = {{}} 180 | for i, member in ipairs(members) do 181 | new_scoremembers[2*i] = member 182 | new_scoremembers[2*i-1] = new_score 183 | end 184 | 185 | if #members > 0 then 186 | -- Remove affected members and add them to the destination. 187 | redis.call('zremrangebyrank', source, 0, #members-1) 188 | redis.call('zadd', destination, unpack(new_scoremembers)) 189 | 190 | -- Perform any "on success" action. 191 | {on_success} 192 | end 193 | 194 | -- Return moved members 195 | return members 196 | """ 197 | 198 | ZPOPPUSH = _ZPOPPUSH_TEMPLATE.format(on_success="") 199 | 200 | # KEYS = { source, destination, remove_from_set, add_to_set } 201 | # ARGV = { score, count, new_score, set_value } 202 | ZPOPPUSH_UPDATE_SETS = _ZPOPPUSH_TEMPLATE.format( 203 | on_success=_ON_SUCCESS_UPDATE_SETS_TEMPLATE.format( 204 | set_value="ARGV[4]", add_to_set="KEYS[4]", remove_from_set="KEYS[3]" 205 | ) 206 | ) 207 | 208 | # ARGV = { score, count, new_score } 209 | ZPOPPUSH_WITHSCORES = """ 210 | local members_scores = redis.call('zrangebyscore', KEYS[1], '-inf', ARGV[1], 'WITHSCORES', 'LIMIT', 0, ARGV[2]) 211 | local new_scoremembers = {} 212 | for i, member in ipairs(members_scores) do 213 | if i % 2 == 1 then -- Just the members (1, 3, 5, ...) 214 | -- Insert scores at 1, 3, 5, ... 215 | new_scoremembers[i] = ARGV[3] 216 | 217 | -- Insert members at 2, 4, 6, ... 218 | new_scoremembers[i+1] = member 219 | end 220 | end 221 | if #members_scores > 0 then 222 | redis.call('zremrangebyrank', KEYS[1], 0, #members_scores/2-1) 223 | redis.call('zadd', KEYS[2], unpack(new_scoremembers)) 224 | end 225 | return members_scores 226 | """ 227 | 228 | # KEYS = { key, other_key } 229 | # ARGV = { member } 230 | SREM_IF_NOT_EXISTS = """ 231 | local exists = redis.call('exists', KEYS[2]) 232 | local result 233 | if exists == 0 then 234 | result = redis.call('srem', KEYS[1], ARGV[1]) 235 | else 236 | result = 0 237 | end 238 | return result 239 | """ 240 | 241 | # KEYS = { del1, [, ..., delN], zset1 [, ..., zsetN] } 242 | # ARGV = { to_delete_count, value } 243 | DELETE_IF_NOT_IN_ZSETS = """ 244 | local found = 0 245 | for i=ARGV[1] + 1,#KEYS do 246 | if redis.call('zscore', KEYS[i], ARGV[2]) then 247 | found = 1 248 | break 249 | end 250 | end 251 | if found == 0 then 252 | return redis.call('del', unpack(KEYS, 1, ARGV[1])) 253 | end 254 | return 0 255 | """ 256 | 257 | # KEYS = { key } 258 | # ARGV = { member } 259 | FAIL_IF_NOT_IN_ZSET = """ 260 | assert(redis.call('zscore', KEYS[1], ARGV[1]), '') 261 | """ 262 | 263 | # KEYS = { } 264 | # ARGV = { key_prefix, time, batch_size } 265 | GET_EXPIRED_TASKS = """ 266 | local key_prefix = ARGV[1] 267 | local time = ARGV[2] 268 | local batch_size = ARGV[3] 269 | local active_queues = redis.call('smembers', key_prefix .. ':' .. 'active') 270 | local result = {} 271 | local result_n = 1 272 | 273 | for i=1, #active_queues do 274 | local queue_name = active_queues[i] 275 | local queue_key = key_prefix .. ':' .. 'active' .. 276 | ':' .. queue_name 277 | 278 | local members = redis.call('zrangebyscore', 279 | queue_key, 0, time, 'LIMIT', 0, batch_size) 280 | 281 | for j=1, #members do 282 | result[result_n] = queue_name 283 | result[result_n + 1] = members[j] 284 | result_n = result_n + 2 285 | end 286 | 287 | batch_size = batch_size - #members 288 | if batch_size <= 0 then 289 | break 290 | end 291 | end 292 | 293 | return result 294 | """ 295 | 296 | 297 | class RedisScripts: 298 | def __init__(self, redis: Redis) -> None: 299 | self.redis = redis 300 | 301 | self._zadd_noupdate = redis.register_script(ZADD_NOUPDATE) 302 | self._zadd_update_existing = redis.register_script( 303 | ZADD_UPDATE_EXISTING 304 | ) 305 | self._zadd_update_min = redis.register_script(ZADD_UPDATE_MIN) 306 | self._zadd_update_max = redis.register_script(ZADD_UPDATE_MAX) 307 | 308 | self._zpoppush = redis.register_script(ZPOPPUSH) 309 | self._zpoppush_update_sets = redis.register_script( 310 | ZPOPPUSH_UPDATE_SETS 311 | ) 312 | self._zpoppush_withscores = redis.register_script(ZPOPPUSH_WITHSCORES) 313 | self._zpoppush_exists_min_update_sets = redis.register_script( 314 | ZPOPPUSH_EXISTS_MIN_UPDATE_SETS 315 | ) 316 | self._zpoppush_exists_ignore_update_sets = redis.register_script( 317 | ZPOPPUSH_EXISTS_IGNORE_UPDATE_SETS 318 | ) 319 | 320 | self._srem_if_not_exists = redis.register_script(SREM_IF_NOT_EXISTS) 321 | 322 | self._delete_if_not_in_zsets = redis.register_script( 323 | DELETE_IF_NOT_IN_ZSETS 324 | ) 325 | 326 | self._fail_if_not_in_zset = redis.register_script(FAIL_IF_NOT_IN_ZSET) 327 | 328 | self._get_expired_tasks = redis.register_script(GET_EXPIRED_TASKS) 329 | 330 | self._move_task = self.register_script_from_file( 331 | "lua/move_task.lua", 332 | include_functions={ 333 | "zadd_noupdate": ZADD_NOUPDATE, 334 | "zadd_update_min": ZADD_UPDATE_MIN, 335 | "srem_if_not_exists": SREM_IF_NOT_EXISTS, 336 | "delete_if_not_in_zsets": DELETE_IF_NOT_IN_ZSETS, 337 | }, 338 | ) 339 | 340 | @property 341 | def can_replicate_commands(self) -> bool: 342 | """ 343 | Whether Redis supports single command replication. 344 | """ 345 | if not hasattr(self, "_can_replicate_commands"): 346 | info = self.redis.info("server") 347 | version_info = info["redis_version"].split(".") 348 | major, minor = int(version_info[0]), int(version_info[1]) 349 | result = major > 3 or major == 3 and minor >= 2 350 | self._can_replicate_commands = result 351 | return self._can_replicate_commands 352 | 353 | def register_script_from_file( 354 | self, filename: str, include_functions: Optional[dict] = None 355 | ) -> Script: 356 | with open( 357 | os.path.join(os.path.dirname(os.path.realpath(__file__)), filename) 358 | ) as f: 359 | script = f.read() 360 | if include_functions: 361 | function_definitions = [] 362 | for func_name in sorted(include_functions.keys()): 363 | function_definitions.append( 364 | LOCAL_FUNC_TEMPLATE.format( 365 | func_name=func_name, 366 | func_body=include_functions[func_name], 367 | ) 368 | ) 369 | script = "\n".join(function_definitions + [script]) 370 | 371 | return self.redis.register_script(script) 372 | 373 | def zadd( 374 | self, 375 | key: str, 376 | score: float, 377 | member: str, 378 | mode: str, 379 | client: Optional[Redis] = None, 380 | ) -> int: 381 | """ 382 | Like ZADD, but supports different score update modes, in case the 383 | member already exists in the ZSET: 384 | - "nx": Don't update the score 385 | - "xx": Only update elements that already exist. Never add elements. 386 | - "min": Use the smaller of the given and existing score 387 | - "max": Use the larger of the given and existing score 388 | """ 389 | if mode == "nx": 390 | f = self._zadd_noupdate 391 | elif mode == "xx": 392 | f = self._zadd_update_existing 393 | elif mode == "min": 394 | f = self._zadd_update_min 395 | elif mode == "max": 396 | f = self._zadd_update_max 397 | else: 398 | raise NotImplementedError('mode "%s" unsupported' % mode) 399 | return f(keys=[key], args=[score, member], client=client) 400 | 401 | def zpoppush( 402 | self, 403 | source: str, 404 | destination: str, 405 | count: int, 406 | score: Optional[Union[float, Literal["+inf"]]], 407 | new_score: float, 408 | client: Optional[Redis] = None, 409 | withscores: bool = False, 410 | on_success: Any = None, 411 | if_exists: Any = None, 412 | ) -> Any: 413 | """ 414 | Pops the first ``count`` members from the ZSET ``source`` and adds them 415 | to the ZSET ``destination`` with a score of ``new_score``. If ``score`` 416 | is not None, only members up to a score of ``score`` are used. Returns 417 | the members that were moved and, if ``withscores`` is True, their 418 | original scores. 419 | 420 | If items were moved, the action defined in ``on_success`` is executed. 421 | The only implemented option is a tuple in the form ('update_sets', 422 | ``set_value``, ``remove_from_set``, ``add_to_set`` 423 | [, ``add_to_set_if_exists``]). 424 | If no items are left in the ``source`` ZSET, the ``set_value`` is 425 | removed from ``remove_from_set``. If any items were moved to the 426 | ``destination`` ZSET, the ``set_value`` is added to ``add_to_set``. If 427 | any items were moved to the ``if_exists_key`` ZSET (see below), the 428 | ``set_value`` is added to the ``add_to_set_if_exists`` set. 429 | 430 | If ``if_exists`` is specified as a tuple ('add', if_exists_key, 431 | if_exists_score, if_exists_mode), then members that are already in the 432 | ``destination`` set will not be returned or updated, but they will be 433 | added to a ZSET ``if_exists_key`` with a score of ``if_exists_score`` 434 | and the given behavior specified in ``if_exists_mode`` for members that 435 | already exist in the ``if_exists_key`` ZSET. ``if_exists_mode`` can be 436 | one of the following: 437 | - "nx": Don't update the score 438 | - "min": Use the smaller of the given and existing score 439 | - "max": Use the larger of the given and existing score 440 | 441 | If ``if_exists`` is specified as a tuple ('noupdate',), then no action 442 | will be taken for members that are already in the ``destination`` ZSET 443 | (their score will not be updated). 444 | """ 445 | if score is None: 446 | score = "+inf" # Include all elements. 447 | if withscores: 448 | if on_success: 449 | raise NotImplementedError() 450 | return self._zpoppush_withscores( 451 | keys=[source, destination], 452 | args=[score, count, new_score], 453 | client=client, 454 | ) 455 | else: 456 | if if_exists and if_exists[0] == "add": 457 | _, if_exists_key, if_exists_score, if_exists_mode = if_exists 458 | if if_exists_mode != "min": 459 | raise NotImplementedError() 460 | 461 | if not on_success or on_success[0] != "update_sets": 462 | raise NotImplementedError() 463 | ( 464 | set_value, 465 | remove_from_set, 466 | add_to_set, 467 | add_to_set_if_exists, 468 | ) = on_success[1:] 469 | 470 | return self._zpoppush_exists_min_update_sets( 471 | keys=[ 472 | source, 473 | destination, 474 | remove_from_set, 475 | add_to_set, 476 | add_to_set_if_exists, 477 | if_exists_key, 478 | ], 479 | args=[score, count, new_score, set_value, if_exists_score], 480 | ) 481 | elif if_exists and if_exists[0] == "noupdate": 482 | if not on_success or on_success[0] != "update_sets": 483 | raise NotImplementedError() 484 | set_value, remove_from_set, add_to_set = on_success[1:] 485 | 486 | return self._zpoppush_exists_ignore_update_sets( 487 | keys=[source, destination, remove_from_set, add_to_set], 488 | args=[score, count, new_score, set_value], 489 | ) 490 | 491 | if on_success: 492 | if on_success[0] != "update_sets": 493 | raise NotImplementedError() 494 | else: 495 | set_value, remove_from_set, add_to_set = on_success[1:] 496 | return self._zpoppush_update_sets( 497 | keys=[ 498 | source, 499 | destination, 500 | remove_from_set, 501 | add_to_set, 502 | ], 503 | args=[score, count, new_score, set_value], 504 | client=client, 505 | ) 506 | else: 507 | return self._zpoppush( 508 | keys=[source, destination], 509 | args=[score, count, new_score], 510 | client=client, 511 | ) 512 | 513 | def srem_if_not_exists( 514 | self, 515 | key: str, 516 | member: str, 517 | other_key: str, 518 | client: Optional[Redis] = None, 519 | ) -> int: 520 | """ 521 | Removes ``member`` from the set ``key`` if ``other_key`` does not 522 | exist (i.e. is empty). Returns the number of removed elements (0 or 1). 523 | """ 524 | return self._srem_if_not_exists( 525 | keys=[key, other_key], args=[member], client=client 526 | ) 527 | 528 | def delete_if_not_in_zsets( 529 | self, 530 | to_delete: List[str], 531 | value: str, 532 | zsets: List[str], 533 | client: Optional[Redis] = None, 534 | ) -> int: 535 | """ 536 | Removes keys in ``to_delete`` only if ``value`` is not a member of any 537 | sorted sets in ``zsets``. Returns the number of removed elements. 538 | """ 539 | return self._delete_if_not_in_zsets( 540 | keys=to_delete + zsets, 541 | args=[len(to_delete), value], 542 | client=client, 543 | ) 544 | 545 | def fail_if_not_in_zset( 546 | self, key: str, member: str, client: Optional[Redis] = None 547 | ) -> None: 548 | """ 549 | Fails with an error containing the string '' if 550 | the given ``member`` is not in the ZSET ``key``. This can be used in 551 | a pipeline to assert that the member is in the ZSET and cancel the 552 | execution otherwise. 553 | """ 554 | self._fail_if_not_in_zset(keys=[key], args=[member], client=client) 555 | 556 | def get_expired_tasks( 557 | self, 558 | key_prefix: str, 559 | time: float, 560 | batch_size: int, 561 | client: Optional[Redis] = None, 562 | ) -> List[Tuple[str, str]]: 563 | """ 564 | Returns a list of expired tasks (older than ``time``) by looking at all 565 | active queues. The list is capped at ``batch_size``. The list contains 566 | tuples (queue, task_id). 567 | """ 568 | result = self._get_expired_tasks( 569 | args=[key_prefix, time, batch_size], client=client 570 | ) 571 | 572 | # [queue1, task1, queue2, task2] -> [(queue1, task1), (queue2, task2)] 573 | return list(zip(result[::2], result[1::2])) 574 | 575 | def move_task( 576 | self, 577 | id: str, 578 | queue: str, 579 | from_state: str, 580 | to_state: Optional[str], 581 | unique: bool, 582 | when: float, 583 | mode: Optional[str], 584 | key_func: Callable[..., str], 585 | publish_queued_tasks: bool, 586 | client: Optional[Redis] = None, 587 | ) -> Any: 588 | """ 589 | Refer to task._move internal helper documentation. 590 | """ 591 | 592 | def _bool_to_str(v: bool) -> str: 593 | return "true" if v else "false" 594 | 595 | def _none_to_empty_str(v: Optional[str]) -> str: 596 | return v or "" 597 | 598 | key_task_id = key_func("task", id) 599 | key_task_id_executions = key_func("task", id, "executions") 600 | key_task_id_executions_count = key_func("task", id, "executions_count") 601 | key_from_state = key_func(from_state) 602 | key_to_state = key_func(to_state) if to_state else "" 603 | key_active_queue = key_func(ACTIVE, queue) 604 | key_queued_queue = key_func(QUEUED, queue) 605 | key_error_queue = key_func(ERROR, queue) 606 | key_scheduled_queue = key_func(SCHEDULED, queue) 607 | key_activity = key_func("activity") 608 | 609 | return self._move_task( 610 | keys=[ 611 | key_task_id, 612 | key_task_id_executions, 613 | key_task_id_executions_count, 614 | key_from_state, 615 | key_to_state, 616 | key_active_queue, 617 | key_queued_queue, 618 | key_error_queue, 619 | key_scheduled_queue, 620 | key_activity, 621 | ], 622 | args=[ 623 | id, 624 | queue, 625 | from_state, 626 | _none_to_empty_str(to_state), 627 | _bool_to_str(unique), 628 | when, 629 | _none_to_empty_str(mode), 630 | _bool_to_str(publish_queued_tasks), 631 | ], 632 | client=client, 633 | ) 634 | -------------------------------------------------------------------------------- /tasktiger/redis_semaphore.py: -------------------------------------------------------------------------------- 1 | """Redis Semaphore lock.""" 2 | 3 | import os 4 | import time 5 | from typing import Optional, Tuple 6 | 7 | from redis import Redis 8 | 9 | SYSTEM_LOCK_ID = "SYSTEM_LOCK" 10 | 11 | 12 | class Semaphore: 13 | """Semaphore lock using Redis ZSET.""" 14 | 15 | def __init__( 16 | self, 17 | redis: Redis, 18 | name: str, 19 | lock_id: str, 20 | timeout: float, 21 | max_locks: int = 1, 22 | ) -> None: 23 | """ 24 | Semaphore lock. 25 | 26 | Semaphore logic is implemented in the lua/semaphore.lua script. 27 | Individual locks within the semaphore are managed inside a ZSET 28 | using scores to track when they expire. 29 | 30 | Arguments: 31 | redis: Redis client 32 | name: Name of lock. Used as ZSET key. 33 | lock_id: Lock ID 34 | timeout: Timeout in seconds 35 | max_locks: Maximum number of locks allowed for this semaphore 36 | """ 37 | 38 | self.redis = redis 39 | self.name = name 40 | self.lock_id = lock_id 41 | self.max_locks = max_locks 42 | self.timeout = timeout 43 | with open( 44 | os.path.join( 45 | os.path.dirname(os.path.realpath(__file__)), 46 | "lua/semaphore.lua", 47 | ) 48 | ) as f: 49 | self._semaphore = self.redis.register_script(f.read()) 50 | 51 | @classmethod 52 | def get_system_lock(cls, redis: Redis, name: str) -> Optional[float]: 53 | """ 54 | Get system lock timeout for the semaphore. 55 | 56 | Arguments: 57 | redis: Redis client 58 | name: Name of lock. Used as ZSET key. 59 | 60 | Returns: Time system lock expires or None if lock does not exist 61 | """ 62 | 63 | return redis.zscore(name, SYSTEM_LOCK_ID) 64 | 65 | @classmethod 66 | def set_system_lock(cls, redis: Redis, name: str, timeout: int) -> None: 67 | """ 68 | Set system lock for the semaphore. 69 | 70 | Sets a system lock that will expire in timeout seconds. This 71 | overrides all other locks. Existing locks cannot be renewed 72 | and no new locks will be permitted until the system lock 73 | expires. 74 | 75 | Arguments: 76 | redis: Redis client 77 | name: Name of lock. Used as ZSET key. 78 | timeout: Timeout in seconds for system lock 79 | """ 80 | 81 | pipeline = redis.pipeline() 82 | pipeline.zadd(name, {SYSTEM_LOCK_ID: time.time() + timeout}) 83 | pipeline.expire( 84 | name, timeout + 10 85 | ) # timeout plus buffer for troubleshooting 86 | pipeline.execute() 87 | 88 | def release(self) -> None: 89 | """Release semaphore.""" 90 | 91 | self.redis.zrem(self.name, self.lock_id) 92 | 93 | def acquire(self) -> Tuple[bool, int]: 94 | """ 95 | Obtain a semaphore lock. 96 | 97 | Returns: Tuple that contains True/False if the lock was acquired and number of 98 | locks in semaphore. 99 | """ 100 | 101 | acquired, locks = self._semaphore( 102 | keys=[self.name], 103 | args=[self.lock_id, self.max_locks, self.timeout, time.time()], 104 | ) 105 | 106 | # Convert Lua boolean returns to Python booleans 107 | acquired = True if acquired == 1 else False 108 | 109 | return acquired, locks 110 | 111 | def renew(self) -> Tuple[bool, int]: 112 | """ 113 | Attempt to renew semaphore. 114 | 115 | Technically this doesn't know the difference between losing the lock 116 | but then successfully getting a new lock versus renewing your lock 117 | before the timeout. Both will return True. 118 | """ 119 | 120 | return self.acquire() 121 | -------------------------------------------------------------------------------- /tasktiger/retry.py: -------------------------------------------------------------------------------- 1 | # The retry logic is documented in the README. 2 | from .exceptions import StopRetry 3 | from .types import RetryStrategy 4 | 5 | 6 | def _fixed(retry: int, delay: float, max_retries: int) -> float: 7 | if retry > max_retries: 8 | raise StopRetry() 9 | return delay 10 | 11 | 12 | def fixed(delay: float, max_retries: int) -> RetryStrategy: 13 | return (_fixed, (delay, max_retries)) 14 | 15 | 16 | def _linear( 17 | retry: int, delay: float, increment: float, max_retries: int 18 | ) -> float: 19 | if retry > max_retries: 20 | raise StopRetry() 21 | return delay + increment * (retry - 1) 22 | 23 | 24 | def linear(delay: float, increment: float, max_retries: int) -> RetryStrategy: 25 | return (_linear, (delay, increment, max_retries)) 26 | 27 | 28 | def _exponential( 29 | retry: int, delay: float, factor: float, max_retries: int 30 | ) -> float: 31 | if retry > max_retries: 32 | raise StopRetry() 33 | return delay * factor ** (retry - 1) 34 | 35 | 36 | def exponential( 37 | delay: float, factor: float, max_retries: int 38 | ) -> RetryStrategy: 39 | return (_exponential, (delay, factor, max_retries)) 40 | -------------------------------------------------------------------------------- /tasktiger/rollbar.py: -------------------------------------------------------------------------------- 1 | import json 2 | from typing import Any 3 | 4 | import rollbar 5 | from rollbar.logger import RollbarHandler 6 | 7 | 8 | class StructlogRollbarHandler(RollbarHandler): 9 | def __init__(self, prefix: str, *args: Any, **kwargs: Any): 10 | """ 11 | Structured rollbar handler. Rollbar messages are prefixed with the 12 | given prefix string. Any other arguments are passed to RollbarHandler. 13 | """ 14 | self.prefix = prefix 15 | super(StructlogRollbarHandler, self).__init__(*args, **kwargs) 16 | 17 | def format_title(self, data: Any) -> str: 18 | # Keys used to construct the title and for grouping purposes. 19 | KEYS = ["event", "func", "exception_name", "queue"] 20 | 21 | def format_field(field: str, value: Any) -> str: 22 | if field == "queue": 23 | return "%s=%s" % (field, value.split(".")[0]) 24 | else: 25 | return "%s=%s" % (field, value) 26 | 27 | return "%s: %s" % ( 28 | self.prefix, 29 | " ".join( 30 | format_field(key, data[key]) for key in KEYS if key in data 31 | ), 32 | ) 33 | 34 | def emit(self, record: Any) -> Any: 35 | level = record.levelname.lower() 36 | if level not in self.SUPPORTED_LEVELS: 37 | return 38 | 39 | if record.levelno < self.notify_level: 40 | return 41 | 42 | try: 43 | data = json.loads(record.msg) 44 | except json.JSONDecodeError: 45 | return super(StructlogRollbarHandler, self).emit(record) 46 | 47 | # Title and grouping 48 | data["title"] = data["fingerprint"] = self.format_title(data) 49 | 50 | uuid = rollbar.report_message( 51 | message=data.pop("traceback", data["title"]), 52 | level=level, 53 | request=rollbar.get_request(), 54 | extra_data={}, 55 | payload_data=data, 56 | ) 57 | 58 | if uuid: 59 | record.rollbar_uuid = uuid 60 | -------------------------------------------------------------------------------- /tasktiger/runner.py: -------------------------------------------------------------------------------- 1 | from typing import TYPE_CHECKING, Any, Dict, List, Type 2 | 3 | from structlog.stdlib import BoundLogger 4 | 5 | from ._internal import import_attribute 6 | from .exceptions import TaskImportError 7 | from .timeouts import UnixSignalDeathPenalty 8 | 9 | if TYPE_CHECKING: 10 | from . import Task, TaskTiger 11 | 12 | 13 | class BaseRunner: 14 | """ 15 | Base implementation of the task runner. 16 | """ 17 | 18 | def __init__(self, tiger: "TaskTiger"): 19 | self.tiger = tiger 20 | 21 | def run_single_task(self, task: "Task", hard_timeout: float) -> None: 22 | """ 23 | Run the given task using the hard timeout in seconds. 24 | 25 | This is called inside of the forked process. 26 | """ 27 | raise NotImplementedError("Single tasks are not supported.") 28 | 29 | def run_batch_tasks( 30 | self, tasks: List["Task"], hard_timeout: float 31 | ) -> None: 32 | """ 33 | Run the given tasks using the hard timeout in seconds. 34 | 35 | This is called inside of the forked process. 36 | """ 37 | raise NotImplementedError("Batch tasks are not supported.") 38 | 39 | def run_eager_task(self, task: "Task") -> None: 40 | """ 41 | Run the task eagerly and return the value. 42 | 43 | Note that the task function could be a batch function. 44 | """ 45 | raise NotImplementedError("Eager tasks are not supported.") 46 | 47 | def on_permanent_error( 48 | self, task: "Task", execution: Dict[str, Any] 49 | ) -> None: 50 | """ 51 | Called if the task fails permanently. 52 | 53 | A task fails permanently if its status is set to ERROR and it is no 54 | longer retried. 55 | 56 | This is called in the main worker process. 57 | """ 58 | 59 | 60 | class DefaultRunner(BaseRunner): 61 | """ 62 | Default implementation of the task runner. 63 | """ 64 | 65 | def run_single_task(self, task: "Task", hard_timeout: float) -> None: 66 | with UnixSignalDeathPenalty(hard_timeout): 67 | task.func(*task.args, **task.kwargs) 68 | 69 | def run_batch_tasks( 70 | self, tasks: List["Task"], hard_timeout: float 71 | ) -> None: 72 | params = [{"args": task.args, "kwargs": task.kwargs} for task in tasks] 73 | func = tasks[0].func 74 | with UnixSignalDeathPenalty(hard_timeout): 75 | func(params) 76 | 77 | def run_eager_task(self, task: "Task") -> None: 78 | func = task.func 79 | is_batch_func = getattr(func, "_task_batch", False) 80 | 81 | if is_batch_func: 82 | return func([{"args": task.args, "kwargs": task.kwargs}]) 83 | else: 84 | return func(*task.args, **task.kwargs) 85 | 86 | 87 | def get_runner_class( 88 | log: BoundLogger, tasks: List["Task"] 89 | ) -> Type[BaseRunner]: 90 | runner_class_paths = {task.serialized_runner_class for task in tasks} 91 | if len(runner_class_paths) > 1: 92 | log.error( 93 | "cannot mix multiple runner classes", 94 | runner_class_paths=", ".join(str(p) for p in runner_class_paths), 95 | ) 96 | raise ValueError("Found multiple runner classes in batch task.") 97 | 98 | runner_class_path = runner_class_paths.pop() 99 | if runner_class_path: 100 | try: 101 | return import_attribute(runner_class_path) 102 | except TaskImportError: 103 | log.error( 104 | "could not import runner class", 105 | runner_class_path=runner_class_path, 106 | ) 107 | raise 108 | return DefaultRunner 109 | -------------------------------------------------------------------------------- /tasktiger/schedule.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | from typing import Callable, Optional, Tuple 3 | 4 | __all__ = ["periodic", "cron_expr"] 5 | 6 | START_DATE = datetime.datetime(2000, 1, 1) 7 | 8 | 9 | def _periodic( 10 | dt: datetime.datetime, 11 | period: int, 12 | start_date: datetime.datetime, 13 | end_date: datetime.datetime, 14 | ) -> Optional[datetime.datetime]: 15 | if end_date and dt >= end_date: 16 | return None 17 | 18 | if dt < start_date: 19 | return start_date 20 | 21 | # Determine the next time the task should be run 22 | delta = dt - start_date 23 | seconds = delta.seconds + delta.days * 86400 24 | runs = seconds // period 25 | next_run = runs + 1 26 | next_date = start_date + datetime.timedelta(seconds=next_run * period) 27 | 28 | # Make sure the time is still within bounds. 29 | if end_date and next_date > end_date: 30 | return None 31 | 32 | return next_date 33 | 34 | 35 | def periodic( 36 | seconds: int = 0, 37 | minutes: int = 0, 38 | hours: int = 0, 39 | days: int = 0, 40 | weeks: int = 0, 41 | start_date: Optional[datetime.datetime] = None, 42 | end_date: Optional[datetime.datetime] = None, 43 | ) -> Tuple[Callable[..., Optional[datetime.datetime]], Tuple]: 44 | """ 45 | Periodic task schedule: Use to schedule a task to run periodically, 46 | starting from start_date (or None to be active immediately) until end_date 47 | (or None to repeat forever). 48 | 49 | The period starts at the given start_date, or on Jan 1st 2000. 50 | 51 | For more details, see README. 52 | """ 53 | period = ( 54 | seconds + minutes * 60 + hours * 3600 + days * 86400 + weeks * 604800 55 | ) 56 | assert period > 0, "Must specify a positive period." 57 | if not start_date: 58 | # Saturday at midnight 59 | start_date = START_DATE 60 | return (_periodic, (period, start_date, end_date)) 61 | 62 | 63 | def _cron_expr( 64 | dt: datetime.datetime, 65 | expr: str, 66 | start_date: datetime.datetime, 67 | end_date: Optional[datetime.datetime] = None, 68 | ) -> Optional[datetime.datetime]: 69 | import croniter # type: ignore 70 | import pytz # type: ignore 71 | 72 | localize = pytz.utc.localize 73 | 74 | if end_date and dt >= end_date: 75 | return None 76 | 77 | if dt < start_date: 78 | return start_date 79 | 80 | assert croniter.croniter.is_valid(expr), "Cron expression is not valid." 81 | 82 | start_date = localize(start_date) 83 | dt = localize(dt) 84 | 85 | next_utc = croniter.croniter(expr, dt).get_next(ret_type=datetime.datetime) 86 | next_utc = next_utc.replace(tzinfo=None) 87 | 88 | # Make sure the time is still within bounds. 89 | if end_date and next_utc > end_date: 90 | return None 91 | 92 | return next_utc 93 | 94 | 95 | def cron_expr( 96 | expr: str, 97 | start_date: Optional[datetime.datetime] = None, 98 | end_date: Optional[datetime.datetime] = None, 99 | ) -> Tuple[Callable[..., Optional[datetime.datetime]], Tuple]: 100 | """ 101 | Periodic task schedule via cron expression: Use to schedule a task to run periodically, 102 | starting from start_date (or None to be active immediately) until end_date 103 | (or None to repeat forever). 104 | 105 | This function behaves similar to the cron jobs, which run with a minimum of 1 minute 106 | granularity. So specifying "* * * * *" expression will the run the task every 107 | minute. 108 | 109 | For more details, see README. 110 | """ 111 | if not start_date: 112 | start_date = START_DATE 113 | return (_cron_expr, (expr, start_date, end_date)) 114 | -------------------------------------------------------------------------------- /tasktiger/stats.py: -------------------------------------------------------------------------------- 1 | import threading 2 | import time 3 | from typing import TYPE_CHECKING, Optional 4 | 5 | from ._internal import g_fork_lock 6 | 7 | if TYPE_CHECKING: 8 | from .worker import Worker 9 | 10 | 11 | class StatsThread(threading.Thread): 12 | def __init__(self, tiger: "Worker") -> None: 13 | super(StatsThread, self).__init__() 14 | self.tiger = tiger 15 | self._stop_event = threading.Event() 16 | 17 | self._task_running = False 18 | self._time_start = time.monotonic() 19 | self._time_busy: float = 0.0 20 | self._task_start_time: Optional[float] = None 21 | self.daemon = True # Exit process if main thread exits unexpectedly 22 | 23 | # Lock that protects stats computations from interleaving. For example, 24 | # we don't want report_task_start() to run at the same time as 25 | # compute_stats(), as it might result in an inconsistent state. 26 | self._computation_lock = threading.Lock() 27 | 28 | def report_task_start(self) -> None: 29 | now = time.monotonic() 30 | with self._computation_lock: 31 | self._task_start_time = now 32 | self._task_running = True 33 | 34 | def report_task_end(self) -> None: 35 | now = time.monotonic() 36 | with self._computation_lock: 37 | assert self._task_start_time is not None 38 | self._time_busy += now - self._task_start_time 39 | self._task_running = False 40 | self._task_start_time = None 41 | 42 | def compute_stats(self) -> None: 43 | now = time.monotonic() 44 | 45 | with self._computation_lock: 46 | time_total = now - self._time_start 47 | time_busy = self._time_busy 48 | self._time_start = now 49 | self._time_busy = 0 50 | if self._task_running: 51 | assert self._task_start_time is not None 52 | time_busy += now - self._task_start_time 53 | self._task_start_time = now 54 | else: 55 | self._task_start_time = None 56 | 57 | if time_total: 58 | utilization = 100.0 / time_total * time_busy 59 | with g_fork_lock: 60 | self.tiger.log.info( 61 | "stats", 62 | time_total=time_total, 63 | time_busy=time_busy, 64 | utilization=utilization, 65 | ) 66 | 67 | def run(self) -> None: 68 | while not self._stop_event.wait(self.tiger.config["STATS_INTERVAL"]): 69 | self.compute_stats() 70 | 71 | def stop(self) -> None: 72 | self._stop_event.set() 73 | -------------------------------------------------------------------------------- /tasktiger/task.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import datetime 3 | import json 4 | import time 5 | from typing import ( 6 | TYPE_CHECKING, 7 | Any, 8 | Callable, 9 | Collection, 10 | Dict, 11 | List, 12 | Optional, 13 | Tuple, 14 | Type, 15 | Union, 16 | ) 17 | 18 | import redis 19 | from structlog.stdlib import BoundLogger 20 | 21 | from ._internal import ( 22 | ERROR, 23 | QUEUED, 24 | SCHEDULED, 25 | g, 26 | gen_id, 27 | gen_unique_id, 28 | get_timestamp, 29 | import_attribute, 30 | serialize_func_name, 31 | serialize_retry_method, 32 | ) 33 | from .exceptions import QueueFullException, TaskImportError, TaskNotFound 34 | from .runner import BaseRunner, get_runner_class 35 | from .types import RetryStrategy 36 | 37 | if TYPE_CHECKING: 38 | from . import TaskTiger 39 | 40 | __all__ = ["Task"] 41 | 42 | 43 | class Task: 44 | def __init__( 45 | self, 46 | tiger: "TaskTiger", 47 | func: Optional[Callable] = None, 48 | args: Optional[Any] = None, 49 | kwargs: Optional[Any] = None, 50 | queue: Optional[str] = None, 51 | hard_timeout: Optional[float] = None, 52 | unique: Optional[bool] = None, 53 | unique_key: Optional[Collection[str]] = None, 54 | lock: Optional[bool] = None, 55 | lock_key: Optional[Collection[str]] = None, 56 | retry: Optional[bool] = None, 57 | retry_on: Optional[Collection[Type[BaseException]]] = None, 58 | retry_method: Optional[ 59 | Union[Callable[[int], float], Tuple[Callable[..., float], Tuple]] 60 | ] = None, 61 | max_queue_size: Optional[int] = None, 62 | max_stored_executions: Optional[int] = None, 63 | runner_class: Optional[Type["BaseRunner"]] = None, 64 | # internal variables 65 | _data: Any = None, 66 | _state: Any = None, 67 | _ts: Any = None, 68 | _executions: Optional[List[Dict[str, Any]]] = None, 69 | ): 70 | """ 71 | Queues a task. See README.rst for an explanation of the options. 72 | """ 73 | 74 | if func and queue is None: 75 | queue = Task.queue_from_function(func, tiger) 76 | 77 | self.tiger = tiger 78 | self._func = func 79 | self._queue = queue 80 | self._state = _state 81 | self._ts = _ts 82 | self._executions = _executions or [] 83 | 84 | # Internal initialization based on raw data. 85 | if _data is not None: 86 | self._data = _data 87 | return 88 | 89 | assert func 90 | 91 | serialized_name = serialize_func_name(func) 92 | 93 | if unique is None: 94 | unique = getattr(func, "_task_unique", False) 95 | 96 | if unique_key is None: 97 | unique_key = getattr(func, "_task_unique_key", None) 98 | 99 | if lock is None: 100 | lock = getattr(func, "_task_lock", False) 101 | 102 | if lock_key is None: 103 | lock_key = getattr(func, "_task_lock_key", None) 104 | 105 | if retry is None: 106 | retry = getattr(func, "_task_retry", False) 107 | 108 | if retry_on is None: 109 | retry_on = getattr(func, "_task_retry_on", None) 110 | 111 | if retry_method is None: 112 | retry_method = getattr(func, "_task_retry_method", None) 113 | 114 | if max_queue_size is None: 115 | max_queue_size = getattr(func, "_task_max_queue_size", None) 116 | 117 | if max_stored_executions is None: 118 | max_stored_executions = getattr( 119 | func, "_task_max_stored_executions", None 120 | ) 121 | 122 | if runner_class is None: 123 | runner_class = getattr(func, "_task_runner_class", None) 124 | 125 | # normalize falsy args/kwargs to empty structures 126 | args = args or [] 127 | kwargs = kwargs or {} 128 | if unique or unique_key: 129 | if unique_key: 130 | task_id = gen_unique_id( 131 | serialized_name, 132 | None, 133 | {key: kwargs.get(key) for key in unique_key}, 134 | ) 135 | else: 136 | task_id = gen_unique_id(serialized_name, args, kwargs) 137 | else: 138 | task_id = gen_id() 139 | 140 | task: Dict[str, Any] = {"id": task_id, "func": serialized_name} 141 | if unique or unique_key: 142 | task["unique"] = True 143 | if unique_key: 144 | task["unique_key"] = unique_key 145 | if lock or lock_key: 146 | task["lock"] = True 147 | if lock_key: 148 | task["lock_key"] = lock_key 149 | if args: 150 | task["args"] = args 151 | if kwargs: 152 | task["kwargs"] = kwargs 153 | if hard_timeout: 154 | task["hard_timeout"] = hard_timeout 155 | if retry or retry_on or retry_method: 156 | if not retry_method: 157 | retry_method = tiger.config["DEFAULT_RETRY_METHOD"] 158 | 159 | task["retry_method"] = serialize_retry_method(retry_method) 160 | if retry_on: 161 | task["retry_on"] = [ 162 | serialize_func_name(cls) for cls in retry_on 163 | ] 164 | if max_queue_size: 165 | task["max_queue_size"] = max_queue_size 166 | if max_stored_executions is not None: 167 | task["max_stored_executions"] = max_stored_executions 168 | if runner_class: 169 | serialized_runner_class = serialize_func_name(runner_class) 170 | task["runner_class"] = serialized_runner_class 171 | 172 | self._data = task 173 | 174 | @property 175 | def id(self) -> str: 176 | return self._data["id"] 177 | 178 | @property 179 | def data(self) -> Dict[str, Any]: 180 | return self._data 181 | 182 | @property 183 | def time_last_queued(self) -> Optional[datetime.datetime]: 184 | timestamp = self._data.get("time_last_queued") 185 | if timestamp is None: 186 | return None 187 | else: 188 | return datetime.datetime.utcfromtimestamp(timestamp) 189 | 190 | @property 191 | def state(self) -> str: 192 | return self._state 193 | 194 | @property 195 | def queue(self) -> str: 196 | assert self._queue 197 | return self._queue 198 | 199 | @property 200 | def serialized_func(self) -> str: 201 | return self._data["func"] 202 | 203 | @property 204 | def lock(self) -> bool: 205 | return self._data.get("lock", False) 206 | 207 | @property 208 | def lock_key(self) -> Optional[str]: 209 | return self._data.get("lock_key") 210 | 211 | @property 212 | def args(self) -> List[Any]: 213 | return self._data.get("args", []) 214 | 215 | @property 216 | def kwargs(self) -> Dict[str, Any]: 217 | return self._data.get("kwargs", {}) 218 | 219 | @property 220 | def hard_timeout(self) -> Optional[float]: 221 | return self._data.get("hard_timeout", None) 222 | 223 | @property 224 | def unique(self) -> bool: 225 | return self._data.get("unique", False) 226 | 227 | @property 228 | def unique_key(self) -> Optional[str]: 229 | return self._data.get("unique_key") 230 | 231 | @property 232 | def retry_method(self) -> Optional[RetryStrategy]: 233 | if "retry_method" in self._data: 234 | retry_func, retry_args = self._data["retry_method"] 235 | return retry_func, retry_args 236 | else: 237 | return None 238 | 239 | @property 240 | def retry_on(self) -> List[str]: 241 | return self._data.get("retry_on") 242 | 243 | def should_retry_on( 244 | self, 245 | exception_class: Type[BaseException], 246 | logger: Optional[BoundLogger] = None, 247 | ) -> bool: 248 | """ 249 | Whether this task should be retried when the given exception occurs. 250 | """ 251 | for n in self.retry_on or []: 252 | try: 253 | if issubclass(exception_class, import_attribute(n)): 254 | return True 255 | except TaskImportError: 256 | if logger: 257 | logger.error( 258 | "should_retry_on could not import class", 259 | exception_name=n, 260 | ) 261 | return False 262 | 263 | @property 264 | def func(self) -> Callable: 265 | if not self._func: 266 | self._func = import_attribute(self.serialized_func) 267 | return self._func 268 | 269 | @property 270 | def max_stored_executions(self) -> Optional[int]: 271 | return self._data.get("max_stored_executions") 272 | 273 | @property 274 | def serialized_runner_class(self) -> str: 275 | return self._data.get("runner_class") 276 | 277 | @property 278 | def ts(self) -> Optional[datetime.datetime]: 279 | """ 280 | The timestamp (datetime) of the task in the queue, or None, if the task 281 | hasn't been queued. 282 | """ 283 | return self._ts 284 | 285 | @property 286 | def executions(self) -> List[Dict[str, Any]]: 287 | return self._executions 288 | 289 | def _move( 290 | self, 291 | from_state: Optional[str] = None, 292 | to_state: Optional[str] = None, 293 | when: Optional[float] = None, 294 | mode: Optional[str] = None, 295 | ) -> None: 296 | """ 297 | Internal helper to move a task from one state to another (e.g. from 298 | QUEUED to DELAYED). The "when" argument indicates the timestamp of the 299 | task in the new state. If no to_state is specified, the task will be 300 | simply removed from the original state. 301 | 302 | The "mode" param can be specified to define how the timestamp in the 303 | new state should be updated and is passed to the ZADD Redis script (see 304 | its documentation for details). 305 | 306 | Raises TaskNotFound if the task is not in the expected state or not in 307 | the expected queue. 308 | """ 309 | 310 | scripts = self.tiger.scripts 311 | 312 | from_state = from_state or self.state 313 | queue = self.queue 314 | 315 | assert from_state 316 | assert queue 317 | 318 | try: 319 | scripts.move_task( 320 | id=self.id, 321 | queue=self.queue, 322 | from_state=from_state, 323 | to_state=to_state, 324 | unique=self.unique, 325 | when=when or time.time(), 326 | mode=mode, 327 | key_func=self.tiger._key, 328 | publish_queued_tasks=self.tiger.config["PUBLISH_QUEUED_TASKS"], 329 | ) 330 | except redis.ResponseError as e: 331 | if "" in e.args[0]: 332 | raise TaskNotFound( 333 | 'Task {} not found in queue "{}" in state "{}".'.format( 334 | self.id, queue, from_state 335 | ) 336 | ) 337 | raise 338 | else: 339 | self._state = to_state 340 | 341 | def execute(self) -> None: 342 | func = self.func 343 | is_batch_func = getattr(func, "_task_batch", False) 344 | 345 | g["current_task_is_batch"] = is_batch_func 346 | g["current_tasks"] = [self] 347 | g["tiger"] = self.tiger 348 | 349 | try: 350 | runner_class = get_runner_class(self.tiger.log, [self]) 351 | runner = runner_class(self.tiger) 352 | return runner.run_eager_task(self) 353 | finally: 354 | g["current_task_is_batch"] = None 355 | g["current_tasks"] = None 356 | g["tiger"] = None 357 | 358 | def delay( 359 | self, 360 | when: Optional[Union[datetime.timedelta, datetime.datetime]] = None, 361 | max_queue_size: Optional[int] = None, 362 | ) -> None: 363 | tiger = self.tiger 364 | 365 | ts = get_timestamp(when) 366 | 367 | now = time.time() 368 | self._data["time_last_queued"] = now 369 | 370 | if max_queue_size is None: 371 | max_queue_size = self._data.get("max_queue_size") 372 | 373 | if not ts or ts <= now: 374 | # Immediately queue if the timestamp is in the past. 375 | ts = now 376 | state = QUEUED 377 | else: 378 | state = SCHEDULED 379 | 380 | # When using ALWAYS_EAGER, make sure we have serialized the task to 381 | # ensure there are no serialization errors. 382 | serialized_task = json.dumps(self._data) 383 | 384 | if max_queue_size: 385 | # This will fail adding a unique task that already is queued but 386 | # the queue size is at the max 387 | queue_size = tiger.get_total_queue_size(self.queue) 388 | if queue_size >= max_queue_size: 389 | raise QueueFullException("Queue size: {}".format(queue_size)) 390 | 391 | if tiger.config["ALWAYS_EAGER"] and state == QUEUED: 392 | return self.execute() 393 | 394 | pipeline = tiger.connection.pipeline() 395 | pipeline.sadd(tiger._key(state), self.queue) 396 | pipeline.set(tiger._key("task", self.id), serialized_task) 397 | # In case of unique tasks, don't update the score. 398 | tiger.scripts.zadd( 399 | tiger._key(state, self.queue), 400 | ts, 401 | self.id, 402 | mode="nx", 403 | client=pipeline, 404 | ) 405 | if state == QUEUED and tiger.config["PUBLISH_QUEUED_TASKS"]: 406 | pipeline.publish(tiger._key("activity"), self.queue) 407 | pipeline.execute() 408 | 409 | self._state = state 410 | self._ts = ts 411 | 412 | def update_scheduled_time( 413 | self, when: Optional[Union[datetime.timedelta, datetime.datetime]] 414 | ) -> None: 415 | """ 416 | Updates a scheduled task's date to the given date. If the task is not 417 | scheduled, a TaskNotFound exception is raised. 418 | """ 419 | tiger = self.tiger 420 | 421 | ts = get_timestamp(when) 422 | assert ts 423 | 424 | pipeline = tiger.connection.pipeline() 425 | key = tiger._key(SCHEDULED, self.queue) 426 | tiger.scripts.zadd(key, ts, self.id, mode="xx", client=pipeline) 427 | pipeline.zscore(key, self.id) 428 | _, score = pipeline.execute() 429 | if not score: 430 | raise TaskNotFound( 431 | 'Task {} not found in queue "{}" in state "{}".'.format( 432 | self.id, self.queue, SCHEDULED 433 | ) 434 | ) 435 | 436 | self._ts = ts 437 | 438 | def __repr__(self) -> str: 439 | return "" % self.func 440 | 441 | @classmethod 442 | def from_id( 443 | cls, 444 | tiger: "TaskTiger", 445 | queue: str, 446 | state: str, 447 | task_id: str, 448 | load_executions: int = 0, 449 | ) -> "Task": 450 | """ 451 | Loads a task with the given ID from the given queue in the given 452 | state. An integer may be passed in the load_executions parameter 453 | to indicate how many executions should be loaded (starting from the 454 | latest). If the task doesn't exist, None is returned. 455 | """ 456 | pipeline = tiger.connection.pipeline() 457 | pipeline.get(tiger._key("task", task_id)) 458 | pipeline.zscore(tiger._key(state, queue), task_id) 459 | if load_executions: 460 | pipeline.lrange( 461 | tiger._key("task", task_id, "executions"), -load_executions, -1 462 | ) 463 | ( 464 | serialized_data, 465 | score, 466 | serialized_executions, 467 | ) = pipeline.execute() 468 | else: 469 | serialized_data, score = pipeline.execute() 470 | serialized_executions = [] 471 | 472 | if serialized_data and score: 473 | data = json.loads(serialized_data) 474 | executions = [json.loads(e) for e in serialized_executions if e] 475 | return Task( 476 | tiger, 477 | queue=queue, 478 | _data=data, 479 | _state=state, 480 | _executions=executions, 481 | _ts=datetime.datetime.utcfromtimestamp(score), 482 | ) 483 | else: 484 | raise TaskNotFound("Task {} not found.".format(task_id)) 485 | 486 | @classmethod 487 | def tasks_from_queue( 488 | cls, 489 | tiger: "TaskTiger", 490 | queue: str, 491 | state: str, 492 | skip: int = 0, 493 | limit: int = 1000, 494 | load_executions: int = 0, 495 | include_not_found: bool = False, 496 | ) -> Tuple[int, List["Task"]]: 497 | """ 498 | Return tasks from a queue. 499 | 500 | Args: 501 | tiger: TaskTiger instance. 502 | queue: Name of the queue. 503 | state: State of the task (QUEUED, ACTIVE, SCHEDULED, ERROR). 504 | limit: Maximum number of tasks to return. 505 | load_executions: Maximum number of executions to load for each task 506 | (starting from the latest). 507 | include_not_found: Whether to include tasks that cannot be loaded. 508 | 509 | Returns: 510 | Tuple with the following information: 511 | * total items in the queue 512 | * tasks from the given queue in the given state, latest first. 513 | """ 514 | 515 | key = tiger._key(state, queue) 516 | pipeline = tiger.connection.pipeline() 517 | pipeline.zcard(key) 518 | pipeline.zrange(key, -limit - skip, -1 - skip, withscores=True) 519 | n, items = pipeline.execute() 520 | 521 | tasks = [] 522 | 523 | if items: 524 | tss = [ 525 | datetime.datetime.utcfromtimestamp(item[1]) for item in items 526 | ] 527 | if load_executions: 528 | pipeline = tiger.connection.pipeline() 529 | pipeline.mget([tiger._key("task", item[0]) for item in items]) 530 | for item in items: 531 | pipeline.lrange( 532 | tiger._key("task", item[0], "executions"), 533 | -load_executions, 534 | -1, 535 | ) 536 | results = pipeline.execute() 537 | 538 | for idx, serialized_data, serialized_executions, ts in zip( 539 | range(len(items)), results[0], results[1:], tss 540 | ): 541 | if serialized_data is None and include_not_found: 542 | data = {"id": items[idx][0]} 543 | else: 544 | data = json.loads(serialized_data) 545 | 546 | executions = [ 547 | json.loads(e) for e in serialized_executions if e 548 | ] 549 | 550 | task = Task( 551 | tiger, 552 | queue=queue, 553 | _data=data, 554 | _state=state, 555 | _ts=ts, 556 | _executions=executions, 557 | ) 558 | 559 | tasks.append(task) 560 | else: 561 | result = tiger.connection.mget( 562 | [tiger._key("task", item[0]) for item in items] 563 | ) 564 | for idx, serialized_data, ts in zip( 565 | range(len(items)), result, tss 566 | ): 567 | if serialized_data is None and include_not_found: 568 | data = {"id": items[idx][0]} 569 | else: 570 | data = json.loads(serialized_data) 571 | 572 | task = Task( 573 | tiger, queue=queue, _data=data, _state=state, _ts=ts 574 | ) 575 | tasks.append(task) 576 | 577 | return n, tasks 578 | 579 | @classmethod 580 | def queue_from_function(cls, func: Any, tiger: "TaskTiger") -> str: 581 | """Get queue from function.""" 582 | return getattr(func, "_task_queue", tiger.config["DEFAULT_QUEUE"]) 583 | 584 | def n_executions(self) -> int: 585 | """ 586 | Queries and returns the number of past task executions. 587 | """ 588 | pipeline = self.tiger.connection.pipeline() 589 | pipeline.exists(self.tiger._key("task", self.id)) 590 | pipeline.get(self.tiger._key("task", self.id, "executions_count")) 591 | 592 | exists, executions_count = pipeline.execute() 593 | if not exists: 594 | raise TaskNotFound("Task {} not found.".format(self.id)) 595 | 596 | return int(executions_count or 0) 597 | 598 | def retry(self) -> None: 599 | """ 600 | Retries a task that's in the error queue. 601 | 602 | Raises TaskNotFound if the task could not be found in the ERROR 603 | queue. 604 | """ 605 | self._move(from_state=ERROR, to_state=QUEUED) 606 | 607 | def cancel(self) -> None: 608 | """ 609 | Cancels a task that is queued in the SCHEDULED queue. 610 | 611 | Raises TaskNotFound if the task could not be found in the SCHEDULED 612 | queue. 613 | """ 614 | self._move(from_state=SCHEDULED) 615 | 616 | def delete(self) -> None: 617 | """ 618 | Removes a task that's in the error queue. 619 | 620 | Raises TaskNotFound if the task could not be found in the ERROR 621 | queue. 622 | """ 623 | self._move(from_state=ERROR) 624 | 625 | def clone(self) -> "Task": 626 | """Returns a clone of the this task""" 627 | return type(self)( 628 | tiger=self.tiger, 629 | func=self.func, 630 | queue=self.queue, 631 | _state=self._state, 632 | _ts=self._ts, 633 | _executions=copy.copy(self._executions), 634 | _data=copy.copy(self._data), 635 | ) 636 | 637 | def _queue_for_next_period(self) -> float: 638 | now = datetime.datetime.utcnow() 639 | schedule = self.func._task_schedule # type: ignore[attr-defined] 640 | if callable(schedule): 641 | schedule_func = schedule 642 | schedule_args = () 643 | else: 644 | schedule_func, schedule_args = schedule 645 | when = schedule_func(now, *schedule_args) 646 | if when: 647 | # recalculate the unique id so that malformed ids don't persist 648 | # between executions 649 | task = self.clone() 650 | task._data["id"] = gen_unique_id( 651 | task.serialized_func, task.args, task.kwargs 652 | ) 653 | task.delay(when=when) 654 | return when 655 | -------------------------------------------------------------------------------- /tasktiger/test_helpers.py: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | from . import TaskTiger 4 | from .task import Task 5 | from .worker import Worker 6 | 7 | __all__ = ["TaskTigerTestMixin"] 8 | 9 | 10 | class TaskTigerTestMixin: 11 | """ 12 | Unit test mixin for tests that use TaskTiger. 13 | """ 14 | 15 | def run_worker( 16 | self, tiger: TaskTiger, raise_on_errors: bool = True, **kwargs: Any 17 | ) -> None: 18 | # A worker run processes queued tasks, and then queues scheduled tasks. 19 | # We therefore need to run the worker twice to execute due scheduled 20 | # tasks. 21 | Worker(tiger, **kwargs).run(once=True) 22 | Worker(tiger, **kwargs).run(once=True) 23 | 24 | # Print any TaskTiger failures for debugging purposes. 25 | prefix = tiger.config["REDIS_PREFIX"] 26 | state = "error" 27 | has_errors = False 28 | for queue in tiger.connection.smembers("{}:{}".format(prefix, state)): 29 | n_tasks, tasks = Task.tasks_from_queue( 30 | tiger, queue, state, load_executions=1 31 | ) 32 | for task in tasks: 33 | print("") 34 | print(task, "failed:") 35 | print(task.executions[0]["traceback"]) 36 | has_errors = True 37 | if has_errors and raise_on_errors: 38 | raise Exception("One or more tasks have failed.") 39 | -------------------------------------------------------------------------------- /tasktiger/timeouts.py: -------------------------------------------------------------------------------- 1 | import signal 2 | from types import TracebackType 3 | from typing import Any, Literal, Optional, Type 4 | 5 | from .exceptions import JobTimeoutException 6 | 7 | 8 | class BaseDeathPenalty: 9 | """Base class to setup job timeouts.""" 10 | 11 | def __init__(self, timeout: float) -> None: 12 | self._timeout = timeout 13 | 14 | def __enter__(self) -> None: 15 | self.setup_death_penalty() 16 | 17 | def __exit__( 18 | self, 19 | type: Optional[Type[BaseException]], 20 | value: Optional[BaseException], 21 | traceback: Optional[TracebackType], 22 | ) -> Literal[False]: 23 | # Always cancel immediately, since we're done 24 | try: 25 | self.cancel_death_penalty() 26 | except JobTimeoutException: 27 | # Weird case: we're done with the with body, but now the alarm is 28 | # fired. We may safely ignore this situation and consider the 29 | # body done. 30 | pass 31 | 32 | # __exit__ may return True to suppress further exception handling. We 33 | # don't want to suppress any exceptions here, since all errors should 34 | # just pass through, JobTimeoutException being handled normally to the 35 | # invoking context. 36 | return False 37 | 38 | def setup_death_penalty(self) -> None: 39 | raise NotImplementedError() 40 | 41 | def cancel_death_penalty(self) -> None: 42 | raise NotImplementedError() 43 | 44 | 45 | class UnixSignalDeathPenalty(BaseDeathPenalty): 46 | def handle_death_penalty(self, signum: int, frame: Any) -> None: 47 | raise JobTimeoutException( 48 | "Job exceeded maximum timeout " 49 | "value (%d seconds)." % self._timeout 50 | ) 51 | 52 | def setup_death_penalty(self) -> None: 53 | """Sets up an alarm signal and a signal handler that raises 54 | a JobTimeoutException after the timeout amount (expressed in 55 | seconds). 56 | """ 57 | signal.signal(signal.SIGALRM, self.handle_death_penalty) 58 | signal.setitimer(signal.ITIMER_REAL, self._timeout) 59 | 60 | def cancel_death_penalty(self) -> None: 61 | """Removes the death penalty alarm and puts back the system into 62 | default signal handling. 63 | """ 64 | signal.alarm(0) 65 | signal.signal(signal.SIGALRM, signal.SIG_DFL) 66 | -------------------------------------------------------------------------------- /tasktiger/types.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Tuple 2 | 3 | RetryStrategy = Tuple[Callable[..., float], Tuple] 4 | -------------------------------------------------------------------------------- /tasktiger/utils.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | REDIS_GLOB_CHARACTER_PATTERN = re.compile(r"([\\?*\[\]])") 4 | 5 | 6 | def redis_glob_escape(value: str) -> str: 7 | return REDIS_GLOB_CHARACTER_PATTERN.sub(r"\\\1", value) 8 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/closeio/tasktiger/ded8463b3e4fc14b0d7019afc374b7d3f6fed24f/tests/__init__.py -------------------------------------------------------------------------------- /tests/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | # How much to delay scheduled tasks for testing purposes. 4 | # Note that on macOS Catalina, when using an unsigned Python version, taskgated 5 | # (com.apple.securityd) needs to approve launching the process. We therefore 6 | # need ample time here (> 0.3s) in order to prevent test failures. 7 | DELAY = 0.4 8 | 9 | # Redis database number which will be wiped and used for the tests 10 | TEST_DB = int(os.environ.get("REDIS_DB", 7)) 11 | 12 | # Redis hostname 13 | REDIS_HOST = os.environ.get("REDIS_HOST", "localhost") 14 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | import pytest 4 | 5 | from .utils import get_tiger 6 | 7 | 8 | @pytest.fixture 9 | def tiger(): 10 | tiger = get_tiger() 11 | redis = tiger.connection 12 | redis.flushdb() 13 | 14 | yield tiger 15 | 16 | redis.flushdb() 17 | redis.close() 18 | # Force disconnect so we don't get Too many open files 19 | redis.connection_pool.disconnect() 20 | 21 | 22 | @pytest.fixture 23 | def redis(tiger): 24 | return tiger.connection 25 | 26 | 27 | @pytest.fixture 28 | def ensure_queues(redis): 29 | def _ensure_queues(queued=None, active=None, error=None, scheduled=None): 30 | expected_queues = { 31 | "queued": {name for name, n in (queued or {}).items() if n}, 32 | "active": {name for name, n in (active or {}).items() if n}, 33 | "error": {name for name, n in (error or {}).items() if n}, 34 | "scheduled": {name for name, n in (scheduled or {}).items() if n}, 35 | } 36 | actual_queues = { 37 | i: redis.smembers("t:{}".format(i)) 38 | for i in ("queued", "active", "error", "scheduled") 39 | } 40 | assert expected_queues == actual_queues 41 | 42 | def _ensure_queue(typ, data): 43 | data = data or {} 44 | ret = {} 45 | for name, n in data.items(): 46 | task_ids = redis.zrange("t:%s:%s" % (typ, name), 0, -1) 47 | assert len(task_ids) == n 48 | ret[name] = [ 49 | json.loads(redis.get("t:task:%s" % task_id)) 50 | for task_id in task_ids 51 | ] 52 | assert [task["id"] for task in ret[name]] == task_ids 53 | return ret 54 | 55 | return { 56 | "queued": _ensure_queue("queued", queued), 57 | "active": _ensure_queue("active", active), 58 | "error": _ensure_queue("error", error), 59 | "scheduled": _ensure_queue("scheduled", scheduled), 60 | } 61 | 62 | return _ensure_queues 63 | -------------------------------------------------------------------------------- /tests/tasks.py: -------------------------------------------------------------------------------- 1 | import json 2 | import time 3 | from math import ceil 4 | 5 | import redis 6 | 7 | from tasktiger import RetryException, TaskTiger 8 | from tasktiger.retry import fixed 9 | from tasktiger.runner import BaseRunner, DefaultRunner 10 | 11 | from .config import DELAY, REDIS_HOST, TEST_DB 12 | from .utils import get_tiger 13 | 14 | LONG_TASK_SIGNAL_KEY = "long_task_ok" 15 | 16 | tiger = get_tiger() 17 | 18 | 19 | def simple_task(): 20 | pass 21 | 22 | 23 | @tiger.task() 24 | def decorated_task(*args, **kwargs): 25 | pass 26 | 27 | 28 | # This decorator below must not contain parenthesis 29 | @tiger.task 30 | def decorated_task_simple_func(*args, **kwargs): 31 | pass 32 | 33 | 34 | def exception_task(): 35 | raise Exception("this failed") 36 | 37 | 38 | def system_exit_task(): 39 | raise SystemExit() 40 | 41 | 42 | @tiger.task(queue="other") 43 | def task_on_other_queue(): 44 | pass 45 | 46 | 47 | def file_args_task(filename, *args, **kwargs): 48 | with open(filename, "w") as f: 49 | f.write(json.dumps({"args": args, "kwargs": kwargs})) 50 | 51 | 52 | @tiger.task(hard_timeout=DELAY) 53 | def long_task_killed(): 54 | time.sleep(DELAY * 2) 55 | 56 | 57 | @tiger.task(hard_timeout=DELAY * 2) 58 | def long_task_ok(): 59 | # Signal task has started 60 | with redis.Redis( 61 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 62 | ) as conn: 63 | conn.lpush(LONG_TASK_SIGNAL_KEY, "1") 64 | time.sleep(DELAY) 65 | 66 | 67 | def wait_for_long_task(): 68 | """Waits for a long task to start.""" 69 | with redis.Redis( 70 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 71 | ) as conn: 72 | result = conn.blpop(LONG_TASK_SIGNAL_KEY, int(ceil(DELAY * 3))) 73 | assert result[1] == "1" 74 | 75 | 76 | @tiger.task(unique=True) 77 | def unique_task(value=None): 78 | with redis.Redis( 79 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 80 | ) as conn: 81 | conn.lpush("unique_task", value) 82 | 83 | 84 | @tiger.task(unique=True) 85 | def unique_exception_task(value=None): 86 | raise Exception("this failed") 87 | 88 | 89 | @tiger.task(unique_key=("a",)) 90 | def unique_key_task(a, b): 91 | pass 92 | 93 | 94 | @tiger.task(lock=True) 95 | def locked_task(key, other=None): 96 | with redis.Redis( 97 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 98 | ) as conn: 99 | data = conn.getset(key, 1) 100 | if data is not None: 101 | raise Exception("task failed, key already set") 102 | time.sleep(DELAY) 103 | conn.delete(key) 104 | 105 | 106 | @tiger.task(queue="batch", batch=True) 107 | def batch_task(params): 108 | with redis.Redis( 109 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 110 | ) as conn: 111 | try: 112 | conn.rpush("batch_task", json.dumps(params)) 113 | except Exception: 114 | pass 115 | if any(p["args"][0] == 10 for p in params if p["args"]): 116 | raise Exception("exception") 117 | 118 | 119 | @tiger.task(queue="batch") 120 | def non_batch_task(arg): 121 | with redis.Redis( 122 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 123 | ) as conn: 124 | conn.rpush("batch_task", arg) 125 | 126 | if arg == 10: 127 | raise Exception("exception") 128 | 129 | 130 | def retry_task(): 131 | raise RetryException() 132 | 133 | 134 | def retry_task_2(): 135 | raise RetryException(method=fixed(DELAY, 1), log_error=False) 136 | 137 | 138 | @tiger.task(retry_method=fixed(DELAY, 1)) 139 | def retry_task_3(): 140 | raise RetryException(log_error=False) 141 | 142 | 143 | def verify_current_task(): 144 | with redis.Redis( 145 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 146 | ) as conn: 147 | try: 148 | tiger.current_tasks 149 | except RuntimeError: 150 | # This is expected (we need to use current_task) 151 | task = tiger.current_task 152 | conn.set("task_id", task.id) 153 | 154 | 155 | @tiger.task(batch=True, queue="batch") 156 | def verify_current_tasks(tasks): 157 | with redis.Redis( 158 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 159 | ) as conn: 160 | try: 161 | tasks = tiger.current_task 162 | except RuntimeError: 163 | # This is expected (we need to use current_tasks) 164 | 165 | tasks = tiger.current_tasks 166 | conn.rpush("task_ids", *[t.id for t in tasks]) 167 | 168 | 169 | def verify_current_serialized_func(): 170 | with redis.Redis( 171 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 172 | ) as conn: 173 | serialized_func = tiger.current_serialized_func 174 | conn.set("serialized_func", serialized_func) 175 | 176 | 177 | @tiger.task(batch=True, queue="batch") 178 | def verify_current_serialized_func_batch(tasks): 179 | with redis.Redis( 180 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 181 | ) as conn: 182 | serialized_func = tiger.current_serialized_func 183 | conn.set("serialized_func", serialized_func) 184 | 185 | 186 | @tiger.task() 187 | def verify_tasktiger_instance(): 188 | # Not necessarily the same object, but the same configuration. 189 | config_1 = dict(TaskTiger.current_instance.config) 190 | config_2 = dict(tiger.config) 191 | 192 | # Changed during the test case, so this may differ. 193 | config_1.pop("ALWAYS_EAGER") 194 | config_2.pop("ALWAYS_EAGER") 195 | 196 | assert config_1 == config_2 197 | 198 | 199 | @tiger.task() 200 | def sleep_task(delay=10): 201 | time.sleep(delay) 202 | 203 | 204 | @tiger.task(hard_timeout=1) 205 | def decorated_task_sleep_timeout(delay=10): 206 | time.sleep(delay) 207 | 208 | 209 | @tiger.task(max_queue_size=1) 210 | def decorated_task_max_queue_size(*args, **kwargs): 211 | pass 212 | 213 | 214 | class StaticTask: 215 | @staticmethod 216 | def task(): 217 | pass 218 | 219 | 220 | class MyRunnerClass(BaseRunner): 221 | def run_single_task(self, task, hard_timeout): 222 | assert self.tiger.config == tiger.config 223 | assert hard_timeout == 300 224 | assert task.func is simple_task 225 | 226 | with redis.Redis( 227 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 228 | ) as conn: 229 | conn.set("task_id", task.id) 230 | 231 | def run_batch_tasks(self, tasks, hard_timeout): 232 | assert self.tiger.config == tiger.config 233 | assert hard_timeout == 300 234 | assert len(tasks) == 2 235 | 236 | with redis.Redis( 237 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 238 | ) as conn: 239 | conn.set("task_args", ",".join(str(t.args[0]) for t in tasks)) 240 | 241 | def run_eager_task(self, task): 242 | return 123 243 | 244 | 245 | class MyErrorRunnerClass(DefaultRunner): 246 | def on_permanent_error(self, task, execution): 247 | assert task.func is exception_task 248 | assert execution["exception_name"] == "builtins:Exception" 249 | with redis.Redis( 250 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 251 | ) as conn: 252 | conn.set("task_id", task.id) 253 | -------------------------------------------------------------------------------- /tests/tasks_periodic.py: -------------------------------------------------------------------------------- 1 | """Periodic test task.""" 2 | 3 | import redis 4 | 5 | from tasktiger.schedule import periodic 6 | 7 | from .config import REDIS_HOST, TEST_DB 8 | from .utils import get_tiger 9 | 10 | tiger = get_tiger() 11 | 12 | 13 | @tiger.task( 14 | schedule=periodic(seconds=1), queue="periodic", retry_on=(ValueError,) 15 | ) 16 | def periodic_task(): 17 | """Periodic task.""" 18 | conn = redis.Redis(host=REDIS_HOST, db=TEST_DB, decode_responses=True) 19 | conn.incr("period_count", 1) 20 | fail = conn.get("fail-periodic-task") 21 | if fail == "retriable": 22 | raise ValueError("retriable failure") 23 | elif fail == "permanent": 24 | raise Exception("permanent failure") 25 | 26 | 27 | @tiger.task(schedule=periodic(seconds=1), queue="periodic_ignore") 28 | def periodic_task_ignore(): 29 | """ 30 | Ignored periodic task. 31 | 32 | This task should never get queued. 33 | """ 34 | pass 35 | -------------------------------------------------------------------------------- /tests/test_context_manager.py: -------------------------------------------------------------------------------- 1 | """Child context manager tests.""" 2 | import redis 3 | 4 | from tasktiger import Worker 5 | 6 | from .config import REDIS_HOST, TEST_DB 7 | from .tasks import exception_task, simple_task 8 | from .test_base import BaseTestCase 9 | 10 | 11 | class ContextManagerTester: 12 | """ 13 | Dummy context manager class. 14 | 15 | Uses Redis to track number of enter/exit calls 16 | """ 17 | 18 | def __init__(self, name): 19 | self.name = name 20 | self.conn = redis.Redis( 21 | host=REDIS_HOST, db=TEST_DB, decode_responses=True 22 | ) 23 | self.conn.set("cm:{}:enter".format(self.name), 0) 24 | self.conn.set("cm:{}:exit".format(self.name), 0) 25 | self.conn.set("cm:{}:exit_with_error".format(self.name), 0) 26 | 27 | def __enter__(self): 28 | self.conn.incr("cm:{}:enter".format(self.name)) 29 | 30 | def __exit__(self, exc_type, exc_val, exc_tb): 31 | self.conn.incr("cm:{}:exit".format(self.name)) 32 | if exc_type is not None: 33 | self.conn.incr("cm:{}:exit_with_error".format(self.name)) 34 | self.conn.close() 35 | 36 | 37 | class TestChildContextManagers(BaseTestCase): 38 | """Child context manager tests.""" 39 | 40 | def _get_context_managers(self, number): 41 | return [ContextManagerTester("cm" + str(i)) for i in range(number)] 42 | 43 | def _test_context_managers(self, num, task, should_fail=False): 44 | cms = self._get_context_managers(num) 45 | 46 | self.tiger.config["CHILD_CONTEXT_MANAGERS"] = cms 47 | self.tiger.delay(task) 48 | Worker(self.tiger).run(once=True) 49 | 50 | for i in range(num): 51 | assert self.conn.get("cm:{}:enter".format(cms[i].name)) == "1" 52 | assert self.conn.get("cm:{}:exit".format(cms[i].name)) == "1" 53 | if should_fail: 54 | assert ( 55 | self.conn.get("cm:{}:exit_with_error".format(cms[i].name)) 56 | == "1" 57 | ) 58 | else: 59 | assert ( 60 | self.conn.get("cm:{}:exit_with_error".format(cms[i].name)) 61 | == "0" 62 | ) 63 | 64 | def test_fixture(self): 65 | cms = self._get_context_managers(1).pop() 66 | with cms: 67 | pass 68 | 69 | assert self.conn.get("cm:{}:enter".format(cms.name)) == "1" 70 | assert self.conn.get("cm:{}:exit".format(cms.name)) == "1" 71 | 72 | def test_single_context_manager(self): 73 | self._test_context_managers(1, simple_task) 74 | self._test_context_managers(1, exception_task, should_fail=True) 75 | 76 | def test_multiple_context_managers(self): 77 | self._test_context_managers(10, simple_task) 78 | self._test_context_managers(10, exception_task, should_fail=True) 79 | -------------------------------------------------------------------------------- /tests/test_lazy_init.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import tempfile 3 | 4 | from tasktiger import TaskTiger, Worker 5 | 6 | from tests.utils import TEST_TIGER_CONFIG, get_redis, setup_structlog 7 | 8 | tiger = TaskTiger(lazy_init=True) 9 | 10 | 11 | @tiger.task 12 | def lazy_task(filename): 13 | with open(filename, "w") as f: 14 | f.write("ok") 15 | 16 | 17 | def test_lazy_init(): 18 | setup_structlog() 19 | tiger.init(connection=get_redis(), config=TEST_TIGER_CONFIG) 20 | tiger.log.setLevel(logging.CRITICAL) 21 | with tempfile.NamedTemporaryFile() as f: 22 | lazy_task.delay(f.name) 23 | Worker(tiger).run(once=True) 24 | assert f.read().decode("utf8") == "ok" 25 | -------------------------------------------------------------------------------- /tests/test_logging.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import structlog 3 | 4 | from tasktiger import TaskTiger, Worker 5 | from tasktiger.logging import tasktiger_processor 6 | 7 | from .test_base import BaseTestCase 8 | from .utils import get_redis, get_tiger 9 | 10 | tiger = get_tiger() 11 | logger = structlog.getLogger("tasktiger") 12 | 13 | 14 | @pytest.fixture(autouse=True) 15 | def restore_structlog_config(): 16 | previous_config = structlog.get_config() 17 | 18 | try: 19 | yield 20 | finally: 21 | structlog.configure(**previous_config) 22 | 23 | 24 | def logging_task(): 25 | log = logger.info("simple task") 26 | 27 | # Confirm tasktiger_processor injected task id and queue name 28 | assert log[1]["task_id"] == tiger.current_task.id 29 | assert log[1]["task_func"] == "tests.test_logging:logging_task" 30 | assert log[1]["queue"] == "foo_qux" 31 | 32 | 33 | class TestLogging(BaseTestCase): 34 | """Test logging.""" 35 | 36 | def test_structlog_processor(self): 37 | # Use ReturnLogger for testing 38 | structlog.configure( 39 | processors=[tasktiger_processor], 40 | context_class=dict, 41 | logger_factory=structlog.ReturnLoggerFactory(), 42 | wrapper_class=structlog.stdlib.BoundLogger, 43 | cache_logger_on_first_use=True, 44 | ) 45 | 46 | # Run a simple task. Logging output is verified in 47 | # the task. 48 | self.tiger.delay(logging_task, queue="foo_qux") 49 | queues = self._ensure_queues(queued={"foo_qux": 1}) 50 | 51 | task = queues["queued"]["foo_qux"][0] 52 | assert task["func"] == "tests.test_logging:logging_task" 53 | 54 | Worker(self.tiger).run(once=True) 55 | self._ensure_queues(queued={"foo_qux": 0}) 56 | assert not self.conn.exists("t:task:%s" % task["id"]) 57 | 58 | 59 | class TestSetupStructlog(BaseTestCase): 60 | def test_setup_structlog_basic(self): 61 | conn = get_redis() 62 | tiger = TaskTiger(connection=conn, setup_structlog=True) 63 | assert tiger 64 | conn.close() 65 | # no errors on init, cool 66 | -------------------------------------------------------------------------------- /tests/test_migrations.py: -------------------------------------------------------------------------------- 1 | import uuid 2 | 3 | import pytest 4 | 5 | from tasktiger.migrations import migrate_executions_count 6 | 7 | from .test_base import BaseTestCase 8 | 9 | 10 | class TestMigrateExecutionsCount(BaseTestCase): 11 | def test_migrate_nothing(self): 12 | migrate_executions_count(self.tiger) 13 | assert self.conn.keys() == [] 14 | 15 | @pytest.mark.parametrize( 16 | "key", 17 | [ 18 | f"foot:task:{uuid.uuid4()}:executions", 19 | f"foo:t:task:{uuid.uuid4()}:executions", 20 | f"t:task:{uuid.uuid4()}:executionsfoo", 21 | f"t:task:{uuid.uuid4()}:executions:foo", 22 | f"t:task:{uuid.uuid4()}", 23 | ], 24 | ) 25 | def test_migrate_ignores_irrelevant_keys(self, key): 26 | self.conn.rpush(key, "{}") 27 | migrate_executions_count(self.tiger) 28 | 29 | assert self.conn.keys() == [key] 30 | 31 | def test_migrate(self): 32 | task_id_1 = uuid.uuid4() 33 | task_id_2 = uuid.uuid4() 34 | 35 | for __ in range(73): 36 | self.conn.rpush(f"t:task:{task_id_1}:executions", "{}") 37 | 38 | for __ in range(35): 39 | self.conn.rpush(f"t:task:{task_id_2}:executions", "{}") 40 | 41 | migrate_executions_count(self.tiger) 42 | assert self.conn.get(f"t:task:{task_id_1}:executions_count") == "73" 43 | assert self.conn.get(f"t:task:{task_id_2}:executions_count") == "35" 44 | 45 | def test_migrate_when_some_tasks_already_migrated(self): 46 | task_id_1 = uuid.uuid4() 47 | task_id_2 = uuid.uuid4() 48 | 49 | for __ in range(73): 50 | self.conn.rpush(f"t:task:{task_id_1}:executions", "{}") 51 | 52 | self.conn.set(f"t:task:{task_id_1}:executions_count", 91) 53 | 54 | for __ in range(35): 55 | self.conn.rpush(f"t:task:{task_id_2}:executions", "{}") 56 | 57 | migrate_executions_count(self.tiger) 58 | assert self.conn.get(f"t:task:{task_id_2}:executions_count") == "35" 59 | 60 | # looks migrated already - left untouched 61 | assert self.conn.get(f"t:task:{task_id_1}:executions_count") == "91" 62 | 63 | def test_migrate_when_counter_is_behind(self): 64 | task_id_1 = uuid.uuid4() 65 | task_id_2 = uuid.uuid4() 66 | 67 | for __ in range(73): 68 | self.conn.rpush(f"t:task:{task_id_1}:executions", "{}") 69 | 70 | self.conn.set(f"t:task:{task_id_1}:executions_count", 10) 71 | 72 | for __ in range(35): 73 | self.conn.rpush(f"t:task:{task_id_2}:executions", "{}") 74 | 75 | migrate_executions_count(self.tiger) 76 | assert self.conn.get(f"t:task:{task_id_2}:executions_count") == "35" 77 | 78 | # updated because the counter value was less than the actual count 79 | assert self.conn.get(f"t:task:{task_id_1}:executions_count") == "73" 80 | -------------------------------------------------------------------------------- /tests/test_periodic.py: -------------------------------------------------------------------------------- 1 | """Periodic task tests.""" 2 | 3 | import datetime 4 | import time 5 | 6 | from tasktiger import Task, Worker, cron_expr, periodic 7 | from tasktiger._internal import ( 8 | QUEUED, 9 | SCHEDULED, 10 | gen_unique_id, 11 | serialize_func_name, 12 | ) 13 | 14 | from .tasks_periodic import periodic_task, tiger 15 | from .test_base import BaseTestCase 16 | from .utils import sleep_until_next_second 17 | 18 | 19 | class TestPeriodicTasks(BaseTestCase): 20 | def test_periodic_schedule(self): 21 | """ 22 | Test the periodic() schedule function. 23 | """ 24 | dt = datetime.datetime(2010, 1, 1) 25 | 26 | f = periodic(seconds=1) 27 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 0, 0, 1) 28 | 29 | f = periodic(minutes=1) 30 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 0, 1) 31 | 32 | f = periodic(hours=1) 33 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 1) 34 | 35 | f = periodic(days=1) 36 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 2) 37 | 38 | f = periodic(weeks=1) 39 | # 2010-01-02 is a Saturday 40 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 2) 41 | 42 | f = periodic(weeks=1, start_date=datetime.datetime(2000, 1, 2)) 43 | # 2000-01-02 and 2010-01-02 are Sundays 44 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 3) 45 | 46 | f = periodic(seconds=1, minutes=2, hours=3, start_date=dt) 47 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 3, 2, 1) 48 | # Make sure we return the start_date if the current date is earlier. 49 | assert f[0](datetime.datetime(1990, 1, 1), *f[1]) == dt 50 | 51 | f = periodic(minutes=1, end_date=dt) 52 | assert f[0]( 53 | datetime.datetime(2009, 12, 31, 23, 58), *f[1] 54 | ) == datetime.datetime(2009, 12, 31, 23, 59) 55 | 56 | f = periodic(minutes=1, end_date=dt) 57 | assert f[0]( 58 | datetime.datetime(2009, 12, 31, 23, 59), *f[1] 59 | ) == datetime.datetime(2010, 1, 1, 0, 0) 60 | 61 | f = periodic(minutes=1, end_date=dt) 62 | assert f[0](datetime.datetime(2010, 1, 1, 0, 0), *f[1]) is None 63 | 64 | f = periodic(minutes=1, end_date=dt) 65 | assert f[0](datetime.datetime(2010, 1, 1, 0, 1), *f[1]) is None 66 | 67 | def test_cron_schedule(self): 68 | """ 69 | Test the cron_expr() schedule function. 70 | """ 71 | dt = datetime.datetime(2010, 1, 1) 72 | 73 | f = cron_expr("* * * * *") 74 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 0, 1) 75 | 76 | f = cron_expr("0 * * * *") 77 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 1) 78 | 79 | f = cron_expr("0 0 * * *") 80 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 2) 81 | 82 | f = cron_expr("0 0 * * 6") 83 | # 2010-01-02 is a Saturday 84 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 2) 85 | 86 | f = cron_expr("0 0 * * 0", start_date=datetime.datetime(2000, 1, 2)) 87 | # 2000-01-02 is a Sunday and 2010-01-02 is a Saturday 88 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 3) 89 | 90 | f = cron_expr("2 3 * * *", start_date=dt) 91 | assert f[0](dt, *f[1]) == datetime.datetime(2010, 1, 1, 3, 2) 92 | # Make sure we return the start_date if the current date is earlier. 93 | assert f[0](datetime.datetime(1990, 1, 1), *f[1]) == dt 94 | 95 | f = cron_expr("* * * * *", end_date=dt) 96 | assert f[0]( 97 | datetime.datetime(2009, 12, 31, 23, 58), *f[1] 98 | ) == datetime.datetime(2009, 12, 31, 23, 59) 99 | 100 | f = cron_expr("* * * * *", end_date=dt) 101 | assert f[0]( 102 | datetime.datetime(2009, 12, 31, 23, 59), *f[1] 103 | ) == datetime.datetime(2010, 1, 1, 0, 0) 104 | 105 | f = cron_expr("* * * * *", end_date=dt) 106 | assert f[0](datetime.datetime(2010, 1, 1, 0, 0), *f[1]) is None 107 | 108 | f = cron_expr("* * * * *", end_date=dt) 109 | assert f[0](datetime.datetime(2010, 1, 1, 0, 1), *f[1]) is None 110 | 111 | def test_periodic_execution(self): 112 | """ 113 | Test periodic task execution. 114 | 115 | Test periodic_task() runs as expected and periodic_task_ignore() 116 | is not queued. 117 | """ 118 | # Sleep until the next second to ensure we have enough time to start 119 | # the worker and get the periodic task queued before the following 120 | # second starts. 121 | sleep_until_next_second() 122 | 123 | # After the first worker run, the periodic task will be queued. 124 | # Note that since periodic tasks register with the Tiger instance, it 125 | # must be the same instance that was used to decorate the task. We 126 | # therefore use `tiger` from the tasks module instead of `self.tiger`. 127 | self._ensure_queues() 128 | Worker(tiger).run(once=True) 129 | 130 | # NOTE: When the worker is started just before the second elapses, 131 | # it's possible that the periodic task is in "queued" state instead 132 | # of "scheduled" to ensure immediate execution. We capture this 133 | # condition by running the task, and retry. 134 | try: 135 | self._ensure_queues(scheduled={"periodic": 1}) 136 | except AssertionError: 137 | Worker(tiger).run(once=True) 138 | self._ensure_queues(scheduled={"periodic": 1}) 139 | assert int(self.conn.get("period_count")) == 1 140 | self.conn.delete("period_count") 141 | 142 | def ensure_run(n): 143 | # Run worker twice (once to move from scheduled to queued, and once 144 | # to execute the task) 145 | Worker(tiger).run(once=True) 146 | self._ensure_queues(queued={"periodic": 1}) 147 | Worker(tiger).run(once=True) 148 | self._ensure_queues(scheduled={"periodic": 1}) 149 | 150 | assert int(self.conn.get("period_count")) == n 151 | 152 | # The task is requeued for the next period 153 | self._ensure_queues(scheduled={"periodic": 1}) 154 | 155 | # Sleep until the next second 156 | sleep_until_next_second() 157 | 158 | ensure_run(1) 159 | 160 | # Within less than a second, the task will be processed again. 161 | time.sleep(1) 162 | 163 | ensure_run(2) 164 | 165 | def test_periodic_execution_unique_ids(self): 166 | """ 167 | Test that periodic tasks generate the same unique ids 168 | 169 | When a periodic task is scheduled initially as part of worker startup 170 | vs re-scheduled from within python the unique id generated should be 171 | the same. If they aren't it could result in duplicate tasks. 172 | """ 173 | # Sleep until the next second 174 | sleep_until_next_second() 175 | 176 | # After the first worker run, the periodic task will be queued. 177 | # Note that since periodic tasks register with the Tiger instance, it 178 | # must be the same instance that was used to decorate the task. We 179 | # therefore use `tiger` from the tasks module instead of `self.tiger`. 180 | self._ensure_queues() 181 | Worker(tiger).run(once=True) 182 | self._ensure_queues(scheduled={"periodic": 1}) 183 | time.sleep(1) 184 | Worker(tiger).run(once=True) 185 | self._ensure_queues(queued={"periodic": 1}) 186 | 187 | # generate the expected unique id 188 | expected_unique_id = gen_unique_id( 189 | serialize_func_name(periodic_task), [], {} 190 | ) 191 | 192 | # pull task out of the queue by id. If found, then the id is correct 193 | task = Task.from_id(tiger, "periodic", QUEUED, expected_unique_id) 194 | assert task is not None 195 | 196 | # execute and reschedule the task 197 | self._ensure_queues(queued={"periodic": 1}) 198 | Worker(tiger).run(once=True) 199 | self._ensure_queues(scheduled={"periodic": 1}) 200 | 201 | # wait for the task to need to be queued 202 | time.sleep(1) 203 | Worker(tiger).run(once=True) 204 | self._ensure_queues(queued={"periodic": 1}) 205 | 206 | # The unique id shouldn't change between executions. Try finding the 207 | # task by id again 208 | task = Task.from_id(tiger, "periodic", QUEUED, expected_unique_id) 209 | assert task is not None 210 | 211 | def test_periodic_execution_unique_ids_manual_scheduling(self): 212 | """ 213 | Periodic tasks should have the same unique ids when manually scheduled 214 | 215 | When a periodic task is scheduled initially as part of worker startup 216 | vs ``.delay``'d manually, the unique id generated should be the same. 217 | If they aren't it could result in duplicate tasks. 218 | """ 219 | # Sleep until the next second 220 | sleep_until_next_second() 221 | 222 | # After the first worker run, the periodic task will be queued. 223 | # Note that since periodic tasks register with the Tiger instance, it 224 | # must be the same instance that was used to decorate the task. We 225 | # therefore use `tiger` from the tasks module instead of `self.tiger`. 226 | self._ensure_queues() 227 | Worker(tiger).run(once=True) 228 | self._ensure_queues(scheduled={"periodic": 1}) 229 | time.sleep(1) 230 | Worker(tiger).run(once=True) 231 | self._ensure_queues(queued={"periodic": 1}) 232 | 233 | # schedule the task manually 234 | periodic_task.delay() 235 | 236 | # make sure a duplicate wasn't scheduled 237 | self._ensure_queues(queued={"periodic": 1}) 238 | 239 | def test_periodic_execution_unique_ids_self_correct(self): 240 | """ 241 | Test that periodic tasks will self-correct unique ids 242 | """ 243 | # Sleep until the next second 244 | sleep_until_next_second() 245 | 246 | # generate the ids 247 | correct_unique_id = gen_unique_id( 248 | serialize_func_name(periodic_task), [], {} 249 | ) 250 | malformed_unique_id = gen_unique_id( 251 | serialize_func_name(periodic_task), None, None 252 | ) 253 | 254 | task = Task(tiger, func=periodic_task) 255 | 256 | # patch the id to something slightly wrong 257 | assert task.id == correct_unique_id 258 | task._data["id"] = malformed_unique_id 259 | assert task.id == malformed_unique_id 260 | 261 | # schedule the task 262 | task.delay() 263 | self._ensure_queues(queued={"periodic": 1}) 264 | 265 | # pull task out of the queue by the malformed id 266 | task = Task.from_id(tiger, "periodic", QUEUED, malformed_unique_id) 267 | assert task is not None 268 | 269 | Worker(tiger).run(once=True) 270 | self._ensure_queues(scheduled={"periodic": 1}) 271 | 272 | # pull task out of the queue by the self-corrected id 273 | task = Task.from_id(tiger, "periodic", SCHEDULED, correct_unique_id) 274 | assert task is not None 275 | 276 | def test_successful_execution_clears_executions_from_retries(self): 277 | """ 278 | Ensure previous executions from retries are cleared after a successful 279 | execution. 280 | """ 281 | sleep_until_next_second() 282 | 283 | # Queue the periodic task. 284 | self._ensure_queues() 285 | Worker(tiger).run(once=True) 286 | 287 | # Prepare to execute the periodic task (as retriable failure). 288 | tiger.connection.set("fail-periodic-task", "retriable") 289 | n_total, tasks = Task.tasks_from_queue(tiger, "periodic", SCHEDULED) 290 | task_id = tasks[0].id 291 | time.sleep(1) 292 | 293 | # Queue the periodic task. 294 | self._ensure_queues(scheduled={"periodic": 1}) 295 | Worker(tiger).run(once=True) 296 | 297 | # Run the failing periodic task. 298 | self._ensure_queues(queued={"periodic": 1}) 299 | Worker(tiger).run(once=True) 300 | 301 | task = Task.from_id( 302 | tiger, "periodic", SCHEDULED, task_id, load_executions=10 303 | ) 304 | assert len(task.executions) == 1 305 | 306 | tiger.connection.delete("fail-periodic-task") 307 | time.sleep(1) 308 | 309 | # Queue the periodic task. 310 | self._ensure_queues(scheduled={"periodic": 1}) 311 | Worker(tiger).run(once=True) 312 | 313 | # Run the successful periodic task. 314 | self._ensure_queues(queued={"periodic": 1}) 315 | Worker(tiger).run(once=True) 316 | 317 | # Ensure we cleared any previous executions. 318 | task = Task.from_id( 319 | tiger, "periodic", SCHEDULED, task_id, load_executions=10 320 | ) 321 | assert len(task.executions) == 0 322 | 323 | def test_successful_execution_doesnt_clear_previous_errors(self): 324 | """ 325 | Ensure previous executions are not cleared if we have had non-retriable 326 | errors. 327 | """ 328 | sleep_until_next_second() 329 | 330 | # Queue the periodic task. 331 | self._ensure_queues() 332 | Worker(tiger).run(once=True) 333 | 334 | # Prepare to execute the periodic task (as permanent failure). 335 | tiger.connection.set("fail-periodic-task", "permanent") 336 | n_total, tasks = Task.tasks_from_queue(tiger, "periodic", SCHEDULED) 337 | task_id = tasks[0].id 338 | time.sleep(1) 339 | 340 | # Queue the periodic task. 341 | self._ensure_queues(scheduled={"periodic": 1}) 342 | Worker(tiger).run(once=True) 343 | 344 | # Run the failing periodic task. 345 | self._ensure_queues(queued={"periodic": 1}) 346 | Worker(tiger).run(once=True) 347 | 348 | task = Task.from_id( 349 | tiger, "periodic", SCHEDULED, task_id, load_executions=10 350 | ) 351 | assert len(task.executions) == 1 352 | 353 | tiger.connection.delete("fail-periodic-task") 354 | time.sleep(1) 355 | 356 | # Queue the periodic task. 357 | self._ensure_queues(scheduled={"periodic": 1}, error={"periodic": 1}) 358 | Worker(tiger).run(once=True) 359 | 360 | # Run the successful periodic task. 361 | self._ensure_queues(queued={"periodic": 1}, error={"periodic": 1}) 362 | Worker(tiger).run(once=True) 363 | 364 | # Ensure we didn't clear previous executions. 365 | task = Task.from_id( 366 | tiger, "periodic", SCHEDULED, task_id, load_executions=10 367 | ) 368 | assert len(task.executions) == 1 369 | -------------------------------------------------------------------------------- /tests/test_queue_size.py: -------------------------------------------------------------------------------- 1 | """Test max queue size limits.""" 2 | 3 | import datetime 4 | import os 5 | import signal 6 | import time 7 | from multiprocessing import Process 8 | 9 | import pytest 10 | 11 | from tasktiger import Task, Worker 12 | from tasktiger.exceptions import QueueFullException 13 | 14 | from .config import DELAY 15 | from .tasks import decorated_task_max_queue_size, simple_task, sleep_task 16 | from .test_base import BaseTestCase 17 | from .utils import external_worker 18 | 19 | 20 | class TestMaxQueue(BaseTestCase): 21 | """TaskTiger test max queue size.""" 22 | 23 | def test_task_simple_delay(self): 24 | """Test enforcing max queue size using delay function.""" 25 | 26 | self.tiger.delay(simple_task, queue="a", max_queue_size=1) 27 | self._ensure_queues(queued={"a": 1}) 28 | 29 | # Queue size would be 2 so it should fail 30 | with pytest.raises(QueueFullException): 31 | self.tiger.delay(simple_task, queue="a", max_queue_size=1) 32 | 33 | # Process first task and then queuing a second should succeed 34 | Worker(self.tiger).run(once=True, force_once=True) 35 | self.tiger.delay(simple_task, queue="a", max_queue_size=1) 36 | self._ensure_queues(queued={"a": 1}) 37 | 38 | def test_task_decorated(self): 39 | """Test max queue size with decorator.""" 40 | 41 | decorated_task_max_queue_size.delay() 42 | self._ensure_queues(queued={"default": 1}) 43 | 44 | with pytest.raises(QueueFullException): 45 | decorated_task_max_queue_size.delay() 46 | 47 | def test_task_all_states(self): 48 | """Test max queue size with tasks in all three states.""" 49 | 50 | # Active 51 | task = Task(self.tiger, sleep_task, queue="a") 52 | task.delay() 53 | self._ensure_queues(queued={"a": 1}) 54 | 55 | # Start a worker and wait until it starts processing. 56 | worker = Process(target=external_worker) 57 | worker.start() 58 | time.sleep(DELAY) 59 | 60 | # Kill the worker while it's still processing the task. 61 | os.kill(worker.pid, signal.SIGKILL) 62 | self._ensure_queues(active={"a": 1}) 63 | 64 | # Scheduled 65 | self.tiger.delay( 66 | simple_task, 67 | queue="a", 68 | max_queue_size=3, 69 | when=datetime.timedelta(seconds=10), 70 | ) 71 | 72 | # Queued 73 | self.tiger.delay(simple_task, queue="a", max_queue_size=3) 74 | 75 | self._ensure_queues( 76 | active={"a": 1}, queued={"a": 1}, scheduled={"a": 1} 77 | ) 78 | 79 | # Should fail to queue task to run immediately 80 | with pytest.raises(QueueFullException): 81 | self.tiger.delay(simple_task, queue="a", max_queue_size=3) 82 | 83 | # Should fail to queue task to run in the future 84 | with pytest.raises(QueueFullException): 85 | self.tiger.delay( 86 | simple_task, 87 | queue="a", 88 | max_queue_size=3, 89 | when=datetime.timedelta(seconds=10), 90 | ) 91 | 92 | 93 | class TestQueueSizes: 94 | @pytest.fixture 95 | def queue_sample_tasks(self, tiger): 96 | tiger.delay(simple_task) 97 | tiger.delay(simple_task) 98 | tiger.delay(simple_task, queue="other") 99 | tiger.delay(simple_task, when=datetime.timedelta(seconds=60)) 100 | 101 | def test_get_total_queue_size(self, tiger, queue_sample_tasks): 102 | assert tiger.get_total_queue_size("other") == 1 103 | assert tiger.get_total_queue_size("default") == 3 104 | 105 | def test_get_queue_sizes(self, tiger, queue_sample_tasks): 106 | assert tiger.get_queue_sizes("default") == { 107 | "active": 0, 108 | "queued": 2, 109 | "scheduled": 1, 110 | } 111 | assert tiger.get_queue_sizes("other") == { 112 | "active": 0, 113 | "queued": 1, 114 | "scheduled": 0, 115 | } 116 | 117 | def test_get_sizes_for_queues_and_states(self, tiger, queue_sample_tasks): 118 | assert tiger.get_sizes_for_queues_and_states( 119 | [ 120 | ("default", "queued"), 121 | ("default", "scheduled"), 122 | ("other", "queued"), 123 | ("other", "scheduled"), 124 | ] 125 | ) == [2, 1, 1, 0] 126 | -------------------------------------------------------------------------------- /tests/test_redis_scripts.py: -------------------------------------------------------------------------------- 1 | from .utils import get_tiger 2 | 3 | 4 | class TestRedisScripts: 5 | def setup_method(self, method): 6 | self.tiger = get_tiger() 7 | self.conn = self.tiger.connection 8 | self.conn.flushdb() 9 | self.scripts = self.tiger.scripts 10 | 11 | def teardown_method(self, method): 12 | self.conn.flushdb() 13 | self.conn.close() 14 | # Force disconnect so we don't get Too many open files 15 | self.conn.connection_pool.disconnect() 16 | 17 | def _test_zadd(self, mode): 18 | self.conn.zadd("z", {"key1": 2}) 19 | self.scripts.zadd("z", 4, "key1", mode=mode) 20 | self.scripts.zadd("z", 3, "key1", mode=mode) 21 | self.scripts.zadd("z", 1, "key2", mode=mode) 22 | self.scripts.zadd("z", 2, "key2", mode=mode) 23 | self.scripts.zadd("z", 0, "key2", mode=mode) 24 | return self.conn.zrange("z", 0, -1, withscores=True) 25 | 26 | def test_zadd_nx(self): 27 | entries = self._test_zadd("nx") 28 | assert entries == [("key2", 1.0), ("key1", 2.0)] 29 | 30 | def test_zadd_xx(self): 31 | entries = self._test_zadd("xx") 32 | assert entries == [("key1", 3.0)] 33 | 34 | def test_zadd_min(self): 35 | entries = self._test_zadd("min") 36 | assert entries == [("key2", 0.0), ("key1", 2.0)] 37 | 38 | def test_zadd_max(self): 39 | entries = self._test_zadd("max") 40 | assert entries == [("key2", 2.0), ("key1", 4.0)] 41 | 42 | def test_zpoppush_1(self): 43 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 44 | result = self.scripts.zpoppush("src", "dst", 3, None, 10) 45 | assert result == ["a", "b", "c"] 46 | 47 | src = self.conn.zrange("src", 0, -1, withscores=True) 48 | assert src == [("d", 4.0)] 49 | 50 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 51 | assert dst == [("a", 10.0), ("b", 10.0), ("c", 10.0)] 52 | 53 | def test_zpoppush_2(self): 54 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 55 | result = self.scripts.zpoppush("src", "dst", 100, None, 10) 56 | assert result == ["a", "b", "c", "d"] 57 | 58 | src = self.conn.zrange("src", 0, -1, withscores=True) 59 | assert src == [] 60 | 61 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 62 | assert dst == [("a", 10.0), ("b", 10.0), ("c", 10.0), ("d", 10.0)] 63 | 64 | def test_zpoppush_3(self): 65 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 66 | result = self.scripts.zpoppush("src", "dst", 3, 2, 10) 67 | assert result == ["a", "b"] 68 | 69 | src = self.conn.zrange("src", 0, -1, withscores=True) 70 | assert src == [("c", 3.0), ("d", 4.0)] 71 | 72 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 73 | assert dst == [("a", 10.0), ("b", 10.0)] 74 | 75 | def test_zpoppush_withscores_1(self): 76 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 77 | result = self.scripts.zpoppush( 78 | "src", "dst", 3, None, 10, withscores=True 79 | ) 80 | assert result == ["a", "1", "b", "2", "c", "3"] 81 | 82 | src = self.conn.zrange("src", 0, -1, withscores=True) 83 | assert src == [("d", 4.0)] 84 | 85 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 86 | assert dst == [("a", 10.0), ("b", 10.0), ("c", 10.0)] 87 | 88 | def test_zpoppush_withscores_2(self): 89 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 90 | result = self.scripts.zpoppush( 91 | "src", "dst", 100, None, 10, withscores=True 92 | ) 93 | assert result == ["a", "1", "b", "2", "c", "3", "d", "4"] 94 | 95 | src = self.conn.zrange("src", 0, -1, withscores=True) 96 | assert src == [] 97 | 98 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 99 | assert dst == [("a", 10.0), ("b", 10.0), ("c", 10.0), ("d", 10.0)] 100 | 101 | def test_zpoppush_withscores_3(self): 102 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 103 | result = self.scripts.zpoppush("src", "dst", 3, 2, 10, withscores=True) 104 | assert result == ["a", "1", "b", "2"] 105 | 106 | src = self.conn.zrange("src", 0, -1, withscores=True) 107 | assert src == [("c", 3.0), ("d", 4.0)] 108 | 109 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 110 | assert dst == [("a", 10.0), ("b", 10.0)] 111 | 112 | def test_zpoppush_on_success_1(self, **kwargs): 113 | """ 114 | 2 out of 4 items moved, so add_set contains "val" 115 | """ 116 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 117 | # Whether the members are in the destination ZSET doesn't make any 118 | # difference here. 119 | self.conn.zadd("dst", {"a": 5, "b": 5}) 120 | self.conn.sadd("remove_set", "val") 121 | result = self.scripts.zpoppush( 122 | "src", 123 | "dst", 124 | count=2, 125 | score=None, 126 | new_score=10, 127 | on_success=("update_sets", "val", "remove_set", "add_set"), 128 | **kwargs 129 | ) 130 | assert result == ["a", "b"] 131 | 132 | src = self.conn.zrange("src", 0, -1, withscores=True) 133 | assert src == [("c", 3.0), ("d", 4.0)] 134 | 135 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 136 | assert dst == [("a", 10.0), ("b", 10.0)] 137 | 138 | assert self.conn.smembers("remove_set") == {"val"} 139 | assert self.conn.smembers("add_set") == {"val"} 140 | 141 | def test_zpoppush_on_success_2(self, **kwargs): 142 | """ 143 | 0 out of 4 items moved, so no sets were changed 144 | """ 145 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 146 | self.conn.sadd("remove_set", "val") 147 | result = self.scripts.zpoppush( 148 | "src", 149 | "dst", 150 | count=2, 151 | score=0, 152 | new_score=10, 153 | on_success=("update_sets", "val", "remove_set", "add_set"), 154 | **kwargs 155 | ) 156 | assert result == [] 157 | 158 | src = self.conn.zrange("src", 0, -1, withscores=True) 159 | assert src == [("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0)] 160 | 161 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 162 | assert dst == [] 163 | 164 | assert self.conn.smembers("remove_set") == {"val"} 165 | assert self.conn.smembers("add_set") == set() 166 | 167 | def test_zpoppush_on_success_3(self, **kwargs): 168 | """ 169 | 4 out of 4 items moved, so both sets were changed 170 | """ 171 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 172 | self.conn.sadd("remove_set", "val") 173 | result = self.scripts.zpoppush( 174 | "src", 175 | "dst", 176 | count=4, 177 | score=None, 178 | new_score=10, 179 | on_success=("update_sets", "val", "remove_set", "add_set"), 180 | **kwargs 181 | ) 182 | assert result == ["a", "b", "c", "d"] 183 | 184 | src = self.conn.zrange("src", 0, -1, withscores=True) 185 | assert src == [] 186 | 187 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 188 | assert dst == [("a", 10), ("b", 10), ("c", 10), ("d", 10)] 189 | 190 | assert self.conn.smembers("remove_set") == set() 191 | assert self.conn.smembers("add_set") == {"val"} 192 | 193 | def test_zpoppush_ignore_if_exists_1(self): 194 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 195 | # Members that are in the destination ZSET are not updated here. 196 | self.conn.zadd("dst", {"a": 5, "b": 5}) 197 | self.conn.sadd("remove_set", "val") 198 | result = self.scripts.zpoppush( 199 | "src", 200 | "dst", 201 | count=2, 202 | score=None, 203 | new_score=10, 204 | on_success=("update_sets", "val", "remove_set", "add_set"), 205 | if_exists=("noupdate",), 206 | ) 207 | assert result == [] 208 | 209 | src = self.conn.zrange("src", 0, -1, withscores=True) 210 | assert src == [("c", 3.0), ("d", 4.0)] 211 | 212 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 213 | assert dst == [("a", 5.0), ("b", 5.0)] 214 | 215 | assert self.conn.smembers("remove_set") == {"val"} 216 | assert self.conn.smembers("add_set") == {"val"} 217 | 218 | def test_zpoppush_ignore_if_exists_2(self): 219 | self.test_zpoppush_on_success_2(if_exists=("noupdate",)) 220 | 221 | def test_zpoppush_ignore_if_exists_3(self): 222 | self.test_zpoppush_on_success_3(if_exists=("noupdate",)) 223 | 224 | def _test_zpoppush_min_if_exists(self, expected_if_exists_score): 225 | self.conn.zadd("src", {"a": 1, "b": 2, "c": 3, "d": 4}) 226 | # Members that are in the destination ZSET are added to the if_exists 227 | # ZSET. 228 | self.conn.zadd("dst", {"a": 5}) 229 | self.conn.sadd("remove_set", "val") 230 | result = self.scripts.zpoppush( 231 | "src", 232 | "dst", 233 | count=2, 234 | score=None, 235 | new_score=10, 236 | on_success=( 237 | "update_sets", 238 | "val", 239 | "remove_set", 240 | "add_set", 241 | "add_set_if_exists", 242 | ), 243 | if_exists=("add", "if_exists", 20, "min"), 244 | ) 245 | assert result == ["b"] 246 | 247 | src = self.conn.zrange("src", 0, -1, withscores=True) 248 | assert src == [("c", 3.0), ("d", 4.0)] 249 | 250 | dst = self.conn.zrange("dst", 0, -1, withscores=True) 251 | assert dst == [("a", 5.0), ("b", 10.0)] 252 | 253 | if_exists = self.conn.zrange("if_exists", 0, -1, withscores=True) 254 | assert if_exists == [("a", expected_if_exists_score)] 255 | 256 | assert self.conn.smembers("remove_set") == {"val"} 257 | assert self.conn.smembers("add_set") == {"val"} 258 | assert self.conn.smembers("add_set_if_exists") == {"val"} 259 | 260 | def test_zpoppush_min_if_exists_1(self): 261 | self._test_zpoppush_min_if_exists(20) 262 | 263 | def test_zpoppush_min_if_exists_2(self): 264 | self.conn.zadd("if_exists", {"a": 10}) 265 | self._test_zpoppush_min_if_exists(10) 266 | 267 | def test_zpoppush_min_if_exists_3(self): 268 | self.conn.zadd("if_exists", {"a": 30}) 269 | self._test_zpoppush_min_if_exists(20) 270 | 271 | def test_srem_if_not_exists_1(self): 272 | self.conn.sadd("set", "member") 273 | result = self.scripts.srem_if_not_exists("set", "member", "other_key") 274 | assert result == 1 275 | assert self.conn.smembers("set") == set() 276 | 277 | def test_srem_if_not_exists_2(self): 278 | self.conn.sadd("set", "member") 279 | self.conn.set("other_key", 0) 280 | result = self.scripts.srem_if_not_exists("set", "member", "other_key") 281 | assert result == 0 282 | assert self.conn.smembers("set") == {"member"} 283 | 284 | def test_delete_if_not_in_zsets_1(self): 285 | self.conn.set("foo", 0) 286 | self.conn.set("bar", 0) 287 | self.conn.set("baz", 0) 288 | 289 | self.conn.zadd("z2", {"other": 0}) 290 | result = self.scripts.delete_if_not_in_zsets( 291 | to_delete=["foo", "bar"], value="member", zsets=["z1", "z2"] 292 | ) 293 | assert result == 2 294 | assert self.conn.exists("foo") == 0 295 | assert self.conn.exists("bar") == 0 296 | assert self.conn.exists("baz") == 1 297 | assert self.conn.exists("z2") == 1 298 | 299 | def test_delete_if_not_in_zsets_2(self): 300 | self.conn.set("foo", 0) 301 | self.conn.set("bar", 0) 302 | 303 | self.conn.zadd("z2", {"member": 0}) 304 | result = self.scripts.delete_if_not_in_zsets( 305 | to_delete=["foo", "bar"], value="member", zsets=["z1", "z2"] 306 | ) 307 | assert result == 0 308 | assert self.conn.exists("foo") == 1 309 | assert self.conn.exists("bar") == 1 310 | assert self.conn.exists("z2") == 1 311 | 312 | def test_delete_if_not_in_zsets_3(self): 313 | self.conn.set("foo", 0) 314 | self.conn.set("bar", 0) 315 | 316 | result = self.scripts.delete_if_not_in_zsets( 317 | to_delete=["foo", "bar"], value="member", zsets=[] 318 | ) 319 | assert result == 2 320 | assert self.conn.exists("foo") == 0 321 | assert self.conn.exists("bar") == 0 322 | 323 | def test_get_expired_tasks(self): 324 | self.conn.sadd("t:active", "q1", "q2", "q3", "q4") 325 | self.conn.zadd("t:active:q1", {"t1": 500}) 326 | self.conn.zadd("t:active:q1", {"t2": 1000}) 327 | self.conn.zadd("t:active:q1", {"t3": 1500}) 328 | self.conn.zadd("t:active:q2", {"t4": 1200}) 329 | self.conn.zadd("t:active:q3", {"t5": 1800}) 330 | self.conn.zadd("t:active:q4", {"t6": 200}) 331 | 332 | expired_task_set = {("q1", "t1"), ("q1", "t2"), ("q4", "t6")} 333 | 334 | result = self.scripts.get_expired_tasks("t", 1000, 10) 335 | assert len(result) == 3 336 | assert set(result) == expired_task_set 337 | 338 | for batch_size in range(1, 4): 339 | result = self.scripts.get_expired_tasks("t", 1000, batch_size) 340 | assert len(result) == batch_size 341 | assert set(result) & expired_task_set == set(result) 342 | -------------------------------------------------------------------------------- /tests/test_semaphore.py: -------------------------------------------------------------------------------- 1 | """Test Redis Semaphore lock.""" 2 | import datetime 3 | import time 4 | 5 | import pytest 6 | from freezefrog import FreezeTime 7 | 8 | from tasktiger.redis_semaphore import Semaphore 9 | 10 | from .utils import get_tiger 11 | 12 | 13 | class TestSemaphore: 14 | """Test Redis Semaphores.""" 15 | 16 | def setup_method(self, method): 17 | """Test setup.""" 18 | 19 | self.tiger = get_tiger() 20 | self.conn = self.tiger.connection 21 | self.conn.flushdb() 22 | 23 | def teardown_method(self, method): 24 | """Test teardown.""" 25 | 26 | self.conn.flushdb() 27 | self.conn.close() 28 | # Force disconnect so we don't get Too many open files 29 | self.conn.connection_pool.disconnect() 30 | 31 | def test_simple_semaphore(self): 32 | """Test semaphore.""" 33 | 34 | semaphore1 = Semaphore( 35 | self.conn, "test_key", "id_1", max_locks=1, timeout=10 36 | ) 37 | semaphore2 = Semaphore( 38 | self.conn, "test_key", "id_2", max_locks=1, timeout=10 39 | ) 40 | 41 | # Get lock and then release 42 | with FreezeTime(datetime.datetime(2014, 1, 1)): 43 | acquired, locks = semaphore1.acquire() 44 | assert acquired 45 | assert locks == 1 46 | semaphore1.release() 47 | 48 | # Get a new lock after releasing old 49 | with FreezeTime(datetime.datetime(2014, 1, 1)): 50 | acquired, locks = semaphore2.acquire() 51 | assert acquired 52 | assert locks == 1 53 | 54 | # Fail getting second lock while still inside time out period 55 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 9)): 56 | acquired, locks = semaphore1.acquire() 57 | assert not acquired 58 | assert locks == 1 59 | 60 | # Successful getting lock after semaphore2 times out 61 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 10)): 62 | acquired, locks = semaphore1.acquire() 63 | assert acquired 64 | assert locks == 1 65 | 66 | def test_multiple_locks(self): 67 | semaphore1 = Semaphore( 68 | self.conn, "test_key", "id_1", max_locks=2, timeout=10 69 | ) 70 | semaphore2 = Semaphore( 71 | self.conn, "test_key", "id_2", max_locks=2, timeout=10 72 | ) 73 | semaphore3 = Semaphore( 74 | self.conn, "test_key", "id_3", max_locks=2, timeout=10 75 | ) 76 | 77 | # First two locks should be acquired 78 | with FreezeTime(datetime.datetime(2014, 1, 1)): 79 | acquired, locks = semaphore1.acquire() 80 | assert acquired 81 | assert locks == 1 82 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 4)): 83 | acquired, locks = semaphore2.acquire() 84 | assert acquired 85 | assert locks == 2 86 | 87 | # Third lock should fail 88 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 6)): 89 | acquired, locks = semaphore3.acquire() 90 | assert not acquired 91 | assert locks == 2 92 | 93 | semaphore2.release() 94 | 95 | # Releasing one of the existing locks should let a new lock succeed 96 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 9)): 97 | acquired, locks = semaphore3.acquire() 98 | assert acquired 99 | assert locks == 2 100 | 101 | def test_semaphores_renew(self): 102 | semaphore1 = Semaphore( 103 | self.conn, "test_key", "id_1", max_locks=1, timeout=10 104 | ) 105 | semaphore2 = Semaphore( 106 | self.conn, "test_key", "id_2", max_locks=1, timeout=10 107 | ) 108 | 109 | with FreezeTime(datetime.datetime(2014, 1, 1)): 110 | acquired, locks = semaphore1.acquire() 111 | assert acquired 112 | assert locks == 1 113 | 114 | # Renew 5 seconds into lock timeout window 115 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 5)): 116 | acquired, locks = semaphore1.renew() 117 | assert acquired 118 | assert locks == 1 119 | 120 | # Fail getting a lock 121 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 14)): 122 | acquired, locks = semaphore2.acquire() 123 | assert not acquired 124 | assert locks == 1 125 | 126 | # Successful getting lock after renewed timeout window passes 127 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 15)): 128 | acquired, locks = semaphore2.acquire() 129 | assert acquired 130 | assert locks == 1 131 | 132 | # Fail renewing 133 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 15)): 134 | acquired, locks = semaphore1.renew() 135 | assert not acquired 136 | assert locks == 1 137 | 138 | # Test system lock shorter and longer than regular lock timeout 139 | @pytest.mark.parametrize("timeout", [8, 30]) 140 | def test_system_lock(self, timeout): 141 | semaphore1 = Semaphore( 142 | self.conn, "test_key", "id_1", max_locks=10, timeout=10 143 | ) 144 | 145 | with FreezeTime(datetime.datetime(2014, 1, 1)): 146 | Semaphore.set_system_lock(self.conn, "test_key", timeout) 147 | ttl = Semaphore.get_system_lock(self.conn, "test_key") 148 | assert ttl == time.time() + timeout 149 | 150 | # Should be blocked by system lock 151 | acquired, locks = semaphore1.acquire() 152 | assert not acquired 153 | assert locks == -1 154 | 155 | # System lock should still block other locks 1 second before it expires 156 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, timeout - 1)): 157 | acquired, locks = semaphore1.acquire() 158 | assert not acquired 159 | assert locks == -1 160 | 161 | # Wait for system lock to expire 162 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, timeout)): 163 | acquired, locks = semaphore1.acquire() 164 | assert acquired 165 | assert locks == 1 166 | -------------------------------------------------------------------------------- /tests/test_stats.py: -------------------------------------------------------------------------------- 1 | import time 2 | from unittest import mock 3 | 4 | import pytest 5 | 6 | from tasktiger.stats import StatsThread 7 | 8 | from tests.utils import get_tiger 9 | 10 | 11 | @pytest.fixture 12 | def tiger(): 13 | t = get_tiger() 14 | t.config["STATS_INTERVAL"] = 0.07 15 | return t 16 | 17 | 18 | def test_start_and_stop(tiger): 19 | stats = StatsThread(tiger) 20 | stats.compute_stats = mock.Mock() 21 | stats.start() 22 | 23 | time.sleep(0.22) 24 | stats.stop() 25 | 26 | assert len(stats.compute_stats.mock_calls) == 3 27 | 28 | # Stats are no longer being collected 29 | time.sleep(0.22) 30 | assert len(stats.compute_stats.mock_calls) == 3 31 | -------------------------------------------------------------------------------- /tests/test_task.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from tasktiger import Task, TaskNotFound 4 | 5 | from .tasks import simple_task 6 | from .utils import get_tiger 7 | 8 | 9 | @pytest.fixture 10 | def tiger(): 11 | return get_tiger() 12 | 13 | 14 | class TestTaskFromId: 15 | @pytest.fixture 16 | def queued_task(self, tiger): 17 | return tiger.delay(simple_task) 18 | 19 | def test_task_found(self, tiger, queued_task): 20 | task = Task.from_id(tiger, "default", "queued", queued_task.id) 21 | assert queued_task.id == task.id 22 | 23 | def test_task_wrong_state(self, tiger, queued_task): 24 | with pytest.raises(TaskNotFound): 25 | Task.from_id(tiger, "default", "active", queued_task.id) 26 | 27 | def test_task_wrong_queue(self, tiger, queued_task): 28 | with pytest.raises(TaskNotFound): 29 | Task.from_id(tiger, "other", "active", queued_task.id) 30 | 31 | 32 | class TestTaskMaxTrackedExecutions: 33 | def test_max_stored_executions_passed_to_tiger_delay(self, tiger): 34 | task = tiger.delay(simple_task, max_stored_executions=17) 35 | assert task.max_stored_executions == 17 36 | 37 | def test_max_stored_executions_passed_to_decorator(self, tiger): 38 | @tiger.task(max_stored_executions=17) 39 | def some_task(): 40 | pass 41 | 42 | task = some_task.delay() 43 | assert task.max_stored_executions == 17 44 | 45 | def test_max_stored_executions_overridden_in_tiger_delay(self, tiger): 46 | @tiger.task(max_stored_executions=17) 47 | def some_task(): 48 | pass 49 | 50 | task = tiger.delay(some_task, max_stored_executions=11) 51 | assert task.max_stored_executions == 11 52 | -------------------------------------------------------------------------------- /tests/test_workers.py: -------------------------------------------------------------------------------- 1 | """Test workers.""" 2 | 3 | import datetime 4 | import time 5 | from multiprocessing import Process 6 | 7 | import pytest 8 | from freezefrog import FreezeTime 9 | 10 | from tasktiger import Task, Worker 11 | from tasktiger._internal import ACTIVE 12 | from tasktiger.executor import SyncExecutor 13 | from tasktiger.worker import LOCK_REDIS_KEY 14 | 15 | from .config import DELAY 16 | from .tasks import ( 17 | exception_task, 18 | long_task_killed, 19 | long_task_ok, 20 | simple_task, 21 | sleep_task, 22 | system_exit_task, 23 | wait_for_long_task, 24 | ) 25 | from .test_base import BaseTestCase 26 | from .utils import external_worker 27 | 28 | 29 | class TestMaxWorkers(BaseTestCase): 30 | """Single Worker Queue tests.""" 31 | 32 | def test_max_workers(self): 33 | """Test Single Worker Queue.""" 34 | 35 | # Queue three tasks 36 | for i in range(0, 3): 37 | task = Task(self.tiger, long_task_ok, queue="a") 38 | task.delay() 39 | self._ensure_queues(queued={"a": 3}) 40 | 41 | # Start two workers and wait until they start processing. 42 | worker1 = Process( 43 | target=external_worker, 44 | kwargs={"worker_kwargs": {"max_workers_per_queue": 2}}, 45 | ) 46 | worker2 = Process( 47 | target=external_worker, 48 | kwargs={"worker_kwargs": {"max_workers_per_queue": 2}}, 49 | ) 50 | worker1.start() 51 | worker2.start() 52 | 53 | # Wait for both tasks to start 54 | wait_for_long_task() 55 | wait_for_long_task() 56 | 57 | # Verify they both are active 58 | self._ensure_queues(active={"a": 2}, queued={"a": 1}) 59 | 60 | # This worker should fail to get the queue lock and exit immediately 61 | worker = Worker(self.tiger) 62 | worker.max_workers_per_queue = 2 63 | worker.run(once=True, force_once=True) 64 | self._ensure_queues(active={"a": 2}, queued={"a": 1}) 65 | # Wait for external workers 66 | worker1.join() 67 | worker2.join() 68 | 69 | def test_single_worker_queue(self): 70 | """ 71 | Test Single Worker Queue. 72 | 73 | Single worker queues are the same as running with MAX_WORKERS_PER_QUEUE 74 | set to 1. 75 | """ 76 | 77 | # Queue two tasks 78 | task = Task(self.tiger, long_task_ok, queue="swq") 79 | task.delay() 80 | task = Task(self.tiger, long_task_ok, queue="swq") 81 | task.delay() 82 | self._ensure_queues(queued={"swq": 2}) 83 | 84 | # Start a worker and wait until it starts processing. 85 | # It should start processing one task and hold a lock on the queue 86 | worker = Process(target=external_worker) 87 | worker.start() 88 | 89 | # Wait for task to start 90 | wait_for_long_task() 91 | 92 | # This worker should fail to get the queue lock and exit immediately 93 | Worker(self.tiger).run(once=True, force_once=True) 94 | self._ensure_queues(active={"swq": 1}, queued={"swq": 1}) 95 | # Wait for external worker 96 | worker.join() 97 | 98 | # Clear out second task 99 | Worker(self.tiger).run(once=True, force_once=True) 100 | self.conn.delete("long_task_ok") 101 | 102 | # Retest using a non-single worker queue 103 | # Queue two tasks 104 | task = Task(self.tiger, long_task_ok, queue="not_swq") 105 | task.delay() 106 | task = Task(self.tiger, long_task_ok, queue="not_swq") 107 | task.delay() 108 | self._ensure_queues(queued={"not_swq": 2}) 109 | 110 | # Start a worker and wait until it starts processing. 111 | # It should start processing one task 112 | worker = Process(target=external_worker) 113 | worker.start() 114 | 115 | # Wait for task to start processing 116 | wait_for_long_task() 117 | 118 | # This worker should process the second task 119 | Worker(self.tiger).run(once=True, force_once=True) 120 | 121 | # Queues should be empty since the first task will have to 122 | # have finished before the second task finishes. 123 | self._ensure_queues() 124 | 125 | worker.join() 126 | 127 | def test_queue_system_lock(self): 128 | """Test queue system lock.""" 129 | 130 | with FreezeTime(datetime.datetime(2014, 1, 1)): 131 | # Queue three tasks 132 | for i in range(0, 3): 133 | task = Task(self.tiger, long_task_ok, queue="a") 134 | task.delay() 135 | self._ensure_queues(queued={"a": 3}) 136 | 137 | # Ensure we can process one 138 | worker = Worker(self.tiger) 139 | worker.max_workers_per_queue = 2 140 | worker.run(once=True, force_once=True) 141 | self._ensure_queues(queued={"a": 2}) 142 | 143 | # Set system lock so no processing should occur for 10 seconds 144 | self.tiger.set_queue_system_lock("a", 10) 145 | 146 | lock_timeout = self.tiger.get_queue_system_lock("a") 147 | assert lock_timeout == time.time() + 10 148 | 149 | # Confirm tasks don't get processed within the system lock timeout 150 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 9)): 151 | worker = Worker(self.tiger) 152 | worker.max_workers_per_queue = 2 153 | worker.run(once=True, force_once=True) 154 | self._ensure_queues(queued={"a": 2}) 155 | 156 | # 10 seconds in the future the lock should have expired 157 | with FreezeTime(datetime.datetime(2014, 1, 1, 0, 0, 10)): 158 | worker = Worker(self.tiger) 159 | worker.max_workers_per_queue = 2 160 | worker.run(once=True, force_once=True) 161 | self._ensure_queues(queued={"a": 1}) 162 | 163 | 164 | class TestSyncExecutorWorker: 165 | def test_success(self, tiger, ensure_queues): 166 | worker = Worker(tiger, executor_class=SyncExecutor) 167 | worker.run(once=True, force_once=True) 168 | 169 | task = Task(tiger, simple_task) 170 | task.delay() 171 | task = Task(tiger, simple_task) 172 | task.delay() 173 | ensure_queues(queued={"default": 2}) 174 | 175 | worker.run(once=True) 176 | ensure_queues() 177 | 178 | def test_handles_exception(self, tiger, ensure_queues): 179 | Task(tiger, exception_task).delay() 180 | worker = Worker(tiger, executor_class=SyncExecutor) 181 | worker.run(once=True, force_once=True) 182 | ensure_queues(error={"default": 1}) 183 | 184 | def test_handles_timeout(self, tiger, ensure_queues): 185 | Task(tiger, long_task_killed).delay() 186 | worker = Worker(tiger, executor_class=SyncExecutor) 187 | # Worker should exit to avoid any inconsistencies. 188 | with pytest.raises(SystemExit): 189 | worker.run(once=True, force_once=True) 190 | ensure_queues(error={"default": 1}) 191 | 192 | def test_heartbeat(self, tiger): 193 | # Test both task heartbeat and lock renewal. 194 | # We set unique=True so the task ID matches the lock key. 195 | task = Task(tiger, sleep_task, lock=True, unique=True) 196 | task.delay() 197 | 198 | # Start a worker and wait until it starts processing. 199 | worker = Process( 200 | target=external_worker, 201 | kwargs={ 202 | "patch_config": {"ACTIVE_TASK_UPDATE_TIMER": DELAY / 2}, 203 | "worker_kwargs": { 204 | # Test queue lock. 205 | "max_workers_per_queue": 1, 206 | "executor_class": SyncExecutor, 207 | }, 208 | }, 209 | ) 210 | worker.start() 211 | 212 | time.sleep(DELAY) 213 | 214 | queue_key = tiger._key(ACTIVE, "default") 215 | queue_lock_key = tiger._key(LOCK_REDIS_KEY, "default") 216 | task_lock_key = tiger._key("lockv2", task.id) 217 | 218 | conn = tiger.connection 219 | 220 | heartbeat_1 = conn.zscore(queue_key, task.id) 221 | queue_lock_1 = conn.zrange(queue_lock_key, 0, -1, withscores=True)[0][ 222 | 1 223 | ] 224 | task_lock_1 = conn.pttl(task_lock_key) 225 | 226 | time.sleep(DELAY / 2) 227 | 228 | heartbeat_2 = conn.zscore(queue_key, task.id) 229 | queue_lock_2 = conn.zrange(queue_lock_key, 0, -1, withscores=True)[0][ 230 | 1 231 | ] 232 | task_lock_2 = conn.pttl(task_lock_key) 233 | 234 | assert heartbeat_2 > heartbeat_1 > 0 235 | assert queue_lock_2 > queue_lock_1 > 0 236 | 237 | # Active task update timeout is 2 * DELAY and we renew every DELAY / 2. 238 | assert task_lock_1 > DELAY 239 | assert task_lock_2 > DELAY 240 | 241 | worker.kill() 242 | 243 | def test_stop_heartbeat_thread_on_unhandled_exception( 244 | self, tiger, ensure_queues 245 | ): 246 | task = Task(tiger, system_exit_task) 247 | task.delay() 248 | 249 | # Start a worker and wait until it starts processing. 250 | worker = Process( 251 | target=external_worker, 252 | kwargs={ 253 | "worker_kwargs": { 254 | "executor_class": SyncExecutor, 255 | }, 256 | }, 257 | ) 258 | worker.start() 259 | 260 | # Ensure process exits and does not hang here. 261 | worker.join() 262 | 263 | # Since SystemExit derives from BaseException and is therefore not 264 | # handled by the executor, the task is still active until it times out 265 | # and gets requeued by another worker. 266 | ensure_queues(active={"default": 1}) 267 | -------------------------------------------------------------------------------- /tests/utils.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import logging 3 | import time 4 | 5 | import redis 6 | import structlog 7 | 8 | from tasktiger import TaskTiger, Worker, fixed 9 | 10 | from .config import DELAY, REDIS_HOST, TEST_DB 11 | 12 | TEST_TIGER_CONFIG = { 13 | # We need this 0 here so we don't pick up scheduled tasks when 14 | # doing a single worker run. 15 | "ACTIVE_TASK_UPDATE_TIMEOUT": 2 * DELAY, 16 | "BATCH_QUEUES": {"batch": 3}, 17 | "DEFAULT_RETRY_METHOD": fixed(DELAY, 2), 18 | "EXCLUDE_QUEUES": ["periodic_ignore"], 19 | "LOCK_RETRY": DELAY * 2.0, 20 | "QUEUE_SCHEDULED_TASKS_TIME": DELAY, 21 | "REQUEUE_EXPIRED_TASKS_INTERVAL": DELAY, 22 | "SELECT_TIMEOUT": 0, 23 | "SINGLE_WORKER_QUEUES": ["swq"], 24 | } 25 | 26 | 27 | class Patch: 28 | """ 29 | Simple context manager to patch a function, e.g.: 30 | 31 | with Patch(module, 'func_name', mocked_func): 32 | module.func_name() # will use mocked_func 33 | module.func_name() # will use the original function 34 | 35 | """ 36 | 37 | def __init__(self, orig_obj, func_name, new_func): 38 | self.orig_obj = orig_obj 39 | self.func_name = func_name 40 | self.new_func = new_func 41 | 42 | def __enter__(self): 43 | self.orig_func = getattr(self.orig_obj, self.func_name) 44 | setattr(self.orig_obj, self.func_name, self.new_func) 45 | 46 | def __exit__(self, *args): 47 | setattr(self.orig_obj, self.func_name, self.orig_func) 48 | 49 | 50 | def setup_structlog(): 51 | structlog.configure( 52 | logger_factory=structlog.stdlib.LoggerFactory(), 53 | wrapper_class=structlog.stdlib.BoundLogger, 54 | ) 55 | logging.basicConfig(format="%(message)s") 56 | 57 | 58 | def get_redis(): 59 | return redis.Redis(host=REDIS_HOST, db=TEST_DB, decode_responses=True) 60 | 61 | 62 | def get_tiger(): 63 | """ 64 | Sets up logging and returns a new tasktiger instance. 65 | """ 66 | setup_structlog() 67 | conn = get_redis() 68 | tiger = TaskTiger(connection=conn, config=TEST_TIGER_CONFIG) 69 | tiger.log.setLevel(logging.CRITICAL) 70 | return tiger 71 | 72 | 73 | def external_worker(n=None, patch_config=None, worker_kwargs=None): 74 | """ 75 | Runs a worker. To be used with multiprocessing.Pool.map. 76 | """ 77 | tiger = get_tiger() 78 | 79 | if patch_config: 80 | tiger.config.update(patch_config) 81 | 82 | if worker_kwargs is None: 83 | worker_kwargs = {} 84 | 85 | Worker(tiger, **worker_kwargs).run(once=True, force_once=True) 86 | 87 | tiger.connection.close() 88 | 89 | 90 | def sleep_until_next_second(): 91 | now = datetime.datetime.utcnow() 92 | time.sleep(1 - now.microsecond / 10.0**6) 93 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py27,py34,py35,py36 3 | 4 | [testenv] 5 | commands=pytest {posargs} 6 | deps=pytest 7 | psutil 8 | freezefrog 9 | --------------------------------------------------------------------------------