Skip to content

Commit 0b09bc5

Browse files
committedJul 21, 2023
wip
1 parent 7ff40b2 commit 0b09bc5

File tree

11 files changed

+604
-2
lines changed

11 files changed

+604
-2
lines changed
 

‎.github/workflows/ci.yml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,3 +50,6 @@ jobs:
5050
- name: Run lint
5151
run: |
5252
poetry run make LINT_FIX=0 lint
53+
- name: Run test
54+
run: |
55+
poetry run make test

‎.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,3 +158,4 @@ cython_debug/
158158
# and can be added to the global gitignore or merged into this file. For a more nuclear
159159
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
160160
#.idea/
161+
/.vscode

‎Makefile

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
LINT_FIX = 1
2+
COVERAGE = 0
23

34
.PHONY: lint_ruff
45
lint_ruff:
@@ -22,3 +23,12 @@ lint_mypy:
2223

2324
.PHONY: lint
2425
lint: lint_ruff lint_black lint_mypy
26+
27+
.PHONY: test
28+
test:
29+
ifeq ($(COVERAGE),0)
30+
pytest tests
31+
else
32+
pytest --no-cov-on-fail --cov=distributed_lock --cov-report=term --cov-report=html --cov-report=xml tests
33+
endif
34+

‎distributed_lock/__init__.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from __future__ import annotations
2+
3+
from distributed_lock.const import DEFAULT_CLUSTER, DEFAULT_LIFETIME, DEFAULT_WAIT
4+
from distributed_lock.exception import (
5+
BadConfigurationError,
6+
DistributedLockError,
7+
DistributedLockException,
8+
NotAcquiredError,
9+
NotAcquiredException,
10+
NotReleasedError,
11+
NotReleasedException,
12+
)
13+
from distributed_lock.sync import AcquiredRessource, DistributedLockClient
14+
15+
__all__ = [
16+
"DEFAULT_CLUSTER",
17+
"DEFAULT_LIFETIME",
18+
"DEFAULT_WAIT",
19+
"AcquiredRessource",
20+
"DistributedLockClient",
21+
"DistributedlockException",
22+
"NotAcquiredError",
23+
"NotReleasedException",
24+
"NotReleasedError",
25+
"NotAcquiredException",
26+
"BadConfigurationError",
27+
"DistributedLockError",
28+
"DistributedLockException",
29+
]

‎distributed_lock/const.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
from __future__ import annotations
2+
3+
DEFAULT_CLUSTER = "europe-free"
4+
DEFAULT_LIFETIME = 3600
5+
DEFAULT_WAIT = 10

‎distributed_lock/exception.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from __future__ import annotations
2+
3+
4+
class DistributedLockException(Exception):
5+
pass
6+
7+
8+
class DistributedLockError(DistributedLockException):
9+
pass
10+
11+
12+
class BadConfigurationError(DistributedLockError):
13+
pass
14+
15+
16+
class NotAcquiredException(DistributedLockException):
17+
pass
18+
19+
20+
class NotReleasedException(DistributedLockException):
21+
pass
22+
23+
24+
class NotReleasedError(DistributedLockError):
25+
pass
26+
27+
28+
class NotAcquiredError(DistributedLockError):
29+
pass

‎distributed_lock/sync.py

Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,240 @@
1+
from __future__ import annotations
2+
3+
import datetime
4+
import logging
5+
import os
6+
import time
7+
from contextlib import contextmanager
8+
from dataclasses import asdict, dataclass, field
9+
from typing import Any
10+
11+
import httpx
12+
13+
from distributed_lock.const import DEFAULT_CLUSTER, DEFAULT_LIFETIME, DEFAULT_WAIT
14+
from distributed_lock.exception import (
15+
BadConfigurationError,
16+
DistributedLockError,
17+
DistributedLockException,
18+
NotAcquiredError,
19+
NotAcquiredException,
20+
NotReleasedError,
21+
NotReleasedException,
22+
)
23+
24+
logger = logging.getLogger("distributed-lock.sync")
25+
26+
27+
def get_cluster() -> str:
28+
if os.environ.get("DLOCK_CLUSTER"):
29+
return os.environ["DLOCK_CLUSTER"].lower().strip()
30+
return DEFAULT_CLUSTER
31+
32+
33+
def get_token() -> str:
34+
if os.environ.get("DLOCK_TOKEN"):
35+
return os.environ["DLOCK_TOKEN"].lower().strip()
36+
raise BadConfigurationError("You must provide a token (or set DLOCK_TOKEN env var)")
37+
38+
39+
def get_tenant_id() -> str:
40+
if os.environ.get("DLOCK_TENANT_ID"):
41+
return os.environ["DLOCK_TENANT_ID"].lower().strip()
42+
raise BadConfigurationError(
43+
"You must provide a tenant_id (or set DLOCK_TENANT_ID env var)"
44+
)
45+
46+
47+
def make_httpx_client() -> httpx.Client:
48+
timeout = httpx.Timeout(connect=10.0, read=65.0, write=10.0, pool=10.0)
49+
return httpx.Client(timeout=timeout)
50+
51+
52+
@dataclass
53+
class AcquiredRessource:
54+
resource: str
55+
lock_id: str
56+
57+
@classmethod
58+
def from_dict(cls, d: dict) -> AcquiredRessource:
59+
for f in ("lock_id", "resource"):
60+
if f not in d:
61+
raise DistributedLockError(f"bad reply from service, missing {f}")
62+
return cls(resource=d["resource"], lock_id=d["lock_id"])
63+
64+
def to_dict(self) -> dict:
65+
return asdict(self)
66+
67+
68+
@dataclass
69+
class DistributedLockClient:
70+
cluster: str = field(default_factory=get_cluster)
71+
token: str = field(default_factory=get_token)
72+
tenant_id: str = field(default_factory=get_tenant_id)
73+
client: httpx.Client = field(default_factory=make_httpx_client)
74+
user_agent: str | None = None
75+
service_wait: int = DEFAULT_WAIT
76+
77+
def get_resource_url(self, resource: str) -> str:
78+
return f"https://{self.cluster}.distributed-lock.com/exclusive_locks/{self.tenant_id}/{resource}"
79+
80+
def get_headers(self) -> dict[str, str]:
81+
return {"Authorization": f"Bearer {self.token}"}
82+
83+
def __del__(self):
84+
self.client.close()
85+
86+
def _acquire(
87+
self,
88+
resource: str,
89+
lifetime: int = DEFAULT_LIFETIME,
90+
user_data: str | None = None,
91+
) -> AcquiredRessource:
92+
body: dict[str, Any] = {"wait": self.service_wait, "lifetime": lifetime}
93+
if self.user_agent:
94+
body["user_agent"] = self.user_agent
95+
if user_data:
96+
body["user_data"] = user_data
97+
url = self.get_resource_url(resource)
98+
logger.debug(f"Try to lock {resource} with url: {url}...")
99+
try:
100+
r = self.client.post(url, json=body, headers=self.get_headers())
101+
except httpx.ConnectTimeout as e:
102+
logger.warning(f"connect timeout error during POST on {url}")
103+
raise NotAcquiredError("timeout during connect") from e
104+
except httpx.ReadTimeout as e:
105+
logger.warning(f"read timeout error during POST on {url}")
106+
raise NotAcquiredError("timeout during read") from e
107+
except httpx.WriteTimeout as e:
108+
logger.warning(f"write timeout error during POST on {url}")
109+
raise NotAcquiredError("timeout during write") from e
110+
except httpx.PoolTimeout as e:
111+
logger.warning("timeout in connection pool")
112+
raise NotAcquiredError("timeout in connection pool") from e
113+
except httpx.HTTPError as e:
114+
logger.warning("generic http error")
115+
raise NotAcquiredError("generic http error") from e
116+
if r.status_code == 409:
117+
logger.info(f"Lock on {resource} NOT acquired")
118+
raise NotAcquiredException()
119+
# FIXME other codes
120+
d = r.json()
121+
logger.info(f"Lock on {resource} acquired")
122+
return AcquiredRessource.from_dict(d)
123+
124+
def acquire_exclusive_lock(
125+
self,
126+
resource: str,
127+
lifetime: int = DEFAULT_LIFETIME,
128+
wait: int = DEFAULT_WAIT,
129+
user_data: str | None = None,
130+
automatic_retry: bool = True,
131+
sleep_after_unsuccessful: float = 1.0,
132+
) -> AcquiredRessource:
133+
before = datetime.datetime.utcnow()
134+
while True:
135+
catched_exception: Exception | None = None
136+
try:
137+
return self._acquire(
138+
resource=resource, lifetime=lifetime, user_data=user_data
139+
)
140+
except DistributedLockError as e:
141+
if not automatic_retry:
142+
raise
143+
catched_exception = e
144+
except DistributedLockException as e:
145+
catched_exception = e
146+
elapsed = (datetime.datetime.utcnow() - before).total_seconds()
147+
if elapsed > wait - sleep_after_unsuccessful:
148+
raise catched_exception
149+
logger.debug(f"wait {sleep_after_unsuccessful}s...")
150+
time.sleep(sleep_after_unsuccessful)
151+
if elapsed + sleep_after_unsuccessful + self.service_wait > wait:
152+
self.service_wait = max(
153+
int(wait - elapsed - sleep_after_unsuccessful), 1
154+
)
155+
156+
def _release(self, resource: str, lock_id: str):
157+
url = self.get_resource_url(resource) + "/" + lock_id
158+
logger.debug(f"Try to unlock {resource} with url: {url}...")
159+
try:
160+
r = self.client.delete(url, headers=self.get_headers())
161+
except httpx.ConnectTimeout as e:
162+
logger.warning(f"connect timeout error during DELETE on {url}")
163+
raise NotReleasedError("timeout during connect") from e
164+
except httpx.ReadTimeout as e:
165+
logger.warning(f"read timeout error during DELTE on {url}")
166+
raise NotReleasedError("timeout during read") from e
167+
except httpx.WriteTimeout as e:
168+
logger.warning(f"write timeout error during DELETE on {url}")
169+
raise NotReleasedError("timeout during write") from e
170+
except httpx.PoolTimeout as e:
171+
logger.warning("timeout in connection pool")
172+
raise NotReleasedError("timeout in connection pool") from e
173+
except httpx.HTTPError as e:
174+
logger.warning("generic http error")
175+
raise NotReleasedError("generic http error") from e
176+
if r.status_code == 409:
177+
logger.warning(
178+
f"Lock on {resource} NOT released (because it's acquired by another lock_id!)"
179+
)
180+
raise NotReleasedException()
181+
if r.status_code == 204:
182+
return
183+
logger.warning(
184+
f"Lock on {resource} NOT released (because of an unexpected status code: {r.status_code})"
185+
)
186+
raise NotReleasedError(f"unexpected status code: {r.status_code}")
187+
188+
def release_exclusive_lock(
189+
self,
190+
resource: str,
191+
lock_id: str,
192+
wait: int = 30,
193+
automatic_retry: bool = True,
194+
sleep_after_unsuccessful: float = 1.0,
195+
):
196+
before = datetime.datetime.utcnow()
197+
while True:
198+
catched_exception = None
199+
try:
200+
return self._release(resource=resource, lock_id=lock_id)
201+
except DistributedLockError as e:
202+
if not automatic_retry:
203+
raise
204+
catched_exception = e
205+
elapsed = (datetime.datetime.utcnow() - before).total_seconds()
206+
if elapsed > wait - sleep_after_unsuccessful:
207+
raise catched_exception
208+
logger.debug(f"wait {sleep_after_unsuccessful}s...")
209+
time.sleep(sleep_after_unsuccessful)
210+
211+
@contextmanager
212+
def exclusive_lock(
213+
self,
214+
resource: str,
215+
lifetime: int = DEFAULT_LIFETIME,
216+
wait: int = DEFAULT_WAIT,
217+
user_data: str | None = None,
218+
automatic_retry: bool = True,
219+
sleep_after_unsuccessful: float = 1.0,
220+
):
221+
ar: AcquiredRessource | None = None
222+
try:
223+
ar = self.acquire_exclusive_lock(
224+
resource=resource,
225+
lifetime=lifetime,
226+
wait=wait,
227+
user_data=user_data,
228+
automatic_retry=automatic_retry,
229+
sleep_after_unsuccessful=sleep_after_unsuccessful,
230+
)
231+
yield
232+
finally:
233+
if ar is not None:
234+
self.release_exclusive_lock(
235+
resource=resource,
236+
lock_id=ar.lock_id,
237+
wait=wait,
238+
automatic_retry=automatic_retry,
239+
sleep_after_unsuccessful=sleep_after_unsuccessful,
240+
)

