-
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
10 changed files
with
285 additions
and
14 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
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,3 +3,19 @@ | |
__author__ = "wh1isper" | ||
__email__ = "[email protected]" | ||
__version__ = "0.3.9.dev0" | ||
|
||
|
||
from .configs import BrqConfig | ||
from .consumer import Consumer | ||
from .decorator import task | ||
from .producer import Producer | ||
from .tools import get_redis_client, get_redis_url | ||
|
||
__all__ = [ | ||
"task", | ||
"Consumer", | ||
"Producer", | ||
"get_redis_client", | ||
"get_redis_url", | ||
"BrqConfig", | ||
] |
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,73 @@ | ||
import os | ||
import uuid | ||
from contextlib import asynccontextmanager | ||
from typing import Any, AsyncGenerator | ||
|
||
import redis.asyncio as redis | ||
from pydantic import BaseModel | ||
from pydantic_settings import BaseSettings, SettingsConfigDict | ||
|
||
from brq.tools import get_redis_client, get_redis_url | ||
|
||
|
||
class Settings(BaseSettings): | ||
model_config = SettingsConfigDict( | ||
env_prefix="brq_", | ||
env_file=".env", | ||
env_nested_delimiter="__", | ||
env_file_encoding="utf-8", | ||
case_sensitive=False, | ||
) | ||
|
||
|
||
class RedisSettingsMixin: | ||
redis_host: str = "localhost" | ||
redis_port: int = 6379 | ||
redis_db: int = 0 | ||
redis_cluster: bool = False | ||
redis_tls: bool = False | ||
redis_username: str = "" | ||
redis_password: str = "" | ||
|
||
@property | ||
def redis_url(self) -> str: | ||
return get_redis_url( | ||
host=self.redis_host, | ||
port=self.redis_port, | ||
db=self.redis_db, | ||
cluster=self.redis_cluster, | ||
tls=self.redis_tls, | ||
username=self.redis_username, | ||
password=self.redis_password, | ||
) | ||
|
||
@asynccontextmanager | ||
async def open_redis_client( | ||
self, | ||
) -> AsyncGenerator[Any, redis.Redis | redis.RedisCluster]: | ||
async with get_redis_client(self.redis_url) as redis_client: | ||
yield redis_client | ||
|
||
|
||
class BrqConfig(Settings, RedisSettingsMixin): | ||
redis_key_prefix: str = "brq" | ||
redis_key_seperator: str = ":" | ||
|
||
producer_max_message_length: int = 1000 | ||
|
||
consumer_group_name: str = "default-workers" | ||
consumer_identifier: str = uuid.uuid4().hex | ||
consumer_count_per_fetch: int = 1 | ||
consumer_block_time: int = 1 | ||
consumer_expire_time: int = 60 * 60 | ||
consumer_process_timeout: int = 60 | ||
consumer_retry_lock_time: int = 300 | ||
consumer_retry_cooldown_time: int = 60 | ||
consumer_enable_enque_deferred_job: bool = True | ||
consumer_enable_reprocess_timeout_job: bool = True | ||
consumer_enable_dead_queue: bool = True | ||
consumer_max_message_len: int = 1000 | ||
consumer_delete_message_after_process: bool = False | ||
consumer_run_parallel: bool = False | ||
|
||
daemon_concurrency: int = 1 |
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,89 @@ | ||
import asyncio | ||
from functools import partial | ||
from typing import Awaitable, Callable | ||
|
||
from anyio import CapacityLimiter | ||
|
||
from brq.configs import BrqConfig | ||
from brq.consumer import Consumer | ||
from brq.daemon import Daemon | ||
from brq.log import logger | ||
from brq.tools import ensure_awaitable | ||
|
||
|
||
class BrqTaskWrapper: | ||
def __init__( | ||
self, | ||
func: Callable | Awaitable, | ||
config: BrqConfig, | ||
register_function_name: str | None = None, | ||
): | ||
self._func = func | ||
self.config = config | ||
self.register_function_name = register_function_name or func.__name__ | ||
|
||
def __call__(self, *args, **kwargs): | ||
return self._func(*args, **kwargs) | ||
|
||
def serve(self): | ||
asyncio.run(self._serve()) | ||
|
||
async def _serve(self): | ||
async with self.config.open_redis_client() as async_redis_client: | ||
awaitable_function = ensure_awaitable( | ||
self._func, | ||
limiter=CapacityLimiter(total_tokens=self.config.daemon_concurrency), | ||
) | ||
consumer_builder = partial( | ||
Consumer, | ||
redis_prefix=self.config.redis_key_prefix, | ||
redis_seperator=self.config.redis_key_seperator, | ||
redis=async_redis_client, | ||
awaitable_function=awaitable_function, | ||
register_function_name=self.register_function_name, | ||
group_name=self.config.consumer_group_name, | ||
consumer_identifier=self.config.consumer_identifier, | ||
count_per_fetch=self.config.consumer_count_per_fetch, | ||
block_time=self.config.consumer_block_time, | ||
expire_time=self.config.consumer_expire_time, | ||
process_timeout=self.config.consumer_process_timeout, | ||
retry_lock_time=self.config.consumer_retry_lock_time, | ||
retry_cooldown_time=self.config.consumer_retry_cooldown_time, | ||
enable_enque_deferred_job=self.config.consumer_enable_enque_deferred_job, | ||
enable_reprocess_timeout_job=self.config.consumer_enable_reprocess_timeout_job, | ||
enable_dead_queue=self.config.consumer_enable_dead_queue, | ||
max_message_len=self.config.consumer_max_message_len, | ||
delete_message_after_process=self.config.consumer_delete_message_after_process, | ||
run_parallel=self.config.consumer_run_parallel, | ||
) | ||
daemon = Daemon(*[consumer_builder() for _ in range(self.config.daemon_concurrency)]) | ||
await daemon.run_forever() | ||
|
||
|
||
def task( | ||
_func=None, | ||
*, | ||
config: BrqConfig | None = None, | ||
register_function_name: str | None = None, | ||
): | ||
if not config: | ||
logger.info("Initializing config from environment variables and .env file") | ||
config = BrqConfig() | ||
else: | ||
logger.info("Using custom config") | ||
config = config | ||
|
||
def _wrapper(func, config, register_function_name): | ||
if not config: | ||
logger.info("Initializing config from environment variables and .env file") | ||
config = BrqConfig() | ||
else: | ||
logger.info("Using custom config") | ||
config = config | ||
|
||
return BrqTaskWrapper(func, config, register_function_name) | ||
|
||
if _func is None: | ||
return _wrapper | ||
else: | ||
return _wrapper(_func, config=config, register_function_name=register_function_name) | ||
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,5 @@ | ||
## Echo example | ||
|
||
Use [start-redis.sh](../../dev/start-redis.sh) to start redis in docker. | ||
|
||
And run `python producer.py` and `python consumer.py` in different terminals at the same time. |
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,14 @@ | ||
from brq import task | ||
|
||
|
||
@task | ||
def echo(message): | ||
print(f"Received message: {message}") | ||
|
||
|
||
if __name__ == "__main__": | ||
# Run the task once, for local debug | ||
# echo("hello") | ||
|
||
# Run as a daemon | ||
echo.serve() |
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,21 @@ | ||
import os | ||
|
||
from brq.configs import BrqConfig | ||
from brq.producer import Producer | ||
|
||
|
||
async def main(): | ||
config = BrqConfig() | ||
async with config.open_redis_client() as async_redis_client: | ||
await Producer( | ||
async_redis_client, | ||
redis_prefix=config.redis_key_prefix, | ||
redis_seperator=config.redis_key_seperator, | ||
max_message_len=config.producer_max_message_length, | ||
).run_job("echo", ["hello"]) | ||
|
||
|
||
if __name__ == "__main__": | ||
import asyncio | ||
|
||
asyncio.run(main()) |
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 |
---|---|---|
@@ -1,3 +1,5 @@ | ||
## Echo example | ||
|
||
Just run `python producer.py` and `python consumer.py` in different terminals at the same time. | ||
Use [start-redis.sh](../../dev/start-redis.sh) to start redis in docker. | ||
|
||
And run `python producer.py` and `python consumer.py` in different terminals at the same time. |
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