├── pals ├── version.py ├── __init__.py ├── core.py └── tests │ └── test_core.py ├── pyp.ini ├── .ci └── pytest.ini ├── .gitignore ├── MANIFEST.in ├── stable-requirements.txt ├── docker-compose.yaml ├── .coveragerc ├── scripts ├── setup-codecov └── hang.py ├── setup.cfg ├── .circleci └── config.yml ├── tox.ini ├── setup.py ├── license.txt ├── changelog.rst └── readme.rst /pals/version.py: -------------------------------------------------------------------------------- 1 | VERSION = '0.4.0' 2 | -------------------------------------------------------------------------------- /pyp.ini: -------------------------------------------------------------------------------- 1 | [pyp] 2 | source_dir = pals 3 | -------------------------------------------------------------------------------- /pals/__init__.py: -------------------------------------------------------------------------------- 1 | from pals.core import * # noqa 2 | -------------------------------------------------------------------------------- /.ci/pytest.ini: -------------------------------------------------------------------------------- 1 | # This file's only purpose is to keep a developer's local pytest.ini from interfering with tox 2 | # tests. 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | MANIFEST 2 | *.egg-info 3 | *.pyc 4 | .coverage 5 | coverage.xml 6 | .tox/ 7 | .cache/ 8 | .ci/test-reports/*.xml 9 | dist 10 | build 11 | .pytest_cache 12 | .vscode 13 | .env 14 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include *.rst *.ini *.txt *.cfg .coveragerc Pipfile* *.yaml 2 | exclude docker-compose.override.yaml 3 | include .ci/* 4 | include .circleci/* 5 | recursive-include scripts *.py setup-codecov 6 | recursive-include pals *.py 7 | -------------------------------------------------------------------------------- /stable-requirements.txt: -------------------------------------------------------------------------------- 1 | attrs==22.1.0 2 | coverage==6.5.0 3 | exceptiongroup==1.0.4 4 | greenlet==2.0.1 5 | iniconfig==1.1.1 6 | packaging==21.3 7 | pluggy==1.0.0 8 | psycopg2-binary==2.9.5 9 | pyparsing==3.0.9 10 | pytest==7.2.0 11 | pytest-cov==4.0.0 12 | SQLAlchemy==1.4.44 13 | tomli==2.0.1 14 | -------------------------------------------------------------------------------- /docker-compose.yaml: -------------------------------------------------------------------------------- 1 | version: '2.1' 2 | services: 3 | pals-pg: 4 | image: postgres:15-alpine 5 | container_name: pals-pg 6 | ports: 7 | - '${PALS_LIB_POSTGRES_IP:-127.0.0.1}:${PALS_LIB_POSTGRES_PORT:-54321}:5432' 8 | environment: 9 | POSTGRES_HOST_AUTH_METHOD: trust 10 | -------------------------------------------------------------------------------- /.coveragerc: -------------------------------------------------------------------------------- 1 | # .coveragerc to control coverage.py 2 | [run] 3 | branch = True 4 | 5 | [report] 6 | # Regexes for lines to exclude from consideration 7 | exclude_lines = 8 | # Have to re-enable the standard pragma 9 | pragma: no cover 10 | 11 | # testing assertions usually used with exceptions 12 | assert False 13 | 14 | # bare raise statements 15 | raise\w*$ 16 | -------------------------------------------------------------------------------- /scripts/setup-codecov: -------------------------------------------------------------------------------- 1 | curl https://keybase.io/codecovsecurity/pgp_keys.asc | gpg --no-default-keyring --keyring trustedkeys.gpg --import 2 | curl -Os https://uploader.codecov.io/latest/linux/codecov 3 | curl -Os https://uploader.codecov.io/latest/linux/codecov.SHA256SUM 4 | curl -Os https://uploader.codecov.io/latest/linux/codecov.SHA256SUM.sig 5 | gpgv codecov.SHA256SUM.sig codecov.SHA256SUM 6 | shasum -a 256 -c codecov.SHA256SUM 7 | chmod +x codecov 8 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal = 1 3 | 4 | # E121 - A line is less indented than it should be for hanging indents. 5 | # E128 - Continuation line under-indented for visual indent. To permit: 6 | # def some_log_function_call_with_params_on_next_line(race_day=race_day, system_id='A', 7 | # system_status='stopped', system_run_id='1') 8 | # E731 - Lambdas should not be assigned to a variable. Instead, they should be defined as functions. 9 | # W503 - Line breaks should occur after the binary operator to keep all variable names aligned. 10 | # Details & examples at: Errors: https://lintlyci.github.io/Flake8Rules/ 11 | [flake8] 12 | max-line-length = 100 13 | exclude=.git,.hg,.tox,dist,doc,*egg,build 14 | ignore=E121,E128,E731,W503 15 | -------------------------------------------------------------------------------- /scripts/hang.py: -------------------------------------------------------------------------------- 1 | """ 2 | If you run this file and monitor connections on your PostgreSQL server, you should see a backend 3 | process dedicated to this lock: 4 | 5 | select 6 | stat.datname 7 | , stat.pid 8 | , stat.usename 9 | , stat.application_name 10 | , stat.state 11 | , locks.mode as lockmode 12 | , stat.query 13 | from pg_stat_activity stat 14 | left join pg_locks locks 15 | on locks.pid = stat.pid 16 | where locktype = 'advisory' 17 | 18 | Then kill the Python process running this file in any way you want and make sure the PostgreSQL 19 | pid with the lock goes away: 20 | 21 | 1) press any key and exit 22 | 2) CTRL-C 23 | 3) kill -9 24 | """ 25 | import os 26 | 27 | import pals 28 | 29 | locker = pals.Locker('pals-hang', 'postgresql://postgres:password@localhost:54321/postgres') 30 | 31 | lock = locker.lock('hang') 32 | lock.acquire() 33 | print('My pid is: ', os.getpid()) 34 | input('Lock acquired, press any key to exit: ') 35 | -------------------------------------------------------------------------------- /.circleci/config.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | jobs: 3 | build: 4 | docker: 5 | - image: level12/python-test-multi 6 | - image: postgres:15 7 | environment: 8 | POSTGRES_USER: postgres 9 | POSTGRES_PASSWORD: password 10 | steps: 11 | - checkout 12 | 13 | - run: 14 | name: folder listing for debugging 15 | command: ls -al 16 | 17 | - run: 18 | name: install latest version of tox 19 | command: python3.10 -m pip install --upgrade --force-reinstall --quiet tox 20 | 21 | - run: 22 | name: version checks 23 | command: | 24 | python3.10 --version 25 | tox --version 26 | 27 | - run: 28 | name: run tox 29 | command: tox 30 | 31 | - store_test_results: 32 | path: .ci/test-reports/ 33 | 34 | - run: 35 | name: push code coverage 36 | command: source scripts/setup-codecov && ./codecov -t bef3273d-1997-4764-9c87-81e7d6a57778 37 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = 3 | py{39,310,311}-{base,stable} 4 | py311-{lowest} 5 | project 6 | 7 | [testenv] 8 | # Ignore all "not installed in testenv" warnings. 9 | whitelist_externals = * 10 | passenv = PALS_DB_URL 11 | 12 | skip_install = true 13 | 14 | recreate = true 15 | deps = 16 | -e .[tests] 17 | stable: -r stable-requirements.txt 18 | lowest: sqlalchemy<2 19 | lowest: psycopg2-binary 20 | commands = 21 | lowest: pip uninstall -y psycopg psycopg-binary 22 | stable: pip uninstall -y psycopg psycopg-binary 23 | pip --version 24 | # Output installed versions to compare with previous test runs in case a dependency's change 25 | # breaks things for our build. 26 | pip freeze 27 | py.test \ 28 | # feed a blank file so that a user's default pytest.ini doesn't get used 29 | -c .ci/pytest.ini \ 30 | -ra \ 31 | --tb native \ 32 | --strict-markers \ 33 | --cov pals \ 34 | --cov-config .coveragerc \ 35 | --cov-report xml \ 36 | --no-cov-on-fail \ 37 | --junit-xml={toxinidir}/.ci/test-reports/{envname}.pytests.xml \ 38 | pals/tests 39 | 40 | 41 | [testenv:project] 42 | basepython = python3.11 43 | skip_install = true 44 | usedevelop = false 45 | deps = 46 | check-manifest 47 | flake8 48 | twine 49 | commands = 50 | pip install -e .[tests] 51 | check-manifest 52 | python setup.py sdist 53 | twine check dist/* 54 | flake8 pals 55 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | from setuptools import setup, find_packages 3 | 4 | cdir = osp.abspath(osp.dirname(__file__)) 5 | README = open(osp.join(cdir, 'readme.rst')).read() 6 | CHANGELOG = open(osp.join(cdir, 'changelog.rst')).read() 7 | 8 | version_fpath = osp.join(cdir, 'pals', 'version.py') 9 | version_globals = {} 10 | with open(version_fpath) as fo: 11 | exec(fo.read(), version_globals) 12 | 13 | setup( 14 | name="PALs", 15 | version=version_globals['VERSION'], 16 | description="Easy distributed locking using PostgreSQL Advisory Locks.", 17 | long_description='\n\n'.join((README, CHANGELOG)), 18 | long_description_content_type='text/x-rst', 19 | author="Randy Syring", 20 | author_email="randy.syring@level12.io", 21 | url='https://github.com/level12/pals', 22 | classifiers=[ 23 | 'Development Status :: 4 - Beta', 24 | 'Intended Audience :: Developers', 25 | 'License :: OSI Approved :: BSD License', 26 | 'Operating System :: OS Independent', 27 | 'Programming Language :: Python :: 3', 28 | 'Programming Language :: Python :: 3.7', 29 | 'Programming Language :: Python :: 3.8', 30 | 'Programming Language :: Python :: 3.9', 31 | 'Programming Language :: Python :: 3.10', 32 | ], 33 | license='BSD-3-Clause', 34 | packages=find_packages(), 35 | zip_safe=False, 36 | include_package_data=True, 37 | install_requires=[ 38 | 'sqlalchemy', 39 | ], 40 | extras_require={ 41 | 'tests': [ 42 | 'pytest', 43 | 'pytest-cov', 44 | 'psycopg[binary]', 45 | ], 46 | } 47 | ) 48 | -------------------------------------------------------------------------------- /license.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2018 by Randy Syring and contributors. 2 | 3 | Some rights reserved. 4 | 5 | Redistribution and use in source and binary forms of the software as well 6 | as documentation, with or without modification, are permitted provided 7 | that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright 10 | notice, this list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above 13 | copyright notice, this list of conditions and the following 14 | disclaimer in the documentation and/or other materials provided 15 | with the distribution. 16 | 17 | * The names of the contributors may not be used to endorse or 18 | promote products derived from this software without specific 19 | prior written permission. 20 | 21 | THIS SOFTWARE AND DOCUMENTATION IS PROVIDED BY THE COPYRIGHT HOLDERS AND 22 | CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT 23 | NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 24 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER 25 | OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, 26 | EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, 27 | PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 28 | PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 29 | LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING 30 | NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 31 | SOFTWARE AND DOCUMENTATION, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH 32 | DAMAGE. 33 | -------------------------------------------------------------------------------- /changelog.rst: -------------------------------------------------------------------------------- 1 | Changelog 2 | ========= 3 | 4 | 0.4.0 released 2024-04-16 5 | ------------------------- 6 | 7 | - support psycopg3, backwards compatible with psycopg2 (thanks to @petr.prikryl) (0139bea_) 8 | - remove python 3.7 and 3.8 from CI, resolve flake8 issue (24bff1f_) 9 | 10 | .. _0139bea: https://github.com/level12/pals/commit/0139bea 11 | .. _24bff1f: https://github.com/level12/pals/commit/24bff1f 12 | 13 | 14 | 0.3.5 released 2023-06-29 15 | ------------------------- 16 | 17 | - fix cursor usage after connection close (thanks to @moser) (ded88e7_) 18 | 19 | .. _ded88e7: https://github.com/level12/pals/commit/ded88e7 20 | 21 | 22 | 0.3.4 released 2023-03-06 23 | ------------------------- 24 | 25 | - support SQLAlchemy 2.0 (6879081_) 26 | 27 | .. _6879081: https://github.com/level12/pals/commit/6879081 28 | 29 | 30 | 0.3.3 released 2023-01-06 31 | ------------------------- 32 | 33 | - add additional info to AcquireFailure exception (6d81db9_) 34 | 35 | .. _6d81db9: https://github.com/level12/pals/commit/6d81db9 36 | 37 | 38 | 0.3.2 released 2021-02-01 39 | ------------------------- 40 | 41 | - Support shared advisory locks (thanks to @absalon-james) (ba2fe21_) 42 | 43 | .. _ba2fe21: https://github.com/level12/pals/commit/ba2fe21 44 | 45 | 46 | 0.3.1 released 2020-09-03 47 | ------------------------- 48 | 49 | - readme: update postgresql link (260bf75_) 50 | - Handle case where a DB connection is returned to the pool which is already closed (5d730c9_) 51 | - Fix a couple of typos in comments (da2b8af_) 52 | - readme improvements (4efba90_) 53 | - CI: fix coverage upload (52daa27_) 54 | - Fix CI: bump CI python to v3.7 and postgres to v11 (23b3028_) 55 | 56 | .. _260bf75: https://github.com/level12/pals/commit/260bf75 57 | .. _5d730c9: https://github.com/level12/pals/commit/5d730c9 58 | .. _da2b8af: https://github.com/level12/pals/commit/da2b8af 59 | .. _4efba90: https://github.com/level12/pals/commit/4efba90 60 | .. _52daa27: https://github.com/level12/pals/commit/52daa27 61 | .. _23b3028: https://github.com/level12/pals/commit/23b3028 62 | 63 | 64 | 0.3.0 released 2019-11-13 65 | ------------------------- 66 | 67 | Enhancements 68 | ~~~~~~~~~~~~ 69 | 70 | - Add acquire timeout and blocking defaults at Locker level (681c3ba_) 71 | - Adjust default lock timeout from 1s to 30s (5a0963b_) 72 | 73 | Project Cleanup 74 | ~~~~~~~~~~~~~~~ 75 | 76 | - adjust flake8 ignore and other tox project warning (ee123fc_) 77 | - fix comment in test (0d8eb98_) 78 | - Additional readme updates (0786766_) 79 | - update locked dependencies (f5743a6_) 80 | - Remove Python 3.5 from CI (b63c71a_) 81 | - Cleaned up the readme code example a bit and added more references (dabb497_) 82 | - Update setup.py to use SPDX license identifier (b811a99_) 83 | - remove Pipefiles (0637f39_) 84 | - move to using piptools for dependency management (af2e91f_) 85 | 86 | .. _ee123fc: https://github.com/level12/pals/commit/ee123fc 87 | .. _681c3ba: https://github.com/level12/pals/commit/681c3ba 88 | .. _5a0963b: https://github.com/level12/pals/commit/5a0963b 89 | .. _0d8eb98: https://github.com/level12/pals/commit/0d8eb98 90 | .. _0786766: https://github.com/level12/pals/commit/0786766 91 | .. _f5743a6: https://github.com/level12/pals/commit/f5743a6 92 | .. _b63c71a: https://github.com/level12/pals/commit/b63c71a 93 | .. _dabb497: https://github.com/level12/pals/commit/dabb497 94 | .. _b811a99: https://github.com/level12/pals/commit/b811a99 95 | .. _0637f39: https://github.com/level12/pals/commit/0637f39 96 | .. _af2e91f: https://github.com/level12/pals/commit/af2e91f 97 | 98 | 99 | 0.2.0 released 2019-03-07 100 | ------------------------- 101 | 102 | - Fix misspelling of "acquire" (737763f_) 103 | 104 | .. _737763f: https://github.com/level12/pals/commit/737763f 105 | 106 | 107 | 0.1.0 released 2019-02-22 108 | ------------------------- 109 | 110 | - Use `lock_timeout` setting to expire blocking calls (d0216ce_) 111 | - fix tox (1b0ffe2_) 112 | - rename to PALs (95d5a3c_) 113 | - improve readme (e8dd6f2_) 114 | - move tests file to better location (a153af5_) 115 | - add flake8 dep (3909c95_) 116 | - fix tests so they work locally too (7102294_) 117 | - get circleci working (28f16d2_) 118 | - suppress exceptions in Lock __del__ (e29c1ce_) 119 | - Add hang.py script (3372ef0_) 120 | - fix packaging stuff, update readme (cebd976_) 121 | - initial commit (871b877_) 122 | 123 | .. _d0216ce: https://github.com/level12/pals/commit/d0216ce 124 | .. _1b0ffe2: https://github.com/level12/pals/commit/1b0ffe2 125 | .. _95d5a3c: https://github.com/level12/pals/commit/95d5a3c 126 | .. _e8dd6f2: https://github.com/level12/pals/commit/e8dd6f2 127 | .. _a153af5: https://github.com/level12/pals/commit/a153af5 128 | .. _3909c95: https://github.com/level12/pals/commit/3909c95 129 | .. _7102294: https://github.com/level12/pals/commit/7102294 130 | .. _28f16d2: https://github.com/level12/pals/commit/28f16d2 131 | .. _e29c1ce: https://github.com/level12/pals/commit/e29c1ce 132 | .. _3372ef0: https://github.com/level12/pals/commit/3372ef0 133 | .. _cebd976: https://github.com/level12/pals/commit/cebd976 134 | .. _871b877: https://github.com/level12/pals/commit/871b877 135 | 136 | -------------------------------------------------------------------------------- /readme.rst: -------------------------------------------------------------------------------- 1 | .. default-role:: code 2 | 3 | PostgreSQL Advisory Locks (PALs) 4 | ################################ 5 | 6 | .. image:: https://circleci.com/gh/level12/pals.svg?style=shield 7 | :target: https://circleci.com/gh/level12/pals 8 | .. image:: https://codecov.io/gh/level12/pals/branch/master/graph/badge.svg 9 | :target: https://codecov.io/gh/level12/pals 10 | 11 | 12 | Introduction 13 | ============ 14 | 15 | PALs makes it easy to use `PostgreSQL Advisory Locks`_ to do distributed application level 16 | locking. 17 | 18 | Do not confuse this type of locking with table or row locking in PostgreSQL. It's not the same 19 | thing. 20 | 21 | Distributed application level locking can be implemented by using Redis, Memcache, ZeroMQ and 22 | others. But for those who are already using PostgreSQL, setup & management of another service is 23 | unnecessary. 24 | 25 | .. _PostgreSQL Advisory Locks: https://www.postgresql.org/docs/current/static/explicit-locking.html#ADVISORY-LOCKS 26 | 27 | 28 | Usage 29 | ======== 30 | 31 | Install with:: 32 | 33 | pip install PALs 34 | 35 | Then usage is as follows: 36 | 37 | .. code:: python 38 | 39 | import datetime as dt 40 | import pals 41 | 42 | # Think of the Locker instance as a Lock factory. 43 | locker = pals.Locker('my-app-name', 'postgresql://user:pass@server/dbname') 44 | 45 | lock1 = locker.lock('my-lock') 46 | lock2 = locker.lock('my-lock') 47 | 48 | # The first acquire works 49 | assert lock1.acquire() is True 50 | 51 | # Non blocking version should fail immediately 52 | assert lock2.acquire(blocking=False) is False 53 | 54 | # Blocking version should fail after a short time 55 | start = dt.datetime.now() 56 | acquired = lock2.acquire(acquire_timeout=300) 57 | waited_ms = duration(start) 58 | 59 | assert acquired is False 60 | assert waited_ms >= 300 and waited_ms < 350 61 | 62 | # Release the lock 63 | lock1.release() 64 | 65 | # Non-blocking usage pattern 66 | if not lock1.acquire(blocking=False): 67 | # Aquire returned False, indicating we did not get the lock. 68 | return 69 | try: 70 | # do your work here 71 | finally: 72 | lock1.release() 73 | 74 | # If you want to block, you can use a context manager: 75 | try: 76 | with lock1: 77 | # Do your work here 78 | pass 79 | except pals.AcquireFailure: 80 | # This indicates the aquire_timeout was reached before the lock could be aquired. 81 | pass 82 | 83 | Docs 84 | ======== 85 | 86 | Just this readme, the code, and tests. It a small project, should be easy to understand. 87 | 88 | Feel free to open an issue with questions. 89 | 90 | Running Tests Locally 91 | ===================== 92 | 93 | Setup Database Connection 94 | ------------------------- 95 | 96 | We have provided a docker-compose file to ease running the tests:: 97 | 98 | $ docker-compose up -d 99 | $ export PALS_DB_URL=postgresql://postgres:password@localhost:54321/postgres 100 | 101 | 102 | Run the Tests 103 | ------------- 104 | 105 | With tox:: 106 | 107 | $ tox 108 | 109 | Or, manually (assuming an activated virtualenv):: 110 | 111 | $ pip install -e .[tests] 112 | $ pytest pals/tests/ 113 | 114 | 115 | Lock Releasing & Expiration 116 | --------------------------- 117 | 118 | Unlike locking systems built on cache services like Memcache and Redis, whose keys can be expired 119 | by the service, there is no faculty for expiring an advisory lock in PostgreSQL. If a client 120 | holds a lock and then sleeps/hangs for mins/hours/days, no other client will be able to get that 121 | lock until the client releases it. This actually seems like a good thing to us, if a lock is 122 | acquired, it should be kept until released. 123 | 124 | But what about accidental failures to release the lock? 125 | 126 | 1. If a developer uses `lock.acquire()` but doesn't later call `lock.release()`? 127 | 2. If code inside a lock accidentally throws an exception (and .release() is not called)? 128 | 3. If the process running the application crashes or the process' server dies? 129 | 130 | PALs helps #1 and #2 above in a few different ways: 131 | 132 | * Locks work as context managers. Use them as much as possible to guarantee a lock is released. 133 | * Locks release their lock when garbage collected. 134 | * PALs uses a dedicated SQLAlchemy connection pool. When a connection is returned to the pool, 135 | either because a connection `.close()` is called or due to garbage collection of the connection, 136 | PALs issues a `pg_advisory_unlock_all()`. It should therefore be impossible for an idle 137 | connection in the pool to ever still be holding a lock. 138 | 139 | Regarding #3 above, `pg_advisory_unlock_all()` is implicitly invoked by PostgreSQL whenever a 140 | connection (a.k.a session) ends, even if the client disconnects ungracefully. So if a process 141 | crashes or otherwise disappears, PostgreSQL should notice and remove all locks held by that 142 | connection/session. 143 | 144 | The possibility could exist that PostgreSQL does not detect a connection has closed and keeps 145 | a lock open indefinitely. However, in manual testing using `scripts/hang.py` no way was found 146 | to end the Python process without PostgreSQL detecting it. 147 | 148 | 149 | See Also 150 | ========== 151 | 152 | * https://vladmihalcea.com/how-do-postgresql-advisory-locks-work/ 153 | * https://github.com/binded/advisory-lock 154 | * https://github.com/vaidik/sherlock 155 | * https://github.com/Xof/django-pglocks 156 | 157 | -------------------------------------------------------------------------------- /pals/core.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import logging 3 | import struct 4 | 5 | import sqlalchemy as sa 6 | 7 | 8 | log = logging.getLogger(__name__) 9 | 10 | __all__ = [ 11 | 'Locker', 12 | 'AcquireFailure', 13 | ] 14 | 15 | 16 | class AcquireFailure(Exception): 17 | def __init__(self, name, message): 18 | self.name = name 19 | self.message = message 20 | 21 | def __str__(self): 22 | return f'Lock acquire failed for "{self.name}". {self.message}.' 23 | 24 | 25 | class Locker: 26 | """ 27 | A Locker instance is intended to be an app-level lock factory. 28 | 29 | It holds the name of the application (so lock names are namespaced and less likely to 30 | collide) and the SQLAlchemy engine instance (and therefore the connection pool). 31 | """ 32 | def __init__(self, app_name, db_url=None, blocking_default=True, acquire_timeout_default=30000, 33 | create_engine_callable=None): 34 | self.app_name = app_name 35 | self.blocking_default = blocking_default 36 | self.acquire_timeout_default = acquire_timeout_default 37 | 38 | if create_engine_callable: 39 | self.engine = create_engine_callable() 40 | else: 41 | self.engine = sa.create_engine(db_url) 42 | 43 | @sa.event.listens_for(self.engine, 'checkin') 44 | def on_conn_checkin(dbapi_connection, connection_record): 45 | """ 46 | This function will be called when a connection is checked back into the connection 47 | pool. That should happen when .close() is called on it or when the connection 48 | proxy goes out of scope and is garbage collected. 49 | """ 50 | if dbapi_connection is None: 51 | # This may occur in rare circumstances where the connection is already closed or an 52 | # error occurred while connecting to the database. In these cases any held locks 53 | # should already be released when the connection terminated. 54 | return 55 | 56 | with dbapi_connection.cursor() as cur: 57 | # If the connection is "closed" we want all locks to be cleaned up since this 58 | # connection is going to be recycled. This step is to take extra care that we don't 59 | # accidentally leave a lock acquired. 60 | cur.execute('select pg_advisory_unlock_all()') 61 | 62 | def _lock_name(self, name): 63 | if self.app_name is None: 64 | return name 65 | 66 | return '{}.{}'.format(self.app_name, name) 67 | 68 | def _lock_num(self, name): 69 | """ 70 | PostgreSQL requires lock ids to be integers. It accepts bigints which gives us 71 | 64 bits to work with. Hash the lock name to an integer. 72 | """ 73 | name = self._lock_name(name) 74 | 75 | name_hash = hashlib.sha1(name.encode('utf-8')) 76 | 77 | # Convert the hash to an integer value in the range of a PostgreSQL bigint 78 | num, = struct.unpack('q', name_hash.digest()[:8]) 79 | 80 | return num 81 | 82 | def lock(self, name, **kwargs): 83 | lock_num = self._lock_num(name) 84 | name = self._lock_name(name) 85 | kwargs.setdefault('blocking', self.blocking_default) 86 | kwargs.setdefault('acquire_timeout', self.acquire_timeout_default) 87 | return Lock(self.engine, lock_num, name, **kwargs) 88 | 89 | 90 | class Lock: 91 | def __init__(self, engine, lock_num, name, blocking=None, acquire_timeout=None, shared=False): 92 | self.engine = engine 93 | self.conn = None 94 | self.lock_num = lock_num 95 | self.name = name 96 | self.blocking = blocking 97 | self.acquire_timeout = acquire_timeout 98 | self.shared_suffix = '_shared' if shared else '' 99 | 100 | def _acquire(self, blocking=None, acquire_timeout=None) -> bool: 101 | blocking = blocking if blocking is not None else self.blocking 102 | acquire_timeout = acquire_timeout or self.acquire_timeout 103 | 104 | if self.conn is None: 105 | self.conn = self.engine.connect() 106 | 107 | if blocking: 108 | timeout_sql = sa.text("select set_config('lock_timeout', :timeout :: text, false)") 109 | with self.conn.begin(): 110 | self.conn.execute(timeout_sql, {'timeout': acquire_timeout}) 111 | 112 | lock_sql = sa.text(f'select pg_advisory_lock{self.shared_suffix}(:lock_num)') 113 | else: 114 | lock_sql = sa.text(f'select pg_try_advisory_lock{self.shared_suffix}(:lock_num)') 115 | 116 | try: 117 | with self.conn.begin(): 118 | result = self.conn.execute(lock_sql, {'lock_num': self.lock_num}) 119 | retval = result.scalar() 120 | log.debug('Lock result was: %r', retval) 121 | # At least on PG 10.6, pg_advisory_lock() returns an empty string 122 | # when it acquires the lock. pg_try_advisory_lock() returns True. 123 | # If pg_try_advisory_lock() fails, it returns False. 124 | if retval in (True, ''): 125 | return True 126 | else: 127 | raise AcquireFailure(self.name, 'result was: {retval}') 128 | except sa.exc.OperationalError as e: 129 | if 'lock timeout' not in str(e): 130 | raise 131 | log.debug('Lock acquire failed due to timeout') 132 | raise AcquireFailure(self.name, 'Failed due to timeout') 133 | 134 | def acquire(self, blocking=None, acquire_timeout=None) -> bool: 135 | try: 136 | return self._acquire(blocking=blocking, acquire_timeout=acquire_timeout) 137 | except AcquireFailure: 138 | return False 139 | 140 | def release(self): 141 | if self.conn is None: 142 | return False 143 | 144 | sql = sa.text(f'select pg_advisory_unlock{self.shared_suffix}(:lock_num)') 145 | with self.conn.begin(): 146 | result = self.conn.execute(sql, {'lock_num': self.lock_num}) 147 | try: 148 | return result.scalar() 149 | finally: 150 | self.conn.close() 151 | self.conn = None 152 | 153 | def __enter__(self): 154 | self._acquire() 155 | return self 156 | 157 | def __exit__(self, exc_type, exc_value, traceback): 158 | self.release() 159 | 160 | def __del__(self): 161 | # Do everything we can to release resources and the connection to avoid accidentally holding 162 | # a lock indefinitely if .release() is forgotten. 163 | try: 164 | self.release() 165 | except Exception: 166 | # Sometimes this will fail if the connection has gone away before the gc runs. Since 167 | # Python is just going to print the exception and we can't do anything about it, 168 | # suppress the exception to keep erroneous noise out of stderr. 169 | pass 170 | try: 171 | self.conn.close() 172 | except Exception: 173 | # ditto 174 | pass 175 | del self.conn 176 | -------------------------------------------------------------------------------- /pals/tests/test_core.py: -------------------------------------------------------------------------------- 1 | import datetime as dt 2 | import gc 3 | import os 4 | import random 5 | import string 6 | import threading 7 | import time 8 | 9 | import pytest 10 | 11 | import pals 12 | 13 | try: 14 | import psycopg # noqa: F401 15 | 16 | db_driver = 'postgresql+psycopg' 17 | except ImportError: 18 | db_driver = 'postgresql' 19 | 20 | 21 | # Default URL will work for CI tests 22 | db_url = os.environ.get( 23 | 'PALS_DB_URL', 24 | f'{db_driver}://postgres:password@localhost/postgres' 25 | ) 26 | 27 | 28 | def random_str(length): 29 | return ''.join(random.choice(string.printable) for _ in range(25)) 30 | 31 | 32 | class TestLocker: 33 | 34 | def test_lock_num_generation(self): 35 | """ 36 | Create 5000 random strings of varying lengths, convert them to their lock number, and 37 | make sure we still have 5000 unique numbers. Given that any application using locking 38 | is probably going to have, at most, locks numbering in the tens or very low hundreds, 39 | this seems like an ok way to test that the method we are using to convert our strings 40 | into numbers is unlikely to have accidental collisions. 41 | """ 42 | locker = pals.Locker('TestLocker', db_url) 43 | 44 | names = [random_str(max(6, x % 25)) for x in range(5000)] 45 | assert len(set(names)) == 5000 46 | nums = [locker._lock_num(name) for name in names] 47 | assert len(set(nums)) == 5000 48 | 49 | def test_locker_defaults(self): 50 | lock = pals.Locker('foo', db_url).lock('a') 51 | assert lock.blocking is True 52 | assert lock.acquire_timeout == 30000 53 | 54 | lock = pals.Locker('bar', db_url, blocking_default=False, acquire_timeout_default=1000) \ 55 | .lock('a') 56 | assert lock.blocking is False 57 | assert lock.acquire_timeout == 1000 58 | 59 | 60 | def duration(started_at): 61 | duration = dt.datetime.now() - started_at 62 | secs = duration.total_seconds() 63 | # in milliseconds 64 | return secs * 1000 65 | 66 | 67 | class TestLock: 68 | locker = pals.Locker('TestLock', db_url, acquire_timeout_default=1000) 69 | 70 | def test_same_lock_fails_acquire(self): 71 | lock1 = self.locker.lock('test_it') 72 | lock2 = self.locker.lock('test_it') 73 | 74 | try: 75 | assert lock1.acquire() is True 76 | 77 | # Non blocking version should fail immediately. Use `is` test to make sure we get the 78 | # correct return value. 79 | assert lock2.acquire(blocking=False) is False 80 | 81 | # Blocking version should fail after a short time 82 | start = dt.datetime.now() 83 | acquired = lock2.acquire(acquire_timeout=300) 84 | waited_ms = duration(start) 85 | 86 | assert acquired is False 87 | assert waited_ms >= 300 and waited_ms < 350 88 | finally: 89 | lock1.release() 90 | 91 | def test_different_lock_name_both_acquire(self): 92 | lock1 = self.locker.lock('test_it') 93 | lock2 = self.locker.lock('test_it2') 94 | 95 | try: 96 | assert lock1.acquire() is True 97 | assert lock2.acquire() is True 98 | finally: 99 | lock1.release() 100 | lock2.release() 101 | 102 | def test_lock_after_release_acquires(self): 103 | lock1 = self.locker.lock('test_it') 104 | lock2 = self.locker.lock('test_it') 105 | 106 | try: 107 | assert lock1.acquire() is True 108 | assert lock1.release() is True 109 | assert lock2.acquire() is True 110 | finally: 111 | lock1.release() 112 | lock2.release() 113 | 114 | def test_class_params_used(self): 115 | """ 116 | If blocking & timeout params are set on the class, make sure they are passed through and 117 | used correctly. 118 | """ 119 | lock1 = self.locker.lock('test_it') 120 | lock2 = self.locker.lock('test_it', blocking=False) 121 | lock3 = self.locker.lock('test_it', acquire_timeout=300) 122 | 123 | try: 124 | assert lock1.acquire() is True 125 | 126 | # Make sure the blocking param applies 127 | acquired = lock2.acquire() 128 | assert acquired is False 129 | 130 | # Make sure the retry params apply 131 | start = dt.datetime.now() 132 | acquired = lock3.acquire() 133 | waited_ms = duration(start) 134 | assert acquired is False 135 | assert waited_ms >= 300 and waited_ms < 350 136 | finally: 137 | lock1.release() 138 | lock2.release() 139 | lock3.release() 140 | 141 | def test_shared_lock(self): 142 | """ 143 | Test shared lock. 144 | """ 145 | shared_lock = self.locker.lock('test_it', shared=True) 146 | another_shared_lock = self.locker.lock('test_it', shared=True) 147 | exclusive_lock = self.locker.lock('test_it', shared=False, blocking=False) 148 | try: 149 | assert shared_lock.acquire() is True 150 | assert another_shared_lock.acquire() is True 151 | 152 | # Assert that an exclusive lock cannot be obtained 153 | assert exclusive_lock.acquire() is False 154 | 155 | # Release one of the shared locks. 156 | shared_lock.release() 157 | 158 | # Assert that the exclusive still cannot be obtained. 159 | assert exclusive_lock.acquire() is False 160 | 161 | # Release the other shared lock. 162 | another_shared_lock.release() 163 | 164 | # Assert that the exclusive lock can now be obtained. 165 | assert exclusive_lock.acquire() is True 166 | finally: 167 | shared_lock.release() 168 | another_shared_lock.release() 169 | exclusive_lock.release() 170 | 171 | def test_context_manager(self): 172 | lock2 = self.locker.lock('test_it', blocking=False) 173 | try: 174 | with self.locker.lock('test_it'): 175 | assert lock2.acquire() is False 176 | 177 | # Outside the lock should have been released and we can get it now 178 | assert lock2.acquire() 179 | finally: 180 | lock2.release() 181 | 182 | def test_context_manager_failure_to_acquire(self): 183 | """ 184 | By default, we want a context manager's failure to acquire to be a hard error so that 185 | a developer doesn't have to remember to explicilty check the return value of acquire 186 | when using a with statement. 187 | """ 188 | lock2 = self.locker.lock('test_it', blocking=False) 189 | assert lock2.acquire() is True 190 | 191 | with pytest.raises( 192 | pals.AcquireFailure, 193 | match='Lock acquire failed for "TestLock.test_it". Failed due to timeout.', 194 | ): 195 | with self.locker.lock("test_it"): 196 | pass # we should never hit this line 197 | 198 | def test_gc_lock_release(self): 199 | lock1 = self.locker.lock('test_it') 200 | lock1.acquire() 201 | del lock1 202 | gc.collect() 203 | 204 | assert self.locker.lock('test_it', blocking=False).acquire() 205 | 206 | def test_contention_on_connection_pool(self): 207 | """ 208 | Given N = size of connection pool, run N + 1 threads that all try to 209 | acquire locks (different ones) at the same time. 210 | 211 | The first N threads should all acquire a lock. The last thread should 212 | wait until one of the first N threads releases their lock, then acquire 213 | it normally. 214 | 215 | This test demonstrates issue #40. 216 | """ 217 | set_size = self.locker.engine.pool.size() + 1 218 | results = [None] * set_size 219 | 220 | def target(n): 221 | try: 222 | with self.locker.lock(f"example-{n}"): 223 | time.sleep(0.05) 224 | results[n] = True 225 | except Exception as e: 226 | results[n] = e 227 | 228 | threads = [threading.Thread(target=target, args=(n,)) for n in range(set_size)] 229 | 230 | for thread in threads: 231 | thread.start() 232 | for thread in threads: 233 | thread.join() 234 | 235 | assert [r for r in results if isinstance(r, Exception)] == [] 236 | --------------------------------------------------------------------------------