-
Notifications
You must be signed in to change notification settings - Fork 24
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Update version to v0.2.24 (#132) Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> Co-authored-by: Khor Shu Heng <khor.heng@gojek.com> Signed-off-by: KeshavSharma <keshav.sharma@gojek.com> * distributed lock for starting/stopping streaming ingestion jobs Signed-off-by: KeshavSharma <keshav.sharma@gojek.com> * unit tests for lock manager Signed-off-by: KeshavSharma <keshav.sharma@gojek.com> * Perform data type conversion automatically (#133) * Perform data type conversion automatically Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> * Use wheel installation for local setup to avoid module not found issue Signed-off-by: Khor Shu Heng <khor.heng@gojek.com> Co-authored-by: Khor Shu Heng <khor.heng@gojek.com> Signed-off-by: KeshavSharma <keshav.sharma@gojek.com> * format python files Signed-off-by: KeshavSharma <keshav.sharma@gojek.com> * install feast-spark from setup.py Signed-off-by: KeshavSharma <keshav.sharma@gojek.com> Co-authored-by: Khor Shu Heng <32997938+khorshuheng@users.noreply.github.com> Co-authored-by: Khor Shu Heng <khor.heng@gojek.com> Co-authored-by: KeshavSharma <keshav.sharma@gojek.com>
- Loading branch information
4 people
committed
Apr 12, 2022
1 parent
7d3aa9d
commit 6499fb2
Showing
6 changed files
with
247 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,115 @@ | ||
| """ | ||
| Classes to manage distributed locks | ||
| """ | ||
| import enum | ||
| import logging | ||
| import secrets | ||
| import time | ||
|
|
||
| import redis | ||
| from redis.exceptions import ConnectionError | ||
|
|
||
| # retries for acquiring lock | ||
| LOCK_ACQUIRE_RETRIES = 3 | ||
| # wait between retries | ||
| LOCK_ACQUIRE_WAIT = 1 | ||
| LOCK_KEY_PREFIX = "lock" | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class JobOperation(enum.Enum): | ||
| """ | ||
| Enum for job operations | ||
| """ | ||
|
|
||
| START = "st" | ||
| CANCEL = "cn" | ||
|
|
||
|
|
||
| class JobOperationLock: | ||
| """ | ||
| Lock for starting and cancelling spark ingestion jobs. | ||
| Implemented as a context manager to automatically release lock after operation. | ||
| Usage: | ||
| with JobOperationLock(job_hash, <start/cancel>): | ||
| client.start_stream_to_online_ingestion(feature_table, [], project=project) | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| redis_host: str, | ||
| redis_port: int, | ||
| lock_expiry: int, | ||
| job_hash: str, | ||
| operation: JobOperation = JobOperation.START, | ||
| ): | ||
| """ | ||
| Init method, initialized redis key for the lock | ||
| Args: | ||
| redis_host: host to redis instance to store locks | ||
| redis_port: port to redis instance to store locks | ||
| lock_expiry: time in seconds for auto releasing lock | ||
| job_hash: job hash string for the job which needs to be operated upon | ||
| operation: operation to be performed <START/CANCEL> | ||
| """ | ||
| self._redis = redis.Redis(host=redis_host, port=redis_port) | ||
| self._lock_expiry = lock_expiry | ||
| self._lock_key = f"{LOCK_KEY_PREFIX}_{operation.value}_{job_hash}" | ||
| self._lock_value = secrets.token_hex(nbytes=8) | ||
|
|
||
| def __enter__(self): | ||
| """ | ||
| Context manager method for setup - acquire lock | ||
| lock_key is a combination of a prefix, job hash and operation(start/cancel) | ||
| lock_value is a randomly generated 8 byte hexadecimal, this is to ensure | ||
| that lock can be deleted only by the agent who created it | ||
| NX option is used only set the key if it does not already exist, | ||
| this will ensure that locks are not overwritten | ||
| EX option is used to set the specified expire time to release the lock automatically after TTL | ||
| """ | ||
| # Retry acquiring lock on connection failures | ||
| retry_attempts = 0 | ||
| while retry_attempts < LOCK_ACQUIRE_RETRIES: | ||
| try: | ||
| if self._redis.set( | ||
| name=self._lock_key, | ||
| value=self._lock_value, | ||
| nx=True, | ||
| ex=self._lock_expiry, | ||
| ): | ||
| return self._lock_value | ||
| else: | ||
| logger.info(f"lock not available: {self._lock_key}") | ||
| return False | ||
| except ConnectionError: | ||
| # wait before attempting to retry | ||
| logger.warning( | ||
| f"connection error while acquiring lock: {self._lock_key}" | ||
| ) | ||
| time.sleep(LOCK_ACQUIRE_WAIT) | ||
| retry_attempts += 1 | ||
| logger.warning(f"Can't acquire lock, backing off: {self._lock_key}") | ||
| return False | ||
|
|
||
| def __exit__(self, *args, **kwargs): | ||
| """ | ||
| context manager method for teardown - release lock | ||
| safe release - delete lock key only if value exists and is same as set by this object | ||
| otherwise rely on auto-release on expiry | ||
| """ | ||
| try: | ||
| lock_value = self._redis.get(self._lock_key) | ||
| if lock_value and lock_value.decode() == self._lock_value: | ||
| self._redis.delete(self._lock_key) | ||
| except ConnectionError: | ||
| logger.warning( | ||
| f"connection error while deleting lock: {self._lock_key}." | ||
| f"rely on auto-release after {self._lock_expiry} seconds" | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,71 @@ | ||
| from unittest.mock import patch | ||
|
|
||
| import pytest | ||
|
|
||
| from feast_spark.lock_manager import JobOperation, JobOperationLock | ||
|
|
||
| job_hash = "dummy_hash" | ||
|
|
||
|
|
||
| class MockRedis: | ||
| def __init__(self, cache=dict()): | ||
| self.cache = cache | ||
|
|
||
| def get(self, name): | ||
| if name in self.cache: | ||
| return self.cache[name] | ||
| return None | ||
|
|
||
| def set(self, name, value, *args, **kwargs): | ||
| if name not in self.cache: | ||
| self.cache[name] = value.encode("utf-8") | ||
| return "OK" | ||
|
|
||
| def delete(self, name): | ||
| if name in self.cache: | ||
| self.cache.pop(name) | ||
| return None | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| def lock_config(): | ||
| return {"redis_host": "localhost", "redis_port": 0, "lock_expiry": 5} | ||
|
|
||
|
|
||
| @patch("redis.Redis") | ||
| def test_lock_manager_context(mock_redis, lock_config): | ||
| mock_redis_connection = MockRedis() | ||
| mock_redis.return_value = mock_redis_connection | ||
| with JobOperationLock( | ||
| job_hash=job_hash, operation=JobOperation.START, **lock_config | ||
| ) as lock: | ||
| # test lock acquired | ||
| assert lock | ||
| # verify lock key in cache | ||
| assert ( | ||
| f"lock_{JobOperation.START.value}_{job_hash}" in mock_redis_connection.cache | ||
| ) | ||
| # verify release | ||
| assert ( | ||
| f"lock_{JobOperation.START.value}_{job_hash}" not in mock_redis_connection.cache | ||
| ) | ||
|
|
||
|
|
||
| @patch("redis.Redis") | ||
| def test_lock_manager_lock_not_available(mock_redis, lock_config): | ||
| cache = {"lock_st_dummy_hash": b"127a32aaf729dc87"} | ||
| mock_redis_connection = MockRedis(cache) | ||
| mock_redis.return_value = mock_redis_connection | ||
| with JobOperationLock( | ||
| job_hash=job_hash, operation=JobOperation.START, **lock_config | ||
| ) as lock: | ||
| # test lock not acquired | ||
| assert not lock | ||
|
|
||
|
|
||
| def test_lock_manager_connection_error(lock_config): | ||
| with JobOperationLock( | ||
| job_hash=job_hash, operation=JobOperation.START, **lock_config | ||
| ) as lock: | ||
| # test lock not acquired | ||
| assert not lock |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters