mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
Merge branch 'litellm_alerts_pr1_ci_backend' into litellm_alerts_reconcile_02_final
This commit is contained in:
commit
99b31f772a
443 changed files with 31060 additions and 8089 deletions
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -53,7 +53,7 @@ jobs:
|
|||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
|
|
|
|||
2
.github/workflows/test-linting.yml
vendored
2
.github/workflows/test-linting.yml
vendored
|
|
@ -21,7 +21,7 @@ jobs:
|
|||
pull-requests: read
|
||||
actions: read
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 30
|
||||
services:
|
||||
redis:
|
||||
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
|
||||
|
|
|
|||
19
.github/workflows/test-unit.yml
vendored
19
.github/workflows/test-unit.yml
vendored
|
|
@ -45,7 +45,7 @@ jobs:
|
|||
id: rust-bridge-key
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }}
|
||||
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml', 'packaging/litellm-proxy/pyproject.toml') }}
|
||||
run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Set up Python
|
||||
|
|
@ -98,6 +98,18 @@ jobs:
|
|||
mkdir -p rust-bridge-dist/litellm/rust_bridge
|
||||
cp litellm/rust_bridge/_native.abi3.so rust-bridge-dist/litellm/rust_bridge/_native.abi3.so
|
||||
|
||||
- name: Build the gateway wheel once for test shards
|
||||
if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true'
|
||||
timeout-minutes: 15
|
||||
working-directory: packaging/litellm-proxy
|
||||
env:
|
||||
RUSTC_WRAPPER: sccache
|
||||
SCCACHE_GHA_ENABLED: "true"
|
||||
SCCACHE_GHA_VERSION: litellm-ci-v1
|
||||
SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }}
|
||||
CARGO_INCREMENTAL: "0"
|
||||
run: uv run --no-project --with maturin==1.15.0 maturin build --profile dev --strip --out ../../rust-bridge-dist
|
||||
|
||||
- name: Upload Rust bridge artifact
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 10
|
||||
|
|
@ -170,6 +182,7 @@ jobs:
|
|||
tests/unit/google_genai
|
||||
tests/unit/router_strategy
|
||||
tests/unit/router_utils
|
||||
tests/unit/redis
|
||||
tests/unit/proxy/common_utils/test_cache_aware_routing.py
|
||||
tests/unit/enterprise/enterprise_callbacks/send_emails
|
||||
tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py
|
||||
|
|
@ -1025,7 +1038,7 @@ jobs:
|
|||
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
else
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
|
||||
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
|
||||
|
|
@ -1487,7 +1500,7 @@ jobs:
|
|||
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
else
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -76,6 +76,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
|||
COPY pyproject.toml uv.lock ./
|
||||
COPY enterprise/pyproject.toml enterprise/
|
||||
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
|
||||
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
|
||||
|
||||
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
|
||||
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
|
|
@ -101,7 +102,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
|
|||
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
|
||||
|
||||
# Install project and workspace packages (fast - deps already cached)
|
||||
RUN uv sync --frozen --no-default-groups --no-editable \
|
||||
RUN uv build --package litellm-proxy --wheel --out-dir /wheels
|
||||
RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -111,6 +113,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
|
|
|||
|
|
@ -74,6 +74,7 @@ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
|||
COPY pyproject.toml uv.lock ./
|
||||
COPY enterprise/pyproject.toml enterprise/
|
||||
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
|
||||
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
|
||||
|
||||
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
|
||||
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
|
|
@ -99,7 +100,8 @@ COPY --from=ui-builder /ui/out/. litellm/proxy/_experimental/out/
|
|||
RUN sed -i 's/\r$//' docker/build_admin_ui.sh && chmod +x docker/build_admin_ui.sh && ./docker/build_admin_ui.sh
|
||||
|
||||
# Install project and workspace packages (fast - deps already cached)
|
||||
RUN uv sync --frozen --no-default-groups --no-editable \
|
||||
RUN uv build --package litellm-proxy --wheel --out-dir /wheels
|
||||
RUN uv sync --frozen --no-install-package litellm-proxy --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -109,6 +111,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm_proxy-*.whl
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ import re
|
|||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
|
|
@ -766,7 +765,6 @@ class CheckBatchCost:
|
|||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
_file_content = await afile_content(
|
||||
file_id=cast(str, raw_output_file_id),
|
||||
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
|
||||
**cast(FileContentCallOptions, credentials),
|
||||
)
|
||||
|
||||
|
|
@ -867,10 +865,7 @@ class CheckBatchCost:
|
|||
error_file_failed_requests: Final = await count_error_file_failed_requests(
|
||||
response,
|
||||
custom_llm_provider=batch_file_provider,
|
||||
litellm_params={
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
},
|
||||
litellm_params=credentials,
|
||||
)
|
||||
batch_result: Final = (
|
||||
output_file_result
|
||||
|
|
@ -1029,11 +1024,7 @@ class CheckBatchCost:
|
|||
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
|
||||
)
|
||||
cost_provider: Final = cast(BatchFileProvider, llm_provider)
|
||||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
file_access: Final = {
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
}
|
||||
file_access: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
output_content: Final = (
|
||||
await fetch_batch_managed_file_content(
|
||||
response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access
|
||||
|
|
|
|||
|
|
@ -127,7 +127,6 @@ class _ManagedFileRetrieve(Protocol):
|
|||
self,
|
||||
*,
|
||||
file_id: str,
|
||||
_litellm_internal_model_credentials: Mapping[str, object] | None = None,
|
||||
**kwargs: Unpack[FileRetrieveCallOptions],
|
||||
) -> OpenAIFileObject: ...
|
||||
|
||||
|
|
@ -400,7 +399,7 @@ class _RouterFileCallKwargs(TypedDict, total=False):
|
|||
|
||||
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
|
||||
_FILE_UPLOAD_ONLY_FIELDS: Final = frozenset({"file", "purpose", "expires_after", "extra_body"})
|
||||
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id", "_litellm_internal_model_credentials"})
|
||||
_FILE_DELETE_RESERVED_FIELDS: Final = frozenset({"model", "file_id"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -459,15 +458,7 @@ async def _delete_provider_file(
|
|||
provider_file_id: str,
|
||||
request_data: Mapping[str, object],
|
||||
) -> None:
|
||||
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data: Final = {
|
||||
**{k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS},
|
||||
**(
|
||||
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
delete_data: Final = {k: v for k, v in request_data.items() if k not in _FILE_DELETE_RESERVED_FIELDS}
|
||||
await llm_router.afile_delete(
|
||||
model=model_id,
|
||||
file_id=provider_file_id,
|
||||
|
|
@ -1160,7 +1151,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
*,
|
||||
provider_file_id: str,
|
||||
call_options: Mapping[str, object],
|
||||
internal_model_credentials: Mapping[str, object] | None = None,
|
||||
) -> OpenAIFileObject:
|
||||
"""Retrieve provider file details with transient retries and SDK retries disabled."""
|
||||
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
|
||||
|
|
@ -1172,14 +1162,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if attempt > 0:
|
||||
await self._sleep(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS[attempt - 1])
|
||||
try:
|
||||
if internal_model_credentials is None:
|
||||
return await retrieve_file(
|
||||
file_id=provider_file_id,
|
||||
**retrieve_options,
|
||||
)
|
||||
return await retrieve_file(
|
||||
file_id=provider_file_id,
|
||||
_litellm_internal_model_credentials=MappingProxyType(dict(internal_model_credentials)),
|
||||
**retrieve_options,
|
||||
)
|
||||
except Exception as error:
|
||||
|
|
@ -1221,7 +1205,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
self._afile_retrieve_with_retries(
|
||||
provider_file_id=provider_file_id,
|
||||
call_options=router_credentials,
|
||||
internal_model_credentials=router_credentials,
|
||||
),
|
||||
timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
@ -2866,15 +2849,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
exception_dict = {}
|
||||
for model_id, provider_file_id in specific_model_file_id_mapping.items():
|
||||
try:
|
||||
# Cloud-storage providers (e.g. Bedrock S3) validate file ids
|
||||
# against the deployment's configured bucket, which they only
|
||||
# trust from this immutable server-side snapshot, never from
|
||||
# request params.
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
if credentials is not None:
|
||||
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
|
||||
else:
|
||||
data.pop("_litellm_internal_model_credentials", None)
|
||||
return cast(
|
||||
HttpxBinaryResponseContent,
|
||||
await llm_router.afile_content(
|
||||
|
|
|
|||
|
|
@ -41,10 +41,14 @@ COPY pyproject.toml README.md rust-toolchain.toml ./
|
|||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
COPY model_prices_and_context_window.json ./
|
||||
COPY packaging/litellm-proxy/pyproject.toml packaging/litellm-proxy/
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
|
||||
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
|
||||
python -c "import maturin; maturin.build_wheel('/wheels')" && \
|
||||
cd packaging/litellm-proxy && \
|
||||
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
|
||||
python -c "import maturin; maturin.build_wheel('/wheels')"
|
||||
|
||||
|
|
@ -85,6 +89,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--mount=type=bind,source=uv.lock,target=uv.lock \
|
||||
--mount=type=bind,source=enterprise/pyproject.toml,target=enterprise/pyproject.toml \
|
||||
--mount=type=bind,source=litellm-proxy-extras/pyproject.toml,target=litellm-proxy-extras/pyproject.toml \
|
||||
--mount=type=bind,source=packaging/litellm-proxy/pyproject.toml,target=packaging/litellm-proxy/pyproject.toml \
|
||||
uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
|
|
@ -98,7 +103,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
COPY . .
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv sync --frozen --no-install-project --no-default-groups --no-editable \
|
||||
uv sync --frozen --no-install-project --no-install-package litellm-proxy --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--group admin-mcp \
|
||||
|
|
@ -108,7 +113,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--python python3.13
|
||||
|
||||
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
|
||||
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
|
||||
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl /wheels/litellm_proxy-*.whl
|
||||
|
||||
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
|
||||
# compiled Rust extension must live next to the source or it is never imported.
|
||||
|
|
|
|||
53
litellm-rust/.agents/skills/rust-python-boundary/SKILL.md
Normal file
53
litellm-rust/.agents/skills/rust-python-boundary/SKILL.md
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
---
|
||||
name: rust-python-boundary
|
||||
description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Python boundary
|
||||
|
||||
Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value
|
||||
|
||||
## The two boundaries
|
||||
|
||||
- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat`
|
||||
- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python
|
||||
- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under
|
||||
|
||||
## Python value semantics
|
||||
|
||||
Reading a value the way Python itself reads it. This lives only at the two boundaries
|
||||
|
||||
- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)`
|
||||
- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false`
|
||||
- `os.environ/` references and `REDIS_<KWARG>` style environment layering
|
||||
|
||||
A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on
|
||||
|
||||
Each typed field keeps where its value came from (the key, `REDIS_<KWARG>` or the fallback name), so an error or warning downstream still names it
|
||||
|
||||
## What stays downstream
|
||||
|
||||
A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps
|
||||
|
||||
- Which target wins, and the defaults redis-py applies
|
||||
- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters
|
||||
- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one
|
||||
|
||||
Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped
|
||||
|
||||
## litellm-python-compat
|
||||
|
||||
The rule is about where config and kwargs get coerced, not about who may depend on the crate
|
||||
|
||||
- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis
|
||||
- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself
|
||||
|
||||
## Existing code
|
||||
|
||||
Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place
|
||||
|
||||
## Tests
|
||||
|
||||
Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce
|
||||
|
|
@ -8,6 +8,10 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust
|
|||
|
||||
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
|
||||
|
||||
For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md)
|
||||
|
||||
The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics
|
||||
|
||||
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
|
||||
|
||||
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`
|
||||
|
|
|
|||
207
litellm-rust/Cargo.lock
generated
207
litellm-rust/Cargo.lock
generated
|
|
@ -8,6 +8,17 @@ version = "2.0.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
|
||||
|
||||
[[package]]
|
||||
name = "aes"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cipher",
|
||||
"cpufeatures 0.2.17",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ahash"
|
||||
version = "0.8.12"
|
||||
|
|
@ -982,6 +993,15 @@ dependencies = [
|
|||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "block-padding"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93"
|
||||
dependencies = [
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bollard"
|
||||
version = "0.20.2"
|
||||
|
|
@ -1134,6 +1154,15 @@ dependencies = [
|
|||
"rustversion",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cbc"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6"
|
||||
dependencies = [
|
||||
"cipher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cc"
|
||||
version = "1.3.0"
|
||||
|
|
@ -1210,6 +1239,16 @@ dependencies = [
|
|||
"half",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cipher"
|
||||
version = "0.4.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
|
||||
dependencies = [
|
||||
"crypto-common 0.1.7",
|
||||
"inout",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clap"
|
||||
version = "4.6.7"
|
||||
|
|
@ -1774,6 +1813,15 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "des"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e"
|
||||
dependencies = [
|
||||
"cipher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.7"
|
||||
|
|
@ -3105,6 +3153,16 @@ dependencies = [
|
|||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "inout"
|
||||
version = "0.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01"
|
||||
dependencies = [
|
||||
"block-padding",
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ipnet"
|
||||
version = "2.12.0"
|
||||
|
|
@ -3663,8 +3721,7 @@ dependencies = [
|
|||
"strum",
|
||||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3767,6 +3824,7 @@ dependencies = [
|
|||
"http-body-util",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-config",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-gateway-forward",
|
||||
|
|
@ -3777,11 +3835,15 @@ dependencies = [
|
|||
"litellm-inference",
|
||||
"litellm-llms",
|
||||
"litellm-redis-client",
|
||||
"litellm-redis-connect",
|
||||
"litellm-redis-resolve",
|
||||
"litellm-redis-testing",
|
||||
"litellm-redis-types",
|
||||
"litellm-secret-value",
|
||||
"litellm-secrets",
|
||||
"litellm-tracing",
|
||||
"rstest",
|
||||
"rustls 0.23.42",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
|
|
@ -3792,6 +3854,7 @@ dependencies = [
|
|||
"tower-sessions-moka-store",
|
||||
"tracing",
|
||||
"uuid",
|
||||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4323,6 +4386,7 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-response",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
|
|
@ -4337,6 +4401,12 @@ dependencies = [
|
|||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-model-catalog",
|
||||
"litellm-python-compat",
|
||||
"litellm-redis-client",
|
||||
"litellm-redis-connect",
|
||||
"litellm-redis-resolve",
|
||||
"litellm-redis-store",
|
||||
"litellm-redis-types",
|
||||
"litellm-router-core",
|
||||
"litellm-router-types",
|
||||
"litellm-secret-value",
|
||||
|
|
@ -4387,23 +4457,108 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-python-compat",
|
||||
"litellm-redis-testing",
|
||||
"litellm-redis-types",
|
||||
"litellm-secret-value",
|
||||
"pem",
|
||||
"pkcs8",
|
||||
"redis",
|
||||
"rstest",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
"socket2 0.6.5",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-connect"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-auth-aws",
|
||||
"litellm-auth-azure",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-auth-types",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-redis-client",
|
||||
"litellm-redis-resolve",
|
||||
"litellm-redis-testing",
|
||||
"litellm-redis-types",
|
||||
"litellm-secret-value",
|
||||
"rstest",
|
||||
"rustls 0.23.42",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-resolve"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-python-compat",
|
||||
"litellm-redis-types",
|
||||
"litellm-secret-value",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"tracing",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-store"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"litellm-cache",
|
||||
"litellm-cache-memory",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-python-compat",
|
||||
"litellm-redis-client",
|
||||
"litellm-redis-testing",
|
||||
"litellm-secret-value",
|
||||
"litellm-tracing",
|
||||
"rstest",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-testing"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"pkcs8",
|
||||
"rcgen",
|
||||
"tempfile",
|
||||
"testcontainers-modules",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-redis-types"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-secret-value",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-router"
|
||||
version = "0.1.0"
|
||||
|
|
@ -5286,6 +5441,16 @@ version = "0.2.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4"
|
||||
|
||||
[[package]]
|
||||
name = "pbkdf2"
|
||||
version = "0.12.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"hmac 0.12.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pem"
|
||||
version = "4.0.0"
|
||||
|
|
@ -5396,6 +5561,22 @@ dependencies = [
|
|||
"spki",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkcs5"
|
||||
version = "0.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6"
|
||||
dependencies = [
|
||||
"aes",
|
||||
"cbc",
|
||||
"der",
|
||||
"des",
|
||||
"pbkdf2",
|
||||
"scrypt",
|
||||
"sha2 0.10.9",
|
||||
"spki",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkcs8"
|
||||
version = "0.10.2"
|
||||
|
|
@ -5403,6 +5584,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
|
||||
dependencies = [
|
||||
"der",
|
||||
"pkcs5",
|
||||
"rand_core 0.6.4",
|
||||
"spki",
|
||||
]
|
||||
|
||||
|
|
@ -6411,6 +6594,15 @@ version = "1.0.23"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
|
||||
|
||||
[[package]]
|
||||
name = "salsa20"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213"
|
||||
dependencies = [
|
||||
"cipher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "same-file"
|
||||
version = "1.0.6"
|
||||
|
|
@ -6474,6 +6666,17 @@ version = "1.2.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||
|
||||
[[package]]
|
||||
name = "scrypt"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f"
|
||||
dependencies = [
|
||||
"pbkdf2",
|
||||
"salsa20",
|
||||
"sha2 0.10.9",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sct"
|
||||
version = "0.7.1"
|
||||
|
|
|
|||
|
|
@ -73,7 +73,11 @@ litellm-host-python = { path = "crates/host-python" }
|
|||
litellm-python-compat = { path = "crates/python-compat" }
|
||||
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
|
||||
litellm-redis-client = { path = "crates/redis-client" }
|
||||
litellm-redis-connect = { path = "crates/redis-connect" }
|
||||
litellm-redis-store = { path = "crates/redis-store" }
|
||||
litellm-redis-testing = { path = "crates/redis-testing" }
|
||||
litellm-redis-resolve = { path = "crates/redis-resolve" }
|
||||
litellm-redis-types = { path = "crates/redis-types" }
|
||||
litellm-cost = { path = "crates/cost" }
|
||||
litellm-model-catalog = { path = "crates/model-catalog" }
|
||||
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }
|
||||
|
|
|
|||
|
|
@ -2,7 +2,10 @@ use std::{
|
|||
cmp::Reverse,
|
||||
collections::{BinaryHeap, HashMap, HashSet},
|
||||
hash::Hash,
|
||||
sync::{Arc, Mutex, MutexGuard},
|
||||
sync::{
|
||||
Arc, Mutex, MutexGuard,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
|
|
@ -31,7 +34,7 @@ struct CacheState<V> {
|
|||
|
||||
pub struct InMemoryCache<V: Clone> {
|
||||
state: Mutex<CacheState<V>>,
|
||||
max_size_in_memory: usize,
|
||||
max_size_in_memory: AtomicUsize,
|
||||
default_ttl: Duration,
|
||||
max_entry_bytes: Option<usize>,
|
||||
measure_value: Option<ValueMeasure<V>>,
|
||||
|
|
@ -74,7 +77,9 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
expirations: HashMap::new(),
|
||||
expiration_heap: BinaryHeap::new(),
|
||||
}),
|
||||
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
|
||||
max_size_in_memory: AtomicUsize::new(
|
||||
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
|
||||
),
|
||||
default_ttl: default_ttl
|
||||
.filter(|ttl| !ttl.is_zero())
|
||||
.unwrap_or(DEFAULT_TTL),
|
||||
|
|
@ -94,7 +99,7 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
value: V,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<CacheWrite, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
if self.max_size_in_memory() == 0 {
|
||||
return Ok(CacheWrite::Disabled);
|
||||
}
|
||||
let now = (self.now)();
|
||||
|
|
@ -129,7 +134,7 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
pub fn evict_cache(&self) -> Result<(), Error> {
|
||||
let now = (self.now)();
|
||||
let mut state = self.lock()?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now, None);
|
||||
Self::evict(&mut state, self.max_size_in_memory(), now, None);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -170,7 +175,16 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
}
|
||||
|
||||
pub fn max_size_in_memory(&self) -> usize {
|
||||
self.max_size_in_memory
|
||||
self.max_size_in_memory.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay
|
||||
/// until the next write evicts down to the new size.
|
||||
pub fn update_max_size_in_memory(&self, max_size_in_memory: Option<usize>) {
|
||||
self.max_size_in_memory.store(
|
||||
max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
|
||||
pub fn max_entry_bytes(&self) -> Option<usize> {
|
||||
|
|
@ -200,6 +214,23 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
.collect())
|
||||
}
|
||||
|
||||
/// Every live entry, for callers that inspect the whole tier.
|
||||
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
|
||||
let now = (self.now)();
|
||||
let state = self.lock()?;
|
||||
Ok(state
|
||||
.values
|
||||
.iter()
|
||||
.filter(|(key, _)| {
|
||||
state
|
||||
.expirations
|
||||
.get(*key)
|
||||
.is_none_or(|expiration| *expiration >= now)
|
||||
})
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let mut state = self.lock()?;
|
||||
Self::remove(&mut state, key);
|
||||
|
|
@ -276,7 +307,7 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
ttl: Option<Duration>,
|
||||
now: Duration,
|
||||
) -> Result<CacheWrite, Error> {
|
||||
Self::evict(state, self.max_size_in_memory, now, None);
|
||||
Self::evict(state, self.max_size_in_memory(), now, None);
|
||||
if !self.check_value_size(&value)? {
|
||||
return Ok(CacheWrite::TooLarge);
|
||||
}
|
||||
|
|
@ -300,12 +331,12 @@ where
|
|||
eligible: &[V],
|
||||
context: ExactCacheContext,
|
||||
) -> Result<V, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
if self.max_size_in_memory() == 0 {
|
||||
return Ok(candidate);
|
||||
}
|
||||
let now = (self.now)();
|
||||
let mut state = self.lock()?;
|
||||
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
|
||||
Self::evict(&mut state, self.max_size_in_memory(), now, Some(key));
|
||||
let existing = state
|
||||
.values
|
||||
.get(key)
|
||||
|
|
@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache<f64> {
|
|||
amount: f64,
|
||||
context: ExactCacheContext,
|
||||
) -> Result<f64, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
if self.max_size_in_memory() == 0 {
|
||||
return Ok(amount);
|
||||
}
|
||||
let now = (self.now)();
|
||||
|
|
@ -408,7 +439,7 @@ where
|
|||
values: Vec<Self::SetValue>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Self::SetResult, Error> {
|
||||
if self.max_size_in_memory == 0 {
|
||||
if self.max_size_in_memory() == 0 {
|
||||
return Ok(values);
|
||||
}
|
||||
let now = (self.now)();
|
||||
|
|
|
|||
3
litellm-rust/crates/config/AGENTS.md
Normal file
3
litellm-rust/crates/config/AGENTS.md
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
## Redis settings
|
||||
|
||||
- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve`
|
||||
|
|
@ -6,17 +6,16 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-secret-value.workspace = true
|
||||
litellm-router-types.workspace = true
|
||||
litellm-secret-value.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
serde_yaml_ng = "0.10.0"
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
tracing.workspace = true
|
||||
veil.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-error.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tempfile.workspace = true
|
||||
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
|
||||
|
|
|
|||
|
|
@ -11,11 +11,4 @@ pub enum Error {
|
|||
path: PathBuf,
|
||||
source: serde_yaml_ng::Error,
|
||||
},
|
||||
#[error("Redis host does not yet support configured setting {0}")]
|
||||
RedisSetting(String),
|
||||
#[error("invalid Redis setting {setting}: {reason}")]
|
||||
RedisInvalid {
|
||||
setting: &'static str,
|
||||
reason: &'static str,
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,10 +14,7 @@ pub use error::Error;
|
|||
pub use litellm_router_types::{LitellmParams, Spelled};
|
||||
pub use mcp::{McpAuth, McpServer, McpTransport};
|
||||
pub use model::Model;
|
||||
pub use redis::{
|
||||
CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings,
|
||||
ResolvedRedis, environment_has_target,
|
||||
};
|
||||
pub use redis::{CacheSettings, RedisSettings};
|
||||
pub use settings::{
|
||||
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
|
||||
TracingStoreSettings,
|
||||
|
|
|
|||
|
|
@ -1,195 +1,21 @@
|
|||
use std::{collections::BTreeMap, fmt, time::Duration};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_router_types::Spelled;
|
||||
use litellm_secret_value::SecretValue;
|
||||
use serde::{Deserialize, Deserializer, de::Error as _};
|
||||
use veil::Redact;
|
||||
|
||||
use crate::{AdditionalFields, Error, Object, Value};
|
||||
use crate::{Object, Value};
|
||||
|
||||
pub type Lookup<'a> = &'a dyn Fn(&str) -> Option<String>;
|
||||
|
||||
const SUPPORTED: &[&str] = &[
|
||||
"url",
|
||||
"host",
|
||||
"port",
|
||||
"username",
|
||||
"password",
|
||||
"ssl",
|
||||
"db",
|
||||
"socket_timeout",
|
||||
"startup_nodes",
|
||||
"namespace",
|
||||
];
|
||||
|
||||
const IGNORED: &[&str] = &[
|
||||
"max_connections",
|
||||
"single_connection_client",
|
||||
"health_check_interval",
|
||||
"socket_keepalive",
|
||||
"socket_keepalive_options",
|
||||
"decode_responses",
|
||||
"encoding",
|
||||
"encoding_errors",
|
||||
"charset",
|
||||
"errors",
|
||||
"client_name",
|
||||
"lib_name",
|
||||
"lib_version",
|
||||
"sentinel_password",
|
||||
"gcp_ssl_ca_certs",
|
||||
"azure_client_id",
|
||||
"azure_tenant_id",
|
||||
"azure_client_secret",
|
||||
"aws_iam_user_name",
|
||||
"aws_iam_cache_name",
|
||||
"aws_iam_region",
|
||||
"aws_iam_serverless",
|
||||
];
|
||||
|
||||
const WARN_IGNORED: &[&str] = &[
|
||||
"max_connections",
|
||||
"health_check_interval",
|
||||
"socket_keepalive",
|
||||
"socket_keepalive_options",
|
||||
];
|
||||
|
||||
const REDIS_PY_KWARGS: &[&str] = &[
|
||||
"aws_iam_auth",
|
||||
"aws_iam_cache_name",
|
||||
"aws_iam_region",
|
||||
"aws_iam_serverless",
|
||||
"aws_iam_user_name",
|
||||
"azure_client_id",
|
||||
"azure_client_secret",
|
||||
"azure_redis_ad_token",
|
||||
"azure_tenant_id",
|
||||
"cache",
|
||||
"cache_config",
|
||||
"charset",
|
||||
"client_name",
|
||||
"db",
|
||||
"decode_responses",
|
||||
"encoding",
|
||||
"encoding_errors",
|
||||
"errors",
|
||||
"event_dispatcher",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
"health_check_interval",
|
||||
"host",
|
||||
"lib_name",
|
||||
"lib_version",
|
||||
"max_connections",
|
||||
"password",
|
||||
"port",
|
||||
"protocol",
|
||||
"redis_connect_func",
|
||||
"retry_on_error",
|
||||
"retry_on_timeout",
|
||||
"single_connection_client",
|
||||
"socket_connect_timeout",
|
||||
"socket_keepalive",
|
||||
"socket_keepalive_options",
|
||||
"socket_timeout",
|
||||
"ssl",
|
||||
"ssl_ca_certs",
|
||||
"ssl_ca_data",
|
||||
"ssl_ca_path",
|
||||
"ssl_cert_reqs",
|
||||
"ssl_certfile",
|
||||
"ssl_check_hostname",
|
||||
"ssl_ciphers",
|
||||
"ssl_keyfile",
|
||||
"ssl_min_version",
|
||||
"ssl_ocsp_context",
|
||||
"ssl_ocsp_expected_cert",
|
||||
"ssl_password",
|
||||
"ssl_validate_ocsp",
|
||||
"ssl_validate_ocsp_stapled",
|
||||
"unix_socket_path",
|
||||
"url",
|
||||
"username",
|
||||
];
|
||||
|
||||
const UNSUPPORTED_EXTRA: &[&str] = &[
|
||||
"credential_provider",
|
||||
"sentinel_nodes",
|
||||
"retry",
|
||||
"connection_pool",
|
||||
];
|
||||
|
||||
const CACHE_ONLY: &[&str] = &[
|
||||
"type",
|
||||
"mode",
|
||||
"ttl",
|
||||
"default_in_memory_ttl",
|
||||
"default_in_redis_ttl",
|
||||
"supported_call_types",
|
||||
"similarity_threshold",
|
||||
"max_messages",
|
||||
"_backend",
|
||||
];
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
/// A Redis section as `config.yaml` spells it, before any surface reads it.
|
||||
#[derive(Clone, Default, Deserialize, Redact)]
|
||||
#[serde(transparent)]
|
||||
pub struct RedisSettings {
|
||||
pub url: Option<SecretValue>,
|
||||
pub host: Option<Spelled<String>>,
|
||||
pub port: Option<Spelled<u16>>,
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
pub ssl: Option<Spelled<bool>>,
|
||||
pub db: Option<Spelled<i64>>,
|
||||
pub socket_timeout: Option<Spelled<f64>>,
|
||||
pub startup_nodes: Option<Value>,
|
||||
pub namespace: Option<Spelled<String>>,
|
||||
pub additional_fields: AdditionalFields,
|
||||
unparsed: DeferredInputs,
|
||||
#[redact(with = "[REDACTED]")]
|
||||
fields: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for RedisSettings {
|
||||
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
let fields = BTreeMap::<String, Value>::deserialize(deserializer)?;
|
||||
let scalar = |key: &str| fields.get(key).and_then(as_text);
|
||||
Ok(Self {
|
||||
url: scalar("url").map(SecretValue::new),
|
||||
host: scalar("host").map(Spelled::Value),
|
||||
port: spelled(scalar("port")),
|
||||
username: scalar("username").map(SecretValue::new),
|
||||
password: scalar("password").map(SecretValue::new),
|
||||
ssl: spelled(scalar("ssl")),
|
||||
db: spelled(scalar("db")),
|
||||
socket_timeout: spelled(scalar("socket_timeout")),
|
||||
startup_nodes: fields
|
||||
.get("startup_nodes")
|
||||
.filter(|value| !value.is_null())
|
||||
.cloned(),
|
||||
namespace: scalar("namespace").map(Spelled::Value),
|
||||
additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)),
|
||||
unparsed: deferred_inputs(&fields),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for RedisSettings {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("RedisSettings")
|
||||
.field("url", &self.url)
|
||||
.field("host", &self.host)
|
||||
.field("port", &self.port)
|
||||
.field("username", &self.username)
|
||||
.field("password", &self.password)
|
||||
.field("ssl", &self.ssl)
|
||||
.field("db", &self.db)
|
||||
.field("socket_timeout", &self.socket_timeout)
|
||||
.field(
|
||||
"startup_nodes",
|
||||
&self.startup_nodes.as_ref().map(|_| "<configured>"),
|
||||
)
|
||||
.field("namespace", &self.namespace)
|
||||
.field("additional_fields", &self.additional_fields.keys())
|
||||
.field("unparsed", &self.unparsed.keys())
|
||||
.finish()
|
||||
impl RedisSettings {
|
||||
pub fn fields(&self) -> &BTreeMap<String, Value> {
|
||||
&self.fields
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct ResolvedRedis {
|
||||
pub endpoint: RedisEndpoint,
|
||||
pub namespace: Option<String>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisSettings;
|
||||
use rstest::rstest;
|
||||
|
||||
impl fmt::Debug for ResolvedRedis {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ResolvedRedis")
|
||||
.field("endpoint", &self.endpoint)
|
||||
.field("namespace", &self.namespace)
|
||||
.field("timeout", &self.timeout)
|
||||
.finish()
|
||||
#[rstest]
|
||||
#[case::scalar(serde_json::json!("credential-value"))]
|
||||
#[case::object(serde_json::json!({"secret": "credential-value"}))]
|
||||
#[case::array(serde_json::json!(["credential-value"]))]
|
||||
fn credentials_are_redacted(#[case] password: serde_json::Value) {
|
||||
let settings: RedisSettings = serde_json::from_value(serde_json::json!({
|
||||
"host": "redis", "password": password,
|
||||
}))
|
||||
.unwrap();
|
||||
let debug = format!("{settings:?}");
|
||||
assert!(!debug.contains("credential-value"));
|
||||
assert!(debug.contains("[REDACTED]"));
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum RedisEndpoint {
|
||||
Url(SecretValue),
|
||||
Host(RedisHost),
|
||||
Cluster(RedisCluster),
|
||||
}
|
||||
|
||||
impl fmt::Debug for RedisEndpoint {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::Url(_) => formatter.debug_tuple("Url").field(&"<redacted>").finish(),
|
||||
Self::Host(host) => host.fmt(formatter),
|
||||
Self::Cluster(cluster) => cluster.fmt(formatter),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct RedisHost {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
pub tls: bool,
|
||||
pub db: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct RedisCluster {
|
||||
pub nodes: Vec<RedisNode>,
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
pub tls: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct RedisNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct Entry {
|
||||
value: Value,
|
||||
source: String,
|
||||
}
|
||||
|
||||
type Kwargs = BTreeMap<String, Entry>;
|
||||
|
||||
impl RedisSettings {
|
||||
pub fn resolve(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
|
||||
self.resolve_with(lookup, false)
|
||||
}
|
||||
|
||||
pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
|
||||
self.resolve_with(lookup, true)
|
||||
}
|
||||
|
||||
fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result<ResolvedRedis, Error> {
|
||||
let block = self.kwargs(cache);
|
||||
let block: Kwargs = block
|
||||
.into_iter()
|
||||
.filter_map(|(key, entry)| {
|
||||
let value = environment_reference(&entry.value, lookup)?;
|
||||
Some((key, Entry { value, ..entry }))
|
||||
})
|
||||
.collect();
|
||||
let environment = environment_kwargs(lookup);
|
||||
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
|
||||
.iter()
|
||||
.any(|key| block.contains_key(*key));
|
||||
let drop_environment_url = names_a_target && !block.contains_key("url");
|
||||
let merged: Kwargs = environment
|
||||
.into_iter()
|
||||
.filter(|(key, _)| !(key == "url" && drop_environment_url))
|
||||
.chain(block)
|
||||
.collect();
|
||||
reject_unsupported(&merged, lookup)?;
|
||||
warn_ignored(&merged);
|
||||
|
||||
let startup_nodes = match merged.get("startup_nodes") {
|
||||
Some(entry) => Some(cluster_nodes(&entry.value)?),
|
||||
None => lookup("REDIS_CLUSTER_NODES")
|
||||
.map(|nodes| cluster_nodes(&Value::String(nodes)))
|
||||
.transpose()?,
|
||||
};
|
||||
let username = secret_value(&merged, "username");
|
||||
let password = secret_value(&merged, "password");
|
||||
let tls = bool_value(&merged, "ssl");
|
||||
let endpoint = match (startup_nodes, text_value(&merged, "url")) {
|
||||
(Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster {
|
||||
nodes,
|
||||
username,
|
||||
password,
|
||||
tls,
|
||||
}),
|
||||
(_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)),
|
||||
_ => RedisEndpoint::Host(RedisHost {
|
||||
host: text_value(&merged, "host").ok_or(Error::RedisInvalid {
|
||||
setting: "host",
|
||||
reason: "either host or url must be specified for redis",
|
||||
})?,
|
||||
port: integer(&merged, "port")?.unwrap_or(6379),
|
||||
username,
|
||||
password,
|
||||
tls,
|
||||
db: integer(&merged, "db")?.unwrap_or(0),
|
||||
}),
|
||||
};
|
||||
let timeout = merged
|
||||
.get("socket_timeout")
|
||||
.and_then(|entry| as_text(&entry.value))
|
||||
.map(|seconds| {
|
||||
seconds
|
||||
.parse::<f64>()
|
||||
.ok()
|
||||
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
|
||||
.ok_or(Error::RedisInvalid {
|
||||
setting: "socket_timeout",
|
||||
reason: "must be a non-negative number of seconds",
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
Ok(ResolvedRedis {
|
||||
endpoint,
|
||||
namespace: text_value(&merged, "namespace"),
|
||||
timeout,
|
||||
})
|
||||
}
|
||||
|
||||
fn kwargs(&self, cache: bool) -> Kwargs {
|
||||
let known = [
|
||||
(
|
||||
"url",
|
||||
self.url.as_ref().map(|value| value.expose().to_owned()),
|
||||
),
|
||||
("host", spelled_text(self.host.as_ref())),
|
||||
("port", spelled_text(self.port.as_ref())),
|
||||
(
|
||||
"username",
|
||||
self.username
|
||||
.as_ref()
|
||||
.map(|value| value.expose().to_owned()),
|
||||
),
|
||||
(
|
||||
"password",
|
||||
self.password
|
||||
.as_ref()
|
||||
.map(|value| value.expose().to_owned()),
|
||||
),
|
||||
("ssl", spelled_text(self.ssl.as_ref())),
|
||||
("db", spelled_text(self.db.as_ref())),
|
||||
("socket_timeout", spelled_text(self.socket_timeout.as_ref())),
|
||||
("namespace", spelled_text(self.namespace.as_ref())),
|
||||
];
|
||||
let known = known.into_iter().filter_map(|(key, value)| {
|
||||
value.map(|value| {
|
||||
(
|
||||
key.to_owned(),
|
||||
Entry {
|
||||
value: Value::String(value),
|
||||
source: key.to_owned(),
|
||||
},
|
||||
)
|
||||
})
|
||||
});
|
||||
let startup_nodes = self.startup_nodes.clone().map(|value| {
|
||||
(
|
||||
"startup_nodes".to_owned(),
|
||||
Entry {
|
||||
value,
|
||||
source: "startup_nodes".to_owned(),
|
||||
},
|
||||
)
|
||||
});
|
||||
let extras = self.additional_fields.iter().filter_map(|(key, value)| {
|
||||
if cache && CACHE_ONLY.contains(&key.as_str()) {
|
||||
return None;
|
||||
}
|
||||
let key = match (cache, key.as_str()) {
|
||||
(true, "redis_startup_nodes") => "startup_nodes".to_owned(),
|
||||
_ => key.clone(),
|
||||
};
|
||||
Some((
|
||||
key.clone(),
|
||||
Entry {
|
||||
value: value.clone(),
|
||||
source: key,
|
||||
},
|
||||
))
|
||||
});
|
||||
let unparsed = self.unparsed.iter().map(|(key, value)| {
|
||||
(
|
||||
key.clone(),
|
||||
Entry {
|
||||
value: value.0.clone(),
|
||||
source: key.clone(),
|
||||
},
|
||||
)
|
||||
});
|
||||
extras
|
||||
.chain(unparsed)
|
||||
.chain(known)
|
||||
.chain(startup_nodes)
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn environment_has_target(lookup: Lookup<'_>) -> bool {
|
||||
[
|
||||
"REDIS_HOST",
|
||||
"REDIS_URL",
|
||||
"REDIS_CLUSTER_NODES",
|
||||
"REDIS_SENTINEL_NODES",
|
||||
]
|
||||
.iter()
|
||||
.any(|name| lookup(name).is_some())
|
||||
}
|
||||
|
||||
fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> {
|
||||
let active_auth = [
|
||||
merged
|
||||
.get("gcp_service_account")
|
||||
.map(|entry| entry.source.as_str())
|
||||
.or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")),
|
||||
text_value(merged, "azure_redis_ad_token")
|
||||
.or_else(|| lookup("REDIS_AZURE_AD_TOKEN"))
|
||||
.filter(|value| value.eq_ignore_ascii_case("true"))
|
||||
.map(|_| {
|
||||
merged
|
||||
.get("azure_redis_ad_token")
|
||||
.map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str())
|
||||
}),
|
||||
bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()),
|
||||
merged
|
||||
.get("sentinel_nodes")
|
||||
.map(|entry| entry.source.as_str())
|
||||
.or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")),
|
||||
];
|
||||
if let Some(setting) = active_auth.into_iter().flatten().next() {
|
||||
return Err(Error::RedisSetting(setting.to_owned()));
|
||||
}
|
||||
match merged
|
||||
.iter()
|
||||
.find(|(key, _)| is_unsupported(key))
|
||||
.map(|(_, entry)| entry.source.clone())
|
||||
{
|
||||
Some(setting) => Err(Error::RedisSetting(setting)),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
fn warn_ignored(merged: &Kwargs) {
|
||||
for setting in WARN_IGNORED {
|
||||
if merged.contains_key(*setting) {
|
||||
tracing::warn!(setting, "Redis setting is accepted but ignored");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_unsupported(key: &str) -> bool {
|
||||
let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key);
|
||||
let handled = SUPPORTED.contains(&key)
|
||||
|| IGNORED.contains(&key)
|
||||
|| matches!(
|
||||
key,
|
||||
"gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth"
|
||||
);
|
||||
known && !handled
|
||||
}
|
||||
|
||||
fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs {
|
||||
REDIS_PY_KWARGS
|
||||
.iter()
|
||||
.filter_map(|key| {
|
||||
let source = format!("REDIS_{}", key.to_ascii_uppercase());
|
||||
let value = lookup(&source)?;
|
||||
Some((
|
||||
(*key).to_owned(),
|
||||
Entry {
|
||||
value: Value::String(value),
|
||||
source,
|
||||
},
|
||||
))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option<Value> {
|
||||
match value {
|
||||
Value::String(text) => match text.strip_prefix("os.environ/") {
|
||||
Some(name) => lookup(name).map(Value::String),
|
||||
None => Some(value.clone()),
|
||||
},
|
||||
Value::Null => None,
|
||||
_ => Some(value.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
fn as_text(value: &Value) -> Option<String> {
|
||||
match value {
|
||||
Value::String(text) => Some(text.clone()),
|
||||
Value::Number(number) => Some(number.to_string()),
|
||||
Value::Bool(flag) => Some(flag.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn text_value(kwargs: &Kwargs, key: &str) -> Option<String> {
|
||||
kwargs.get(key).and_then(|entry| as_text(&entry.value))
|
||||
}
|
||||
|
||||
fn secret_value(kwargs: &Kwargs, key: &str) -> Option<SecretValue> {
|
||||
text_value(kwargs, key).map(SecretValue::new)
|
||||
}
|
||||
|
||||
fn bool_value(kwargs: &Kwargs, key: &str) -> bool {
|
||||
match kwargs.get(key).map(|entry| &entry.value) {
|
||||
Some(Value::Bool(flag)) => *flag,
|
||||
Some(value) => as_text(value)
|
||||
.is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")),
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn integer<T: std::str::FromStr>(kwargs: &Kwargs, key: &'static str) -> Result<Option<T>, Error> {
|
||||
kwargs
|
||||
.get(key)
|
||||
.and_then(|entry| as_text(&entry.value))
|
||||
.map(|text| {
|
||||
text.parse().map_err(|_| Error::RedisInvalid {
|
||||
setting: key,
|
||||
reason: "must be an integer",
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
const INVALID_NODES: Error = Error::RedisInvalid {
|
||||
setting: "startup_nodes",
|
||||
reason: "must be a list of {host, port} objects",
|
||||
};
|
||||
|
||||
fn cluster_nodes(value: &Value) -> Result<Vec<RedisNode>, Error> {
|
||||
match value {
|
||||
Value::String(json) => {
|
||||
let nodes: Vec<serde_json::Map<String, serde_json::Value>> =
|
||||
serde_json::from_str(json).map_err(|_| INVALID_NODES)?;
|
||||
nodes
|
||||
.iter()
|
||||
.map(|node| {
|
||||
cluster_node(
|
||||
node.get("host").and_then(json_text),
|
||||
node.get("port").and_then(json_text),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
Value::Sequence(nodes) => nodes
|
||||
.iter()
|
||||
.map(|node| {
|
||||
cluster_node(
|
||||
node.get("host").and_then(as_text),
|
||||
node.get("port").and_then(as_text),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
_ => Err(INVALID_NODES),
|
||||
}
|
||||
}
|
||||
|
||||
fn cluster_node(host: Option<String>, port: Option<String>) -> Result<RedisNode, Error> {
|
||||
match (host, port.and_then(|port| port.parse().ok())) {
|
||||
(Some(host), Some(port)) => Ok(RedisNode { host, port }),
|
||||
_ => Err(INVALID_NODES),
|
||||
}
|
||||
}
|
||||
|
||||
fn json_text(value: &serde_json::Value) -> Option<String> {
|
||||
match value {
|
||||
serde_json::Value::String(text) => Some(text.clone()),
|
||||
serde_json::Value::Number(number) => Some(number.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn spelled_text<T: ToString>(value: Option<&Spelled<T>>) -> Option<String> {
|
||||
value.map(|value| match value {
|
||||
Spelled::Value(value) => value.to_string(),
|
||||
Spelled::Text(value) => value.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn spelled<T: std::str::FromStr>(text: Option<String>) -> Option<Spelled<T>> {
|
||||
text.map(|text| match text.parse() {
|
||||
Ok(value) => Spelled::Value(value),
|
||||
Err(_) => Spelled::Text(text),
|
||||
})
|
||||
}
|
||||
|
||||
fn cloned(
|
||||
fields: &BTreeMap<String, Value>,
|
||||
keep: impl Fn(&str, &Value) -> bool,
|
||||
) -> AdditionalFields {
|
||||
fields
|
||||
.iter()
|
||||
.filter(|(key, value)| keep(key, value))
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct DeferredInput(Value);
|
||||
|
||||
type DeferredInputs = BTreeMap<String, DeferredInput>;
|
||||
|
||||
fn deferred_inputs(fields: &BTreeMap<String, Value>) -> DeferredInputs {
|
||||
fields
|
||||
.iter()
|
||||
.filter(|(key, value)| is_unparsed(key, value))
|
||||
.map(|(key, value)| (key.clone(), DeferredInput(value.clone())))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn is_unparsed(key: &str, value: &Value) -> bool {
|
||||
SUPPORTED.contains(&key)
|
||||
&& match key {
|
||||
"startup_nodes" => value.is_null(),
|
||||
_ => as_text(value).is_none(),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() {
|
|||
fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) {
|
||||
assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
|
||||
let config = Config::from_yaml(
|
||||
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
config.litellm_settings.cache_params,
|
||||
Some(litellm_config::CacheSettings::Redis(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn non_redis_cache_settings_remain_opaque() {
|
||||
let config = Config::from_yaml(
|
||||
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
config.litellm_settings.cache_params,
|
||||
Some(litellm_config::CacheSettings::Other(_))
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,405 +0,0 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
fmt::Write,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis};
|
||||
use rstest::rstest;
|
||||
use tracing::{Event, Subscriber, field::Visit};
|
||||
use tracing_subscriber::{Layer, layer::Context, prelude::*};
|
||||
|
||||
fn resolve(
|
||||
yaml: &str,
|
||||
environment: &[(&str, &str)],
|
||||
) -> Result<litellm_config::ResolvedRedis, Error> {
|
||||
let environment: HashMap<String, String> = environment
|
||||
.iter()
|
||||
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
|
||||
.collect();
|
||||
let config = Config::from_yaml(yaml).unwrap();
|
||||
config
|
||||
.general_settings
|
||||
.coordination_redis
|
||||
.as_ref()
|
||||
.unwrap_or(&RedisSettings::default())
|
||||
.resolve(&|name| environment.get(name).cloned())
|
||||
}
|
||||
|
||||
fn resolve_cache(yaml: &str) -> Result<ResolvedRedis, Error> {
|
||||
let config = Config::from_yaml(yaml).unwrap();
|
||||
match config.litellm_settings.cache_params.unwrap() {
|
||||
CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None),
|
||||
CacheSettings::Other(_) => panic!("expected a Redis cache"),
|
||||
}
|
||||
}
|
||||
|
||||
fn cluster_ports(resolved: &ResolvedRedis) -> Vec<u16> {
|
||||
match &resolved.endpoint {
|
||||
RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
|
||||
endpoint => panic!("expected a cluster, got {endpoint:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::host(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
|
||||
&[],
|
||||
"host:redis:6380:2:true:app:1500"
|
||||
)]
|
||||
#[case::environment_url(
|
||||
"model_list: []",
|
||||
&[("REDIS_URL", "redis://environment:6379")],
|
||||
"url:redis://environment:6379"
|
||||
)]
|
||||
#[case::explicit_host_drops_environment_url(
|
||||
"general_settings:\n coordination_redis:\n host: configured",
|
||||
&[("REDIS_URL", "redis://environment:6379")],
|
||||
"host:configured:6379:0:false::none"
|
||||
)]
|
||||
#[case::yaml_cluster(
|
||||
"general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"",
|
||||
&[],
|
||||
"cluster:node-a:7000,node-b:7001"
|
||||
)]
|
||||
#[case::json_environment_cluster(
|
||||
"model_list: []",
|
||||
&[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)],
|
||||
"cluster:node-a:7000,node-b:7001"
|
||||
)]
|
||||
#[case::missing_reference_keeps_environment_fallback(
|
||||
"general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST",
|
||||
&[("REDIS_HOST", "environment")],
|
||||
"host:environment:6379:0:false::none"
|
||||
)]
|
||||
#[case::configured_values_beat_the_same_environment_values(
|
||||
"general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3",
|
||||
&[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")],
|
||||
"host:configured:6380:3:false::none"
|
||||
)]
|
||||
#[case::environment_url_beats_the_environment_host(
|
||||
"model_list: []",
|
||||
&[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")],
|
||||
"url:redis://environment:6379"
|
||||
)]
|
||||
#[case::startup_nodes_beat_a_url(
|
||||
"general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000",
|
||||
&[],
|
||||
"cluster:node-a:7000"
|
||||
)]
|
||||
#[case::environment_references_resolve_the_host_and_port(
|
||||
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
|
||||
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
|
||||
"host:referenced:6400:0:false::none"
|
||||
)]
|
||||
#[case::a_null_host_falls_back_to_the_environment(
|
||||
"general_settings:\n coordination_redis:\n host: null",
|
||||
&[("REDIS_HOST", "environment")],
|
||||
"host:environment:6379:0:false::none"
|
||||
)]
|
||||
fn resolves_redis_through_the_parsed_config(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let resolved = resolve(yaml, environment).unwrap();
|
||||
let description = match resolved.endpoint {
|
||||
RedisEndpoint::Url(url) => format!("url:{}", url.expose()),
|
||||
RedisEndpoint::Host(host) => format!(
|
||||
"host:{}:{}:{}:{}:{}:{}",
|
||||
host.host,
|
||||
host.port,
|
||||
host.db,
|
||||
host.tls,
|
||||
resolved.namespace.unwrap_or_default(),
|
||||
resolved
|
||||
.timeout
|
||||
.map(|timeout| timeout.as_millis().to_string())
|
||||
.unwrap_or_else(|| "none".to_owned())
|
||||
),
|
||||
RedisEndpoint::Cluster(cluster) => format!(
|
||||
"cluster:{}",
|
||||
cluster
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|node| format!("{}:{}", node.host, node.port))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
),
|
||||
};
|
||||
|
||||
assert_eq!(description, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::custom_tls_file(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem",
|
||||
&[],
|
||||
"ssl_ca_certs"
|
||||
)]
|
||||
#[case::gcp_auth_from_config(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account",
|
||||
&[],
|
||||
"gcp_service_account"
|
||||
)]
|
||||
#[case::gcp_auth_from_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")],
|
||||
"REDIS_GCP_SERVICE_ACCOUNT"
|
||||
)]
|
||||
#[case::azure_auth_from_config(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true",
|
||||
&[],
|
||||
"azure_redis_ad_token"
|
||||
)]
|
||||
#[case::azure_auth_from_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")],
|
||||
"REDIS_AZURE_AD_TOKEN"
|
||||
)]
|
||||
#[case::aws_auth_from_config(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true",
|
||||
&[],
|
||||
"aws_iam_auth"
|
||||
)]
|
||||
#[case::aws_auth_from_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")],
|
||||
"REDIS_AWS_IAM_AUTH"
|
||||
)]
|
||||
#[case::sentinel_from_config(
|
||||
"general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]",
|
||||
&[],
|
||||
"sentinel_nodes"
|
||||
)]
|
||||
#[case::sentinel_from_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")],
|
||||
"REDIS_SENTINEL_NODES"
|
||||
)]
|
||||
fn rejected_settings_name_their_source(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(matches!(error, Error::RedisSetting(setting) if setting == expected));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_explicit_host_keeps_environment_credentials() {
|
||||
let resolved = resolve(
|
||||
"general_settings:\n coordination_redis:\n host: configured",
|
||||
&[
|
||||
("REDIS_PASSWORD", "environment-secret"),
|
||||
("REDIS_SSL", "True"),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let RedisEndpoint::Host(host) = resolved.endpoint else {
|
||||
panic!("expected a host endpoint");
|
||||
};
|
||||
assert_eq!(
|
||||
host.password.as_ref().map(|password| password.expose()),
|
||||
Some("environment-secret")
|
||||
);
|
||||
assert!(host.tls);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
|
||||
#[case::a_structured_host_with_an_environment_target(
|
||||
"general_settings:\n coordination_redis:\n host: []",
|
||||
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
|
||||
)]
|
||||
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::RedisInvalid {
|
||||
setting: "host",
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_cache_alias_names_cluster_nodes() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cluster_ports(&resolved), vec![7000]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn canonical_startup_nodes_beat_the_cache_alias() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cluster_ports(&resolved), vec![7000]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rejected_environment_settings_never_print_the_value() {
|
||||
let error = resolve(
|
||||
"model_list: []",
|
||||
&[
|
||||
("REDIS_HOST", "redis"),
|
||||
("REDIS_SSL_CERT_REQS", "private-option-value"),
|
||||
],
|
||||
)
|
||||
.unwrap_err();
|
||||
let message = error.to_string();
|
||||
|
||||
assert!(message.contains("REDIS_SSL_CERT_REQS"));
|
||||
assert!(!message.contains("private-option-value"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::negative("-1")]
|
||||
#[case::not_a_number("slow")]
|
||||
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
|
||||
let yaml = format!(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
|
||||
);
|
||||
let error = resolve(&yaml, &[]).unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::RedisInvalid {
|
||||
setting: "socket_timeout",
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn defaults_socket_timeout_for_the_gateway() {
|
||||
let resolved = resolve(
|
||||
"general_settings:\n coordination_redis:\n host: redis",
|
||||
&[],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resolved.timeout, None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn debug_redacts_redis_secrets_and_url() {
|
||||
let config = Config::from_yaml(
|
||||
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password",
|
||||
)
|
||||
.unwrap();
|
||||
let settings = config.general_settings.coordination_redis.unwrap();
|
||||
let resolved = settings.resolve(&|_| None).unwrap();
|
||||
let debug = format!("{settings:?} {resolved:?}");
|
||||
|
||||
assert!(!debug.contains("redis://user:password"));
|
||||
assert!(!debug.contains("credential-user"));
|
||||
assert!(!debug.contains("credential-password"));
|
||||
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
|
||||
let config = Config::from_yaml(
|
||||
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
config.litellm_settings.cache_params,
|
||||
Some(litellm_config::CacheSettings::Redis(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn non_redis_cache_settings_remain_opaque() {
|
||||
let config = Config::from_yaml(
|
||||
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
config.litellm_settings.cache_params,
|
||||
Some(litellm_config::CacheSettings::Other(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn legacy_unknown_settings_do_not_change_resolution() {
|
||||
let resolved = resolve(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
|
||||
&[],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
|
||||
}
|
||||
|
||||
struct WarningCapture(Arc<Mutex<Vec<String>>>);
|
||||
|
||||
impl<S: Subscriber> Layer<S> for WarningCapture {
|
||||
fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
|
||||
let mut fields = Fields::default();
|
||||
event.record(&mut fields);
|
||||
self.0.lock().unwrap().push(fields.0);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Fields(String);
|
||||
|
||||
impl Visit for Fields {
|
||||
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
|
||||
write!(self.0, "{}={value:?} ", field.name()).unwrap();
|
||||
}
|
||||
|
||||
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
|
||||
write!(self.0, "{}={value} ", field.name()).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn ignored_setting_warning_names_the_setting_without_its_value() {
|
||||
let events = Arc::new(Mutex::new(Vec::new()));
|
||||
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
resolve(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
|
||||
&[],
|
||||
)
|
||||
.unwrap();
|
||||
});
|
||||
let warnings = events.lock().unwrap().join(" ");
|
||||
|
||||
assert!(warnings.contains("max_connections"));
|
||||
assert!(!warnings.contains("47"));
|
||||
}
|
||||
|
|
@ -5,8 +5,11 @@
|
|||
- Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate
|
||||
|
||||
## Redis
|
||||
- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
|
||||
- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
|
||||
- Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands
|
||||
- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped
|
||||
- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME`
|
||||
- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag
|
||||
- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it
|
||||
- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md`
|
||||
- Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python
|
||||
- `/health/readiness` never touches Redis, matching the proxy
|
||||
|
|
|
|||
|
|
@ -20,9 +20,14 @@ futures-util.workspace = true
|
|||
litellm-gateway-auth.workspace = true
|
||||
litellm-gateway-ui.workspace = true
|
||||
litellm-config.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
litellm-error.workspace = true
|
||||
litellm-circuit-breaker.workspace = true
|
||||
litellm-redis-client.workspace = true
|
||||
litellm-redis-connect.workspace = true
|
||||
litellm-redis-resolve.workspace = true
|
||||
litellm-redis-types.workspace = true
|
||||
rustls.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-secrets.workspace = true
|
||||
|
|
@ -32,10 +37,12 @@ serde.workspace = true
|
|||
tracing.workspace = true
|
||||
tokio = { workspace = true, features = ["signal"] }
|
||||
uuid.workspace = true
|
||||
veil.workspace = true
|
||||
tower-sessions-moka-store = "0.15.0"
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-redis-testing.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
base64.workspace = true
|
||||
rstest.workspace = true
|
||||
tempfile.workspace = true
|
||||
|
|
|
|||
|
|
@ -12,13 +12,12 @@ pub enum Error {
|
|||
Mcp(#[from] litellm_gateway_mcp::ConnectError),
|
||||
#[error(transparent)]
|
||||
Config(#[from] litellm_config::Error),
|
||||
#[error("invalid Redis setting {setting}: {reason}")]
|
||||
RedisInvalid {
|
||||
setting: &'static str,
|
||||
reason: &'static str,
|
||||
},
|
||||
#[error(transparent)]
|
||||
RedisConnect(#[from] litellm_redis_client::ConnectError),
|
||||
Redis(#[from] litellm_redis_resolve::Error),
|
||||
#[error("Redis setting {setting} is not plain data")]
|
||||
RedisNotPlainData { setting: String },
|
||||
#[error(transparent)]
|
||||
RedisConnect(#[from] litellm_redis_connect::Error),
|
||||
#[error("MCP requires working gateway authentication")]
|
||||
Auth(#[from] litellm_gateway_auth::Error),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
|
|||
Ok(Arc::new(Gateway::new(resources, http, secrets, router)?))
|
||||
}
|
||||
|
||||
pub async fn build_redis(config: &Config) -> Result<RedisClients, Error> {
|
||||
let environment = secrets::environment_values(&config.environment_variables)?;
|
||||
let lookup = |name: &str| {
|
||||
pub fn install_crypto_provider() {
|
||||
if rustls::crypto::CryptoProvider::get_default().is_none() {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn build_redis(
|
||||
config: &Config,
|
||||
pool: &HttpClientPool,
|
||||
http: &litellm_http::HttpClientConfig,
|
||||
) -> Result<RedisClients, Error> {
|
||||
let environment = Arc::new(secrets::environment_values(&config.environment_variables)?);
|
||||
let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| {
|
||||
environment
|
||||
.get(name)
|
||||
.map(|value| value.expose().to_owned())
|
||||
.or_else(|| std::env::var(name).ok())
|
||||
});
|
||||
let resources = redis::Resources {
|
||||
http: pool.client(http, ClientVariant::Provider)?,
|
||||
lookup: lookup.clone(),
|
||||
};
|
||||
redis::connect(config, &lookup).await
|
||||
redis::connect(config, &*lookup, &resources).await
|
||||
}
|
||||
|
||||
pub async fn build_mcp(
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ impl Sink for StderrSink {
|
|||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn Error>> {
|
||||
litellm_gateway::install_crypto_provider();
|
||||
if let Some(forward) = litellm_gateway_forward::Settings::from_env()? {
|
||||
let log: LogSettings = settings::from_iter(std::env::vars_os())?;
|
||||
Logger::new(StderrSink {
|
||||
|
|
@ -55,7 +56,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
Logger::new(StderrSink { level }).install_global()?;
|
||||
let config = Config::load(settings.litellm_config)?;
|
||||
let inference = litellm_gateway::build_inference(&config)?;
|
||||
let _redis = litellm_gateway::build_redis(&config).await?;
|
||||
let _redis =
|
||||
litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?;
|
||||
let ui = match std::env::var_os("LITELLM_UI_PATH") {
|
||||
Some(directory) => {
|
||||
let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?;
|
||||
|
|
|
|||
144
litellm-rust/crates/gateway/src/redis/configured.rs
Normal file
144
litellm-rust/crates/gateway/src/redis/configured.rs
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_config::{RedisSettings, Value};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_redis_resolve::{Expected, Kwargs};
|
||||
use litellm_redis_types::{ProviderHandle, Resolution};
|
||||
use veil::Redact;
|
||||
|
||||
use super::{layers, secrets::SecretSnapshot};
|
||||
use crate::Error;
|
||||
|
||||
const CACHE_ONLY: &[&str] = &[
|
||||
"type",
|
||||
"mode",
|
||||
"ttl",
|
||||
"default_in_memory_ttl",
|
||||
"default_in_redis_ttl",
|
||||
"supported_call_types",
|
||||
"similarity_threshold",
|
||||
"max_messages",
|
||||
"_backend",
|
||||
];
|
||||
|
||||
/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets.
|
||||
#[derive(Clone, Redact)]
|
||||
pub(crate) struct ConfiguredRedis {
|
||||
#[redact(with = "[REDACTED]")]
|
||||
pub(super) fields: BTreeMap<String, Value>,
|
||||
pub(super) nulls: Nulls,
|
||||
pub(super) credential_provider: Option<ProviderHandle>,
|
||||
}
|
||||
|
||||
/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub(super) enum Nulls {
|
||||
/// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`.
|
||||
Dropped,
|
||||
/// `Cache` drops a null only for the parameters it names and passes any other key on.
|
||||
KeptUnlessNamedByCache,
|
||||
/// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives.
|
||||
#[cfg(test)]
|
||||
Kept,
|
||||
}
|
||||
|
||||
impl ConfiguredRedis {
|
||||
/// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first.
|
||||
/// An `os.environ/` port is only known after merging, where it is read like any other port.
|
||||
pub(crate) fn coordination(
|
||||
settings: &RedisSettings,
|
||||
) -> Result<Self, litellm_redis_resolve::Error> {
|
||||
if let Some(port) = settings.fields().get("port").and_then(text)
|
||||
&& !port.starts_with("os.environ/")
|
||||
&& port.parse::<u16>().is_err()
|
||||
{
|
||||
return Err(litellm_redis_resolve::Error::Invalid {
|
||||
setting: "port".to_owned(),
|
||||
expected: Expected::Integer,
|
||||
});
|
||||
}
|
||||
Ok(Self {
|
||||
fields: settings.fields().clone(),
|
||||
nulls: Nulls::Dropped,
|
||||
credential_provider: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`.
|
||||
pub(crate) fn cache_params(settings: &RedisSettings) -> Self {
|
||||
let fields = settings.fields();
|
||||
let alias = fields
|
||||
.get("redis_startup_nodes")
|
||||
.filter(|_| !fields.contains_key("startup_nodes"))
|
||||
.map(|nodes| ("startup_nodes".to_owned(), nodes.clone()));
|
||||
let fields = fields
|
||||
.iter()
|
||||
.filter(|(key, _)| {
|
||||
key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str())
|
||||
})
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.chain(alias)
|
||||
.collect();
|
||||
Self {
|
||||
fields,
|
||||
nulls: Nulls::KeptUnlessNamedByCache,
|
||||
credential_provider: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `RedisCache.redis_kwargs`, the dict every Python surface ends in.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn redis_cache_kwargs(fields: BTreeMap<String, Value>) -> Self {
|
||||
Self {
|
||||
fields,
|
||||
nulls: Nulls::Kept,
|
||||
credential_provider: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The handle the host holds for the `credential_provider` object in these settings.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self {
|
||||
Self {
|
||||
credential_provider: Some(handle),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
/// No configured settings, so only `REDIS_*` and the other environment fallbacks apply.
|
||||
pub(crate) fn environment_only() -> Self {
|
||||
Self {
|
||||
fields: BTreeMap::new(),
|
||||
nulls: Nulls::Dropped,
|
||||
credential_provider: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every
|
||||
/// `get_secret` fallback filled.
|
||||
pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result<Kwargs, Error> {
|
||||
layers::merged(self, secrets)
|
||||
}
|
||||
|
||||
/// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure.
|
||||
pub(crate) fn resolve(
|
||||
&self,
|
||||
secrets: &SecretSnapshot,
|
||||
environment: &dyn Lookup,
|
||||
) -> Result<Resolution, Error> {
|
||||
Ok(litellm_redis_resolve::interpret(
|
||||
&self.merged(secrets)?,
|
||||
self.credential_provider,
|
||||
environment.get("REDIS_USERNAME").as_deref(),
|
||||
)?)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn text(value: &Value) -> Option<String> {
|
||||
match value {
|
||||
Value::String(text) => Some(text.clone()),
|
||||
Value::Number(number) => Some(number.to_string()),
|
||||
Value::Bool(flag) => Some(flag.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
226
litellm-rust/crates/gateway/src/redis/layers.rs
Normal file
226
litellm-rust/crates/gateway/src/redis/layers.rs
Normal file
|
|
@ -0,0 +1,226 @@
|
|||
use litellm_config::Value;
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting};
|
||||
|
||||
use super::configured::{ConfiguredRedis, Nulls};
|
||||
use crate::Error;
|
||||
|
||||
const NAMED_BY_CACHE: &[&str] = &[
|
||||
"host",
|
||||
"port",
|
||||
"password",
|
||||
"socket_timeout",
|
||||
"startup_nodes",
|
||||
"namespace",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
];
|
||||
|
||||
/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and
|
||||
/// `get_secret_str` treats that bool as unset.
|
||||
#[derive(Clone, Copy)]
|
||||
enum Read {
|
||||
Secret,
|
||||
SecretStr,
|
||||
}
|
||||
|
||||
const TARGETS: &[(&str, Read)] = &[
|
||||
("startup_nodes", Read::Secret),
|
||||
("sentinel_nodes", Read::Secret),
|
||||
("sentinel_password", Read::SecretStr),
|
||||
("service_name", Read::Secret),
|
||||
];
|
||||
|
||||
const IDENTITIES: &[(&str, Read)] = &[
|
||||
("gcp_service_account", Read::SecretStr),
|
||||
("gcp_ssl_ca_certs", Read::SecretStr),
|
||||
("azure_redis_ad_token", Read::Secret),
|
||||
];
|
||||
|
||||
const AZURE: &[(&str, Read)] = &[
|
||||
("azure_client_id", Read::SecretStr),
|
||||
("azure_tenant_id", Read::SecretStr),
|
||||
("azure_client_secret", Read::SecretStr),
|
||||
];
|
||||
|
||||
const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)];
|
||||
|
||||
/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}`
|
||||
/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is
|
||||
/// an explicit `None` that hid its environment value.
|
||||
pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Kwargs, Error> {
|
||||
let block: Kwargs = configured
|
||||
.fields
|
||||
.iter()
|
||||
.map(|(key, value)| (key, environment_reference(value, lookup)))
|
||||
.filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key))
|
||||
.map(|(key, value)| {
|
||||
let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData {
|
||||
setting: key.clone(),
|
||||
})?;
|
||||
let entry = Entry {
|
||||
value,
|
||||
source: key.clone(),
|
||||
};
|
||||
Ok((key.clone(), entry))
|
||||
})
|
||||
.collect::<Result<_, Error>>()?;
|
||||
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
|
||||
.iter()
|
||||
.any(|key| is_set(&block, key));
|
||||
let drop_environment_url = names_a_target && !is_set(&block, "url");
|
||||
let layered: Kwargs = environment(lookup)
|
||||
.filter(|(key, _)| !(key == "url" && drop_environment_url))
|
||||
.chain(block)
|
||||
.collect();
|
||||
let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup);
|
||||
Ok(layered.into_iter().chain(fallbacks).collect())
|
||||
}
|
||||
|
||||
// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks
|
||||
// for is never filled in
|
||||
fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs {
|
||||
let fill = |kwargs: &[(&str, Read)]| -> Kwargs {
|
||||
kwargs
|
||||
.iter()
|
||||
.filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup))
|
||||
.collect()
|
||||
};
|
||||
let targets = fill(TARGETS);
|
||||
if has_provider || is_set(layered, "credential_provider") {
|
||||
return targets;
|
||||
}
|
||||
let identities = fill(IDENTITIES);
|
||||
let value = |kwarg: &str| {
|
||||
layered
|
||||
.get(kwarg)
|
||||
.filter(|entry| is_present(&entry.value))
|
||||
.or_else(|| identities.get(kwarg))
|
||||
.map(|entry| &entry.value)
|
||||
};
|
||||
let gcp = value("gcp_service_account").is_some();
|
||||
let azure = value("azure_redis_ad_token")
|
||||
.is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true")));
|
||||
let aws = layered
|
||||
.get("aws_iam_auth")
|
||||
.is_some_and(|entry| coerces_to_true(&entry.value));
|
||||
let azure_settings = match azure && !gcp {
|
||||
true => fill(AZURE),
|
||||
false => Kwargs::new(),
|
||||
};
|
||||
let aws_settings = match aws && !gcp && !azure {
|
||||
true => fill(AWS),
|
||||
false => Kwargs::new(),
|
||||
};
|
||||
targets
|
||||
.into_iter()
|
||||
.chain(identities)
|
||||
.chain(azure_settings)
|
||||
.chain(aws_settings)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn fallback(
|
||||
layered: &Kwargs,
|
||||
kwarg: &str,
|
||||
read: Read,
|
||||
lookup: &dyn Lookup,
|
||||
) -> Option<(String, Entry)> {
|
||||
if layered
|
||||
.get(kwarg)
|
||||
.is_some_and(|entry| is_present(&entry.value))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback);
|
||||
names.iter().find_map(|name| {
|
||||
let value = lookup
|
||||
.get(name)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(secret)?;
|
||||
match (read, &value) {
|
||||
(Read::SecretStr, serde_json::Value::Bool(_)) => None,
|
||||
_ => {
|
||||
let entry = Entry {
|
||||
value,
|
||||
source: (*name).to_owned(),
|
||||
};
|
||||
Some((kwarg.to_owned(), entry))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// get_secret reads an environment value of true or false, in any case and padding, as a bool
|
||||
fn secret_flag(value: &str) -> Option<bool> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"true" => Some(true),
|
||||
"false" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn secret(value: String) -> serde_json::Value {
|
||||
secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool)
|
||||
}
|
||||
|
||||
// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness
|
||||
fn coerces_to_true(value: &serde_json::Value) -> bool {
|
||||
match value {
|
||||
serde_json::Value::String(text) => {
|
||||
matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")
|
||||
}
|
||||
serde_json::Value::Null => false,
|
||||
serde_json::Value::Bool(flag) => *flag,
|
||||
serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0),
|
||||
serde_json::Value::Array(items) => !items.is_empty(),
|
||||
serde_json::Value::Object(entries) => !entries.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
fn text(value: &serde_json::Value) -> Option<String> {
|
||||
match value {
|
||||
serde_json::Value::String(text) => Some(text.clone()),
|
||||
serde_json::Value::Number(number) => Some(number.to_string()),
|
||||
serde_json::Value::Bool(flag) => Some(flag.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn keeps_null(nulls: Nulls, key: &str) -> bool {
|
||||
match nulls {
|
||||
Nulls::Dropped => false,
|
||||
Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key),
|
||||
#[cfg(test)]
|
||||
Nulls::Kept => true,
|
||||
}
|
||||
}
|
||||
|
||||
fn environment(lookup: &dyn Lookup) -> impl Iterator<Item = (String, Entry)> + '_ {
|
||||
SETTINGS.iter().filter_map(|setting| {
|
||||
let source = setting.environment_name()?;
|
||||
let value = lookup.get(&source)?;
|
||||
let entry = Entry {
|
||||
value: secret(value),
|
||||
source,
|
||||
};
|
||||
Some((setting.kwarg.to_owned(), entry))
|
||||
})
|
||||
}
|
||||
|
||||
// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other
|
||||
fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value {
|
||||
match value {
|
||||
Value::String(text) => match text.strip_prefix("os.environ/") {
|
||||
Some(name) => lookup.get(name).map_or(Value::Null, |value| {
|
||||
secret_flag(&value).map_or(Value::String(value), Value::Bool)
|
||||
}),
|
||||
None => value.clone(),
|
||||
},
|
||||
_ => value.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_set(kwargs: &Kwargs, key: &str) -> bool {
|
||||
kwargs.get(key).is_some_and(|entry| !entry.value.is_null())
|
||||
}
|
||||
|
|
@ -1,19 +1,31 @@
|
|||
use std::time::Duration;
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State};
|
||||
use litellm_config::{
|
||||
CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value,
|
||||
environment_has_target,
|
||||
};
|
||||
use litellm_redis_client::{
|
||||
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient,
|
||||
RedisConfig,
|
||||
};
|
||||
use litellm_circuit_breaker::Config as BreakerConfig;
|
||||
use litellm_config::{CacheSettings, Config, RedisSettings, Value};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_redis_client::{RedisClient, RedisEvents};
|
||||
use litellm_redis_connect::{HostServices, SharedEnvironment, TracingEvents};
|
||||
use litellm_redis_resolve::Expected;
|
||||
use litellm_redis_types::Resolution;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
mod configured;
|
||||
mod layers;
|
||||
mod secrets;
|
||||
#[cfg(test)]
|
||||
mod surface_tests;
|
||||
|
||||
use configured::ConfiguredRedis;
|
||||
use secrets::SecretSnapshot;
|
||||
|
||||
pub(crate) struct Resources {
|
||||
pub(crate) http: litellm_http::Client,
|
||||
pub(crate) lookup: SharedEnvironment,
|
||||
}
|
||||
|
||||
const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
type Params = ResolvedRedis;
|
||||
type Params = Resolution;
|
||||
|
||||
pub struct RedisClients {
|
||||
pub coordination: Option<RedisClient>,
|
||||
|
|
@ -34,12 +46,15 @@ enum Target {
|
|||
Environment(Params),
|
||||
}
|
||||
|
||||
pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
|
||||
pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result<Plan, Error> {
|
||||
let response_cache = redis_cache_params(config)
|
||||
.map(|settings| settings.resolve_cache(lookup))
|
||||
.map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup))
|
||||
.transpose()?;
|
||||
let coordination = match &config.general_settings.coordination_redis {
|
||||
Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)),
|
||||
Some(settings) => Some(Target::Configured(resolve(
|
||||
&ConfiguredRedis::coordination(settings)?,
|
||||
lookup,
|
||||
)?)),
|
||||
None => match (&response_cache, config.litellm_settings.cache == Some(true)) {
|
||||
(Some(cache), _) => Some(Target::Configured(cache.clone())),
|
||||
(None, true) => environment(lookup)?.map(Target::Configured),
|
||||
|
|
@ -57,12 +72,29 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
|
|||
})
|
||||
}
|
||||
|
||||
pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients, Error> {
|
||||
pub(crate) async fn connect(
|
||||
config: &Config,
|
||||
lookup: &dyn Lookup,
|
||||
resources: &Resources,
|
||||
) -> Result<RedisClients, Error> {
|
||||
let plan = plan(config, lookup)?;
|
||||
let breaker = breaker_config(lookup)?;
|
||||
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
|
||||
let client = |params: Params| -> Result<RedisClient, Error> {
|
||||
Ok(litellm_redis_connect::connect(
|
||||
params,
|
||||
HostServices {
|
||||
events: Some(events.clone()),
|
||||
http: resources.http.clone(),
|
||||
environment: resources.lookup.clone(),
|
||||
credential_provider: None,
|
||||
},
|
||||
breaker,
|
||||
)?)
|
||||
};
|
||||
let coordination = match plan.coordination {
|
||||
Some(Target::Configured(params)) => Some(client(params, breaker)?),
|
||||
Some(Target::Environment(params)) => match client(params, breaker) {
|
||||
Some(Target::Configured(params)) => Some(client(params)?),
|
||||
Some(Target::Environment(params)) => match client(params) {
|
||||
Ok(client) => reachable(client).await,
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
|
|
@ -75,13 +107,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients
|
|||
};
|
||||
let transaction_buffer = match (&coordination, plan.transaction_buffer_fallback) {
|
||||
(Some(coordination), _) => Some(coordination.clone()),
|
||||
(None, Some(params)) => Some(client(params, breaker)?),
|
||||
(None, Some(params)) => Some(client(params)?),
|
||||
(None, None) => None,
|
||||
};
|
||||
let response_cache = plan
|
||||
.response_cache
|
||||
.map(|params| client(params, breaker))
|
||||
.transpose()?;
|
||||
let response_cache = plan.response_cache.map(client).transpose()?;
|
||||
Ok(RedisClients {
|
||||
coordination,
|
||||
transaction_buffer,
|
||||
|
|
@ -98,17 +127,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> {
|
|||
}
|
||||
}
|
||||
|
||||
fn environment(lookup: Lookup<'_>) -> Result<Option<Params>, Error> {
|
||||
fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result<Params, Error> {
|
||||
let secrets = SecretSnapshot::read(configured.secret_names(), lookup);
|
||||
let resolution = configured.resolve(&secrets, lookup)?;
|
||||
for ignored in &resolution.ignored {
|
||||
tracing::warn!(
|
||||
"Redis setting {} is not supported by the Rust client and is ignored: {}",
|
||||
ignored.setting,
|
||||
ignored.consequence
|
||||
);
|
||||
}
|
||||
Ok(resolution)
|
||||
}
|
||||
|
||||
fn environment(lookup: &dyn Lookup) -> Result<Option<Params>, Error> {
|
||||
match environment_has_target(lookup) {
|
||||
true => RedisSettings::default()
|
||||
.resolve(lookup)
|
||||
.map(Some)
|
||||
.map_err(Into::into),
|
||||
true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)),
|
||||
false => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_environment(lookup: Lookup<'_>) -> Option<Params> {
|
||||
fn environment_has_target(lookup: &dyn Lookup) -> bool {
|
||||
[
|
||||
"REDIS_HOST",
|
||||
"REDIS_URL",
|
||||
"REDIS_CLUSTER_NODES",
|
||||
"REDIS_SENTINEL_NODES",
|
||||
]
|
||||
.iter()
|
||||
.any(|name| lookup.get(name).is_some())
|
||||
}
|
||||
|
||||
fn optional_environment(lookup: &dyn Lookup) -> Option<Params> {
|
||||
match environment(lookup) {
|
||||
Ok(params) => params,
|
||||
Err(error) => {
|
||||
|
|
@ -149,122 +199,101 @@ async fn reachable(client: RedisClient) -> Option<RedisClient> {
|
|||
}
|
||||
}
|
||||
|
||||
fn client(params: Params, breaker: Option<BreakerConfig>) -> Result<RedisClient, Error> {
|
||||
tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured");
|
||||
let endpoint = match params.endpoint {
|
||||
RedisEndpoint::Url(url) => Endpoint::Url(url),
|
||||
RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint {
|
||||
host: host.host,
|
||||
port: host.port,
|
||||
username: host.username,
|
||||
password: host.password,
|
||||
tls: host.tls,
|
||||
db: host.db,
|
||||
}),
|
||||
RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
|
||||
nodes: cluster
|
||||
.nodes
|
||||
.into_iter()
|
||||
.map(|node| ClusterNode {
|
||||
host: node.host,
|
||||
port: node.port,
|
||||
})
|
||||
.collect(),
|
||||
username: cluster.username,
|
||||
password: cluster.password,
|
||||
tls: cluster.tls,
|
||||
}),
|
||||
};
|
||||
Ok(RedisClient::new(RedisConfig {
|
||||
endpoint,
|
||||
namespace: params.namespace.map(Namespace::new).unwrap_or_default(),
|
||||
timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT),
|
||||
breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))),
|
||||
})?)
|
||||
fn parse_seconds(text: &str) -> Option<Duration> {
|
||||
text.parse::<f64>()
|
||||
.ok()
|
||||
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
|
||||
}
|
||||
|
||||
fn log_transition(transition: &litellm_circuit_breaker::Transition<State>) {
|
||||
match transition.to {
|
||||
State::Open { .. } => tracing::warn!(
|
||||
generation = transition.generation,
|
||||
"redis circuit breaker opened, fast-failing redis calls"
|
||||
),
|
||||
State::HalfOpen => tracing::info!(
|
||||
generation = transition.generation,
|
||||
"redis circuit breaker half open, probing redis"
|
||||
),
|
||||
State::Closed { .. } => tracing::info!(
|
||||
generation = transition.generation,
|
||||
"redis circuit breaker closed, redis recovered"
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn breaker_config(lookup: Lookup<'_>) -> Result<Option<BreakerConfig>, Error> {
|
||||
let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED")
|
||||
fn breaker_config(lookup: &dyn Lookup) -> Result<Option<BreakerConfig>, Error> {
|
||||
let enabled = lookup
|
||||
.get("REDIS_CIRCUIT_BREAKER_ENABLED")
|
||||
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
|
||||
if !enabled {
|
||||
return Ok(None);
|
||||
}
|
||||
let defaults = BreakerConfig::default();
|
||||
let seconds = |name: &'static str, default: Duration| -> Result<Duration, Error> {
|
||||
lookup(name)
|
||||
.map(|value| {
|
||||
value
|
||||
.parse::<f64>()
|
||||
.ok()
|
||||
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
|
||||
.ok_or(Error::RedisInvalid {
|
||||
setting: name,
|
||||
reason: "must be a non-negative number of seconds",
|
||||
})
|
||||
})
|
||||
.unwrap_or(Ok(default))
|
||||
};
|
||||
let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD")
|
||||
Ok(Some(BreakerConfig {
|
||||
failure_threshold: environment_value(
|
||||
lookup,
|
||||
"REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
|
||||
Expected::Integer,
|
||||
|text| text.parse().ok(),
|
||||
)?
|
||||
.unwrap_or(defaults.failure_threshold),
|
||||
recovery_timeout: environment_value(
|
||||
lookup,
|
||||
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
|
||||
Expected::Seconds,
|
||||
parse_seconds,
|
||||
)?
|
||||
.unwrap_or(defaults.recovery_timeout),
|
||||
timeout_min_duration: environment_value(
|
||||
lookup,
|
||||
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
|
||||
Expected::Seconds,
|
||||
parse_seconds,
|
||||
)?
|
||||
.unwrap_or(defaults.timeout_min_duration),
|
||||
}))
|
||||
}
|
||||
|
||||
fn environment_value<T>(
|
||||
lookup: &dyn Lookup,
|
||||
name: &str,
|
||||
expected: Expected,
|
||||
parse: impl Fn(&str) -> Option<T>,
|
||||
) -> Result<Option<T>, litellm_redis_resolve::Error> {
|
||||
lookup
|
||||
.get(name)
|
||||
.map(|value| {
|
||||
value.parse().map_err(|_| Error::RedisInvalid {
|
||||
setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
|
||||
reason: "must be an integer",
|
||||
parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid {
|
||||
setting: name.to_owned(),
|
||||
expected,
|
||||
})
|
||||
})
|
||||
.unwrap_or(Ok(defaults.failure_threshold))?;
|
||||
Ok(Some(BreakerConfig {
|
||||
failure_threshold,
|
||||
recovery_timeout: seconds(
|
||||
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
|
||||
defaults.recovery_timeout,
|
||||
)?,
|
||||
timeout_min_duration: seconds(
|
||||
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
|
||||
defaults.timeout_min_duration,
|
||||
)?,
|
||||
}))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use std::{collections::HashMap, time::Duration};
|
||||
|
||||
use litellm_config::{Config, RedisEndpoint, RedisHost};
|
||||
use litellm_config::Config;
|
||||
use litellm_redis_client::{Standalone, StaticAuth, commands};
|
||||
use litellm_redis_testing::RedisServer;
|
||||
use litellm_redis_types::{Auth, HealthCheck};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Params, Plan, Target, connect, plan};
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{Params, Plan, Resources, Target, plan};
|
||||
use crate::Error;
|
||||
|
||||
async fn connect(config: &Config, lookup: &dyn Lookup) -> Result<super::RedisClients, Error> {
|
||||
let resources = Resources {
|
||||
http: litellm_http::Client::for_test(),
|
||||
lookup: Arc::new(|_: &str| None),
|
||||
};
|
||||
super::connect(config, lookup, &resources).await
|
||||
}
|
||||
|
||||
fn host(name: &str, port: u16) -> Params {
|
||||
Params {
|
||||
endpoint: RedisEndpoint::Host(RedisHost {
|
||||
target: litellm_redis_client::Target::Standalone(Standalone {
|
||||
host: name.to_owned(),
|
||||
port,
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
tls: None,
|
||||
db: 0,
|
||||
}),
|
||||
auth: Auth::Static(StaticAuth::default()),
|
||||
namespace: None,
|
||||
timeout: None,
|
||||
connect_timeout: None,
|
||||
health_check: HealthCheck::Unset,
|
||||
keepalive: true,
|
||||
ignored: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -273,7 +302,7 @@ mod tests {
|
|||
.iter()
|
||||
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
|
||||
.collect();
|
||||
plan(&Config::from_yaml(yaml).unwrap(), &|name| {
|
||||
plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| {
|
||||
environment.get(name).cloned()
|
||||
})
|
||||
}
|
||||
|
|
@ -350,36 +379,44 @@ mod tests {
|
|||
port_of(&server)
|
||||
))
|
||||
.unwrap();
|
||||
let clients = connect(&config, &|_| None).await.unwrap();
|
||||
let clients = connect(&config, &|_: &str| None).await.unwrap();
|
||||
let coordination = clients.coordination.unwrap();
|
||||
let buffer = clients.transaction_buffer.unwrap();
|
||||
let cache = clients.response_cache.unwrap();
|
||||
let unnamespaced = connect(&coordination_config(&server), &|_| None)
|
||||
let unnamespaced = connect(&coordination_config(&server), &|_: &str| None)
|
||||
.await
|
||||
.unwrap()
|
||||
.coordination
|
||||
.unwrap();
|
||||
|
||||
coordination
|
||||
.set("rate-limit", "1", Duration::from_secs(60))
|
||||
.send(commands::set("rate-limit", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
buffer
|
||||
.set("spend", "2", Duration::from_secs(60))
|
||||
.send(commands::set("spend", "2", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
cache.get::<String>("rate-limit").await.unwrap().as_deref(),
|
||||
cache
|
||||
.send(commands::get::<String>("rate-limit"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("1")
|
||||
);
|
||||
assert_eq!(
|
||||
cache.get::<String>("spend").await.unwrap().as_deref(),
|
||||
cache
|
||||
.send(commands::get::<String>("spend"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("2")
|
||||
);
|
||||
assert_eq!(
|
||||
unnamespaced
|
||||
.get::<String>("responses:rate-limit")
|
||||
.send(commands::get::<String>("responses:rate-limit"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
|
|
@ -387,7 +424,7 @@ mod tests {
|
|||
);
|
||||
assert_eq!(
|
||||
unnamespaced
|
||||
.get::<String>("responses:spend")
|
||||
.send(commands::get::<String>("responses:spend"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
|
|
@ -395,16 +432,21 @@ mod tests {
|
|||
);
|
||||
assert!(
|
||||
unnamespaced
|
||||
.get::<String>("rate-limit")
|
||||
.send(commands::get::<String>("rate-limit"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
unnamespaced
|
||||
.send(commands::get::<String>("spend"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(unnamespaced.get::<String>("spend").await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
|
||||
#[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])]
|
||||
fn redis_migration_regression_malformed_optional_environment_is_absent(
|
||||
#[case] environment: &[(&str, &str)],
|
||||
|
|
@ -413,54 +455,45 @@ mod tests {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
fn redis_migration_regression_explicit_malformed_setting_still_fails() {
|
||||
#[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
|
||||
#[case::transaction_buffer_environment(
|
||||
"general_settings:\n use_redis_transaction_buffer: true",
|
||||
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")]
|
||||
)]
|
||||
#[case::response_cache(
|
||||
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port",
|
||||
&[]
|
||||
)]
|
||||
fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
) {
|
||||
let plan = resolved(yaml, environment).unwrap();
|
||||
|
||||
let target = match (plan.coordination, plan.transaction_buffer_fallback) {
|
||||
(Some(Target::Configured(params) | Target::Environment(params)), _)
|
||||
| (None, Some(params)) => params,
|
||||
(None, None) => panic!("expected a Redis target"),
|
||||
};
|
||||
assert_eq!(target.target, host("env", 6379).target);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::port("port: not-a-port", "port")]
|
||||
#[case::ssl("ssl: maybe", "ssl")]
|
||||
fn a_malformed_coordination_block_fails_like_the_pydantic_model(
|
||||
#[case] setting: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let error = resolved(
|
||||
"general_settings:\n coordination_redis:\n host: coord\n port: not-a-port",
|
||||
&format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"),
|
||||
&[],
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Config(litellm_config::Error::RedisInvalid {
|
||||
setting: "port",
|
||||
..
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() {
|
||||
let error = resolved(
|
||||
"general_settings:\n use_redis_transaction_buffer: true",
|
||||
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")],
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Config(litellm_config::Error::RedisInvalid {
|
||||
setting: "port",
|
||||
..
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn redis_migration_regression_invalid_response_cache_still_fails() {
|
||||
let error = resolved(
|
||||
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port",
|
||||
&[],
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Config(litellm_config::Error::RedisInvalid {
|
||||
setting: "port",
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(
|
||||
matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -468,7 +501,7 @@ mod tests {
|
|||
async fn redis_migration_regression_invalid_optional_environment_url_is_absent() {
|
||||
let config = Config::from_yaml("model_list: []").unwrap();
|
||||
|
||||
let clients = connect(&config, &|name| {
|
||||
let clients = connect(&config, &|name: &str| {
|
||||
(name == "REDIS_URL").then(|| "not a redis URL".to_owned())
|
||||
})
|
||||
.await
|
||||
|
|
@ -506,7 +539,7 @@ mod tests {
|
|||
let server = RedisServer::start().await.unwrap();
|
||||
server.stop().await.unwrap();
|
||||
|
||||
let clients = connect(&coordination_config(&server), &|_| None)
|
||||
let clients = connect(&coordination_config(&server), &|_: &str| None)
|
||||
.await
|
||||
.unwrap();
|
||||
let coordination = clients.coordination.unwrap();
|
||||
|
|
@ -521,12 +554,12 @@ mod tests {
|
|||
#[tokio::test]
|
||||
async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let clients = connect(&coordination_config(&server), &|_| None)
|
||||
let clients = connect(&coordination_config(&server), &|_: &str| None)
|
||||
.await
|
||||
.unwrap();
|
||||
let coordination = clients.coordination.unwrap();
|
||||
coordination
|
||||
.set("written", "1", Duration::from_secs(60))
|
||||
.send(commands::set("written", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
|
@ -537,6 +570,30 @@ mod tests {
|
|||
assert!(buffer.ping().await.is_ok());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_environment_only_tls_redis_is_reached_with_its_ca() {
|
||||
crate::install_crypto_provider();
|
||||
let server = litellm_redis_testing::RedisTlsServer::start(false)
|
||||
.await
|
||||
.unwrap();
|
||||
let ca = server.files().ca.display().to_string();
|
||||
let port = server.port().to_string();
|
||||
let config = Config::from_yaml("model_list: []").unwrap();
|
||||
|
||||
let clients = connect(&config, &|name: &str| match name {
|
||||
"REDIS_HOST" => Some(server.host().to_owned()),
|
||||
"REDIS_PORT" => Some(port.clone()),
|
||||
"REDIS_SSL" => Some("true".to_owned()),
|
||||
"REDIS_SSL_CA_CERTS" => Some(ca.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(clients.coordination.is_some());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() {
|
||||
|
|
@ -545,7 +602,7 @@ mod tests {
|
|||
server.stop().await.unwrap();
|
||||
let config = Config::from_yaml("model_list: []").unwrap();
|
||||
|
||||
let clients = connect(&config, &|name| match name {
|
||||
let clients = connect(&config, &|name: &str| match name {
|
||||
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
|
||||
"REDIS_PORT" => Some(port.clone()),
|
||||
_ => None,
|
||||
|
|
@ -564,7 +621,7 @@ mod tests {
|
|||
let port = port_of(&server);
|
||||
let config = Config::from_yaml("model_list: []").unwrap();
|
||||
|
||||
let clients = connect(&config, &|name| match name {
|
||||
let clients = connect(&config, &|name: &str| match name {
|
||||
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
|
||||
"REDIS_PORT" => Some(port.clone()),
|
||||
_ => None,
|
||||
61
litellm-rust/crates/gateway/src/redis/secrets.rs
Normal file
61
litellm-rust/crates/gateway/src/redis/secrets.rs
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use litellm_config::Value;
|
||||
use litellm_redis_resolve::SETTINGS;
|
||||
|
||||
use super::configured::ConfiguredRedis;
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub(crate) struct SecretSnapshot(BTreeMap<String, Option<SecretValue>>);
|
||||
|
||||
impl SecretSnapshot {
|
||||
pub(crate) fn read(names: BTreeSet<String>, lookup: &dyn Lookup) -> Self {
|
||||
Self(
|
||||
names
|
||||
.into_iter()
|
||||
.map(|name| {
|
||||
let value = lookup.get(&name).map(SecretValue::new);
|
||||
(name, value)
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Lookup for SecretSnapshot {
|
||||
fn get(&self, name: &str) -> Option<String> {
|
||||
match self.0.get(name) {
|
||||
Some(value) => value.as_ref().map(|value| value.expose().to_owned()),
|
||||
None => {
|
||||
debug_assert!(
|
||||
false,
|
||||
"{name} is missing from ConfiguredRedis::secret_names"
|
||||
);
|
||||
tracing::error!(
|
||||
name,
|
||||
"a Redis secret was read without being in the snapshot"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ConfiguredRedis {
|
||||
pub(crate) fn secret_names(&self) -> BTreeSet<String> {
|
||||
let references = self.fields.values().filter_map(|value| match value {
|
||||
Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned),
|
||||
_ => None,
|
||||
});
|
||||
let declared = SETTINGS.iter().flat_map(|setting| {
|
||||
setting
|
||||
.environment_name()
|
||||
.into_iter()
|
||||
.chain(setting.fallback.iter().map(|name| (*name).to_owned()))
|
||||
});
|
||||
declared.chain(references).collect()
|
||||
}
|
||||
}
|
||||
964
litellm-rust/crates/gateway/src/redis/surface_tests.rs
Normal file
964
litellm-rust/crates/gateway/src/redis/surface_tests.rs
Normal file
|
|
@ -0,0 +1,964 @@
|
|||
use std::collections::{BTreeMap, HashMap};
|
||||
|
||||
use litellm_config::{CacheSettings, Config, Value};
|
||||
use litellm_redis_resolve::Expected;
|
||||
use litellm_redis_types::{
|
||||
Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions,
|
||||
};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{configured::ConfiguredRedis, secrets::SecretSnapshot};
|
||||
use crate::Error;
|
||||
|
||||
fn resolution(
|
||||
configured: &ConfiguredRedis,
|
||||
environment: &HashMap<String, String>,
|
||||
) -> Result<Resolution, Error> {
|
||||
let lookup = |name: &str| environment.get(name).cloned();
|
||||
configured.resolve(
|
||||
&SecretSnapshot::read(configured.secret_names(), &lookup),
|
||||
&lookup,
|
||||
)
|
||||
}
|
||||
|
||||
const PROVIDER: ProviderHandle = ProviderHandle::new(7);
|
||||
|
||||
fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis {
|
||||
let has_provider = kwargs
|
||||
.get("credential_provider")
|
||||
.is_some_and(|provider| !provider.is_null());
|
||||
let configured = ConfiguredRedis::redis_cache_kwargs(
|
||||
serde_json::from_value::<BTreeMap<String, Value>>(kwargs).unwrap(),
|
||||
);
|
||||
match has_provider {
|
||||
true => configured.with_credential_provider(PROVIDER),
|
||||
false => configured,
|
||||
}
|
||||
}
|
||||
|
||||
fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
|
||||
let environment: HashMap<String, String> = environment
|
||||
.iter()
|
||||
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
|
||||
.collect();
|
||||
let config = Config::from_yaml(yaml).unwrap();
|
||||
let configured = config
|
||||
.general_settings
|
||||
.coordination_redis
|
||||
.unwrap_or_default();
|
||||
let configured = ConfiguredRedis::coordination(&configured)?;
|
||||
resolution(&configured, &environment)
|
||||
}
|
||||
|
||||
fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
|
||||
coordination_resolution(yaml, environment)
|
||||
}
|
||||
|
||||
fn resolve_cache(yaml: &str) -> Result<Resolution, Error> {
|
||||
resolve_cache_in(yaml, &[])
|
||||
}
|
||||
|
||||
fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result<Resolution, Error> {
|
||||
let lookup = |name: &str| {
|
||||
environment
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| (*value).to_owned())
|
||||
};
|
||||
let config = Config::from_yaml(yaml).unwrap();
|
||||
let configured = match config.litellm_settings.cache_params.unwrap() {
|
||||
CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings),
|
||||
CacheSettings::Other(_) => panic!("expected a Redis cache"),
|
||||
};
|
||||
configured.resolve(
|
||||
&SecretSnapshot::read(configured.secret_names(), &lookup),
|
||||
&lookup,
|
||||
)
|
||||
}
|
||||
|
||||
fn cluster_ports(resolved: &Resolution) -> Vec<u16> {
|
||||
match &resolved.target {
|
||||
Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
|
||||
target => panic!("expected a cluster, got {target:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::host(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
|
||||
&[],
|
||||
"host:redis:6380:2:true:app:1500"
|
||||
)]
|
||||
#[case::environment_references_resolve_the_host_and_port(
|
||||
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
|
||||
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
|
||||
"host:referenced:6400:0:false::none"
|
||||
)]
|
||||
#[case::a_null_host_falls_back_to_the_environment(
|
||||
"general_settings:\n coordination_redis:\n host: null",
|
||||
&[("REDIS_HOST", "environment")],
|
||||
"host:environment:6379:0:false::none"
|
||||
)]
|
||||
fn resolves_redis_through_the_parsed_config(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let resolved = resolve(yaml, environment).unwrap();
|
||||
let description = match resolved.target {
|
||||
Target::Unix(socket) => format!("unix:{}", socket.path.display()),
|
||||
Target::Url(url) => format!("url:{}", url.url.expose()),
|
||||
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
|
||||
Target::Standalone(host) => format!(
|
||||
"host:{}:{}:{}:{}:{}:{}",
|
||||
host.host,
|
||||
host.port,
|
||||
host.db,
|
||||
host.tls.is_some(),
|
||||
resolved.namespace.unwrap_or_default(),
|
||||
resolved
|
||||
.timeout
|
||||
.map(|timeout| timeout.as_millis().to_string())
|
||||
.unwrap_or_else(|| "none".to_owned())
|
||||
),
|
||||
Target::Cluster(cluster) => format!(
|
||||
"cluster:{}",
|
||||
cluster
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|node| format!("{}:{}", node.host, node.port))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
),
|
||||
};
|
||||
|
||||
assert_eq!(description, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::tls_ciphers(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH",
|
||||
&[],
|
||||
"ssl_ciphers"
|
||||
)]
|
||||
#[case::ssl_ciphers_from_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")],
|
||||
"REDIS_SSL_CIPHERS"
|
||||
)]
|
||||
#[case::connection_pool(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared",
|
||||
&[],
|
||||
"connection_pool"
|
||||
)]
|
||||
fn unsupported_settings_resolve_and_are_named_by_their_source(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let resolution = coordination_resolution(yaml, environment).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolution.target,
|
||||
Target::Standalone(litellm_redis_types::Standalone {
|
||||
host: "redis".into(),
|
||||
port: 6379,
|
||||
db: 0,
|
||||
tls: None,
|
||||
})
|
||||
);
|
||||
let [ignored] = resolution.ignored.as_slice() else {
|
||||
panic!("expected one ignored setting, got {:?}", resolution.ignored);
|
||||
};
|
||||
assert_eq!(ignored.setting, expected);
|
||||
assert!(!ignored.consequence.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::kwarg(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true",
|
||||
&[],
|
||||
"ssl_validate_ocsp"
|
||||
)]
|
||||
#[case::stapled_from_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")],
|
||||
"REDIS_SSL_VALIDATE_OCSP_STAPLED"
|
||||
)]
|
||||
#[case::rediss_url_query(
|
||||
"general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true",
|
||||
&[],
|
||||
"url ?ssl_validate_ocsp"
|
||||
)]
|
||||
fn an_ocsp_check_fails_instead_of_connecting_without_it(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] setting: &str,
|
||||
) {
|
||||
let error = coordination_resolution(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")]
|
||||
#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")]
|
||||
fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) {
|
||||
let yaml = format!("general_settings:\n coordination_redis:\n{block}");
|
||||
|
||||
let resolution = coordination_resolution(&yaml, &[]).unwrap();
|
||||
|
||||
assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored);
|
||||
}
|
||||
|
||||
fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth {
|
||||
Auth::Static(StaticAuth {
|
||||
username: username.map(litellm_secret_value::SecretValue::new),
|
||||
password: password.map(litellm_secret_value::SecretValue::new),
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_explicit_host_keeps_environment_credentials() {
|
||||
let resolved = resolve(
|
||||
"general_settings:\n coordination_redis:\n host: configured",
|
||||
&[
|
||||
("REDIS_PASSWORD", "environment-secret"),
|
||||
("REDIS_SSL", "True"),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resolved.auth, static_auth(None, Some("environment-secret")));
|
||||
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some()));
|
||||
}
|
||||
|
||||
fn azure(username: Option<&str>) -> Auth {
|
||||
Auth::AzureAd(AzureAdAuth {
|
||||
client_id: Some("client".to_owned()),
|
||||
tenant_id: None,
|
||||
client_secret: None,
|
||||
username: username.map(str::to_owned),
|
||||
})
|
||||
}
|
||||
|
||||
fn aws(serverless: bool) -> Auth {
|
||||
Auth::AwsIam(AwsIamAuth {
|
||||
user_name: "iam-user".to_owned(),
|
||||
cache_name: "cache-1".to_owned(),
|
||||
region: "us-east-1".to_owned(),
|
||||
serverless,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::static_credentials(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass",
|
||||
&[],
|
||||
static_auth(Some("user"), Some("pass"))
|
||||
)]
|
||||
#[case::gcp_beats_azure_and_aws(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true",
|
||||
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")],
|
||||
Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() }
|
||||
)]
|
||||
#[case::azure_beats_aws_and_takes_the_environment_username(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true",
|
||||
&[("REDIS_USERNAME", "environment-user")],
|
||||
azure(Some("environment-user"))
|
||||
)]
|
||||
#[case::azure_without_an_environment_username_sends_only_the_token(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client",
|
||||
&[],
|
||||
azure(None)
|
||||
)]
|
||||
#[case::azure_needs_the_literal_true(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"",
|
||||
&[],
|
||||
static_auth(None, None)
|
||||
)]
|
||||
#[case::aws_with_a_region_from_the_environment(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1",
|
||||
&[("AWS_DEFAULT_REGION", "us-east-1")],
|
||||
aws(false)
|
||||
)]
|
||||
#[case::aws_serverless_over_a_rediss_url(
|
||||
"general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true",
|
||||
&[],
|
||||
aws(true)
|
||||
)]
|
||||
fn auth_follows_python_precedence(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: Auth,
|
||||
) {
|
||||
assert_eq!(resolve(yaml, environment).unwrap().auth, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::coordination(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object",
|
||||
&[]
|
||||
)]
|
||||
#[case::without_a_target(
|
||||
"general_settings:\n coordination_redis:\n credential_provider: host-object",
|
||||
&[("REDIS_HOST", "redis")]
|
||||
)]
|
||||
fn a_credential_provider_without_a_host_handle_fails(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
error,
|
||||
Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle)
|
||||
),
|
||||
"{error}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::without_tls(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
|
||||
"requires TLS"
|
||||
)]
|
||||
#[case::a_plain_url_ignores_ssl(
|
||||
"general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r",
|
||||
"requires TLS"
|
||||
)]
|
||||
#[case::missing_settings(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c",
|
||||
"requires: aws_iam_user_name, aws_iam_region"
|
||||
)]
|
||||
fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) {
|
||||
let error = resolve(yaml, &[]).unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains(message), "{error}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::defaults_verify_the_certificate_but_not_the_hostname(
|
||||
" ssl: true",
|
||||
&[],
|
||||
Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
|
||||
)]
|
||||
#[case::every_option(
|
||||
" ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true",
|
||||
&[],
|
||||
Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true })
|
||||
)]
|
||||
#[case::options_from_the_environment(
|
||||
" ssl: true",
|
||||
&[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")],
|
||||
Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false })
|
||||
)]
|
||||
#[case::gcp_ca_certs_replace_the_configured_ones(
|
||||
" ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem",
|
||||
&[],
|
||||
Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false })
|
||||
)]
|
||||
#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)]
|
||||
fn tls_options_resolve_like_redis_py(
|
||||
#[case] block: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: Option<TlsOptions>,
|
||||
) {
|
||||
let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}");
|
||||
|
||||
let resolved = resolve(&yaml, environment).unwrap();
|
||||
|
||||
assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_unknown_cert_requirement_fails() {
|
||||
let error = resolve(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict",
|
||||
&[],
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs"
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::from_config(
|
||||
"general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2",
|
||||
&[]
|
||||
)]
|
||||
#[case::from_the_environment(
|
||||
"model_list: []",
|
||||
&[
|
||||
("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#),
|
||||
("REDIS_SERVICE_NAME", "primary"),
|
||||
("REDIS_SENTINEL_PASSWORD", "sentinel-secret"),
|
||||
("REDIS_DB", "2"),
|
||||
]
|
||||
)]
|
||||
fn sentinel_nodes_resolve_to_a_sentinel_endpoint(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
) {
|
||||
let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else {
|
||||
panic!("expected a sentinel endpoint");
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
sentinel.monitors,
|
||||
vec![
|
||||
Node {
|
||||
host: "sentinel-a".to_owned(),
|
||||
port: 26379
|
||||
},
|
||||
Node {
|
||||
host: "sentinel-b".to_owned(),
|
||||
port: 26380
|
||||
},
|
||||
]
|
||||
);
|
||||
assert_eq!(sentinel.service_name, "primary");
|
||||
assert_eq!(
|
||||
sentinel.monitor_auth,
|
||||
Some(StaticAuth {
|
||||
username: None,
|
||||
password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")),
|
||||
})
|
||||
);
|
||||
assert_eq!(sentinel.db, 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::a_url_beats_sentinels(
|
||||
"general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
|
||||
"url:redis://single:6379"
|
||||
)]
|
||||
fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) {
|
||||
let resolved = resolve(yaml, &[]).unwrap();
|
||||
let description = match resolved.target {
|
||||
Target::Url(url) => format!("url:{}", url.url.expose()),
|
||||
Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
};
|
||||
|
||||
assert_eq!(description, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn sentinels_without_a_service_name_fail() {
|
||||
let error = resolve(
|
||||
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]",
|
||||
&[],
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
error,
|
||||
Error::Redis(litellm_redis_resolve::Error::NoSentinelService)
|
||||
),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configured_host(
|
||||
"general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
|
||||
&[],
|
||||
"host"
|
||||
)]
|
||||
#[case::environment_port(
|
||||
"general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary",
|
||||
&[("REDIS_PORT", "6380")],
|
||||
"REDIS_PORT"
|
||||
)]
|
||||
fn a_host_or_port_cannot_reach_sentinel(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] source: &str,
|
||||
) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_pairs(
|
||||
"general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary",
|
||||
&[],
|
||||
"sentinel_nodes"
|
||||
)]
|
||||
#[case::from_the_environment(
|
||||
"model_list: []",
|
||||
&[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")],
|
||||
"REDIS_SENTINEL_NODES"
|
||||
)]
|
||||
fn invalid_sentinel_nodes_name_their_source(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] source: &str,
|
||||
) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")]
|
||||
#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")]
|
||||
#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")]
|
||||
fn invalid_values_name_their_source(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] source: &str,
|
||||
) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)]
|
||||
#[case::explicit_connect_timeout(
|
||||
" socket_timeout: 1.5\n socket_connect_timeout: 0.25",
|
||||
Some(1500),
|
||||
Some(250)
|
||||
)]
|
||||
fn connect_timeout_is_resolved_separately(
|
||||
#[case] block: &str,
|
||||
#[case] timeout_ms: Option<u128>,
|
||||
#[case] connect_ms: Option<u128>,
|
||||
) {
|
||||
let resolved = resolve(
|
||||
&format!("general_settings:\n coordination_redis:\n host: redis\n{block}"),
|
||||
&[],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
resolved.timeout.map(|timeout| timeout.as_millis()),
|
||||
timeout_ms
|
||||
);
|
||||
assert_eq!(
|
||||
resolved.connect_timeout.map(|timeout| timeout.as_millis()),
|
||||
connect_ms
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
|
||||
#[case::a_structured_host_with_an_environment_target(
|
||||
"general_settings:\n coordination_redis:\n host: []",
|
||||
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
|
||||
)]
|
||||
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
|
||||
let error = resolve(yaml, environment).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_cache_alias_names_cluster_nodes() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cluster_ports(&resolved), vec![7000]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn canonical_startup_nodes_beat_the_cache_alias() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cluster_ports(&resolved), vec![7000]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(resolved.target, Target::Standalone(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
|
||||
let resolved = resolve_cache(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(resolved.target, Target::Standalone(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn unsupported_environment_settings_never_carry_the_value() {
|
||||
let resolution = coordination_resolution(
|
||||
"model_list: []",
|
||||
&[
|
||||
("REDIS_HOST", "redis"),
|
||||
("REDIS_SSL_CIPHERS", "private-option-value"),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let ignored = format!("{:?}", resolution.ignored);
|
||||
|
||||
assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}");
|
||||
assert!(!ignored.contains("private-option-value"), "{ignored}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::negative("-1")]
|
||||
#[case::infinite("inf")]
|
||||
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
|
||||
let yaml = format!(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
|
||||
);
|
||||
let error = resolve(&yaml, &[]).unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout"
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn debug_redacts_redis_secrets_and_url() {
|
||||
let config = Config::from_yaml(
|
||||
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}",
|
||||
)
|
||||
.unwrap();
|
||||
let settings = config.general_settings.coordination_redis.unwrap();
|
||||
let configured = ConfiguredRedis::coordination(&settings).unwrap();
|
||||
let resolved = resolution(&configured, &HashMap::new()).unwrap();
|
||||
let debug = format!("{settings:?} {configured:?} {resolved:?}");
|
||||
|
||||
assert!(!debug.contains("redis://user:password"));
|
||||
assert!(!debug.contains("credential-user"));
|
||||
assert!(!debug.contains("credential-password"));
|
||||
assert!(!debug.contains("credential-passthrough"));
|
||||
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn legacy_unknown_settings_do_not_change_resolution() {
|
||||
let resolved = resolve(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
|
||||
&[],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(resolved.target, Target::Standalone(_)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configured(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
|
||||
&[],
|
||||
"max_connections"
|
||||
)]
|
||||
#[case::from_the_environment(
|
||||
"general_settings:\n coordination_redis:\n host: redis",
|
||||
&[("REDIS_MAX_CONNECTIONS", "47")],
|
||||
"REDIS_MAX_CONNECTIONS"
|
||||
)]
|
||||
fn ignored_settings_are_named_by_their_source(
|
||||
#[case] yaml: &str,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] source: &str,
|
||||
) {
|
||||
let resolution = coordination_resolution(yaml, environment).unwrap();
|
||||
|
||||
let settings: Vec<&str> = resolution
|
||||
.ignored
|
||||
.iter()
|
||||
.map(|ignored| ignored.setting.as_str())
|
||||
.collect();
|
||||
assert_eq!(settings, vec![source]);
|
||||
}
|
||||
|
||||
fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot {
|
||||
let values: HashMap<&str, &str> = values.iter().copied().collect();
|
||||
SecretSnapshot::read(
|
||||
ConfiguredRedis::environment_only().secret_names(),
|
||||
&|name: &str| values.get(name).map(|value| (*value).to_owned()),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn secret_manager_values_reach_the_resolver_through_the_snapshot() {
|
||||
let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"}));
|
||||
let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name {
|
||||
"COORD_HOST" => Some("from-the-manager".to_owned()),
|
||||
"REDIS_PASSWORD" => Some("manager-secret".to_owned()),
|
||||
_ => None,
|
||||
});
|
||||
|
||||
let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager")
|
||||
);
|
||||
assert_eq!(resolved.auth, static_auth(None, Some("manager-secret")));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))]
|
||||
#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)]
|
||||
fn the_azure_username_comes_from_the_process_environment_like_python(
|
||||
#[case] secrets: &[(&str, &str)],
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
|
||||
let environment: HashMap<&str, &str> = environment.iter().copied().collect();
|
||||
|
||||
let resolved = settings
|
||||
.resolve(&snapshot(secrets), &|name: &str| {
|
||||
environment.get(name).map(|value| (*value).to_owned())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
|
||||
"{:?}",
|
||||
resolved.auth
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::coordination_drops_the_null(
|
||||
"general_settings:\n coordination_redis:\n host: redis\n ssl: null",
|
||||
false,
|
||||
true
|
||||
)]
|
||||
#[case::cache_passes_an_unnamed_null_on(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null",
|
||||
true,
|
||||
false
|
||||
)]
|
||||
#[case::cache_drops_a_null_it_names(
|
||||
"litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true",
|
||||
true,
|
||||
true
|
||||
)]
|
||||
fn a_null_overrides_the_environment_only_where_python_passes_it_on(
|
||||
#[case] yaml: &str,
|
||||
#[case] cache: bool,
|
||||
#[case] tls: bool,
|
||||
) {
|
||||
let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")];
|
||||
let resolved = match cache {
|
||||
true => resolve_cache_in(yaml, &environment),
|
||||
false => resolve(yaml, &environment),
|
||||
}
|
||||
.unwrap();
|
||||
|
||||
let Target::Standalone(host) = resolved.target else {
|
||||
panic!("expected a standalone target, got {:?}", resolved.target);
|
||||
};
|
||||
assert_eq!(host.tls.is_some(), tls);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty_cluster_nodes_use_the_environment(
|
||||
json!({"startup_nodes": []}),
|
||||
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)],
|
||||
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}),
|
||||
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
|
||||
)]
|
||||
#[case::empty_sentinel_nodes_use_the_environment(
|
||||
json!({"sentinel_nodes": [], "service_name": "primary"}),
|
||||
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)],
|
||||
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}),
|
||||
&[("sentinel_nodes", "REDIS_SENTINEL_NODES")],
|
||||
)]
|
||||
#[case::empty_service_name_uses_the_environment(
|
||||
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}),
|
||||
&[("REDIS_SERVICE_NAME", "primary")],
|
||||
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
|
||||
&[("service_name", "REDIS_SERVICE_NAME")],
|
||||
)]
|
||||
#[case::environment_host_port_and_password(
|
||||
json!({}),
|
||||
&[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")],
|
||||
json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}),
|
||||
&[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")],
|
||||
)]
|
||||
#[case::padded_environment_numbers_parse_like_python_int(
|
||||
json!({}),
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")],
|
||||
json!({"port": "6_380", "host": "redis", "db": "3\n"}),
|
||||
&[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
|
||||
)]
|
||||
#[case::socket_timeout_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SOCKET_TIMEOUT", "0.75")],
|
||||
json!({"socket_timeout": "0.75", "host": "redis"}),
|
||||
&[("socket_timeout", "REDIS_SOCKET_TIMEOUT")],
|
||||
)]
|
||||
#[case::environment_url(
|
||||
json!({}),
|
||||
&[("REDIS_URL", "redis://environment:6379/1")],
|
||||
json!({"url": "redis://environment:6379/1"}),
|
||||
&[("url", "REDIS_URL")],
|
||||
)]
|
||||
#[case::environment_url_beats_the_environment_host(
|
||||
json!({}),
|
||||
&[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")],
|
||||
json!({"host": "ignored", "url": "redis://from-url:6379"}),
|
||||
&[("host", "REDIS_HOST"), ("url", "REDIS_URL")],
|
||||
)]
|
||||
#[case::an_explicit_host_drops_the_environment_url(
|
||||
json!({"host": "configured"}),
|
||||
&[("REDIS_URL", "redis://environment:6379")],
|
||||
json!({"host": "configured"}),
|
||||
&[],
|
||||
)]
|
||||
#[case::environment_cluster_nodes(
|
||||
json!({}),
|
||||
&[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)],
|
||||
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}),
|
||||
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
|
||||
)]
|
||||
#[case::sentinel_from_the_environment(
|
||||
json!({}),
|
||||
&[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")],
|
||||
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}),
|
||||
&[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")],
|
||||
)]
|
||||
#[case::an_explicit_credential_provider_beats_cloud_auth(
|
||||
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
|
||||
&[("REDIS_AZURE_AD_TOKEN", "true")],
|
||||
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
|
||||
&[],
|
||||
)]
|
||||
#[case::gcp_beats_azure_and_aws(
|
||||
json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
|
||||
&[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")],
|
||||
json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
|
||||
&[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")],
|
||||
)]
|
||||
#[case::azure_beats_aws_and_takes_the_environment_username(
|
||||
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
|
||||
&[("REDIS_USERNAME", "environment-user")],
|
||||
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
|
||||
&[],
|
||||
)]
|
||||
#[case::an_empty_environment_username_sends_only_the_token(
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}),
|
||||
&[("REDIS_USERNAME", "")],
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}),
|
||||
&[("username", "REDIS_USERNAME")],
|
||||
)]
|
||||
#[case::aws_iam_over_tls_with_the_default_region(
|
||||
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}),
|
||||
&[("AWS_DEFAULT_REGION", "us-east-1")],
|
||||
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}),
|
||||
&[("aws_iam_region", "AWS_DEFAULT_REGION")],
|
||||
)]
|
||||
#[case::aws_region_beats_the_default_region(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}),
|
||||
&[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")],
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}),
|
||||
&[("aws_iam_region", "AWS_REGION")],
|
||||
)]
|
||||
#[case::a_null_cert_requirement_skips_verification(
|
||||
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
|
||||
&[("REDIS_SSL_CERT_REQS", "required")],
|
||||
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
|
||||
&[],
|
||||
)]
|
||||
#[case::a_cert_requirement_spelled_like_python_none_fails(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")],
|
||||
json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}),
|
||||
&[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
|
||||
)]
|
||||
#[case::tls_options_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")],
|
||||
json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}),
|
||||
&[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
|
||||
)]
|
||||
#[case::environment_references_resolve(
|
||||
json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}),
|
||||
&[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")],
|
||||
json!({"host": "referenced", "password": "referenced-secret"}),
|
||||
&[],
|
||||
)]
|
||||
#[case::malformed_port_from_the_environment(
|
||||
json!({}),
|
||||
&[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")],
|
||||
json!({"port": "abc", "host": "redis"}),
|
||||
&[("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
|
||||
)]
|
||||
#[case::retry_on_error_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_RETRY_ON_ERROR", "true")],
|
||||
json!({"retry_on_error": true, "host": "redis"}),
|
||||
&[("retry_on_error", "REDIS_RETRY_ON_ERROR")],
|
||||
)]
|
||||
#[case::ssl_ca_path_from_the_environment(
|
||||
json!({"host": "redis"}),
|
||||
&[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")],
|
||||
json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}),
|
||||
&[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")],
|
||||
)]
|
||||
fn merges_like_python_redis(
|
||||
#[case] kwargs: serde_json::Value,
|
||||
#[case] environment: &[(&str, &str)],
|
||||
#[case] merged: serde_json::Value,
|
||||
#[case] sources: &[(&str, &str)],
|
||||
) {
|
||||
let configured = redis_cache_kwargs(kwargs);
|
||||
let lookup = |name: &str| {
|
||||
environment
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| (*value).to_owned())
|
||||
};
|
||||
|
||||
let kwargs = configured
|
||||
.merged(&SecretSnapshot::read(configured.secret_names(), &lookup))
|
||||
.unwrap();
|
||||
|
||||
let values: serde_json::Map<String, serde_json::Value> = kwargs
|
||||
.iter()
|
||||
.map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone()))
|
||||
.collect();
|
||||
let renamed: Vec<(&str, &str)> = kwargs
|
||||
.iter()
|
||||
.filter(|(kwarg, entry)| *kwarg != &entry.source)
|
||||
.map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str()))
|
||||
.collect();
|
||||
assert_eq!(serde_json::Value::Object(values), merged);
|
||||
assert_eq!(renamed, sources);
|
||||
}
|
||||
|
|
@ -37,7 +37,7 @@ pub use owned::PythonOwned;
|
|||
pub use runtime::{
|
||||
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
|
||||
ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync,
|
||||
run_sync_value, runtime_started,
|
||||
run_sync_value, runtime_handle, runtime_started,
|
||||
};
|
||||
pub use services::PythonHostCalls;
|
||||
|
||||
|
|
|
|||
|
|
@ -68,6 +68,11 @@ fn runtime() -> PyResult<&'static Runtime> {
|
|||
Ok(pyo3_async_runtimes::tokio::get_runtime())
|
||||
}
|
||||
|
||||
/// The native runtime, for work that must outlive the call that schedules it.
|
||||
pub fn runtime_handle() -> PyResult<Handle> {
|
||||
Ok(runtime()?.handle().clone())
|
||||
}
|
||||
|
||||
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
|
||||
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
|
||||
where
|
||||
|
|
|
|||
|
|
@ -30,6 +30,14 @@ bytes.workspace = true
|
|||
futures-util.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-error.workspace = true
|
||||
litellm-circuit-breaker.workspace = true
|
||||
litellm-redis-client.workspace = true
|
||||
litellm-redis-connect.workspace = true
|
||||
litellm-redis-store.workspace = true
|
||||
litellm-redis-resolve.workspace = true
|
||||
litellm-python-compat.workspace = true
|
||||
litellm-redis-types.workspace = true
|
||||
tracing.workspace = true
|
||||
litellm-cache-response.workspace = true
|
||||
serde.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ mod marshal;
|
|||
mod model_registry;
|
||||
mod preflight;
|
||||
mod python_settings;
|
||||
mod redis;
|
||||
mod router;
|
||||
mod routes;
|
||||
mod secrets;
|
||||
|
|
@ -37,6 +38,14 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::model_registry::bundled_model_cost_map;
|
||||
#[pymodule_export]
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
#[pymodule_export]
|
||||
use crate::redis::client::native_redis_breaker_states;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::request::{NativeRequestRedis, NativeSend};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::{NativePending, NativeRedisStoreError};
|
||||
#[pymodule_export]
|
||||
use crate::router::{RouterCall, RouterCoreHandle};
|
||||
#[pymodule_export]
|
||||
use crate::routes::audio_transcription::{atranscription, transcription};
|
||||
|
|
@ -111,6 +120,12 @@ mod tests {
|
|||
"aresponses",
|
||||
"ResponsesWebSocketConnection",
|
||||
"NativeDiagnosticProcessor",
|
||||
"NativeRedisClient",
|
||||
"NativeRequestRedis",
|
||||
"NativeSend",
|
||||
"NativePending",
|
||||
"NativeRedisStoreError",
|
||||
"native_redis_breaker_states",
|
||||
"NativeClickHouseSpendConfig",
|
||||
"NativeClickHouseSpendStorage",
|
||||
"NativeDatabase",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,47 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_host_python::{PythonContext, attach_blocking};
|
||||
use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.host.redis_credentials";
|
||||
|
||||
/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host
|
||||
/// module the way redis-py's async connection asks it.
|
||||
pub(crate) struct PythonCredentialSource {
|
||||
provider: Arc<Py<PyAny>>,
|
||||
context: PythonContext,
|
||||
}
|
||||
|
||||
impl PythonCredentialSource {
|
||||
pub(crate) fn new(provider: Py<PyAny>, context: PythonContext) -> Self {
|
||||
Self {
|
||||
provider: Arc::new(provider),
|
||||
context,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CredentialSource for PythonCredentialSource {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
let provider = Arc::clone(&self.provider);
|
||||
let context = self.context.clone();
|
||||
Box::pin(async move {
|
||||
let (username, password) = attach_blocking(context, move |py| {
|
||||
py.import(MODULE)?
|
||||
.getattr("credentials")?
|
||||
.call1((provider.bind(py),))?
|
||||
.extract::<(Option<String>, String)>()
|
||||
})
|
||||
.await
|
||||
.and_then(|credentials| credentials)
|
||||
.map_err(Error::credential(CredentialStage::Provider))?;
|
||||
Ok(Expiring {
|
||||
username,
|
||||
password: SecretValue::new(password),
|
||||
expires_at: None,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
385
litellm-rust/crates/python-bridge/src/redis/client/mod.rs
Normal file
385
litellm-rust/crates/python-bridge/src/redis/client/mod.rs
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
mod credentials;
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::{Arc, LazyLock, Mutex, PoisonError},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::Config as BreakerConfig;
|
||||
use litellm_core_utils::settings::ProcessEnvironment;
|
||||
use litellm_host_python::PythonContext;
|
||||
use litellm_python_compat::number::{self, SecondsError};
|
||||
use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient};
|
||||
use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment};
|
||||
use litellm_redis_resolve::{Kwargs, PythonException};
|
||||
use litellm_redis_types::{ProviderHandle, Resolution};
|
||||
use pyo3::{
|
||||
exceptions::{PyOverflowError, PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
sync::MutexExt,
|
||||
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple},
|
||||
};
|
||||
|
||||
use self::credentials::PythonCredentialSource;
|
||||
use crate::redis::events::{ServiceLoggingEvents, state_label};
|
||||
|
||||
pyo3::import_exception!(redis.exceptions, RedisClusterException);
|
||||
|
||||
/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way
|
||||
/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the
|
||||
/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object
|
||||
/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle.
|
||||
pub(crate) struct RedisSpec {
|
||||
kwargs: Kwargs,
|
||||
credential_provider: Option<Py<PyAny>>,
|
||||
}
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let py = kwargs.py();
|
||||
let plain = PyDict::new(py);
|
||||
let mut host_objects = Vec::new();
|
||||
for (name, value) in kwargs.cast::<PyDict>()?.iter() {
|
||||
match value.is_none() || is_plain(&value) {
|
||||
true => plain.set_item(&name, value)?,
|
||||
false => host_objects.push(name.extract::<String>()?),
|
||||
}
|
||||
}
|
||||
let credential_provider = kwargs
|
||||
.cast::<PyDict>()?
|
||||
.get_item("credential_provider")?
|
||||
.filter(|provider| !provider.is_none() && !is_plain(provider))
|
||||
.map(Bound::unbind);
|
||||
let fields: BTreeMap<String, serde_json::Value> =
|
||||
litellm_host_python::from_py_argument(plain.as_any())?;
|
||||
let objects = host_objects
|
||||
.into_iter()
|
||||
.map(|name| (name, serde_json::Value::Bool(true)));
|
||||
Ok(Self {
|
||||
kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)),
|
||||
credential_provider,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// `litellm.constants`' `REDIS_CIRCUIT_BREAKER_*`, which Python reads from the environment once.
|
||||
pub(crate) struct RedisBreaker(BreakerConfig);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for RedisBreaker {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(breaker: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let (failure_threshold, recovery_timeout, timeout_min_duration): (u32, f64, f64) =
|
||||
breaker.extract()?;
|
||||
let seconds = |seconds: f64| {
|
||||
Duration::try_from_secs_f64(seconds)
|
||||
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))
|
||||
};
|
||||
Ok(Self(BreakerConfig {
|
||||
failure_threshold,
|
||||
recovery_timeout: seconds(recovery_timeout)?,
|
||||
timeout_min_duration: seconds(timeout_min_duration)?,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
|
||||
value.is_instance_of::<PyString>()
|
||||
|| value.is_instance_of::<PyBool>()
|
||||
|| value.is_instance_of::<PyInt>()
|
||||
|| value.is_instance_of::<PyFloat>()
|
||||
|| value.is_instance_of::<PyList>()
|
||||
|| value.is_instance_of::<PyTuple>()
|
||||
|| value.is_instance_of::<PyDict>()
|
||||
}
|
||||
|
||||
/// A client the process built, held weakly so a discarded `Router` drops its client and provider.
|
||||
struct Shared {
|
||||
resolved: Resolution,
|
||||
client: WeakRedisClient,
|
||||
}
|
||||
|
||||
static CLIENTS: LazyLock<Mutex<Vec<Shared>>> = LazyLock::new(|| Mutex::new(Vec::new()));
|
||||
|
||||
fn live_clients(py: Python<'_>) -> std::sync::MutexGuard<'static, Vec<Shared>> {
|
||||
let mut clients = CLIENTS
|
||||
.lock_py_attached(py)
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
clients.retain(|shared| shared.client.upgrade().is_some());
|
||||
clients
|
||||
}
|
||||
|
||||
/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool
|
||||
/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on
|
||||
/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error.
|
||||
fn shared_client(
|
||||
py: Python<'_>,
|
||||
spec: RedisSpec,
|
||||
namespace: Option<String>,
|
||||
breaker: Option<RedisBreaker>,
|
||||
) -> PyResult<Option<RedisClient>> {
|
||||
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
|
||||
let RedisSpec {
|
||||
kwargs,
|
||||
credential_provider,
|
||||
} = spec;
|
||||
// The registry entry's client keeps the provider alive, so no live entry shares its address
|
||||
let handle = credential_provider
|
||||
.as_ref()
|
||||
.map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64));
|
||||
let logger = crate::logger::capture(py);
|
||||
let environ = py.import("os")?.getattr("environ")?;
|
||||
let resolved = match logger.scope(|| resolve(&kwargs, handle, &environ))? {
|
||||
Ok(resolution) => {
|
||||
logger.scope(|| {
|
||||
for ignored in &resolution.ignored {
|
||||
litellm_tracing::warn!(
|
||||
"Redis setting {} is not supported by the Rust client and is ignored: {}",
|
||||
ignored.setting,
|
||||
ignored.consequence
|
||||
);
|
||||
}
|
||||
});
|
||||
Resolution {
|
||||
namespace,
|
||||
..resolution
|
||||
}
|
||||
}
|
||||
Err(error) => return unresolved(py, &error),
|
||||
};
|
||||
let runtime = litellm_host_python::runtime_handle()?;
|
||||
let events: Arc<dyn RedisEvents> = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?);
|
||||
let host = HostServices {
|
||||
events: Some(events),
|
||||
http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
|
||||
environment,
|
||||
credential_provider: credential_provider
|
||||
.as_ref()
|
||||
.map(|provider| -> PyResult<Arc<dyn CredentialSource>> {
|
||||
Ok(Arc::new(PythonCredentialSource::new(
|
||||
provider.clone_ref(py),
|
||||
PythonContext::capture(py)?,
|
||||
)))
|
||||
})
|
||||
.transpose()?,
|
||||
};
|
||||
let mut clients = live_clients(py);
|
||||
if let Some(client) = clients
|
||||
.iter()
|
||||
.filter(|shared| shared.resolved == resolved)
|
||||
.find_map(|shared| shared.client.upgrade())
|
||||
{
|
||||
return Ok(Some(client));
|
||||
}
|
||||
let connected = {
|
||||
let _runtime = runtime.enter();
|
||||
litellm_redis_connect::connect(resolved.clone(), host, breaker.map(|breaker| breaker.0))
|
||||
};
|
||||
let client = match connected {
|
||||
Ok(client) => client,
|
||||
Err(error) => {
|
||||
drop(clients);
|
||||
return Ok(unconnected(py, &error));
|
||||
}
|
||||
};
|
||||
clients.push(Shared {
|
||||
resolved,
|
||||
client: client.downgrade(),
|
||||
});
|
||||
Ok(Some(client))
|
||||
}
|
||||
|
||||
/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store
|
||||
/// never sees the redis-py kwargs or the namespace.
|
||||
#[pyclass(frozen, name = "NativeRedisClient")]
|
||||
pub(crate) struct NativeRedisClient {
|
||||
pub(crate) client: RedisClient,
|
||||
/// `RedisCache.default_ttl`, the TTL a store's write takes when its caller gives none.
|
||||
pub(crate) default_ttl: Duration,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRedisClient {
|
||||
/// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory.
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (redis_kwargs, *, default_ttl_seconds, namespace = None, breaker = None))]
|
||||
fn connect(
|
||||
py: Python<'_>,
|
||||
redis_kwargs: RedisSpec,
|
||||
default_ttl_seconds: f64,
|
||||
namespace: Option<String>,
|
||||
breaker: Option<RedisBreaker>,
|
||||
) -> PyResult<Option<Self>> {
|
||||
let default_ttl = number::seconds(default_ttl_seconds).map_err(seconds_error)?;
|
||||
Ok(
|
||||
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
|
||||
client,
|
||||
default_ttl,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn default_ttl_seconds(&self) -> f64 {
|
||||
self.default_ttl.as_secs_f64()
|
||||
}
|
||||
}
|
||||
|
||||
/// The breaker state of every native Redis client this process built, so the Prometheus state
|
||||
/// gauge counts them next to the Python `RedisCache` breakers.
|
||||
#[pyfunction]
|
||||
pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> {
|
||||
live_clients(py)
|
||||
.iter()
|
||||
.filter_map(|shared| shared.client.upgrade())
|
||||
.filter_map(|client| client.breaker_state())
|
||||
.map(|state| state_label(&state))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth
|
||||
/// reads `os.environ`.
|
||||
fn resolve(
|
||||
kwargs: &Kwargs,
|
||||
provider: Option<ProviderHandle>,
|
||||
environ: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Result<Resolution, litellm_redis_resolve::Error>> {
|
||||
let redis_username: Option<String> = environ
|
||||
.call_method1("get", ("REDIS_USERNAME",))?
|
||||
.extract()?;
|
||||
Ok(litellm_redis_resolve::interpret(
|
||||
kwargs,
|
||||
provider,
|
||||
redis_username.as_deref(),
|
||||
))
|
||||
}
|
||||
|
||||
fn unresolved(
|
||||
py: Python<'_>,
|
||||
error: &litellm_redis_resolve::Error,
|
||||
) -> PyResult<Option<RedisClient>> {
|
||||
let message = || litellm_error::Report::new(error).to_string();
|
||||
match error.raised_building_the_python_client() {
|
||||
Some(PythonException::ValueError) => Err(PyValueError::new_err(message())),
|
||||
Some(PythonException::TypeError) => Err(PyTypeError::new_err(message())),
|
||||
Some(PythonException::RedisClusterException) => {
|
||||
Err(RedisClusterException::new_err(message()))
|
||||
}
|
||||
None => Ok(memory_only(py, error)),
|
||||
}
|
||||
}
|
||||
|
||||
// redis-py reads TLS files and fetches credentials only once a command connects, so none of these fails building
|
||||
// its client
|
||||
fn unconnected(py: Python<'_>, error: &litellm_redis_connect::Error) -> Option<RedisClient> {
|
||||
match error {
|
||||
litellm_redis_connect::Error::Connect(_)
|
||||
| litellm_redis_connect::Error::NoHostCredentialProvider
|
||||
| litellm_redis_connect::Error::ProcessEnvironmentRequired { .. }
|
||||
| litellm_redis_connect::Error::Credential { .. }
|
||||
| litellm_redis_connect::Error::CredentialTimeout { .. }
|
||||
| litellm_redis_connect::Error::ExpiredCredential
|
||||
| litellm_redis_connect::Error::MissingCredential { .. } => memory_only(py, error),
|
||||
}
|
||||
}
|
||||
|
||||
fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option<RedisClient> {
|
||||
let report = litellm_error::Report::new(error).to_string();
|
||||
crate::logger::capture(py).scope(|| {
|
||||
litellm_tracing::error!(
|
||||
"Redis is not usable from the Rust stores, their state stays in memory: {report}"
|
||||
)
|
||||
});
|
||||
None
|
||||
}
|
||||
|
||||
fn seconds_error(error: SecondsError) -> PyErr {
|
||||
match error {
|
||||
SecondsError::Overflow => PyOverflowError::new_err(error.to_string()),
|
||||
SecondsError::NotANumber => PyValueError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::ffi::CStr;
|
||||
|
||||
use litellm_redis_types::Auth;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn kwargs(fields: serde_json::Value) -> Kwargs {
|
||||
let fields: BTreeMap<String, serde_json::Value> = serde_json::from_value(fields).unwrap();
|
||||
litellm_redis_resolve::kwargs(fields)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::set(c"environ = {'REDIS_USERNAME': 'azure-user'}", Some("azure-user"))]
|
||||
#[case::empty_is_unset(c"environ = {'REDIS_USERNAME': ''}", None)]
|
||||
fn azure_auth_takes_its_username_from_the_environ_like_python(
|
||||
#[case] source: &CStr,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(source, Some(&locals), Some(&locals)).unwrap();
|
||||
let environ = locals.get_item("environ").unwrap().unwrap();
|
||||
let kwargs = kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"}));
|
||||
|
||||
let resolved = resolve(&kwargs, None, &environ).unwrap().unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected),
|
||||
"{:?}",
|
||||
resolved.auth
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_target(json!({"port": 6379}), Some("ValueError"))]
|
||||
#[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))]
|
||||
#[case::fails_only_once_a_command_connects(
|
||||
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": "None"}),
|
||||
None
|
||||
)]
|
||||
fn a_config_raises_where_redis_py_raises_building_its_client(
|
||||
#[case] fields: serde_json::Value,
|
||||
#[case] raised: Option<&str>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = litellm_redis_resolve::interpret(&kwargs(fields), None, None).unwrap_err();
|
||||
|
||||
let outcome = unresolved(py, &error)
|
||||
.map(|client| client.is_some())
|
||||
.map_err(|error| error.get_type(py).name().unwrap().to_string());
|
||||
|
||||
assert_eq!(
|
||||
outcome,
|
||||
raised.map_or(Ok(false), |class| Err(class.to_owned()))
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::overflow(f64::INFINITY, "OverflowError")]
|
||||
#[case::not_a_number(f64::NAN, "ValueError")]
|
||||
fn a_default_ttl_python_cannot_make_an_int_raises_like_int(
|
||||
#[case] seconds: f64,
|
||||
#[case] raised: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = number::seconds(seconds).map_err(seconds_error).unwrap_err();
|
||||
|
||||
assert_eq!(error.get_type(py).name().unwrap().to_string(), raised);
|
||||
});
|
||||
}
|
||||
}
|
||||
391
litellm-rust/crates/python-bridge/src/redis/events.rs
Normal file
391
litellm-rust/crates/python-bridge/src/redis/events.rs
Normal file
|
|
@ -0,0 +1,391 @@
|
|||
use std::{
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::{State, Transition};
|
||||
use litellm_host_python::{PythonContext, attach_blocking};
|
||||
use litellm_redis_client::{
|
||||
CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged,
|
||||
};
|
||||
use pyo3::{prelude::*, types::PyModule};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
/// The asyncio loop a Python caller ran on. `RedisCache` scheduled its service events there with
|
||||
/// `asyncio.create_task`, so the native events go there too instead of a new loop per event.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct EventLoop(Arc<Py<PyAny>>);
|
||||
|
||||
impl EventLoop {
|
||||
pub(crate) fn running(py: Python<'_>) -> Option<Self> {
|
||||
pyo3_async_runtimes::get_running_loop(py)
|
||||
.ok()
|
||||
.map(|event_loop| Self(Arc::new(event_loop.unbind())))
|
||||
}
|
||||
}
|
||||
|
||||
/// The Python caller of one store operation. Its context carries `with_service_target` and the
|
||||
/// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep
|
||||
/// their series.
|
||||
pub(crate) struct RedisCall {
|
||||
pub(crate) context: PythonContext,
|
||||
pub(crate) event_loop: Option<EventLoop>,
|
||||
pub(crate) call_type: &'static str,
|
||||
}
|
||||
|
||||
impl RedisCall {
|
||||
pub(crate) fn tag(self) -> CallTag {
|
||||
Arc::new(self)
|
||||
}
|
||||
|
||||
pub(crate) async fn scope<F: Future>(self, operation: F) -> F::Output {
|
||||
tagged(self.tag(), operation).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Reports every native Redis call and breaker transition to `ServiceLogging`.
|
||||
pub(crate) struct ServiceLoggingEvents {
|
||||
record_call: Arc<Py<PyAny>>,
|
||||
record_breaker: Arc<Py<PyAny>>,
|
||||
context: PythonContext,
|
||||
/// The loop the client was built on, for calls whose caller had none running.
|
||||
event_loop: Option<EventLoop>,
|
||||
runtime: Handle,
|
||||
}
|
||||
|
||||
impl ServiceLoggingEvents {
|
||||
pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
|
||||
const MODULE: &str = "litellm.rust_bridge.host.redis_events";
|
||||
Self::from_module(py, &py.import(MODULE)?, runtime)
|
||||
}
|
||||
|
||||
fn from_module(
|
||||
py: Python<'_>,
|
||||
module: &Bound<'_, PyModule>,
|
||||
runtime: Handle,
|
||||
) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
record_call: Arc::new(module.getattr("record_redis_call")?.unbind()),
|
||||
record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()),
|
||||
context: PythonContext::capture(py)?,
|
||||
event_loop: EventLoop::running(py),
|
||||
runtime,
|
||||
})
|
||||
}
|
||||
|
||||
fn report(&self, context: PythonContext, record: Arc<Py<PyAny>>, args: ReportArgs) {
|
||||
self.runtime.spawn(async move {
|
||||
let reported = attach_blocking(context, move |py| match args {
|
||||
ReportArgs::Call(call) => record
|
||||
.bind(py)
|
||||
.call1((
|
||||
call.event_loop
|
||||
.as_ref()
|
||||
.map(|event_loop| event_loop.0.bind(py)),
|
||||
call.call_type,
|
||||
call.start_time,
|
||||
call.end_time,
|
||||
call.failure,
|
||||
))
|
||||
.map(drop),
|
||||
ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop),
|
||||
})
|
||||
.await
|
||||
.and_then(|reported| reported);
|
||||
if let Err(error) = reported {
|
||||
tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// One call as `record_redis_call` takes it, its times in `time.time()` seconds.
|
||||
struct CallReport {
|
||||
event_loop: Option<EventLoop>,
|
||||
call_type: &'static str,
|
||||
start_time: f64,
|
||||
end_time: f64,
|
||||
failure: Option<&'static str>,
|
||||
}
|
||||
|
||||
enum ReportArgs {
|
||||
Call(CallReport),
|
||||
Breaker(&'static str),
|
||||
}
|
||||
|
||||
/// The wall-clock start and end of a call that just finished after `duration`.
|
||||
fn wall_clock_span(duration: Duration) -> (f64, f64) {
|
||||
let end = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default();
|
||||
(
|
||||
end.saturating_sub(duration).as_secs_f64(),
|
||||
end.as_secs_f64(),
|
||||
)
|
||||
}
|
||||
|
||||
fn operation_call_type(operation: Operation) -> &'static str {
|
||||
match operation {
|
||||
Operation::Command(name) => name,
|
||||
Operation::Script => "eval",
|
||||
Operation::Pipeline { .. } => "pipeline",
|
||||
}
|
||||
}
|
||||
|
||||
fn failure_label(outcome: Outcome) -> Option<&'static str> {
|
||||
match outcome {
|
||||
Outcome::Succeeded => None,
|
||||
Outcome::Failed(FailureKind::Timeout) => Some("timeout"),
|
||||
Outcome::Failed(FailureKind::Transport) => Some("transport"),
|
||||
Outcome::Failed(FailureKind::Server) => Some("server"),
|
||||
Outcome::Failed(FailureKind::Decode) => Some("decode"),
|
||||
Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn state_label(state: &State) -> &'static str {
|
||||
match state {
|
||||
State::Closed { .. } => "closed",
|
||||
State::Open { .. } => "open",
|
||||
State::HalfOpen => "half_open",
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisEvents for ServiceLoggingEvents {
|
||||
fn call(&self, event: &CallEvent) {
|
||||
let scoped = current_tag().and_then(|tag| tag.downcast::<RedisCall>().ok());
|
||||
let (context, event_loop, call_type) = match scoped {
|
||||
Some(call) => (
|
||||
call.context.clone(),
|
||||
call.event_loop.clone().or_else(|| self.event_loop.clone()),
|
||||
call.call_type,
|
||||
),
|
||||
None => (
|
||||
self.context.clone(),
|
||||
self.event_loop.clone(),
|
||||
operation_call_type(event.operation),
|
||||
),
|
||||
};
|
||||
let (start_time, end_time) = wall_clock_span(event.duration);
|
||||
self.report(
|
||||
context,
|
||||
self.record_call.clone(),
|
||||
ReportArgs::Call(CallReport {
|
||||
event_loop,
|
||||
call_type,
|
||||
start_time,
|
||||
end_time,
|
||||
failure: failure_label(event.outcome),
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
fn breaker(&self, transition: &Transition<State>) {
|
||||
self.report(
|
||||
self.context.clone(),
|
||||
self.record_breaker.clone(),
|
||||
ReportArgs::Breaker(state_label(&transition.to)),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_host_python::PythonContext;
|
||||
use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents};
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyCFunction, PyModule, PyString},
|
||||
};
|
||||
use rstest::rstest;
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
use super::{EventLoop, RedisCall, ServiceLoggingEvents};
|
||||
|
||||
const RECORDER: &str = "
|
||||
import contextvars, time
|
||||
target = contextvars.ContextVar('target', default=None)
|
||||
seen = []
|
||||
def record_redis_call(event_loop, call_type, start_time, end_time, failure):
|
||||
just_now = abs(time.time() - end_time) < 60
|
||||
seen.append((event_loop, call_type, round(end_time - start_time, 3), just_now, failure, target.get()))
|
||||
def record_redis_breaker(state):
|
||||
seen.append(('breaker', state))
|
||||
";
|
||||
|
||||
fn stand_in_loop(py: Python<'_>, name: &str) -> EventLoop {
|
||||
EventLoop(Arc::new(PyString::new(py, name).into_any().unbind()))
|
||||
}
|
||||
|
||||
fn recorder(name: &str, client_loop: Option<&str>) -> (ServiceLoggingEvents, Py<PyModule>) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap();
|
||||
let module =
|
||||
PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name)
|
||||
.unwrap();
|
||||
let events = ServiceLoggingEvents::from_module(py, &module, Handle::current()).unwrap();
|
||||
(
|
||||
ServiceLoggingEvents {
|
||||
event_loop: client_loop.map(|client_loop| stand_in_loop(py, client_loop)),
|
||||
..events
|
||||
},
|
||||
module.unbind(),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
async fn seen(module: &Py<PyModule>, count: usize) -> String {
|
||||
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
|
||||
loop {
|
||||
let seen: (usize, String) = Python::attach(|py| {
|
||||
let seen = module.bind(py).getattr("seen").unwrap();
|
||||
(seen.len().unwrap(), seen.repr().unwrap().to_string())
|
||||
});
|
||||
if seen.0 >= count || tokio::time::Instant::now() > deadline {
|
||||
return seen.1;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn scoped_call(namespace: &Py<PyModule>, caller_loop: Option<&str>) -> RedisCall {
|
||||
Python::attach(|py| {
|
||||
namespace
|
||||
.bind(py)
|
||||
.getattr("target")
|
||||
.unwrap()
|
||||
.call_method1("set", ("spend_counter",))
|
||||
.unwrap();
|
||||
RedisCall {
|
||||
context: PythonContext::capture(py).unwrap(),
|
||||
event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)),
|
||||
call_type: "async_increment",
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::on_the_callers_loop(Some("caller loop"), "caller loop")]
|
||||
#[case::on_the_clients_loop_when_the_caller_had_none(None, "client loop")]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn a_scoped_call_reports_the_store_labels_and_its_times_in_the_callers_context(
|
||||
#[case] caller_loop: Option<&str>,
|
||||
#[case] expected_loop: &str,
|
||||
) {
|
||||
let (events, namespace) = recorder(
|
||||
expected_loop.replace(' ', "_").as_str(),
|
||||
Some("client loop"),
|
||||
);
|
||||
let call = scoped_call(&namespace, caller_loop);
|
||||
let event = CallEvent {
|
||||
operation: Operation::Command("INCRBYFLOAT"),
|
||||
duration: Duration::from_millis(250),
|
||||
outcome: Outcome::Failed(FailureKind::Timeout),
|
||||
};
|
||||
|
||||
call.scope(async { events.call(&event) }).await;
|
||||
|
||||
assert_eq!(
|
||||
seen(&namespace, 1).await,
|
||||
format!(
|
||||
"[('{expected_loop}', 'async_increment', 0.25, True, 'timeout', 'spend_counter')]"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")]
|
||||
#[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn breaker_transitions_report_python_state_names(
|
||||
#[case] to: litellm_circuit_breaker::State,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let (events, namespace) = recorder(expected, None);
|
||||
|
||||
events.breaker(&litellm_circuit_breaker::Transition {
|
||||
from: litellm_circuit_breaker::State::HalfOpen,
|
||||
to,
|
||||
generation: 1,
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
seen(&namespace, 1).await,
|
||||
format!("[('breaker', '{expected}')]")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() {
|
||||
let (events, namespace) = recorder("plain_thread", None);
|
||||
|
||||
std::thread::spawn(move || {
|
||||
events.breaker(&litellm_circuit_breaker::Transition {
|
||||
from: litellm_circuit_breaker::State::HalfOpen,
|
||||
to: litellm_circuit_breaker::State::HalfOpen,
|
||||
generation: 1,
|
||||
});
|
||||
})
|
||||
.join()
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn an_unscoped_call_falls_back_to_the_command_name_and_the_clients_loop() {
|
||||
let (events, namespace) = recorder("unscoped", None);
|
||||
|
||||
events.call(&CallEvent {
|
||||
operation: Operation::Pipeline { ops: 3 },
|
||||
duration: Duration::from_millis(5),
|
||||
outcome: Outcome::Succeeded,
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
seen(&namespace, 1).await,
|
||||
"[(None, 'pipeline', 0.005, True, None, None)]"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn running_is_the_loop_the_caller_runs_on_and_none_outside_one() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let module = PyModule::from_code(
|
||||
py,
|
||||
c"import asyncio\nasync def probe_inside(probe):\n return probe() is asyncio.get_running_loop()\n",
|
||||
c"event_loop_probe",
|
||||
c"event_loop_probe",
|
||||
)
|
||||
.unwrap();
|
||||
let probe = PyCFunction::new_closure(py, None, None, |_, _| {
|
||||
Python::attach(|py| -> PyResult<Option<Py<PyAny>>> {
|
||||
Ok(EventLoop::running(py).map(|event_loop| event_loop.0.clone_ref(py)))
|
||||
})
|
||||
})
|
||||
.unwrap();
|
||||
let coroutine = module
|
||||
.getattr("probe_inside")
|
||||
.unwrap()
|
||||
.call1((probe,))
|
||||
.unwrap();
|
||||
|
||||
let inside: bool = py
|
||||
.import("asyncio")
|
||||
.unwrap()
|
||||
.call_method1("run", (coroutine,))
|
||||
.unwrap()
|
||||
.extract()
|
||||
.unwrap();
|
||||
|
||||
assert!(inside);
|
||||
assert!(EventLoop::running(py).is_none());
|
||||
});
|
||||
}
|
||||
}
|
||||
11
litellm-rust/crates/python-bridge/src/redis/mod.rs
Normal file
11
litellm-rust/crates/python-bridge/src/redis/mod.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
#[expect(
|
||||
dead_code,
|
||||
reason = "the native stores that use these helpers land in the next PR"
|
||||
)]
|
||||
pub(crate) mod client;
|
||||
mod events;
|
||||
#[expect(
|
||||
dead_code,
|
||||
reason = "the native stores that use these helpers land in the next PR"
|
||||
)]
|
||||
pub(crate) mod store;
|
||||
248
litellm-rust/crates/python-bridge/src/redis/store/mod.rs
Normal file
248
litellm-rust/crates/python-bridge/src/redis/store/mod.rs
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
pub(crate) mod request;
|
||||
mod value;
|
||||
|
||||
use value::PyValue;
|
||||
|
||||
use std::{
|
||||
pin::Pin,
|
||||
sync::{Mutex, PoisonError},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_host_python::PythonContext;
|
||||
use litellm_python_compat::number::{SecondsError, seconds};
|
||||
use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote};
|
||||
use pyo3::{
|
||||
IntoPyObjectExt,
|
||||
exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{PyBool, PyFloat, PyInt},
|
||||
};
|
||||
|
||||
use crate::execution::{run_async_value, run_sync_value};
|
||||
use crate::redis::{
|
||||
client::NativeRedisClient,
|
||||
events::{EventLoop, RedisCall},
|
||||
};
|
||||
|
||||
/// The labels `RedisCache` reports for this operation, in the caller's context so its
|
||||
/// `service_target` reaches the service event, dispatched on the loop the caller runs on.
|
||||
fn call(py: Python<'_>, call_type: &'static str) -> PyResult<RedisCall> {
|
||||
Ok(RedisCall {
|
||||
context: PythonContext::capture(py)?,
|
||||
event_loop: EventLoop::running(py),
|
||||
call_type,
|
||||
})
|
||||
}
|
||||
|
||||
/// Runs one Redis operation as a Python awaitable, reported to `ServiceLogging` as `call_type`.
|
||||
fn redis_async<'py, T, E: 'static>(
|
||||
py: Python<'py>,
|
||||
call_type: &'static str,
|
||||
error: fn(E) -> PyErr,
|
||||
operation: impl Future<Output = Result<T, E>> + Send + 'static,
|
||||
) -> PyResult<Bound<'py, PyAny>>
|
||||
where
|
||||
T: for<'a> IntoPyObject<'a> + Send + 'static,
|
||||
{
|
||||
let scope = call(py, call_type)?;
|
||||
run_async_value(
|
||||
py,
|
||||
async move { scope.scope(operation).await.map_err(error) },
|
||||
)
|
||||
}
|
||||
|
||||
/// Runs one Redis operation with the GIL released, for the router's sync paths.
|
||||
fn redis_sync<T, E: 'static>(
|
||||
py: Python<'_>,
|
||||
call_type: &'static str,
|
||||
error: fn(E) -> PyErr,
|
||||
operation: impl Future<Output = Result<T, E>> + Send + 'static,
|
||||
) -> PyResult<T>
|
||||
where
|
||||
T: Send + 'static,
|
||||
{
|
||||
let scope = call(py, call_type)?;
|
||||
run_sync_value(
|
||||
py,
|
||||
async move { scope.scope(operation).await.map_err(error) },
|
||||
)
|
||||
}
|
||||
|
||||
/// A TTL in seconds, raising what `int(seconds)` raises for infinity and NaN.
|
||||
fn duration(value: f64) -> PyResult<Duration> {
|
||||
seconds(value).map_err(|error| match error {
|
||||
SecondsError::Overflow => PyOverflowError::new_err(error.to_string()),
|
||||
SecondsError::NotANumber => PyValueError::new_err(error.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
fn ttl(value: Option<f64>) -> PyResult<Option<Duration>> {
|
||||
value.map(duration).transpose()
|
||||
}
|
||||
|
||||
fn remote(client: &NativeRedisClient) -> Remote {
|
||||
Remote {
|
||||
client: client.client.clone(),
|
||||
default_ttl: client.default_ttl,
|
||||
}
|
||||
}
|
||||
|
||||
type Resolve = Pin<Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>>;
|
||||
|
||||
/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone
|
||||
/// out and answers the call, once.
|
||||
#[pyclass(frozen, name = "NativePending")]
|
||||
pub(crate) struct NativePending {
|
||||
resolve: Mutex<Option<Resolve>>,
|
||||
call_type: &'static str,
|
||||
}
|
||||
|
||||
impl NativePending {
|
||||
pub(crate) fn new<T, E: 'static>(
|
||||
call_type: &'static str,
|
||||
error: fn(E) -> PyErr,
|
||||
reply: impl Future<Output = Result<T, E>> + Send + 'static,
|
||||
) -> Self
|
||||
where
|
||||
T: for<'a> IntoPyObject<'a> + Send + 'static,
|
||||
{
|
||||
let resolve = async move {
|
||||
let reply = reply.await.map_err(error)?;
|
||||
Python::attach(|py| reply.into_py_any(py))
|
||||
};
|
||||
Self {
|
||||
resolve: Mutex::new(Some(Box::pin(resolve))),
|
||||
call_type,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativePending {
|
||||
fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let resolve = self
|
||||
.resolve
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.take()
|
||||
.ok_or_else(|| PyRuntimeError::new_err("this Redis call was already resolved"))?;
|
||||
redis_async(py, self.call_type, std::convert::identity, resolve)
|
||||
}
|
||||
}
|
||||
|
||||
pyo3::create_exception!(
|
||||
_native,
|
||||
NativeRedisStoreError,
|
||||
pyo3::exceptions::PyException,
|
||||
"A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind."
|
||||
);
|
||||
|
||||
/// What every native store raises, its kind telling a caller which failures it may drop.
|
||||
fn kind_error(error: Error) -> PyErr {
|
||||
let kind: &'static str = error.kind().into();
|
||||
NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string()))
|
||||
}
|
||||
|
||||
/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it.
|
||||
fn pending<T, R>(
|
||||
call_type: &'static str,
|
||||
error: fn(Error) -> PyErr,
|
||||
declared: Pending<T>,
|
||||
reply: fn(T) -> R,
|
||||
) -> NativePending
|
||||
where
|
||||
T: Send + 'static,
|
||||
R: for<'a> IntoPyObject<'a> + Send + 'static,
|
||||
{
|
||||
NativePending::new(call_type, error, async move { declared.await.map(reply) })
|
||||
}
|
||||
|
||||
/// A Python `int` or `float` across the bridge, keeping which one it is.
|
||||
pub(crate) struct PyNumber(Number);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
if value.is_instance_of::<PyInt>() && !value.is_instance_of::<PyBool>() {
|
||||
return Ok(Self(Number::Int(value.extract()?)));
|
||||
}
|
||||
if value.is_instance_of::<PyFloat>() {
|
||||
return Ok(Self(Number::Float(value.extract()?)));
|
||||
}
|
||||
Err(PyTypeError::new_err("expected an int or a float"))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PyNumber {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, PyAny>;
|
||||
type Error = PyErr;
|
||||
|
||||
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
|
||||
Ok(match self.0 {
|
||||
Number::Int(value) => value.into_pyobject(py)?.into_any(),
|
||||
Number::Float(value) => value.into_pyobject(py)?.into_any(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue`
|
||||
/// so the bridge never names a field. A shape the store refuses is `ValueError`.
|
||||
pub(crate) struct PyData<T>(pub(crate) T);
|
||||
|
||||
impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData<T> {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let PyValue(value) = value.extract()?;
|
||||
T::from_value(&value)
|
||||
.map(Self)
|
||||
.map_err(|error| PyValueError::new_err(error.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py, T: IntoValue> IntoPyObject<'py> for PyData<T> {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, PyAny>;
|
||||
type Error = PyErr;
|
||||
|
||||
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
|
||||
PyValue(self.0.into_value()).into_pyobject(py)
|
||||
}
|
||||
}
|
||||
|
||||
fn unwrap_all<T>(values: Vec<PyData<T>>) -> Vec<T> {
|
||||
values.into_iter().map(|value| value.0).collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::ttl;
|
||||
|
||||
#[rstest]
|
||||
#[case::infinity(f64::INFINITY, "OverflowError")]
|
||||
#[case::nan(f64::NAN, "ValueError")]
|
||||
fn a_ttl_int_refuses_raises_what_int_raises(#[case] seconds: f64, #[case] raised: &str) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = ttl(Some(seconds)).unwrap_err();
|
||||
|
||||
assert_eq!(error.get_type(py).name().unwrap().to_string(), raised);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::fraction(Some(1.5), Some(Duration::from_millis(1500)))]
|
||||
#[case::negative_clamps_to_zero(Some(-1.0), Some(Duration::ZERO))]
|
||||
#[case::absent(None, None)]
|
||||
fn a_finite_ttl_converts(#[case] seconds: Option<f64>, #[case] expected: Option<Duration>) {
|
||||
assert_eq!(ttl(seconds).unwrap(), expected);
|
||||
}
|
||||
}
|
||||
60
litellm-rust/crates/python-bridge/src/redis/store/request.rs
Normal file
60
litellm-rust/crates/python-bridge/src/redis/store/request.rs
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
use litellm_redis_store::{RequestRedis, Send};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::{call, duration};
|
||||
use crate::execution::run_async_value;
|
||||
|
||||
/// One request's native Redis batches. Python builds one per request scope and passes it to each
|
||||
/// request-aware store call.
|
||||
#[pyclass(frozen, name = "NativeRequestRedis")]
|
||||
pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRequestRedis {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult<Self> {
|
||||
let runtime = litellm_host_python::runtime_handle()?;
|
||||
let deadline = duration(post_call_deadline_seconds)?;
|
||||
let deadline_tag = call(py, "post_call_redis_batch")?.tag();
|
||||
Ok(Self(crate::logger::capture(py).scope(|| {
|
||||
RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag))
|
||||
})))
|
||||
}
|
||||
|
||||
fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request = self.0.clone();
|
||||
let scope = call(py, "request_redis_batch")?;
|
||||
run_async_value(py, async move {
|
||||
scope.scope(request.flush_pre()).await;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn flush_post<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let request = self.0.clone();
|
||||
let scope = call(py, "post_call_redis_batch")?;
|
||||
run_async_value(py, async move {
|
||||
scope.scope(request.flush_post()).await;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it
|
||||
/// needs, so there is no phase string and no phase without a request.
|
||||
#[pyclass(frozen, name = "NativeSend")]
|
||||
pub(crate) enum NativeSend {
|
||||
Now(),
|
||||
PreCall(Py<NativeRequestRedis>),
|
||||
PostCall(Py<NativeRequestRedis>),
|
||||
}
|
||||
|
||||
impl NativeSend {
|
||||
pub(crate) fn send(&self) -> Send<'_> {
|
||||
match self {
|
||||
Self::Now() => Send::Now,
|
||||
Self::PreCall(request) => Send::PreCall(&request.get().0),
|
||||
Self::PostCall(request) => Send::PostCall(&request.get().0),
|
||||
}
|
||||
}
|
||||
}
|
||||
125
litellm-rust/crates/python-bridge/src/redis/store/value.rs
Normal file
125
litellm-rust/crates/python-bridge/src/redis/store/value.rs
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
use litellm_python_compat::{BigInt, MAX_DEPTH, Value};
|
||||
use pyo3::{
|
||||
exceptions::{PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{
|
||||
PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString,
|
||||
PyTuple,
|
||||
},
|
||||
};
|
||||
|
||||
/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge.
|
||||
/// Anything else, a live object included, is refused with `TypeError`.
|
||||
pub(crate) struct PyValue(pub(crate) Value);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyValue {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
to_value(&value, 0).map(Self)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PyValue {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, PyAny>;
|
||||
type Error = PyErr;
|
||||
|
||||
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
|
||||
to_python(py, self.0)
|
||||
}
|
||||
}
|
||||
|
||||
fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult<Value> {
|
||||
if depth > MAX_DEPTH {
|
||||
return Err(PyValueError::new_err("value is nested too deeply"));
|
||||
}
|
||||
let nested = |items: Bound<'_, PyAny>| -> PyResult<Vec<Value>> {
|
||||
items
|
||||
.try_iter()?
|
||||
.map(|item| to_value(&item?, depth + 1))
|
||||
.collect()
|
||||
};
|
||||
if value.is_none() {
|
||||
return Ok(Value::None);
|
||||
}
|
||||
if let Ok(flag) = value.cast::<PyBool>() {
|
||||
return Ok(Value::Bool(flag.is_true()));
|
||||
}
|
||||
if value.is_instance_of::<PyInt>() {
|
||||
return Ok(Value::Int(match value.extract::<i64>() {
|
||||
Ok(small) => small.into(),
|
||||
Err(_) => value
|
||||
.str()?
|
||||
.to_str()?
|
||||
.parse::<BigInt>()
|
||||
.map_err(|error| PyValueError::new_err(error.to_string()))?,
|
||||
}));
|
||||
}
|
||||
if let Ok(number) = value.cast::<PyFloat>() {
|
||||
return Ok(Value::Float(number.value()));
|
||||
}
|
||||
if let Ok(text) = value.cast::<PyString>() {
|
||||
return Ok(Value::Str(text.to_str()?.to_owned()));
|
||||
}
|
||||
if let Ok(bytes) = value.cast::<PyBytes>() {
|
||||
return Ok(Value::Bytes(bytes.as_bytes().to_vec()));
|
||||
}
|
||||
if let Ok(complex) = value.cast::<PyComplex>() {
|
||||
return Ok(Value::Complex {
|
||||
re: complex.real(),
|
||||
im: complex.imag(),
|
||||
});
|
||||
}
|
||||
if let Ok(entries) = value.cast::<PyDict>() {
|
||||
return entries
|
||||
.iter()
|
||||
.map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?)))
|
||||
.collect::<PyResult<_>>()
|
||||
.map(Value::Dict);
|
||||
}
|
||||
if value.is_instance_of::<PyList>() {
|
||||
return nested(value.clone()).map(Value::List);
|
||||
}
|
||||
if value.is_instance_of::<PyTuple>() {
|
||||
return nested(value.clone()).map(Value::Tuple);
|
||||
}
|
||||
if value.is_instance_of::<PySet>() || value.is_instance_of::<PyFrozenSet>() {
|
||||
return nested(value.clone()).map(Value::Set);
|
||||
}
|
||||
Err(PyTypeError::new_err(format!(
|
||||
"{} is not a plain data value",
|
||||
value.get_type().name()?
|
||||
)))
|
||||
}
|
||||
|
||||
fn to_python(py: Python<'_>, value: Value) -> PyResult<Bound<'_, PyAny>> {
|
||||
let all = |items: Vec<Value>| -> PyResult<Vec<Bound<'_, PyAny>>> {
|
||||
items.into_iter().map(|item| to_python(py, item)).collect()
|
||||
};
|
||||
Ok(match value {
|
||||
Value::None => py.None().into_bound(py),
|
||||
Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(),
|
||||
Value::Int(number) => match i64::try_from(&number) {
|
||||
Ok(small) => small.into_pyobject(py)?.into_any(),
|
||||
Err(_) => py
|
||||
.get_type::<PyInt>()
|
||||
.call1((number.to_string(),))?
|
||||
.into_any(),
|
||||
},
|
||||
Value::Float(number) => number.into_pyobject(py)?.into_any(),
|
||||
Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(),
|
||||
Value::Str(text) => text.into_pyobject(py)?.into_any(),
|
||||
Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(),
|
||||
Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(),
|
||||
Value::List(items) => PyList::new(py, all(items)?)?.into_any(),
|
||||
Value::Set(items) => PySet::new(py, all(items)?)?.into_any(),
|
||||
Value::Dict(entries) => {
|
||||
let dict = PyDict::new(py);
|
||||
for (key, item) in entries {
|
||||
dict.set_item(to_python(py, key)?, to_python(py, item)?)?;
|
||||
}
|
||||
dict.into_any()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -10,7 +10,7 @@
|
|||
use std::{hint::black_box, time::Duration};
|
||||
|
||||
use criterion::{
|
||||
BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
|
||||
BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
|
||||
measurement::WallTime,
|
||||
};
|
||||
use litellm_python_compat::{Value, json, literal::literal_eval, pickle, repr::repr};
|
||||
|
|
@ -72,14 +72,9 @@ fn formats(c: &mut Criterion) {
|
|||
let mut group = c.benchmark_group("cached_completion");
|
||||
text_formats(&mut group, &text, &value);
|
||||
binary_formats(&mut group, &value, &pickled);
|
||||
// `from_json` consumes its input, so each iteration gets a freshly parsed one.
|
||||
group.throughput(Throughput::Bytes(dumped.len() as u64));
|
||||
group.bench_function("from_json", |bencher| {
|
||||
bencher.iter_batched(
|
||||
|| serde_json::from_str::<serde_json::Value>(&dumped).expect("dumps output parses"),
|
||||
json::from_json,
|
||||
BatchSize::SmallInput,
|
||||
)
|
||||
group.bench_function("json_loads", |bencher| {
|
||||
bencher.iter(|| json::loads(black_box(&dumped)))
|
||||
});
|
||||
group.finish();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ use crate::MAX_DEPTH;
|
|||
pub enum Error {
|
||||
#[error("malformed Python literal at byte {0}")]
|
||||
InvalidLiteral(usize),
|
||||
#[error("malformed JSON at byte {0}")]
|
||||
InvalidJson(usize),
|
||||
#[error("unhashable type: '{0}'")]
|
||||
Unhashable(&'static str),
|
||||
#[error("value nests deeper than {MAX_DEPTH} levels")]
|
||||
|
|
|
|||
|
|
@ -1,14 +1,19 @@
|
|||
//! `json.dumps` with CPython's default options, and the JSON value `json.loads` returns.
|
||||
//! `json.dumps` with CPython's default options, and `json.loads` with its default decoder.
|
||||
//!
|
||||
//! Defaults are `ensure_ascii=True`, `allow_nan=True`, separators `(", ", ": ")`, and no key
|
||||
//! sorting. Python's `json.loads` is mapped by [`from_json`]; a `serde_json` number that does
|
||||
//! not fit `i64` or `u64` arrives as a float, where Python would keep an `int`.
|
||||
//! sorting on the way out, and `strict=True` with `NaN`, `Infinity` and `-Infinity` read as floats
|
||||
//! on the way in.
|
||||
|
||||
use std::fmt::Write;
|
||||
use std::{
|
||||
collections::{HashMap, hash_map::Entry},
|
||||
fmt::Write,
|
||||
str::FromStr,
|
||||
};
|
||||
|
||||
use num_bigint::BigInt;
|
||||
use serde_json::{Map, Number};
|
||||
|
||||
use crate::{Error, Value, repr::float_repr};
|
||||
use crate::{Error, MAX_DEPTH, Value, repr::float_repr};
|
||||
|
||||
/// `json.dumps(value)`.
|
||||
pub fn dumps(value: &Value) -> Result<String, Error> {
|
||||
|
|
@ -44,6 +49,260 @@ pub fn from_json(value: serde_json::Value) -> Value {
|
|||
}
|
||||
}
|
||||
|
||||
/// `json.loads(text)`. Integers keep full precision, and a repeated key keeps its first position
|
||||
/// with its last value, as `dict(pairs)` does. An escape that leaves a lone surrogate fails, since
|
||||
/// a Rust `String` cannot hold one.
|
||||
pub fn loads(text: &str) -> Result<Value, Error> {
|
||||
let mut decoder = Decoder {
|
||||
text,
|
||||
bytes: text.as_bytes(),
|
||||
pos: 0,
|
||||
depth: 0,
|
||||
};
|
||||
decoder.skip_whitespace();
|
||||
let value = decoder.value()?;
|
||||
decoder.skip_whitespace();
|
||||
if decoder.pos != decoder.bytes.len() {
|
||||
return Err(decoder.error());
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
/// CPython's `json.scanner` and `scanstring`, one pass over the bytes.
|
||||
struct Decoder<'a> {
|
||||
text: &'a str,
|
||||
bytes: &'a [u8],
|
||||
pos: usize,
|
||||
depth: usize,
|
||||
}
|
||||
|
||||
impl Decoder<'_> {
|
||||
fn error(&self) -> Error {
|
||||
Error::InvalidJson(self.pos)
|
||||
}
|
||||
|
||||
fn peek(&self) -> Option<u8> {
|
||||
self.bytes.get(self.pos).copied()
|
||||
}
|
||||
|
||||
fn eat(&mut self, token: &str) -> bool {
|
||||
let found = self.bytes[self.pos..].starts_with(token.as_bytes());
|
||||
if found {
|
||||
self.pos += token.len();
|
||||
}
|
||||
found
|
||||
}
|
||||
|
||||
fn expect(&mut self, token: &str) -> Result<(), Error> {
|
||||
if self.eat(token) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(self.error())
|
||||
}
|
||||
}
|
||||
|
||||
fn skip_whitespace(&mut self) {
|
||||
while matches!(self.peek(), Some(b' ' | b'\t' | b'\n' | b'\r')) {
|
||||
self.pos += 1;
|
||||
}
|
||||
}
|
||||
|
||||
fn value(&mut self) -> Result<Value, Error> {
|
||||
match self.peek() {
|
||||
Some(b'"') => self.string().map(Value::Str),
|
||||
Some(b'[') => self.nested(Self::array),
|
||||
Some(b'{') => self.nested(Self::object),
|
||||
_ if self.eat("null") => Ok(Value::None),
|
||||
_ if self.eat("true") => Ok(Value::Bool(true)),
|
||||
_ if self.eat("false") => Ok(Value::Bool(false)),
|
||||
_ if self.eat("NaN") => Ok(Value::Float(f64::NAN)),
|
||||
_ if self.eat("Infinity") => Ok(Value::Float(f64::INFINITY)),
|
||||
_ if self.eat("-Infinity") => Ok(Value::Float(f64::NEG_INFINITY)),
|
||||
_ => self.number(),
|
||||
}
|
||||
}
|
||||
|
||||
fn nested(&mut self, parse: fn(&mut Self) -> Result<Value, Error>) -> Result<Value, Error> {
|
||||
self.depth += 1;
|
||||
if self.depth > MAX_DEPTH {
|
||||
return Err(Error::TooDeep);
|
||||
}
|
||||
let value = parse(self)?;
|
||||
self.depth -= 1;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn array(&mut self) -> Result<Value, Error> {
|
||||
self.expect("[")?;
|
||||
self.skip_whitespace();
|
||||
let mut values = Vec::new();
|
||||
if self.eat("]") {
|
||||
return Ok(Value::List(values));
|
||||
}
|
||||
loop {
|
||||
values.push(self.value()?);
|
||||
self.skip_whitespace();
|
||||
if self.eat("]") {
|
||||
return Ok(Value::List(values));
|
||||
}
|
||||
self.expect(",")?;
|
||||
self.skip_whitespace();
|
||||
}
|
||||
}
|
||||
|
||||
fn object(&mut self) -> Result<Value, Error> {
|
||||
self.expect("{")?;
|
||||
self.skip_whitespace();
|
||||
let mut entries = Vec::new();
|
||||
let mut positions = HashMap::<_, usize>::new();
|
||||
if self.eat("}") {
|
||||
return Ok(Value::Dict(entries));
|
||||
}
|
||||
loop {
|
||||
if self.peek() != Some(b'"') {
|
||||
return Err(self.error());
|
||||
}
|
||||
let key = self.string()?;
|
||||
self.skip_whitespace();
|
||||
self.expect(":")?;
|
||||
self.skip_whitespace();
|
||||
let value = self.value()?;
|
||||
match positions.entry(key) {
|
||||
Entry::Occupied(position) => entries[*position.get()].1 = value,
|
||||
Entry::Vacant(position) => {
|
||||
entries.push((Value::Str(position.key().clone()), value));
|
||||
position.insert(entries.len() - 1);
|
||||
}
|
||||
}
|
||||
self.skip_whitespace();
|
||||
if self.eat("}") {
|
||||
return Ok(Value::Dict(entries));
|
||||
}
|
||||
self.expect(",")?;
|
||||
self.skip_whitespace();
|
||||
}
|
||||
}
|
||||
|
||||
/// `scanstring` with `strict=True`: control characters must be escaped.
|
||||
fn string(&mut self) -> Result<String, Error> {
|
||||
self.expect("\"")?;
|
||||
let mut out = String::new();
|
||||
loop {
|
||||
let start = self.pos;
|
||||
while self
|
||||
.peek()
|
||||
.is_some_and(|byte| byte != b'"' && byte != b'\\' && byte >= 0x20)
|
||||
{
|
||||
self.pos += 1;
|
||||
}
|
||||
out.push_str(&self.text[start..self.pos]);
|
||||
match self.peek() {
|
||||
Some(b'"') => {
|
||||
self.pos += 1;
|
||||
return Ok(out);
|
||||
}
|
||||
Some(b'\\') => {
|
||||
self.pos += 1;
|
||||
out.push(self.escape()?);
|
||||
}
|
||||
_ => return Err(self.error()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn escape(&mut self) -> Result<char, Error> {
|
||||
let escaped = match self.peek() {
|
||||
Some(b'"') => '"',
|
||||
Some(b'\\') => '\\',
|
||||
Some(b'/') => '/',
|
||||
Some(b'b') => '\u{8}',
|
||||
Some(b'f') => '\u{c}',
|
||||
Some(b'n') => '\n',
|
||||
Some(b'r') => '\r',
|
||||
Some(b't') => '\t',
|
||||
Some(b'u') => {
|
||||
self.pos += 1;
|
||||
return self.unicode_escape();
|
||||
}
|
||||
_ => return Err(self.error()),
|
||||
};
|
||||
self.pos += 1;
|
||||
Ok(escaped)
|
||||
}
|
||||
|
||||
/// `\uXXXX`, joining a high surrogate with the low one escaped right after it.
|
||||
fn unicode_escape(&mut self) -> Result<char, Error> {
|
||||
let unit = self.hex4()?;
|
||||
let code = if (0xd800..0xdc00).contains(&unit) {
|
||||
self.expect("\\u")?;
|
||||
let low = self.hex4()?;
|
||||
if !(0xdc00..0xe000).contains(&low) {
|
||||
return Err(self.error());
|
||||
}
|
||||
0x10000 + ((unit - 0xd800) << 10) + (low - 0xdc00)
|
||||
} else {
|
||||
unit
|
||||
};
|
||||
char::from_u32(code).ok_or_else(|| self.error())
|
||||
}
|
||||
|
||||
fn hex4(&mut self) -> Result<u32, Error> {
|
||||
let digits = self
|
||||
.text
|
||||
.get(self.pos..self.pos + 4)
|
||||
.filter(|digits| digits.bytes().all(|byte| byte.is_ascii_hexdigit()))
|
||||
.ok_or_else(|| self.error())?;
|
||||
self.pos += 4;
|
||||
u32::from_str_radix(digits, 16).map_err(|_| self.error())
|
||||
}
|
||||
|
||||
/// `NUMBER_RE`: `(-?(?:0|[1-9][0-9]*))(\.[0-9]+)?([eE][-+]?[0-9]+)?`, an `int` unless a
|
||||
/// fraction or exponent matched.
|
||||
fn number(&mut self) -> Result<Value, Error> {
|
||||
let start = self.pos;
|
||||
self.eat("-");
|
||||
match self.peek() {
|
||||
Some(b'0') => self.pos += 1,
|
||||
Some(b'1'..=b'9') => self.digits(),
|
||||
_ => return Err(Error::InvalidJson(start)),
|
||||
}
|
||||
let integer_end = self.pos;
|
||||
if self.peek() == Some(b'.') && self.bytes.get(self.pos + 1).is_some_and(u8::is_ascii_digit)
|
||||
{
|
||||
self.pos += 1;
|
||||
self.digits();
|
||||
}
|
||||
if matches!(self.peek(), Some(b'e' | b'E')) {
|
||||
let exponent = self.pos;
|
||||
self.pos += 1;
|
||||
if matches!(self.peek(), Some(b'+' | b'-')) {
|
||||
self.pos += 1;
|
||||
}
|
||||
if self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
|
||||
self.digits();
|
||||
} else {
|
||||
self.pos = exponent;
|
||||
}
|
||||
}
|
||||
let text = &self.text[start..self.pos];
|
||||
if self.pos == integer_end {
|
||||
BigInt::from_str(text)
|
||||
.map(Value::Int)
|
||||
.map_err(|_| Error::InvalidJson(start))
|
||||
} else {
|
||||
f64::from_str(text)
|
||||
.map(Value::Float)
|
||||
.map_err(|_| Error::InvalidJson(start))
|
||||
}
|
||||
}
|
||||
|
||||
fn digits(&mut self) {
|
||||
while self.peek().is_some_and(|byte| byte.is_ascii_digit()) {
|
||||
self.pos += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `json.loads(json.dumps(value))` as a `serde_json` value: tuples become arrays and dict
|
||||
/// keys are coerced to strings as `json.dumps` does. Non-finite floats, which Python writes
|
||||
/// as `NaN` and `Infinity`, have no `serde_json` form and fail with [`Error::NonFiniteFloat`].
|
||||
|
|
@ -166,3 +425,86 @@ fn write_string(out: &mut String, text: &str) {
|
|||
}
|
||||
out.push('"');
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{dumps, loads};
|
||||
use crate::{Error, Value, repr::repr};
|
||||
|
||||
// Expected values are the repr of CPython 3.14 json.loads(text)
|
||||
#[rstest]
|
||||
#[case::nan("NaN", "nan")]
|
||||
#[case::infinity("Infinity", "inf")]
|
||||
#[case::negative_infinity("-Infinity", "-inf")]
|
||||
#[case::non_finite_in_a_list("[NaN, Infinity, -Infinity]", "[nan, inf, -inf]")]
|
||||
#[case::beyond_u64("18446744073709551616", "18446744073709551616")]
|
||||
#[case::exponent_overflow("1e400", "inf")]
|
||||
#[case::upper_exponent("1E+2", "100.0")]
|
||||
#[case::negative_zero_int("-0", "0")]
|
||||
#[case::negative_zero_float("-0.0", "-0.0")]
|
||||
#[case::repeated_key(r#"{"a": 1, "b": 2, "a": 3}"#, "{'a': 3, 'b': 2}")]
|
||||
#[case::surrogate_pair(r#""😀""#, "'😀'")]
|
||||
#[case::upper_hex_surrogate_pair(r#""😀""#, "'😀'")]
|
||||
#[case::escapes(
|
||||
r#"{"k": [true, false, null, 1.5, "é\/"]}"#,
|
||||
"{'k': [True, False, None, 1.5, 'é/']}"
|
||||
)]
|
||||
#[case::delete_is_not_a_control_character("\"\x7f\"", "'\\x7f'")]
|
||||
#[case::surrounding_whitespace(" \n[1]\t", "[1]")]
|
||||
#[case::empty_containers("[{}, []]", "[{}, []]")]
|
||||
fn loads_reads_like_python(#[case] text: &str, #[case] expected: &str) {
|
||||
assert_eq!(repr(&loads(text).unwrap()), expected);
|
||||
}
|
||||
|
||||
// CPython 3.14 json.loads raises JSONDecodeError on each
|
||||
#[rstest]
|
||||
#[case::lowercase_nan("nan")]
|
||||
#[case::signed_nan("-NaN")]
|
||||
#[case::plus_sign("+1")]
|
||||
#[case::bare_minus("-")]
|
||||
#[case::leading_zero("01")]
|
||||
#[case::bare_point("1.")]
|
||||
#[case::bare_exponent("1e")]
|
||||
#[case::trailing_data("Infinity1")]
|
||||
#[case::trailing_comma("[1,]")]
|
||||
#[case::missing_comma("[1 2]")]
|
||||
#[case::missing_colon(r#"{"a" 1}"#)]
|
||||
#[case::unquoted_key("{1: 2}")]
|
||||
#[case::raw_newline_in_string("\"a\nb\"")]
|
||||
#[case::unknown_escape(r#""\x41""#)]
|
||||
#[case::empty("")]
|
||||
#[case::byte_order_mark("\u{feff}1")]
|
||||
fn loads_rejects_what_python_rejects(#[case] text: &str) {
|
||||
assert!(matches!(loads(text), Err(Error::InvalidJson(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_lone_surrogate_is_a_known_limit() {
|
||||
assert!(matches!(loads(r#""\ud800""#), Err(Error::InvalidJson(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn loads_reads_back_the_non_finite_floats_dumps_writes() {
|
||||
let written = dumps(&Value::List(vec![
|
||||
Value::Float(f64::NAN),
|
||||
Value::Float(f64::INFINITY),
|
||||
Value::Float(f64::NEG_INFINITY),
|
||||
]))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(repr(&loads(&written).unwrap()), "[nan, inf, -inf]");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn loads_rejects_nesting_past_max_depth() {
|
||||
let text = format!(
|
||||
"{}{}",
|
||||
"[".repeat(crate::MAX_DEPTH + 1),
|
||||
"]".repeat(crate::MAX_DEPTH + 1)
|
||||
);
|
||||
|
||||
assert!(matches!(loads(&text), Err(Error::TooDeep)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,9 +7,10 @@
|
|||
//! |---|---|
|
||||
//! | [`literal`] | `ast.literal_eval(text)` |
|
||||
//! | [`repr`] | `repr(value)` and `str(value)` |
|
||||
//! | [`json`] | `json.dumps(value)`, and `json.loads(json.dumps(value))` as a JSON value |
|
||||
//! | [`json`] | `json.dumps(value)` and `json.loads(text)` |
|
||||
//! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data |
|
||||
//! | [`truthy`] | `bool(value)` |
|
||||
//! | [`number`] | `int(text)`, `float(text)`, and `int(seconds)` as a `Duration` |
|
||||
//!
|
||||
//! [`Value`] is the closed data model these formats share. Live Python objects
|
||||
//! (descriptors, `__bool__`, `__str__`, callbacks) are out of scope: those belong to the
|
||||
|
|
@ -20,10 +21,12 @@
|
|||
//! (string hashes are randomized per process). Sets keep their literal order.
|
||||
//! - `str` values are Rust `String`s, so lone surrogates cannot be represented.
|
||||
//! - [`pickle::loads`] decodes `tuple`, `set`, and `frozenset` as lists.
|
||||
//! - [`number`] reads ASCII digits only, while Python's `int` and `float` accept any Unicode decimal digit.
|
||||
|
||||
mod error;
|
||||
pub mod json;
|
||||
pub mod literal;
|
||||
pub mod number;
|
||||
pub mod pickle;
|
||||
pub mod repr;
|
||||
pub mod truthy;
|
||||
|
|
|
|||
121
litellm-rust/crates/python-compat/src/number.rs
Normal file
121
litellm-rust/crates/python-compat/src/number.rs
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
use std::{str::FromStr, time::Duration};
|
||||
|
||||
use num_bigint::BigInt;
|
||||
|
||||
/// `int(text)` in base 10, or `None` where Python raises `ValueError`.
|
||||
pub fn int(text: &str) -> Option<BigInt> {
|
||||
parse(text)
|
||||
}
|
||||
|
||||
/// `float(text)`, or `None` where Python raises `ValueError`.
|
||||
pub fn float(text: &str) -> Option<f64> {
|
||||
parse(text)
|
||||
}
|
||||
|
||||
/// Why [`seconds`] refused a value, as `int(value)` would. Messages are CPython's.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum SecondsError {
|
||||
/// `OverflowError`. Finite values past `Duration::MAX` land here too, where Python builds a big int.
|
||||
#[error("cannot convert float infinity to integer")]
|
||||
Overflow,
|
||||
/// `ValueError`.
|
||||
#[error("cannot convert float NaN to integer")]
|
||||
NotANumber,
|
||||
}
|
||||
|
||||
/// A float of seconds as a `Duration`, refusing what `int(value)` refuses. Finite negative values clamp to zero.
|
||||
pub fn seconds(value: f64) -> Result<Duration, SecondsError> {
|
||||
match value {
|
||||
value if value.is_nan() => Err(SecondsError::NotANumber),
|
||||
value if value.is_infinite() => Err(SecondsError::Overflow),
|
||||
value => Duration::try_from_secs_f64(value.max(0.0)).map_err(|_| SecondsError::Overflow),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse<N: FromStr>(text: &str) -> Option<N> {
|
||||
let text = text.trim();
|
||||
let digit_at = |index: Option<usize>| {
|
||||
index
|
||||
.and_then(|index| text.as_bytes().get(index))
|
||||
.is_some_and(u8::is_ascii_digit)
|
||||
};
|
||||
let underscores_between_digits = text
|
||||
.match_indices('_')
|
||||
.all(|(index, _)| digit_at(index.checked_sub(1)) && digit_at(Some(index + 1)));
|
||||
if !underscores_between_digits {
|
||||
return None;
|
||||
}
|
||||
text.replace('_', "").parse().ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use num_bigint::BigInt;
|
||||
use rstest::rstest;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{SecondsError, float, int, seconds};
|
||||
|
||||
// Expected values are CPython 3.14 int(text)
|
||||
#[rstest]
|
||||
#[case::trailing_newline("2\n", Some(2))]
|
||||
#[case::surrounding_spaces_and_underscore(" 6_380 ", Some(6380))]
|
||||
#[case::plus_sign("+5", Some(5))]
|
||||
#[case::negative_zero("-0", Some(0))]
|
||||
#[case::leading_zeros("007", Some(7))]
|
||||
#[case::unicode_spaces("\u{a0}7\u{2003}", Some(7))]
|
||||
#[case::beyond_i64("99999999999999999999", Some(99_999_999_999_999_999_999_i128))]
|
||||
#[case::double_underscore("1__0", None)]
|
||||
#[case::leading_underscore("_1", None)]
|
||||
#[case::trailing_underscore("1_", None)]
|
||||
#[case::decimal_point("5.0", None)]
|
||||
#[case::information_separator("\x1c7", None)]
|
||||
#[case::empty("", None)]
|
||||
#[case::sign_only("+", None)]
|
||||
#[case::non_ascii_digit_is_a_known_limit("\u{667}", None)]
|
||||
fn int_parses_like_python(#[case] text: &str, #[case] expected: Option<i128>) {
|
||||
assert_eq!(int(text), expected.map(BigInt::from));
|
||||
}
|
||||
|
||||
// Expected values are CPython 3.14 float(text)
|
||||
#[rstest]
|
||||
#[case::underscore_and_newline(" 1_0.5\n", Some(10.5))]
|
||||
#[case::exponent("1e3", Some(1000.0))]
|
||||
#[case::leading_point(".5", Some(0.5))]
|
||||
#[case::trailing_point("5.", Some(5.0))]
|
||||
#[case::plus_sign("+1.5", Some(1.5))]
|
||||
#[case::infinity_word("Infinity", Some(f64::INFINITY))]
|
||||
#[case::infinity_with_space("infinity ", Some(f64::INFINITY))]
|
||||
#[case::underscore_before_exponent("1_e3", None)]
|
||||
#[case::underscore_after_exponent("1e_3", None)]
|
||||
#[case::underscore_after_point("1._5", None)]
|
||||
#[case::underscore_after_word("inf_", None)]
|
||||
#[case::hexadecimal("0x10", None)]
|
||||
#[case::bare_exponent("1.5e", None)]
|
||||
fn float_parses_like_python(#[case] text: &str, #[case] expected: Option<f64>) {
|
||||
assert_eq!(float(text), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn float_reads_a_signed_nan_like_python() {
|
||||
assert!(float("-NaN").is_some_and(f64::is_nan));
|
||||
}
|
||||
|
||||
// CPython 3.14: int(inf) and int(-inf) raise OverflowError, int(nan) raises ValueError
|
||||
#[rstest]
|
||||
#[case::fraction(1.5, Ok(Duration::from_millis(1500)))]
|
||||
#[case::zero(0.0, Ok(Duration::ZERO))]
|
||||
#[case::negative_clamps_to_zero(-3.0, Ok(Duration::ZERO))]
|
||||
#[case::infinity(f64::INFINITY, Err(SecondsError::Overflow))]
|
||||
#[case::negative_infinity(f64::NEG_INFINITY, Err(SecondsError::Overflow))]
|
||||
#[case::past_duration_max(1e20, Err(SecondsError::Overflow))]
|
||||
#[case::nan(f64::NAN, Err(SecondsError::NotANumber))]
|
||||
#[case::negative_nan(-f64::NAN, Err(SecondsError::NotANumber))]
|
||||
fn seconds_refuses_what_int_refuses(
|
||||
#[case] value: f64,
|
||||
#[case] expected: Result<Duration, SecondsError>,
|
||||
) {
|
||||
assert_eq!(seconds(value), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -76,6 +76,11 @@ const KNOWN: &[(&str, &str, &str)] = &[
|
|||
"literal_eval",
|
||||
"deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it",
|
||||
),
|
||||
(
|
||||
"nested_150",
|
||||
"json.loads",
|
||||
"deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it",
|
||||
),
|
||||
(
|
||||
"nested_150",
|
||||
"pickle.loads",
|
||||
|
|
@ -341,3 +346,19 @@ fn to_json_matches_python_json_round_trip(fixtures: &Fixtures) {
|
|||
}
|
||||
mismatches.finish();
|
||||
}
|
||||
|
||||
/// CPython gives back the same text from `json.dumps(json.loads(text))` for every text its
|
||||
/// `json.dumps` wrote here, non-finite floats included, so `loads` must too.
|
||||
#[rstest]
|
||||
fn loads_reads_back_what_python_json_dumps_wrote(fixtures: &Fixtures) {
|
||||
let mut mismatches = Mismatches::default();
|
||||
for row in &fixtures.rows {
|
||||
let Some(text) = &row.json else { continue };
|
||||
let actual = match json::loads(text).and_then(|value| json::dumps(&value)) {
|
||||
Ok(text) => text,
|
||||
Err(error) => format!("error: {error}"),
|
||||
};
|
||||
mismatches.check(row, "json.loads", text, &actual);
|
||||
}
|
||||
mismatches.finish();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
- The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace
|
||||
- It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default
|
||||
- Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route
|
||||
- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`)
|
||||
- Every connection sets TCP_NODELAY, like redis-py's `socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)`. Set every redis-rs timeout and retry explicitly, citing the redis-py 5.3.1 or `litellm/` source next to each constant. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's. `Timeouts` takes `socket_timeout` and `socket_connect_timeout`, and the connect timeout falls back to the socket timeout like redis-py. Standalone reconnects on the next command with no retry of the failed one, matching redis-py's `Retry(NoBackoff(), 0)`. redis-rs keeps a socket whose command timed out, so `Renewing` in `connection.rs` swaps in a fresh `ConnectionManager`, or a fresh cluster connection like redis-py's cluster `aclose`, after a timeout. Without it a socket that went silent times out every command until TCP gives up on it. Cluster keeps redis-py's 3 `cluster_error_retry_attempts` and waits a flat 250ms between them, which is redis-py's CLUSTERDOWN sleep. redis-py's jittered per-node reconnect backoff has no redis-rs equivalent. The retry set also differs, redis-rs retries LOADING and READONLY internally, while redis-py 5.3.1 raises them. Do not claim exact retry parity
|
||||
- Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values
|
||||
|
||||
## Errors
|
||||
- `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it
|
||||
- Classify in `From<RedisError> for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error
|
||||
- Classify with `Error::classify(error, deployment)`, never a bare `From<RedisError>`, because redis-py (5.3.1) decides by topology. `Transport` is what Python's `_redis_health_error_types` counts: LOADING everywhere, CLUSTERDOWN and MASTERDOWN only on cluster (only `ClusterParser` maps them), READONLY only under Sentinel (where `SentinelManagedConnection` turns it into `ConnectionError` and rediscovers), plus NOAUTH, WRONGPASS and the `ERR` messages redis-py maps to `ConnectionError`, matched exactly like `parse_error`. TRYAGAIN is `Server`, since the cluster client already retried it. Each code gets its own rstest case built from a parsed RESP error
|
||||
- A `Server` error never drops the connection
|
||||
- `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome
|
||||
- Never swallow an error or turn it into a miss. The caller picks the fallback
|
||||
|
|
@ -17,27 +17,61 @@
|
|||
- Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it
|
||||
- `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success
|
||||
- An open breaker returns `Error::BreakerOpen` without touching the socket
|
||||
- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count
|
||||
|
||||
## Namespace
|
||||
- Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace`
|
||||
- The rule never changes, because Python and Rust pods share keys in the same Redis
|
||||
|
||||
## Commands
|
||||
- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
|
||||
- A command is a value, one constructor in `commands.rs` returning `Command<T>`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::with` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
|
||||
- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate
|
||||
- On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test
|
||||
- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`)
|
||||
- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same
|
||||
- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds
|
||||
- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat`
|
||||
- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ
|
||||
|
||||
## Scripts and pipelines
|
||||
- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production
|
||||
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async`
|
||||
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python
|
||||
- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice
|
||||
- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared<T>`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s
|
||||
- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip
|
||||
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order
|
||||
- Pipeline ops take one key each. A multi-key op would span slots inside one group
|
||||
- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once
|
||||
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed
|
||||
- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed
|
||||
- Command constructors live in `commands.rs` and every executor sends them, so a command is written once
|
||||
|
||||
## Tests
|
||||
## Pub/sub
|
||||
- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel
|
||||
- Channel names never get the namespace. Python builds `{namespace}:{channel}` itself, so the store does too
|
||||
- `publish` is a command like any other, guarded by the breaker and reported to events
|
||||
- `subscribe` returns a `Stream` that opens its own RESP2 `PubSub` connection when first polled, like redis-py's `PubSub`, to the endpoint, the first reachable cluster node, or the master Sentinel names at that moment. It never touches the breaker or events, like Python's subscriber, which reads off its own connection
|
||||
- After a lost connection or a failed subscribe it waits and resubscribes, 5s doubling to 60s and reset by a successful subscribe, like `AuthCacheInvalidationSubscriber._run`. Messages published while it is down are lost. The connection has no idle ping, so a socket that dies silently is only noticed by TCP
|
||||
- redis-rs's `PubSub` takes no credentials provider, so a provider is asked for its current credentials once when the subscription connects, which is when redis-py's pubsub connection asks it. The connect is bounded by the connect timeout
|
||||
|
||||
## Batches
|
||||
- `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment`
|
||||
- `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`, and answers one value per key it was given, in that order. Use it instead of one `get` per key whenever Python reads several keys in one call
|
||||
- It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local
|
||||
- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared<T>` is the one reply handle on every route
|
||||
- Each op returns a `Declared<T>` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once
|
||||
- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis
|
||||
- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client
|
||||
- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
|
||||
- Tests that only need the public API live in `tests/<module>.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`)
|
||||
- Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection
|
||||
- Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning)
|
||||
- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables
|
||||
- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables. A socket that never answers again comes from `TcpProxy::blackhole_open_connections`, since `thaw` revives the same socket
|
||||
- Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline
|
||||
|
||||
## TLS, credentials and Sentinel
|
||||
- `TlsOptions` follows redis-py 5.3.1. It verifies the certificate by default, skips the hostname check by default (`ssl_check_hostname=False`), and trusts the system roots plus `ca_certs`
|
||||
- rustls only reads unencrypted keys, so a PKCS#8 `ENCRYPTED PRIVATE KEY` is decrypted with `key_password` before it reaches redis-rs. A legacy `Proc-Type: 4,ENCRYPTED` PEM key is refused with `KeyError::LegacyEncryption` rather than decrypted by hand-rolled OpenSSL key derivation
|
||||
- redis-rs loads the system trust store for every TLS connection and fails if that fails, even with `ca_certs` set. Where the macOS keychain is blocked (agent sandboxes), run the TLS tests with `SSL_CERT_FILE=/etc/ssl/cert.pem`
|
||||
- `Credentials::Static` are sent at connect, and URL userinfo wins over them, like redis-py's `from_url`
|
||||
- A url's `?protocol=` is ignored and every connection speaks RESP2, which every reply decode assumes
|
||||
- `Credentials::Provider` wraps a redis-rs `StreamingCredentialsProvider`. redis-rs takes the first item at connect and re-`AUTH`s live connections on every later item, on each cluster node. Providers push, redis-rs never polls. With a provider, redis-rs overwrites URL userinfo at connect, which matches Python stripping `user:pass@` from the URL
|
||||
- Sentinel discovers the master on each new connection and drops it on a transport failure or a READONLY reply, to one command or inside a pipeline, like redis-py's `SentinelConnectionPool`
|
||||
- redis-rs sends `ROLE` to the master a Sentinel names before using it, which redis-py never does. That probe gets the data credentials (a provider's current pair) and TLS, so the Sentinel client is built on each discovery. redis-rs gives the probe no way to skip the hostname check, and drops the client certificate when the certificate is not verified, so a TLS master whose certificate does not name the announced address, or that requires a client certificate while verification is off, fails discovery where Python connects
|
||||
- Sentinel monitors get `sentinel_password`, the data username when that password is set, and the TLS mode and certificates. Python's `_sentinel_auth_kwargs` forwards every other connection kwarg too, and it sends the username even without a sentinel password
|
||||
- A `Debug` impl never shows a secret. Follow `.agents/skills/rust-redaction`
|
||||
|
|
|
|||
|
|
@ -6,13 +6,28 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
tracing.workspace = true
|
||||
litellm-circuit-breaker.workspace = true
|
||||
litellm-secret-value.workspace = true
|
||||
litellm-python-compat.workspace = true
|
||||
litellm-redis-types.workspace = true
|
||||
futures-util.workspace = true
|
||||
redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] }
|
||||
redis = { workspace = true, features = [
|
||||
"cluster-async",
|
||||
"connection-manager",
|
||||
"sentinel",
|
||||
"tls-rustls-insecure",
|
||||
"token-based-authentication",
|
||||
"tokio-rustls-comp",
|
||||
] }
|
||||
pem = "4.0.0"
|
||||
pkcs8 = { version = "0.10.2", features = ["3des", "encryption", "std"] }
|
||||
rustls-native-certs.workspace = true
|
||||
socket2 = "0.6.5"
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["rt", "sync", "time"] }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-redis-testing.workspace = true
|
||||
rustls.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
147
litellm-rust/crates/redis-client/src/batch.rs
Normal file
147
litellm-rust/crates/redis-client/src/batch.rs
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
use std::{
|
||||
io,
|
||||
sync::{Arc, Mutex, MutexGuard},
|
||||
};
|
||||
|
||||
use futures_util::future::join_all;
|
||||
use redis::{RedisError, Value};
|
||||
use tokio::sync::oneshot::{self, error::TryRecvError};
|
||||
|
||||
use crate::{Command, Error, Pipeline, RedisClient, commands::Op};
|
||||
|
||||
/// Ops declared by many callers and sent together, like litellm's `RedisBatch`: one pipeline on a
|
||||
/// standalone or Sentinel Redis, and each op on its own on a cluster or in an alone batch.
|
||||
#[derive(Clone)]
|
||||
pub struct RedisBatch(Arc<Shared>);
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Mode {
|
||||
/// litellm's `RedisBatch`: one pipeline, or each op on its own on a cluster.
|
||||
Batch,
|
||||
/// Each op on its own, the way the `RedisCache` method behind it sends it.
|
||||
Alone,
|
||||
/// One redis-py pipeline on every topology, split by slot on a cluster.
|
||||
Pipeline,
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
client: RedisClient,
|
||||
mode: Mode,
|
||||
pending: Mutex<Pending>,
|
||||
flushing: tokio::sync::Mutex<()>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Pending {
|
||||
ops: Vec<Op>,
|
||||
replies: Vec<oneshot::Sender<Result<Value, Error>>>,
|
||||
}
|
||||
|
||||
type Decode<T> = Box<dyn FnOnce(Value) -> Result<T, Error> + Send>;
|
||||
|
||||
#[must_use]
|
||||
pub struct Declared<T> {
|
||||
batch: RedisBatch,
|
||||
reply: oneshot::Receiver<Result<Value, Error>>,
|
||||
decode: Decode<T>,
|
||||
}
|
||||
|
||||
impl RedisBatch {
|
||||
pub(crate) fn new(client: RedisClient) -> Self {
|
||||
Self::build(client, Mode::Batch)
|
||||
}
|
||||
|
||||
/// A batch whose every op goes out on its own, the way the `RedisCache` method behind it sends
|
||||
/// it, so a fresh one is "send it now": the first `resolve` sends it.
|
||||
pub fn alone(client: RedisClient) -> Self {
|
||||
Self::build(client, Mode::Alone)
|
||||
}
|
||||
|
||||
/// A batch sent as one redis-py pipeline, like `async_increment_pipeline`, on every topology.
|
||||
pub fn pipelined(client: RedisClient) -> Self {
|
||||
Self::build(client, Mode::Pipeline)
|
||||
}
|
||||
|
||||
fn build(client: RedisClient, mode: Mode) -> Self {
|
||||
Self(Arc::new(Shared {
|
||||
client,
|
||||
mode,
|
||||
pending: Mutex::new(Pending::default()),
|
||||
flushing: tokio::sync::Mutex::new(()),
|
||||
}))
|
||||
}
|
||||
|
||||
pub fn add<T>(&self, command: Command<T>) -> Declared<T> {
|
||||
let (op, decode) = command.prepare(self.0.client.namespace());
|
||||
let (sender, reply) = oneshot::channel();
|
||||
let mut pending = self.pending();
|
||||
pending.ops.push(op);
|
||||
pending.replies.push(sender);
|
||||
Declared {
|
||||
batch: self.clone(),
|
||||
reply,
|
||||
decode,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn flush(&self) {
|
||||
let _flushing = self.0.flushing.lock().await;
|
||||
let Pending { ops, replies } = std::mem::take(&mut *self.pending());
|
||||
if replies.is_empty() {
|
||||
return;
|
||||
}
|
||||
let client = &self.0.client;
|
||||
let alone = match self.0.mode {
|
||||
Mode::Alone => true,
|
||||
Mode::Batch => client.is_cluster(),
|
||||
Mode::Pipeline => false,
|
||||
};
|
||||
let results = match alone {
|
||||
true => join_all(ops.into_iter().map(|op| client.run_alone(op))).await,
|
||||
false => {
|
||||
client
|
||||
.run(&Pipeline::of(client.namespace().clone(), ops))
|
||||
.await
|
||||
}
|
||||
};
|
||||
for (reply, result) in replies.into_iter().zip(results) {
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
}
|
||||
|
||||
fn pending(&self) -> MutexGuard<'_, Pending> {
|
||||
self.0
|
||||
.pending
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Declared<T> {
|
||||
pub async fn resolve(mut self) -> Result<T, Error> {
|
||||
match self.reply.try_recv() {
|
||||
Ok(reply) => return reply.and_then(self.decode),
|
||||
Err(TryRecvError::Closed) => return Err(abandoned()),
|
||||
Err(TryRecvError::Empty) => {}
|
||||
}
|
||||
self.batch.flush().await;
|
||||
self.settled().await
|
||||
}
|
||||
|
||||
pub async fn settled(self) -> Result<T, Error> {
|
||||
let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned()));
|
||||
reply.and_then(self.decode)
|
||||
}
|
||||
|
||||
/// Every reply in order, flushing their batch once if nobody has yet.
|
||||
pub async fn all(declared: impl IntoIterator<Item = Self>) -> Vec<Result<T, Error>> {
|
||||
join_all(declared.into_iter().map(Self::resolve)).await
|
||||
}
|
||||
}
|
||||
|
||||
fn abandoned() -> Error {
|
||||
Error::Transport(RedisError::from(io::Error::new(
|
||||
io::ErrorKind::Interrupted,
|
||||
"the batch flush carrying this op was cancelled",
|
||||
)))
|
||||
}
|
||||
|
|
@ -1,21 +1,43 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
sync::{Arc, Mutex, PoisonError, Weak},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::{Admission, Breaker, Failure};
|
||||
use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value};
|
||||
use futures_util::stream::BoxStream;
|
||||
use litellm_circuit_breaker::{Admission, Breaker, Failure, State};
|
||||
use redis::{Cmd, RedisError, Value};
|
||||
|
||||
use crate::{
|
||||
ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands,
|
||||
CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline,
|
||||
RedisBatch, RedisConfig, RedisEvents,
|
||||
commands::{self, Op, Plan, Trips},
|
||||
connection::Connection,
|
||||
pubsub::Subscriber,
|
||||
};
|
||||
|
||||
/// One Redis client. Clones share the connection, the breaker and the subscriber.
|
||||
#[derive(Clone)]
|
||||
pub struct RedisClient {
|
||||
pub(crate) connection: Connection,
|
||||
pub struct RedisClient(Arc<Shared>);
|
||||
|
||||
/// A handle that does not keep the client alive, for a registry that must not outlive its users.
|
||||
#[derive(Clone)]
|
||||
pub struct WeakRedisClient(Weak<Shared>);
|
||||
|
||||
impl WeakRedisClient {
|
||||
pub fn upgrade(&self) -> Option<RedisClient> {
|
||||
self.0.upgrade().map(RedisClient)
|
||||
}
|
||||
}
|
||||
|
||||
struct Shared {
|
||||
connection: Connection,
|
||||
subscriber: Subscriber,
|
||||
namespace: Namespace,
|
||||
breaker: Option<Breaker>,
|
||||
events: Option<Arc<dyn RedisEvents>>,
|
||||
health_check: Option<Duration>,
|
||||
last_used: Mutex<Instant>,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
|
|
@ -27,122 +49,157 @@ pub enum Ttl {
|
|||
|
||||
impl RedisClient {
|
||||
pub fn new(config: RedisConfig) -> Result<Self, ConnectError> {
|
||||
let connection = match &config.endpoint {
|
||||
Endpoint::Url(url) => {
|
||||
Connection::standalone(url.expose().into_connection_info()?, config.timeout)?
|
||||
}
|
||||
Endpoint::Host(host) => {
|
||||
Connection::standalone(host.connection_info()?, config.timeout)?
|
||||
}
|
||||
Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?,
|
||||
};
|
||||
Ok(Self {
|
||||
let connection = Connection::new(
|
||||
&config.target,
|
||||
&config.credentials,
|
||||
config.timeouts,
|
||||
config.keepalive,
|
||||
)?;
|
||||
Ok(Self(Arc::new(Shared {
|
||||
subscriber: Subscriber::new(
|
||||
&config.target,
|
||||
&config.credentials,
|
||||
config.timeouts,
|
||||
config.keepalive,
|
||||
&connection,
|
||||
)?,
|
||||
connection,
|
||||
namespace: config.namespace,
|
||||
breaker: config.breaker,
|
||||
})
|
||||
events: config.events,
|
||||
health_check: config.health_check,
|
||||
last_used: Mutex::new(Instant::now()),
|
||||
})))
|
||||
}
|
||||
|
||||
pub fn downgrade(&self) -> WeakRedisClient {
|
||||
WeakRedisClient(Arc::downgrade(&self.0))
|
||||
}
|
||||
|
||||
/// The breaker's current state, or none when the breaker is disabled.
|
||||
pub fn breaker_state(&self) -> Option<State> {
|
||||
self.0.breaker.as_ref().map(Breaker::state)
|
||||
}
|
||||
|
||||
pub async fn ping(&self) -> Result<(), Error> {
|
||||
self.query(&commands::ping()).await
|
||||
}
|
||||
|
||||
pub async fn get<T: FromRedisValue>(&self, key: &str) -> Result<Option<T>, Error> {
|
||||
self.query(&commands::get(self.key(key))).await
|
||||
}
|
||||
|
||||
pub async fn mget<T: FromRedisValue>(&self, keys: &[&str]) -> Result<Vec<Option<T>>, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.query(&commands::mget(self.keys(keys))).await
|
||||
}
|
||||
|
||||
pub async fn set(
|
||||
&self,
|
||||
key: &str,
|
||||
value: impl ToSingleRedisArg,
|
||||
ttl: Duration,
|
||||
) -> Result<(), Error> {
|
||||
self.query(&commands::set(self.key(key), value, ttl)).await
|
||||
}
|
||||
|
||||
pub async fn set_nx(
|
||||
&self,
|
||||
key: &str,
|
||||
value: impl ToSingleRedisArg,
|
||||
ttl: Duration,
|
||||
) -> Result<bool, Error> {
|
||||
let reply: Option<String> = self
|
||||
.query(&commands::set_nx(self.key(key), value, ttl))
|
||||
.await?;
|
||||
Ok(reply.is_some())
|
||||
}
|
||||
|
||||
pub async fn del(&self, keys: &[&str]) -> Result<usize, Error> {
|
||||
if keys.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
self.query(&commands::del(self.keys(keys))).await
|
||||
}
|
||||
|
||||
pub async fn expire(&self, key: &str, ttl: Duration) -> Result<bool, Error> {
|
||||
self.query(&commands::expire(self.key(key), ttl)).await
|
||||
}
|
||||
|
||||
pub async fn ttl(&self, key: &str) -> Result<Ttl, Error> {
|
||||
let reply: i64 = self.query(&commands::ttl(self.key(key))).await?;
|
||||
Ok(match reply {
|
||||
-2 => Ttl::Missing,
|
||||
-1 => Ttl::Persistent,
|
||||
millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn incr_by(&self, key: &str, delta: i64) -> Result<i64, Error> {
|
||||
self.query(&commands::incr_by(self.key(key), delta)).await
|
||||
}
|
||||
|
||||
pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result<f64, Error> {
|
||||
self.query(&commands::incr_by_float(self.key(key), delta))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result<usize, Error> {
|
||||
self.query(&commands::rpush(self.key(key), values)).await
|
||||
}
|
||||
|
||||
pub async fn eval<T: FromRedisValue>(
|
||||
&self,
|
||||
script: &Script,
|
||||
keys: &[&str],
|
||||
args: impl ToRedisArgs,
|
||||
) -> Result<T, Error> {
|
||||
let keys = self.keys(keys);
|
||||
let args = args.to_redis_args();
|
||||
self.guarded(async {
|
||||
Ok(script
|
||||
.invocation(&keys, &args)
|
||||
.invoke_async(&mut self.connection.clone())
|
||||
.await?)
|
||||
self.guarded(Operation::Command("PING"), async {
|
||||
commands::ping()
|
||||
.query_async(&mut self.0.connection.clone())
|
||||
.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Sends one command on its own, the way the `RedisCache` method behind it sends it.
|
||||
pub async fn send<T>(&self, command: Command<T>) -> Result<T, Error> {
|
||||
let (op, decode) = command.prepare(&self.0.namespace);
|
||||
self.run_alone(op).await.and_then(decode)
|
||||
}
|
||||
|
||||
pub(crate) async fn run_alone(&self, op: Op) -> Result<Value, Error> {
|
||||
let Op {
|
||||
operation, plan, ..
|
||||
} = op;
|
||||
self.guarded(operation, async {
|
||||
let mut connection = self.0.connection.clone();
|
||||
match plan {
|
||||
Plan::Reply(value) => Ok(value),
|
||||
Plan::Commands(commands) => match <[Cmd; 1]>::try_from(commands) {
|
||||
Ok([command]) => command.query_async(&mut connection).await,
|
||||
Err(commands) => {
|
||||
let mut pipeline = redis::pipe();
|
||||
for command in commands {
|
||||
pipeline.add_command(command);
|
||||
}
|
||||
let replies: Vec<Value> = pipeline.query_async(&mut connection).await?;
|
||||
Ok(replies.into_iter().next().unwrap_or(Value::Nil))
|
||||
}
|
||||
},
|
||||
Plan::Script(call) => {
|
||||
call.script
|
||||
.invocation(&call.keys, &call.args)
|
||||
.invoke_async(&mut connection)
|
||||
.await
|
||||
}
|
||||
Plan::KeepingTtl {
|
||||
increment,
|
||||
key,
|
||||
ttl,
|
||||
trips,
|
||||
} => {
|
||||
let (value, current): (Value, i64) = match trips {
|
||||
Trips::Shared => {
|
||||
redis::pipe()
|
||||
.add_command(increment)
|
||||
.add_command(commands::ttl_cmd(&key))
|
||||
.query_async(&mut connection)
|
||||
.await?
|
||||
}
|
||||
Trips::Separate => (
|
||||
increment.query_async(&mut connection).await?,
|
||||
commands::ttl_cmd(&key).query_async(&mut connection).await?,
|
||||
),
|
||||
};
|
||||
if current == -1 {
|
||||
redis::cmd("EXPIRE")
|
||||
.arg(&key)
|
||||
.arg(ttl.as_secs())
|
||||
.exec_async(&mut connection)
|
||||
.await?;
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Messages published on `channel`, on a connection of its own. After the connection is lost
|
||||
/// it reconnects and resubscribes, and messages published meanwhile are not delivered.
|
||||
pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec<u8>> {
|
||||
self.0.subscriber.clone().messages(channel.to_owned())
|
||||
}
|
||||
|
||||
pub fn pipeline(&self) -> Pipeline {
|
||||
Pipeline::new(self.namespace.clone())
|
||||
Pipeline::new(self.0.namespace.clone())
|
||||
}
|
||||
|
||||
pub(crate) fn namespace(&self) -> &Namespace {
|
||||
&self.0.namespace
|
||||
}
|
||||
|
||||
pub fn is_cluster(&self) -> bool {
|
||||
matches!(self.0.connection, Connection::Cluster(_))
|
||||
}
|
||||
|
||||
pub fn same_client(&self, other: &Self) -> bool {
|
||||
Arc::ptr_eq(&self.0, &other.0)
|
||||
}
|
||||
|
||||
pub fn batch(&self) -> RedisBatch {
|
||||
RedisBatch::new(self.clone())
|
||||
}
|
||||
|
||||
pub async fn run(&self, pipeline: &Pipeline) -> Vec<Result<Value, Error>> {
|
||||
let operation = Operation::Pipeline {
|
||||
ops: pipeline.len(),
|
||||
};
|
||||
let started = Instant::now();
|
||||
let admission = match self.admit() {
|
||||
Ok(admission) => admission,
|
||||
Err(_) => {
|
||||
self.observe(
|
||||
operation,
|
||||
started,
|
||||
Outcome::Failed(FailureKind::BreakerOpen),
|
||||
);
|
||||
return (0..pipeline.len())
|
||||
.map(|_| Err(Error::BreakerOpen))
|
||||
.collect();
|
||||
}
|
||||
};
|
||||
let results = pipeline.run(&self.connection).await;
|
||||
self.check_health().await;
|
||||
let results = pipeline.run(&self.0.connection).await;
|
||||
let transport = results.iter().find_map(|result| {
|
||||
result
|
||||
.as_ref()
|
||||
|
|
@ -150,31 +207,76 @@ impl RedisClient {
|
|||
.filter(|error| matches!(error, Error::Transport(_)))
|
||||
});
|
||||
settle(admission, transport);
|
||||
let failure = transport.or_else(|| results.iter().find_map(|result| result.as_ref().err()));
|
||||
let outcome = failure.map_or(Outcome::Succeeded, |error| {
|
||||
Outcome::Failed(FailureKind::of(error))
|
||||
});
|
||||
self.observe(operation, started, outcome);
|
||||
results
|
||||
}
|
||||
|
||||
fn key(&self, key: &str) -> String {
|
||||
self.namespace.apply(key)
|
||||
}
|
||||
|
||||
fn keys(&self, keys: &[&str]) -> Vec<String> {
|
||||
keys.iter().map(|key| self.key(key)).collect()
|
||||
}
|
||||
|
||||
async fn query<T: FromRedisValue>(&self, command: &Cmd) -> Result<T, Error> {
|
||||
self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) })
|
||||
.await
|
||||
}
|
||||
|
||||
async fn guarded<T>(&self, call: impl Future<Output = Result<T, Error>>) -> Result<T, Error> {
|
||||
let admission = self.admit()?;
|
||||
let result = call.await;
|
||||
settle(admission, result.as_ref().err());
|
||||
async fn guarded<T>(
|
||||
&self,
|
||||
operation: Operation,
|
||||
call: impl Future<Output = Result<T, RedisError>>,
|
||||
) -> Result<T, Error> {
|
||||
let started = Instant::now();
|
||||
let result = match self.admit() {
|
||||
Ok(admission) => {
|
||||
self.check_health().await;
|
||||
let result = call
|
||||
.await
|
||||
.map_err(|error| Error::classify(error, self.0.connection.deployment()));
|
||||
settle(admission, result.as_ref().err());
|
||||
result
|
||||
}
|
||||
Err(error) => Err(error),
|
||||
};
|
||||
self.observe(operation, started, Outcome::of(&result));
|
||||
result
|
||||
}
|
||||
|
||||
fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) {
|
||||
*self
|
||||
.0
|
||||
.last_used
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner) = Instant::now();
|
||||
if let Some(events) = &self.0.events {
|
||||
events.call(&CallEvent {
|
||||
operation,
|
||||
duration: started.elapsed(),
|
||||
outcome,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// redis-py pings a connection that sat idle longer than `health_check_interval` before
|
||||
/// sending the next command on it. A ping that times out or finds the socket dropped leaves
|
||||
/// that connection behind, as a command would, so the command then goes out on a fresh one,
|
||||
/// where redis-py 5.3.1 disconnects and fails the command with the ping's error. The ping is
|
||||
/// not a call of the caller's and never reaches the breaker or the events.
|
||||
async fn check_health(&self) {
|
||||
let Some(interval) = self.0.health_check else {
|
||||
return;
|
||||
};
|
||||
let idle = self
|
||||
.0
|
||||
.last_used
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.elapsed();
|
||||
if idle < interval {
|
||||
return;
|
||||
}
|
||||
if let Err(error) = self.0.connection.ping_every_node().await {
|
||||
tracing::debug!(%error, "idle Redis connection failed its health check, replacing it");
|
||||
}
|
||||
}
|
||||
|
||||
fn admit(&self) -> Result<Option<Admission>, Error> {
|
||||
self.breaker
|
||||
self.0
|
||||
.breaker
|
||||
.as_ref()
|
||||
.map(|breaker| {
|
||||
breaker
|
||||
|
|
@ -198,455 +300,3 @@ fn settle(admission: Option<Admission>, error: Option<&Error>) {
|
|||
Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use litellm_redis_testing::RedisServer;
|
||||
use litellm_secret_value::SecretValue;
|
||||
use rstest::rstest;
|
||||
|
||||
use litellm_circuit_breaker::{Breaker, Config, State};
|
||||
|
||||
use super::{RedisClient, Ttl};
|
||||
use crate::{
|
||||
DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig,
|
||||
test_support::{Server, Topology, client_id},
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn namespaces_isolate_keys_on_the_server(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let a = server.client("a", DEFAULT_TIMEOUT).await;
|
||||
let b = server.client("b", DEFAULT_TIMEOUT).await;
|
||||
let bare = server.client("", DEFAULT_TIMEOUT).await;
|
||||
|
||||
a.set("x", "1", Duration::from_secs(60)).await.unwrap();
|
||||
|
||||
assert_eq!(a.get::<String>("x").await.unwrap().as_deref(), Some("1"));
|
||||
assert_eq!(a.get::<String>("a:x").await.unwrap().as_deref(), Some("1"));
|
||||
assert_eq!(
|
||||
bare.get::<String>("a:x").await.unwrap().as_deref(),
|
||||
Some("1")
|
||||
);
|
||||
assert_eq!(b.get::<String>("x").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_rejected_command_keeps_the_connection(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
client.rpush("list", "item").await.unwrap();
|
||||
let before = client_id(&client, "list").await;
|
||||
|
||||
let rejected = client.incr_by("list", 1).await;
|
||||
|
||||
assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}");
|
||||
assert_eq!(client.rpush("list", "next").await.unwrap(), 2);
|
||||
assert_eq!(client_id(&client, "list").await, before);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
client
|
||||
.set("float", "15.0", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let decoded = client.get::<i64>("float").await;
|
||||
|
||||
assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}");
|
||||
assert_eq!(client.get::<i64>("missing").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", Duration::from_millis(500)).await;
|
||||
client
|
||||
.set("first", "1", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.set("second", "2", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
server.freeze().await;
|
||||
let timed_out = client.get::<String>("first").await;
|
||||
server.thaw().await;
|
||||
|
||||
assert!(
|
||||
matches!(timed_out, Err(Error::Transport(_))),
|
||||
"{timed_out:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
client.get::<String>("second").await.unwrap().as_deref(),
|
||||
Some("2")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_same_client_recovers_after_redis_restarts(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", Duration::from_millis(500)).await;
|
||||
|
||||
server.restart().await;
|
||||
|
||||
let deadline = Instant::now() + Duration::from_secs(30);
|
||||
let recovered = loop {
|
||||
match client.incr_by("after_restart", 1).await {
|
||||
Ok(value) => break Some(value),
|
||||
Err(_) if Instant::now() < deadline => {
|
||||
tokio::time::sleep(Duration::from_millis(200)).await
|
||||
}
|
||||
Err(_) => break None,
|
||||
}
|
||||
};
|
||||
assert_eq!(recovered, Some(1));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_client_created_while_redis_is_down_works_once_redis_is_up(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
server.stop().await;
|
||||
|
||||
let client = server.client("", Duration::from_millis(500)).await;
|
||||
let while_down = client.ping().await;
|
||||
server.resume().await;
|
||||
|
||||
assert!(
|
||||
matches!(while_down, Err(Error::Transport(_))),
|
||||
"{while_down:?}"
|
||||
);
|
||||
let deadline = Instant::now() + Duration::from_secs(30);
|
||||
let recovered = loop {
|
||||
match client.ping().await {
|
||||
Ok(()) => break true,
|
||||
Err(_) if Instant::now() < deadline => {
|
||||
tokio::time::sleep(Duration::from_millis(200)).await
|
||||
}
|
||||
Err(_) => break false,
|
||||
}
|
||||
};
|
||||
assert!(recovered);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn empty_batches_succeed_without_asking_redis(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
|
||||
assert_eq!(
|
||||
client.mget::<String>(&[]).await.unwrap(),
|
||||
Vec::<Option<String>>::new()
|
||||
);
|
||||
assert_eq!(client.del(&[]).await.unwrap(), 0);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mget_and_del_span_slots_in_input_order(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
client.set("a", "1", Duration::from_secs(60)).await.unwrap();
|
||||
client.set("c", "3", Duration::from_secs(60)).await.unwrap();
|
||||
|
||||
let values = client.mget::<String>(&["c", "b", "a"]).await.unwrap();
|
||||
let deleted = client.del(&["a", "b", "c"]).await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![Some("3".to_owned()), None, Some("1".to_owned())]
|
||||
);
|
||||
assert_eq!(deleted, 2);
|
||||
assert_eq!(
|
||||
client.mget::<String>(&["a", "c"]).await.unwrap(),
|
||||
vec![None, None]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn set_nx_only_writes_a_missing_key(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
|
||||
let first = client
|
||||
.set_nx("lock", "owner-1", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
let second = client
|
||||
.set_nx("lock", "owner-2", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(first);
|
||||
assert!(!second);
|
||||
assert_eq!(
|
||||
client.get::<String>("lock").await.unwrap().as_deref(),
|
||||
Some("owner-1")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ttl_reports_missing_persistent_and_remaining_time(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
client.incr_by("persistent", 1).await.unwrap();
|
||||
client
|
||||
.set("expiring", "1", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
client.incr_by_float("expired_later", 0.5).await.unwrap();
|
||||
assert!(
|
||||
client
|
||||
.expire("expired_later", Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing);
|
||||
assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent);
|
||||
assert!(matches!(
|
||||
client.ttl("expiring").await.unwrap(),
|
||||
Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60)
|
||||
));
|
||||
assert!(matches!(
|
||||
client.ttl("expired_later").await.unwrap(),
|
||||
Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30)
|
||||
));
|
||||
assert!(
|
||||
!client
|
||||
.expire("missing", Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn incr_by_float_accumulates_fractions(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
|
||||
client.incr_by_float("spend", 0.25).await.unwrap();
|
||||
|
||||
assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn concurrent_commands_on_one_cloned_client_all_apply(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
|
||||
let tasks: Vec<_> = (0..100)
|
||||
.map(|_| {
|
||||
let client = client.clone();
|
||||
tokio::spawn(async move {
|
||||
for _ in 0..10 {
|
||||
client.incr_by("shared", 1).await.unwrap();
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
for task in tasks {
|
||||
task.await.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(client.get::<i64>("shared").await.unwrap(), Some(1000));
|
||||
}
|
||||
|
||||
fn breaker(failure_threshold: u32) -> Breaker {
|
||||
Breaker::new(
|
||||
Config {
|
||||
failure_threshold,
|
||||
recovery_timeout: Duration::from_millis(200),
|
||||
timeout_min_duration: Duration::ZERO,
|
||||
},
|
||||
Box::new(|_| {}),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let breaker = breaker(2);
|
||||
let client = server
|
||||
.guarded_client("", Duration::from_millis(300), Some(breaker.clone()))
|
||||
.await;
|
||||
|
||||
server.freeze().await;
|
||||
let first = client.get::<String>("k").await;
|
||||
let second = client.get::<String>("k").await;
|
||||
let started = Instant::now();
|
||||
let third = client.get::<String>("k").await;
|
||||
let rejected_in = started.elapsed();
|
||||
server.thaw().await;
|
||||
|
||||
assert!(matches!(first, Err(Error::Transport(_))), "{first:?}");
|
||||
assert!(matches!(second, Err(Error::Transport(_))), "{second:?}");
|
||||
assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}");
|
||||
assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}");
|
||||
|
||||
let deadline = Instant::now() + Duration::from_secs(30);
|
||||
let recovered = loop {
|
||||
match client.get::<String>("k").await {
|
||||
Ok(value) => break Some(value),
|
||||
Err(_) if Instant::now() < deadline => {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await
|
||||
}
|
||||
Err(_) => break None,
|
||||
}
|
||||
};
|
||||
assert_eq!(recovered, Some(None));
|
||||
assert!(
|
||||
matches!(breaker.state(), State::Closed { .. }),
|
||||
"{:?}",
|
||||
breaker.state()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rejected_commands_never_open_the_breaker(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let breaker = breaker(2);
|
||||
let client = server
|
||||
.guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone()))
|
||||
.await;
|
||||
client.rpush("list", "item").await.unwrap();
|
||||
|
||||
let rejected = [
|
||||
client.incr_by("list", 1).await,
|
||||
client.incr_by("list", 1).await,
|
||||
client.incr_by("list", 1).await,
|
||||
];
|
||||
|
||||
assert!(
|
||||
matches!(rejected[2], Err(Error::Server(_))),
|
||||
"{:?}",
|
||||
rejected[2]
|
||||
);
|
||||
assert!(
|
||||
matches!(breaker.state(), State::Closed { hard: 0, .. }),
|
||||
"{:?}",
|
||||
breaker.state()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_open_breaker_fails_every_pipeline_op_without_sending(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server
|
||||
.guarded_client("", Duration::from_millis(300), Some(breaker(1)))
|
||||
.await;
|
||||
let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1);
|
||||
|
||||
server.freeze().await;
|
||||
let failed = client.run(&pipeline).await;
|
||||
let rejected = client.run(&pipeline).await;
|
||||
server.thaw().await;
|
||||
|
||||
assert!(
|
||||
matches!(failed[0], Err(Error::Transport(_))),
|
||||
"{:?}",
|
||||
failed[0]
|
||||
);
|
||||
assert!(
|
||||
matches!(rejected[0], Err(Error::BreakerOpen)),
|
||||
"{:?}",
|
||||
rejected[0]
|
||||
);
|
||||
assert!(
|
||||
matches!(rejected[1], Err(Error::BreakerOpen)),
|
||||
"{:?}",
|
||||
rejected[1]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_host_endpoint_connects_to_its_database() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let port = server.url().rsplit(':').next().unwrap().parse().unwrap();
|
||||
let connect = |endpoint| {
|
||||
RedisClient::new(RedisConfig {
|
||||
endpoint,
|
||||
namespace: Namespace::default(),
|
||||
timeout: DEFAULT_TIMEOUT,
|
||||
breaker: None,
|
||||
})
|
||||
};
|
||||
let in_db_3 = connect(Endpoint::Host(HostEndpoint {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
db: 3,
|
||||
}))
|
||||
.unwrap();
|
||||
let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap();
|
||||
let via_url = connect(Endpoint::Url(SecretValue::new(format!(
|
||||
"{}/3",
|
||||
server.url()
|
||||
))))
|
||||
.unwrap();
|
||||
|
||||
in_db_3
|
||||
.set("k", "3", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(in_db_0.get::<String>("k").await.unwrap(), None);
|
||||
assert_eq!(
|
||||
via_url.get::<String>("k").await.unwrap().as_deref(),
|
||||
Some("3")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,62 +1,407 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use redis::{Cmd, ToRedisArgs, ToSingleRedisArg};
|
||||
use litellm_python_compat::repr::float_repr;
|
||||
use redis::{
|
||||
Cmd, ErrorKind, FromRedisValue, RedisError, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value,
|
||||
};
|
||||
|
||||
use crate::{Error, Namespace, Operation, Script, Ttl, script::ScriptCall};
|
||||
|
||||
/// How an increment treats the key's TTL, the two ways `RedisCache` increments with one.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Expiry {
|
||||
/// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write.
|
||||
Refresh(Duration),
|
||||
/// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none.
|
||||
/// Pipelined it is lowered to [`Expiry::Refresh`], like `RedisBatch.increment`.
|
||||
KeepExisting(Duration),
|
||||
}
|
||||
|
||||
/// A number written the way redis-py encodes it: an `int` with `str()`, a `float` with `repr()`.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub enum Amount {
|
||||
Int(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl ToRedisArgs for Amount {
|
||||
fn write_redis_args<W: ?Sized + RedisWrite>(&self, out: &mut W) {
|
||||
match self {
|
||||
Self::Int(value) => out.write_arg(value.to_string().as_bytes()),
|
||||
Self::Float(value) => out.write_arg(float_repr(*value).as_bytes()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ToSingleRedisArg for Amount {}
|
||||
|
||||
type Build = Box<dyn FnOnce(Vec<String>) -> Plan + Send>;
|
||||
type Decode<T> = Box<dyn FnOnce(Value) -> Result<T, Error> + Send>;
|
||||
|
||||
/// One Redis command as a value: what it sends alone and pipelined, and how its reply decodes. Every
|
||||
/// executor (`RedisClient::send`, `Pipeline::with`, `RedisBatch::add`) applies its namespace to the
|
||||
/// keys once.
|
||||
#[must_use]
|
||||
pub struct Command<T> {
|
||||
operation: Operation,
|
||||
keys: Vec<String>,
|
||||
idempotent: bool,
|
||||
build: Build,
|
||||
decode: Decode<T>,
|
||||
}
|
||||
|
||||
/// What a command sends once its keys carry the namespace. Only the cases whose alone and pipelined
|
||||
/// forms differ get a variant of their own.
|
||||
pub(crate) enum Plan {
|
||||
/// Answered without a round trip, like an empty `MGET` or `DEL`.
|
||||
Reply(Value),
|
||||
/// Sent together; the reply is the first command's, failing if any failed.
|
||||
Commands(Vec<Cmd>),
|
||||
/// Alone through `invoke_async`, pipelined as `EVALSHA` resent once on `NOSCRIPT`.
|
||||
Script(ScriptCall),
|
||||
/// Alone the increment reads the `TTL` and sends `EXPIRE` only to a key without one, pipelined it
|
||||
/// is the increment then `EXPIRE`.
|
||||
KeepingTtl {
|
||||
increment: Cmd,
|
||||
key: String,
|
||||
ttl: Duration,
|
||||
trips: Trips,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(crate) enum Trips {
|
||||
/// `async_increment`: the increment and its `TTL` in one round trip.
|
||||
Shared,
|
||||
/// `increment_cache`: the increment and its `TTL` each on its own round trip.
|
||||
Separate,
|
||||
}
|
||||
|
||||
/// A command with the namespace applied, ready for an executor to send.
|
||||
pub(crate) struct Op {
|
||||
pub(crate) operation: Operation,
|
||||
pub(crate) key: String,
|
||||
pub(crate) idempotent: bool,
|
||||
pub(crate) plan: Plan,
|
||||
}
|
||||
|
||||
impl Plan {
|
||||
/// What a pipeline sends for this op.
|
||||
pub(crate) fn pipelined(&self) -> Vec<Cmd> {
|
||||
match self {
|
||||
Self::Reply(_) => Vec::new(),
|
||||
Self::Commands(commands) => commands.clone(),
|
||||
Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)],
|
||||
Self::KeepingTtl {
|
||||
increment,
|
||||
key,
|
||||
ttl,
|
||||
..
|
||||
} => vec![increment.clone(), expire_cmd(key, *ttl)],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Command<T> {
|
||||
fn new(
|
||||
operation: Operation,
|
||||
keys: Vec<String>,
|
||||
idempotent: bool,
|
||||
build: impl FnOnce(Vec<String>) -> Plan + Send + 'static,
|
||||
decode: impl FnOnce(Value) -> Result<T, Error> + Send + 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
operation,
|
||||
keys,
|
||||
idempotent,
|
||||
build: Box::new(build),
|
||||
decode: Box::new(decode),
|
||||
}
|
||||
}
|
||||
|
||||
/// Attaches a codec over the decoded reply, such as a store's value format.
|
||||
pub fn map<U: 'static>(self, decode: fn(T) -> Result<U, Error>) -> Command<U>
|
||||
where
|
||||
T: 'static,
|
||||
{
|
||||
let inner = self.decode;
|
||||
Command {
|
||||
operation: self.operation,
|
||||
keys: self.keys,
|
||||
idempotent: self.idempotent,
|
||||
build: self.build,
|
||||
decode: Box::new(move |value| inner(value).and_then(decode)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn prepare(self, namespace: &Namespace) -> (Op, Decode<T>) {
|
||||
let keys: Vec<String> = self.keys.iter().map(|key| namespace.apply(key)).collect();
|
||||
let key = keys.first().cloned().unwrap_or_default();
|
||||
let op = Op {
|
||||
operation: self.operation,
|
||||
key,
|
||||
idempotent: self.idempotent,
|
||||
plan: (self.build)(keys),
|
||||
};
|
||||
(op, self.decode)
|
||||
}
|
||||
}
|
||||
|
||||
fn decode<T: FromRedisValue>(value: Value) -> Result<T, Error> {
|
||||
T::from_redis_value(value).map_err(|error| Error::Decode(error.into()))
|
||||
}
|
||||
|
||||
fn owned(keys: &[&str]) -> Vec<String> {
|
||||
keys.iter().map(|key| (*key).to_owned()).collect()
|
||||
}
|
||||
|
||||
fn single(keys: Vec<String>) -> String {
|
||||
keys.into_iter().next().unwrap_or_default()
|
||||
}
|
||||
|
||||
fn expire_cmd(key: &str, ttl: Duration) -> Cmd {
|
||||
redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn ttl_cmd(key: &str) -> Cmd {
|
||||
redis::cmd("TTL").arg(key).to_owned()
|
||||
}
|
||||
|
||||
/// A reply whose shape a codec over it refuses, as `what` describes.
|
||||
pub fn unexpected_reply(what: &'static str) -> Error {
|
||||
Error::Decode(RedisError::from((ErrorKind::UnexpectedReturnType, what)))
|
||||
}
|
||||
|
||||
pub fn get<T: FromRedisValue + 'static>(key: &str) -> Command<Option<T>> {
|
||||
Command::new(
|
||||
Operation::Command("GET"),
|
||||
vec![key.to_owned()],
|
||||
true,
|
||||
|keys| Plan::Commands(vec![redis::cmd("GET").arg(single(keys)).to_owned()]),
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`, answering one
|
||||
/// value per key of `keys` in their order. No keys is no round trip, since Redis rejects an empty
|
||||
/// `MGET`. Pipelined, every key must share a hash slot on a cluster, and a batch never pipelines
|
||||
/// there.
|
||||
pub fn mget<T: FromRedisValue + 'static>(keys: &[&str]) -> Command<Vec<Option<T>>> {
|
||||
let unique = keys.iter().fold(Vec::<String>::new(), |mut unique, key| {
|
||||
if !unique.iter().any(|seen| seen == key) {
|
||||
unique.push((*key).to_owned());
|
||||
}
|
||||
unique
|
||||
});
|
||||
let positions: Vec<usize> = keys
|
||||
.iter()
|
||||
.filter_map(|key| unique.iter().position(|seen| seen == key))
|
||||
.collect();
|
||||
Command::new(
|
||||
Operation::Command("MGET"),
|
||||
unique,
|
||||
true,
|
||||
|keys| match keys.is_empty() {
|
||||
true => Plan::Reply(Value::Array(Vec::new())),
|
||||
false => Plan::Commands(vec![redis::cmd("MGET").arg(keys).to_owned()]),
|
||||
},
|
||||
move |reply| {
|
||||
let values: Vec<Value> = decode(reply)?;
|
||||
let in_key_order = positions
|
||||
.iter()
|
||||
.map(|&position| values.get(position).cloned())
|
||||
.collect::<Option<Vec<Value>>>()
|
||||
.ok_or_else(|| unexpected_reply("MGET answered fewer values than keys"))?;
|
||||
decode(Value::Array(in_key_order))
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())`
|
||||
pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()> {
|
||||
let value = value.to_redis_args().concat();
|
||||
Command::new(
|
||||
Operation::Command("SET"),
|
||||
vec![key.to_owned()],
|
||||
true,
|
||||
move |keys| {
|
||||
Plan::Commands(vec![
|
||||
redis::cmd("SET")
|
||||
.arg(single(keys))
|
||||
.arg(value)
|
||||
.arg("EX")
|
||||
.arg(ttl.as_secs())
|
||||
.to_owned(),
|
||||
])
|
||||
},
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
/// `SET ... EX ... NX`, `EX` before `NX` as redis-py builds it. `true` when this call wrote the key.
|
||||
pub fn set_nx(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<bool> {
|
||||
let value = value.to_redis_args().concat();
|
||||
Command::new(
|
||||
Operation::Command("SET"),
|
||||
vec![key.to_owned()],
|
||||
false,
|
||||
move |keys| {
|
||||
Plan::Commands(vec![
|
||||
redis::cmd("SET")
|
||||
.arg(single(keys))
|
||||
.arg(value)
|
||||
.arg("EX")
|
||||
.arg(ttl.as_secs())
|
||||
.arg("NX")
|
||||
.to_owned(),
|
||||
])
|
||||
},
|
||||
|value| decode::<Option<String>>(value).map(|reply| reply.is_some()),
|
||||
)
|
||||
}
|
||||
|
||||
/// One multi-key `DEL`, no round trip for no keys. Pipelined, every key must share a hash slot.
|
||||
pub fn del(keys: &[&str]) -> Command<usize> {
|
||||
Command::new(
|
||||
Operation::Command("DEL"),
|
||||
owned(keys),
|
||||
true,
|
||||
|keys| match keys.is_empty() {
|
||||
true => Plan::Reply(Value::Int(0)),
|
||||
false => Plan::Commands(vec![redis::cmd("DEL").arg(keys).to_owned()]),
|
||||
},
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn expire(key: &str, ttl: Duration) -> Command<bool> {
|
||||
Command::new(
|
||||
Operation::Command("EXPIRE"),
|
||||
vec![key.to_owned()],
|
||||
true,
|
||||
move |keys| Plan::Commands(vec![expire_cmd(&single(keys), ttl)]),
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn ttl(key: &str) -> Command<Ttl> {
|
||||
Command::new(
|
||||
Operation::Command("TTL"),
|
||||
vec![key.to_owned()],
|
||||
true,
|
||||
|keys| Plan::Commands(vec![ttl_cmd(&single(keys))]),
|
||||
|value| {
|
||||
Ok(match decode::<i64>(value)? {
|
||||
-2 => Ttl::Missing,
|
||||
-1 => Ttl::Persistent,
|
||||
seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())),
|
||||
})
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no TTL,
|
||||
/// each its own round trip.
|
||||
pub fn incr_by(key: &str, delta: i64, ttl: Duration) -> Command<i64> {
|
||||
Command::new(
|
||||
Operation::Command("INCRBY"),
|
||||
vec![key.to_owned()],
|
||||
false,
|
||||
move |keys| {
|
||||
let key = single(keys);
|
||||
Plan::KeepingTtl {
|
||||
increment: redis::cmd("INCRBY").arg(&key).arg(delta).to_owned(),
|
||||
key,
|
||||
ttl,
|
||||
trips: Trips::Separate,
|
||||
}
|
||||
},
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
/// `INCRBYFLOAT`, with the TTL handled as `expiry` says, or left alone without one. Alone it is
|
||||
/// `RedisCache.async_increment`, pipelined `RedisBatch.increment`.
|
||||
pub fn incr_by_float(key: &str, amount: Amount, expiry: Option<Expiry>) -> Command<f64> {
|
||||
Command::new(
|
||||
Operation::Command("INCRBYFLOAT"),
|
||||
vec![key.to_owned()],
|
||||
false,
|
||||
move |keys| {
|
||||
let key = single(keys);
|
||||
let increment = redis::cmd("INCRBYFLOAT").arg(&key).arg(amount).to_owned();
|
||||
match expiry {
|
||||
None => Plan::Commands(vec![increment]),
|
||||
Some(Expiry::Refresh(ttl)) => {
|
||||
Plan::Commands(vec![increment, expire_cmd(&key, ttl)])
|
||||
}
|
||||
Some(Expiry::KeepExisting(ttl)) => Plan::KeepingTtl {
|
||||
increment,
|
||||
key,
|
||||
ttl,
|
||||
trips: Trips::Shared,
|
||||
},
|
||||
}
|
||||
},
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
/// A registered script, like redis-py's `Script`. Every key it touches must be in `keys`.
|
||||
pub fn eval<T: FromRedisValue + 'static>(
|
||||
script: &Script,
|
||||
keys: &[&str],
|
||||
args: impl ToRedisArgs,
|
||||
) -> Command<T> {
|
||||
let script = script.clone();
|
||||
let args = args.to_redis_args();
|
||||
Command::new(
|
||||
Operation::Script,
|
||||
owned(keys),
|
||||
false,
|
||||
move |keys| Plan::Script(ScriptCall { script, keys, args }),
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the scripts
|
||||
/// `RedisCache` runs that way rather than through `register_script`.
|
||||
pub fn eval_source<T: FromRedisValue + 'static>(
|
||||
source: &'static str,
|
||||
keys: &[&str],
|
||||
args: impl ToRedisArgs,
|
||||
) -> Command<T> {
|
||||
let args = args.to_redis_args();
|
||||
Command::new(
|
||||
Operation::Script,
|
||||
owned(keys),
|
||||
false,
|
||||
move |keys| {
|
||||
Plan::Commands(vec![
|
||||
redis::cmd("EVAL")
|
||||
.arg(source)
|
||||
.arg(keys.len())
|
||||
.arg(keys)
|
||||
.arg(args)
|
||||
.to_owned(),
|
||||
])
|
||||
},
|
||||
decode,
|
||||
)
|
||||
}
|
||||
|
||||
/// The channel is sent as given, without the namespace, like redis-py's `publish`.
|
||||
pub fn publish(channel: &str, payload: &[u8]) -> Command<()> {
|
||||
let command = redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned();
|
||||
Command::new(
|
||||
Operation::Command("PUBLISH"),
|
||||
Vec::new(),
|
||||
false,
|
||||
move |_| Plan::Commands(vec![command]),
|
||||
|_| Ok(()),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn ping() -> Cmd {
|
||||
redis::cmd("PING")
|
||||
}
|
||||
|
||||
pub(crate) fn get(key: String) -> Cmd {
|
||||
redis::cmd("GET").arg(key).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn mget(keys: Vec<String>) -> Cmd {
|
||||
redis::cmd("MGET").arg(keys).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
|
||||
redis::cmd("SET")
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg("PX")
|
||||
.arg(millis(ttl))
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
|
||||
redis::cmd("SET")
|
||||
.arg(key)
|
||||
.arg(value)
|
||||
.arg("NX")
|
||||
.arg("PX")
|
||||
.arg(millis(ttl))
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn del(keys: Vec<String>) -> Cmd {
|
||||
redis::cmd("DEL").arg(keys).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn expire(key: String, ttl: Duration) -> Cmd {
|
||||
redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn ttl(key: String) -> Cmd {
|
||||
redis::cmd("PTTL").arg(key).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn incr_by(key: String, delta: i64) -> Cmd {
|
||||
redis::cmd("INCRBY").arg(key).arg(delta).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd {
|
||||
redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd {
|
||||
redis::cmd("RPUSH").arg(key).arg(values).to_owned()
|
||||
}
|
||||
|
||||
fn millis(duration: Duration) -> u64 {
|
||||
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,118 +1,458 @@
|
|||
use std::time::Duration;
|
||||
use std::{fmt, path::PathBuf, pin::Pin, sync::Arc, time::Duration};
|
||||
|
||||
use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult};
|
||||
use futures_util::{Stream, StreamExt};
|
||||
|
||||
use redis::{
|
||||
BasicAuth, ClientTlsConfig, ConnectionAddr, ConnectionInfo, ErrorKind, IntoConnectionInfo,
|
||||
ProtocolVersion, RedisConnectionInfo, RedisError, RedisResult, StreamingCredentialsProvider,
|
||||
TlsCertificates, TlsMode, io::tcp::TcpSettings,
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::Breaker;
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use crate::Namespace;
|
||||
use pkcs8::EncryptedPrivateKeyInfo;
|
||||
|
||||
use crate::{ConnectError, KeyError, Namespace, RedisEvents};
|
||||
use litellm_redis_types::{Cluster, RedisUrl, Target, TlsOptions};
|
||||
|
||||
// litellm/constants.py REDIS_SOCKET_TIMEOUT, which _coerce_redis_kwargs_types sets when socket_timeout is unset
|
||||
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
// below the idle timeouts of common cloud load balancers (AWS NLB 350s, Azure Load Balancer 240s), which the
|
||||
// OS default of two hours never probes within
|
||||
const KEEPALIVE_IDLE: Duration = Duration::from_secs(60);
|
||||
|
||||
pub struct RedisConfig {
|
||||
pub endpoint: Endpoint,
|
||||
pub target: Target,
|
||||
pub credentials: Credentials,
|
||||
pub namespace: Namespace,
|
||||
pub timeout: Duration,
|
||||
pub timeouts: Timeouts,
|
||||
pub breaker: Option<Breaker>,
|
||||
pub events: Option<Arc<dyn RedisEvents>>,
|
||||
/// Ping every node before the next command once the client sat idle this long, like
|
||||
/// redis-py's `health_check_interval`, so a connection the server dropped silently is
|
||||
/// replaced before a real command pays for it.
|
||||
pub health_check: Option<Duration>,
|
||||
/// Probe a connection with TCP keepalive once it sat idle, like redis-py's `socket_keepalive`.
|
||||
pub keepalive: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Endpoint {
|
||||
Url(SecretValue),
|
||||
Host(HostEndpoint),
|
||||
Cluster(ClusterEndpoint),
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct Timeouts {
|
||||
pub response: Duration,
|
||||
pub connect: Duration,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct HostEndpoint {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
pub tls: bool,
|
||||
pub db: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct ClusterEndpoint {
|
||||
pub nodes: Vec<ClusterNode>,
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
pub tls: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct ClusterNode {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
impl ClusterEndpoint {
|
||||
pub(crate) fn addrs(&self) -> impl Iterator<Item = ConnectionAddr> + '_ {
|
||||
self.nodes
|
||||
.iter()
|
||||
.map(|node| addr(&node.host, node.port, self.tls))
|
||||
impl Timeouts {
|
||||
pub fn new(socket_timeout: Duration, connect_timeout: Option<Duration>) -> Self {
|
||||
Self {
|
||||
response: socket_timeout,
|
||||
connect: connect_timeout.unwrap_or(socket_timeout),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HostEndpoint {
|
||||
pub(crate) fn connection_info(&self) -> RedisResult<ConnectionInfo> {
|
||||
let settings = RedisConnectionInfo::default().set_db(self.db);
|
||||
let settings = match &self.username {
|
||||
Some(username) => settings.set_username(username.expose()),
|
||||
None => settings,
|
||||
};
|
||||
let settings = match &self.password {
|
||||
Some(password) => settings.set_password(password.expose()),
|
||||
None => settings,
|
||||
};
|
||||
Ok(addr(&self.host, self.port, self.tls)
|
||||
.into_connection_info()?
|
||||
.set_redis_settings(settings))
|
||||
impl Default for Timeouts {
|
||||
fn default() -> Self {
|
||||
Self::new(DEFAULT_TIMEOUT, None)
|
||||
}
|
||||
}
|
||||
|
||||
fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr {
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Credentials {
|
||||
Static {
|
||||
username: Option<SecretValue>,
|
||||
password: Option<SecretValue>,
|
||||
},
|
||||
Provider(CredentialsProvider),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct CredentialsProvider(pub Arc<dyn StreamingCredentialsProvider>);
|
||||
|
||||
impl fmt::Debug for CredentialsProvider {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("CredentialsProvider")
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamingCredentialsProvider for CredentialsProvider {
|
||||
fn subscribe(&self) -> Pin<Box<dyn Stream<Item = RedisResult<BasicAuth>> + Send + 'static>> {
|
||||
self.0.subscribe()
|
||||
}
|
||||
}
|
||||
|
||||
impl Credentials {
|
||||
pub const NONE: Self = Self::Static {
|
||||
username: None,
|
||||
password: None,
|
||||
};
|
||||
}
|
||||
|
||||
pub(crate) fn tls_mode(tls: &TlsOptions) -> TlsMode {
|
||||
match tls.verify_certificate {
|
||||
true => TlsMode::Secure,
|
||||
false => TlsMode::Insecure,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn tls_certificates(tls: &TlsOptions) -> Result<TlsCertificates, ConnectError> {
|
||||
let client_tls = match (&tls.certfile, &tls.keyfile) {
|
||||
(Some(certfile), keyfile) => Some(ClientTlsConfig {
|
||||
client_cert: read(certfile)?,
|
||||
client_key: client_key(
|
||||
keyfile.as_ref().unwrap_or(certfile),
|
||||
tls.key_password.as_ref(),
|
||||
)?,
|
||||
}),
|
||||
(None, _) => None,
|
||||
};
|
||||
let from_file: Option<Vec<u8>> = tls.ca_certs.as_ref().map(read).transpose()?;
|
||||
let from_directory: Vec<Vec<u8>> = tls
|
||||
.ca_path
|
||||
.as_ref()
|
||||
.map(read_directory)
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
let configured: Vec<Vec<u8>> = from_file
|
||||
.into_iter()
|
||||
.chain(tls.ca_data.iter().map(|pem| pem.clone().into_bytes()))
|
||||
.chain(from_directory)
|
||||
.collect();
|
||||
Ok(TlsCertificates {
|
||||
client_tls,
|
||||
root_cert: (!configured.is_empty()).then(|| with_system_roots(configured.join(&b'\n'))),
|
||||
})
|
||||
}
|
||||
|
||||
/// Every regular file in `path`, in name order, as redis-py hands `ssl_ca_path` to
|
||||
/// `load_verify_locations(capath=...)`
|
||||
fn read_directory(path: &PathBuf) -> Result<Vec<Vec<u8>>, ConnectError> {
|
||||
let tls_file = |source| ConnectError::TlsFile {
|
||||
path: path.display().to_string(),
|
||||
source,
|
||||
};
|
||||
let mut files: Vec<PathBuf> = std::fs::read_dir(path)
|
||||
.map_err(tls_file)?
|
||||
.map(|entry| entry.map(|entry| entry.path()))
|
||||
.collect::<Result<_, _>>()
|
||||
.map_err(tls_file)?;
|
||||
files.sort();
|
||||
files
|
||||
.iter()
|
||||
.filter(|file| file.is_file())
|
||||
.map(read)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn apply_tls(
|
||||
tls: &TlsOptions,
|
||||
info: ConnectionInfo,
|
||||
) -> Result<ConnectionInfo, ConnectError> {
|
||||
let tls_info = redis::Client::build_with_tls(info, tls_certificates(tls)?)?
|
||||
.get_connection_info()
|
||||
.clone();
|
||||
let mut addr = match tls_info.addr().clone() {
|
||||
ConnectionAddr::TcpTls {
|
||||
host,
|
||||
port,
|
||||
tls_params,
|
||||
..
|
||||
} => ConnectionAddr::TcpTls {
|
||||
host,
|
||||
port,
|
||||
insecure: !tls.verify_certificate,
|
||||
tls_params,
|
||||
},
|
||||
addr => addr,
|
||||
};
|
||||
addr.set_danger_accept_invalid_hostnames(!tls.check_hostname);
|
||||
Ok(tls_info.set_addr(addr))
|
||||
}
|
||||
|
||||
fn with_system_roots(configured: Vec<u8>) -> Vec<u8> {
|
||||
let system = rustls_native_certs::load_native_certs();
|
||||
let system_pem = system
|
||||
.certs
|
||||
.iter()
|
||||
.map(|cert| pem::encode(&pem::Pem::new("CERTIFICATE", cert.to_vec())));
|
||||
configured
|
||||
.into_iter()
|
||||
.chain(b"\n".iter().copied())
|
||||
.chain(system_pem.flat_map(String::into_bytes))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The key as unencrypted PEM, decrypting a PKCS#8 `ENCRYPTED PRIVATE KEY` with `ssl_password` the way
|
||||
/// `load_cert_chain(password=...)` does, since rustls only reads unencrypted keys. Like OpenSSL, a password next to
|
||||
/// an unencrypted key is unused
|
||||
fn client_key(path: &PathBuf, password: Option<&SecretValue>) -> Result<Vec<u8>, ConnectError> {
|
||||
let contents = read(path)?;
|
||||
let key_error = |reason| ConnectError::TlsKey {
|
||||
path: path.display().to_string(),
|
||||
reason,
|
||||
};
|
||||
let Ok(sections) = pem::parse_many(&contents) else {
|
||||
return Ok(contents);
|
||||
};
|
||||
let encrypted = sections
|
||||
.iter()
|
||||
.find(|section| section.tag() == "ENCRYPTED PRIVATE KEY");
|
||||
match (encrypted, password) {
|
||||
(Some(_), None) => Err(key_error(KeyError::MissingPassword)),
|
||||
(Some(section), Some(password)) => {
|
||||
let decrypted = EncryptedPrivateKeyInfo::try_from(section.contents())
|
||||
.and_then(|info| info.decrypt(password.expose()))
|
||||
.map_err(|source| key_error(KeyError::Decrypt(source)))?;
|
||||
Ok(pem::encode(&pem::Pem::new("PRIVATE KEY", decrypted.as_bytes())).into_bytes())
|
||||
}
|
||||
(None, _) if sections.iter().any(is_legacy_encrypted) => {
|
||||
Err(key_error(KeyError::LegacyEncryption))
|
||||
}
|
||||
(None, _) => Ok(contents),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_legacy_encrypted(section: &pem::Pem) -> bool {
|
||||
section
|
||||
.headers()
|
||||
.get("Proc-Type")
|
||||
.is_some_and(|proc_type| proc_type.contains("ENCRYPTED"))
|
||||
}
|
||||
|
||||
fn read(path: &PathBuf) -> Result<Vec<u8>, ConnectError> {
|
||||
std::fs::read(path).map_err(|source| ConnectError::TlsFile {
|
||||
path: path.display().to_string(),
|
||||
source,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn addr(host: &str, port: u16, tls: Option<&TlsOptions>) -> ConnectionAddr {
|
||||
match tls {
|
||||
true => ConnectionAddr::TcpTls {
|
||||
Some(tls) => ConnectionAddr::TcpTls {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
insecure: false,
|
||||
insecure: !tls.verify_certificate,
|
||||
tls_params: None,
|
||||
},
|
||||
false => ConnectionAddr::Tcp(host.to_owned(), port),
|
||||
None => ConnectionAddr::Tcp(host.to_owned(), port),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn cluster_addrs(cluster: &Cluster) -> impl Iterator<Item = ConnectionAddr> + '_ {
|
||||
cluster
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|node| addr(&node.host, node.port, cluster.tls.as_ref()))
|
||||
}
|
||||
|
||||
pub(crate) fn tcp_settings(keepalive: bool) -> TcpSettings {
|
||||
let settings = TcpSettings::default().set_nodelay(true);
|
||||
match keepalive {
|
||||
true => settings.set_keepalive(socket2::TcpKeepalive::new().with_time(KEEPALIVE_IDLE)),
|
||||
false => settings,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn connection_info(
|
||||
target: &Target,
|
||||
credentials: &Credentials,
|
||||
keepalive: bool,
|
||||
) -> Result<ConnectionInfo, ConnectError> {
|
||||
let (info, tls) = match target {
|
||||
Target::Url(url) => (url_info(url, credentials)?, url.tls.as_ref()),
|
||||
Target::Standalone(host) => {
|
||||
let info = addr(&host.host, host.port, host.tls.as_ref()).into_connection_info()?;
|
||||
let settings = info.redis_settings().clone().set_db(host.db);
|
||||
(info.set_redis_settings(settings), host.tls.as_ref())
|
||||
}
|
||||
Target::Unix(socket) => {
|
||||
let info = ConnectionAddr::Unix(socket.path.clone()).into_connection_info()?;
|
||||
let settings = info.redis_settings().clone().set_db(socket.db);
|
||||
(info.set_redis_settings(settings), None)
|
||||
}
|
||||
Target::Cluster(_) | Target::Sentinel(_) => {
|
||||
unreachable!("cluster and sentinel endpoints build their own connection info")
|
||||
}
|
||||
};
|
||||
let info = with_static_credentials(info, credentials).set_tcp_settings(tcp_settings(keepalive));
|
||||
match (tls, info.addr()) {
|
||||
(Some(tls), ConnectionAddr::TcpTls { .. }) => apply_tls(tls, info),
|
||||
_ => Ok(info),
|
||||
}
|
||||
}
|
||||
|
||||
// parse_url reads ?db= over the path, and the url's userinfo over ?username= and ?password=. A credential
|
||||
// provider supersedes both, as it does in redis-py. ?protocol=3 would switch redis-rs to RESP3, which no
|
||||
// reply decode here expects, so the protocol is ignored as the resolver warns
|
||||
fn url_info(url: &RedisUrl, credentials: &Credentials) -> Result<ConnectionInfo, ConnectError> {
|
||||
let info = url.url.expose().into_connection_info()?;
|
||||
let settings = info
|
||||
.redis_settings()
|
||||
.clone()
|
||||
.set_protocol(ProtocolVersion::RESP2);
|
||||
let settings = match url.db {
|
||||
Some(db) => settings.set_db(db),
|
||||
None => settings,
|
||||
};
|
||||
let info = info.set_redis_settings(settings);
|
||||
Ok(match credentials {
|
||||
Credentials::Static { .. } => with_static_credentials(
|
||||
info,
|
||||
&Credentials::Static {
|
||||
username: url.username.clone(),
|
||||
password: url.password.clone(),
|
||||
},
|
||||
),
|
||||
Credentials::Provider(_) => info,
|
||||
})
|
||||
}
|
||||
|
||||
fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> ConnectionInfo {
|
||||
let settings = static_credentials(info.redis_settings().clone(), credentials);
|
||||
info.set_redis_settings(settings)
|
||||
}
|
||||
|
||||
fn static_credentials(
|
||||
settings: RedisConnectionInfo,
|
||||
credentials: &Credentials,
|
||||
) -> RedisConnectionInfo {
|
||||
let Credentials::Static { username, password } = credentials else {
|
||||
return settings;
|
||||
};
|
||||
let settings = match (username, settings.username()) {
|
||||
(Some(username), None) => settings.set_username(username.expose()),
|
||||
_ => settings,
|
||||
};
|
||||
match (password, settings.password()) {
|
||||
(Some(password), None) => settings.set_password(password.expose()),
|
||||
_ => settings,
|
||||
}
|
||||
}
|
||||
|
||||
/// `settings` with the credentials a connection opened now sends. A provider is asked once for its
|
||||
/// current pair, which is when redis-py asks its `credential_provider`.
|
||||
pub(crate) async fn authenticated(
|
||||
settings: RedisConnectionInfo,
|
||||
credentials: &Credentials,
|
||||
) -> RedisResult<RedisConnectionInfo> {
|
||||
let Credentials::Provider(provider) = credentials else {
|
||||
return Ok(static_credentials(settings, credentials));
|
||||
};
|
||||
let current = provider.subscribe().next().await.ok_or_else(|| {
|
||||
RedisError::from((
|
||||
ErrorKind::AuthenticationFailed,
|
||||
"the Redis credential provider ended before giving credentials",
|
||||
))
|
||||
})??;
|
||||
Ok(settings
|
||||
.set_username(current.username())
|
||||
.set_password(current.password()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_secret_value::SecretValue;
|
||||
use litellm_redis_types::{RedisUrl, UnixSocket};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{ClusterEndpoint, Endpoint, HostEndpoint};
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::url(Endpoint::Url(SecretValue::new("redis://credential-value@cache:6379/0")))]
|
||||
#[case::host(
|
||||
Endpoint::Host(HostEndpoint {
|
||||
host: "cache".to_owned(),
|
||||
port: 6379,
|
||||
username: Some(SecretValue::new("credential-value")),
|
||||
password: Some(SecretValue::new("credential-value")),
|
||||
tls: false,
|
||||
db: 0,
|
||||
})
|
||||
)]
|
||||
#[case::cluster(Endpoint::Cluster(ClusterEndpoint {
|
||||
nodes: Vec::new(),
|
||||
username: Some(SecretValue::new("credential-value")),
|
||||
password: Some(SecretValue::new("credential-value")),
|
||||
tls: false,
|
||||
#[case::host(Target::Standalone(crate::Standalone { host: "redis".into(), port: 6379, tls: None, db: 0 }))]
|
||||
#[case::url(Target::Url(RedisUrl {
|
||||
url: SecretValue::new("redis://cache:6379".to_owned()),
|
||||
db: None,
|
||||
username: None,
|
||||
password: None,
|
||||
tls: None,
|
||||
}))]
|
||||
fn debug_output_never_shows_the_secret(#[case] endpoint: Endpoint) {
|
||||
let debug = format!("{endpoint:?}");
|
||||
fn keepalive_probes_an_idle_socket_only_when_enabled(#[case] target: Target) {
|
||||
let probed = |keepalive| {
|
||||
connection_info(&target, &Credentials::NONE, keepalive)
|
||||
.unwrap()
|
||||
.tcp_settings()
|
||||
.keepalive()
|
||||
.is_some()
|
||||
};
|
||||
|
||||
assert!(!debug.contains("credential-value"), "{debug}");
|
||||
assert!(debug.contains("[REDACTED]"), "{debug}");
|
||||
assert!(probed(true));
|
||||
assert!(!probed(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_unix_socket_replaces_host_port_and_tls_like_redis_py() {
|
||||
let target = Target::Unix(UnixSocket {
|
||||
path: "/tmp/redis.sock".into(),
|
||||
db: 3,
|
||||
});
|
||||
|
||||
let info = connection_info(&target, &Credentials::NONE, true).unwrap();
|
||||
|
||||
assert_eq!(*info.addr(), ConnectionAddr::Unix("/tmp/redis.sock".into()));
|
||||
assert_eq!(info.redis_settings().db(), 3);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::the_query_db_wins_over_the_path("redis://cache:6379/1", Some(3), None, 3, None)]
|
||||
#[case::the_path_db_without_a_query_db("redis://cache:6379/1", None, None, 1, None)]
|
||||
#[case::the_query_password_fills_a_missing_userinfo(
|
||||
"redis://cache:6379",
|
||||
None,
|
||||
Some("query"),
|
||||
0,
|
||||
Some("query")
|
||||
)]
|
||||
#[case::the_userinfo_wins_over_the_query_password(
|
||||
"redis://:userinfo@cache:6379",
|
||||
None,
|
||||
Some("query"),
|
||||
0,
|
||||
Some("userinfo")
|
||||
)]
|
||||
fn a_url_reads_its_query_like_parse_url(
|
||||
#[case] url: &str,
|
||||
#[case] db: Option<i64>,
|
||||
#[case] password: Option<&str>,
|
||||
#[case] expected_db: i64,
|
||||
#[case] expected_password: Option<&str>,
|
||||
) {
|
||||
let target = Target::Url(RedisUrl {
|
||||
url: SecretValue::new(url),
|
||||
tls: None,
|
||||
db,
|
||||
username: None,
|
||||
password: password.map(SecretValue::new),
|
||||
});
|
||||
let static_password = Credentials::Static {
|
||||
username: None,
|
||||
password: Some(SecretValue::new("static")),
|
||||
};
|
||||
|
||||
let info = connection_info(&target, &static_password, true).unwrap();
|
||||
|
||||
assert_eq!(info.redis_settings().db(), expected_db);
|
||||
assert_eq!(
|
||||
info.redis_settings().password(),
|
||||
expected_password.or(Some("static"))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::by_number("redis://cache:6379/?protocol=3")]
|
||||
#[case::by_name("redis://cache:6379/?protocol=resp3")]
|
||||
fn a_url_never_switches_the_protocol_from_resp2(#[case] url: &str) {
|
||||
let target = Target::Url(RedisUrl {
|
||||
url: SecretValue::new(url),
|
||||
tls: None,
|
||||
db: None,
|
||||
username: None,
|
||||
password: None,
|
||||
});
|
||||
|
||||
let info = connection_info(&target, &Credentials::NONE, true).unwrap();
|
||||
|
||||
assert_eq!(info.redis_settings().protocol(), ProtocolVersion::RESP2);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,53 +1,161 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use redis::{
|
||||
Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value,
|
||||
aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig},
|
||||
cluster::{ClusterClientBuilder, ClusterConfig},
|
||||
cluster_async::ClusterConnection,
|
||||
use std::{
|
||||
sync::{
|
||||
Arc, Mutex, PoisonError,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use crate::ClusterEndpoint;
|
||||
use redis::{
|
||||
AsyncConnectionConfig, Cmd, ConnectionAddr, ConnectionInfo, ErrorKind, Pipeline,
|
||||
RedisConnectionInfo, RedisError, RedisFuture, RedisResult, ServerErrorKind, TlsCertificates,
|
||||
TlsMode, Value,
|
||||
aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig, MultiplexedConnection},
|
||||
cluster::{ClusterClientBuilder, ClusterConfig},
|
||||
cluster_async::ClusterConnection,
|
||||
cluster_routing::{MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo},
|
||||
sentinel::{SentinelClient, SentinelClientBuilder, SentinelServerType},
|
||||
};
|
||||
|
||||
const RETRIES: u32 = 3;
|
||||
use crate::{
|
||||
Cluster, ConnectError, Credentials, Sentinel, Target, Timeouts,
|
||||
config::{self, addr},
|
||||
error::Deployment,
|
||||
};
|
||||
|
||||
// redis-py 5.3.1 redis/asyncio/connection.py: without retry options a connection gets Retry(NoBackoff(), 0)
|
||||
const RECONNECT_RETRIES: usize = 0;
|
||||
// redis-py 5.3.1 redis/asyncio/cluster.py: cluster_error_retry_attempts=3
|
||||
const CLUSTER_RETRIES: u32 = 3;
|
||||
// redis-py 5.3.1 redis/asyncio/cluster.py sleeps a flat 0.25s before retrying after CLUSTERDOWN.
|
||||
// redis-rs draws the wait from [min, max), so a 1ms window pins it to 250ms
|
||||
const CLUSTER_RETRY_WAIT_MS: u64 = 250;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum Connection {
|
||||
Standalone(ConnectionManager),
|
||||
Cluster(ClusterConnection),
|
||||
Standalone(Renewing<ConnectionManager>),
|
||||
Cluster(Renewing<ClusterConnection>),
|
||||
Sentinel(SentinelConnection),
|
||||
}
|
||||
|
||||
impl Connection {
|
||||
pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult<Self> {
|
||||
let config = ConnectionManagerConfig::new()
|
||||
.set_connection_timeout(Some(timeout))
|
||||
.set_response_timeout(Some(timeout))
|
||||
.set_number_of_retries(RETRIES as usize);
|
||||
let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?;
|
||||
Ok(Self::Standalone(manager))
|
||||
pub(crate) fn new(
|
||||
target: &Target,
|
||||
credentials: &Credentials,
|
||||
timeouts: Timeouts,
|
||||
keepalive: bool,
|
||||
) -> Result<Self, ConnectError> {
|
||||
match target {
|
||||
Target::Url(_) | Target::Standalone(_) | Target::Unix(_) => Self::standalone(
|
||||
config::connection_info(target, credentials, keepalive)?,
|
||||
credentials,
|
||||
timeouts,
|
||||
),
|
||||
Target::Cluster(cluster) => Self::cluster(cluster, credentials, timeouts, keepalive),
|
||||
Target::Sentinel(sentinel) => Ok(Self::Sentinel(SentinelConnection::new(
|
||||
sentinel,
|
||||
credentials,
|
||||
timeouts,
|
||||
keepalive,
|
||||
)?)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult<Self> {
|
||||
let builder = ClusterClientBuilder::new(endpoint.addrs())
|
||||
.connection_timeout(timeout)
|
||||
.response_timeout(timeout)
|
||||
.retries(RETRIES);
|
||||
let builder = match &endpoint.username {
|
||||
Some(username) => builder.username(username.expose()),
|
||||
pub(crate) fn deployment(&self) -> Deployment {
|
||||
match self {
|
||||
Self::Standalone(_) => Deployment::Standalone,
|
||||
Self::Cluster(_) => Deployment::Cluster,
|
||||
Self::Sentinel(_) => Deployment::Sentinel,
|
||||
}
|
||||
}
|
||||
|
||||
fn standalone(
|
||||
info: ConnectionInfo,
|
||||
credentials: &Credentials,
|
||||
timeouts: Timeouts,
|
||||
) -> Result<Self, ConnectError> {
|
||||
let config = ConnectionManagerConfig::new()
|
||||
.set_connection_timeout(Some(timeouts.connect))
|
||||
.set_response_timeout(Some(timeouts.response))
|
||||
.set_number_of_retries(RECONNECT_RETRIES);
|
||||
let config = match credentials {
|
||||
Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()),
|
||||
Credentials::Static { .. } => config,
|
||||
};
|
||||
let client = redis::Client::open(info)?;
|
||||
Ok(Self::Standalone(Renewing::new(move || {
|
||||
ConnectionManager::new_lazy_with_config(client.clone(), config.clone())
|
||||
})?))
|
||||
}
|
||||
|
||||
fn cluster(
|
||||
endpoint: &Cluster,
|
||||
credentials: &Credentials,
|
||||
timeouts: Timeouts,
|
||||
keepalive: bool,
|
||||
) -> Result<Self, ConnectError> {
|
||||
let builder = ClusterClientBuilder::new(config::cluster_addrs(endpoint))
|
||||
.tcp_settings(config::tcp_settings(keepalive))
|
||||
.connection_timeout(timeouts.connect)
|
||||
.response_timeout(timeouts.response)
|
||||
.retries(CLUSTER_RETRIES)
|
||||
.retry_wait_formula(0, 2)
|
||||
.min_retry_wait(CLUSTER_RETRY_WAIT_MS)
|
||||
.max_retry_wait(CLUSTER_RETRY_WAIT_MS + 1);
|
||||
let builder = match &endpoint.tls {
|
||||
Some(tls) => builder
|
||||
.tls(config::tls_mode(tls))
|
||||
.certs(config::tls_certificates(tls)?)
|
||||
.danger_accept_invalid_hostnames(!tls.check_hostname),
|
||||
None => builder,
|
||||
};
|
||||
let builder = match &endpoint.password {
|
||||
Some(password) => builder.password(password.expose()),
|
||||
None => builder,
|
||||
let builder = match credentials {
|
||||
Credentials::Static { username, password } => {
|
||||
let builder = match username {
|
||||
Some(username) => builder.username(username.expose()),
|
||||
None => builder,
|
||||
};
|
||||
match password {
|
||||
Some(password) => builder.password(password.expose()),
|
||||
None => builder,
|
||||
}
|
||||
}
|
||||
Credentials::Provider(provider) => builder.set_credentials_provider(provider.clone()),
|
||||
};
|
||||
let config = ClusterConfig::new()
|
||||
.set_connection_timeout(timeout)
|
||||
.set_response_timeout(timeout);
|
||||
Ok(Self::Cluster(
|
||||
builder
|
||||
.build()?
|
||||
.get_pending_async_connection_with_config(config),
|
||||
))
|
||||
// The builder's response timeout bounds the whole request, but with_config overwrites each node's
|
||||
// response timeout with this one even when unset, so it has to be repeated here
|
||||
let config = ClusterConfig::new().set_response_timeout(timeouts.response);
|
||||
let client = builder.build()?;
|
||||
Ok(Self::Cluster(Renewing::new(move || {
|
||||
Ok(client.get_pending_async_connection_with_config(config.clone()))
|
||||
})?))
|
||||
}
|
||||
}
|
||||
|
||||
impl Connection {
|
||||
/// `PING` on the one multiplexed connection, or on every cluster node, so a dead socket is
|
||||
/// replaced before a real command uses it.
|
||||
pub(crate) async fn ping_every_node(&self) -> RedisResult<()> {
|
||||
let ping = redis::cmd("PING");
|
||||
match self {
|
||||
Self::Cluster(renewing) => {
|
||||
let (generation, mut connection) = renewing.current();
|
||||
let reply = connection
|
||||
.route_command(
|
||||
ping,
|
||||
RoutingInfo::MultiNode((
|
||||
MultipleNodeRoutingInfo::AllNodes,
|
||||
Some(ResponsePolicy::AllSucceeded),
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
renewing.settle(generation, &reply);
|
||||
reply.map(drop)
|
||||
}
|
||||
Self::Standalone(_) | Self::Sentinel(_) => {
|
||||
ping.query_async::<Value>(&mut self.clone()).await.map(drop)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -56,6 +164,7 @@ impl ConnectionLike for Connection {
|
|||
match self {
|
||||
Self::Standalone(connection) => connection.req_packed_command(cmd),
|
||||
Self::Cluster(connection) => connection.req_packed_command(cmd),
|
||||
Self::Sentinel(connection) => connection.req_packed_command(cmd),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -68,6 +177,7 @@ impl ConnectionLike for Connection {
|
|||
match self {
|
||||
Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count),
|
||||
Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count),
|
||||
Self::Sentinel(connection) => connection.req_packed_commands(pipeline, offset, count),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -75,6 +185,441 @@ impl ConnectionLike for Connection {
|
|||
match self {
|
||||
Self::Standalone(connection) => connection.get_db(),
|
||||
Self::Cluster(connection) => connection.get_db(),
|
||||
Self::Sentinel(connection) => connection.get_db(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A connection swapped for a fresh one once a command on it times out, like redis-py disconnecting
|
||||
/// a connection whose read timed out. redis-rs keeps the socket after a timeout and the cluster
|
||||
/// client retries on it, so a socket that went silent would time out every command until TCP gives
|
||||
/// up on it.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Renewing<C>(Arc<RenewingState<C>>);
|
||||
|
||||
type Open<C> = Box<dyn Fn() -> RedisResult<C> + Send + Sync>;
|
||||
|
||||
struct RenewingState<C> {
|
||||
open: Open<C>,
|
||||
current: Mutex<(u64, C)>,
|
||||
}
|
||||
|
||||
impl<C: Clone> Renewing<C> {
|
||||
fn new(open: impl Fn() -> RedisResult<C> + Send + Sync + 'static) -> RedisResult<Self> {
|
||||
let first = open()?;
|
||||
Ok(Self(Arc::new(RenewingState {
|
||||
open: Box::new(open),
|
||||
current: Mutex::new((0, first)),
|
||||
})))
|
||||
}
|
||||
|
||||
fn current(&self) -> (u64, C) {
|
||||
self.0
|
||||
.current
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Replaces the connection of `generation` when its reply timed out, unless another timeout
|
||||
/// already replaced it.
|
||||
fn settle<T>(&self, generation: u64, reply: &RedisResult<T>) {
|
||||
if !reply.as_ref().is_err_and(RedisError::is_timeout) {
|
||||
return;
|
||||
}
|
||||
let mut current = self
|
||||
.0
|
||||
.current
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
if current.0 != generation {
|
||||
return;
|
||||
}
|
||||
match (self.0.open)() {
|
||||
Ok(fresh) => *current = (generation + 1, fresh),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "replacing a timed-out Redis connection failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<C: ConnectionLike + Clone + Send> ConnectionLike for Renewing<C> {
|
||||
fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
|
||||
Box::pin(async move {
|
||||
let (generation, mut connection) = self.current();
|
||||
let reply = connection.req_packed_command(cmd).await;
|
||||
self.settle(generation, &reply);
|
||||
reply
|
||||
})
|
||||
}
|
||||
|
||||
fn req_packed_commands<'a>(
|
||||
&'a mut self,
|
||||
pipeline: &'a Pipeline,
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> RedisFuture<'a, Vec<Value>> {
|
||||
Box::pin(async move {
|
||||
let (generation, mut connection) = self.current();
|
||||
let replies = connection
|
||||
.req_packed_commands(pipeline, offset, count)
|
||||
.await;
|
||||
self.settle(generation, &replies);
|
||||
replies
|
||||
})
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
self.current().1.get_db()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct SentinelConnection(Arc<SentinelState>);
|
||||
|
||||
struct SentinelState {
|
||||
discovery: tokio::sync::Mutex<()>,
|
||||
current: Mutex<Option<(u64, MultiplexedConnection)>>,
|
||||
generations: AtomicU64,
|
||||
endpoint: Sentinel,
|
||||
credentials: Credentials,
|
||||
tls: Option<(TlsMode, TlsCertificates)>,
|
||||
config: AsyncConnectionConfig,
|
||||
connect_timeout: Duration,
|
||||
keepalive: bool,
|
||||
}
|
||||
|
||||
impl SentinelConnection {
|
||||
fn new(
|
||||
endpoint: &Sentinel,
|
||||
credentials: &Credentials,
|
||||
timeouts: Timeouts,
|
||||
keepalive: bool,
|
||||
) -> Result<Self, ConnectError> {
|
||||
let tls = match &endpoint.tls {
|
||||
Some(tls) => Some((config::tls_mode(tls), config::tls_certificates(tls)?)),
|
||||
None => None,
|
||||
};
|
||||
let state = SentinelState {
|
||||
discovery: tokio::sync::Mutex::new(()),
|
||||
current: Mutex::new(None),
|
||||
generations: AtomicU64::new(0),
|
||||
endpoint: endpoint.clone(),
|
||||
credentials: credentials.clone(),
|
||||
tls,
|
||||
config: async_config(credentials, timeouts),
|
||||
connect_timeout: timeouts.connect,
|
||||
keepalive,
|
||||
};
|
||||
state.sentinel_client(&RedisConnectionInfo::default())?;
|
||||
Ok(Self(Arc::new(state)))
|
||||
}
|
||||
|
||||
fn cached(&self) -> Option<(u64, MultiplexedConnection)> {
|
||||
self.0
|
||||
.current
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
.clone()
|
||||
}
|
||||
|
||||
async fn connection(&self) -> RedisResult<(u64, MultiplexedConnection)> {
|
||||
if let Some(current) = self.cached() {
|
||||
return Ok(current);
|
||||
}
|
||||
let _discovering = self.0.discovery.lock().await;
|
||||
if let Some(current) = self.cached() {
|
||||
return Ok(current);
|
||||
}
|
||||
let connection = self
|
||||
.discover()
|
||||
.await?
|
||||
.get_multiplexed_async_connection_with_config(&self.0.config)
|
||||
.await?;
|
||||
let mut current = self
|
||||
.0
|
||||
.current
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
let generation = self.0.generations.fetch_add(1, Ordering::Relaxed);
|
||||
*current = Some((generation, connection.clone()));
|
||||
Ok((generation, connection))
|
||||
}
|
||||
|
||||
pub(crate) async fn master(&self) -> RedisResult<redis::Client> {
|
||||
self.discover().await
|
||||
}
|
||||
|
||||
async fn discover(&self) -> RedisResult<redis::Client> {
|
||||
let master = tokio::time::timeout(self.0.connect_timeout, self.ask_sentinels())
|
||||
.await
|
||||
.map_err(|_| timed_out("asking the Sentinels for the master timed out"))??;
|
||||
self.master_client(master.get_connection_info().addr())
|
||||
}
|
||||
|
||||
/// redis-rs sends `ROLE` to the master a Sentinel names before trusting it, so that probe gets
|
||||
/// the data credentials and TLS a connection opened now would use. The Sentinel client is built
|
||||
/// per discovery since a provider's credentials change after `new`.
|
||||
async fn ask_sentinels(&self) -> RedisResult<redis::Client> {
|
||||
let data_auth =
|
||||
config::authenticated(RedisConnectionInfo::default(), &self.0.credentials).await?;
|
||||
self.0.sentinel_client(&data_auth)?.async_get_client().await
|
||||
}
|
||||
|
||||
fn master_client(&self, master: &ConnectionAddr) -> Result<redis::Client, RedisError> {
|
||||
let (host, port) = match master {
|
||||
ConnectionAddr::Tcp(host, port) | ConnectionAddr::TcpTls { host, port, .. } => {
|
||||
(host.as_str(), *port)
|
||||
}
|
||||
_ => {
|
||||
return Err(RedisError::from((
|
||||
ErrorKind::InvalidClientConfig,
|
||||
"Sentinel named a master that is not a TCP address",
|
||||
)));
|
||||
}
|
||||
};
|
||||
let endpoint = Target::Standalone(crate::Standalone {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
tls: self.0.endpoint.tls.clone(),
|
||||
db: self.0.endpoint.db,
|
||||
});
|
||||
let info = config::connection_info(&endpoint, &self.0.credentials, self.0.keepalive)
|
||||
.map_err(|error| match error {
|
||||
ConnectError::Redis(error) => error,
|
||||
ConnectError::TlsFile { path, .. } => RedisError::from((
|
||||
ErrorKind::InvalidClientConfig,
|
||||
"reading a Redis TLS file",
|
||||
path,
|
||||
)),
|
||||
ConnectError::TlsKey { path, reason } => RedisError::from((
|
||||
ErrorKind::InvalidClientConfig,
|
||||
"reading the Redis TLS key",
|
||||
format!("{path}: {reason}"),
|
||||
)),
|
||||
})?;
|
||||
redis::Client::open(info)
|
||||
}
|
||||
|
||||
fn forget(&self, generation: u64) {
|
||||
let mut current = self
|
||||
.0
|
||||
.current
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
if current
|
||||
.as_ref()
|
||||
.is_some_and(|(cached, _)| *cached == generation)
|
||||
{
|
||||
*current = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SentinelState {
|
||||
fn sentinel_client(&self, data_auth: &RedisConnectionInfo) -> RedisResult<SentinelClient> {
|
||||
let endpoint = &self.endpoint;
|
||||
let sentinels = endpoint
|
||||
.monitors
|
||||
.iter()
|
||||
.map(|node| addr(&node.host, node.port, endpoint.tls.as_ref()));
|
||||
let builder = SentinelClientBuilder::new(
|
||||
sentinels,
|
||||
&endpoint.service_name,
|
||||
SentinelServerType::Master,
|
||||
)?;
|
||||
let (username, password) = match &endpoint.monitor_auth {
|
||||
Some(auth) => (auth.username.as_ref(), auth.password.as_ref()),
|
||||
None => (None, None),
|
||||
};
|
||||
let builder = match username {
|
||||
Some(username) => builder.set_client_to_sentinel_username(username.expose()),
|
||||
None => builder,
|
||||
};
|
||||
let builder = match password {
|
||||
Some(password) => builder.set_client_to_sentinel_password(password.expose()),
|
||||
None => builder,
|
||||
};
|
||||
let builder = match data_auth.username() {
|
||||
Some(username) => builder.set_client_to_redis_username(username),
|
||||
None => builder,
|
||||
};
|
||||
let builder = match data_auth.password() {
|
||||
Some(password) => builder.set_client_to_redis_password(password),
|
||||
None => builder,
|
||||
};
|
||||
let builder = match &self.tls {
|
||||
Some((mode, certificates)) => builder
|
||||
.set_client_to_sentinel_tls_mode(*mode)
|
||||
.set_client_to_sentinel_certificates(certificates.clone())
|
||||
.set_client_to_redis_tls_mode(*mode)
|
||||
.set_client_to_redis_certificates(certificates.clone()),
|
||||
None => builder,
|
||||
};
|
||||
builder.build()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn async_config(credentials: &Credentials, timeouts: Timeouts) -> AsyncConnectionConfig {
|
||||
let config = AsyncConnectionConfig::new()
|
||||
.set_connection_timeout(Some(timeouts.connect))
|
||||
.set_response_timeout(Some(timeouts.response));
|
||||
match credentials {
|
||||
Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()),
|
||||
Credentials::Static { .. } => config,
|
||||
}
|
||||
}
|
||||
|
||||
fn timed_out(message: &'static str) -> RedisError {
|
||||
RedisError::from(std::io::Error::new(std::io::ErrorKind::TimedOut, message))
|
||||
}
|
||||
|
||||
fn needs_rediscovery(error: &RedisError) -> bool {
|
||||
matches!(
|
||||
crate::Error::classify(error.clone(), Deployment::Sentinel),
|
||||
crate::Error::Transport(_)
|
||||
)
|
||||
}
|
||||
|
||||
fn replied_read_only(values: &[Value]) -> bool {
|
||||
values.iter().any(|value| {
|
||||
matches!(value, Value::ServerError(error) if error.kind() == Some(ServerErrorKind::ReadOnly))
|
||||
})
|
||||
}
|
||||
|
||||
impl ConnectionLike for SentinelConnection {
|
||||
fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
|
||||
Box::pin(async move {
|
||||
let (generation, mut connection) = self.connection().await?;
|
||||
let reply = connection.req_packed_command(cmd).await;
|
||||
let rediscover = match &reply {
|
||||
Ok(value) => replied_read_only(std::slice::from_ref(value)),
|
||||
Err(error) => needs_rediscovery(error),
|
||||
};
|
||||
if rediscover {
|
||||
self.forget(generation);
|
||||
}
|
||||
reply
|
||||
})
|
||||
}
|
||||
|
||||
fn req_packed_commands<'a>(
|
||||
&'a mut self,
|
||||
pipeline: &'a Pipeline,
|
||||
offset: usize,
|
||||
count: usize,
|
||||
) -> RedisFuture<'a, Vec<Value>> {
|
||||
Box::pin(async move {
|
||||
let (generation, mut connection) = self.connection().await?;
|
||||
let replies = connection
|
||||
.req_packed_commands(pipeline, offset, count)
|
||||
.await;
|
||||
let rediscover = match &replies {
|
||||
Ok(values) => replied_read_only(values),
|
||||
Err(error) => needs_rediscovery(error),
|
||||
};
|
||||
if rediscover {
|
||||
self.forget(generation);
|
||||
}
|
||||
replies
|
||||
})
|
||||
}
|
||||
|
||||
fn get_db(&self) -> i64 {
|
||||
self.0.endpoint.db
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::Node;
|
||||
use litellm_redis_testing::{RedisSentinel, SERVICE_NAME};
|
||||
|
||||
fn endpoint(server: &RedisSentinel) -> Sentinel {
|
||||
let (host, port) = server.sentinel();
|
||||
Sentinel {
|
||||
monitors: vec![Node {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
}],
|
||||
service_name: SERVICE_NAME.to_owned(),
|
||||
monitor_auth: None,
|
||||
tls: None,
|
||||
db: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn a_late_failure_does_not_discard_the_reconnected_sentinel_master() {
|
||||
let server = RedisSentinel::start().await.unwrap();
|
||||
let mut client = SentinelConnection::new(
|
||||
&endpoint(&server),
|
||||
&Credentials::NONE,
|
||||
Timeouts::default(),
|
||||
true,
|
||||
)
|
||||
.unwrap();
|
||||
let (old_generation, _) = client.connection().await.unwrap();
|
||||
let original: i64 = redis::cmd("CLIENT")
|
||||
.arg("ID")
|
||||
.query_async(&mut client)
|
||||
.await
|
||||
.unwrap();
|
||||
client.forget(old_generation);
|
||||
let replacement: i64 = redis::cmd("CLIENT")
|
||||
.arg("ID")
|
||||
.query_async(&mut client)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_ne!(original, replacement);
|
||||
|
||||
client.forget(old_generation);
|
||||
let after_late_failure: i64 = redis::cmd("CLIENT")
|
||||
.arg("ID")
|
||||
.query_async(&mut client)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(replacement, after_late_failure);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn a_read_only_reply_to_one_command_rediscovers_the_master() {
|
||||
let server = RedisSentinel::start().await.unwrap();
|
||||
let mut client = SentinelConnection::new(
|
||||
&endpoint(&server),
|
||||
&Credentials::NONE,
|
||||
Timeouts::default(),
|
||||
true,
|
||||
)
|
||||
.unwrap();
|
||||
let replica = redis::Client::open(format!("redis://127.0.0.1:{}", server.replica_port()))
|
||||
.unwrap()
|
||||
.get_multiplexed_async_connection()
|
||||
.await
|
||||
.unwrap();
|
||||
let stale = client.0.generations.fetch_add(1, Ordering::Relaxed);
|
||||
*client.0.current.lock().unwrap() = Some((stale, replica));
|
||||
|
||||
let on_replica = redis::cmd("SET")
|
||||
.arg("k")
|
||||
.arg("1")
|
||||
.exec_async(&mut client)
|
||||
.await;
|
||||
let on_master = redis::cmd("SET")
|
||||
.arg("k")
|
||||
.arg("1")
|
||||
.exec_async(&mut client)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
on_replica.unwrap_err().kind(),
|
||||
ErrorKind::Server(ServerErrorKind::ReadOnly)
|
||||
);
|
||||
on_master.unwrap();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,32 +13,93 @@ pub enum Error {
|
|||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("connecting to Redis")]
|
||||
pub struct ConnectError(#[from] RedisError);
|
||||
pub enum ConnectError {
|
||||
#[error("connecting to Redis")]
|
||||
Redis(#[from] RedisError),
|
||||
#[error("reading the Redis TLS file {path}")]
|
||||
TlsFile {
|
||||
path: String,
|
||||
#[source]
|
||||
source: std::io::Error,
|
||||
},
|
||||
#[error("reading the Redis TLS key {path}")]
|
||||
TlsKey {
|
||||
path: String,
|
||||
#[source]
|
||||
reason: KeyError,
|
||||
},
|
||||
}
|
||||
|
||||
impl From<RedisError> for Error {
|
||||
fn from(error: RedisError) -> Self {
|
||||
match error.kind() {
|
||||
ErrorKind::Server(
|
||||
ServerErrorKind::BusyLoading
|
||||
| ServerErrorKind::ClusterDown
|
||||
| ServerErrorKind::MasterDown
|
||||
| ServerErrorKind::TryAgain
|
||||
| ServerErrorKind::ReadOnly,
|
||||
) => Self::Transport(error),
|
||||
ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error),
|
||||
ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error),
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum KeyError {
|
||||
#[error("the key is encrypted, so ssl_password is required")]
|
||||
MissingPassword,
|
||||
#[error("ssl_password does not decrypt the key, or the key uses an unsupported cipher")]
|
||||
Decrypt(#[source] pkcs8::Error),
|
||||
#[error(
|
||||
"the key uses legacy OpenSSL PEM encryption, which Rust cannot read. Convert it to PKCS#8 with `openssl pkcs8 -topk8`"
|
||||
)]
|
||||
LegacyEncryption,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum Deployment {
|
||||
Standalone,
|
||||
Cluster,
|
||||
Sentinel,
|
||||
}
|
||||
|
||||
// redis-py 5.3.1 redis/_parsers/base.py EXCEPTION_CLASSES["ERR"] entries that subclass ConnectionError
|
||||
const CONNECTION_ERR_MESSAGES: &[&str] = &[
|
||||
"max number of clients reached",
|
||||
"invalid password",
|
||||
"Client sent AUTH, but no password is set",
|
||||
"AUTH <password> called without any password configured for the default user. Are you sure your configuration is correct?",
|
||||
];
|
||||
|
||||
impl Error {
|
||||
pub(crate) fn classify(error: RedisError, deployment: Deployment) -> Self {
|
||||
match (error.kind(), deployment) {
|
||||
(ErrorKind::Parse | ErrorKind::UnexpectedReturnType, _) => Self::Decode(error),
|
||||
(ErrorKind::Server(kind), _) if is_unavailable(kind, deployment) => {
|
||||
Self::Transport(error)
|
||||
}
|
||||
(ErrorKind::Server(ServerErrorKind::ResponseError), _)
|
||||
if error
|
||||
.detail()
|
||||
.is_some_and(|detail| CONNECTION_ERR_MESSAGES.contains(&detail)) =>
|
||||
{
|
||||
Self::Transport(error)
|
||||
}
|
||||
(ErrorKind::Extension, _) if matches!(error.code(), Some("NOAUTH" | "WRONGPASS")) => {
|
||||
Self::Transport(error)
|
||||
}
|
||||
(ErrorKind::Server(_) | ErrorKind::Extension | ErrorKind::InvalidClientConfig, _) => {
|
||||
Self::Server(error)
|
||||
}
|
||||
_ => Self::Transport(error),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn is_unavailable(kind: ServerErrorKind, deployment: Deployment) -> bool {
|
||||
matches!(
|
||||
(kind, deployment),
|
||||
(ServerErrorKind::BusyLoading, _)
|
||||
| (
|
||||
ServerErrorKind::ClusterDown | ServerErrorKind::MasterDown,
|
||||
Deployment::Cluster
|
||||
)
|
||||
| (ServerErrorKind::ReadOnly, Deployment::Sentinel)
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use redis::{RedisError, Value};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::Error;
|
||||
use super::{Deployment, Error};
|
||||
|
||||
fn reply(resp: &str) -> RedisError {
|
||||
match redis::parse_redis_value(resp.as_bytes()).unwrap() {
|
||||
|
|
@ -47,23 +108,101 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
fn is_transport(resp: &str, deployment: Deployment) -> bool {
|
||||
match Error::classify(reply(resp), deployment) {
|
||||
Error::Transport(_) => true,
|
||||
Error::Server(_) => false,
|
||||
other => panic!("{resp:?} classified as {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")]
|
||||
#[case::no_script("-NOSCRIPT No matching script\r\n")]
|
||||
#[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")]
|
||||
#[case::generic("-ERR value is not an integer or out of range\r\n")]
|
||||
#[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")]
|
||||
fn rejections_are_server_errors(#[case] resp: &str) {
|
||||
assert!(matches!(Error::from(reply(resp)), Error::Server(_)));
|
||||
#[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")]
|
||||
#[case::out_of_memory("-OOM command not allowed when used memory > 'maxmemory'.\r\n")]
|
||||
#[case::auth_arity("-ERR wrong number of arguments for 'auth' command\r\n")]
|
||||
#[case::max_clients_with_suffix("-ERR max number of clients reached, try later\r\n")]
|
||||
fn rejections_are_server_errors_on_every_deployment(
|
||||
#[case] resp: &str,
|
||||
#[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)]
|
||||
deployment: Deployment,
|
||||
) {
|
||||
assert!(!is_transport(resp, deployment));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::loading("-LOADING Redis is loading the dataset in memory\r\n")]
|
||||
#[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")]
|
||||
#[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")]
|
||||
#[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")]
|
||||
#[case::read_only("-READONLY You can't write against a read only replica.\r\n")]
|
||||
fn unavailable_nodes_are_transport_errors(#[case] resp: &str) {
|
||||
assert!(matches!(Error::from(reply(resp)), Error::Transport(_)));
|
||||
#[case::no_auth("-NOAUTH Authentication required.\r\n")]
|
||||
#[case::wrong_pass("-WRONGPASS invalid username-password pair or user is disabled.\r\n")]
|
||||
#[case::max_clients("-ERR max number of clients reached\r\n")]
|
||||
#[case::invalid_password("-ERR invalid password\r\n")]
|
||||
#[case::no_password_set("-ERR Client sent AUTH, but no password is set\r\n")]
|
||||
#[case::no_default_password(
|
||||
"-ERR AUTH <password> called without any password configured for the default user. Are you sure your configuration is correct?\r\n"
|
||||
)]
|
||||
fn unavailable_nodes_are_transport_errors_on_every_deployment(
|
||||
#[case] resp: &str,
|
||||
#[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)]
|
||||
deployment: Deployment,
|
||||
) {
|
||||
assert!(is_transport(resp, deployment));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::master_down_on_standalone(
|
||||
"-MASTERDOWN Link with MASTER is down\r\n",
|
||||
Deployment::Standalone,
|
||||
false
|
||||
)]
|
||||
#[case::master_down_on_cluster(
|
||||
"-MASTERDOWN Link with MASTER is down\r\n",
|
||||
Deployment::Cluster,
|
||||
true
|
||||
)]
|
||||
#[case::cluster_down_on_standalone(
|
||||
"-CLUSTERDOWN The cluster is down\r\n",
|
||||
Deployment::Standalone,
|
||||
false
|
||||
)]
|
||||
#[case::cluster_down_on_cluster(
|
||||
"-CLUSTERDOWN The cluster is down\r\n",
|
||||
Deployment::Cluster,
|
||||
true
|
||||
)]
|
||||
#[case::read_only_on_standalone(
|
||||
"-READONLY You can't write against a read only replica.\r\n",
|
||||
Deployment::Standalone,
|
||||
false
|
||||
)]
|
||||
#[case::read_only_on_cluster(
|
||||
"-READONLY You can't write against a read only replica.\r\n",
|
||||
Deployment::Cluster,
|
||||
false
|
||||
)]
|
||||
#[case::master_down_on_sentinel(
|
||||
"-MASTERDOWN Link with MASTER is down\r\n",
|
||||
Deployment::Sentinel,
|
||||
false
|
||||
)]
|
||||
#[case::cluster_down_on_sentinel(
|
||||
"-CLUSTERDOWN The cluster is down\r\n",
|
||||
Deployment::Sentinel,
|
||||
false
|
||||
)]
|
||||
#[case::read_only_on_sentinel(
|
||||
"-READONLY You can't write against a read only replica.\r\n",
|
||||
Deployment::Sentinel,
|
||||
true
|
||||
)]
|
||||
fn topology_decides_like_redis_py(
|
||||
#[case] resp: &str,
|
||||
#[case] deployment: Deployment,
|
||||
#[case] transport: bool,
|
||||
) {
|
||||
assert_eq!(is_transport(resp, deployment), transport);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
78
litellm-rust/crates/redis-client/src/events.rs
Normal file
78
litellm-rust/crates/redis-client/src/events.rs
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
use std::{any::Any, future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_circuit_breaker::{State, Transition};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
/// What the host attached to the calls made inside `tagged`, opaque to the client.
|
||||
pub type CallTag = Arc<dyn Any + Send + Sync>;
|
||||
|
||||
tokio::task_local! {
|
||||
static TAG: CallTag;
|
||||
}
|
||||
|
||||
pub async fn tagged<F: Future>(tag: CallTag, operation: F) -> F::Output {
|
||||
TAG.scope(tag, operation).await
|
||||
}
|
||||
|
||||
/// The tag of the `tagged` scope this task runs in. `RedisEvents::call` runs on the calling task,
|
||||
/// so a sink reads its call's tag here.
|
||||
pub fn current_tag() -> Option<CallTag> {
|
||||
TAG.try_with(Arc::clone).ok()
|
||||
}
|
||||
|
||||
pub trait RedisEvents: Send + Sync {
|
||||
fn call(&self, event: &CallEvent);
|
||||
|
||||
fn breaker(&self, transition: &Transition<State>);
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Operation {
|
||||
Command(&'static str),
|
||||
Script,
|
||||
Pipeline { ops: usize },
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Outcome {
|
||||
Succeeded,
|
||||
Failed(FailureKind),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum FailureKind {
|
||||
Timeout,
|
||||
Transport,
|
||||
Server,
|
||||
Decode,
|
||||
BreakerOpen,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub struct CallEvent {
|
||||
pub operation: Operation,
|
||||
pub duration: Duration,
|
||||
pub outcome: Outcome,
|
||||
}
|
||||
|
||||
impl FailureKind {
|
||||
pub fn of(error: &Error) -> Self {
|
||||
match error {
|
||||
Error::Transport(error) if error.is_timeout() => Self::Timeout,
|
||||
Error::Transport(_) => Self::Transport,
|
||||
Error::Server(_) => Self::Server,
|
||||
Error::Decode(_) => Self::Decode,
|
||||
Error::BreakerOpen => Self::BreakerOpen,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Outcome {
|
||||
pub(crate) fn of<T>(result: &Result<T, Error>) -> Self {
|
||||
match result {
|
||||
Ok(_) => Self::Succeeded,
|
||||
Err(error) => Self::Failed(FailureKind::of(error)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,20 +1,30 @@
|
|||
mod batch;
|
||||
mod client;
|
||||
mod commands;
|
||||
pub mod commands;
|
||||
mod config;
|
||||
mod connection;
|
||||
mod error;
|
||||
mod events;
|
||||
mod namespace;
|
||||
mod pipeline;
|
||||
mod pubsub;
|
||||
mod script;
|
||||
#[cfg(test)]
|
||||
mod test_support;
|
||||
|
||||
pub use client::{RedisClient, Ttl};
|
||||
pub use config::{
|
||||
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig,
|
||||
pub use batch::{Declared, RedisBatch};
|
||||
pub use client::{RedisClient, Ttl, WeakRedisClient};
|
||||
pub use commands::{Amount, Command, Expiry};
|
||||
pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts};
|
||||
pub use error::{ConnectError, Error, KeyError};
|
||||
pub use events::{
|
||||
CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged,
|
||||
};
|
||||
pub use litellm_redis_types::{
|
||||
Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket,
|
||||
};
|
||||
pub use error::{ConnectError, Error};
|
||||
pub use namespace::Namespace;
|
||||
pub use pipeline::Pipeline;
|
||||
pub use redis::Value;
|
||||
pub use redis::{
|
||||
BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value,
|
||||
cluster_routing::Slot,
|
||||
};
|
||||
pub use script::Script;
|
||||
|
|
|
|||
|
|
@ -1,30 +1,20 @@
|
|||
use std::{collections::HashMap, time::Duration};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use futures_util::future::join_all;
|
||||
use redis::{
|
||||
Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value,
|
||||
cluster_routing::Slot,
|
||||
};
|
||||
use redis::{ErrorKind, RedisError, RedisResult, ServerErrorKind, Value, cluster_routing::Slot};
|
||||
|
||||
use crate::{Error, Namespace, Script, commands, connection::Connection};
|
||||
use crate::{
|
||||
Command, Error, Namespace,
|
||||
commands::{Op, Plan},
|
||||
connection::Connection,
|
||||
script::ScriptCall,
|
||||
};
|
||||
|
||||
pub struct Pipeline {
|
||||
namespace: Namespace,
|
||||
ops: Vec<Op>,
|
||||
}
|
||||
|
||||
struct Op {
|
||||
command: Cmd,
|
||||
key: String,
|
||||
script: Option<ScriptCall>,
|
||||
}
|
||||
|
||||
struct ScriptCall {
|
||||
script: Script,
|
||||
keys: Vec<String>,
|
||||
args: Vec<Vec<u8>>,
|
||||
}
|
||||
|
||||
type Settled = (usize, Result<Value, Error>);
|
||||
|
||||
impl Pipeline {
|
||||
|
|
@ -35,60 +25,13 @@ impl Pipeline {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn get(self, key: &str) -> Self {
|
||||
let key = self.namespace.apply(key);
|
||||
self.command(commands::get(key.clone()), key)
|
||||
pub(crate) fn of(namespace: Namespace, ops: Vec<Op>) -> Self {
|
||||
Self { namespace, ops }
|
||||
}
|
||||
|
||||
pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self {
|
||||
let key = self.namespace.apply(key);
|
||||
self.command(commands::set(key.clone(), value, ttl), key)
|
||||
}
|
||||
|
||||
pub fn del(self, key: &str) -> Self {
|
||||
let key = self.namespace.apply(key);
|
||||
self.command(commands::del(vec![key.clone()]), key)
|
||||
}
|
||||
|
||||
pub fn expire(self, key: &str, ttl: Duration) -> Self {
|
||||
let key = self.namespace.apply(key);
|
||||
self.command(commands::expire(key.clone(), ttl), key)
|
||||
}
|
||||
|
||||
pub fn incr_by(self, key: &str, delta: i64) -> Self {
|
||||
let key = self.namespace.apply(key);
|
||||
self.command(commands::incr_by(key.clone(), delta), key)
|
||||
}
|
||||
|
||||
pub fn incr_by_float(self, key: &str, delta: f64) -> Self {
|
||||
let key = self.namespace.apply(key);
|
||||
self.command(commands::incr_by_float(key.clone(), delta), key)
|
||||
}
|
||||
|
||||
pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self {
|
||||
let keys: Vec<String> = keys.iter().map(|key| self.namespace.apply(key)).collect();
|
||||
let args = args.to_redis_args();
|
||||
let op = Op {
|
||||
command: script.evalsha(&keys, &args),
|
||||
key: keys.first().cloned().unwrap_or_default(),
|
||||
script: Some(ScriptCall {
|
||||
script: script.clone(),
|
||||
keys,
|
||||
args,
|
||||
}),
|
||||
};
|
||||
self.push(op)
|
||||
}
|
||||
|
||||
fn command(self, command: Cmd, key: String) -> Self {
|
||||
self.push(Op {
|
||||
command,
|
||||
key,
|
||||
script: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn push(mut self, op: Op) -> Self {
|
||||
/// Appends one command in its pipelined form; `RedisClient::run` answers it in order, undecoded.
|
||||
pub fn with<T>(mut self, command: Command<T>) -> Self {
|
||||
let (op, _) = command.prepare(&self.namespace);
|
||||
self.ops.push(op);
|
||||
self
|
||||
}
|
||||
|
|
@ -102,7 +45,7 @@ impl Pipeline {
|
|||
let settled = join_all(
|
||||
groups
|
||||
.iter()
|
||||
.map(|group| self.run_group(connection.clone(), group)),
|
||||
.map(|batches| self.run_in_order(connection, batches)),
|
||||
)
|
||||
.await;
|
||||
let mut settled: Vec<Settled> = settled.into_iter().flatten().collect();
|
||||
|
|
@ -110,183 +53,131 @@ impl Pipeline {
|
|||
settled.into_iter().map(|(_, result)| result).collect()
|
||||
}
|
||||
|
||||
fn groups(&self, connection: &Connection) -> Vec<Vec<usize>> {
|
||||
async fn run_in_order(&self, connection: &Connection, batches: &[Vec<usize>]) -> Vec<Settled> {
|
||||
let mut settled = Vec::with_capacity(batches.iter().map(Vec::len).sum());
|
||||
for batch in batches {
|
||||
settled.extend(self.run_group(connection.clone(), batch).await);
|
||||
}
|
||||
settled
|
||||
}
|
||||
|
||||
fn groups(&self, connection: &Connection) -> Vec<Vec<Vec<usize>>> {
|
||||
match connection {
|
||||
_ if self.ops.is_empty() => Vec::new(),
|
||||
Connection::Standalone(_) => vec![(0..self.ops.len()).collect()],
|
||||
Connection::Standalone(_) | Connection::Sentinel(_) => {
|
||||
vec![vec![(0..self.ops.len()).collect()]]
|
||||
}
|
||||
Connection::Cluster(_) => {
|
||||
let mut by_slot: HashMap<Slot, Vec<usize>> = HashMap::new();
|
||||
for (index, op) in self.ops.iter().enumerate() {
|
||||
for index in 0..self.ops.len() {
|
||||
by_slot
|
||||
.entry(Slot::for_key(&op.key))
|
||||
.entry(Slot::for_key(&self.ops[index].key))
|
||||
.or_default()
|
||||
.push(index);
|
||||
}
|
||||
by_slot.into_values().collect()
|
||||
by_slot
|
||||
.into_values()
|
||||
.map(|indices| self.retry_safe_batches(indices))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn retry_safe_batches(&self, indices: Vec<usize>) -> Vec<Vec<usize>> {
|
||||
indices
|
||||
.chunk_by(|&left, &right| self.ops[left].idempotent && self.ops[right].idempotent)
|
||||
.map(<[usize]>::to_vec)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A script the server had not cached ran nothing, so it alone runs again, the way litellm's
|
||||
/// `RedisBatch` reruns it through redis-py's `Script`: `EVALSHA`, then `SCRIPT LOAD` and
|
||||
/// `EVALSHA` on `NOSCRIPT`.
|
||||
async fn settle_script(
|
||||
&self,
|
||||
connection: Connection,
|
||||
index: usize,
|
||||
reply: RedisResult<Value>,
|
||||
) -> (usize, RedisResult<Value>) {
|
||||
match (&self.ops[index].plan, reply) {
|
||||
(Plan::Script(call), Err(error)) if is_missing_script(&error) => {
|
||||
(index, rerun_missing_script(connection, call).await)
|
||||
}
|
||||
(_, reply) => (index, reply),
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec<Settled> {
|
||||
let deployment = connection.deployment();
|
||||
let mut pipeline = redis::pipe();
|
||||
pipeline.ignore_errors();
|
||||
for &index in indices {
|
||||
pipeline.add_command(self.ops[index].command.clone());
|
||||
}
|
||||
let replies: RedisResult<Vec<RedisResult<Value>>> =
|
||||
pipeline.query_async(&mut connection).await;
|
||||
let widths: Vec<usize> = indices
|
||||
.iter()
|
||||
.map(|&index| {
|
||||
let commands = self.ops[index].plan.pipelined();
|
||||
let width = commands.len();
|
||||
for command in commands {
|
||||
pipeline.add_command(command);
|
||||
}
|
||||
width
|
||||
})
|
||||
.collect();
|
||||
let replies: RedisResult<Vec<RedisResult<Value>>> = match pipeline.is_empty() {
|
||||
true => Ok(Vec::new()),
|
||||
false => pipeline.query_async(&mut connection).await,
|
||||
};
|
||||
let replies = match replies {
|
||||
Ok(replies) => replies,
|
||||
Err(error) => {
|
||||
return indices
|
||||
.iter()
|
||||
.map(|&index| (index, Err(Error::from(error.clone()))))
|
||||
.map(|&index| (index, Err(Error::classify(error.clone(), deployment))))
|
||||
.collect();
|
||||
}
|
||||
};
|
||||
let mut settled = Vec::with_capacity(indices.len());
|
||||
for (&index, reply) in indices.iter().zip(replies) {
|
||||
settled.push((index, self.settle(&mut connection, index, reply).await));
|
||||
}
|
||||
settled
|
||||
}
|
||||
|
||||
async fn settle(
|
||||
&self,
|
||||
connection: &mut Connection,
|
||||
index: usize,
|
||||
reply: RedisResult<Value>,
|
||||
) -> Result<Value, Error> {
|
||||
match (reply, &self.ops[index].script) {
|
||||
(Err(error), Some(call))
|
||||
if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) =>
|
||||
{
|
||||
Ok(call
|
||||
.script
|
||||
.invocation(&call.keys, &call.args)
|
||||
.invoke_async(connection)
|
||||
.await?)
|
||||
}
|
||||
(reply, _) => Ok(reply?),
|
||||
}
|
||||
let mut replies = replies.into_iter();
|
||||
let per_op: Vec<(usize, RedisResult<Value>)> = indices
|
||||
.iter()
|
||||
.zip(widths)
|
||||
.map(|(&index, width)| match &self.ops[index].plan {
|
||||
Plan::Reply(value) => (index, Ok(value.clone())),
|
||||
_ => (index, first_reply(replies.by_ref().take(width))),
|
||||
})
|
||||
.collect();
|
||||
join_all(
|
||||
per_op
|
||||
.into_iter()
|
||||
.map(|(index, reply)| self.settle_script(connection.clone(), index, reply)),
|
||||
)
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|(index, reply)| {
|
||||
(
|
||||
index,
|
||||
reply.map_err(|error| Error::classify(error, deployment)),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use redis::Value;
|
||||
use rstest::rstest;
|
||||
|
||||
use crate::{
|
||||
DEFAULT_TIMEOUT, Error, Script,
|
||||
test_support::{Server, Topology, flush_scripts},
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn each_op_settles_on_its_own_in_input_order(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
client
|
||||
.set("text", "not a number", Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
let pipeline = client
|
||||
.pipeline()
|
||||
.incr_by("a", 1)
|
||||
.incr_by("text", 1)
|
||||
.incr_by("c", 5)
|
||||
.set("b", "2", Duration::from_secs(60))
|
||||
.get("b")
|
||||
.incr_by("a", 1)
|
||||
.get("missing");
|
||||
|
||||
let results = client.run(&pipeline).await;
|
||||
|
||||
assert_eq!(results.len(), 7);
|
||||
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1));
|
||||
assert!(
|
||||
matches!(results[1], Err(Error::Server(_))),
|
||||
"{:?}",
|
||||
results[1]
|
||||
);
|
||||
assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5));
|
||||
assert_eq!(results[3].as_ref().unwrap(), &Value::Okay);
|
||||
assert_eq!(
|
||||
results[4].as_ref().unwrap(),
|
||||
&Value::BulkString(b"2".to_vec())
|
||||
);
|
||||
assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2));
|
||||
assert_eq!(results[6].as_ref().unwrap(), &Value::Nil);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_flushed_script_in_a_pipeline_reruns_alone(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("ns", DEFAULT_TIMEOUT).await;
|
||||
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
|
||||
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
|
||||
flush_scripts(&client).await;
|
||||
let pipeline = client
|
||||
.pipeline()
|
||||
.eval(&incr, &["counter"], 2)
|
||||
.incr_by("other", 7);
|
||||
|
||||
let results = client.run(&pipeline).await;
|
||||
|
||||
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3));
|
||||
assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7));
|
||||
assert_eq!(client.get::<i64>("counter").await.unwrap(), Some(3));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_unreachable_server_fails_every_op_as_transport(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", Duration::from_millis(500)).await;
|
||||
let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1);
|
||||
|
||||
server.freeze().await;
|
||||
let results = client.run(&pipeline).await;
|
||||
server.thaw().await;
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
assert!(
|
||||
matches!(results[0], Err(Error::Transport(_))),
|
||||
"{:?}",
|
||||
results[0]
|
||||
);
|
||||
assert!(
|
||||
matches!(results[1], Err(Error::Transport(_))),
|
||||
"{:?}",
|
||||
results[1]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_empty_pipeline_returns_without_asking_redis(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", Duration::from_secs(2)).await;
|
||||
|
||||
server.freeze().await;
|
||||
let started = Instant::now();
|
||||
let results = client.run(&client.pipeline()).await;
|
||||
let elapsed = started.elapsed();
|
||||
server.thaw().await;
|
||||
|
||||
assert!(results.is_empty());
|
||||
assert!(elapsed < Duration::from_secs(1), "{elapsed:?}");
|
||||
}
|
||||
async fn rerun_missing_script(mut connection: Connection, call: &ScriptCall) -> RedisResult<Value> {
|
||||
call.script
|
||||
.invocation(&call.keys, &call.args)
|
||||
.invoke_async(&mut connection)
|
||||
.await
|
||||
}
|
||||
|
||||
fn is_missing_script(error: &RedisError) -> bool {
|
||||
error.kind() == ErrorKind::Server(ServerErrorKind::NoScript)
|
||||
}
|
||||
|
||||
// RedisBatch settles an op from its first reply and fails it on any reply that failed
|
||||
fn first_reply(replies: impl Iterator<Item = RedisResult<Value>>) -> RedisResult<Value> {
|
||||
replies
|
||||
.reduce(|first, next| match (first, next) {
|
||||
(Err(error), _) | (Ok(_), Err(error)) => Err(error),
|
||||
(first, Ok(_)) => first,
|
||||
})
|
||||
.unwrap_or(Ok(Value::Nil))
|
||||
}
|
||||
|
|
|
|||
225
litellm-rust/crates/redis-client/src/pubsub.rs
Normal file
225
litellm-rust/crates/redis-client/src/pubsub.rs
Normal file
|
|
@ -0,0 +1,225 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use redis::{ConnectionInfo, RedisResult, aio::PubSubStream};
|
||||
|
||||
use crate::{
|
||||
ConnectError, Credentials, Standalone, Target, Timeouts, config,
|
||||
connection::{Connection, SentinelConnection},
|
||||
};
|
||||
|
||||
// litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS
|
||||
const BACKOFF_INITIAL: Duration = Duration::from_secs(5);
|
||||
const BACKOFF_MAX: Duration = Duration::from_secs(60);
|
||||
|
||||
/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends
|
||||
/// `SUBSCRIBE`, authenticated once when it connects.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct Subscriber {
|
||||
source: Source,
|
||||
credentials: Credentials,
|
||||
connect_timeout: Duration,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Source {
|
||||
Nodes(Arc<[ConnectionInfo]>),
|
||||
Sentinel(SentinelConnection),
|
||||
}
|
||||
|
||||
impl Subscriber {
|
||||
pub(crate) fn new(
|
||||
endpoint: &Target,
|
||||
credentials: &Credentials,
|
||||
timeouts: Timeouts,
|
||||
keepalive: bool,
|
||||
connection: &Connection,
|
||||
) -> Result<Self, ConnectError> {
|
||||
let source = match (endpoint, connection) {
|
||||
(_, Connection::Sentinel(sentinel)) => Source::Sentinel(sentinel.clone()),
|
||||
(Target::Cluster(cluster), _) => Source::Nodes(
|
||||
cluster
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|node| {
|
||||
let node = Target::Standalone(Standalone {
|
||||
host: node.host.clone(),
|
||||
port: node.port,
|
||||
tls: cluster.tls.clone(),
|
||||
db: 0,
|
||||
});
|
||||
config::connection_info(&node, credentials, keepalive)
|
||||
})
|
||||
.collect::<Result<_, ConnectError>>()?,
|
||||
),
|
||||
_ => Source::Nodes(Arc::new([config::connection_info(
|
||||
endpoint,
|
||||
credentials,
|
||||
keepalive,
|
||||
)?])),
|
||||
};
|
||||
Ok(Self {
|
||||
source,
|
||||
credentials: credentials.clone(),
|
||||
connect_timeout: timeouts.connect,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec<u8>> {
|
||||
futures_util::stream::unfold(
|
||||
(self, channel, None::<PubSubStream>, BACKOFF_INITIAL),
|
||||
|(subscriber, channel, live, backoff)| async move {
|
||||
let mut live = live;
|
||||
let mut backoff = backoff;
|
||||
loop {
|
||||
let Some(current) = live.as_mut() else {
|
||||
match subscriber.subscribe(&channel).await {
|
||||
Ok(subscribed) => {
|
||||
live = Some(subscribed);
|
||||
backoff = BACKOFF_INITIAL;
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying");
|
||||
tokio::time::sleep(backoff).await;
|
||||
backoff = (backoff * 2).min(BACKOFF_MAX);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
};
|
||||
match current.next().await {
|
||||
Some(message) => {
|
||||
let payload = message.get_payload_bytes().to_vec();
|
||||
return Some((payload, (subscriber, channel, live, backoff)));
|
||||
}
|
||||
None => live = None,
|
||||
}
|
||||
tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing");
|
||||
tokio::time::sleep(backoff).await;
|
||||
backoff = (backoff * 2).min(BACKOFF_MAX);
|
||||
}
|
||||
},
|
||||
)
|
||||
.boxed()
|
||||
}
|
||||
|
||||
async fn subscribe(&self, channel: &str) -> RedisResult<PubSubStream> {
|
||||
match &self.source {
|
||||
Source::Sentinel(sentinel) => {
|
||||
let master = sentinel.master().await?;
|
||||
self.subscribe_on(master.get_connection_info().clone(), channel)
|
||||
.await
|
||||
}
|
||||
Source::Nodes(nodes) => {
|
||||
let mut last = None;
|
||||
for node in nodes.iter() {
|
||||
match self.subscribe_on(node.clone(), channel).await {
|
||||
Ok(live) => return Ok(live),
|
||||
Err(error) => last = Some(error),
|
||||
}
|
||||
}
|
||||
Err(last.expect("an endpoint always has at least one node"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult<PubSubStream> {
|
||||
let settings =
|
||||
config::authenticated(info.redis_settings().clone(), &self.credentials).await?;
|
||||
let client = redis::Client::open(info.set_redis_settings(settings))?;
|
||||
let mut pubsub = tokio::time::timeout(self.connect_timeout, client.get_async_pubsub())
|
||||
.await
|
||||
.map_err(|_| {
|
||||
redis::RedisError::from(std::io::Error::new(
|
||||
std::io::ErrorKind::TimedOut,
|
||||
"connecting the Redis subscription timed out",
|
||||
))
|
||||
})??;
|
||||
subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?;
|
||||
Ok(pubsub.into_on_message())
|
||||
}
|
||||
}
|
||||
|
||||
async fn subscribe_connected(
|
||||
pubsub: &mut redis::aio::PubSub,
|
||||
channel: &str,
|
||||
connect_timeout: Duration,
|
||||
) -> RedisResult<()> {
|
||||
tokio::time::timeout(connect_timeout, pubsub.subscribe(channel))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
redis::RedisError::from(std::io::Error::new(
|
||||
std::io::ErrorKind::TimedOut,
|
||||
"subscribing to the Redis channel timed out",
|
||||
))
|
||||
})?
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_redis_testing::{RedisServer, TcpProxy};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::subscribe_connected;
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_stalled_pubsub_subscribe_times_out_and_a_fresh_pubsub_recovers() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let proxy = TcpProxy::start(server.port()).await.unwrap();
|
||||
let client = redis::Client::open(format!("redis://127.0.0.1:{}", proxy.port())).unwrap();
|
||||
let mut pubsub = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub())
|
||||
.await
|
||||
.expect("connecting the initial pubsub should finish")
|
||||
.unwrap();
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
subscribe_connected(&mut pubsub, "existing", Duration::from_secs(1)),
|
||||
)
|
||||
.await
|
||||
.expect("the initial subscription acknowledgement should arrive")
|
||||
.unwrap();
|
||||
|
||||
proxy.blackhole_open_connections();
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(1),
|
||||
subscribe_connected(&mut pubsub, "stalled", Duration::from_millis(100)),
|
||||
)
|
||||
.await
|
||||
.expect("the stalled subscription should return before the safety deadline");
|
||||
assert!(
|
||||
matches!(&result, Err(error) if error.is_timeout()),
|
||||
"{result:?}"
|
||||
);
|
||||
drop(pubsub);
|
||||
|
||||
let mut recovered = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub())
|
||||
.await
|
||||
.expect("a fresh pubsub connection should finish")
|
||||
.unwrap();
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
subscribe_connected(&mut recovered, "recovered", Duration::from_secs(1)),
|
||||
)
|
||||
.await
|
||||
.expect("the fresh subscription acknowledgement should arrive")
|
||||
.unwrap();
|
||||
|
||||
let publisher = redis::Client::open(server.url()).unwrap();
|
||||
let mut connection = publisher.get_multiplexed_async_connection().await.unwrap();
|
||||
let subscriber_count: i64 = redis::cmd("PUBLISH")
|
||||
.arg("recovered")
|
||||
.arg("payload")
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(subscriber_count, 1);
|
||||
let message = tokio::time::timeout(Duration::from_secs(2), recovered.on_message().next())
|
||||
.await
|
||||
.expect("the fresh subscription should receive the published payload")
|
||||
.unwrap();
|
||||
assert_eq!(message.get_payload::<String>().unwrap(), "payload");
|
||||
}
|
||||
}
|
||||
|
|
@ -1,13 +1,18 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use redis::{Cmd, ScriptInvocation};
|
||||
|
||||
/// A Lua script, cheap to clone into every call that runs it.
|
||||
#[derive(Clone)]
|
||||
pub struct Script(redis::Script);
|
||||
pub struct Script(Arc<redis::Script>);
|
||||
|
||||
impl Script {
|
||||
pub fn new(source: &str) -> Self {
|
||||
Self(redis::Script::new(source))
|
||||
Self(Arc::new(redis::Script::new(source)))
|
||||
}
|
||||
|
||||
/// The `EVALSHA` a pipeline sends, like redis-py. A `NOSCRIPT` reply is settled by
|
||||
/// [`Script::invocation`] alone, which loads the script and sends it again.
|
||||
pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec<u8>]) -> Cmd {
|
||||
redis::cmd("EVALSHA")
|
||||
.arg(self.0.get_hash())
|
||||
|
|
@ -24,65 +29,10 @@ impl Script {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::Script;
|
||||
use crate::{
|
||||
DEFAULT_TIMEOUT, Error,
|
||||
test_support::{Server, Topology, flush_scripts},
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn script_keys_are_namespaced(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("a", DEFAULT_TIMEOUT).await;
|
||||
let bare = server.client("", DEFAULT_TIMEOUT).await;
|
||||
let read = Script::new("return redis.call('GET', KEYS[1])");
|
||||
bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap();
|
||||
|
||||
let value: Option<String> = client
|
||||
.eval(&read, &["x"], Vec::<String>::new())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(value.as_deref(), Some("1"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failing_script_is_a_server_error_and_keeps_the_connection(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
|
||||
client.rpush("list", "item").await.unwrap();
|
||||
|
||||
let failed = client.eval::<i64>(&incr, &["list"], 1).await;
|
||||
|
||||
assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}");
|
||||
assert_eq!(client.eval::<i64>(&incr, &["counter"], 3).await.unwrap(), 3);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_flushed_script_is_loaded_again(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = Server::start(topology).await;
|
||||
let client = server.client("", DEFAULT_TIMEOUT).await;
|
||||
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
|
||||
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
|
||||
|
||||
flush_scripts(&client).await;
|
||||
|
||||
assert_eq!(client.eval::<i64>(&incr, &["counter"], 1).await.unwrap(), 2);
|
||||
}
|
||||
/// A script call in a pipeline, kept apart from plain commands because it is the one entry a
|
||||
/// pipeline may send again.
|
||||
pub(crate) struct ScriptCall {
|
||||
pub(crate) script: Script,
|
||||
pub(crate) keys: Vec<String>,
|
||||
pub(crate) args: Vec<Vec<u8>>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,126 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_circuit_breaker::Breaker;
|
||||
use litellm_redis_testing::{RedisCluster, RedisServer};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr};
|
||||
|
||||
use crate::{
|
||||
ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig,
|
||||
connection::Connection,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub(crate) enum Topology {
|
||||
Standalone,
|
||||
Cluster,
|
||||
}
|
||||
|
||||
pub(crate) enum Server {
|
||||
Standalone(RedisServer),
|
||||
Cluster(RedisCluster),
|
||||
}
|
||||
|
||||
impl Server {
|
||||
pub(crate) async fn start(topology: Topology) -> Self {
|
||||
match topology {
|
||||
Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()),
|
||||
Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn endpoint(&self) -> Endpoint {
|
||||
match self {
|
||||
Self::Standalone(server) => Endpoint::Url(SecretValue::new(server.url().to_owned())),
|
||||
Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
|
||||
nodes: cluster
|
||||
.nodes()
|
||||
.map(|(host, port)| ClusterNode {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
})
|
||||
.collect(),
|
||||
username: None,
|
||||
password: None,
|
||||
tls: false,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient {
|
||||
self.guarded_client(namespace, timeout, None).await
|
||||
}
|
||||
|
||||
pub(crate) async fn guarded_client(
|
||||
&self,
|
||||
namespace: &str,
|
||||
timeout: Duration,
|
||||
breaker: Option<Breaker>,
|
||||
) -> RedisClient {
|
||||
RedisClient::new(RedisConfig {
|
||||
endpoint: self.endpoint(),
|
||||
namespace: Namespace::new(namespace),
|
||||
timeout,
|
||||
breaker,
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub(crate) async fn restart(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.restart().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.restart().await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn stop(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.stop().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.stop().await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resume(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.resume().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.resume().await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn freeze(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.freeze().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.freeze().await.unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn thaw(&self) {
|
||||
match self {
|
||||
Self::Standalone(server) => server.thaw().await.unwrap(),
|
||||
Self::Cluster(cluster) => cluster.thaw().await.unwrap(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn client_id(client: &RedisClient, key: &str) -> i64 {
|
||||
let command = redis::cmd("CLIENT").arg("ID").to_owned();
|
||||
match client.connection.clone() {
|
||||
Connection::Standalone(mut connection) => {
|
||||
command.query_async(&mut connection).await.unwrap()
|
||||
}
|
||||
Connection::Cluster(mut connection) => {
|
||||
let node = Route::with_key(key, SlotAddr::Master);
|
||||
let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node));
|
||||
redis::from_redis_value(connection.route_command(command, routing).await.unwrap())
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn flush_scripts(client: &RedisClient) {
|
||||
redis::cmd("SCRIPT")
|
||||
.arg("FLUSH")
|
||||
.query_async::<()>(&mut client.connection.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
330
litellm-rust/crates/redis-client/tests/batch.rs
Normal file
330
litellm-rust/crates/redis-client/tests/batch.rs
Normal file
|
|
@ -0,0 +1,330 @@
|
|||
mod support;
|
||||
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::{State, Transition};
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{
|
||||
Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Declared, Error, Expiry, Namespace, Operation,
|
||||
RedisBatch, RedisClient, RedisConfig, RedisEvents, Script, Timeouts, Ttl,
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{connect, endpoint};
|
||||
|
||||
const MINUTE: Option<Expiry> = Some(Expiry::KeepExisting(Duration::from_secs(60)));
|
||||
|
||||
// RedisBatch pipelines a standalone flush and runs each op on its own on a cluster
|
||||
fn expected_trips(topology: Topology, alone: &[Operation]) -> Vec<Operation> {
|
||||
match topology {
|
||||
Topology::Standalone => vec![Operation::Pipeline { ops: alone.len() }],
|
||||
Topology::Cluster => alone.to_vec(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RoundTrips(Mutex<Vec<Operation>>);
|
||||
|
||||
impl RoundTrips {
|
||||
fn seen(&self) -> Vec<Operation> {
|
||||
self.0.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisEvents for RoundTrips {
|
||||
fn call(&self, event: &CallEvent) {
|
||||
self.0.lock().unwrap().push(event.operation);
|
||||
}
|
||||
|
||||
fn breaker(&self, _transition: &Transition<State>) {}
|
||||
}
|
||||
|
||||
fn observed(server: &RedisDeployment) -> (RedisClient, Arc<RoundTrips>) {
|
||||
let trips = Arc::new(RoundTrips::default());
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: endpoint(server),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: Namespace::new("ns"),
|
||||
timeouts: Timeouts::default(),
|
||||
breaker: None,
|
||||
events: Some(trips.clone()),
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
(client, trips)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn declared_ops_wait_for_the_flush_and_share_one_round_trip(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = observed(&server);
|
||||
let batch = client.batch();
|
||||
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(3), MINUTE));
|
||||
let total = batch.add(commands::incr_by_float("total", Amount::Float(1.5), MINUTE));
|
||||
let written = batch.add(commands::set("other", "v", Duration::from_secs(60)));
|
||||
let read = batch.add(commands::get::<String>("value"));
|
||||
|
||||
let outside = connect(&server, "ns", DEFAULT_TIMEOUT);
|
||||
assert_eq!(
|
||||
outside.send(commands::get::<i64>("counter")).await.unwrap(),
|
||||
None
|
||||
);
|
||||
assert!(trips.seen().is_empty());
|
||||
|
||||
batch.flush().await;
|
||||
|
||||
let expected = expected_trips(
|
||||
topology,
|
||||
&[
|
||||
Operation::Command("INCRBYFLOAT"),
|
||||
Operation::Command("INCRBYFLOAT"),
|
||||
Operation::Command("SET"),
|
||||
Operation::Command("GET"),
|
||||
],
|
||||
);
|
||||
let mut seen = trips.seen();
|
||||
seen.sort_by_key(|operation| format!("{operation:?}"));
|
||||
let mut sorted = expected.clone();
|
||||
sorted.sort_by_key(|operation| format!("{operation:?}"));
|
||||
assert_eq!(seen, sorted);
|
||||
assert_eq!(counter.resolve().await.unwrap(), 3.0);
|
||||
assert_eq!(total.resolve().await.unwrap(), 1.5);
|
||||
written.resolve().await.unwrap();
|
||||
assert_eq!(read.resolve().await.unwrap(), None);
|
||||
assert_eq!(trips.seen().len(), expected.len());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn resolving_before_a_flush_sends_everything_declared_so_far(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = observed(&server);
|
||||
let batch = client.batch();
|
||||
let first = batch.add(commands::incr_by_float("a", Amount::Int(1), MINUTE));
|
||||
let second = batch.add(commands::incr_by_float("b", Amount::Int(2), MINUTE));
|
||||
let increment = Operation::Command("INCRBYFLOAT");
|
||||
|
||||
assert_eq!(second.resolve().await.unwrap(), 2.0);
|
||||
let late = batch.add(commands::incr_by_float("a", Amount::Int(10), MINUTE));
|
||||
assert_eq!(first.resolve().await.unwrap(), 1.0);
|
||||
assert_eq!(
|
||||
trips.seen(),
|
||||
expected_trips(topology, &[increment, increment])
|
||||
);
|
||||
|
||||
assert_eq!(late.resolve().await.unwrap(), 11.0);
|
||||
let both: Vec<Operation> = expected_trips(topology, &[increment, increment])
|
||||
.into_iter()
|
||||
.chain(expected_trips(topology, &[increment]))
|
||||
.collect();
|
||||
assert_eq!(trips.seen(), both);
|
||||
batch.flush().await;
|
||||
assert_eq!(trips.seen().len(), both.len());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_pipelined_batch_is_one_pipeline_on_every_topology(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = observed(&server);
|
||||
let batch = RedisBatch::pipelined(client);
|
||||
let declared =
|
||||
["a", "b", "c"].map(|key| batch.add(commands::incr_by_float(key, Amount::Int(2), MINUTE)));
|
||||
|
||||
let totals = Declared::all(declared).await;
|
||||
|
||||
assert_eq!(
|
||||
totals.into_iter().map(Result::unwrap).collect::<Vec<_>>(),
|
||||
[2.0; 3]
|
||||
);
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn each_declared_op_gets_its_own_typed_result(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
client
|
||||
.send(commands::set(
|
||||
"text",
|
||||
"not a number",
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = client.batch();
|
||||
let rejected = batch.add(commands::incr_by_float("text", Amount::Int(1), MINUTE));
|
||||
let undecodable = batch.add(commands::get::<i64>("text"));
|
||||
let missing = batch.add(commands::get::<String>("missing"));
|
||||
let script = batch.add(commands::eval::<Vec<String>>(
|
||||
&Script::new("return {KEYS[1], ARGV[1]}"),
|
||||
&["k"],
|
||||
"arg",
|
||||
));
|
||||
let removed = batch.add(commands::del(&["text"]));
|
||||
|
||||
batch.flush().await;
|
||||
|
||||
assert!(matches!(rejected.resolve().await, Err(Error::Server(_))));
|
||||
assert!(matches!(undecodable.resolve().await, Err(Error::Decode(_))));
|
||||
assert_eq!(missing.resolve().await.unwrap(), None);
|
||||
assert_eq!(script.resolve().await.unwrap(), ["k", "arg"]);
|
||||
assert_eq!(removed.resolve().await.unwrap(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_op_whose_flush_was_cancelled_resolves_as_transport(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_secs(5));
|
||||
let batch = client.batch();
|
||||
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE));
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await;
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert!(cancelled.is_err());
|
||||
assert!(matches!(counter.resolve().await, Err(Error::Transport(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = observed(&server);
|
||||
let batch = client.batch();
|
||||
let counter = tokio::spawn(
|
||||
batch
|
||||
.add(commands::incr_by_float("counter", Amount::Int(4), MINUTE))
|
||||
.settled(),
|
||||
);
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
assert!(!counter.is_finished());
|
||||
assert!(trips.seen().is_empty());
|
||||
|
||||
batch.flush().await;
|
||||
|
||||
assert_eq!(counter.await.unwrap().unwrap(), 4.0);
|
||||
assert_eq!(
|
||||
trips.seen(),
|
||||
expected_trips(topology, &[Operation::Command("INCRBYFLOAT")])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() {
|
||||
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_secs(5));
|
||||
let batch = client.batch();
|
||||
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE));
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await;
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert!(cancelled.is_err());
|
||||
assert!(matches!(counter.settled().await, Err(Error::Transport(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn clones_are_the_same_client_and_a_second_connection_is_not() {
|
||||
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let twin = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
|
||||
assert!(client.same_client(&client.clone()));
|
||||
assert!(!client.same_client(&twin));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn mget_reads_each_key_once_and_answers_every_key_in_order(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
client
|
||||
.send(commands::set("a", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
client
|
||||
.send(commands::set("c", "3", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = client.batch();
|
||||
|
||||
let values = batch.add(commands::mget::<String>(&["c", "missing", "a", "c"]));
|
||||
batch.flush().await;
|
||||
|
||||
assert_eq!(
|
||||
values.resolve().await.unwrap(),
|
||||
[
|
||||
Some("3".to_owned()),
|
||||
None,
|
||||
Some("1".to_owned()),
|
||||
Some("3".to_owned())
|
||||
]
|
||||
);
|
||||
let stats = support::command_stats(&server).await;
|
||||
assert!(support::calls(&stats, "mget") > 0, "{stats}");
|
||||
assert_eq!(support::calls(&stats, "get"), 0, "{stats}");
|
||||
assert_eq!(support::keys_read(&server).await, 3);
|
||||
}
|
||||
|
||||
// RedisBatch's _Increment refreshes the TTL in a pipeline, while its cluster path calls
|
||||
// async_increment, which keeps a TTL the key already has
|
||||
#[rstest]
|
||||
#[case::pipelined(Topology::Standalone, false)]
|
||||
#[case::alone_on_a_cluster(Topology::Cluster, true)]
|
||||
#[tokio::test]
|
||||
async fn a_batched_increment_treats_the_ttl_like_redis_batch(
|
||||
#[case] topology: Topology,
|
||||
#[case] keeps_the_ttl: bool,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
client
|
||||
.send(commands::set("counter", "1", Duration::from_secs(30)))
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = client.batch();
|
||||
|
||||
let counter = batch.add(commands::incr_by_float(
|
||||
"counter",
|
||||
Amount::Int(1),
|
||||
Some(Expiry::KeepExisting(Duration::from_secs(600))),
|
||||
));
|
||||
batch.flush().await;
|
||||
|
||||
assert_eq!(counter.resolve().await.unwrap(), 2.0);
|
||||
let Ttl::Expires(left) = client.send(commands::ttl("counter")).await.unwrap() else {
|
||||
panic!("expected an expiring counter")
|
||||
};
|
||||
assert_eq!(left <= Duration::from_secs(30), keeps_the_ttl, "{left:?}");
|
||||
}
|
||||
1039
litellm-rust/crates/redis-client/tests/client.rs
Normal file
1039
litellm-rust/crates/redis-client/tests/client.rs
Normal file
File diff suppressed because it is too large
Load diff
290
litellm-rust/crates/redis-client/tests/config.rs
Normal file
290
litellm-rust/crates/redis-client/tests/config.rs
Normal file
|
|
@ -0,0 +1,290 @@
|
|||
use rstest::rstest;
|
||||
|
||||
use litellm_redis_client::{Credentials, RedisUrl, Sentinel, StaticAuth, Target, TlsOptions};
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
#[rstest]
|
||||
#[case::url(format!("{:?}", Target::Url(RedisUrl { url: SecretValue::new("redis://user:credential-value@cache:6379/0"), tls: None, db: None, username: None, password: None })))]
|
||||
#[case::static_password(format!("{:?}", Credentials::Static { username: None, password: Some(SecretValue::new("credential-value")) }))]
|
||||
#[case::sentinel_password(format!("{:?}", Target::Sentinel(Sentinel {
|
||||
monitors: Vec::new(),
|
||||
service_name: "primary".to_owned(),
|
||||
monitor_auth: Some(StaticAuth { username: None, password: Some(SecretValue::new("credential-value")) }),
|
||||
tls: None,
|
||||
db: 0,
|
||||
})))]
|
||||
fn debug_output_never_shows_the_secret(#[case] debug: String) {
|
||||
assert!(!debug.contains("credential-value"), "{debug}");
|
||||
assert!(debug.contains("[REDACTED]"), "{debug}");
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Ca {
|
||||
Trusted,
|
||||
TrustedData,
|
||||
TrustedDirectory,
|
||||
SystemRoots,
|
||||
}
|
||||
|
||||
fn directory_holding(ca: &std::path::Path) -> std::path::PathBuf {
|
||||
let directory = std::env::temp_dir().join(format!("litellm-ca-dir-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&directory).unwrap();
|
||||
std::fs::copy(ca, directory.join("ca.pem")).unwrap();
|
||||
directory
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum ClientKey {
|
||||
Absent,
|
||||
Plain,
|
||||
Encrypted { password: Option<&'static str> },
|
||||
}
|
||||
|
||||
async fn tls_server(require_client_certificate: bool) -> litellm_redis_testing::RedisTlsServer {
|
||||
if rustls::crypto::CryptoProvider::get_default().is_none() {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
}
|
||||
litellm_redis_testing::RedisTlsServer::start(require_client_certificate)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn tls_client(
|
||||
server: &litellm_redis_testing::RedisTlsServer,
|
||||
host: &str,
|
||||
tls: TlsOptions,
|
||||
) -> Result<litellm_redis_client::RedisClient, litellm_redis_client::ConnectError> {
|
||||
litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig {
|
||||
target: Target::Standalone(litellm_redis_client::Standalone {
|
||||
host: host.to_owned(),
|
||||
port: server.port(),
|
||||
tls: Some(tls),
|
||||
db: 0,
|
||||
}),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: litellm_redis_client::Namespace::default(),
|
||||
timeouts: litellm_redis_client::Timeouts::default(),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
}
|
||||
|
||||
fn tls_options(
|
||||
files: &litellm_redis_testing::TlsFiles,
|
||||
ca: Ca,
|
||||
verify_certificate: bool,
|
||||
check_hostname: bool,
|
||||
client_key: ClientKey,
|
||||
) -> TlsOptions {
|
||||
let (keyfile, key_password) = match client_key {
|
||||
ClientKey::Absent => (None, None),
|
||||
ClientKey::Plain => (Some(files.client_key.clone()), None),
|
||||
ClientKey::Encrypted { password } => (
|
||||
Some(files.encrypted_client_key.clone()),
|
||||
password.map(SecretValue::new),
|
||||
),
|
||||
};
|
||||
TlsOptions {
|
||||
ca_certs: matches!(ca, Ca::Trusted).then(|| files.ca.clone()),
|
||||
ca_data: matches!(ca, Ca::TrustedData).then(|| std::fs::read_to_string(&files.ca).unwrap()),
|
||||
ca_path: matches!(ca, Ca::TrustedDirectory).then(|| directory_holding(&files.ca)),
|
||||
certfile: keyfile.is_some().then(|| files.client_cert.clone()),
|
||||
keyfile,
|
||||
key_password,
|
||||
verify_certificate,
|
||||
check_hostname,
|
||||
}
|
||||
}
|
||||
|
||||
async fn ping_over_tls(
|
||||
require_client_certificate: bool,
|
||||
host: &str,
|
||||
ca: Ca,
|
||||
verify_certificate: bool,
|
||||
check_hostname: bool,
|
||||
client_key: ClientKey,
|
||||
) -> Result<(), litellm_redis_client::Error> {
|
||||
let server = tls_server(require_client_certificate).await;
|
||||
let tls = tls_options(
|
||||
server.files(),
|
||||
ca,
|
||||
verify_certificate,
|
||||
check_hostname,
|
||||
client_key,
|
||||
);
|
||||
tls_client(&server, host, tls).unwrap().ping().await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trusted_ca_without_hostname_check_like_redis_py(
|
||||
"127.0.0.1",
|
||||
Ca::Trusted,
|
||||
true,
|
||||
false,
|
||||
true
|
||||
)]
|
||||
#[case::trusted_ca_as_pem_text_like_ssl_ca_data("127.0.0.1", Ca::TrustedData, true, false, true)]
|
||||
#[case::trusted_ca_in_a_directory_like_ssl_ca_path(
|
||||
"127.0.0.1",
|
||||
Ca::TrustedDirectory,
|
||||
true,
|
||||
false,
|
||||
true
|
||||
)]
|
||||
#[case::hostname_check_with_a_matching_name("localhost", Ca::Trusted, true, true, true)]
|
||||
#[case::hostname_check_against_an_unlisted_address("127.0.0.1", Ca::Trusted, true, true, false)]
|
||||
#[case::an_untrusted_ca_is_refused("127.0.0.1", Ca::SystemRoots, true, false, false)]
|
||||
#[case::cert_reqs_none_skips_verification("127.0.0.1", Ca::SystemRoots, false, false, true)]
|
||||
#[tokio::test]
|
||||
async fn tls_options_decide_which_servers_are_trusted(
|
||||
#[case] host: &str,
|
||||
#[case] ca: Ca,
|
||||
#[case] verify_certificate: bool,
|
||||
#[case] check_hostname: bool,
|
||||
#[case] connects: bool,
|
||||
) {
|
||||
let result = ping_over_tls(
|
||||
false,
|
||||
host,
|
||||
ca,
|
||||
verify_certificate,
|
||||
check_hostname,
|
||||
ClientKey::Absent,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(result.is_ok(), connects, "{result:?}");
|
||||
if let Err(error) = result {
|
||||
assert!(
|
||||
matches!(error, litellm_redis_client::Error::Transport(_)),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::with_the_client_certificate(ClientKey::Plain, true)]
|
||||
#[case::with_an_encrypted_key_and_its_ssl_password(
|
||||
ClientKey::Encrypted { password: Some(litellm_redis_testing::CLIENT_KEY_PASSWORD) },
|
||||
true
|
||||
)]
|
||||
#[case::without_a_client_certificate(ClientKey::Absent, false)]
|
||||
#[tokio::test]
|
||||
async fn mutual_tls_sends_the_configured_client_certificate(
|
||||
#[case] client_key: ClientKey,
|
||||
#[case] connects: bool,
|
||||
) {
|
||||
let result = ping_over_tls(true, "127.0.0.1", Ca::Trusted, true, false, client_key).await;
|
||||
|
||||
assert_eq!(result.is_ok(), connects, "{result:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::without_ssl_password(None, "ssl_password is required")]
|
||||
#[case::with_the_wrong_ssl_password(Some("wrong-password"), "does not decrypt the key")]
|
||||
#[tokio::test]
|
||||
async fn an_encrypted_key_that_cannot_be_decrypted_fails_before_connecting(
|
||||
#[case] password: Option<&'static str>,
|
||||
#[case] reason: &str,
|
||||
) {
|
||||
let server = tls_server(true).await;
|
||||
let tls = tls_options(
|
||||
server.files(),
|
||||
Ca::Trusted,
|
||||
true,
|
||||
false,
|
||||
ClientKey::Encrypted { password },
|
||||
);
|
||||
|
||||
let error = tls_client(&server, "127.0.0.1", tls).err().unwrap();
|
||||
|
||||
let litellm_redis_client::ConnectError::TlsKey {
|
||||
path,
|
||||
reason: cause,
|
||||
} = &error
|
||||
else {
|
||||
panic!("expected a TLS key error, got {error:?}");
|
||||
};
|
||||
assert_eq!(
|
||||
*path,
|
||||
server.files().encrypted_client_key.display().to_string()
|
||||
);
|
||||
assert!(cause.to_string().contains(reason), "{cause}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_missing_tls_file_names_the_path() {
|
||||
let tls = TlsOptions {
|
||||
ca_certs: Some("/nonexistent/ca.pem".into()),
|
||||
..TlsOptions::verified()
|
||||
};
|
||||
|
||||
let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig {
|
||||
target: Target::Standalone(litellm_redis_client::Standalone {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 6379,
|
||||
tls: Some(tls),
|
||||
db: 0,
|
||||
}),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: litellm_redis_client::Namespace::default(),
|
||||
timeouts: litellm_redis_client::Timeouts::default(),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.err()
|
||||
.unwrap();
|
||||
|
||||
assert!(error.to_string().contains("/nonexistent/ca.pem"), "{error}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_legacy_encrypted_pem_key_is_refused_with_a_conversion_hint() {
|
||||
let directory = std::env::temp_dir().join(format!("litellm-legacy-key-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&directory).unwrap();
|
||||
let keyfile = directory.join("client.key");
|
||||
std::fs::write(
|
||||
&keyfile,
|
||||
"-----BEGIN RSA PRIVATE KEY-----\nProc-Type: 4,ENCRYPTED\nDEK-Info: AES-256-CBC,00112233445566778899AABBCCDDEEFF\n\nAAAA\n-----END RSA PRIVATE KEY-----\n",
|
||||
)
|
||||
.unwrap();
|
||||
let tls = TlsOptions {
|
||||
certfile: Some(keyfile.clone()),
|
||||
keyfile: Some(keyfile),
|
||||
key_password: Some(SecretValue::new("legacy-password")),
|
||||
..TlsOptions::verified()
|
||||
};
|
||||
|
||||
let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig {
|
||||
target: Target::Standalone(litellm_redis_client::Standalone {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 6379,
|
||||
tls: Some(tls),
|
||||
db: 0,
|
||||
}),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: litellm_redis_client::Namespace::default(),
|
||||
timeouts: litellm_redis_client::Timeouts::default(),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.err()
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(
|
||||
&error,
|
||||
litellm_redis_client::ConnectError::TlsKey {
|
||||
reason: litellm_redis_client::KeyError::LegacyEncryption,
|
||||
..
|
||||
}
|
||||
),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
329
litellm-rust/crates/redis-client/tests/connection.rs
Normal file
329
litellm-rust/crates/redis-client/tests/connection.rs
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
mod support;
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_testing::{
|
||||
ClientConnection, RedisDeployment, RedisSentinel, RedisServer, SERVICE_NAME, Topology,
|
||||
eventually,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{
|
||||
Credentials, CredentialsProvider, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel,
|
||||
Standalone, Target, Timeouts,
|
||||
};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use support::{Rotating, connect_as, connections, proxied};
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect_uses_the_latest(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&[
|
||||
"ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&[
|
||||
"ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
let provider = Arc::new(Rotating::new("alice", "first"));
|
||||
let client = connect_as(
|
||||
&server,
|
||||
Credentials::Provider(CredentialsProvider(provider.clone())),
|
||||
);
|
||||
client
|
||||
.send(commands::set("k", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
let before = connections(&server).await;
|
||||
|
||||
provider.rotate("bob", "second");
|
||||
let rotated = eventually(async || {
|
||||
let now = connections(&server).await;
|
||||
(users(&now) == BTreeSet::from(["bob"])).then_some(now)
|
||||
})
|
||||
.await;
|
||||
server
|
||||
.cli_on_every_node(&["ACL", "DELUSER", "alice"])
|
||||
.await
|
||||
.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"])
|
||||
.await
|
||||
.unwrap();
|
||||
let after_reconnect = eventually(async || {
|
||||
client
|
||||
.send(commands::get::<String>("k"))
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
})
|
||||
.await;
|
||||
|
||||
let reconnected = connections(&server).await;
|
||||
|
||||
assert_eq!(users(&before), BTreeSet::from(["alice"]));
|
||||
assert_eq!(ids(&rotated), ids(&before));
|
||||
assert_eq!(after_reconnect, "1");
|
||||
assert!(!reconnected.is_empty());
|
||||
assert_eq!(users(&reconnected), BTreeSet::from(["bob"]));
|
||||
}
|
||||
|
||||
fn users(connections: &[ClientConnection]) -> BTreeSet<&str> {
|
||||
connections
|
||||
.iter()
|
||||
.map(|connection| connection.user.as_str())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn ids(connections: &[ClientConnection]) -> Vec<(u16, u64)> {
|
||||
connections
|
||||
.iter()
|
||||
.map(|connection| (connection.node, connection.id))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_out_of_range_database_is_a_server_rejection() {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
let port = server.url().rsplit(':').next().unwrap().parse().unwrap();
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: Target::Standalone(Standalone {
|
||||
host: "127.0.0.1".into(),
|
||||
port,
|
||||
tls: None,
|
||||
db: 999,
|
||||
}),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::new(Duration::from_secs(2), None),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
let result = client.ping().await;
|
||||
assert!(matches!(result, Err(Error::Server(_))), "{result:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_rejected_credential_is_a_transport_failure() {
|
||||
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"])
|
||||
.await
|
||||
.unwrap();
|
||||
let with_password = |password: &str| {
|
||||
connect_as(
|
||||
&server,
|
||||
Credentials::Static {
|
||||
username: None,
|
||||
password: Some(SecretValue::new(password)),
|
||||
},
|
||||
)
|
||||
};
|
||||
with_password("right").ping().await.unwrap();
|
||||
|
||||
let rejected = with_password("wrong").ping().await;
|
||||
|
||||
assert!(matches!(rejected, Err(Error::Transport(_))), "{rejected:?}");
|
||||
}
|
||||
|
||||
fn sentinel_client(sentinel: &RedisSentinel, credentials: Credentials) -> RedisClient {
|
||||
let (host, port) = sentinel.sentinel();
|
||||
RedisClient::new(RedisConfig {
|
||||
target: Target::Sentinel(Sentinel {
|
||||
monitors: vec![Node {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
}],
|
||||
service_name: SERVICE_NAME.to_owned(),
|
||||
monitor_auth: None,
|
||||
tls: None,
|
||||
db: 0,
|
||||
}),
|
||||
credentials,
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::new(Duration::from_secs(1), None),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn sentinel_writes_follow_the_master_through_a_failover() {
|
||||
let sentinel = RedisSentinel::start().await.unwrap();
|
||||
let client = sentinel_client(&sentinel, Credentials::NONE);
|
||||
client
|
||||
.send(commands::set("before", "1", Duration::from_secs(600)))
|
||||
.await
|
||||
.unwrap();
|
||||
eventually(async || {
|
||||
let replica =
|
||||
redis::Client::open(format!("redis://127.0.0.1:{}", sentinel.replica_port())).ok()?;
|
||||
let mut replica = replica.get_multiplexed_async_connection().await.ok()?;
|
||||
let value: Option<String> = redis::cmd("GET")
|
||||
.arg("before")
|
||||
.query_async(&mut replica)
|
||||
.await
|
||||
.ok()?;
|
||||
value
|
||||
})
|
||||
.await;
|
||||
|
||||
let promoted = sentinel.failover().await.unwrap();
|
||||
eventually(async || {
|
||||
client
|
||||
.send(commands::set("after", "2", Duration::from_secs(600)))
|
||||
.await
|
||||
.ok()
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(promoted, sentinel.replica_port());
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>("before"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("1")
|
||||
);
|
||||
let promoted = redis::Client::open(format!("redis://127.0.0.1:{promoted}")).unwrap();
|
||||
let mut promoted = promoted.get_multiplexed_async_connection().await.unwrap();
|
||||
let after: Option<String> = redis::cmd("GET")
|
||||
.arg("after")
|
||||
.query_async(&mut promoted)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(after.as_deref(), Some("2"));
|
||||
}
|
||||
|
||||
const DATA_PASSWORD: &str = "data-secret";
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn sentinel_reaches_a_master_that_requires_a_password() {
|
||||
let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD)
|
||||
.await
|
||||
.unwrap();
|
||||
let client = sentinel_client(
|
||||
&sentinel,
|
||||
Credentials::Static {
|
||||
username: None,
|
||||
password: Some(SecretValue::new(DATA_PASSWORD)),
|
||||
},
|
||||
);
|
||||
|
||||
client
|
||||
.send(commands::set("k", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>("k"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("1")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn sentinel_rediscovers_with_the_credentials_the_provider_holds_now() {
|
||||
let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD)
|
||||
.await
|
||||
.unwrap();
|
||||
sentinel
|
||||
.cli_on_data_nodes(&[
|
||||
"ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
sentinel
|
||||
.cli_on_data_nodes(&[
|
||||
"ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
let provider = Arc::new(Rotating::new("alice", "first"));
|
||||
let client = sentinel_client(
|
||||
&sentinel,
|
||||
Credentials::Provider(CredentialsProvider(provider.clone())),
|
||||
);
|
||||
client
|
||||
.send(commands::set("k", "1", Duration::from_secs(600)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
provider.rotate("bob", "second");
|
||||
sentinel
|
||||
.cli_on_data_nodes(&["ACL", "DELUSER", "alice"])
|
||||
.await
|
||||
.unwrap();
|
||||
sentinel
|
||||
.cli_on_data_nodes(&["CLIENT", "KILL", "USER", "bob"])
|
||||
.await
|
||||
.unwrap();
|
||||
let after_rediscovery = eventually(async || {
|
||||
client
|
||||
.send(commands::get::<String>("k"))
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(after_rediscovery, "1");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_timed_out_command_replaces_a_connection_that_stopped_answering(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let proxies = proxied(&server).await;
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: proxies.target.clone(),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::new(Duration::from_millis(500), None),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
client
|
||||
.send(commands::set("k", "1", Duration::from_secs(600)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
proxies.blackhole_open_connections();
|
||||
let stalled = client.send(commands::get::<String>("k")).await;
|
||||
let next = client.send(commands::get::<String>("k")).await;
|
||||
|
||||
assert!(
|
||||
matches!(&stalled, Err(Error::Transport(error)) if error.is_timeout()),
|
||||
"{stalled:?}"
|
||||
);
|
||||
assert_eq!(next.unwrap().as_deref(), Some("1"));
|
||||
}
|
||||
362
litellm-rust/crates/redis-client/tests/pipeline.rs
Normal file
362
litellm-rust/crates/redis-client/tests/pipeline.rs
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
mod support;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use litellm_redis_client::Value;
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
|
||||
use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Expiry, Script};
|
||||
use support::connect;
|
||||
|
||||
async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) {
|
||||
let (_, first) = cluster.nodes().next().unwrap();
|
||||
let slot = cluster
|
||||
.cli_on_node(first, &["CLUSTER", "KEYSLOT", key])
|
||||
.await
|
||||
.unwrap();
|
||||
let slot = slot.trim();
|
||||
let number: u16 = slot.parse().unwrap();
|
||||
let nodes = cluster
|
||||
.cli_on_node(first, &["CLUSTER", "NODES"])
|
||||
.await
|
||||
.unwrap();
|
||||
let owner = nodes.lines().find(|line| line.split_whitespace().skip(8).any(|range| {
|
||||
let (start, end) = range.split_once('-').unwrap_or((range, range));
|
||||
matches!((start.parse::<u16>(), end.parse::<u16>()), (Ok(start), Ok(end)) if (start..=end).contains(&number))
|
||||
})).unwrap();
|
||||
let target = nodes.lines().find(|line| *line != owner).unwrap();
|
||||
let address = |line: &str| {
|
||||
let value = line.split_whitespace().nth(1).unwrap();
|
||||
value
|
||||
.split_once('@')
|
||||
.unwrap()
|
||||
.0
|
||||
.rsplit_once(':')
|
||||
.unwrap()
|
||||
.1
|
||||
.parse::<u16>()
|
||||
.unwrap()
|
||||
};
|
||||
let source_port = address(owner);
|
||||
let target_port = address(target);
|
||||
let source_id = owner.split_whitespace().next().unwrap();
|
||||
let target_id = target.split_whitespace().next().unwrap();
|
||||
assert_eq!(
|
||||
cluster
|
||||
.cli_on_node(source_port, &["GET", key])
|
||||
.await
|
||||
.unwrap()
|
||||
.trim(),
|
||||
"0"
|
||||
);
|
||||
assert_eq!(
|
||||
cluster
|
||||
.cli_on_node(
|
||||
target_port,
|
||||
&["CLUSTER", "SETSLOT", slot, "IMPORTING", source_id]
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.trim(),
|
||||
"OK"
|
||||
);
|
||||
assert_eq!(
|
||||
cluster
|
||||
.cli_on_node(
|
||||
source_port,
|
||||
&["CLUSTER", "SETSLOT", slot, "MIGRATING", target_id]
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.trim(),
|
||||
"OK"
|
||||
);
|
||||
let state = cluster
|
||||
.cli_on_node(source_port, &["CLUSTER", "NODES"])
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
state.contains(&format!("[{slot}->-{target_id}]")),
|
||||
"{state}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() {
|
||||
let cluster = litellm_redis_testing::RedisCluster::start().await.unwrap();
|
||||
let server = RedisDeployment::Cluster(cluster);
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let counter = "{migration}:counter";
|
||||
client
|
||||
.send(commands::set(counter, 0, Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
let RedisDeployment::Cluster(cluster) = &server else {
|
||||
unreachable!()
|
||||
};
|
||||
begin_migration(cluster, counter).await;
|
||||
let both_keys = Script::new("return redis.call('MGET', KEYS[1], KEYS[2])");
|
||||
|
||||
let results = client
|
||||
.run(
|
||||
&client
|
||||
.pipeline()
|
||||
.with(commands::incr_by_float(
|
||||
counter,
|
||||
Amount::Int(1),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::eval::<Value>(
|
||||
&both_keys,
|
||||
&[counter, "{migration}:missing"],
|
||||
Vec::<String>::new(),
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
client.send(commands::get::<i64>(counter)).await.unwrap(),
|
||||
Some(1)
|
||||
);
|
||||
assert_eq!(
|
||||
results[0].as_ref().unwrap(),
|
||||
&Value::BulkString(b"1".to_vec())
|
||||
);
|
||||
assert!(
|
||||
matches!(&results[1], Err(Error::Server(error)) if error.code() == Some("TRYAGAIN")),
|
||||
"{:?}",
|
||||
results[1]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn each_op_settles_on_its_own_in_input_order(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
client
|
||||
.send(commands::set(
|
||||
"text",
|
||||
"not a number",
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let pipeline = client
|
||||
.pipeline()
|
||||
.with(commands::incr_by_float(
|
||||
"a",
|
||||
Amount::Int(1),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::incr_by_float(
|
||||
"text",
|
||||
Amount::Int(1),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::incr_by_float(
|
||||
"c",
|
||||
Amount::Int(5),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::set("b", "2", Duration::from_secs(60)))
|
||||
.with(commands::get::<Value>("b"))
|
||||
.with(commands::incr_by_float(
|
||||
"a",
|
||||
Amount::Int(1),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::get::<Value>("missing"));
|
||||
|
||||
let results = client.run(&pipeline).await;
|
||||
|
||||
assert_eq!(results.len(), 7);
|
||||
assert_eq!(
|
||||
results[0].as_ref().unwrap(),
|
||||
&Value::BulkString(b"1".to_vec())
|
||||
);
|
||||
assert!(
|
||||
matches!(results[1], Err(Error::Server(_))),
|
||||
"{:?}",
|
||||
results[1]
|
||||
);
|
||||
assert_eq!(
|
||||
results[2].as_ref().unwrap(),
|
||||
&Value::BulkString(b"5".to_vec())
|
||||
);
|
||||
assert_eq!(results[3].as_ref().unwrap(), &Value::Okay);
|
||||
assert_eq!(
|
||||
results[4].as_ref().unwrap(),
|
||||
&Value::BulkString(b"2".to_vec())
|
||||
);
|
||||
assert_eq!(
|
||||
results[5].as_ref().unwrap(),
|
||||
&Value::BulkString(b"2".to_vec())
|
||||
);
|
||||
assert_eq!(results[6].as_ref().unwrap(), &Value::Nil);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::cached(true, Value::BulkString(b"2".to_vec()))]
|
||||
#[case::lost_by_the_server(false, Value::Nil)]
|
||||
#[tokio::test]
|
||||
async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py(
|
||||
#[case] cached: bool,
|
||||
#[case] read_in_the_same_pipeline: Value,
|
||||
) {
|
||||
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let client = connect(&server, "ns", DEFAULT_TIMEOUT);
|
||||
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
|
||||
client
|
||||
.send(commands::eval::<i64>(&incr, &["warmup"], 1))
|
||||
.await
|
||||
.unwrap();
|
||||
if !cached {
|
||||
server
|
||||
.cli_on_every_node(&["SCRIPT", "FLUSH"])
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let pipeline = client
|
||||
.pipeline()
|
||||
.with(commands::eval::<Value>(&incr, &["counter"], 2))
|
||||
.with(commands::get::<Value>("counter"));
|
||||
|
||||
let results = client.run(&pipeline).await;
|
||||
|
||||
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2));
|
||||
assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline);
|
||||
assert_eq!(
|
||||
client.send(commands::get::<i64>("counter")).await.unwrap(),
|
||||
Some(2)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_unreachable_server_fails_every_op_as_transport(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_millis(500));
|
||||
let pipeline = client
|
||||
.pipeline()
|
||||
.with(commands::incr_by_float(
|
||||
"a",
|
||||
Amount::Int(1),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::incr_by_float(
|
||||
"c",
|
||||
Amount::Int(1),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
));
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let results = client.run(&pipeline).await;
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
assert!(
|
||||
matches!(results[0], Err(Error::Transport(_))),
|
||||
"{:?}",
|
||||
results[0]
|
||||
);
|
||||
assert!(
|
||||
matches!(results[1], Err(Error::Transport(_))),
|
||||
"{:?}",
|
||||
results[1]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_empty_pipeline_returns_without_asking_redis(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_secs(2));
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let started = Instant::now();
|
||||
let results = client.run(&client.pipeline()).await;
|
||||
let elapsed = started.elapsed();
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert!(results.is_empty());
|
||||
assert!(elapsed < Duration::from_secs(1), "{elapsed:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_lost(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&deployment, "scripts", DEFAULT_TIMEOUT);
|
||||
let count = Script::new("return redis.call('INCR', KEYS[1])");
|
||||
deployment
|
||||
.cli_on_every_node(&["SCRIPT", "FLUSH"])
|
||||
.await
|
||||
.unwrap();
|
||||
deployment
|
||||
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let first = client
|
||||
.run(
|
||||
&client
|
||||
.pipeline()
|
||||
.with(commands::incr_by_float(
|
||||
"{a}spend",
|
||||
Amount::Float(0.5),
|
||||
Some(Expiry::Refresh(Duration::from_secs(60))),
|
||||
))
|
||||
.with(commands::eval::<Value>(
|
||||
&count,
|
||||
&["{a}"],
|
||||
Vec::<String>::new(),
|
||||
))
|
||||
.with(commands::eval::<Value>(
|
||||
&count,
|
||||
&["b"],
|
||||
Vec::<String>::new(),
|
||||
)),
|
||||
)
|
||||
.await;
|
||||
let second = client
|
||||
.run(&client.pipeline().with(commands::eval::<Value>(
|
||||
&count,
|
||||
&["{a}"],
|
||||
Vec::<String>::new(),
|
||||
)))
|
||||
.await;
|
||||
|
||||
let replies: Vec<Value> = first
|
||||
.into_iter()
|
||||
.chain(second)
|
||||
.map(Result::unwrap)
|
||||
.collect();
|
||||
assert_eq!(
|
||||
replies,
|
||||
[
|
||||
Value::BulkString(b"0.5".to_vec()),
|
||||
Value::Int(1),
|
||||
Value::Int(1),
|
||||
Value::Int(2)
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
client.send(commands::get::<f64>("{a}spend")).await.unwrap(),
|
||||
Some(0.5)
|
||||
);
|
||||
let stats = support::command_stats(&deployment).await;
|
||||
assert!(support::calls(&stats, "evalsha") > 0, "{stats}");
|
||||
assert_eq!(support::calls(&stats, "eval"), 0, "{stats}");
|
||||
}
|
||||
189
litellm-rust/crates/redis-client/tests/pubsub.rs
Normal file
189
litellm-rust/crates/redis-client/tests/pubsub.rs
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
mod support;
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_circuit_breaker::{Breaker, Config};
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{
|
||||
Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel,
|
||||
Target, Timeouts,
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually};
|
||||
use rstest::rstest;
|
||||
use support::{connect, connect_guarded};
|
||||
use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel};
|
||||
|
||||
const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation";
|
||||
|
||||
fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver<Vec<u8>> {
|
||||
let (sender, receiver) = unbounded_channel();
|
||||
let mut messages = client.subscribe(channel);
|
||||
tokio::spawn(async move {
|
||||
while let Some(message) = messages.next().await {
|
||||
if sender.send(message).is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
receiver
|
||||
}
|
||||
|
||||
async fn delivered(
|
||||
publisher: &RedisClient,
|
||||
channel: &str,
|
||||
received: &mut UnboundedReceiver<Vec<u8>>,
|
||||
payload: &[u8],
|
||||
) {
|
||||
eventually(async || {
|
||||
publisher
|
||||
.send(commands::publish(channel, payload))
|
||||
.await
|
||||
.ok()?;
|
||||
let deadline = tokio::time::Instant::now() + Duration::from_millis(200);
|
||||
while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await {
|
||||
if message == payload {
|
||||
return Some(());
|
||||
}
|
||||
}
|
||||
None
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_subscriber_receives_raw_bytes_on_the_channel_name_without_the_namespace(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let publisher = connect(&server, "ns", DEFAULT_TIMEOUT);
|
||||
let plain = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let mut received = listen(&plain, CHANNEL);
|
||||
|
||||
delivered(
|
||||
&publisher,
|
||||
CHANNEL,
|
||||
&mut received,
|
||||
b"{\"cache_key\":\"k\"}\xff",
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_subscriber_resubscribes_after_redis_restarts(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", Duration::from_millis(500));
|
||||
let mut received = listen(&client, CHANNEL);
|
||||
delivered(&client, CHANNEL, &mut received, b"before").await;
|
||||
|
||||
server.restart().await.unwrap();
|
||||
|
||||
delivered(&client, CHANNEL, &mut received, b"after").await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_sentinel_subscriber_receives_from_the_master() {
|
||||
let server = RedisSentinel::start().await.unwrap();
|
||||
let (host, port) = server.sentinel();
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: Target::Sentinel(Sentinel {
|
||||
monitors: vec![Node {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
}],
|
||||
service_name: SERVICE_NAME.to_owned(),
|
||||
monitor_auth: None,
|
||||
tls: None,
|
||||
db: 0,
|
||||
}),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::default(),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
let mut received = listen(&client, CHANNEL);
|
||||
|
||||
delivered(&client, CHANNEL, &mut received, b"payload").await;
|
||||
}
|
||||
|
||||
async fn publish_calls(server: &RedisDeployment) -> Vec<String> {
|
||||
let replies = match server {
|
||||
RedisDeployment::Standalone(server) => {
|
||||
vec![server.cli(&["INFO", "commandstats"]).await.unwrap()]
|
||||
}
|
||||
RedisDeployment::Cluster(cluster) => cluster
|
||||
.cli_outputs(&["INFO", "commandstats"])
|
||||
.await
|
||||
.unwrap(),
|
||||
};
|
||||
replies
|
||||
.iter()
|
||||
.flat_map(|reply| reply.lines())
|
||||
.filter(|line| line.starts_with("cmdstat_publish:"))
|
||||
.map(str::to_owned)
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_open_breaker_rejects_a_publish_without_sending_it(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let breaker = Breaker::new(
|
||||
Config {
|
||||
failure_threshold: 2,
|
||||
recovery_timeout: Duration::from_secs(60),
|
||||
timeout_min_duration: Duration::ZERO,
|
||||
},
|
||||
Box::new(|_| {}),
|
||||
);
|
||||
let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker));
|
||||
|
||||
server.freeze().await.unwrap();
|
||||
let opened = [client.ping().await, client.ping().await];
|
||||
let started = Instant::now();
|
||||
let rejected = client.send(commands::publish(CHANNEL, b"payload")).await;
|
||||
let rejected_in = started.elapsed();
|
||||
server.thaw().await.unwrap();
|
||||
|
||||
assert!(
|
||||
opened
|
||||
.iter()
|
||||
.all(|ping| matches!(ping, Err(Error::Transport(_)))),
|
||||
"{opened:?}"
|
||||
);
|
||||
assert!(matches!(rejected, Err(Error::BreakerOpen)), "{rejected:?}");
|
||||
assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}");
|
||||
assert_eq!(publish_calls(&server).await, Vec::<String>::new());
|
||||
}
|
||||
|
||||
// redis-py's PubSub subscribes on a RESP2 connection, so it never sends HELLO
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_subscription_speaks_resp2_like_redis_py(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
|
||||
.await
|
||||
.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let mut received = listen(&client, CHANNEL);
|
||||
|
||||
delivered(&client, CHANNEL, &mut received, b"evict").await;
|
||||
|
||||
let stats = support::command_stats(&server).await;
|
||||
assert!(support::calls(&stats, "subscribe") > 0, "{stats}");
|
||||
assert_eq!(support::calls(&stats, "hello"), 0, "{stats}");
|
||||
}
|
||||
133
litellm-rust/crates/redis-client/tests/script.rs
Normal file
133
litellm-rust/crates/redis-client/tests/script.rs
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
|
||||
use litellm_redis_client::Script;
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl};
|
||||
use support::connect;
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn script_keys_are_namespaced(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "a", DEFAULT_TIMEOUT);
|
||||
let bare = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let read = Script::new("return redis.call('GET', KEYS[1])");
|
||||
bare.send(commands::set("a:x", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let value: Option<String> = client
|
||||
.send(commands::eval(&read, &["x"], Vec::<String>::new()))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(value.as_deref(), Some("1"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failing_script_is_a_server_error_and_keeps_the_connection(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
|
||||
client
|
||||
.send(commands::set("list", "item", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let failed = client
|
||||
.send(commands::eval::<i64>(&incr, &["list"], 1))
|
||||
.await;
|
||||
|
||||
assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}");
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::eval::<i64>(&incr, &["counter"], 3))
|
||||
.await
|
||||
.unwrap(),
|
||||
3
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_flushed_script_is_loaded_again(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "", DEFAULT_TIMEOUT);
|
||||
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
|
||||
client
|
||||
.send(commands::eval::<i64>(&incr, &["counter"], 1))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
server
|
||||
.cli_on_every_node(&["SCRIPT", "FLUSH"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::eval::<i64>(&incr, &["counter"], 1))
|
||||
.await
|
||||
.unwrap(),
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \
|
||||
if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end \
|
||||
if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end \
|
||||
return count";
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_call(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
#[values(false, true)] pipelined: bool,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
let client = connect(&server, "ns", DEFAULT_TIMEOUT);
|
||||
let floor = Script::new(INCREMENT_WITH_FLOOR);
|
||||
server
|
||||
.cli_on_every_node(&["SCRIPT", "FLUSH"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let count = match pipelined {
|
||||
false => client
|
||||
.send(commands::eval::<i64>(&floor, &["slots"], (-5, 60)))
|
||||
.await
|
||||
.unwrap(),
|
||||
true => {
|
||||
let results = client
|
||||
.run(&client.pipeline().with(commands::eval::<redis::Value>(
|
||||
&floor,
|
||||
&["slots"],
|
||||
(-5, 60),
|
||||
)))
|
||||
.await;
|
||||
redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap()
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(count, 0);
|
||||
assert_eq!(
|
||||
client.send(commands::get::<i64>("slots")).await.unwrap(),
|
||||
Some(0)
|
||||
);
|
||||
assert!(matches!(
|
||||
client.send(commands::ttl("slots")).await.unwrap(),
|
||||
Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60)
|
||||
));
|
||||
}
|
||||
206
litellm-rust/crates/redis-client/tests/support/mod.rs
Normal file
206
litellm-rust/crates/redis-client/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,206 @@
|
|||
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
|
||||
|
||||
use std::{pin::Pin, time::Duration};
|
||||
|
||||
use futures_util::Stream;
|
||||
use litellm_circuit_breaker::Breaker;
|
||||
use litellm_redis_client::{
|
||||
BasicAuth, Cluster, Credentials, Expiry, Namespace, Node, RedisClient, RedisConfig, RedisError,
|
||||
RedisUrl, Standalone, StreamingCredentialsProvider, Target, Timeouts,
|
||||
};
|
||||
use litellm_redis_testing::{ClientConnection, RedisDeployment, TcpProxy};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use tokio::sync::watch;
|
||||
|
||||
pub const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60));
|
||||
|
||||
pub fn endpoint(deployment: &RedisDeployment) -> Target {
|
||||
match deployment {
|
||||
RedisDeployment::Standalone(server) => Target::Url(RedisUrl {
|
||||
url: SecretValue::new(server.url()),
|
||||
tls: None,
|
||||
db: None,
|
||||
username: None,
|
||||
password: None,
|
||||
}),
|
||||
RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster {
|
||||
nodes: cluster
|
||||
.nodes()
|
||||
.map(|(host, port)| Node {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
})
|
||||
.collect(),
|
||||
tls: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn connect(deployment: &RedisDeployment, namespace: &str, timeout: Duration) -> RedisClient {
|
||||
connect_guarded(deployment, namespace, timeout, None)
|
||||
}
|
||||
|
||||
pub fn connect_guarded(
|
||||
deployment: &RedisDeployment,
|
||||
namespace: &str,
|
||||
timeout: Duration,
|
||||
breaker: Option<Breaker>,
|
||||
) -> RedisClient {
|
||||
RedisClient::new(RedisConfig {
|
||||
target: endpoint(deployment),
|
||||
credentials: Credentials::NONE,
|
||||
namespace: Namespace::new(namespace),
|
||||
timeouts: Timeouts::new(timeout, None),
|
||||
breaker,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> RedisClient {
|
||||
RedisClient::new(RedisConfig {
|
||||
target: endpoint(deployment),
|
||||
credentials,
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::new(Duration::from_secs(2), None),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub async fn connections(deployment: &RedisDeployment) -> Vec<ClientConnection> {
|
||||
deployment.clients().await.unwrap()
|
||||
}
|
||||
|
||||
/// A `TcpProxy` in front of every node, and a target that reaches the nodes only through them. A
|
||||
/// cluster announces the proxy ports, so its slot map and redirects lead there too.
|
||||
pub struct Proxied {
|
||||
proxies: Vec<TcpProxy>,
|
||||
pub target: Target,
|
||||
}
|
||||
|
||||
impl Proxied {
|
||||
pub fn blackhole_open_connections(&self) {
|
||||
self.proxies
|
||||
.iter()
|
||||
.for_each(TcpProxy::blackhole_open_connections);
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn proxied(deployment: &RedisDeployment) -> Proxied {
|
||||
match deployment {
|
||||
RedisDeployment::Standalone(server) => {
|
||||
let proxy = TcpProxy::start(server.port()).await.unwrap();
|
||||
let target = Target::Standalone(Standalone {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: proxy.port(),
|
||||
tls: None,
|
||||
db: 0,
|
||||
});
|
||||
Proxied {
|
||||
proxies: vec![proxy],
|
||||
target,
|
||||
}
|
||||
}
|
||||
RedisDeployment::Cluster(cluster) => {
|
||||
let mut proxies = Vec::new();
|
||||
for (_, port) in cluster.nodes() {
|
||||
proxies.push(TcpProxy::start(port).await.unwrap());
|
||||
}
|
||||
let ports: [u16; 3] = std::array::from_fn(|node| proxies[node].port());
|
||||
cluster.announce_ports(ports).await.unwrap();
|
||||
let target = Target::Cluster(Cluster {
|
||||
nodes: ports
|
||||
.iter()
|
||||
.map(|&port| Node {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port,
|
||||
})
|
||||
.collect(),
|
||||
tls: None,
|
||||
});
|
||||
Proxied { proxies, target }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `INFO <section>` from every node, concatenated.
|
||||
async fn info(deployment: &RedisDeployment, section: &str) -> String {
|
||||
match deployment {
|
||||
RedisDeployment::Standalone(server) => server.cli(&["INFO", section]).await.unwrap(),
|
||||
RedisDeployment::Cluster(cluster) => {
|
||||
let mut stats = String::new();
|
||||
for (_, port) in cluster.nodes() {
|
||||
stats.push_str(&cluster.cli_on_node(port, &["INFO", section]).await.unwrap());
|
||||
}
|
||||
stats
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `INFO commandstats` from every node, so a test can count what was actually sent.
|
||||
pub async fn command_stats(deployment: &RedisDeployment) -> String {
|
||||
info(deployment, "commandstats").await
|
||||
}
|
||||
|
||||
/// Keys looked up on every node since the last `CONFIG RESETSTAT`, found or not.
|
||||
pub async fn keys_read(deployment: &RedisDeployment) -> u64 {
|
||||
let stats = info(deployment, "stats").await;
|
||||
["keyspace_hits:", "keyspace_misses:"]
|
||||
.iter()
|
||||
.flat_map(|field| {
|
||||
stats
|
||||
.lines()
|
||||
.filter_map(move |line| line.trim().strip_prefix(field)?.parse::<u64>().ok())
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// How many times `command` reached a server, summed across the nodes in `stats`.
|
||||
pub fn calls(stats: &str, command: &str) -> u64 {
|
||||
let prefix = format!("cmdstat_{command}:calls=");
|
||||
stats
|
||||
.lines()
|
||||
.filter_map(|line| line.trim().strip_prefix(&prefix))
|
||||
.filter_map(|rest| rest.split(',').next()?.parse::<u64>().ok())
|
||||
.sum()
|
||||
}
|
||||
|
||||
pub struct Rotating(watch::Sender<BasicAuth>);
|
||||
|
||||
impl Rotating {
|
||||
pub fn new(username: &str, password: &str) -> Self {
|
||||
Self(watch::Sender::new(BasicAuth::new(
|
||||
username.to_owned(),
|
||||
password.to_owned(),
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn rotate(&self, username: &str, password: &str) {
|
||||
self.0
|
||||
.send_replace(BasicAuth::new(username.to_owned(), password.to_owned()));
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamingCredentialsProvider for Rotating {
|
||||
fn subscribe(
|
||||
&self,
|
||||
) -> Pin<Box<dyn Stream<Item = Result<BasicAuth, RedisError>> + Send + 'static>> {
|
||||
let receiver = self.0.subscribe();
|
||||
Box::pin(futures_util::stream::unfold(
|
||||
(receiver, true),
|
||||
|(mut receiver, first)| async move {
|
||||
if !first {
|
||||
receiver.changed().await.ok()?;
|
||||
}
|
||||
let credentials = receiver.borrow_and_update().clone();
|
||||
Some((Ok(credentials), (receiver, false)))
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
10
litellm-rust/crates/redis-connect/AGENTS.md
Normal file
10
litellm-rust/crates/redis-connect/AGENTS.md
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
- Turns a `litellm_redis_types::Resolution` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `redis-client` stays cloud-agnostic, so credential providers and cloud token sources are built here
|
||||
- The host passes only what it alone can build in `HostServices`: the event sink, the HTTP client and the environment lookup. The breaker config comes from `Resolution::breaker`, which `config` reads from `REDIS_CIRCUIT_BREAKER_*` with Python's fallback, so a bad value keeps the default instead of failing connect. Each cloud source asks its auth crate for one explicit identity, the way the Python SDK builds it: `AwsAuthService::default_chain_credentials` (botocore's default chain), `CredentialSource::application_default` (`google.auth.default`) and `AzureIdentity` (the class `_build_azure_credential` picks). LiteLLM's provider settings such as `AWS_ROLE_NAME`, `AZURE_AD_TOKEN` or `VERTEXAI_CREDENTIALS` belong to Bedrock, Azure OpenAI and Vertex and never reach those constructors. `Auth::Provider` (a handle for a Python `credential_provider` object) fails at connect with an error naming the key, since no host can supply one yet
|
||||
- Every expiring credential goes through the one refresh loop in `refresh.rs`. A cloud source only fetches. Never add a second loop or let a source sleep. The expiry deadline lives only in the published `Latest::Ready`, so expiring it leaves `Ready` and cannot fire again. Never keep a copy of the deadline in the loop
|
||||
- AWS SDK identity settings must be in the process environment. `aws-config` cannot take the host lookup, so config-only overrides for profiles, role/web-identity settings, credentials, files and metadata sources fail before connecting rather than silently using another identity. Equal process/config values are allowed. Bedrock-only settings remain unrelated to Redis
|
||||
- redis-rs subscribes once for the connection handshake and again for re-authentication, so a new connection sends a redundant initial AUTH. Unchanged refreshes do not trigger another AUTH. Do not skip a provider's first or alternating subscriptions, because concurrent connections and retries each need the current token. Removing the initial duplicate requires an upstream connection-level fix
|
||||
- The loop starts on the first `subscribe()`, which redis-rs calls inside its runtime, so `connect` is safe off a runtime
|
||||
- A token-only credential (Azure without `REDIS_USERNAME`, GCP) is sent as `AUTH default <token>`, because redis-rs always sends a username with a credentials provider. On Redis 6+ that is the same as Python's `AUTH <token>`
|
||||
- Tests for the refresh loop use a fake source and real Redis ACL users. Cloud sources are proven against the live service, unit tests only pin request shapes we own
|
||||
- Rotation tests must retain the same server-side connection IDs across re-AUTH, then force a reconnect after removing the old user and verify the new identity still works
|
||||
- Live cloud tests live in `tests/live/`. Set `LITELLM_REDIS_AWS_LIVE_CONFIG`, `LITELLM_REDIS_AZURE_LIVE_CONFIG` or `LITELLM_REDIS_GCP_LIVE_CONFIG` to a JSON file containing `redis` settings and `hold_seconds`. Set the waiting period from the deployment's token lifetime when checking expiry. Optional fixtures return early when unset, and configured tests run under plain `cargo test`. An unconfigured run never counts as live proof
|
||||
39
litellm-rust/crates/redis-connect/Cargo.toml
Normal file
39
litellm-rust/crates/redis-connect/Cargo.toml
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
[package]
|
||||
name = "litellm-redis-connect"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-auth-aws.workspace = true
|
||||
litellm-auth-azure.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-error.workspace = true
|
||||
litellm-secret-value.workspace = true
|
||||
strum.workspace = true
|
||||
litellm-circuit-breaker.workspace = true
|
||||
litellm-redis-types.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
litellm-http.workspace = true
|
||||
rustls.workspace = true
|
||||
litellm-redis-client.workspace = true
|
||||
futures-util.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
veil.workspace = true
|
||||
time.workspace = true
|
||||
url.workspace = true
|
||||
tokio = { workspace = true, features = ["sync", "time", "macros"] }
|
||||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-redis-resolve.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-redis-testing.workspace = true
|
||||
rustls.workspace = true
|
||||
rstest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] }
|
||||
154
litellm-rust/crates/redis-connect/src/connect.rs
Normal file
154
litellm-rust/crates/redis-connect/src/connect.rs
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_redis_client::{
|
||||
Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig,
|
||||
RedisEvents, Target, Timeouts,
|
||||
};
|
||||
use litellm_redis_types::{Auth, HealthCheck, Resolution, StaticAuth};
|
||||
|
||||
use crate::{
|
||||
CredentialSource, Error, RefreshingCredentials,
|
||||
sources::{AzureAd, ElastiCacheIam, GcpIam},
|
||||
};
|
||||
|
||||
pub type SharedEnvironment = Arc<dyn Lookup + Send + Sync>;
|
||||
|
||||
pub struct HostServices {
|
||||
pub events: Option<Arc<dyn RedisEvents>>,
|
||||
pub http: litellm_http::Client,
|
||||
pub environment: SharedEnvironment,
|
||||
pub credential_provider: Option<Arc<dyn CredentialSource>>,
|
||||
}
|
||||
|
||||
// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only
|
||||
const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25);
|
||||
|
||||
pub fn connect(
|
||||
resolved: Resolution,
|
||||
host: HostServices,
|
||||
breaker: Option<BreakerConfig>,
|
||||
) -> Result<RedisClient, Error> {
|
||||
if rustls::crypto::CryptoProvider::get_default().is_none() {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
}
|
||||
tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured");
|
||||
let credentials = credentials(resolved.auth, &host)?;
|
||||
let events = host.events;
|
||||
let breaker = breaker.map(|config| {
|
||||
let observer = events.clone();
|
||||
Breaker::new(
|
||||
config,
|
||||
Box::new(move |transition| {
|
||||
if let Some(events) = &observer {
|
||||
events.breaker(transition);
|
||||
}
|
||||
}),
|
||||
)
|
||||
});
|
||||
let health_check = health_check(&resolved.target, resolved.health_check);
|
||||
Ok(RedisClient::new(RedisConfig {
|
||||
target: resolved.target,
|
||||
credentials,
|
||||
namespace: resolved.namespace.map(Namespace::new).unwrap_or_default(),
|
||||
timeouts: Timeouts::new(
|
||||
resolved.timeout.unwrap_or(DEFAULT_TIMEOUT),
|
||||
resolved.connect_timeout,
|
||||
),
|
||||
breaker,
|
||||
events,
|
||||
health_check,
|
||||
keepalive: resolved.keepalive,
|
||||
})?)
|
||||
}
|
||||
|
||||
fn health_check(target: &Target, configured: HealthCheck) -> Option<Duration> {
|
||||
match (target, configured) {
|
||||
(_, HealthCheck::Every(interval)) => Some(interval),
|
||||
(_, HealthCheck::Disabled) => None,
|
||||
(Target::Cluster(_), HealthCheck::Unset) => Some(CLUSTER_HEALTH_CHECK_INTERVAL),
|
||||
(_, HealthCheck::Unset) => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn credentials(auth: Auth, host: &HostServices) -> Result<Credentials, Error> {
|
||||
let refreshing = |provider: RefreshingCredentials| {
|
||||
Credentials::Provider(CredentialsProvider(Arc::new(provider)))
|
||||
};
|
||||
Ok(match auth {
|
||||
Auth::Static(StaticAuth { username, password }) => {
|
||||
Credentials::Static { username, password }
|
||||
}
|
||||
Auth::Provider(_) => match &host.credential_provider {
|
||||
Some(source) => refreshing(RefreshingCredentials::new(Arc::clone(source))),
|
||||
None => return Err(Error::NoHostCredentialProvider),
|
||||
},
|
||||
Auth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new(
|
||||
auth,
|
||||
host.environment.clone(),
|
||||
))),
|
||||
Auth::AwsIam(auth) => refreshing(RefreshingCredentials::new(ElastiCacheIam::new(
|
||||
auth,
|
||||
&*host.environment,
|
||||
)?)),
|
||||
Auth::GcpIam { service_account } => refreshing(RefreshingCredentials::new(GcpIam::new(
|
||||
host.http.clone(),
|
||||
service_account,
|
||||
host.environment.clone(),
|
||||
))),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_redis_client::{Cluster, Node, Standalone};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn cluster() -> Target {
|
||||
Target::Cluster(Cluster {
|
||||
nodes: vec![Node {
|
||||
host: "node".to_owned(),
|
||||
port: 7000,
|
||||
}],
|
||||
tls: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn standalone() -> Target {
|
||||
Target::Standalone(Standalone {
|
||||
host: "redis".to_owned(),
|
||||
port: 6379,
|
||||
tls: None,
|
||||
db: 0,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::an_unset_cluster_takes_the_python_default(
|
||||
cluster(),
|
||||
HealthCheck::Unset,
|
||||
Some(CLUSTER_HEALTH_CHECK_INTERVAL)
|
||||
)]
|
||||
#[case::zero_keeps_a_cluster_off(cluster(), HealthCheck::Disabled, None)]
|
||||
#[case::an_unset_standalone_stays_off(standalone(), HealthCheck::Unset, None)]
|
||||
#[case::a_configured_interval_wins(
|
||||
cluster(),
|
||||
HealthCheck::Every(Duration::from_secs(5)),
|
||||
Some(Duration::from_secs(5))
|
||||
)]
|
||||
#[case::a_negative_interval_checks_every_command(
|
||||
standalone(),
|
||||
HealthCheck::Every(Duration::ZERO),
|
||||
Some(Duration::ZERO)
|
||||
)]
|
||||
fn health_checks_follow_redis_py(
|
||||
#[case] target: Target,
|
||||
#[case] configured: HealthCheck,
|
||||
#[case] expected: Option<Duration>,
|
||||
) {
|
||||
assert_eq!(health_check(&target, configured), expected);
|
||||
}
|
||||
}
|
||||
102
litellm-rust/crates/redis-connect/src/error.rs
Normal file
102
litellm-rust/crates/redis-connect/src/error.rs
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
#[derive(veil::Redact, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error(transparent)]
|
||||
Connect(#[from] litellm_redis_client::ConnectError),
|
||||
#[error(
|
||||
"Redis setting credential_provider names a Python object, which this host cannot supply"
|
||||
)]
|
||||
NoHostCredentialProvider,
|
||||
#[error(
|
||||
"Redis credential setting {setting} must be set in the process environment, not only environment_variables"
|
||||
)]
|
||||
ProcessEnvironmentRequired { setting: &'static str },
|
||||
#[error("{stage} failed{}", http_status_suffix(*http_status))]
|
||||
Credential {
|
||||
stage: CredentialStage,
|
||||
http_status: Option<u16>,
|
||||
#[source]
|
||||
#[redact(with = "[REDACTED]")]
|
||||
cause: litellm_error::Source,
|
||||
},
|
||||
#[error("credential acquisition exceeded {timeout:?}")]
|
||||
CredentialTimeout {
|
||||
timeout: std::time::Duration,
|
||||
#[source]
|
||||
cause: tokio::time::error::Elapsed,
|
||||
},
|
||||
#[error("Redis credential has expired")]
|
||||
ExpiredCredential,
|
||||
#[error("no {provider} credential is available for Redis")]
|
||||
MissingCredential { provider: &'static str },
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, strum::Display, strum::VariantArray)]
|
||||
pub enum CredentialStage {
|
||||
#[strum(serialize = "caller identity")]
|
||||
Identity,
|
||||
#[strum(serialize = "Redis token signing")]
|
||||
Signing,
|
||||
#[strum(serialize = "Redis token exchange")]
|
||||
Exchange,
|
||||
#[strum(serialize = "Redis token response")]
|
||||
Response,
|
||||
#[strum(serialize = "Redis credential_provider")]
|
||||
Provider,
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub fn credential<E: std::error::Error + Send + Sync + 'static>(
|
||||
stage: CredentialStage,
|
||||
) -> impl FnOnce(E) -> Self {
|
||||
move |cause| Self::Credential {
|
||||
stage,
|
||||
http_status: None,
|
||||
cause: litellm_error::Source::new(cause),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn credential_response(
|
||||
stage: CredentialStage,
|
||||
status: u16,
|
||||
) -> impl FnOnce(litellm_http::Error) -> Self {
|
||||
move |cause| Self::Credential {
|
||||
stage,
|
||||
http_status: Some(status),
|
||||
cause: litellm_error::Source::new(cause),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn upstream_status(&self) -> Option<u16> {
|
||||
match self {
|
||||
Self::Credential { http_status, .. } => *http_status,
|
||||
Self::Connect(_)
|
||||
| Self::NoHostCredentialProvider
|
||||
| Self::ProcessEnvironmentRequired { .. }
|
||||
| Self::CredentialTimeout { .. }
|
||||
| Self::ExpiredCredential
|
||||
| Self::MissingCredential { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn kind(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Connect(_) => "connection",
|
||||
Self::NoHostCredentialProvider => "host-credential-provider",
|
||||
Self::ProcessEnvironmentRequired { .. } => "process-environment-required",
|
||||
Self::Credential { stage, .. } => match stage {
|
||||
CredentialStage::Identity => "caller-identity",
|
||||
CredentialStage::Signing => "token-signing",
|
||||
CredentialStage::Exchange => "token-exchange",
|
||||
CredentialStage::Response => "token-response",
|
||||
CredentialStage::Provider => "credential-provider",
|
||||
},
|
||||
Self::CredentialTimeout { .. } => "acquisition-timeout",
|
||||
Self::ExpiredCredential => "expired-credential",
|
||||
Self::MissingCredential { .. } => "missing-credential",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn http_status_suffix(status: Option<u16>) -> String {
|
||||
status.map_or_else(String::new, |status| format!(" (HTTP {status})"))
|
||||
}
|
||||
129
litellm-rust/crates/redis-connect/src/events.rs
Normal file
129
litellm-rust/crates/redis-connect/src/events.rs
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
use std::{
|
||||
sync::Mutex,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use litellm_circuit_breaker::{State, Transition};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_redis_client::{CallEvent, FailureKind, Outcome, RedisEvents};
|
||||
|
||||
const DEFAULT_TIMEOUT_LOG_INTERVAL: Duration = Duration::from_secs(5);
|
||||
|
||||
pub struct TracingEvents {
|
||||
throttle: TimeoutLogThrottle,
|
||||
}
|
||||
|
||||
impl TracingEvents {
|
||||
pub fn new(lookup: &dyn Lookup) -> Self {
|
||||
let interval = lookup
|
||||
.get("REDIS_TIMEOUT_LOG_INTERVAL")
|
||||
.and_then(|seconds| seconds.parse::<f64>().ok())
|
||||
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
|
||||
.unwrap_or(DEFAULT_TIMEOUT_LOG_INTERVAL);
|
||||
Self {
|
||||
throttle: TimeoutLogThrottle::new(interval),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RedisEvents for TracingEvents {
|
||||
fn call(&self, event: &CallEvent) {
|
||||
let Outcome::Failed(kind) = event.outcome else {
|
||||
return;
|
||||
};
|
||||
match kind {
|
||||
FailureKind::BreakerOpen => {
|
||||
tracing::debug!(operation = ?event.operation, "redis call rejected by the open circuit breaker")
|
||||
}
|
||||
FailureKind::Timeout => match self.throttle.admit(Instant::now()) {
|
||||
None => {
|
||||
tracing::debug!(operation = ?event.operation, duration = ?event.duration, "redis call timed out")
|
||||
}
|
||||
Some(0) => {
|
||||
tracing::error!(operation = ?event.operation, duration = ?event.duration, "redis call timed out")
|
||||
}
|
||||
Some(suppressed) => tracing::error!(
|
||||
operation = ?event.operation,
|
||||
duration = ?event.duration,
|
||||
suppressed,
|
||||
"redis call timed out, more redis timeouts since the previous line were logged at debug"
|
||||
),
|
||||
},
|
||||
FailureKind::Transport | FailureKind::Server | FailureKind::Decode => {
|
||||
tracing::error!(operation = ?event.operation, failure = ?kind, "redis call failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn breaker(&self, transition: &Transition<State>) {
|
||||
match transition.to {
|
||||
State::Open { .. } => tracing::warn!(
|
||||
generation = transition.generation,
|
||||
"redis circuit breaker opened, fast-failing redis calls"
|
||||
),
|
||||
State::HalfOpen => tracing::info!(
|
||||
generation = transition.generation,
|
||||
"redis circuit breaker half open, probing redis"
|
||||
),
|
||||
State::Closed { .. } => tracing::info!(
|
||||
generation = transition.generation,
|
||||
"redis circuit breaker closed, redis recovered"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct TimeoutLogThrottle {
|
||||
interval: Duration,
|
||||
state: Mutex<(Option<Instant>, u64)>,
|
||||
}
|
||||
|
||||
impl TimeoutLogThrottle {
|
||||
fn new(interval: Duration) -> Self {
|
||||
Self {
|
||||
interval,
|
||||
state: Mutex::new((None, 0)),
|
||||
}
|
||||
}
|
||||
|
||||
fn admit(&self, now: Instant) -> Option<u64> {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
let (last_logged_at, suppressed) = &mut *state;
|
||||
if last_logged_at.is_some_and(|last| now.duration_since(last) < self.interval) {
|
||||
*suppressed += 1;
|
||||
return None;
|
||||
}
|
||||
*last_logged_at = Some(now);
|
||||
Some(std::mem::take(suppressed))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::TimeoutLogThrottle;
|
||||
|
||||
#[rstest]
|
||||
fn admits_one_line_per_interval_and_counts_the_rest() {
|
||||
let throttle = TimeoutLogThrottle::new(Duration::from_secs(5));
|
||||
let start = Instant::now();
|
||||
let at = |seconds: u64| start + Duration::from_secs(seconds);
|
||||
|
||||
let admitted = [
|
||||
throttle.admit(at(0)),
|
||||
throttle.admit(at(1)),
|
||||
throttle.admit(at(4)),
|
||||
throttle.admit(at(5)),
|
||||
throttle.admit(at(6)),
|
||||
throttle.admit(at(20)),
|
||||
];
|
||||
|
||||
assert_eq!(admitted, [Some(0), None, None, Some(2), None, Some(1)]);
|
||||
}
|
||||
}
|
||||
11
litellm-rust/crates/redis-connect/src/lib.rs
Normal file
11
litellm-rust/crates/redis-connect/src/lib.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
mod connect;
|
||||
mod error;
|
||||
mod events;
|
||||
mod refresh;
|
||||
mod sources;
|
||||
|
||||
pub use connect::{HostServices, SharedEnvironment, connect};
|
||||
pub use error::{CredentialStage, Error};
|
||||
pub use events::TracingEvents;
|
||||
pub(crate) use refresh::RefreshingCredentials;
|
||||
pub use refresh::{CredentialSource, Expiring, SourceFuture};
|
||||
783
litellm-rust/crates/redis-connect/src/refresh.rs
Normal file
783
litellm-rust/crates/redis-connect/src/refresh.rs
Normal file
|
|
@ -0,0 +1,783 @@
|
|||
use std::{
|
||||
future::{Future, pending},
|
||||
pin::Pin,
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use futures_util::{Stream, stream};
|
||||
use litellm_redis_client::{BasicAuth, RedisError, StreamingCredentialsProvider};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use tokio::{sync::watch, time::Instant};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
const TOKEN_ONLY_USERNAME: &str = "default";
|
||||
const UNKNOWN_LIFETIME: Duration = Duration::from_secs(300);
|
||||
const MIN_REFRESH: Duration = Duration::from_secs(1);
|
||||
const ACQUISITION_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const FIRST_RETRY: Duration = Duration::from_secs(1);
|
||||
const MAX_RETRY: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct Expiring {
|
||||
pub username: Option<String>,
|
||||
pub password: SecretValue,
|
||||
pub expires_at: Option<SystemTime>,
|
||||
}
|
||||
|
||||
pub type SourceFuture<'a> = Pin<Box<dyn Future<Output = Result<Expiring, Error>> + Send + 'a>>;
|
||||
|
||||
pub trait CredentialSource: Send + Sync + 'static {
|
||||
fn fetch(&self) -> SourceFuture<'_>;
|
||||
}
|
||||
|
||||
impl CredentialSource for Arc<dyn CredentialSource> {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
(**self).fetch()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Latest {
|
||||
Pending,
|
||||
Ready {
|
||||
username: String,
|
||||
password: SecretValue,
|
||||
expires_at: Option<Instant>,
|
||||
},
|
||||
Unavailable(Arc<Error>),
|
||||
}
|
||||
|
||||
impl Latest {
|
||||
fn is_same_token(&self, other: &Self) -> bool {
|
||||
matches!(
|
||||
(self, other),
|
||||
(
|
||||
Self::Ready { username, password, .. },
|
||||
Self::Ready { username: other_username, password: other_password, .. },
|
||||
) if username == other_username && password == other_password
|
||||
)
|
||||
}
|
||||
|
||||
fn expires_at(&self) -> Option<Instant> {
|
||||
match self {
|
||||
Self::Ready { expires_at, .. } => *expires_at,
|
||||
Self::Pending | Self::Unavailable(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type Start = Box<dyn FnOnce() + Send>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RefreshingCredentials {
|
||||
latest: watch::Receiver<Latest>,
|
||||
start: Arc<Mutex<Option<Start>>>,
|
||||
}
|
||||
|
||||
impl RefreshingCredentials {
|
||||
pub fn new(source: impl CredentialSource) -> Self {
|
||||
let (sender, latest) = watch::channel(Latest::Pending);
|
||||
let start: Start = Box::new(move || drop(tokio::spawn(refresh(source, sender))));
|
||||
Self {
|
||||
latest,
|
||||
start: Arc::new(Mutex::new(Some(start))),
|
||||
}
|
||||
}
|
||||
|
||||
fn start_once(&self) {
|
||||
let start = self
|
||||
.start
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
.take();
|
||||
if let Some(start) = start {
|
||||
start();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamingCredentialsProvider for RefreshingCredentials {
|
||||
fn subscribe(
|
||||
&self,
|
||||
) -> Pin<Box<dyn Stream<Item = Result<BasicAuth, RedisError>> + Send + 'static>> {
|
||||
self.start_once();
|
||||
Box::pin(stream::unfold(
|
||||
(self.latest.clone(), true),
|
||||
|(mut receiver, first)| async move {
|
||||
if !first {
|
||||
receiver.changed().await.ok()?;
|
||||
}
|
||||
let latest = receiver
|
||||
.wait_for(|latest| !matches!(latest, Latest::Pending))
|
||||
.await
|
||||
.ok()?
|
||||
.clone();
|
||||
let item = match latest {
|
||||
Latest::Ready {
|
||||
username, password, ..
|
||||
} => Ok(BasicAuth::new(username, password.expose().to_owned())),
|
||||
Latest::Unavailable(error) => Err(public_error(&error)),
|
||||
Latest::Pending => unreachable!("wait_for skips Pending"),
|
||||
};
|
||||
Some((item, (receiver, false)))
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn public_error(error: &Error) -> RedisError {
|
||||
RedisError::from((
|
||||
litellm_redis_client::RedisErrorKind::AuthenticationFailed,
|
||||
"fetching Redis credentials failed",
|
||||
error.to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
struct Fresh {
|
||||
latest: Latest,
|
||||
refresh_in: Duration,
|
||||
}
|
||||
|
||||
fn fresh(credentials: Expiring, now: SystemTime) -> Result<Fresh, Error> {
|
||||
let remaining = credentials
|
||||
.expires_at
|
||||
.map(|expires_at| expires_at.duration_since(now).unwrap_or_default());
|
||||
if remaining == Some(Duration::ZERO) {
|
||||
return Err(Error::ExpiredCredential);
|
||||
}
|
||||
Ok(Fresh {
|
||||
latest: Latest::Ready {
|
||||
username: credentials
|
||||
.username
|
||||
.unwrap_or_else(|| TOKEN_ONLY_USERNAME.to_owned()),
|
||||
password: credentials.password,
|
||||
expires_at: remaining.map(|remaining| Instant::now() + remaining),
|
||||
},
|
||||
refresh_in: refresh_after(credentials.expires_at, now),
|
||||
})
|
||||
}
|
||||
|
||||
async fn until_expiry<T>(
|
||||
future: impl Future<Output = T>,
|
||||
sender: &watch::Sender<Latest>,
|
||||
) -> Option<T> {
|
||||
tokio::pin!(future);
|
||||
loop {
|
||||
let expires_at = sender.borrow().expires_at();
|
||||
let expiration = async {
|
||||
match expires_at {
|
||||
Some(deadline) => tokio::time::sleep_until(deadline).await,
|
||||
None => pending().await,
|
||||
}
|
||||
};
|
||||
tokio::select! {
|
||||
biased;
|
||||
() = sender.closed() => return None,
|
||||
() = expiration => {
|
||||
sender.send_replace(Latest::Unavailable(Arc::new(Error::ExpiredCredential)));
|
||||
}
|
||||
result = &mut future => return Some(result),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn refresh(source: impl CredentialSource, sender: watch::Sender<Latest>) {
|
||||
let mut retry = FIRST_RETRY;
|
||||
loop {
|
||||
let acquisition = async {
|
||||
tokio::time::timeout(ACQUISITION_TIMEOUT, source.fetch())
|
||||
.await
|
||||
.map_err(|cause| Error::CredentialTimeout {
|
||||
timeout: ACQUISITION_TIMEOUT,
|
||||
cause,
|
||||
})?
|
||||
.and_then(|credentials| fresh(credentials, SystemTime::now()))
|
||||
};
|
||||
let Some(result) = until_expiry(acquisition, &sender).await else {
|
||||
return;
|
||||
};
|
||||
let wait = match result {
|
||||
Ok(Fresh { latest, refresh_in }) => {
|
||||
retry = FIRST_RETRY;
|
||||
sender.send_if_modified(|current| {
|
||||
let changed = !current.is_same_token(&latest);
|
||||
*current = latest;
|
||||
changed
|
||||
});
|
||||
refresh_in
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(failure = error.kind(), upstream_status = error.upstream_status(), retry_in = ?retry, "fetching Redis credentials failed");
|
||||
if !matches!(*sender.borrow(), Latest::Ready { .. }) {
|
||||
sender.send_replace(Latest::Unavailable(Arc::new(error)));
|
||||
}
|
||||
let wait = retry;
|
||||
retry = (retry * 2).min(MAX_RETRY);
|
||||
wait
|
||||
}
|
||||
};
|
||||
if until_expiry(tokio::time::sleep(wait), &sender)
|
||||
.await
|
||||
.is_none()
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_after(expires_at: Option<SystemTime>, now: SystemTime) -> Duration {
|
||||
let remaining = expires_at
|
||||
.map(|expires_at| expires_at.duration_since(now).unwrap_or_default())
|
||||
.unwrap_or(UNKNOWN_LIFETIME);
|
||||
(remaining * 3 / 4).max(MIN_REFRESH)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_redis_client::commands;
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_redis_client::StreamingCredentialsProvider;
|
||||
use litellm_secret_value::SecretValue;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after};
|
||||
use litellm_redis_client::{
|
||||
Cluster, Credentials, CredentialsProvider, Namespace, Node, RedisClient, RedisConfig,
|
||||
RedisUrl, Sentinel, Target, Timeouts,
|
||||
};
|
||||
use litellm_redis_testing::{
|
||||
ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually,
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[case::three_quarters_of_the_lifetime(
|
||||
Some(Duration::from_secs(900)),
|
||||
Duration::from_secs(675)
|
||||
)]
|
||||
#[case::already_expired(Some(Duration::ZERO), Duration::from_secs(1))]
|
||||
#[case::unknown_expiry(None, Duration::from_secs(225))]
|
||||
fn refreshes_well_before_expiry(
|
||||
#[case] lifetime: Option<Duration>,
|
||||
#[case] expected: Duration,
|
||||
) {
|
||||
let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000);
|
||||
|
||||
assert_eq!(
|
||||
refresh_after(lifetime.map(|lifetime| now + lifetime), now),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
struct Scripted {
|
||||
replies: Mutex<Vec<Result<Expiring, &'static str>>>,
|
||||
calls: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl CredentialSource for Scripted {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
let reply = self.replies.lock().unwrap().remove(0);
|
||||
Box::pin(async move {
|
||||
reply.map_err(|reason| {
|
||||
crate::Error::credential(crate::error::CredentialStage::Identity)(
|
||||
std::io::Error::other(reason),
|
||||
)
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn token(password: &str, lifetime: Duration) -> Result<Expiring, &'static str> {
|
||||
Ok(Expiring {
|
||||
username: None,
|
||||
password: SecretValue::new(password),
|
||||
expires_at: Some(SystemTime::now() + lifetime),
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn pushes_each_refresh_and_keeps_the_last_token_through_failures() {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![
|
||||
token("first", Duration::from_secs(8)),
|
||||
Err("token endpoint down"),
|
||||
token("second", Duration::from_secs(3600)),
|
||||
]),
|
||||
calls: calls.clone(),
|
||||
});
|
||||
let mut stream = provider.subscribe();
|
||||
|
||||
let first = stream.next().await.unwrap().unwrap();
|
||||
let second = stream.next().await.unwrap().unwrap();
|
||||
let late = provider.subscribe().next().await.unwrap().unwrap();
|
||||
|
||||
assert_eq!((first.username(), first.password()), ("default", "first"));
|
||||
assert_eq!(second.password(), "second");
|
||||
assert_eq!(late.password(), "second");
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn building_outside_a_runtime_fetches_nothing_until_a_connection_subscribes() {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![token("first", Duration::from_secs(3600))]),
|
||||
calls: calls.clone(),
|
||||
});
|
||||
let runtime = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let before = calls.load(Ordering::SeqCst);
|
||||
let first = runtime.block_on(async { provider.subscribe().next().await.unwrap().unwrap() });
|
||||
|
||||
assert_eq!(before, 0);
|
||||
assert_eq!(first.password(), "first");
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn a_failure_before_any_token_fails_the_connection_instead_of_hanging() {
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![
|
||||
Err("no identity"),
|
||||
token("late", Duration::from_secs(3600)),
|
||||
]),
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let mut stream = provider.subscribe();
|
||||
|
||||
assert!(stream.next().await.unwrap().is_err());
|
||||
assert_eq!(stream.next().await.unwrap().unwrap().password(), "late");
|
||||
}
|
||||
|
||||
struct Blocking {
|
||||
first: Mutex<Option<Expiring>>,
|
||||
started: Arc<tokio::sync::Notify>,
|
||||
cancelled: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
struct Cancellation(Arc<AtomicUsize>);
|
||||
|
||||
impl Drop for Cancellation {
|
||||
fn drop(&mut self) {
|
||||
self.0.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
impl CredentialSource for Blocking {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
let first = self.first.lock().unwrap().take();
|
||||
Box::pin(async move {
|
||||
if let Some(first) = first {
|
||||
return Ok(first);
|
||||
}
|
||||
let _guard = Cancellation(self.cancelled.clone());
|
||||
self.started.notify_one();
|
||||
std::future::pending().await
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn a_refresh_that_returns_the_same_token_is_not_pushed_again() {
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![
|
||||
token("same", Duration::from_secs(8)),
|
||||
token("same", Duration::from_secs(3600)),
|
||||
]),
|
||||
calls: calls.clone(),
|
||||
});
|
||||
let mut stream = provider.subscribe();
|
||||
assert_eq!(stream.next().await.unwrap().unwrap().password(), "same");
|
||||
|
||||
let pushed = tokio::time::timeout(Duration::from_secs(60), stream.next()).await;
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 2);
|
||||
assert!(pushed.is_err(), "an unchanged token was pushed again");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn expiration_is_published_while_refresh_is_stalled() {
|
||||
let started = Arc::new(tokio::sync::Notify::new());
|
||||
let provider = RefreshingCredentials::new(Blocking {
|
||||
first: Mutex::new(Some(token("first", Duration::from_secs(4)).unwrap())),
|
||||
started: started.clone(),
|
||||
cancelled: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let mut stream = provider.subscribe();
|
||||
assert_eq!(stream.next().await.unwrap().unwrap().password(), "first");
|
||||
let start = tokio::time::Instant::now();
|
||||
started.notified().await;
|
||||
let expired = stream.next().await.unwrap().unwrap_err();
|
||||
assert!(start.elapsed() <= Duration::from_secs(5));
|
||||
assert!(expired.to_string().contains("expired"));
|
||||
let fresh = provider.subscribe().next().await.unwrap().unwrap_err();
|
||||
assert!(fresh.to_string().contains("expired"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn expiration_is_published_once_and_later_failures_keep_their_cause() {
|
||||
let replies = std::iter::once(token("first", Duration::from_secs(8)))
|
||||
.chain(std::iter::repeat_n(Err("token endpoint down"), 16))
|
||||
.collect();
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(replies),
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let mut stream = provider.subscribe();
|
||||
assert_eq!(stream.next().await.unwrap().unwrap().password(), "first");
|
||||
|
||||
let failures: Vec<String> = stream
|
||||
.take(5)
|
||||
.map(|item| item.unwrap_err().to_string())
|
||||
.collect()
|
||||
.await;
|
||||
|
||||
assert!(failures[0].contains("expired"), "{failures:?}");
|
||||
assert!(
|
||||
failures[1..]
|
||||
.iter()
|
||||
.all(|failure| failure.contains("caller identity") && !failure.contains("expired")),
|
||||
"{failures:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn initial_acquisition_has_a_bounded_wait() {
|
||||
let provider = RefreshingCredentials::new(Blocking {
|
||||
first: Mutex::new(None),
|
||||
started: Arc::new(tokio::sync::Notify::new()),
|
||||
cancelled: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let error = provider.subscribe().next().await.unwrap().unwrap_err();
|
||||
assert!(error.to_string().contains("exceeded"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn dropping_all_subscribers_cancels_inflight_acquisition() {
|
||||
let started = Arc::new(tokio::sync::Notify::new());
|
||||
let cancelled = Arc::new(AtomicUsize::new(0));
|
||||
let provider = RefreshingCredentials::new(Blocking {
|
||||
first: Mutex::new(None),
|
||||
started: started.clone(),
|
||||
cancelled: cancelled.clone(),
|
||||
});
|
||||
let stream = provider.subscribe();
|
||||
started.notified().await;
|
||||
drop(stream);
|
||||
drop(provider);
|
||||
tokio::task::yield_now().await;
|
||||
assert_eq!(cancelled.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn an_expired_fetch_is_rejected_and_a_later_valid_token_recovers() {
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![
|
||||
Ok(Expiring {
|
||||
username: None,
|
||||
password: SecretValue::new("expired"),
|
||||
expires_at: Some(SystemTime::UNIX_EPOCH),
|
||||
}),
|
||||
token("recovered", Duration::from_secs(3600)),
|
||||
]),
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let mut stream = provider.subscribe();
|
||||
assert!(
|
||||
stream
|
||||
.next()
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("expired")
|
||||
);
|
||||
assert_eq!(
|
||||
stream.next().await.unwrap().unwrap().password(),
|
||||
"recovered"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn credential_failure_keeps_its_typed_cause_but_hides_it_at_the_redis_boundary() {
|
||||
let error = crate::Error::credential(crate::error::CredentialStage::Identity)(
|
||||
std::io::Error::new(std::io::ErrorKind::PermissionDenied, "credential-value"),
|
||||
);
|
||||
assert_eq!(
|
||||
litellm_error::find::<std::io::Error>(&error)
|
||||
.unwrap()
|
||||
.kind(),
|
||||
std::io::ErrorKind::PermissionDenied
|
||||
);
|
||||
let debug = format!("{error:?}");
|
||||
assert!(!debug.contains("credential-value"));
|
||||
assert!(debug.contains("[REDACTED]"));
|
||||
let public = super::public_error(&error).to_string();
|
||||
assert!(public.contains("caller identity"));
|
||||
assert!(!public.contains("credential-value"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn credential_http_failures_expose_status_without_exposing_the_response_body() {
|
||||
let error = crate::Error::credential_response(crate::CredentialStage::Exchange, 503)(
|
||||
litellm_http::Error::upstream(503, "credential-value".into()),
|
||||
);
|
||||
let public = super::public_error(&error).to_string();
|
||||
assert!(public.contains("HTTP 503"));
|
||||
assert!(!public.contains("credential-value"));
|
||||
assert_eq!(
|
||||
litellm_error::find::<litellm_http::Error>(&error)
|
||||
.unwrap()
|
||||
.upstream_response()
|
||||
.unwrap()
|
||||
.0,
|
||||
503
|
||||
);
|
||||
}
|
||||
|
||||
fn endpoint(deployment: &RedisDeployment) -> Target {
|
||||
match deployment {
|
||||
RedisDeployment::Standalone(server) => Target::Url(RedisUrl {
|
||||
url: SecretValue::new(server.url()),
|
||||
tls: None,
|
||||
db: None,
|
||||
username: None,
|
||||
password: None,
|
||||
}),
|
||||
RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster {
|
||||
nodes: cluster
|
||||
.nodes()
|
||||
.map(|(host, port)| Node {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
})
|
||||
.collect(),
|
||||
tls: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn user(username: &str, password: &str, lifetime: Duration) -> Result<Expiring, &'static str> {
|
||||
Ok(Expiring {
|
||||
username: Some(username.to_owned()),
|
||||
password: SecretValue::new(password),
|
||||
expires_at: Some(SystemTime::now() + lifetime),
|
||||
})
|
||||
}
|
||||
|
||||
fn connection_users(connections: &[ClientConnection]) -> BTreeSet<&str> {
|
||||
connections
|
||||
.iter()
|
||||
.map(|connection| connection.user.as_str())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn connection_ids(connections: &[ClientConnection]) -> BTreeSet<(u16, u64)> {
|
||||
connections
|
||||
.iter()
|
||||
.map(|connection| (connection.node, connection.id))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_refreshed_token_reauthenticates_live_connections_before_the_old_one_expires(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&[
|
||||
"ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&[
|
||||
"ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![
|
||||
user("alice", "first", Duration::from_secs(2)),
|
||||
user("bob", "second", Duration::from_secs(3600)),
|
||||
]),
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: endpoint(&server),
|
||||
credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))),
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::default(),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
client
|
||||
.send(commands::set("k", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
let before = server.clients().await.unwrap();
|
||||
|
||||
let after = eventually(async || {
|
||||
let connections = server.clients().await.ok()?;
|
||||
(!connections.is_empty() && connection_users(&connections) == BTreeSet::from(["bob"]))
|
||||
.then_some(connections)
|
||||
})
|
||||
.await;
|
||||
server
|
||||
.cli_on_every_node(&["ACL", "DELUSER", "alice"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!before.is_empty());
|
||||
assert_eq!(connection_users(&before), BTreeSet::from(["alice"]));
|
||||
assert_eq!(connection_ids(&after), connection_ids(&before));
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>("k"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("1")
|
||||
);
|
||||
server
|
||||
.cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"])
|
||||
.await
|
||||
.unwrap();
|
||||
let value = eventually(async || {
|
||||
client
|
||||
.send(commands::get::<String>("k"))
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
})
|
||||
.await;
|
||||
let reconnected = server.clients().await.unwrap();
|
||||
|
||||
assert_eq!(value, "1");
|
||||
assert!(!reconnected.is_empty());
|
||||
assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"]));
|
||||
assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after)));
|
||||
}
|
||||
|
||||
/// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the
|
||||
/// client presents to the master Sentinel promotes, not the one it connected with first.
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_refreshed_user_reaches_the_master_sentinel_promotes() {
|
||||
let sentinel = RedisSentinel::start().await.unwrap();
|
||||
sentinel
|
||||
.cli_on_data_nodes(&[
|
||||
"ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
sentinel
|
||||
.cli_on_data_nodes(&[
|
||||
"ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all",
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
let provider = RefreshingCredentials::new(Scripted {
|
||||
replies: Mutex::new(vec![
|
||||
user("alice", "first", Duration::from_secs(2)),
|
||||
user("bob", "second", Duration::from_secs(3600)),
|
||||
]),
|
||||
calls: Arc::new(AtomicUsize::new(0)),
|
||||
});
|
||||
let (host, port) = sentinel.sentinel();
|
||||
let client = RedisClient::new(RedisConfig {
|
||||
target: Target::Sentinel(Sentinel {
|
||||
monitors: vec![Node {
|
||||
host: host.to_owned(),
|
||||
port,
|
||||
}],
|
||||
service_name: SERVICE_NAME.to_owned(),
|
||||
monitor_auth: None,
|
||||
tls: None,
|
||||
db: 0,
|
||||
}),
|
||||
credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))),
|
||||
namespace: Namespace::default(),
|
||||
timeouts: Timeouts::new(Duration::from_secs(1), None),
|
||||
breaker: None,
|
||||
events: None,
|
||||
health_check: None,
|
||||
keepalive: true,
|
||||
})
|
||||
.unwrap();
|
||||
client
|
||||
.send(commands::set("before", "1", Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
let master = sentinel.master_port().await.unwrap();
|
||||
let app_users = |connections: &[ClientConnection]| -> BTreeSet<String> {
|
||||
connections
|
||||
.iter()
|
||||
.filter(|connection| connection.user != "default")
|
||||
.map(|connection| connection.user.clone())
|
||||
.collect()
|
||||
};
|
||||
eventually(async || {
|
||||
let connections = sentinel.clients_on(master).await.ok()?;
|
||||
(app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(())
|
||||
})
|
||||
.await;
|
||||
sentinel
|
||||
.cli_on_data_nodes(&["ACL", "DELUSER", "alice"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let promoted = sentinel.failover().await.unwrap();
|
||||
eventually(async || {
|
||||
client
|
||||
.send(commands::set("after", "2", Duration::from_secs(60)))
|
||||
.await
|
||||
.ok()
|
||||
})
|
||||
.await;
|
||||
let on_promoted = sentinel.clients_on(promoted).await.unwrap();
|
||||
|
||||
assert_eq!(promoted, sentinel.replica_port());
|
||||
assert_eq!(
|
||||
app_users(&on_promoted),
|
||||
BTreeSet::from(["bob".to_owned()]),
|
||||
"{on_promoted:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>("after"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("2")
|
||||
);
|
||||
}
|
||||
}
|
||||
188
litellm-rust/crates/redis-connect/src/sources/aws.rs
Normal file
188
litellm-rust/crates/redis-connect/src/sources/aws.rs
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use litellm_auth_aws::{AwsAuthService, SigningScope, presign_url};
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_redis_types::AwsIamAuth;
|
||||
|
||||
use crate::Error;
|
||||
use crate::error::CredentialStage;
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use crate::{CredentialSource, Expiring, SourceFuture};
|
||||
|
||||
const SERVICE: &str = "elasticache";
|
||||
const TOKEN_LIFETIME: Duration = Duration::from_secs(900);
|
||||
|
||||
const SDK_IDENTITY_ENVIRONMENT: &[&str] = &[
|
||||
"AWS_ACCESS_KEY_ID",
|
||||
"AWS_SECRET_ACCESS_KEY",
|
||||
"AWS_SESSION_TOKEN",
|
||||
"AWS_PROFILE",
|
||||
"AWS_DEFAULT_PROFILE",
|
||||
"AWS_SHARED_CREDENTIALS_FILE",
|
||||
"AWS_CONFIG_FILE",
|
||||
"AWS_WEB_IDENTITY_TOKEN_FILE",
|
||||
"AWS_ROLE_ARN",
|
||||
"AWS_ROLE_SESSION_NAME",
|
||||
"AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
|
||||
"AWS_CONTAINER_CREDENTIALS_FULL_URI",
|
||||
"AWS_CONTAINER_AUTHORIZATION_TOKEN",
|
||||
"AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE",
|
||||
"AWS_EC2_METADATA_DISABLED",
|
||||
"AWS_EC2_METADATA_SERVICE_ENDPOINT",
|
||||
"AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE",
|
||||
"AWS_ENDPOINT_URL",
|
||||
"AWS_ENDPOINT_URL_STS",
|
||||
];
|
||||
|
||||
pub(crate) struct ElastiCacheIam {
|
||||
service: AwsAuthService,
|
||||
auth: AwsIamAuth,
|
||||
}
|
||||
|
||||
impl ElastiCacheIam {
|
||||
pub(crate) fn new(auth: AwsIamAuth, environment: &dyn Lookup) -> Result<Self, Error> {
|
||||
validate_environment(environment, &|name: &str| std::env::var(name).ok())?;
|
||||
Ok(Self {
|
||||
service: AwsAuthService::default(),
|
||||
auth,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_environment(environment: &dyn Lookup, process: &dyn Lookup) -> Result<(), Error> {
|
||||
match SDK_IDENTITY_ENVIRONMENT.iter().find(|name| {
|
||||
environment.get(name).is_some_and(|value| {
|
||||
Some(SecretValue::new(value)) != process.get(name).map(SecretValue::new)
|
||||
})
|
||||
}) {
|
||||
Some(setting) => Err(Error::ProcessEnvironmentRequired { setting }),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn connect_url(auth: &AwsIamAuth) -> String {
|
||||
let query = [
|
||||
Some(("Action", "connect")),
|
||||
Some(("User", auth.user_name.as_str())),
|
||||
auth.serverless
|
||||
.then_some(("ResourceType", "ServerlessCache")),
|
||||
];
|
||||
let mut url = String::from("https://");
|
||||
url.push_str(&auth.cache_name.to_lowercase());
|
||||
url.push_str("/?");
|
||||
url.push_str(
|
||||
&url::form_urlencoded::Serializer::new(String::new())
|
||||
.extend_pairs(query.into_iter().flatten())
|
||||
.finish(),
|
||||
);
|
||||
url
|
||||
}
|
||||
|
||||
impl CredentialSource for ElastiCacheIam {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
Box::pin(async move {
|
||||
let credentials = self
|
||||
.service
|
||||
.default_chain_credentials(&self.auth.region)
|
||||
.await
|
||||
.map_err(Error::credential(CredentialStage::Identity))?;
|
||||
let now = SystemTime::now();
|
||||
let signed = presign_url(
|
||||
&connect_url(&self.auth),
|
||||
TOKEN_LIFETIME,
|
||||
SigningScope {
|
||||
region: &self.auth.region,
|
||||
service: SERVICE,
|
||||
credentials: &credentials,
|
||||
time: now,
|
||||
},
|
||||
)
|
||||
.map_err(Error::credential(CredentialStage::Signing))?;
|
||||
let token_expiry = now + TOKEN_LIFETIME;
|
||||
Ok(Expiring {
|
||||
username: Some(self.auth.user_name.clone()),
|
||||
password: SecretValue::new(signed.trim_start_matches("https://")),
|
||||
expires_at: Some(
|
||||
credentials
|
||||
.expiry()
|
||||
.map_or(token_expiry, |expiry| expiry.min(token_expiry)),
|
||||
),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_redis_types::AwsIamAuth;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{connect_url, validate_environment};
|
||||
|
||||
#[rstest]
|
||||
#[case::profile("AWS_PROFILE")]
|
||||
#[case::role("AWS_ROLE_ARN")]
|
||||
#[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")]
|
||||
#[case::source_key("AWS_ACCESS_KEY_ID")]
|
||||
#[case::source_secret("AWS_SECRET_ACCESS_KEY")]
|
||||
#[case::session("AWS_SESSION_TOKEN")]
|
||||
#[case::profile_path("AWS_SHARED_CREDENTIALS_FILE")]
|
||||
#[case::container("AWS_CONTAINER_CREDENTIALS_FULL_URI")]
|
||||
fn configured_identity_overrides_cannot_silently_use_an_ambient_identity(
|
||||
#[case] setting: &'static str,
|
||||
) {
|
||||
let configured = |name: &str| (name == setting).then(|| "credential-value".to_owned());
|
||||
let process = |name: &str| (name == setting).then(|| "ambient".to_owned());
|
||||
let error = validate_environment(&configured, &process).unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
crate::Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting
|
||||
));
|
||||
assert!(!error.to_string().contains("credential-value"));
|
||||
assert!(!format!("{error:?}").contains("credential-value"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::process_environment("AWS_PROFILE", Some("profile"), Some("profile"))]
|
||||
#[case::ambient_only("AWS_PROFILE", None, Some("profile"))]
|
||||
#[case::unconfigured("AWS_PROFILE", None, None)]
|
||||
#[case::unrelated_bedrock_setting("AWS_ROLE_NAME", Some("bedrock"), None)]
|
||||
fn unchanged_or_unrelated_settings_preserve_the_sdk_chain(
|
||||
#[case] setting: &str,
|
||||
#[case] configured: Option<&str>,
|
||||
#[case] ambient: Option<&str>,
|
||||
) {
|
||||
let configured = |name: &str| {
|
||||
(name == setting)
|
||||
.then_some(configured)
|
||||
.flatten()
|
||||
.map(str::to_owned)
|
||||
};
|
||||
let ambient = |name: &str| {
|
||||
(name == setting)
|
||||
.then_some(ambient)
|
||||
.flatten()
|
||||
.map(str::to_owned)
|
||||
};
|
||||
assert!(validate_environment(&configured, &ambient).is_ok());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::provisioned(false, "https://cache-1/?Action=connect&User=iam+user%2F1")]
|
||||
#[case::serverless(
|
||||
true,
|
||||
"https://cache-1/?Action=connect&User=iam+user%2F1&ResourceType=ServerlessCache"
|
||||
)]
|
||||
fn the_connect_url_matches_python_urlencode(#[case] serverless: bool, #[case] expected: &str) {
|
||||
let auth = AwsIamAuth {
|
||||
user_name: "iam user/1".to_owned(),
|
||||
cache_name: "Cache-1".to_owned(),
|
||||
region: "us-east-1".to_owned(),
|
||||
serverless,
|
||||
};
|
||||
|
||||
assert_eq!(connect_url(&auth), expected);
|
||||
}
|
||||
}
|
||||
64
litellm-rust/crates/redis-connect/src/sources/azure.rs
Normal file
64
litellm-rust/crates/redis-connect/src/sources/azure.rs
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
use litellm_auth_azure::{AzureAuthService, AzureIdentity};
|
||||
use litellm_auth_types::ResolvedCredential;
|
||||
use litellm_redis_types::AzureAdAuth;
|
||||
|
||||
use crate::Error;
|
||||
use crate::error::CredentialStage;
|
||||
|
||||
use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture};
|
||||
|
||||
const AZURE_REDIS_SCOPE: &str = "https://redis.azure.com/.default";
|
||||
|
||||
pub(crate) struct AzureAd {
|
||||
service: AzureAuthService,
|
||||
identity: AzureIdentity,
|
||||
username: Option<String>,
|
||||
lookup: SharedEnvironment,
|
||||
}
|
||||
|
||||
impl AzureAd {
|
||||
pub(crate) fn new(auth: AzureAdAuth, lookup: SharedEnvironment) -> Self {
|
||||
// _build_azure_credential in litellm/_redis.py picks the credential class the same way
|
||||
let identity = match (auth.client_id, auth.tenant_id, auth.client_secret) {
|
||||
(Some(client_id), Some(tenant_id), Some(client_secret)) => {
|
||||
AzureIdentity::ClientSecret {
|
||||
tenant_id,
|
||||
client_id,
|
||||
client_secret,
|
||||
}
|
||||
}
|
||||
(Some(client_id), _, _) => AzureIdentity::ManagedIdentity { client_id },
|
||||
(None, _, _) => AzureIdentity::Default,
|
||||
};
|
||||
Self {
|
||||
service: AzureAuthService::default(),
|
||||
identity,
|
||||
username: auth.username,
|
||||
lookup,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CredentialSource for AzureAd {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
Box::pin(async move {
|
||||
let credential = self
|
||||
.service
|
||||
.identity_token(&self.identity, AZURE_REDIS_SCOPE, &|name| {
|
||||
self.lookup.get(name)
|
||||
})
|
||||
.await
|
||||
.map_err(Error::credential(CredentialStage::Identity))?
|
||||
.ok_or(Error::MissingCredential { provider: "Azure" })?;
|
||||
let (token, expires_at) = match credential {
|
||||
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
|
||||
ResolvedCredential::Static(token) => (token, None),
|
||||
};
|
||||
Ok(Expiring {
|
||||
username: self.username.clone(),
|
||||
password: token,
|
||||
expires_at,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
126
litellm-rust/crates/redis-connect/src/sources/gcp.rs
Normal file
126
litellm-rust/crates/redis-connect/src/sources/gcp.rs
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
use std::time::SystemTime;
|
||||
|
||||
use litellm_auth_gcp::{CredentialSource as GcpCredentialSource, GcpAuth};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::Error;
|
||||
use crate::error::CredentialStage;
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture};
|
||||
|
||||
const IAM_CREDENTIALS: &str = "https://iamcredentials.googleapis.com/v1";
|
||||
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
|
||||
|
||||
pub(crate) struct GcpIam {
|
||||
http: litellm_http::Client,
|
||||
identity: GcpAuth,
|
||||
service_account: String,
|
||||
lookup: SharedEnvironment,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccessToken {
|
||||
access_token: String,
|
||||
expire_time: String,
|
||||
}
|
||||
|
||||
impl GcpIam {
|
||||
pub(crate) fn new(
|
||||
http: litellm_http::Client,
|
||||
service_account: String,
|
||||
lookup: SharedEnvironment,
|
||||
) -> Self {
|
||||
Self {
|
||||
http,
|
||||
identity: GcpAuth::default(),
|
||||
service_account,
|
||||
lookup,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CredentialSource for GcpIam {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
Box::pin(async move {
|
||||
let lookup = |name: &str| self.lookup.get(name);
|
||||
let caller = self
|
||||
.identity
|
||||
.access_token(GcpCredentialSource::application_default(&lookup))
|
||||
.await
|
||||
.map_err(Error::credential(CredentialStage::Identity))?;
|
||||
let response = self
|
||||
.http
|
||||
.post(format!(
|
||||
"{IAM_CREDENTIALS}/{}:generateAccessToken",
|
||||
self.service_account
|
||||
))
|
||||
.bearer_auth(caller.expose())
|
||||
.json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] }))
|
||||
.send()
|
||||
.await
|
||||
.map_err(Error::credential(CredentialStage::Exchange))?;
|
||||
let status = response.status().as_u16();
|
||||
let reply: AccessToken = response
|
||||
.error_for_status()
|
||||
.map_err(litellm_http::Error::from)
|
||||
.map_err(Error::credential_response(
|
||||
CredentialStage::Exchange,
|
||||
status,
|
||||
))?
|
||||
.json()
|
||||
.await
|
||||
.map_err(litellm_http::Error::from)
|
||||
.map_err(Error::credential(CredentialStage::Response))?;
|
||||
let expires_at = time::OffsetDateTime::parse(
|
||||
&reply.expire_time,
|
||||
&time::format_description::well_known::Rfc3339,
|
||||
)
|
||||
.map(SystemTime::from)
|
||||
.map_err(Error::credential(CredentialStage::Response))?;
|
||||
Ok(Expiring {
|
||||
username: None,
|
||||
password: SecretValue::new(reply.access_token),
|
||||
expires_at: Some(expires_at),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::GcpIam;
|
||||
use crate::CredentialSource;
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() {
|
||||
let environment = |name: &str| {
|
||||
match name {
|
||||
"VERTEXAI_CREDENTIALS" => {
|
||||
Some(r#"{"type":"service_account","project_id":"vertex"}"#)
|
||||
}
|
||||
"GOOGLE_APPLICATION_CREDENTIALS" => Some("/nonexistent/redis-adc.json"),
|
||||
_ => None,
|
||||
}
|
||||
.map(str::to_owned)
|
||||
};
|
||||
let source = GcpIam::new(
|
||||
litellm_http::Client::for_test(),
|
||||
"projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com".to_owned(),
|
||||
std::sync::Arc::new(environment),
|
||||
);
|
||||
|
||||
let Err(error) = source.fetch().await else {
|
||||
panic!("a missing credentials file must fail");
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
litellm_error::find::<std::io::Error>(&error).map(std::io::Error::kind),
|
||||
Some(std::io::ErrorKind::NotFound)
|
||||
);
|
||||
}
|
||||
}
|
||||
7
litellm-rust/crates/redis-connect/src/sources/mod.rs
Normal file
7
litellm-rust/crates/redis-connect/src/sources/mod.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
mod aws;
|
||||
mod azure;
|
||||
mod gcp;
|
||||
|
||||
pub(crate) use aws::ElastiCacheIam;
|
||||
pub(crate) use azure::AzureAd;
|
||||
pub(crate) use gcp::GcpIam;
|
||||
230
litellm-rust/crates/redis-connect/tests/connect.rs
Normal file
230
litellm-rust/crates/redis-connect/tests/connect.rs
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
use litellm_redis_client::commands;
|
||||
use std::{collections::BTreeMap, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_connect::{
|
||||
CredentialSource, Error, Expiring, HostServices, SourceFuture, connect,
|
||||
};
|
||||
use litellm_redis_testing::{
|
||||
RedisDeployment, RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, Topology, eventually,
|
||||
};
|
||||
use litellm_redis_types::{ProviderHandle, Resolution};
|
||||
use litellm_secret_value::SecretValue;
|
||||
use rstest::rstest;
|
||||
|
||||
// a host hands its credential_provider object over as a handle, so these tests do the same
|
||||
fn resolved(kwargs: serde_json::Value) -> Resolution {
|
||||
let provider = kwargs
|
||||
.get("credential_provider")
|
||||
.map(|_| ProviderHandle::new(1));
|
||||
let kwargs = serde_json::from_value::<BTreeMap<String, serde_json::Value>>(kwargs).unwrap();
|
||||
litellm_redis_resolve::interpret(&litellm_redis_resolve::kwargs(kwargs), provider, None)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn services() -> HostServices {
|
||||
HostServices {
|
||||
events: None,
|
||||
http: litellm_http::Client::for_test(),
|
||||
environment: Arc::new(|_: &str| None),
|
||||
credential_provider: None,
|
||||
}
|
||||
}
|
||||
|
||||
struct HostPassword(&'static str);
|
||||
|
||||
impl CredentialSource for HostPassword {
|
||||
fn fetch(&self) -> SourceFuture<'_> {
|
||||
Box::pin(async move {
|
||||
Ok(Expiring {
|
||||
username: None,
|
||||
password: SecretValue::new(self.0),
|
||||
expires_at: None,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn port(url: &str) -> u16 {
|
||||
url.rsplit(':').next().unwrap().parse().unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trusted_ca(true, true)]
|
||||
#[case::system_roots_only(false, false)]
|
||||
#[tokio::test]
|
||||
async fn python_tls_kwargs_reach_the_server(#[case] with_ca: bool, #[case] connects: bool) {
|
||||
if rustls::crypto::CryptoProvider::get_default().is_none() {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
}
|
||||
let server = RedisTlsServer::start(false).await.unwrap();
|
||||
let kwargs = serde_json::json!({
|
||||
"host": server.host(),
|
||||
"port": server.port(),
|
||||
"ssl": true,
|
||||
"ssl_ca_certs": with_ca.then(|| server.files().ca.display().to_string()),
|
||||
"socket_timeout": 2,
|
||||
});
|
||||
|
||||
let pinged = connect(resolved(kwargs), services(), None)
|
||||
.unwrap()
|
||||
.ping()
|
||||
.await;
|
||||
|
||||
assert_eq!(pinged.is_ok(), connects, "{pinged:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::right_password("right", true)]
|
||||
#[case::wrong_password("wrong", false)]
|
||||
#[tokio::test]
|
||||
async fn static_credentials_authenticate(
|
||||
#[case] password: &str,
|
||||
#[case] connects: bool,
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let server = RedisDeployment::start(topology).await.unwrap();
|
||||
server
|
||||
.cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"])
|
||||
.await
|
||||
.unwrap();
|
||||
let kwargs = match &server {
|
||||
RedisDeployment::Standalone(server) => serde_json::json!({
|
||||
"host": "127.0.0.1",
|
||||
"port": port(server.url()),
|
||||
"password": password,
|
||||
"socket_timeout": 2,
|
||||
}),
|
||||
RedisDeployment::Cluster(cluster) => serde_json::json!({
|
||||
"startup_nodes": cluster.nodes().map(|(host, port)| serde_json::json!({
|
||||
"host": host,
|
||||
"port": port,
|
||||
})).collect::<Vec<_>>(),
|
||||
"password": password,
|
||||
"socket_timeout": 2,
|
||||
}),
|
||||
};
|
||||
|
||||
let pinged = connect(resolved(kwargs), services(), None)
|
||||
.unwrap()
|
||||
.ping()
|
||||
.await;
|
||||
|
||||
assert_eq!(pinged.is_ok(), connects, "{pinged:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn python_sentinel_kwargs_follow_the_master_through_a_failover() {
|
||||
let sentinel = RedisSentinel::start().await.unwrap();
|
||||
let (host, sentinel_port) = sentinel.sentinel();
|
||||
let kwargs = serde_json::json!({
|
||||
"sentinel_nodes": [[host, sentinel_port]],
|
||||
"service_name": SERVICE_NAME,
|
||||
"socket_timeout": 1,
|
||||
});
|
||||
let client = connect(resolved(kwargs), services(), None).unwrap();
|
||||
client
|
||||
.send(commands::set("before", "1", Duration::from_secs(600)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let promoted = sentinel.failover().await.unwrap();
|
||||
eventually(async || {
|
||||
client
|
||||
.send(commands::set("after", "2", Duration::from_secs(600)))
|
||||
.await
|
||||
.ok()
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(promoted, sentinel.replica_port());
|
||||
assert_eq!(sentinel.master_port().await.unwrap(), promoted);
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>("after"))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("2")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_python_credential_provider_without_a_host_source_fails_at_boot() {
|
||||
let error = connect(
|
||||
resolved(serde_json::json!({"host": "redis", "credential_provider": true})),
|
||||
services(),
|
||||
None,
|
||||
)
|
||||
.err()
|
||||
.unwrap();
|
||||
|
||||
assert!(
|
||||
matches!(error, Error::NoHostCredentialProvider),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::profile("AWS_PROFILE")]
|
||||
#[case::role("AWS_ROLE_ARN")]
|
||||
#[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")]
|
||||
fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static str) {
|
||||
let configured = format!("{}-override", std::env::var(setting).unwrap_or_default());
|
||||
let host = HostServices {
|
||||
environment: Arc::new(move |name: &str| (name == setting).then(|| configured.clone())),
|
||||
..services()
|
||||
};
|
||||
let error = connect(
|
||||
resolved(serde_json::json!({
|
||||
"host": "redis", "ssl": true, "aws_iam_auth": true,
|
||||
"aws_iam_user_name": "user", "aws_iam_cache_name": "cache",
|
||||
"aws_iam_region": "us-east-1",
|
||||
})),
|
||||
host,
|
||||
None,
|
||||
)
|
||||
.err()
|
||||
.expect("must not select ambient credentials");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::right_password("right", true)]
|
||||
#[case::wrong_password("wrong", false)]
|
||||
#[tokio::test]
|
||||
async fn a_python_credential_provider_authenticates_through_the_host_source(
|
||||
#[case] password: &'static str,
|
||||
#[case] connects: bool,
|
||||
) {
|
||||
let server = RedisServer::start().await.unwrap();
|
||||
server
|
||||
.cli(&["ACL", "SETUSER", "default", "on", ">right"])
|
||||
.await
|
||||
.unwrap();
|
||||
let kwargs = serde_json::json!({
|
||||
"host": "127.0.0.1",
|
||||
"port": port(server.url()),
|
||||
"password": "ignored",
|
||||
"credential_provider": true,
|
||||
"socket_timeout": 2,
|
||||
});
|
||||
|
||||
let pinged = connect(
|
||||
resolved(kwargs),
|
||||
HostServices {
|
||||
credential_provider: Some(Arc::new(HostPassword(password))),
|
||||
..services()
|
||||
},
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
.ping()
|
||||
.await;
|
||||
|
||||
assert_eq!(pinged.is_ok(), connects, "{pinged:?}");
|
||||
}
|
||||
17
litellm-rust/crates/redis-connect/tests/live/aws.rs
Normal file
17
litellm-rust/crates/redis-connect/tests/live/aws.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use rstest::{fixture, rstest};
|
||||
|
||||
use super::{LiveConfig, config, roundtrip};
|
||||
|
||||
#[fixture]
|
||||
fn live_config() -> Option<LiveConfig> {
|
||||
config("LITELLM_REDIS_AWS_LIVE_CONFIG")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option<LiveConfig>) {
|
||||
let Some(config) = live_config else {
|
||||
return;
|
||||
};
|
||||
roundtrip(config).await;
|
||||
}
|
||||
17
litellm-rust/crates/redis-connect/tests/live/azure.rs
Normal file
17
litellm-rust/crates/redis-connect/tests/live/azure.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use rstest::{fixture, rstest};
|
||||
|
||||
use super::{LiveConfig, config, roundtrip};
|
||||
|
||||
#[fixture]
|
||||
fn live_config() -> Option<LiveConfig> {
|
||||
config("LITELLM_REDIS_AZURE_LIVE_CONFIG")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option<LiveConfig>) {
|
||||
let Some(config) = live_config else {
|
||||
return;
|
||||
};
|
||||
roundtrip(config).await;
|
||||
}
|
||||
17
litellm-rust/crates/redis-connect/tests/live/gcp.rs
Normal file
17
litellm-rust/crates/redis-connect/tests/live/gcp.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use rstest::{fixture, rstest};
|
||||
|
||||
use super::{LiveConfig, config, roundtrip};
|
||||
|
||||
#[fixture]
|
||||
fn live_config() -> Option<LiveConfig> {
|
||||
config("LITELLM_REDIS_GCP_LIVE_CONFIG")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option<LiveConfig>) {
|
||||
let Some(config) = live_config else {
|
||||
return;
|
||||
};
|
||||
roundtrip(config).await;
|
||||
}
|
||||
111
litellm-rust/crates/redis-connect/tests/live/main.rs
Normal file
111
litellm-rust/crates/redis-connect/tests/live/main.rs
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
mod aws;
|
||||
mod azure;
|
||||
mod gcp;
|
||||
|
||||
use std::{collections::BTreeMap, path::PathBuf, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_http::{
|
||||
ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution,
|
||||
media::PublicDnsResolver,
|
||||
};
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_connect::{HostServices, connect};
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct LiveConfig {
|
||||
redis: BTreeMap<String, serde_json::Value>,
|
||||
hold_seconds: u64,
|
||||
}
|
||||
|
||||
fn config(variable: &str) -> Option<LiveConfig> {
|
||||
let path = PathBuf::from(std::env::var_os(variable)?);
|
||||
Some(
|
||||
serde_json::from_slice(&std::fs::read(path).expect("read live Redis config"))
|
||||
.expect("parse live Redis config"),
|
||||
)
|
||||
}
|
||||
|
||||
async fn roundtrip(config: LiveConfig) {
|
||||
if rustls::crypto::CryptoProvider::get_default().is_none() {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
}
|
||||
let lookup = |name: &str| std::env::var(name).ok();
|
||||
let resolved = litellm_redis_resolve::interpret(
|
||||
&litellm_redis_resolve::kwargs(config.redis),
|
||||
None,
|
||||
lookup("REDIS_USERNAME").as_deref(),
|
||||
)
|
||||
.expect("resolve live Redis config");
|
||||
let http_settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&lookup)]);
|
||||
let pool = HttpClientPool::new(Arc::new(PublicDnsResolver));
|
||||
let http = pool
|
||||
.client(
|
||||
&Resolution::from(&http_settings).config,
|
||||
ClientVariant::Provider,
|
||||
)
|
||||
.expect("build production HTTP client");
|
||||
let services = || HostServices {
|
||||
events: None,
|
||||
http: http.clone(),
|
||||
environment: Arc::new(|name: &str| std::env::var(name).ok()),
|
||||
credential_provider: None,
|
||||
};
|
||||
let client = connect(resolved.clone(), services(), None).expect("build live Redis client");
|
||||
let marker = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos();
|
||||
let key = format!("litellm:redis-connect:live:{}:{marker}", std::process::id());
|
||||
let ttl = Duration::from_secs(config.hold_seconds + 120);
|
||||
|
||||
client.ping().await.expect("authenticate and ping");
|
||||
client
|
||||
.send(commands::set(&key, "before-refresh", ttl))
|
||||
.await
|
||||
.expect("set before refresh");
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>(&key))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("before-refresh")
|
||||
);
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(config.hold_seconds)).await;
|
||||
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>(&key))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("before-refresh")
|
||||
);
|
||||
let fresh = connect(resolved, services(), None).expect("build a fresh live Redis client");
|
||||
assert_eq!(
|
||||
fresh
|
||||
.send(commands::get::<String>(&key))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("before-refresh")
|
||||
);
|
||||
fresh
|
||||
.send(commands::set(&key, "after-refresh", ttl))
|
||||
.await
|
||||
.expect("set after refresh");
|
||||
assert_eq!(
|
||||
client
|
||||
.send(commands::get::<String>(&key))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some("after-refresh")
|
||||
);
|
||||
client
|
||||
.send(commands::del(&[&key]))
|
||||
.await
|
||||
.expect("delete live test key");
|
||||
}
|
||||
31
litellm-rust/crates/redis-resolve/AGENTS.md
Normal file
31
litellm-rust/crates/redis-resolve/AGENTS.md
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
## Interpreting merged kwargs
|
||||
|
||||
- The boundary is merged redis-py kwargs in, `litellm_redis_types::Resolution` out. `interpret` takes the dict a host already merged the way `_merge_redis_kwargs` in `litellm/_redis.py` does, with its `REDIS_*` values and `get_secret` fallbacks filled, and stands in for redis-py's constructors and `from_url`. It never reads the environment or a secret, so both hosts share it. Every `Entry` keeps the source a warning or error names, which `kwargs` sets to the kwarg name and the gateway merge sets to `REDIS_<KWARG>` or the fallback name
|
||||
- The `Resolution` holds the target, the auth and the timeouts Rust models, and every setting the Rust client ignores as an `Ignored` with its source and consequence, which the host logs. A host passes its `credential_provider` object as a `ProviderHandle`, and a `credential_provider` key without one is `Error::ProviderWithoutHandle`. Python's Azure AD auth reads `REDIS_USERNAME` from `os.environ` instead of the merged kwargs, so it is a required argument of `interpret` that each host reads from its own process environment, the gateway through its `Lookup` and python-bridge through Python's `os.environ`. An empty one is unset, like `or None`
|
||||
- `options.rs` reads the kwargs into the typed `Options` and collects the ignored ones. `endpoint.rs`, `auth.rs` and `tls.rs` pick the target, the credentials and TLS from `Options` alone
|
||||
- Rust never introspects redis-py. Every Redis kwarg redis-py or `litellm/_redis.py` accepts is one row in `SETTINGS` in `src/keys.rs`, and that row is the whole answer for the kwarg: its `Handling` (`Read`, `Warn` with the consequence the warning names, or `Reject`), whether a url query carries it (`FromUrl`), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_<KWARG>`, and the names `_merge_redis_kwargs` falls back to. The last two columns are only read by the gateway merge, and they stay on the row so there is one table. A `Read` row is a field of `Options`, which an inline test pins. A kwarg with no row is not a Redis setting, so it is dropped without a warning
|
||||
- `redis-py-kwargs.json` holds the Python facts each row answers: which redis-py versions name the kwarg in which signature, and whether Python reads its `REDIS_*` name under the locked redis-py. Never edit it by hand. `redis_py_kwargs.py` next to it regenerates it, and its docstring says how to run it. A kwarg litellm reads that no signature or `REDIS_*` name covers goes in its `LITELLM_ONLY`, and a new redis-py version in its `VERSIONS` and the `redis-compat` CI matrix together. `test_redis_py_kwargs_catalog_matches_the_installed_redis_py` in `tests/unit/test_redis.py` checks it against each version the `redis-compat` CI matrix installs, and `tests/catalog.rs` fails until `SETTINGS` has exactly one row per kwarg in it with the same `environment`. A redis-py bump that adds a kwarg therefore cannot land without a row deciding how Rust handles it
|
||||
- The contract is `resolves_like_python_redis` in `tests/interpret.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Each case holds the merged dict, the sources the merge gave it, the `REDIS_USERNAME` the host passes, and the expected outcome written next to it. `Resolves` is the view both sides build. `Refuses` is a config Rust refuses and Python connects with, a disagreement on purpose. `Fails` is a config redis-py 5.3.1 already raises on, building the client or its first connection, so Rust failing changes no deployment that works today. Its Rust side names the class `Error::raised_building_the_python_client` answers, `None` when Python only raises once a command connects, and its Python mirror `_Fails` asserts the class and message redis-py raises. The Python side asserts that `_merge_redis_kwargs` returns that merged dict and reads the effective redis-py 5.3.1 connection objects without connecting, and `merges_like_python_redis` in `gateway/src/redis/surface_tests.rs` pins the same merged dicts from the Rust merge. A new rule is a new case in every file, never a generated fixture. Python and Rust pods share one Redis during the switch, so a change here is a change to a wire contract
|
||||
- A url carries its own TLS. redis-py's url allowlist is the plain `Connection` signature, so every `ssl_*` setting outside the url is dropped, and a credential provider strips the url's userinfo. A `rediss://` url's `ssl_cert_reqs`, `ssl_ca_certs`, `ssl_ca_data`, `ssl_ca_path`, `ssl_certfile`, `ssl_keyfile`, `ssl_password` and `ssl_check_hostname` query parameters are read like `parse_url` reads them
|
||||
- TLS options and `db` are only read for the endpoint that uses them, so a bad `ssl_cert_reqs` never fails a url target. A cluster with a non-zero `db`, and Sentinel with a `host` or `port`, are errors because redis-py raises building those clients
|
||||
- An empty value (`""`, `[]`, `{}`) is the same as an unset one, like Python's `x or ...`. A null is unset too, except `ssl_cert_reqs`, where redis-py reads `None` as `CERT_NONE`, so a null skips certificate verification
|
||||
- `ssl_cert_reqs` takes `none`, `optional` and `required` case-sensitively like redis-py's lookup, so `None` (what `REDIS_SSL_CERT_REQS=None` gives) is `Error::Invalid`, and an `ssl.VerifyMode` int, where 0 skips verification and 1 or 2 verifies. `optional` verifies, because rustls has no optional mode
|
||||
- Precedence matches Python. The target is cluster, then URL, then Sentinel (`sentinel_nodes` plus `service_name`), then host. Auth is a `credential_provider` handle, then GCP, then Azure, then AWS, and the first one wins with a warning
|
||||
- Numbers are read with `litellm_python_compat::number`, so whitespace and digit underscores parse like Python's `int()` and `float()`. A value they cannot parse (`port: abc`, `socket_timeout: slow`) is dropped with a warning and the default is used, like `_coerce_redis_kwargs_types`. A number that parses but is out of range (negative seconds, a port above 65535) fails with `Error::Invalid`
|
||||
- Flags (`ssl`, `ssl_check_hostname`, `aws_iam_auth`, `aws_iam_serverless`) are read one way on every surface. `true`, `1` and `yes` are true and `false`, `0`, `no`, `off`, `f` and `n` are false, the text redis-py's `_str_to_bool` and the proxy's pydantic models agree on. Any other text is `Error::Invalid`, so a TLS setting is never guessed
|
||||
- The resolver fails on a setting it cannot honor only when that setting is a security check. A `Warn` row is dropped and returned as an `Ignored` naming its source (`ssl_ciphers` or `REDIS_SSL_CIPHERS`) and the consequence on its row, because a Python deployment that worked must keep working when its stores move to Rust. Only a value that is malformed (`Error::Invalid` or `Error::UrlQuery`), a target that is missing or a `Reject` row that asks for its check fails. Promoting a `Warn` row to `Read` is then a pure improvement with no behavior cliff, while moving one to `Reject` can fail a config that connects today
|
||||
- A `Warn` row with a `met` check is dropped silently when its value asks for nothing more than Rust does, so `ssl_min_version` of TLS 1.2 or lower never warns
|
||||
- A `Reject` row is a security check rustls cannot perform that Python's locked async client already refuses with a `TypeError`, so failing closed breaks no proxy that runs today. `tests/catalog.rs` fails if the locked redis-py's async, cluster or connection signature accepts a `Reject` row, so a redis-py bump forces the decision again. A row's `Asked` says which values ask for the check, from the kwargs, `REDIS_*` or a url query, and those are `Error::Unenforceable`
|
||||
- `ssl_validate_ocsp` and `ssl_validate_ocsp_stapled` ask when true, because rustls has no client-side OCSP. False asks for nothing and is dropped without a warning
|
||||
- `ssl_include_verify_flags` and `ssl_exclude_verify_flags` ask when set, because rustls takes no X.509 verify flags such as CRL checking. An empty list asks for nothing
|
||||
- `ssl_ca_data` and `ssl_ca_path` are read into `TlsOptions` and trusted next to `ssl_ca_certs`. `ssl_password` is read into `TlsOptions::key_password`, which `redis-client` uses to decrypt a PKCS#8 key like the sync client's `load_cert_chain(password=...)`. redis-py 5.3.1's async client refuses it, so this is a superset of the path the proxy uses, like `ssl_ca_path`
|
||||
- `unix_socket_path` resolves to `Target::Unix`, which like redis-py's `UnixDomainSocketConnection` replaces host, port and TLS. Python still requires `host` next to it, so Rust does too
|
||||
- A url's `?db=`, `?username=`, `?password=`, `?socket_timeout=`, `?socket_connect_timeout=` and `?health_check_interval=` are read like `parse_url` reads them, so `?db=` wins over the path, the userinfo wins over `?username=` and `?password=`, and the timeouts and health check win over the kwargs like `from_url`. A repeated parameter takes its first value like `parse_qs`. A number `int()` or `float()` cannot read is `Error::UrlQuery`, as `parse_url` raises `ValueError`. Any other query parameter, including one with no row, is ignored naming `url ?<name>`, because redis-py turns each one into a client kwarg. Which ones a url carries is the `FromUrl` column
|
||||
- `socket_keepalive` is a flag that defaults to true on every target, like `_get_redis_cluster_kwargs` defaults it for a cluster. Python leaves it off for a standalone or Sentinel client, so an unset value there is a difference on purpose, because without keepalive a load balancer drops an idle connection silently. False turns it off. `socket_keepalive_options` is warned, and `redis-client` probes after its own idle time instead
|
||||
- `health_check_interval` is read as an `int` of seconds like `_coerce_redis_kwargs_types` into `HealthCheck`. Zero is `Disabled` and stays off on every target, a negative interval checks before every command like redis-py, and only `Unset` takes a default. `redis-connect` gives an unset cluster Python's 25s default and `redis-client` pings every node before the next command once the client sat idle that long, outside the breaker and the call events
|
||||
- `Error::raised_building_the_python_client` is how python-bridge treats a config Rust cannot use. Where Python raises building its client, the bridge raises the same class from `NativeRedisClient.connect`. Where Python builds the client and only fails once a command connects, which `RedisCache` catches and logs on every call, the stores stay on memory and the bridge logs an error
|
||||
- Cases where Rust and Python disagree on purpose are cases in both contract tests with a different expected outcome on each side
|
||||
- `ssl_ca_path`: redis-py 5.3.1's async `Redis` has no such parameter and raises `TypeError` building the client, while Rust trusts the directory like `load_verify_locations(capath=...)`
|
||||
- `on`, `t` and `y` for a flag are `Error::Invalid` on every surface. pydantic reads them as true for `coordination_redis` and `_str_to_bool` reads them as false everywhere else, so Rust refuses to pick one. Other unrecognized text is false in Python and an error here
|
||||
- `ssl_cert_reqs` as an `ssl.VerifyMode` int: redis-py 5.3.1's sync client takes it, while its async `RedisSSLContext` leaves `cert_reqs` unset and raises `AttributeError` once a command connects. Rust reads it like the sync client
|
||||
- Empty `sentinel_nodes` with no `REDIS_SENTINEL_NODES` is no target at all, so it fails as "either host or url must be specified". redis-py keeps the empty list and fails building the Sentinel client
|
||||
20
litellm-rust/crates/redis-resolve/Cargo.toml
Normal file
20
litellm-rust/crates/redis-resolve/Cargo.toml
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
[package]
|
||||
name = "litellm-redis-resolve"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-python-compat.workspace = true
|
||||
litellm-redis-types.workspace = true
|
||||
litellm-secret-value.workspace = true
|
||||
serde_json.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
tracing.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
serde.workspace = true
|
||||
118
litellm-rust/crates/redis-resolve/redis-py-kwargs.json
Normal file
118
litellm-rust/crates/redis-resolve/redis-py-kwargs.json
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
{
|
||||
"meta": {
|
||||
"description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here",
|
||||
"redis_py_versions": [
|
||||
"5.3.1",
|
||||
"6.4.0",
|
||||
"7.4.1",
|
||||
"8.0.1"
|
||||
],
|
||||
"python": "3.12",
|
||||
"method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self",
|
||||
"signatures": {
|
||||
"sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_<KWARG> environment mapping",
|
||||
"async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built",
|
||||
"connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs",
|
||||
"cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs"
|
||||
},
|
||||
"environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_<ARG> under the locked redis-py",
|
||||
"redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it",
|
||||
"checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it",
|
||||
"redis_py_locked": "5.3.1"
|
||||
},
|
||||
"args": {
|
||||
"address_remap": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"auto_close_connection_pool": {"environment": false, "redis_py": {"5.3.1": ["async"], "6.4.0": ["async"], "7.4.1": ["async"], "8.0.1": ["async"]}},
|
||||
"aws_iam_auth": {"environment": true, "redis_py": {}},
|
||||
"aws_iam_cache_name": {"environment": true, "redis_py": {}},
|
||||
"aws_iam_region": {"environment": true, "redis_py": {}},
|
||||
"aws_iam_serverless": {"environment": true, "redis_py": {}},
|
||||
"aws_iam_user_name": {"environment": true, "redis_py": {}},
|
||||
"azure_client_id": {"environment": true, "redis_py": {}},
|
||||
"azure_client_secret": {"environment": true, "redis_py": {}},
|
||||
"azure_redis_ad_token": {"environment": true, "redis_py": {}},
|
||||
"azure_tenant_id": {"environment": true, "redis_py": {}},
|
||||
"cache": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}},
|
||||
"cache_config": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}},
|
||||
"charset": {"environment": true, "redis_py": {"5.3.1": ["sync"]}},
|
||||
"client_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"cluster_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"command_packer": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"connection_class": {"environment": false, "redis_py": {}},
|
||||
"connection_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"]}},
|
||||
"connection_pool": {"environment": false, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}},
|
||||
"credential_provider": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"db": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"decode_responses": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"driver_info": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"dynamic_startup_nodes": {"environment": false, "redis_py": {"6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"encoder_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"encoding": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"encoding_errors": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"errors": {"environment": true, "redis_py": {"5.3.1": ["sync"]}},
|
||||
"event_dispatcher": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"gcp_service_account": {"environment": true, "redis_py": {}},
|
||||
"gcp_ssl_ca_certs": {"environment": true, "redis_py": {}},
|
||||
"health_check_interval": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"host": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"legacy_responses": {"environment": false, "redis_py": {"8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"lib_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"lib_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"load_balancing_strategy": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"maint_notifications_config": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}},
|
||||
"maint_notifications_pool_handler": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"maintenance_notification_hash": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"maintenance_state": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"max_connections": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"namespace": {"environment": false, "redis_py": {}},
|
||||
"orig_host_address": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"orig_socket_connect_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"orig_socket_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"oss_cluster_maint_notifications_handler": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}},
|
||||
"parser": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"parser_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"password": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"path": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"policy_resolver": {"environment": false, "redis_py": {"7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"port": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"protocol": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"read_from_replicas": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"redis_connect_func": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}},
|
||||
"reinitialize_steps": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"require_full_coverage": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"retry": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"retry_on_error": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"retry_on_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}},
|
||||
"sentinel_nodes": {"environment": false, "redis_py": {}},
|
||||
"sentinel_password": {"environment": false, "redis_py": {}},
|
||||
"service_name": {"environment": false, "redis_py": {}},
|
||||
"single_connection_client": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}},
|
||||
"socket_connect_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"socket_keepalive": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"socket_keepalive_options": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"socket_read_size": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"socket_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}},
|
||||
"socket_type": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}},
|
||||
"ssl": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_ca_certs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_ca_data": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_ca_path": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}},
|
||||
"ssl_cert_reqs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_certfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_check_hostname": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_ciphers": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_exclude_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_include_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_keyfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_min_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}},
|
||||
"ssl_ocsp_context": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}},
|
||||
"ssl_ocsp_expected_cert": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}},
|
||||
"ssl_password": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}},
|
||||
"ssl_validate_ocsp": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}},
|
||||
"ssl_validate_ocsp_stapled": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}},
|
||||
"startup_nodes": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}},
|
||||
"unix_socket_path": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}},
|
||||
"url": {"environment": true, "redis_py": {}},
|
||||
"username": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}
|
||||
}
|
||||
}
|
||||
110
litellm-rust/crates/redis-resolve/redis_py_kwargs.py
Normal file
110
litellm-rust/crates/redis-resolve/redis_py_kwargs.py
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
"""Regenerates redis-py-kwargs.json, the Python facts SETTINGS in src/keys.rs answers.
|
||||
|
||||
Run it from the repository root in the project environment, which has the locked redis-py and litellm:
|
||||
|
||||
uv run python litellm-rust/crates/redis-resolve/redis_py_kwargs.py
|
||||
|
||||
It reads every other redis-py version's signatures in a throwaway environment through
|
||||
`uv run --no-project --with redis==<version>`, so nothing else is installed. Then update SETTINGS until
|
||||
tests/catalog.rs passes.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
VERSIONS: Final = ("5.3.1", "6.4.0", "7.4.1", "8.0.1")
|
||||
PYTHON: Final = "3.12"
|
||||
CATALOG: Final = Path(__file__).with_name("redis-py-kwargs.json")
|
||||
# Keys litellm/_redis.py or its callers read that neither a redis-py signature nor the REDIS_* mapping names
|
||||
LITELLM_ONLY: Final = frozenset(
|
||||
{"connection_class", "namespace", "sentinel_nodes", "sentinel_password", "service_name"}
|
||||
)
|
||||
|
||||
META: Final = {
|
||||
"description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here",
|
||||
"redis_py_versions": list(VERSIONS),
|
||||
"python": PYTHON,
|
||||
"method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self",
|
||||
"signatures": {
|
||||
"sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_<KWARG> environment mapping",
|
||||
"async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built",
|
||||
"connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs",
|
||||
"cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs",
|
||||
},
|
||||
"environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_<ARG> under the locked redis-py",
|
||||
"redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it",
|
||||
"checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it",
|
||||
}
|
||||
|
||||
|
||||
def _named_args(*classes: type) -> frozenset[str]:
|
||||
return frozenset(
|
||||
name
|
||||
for cls in classes
|
||||
for spec in (inspect.getfullargspec(inspect.unwrap(cls.__init__)),)
|
||||
for name in spec.args + spec.kwonlyargs
|
||||
if name != "self"
|
||||
)
|
||||
|
||||
|
||||
def arg_signatures() -> dict[str, list[str]]:
|
||||
"""Each kwarg the installed redis-py names, with the signatures that name it."""
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
|
||||
connection_classes: Final = (*redis.Connection.__mro__, *async_redis.Connection.__mro__)
|
||||
signatures: Final = {
|
||||
"sync": _named_args(redis.Redis),
|
||||
"async": _named_args(async_redis.Redis),
|
||||
"connection": _named_args(*(cls for cls in connection_classes if cls is not object)),
|
||||
"cluster": _named_args(async_redis.RedisCluster),
|
||||
}
|
||||
names: Final = sorted(frozenset().union(*signatures.values()))
|
||||
return {name: [kind for kind, named in signatures.items() if name in named] for name in names}
|
||||
|
||||
|
||||
def _arg_signatures_of(version: str) -> Mapping[str, list[str]]:
|
||||
command: Final = ("uv", "run", "--no-project", "--python", PYTHON, "--with", f"redis=={version}")
|
||||
script: Final = ("python", "-I", __file__, "--arg-signatures")
|
||||
return json.loads(subprocess.run((*command, *script), check=True, capture_output=True, text=True).stdout)
|
||||
|
||||
|
||||
def catalog() -> dict[str, object]:
|
||||
import redis
|
||||
|
||||
from litellm._redis import _get_redis_env_kwarg_mapping
|
||||
|
||||
environment: Final = frozenset(_get_redis_env_kwarg_mapping().values())
|
||||
by_version: Final = {version: _arg_signatures_of(version) for version in VERSIONS}
|
||||
names: Final = sorted(LITELLM_ONLY | environment | frozenset().union(*by_version.values()))
|
||||
return {
|
||||
"meta": {**META, "redis_py_locked": redis.__version__},
|
||||
"args": {
|
||||
name: {
|
||||
"environment": name in environment,
|
||||
"redis_py": {version: args[name] for version, args in by_version.items() if name in args},
|
||||
}
|
||||
for name in names
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# one line per arg, so a version bump reads as a diff of the args it touched
|
||||
def render(catalog: Mapping[str, object]) -> str:
|
||||
meta: Final = json.dumps(catalog["meta"], indent=2).replace("\n", "\n ")
|
||||
args: Final = catalog["args"]
|
||||
assert isinstance(args, dict)
|
||||
rows: Final = ",\n".join(f" {json.dumps(name)}: {json.dumps(arg)}" for name, arg in args.items())
|
||||
return f'{{\n "meta": {meta},\n "args": {{\n{rows}\n }}\n}}\n'
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if sys.argv[1:] == ["--arg-signatures"]:
|
||||
sys.stdout.write(json.dumps(arg_signatures()))
|
||||
else:
|
||||
CATALOG.write_text(render(catalog()))
|
||||
95
litellm-rust/crates/redis-resolve/src/auth.rs
Normal file
95
litellm-rust/crates/redis-resolve/src/auth.rs
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
use litellm_redis_types::{Auth, AwsIamAuth, AzureAdAuth, ProviderHandle, StaticAuth};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
options::{self, AwsOptions, Options, flag, secret, text},
|
||||
};
|
||||
|
||||
pub(super) fn redis_auth(
|
||||
options: &Options,
|
||||
provider: Option<ProviderHandle>,
|
||||
redis_username: Option<&str>,
|
||||
) -> Result<Auth, Error> {
|
||||
match (provider, &options.credential_provider) {
|
||||
(Some(handle), _) => return Ok(Auth::Provider(handle)),
|
||||
(None, Some(_)) => return Err(Error::ProviderWithoutHandle),
|
||||
(None, None) => {}
|
||||
}
|
||||
let gcp = text(options.gcp.service_account.as_ref());
|
||||
let azure = text(options.azure.ad_token.as_ref())
|
||||
.is_some_and(|enabled| enabled.eq_ignore_ascii_case("true"));
|
||||
let aws = flag(options.aws.auth.as_ref())?;
|
||||
match (gcp, azure, aws) {
|
||||
(Some(service_account), azure, aws) => {
|
||||
if azure {
|
||||
tracing::warn!(
|
||||
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration."
|
||||
);
|
||||
}
|
||||
if aws {
|
||||
tracing::warn!(
|
||||
"Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration."
|
||||
);
|
||||
}
|
||||
Ok(Auth::GcpIam { service_account })
|
||||
}
|
||||
(None, true, aws) => {
|
||||
if aws {
|
||||
tracing::warn!(
|
||||
"Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using Azure AD. Remove one to avoid misconfiguration."
|
||||
);
|
||||
}
|
||||
Ok(Auth::AzureAd(AzureAdAuth {
|
||||
client_id: text(options.azure.client_id.as_ref()),
|
||||
tenant_id: text(options.azure.tenant_id.as_ref()),
|
||||
client_secret: secret(options.azure.client_secret.as_ref()),
|
||||
username: redis_username
|
||||
.filter(|username| !username.is_empty())
|
||||
.map(str::to_owned),
|
||||
}))
|
||||
}
|
||||
(None, false, true) => aws_iam(&options.aws, uses_tls(options)?).map(Auth::AwsIam),
|
||||
(None, false, false) => Ok(Auth::Static(StaticAuth {
|
||||
username: secret(options.username.as_ref()),
|
||||
password: secret(options.password.as_ref()),
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
// _uses_tls: a url's scheme decides unless cluster nodes are set, otherwise the ssl kwarg does
|
||||
fn uses_tls(options: &Options) -> Result<bool, Error> {
|
||||
match (&options.startup_nodes, text(options.url.as_ref())) {
|
||||
(None, Some(url)) => Ok(url
|
||||
.split_once("://")
|
||||
.is_some_and(|(scheme, _)| scheme.eq_ignore_ascii_case("rediss"))),
|
||||
_ => flag(options.ssl.as_ref()),
|
||||
}
|
||||
}
|
||||
|
||||
fn aws_iam(aws: &AwsOptions, uses_tls: bool) -> Result<AwsIamAuth, Error> {
|
||||
if !uses_tls {
|
||||
return Err(Error::AwsIamWithoutTls);
|
||||
}
|
||||
match (
|
||||
text(aws.user_name.as_ref()),
|
||||
text(aws.cache_name.as_ref()),
|
||||
text(aws.region.as_ref()),
|
||||
) {
|
||||
(Some(user_name), Some(cache_name), Some(region)) => Ok(AwsIamAuth {
|
||||
user_name,
|
||||
cache_name,
|
||||
region,
|
||||
serverless: options::flag(aws.serverless.as_ref())?,
|
||||
}),
|
||||
(user_name, cache_name, region) => Err(Error::AwsIamIncomplete {
|
||||
missing: [
|
||||
("aws_iam_user_name", user_name.is_none()),
|
||||
("aws_iam_cache_name", cache_name.is_none()),
|
||||
("aws_iam_region", region.is_none()),
|
||||
]
|
||||
.into_iter()
|
||||
.filter_map(|(name, missing)| missing.then_some(name))
|
||||
.collect(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
180
litellm-rust/crates/redis-resolve/src/endpoint.rs
Normal file
180
litellm-rust/crates/redis-resolve/src/endpoint.rs
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
use litellm_redis_types::{
|
||||
Auth, Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket,
|
||||
};
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
Entry, Error, Expected,
|
||||
options::{Options, as_text, flag, integer, invalid, secret, text},
|
||||
tls::{self, QuerySettings},
|
||||
};
|
||||
|
||||
/// The target, first match wins like `get_redis_client`: cluster, url, Sentinel, then host. A
|
||||
/// url also returns the settings its query carries.
|
||||
pub(super) fn target(options: &Options, auth: &Auth) -> Result<(Target, QuerySettings), Error> {
|
||||
if let Some(nodes) = cluster_nodes(options.startup_nodes.as_ref())? {
|
||||
reject_cluster_db(options.db.as_ref())?;
|
||||
let cluster = Cluster {
|
||||
nodes,
|
||||
tls: tls_for(options, auth)?,
|
||||
};
|
||||
return Ok((Target::Cluster(cluster), QuerySettings::default()));
|
||||
}
|
||||
if let Some(url) = text(options.url.as_ref()) {
|
||||
let query = tls::url_query(&url)?;
|
||||
let url = RedisUrl {
|
||||
url: SecretValue::new(match auth {
|
||||
Auth::Static(_) => url,
|
||||
_ => without_userinfo(&url),
|
||||
}),
|
||||
tls: query.tls,
|
||||
db: query.db,
|
||||
username: query.username,
|
||||
password: query.password,
|
||||
};
|
||||
return Ok((Target::Url(url), query.settings));
|
||||
}
|
||||
if let Some(sentinels) = options.sentinel.nodes.as_ref() {
|
||||
let sentinel = sentinel(options, auth, sentinels)?;
|
||||
return Ok((Target::Sentinel(sentinel), QuerySettings::default()));
|
||||
}
|
||||
let host = text(options.host.as_ref()).ok_or(Error::NoTarget)?;
|
||||
let target = match text(options.unix_socket_path.as_ref()) {
|
||||
Some(path) => Target::Unix(UnixSocket {
|
||||
path: path.into(),
|
||||
db: db(options)?,
|
||||
}),
|
||||
None => Target::Standalone(Standalone {
|
||||
host,
|
||||
port: integer(options.port.as_ref())?.unwrap_or(6379),
|
||||
tls: tls_for(options, auth)?,
|
||||
db: db(options)?,
|
||||
}),
|
||||
};
|
||||
Ok((target, QuerySettings::default()))
|
||||
}
|
||||
|
||||
fn tls_for(options: &Options, auth: &Auth) -> Result<Option<TlsOptions>, Error> {
|
||||
if !flag(options.ssl.as_ref())? {
|
||||
return Ok(None);
|
||||
}
|
||||
let gcp_ca_certs = match auth {
|
||||
Auth::GcpIam { .. } => options.gcp.ssl_ca_certs.as_ref(),
|
||||
_ => None,
|
||||
};
|
||||
tls::tls_options(&options.tls, gcp_ca_certs).map(Some)
|
||||
}
|
||||
|
||||
fn db(options: &Options) -> Result<i64, Error> {
|
||||
integer(options.db.as_ref()).map(|db| db.unwrap_or(0))
|
||||
}
|
||||
|
||||
fn sentinel(options: &Options, auth: &Auth, sentinels: &Entry) -> Result<Sentinel, Error> {
|
||||
let monitors = sentinel_nodes(&sentinels.value)
|
||||
.ok_or_else(|| invalid(&sentinels.source, Expected::SentinelNodes))?;
|
||||
let service_name =
|
||||
text(options.sentinel.service_name.as_ref()).ok_or(Error::NoSentinelService)?;
|
||||
// redis.asyncio.Sentinel passes the connection kwargs to each monitor's Redis(host=..., port=...), so a
|
||||
// host or port setting raises TypeError when litellm builds the client
|
||||
if let Some(entry) = options.host.as_ref().or(options.port.as_ref()) {
|
||||
return Err(Error::SentinelWithHost {
|
||||
setting: entry.source.clone(),
|
||||
});
|
||||
}
|
||||
Ok(Sentinel {
|
||||
monitors,
|
||||
service_name,
|
||||
monitor_auth: monitor_auth(options, auth),
|
||||
tls: tls_for(options, auth)?,
|
||||
db: db(options)?,
|
||||
})
|
||||
}
|
||||
|
||||
// _sentinel_auth_kwargs gives the monitors the data node's connection kwargs with sentinel_password as the
|
||||
// password, so a static username reaches them too
|
||||
fn monitor_auth(options: &Options, auth: &Auth) -> Option<StaticAuth> {
|
||||
let password = secret(options.sentinel.password.as_ref())?;
|
||||
let username = match auth {
|
||||
Auth::Static(StaticAuth { username, .. }) => username.clone(),
|
||||
_ => None,
|
||||
};
|
||||
Some(StaticAuth {
|
||||
username,
|
||||
password: Some(password),
|
||||
})
|
||||
}
|
||||
|
||||
fn cluster_nodes(entry: Option<&Entry>) -> Result<Option<Vec<Node>>, Error> {
|
||||
let Some(entry) = entry else {
|
||||
return Ok(None);
|
||||
};
|
||||
let nodes = nodes(&entry.value, |node| {
|
||||
node_address(
|
||||
node.get("host").and_then(as_text),
|
||||
node.get("port").and_then(as_text),
|
||||
)
|
||||
})
|
||||
.ok_or_else(|| invalid(&entry.source, Expected::StartupNodes))?;
|
||||
Ok((!nodes.is_empty()).then_some(nodes))
|
||||
}
|
||||
|
||||
fn sentinel_nodes(value: &Value) -> Option<Vec<Node>> {
|
||||
let nodes = nodes(value, |pair| match pair {
|
||||
Value::Array(pair) if pair.len() == 2 => node_address(as_text(&pair[0]), as_text(&pair[1])),
|
||||
_ => None,
|
||||
})?;
|
||||
(!nodes.is_empty()).then_some(nodes)
|
||||
}
|
||||
|
||||
fn nodes(value: &Value, parse: impl Fn(&Value) -> Option<Node>) -> Option<Vec<Node>> {
|
||||
match value {
|
||||
Value::String(json) => serde_json::from_str::<Vec<Value>>(json)
|
||||
.ok()?
|
||||
.iter()
|
||||
.map(parse)
|
||||
.collect(),
|
||||
Value::Array(nodes) => nodes.iter().map(parse).collect(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn node_address(host: Option<String>, port: Option<String>) -> Option<Node> {
|
||||
Some(Node {
|
||||
host: host?,
|
||||
port: port?.parse().ok()?,
|
||||
})
|
||||
}
|
||||
|
||||
// redis-py's RedisCluster raises unless db is 0 or None
|
||||
fn reject_cluster_db(db: Option<&Entry>) -> Result<(), Error> {
|
||||
let Some(entry) = db else {
|
||||
return Ok(());
|
||||
};
|
||||
let is_zero = match &entry.value {
|
||||
Value::String(_) => matches!(integer::<i64>(Some(entry)), Ok(None | Some(0))),
|
||||
Value::Bool(flag) => !flag,
|
||||
Value::Number(number) => number.as_f64() == Some(0.0),
|
||||
_ => false,
|
||||
};
|
||||
match is_zero {
|
||||
true => Ok(()),
|
||||
false => Err(Error::DbInCluster {
|
||||
setting: entry.source.clone(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// _get_redis_url_kwargs keeps only plain Connection kwargs, so redis-py drops every ssl_* setting for a url
|
||||
// and a credential provider replaces the url's userinfo, which _url_without_userinfo strips
|
||||
fn without_userinfo(url: &str) -> String {
|
||||
let Some((scheme, rest)) = url.split_once("://") else {
|
||||
return url.to_owned();
|
||||
};
|
||||
let (authority, tail) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len()));
|
||||
let host = authority
|
||||
.rsplit_once('@')
|
||||
.map_or(authority, |(_, host)| host);
|
||||
format!("{scheme}://{host}{tail}")
|
||||
}
|
||||
80
litellm-rust/crates/redis-resolve/src/error.rs
Normal file
80
litellm-rust/crates/redis-resolve/src/error.rs
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("invalid Redis setting {setting}: expected {expected}")]
|
||||
Invalid { setting: String, expected: Expected },
|
||||
#[error("invalid Redis setting {setting}: redis-py's parse_url cannot read it as a number")]
|
||||
UrlQuery { setting: String },
|
||||
#[error("either host or url must be specified for redis")]
|
||||
NoTarget,
|
||||
#[error("service_name is required with sentinel_nodes")]
|
||||
NoSentinelService,
|
||||
#[error("{setting} cannot be combined with sentinel_nodes, which redis-py rejects at startup")]
|
||||
SentinelWithHost { setting: String },
|
||||
#[error("{setting} must be 0 in cluster mode")]
|
||||
DbInCluster { setting: String },
|
||||
#[error("credential_provider is a host object, and the host passed no handle for it")]
|
||||
ProviderWithoutHandle,
|
||||
#[error(
|
||||
"{setting} asks for {check}, which rustls cannot perform, so Redis is not connected without them"
|
||||
)]
|
||||
Unenforceable {
|
||||
setting: String,
|
||||
check: &'static str,
|
||||
},
|
||||
#[error("AWS ElastiCache IAM Redis authentication requires TLS")]
|
||||
AwsIamWithoutTls,
|
||||
#[error("AWS ElastiCache IAM Redis authentication requires: {}", missing.join(", "))]
|
||||
AwsIamIncomplete { missing: Vec<&'static str> },
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)]
|
||||
pub enum Expected {
|
||||
#[strum(serialize = "an integer")]
|
||||
Integer,
|
||||
#[strum(serialize = "a boolean")]
|
||||
Boolean,
|
||||
#[strum(serialize = "a non-negative number of seconds")]
|
||||
Seconds,
|
||||
#[strum(serialize = "none, optional, required, null or an ssl.VerifyMode value")]
|
||||
CertRequirement,
|
||||
#[strum(serialize = "a list of objects with host and port")]
|
||||
StartupNodes,
|
||||
#[strum(serialize = "a non-empty list of [host, port] pairs")]
|
||||
SentinelNodes,
|
||||
}
|
||||
|
||||
/// The exception class redis-py 5.3.1 raises building its client.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum PythonException {
|
||||
ValueError,
|
||||
TypeError,
|
||||
RedisClusterException,
|
||||
}
|
||||
|
||||
impl Error {
|
||||
/// What `_get_redis_client_logic` and redis-py's async client constructors raise on the same kwargs, or `None`
|
||||
/// when Python builds the client and fails only once a command opens a connection, or never fails.
|
||||
pub fn raised_building_the_python_client(&self) -> Option<PythonException> {
|
||||
match self {
|
||||
Self::Invalid { expected, .. } => match expected {
|
||||
Expected::StartupNodes | Expected::SentinelNodes => {
|
||||
Some(PythonException::ValueError)
|
||||
}
|
||||
Expected::Integer
|
||||
| Expected::Boolean
|
||||
| Expected::Seconds
|
||||
| Expected::CertRequirement => None,
|
||||
},
|
||||
Self::UrlQuery { .. }
|
||||
| Self::NoTarget
|
||||
| Self::NoSentinelService
|
||||
| Self::AwsIamWithoutTls
|
||||
| Self::AwsIamIncomplete { .. } => Some(PythonException::ValueError),
|
||||
Self::SentinelWithHost { .. } | Self::Unenforceable { .. } => {
|
||||
Some(PythonException::TypeError)
|
||||
}
|
||||
Self::DbInCluster { .. } => Some(PythonException::RedisClusterException),
|
||||
Self::ProviderWithoutHandle => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
304
litellm-rust/crates/redis-resolve/src/keys.rs
Normal file
304
litellm-rust/crates/redis-resolve/src/keys.rs
Normal file
|
|
@ -0,0 +1,304 @@
|
|||
use litellm_redis_types::Ignored;
|
||||
use serde_json::Value;
|
||||
|
||||
/// How the Rust client honors a redis-py kwarg.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum Handling {
|
||||
/// Modeled by the resolver into a field of `Options`.
|
||||
Read,
|
||||
/// Dropped and returned as an `Ignored` naming `consequence`, unless `met` says the value asks for nothing the
|
||||
/// Rust client does not already do.
|
||||
Warn {
|
||||
consequence: &'static str,
|
||||
met: Option<fn(&Value) -> bool>,
|
||||
},
|
||||
/// A security check the Rust client cannot perform and Python's async client already refuses, so asking for it
|
||||
/// fails instead of connecting without it.
|
||||
Reject { check: &'static str, asked: Asked },
|
||||
}
|
||||
|
||||
/// Which values of a `Reject` row ask for its check.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Asked {
|
||||
/// A flag read as true. False asks for nothing.
|
||||
WhenTrue,
|
||||
/// Any present value, like a list of flags. An empty one asks for nothing.
|
||||
WhenSet,
|
||||
}
|
||||
|
||||
/// Whether redis-py's `parse_url` reads the setting from a url's query and the Rust resolver does too.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum FromUrl {
|
||||
No,
|
||||
Any,
|
||||
Rediss,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct Setting {
|
||||
pub kwarg: &'static str,
|
||||
pub handling: Handling,
|
||||
pub from_url: FromUrl,
|
||||
/// `_get_redis_env_kwarg_mapping` reads it as `REDIS_<KWARG>`.
|
||||
pub environment: bool,
|
||||
/// The names `_merge_redis_kwargs` reads with `get_secret` when the merged kwargs lack the setting.
|
||||
pub fallback: &'static [&'static str],
|
||||
}
|
||||
|
||||
impl Setting {
|
||||
const fn new(kwarg: &'static str, handling: Handling) -> Self {
|
||||
Self {
|
||||
kwarg,
|
||||
handling,
|
||||
from_url: FromUrl::No,
|
||||
environment: false,
|
||||
fallback: &[],
|
||||
}
|
||||
}
|
||||
|
||||
const fn env(self) -> Self {
|
||||
Self {
|
||||
environment: true,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
const fn in_url(self, from_url: FromUrl) -> Self {
|
||||
Self { from_url, ..self }
|
||||
}
|
||||
|
||||
const fn or_secrets(self, fallback: &'static [&'static str]) -> Self {
|
||||
Self { fallback, ..self }
|
||||
}
|
||||
|
||||
pub fn environment_name(&self) -> Option<String> {
|
||||
self.environment
|
||||
.then(|| format!("REDIS_{}", self.kwarg.to_ascii_uppercase()))
|
||||
}
|
||||
|
||||
pub fn is_read(&self) -> bool {
|
||||
matches!(self.handling, Handling::Read)
|
||||
}
|
||||
|
||||
/// The warning a set value earns, or `None` when the setting is read, refused or already met.
|
||||
pub(crate) fn warning(&self, value: &Value) -> Option<&'static str> {
|
||||
match self.handling {
|
||||
Handling::Warn { consequence, met } if !met.is_some_and(|met| met(value)) => {
|
||||
Some(consequence)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const fn read(kwarg: &'static str) -> Setting {
|
||||
Setting::new(kwarg, Handling::Read)
|
||||
}
|
||||
|
||||
const fn warn(kwarg: &'static str, consequence: &'static str) -> Setting {
|
||||
Setting::new(
|
||||
kwarg,
|
||||
Handling::Warn {
|
||||
consequence,
|
||||
met: None,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
const fn warn_unless(
|
||||
kwarg: &'static str,
|
||||
consequence: &'static str,
|
||||
met: fn(&Value) -> bool,
|
||||
) -> Setting {
|
||||
Setting::new(
|
||||
kwarg,
|
||||
Handling::Warn {
|
||||
consequence,
|
||||
met: Some(met),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
const fn reject(kwarg: &'static str, check: &'static str, asked: Asked) -> Setting {
|
||||
Setting::new(kwarg, Handling::Reject { check, asked })
|
||||
}
|
||||
|
||||
const NO_EQUIVALENT: &str = "the Rust client has no equivalent";
|
||||
const OWN_STACK: &str = "the Rust client uses its own pool, retries and events";
|
||||
const OWN_CONNECTION: &str = "it tunes redis-py's own connection, and the Rust client has its own";
|
||||
const OWN_CLUSTER: &str = "the Rust cluster client discovers, routes and retries on its own";
|
||||
const MAINTENANCE: &str = "the Rust client does not follow server maintenance notifications";
|
||||
const RUSTLS: &str = "rustls negotiates TLS 1.2+ with its own ECDHE AEAD suites";
|
||||
const REVOCATION: &str = "OCSP revocation checks";
|
||||
const OCSP: &str = "it only configures ssl_validate_ocsp, which the Rust client refuses when set";
|
||||
const VERIFY_FLAGS: &str = "X.509 verify flags such as CRL checking";
|
||||
|
||||
/// Every Redis kwarg redis-py or `litellm/_redis.py` accepts, one row each. `redis-py-kwargs.json` records which
|
||||
/// redis-py versions and `REDIS_*` names accept them, and `tests/catalog.rs` keeps the two in step.
|
||||
pub const SETTINGS: &[Setting] = &[
|
||||
warn("address_remap", OWN_CLUSTER),
|
||||
warn("auto_close_connection_pool", OWN_STACK),
|
||||
read("aws_iam_auth").env(),
|
||||
read("aws_iam_cache_name").env(),
|
||||
read("aws_iam_region")
|
||||
.env()
|
||||
.or_secrets(&["AWS_REGION", "AWS_DEFAULT_REGION"]),
|
||||
read("aws_iam_serverless").env(),
|
||||
read("aws_iam_user_name").env(),
|
||||
read("azure_client_id")
|
||||
.env()
|
||||
.or_secrets(&["AZURE_CLIENT_ID"]),
|
||||
read("azure_client_secret")
|
||||
.env()
|
||||
.or_secrets(&["AZURE_CLIENT_SECRET"]),
|
||||
read("azure_redis_ad_token")
|
||||
.env()
|
||||
.or_secrets(&["REDIS_AZURE_AD_TOKEN"]),
|
||||
read("azure_tenant_id")
|
||||
.env()
|
||||
.or_secrets(&["AZURE_TENANT_ID"]),
|
||||
warn("cache", NO_EQUIVALENT).env(),
|
||||
warn("cache_config", NO_EQUIVALENT).env(),
|
||||
warn("charset", NO_EQUIVALENT).env(),
|
||||
warn("client_name", NO_EQUIVALENT).env(),
|
||||
warn("cluster_error_retry_attempts", OWN_CLUSTER),
|
||||
warn("command_packer", OWN_CONNECTION),
|
||||
warn("connection_class", OWN_STACK),
|
||||
warn("connection_error_retry_attempts", OWN_CLUSTER),
|
||||
warn("connection_pool", OWN_STACK),
|
||||
read("credential_provider"),
|
||||
read("db").env().in_url(FromUrl::Any),
|
||||
warn("decode_responses", NO_EQUIVALENT).env(),
|
||||
warn("driver_info", NO_EQUIVALENT),
|
||||
warn("dynamic_startup_nodes", OWN_CLUSTER),
|
||||
warn("encoder_class", OWN_CONNECTION),
|
||||
warn("encoding", NO_EQUIVALENT).env(),
|
||||
warn("encoding_errors", NO_EQUIVALENT).env(),
|
||||
warn("errors", NO_EQUIVALENT).env(),
|
||||
warn("event_dispatcher", OWN_STACK).env(),
|
||||
read("gcp_service_account")
|
||||
.env()
|
||||
.or_secrets(&["REDIS_GCP_SERVICE_ACCOUNT"]),
|
||||
read("gcp_ssl_ca_certs")
|
||||
.env()
|
||||
.or_secrets(&["REDIS_GCP_SSL_CA_CERTS"]),
|
||||
read("health_check_interval").env().in_url(FromUrl::Any),
|
||||
read("host").env(),
|
||||
warn("legacy_responses", NO_EQUIVALENT),
|
||||
warn("lib_name", NO_EQUIVALENT).env(),
|
||||
warn("lib_version", NO_EQUIVALENT).env(),
|
||||
warn("load_balancing_strategy", OWN_CLUSTER),
|
||||
warn("maint_notifications_config", MAINTENANCE),
|
||||
warn("maint_notifications_pool_handler", MAINTENANCE),
|
||||
warn("maintenance_notification_hash", MAINTENANCE),
|
||||
warn("maintenance_state", MAINTENANCE),
|
||||
warn("max_connections", "the Rust client has no connection cap").env(),
|
||||
read("namespace"),
|
||||
warn("orig_host_address", MAINTENANCE),
|
||||
warn("orig_socket_connect_timeout", MAINTENANCE),
|
||||
warn("orig_socket_timeout", MAINTENANCE),
|
||||
warn("oss_cluster_maint_notifications_handler", MAINTENANCE),
|
||||
warn("parser", OWN_CONNECTION),
|
||||
warn("parser_class", OWN_CONNECTION),
|
||||
read("password").env().in_url(FromUrl::Any),
|
||||
warn("path", OWN_CLUSTER),
|
||||
warn("policy_resolver", OWN_CLUSTER),
|
||||
read("port").env(),
|
||||
warn("protocol", NO_EQUIVALENT).env(),
|
||||
warn("read_from_replicas", OWN_CLUSTER),
|
||||
warn(
|
||||
"redis_connect_func",
|
||||
"the Python object is not callable from Rust, so static username and password are used",
|
||||
)
|
||||
.env(),
|
||||
warn("reinitialize_steps", OWN_CLUSTER),
|
||||
warn("require_full_coverage", OWN_CLUSTER),
|
||||
warn("retry", OWN_STACK),
|
||||
warn("retry_on_error", NO_EQUIVALENT).env(),
|
||||
warn("retry_on_timeout", NO_EQUIVALENT).env(),
|
||||
read("sentinel_nodes").or_secrets(&["REDIS_SENTINEL_NODES"]),
|
||||
read("sentinel_password").or_secrets(&["REDIS_SENTINEL_PASSWORD"]),
|
||||
read("service_name").or_secrets(&["REDIS_SERVICE_NAME"]),
|
||||
warn("single_connection_client", NO_EQUIVALENT).env(),
|
||||
read("socket_connect_timeout").env().in_url(FromUrl::Any),
|
||||
read("socket_keepalive").env(),
|
||||
warn("socket_keepalive_options", NO_EQUIVALENT).env(),
|
||||
warn("socket_read_size", OWN_CONNECTION),
|
||||
read("socket_timeout").env().in_url(FromUrl::Any),
|
||||
warn("socket_type", OWN_CONNECTION),
|
||||
read("ssl").env(),
|
||||
read("ssl_ca_certs").env().in_url(FromUrl::Rediss),
|
||||
read("ssl_ca_data").env().in_url(FromUrl::Rediss),
|
||||
read("ssl_ca_path").env().in_url(FromUrl::Rediss),
|
||||
read("ssl_cert_reqs").env().in_url(FromUrl::Rediss),
|
||||
read("ssl_certfile").env().in_url(FromUrl::Rediss),
|
||||
read("ssl_check_hostname").env().in_url(FromUrl::Rediss),
|
||||
warn("ssl_ciphers", RUSTLS).env(),
|
||||
reject("ssl_exclude_verify_flags", VERIFY_FLAGS, Asked::WhenSet),
|
||||
reject("ssl_include_verify_flags", VERIFY_FLAGS, Asked::WhenSet),
|
||||
read("ssl_keyfile").env().in_url(FromUrl::Rediss),
|
||||
warn_unless("ssl_min_version", RUSTLS, is_met_by_rustls).env(),
|
||||
warn("ssl_ocsp_context", OCSP).env(),
|
||||
warn("ssl_ocsp_expected_cert", OCSP).env(),
|
||||
read("ssl_password").env().in_url(FromUrl::Rediss),
|
||||
reject("ssl_validate_ocsp", REVOCATION, Asked::WhenTrue).env(),
|
||||
reject("ssl_validate_ocsp_stapled", REVOCATION, Asked::WhenTrue).env(),
|
||||
read("startup_nodes").or_secrets(&["REDIS_CLUSTER_NODES"]),
|
||||
read("unix_socket_path").env(),
|
||||
read("url").env(),
|
||||
read("username").env().in_url(FromUrl::Any),
|
||||
];
|
||||
|
||||
// rustls only speaks TLS 1.2 and 1.3, so a minimum of 1.2 or lower is already met. redis-py takes an
|
||||
// ssl.TLSVersion, which reaches Rust as its name or its value
|
||||
fn is_met_by_rustls(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Number(number) => matches!(number.as_i64(), Some(-2 | 769 | 770 | 771)),
|
||||
Value::String(text) => matches!(
|
||||
text.strip_prefix("TLSVersion.")
|
||||
.unwrap_or(text)
|
||||
.to_ascii_uppercase()
|
||||
.as_str(),
|
||||
"MINIMUM_SUPPORTED" | "TLSV1" | "TLSV1_1" | "TLSV1_2"
|
||||
),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn setting(kwarg: &str) -> Option<&'static Setting> {
|
||||
SETTINGS.iter().find(|setting| setting.kwarg == kwarg)
|
||||
}
|
||||
|
||||
/// The check a `Reject` row stands for and which values ask for it.
|
||||
pub(crate) fn refusal(kwarg: &str) -> Option<(&'static str, Asked)> {
|
||||
match setting(kwarg)?.handling {
|
||||
Handling::Reject { check, asked } => Some((check, asked)),
|
||||
Handling::Read | Handling::Warn { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// A warned setting the kwargs set, with what dropping it means. A kwarg with no row is not a Redis setting, so it
|
||||
/// is dropped without a warning.
|
||||
pub(crate) fn ignored(source: &str, kwarg: &str, value: &Value) -> Option<Ignored> {
|
||||
let consequence = setting(kwarg)?.warning(value)?;
|
||||
Some(Ignored {
|
||||
setting: source.to_owned(),
|
||||
consequence,
|
||||
})
|
||||
}
|
||||
|
||||
/// A url query parameter the resolver does not read. redis-py turns every one into a client kwarg, so even one with
|
||||
/// no row is reported.
|
||||
pub(crate) fn ignored_query(source: &str, name: &str, value: &Value) -> Option<Ignored> {
|
||||
let consequence = match setting(name) {
|
||||
None => NO_EQUIVALENT,
|
||||
Some(setting) if setting.is_read() => {
|
||||
"the Rust client does not read it from this url's query"
|
||||
}
|
||||
Some(setting) => setting.warning(value)?,
|
||||
};
|
||||
Some(Ignored {
|
||||
setting: source.to_owned(),
|
||||
consequence,
|
||||
})
|
||||
}
|
||||
95
litellm-rust/crates/redis-resolve/src/lib.rs
Normal file
95
litellm-rust/crates/redis-resolve/src/lib.rs
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use litellm_redis_types::{Auth, HealthCheck, ProviderHandle, Resolution, StaticAuth, Target};
|
||||
use serde_json::Value;
|
||||
|
||||
mod auth;
|
||||
mod endpoint;
|
||||
mod error;
|
||||
mod keys;
|
||||
mod options;
|
||||
mod tls;
|
||||
|
||||
pub use error::{Error, Expected, PythonException};
|
||||
pub use keys::{Asked, FromUrl, Handling, SETTINGS, Setting, setting};
|
||||
pub use options::is_present;
|
||||
|
||||
use options::Options;
|
||||
|
||||
/// One merged redis-py kwarg and the name a warning or error reports it under.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Entry {
|
||||
pub value: Value,
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
pub type Kwargs = BTreeMap<String, Entry>;
|
||||
|
||||
/// Kwargs whose every entry is reported under its own kwarg name.
|
||||
pub fn kwargs(fields: impl IntoIterator<Item = (String, Value)>) -> Kwargs {
|
||||
fields
|
||||
.into_iter()
|
||||
.map(|(kwarg, value)| {
|
||||
let entry = Entry {
|
||||
value,
|
||||
source: kwarg.clone(),
|
||||
};
|
||||
(kwarg, entry)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// What redis-py's constructors make of kwargs a host already merged the way `_get_redis_client_logic` does. A
|
||||
/// host passes its `credential_provider` object as `provider`, and `REDIS_USERNAME` as its process environment
|
||||
/// holds it, which Python's Azure AD auth reads from `os.environ` instead of the merged kwargs.
|
||||
pub fn interpret(
|
||||
kwargs: &Kwargs,
|
||||
provider: Option<ProviderHandle>,
|
||||
redis_username: Option<&str>,
|
||||
) -> Result<Resolution, Error> {
|
||||
let (options, ignored) = Options::read(kwargs)?;
|
||||
let auth = auth::redis_auth(&options, provider, redis_username)?;
|
||||
let (target, query) = endpoint::target(&options, &auth)?;
|
||||
let auth = match (&target, auth) {
|
||||
(Target::Url(_), Auth::Static(StaticAuth { username, .. })) => Auth::Static(StaticAuth {
|
||||
username,
|
||||
password: None,
|
||||
}),
|
||||
(_, auth) => auth,
|
||||
};
|
||||
let timeout = query.timeout.map_or_else(
|
||||
|| options::seconds(options.socket_timeout.as_ref()),
|
||||
|timeout| Ok(Some(timeout)),
|
||||
)?;
|
||||
let connect_timeout = query.connect_timeout.map_or_else(
|
||||
|| options::seconds(options.socket_connect_timeout.as_ref()),
|
||||
|timeout| Ok(Some(timeout)),
|
||||
)?;
|
||||
let health_check_interval = query.health_check_interval.map_or_else(
|
||||
|| options::integer::<i64>(options.health_check_interval.as_ref()),
|
||||
|interval| Ok(Some(interval)),
|
||||
)?;
|
||||
Ok(Resolution {
|
||||
target,
|
||||
auth,
|
||||
namespace: options::text(options.namespace.as_ref()),
|
||||
timeout,
|
||||
connect_timeout,
|
||||
health_check: health_check(health_check_interval),
|
||||
keepalive: options
|
||||
.socket_keepalive
|
||||
.as_ref()
|
||||
.map_or(Ok(true), |entry| options::flag(Some(entry)))?,
|
||||
ignored: ignored.into_iter().chain(query.ignored).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
// redis-py checks health only when the interval is truthy, so zero turns it off and a negative interval checks
|
||||
// before every command
|
||||
fn health_check(seconds: Option<i64>) -> HealthCheck {
|
||||
match seconds {
|
||||
None => HealthCheck::Unset,
|
||||
Some(0) => HealthCheck::Disabled,
|
||||
Some(seconds) => HealthCheck::Every(Duration::from_secs(seconds.try_into().unwrap_or(0))),
|
||||
}
|
||||
}
|
||||
281
litellm-rust/crates/redis-resolve/src/options.rs
Normal file
281
litellm-rust/crates/redis-resolve/src/options.rs
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::{BigInt, number};
|
||||
use litellm_redis_types::Ignored;
|
||||
use litellm_secret_value::SecretValue;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
Entry, Error, Expected, Kwargs,
|
||||
keys::{self, Asked},
|
||||
};
|
||||
|
||||
/// Every setting the resolver reads. A setting is `None` when it is unset, null or empty, except that a null
|
||||
/// `tls.cert_reqs` is kept because redis-py reads `ssl_cert_reqs=None` as `CERT_NONE`.
|
||||
pub(super) struct Options {
|
||||
pub(super) url: Option<Entry>,
|
||||
pub(super) host: Option<Entry>,
|
||||
pub(super) port: Option<Entry>,
|
||||
pub(super) db: Option<Entry>,
|
||||
pub(super) unix_socket_path: Option<Entry>,
|
||||
pub(super) startup_nodes: Option<Entry>,
|
||||
pub(super) sentinel: SentinelOptions,
|
||||
pub(super) username: Option<Entry>,
|
||||
pub(super) password: Option<Entry>,
|
||||
pub(super) credential_provider: Option<Entry>,
|
||||
pub(super) gcp: GcpOptions,
|
||||
pub(super) azure: AzureOptions,
|
||||
pub(super) aws: AwsOptions,
|
||||
pub(super) ssl: Option<Entry>,
|
||||
pub(super) tls: TlsInput,
|
||||
pub(super) namespace: Option<Entry>,
|
||||
pub(super) socket_timeout: Option<Entry>,
|
||||
pub(super) socket_connect_timeout: Option<Entry>,
|
||||
pub(super) health_check_interval: Option<Entry>,
|
||||
pub(super) socket_keepalive: Option<Entry>,
|
||||
}
|
||||
|
||||
pub(super) struct SentinelOptions {
|
||||
pub(super) nodes: Option<Entry>,
|
||||
pub(super) service_name: Option<Entry>,
|
||||
pub(super) password: Option<Entry>,
|
||||
}
|
||||
|
||||
pub(super) struct GcpOptions {
|
||||
pub(super) service_account: Option<Entry>,
|
||||
pub(super) ssl_ca_certs: Option<Entry>,
|
||||
}
|
||||
|
||||
pub(super) struct AzureOptions {
|
||||
pub(super) ad_token: Option<Entry>,
|
||||
pub(super) client_id: Option<Entry>,
|
||||
pub(super) tenant_id: Option<Entry>,
|
||||
pub(super) client_secret: Option<Entry>,
|
||||
}
|
||||
|
||||
pub(super) struct AwsOptions {
|
||||
pub(super) auth: Option<Entry>,
|
||||
pub(super) user_name: Option<Entry>,
|
||||
pub(super) cache_name: Option<Entry>,
|
||||
pub(super) region: Option<Entry>,
|
||||
pub(super) serverless: Option<Entry>,
|
||||
}
|
||||
|
||||
/// The TLS settings, read from the kwargs or from a `rediss://` url's query.
|
||||
pub(super) struct TlsInput {
|
||||
pub(super) cert_reqs: Option<Entry>,
|
||||
pub(super) ca_certs: Option<Entry>,
|
||||
pub(super) ca_data: Option<Entry>,
|
||||
pub(super) ca_path: Option<Entry>,
|
||||
pub(super) certfile: Option<Entry>,
|
||||
pub(super) keyfile: Option<Entry>,
|
||||
pub(super) key_password: Option<Entry>,
|
||||
pub(super) check_hostname: Option<Entry>,
|
||||
}
|
||||
|
||||
/// The settings the kwargs set that the Rust client drops, once none of them asks for a check it cannot perform.
|
||||
fn ignored(kwargs: &Kwargs) -> Result<Vec<Ignored>, Error> {
|
||||
refuse_unenforceable(kwargs.iter())?;
|
||||
Ok(kwargs
|
||||
.iter()
|
||||
.filter_map(|(key, entry)| keys::ignored(&entry.source, key, &entry.value))
|
||||
.collect())
|
||||
}
|
||||
|
||||
impl Options {
|
||||
/// The typed settings and the ones the Rust client drops.
|
||||
pub(super) fn read(kwargs: &Kwargs) -> Result<(Self, Vec<Ignored>), Error> {
|
||||
let ignored = ignored(kwargs)?;
|
||||
let options = Self::from_lookup(|key| kwargs.get(key).cloned());
|
||||
Ok((options, ignored))
|
||||
}
|
||||
|
||||
fn from_lookup(lookup: impl Fn(&'static str) -> Option<Entry>) -> Self {
|
||||
let get = |key| lookup(key).filter(|entry| is_present(&entry.value));
|
||||
Self {
|
||||
url: get("url"),
|
||||
host: get("host"),
|
||||
port: get("port"),
|
||||
db: get("db"),
|
||||
unix_socket_path: get("unix_socket_path"),
|
||||
startup_nodes: get("startup_nodes"),
|
||||
sentinel: SentinelOptions {
|
||||
nodes: get("sentinel_nodes"),
|
||||
service_name: get("service_name"),
|
||||
password: get("sentinel_password"),
|
||||
},
|
||||
username: get("username"),
|
||||
password: get("password"),
|
||||
credential_provider: get("credential_provider"),
|
||||
gcp: GcpOptions {
|
||||
service_account: get("gcp_service_account"),
|
||||
ssl_ca_certs: get("gcp_ssl_ca_certs"),
|
||||
},
|
||||
azure: AzureOptions {
|
||||
ad_token: get("azure_redis_ad_token"),
|
||||
client_id: get("azure_client_id"),
|
||||
tenant_id: get("azure_tenant_id"),
|
||||
client_secret: get("azure_client_secret"),
|
||||
},
|
||||
aws: AwsOptions {
|
||||
auth: get("aws_iam_auth"),
|
||||
user_name: get("aws_iam_user_name"),
|
||||
cache_name: get("aws_iam_cache_name"),
|
||||
region: get("aws_iam_region"),
|
||||
serverless: get("aws_iam_serverless"),
|
||||
},
|
||||
ssl: get("ssl"),
|
||||
tls: TlsInput {
|
||||
cert_reqs: lookup("ssl_cert_reqs")
|
||||
.filter(|entry| entry.value.is_null() || is_present(&entry.value)),
|
||||
ca_certs: get("ssl_ca_certs"),
|
||||
ca_data: get("ssl_ca_data"),
|
||||
ca_path: get("ssl_ca_path"),
|
||||
certfile: get("ssl_certfile"),
|
||||
keyfile: get("ssl_keyfile"),
|
||||
key_password: get("ssl_password"),
|
||||
check_hostname: get("ssl_check_hostname"),
|
||||
},
|
||||
namespace: get("namespace"),
|
||||
socket_timeout: get("socket_timeout"),
|
||||
socket_connect_timeout: get("socket_connect_timeout"),
|
||||
health_check_interval: get("health_check_interval"),
|
||||
socket_keepalive: get("socket_keepalive"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A `Reject` row whose value asks for its check fails, because Rust cannot perform it.
|
||||
pub(super) fn refuse_unenforceable<'a>(
|
||||
entries: impl IntoIterator<Item = (&'a String, &'a Entry)>,
|
||||
) -> Result<(), Error> {
|
||||
entries
|
||||
.into_iter()
|
||||
.filter_map(|(key, entry)| Some((keys::refusal(key)?, entry)))
|
||||
.try_for_each(|((check, asked), entry)| {
|
||||
let asks = match asked {
|
||||
Asked::WhenTrue => flag(Some(entry))?,
|
||||
Asked::WhenSet => is_present(&entry.value),
|
||||
};
|
||||
match asks {
|
||||
true => Err(Error::Unenforceable {
|
||||
setting: entry.source.clone(),
|
||||
check,
|
||||
}),
|
||||
false => Ok(()),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Python reads most of these with `x or fallback`, so an empty value counts as unset everywhere
|
||||
pub fn is_present(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Null => false,
|
||||
Value::String(text) => !text.is_empty(),
|
||||
Value::Array(items) => !items.is_empty(),
|
||||
Value::Object(entries) => !entries.is_empty(),
|
||||
Value::Bool(_) | Value::Number(_) => true,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn invalid(setting: &str, expected: Expected) -> Error {
|
||||
Error::Invalid {
|
||||
setting: setting.to_owned(),
|
||||
expected,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn as_text(value: &Value) -> Option<String> {
|
||||
match value {
|
||||
Value::String(text) => Some(text.clone()),
|
||||
Value::Number(number) => Some(number.to_string()),
|
||||
Value::Bool(flag) => Some(flag.to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn text(entry: Option<&Entry>) -> Option<String> {
|
||||
entry.and_then(|entry| as_text(&entry.value))
|
||||
}
|
||||
|
||||
pub(super) fn secret(entry: Option<&Entry>) -> Option<SecretValue> {
|
||||
text(entry).map(SecretValue::new)
|
||||
}
|
||||
|
||||
/// A flag every Python surface reads the same way. `_str_to_bool` takes only `true`, `1` and `yes`
|
||||
/// as true while the proxy's pydantic models also take `on`, `t` and `y`, so text they disagree on,
|
||||
/// or text neither accepts, is an error instead of a guess.
|
||||
pub(super) fn flag(entry: Option<&Entry>) -> Result<bool, Error> {
|
||||
let Some(entry) = entry else {
|
||||
return Ok(false);
|
||||
};
|
||||
match &entry.value {
|
||||
Value::Bool(flag) => Ok(*flag),
|
||||
Value::Number(number) => Ok(number.as_f64().is_some_and(|number| number != 0.0)),
|
||||
Value::String(text) => match text.to_ascii_lowercase().as_str() {
|
||||
"true" | "1" | "yes" => Ok(true),
|
||||
"false" | "0" | "no" | "off" | "f" | "n" => Ok(false),
|
||||
_ => Err(invalid(&entry.source, Expected::Boolean)),
|
||||
},
|
||||
_ => Err(invalid(&entry.source, Expected::Boolean)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn integer<T: TryFrom<BigInt>>(entry: Option<&Entry>) -> Result<Option<T>, Error> {
|
||||
coerced(entry, Expected::Integer, number::int, |number| {
|
||||
T::try_from(number).ok()
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn seconds(entry: Option<&Entry>) -> Result<Option<Duration>, Error> {
|
||||
coerced(entry, Expected::Seconds, number::float, |seconds| {
|
||||
Duration::try_from_secs_f64(seconds).ok()
|
||||
})
|
||||
}
|
||||
|
||||
fn coerced<N, T>(
|
||||
entry: Option<&Entry>,
|
||||
expected: Expected,
|
||||
parse: impl Fn(&str) -> Option<N>,
|
||||
convert: impl Fn(N) -> Option<T>,
|
||||
) -> Result<Option<T>, Error> {
|
||||
let Some((entry, text)) = entry.and_then(|entry| Some((entry, as_text(&entry.value)?))) else {
|
||||
return Ok(None);
|
||||
};
|
||||
match parse(&text) {
|
||||
Some(number) => convert(number)
|
||||
.map(Some)
|
||||
.ok_or_else(|| invalid(&entry.source, expected)),
|
||||
None => {
|
||||
tracing::warn!(setting = %entry.source, "Redis setting is not a number, so the default is used like redis-py");
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{cell::RefCell, collections::BTreeSet};
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::Options;
|
||||
use crate::SETTINGS;
|
||||
|
||||
// a Read row that Options never asks for would be neither modeled nor passed through nor warned about
|
||||
#[rstest]
|
||||
fn options_ask_for_exactly_the_read_rows() {
|
||||
let asked = RefCell::new(BTreeSet::new());
|
||||
let _ = Options::from_lookup(|kwarg| {
|
||||
asked.borrow_mut().insert(kwarg);
|
||||
None
|
||||
});
|
||||
|
||||
let read: BTreeSet<&str> = SETTINGS
|
||||
.iter()
|
||||
.filter(|setting| setting.is_read())
|
||||
.map(|setting| setting.kwarg)
|
||||
.collect();
|
||||
assert_eq!(asked.into_inner(), read);
|
||||
}
|
||||
}
|
||||
189
litellm-rust/crates/redis-resolve/src/tls.rs
Normal file
189
litellm-rust/crates/redis-resolve/src/tls.rs
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::number;
|
||||
use litellm_redis_types::{Ignored, TlsOptions};
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
Entry, Error, Expected,
|
||||
keys::{self, FromUrl},
|
||||
options::{self, TlsInput, flag, invalid, secret, text},
|
||||
};
|
||||
|
||||
/// `gcp_ca_certs` replaces `ssl_ca_certs` when GCP IAM is on, as `_get_redis_client_logic` does.
|
||||
pub(super) fn tls_options(
|
||||
input: &TlsInput,
|
||||
gcp_ca_certs: Option<&Entry>,
|
||||
) -> Result<TlsOptions, Error> {
|
||||
Ok(TlsOptions {
|
||||
ca_certs: text(gcp_ca_certs.or(input.ca_certs.as_ref())).map(Into::into),
|
||||
ca_data: text(input.ca_data.as_ref()),
|
||||
ca_path: text(input.ca_path.as_ref()).map(Into::into),
|
||||
certfile: text(input.certfile.as_ref()).map(Into::into),
|
||||
keyfile: text(input.keyfile.as_ref()).map(Into::into),
|
||||
key_password: secret(input.key_password.as_ref()),
|
||||
verify_certificate: verify_certificate(input.cert_reqs.as_ref())?,
|
||||
check_hostname: flag(input.check_hostname.as_ref())?,
|
||||
})
|
||||
}
|
||||
|
||||
// RedisSSLContext looks text up case-sensitively and takes an ssl.VerifyMode as is
|
||||
fn verify_certificate(cert_reqs: Option<&Entry>) -> Result<bool, Error> {
|
||||
let Some(entry) = cert_reqs else {
|
||||
return Ok(true);
|
||||
};
|
||||
match &entry.value {
|
||||
// Explicit None intentionally means CERT_NONE, matching redis-py rather than the unset default.
|
||||
Value::Null => Ok(false),
|
||||
Value::String(text) if text == "none" => Ok(false),
|
||||
Value::String(text) if text == "optional" || text == "required" => Ok(true),
|
||||
Value::Number(mode) => match mode.as_i64() {
|
||||
Some(0) => Ok(false),
|
||||
Some(1 | 2) => Ok(true),
|
||||
_ => Err(invalid(&entry.source, Expected::CertRequirement)),
|
||||
},
|
||||
_ => Err(invalid(&entry.source, Expected::CertRequirement)),
|
||||
}
|
||||
}
|
||||
|
||||
/// What redis-py's `parse_url` reads from a url's query on top of the url itself.
|
||||
pub(super) struct UrlQuery {
|
||||
pub(super) tls: Option<TlsOptions>,
|
||||
pub(super) db: Option<i64>,
|
||||
pub(super) username: Option<SecretValue>,
|
||||
pub(super) password: Option<SecretValue>,
|
||||
pub(super) settings: QuerySettings,
|
||||
}
|
||||
|
||||
/// The connection settings a url's query sets, which win over the kwargs because `from_url` applies them last.
|
||||
#[derive(Default)]
|
||||
pub(super) struct QuerySettings {
|
||||
pub(super) timeout: Option<Duration>,
|
||||
pub(super) connect_timeout: Option<Duration>,
|
||||
pub(super) health_check_interval: Option<i64>,
|
||||
pub(super) ignored: Vec<Ignored>,
|
||||
}
|
||||
|
||||
pub(super) fn url_query(url: &str) -> Result<UrlQuery, Error> {
|
||||
let Ok(url) = url::Url::parse(url) else {
|
||||
return Ok(UrlQuery {
|
||||
tls: None,
|
||||
db: None,
|
||||
username: None,
|
||||
password: None,
|
||||
settings: QuerySettings::default(),
|
||||
});
|
||||
};
|
||||
let tls_scheme = url.scheme().eq_ignore_ascii_case("rediss");
|
||||
let pairs: Vec<(String, Entry)> = url
|
||||
.query_pairs()
|
||||
.filter(|(_, value)| !value.is_empty())
|
||||
.map(|(name, value)| {
|
||||
(
|
||||
name.to_string(),
|
||||
Entry {
|
||||
value: Value::String(value.into_owned()),
|
||||
source: format!("url ?{name}"),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
// parse_qs keeps every value of a repeated parameter and parse_url reads the first
|
||||
let query: Vec<(String, Entry)> = pairs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(index, (name, _))| !pairs[..*index].iter().any(|(earlier, _)| earlier == name))
|
||||
.map(|(_, pair)| pair.clone())
|
||||
.collect();
|
||||
options::refuse_unenforceable(query.iter().map(|(name, entry)| (name, entry)))?;
|
||||
let (query, ignored): (Vec<_>, Vec<_>) = query.into_iter().partition(|(name, _)| {
|
||||
keys::setting(name).is_some_and(|setting| match setting.from_url {
|
||||
FromUrl::Any => true,
|
||||
FromUrl::Rediss => tls_scheme,
|
||||
FromUrl::No => false,
|
||||
})
|
||||
});
|
||||
let get = |key: &str| {
|
||||
query
|
||||
.iter()
|
||||
.find(|(name, _)| name == key)
|
||||
.map(|(_, entry)| entry.clone())
|
||||
};
|
||||
let seconds = |key: &str| {
|
||||
query_number(
|
||||
get(key).as_ref(),
|
||||
number::float,
|
||||
Expected::Seconds,
|
||||
|seconds| Duration::try_from_secs_f64(seconds).ok(),
|
||||
)
|
||||
};
|
||||
let integer = |key: &str| {
|
||||
query_number(
|
||||
get(key).as_ref(),
|
||||
number::int,
|
||||
Expected::Integer,
|
||||
|number| i64::try_from(number).ok(),
|
||||
)
|
||||
};
|
||||
let db = integer("db")?;
|
||||
let settings = QuerySettings {
|
||||
timeout: seconds("socket_timeout")?,
|
||||
connect_timeout: seconds("socket_connect_timeout")?,
|
||||
health_check_interval: integer("health_check_interval")?,
|
||||
ignored: ignored
|
||||
.iter()
|
||||
.filter_map(|(name, entry)| keys::ignored_query(&entry.source, name, &entry.value))
|
||||
.collect(),
|
||||
};
|
||||
let tls = match tls_scheme {
|
||||
true => Some(TlsOptions {
|
||||
ca_certs: text(get("ssl_ca_certs").as_ref()).map(Into::into),
|
||||
ca_data: text(get("ssl_ca_data").as_ref()),
|
||||
ca_path: text(get("ssl_ca_path").as_ref()).map(Into::into),
|
||||
certfile: text(get("ssl_certfile").as_ref()).map(Into::into),
|
||||
keyfile: text(get("ssl_keyfile").as_ref()).map(Into::into),
|
||||
key_password: secret(get("ssl_password").as_ref()),
|
||||
verify_certificate: verify_certificate(get("ssl_cert_reqs").as_ref())?,
|
||||
check_hostname: text(get("ssl_check_hostname").as_ref())
|
||||
.is_some_and(|text| url_bool(&text)),
|
||||
}),
|
||||
false => None,
|
||||
};
|
||||
Ok(UrlQuery {
|
||||
tls,
|
||||
db,
|
||||
username: secret(get("username").as_ref()),
|
||||
password: secret(get("password").as_ref()),
|
||||
settings,
|
||||
})
|
||||
}
|
||||
|
||||
// parse_url reads these with int() or float() and raises ValueError building the client on anything else
|
||||
fn query_number<N, T>(
|
||||
entry: Option<&Entry>,
|
||||
parse: impl Fn(&str) -> Option<N>,
|
||||
expected: Expected,
|
||||
convert: impl Fn(N) -> Option<T>,
|
||||
) -> Result<Option<T>, Error> {
|
||||
let Some(entry) = entry else {
|
||||
return Ok(None);
|
||||
};
|
||||
let number = text(Some(entry))
|
||||
.and_then(|text| parse(&text))
|
||||
.ok_or_else(|| Error::UrlQuery {
|
||||
setting: entry.source.clone(),
|
||||
})?;
|
||||
convert(number)
|
||||
.map(Some)
|
||||
.ok_or_else(|| invalid(&entry.source, expected))
|
||||
}
|
||||
|
||||
// redis-py's to_bool: only 0, F, FALSE, N and NO (any case) are false
|
||||
fn url_bool(text: &str) -> bool {
|
||||
!matches!(
|
||||
text.to_ascii_uppercase().as_str(),
|
||||
"" | "0" | "F" | "FALSE" | "N" | "NO"
|
||||
)
|
||||
}
|
||||
129
litellm-rust/crates/redis-resolve/tests/catalog.rs
Normal file
129
litellm-rust/crates/redis-resolve/tests/catalog.rs
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use litellm_redis_resolve::{Asked, Error, Handling, SETTINGS, Setting, interpret, kwargs};
|
||||
use rstest::rstest;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Catalog {
|
||||
meta: Meta,
|
||||
args: BTreeMap<String, Arg>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Meta {
|
||||
redis_py_locked: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Arg {
|
||||
environment: bool,
|
||||
redis_py: BTreeMap<String, Vec<String>>,
|
||||
}
|
||||
|
||||
// the signatures of the clients the proxy builds, where an unknown kwarg raises TypeError
|
||||
const PROXY_SIGNATURES: &[&str] = &["async", "cluster", "connection"];
|
||||
|
||||
fn rejected() -> impl Iterator<Item = (&'static Setting, Asked)> {
|
||||
SETTINGS
|
||||
.iter()
|
||||
.filter_map(|setting| match setting.handling {
|
||||
Handling::Reject { asked, .. } => Some((setting, asked)),
|
||||
Handling::Read | Handling::Warn { .. } => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn catalog() -> Catalog {
|
||||
serde_json::from_str(include_str!("../redis-py-kwargs.json")).unwrap()
|
||||
}
|
||||
|
||||
fn host_with(kwarg: &str, value: serde_json::Value) -> litellm_redis_resolve::Kwargs {
|
||||
kwargs([
|
||||
("host".to_owned(), json!("redis")),
|
||||
(kwarg.to_owned(), value),
|
||||
])
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn every_catalogued_kwarg_has_exactly_one_settings_row() {
|
||||
let rows: Vec<&str> = SETTINGS.iter().map(|setting| setting.kwarg).collect();
|
||||
let unique: BTreeSet<&str> = rows.iter().copied().collect();
|
||||
assert_eq!(rows.len(), unique.len(), "a kwarg has two SETTINGS rows");
|
||||
|
||||
let catalog = catalog();
|
||||
let catalogued: BTreeSet<&str> = catalog.args.keys().map(String::as_str).collect();
|
||||
assert_eq!(unique, catalogued);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn settings_read_redis_environment_names_exactly_where_python_does() {
|
||||
let catalog = catalog();
|
||||
|
||||
let differing: Vec<&str> = SETTINGS
|
||||
.iter()
|
||||
.filter(|setting| {
|
||||
catalog
|
||||
.args
|
||||
.get(setting.kwarg)
|
||||
.is_some_and(|arg| arg.environment != setting.environment)
|
||||
})
|
||||
.map(|setting| setting.kwarg)
|
||||
.collect();
|
||||
assert!(differing.is_empty(), "{differing:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn every_warned_setting_is_reported() {
|
||||
let warned = SETTINGS
|
||||
.iter()
|
||||
.filter(|setting| matches!(setting.handling, Handling::Warn { met: None, .. }));
|
||||
|
||||
for setting in warned {
|
||||
let resolved = interpret(&host_with(setting.kwarg, json!("set")), None, None).unwrap();
|
||||
|
||||
let reported: Vec<&str> = resolved
|
||||
.ignored
|
||||
.iter()
|
||||
.map(|ignored| ignored.setting.as_str())
|
||||
.collect();
|
||||
assert_eq!(reported, vec![setting.kwarg]);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn every_rejected_setting_fails_when_asked_for_and_not_otherwise() {
|
||||
for (setting, asked) in rejected() {
|
||||
let (asking, not_asking) = match asked {
|
||||
Asked::WhenTrue => (json!(true), json!(false)),
|
||||
Asked::WhenSet => (json!(["VERIFY_CRL_CHECK_LEAF"]), json!([])),
|
||||
};
|
||||
|
||||
let error = interpret(&host_with(setting.kwarg, asking), None, None).unwrap_err();
|
||||
assert!(
|
||||
matches!(&error, Error::Unenforceable { setting: source, .. } if source == setting.kwarg),
|
||||
"{error:?}"
|
||||
);
|
||||
assert!(interpret(&host_with(setting.kwarg, not_asking), None, None).is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn python_already_refuses_every_rejected_setting() {
|
||||
let catalog = catalog();
|
||||
|
||||
let accepted: Vec<&str> = rejected()
|
||||
.map(|(setting, _)| setting.kwarg)
|
||||
.filter(|kwarg| {
|
||||
catalog.args[*kwarg]
|
||||
.redis_py
|
||||
.get(&catalog.meta.redis_py_locked)
|
||||
.is_some_and(|signatures| {
|
||||
signatures
|
||||
.iter()
|
||||
.any(|signature| PROXY_SIGNATURES.contains(&signature.as_str()))
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
assert!(accepted.is_empty(), "{accepted:?}");
|
||||
}
|
||||
790
litellm-rust/crates/redis-resolve/tests/interpret.rs
Normal file
790
litellm-rust/crates/redis-resolve/tests/interpret.rs
Normal file
|
|
@ -0,0 +1,790 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_redis_resolve::{Entry, Error, Expected, Kwargs, PythonException, interpret, kwargs};
|
||||
use litellm_redis_types::{
|
||||
Auth, HealthCheck, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions,
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use Outcome::{Fails, Refuses, Resolves};
|
||||
use PythonException::{RedisClusterException, TypeError, ValueError};
|
||||
|
||||
const PROVIDER: ProviderHandle = ProviderHandle::new(7);
|
||||
|
||||
fn merged(fields: serde_json::Value, sources: &[(&str, &str)]) -> Kwargs {
|
||||
let fields = serde_json::from_value::<BTreeMap<String, serde_json::Value>>(fields).unwrap();
|
||||
kwargs(fields)
|
||||
.into_iter()
|
||||
.map(|(kwarg, entry)| {
|
||||
let source = sources
|
||||
.iter()
|
||||
.find(|(named, _)| *named == kwarg)
|
||||
.map_or(entry.source, |(_, source)| (*source).to_owned());
|
||||
(kwarg, Entry { source, ..entry })
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// a host hands its credential_provider object over as a handle, so the cases do the same
|
||||
fn interpret_merged(
|
||||
fields: serde_json::Value,
|
||||
sources: &[(&str, &str)],
|
||||
redis_username: Option<&str>,
|
||||
) -> Result<Resolution, Error> {
|
||||
let has_provider = fields
|
||||
.get("credential_provider")
|
||||
.is_some_and(|provider| !provider.is_null());
|
||||
interpret(
|
||||
&merged(fields, sources),
|
||||
has_provider.then_some(PROVIDER),
|
||||
redis_username,
|
||||
)
|
||||
}
|
||||
|
||||
fn interpret_kwargs(fields: serde_json::Value) -> Result<Resolution, Error> {
|
||||
interpret_merged(fields, &[], None)
|
||||
}
|
||||
|
||||
/// What a contract case expects of the resolver, next to what its Python mirror expects of redis-py.
|
||||
enum Outcome {
|
||||
Resolves(serde_json::Value),
|
||||
/// Rust refuses a config Python connects with, a disagreement on purpose.
|
||||
Refuses(&'static str),
|
||||
/// redis-py 5.3.1 already raises on the config, building the client or its first connection, so failing
|
||||
/// here breaks no deployment that works today. The class is what it raises building the client, and `None`
|
||||
/// when it only raises once a command connects.
|
||||
Fails(&'static str, Option<PythonException>),
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_provider_handle_beats_every_cloud_auth_and_static_credentials() {
|
||||
let resolved = interpret_kwargs(json!({
|
||||
"host": "redis", "password": "pass", "credential_provider": "host-object",
|
||||
"gcp_service_account": "account", "aws_iam_auth": true, "azure_redis_ad_token": "true",
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resolved.auth, Auth::Provider(PROVIDER));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_rediss_url_reads_ssl_password_from_its_query() {
|
||||
let resolved =
|
||||
interpret_kwargs(json!({"url": "rediss://cache:6380/0?ssl_password=query-key-secret"}))
|
||||
.unwrap();
|
||||
|
||||
let Target::Url(parsed) = resolved.target else {
|
||||
panic!("expected a url target, got {:?}", resolved.target);
|
||||
};
|
||||
let key_password = parsed.tls.and_then(|tls| tls.key_password);
|
||||
assert_eq!(expose(key_password.as_ref()), Some("query-key-secret"));
|
||||
assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::read_on_a_plain_url(
|
||||
"redis://cache:6379/1?db=3&username=query-user&password=query-secret",
|
||||
Some(3),
|
||||
Some("query-user"),
|
||||
Some("query-secret")
|
||||
)]
|
||||
#[case::absent("rediss://cache:6380/1", None, None, None)]
|
||||
fn a_url_query_carries_db_and_credentials_like_parse_url(
|
||||
#[case] url: &str,
|
||||
#[case] db: Option<i64>,
|
||||
#[case] username: Option<&str>,
|
||||
#[case] password: Option<&str>,
|
||||
) {
|
||||
let resolved = interpret_kwargs(json!({"url": url})).unwrap();
|
||||
|
||||
let Target::Url(parsed) = resolved.target else {
|
||||
panic!("expected a url target, got {:?}", resolved.target);
|
||||
};
|
||||
assert_eq!(parsed.db, db);
|
||||
assert_eq!(expose(parsed.username.as_ref()), username);
|
||||
assert_eq!(expose(parsed.password.as_ref()), password);
|
||||
assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_url_query_db_that_is_not_an_integer_fails_like_parse_url() {
|
||||
let error = interpret_kwargs(json!({"url": "redis://cache:6379?db=first"})).unwrap_err();
|
||||
|
||||
assert!(matches!(error, Error::UrlQuery { setting } if setting == "url ?db"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::tls_1_2_is_already_met(json!("TLSv1_2"), false)]
|
||||
#[case::the_enum_name(json!("TLSVersion.TLSv1_2"), false)]
|
||||
#[case::the_enum_value(json!(771), false)]
|
||||
#[case::tls_1_3_is_not_enforced(json!("TLSv1_3"), true)]
|
||||
fn ssl_min_version_warns_only_when_rustls_cannot_meet_it(
|
||||
#[case] version: serde_json::Value,
|
||||
#[case] warns: bool,
|
||||
) {
|
||||
let resolved =
|
||||
interpret_kwargs(json!({"host": "redis", "ssl": true, "ssl_min_version": version}))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
!resolved.ignored.is_empty(),
|
||||
warns,
|
||||
"{:?}",
|
||||
resolved.ignored
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unset(json!({"host": "redis"}), true)]
|
||||
#[case::empty_like_unset(json!({"host": "redis", "socket_keepalive": ""}), true)]
|
||||
#[case::enabled(json!({"host": "redis", "socket_keepalive": "yes"}), true)]
|
||||
#[case::disabled(json!({"host": "redis", "socket_keepalive": false}), false)]
|
||||
#[case::disabled_as_text(json!({"host": "redis", "socket_keepalive": "0"}), false)]
|
||||
#[case::disabled_on_a_cluster(json!({"startup_nodes": [{"host": "node-a", "port": 7000}], "socket_keepalive": "false"}), false)]
|
||||
fn keepalive_stays_on_unless_socket_keepalive_turns_it_off(
|
||||
#[case] kwargs: serde_json::Value,
|
||||
#[case] keepalive: bool,
|
||||
) {
|
||||
let resolved = interpret_kwargs(kwargs).unwrap();
|
||||
|
||||
assert_eq!(resolved.keepalive, keepalive);
|
||||
assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_socket_keepalive_no_python_surface_agrees_on_is_invalid() {
|
||||
let error = interpret_kwargs(json!({"host": "redis", "socket_keepalive": "on"})).unwrap_err();
|
||||
|
||||
assert!(
|
||||
matches!(&error, Error::Invalid { setting, expected: Expected::Boolean } if setting == "socket_keepalive"),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
fn tls_view(tls: Option<&TlsOptions>) -> serde_json::Value {
|
||||
tls.map_or(serde_json::Value::Null, |tls| {
|
||||
serde_json::json!({
|
||||
"ca_certs": tls.ca_certs,
|
||||
"ca_data": tls.ca_data,
|
||||
"ca_path": tls.ca_path,
|
||||
"certfile": tls.certfile,
|
||||
"keyfile": tls.keyfile,
|
||||
"verify_certificate": tls.verify_certificate,
|
||||
"check_hostname": tls.check_hostname,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn nodes_view(nodes: &[Node]) -> Vec<String> {
|
||||
nodes
|
||||
.iter()
|
||||
.map(|node| format!("{}:{}", node.host, node.port))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> {
|
||||
value.map(litellm_secret_value::SecretValue::expose)
|
||||
}
|
||||
|
||||
fn view(resolved: &Resolution) -> serde_json::Value {
|
||||
let target = match &resolved.target {
|
||||
Target::Unix(socket) => serde_json::json!({
|
||||
"mode": "unix", "path": socket.path, "tls": null,
|
||||
}),
|
||||
Target::Standalone(host) => serde_json::json!({
|
||||
"mode": "host",
|
||||
"host": host.host,
|
||||
"port": host.port,
|
||||
"db": host.db,
|
||||
"tls": tls_view(host.tls.as_ref()),
|
||||
}),
|
||||
Target::Url(url) => serde_json::json!({
|
||||
"mode": "url",
|
||||
"url": url.url.expose(),
|
||||
"tls": tls_view(url.tls.as_ref()),
|
||||
}),
|
||||
Target::Cluster(cluster) => serde_json::json!({
|
||||
"mode": "cluster",
|
||||
"nodes": nodes_view(&cluster.nodes),
|
||||
"tls": tls_view(cluster.tls.as_ref()),
|
||||
}),
|
||||
Target::Sentinel(sentinel) => serde_json::json!({
|
||||
"mode": "sentinel",
|
||||
"nodes": nodes_view(&sentinel.monitors),
|
||||
"service_name": sentinel.service_name,
|
||||
"sentinel_password": expose(sentinel.monitor_auth.as_ref().and_then(|auth| auth.password.as_ref())),
|
||||
"db": sentinel.db,
|
||||
"tls": tls_view(sentinel.tls.as_ref()),
|
||||
}),
|
||||
};
|
||||
let auth = match &resolved.auth {
|
||||
Auth::Static(StaticAuth { username, password }) => serde_json::json!({
|
||||
"kind": "static",
|
||||
"username": expose(username.as_ref()),
|
||||
"password": expose(password.as_ref()),
|
||||
}),
|
||||
Auth::Provider(_) => serde_json::json!({"kind": "host"}),
|
||||
Auth::GcpIam { service_account } => {
|
||||
serde_json::json!({"kind": "gcp_iam", "service_account": service_account})
|
||||
}
|
||||
Auth::AzureAd(azure) => {
|
||||
serde_json::json!({"kind": "azure_ad", "username": azure.username})
|
||||
}
|
||||
Auth::AwsIam(aws) => serde_json::json!({
|
||||
"kind": "aws_iam",
|
||||
"user_name": aws.user_name,
|
||||
"cache_name": aws.cache_name,
|
||||
"region": aws.region,
|
||||
"serverless": aws.serverless,
|
||||
}),
|
||||
};
|
||||
let serde_json::Value::Object(target) = target else {
|
||||
unreachable!("every target view is an object")
|
||||
};
|
||||
let rest = serde_json::json!({
|
||||
"auth": auth,
|
||||
"socket_timeout": resolved.timeout.map_or(5.0, |timeout| timeout.as_secs_f64()),
|
||||
"socket_connect_timeout": resolved.connect_timeout.or(resolved.timeout).map_or(5.0, |timeout| timeout.as_secs_f64()),
|
||||
"health_check_interval": match (&resolved.target, resolved.health_check) {
|
||||
(_, HealthCheck::Every(interval)) => Some(interval.as_secs_f64()),
|
||||
(Target::Cluster(_), HealthCheck::Unset) => Some(25.0),
|
||||
(_, HealthCheck::Unset | HealthCheck::Disabled) => None,
|
||||
},
|
||||
});
|
||||
let serde_json::Value::Object(rest) = rest else {
|
||||
unreachable!("the view is an object")
|
||||
};
|
||||
serde_json::Value::Object(target.into_iter().chain(rest).collect())
|
||||
}
|
||||
|
||||
fn expected_view(fields: serde_json::Value) -> serde_json::Value {
|
||||
let serde_json::Value::Object(fields) = fields else {
|
||||
unreachable!("expected views are objects")
|
||||
};
|
||||
let defaults = json!({
|
||||
"auth": {"kind": "static", "username": null, "password": null},
|
||||
"tls": null,
|
||||
"socket_timeout": 5.0,
|
||||
"socket_connect_timeout": 5.0,
|
||||
"health_check_interval": null,
|
||||
});
|
||||
let serde_json::Value::Object(defaults) = defaults else {
|
||||
unreachable!("the defaults are an object")
|
||||
};
|
||||
serde_json::Value::Object(defaults.into_iter().chain(fields).collect())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::host_defaults(
|
||||
json!({"host": "redis"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::empty_cluster_nodes_use_the_environment(
|
||||
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}),
|
||||
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::empty_sentinel_nodes_use_the_environment(
|
||||
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}),
|
||||
&[("sentinel_nodes", "REDIS_SENTINEL_NODES")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))),
|
||||
)]
|
||||
#[case::empty_service_name_uses_the_environment(
|
||||
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
|
||||
&[("service_name", "REDIS_SERVICE_NAME")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))),
|
||||
)]
|
||||
#[case::plaintext_ignores_tls_requirements(
|
||||
json!({"host": "redis", "ssl_cert_reqs": "strict"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::an_ambiguous_flag_is_refused(
|
||||
json!({"host": "redis", "ssl": "on"}),
|
||||
&[],
|
||||
None,
|
||||
Refuses("invalid Redis setting ssl: expected a boolean"),
|
||||
)]
|
||||
#[case::empty_sentinel_nodes_without_the_environment_name_no_target(
|
||||
json!({"sentinel_nodes": [], "service_name": "primary"}),
|
||||
&[],
|
||||
None,
|
||||
Fails("either host or url must be specified for redis", Some(ValueError)),
|
||||
)]
|
||||
#[case::ssl_ca_path_is_unsupported(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_path": "/ca"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/ca", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::ssl_validate_ocsp_is_unsupported(
|
||||
json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}),
|
||||
&[],
|
||||
None,
|
||||
Fails("ssl_validate_ocsp asks for OCSP revocation checks, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)),
|
||||
)]
|
||||
#[case::ssl_include_verify_flags_is_unsupported(
|
||||
json!({"host": "redis", "ssl": true, "ssl_include_verify_flags": ["VERIFY_CRL_CHECK_LEAF"]}),
|
||||
&[],
|
||||
None,
|
||||
Fails("ssl_include_verify_flags asks for X.509 verify flags such as CRL checking, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)),
|
||||
)]
|
||||
#[case::ssl_password_is_unsupported(
|
||||
json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_password": "key-secret"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::host_with_every_basic_setting(
|
||||
json!({"db": 2, "host": "redis", "password": "pass", "port": 6380, "socket_timeout": 1.5, "ssl": true, "username": "user"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 2, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": "pass"}, "socket_timeout": 1.5, "socket_connect_timeout": 1.5}))),
|
||||
)]
|
||||
#[case::environment_host_port_and_password(
|
||||
json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}),
|
||||
&[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "environment", "port": 6390, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": null, "password": "environment-secret"}}))),
|
||||
)]
|
||||
#[case::string_numbers_are_coerced(
|
||||
json!({"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 3, "socket_timeout": 2.5, "socket_connect_timeout": 2.5}))),
|
||||
)]
|
||||
#[case::padded_and_underscored_numbers_parse_like_python_int_and_float(
|
||||
json!({"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"db": 2, "host": "redis", "mode": "host", "port": 6380, "socket_timeout": 10.5, "socket_connect_timeout": 10.5}))),
|
||||
)]
|
||||
#[case::padded_environment_numbers_parse_like_python_int(
|
||||
json!({"port": "6_380", "host": "redis", "db": "3\n"}),
|
||||
&[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"db": 3, "host": "redis", "mode": "host", "port": 6380}))),
|
||||
)]
|
||||
#[case::socket_timeout_from_the_environment(
|
||||
json!({"socket_timeout": "0.75", "host": "redis"}),
|
||||
&[("socket_timeout", "REDIS_SOCKET_TIMEOUT")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 0.75, "socket_connect_timeout": 0.75}))),
|
||||
)]
|
||||
#[case::environment_url(
|
||||
json!({"url": "redis://environment:6379/1"}),
|
||||
&[("url", "REDIS_URL")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "redis://environment:6379/1"}))),
|
||||
)]
|
||||
#[case::environment_url_beats_the_environment_host(
|
||||
json!({"host": "ignored", "url": "redis://from-url:6379"}),
|
||||
&[("host", "REDIS_HOST"), ("url", "REDIS_URL")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "redis://from-url:6379"}))),
|
||||
)]
|
||||
#[case::an_explicit_host_drops_the_environment_url(
|
||||
json!({"host": "configured"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "configured", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::a_url_drops_host_db_and_password_and_keeps_the_username(
|
||||
json!({"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": null}}))),
|
||||
)]
|
||||
#[case::a_url_ignores_tls_settings_outside_the_url(
|
||||
json!({"ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true, "url": "rediss://cache:6380"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::startup_nodes_beat_a_url(
|
||||
json!({"password": "cluster-secret", "startup_nodes": [{"host": "node-a", "port": 7000}], "url": "redis://single:6379"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "auth": {"kind": "static", "username": null, "password": "cluster-secret"}, "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::environment_cluster_nodes(
|
||||
json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}),
|
||||
&[("startup_nodes", "REDIS_CLUSTER_NODES")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::cluster_tls_options(
|
||||
json!({"ssl": "true", "ssl_ca_certs": "/ca.pem", "ssl_check_hostname": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": true}, "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::sentinel_from_settings(
|
||||
json!({"password": "data-secret", "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], "sentinel_password": "sentinel-secret", "service_name": "primary"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0, "auth": {"kind": "static", "username": null, "password": "data-secret"}}))),
|
||||
)]
|
||||
#[case::sentinel_from_the_environment(
|
||||
json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}),
|
||||
&[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0}))),
|
||||
)]
|
||||
#[case::a_host_breaks_sentinel(
|
||||
json!({"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}),
|
||||
&[],
|
||||
None,
|
||||
Fails("host cannot be combined with sentinel_nodes, which redis-py rejects at startup", Some(TypeError)),
|
||||
)]
|
||||
#[case::a_url_beats_sentinels(
|
||||
json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary", "url": "redis://single:6379"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "redis://single:6379"}))),
|
||||
)]
|
||||
#[case::an_explicit_credential_provider_beats_cloud_auth(
|
||||
json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "host"}}))),
|
||||
)]
|
||||
#[case::gcp_beats_azure_and_aws(
|
||||
json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}),
|
||||
&[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}}))),
|
||||
)]
|
||||
#[case::gcp_ca_certs_replace_the_configured_ones_under_ssl(
|
||||
json!({"gcp_service_account": "account", "gcp_ssl_ca_certs": "/gcp-ca.pem", "host": "redis", "ssl": true, "ssl_ca_certs": "/configured-ca.pem"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/gcp-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "gcp_iam", "service_account": "account"}}))),
|
||||
)]
|
||||
#[case::azure_beats_aws_and_takes_the_environment_username(
|
||||
json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}),
|
||||
&[],
|
||||
Some("environment-user"),
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": "environment-user"}}))),
|
||||
)]
|
||||
#[case::an_empty_environment_username_sends_only_the_token(
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}),
|
||||
&[("username", "REDIS_USERNAME")],
|
||||
Some(""),
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))),
|
||||
)]
|
||||
#[case::azure_without_a_username_sends_only_the_token(
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))),
|
||||
)]
|
||||
#[case::azure_needs_the_literal_true(
|
||||
json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": "1", "azure_tenant_id": "tenant", "host": "redis"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::aws_iam_over_tls_with_the_default_region(
|
||||
json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}),
|
||||
&[("aws_iam_region", "AWS_DEFAULT_REGION")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "us-east-1", "serverless": false}}))),
|
||||
)]
|
||||
#[case::aws_region_beats_the_default_region(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}),
|
||||
&[("aws_iam_region", "AWS_REGION")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-central-1", "serverless": false}}))),
|
||||
)]
|
||||
#[case::aws_iam_serverless_over_a_rediss_url_drops_its_userinfo(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "eu-west-1", "aws_iam_serverless": true, "aws_iam_user_name": "iam-user", "url": "rediss://user:pass@cache-1:6379"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache-1:6379", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-west-1", "serverless": true}}))),
|
||||
)]
|
||||
#[case::aws_iam_without_tls_fails(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "us-east-1", "aws_iam_user_name": "iam-user", "host": "cache-1"}),
|
||||
&[],
|
||||
None,
|
||||
Fails("AWS ElastiCache IAM Redis authentication requires TLS", Some(ValueError)),
|
||||
)]
|
||||
#[case::aws_iam_without_a_region_fails(
|
||||
json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}),
|
||||
&[],
|
||||
None,
|
||||
Fails("AWS ElastiCache IAM Redis authentication requires: aws_iam_region", Some(ValueError)),
|
||||
)]
|
||||
#[case::tls_options_from_settings(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))),
|
||||
)]
|
||||
#[case::a_null_cert_requirement_skips_verification(
|
||||
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::a_verify_mode_of_none_skips_verification(
|
||||
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 0}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::a_verify_mode_of_optional_verifies(
|
||||
json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 1}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::a_cert_requirement_spelled_like_python_none_fails(
|
||||
json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}),
|
||||
&[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
|
||||
None,
|
||||
Fails("invalid Redis setting REDIS_SSL_CERT_REQS: expected none, optional, required, null or an ssl.VerifyMode value", None),
|
||||
)]
|
||||
#[case::tls_options_from_the_environment(
|
||||
json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}),
|
||||
&[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/env-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::connect_timeout_is_separate_from_the_socket_timeout(
|
||||
json!({"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 2.5, "socket_connect_timeout": 0.25}))),
|
||||
)]
|
||||
#[case::environment_references_resolve(
|
||||
json!({"host": "referenced", "password": "referenced-secret"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "referenced", "port": 6379, "db": 0, "auth": {"kind": "static", "username": null, "password": "referenced-secret"}}))),
|
||||
)]
|
||||
#[case::no_target_fails(
|
||||
json!({"port": 6379}),
|
||||
&[],
|
||||
None,
|
||||
Fails("either host or url must be specified for redis", Some(ValueError)),
|
||||
)]
|
||||
#[case::client_certificates(
|
||||
json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": "/client.key", "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::a_rediss_url_takes_tls_options_from_its_query(
|
||||
json!({"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))),
|
||||
)]
|
||||
#[case::a_url_query_carries_socket_timeout(
|
||||
json!({"url": "redis://cache:6379/0?socket_timeout=2"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))),
|
||||
)]
|
||||
#[case::url_query_settings_win_over_the_kwargs(
|
||||
json!({"health_check_interval": 4, "socket_connect_timeout": 7, "socket_timeout": 1, "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", "socket_timeout": 3.0, "socket_connect_timeout": 0.5, "health_check_interval": 9.0}))),
|
||||
)]
|
||||
#[case::a_repeated_url_query_parameter_takes_its_first_value(
|
||||
json!({"url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))),
|
||||
)]
|
||||
#[case::a_url_query_number_parse_url_cannot_read(
|
||||
json!({"url": "redis://cache:6379/0?socket_timeout=slow"}),
|
||||
&[],
|
||||
None,
|
||||
Fails("invalid Redis setting url ?socket_timeout: redis-py's parse_url cannot read it as a number", Some(ValueError)),
|
||||
)]
|
||||
#[case::tls_settings_outside_a_url_never_fail_it(
|
||||
json!({"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::db_is_not_read_for_a_cluster(
|
||||
json!({"db": 99_999_999_999_i64, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Fails("db must be 0 in cluster mode", Some(RedisClusterException)),
|
||||
)]
|
||||
#[case::cluster_db_true(
|
||||
json!({"db": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Fails("db must be 0 in cluster mode", Some(RedisClusterException)),
|
||||
)]
|
||||
#[case::cluster_db_false(
|
||||
json!({"db": false, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::cluster_db_fraction(
|
||||
json!({"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Fails("db must be 0 in cluster mode", Some(RedisClusterException)),
|
||||
)]
|
||||
#[case::cluster_db_float_zero(
|
||||
json!({"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::cluster_db_nonempty_list(
|
||||
json!({"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Fails("db must be 0 in cluster mode", Some(RedisClusterException)),
|
||||
)]
|
||||
#[case::cluster_db_empty_list(
|
||||
json!({"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))),
|
||||
)]
|
||||
#[case::ssl_ciphers(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ciphers": "HIGH"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::retry_on_timeout(
|
||||
json!({"host": "redis", "retry_on_timeout": true}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::malformed_port(
|
||||
json!({"host": "redis", "port": "abc"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::malformed_db(
|
||||
json!({"db": "first", "host": "redis"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::socket_timeout_not_a_number(
|
||||
json!({"host": "redis", "socket_timeout": "slow"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::malformed_port_from_the_environment(
|
||||
json!({"port": "abc", "host": "redis"}),
|
||||
&[("host", "REDIS_HOST"), ("port", "REDIS_PORT")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::protocol_3(
|
||||
json!({"host": "redis", "protocol": 3}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::retry_on_error_from_the_environment(
|
||||
json!({"retry_on_error": true, "host": "redis"}),
|
||||
&[("retry_on_error", "REDIS_RETRY_ON_ERROR")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::unix_socket_path(
|
||||
json!({"host": "redis", "unix_socket_path": "/tmp/redis.sock"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))),
|
||||
)]
|
||||
#[case::unix_socket_path_keeps_tls_settings_out_of_the_target(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))),
|
||||
)]
|
||||
#[case::ssl_ca_data(
|
||||
json!({"host": "redis", "ssl": true, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::ssl_ca_path_from_the_environment(
|
||||
json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}),
|
||||
&[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/etc/ssl/private-cas", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))),
|
||||
)]
|
||||
#[case::health_check_interval_is_whole_seconds(
|
||||
json!({"health_check_interval": "7", "host": "redis"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "health_check_interval": 7.0}))),
|
||||
)]
|
||||
#[case::health_check_interval_zero_is_off(
|
||||
json!({"health_check_interval": 0, "host": "redis"}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))),
|
||||
)]
|
||||
#[case::health_check_interval_zero_keeps_a_cluster_off(
|
||||
json!({"health_check_interval": 0, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))),
|
||||
)]
|
||||
#[case::a_cluster_rejects_a_nonzero_db(
|
||||
json!({"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}),
|
||||
&[],
|
||||
None,
|
||||
Fails("db must be 0 in cluster mode", Some(RedisClusterException)),
|
||||
)]
|
||||
#[case::unix_socket_path_without_a_host_fails_like_python(
|
||||
json!({"unix_socket_path": "/tmp/redis.sock"}),
|
||||
&[],
|
||||
None,
|
||||
Fails("either host or url must be specified for redis", Some(ValueError)),
|
||||
)]
|
||||
fn resolves_like_python_redis(
|
||||
#[case] merged: serde_json::Value,
|
||||
#[case] sources: &[(&str, &str)],
|
||||
#[case] redis_username: Option<&str>,
|
||||
#[case] expected: Outcome,
|
||||
) {
|
||||
let outcome = interpret_merged(merged, sources, redis_username);
|
||||
let rendered = outcome
|
||||
.as_ref()
|
||||
.map(view)
|
||||
.map_err(|error| error.to_string());
|
||||
|
||||
match expected {
|
||||
Resolves(expected) => assert_eq!(rendered, Ok(expected)),
|
||||
Refuses(message) => assert_eq!(rendered, Err(message.to_owned())),
|
||||
Fails(message, raised) => {
|
||||
assert_eq!(rendered, Err(message.to_owned()));
|
||||
assert_eq!(
|
||||
outcome
|
||||
.err()
|
||||
.and_then(|error| error.raised_building_the_python_client()),
|
||||
raised
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
13
litellm-rust/crates/redis-store/AGENTS.md
Normal file
13
litellm-rust/crates/redis-store/AGENTS.md
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (Python-facing values with their `FromValue` / `IntoValue`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
|
||||
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on
|
||||
- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s
|
||||
- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec
|
||||
- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc<RemoteSlot>`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`
|
||||
- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect
|
||||
- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not
|
||||
- Key suffixes and wire spellings derive from strum on the family's enums, and a type whose Python side is plain data implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`, so the bridge never names a field or a spelling
|
||||
- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together
|
||||
- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter
|
||||
- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here
|
||||
- Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect
|
||||
- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages
|
||||
25
litellm-rust/crates/redis-store/Cargo.toml
Normal file
25
litellm-rust/crates/redis-store/Cargo.toml
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
[package]
|
||||
name = "litellm-redis-store"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
futures-util.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-python-compat.workspace = true
|
||||
litellm-redis-client.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] }
|
||||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-circuit-breaker.workspace = true
|
||||
litellm-redis-testing.workspace = true
|
||||
litellm-secret-value.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue