Skip to content

WIP: refactor worker DB data preloading #403

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
krokicki wants to merge 1 commit into
base: main
Choose a base branch
from
Open

WIP: refactor worker DB data preloading #403

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions fileglancer/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,12 +365,19 @@ async def _worker_exec(username: str, action: str, **kwargs):
raise HTTPException(status_code=e.status_code, detail=str(e))
else:
# CLI mode: run action directly in-process (single-user, no setuid)
from fileglancer.user_worker import _ACTIONS, WorkerContext, LocalDbProxy
from fileglancer.user_worker import (
_ACTIONS,
WorkerContext,
LocalDbProxy,
prepare_worker_request,
)
handler = _ACTIONS.get(action)
if handler is None:
raise HTTPException(status_code=500, detail=f"Unknown action: {action}")
ctx = WorkerContext(username=username, db=LocalDbProxy(settings.db_url))
request = {"action": action, **kwargs}
db_proxy = LocalDbProxy(settings.db_url)
ctx = WorkerContext(username=username)
request = prepare_worker_request(
{"action": action, **kwargs}, username, db_proxy)
try:
result = handler(request, ctx)
except Exception as e:
Expand Down
161 changes: 100 additions & 61 deletions fileglancer/user_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,19 +147,16 @@ def decorator(fn):


# ---------------------------------------------------------------------------
# DB proxy
# Parent-prepared DB data
# ---------------------------------------------------------------------------
#
# Worker subprocesses run as the (untrusted) target user, so we don't give
# them the database URL. Instead, action handlers go through a DbProxy that
# reverse-RPCs back to the parent over the same socket. The parent runs the
# query with full credentials and returns the result.
# them the database URL. The parent process reads the small amount of DB data
# an action needs before sending the request to the worker and includes it in
# the request payload. Workers never send reverse DB RPCs.
#
# In dev/test mode (no subprocess), the same handlers use LocalDbProxy which
# calls the database functions directly.
#
# DB_METHODS is the whitelist of methods both proxies expose; the parent's
# inbound dispatch only accepts these names.
# calls the database functions directly to prepare the same request payload.

from types import SimpleNamespace

Expand All @@ -183,10 +180,11 @@ def _job_db_to_dict(j) -> dict:


class LocalDbProxy:
"""DbProxy backed by a real database connection.
"""Small parent-side DB facade used to prepare worker requests.

Used in dev/test mode (in-process) and on the parent side as the
backend for inbound db_request messages from worker subprocesses.
Used by WorkerPool before dispatching to a subprocess and by CLI/dev mode
before calling action handlers in-process. It is never constructed inside
the user worker subprocess.
"""

def __init__(self, db_url: str):
Expand All @@ -206,52 +204,93 @@ def get_job(self, job_id: int, username: str):
return SimpleNamespace(**_job_db_to_dict(j))


class RpcDbProxy:
"""DbProxy that reverse-RPCs each call back to the parent over the socket."""

def __init__(self, sock: socket.socket):
self.sock = sock
# Actions that need the file-share-path list to resolve fsp_name, validate
# containment, or build redirect/browse-link metadata.
ACTIONS_REQUIRING_FILE_SHARE_PATHS = frozenset({
"list_dir",
"list_dir_paged",
"get_file_info",
"check_binary",
"open_file",
"head_file",
"create_dir",
"create_file",
"rename",
"delete",
"chmod",
"update_file",
"validate_proxied_path",
"validate_paths",
"create_dirs",
"get_profile",
"submit",
})

# Actions that need a minimal job row. get_service_urls already receives a
# caller-built list of job dicts to batch many jobs into one worker round-trip.
ACTIONS_REQUIRING_JOB = frozenset({
"get_job_file",
"delete_job_work_dir",
"get_service_url",
})


def _serialize_file_share_paths(paths) -> list[dict]:
return [fsp.model_dump(mode="json") for fsp in paths]


def _serialize_job_for_worker(job) -> Optional[dict]:
return None if job is None else vars(job)


def prepare_worker_request(request: dict, username: str, db_proxy: LocalDbProxy) -> dict:
"""Return a request enriched with all DB-derived data the worker needs.

The worker protocol is one request / one response. Any DB lookups required
by an action happen here, in the parent process, before the request is sent
to the untrusted user-context subprocess.
"""
action_name = request.get("action")
enriched = dict(request)

def _call(self, method: str, **kwargs):
_send(self.sock, {"_kind": "db_request", "method": method, "kwargs": kwargs})
resp = _recv(self.sock)
if resp.get("_kind") != "db_response":
raise RuntimeError(f"Expected db_response, got: {resp!r}")
if not resp.get("ok"):
raise RuntimeError(resp.get("error", "DB request failed"))
return resp.get("result")
if action_name in ACTIONS_REQUIRING_FILE_SHARE_PATHS:
enriched["file_share_paths"] = _serialize_file_share_paths(
db_proxy.get_file_share_paths()
)

def get_file_share_paths(self):
from fileglancer.model import FileSharePath
rows = self._call("get_file_share_paths") or []
return [FileSharePath(**r) for r in rows]
if action_name in ACTIONS_REQUIRING_JOB:
job_id = enriched["job_id"]
enriched["job"] = _serialize_job_for_worker(
db_proxy.get_job(job_id, username)
)

def get_job(self, job_id: int, username: str):
result = self._call("get_job", job_id=job_id, username=username)
return SimpleNamespace(**result) if result else None
return enriched


# Whitelist of method names the parent will dispatch; used by worker_pool
# when handling inbound db_request messages.
DB_METHODS = frozenset({"get_file_share_paths", "get_job"})
def _file_share_paths_from_request(request: dict) -> list:
"""Deserialize parent-provided file-share-path payload."""
from fileglancer.model import FileSharePath

try:
rows = request["file_share_paths"]
except KeyError as e:
raise RuntimeError(
f"Worker request for {request.get('action')} is missing "
"parent-provided file_share_paths"
) from e
return [FileSharePath(**row) for row in (rows or [])]

def serialize_db_result(method: str, value):
"""Convert a LocalDbProxy result into a JSON-serializable form.

Called by the parent before sending a db_response back to the worker.
Keeps the wire format consistent regardless of which backend produced
the value.
"""
if value is None:
return None
if method == "get_file_share_paths":
# value is a list of FileSharePath models
return [fsp.model_dump(mode="json") for fsp in value]
if method == "get_job":
# value is a SimpleNamespace; vars() gives the underlying dict
return vars(value)
raise ValueError(f"Unknown db method: {method}")
def _job_from_request(request: dict):
"""Deserialize parent-provided minimal job payload."""
try:
row = request["job"]
except KeyError as e:
raise RuntimeError(
f"Worker request for {request.get('action')} is missing "
"parent-provided job"
) from e
return SimpleNamespace(**row) if row else None


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -338,7 +377,7 @@ def with_filestore(fn):
"""
@functools.wraps(fn)
def wrapper(request: dict, ctx: WorkerContext) -> dict:
fsps = ctx.db.get_file_share_paths()
fsps = _file_share_paths_from_request(request)
filestore, error_response = _get_filestore(request["fsp_name"], fsps)
if filestore is None:
return error_response
Expand Down Expand Up @@ -655,7 +694,7 @@ def _action_validate_paths(request: dict, ctx: WorkerContext) -> dict:
# Expected type per key ('file' or 'directory'): when the path exists, its
# type must match.
types = request.get("types") or {}
fsps = ctx.db.get_file_share_paths()
fsps = _file_share_paths_from_request(request)
errors = {}
for param_key, path_value in paths.items():
error = validate_path_in_filestore(
Expand All @@ -681,7 +720,7 @@ def _action_create_dirs(request: dict, ctx: WorkerContext) -> dict:
from fileglancer.apps.command import validate_path_in_filestore

paths = request["paths"]
fsps = ctx.db.get_file_share_paths()
fsps = _file_share_paths_from_request(request)
errors = {}
for key, path_value in paths.items():
# Containment check without the exists/readable check — the directory is
Expand All @@ -702,7 +741,7 @@ def _action_create_dirs(request: dict, ctx: WorkerContext) -> dict:
def _action_get_profile(request: dict, ctx: WorkerContext) -> dict:
"""Get user profile information."""
username = ctx.username
paths = ctx.db.get_file_share_paths()
paths = _file_share_paths_from_request(request)

home_fsp = next((fsp for fsp in paths if fsp.mount_path in ('~', '~/')), None)
if home_fsp:
Expand Down Expand Up @@ -774,7 +813,7 @@ def _action_get_job_file(request: dict, ctx: WorkerContext) -> dict:
job_id = request["job_id"]
file_type = request["file_type"]

db_job = ctx.db.get_job(job_id, ctx.username)
db_job = _job_from_request(request)
if db_job is None:
return {"error": f"Job {job_id} not found", "status_code": 404}

Expand All @@ -791,7 +830,7 @@ def _action_delete_job_work_dir(request: dict, ctx: WorkerContext) -> dict:
from fileglancer.apps.jobfiles import delete_job_work_dir

job_id = request["job_id"]
db_job = ctx.db.get_job(job_id, ctx.username)
db_job = _job_from_request(request)
if db_job is None:
return {"error": f"Job {job_id} not found", "status_code": 404}
if not db.is_terminal_job_status(db_job.status):
Expand Down Expand Up @@ -936,7 +975,7 @@ def _action_get_service_url(request: dict, ctx: WorkerContext) -> dict:
from fileglancer.apps.jobfiles import get_service_url, get_service_phase

job_id = request["job_id"]
db_job = ctx.db.get_job(job_id, ctx.username)
db_job = _job_from_request(request)
if db_job is None:
return {"error": f"Job {job_id} not found", "status_code": 404}
return {"service_url": get_service_url(db_job), "phase": get_service_phase(db_job)}
Expand Down Expand Up @@ -1166,7 +1205,7 @@ def _action_submit(request: dict, ctx: WorkerContext) -> dict:
work_dir_subpath = None
try:
from fileglancer.database import find_fsp_in_paths
match = find_fsp_in_paths(ctx.db.get_file_share_paths(), str(work_dir))
match = find_fsp_in_paths(_file_share_paths_from_request(request), str(work_dir))
if match:
work_dir_fsp_name = match[0].name
work_dir_subpath = match[1]
Expand Down Expand Up @@ -1356,7 +1395,7 @@ def _action_remote_heads(request: dict, ctx: WorkerContext) -> dict:
class WorkerContext:
"""Holds per-worker state."""

def __init__(self, username: str, db):
def __init__(self, username: str, db=None):
self.username = username
self.db = db

Expand Down Expand Up @@ -1405,9 +1444,9 @@ def emit(self, record: logging.LogRecord) -> None:
except KeyError:
username = str(uid)

# Worker subprocess never gets DB credentials; all DB access goes back
# through the parent over the same socket via RpcDbProxy.
ctx = WorkerContext(username=username, db=RpcDbProxy(sock))
# Worker subprocess never gets DB credentials. Any DB-derived data needed
# by an action arrives already serialized in that action's request.
ctx = WorkerContext(username=username)

logger.info(
f"Worker started for {username} "
Expand Down
55 changes: 16 additions & 39 deletions fileglancer/worker_pool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def _build_worker_env(base_env: dict, home_dir: str, log_level: str,
unconditionally — even if an operator lists it in passthrough — so
Fileglancer secrets (session key, Okta/Atlassian secrets, DB URL) can never
reach the user. The worker gets its config over the IPC socket instead (DB
access via RpcDbProxy).
data is attached to the request by the parent).

A passthrough entry ending in ``_`` is treated as a prefix; otherwise it is
an exact variable name.
Expand Down Expand Up @@ -162,7 +162,7 @@ def __init__(self, username: str, process: subprocess.Popen,
self.username = username
self.process = process
self.sock = sock
self.db_proxy = db_proxy # LocalDbProxy used to satisfy worker db_requests
self.db_proxy = db_proxy # LocalDbProxy used to prepare worker requests
self.last_activity = time.monotonic()
self._busy = False
self._lock = asyncio.Lock() # serialize requests to the worker
Expand Down Expand Up @@ -227,20 +227,23 @@ def _send_and_recv(self, request: dict, timeout: float = _DEFAULT_REQUEST_TIMEOU
request (requests are serialized under the per-worker lock, so this is
safe) so that a long git clone/snapshot isn't misread as a dead worker.

Loops on receive: any inbound ``_kind == "db_request"`` message is a
reverse-RPC from the worker (which has no DB credentials) asking the
parent to run a DB query on its behalf. We dispatch it, send back a
``db_response``, and keep reading. Anything else is the action result.
Any DB data the worker needs is attached to the request before it is
sent. The worker never gets DB credentials and never sends reverse DB
RPCs; this remains a simple one request / one response exchange.
"""
from fileglancer.user_worker import prepare_worker_request

self.sock.settimeout(timeout)
try:
request = prepare_worker_request(request, self.username, self.db_proxy)
except Exception as e:
logger.exception(
f"Failed to prepare {request.get('action')} request for "
f"{self.username}"
)
raise WorkerError(f"Failed to prepare worker request: {e}") from e
self._send_msg(request)

while True:
response = self._recv_msg()
if response.get("_kind") == "db_request":
self._handle_db_request(response)
continue
return response
return self._recv_msg()

def _send_msg(self, msg: dict):
"""Send a length-prefixed JSON message."""
Expand Down Expand Up @@ -308,32 +311,6 @@ def _recv_msg(self) -> dict:

return response

def _handle_db_request(self, request: dict):
"""Run a DB query on behalf of the worker and send the result back."""
from fileglancer.user_worker import DB_METHODS, serialize_db_result

method = request.get("method")
kwargs = request.get("kwargs", {}) or {}
if method not in DB_METHODS:
self._send_msg({
"_kind": "db_response",
"ok": False,
"error": f"Unknown db method: {method}",
})
return

try:
value = getattr(self.db_proxy, method)(**kwargs)
result = serialize_db_result(method, value)
self._send_msg({"_kind": "db_response", "ok": True, "result": result})
except Exception as e:
logger.exception(f"db_request {method} for {self.username} failed")
self._send_msg({
"_kind": "db_response",
"ok": False,
"error": f"{type(e).__name__}: {e}",
})

async def shutdown(self, timeout: float = 5.0):
"""Ask the worker to shut down gracefully, then force-kill if needed."""
if not self.is_alive:
Expand Down
Loading
Loading