‎poetry.lock

Lines changed: 111 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
[tool.poetry]
2-
name = "distributed-lock"
2+
name = "distributed-lock-client"
33
version = "0.1.0"
44
description = ""
55
authors = ["fab-dlock <fab@distributed-lock.com>"]
@@ -17,7 +17,27 @@ black = "^23.7.0"
1717
ruff = "^0.0.278"
1818
mypy = "^1.4.1"
1919
pytest = "^7.4.0"
20+
respx = "^0.20.2"
21+
pytest-cov = "^4.1.0"
2022

2123
[build-system]
2224
requires = ["poetry-core"]
2325
build-backend = "poetry.core.masonry.api"
26+
27+
[tool.ruff]
28+
# Enable Pyflakes `E` and `F` codes by default.
29+
select = ["E", "F", "W", "N", "UP", "B", "I", "PL", "RUF"]
30+
ignore = [
31+
"E501",
32+
"PLR2004",
33+
"PLR0913",
34+
"PLW0603",
35+
"N805",
36+
"N818"
37+
]
38+
line-length = 88
39+
target-version = "py38"
40+
extend-exclude = []
41+
42+
[tool.ruff.isort]
43+
required-imports = ["from __future__ import annotations"]

‎tests/test_misc.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
from __future__ import annotations
2+
3+
import os
4+
from unittest import mock
5+
6+
import pytest
7+
8+
from distributed_lock import DEFAULT_CLUSTER, BadConfigurationError
9+
from distributed_lock.sync import (
10+
get_cluster,
11+
get_tenant_id,
12+
get_token,
13+
)
14+
15+
16+
@mock.patch.dict(os.environ, {"DLOCK_CLUSTER": "foo"}, clear=True)
17+
def test_get_cluster():
18+
assert get_cluster() == "foo"
19+
20+
21+
@mock.patch.dict(os.environ, {}, clear=True)
22+
def test_get_cluster2():
23+
assert get_cluster() == DEFAULT_CLUSTER
24+
25+
26+
@mock.patch.dict(os.environ, {"DLOCK_TOKEN": "foo"}, clear=True)
27+
def test_get_token():
28+
assert get_token() == "foo"
29+
30+
31+
@mock.patch.dict(os.environ, {}, clear=True)
32+
def test_get_token2():
33+
with pytest.raises(BadConfigurationError):
34+
get_token()
35+
36+
37+
@mock.patch.dict(os.environ, {"DLOCK_TENANT_ID": "foo"}, clear=True)
38+
def test_get_tenant_id():
39+
assert get_tenant_id() == "foo"
40+
41+
42+
@mock.patch.dict(os.environ, {}, clear=True)
43+
def test_get_tenant_id2():
44+
with pytest.raises(BadConfigurationError):
45+
get_tenant_id()

‎tests/test_sync.py

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
from __future__ import annotations
2+
3+
import json
4+
import os
5+
from unittest import mock
6+
7+
import httpx
8+
import pytest
9+
10+
from distributed_lock import (
11+
DEFAULT_LIFETIME,
12+
DEFAULT_WAIT,
13+
AcquiredRessource,
14+
DistributedLockClient,
15+
NotAcquiredException,
16+
NotReleasedError,
17+
)
18+
19+
MOCKED_ENVIRON = {
20+
"DLOCK_CLUSTER": "cluster",
21+
"DLOCK_TENANT_ID": "tenant_id",
22+
"DLOCK_TOKEN": "token",
23+
}
24+
AR = AcquiredRessource(lock_id="1234", resource="bar")
25+
26+
27+
@mock.patch.dict(os.environ, MOCKED_ENVIRON, clear=True)
28+
def test_resource_url():
29+
x = DistributedLockClient()
30+
assert (
31+
x.get_resource_url("bar")
32+
== "https://cluster.distributed-lock.com/exclusive_locks/tenant_id/bar"
33+
)
34+
35+
36+
@mock.patch.dict(os.environ, MOCKED_ENVIRON, clear=True)
37+
@pytest.mark.respx(base_url="https://cluster.distributed-lock.com")
38+
def test_acquire(respx_mock):
39+
respx_mock.post("/exclusive_locks/tenant_id/bar").mock(
40+
return_value=httpx.Response(201, json=AR.to_dict())
41+
)
42+
x = DistributedLockClient()
43+
ar = x.acquire_exclusive_lock("bar")
44+
assert len(respx_mock.calls) == 1
45+
body = json.loads(respx_mock.calls.last.request.content.decode("utf8"))
46+
assert body["wait"] == DEFAULT_WAIT
47+
assert body["lifetime"] == DEFAULT_LIFETIME
48+
headers = respx_mock.calls.last.request.headers
49+
assert headers["host"] == "cluster.distributed-lock.com"
50+
assert headers["authorization"] == "Bearer token"
51+
assert headers["content-type"] == "application/json"
52+
assert ar is not None
53+
assert ar.lock_id == AR.lock_id
54+
assert ar.resource == AR.resource
55+
56+
57+
@mock.patch.dict(os.environ, MOCKED_ENVIRON, clear=True)
58+
@pytest.mark.respx(base_url="https://cluster.distributed-lock.com")
59+
def test_not_acquired(respx_mock):
60+
respx_mock.post("/exclusive_locks/tenant_id/bar").mock(
61+
return_value=httpx.Response(409)
62+
)
63+
x = DistributedLockClient()
64+
with pytest.raises(NotAcquiredException):
65+
x.acquire_exclusive_lock("bar", wait=3)
66+
assert len(respx_mock.calls) > 1
67+
68+
69+
@mock.patch.dict(os.environ, MOCKED_ENVIRON, clear=True)
70+
@pytest.mark.respx(base_url="https://cluster.distributed-lock.com")
71+
def test_release(respx_mock):
72+
respx_mock.delete("/exclusive_locks/tenant_id/bar/1234").mock(
73+
return_value=httpx.Response(204)
74+
)
75+
x = DistributedLockClient()
76+
x.release_exclusive_lock("bar", "1234")
77+
assert len(respx_mock.calls) == 1
78+
headers = respx_mock.calls.last.request.headers
79+
assert headers["host"] == "cluster.distributed-lock.com"
80+
assert headers["authorization"] == "Bearer token"
81+
82+
83+
@mock.patch.dict(os.environ, MOCKED_ENVIRON, clear=True)
84+
@pytest.mark.respx(base_url="https://cluster.distributed-lock.com")
85+
def test_not_released(respx_mock):
86+
respx_mock.delete("/exclusive_locks/tenant_id/bar/1234").mock(
87+
return_value=httpx.Response(500)
88+
)
89+
x = DistributedLockClient()
90+
with pytest.raises(NotReleasedError):
91+
x.release_exclusive_lock("bar", lock_id="1234", wait=3)
92+
assert len(respx_mock.calls) > 1
93+
headers = respx_mock.calls.last.request.headers
94+
assert headers["host"] == "cluster.distributed-lock.com"
95+
assert headers["authorization"] == "Bearer token"
96+
97+
98+
@mock.patch.dict(os.environ, MOCKED_ENVIRON, clear=True)
99+
@pytest.mark.respx(base_url="https://cluster.distributed-lock.com")
100+
def test_context_manager(respx_mock):
101+
respx_mock.post("/exclusive_locks/tenant_id/bar").mock(
102+
return_value=httpx.Response(201, json=AR.to_dict())
103+
)
104+
respx_mock.delete("/exclusive_locks/tenant_id/bar/1234").mock(
105+
return_value=httpx.Response(204)
106+
)
107+
x = DistributedLockClient()
108+
with x.exclusive_lock("bar"):
109+
pass
110+
assert len(respx_mock.calls) == 2

0 commit comments

Comments
 (0)
Please sign in to comment.