You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
266 lines
11 KiB
266 lines
11 KiB
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.exceptions import ValidationError
|
|
from app.core.ids import new_id
|
|
from app.models.entities import Backend, PlacementPolicy
|
|
from app.repositories.backend_repository import BackendRepository
|
|
from app.repositories.placement_policy_repository import PlacementPolicyRepository
|
|
|
|
DEFAULT_PLACEMENT_POLICIES = {
|
|
"document": {
|
|
"require_local": True,
|
|
"stable_replica_count": 1,
|
|
"opportunistic_replica_count": 0,
|
|
"preferred_backend_ids": [],
|
|
"excluded_backend_ids": [],
|
|
"max_non_local_size_bytes": None,
|
|
},
|
|
"image": {
|
|
"require_local": True,
|
|
"stable_replica_count": 1,
|
|
"opportunistic_replica_count": 0,
|
|
"preferred_backend_ids": [],
|
|
"excluded_backend_ids": [],
|
|
"max_non_local_size_bytes": None,
|
|
},
|
|
"video": {
|
|
"require_local": True,
|
|
"stable_replica_count": 1,
|
|
"opportunistic_replica_count": 0,
|
|
"preferred_backend_ids": [],
|
|
"excluded_backend_ids": [],
|
|
"max_non_local_size_bytes": None,
|
|
},
|
|
"other": {
|
|
"require_local": True,
|
|
"stable_replica_count": 1,
|
|
"opportunistic_replica_count": 0,
|
|
"preferred_backend_ids": [],
|
|
"excluded_backend_ids": [],
|
|
"max_non_local_size_bytes": None,
|
|
},
|
|
}
|
|
ALLOWED_FILE_CLASSES = frozenset(DEFAULT_PLACEMENT_POLICIES.keys())
|
|
|
|
|
|
@dataclass
|
|
class PlacementDecision:
|
|
file_class: str
|
|
policy: PlacementPolicy
|
|
backends: list[Backend]
|
|
non_local_allowed: bool
|
|
|
|
|
|
class PlacementPolicyValidationError(ValidationError):
|
|
code = "placement_policy_invalid"
|
|
message = "Placement policy is invalid."
|
|
|
|
|
|
class PolicyService:
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self.session = session
|
|
self.backend_repository = BackendRepository(session)
|
|
self.repository = PlacementPolicyRepository(session)
|
|
|
|
async def ensure_default_policies(self) -> None:
|
|
for file_class, values in DEFAULT_PLACEMENT_POLICIES.items():
|
|
existing = await self.repository.get_by_file_class(file_class)
|
|
if existing is not None:
|
|
continue
|
|
policy = PlacementPolicy(
|
|
id=new_id("plc"),
|
|
file_class=file_class,
|
|
require_local=values["require_local"],
|
|
stable_replica_count=values["stable_replica_count"],
|
|
opportunistic_replica_count=values["opportunistic_replica_count"],
|
|
preferred_backend_ids_json=json.dumps(values["preferred_backend_ids"]),
|
|
excluded_backend_ids_json=json.dumps(values["excluded_backend_ids"]),
|
|
max_non_local_size_bytes=values["max_non_local_size_bytes"],
|
|
)
|
|
await self.repository.create(policy)
|
|
await self.session.commit()
|
|
|
|
async def list_policies(self) -> list[PlacementPolicy]:
|
|
await self.ensure_default_policies()
|
|
return await self.repository.list_all()
|
|
|
|
async def upsert_policy(
|
|
self,
|
|
*,
|
|
file_class: str,
|
|
require_local: bool,
|
|
stable_replica_count: int,
|
|
opportunistic_replica_count: int,
|
|
preferred_backend_ids: list[str],
|
|
excluded_backend_ids: list[str],
|
|
max_non_local_size_bytes: int | None,
|
|
) -> PlacementPolicy:
|
|
await self.ensure_default_policies()
|
|
self._validate_file_class(file_class)
|
|
await self._validate_policy_inputs(
|
|
require_local=require_local,
|
|
stable_replica_count=stable_replica_count,
|
|
opportunistic_replica_count=opportunistic_replica_count,
|
|
preferred_backend_ids=preferred_backend_ids,
|
|
excluded_backend_ids=excluded_backend_ids,
|
|
max_non_local_size_bytes=max_non_local_size_bytes,
|
|
)
|
|
existing = await self.repository.get_by_file_class(file_class)
|
|
if existing is None:
|
|
existing = PlacementPolicy(
|
|
id=new_id("plc"),
|
|
file_class=file_class,
|
|
require_local=require_local,
|
|
stable_replica_count=stable_replica_count,
|
|
opportunistic_replica_count=opportunistic_replica_count,
|
|
preferred_backend_ids_json=json.dumps(preferred_backend_ids),
|
|
excluded_backend_ids_json=json.dumps(excluded_backend_ids),
|
|
max_non_local_size_bytes=max_non_local_size_bytes,
|
|
)
|
|
await self.repository.create(existing)
|
|
else:
|
|
existing.require_local = require_local
|
|
existing.stable_replica_count = stable_replica_count
|
|
existing.opportunistic_replica_count = opportunistic_replica_count
|
|
existing.preferred_backend_ids_json = json.dumps(preferred_backend_ids)
|
|
existing.excluded_backend_ids_json = json.dumps(excluded_backend_ids)
|
|
existing.max_non_local_size_bytes = max_non_local_size_bytes
|
|
await self.repository.save(existing)
|
|
await self.session.commit()
|
|
return existing
|
|
|
|
async def get_placement_decision(self, mime_type: str | None, size_bytes: int | None = None) -> PlacementDecision:
|
|
await self.ensure_default_policies()
|
|
file_class = self.classify_mime_type(mime_type)
|
|
policy = await self.repository.get_by_file_class(file_class)
|
|
if policy is None:
|
|
raise ValueError(f"Placement policy is missing for file class: {file_class}")
|
|
|
|
preferred_backend_ids = self._decode_backend_ids(policy.preferred_backend_ids_json)
|
|
excluded_backend_ids = set(self._decode_backend_ids(policy.excluded_backend_ids_json))
|
|
backends = await self.backend_repository.list_enabled()
|
|
eligible_backends = [backend for backend in backends if backend.id not in excluded_backend_ids]
|
|
local_backends = self._sort_backends(
|
|
[backend for backend in eligible_backends if backend.type == "local_directory"],
|
|
preferred_backend_ids,
|
|
)
|
|
stable_backends = [
|
|
backend
|
|
for backend in eligible_backends
|
|
if backend.type != "local_directory" and backend.stability_class == "stable"
|
|
]
|
|
opportunistic_backends = [
|
|
backend
|
|
for backend in eligible_backends
|
|
if backend.type != "local_directory" and backend.stability_class == "opportunistic"
|
|
]
|
|
stable_backends = self._sort_backends(stable_backends, preferred_backend_ids)
|
|
opportunistic_backends = self._sort_backends(opportunistic_backends, preferred_backend_ids)
|
|
|
|
selected: list[Backend] = []
|
|
if policy.require_local and local_backends:
|
|
selected.append(local_backends[0])
|
|
|
|
non_local_allowed = (
|
|
policy.max_non_local_size_bytes is None
|
|
or size_bytes is None
|
|
or size_bytes <= policy.max_non_local_size_bytes
|
|
)
|
|
if non_local_allowed:
|
|
selected.extend(stable_backends[: policy.stable_replica_count])
|
|
selected.extend(opportunistic_backends[: policy.opportunistic_replica_count])
|
|
|
|
deduped: list[Backend] = []
|
|
seen = set()
|
|
for backend in selected:
|
|
if backend.id in seen:
|
|
continue
|
|
seen.add(backend.id)
|
|
deduped.append(backend)
|
|
|
|
return PlacementDecision(
|
|
file_class=file_class,
|
|
policy=policy,
|
|
backends=deduped,
|
|
non_local_allowed=non_local_allowed,
|
|
)
|
|
|
|
@staticmethod
|
|
def classify_mime_type(mime_type: str | None) -> str:
|
|
if mime_type is None:
|
|
return "other"
|
|
if mime_type.startswith("image/"):
|
|
return "image"
|
|
if mime_type.startswith("video/"):
|
|
return "video"
|
|
if mime_type == "application/pdf" or mime_type.startswith("text/"):
|
|
return "document"
|
|
return "other"
|
|
|
|
@staticmethod
|
|
def _decode_backend_ids(raw_value: str | None) -> list[str]:
|
|
if not raw_value:
|
|
return []
|
|
value = json.loads(raw_value)
|
|
if not isinstance(value, list):
|
|
return []
|
|
return [item for item in value if isinstance(item, str) and item]
|
|
|
|
@staticmethod
|
|
def _sort_backends(backends: list[Backend], preferred_backend_ids: list[str]) -> list[Backend]:
|
|
preferred_positions = {backend_id: index for index, backend_id in enumerate(preferred_backend_ids)}
|
|
return sorted(
|
|
backends,
|
|
key=lambda backend: (
|
|
0 if backend.id in preferred_positions else 1,
|
|
preferred_positions.get(backend.id, 0),
|
|
-backend.write_priority,
|
|
-backend.read_priority,
|
|
backend.name,
|
|
),
|
|
)
|
|
|
|
@staticmethod
|
|
def _validate_file_class(file_class: str) -> None:
|
|
if file_class not in ALLOWED_FILE_CLASSES:
|
|
raise PlacementPolicyValidationError(
|
|
"Unsupported file class. Allowed values: " + ", ".join(sorted(ALLOWED_FILE_CLASSES))
|
|
)
|
|
|
|
async def _validate_policy_inputs(
|
|
self,
|
|
*,
|
|
require_local: bool,
|
|
stable_replica_count: int,
|
|
opportunistic_replica_count: int,
|
|
preferred_backend_ids: list[str],
|
|
excluded_backend_ids: list[str],
|
|
max_non_local_size_bytes: int | None,
|
|
) -> None:
|
|
del require_local
|
|
if stable_replica_count < 0 or opportunistic_replica_count < 0:
|
|
raise PlacementPolicyValidationError("Replica counts must be zero or greater.")
|
|
if max_non_local_size_bytes is not None and max_non_local_size_bytes <= 0:
|
|
raise PlacementPolicyValidationError("max_non_local_size_bytes must be greater than zero.")
|
|
|
|
preferred_set = set(preferred_backend_ids)
|
|
excluded_set = set(excluded_backend_ids)
|
|
if len(preferred_set) != len(preferred_backend_ids):
|
|
raise PlacementPolicyValidationError("preferred_backend_ids must not contain duplicates.")
|
|
if len(excluded_set) != len(excluded_backend_ids):
|
|
raise PlacementPolicyValidationError("excluded_backend_ids must not contain duplicates.")
|
|
overlap = preferred_set & excluded_set
|
|
if overlap:
|
|
raise PlacementPolicyValidationError("A backend cannot be both preferred and excluded.")
|
|
|
|
known_backend_ids = {backend.id for backend in await self.backend_repository.list_all()}
|
|
unknown_backend_ids = sorted((preferred_set | excluded_set) - known_backend_ids)
|
|
if unknown_backend_ids:
|
|
raise PlacementPolicyValidationError(
|
|
"Unknown backend ids referenced by placement policy: " + ", ".join(unknown_backend_ids)
|
|
)
|
|
|