From f4bf7a4ae27df0394635d08c6d7d9a4ff5919907 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 13:52:06 +0530 Subject: [PATCH 01/24] perf(dashboard): disable refetch-on-focus, tune staleTime --- dashboard/src/features/circuit-breakers/hooks.ts | 1 + dashboard/src/features/resources/hooks.ts | 1 + dashboard/src/features/tasks/hooks.ts | 3 +++ dashboard/src/features/workers/hooks.ts | 3 +++ dashboard/src/lib/query-client.ts | 4 +++- 5 files changed, 11 insertions(+), 1 deletion(-) diff --git a/dashboard/src/features/circuit-breakers/hooks.ts b/dashboard/src/features/circuit-breakers/hooks.ts index 78643284..14375dbf 100644 --- a/dashboard/src/features/circuit-breakers/hooks.ts +++ b/dashboard/src/features/circuit-breakers/hooks.ts @@ -6,6 +6,7 @@ export function circuitBreakersQuery() { return queryOptions({ queryKey: ["circuit-breakers"], queryFn: ({ signal }) => fetchCircuitBreakers(signal), + staleTime: 60_000, }); } diff --git a/dashboard/src/features/resources/hooks.ts b/dashboard/src/features/resources/hooks.ts index a79e57a1..c5c7e3e1 100644 --- a/dashboard/src/features/resources/hooks.ts +++ b/dashboard/src/features/resources/hooks.ts @@ -6,6 +6,7 @@ export function resourcesQuery() { return queryOptions({ queryKey: ["resources"], queryFn: ({ signal }) => fetchResources(signal), + staleTime: 30_000, }); } diff --git a/dashboard/src/features/tasks/hooks.ts b/dashboard/src/features/tasks/hooks.ts index 2e91188f..35a1cd3c 100644 --- a/dashboard/src/features/tasks/hooks.ts +++ b/dashboard/src/features/tasks/hooks.ts @@ -25,6 +25,8 @@ export function tasksQuery() { return queryOptions({ queryKey: TASKS_KEY, queryFn: ({ signal }) => listTasks(signal), + // Task registry rarely changes between deploys. + staleTime: 60_000, }); } @@ -32,6 +34,7 @@ export function queuesQuery() { return queryOptions({ queryKey: QUEUES_KEY, queryFn: ({ signal }) => listQueues(signal), + staleTime: 60_000, }); } diff --git a/dashboard/src/features/workers/hooks.ts b/dashboard/src/features/workers/hooks.ts index 6668eb00..55939d88 100644 --- a/dashboard/src/features/workers/hooks.ts +++ b/dashboard/src/features/workers/hooks.ts @@ -6,6 +6,9 @@ export function workersQuery() { return queryOptions({ queryKey: ["workers"], queryFn: ({ signal }) => fetchWorkers(signal), + // Slow-moving: poll on the user interval, but don't refetch on remount + // within the window. + staleTime: 30_000, }); } diff --git a/dashboard/src/lib/query-client.ts b/dashboard/src/lib/query-client.ts index 9d50a02d..c2f3e0f9 100644 --- a/dashboard/src/lib/query-client.ts +++ b/dashboard/src/lib/query-client.ts @@ -13,7 +13,9 @@ export function createQueryClient(): QueryClient { queries: { staleTime: 5_000, gcTime: 5 * 60_000, - refetchOnWindowFocus: true, + // Interval polling already keeps data fresh; focus-refetch would fire + // every active query at once on tab refocus (a thundering-herd burst). + refetchOnWindowFocus: false, retry: (failureCount, error) => { const status = (error as { status?: number }).status; if (status && status >= 400 && status < 500) return false; From 669d6ed9447a32958f66537a50e3cbf6809c8339 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 13:52:12 +0530 Subject: [PATCH 02/24] perf(dashboard): memoize DataTable + stabilize row callbacks --- dashboard/src/components/ui/data-table.tsx | 10 ++++++++-- .../src/features/jobs/components/job-table.tsx | 15 +++++++-------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/dashboard/src/components/ui/data-table.tsx b/dashboard/src/components/ui/data-table.tsx index 9e9ae81a..4b078c15 100644 --- a/dashboard/src/components/ui/data-table.tsx +++ b/dashboard/src/components/ui/data-table.tsx @@ -8,7 +8,7 @@ import { useReactTable, } from "@tanstack/react-table"; import { ArrowDown, ArrowUp, ChevronsUpDown } from "lucide-react"; -import { type ReactNode, useState } from "react"; +import { memo, type ReactNode, useState } from "react"; import { cn } from "@/lib/cn"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "./table"; @@ -22,7 +22,7 @@ interface DataTableProps { initialSorting?: SortingState; } -export function DataTable({ +function DataTableImpl({ columns, data, empty, @@ -126,3 +126,9 @@ export function DataTable({ ); } + +// Memoized so unrelated parent re-renders (e.g. the polling clock) don't +// re-render the whole table. The cast preserves the generic call signature +// that React.memo otherwise erases. Relies on callers passing stable +// `columns`/`onRowClick` (via useMemo/useCallback) to actually skip renders. +export const DataTable = memo(DataTableImpl) as typeof DataTableImpl; diff --git a/dashboard/src/features/jobs/components/job-table.tsx b/dashboard/src/features/jobs/components/job-table.tsx index 7710f8ba..37274dc9 100644 --- a/dashboard/src/features/jobs/components/job-table.tsx +++ b/dashboard/src/features/jobs/components/job-table.tsx @@ -1,6 +1,6 @@ import { useNavigate } from "@tanstack/react-router"; import type { ColumnDef } from "@tanstack/react-table"; -import { useMemo } from "react"; +import { useCallback, useMemo } from "react"; import { Badge, DataTable, @@ -24,7 +24,11 @@ interface JobTableProps { export function JobTable({ jobs, loading, error, onRetry }: JobTableProps) { const navigate = useNavigate(); - const showErrorColumn = !!jobs?.some((j) => j.error); + const showErrorColumn = useMemo(() => !!jobs?.some((j) => j.error), [jobs]); + const handleRowClick = useCallback( + (job: Job) => navigate({ to: "/jobs/$id", params: { id: job.id } }), + [navigate], + ); const columns = useMemo[]>(() => { const base: ColumnDef[] = [ @@ -128,11 +132,6 @@ export function JobTable({ jobs, loading, error, onRetry }: JobTableProps) { } return ( - j.id} - onRowClick={(job) => navigate({ to: "/jobs/$id", params: { id: job.id } })} - /> + j.id} onRowClick={handleRowClick} /> ); } From 5ff7dae68358ee6f816381d6986612b7e3b8393c Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 13:52:17 +0530 Subject: [PATCH 03/24] perf(dashboard): memoize sparkline geometry, unique gradient id --- .../components/throughput-sparkline.tsx | 51 ++++++++++--------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/dashboard/src/features/overview/components/throughput-sparkline.tsx b/dashboard/src/features/overview/components/throughput-sparkline.tsx index 7f90847c..47f4fd9f 100644 --- a/dashboard/src/features/overview/components/throughput-sparkline.tsx +++ b/dashboard/src/features/overview/components/throughput-sparkline.tsx @@ -1,4 +1,4 @@ -import { useState } from "react"; +import { useId, useMemo, useState } from "react"; import { Card, CardContent, CardHeader, CardTitle, ErrorState, Skeleton } from "@/components/ui"; import type { TimeseriesBucket } from "@/lib/api-types"; import { formatCount } from "@/lib/number"; @@ -17,9 +17,9 @@ export function ThroughputSparkline({ error, onRetry, }: ThroughputSparklineProps) { - const points = buckets ?? []; - const total = points.reduce((sum, b) => sum + b.count, 0); - const peak = points.reduce((max, b) => Math.max(max, b.count), 0); + const points = useMemo(() => buckets ?? [], [buckets]); + const total = useMemo(() => points.reduce((sum, b) => sum + b.count, 0), [points]); + const peak = useMemo(() => points.reduce((max, b) => Math.max(max, b.count), 0), [points]); if (error) { return ( @@ -54,29 +54,30 @@ export function ThroughputSparkline({ ); } +const SPARK_WIDTH = 800; +const SPARK_HEIGHT = 80; + function Sparkline({ buckets }: { buckets: TimeseriesBucket[] }) { - const width = 800; - const height = 80; - const maxCount = Math.max(1, ...buckets.map((b) => b.count)); - const step = buckets.length > 1 ? width / (buckets.length - 1) : 0; const [hoverIdx, setHoverIdx] = useState(null); + // Unique per instance so multiple sparklines don't share one gradient . + const gradientId = `sparkline-fill-${useId().replace(/:/g, "")}`; - const areaPath = buckets - .map((b, i) => { - const x = i * step; - const y = height - (b.count / maxCount) * height; - return `${i === 0 ? "M" : "L"} ${x.toFixed(1)} ${y.toFixed(1)}`; - }) - .concat([`L ${width} ${height}`, `L 0 ${height}`, "Z"]) - .join(" "); - - const linePath = buckets - .map((b, i) => { + // Geometry depends only on the data — recompute on new buckets, not on hover. + const { areaPath, linePath } = useMemo(() => { + const maxCount = Math.max(1, ...buckets.map((b) => b.count)); + const step = buckets.length > 1 ? SPARK_WIDTH / (buckets.length - 1) : 0; + const points = buckets.map((b, i) => { const x = i * step; - const y = height - (b.count / maxCount) * height; + const y = SPARK_HEIGHT - (b.count / maxCount) * SPARK_HEIGHT; return `${i === 0 ? "M" : "L"} ${x.toFixed(1)} ${y.toFixed(1)}`; - }) - .join(" "); + }); + return { + areaPath: points + .concat([`L ${SPARK_WIDTH} ${SPARK_HEIGHT}`, `L 0 ${SPARK_HEIGHT}`, "Z"]) + .join(" "), + linePath: points.join(" "), + }; + }, [buckets]); const startLabel = buckets[0] ? formatRelative(buckets[0].timestamp) : ""; const endLabel = "now"; @@ -103,19 +104,19 @@ function Sparkline({ buckets }: { buckets: TimeseriesBucket[] }) { onMouseLeave={() => setHoverIdx(null)} > - + - + Date: Fri, 29 May 2026 13:52:23 +0530 Subject: [PATCH 04/24] perf(dashboard): throttle LastRefreshed re-render cadence --- .../src/components/layout/last-refreshed.tsx | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/dashboard/src/components/layout/last-refreshed.tsx b/dashboard/src/components/layout/last-refreshed.tsx index 8e49a67c..70c9d366 100644 --- a/dashboard/src/components/layout/last-refreshed.tsx +++ b/dashboard/src/components/layout/last-refreshed.tsx @@ -17,10 +17,23 @@ export function LastRefreshed({ className }: { className?: string }) { const { lastRefreshedAt, isFetching } = useLastRefreshed(); const [, setTick] = useState(0); + // Re-render only as often as the label can actually change: every second + // while the "Ns ago" portion ticks, then minute- and hour-granularity once + // the timestamp ages. Avoids a fixed 1s re-render forever. useEffect(() => { - const id = setInterval(() => setTick((n) => n + 1), 1000); - return () => clearInterval(id); - }, []); + if (isFetching) return; + let id: ReturnType; + const schedule = () => { + const age = Date.now() - lastRefreshedAt; + const delay = age < 60_000 ? 1_000 : age < 3_600_000 ? 60_000 : 3_600_000; + id = setTimeout(() => { + setTick((n) => n + 1); + schedule(); + }, delay); + }; + schedule(); + return () => clearTimeout(id); + }, [isFetching, lastRefreshedAt]); const label = isFetching ? "Refreshing…" : `Updated ${formatAgo(Date.now() - lastRefreshedAt)}`; From 95d83d82080b19f1645ef1546e68d933aa5d8d87 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 13:52:28 +0530 Subject: [PATCH 05/24] perf(dashboard): memoize job integration links --- .../jobs/components/job-overview-tab.tsx | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/dashboard/src/features/jobs/components/job-overview-tab.tsx b/dashboard/src/features/jobs/components/job-overview-tab.tsx index 203163ac..25ca2be6 100644 --- a/dashboard/src/features/jobs/components/job-overview-tab.tsx +++ b/dashboard/src/features/jobs/components/job-overview-tab.tsx @@ -1,5 +1,5 @@ import { ExternalLink as ExternalLinkIcon } from "lucide-react"; -import type { ReactNode } from "react"; +import { type ReactNode, useMemo } from "react"; import { buttonVariants, Card, CardContent, CardHeader, CardTitle } from "@/components/ui"; import { applyJobContext, useIntegrations } from "@/features/settings"; import type { Job } from "@/lib/api-types"; @@ -208,13 +208,17 @@ function NotesCard({ raw }: { raw: string | null }) { */ function JobIntegrations({ job }: { job: Job }) { const integrations = useIntegrations(); - const links = ( - [ - { label: "Open in Grafana", href: integrations.grafana }, - { label: "Open in Sentry", href: integrations.sentry }, - { label: "Open in OTel", href: integrations.otel }, - ] as const - ).filter((entry) => entry.href); + const links = useMemo( + () => + ( + [ + { label: "Open in Grafana", href: integrations.grafana }, + { label: "Open in Sentry", href: integrations.sentry }, + { label: "Open in OTel", href: integrations.otel }, + ] as const + ).filter((entry) => entry.href), + [integrations.grafana, integrations.sentry, integrations.otel], + ); if (links.length === 0) return null; From 2eceb87d37f1188943d5b28936d98cc726e5e932 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 13:52:33 +0530 Subject: [PATCH 06/24] perf(dashboard): lazy-load job detail tab queries --- dashboard/src/routes/jobs/$id.tsx | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/dashboard/src/routes/jobs/$id.tsx b/dashboard/src/routes/jobs/$id.tsx index 5145e324..82d04880 100644 --- a/dashboard/src/routes/jobs/$id.tsx +++ b/dashboard/src/routes/jobs/$id.tsx @@ -1,4 +1,5 @@ import { createFileRoute } from "@tanstack/react-router"; +import { useState } from "react"; import { PageHeader } from "@/components/layout"; import { Badge, @@ -33,11 +34,13 @@ export const Route = createFileRoute("/jobs/$id")({ function JobDetailPage() { const { id } = Route.useParams(); + const [tab, setTab] = useState("overview"); const job = useJob(id); - const logs = useJobLogs(id); - const errors = useJobErrors(id); - const replays = useReplayHistory(id); - const dag = useJobDag(id); + // Each tab's data (and its polling) only fetches while that tab is open. + const logs = useJobLogs(id, tab === "logs"); + const errors = useJobErrors(id, tab === "errors"); + const replays = useReplayHistory(id, tab === "replays"); + const dag = useJobDag(id, tab === "dag"); if (job.isLoading && !job.data) { return ( @@ -85,7 +88,7 @@ function JobDetailPage() { } /> - + Overview Logs From 16c9b35e09df303338aaa6514d1d4a619269f999 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 13:52:39 +0530 Subject: [PATCH 07/24] fix(dashboard): read theme from localStorage once --- dashboard/src/providers/theme-provider.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dashboard/src/providers/theme-provider.tsx b/dashboard/src/providers/theme-provider.tsx index 58f1bb06..c4247c4a 100644 --- a/dashboard/src/providers/theme-provider.tsx +++ b/dashboard/src/providers/theme-provider.tsx @@ -31,7 +31,7 @@ function applyTheme(resolved: ResolvedTheme) { export function ThemeProvider({ children }: { children: ReactNode }) { const [theme, setThemeState] = useState(() => readStored()); const [resolved, setResolved] = useState(() => - readStored() === "system" ? readSystem() : (readStored() as ResolvedTheme), + theme === "system" ? readSystem() : theme, ); useEffect(() => { From e566e7de0cd7269984527d565030d29f503dc5a3 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:15:16 +0530 Subject: [PATCH 08/24] perf: trim enqueue and middleware-dispatch hot paths --- py_src/taskito/app.py | 207 +++++++++++++--------- py_src/taskito/mixins/decorators.py | 33 +++- py_src/taskito/mixins/middleware_admin.py | 18 +- py_src/taskito/notes.py | 4 +- tests/core/test_middleware_chain_cache.py | 54 ++++++ 5 files changed, 222 insertions(+), 94 deletions(-) create mode 100644 tests/core/test_middleware_chain_cache.py diff --git a/py_src/taskito/app.py b/py_src/taskito/app.py index 8cbc2df9..b2189ea0 100644 --- a/py_src/taskito/app.py +++ b/py_src/taskito/app.py @@ -226,6 +226,12 @@ def __init__( self._batch_accumulator: BatchAccumulator | None = None self._global_middleware: list[TaskMiddleware] = middleware or [] self._task_middleware: dict[str, list[TaskMiddleware]] = {} + # Per-task middleware-chain cache (chain, disable_version, computed_at). + # Bumping ``_mw_disable_version`` invalidates same-process readers + # instantly; the TTL bounds cross-process staleness. See + # ``_get_middleware_chain``. + self._mw_chain_cache: dict[str, tuple[list[TaskMiddleware], int, float]] = {} + self._mw_disable_version: int = 0 self._task_retry_filters: dict[str, dict[str, list[type[Exception]]]] = {} self._init_predicate_state() self._drain_timeout = drain_timeout @@ -434,44 +440,47 @@ def enqueue( final_args = args final_kwargs = kwargs or {} - # Run on_enqueue middleware hook (options dict is mutable) - enqueue_options: dict[str, Any] = { - "priority": priority, - "delay": delay, - "queue": queue, - "max_retries": max_retries, - "timeout": timeout, - "unique_key": unique_key, - "metadata": metadata, - "notes": notes, - "depends_on": depends_on, - "expires": expires, - "result_ttl": result_ttl, - "idempotency_key": idempotency_key, - "idempotent": idempotent, - } - for mw in self._global_middleware: - if not mw._should_apply(None, task_name=task_name): - continue - try: - mw.on_enqueue(task_name, final_args, final_kwargs, enqueue_options) - except Exception: - logger.exception("middleware on_enqueue() error") - - # Apply any middleware mutations back - priority = enqueue_options.get("priority") - delay = enqueue_options.get("delay") - queue = enqueue_options.get("queue") - max_retries = enqueue_options.get("max_retries") - timeout = enqueue_options.get("timeout") - unique_key = enqueue_options.get("unique_key") - metadata = enqueue_options.get("metadata") - notes = enqueue_options.get("notes") - depends_on = enqueue_options.get("depends_on") - expires = enqueue_options.get("expires") - result_ttl = enqueue_options.get("result_ttl") - idempotency_key = enqueue_options.get("idempotency_key") - idempotent = enqueue_options.get("idempotent") + # Run on_enqueue middleware hook (options dict is mutable). Skip the + # whole build/read-back when no middleware is registered — the common + # hot path — to avoid a needless dict allocation and 13 reads per call. + if self._global_middleware: + enqueue_options: dict[str, Any] = { + "priority": priority, + "delay": delay, + "queue": queue, + "max_retries": max_retries, + "timeout": timeout, + "unique_key": unique_key, + "metadata": metadata, + "notes": notes, + "depends_on": depends_on, + "expires": expires, + "result_ttl": result_ttl, + "idempotency_key": idempotency_key, + "idempotent": idempotent, + } + for mw in self._global_middleware: + if not mw._should_apply(None, task_name=task_name): + continue + try: + mw.on_enqueue(task_name, final_args, final_kwargs, enqueue_options) + except Exception: + logger.exception("middleware on_enqueue() error") + + # Apply any middleware mutations back + priority = enqueue_options.get("priority") + delay = enqueue_options.get("delay") + queue = enqueue_options.get("queue") + max_retries = enqueue_options.get("max_retries") + timeout = enqueue_options.get("timeout") + unique_key = enqueue_options.get("unique_key") + metadata = enqueue_options.get("metadata") + notes = enqueue_options.get("notes") + depends_on = enqueue_options.get("depends_on") + expires = enqueue_options.get("expires") + result_ttl = enqueue_options.get("result_ttl") + idempotency_key = enqueue_options.get("idempotency_key") + idempotent = enqueue_options.get("idempotent") # Validation runs *after* middleware so a mutating hook still gets # the chance to reshape notes before we reject them. @@ -624,51 +633,79 @@ def enqueue_many( ) kw_list = kwargs_list or [{}] * count - # Build a per-job options dict so on_enqueue middleware can mutate - # priority/queue/delay/etc. on a per-job basis before the batch is - # committed. The dispatch must happen BEFORE enqueue_batch — running - # it after (as the previous implementation did) made mutations - # impossible to apply. - per_job_options: list[dict[str, Any]] = [ - { - "priority": priority, - "queue": queue, - "max_retries": max_retries, - "timeout": timeout, - "delay": (delay_list[i] if delay_list is not None else delay), - "unique_key": (unique_keys[i] if unique_keys is not None else None), - "metadata": (metadata_list[i] if metadata_list is not None else metadata), - "notes": (notes_list[i] if notes_list is not None else notes), - "expires": (expires_list[i] if expires_list is not None else expires), - "result_ttl": (result_ttl_list[i] if result_ttl_list is not None else result_ttl), - "idempotency_key": (idempotency_keys[i] if idempotency_keys is not None else None), - "idempotent": idempotent, - } - for i in range(count) - ] - chain = self._get_middleware_chain(task_name) - for i in range(count): - for mw in chain: - if not mw._should_apply(None, task_name=task_name): - continue - try: - mw.on_enqueue(task_name, args_list[i], kw_list[i], per_job_options[i]) - except Exception: - logger.exception("middleware on_enqueue() error") - - # Read mutated per-job options back into the per-job lists passed to - # the Rust batch enqueue. `None` entries are forwarded so the Rust - # side falls back to its defaults. - queues_list = [opt["queue"] or "default" for opt in per_job_options] - priorities_list = [opt["priority"] for opt in per_job_options] - retries_list = [opt["max_retries"] for opt in per_job_options] - timeouts_list = [opt["timeout"] for opt in per_job_options] - delays = [opt["delay"] for opt in per_job_options] - metas = [opt["metadata"] for opt in per_job_options] - notes_encoded = [validate_and_encode_notes(opt["notes"]) for opt in per_job_options] - exp_list = [opt["expires"] for opt in per_job_options] - ttl_list = [opt["result_ttl"] for opt in per_job_options] + # ``_should_apply`` depends only on (task_name, middleware), which is + # constant across the batch — evaluate it once, not per job. [B4] + applicable = [mw for mw in chain if mw._should_apply(None, task_name=task_name)] + + if applicable: + # Middleware may mutate priority/queue/delay/etc. per job, so build + # the mutable per-job options dicts and dispatch ``on_enqueue`` + # BEFORE enqueue_batch so mutations are applied. + per_job_options: list[dict[str, Any]] = [ + { + "priority": priority, + "queue": queue, + "max_retries": max_retries, + "timeout": timeout, + "delay": (delay_list[i] if delay_list is not None else delay), + "unique_key": (unique_keys[i] if unique_keys is not None else None), + "metadata": (metadata_list[i] if metadata_list is not None else metadata), + "notes": (notes_list[i] if notes_list is not None else notes), + "expires": (expires_list[i] if expires_list is not None else expires), + "result_ttl": ( + result_ttl_list[i] if result_ttl_list is not None else result_ttl + ), + "idempotency_key": ( + idempotency_keys[i] if idempotency_keys is not None else None + ), + "idempotent": idempotent, + } + for i in range(count) + ] + for i in range(count): + for mw in applicable: + try: + mw.on_enqueue(task_name, args_list[i], kw_list[i], per_job_options[i]) + except Exception: + logger.exception("middleware on_enqueue() error") + + queues_list = [opt["queue"] or "default" for opt in per_job_options] + priorities_list = [opt["priority"] for opt in per_job_options] + retries_list = [opt["max_retries"] for opt in per_job_options] + timeouts_list = [opt["timeout"] for opt in per_job_options] + delays = [opt["delay"] for opt in per_job_options] + metas = [opt["metadata"] for opt in per_job_options] + notes_encoded = [validate_and_encode_notes(opt["notes"]) for opt in per_job_options] + exp_list = [opt["expires"] for opt in per_job_options] + ttl_list = [opt["result_ttl"] for opt in per_job_options] + uk_src = [opt["unique_key"] for opt in per_job_options] + ik_src = [opt["idempotency_key"] for opt in per_job_options] + idem_src = [opt["idempotent"] for opt in per_job_options] + else: + # No applicable middleware — build the parallel lists the Rust batch + # enqueue wants directly, skipping the per-job dict layer. [B3] + # Copy any per-job input lists so a later predicate mutation can't + # corrupt the caller's arguments. + default_queue = queue or "default" + queues_list = [default_queue] * count + priorities_list = [priority] * count + retries_list = [max_retries] * count + timeouts_list = [timeout] * count + delays = list(delay_list) if delay_list is not None else [delay] * count + metas = list(metadata_list) if metadata_list is not None else [metadata] * count + if notes_list is None and notes is None: + notes_encoded = [None] * count # no per-job validate call needed + else: + notes_source = notes_list if notes_list is not None else [notes] * count + notes_encoded = [validate_and_encode_notes(n) for n in notes_source] + exp_list = list(expires_list) if expires_list is not None else [expires] * count + ttl_list = ( + list(result_ttl_list) if result_ttl_list is not None else [result_ttl] * count + ) + uk_src = list(unique_keys) if unique_keys is not None else [None] * count + ik_src = list(idempotency_keys) if idempotency_keys is not None else [None] * count + idem_src = [idempotent] * count task_serializer = self._get_serializer(task_name) if self._interceptor is not None: @@ -687,9 +724,9 @@ def enqueue_many( self._resolve_unique_key( task_name=task_name, payload=payloads[i], - unique_key=per_job_options[i]["unique_key"], - idempotency_key=per_job_options[i]["idempotency_key"], - idempotent=per_job_options[i]["idempotent"], + unique_key=uk_src[i], + idempotency_key=ik_src[i], + idempotent=idem_src[i], ) for i in range(count) ] diff --git a/py_src/taskito/mixins/decorators.py b/py_src/taskito/mixins/decorators.py index 4ef6daab..8db59335 100644 --- a/py_src/taskito/mixins/decorators.py +++ b/py_src/taskito/mixins/decorators.py @@ -53,6 +53,10 @@ # `CELERYD_TASK_LOG_FORMAT` truncation in spirit. _MAX_RESULT_REPR = 80 +# How long a cached middleware chain stays valid without a version bump. Bounds +# the worst-case lag for an out-of-process dashboard disable change. +_MW_CHAIN_TTL = 1.0 + def _safe_result_repr(value: Any) -> str: """Render a task return value for the success log, bounded and crash-proof.""" @@ -123,19 +127,38 @@ class QueueDecoratorMixin: # serializer. Declared here so mypy sees it through the mixin. _apply_dispatch_predicate: Callable[..., None] + # Middleware-chain cache state (initialized in ``Queue.__init__``). + _mw_chain_cache: dict[str, tuple[list[TaskMiddleware], int, float]] + _mw_disable_version: int + def _get_middleware_chain(self, task_name: str) -> list[TaskMiddleware]: """Get the combined global + per-task middleware list, minus any - middleware the operator has disabled for this task from the dashboard.""" + middleware the operator has disabled for this task from the dashboard. + + Cached per task to avoid a synchronous storage read on every job + dispatch. The cache is invalidated immediately on same-process disable + changes (via ``_mw_disable_version``) and expires after + ``_MW_CHAIN_TTL`` so out-of-process changes still take effect promptly. + """ + version = self._mw_disable_version + cached = self._mw_chain_cache.get(task_name) + if cached is not None: + chain, cached_version, computed_at = cached + if cached_version == version and time.monotonic() - computed_at < _MW_CHAIN_TTL: + return chain + per_task = self._task_middleware.get(task_name, []) chain = self._global_middleware + per_task try: disabled = MiddlewareDisableStore(self).get_for(task_name) # type: ignore[arg-type] except Exception: # pragma: no cover - storage read failure is non-fatal disabled = [] - if not disabled: - return chain - disabled_set = set(disabled) - return [mw for mw in chain if getattr(mw, "name", "") not in disabled_set] + if disabled: + disabled_set = set(disabled) + chain = [mw for mw in chain if getattr(mw, "name", "") not in disabled_set] + + self._mw_chain_cache[task_name] = (chain, version, time.monotonic()) + return chain def _wrap_task( self, fn: Callable, task_name: str, soft_timeout: float | None = None diff --git a/py_src/taskito/mixins/middleware_admin.py b/py_src/taskito/mixins/middleware_admin.py index dd9af800..aa5b6c0c 100644 --- a/py_src/taskito/mixins/middleware_admin.py +++ b/py_src/taskito/mixins/middleware_admin.py @@ -15,6 +15,12 @@ class QueueMiddlewareAdminMixin: _global_middleware: list[TaskMiddleware] _task_middleware: dict[str, list[TaskMiddleware]] + _mw_disable_version: int + + def _bump_mw_disable_version(self) -> None: + """Invalidate the per-task middleware-chain cache for same-process + readers after a disable-list change.""" + self._mw_disable_version += 1 # ── Discovery ────────────────────────────────────────────────── @@ -57,14 +63,20 @@ def get_disabled_middleware_for(self, task_name: str) -> list[str]: return MiddlewareDisableStore(self).get_for(task_name) # type: ignore[arg-type] def disable_middleware_for_task(self, task_name: str, mw_name: str) -> list[str]: - return MiddlewareDisableStore(self).set_disabled( # type: ignore[arg-type] + result = MiddlewareDisableStore(self).set_disabled( # type: ignore[arg-type] task_name, mw_name, disabled=True ) + self._bump_mw_disable_version() + return result def enable_middleware_for_task(self, task_name: str, mw_name: str) -> list[str]: - return MiddlewareDisableStore(self).set_disabled( # type: ignore[arg-type] + result = MiddlewareDisableStore(self).set_disabled( # type: ignore[arg-type] task_name, mw_name, disabled=False ) + self._bump_mw_disable_version() + return result def clear_middleware_disables(self, task_name: str) -> bool: - return MiddlewareDisableStore(self).clear_for(task_name) # type: ignore[arg-type] + result = MiddlewareDisableStore(self).clear_for(task_name) # type: ignore[arg-type] + self._bump_mw_disable_version() + return result diff --git a/py_src/taskito/notes.py b/py_src/taskito/notes.py index 6f9fcb4b..448043c3 100644 --- a/py_src/taskito/notes.py +++ b/py_src/taskito/notes.py @@ -74,7 +74,9 @@ def validate_and_encode_notes(notes: dict[str, Any] | None) -> str | None: _validate_value(key, value, depth=1) encoded = json.dumps(notes, sort_keys=True, ensure_ascii=False, separators=(",", ":")) - encoded_bytes = len(encoded.encode("utf-8")) + # For the common all-ASCII payload, byte length == char length, so skip the + # throwaway UTF-8 encode just to measure size. + encoded_bytes = len(encoded) if encoded.isascii() else len(encoded.encode("utf-8")) if encoded_bytes > MAX_NOTE_BYTES: raise NotesValidationError( f"encoded notes are {encoded_bytes} bytes, exceeds limit of {MAX_NOTE_BYTES} bytes" diff --git a/tests/core/test_middleware_chain_cache.py b/tests/core/test_middleware_chain_cache.py new file mode 100644 index 00000000..70622885 --- /dev/null +++ b/tests/core/test_middleware_chain_cache.py @@ -0,0 +1,54 @@ +"""Regression tests for the per-task middleware-chain cache. + +``_get_middleware_chain`` runs on every job dispatch, so it must not hit +storage (to read the dashboard disable list) on every call. The chain is +cached per task and invalidated immediately on same-process disable changes +via a version counter. +""" + +from __future__ import annotations + +from pathlib import Path + +from taskito import Queue + + +def test_middleware_chain_caches_disable_reads(tmp_path: Path) -> None: + """Repeated dispatches read the disable list from storage at most once + within the TTL window.""" + queue = Queue(db_path=str(tmp_path / "cache.db")) + + reads = {"count": 0} + real_get_setting = queue.get_setting + + def counting_get_setting(key: str) -> str | None: + reads["count"] += 1 + return real_get_setting(key) + + queue.get_setting = counting_get_setting # type: ignore[method-assign] + + for _ in range(50): + assert queue._get_middleware_chain("some.task") == [] + assert reads["count"] == 1, "disable list should be read once, then cached" + + +def test_disable_version_bump_invalidates_cache(tmp_path: Path) -> None: + """A same-process disable change re-reads the disable list on the next call.""" + queue = Queue(db_path=str(tmp_path / "cache.db")) + + reads = {"count": 0} + real_get_setting = queue.get_setting + + def counting_get_setting(key: str) -> str | None: + reads["count"] += 1 + return real_get_setting(key) + + queue.get_setting = counting_get_setting # type: ignore[method-assign] + + queue._get_middleware_chain("some.task") + assert reads["count"] == 1 + + # Simulate a dashboard disable-list change. + queue._bump_mw_disable_version() + queue._get_middleware_chain("some.task") + assert reads["count"] == 2, "version bump should invalidate the cached chain" From fef247c24c0f17f07ac2ebff240c6196a5f7c7da Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:15:20 +0530 Subject: [PATCH 09/24] fix(async): honor per-task serializer in async executor --- py_src/taskito/async_support/executor.py | 4 +- tests/worker/test_native_async.py | 54 ++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/py_src/taskito/async_support/executor.py b/py_src/taskito/async_support/executor.py index f25e13d1..b2267baf 100644 --- a/py_src/taskito/async_support/executor.py +++ b/py_src/taskito/async_support/executor.py @@ -95,8 +95,10 @@ async def _execute( completed_mw: list[Any] = [] try: - args, kwargs = cloudpickle.loads(payload_bytes) queue = self._queue_ref + # Honor the per-task serializer (matches the sync worker path); + # a hardcoded cloudpickle.loads would ignore @task(serializer=...). + args, kwargs = queue._deserialize_payload(task_name, payload_bytes) # Worker-dispatch predicate gate (raw args, pre-reconstruction). if task_name in queue._task_predicates: diff --git a/tests/worker/test_native_async.py b/tests/worker/test_native_async.py index f6fa7d1c..c90ca8ab 100644 --- a/tests/worker/test_native_async.py +++ b/tests/worker/test_native_async.py @@ -151,6 +151,7 @@ class FakeWrapper: registry: dict[str, Any] = {"test_mod.my_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -193,6 +194,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.failing_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -234,6 +236,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.cancelling_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -272,6 +275,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.flaky_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -322,6 +326,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.slow_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -377,6 +382,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.simple_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -417,6 +423,7 @@ class FakeWrapper: fake_db = "fake-conn" queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -462,6 +469,7 @@ class FakeWrapper: registry: dict[str, Any] = {"mod.ctx_task": FakeWrapper()} queue_ref = MagicMock() + queue_ref._deserialize_payload.side_effect = lambda _name, data: cloudpickle.loads(data) queue_ref._interceptor = None queue_ref._proxy_registry = None queue_ref._test_mode_active = False @@ -492,3 +500,49 @@ def test_async_concurrency_default(tmp_path: Path) -> None: """Default async_concurrency is 100.""" queue = Queue(db_path=str(tmp_path / "test.db")) assert queue._async_concurrency == 100 + + +def test_async_executor_honors_per_task_serializer(poll_until: PollUntil) -> None: + """Regression: the async executor deserializes via + ``queue._deserialize_payload`` (honoring ``@task(serializer=...)``) rather + than a hardcoded ``cloudpickle.loads``.""" + import cloudpickle + + from taskito.async_support.executor import AsyncTaskExecutor + from taskito.serializers import JsonSerializer + + sender = MagicMock() + serializer = JsonSerializer() + + async def add(x: int, y: int) -> int: + return x + y + + class FakeWrapper: + _taskito_async_fn = staticmethod(add) + + registry: dict[str, Any] = {"mod.add": FakeWrapper()} + + queue_ref = MagicMock() + # Route deserialization through a NON-cloudpickle serializer; a hardcoded + # cloudpickle.loads would raise on this JSON payload. + queue_ref._deserialize_payload.side_effect = lambda _name, data: serializer.loads(data) + queue_ref._interceptor = None + queue_ref._proxy_registry = None + queue_ref._test_mode_active = False + queue_ref._resource_runtime = None + queue_ref._task_inject_map = {} + queue_ref._task_retry_filters = {} + queue_ref._get_middleware_chain.return_value = [] + queue_ref._proxy_metrics = None + + executor = AsyncTaskExecutor(sender, registry, queue_ref, max_concurrency=10) + executor.start() + + payload = serializer.dumps(([2, 3], {})) + executor.submit_job("ser-job", "mod.add", payload, 0, 3, "default") + poll_until(lambda: sender.report_success.called, message="ser-job result not reported") + executor.stop() + + queue_ref._deserialize_payload.assert_called_once_with("mod.add", payload) + result = cloudpickle.loads(sender.report_success.call_args[0][2]) + assert result == 5 From 68490e7f791f379926705a9ca1a893a450b95ed0 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:15:25 +0530 Subject: [PATCH 10/24] perf(proxies): reuse one thread pool for reconstruction --- py_src/taskito/proxies/reconstruct.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/py_src/taskito/proxies/reconstruct.py b/py_src/taskito/proxies/reconstruct.py index 31a3d79d..8590b835 100644 --- a/py_src/taskito/proxies/reconstruct.py +++ b/py_src/taskito/proxies/reconstruct.py @@ -2,6 +2,7 @@ from __future__ import annotations +import atexit import logging import time from concurrent.futures import ( @@ -26,6 +27,12 @@ _PROXY_KEY = "__taskito_proxy__" _REF_KEY = "__taskito_ref__" +# Shared, process-wide pool for timeout-bounded proxy reconstruction. A fresh +# ThreadPoolExecutor per reconstructed argument would spawn and tear down an OS +# thread on the worker hot path; one small pool serves all reconstructions. +_RECONSTRUCT_POOL = ThreadPoolExecutor(max_workers=4, thread_name_prefix="taskito-proxy") +atexit.register(_RECONSTRUCT_POOL.shutdown, wait=False) + def reconstruct_proxies( args: tuple, @@ -170,14 +177,13 @@ def _reconstruct_one( start = time.monotonic() try: if max_timeout > 0: - with ThreadPoolExecutor(max_workers=1) as pool: - future = pool.submit(handler.reconstruct, recipe, version) - try: - obj = future.result(timeout=max_timeout) - except FuturesTimeout: - raise ProxyReconstructionError( - f"Reconstruction of '{handler_name}' timed out after {max_timeout}s" - ) from None + future = _RECONSTRUCT_POOL.submit(handler.reconstruct, recipe, version) + try: + obj = future.result(timeout=max_timeout) + except FuturesTimeout: + raise ProxyReconstructionError( + f"Reconstruction of '{handler_name}' timed out after {max_timeout}s" + ) from None else: obj = handler.reconstruct(recipe, version) except ProxyReconstructionError: From cdd139c8c51b70531d6fb5afb9ecf9bc1d27f0b0 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:15:29 +0530 Subject: [PATCH 11/24] perf(webhooks): skip lock and list copy when no subscribers --- py_src/taskito/webhooks.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/py_src/taskito/webhooks.py b/py_src/taskito/webhooks.py index 1afd2a43..b26d30e9 100644 --- a/py_src/taskito/webhooks.py +++ b/py_src/taskito/webhooks.py @@ -128,6 +128,11 @@ def add_webhook( def notify(self, event_type: EventType, payload: dict[str, Any]) -> None: """Queue an event for delivery to matching webhooks.""" + # Fast path: no subscriptions → skip the lock and list copy entirely. + # The unlocked read is a benign race; a webhook added concurrently is + # picked up by the next event. + if not self._webhooks: + return with self._lock: webhooks = list(self._webhooks) task_name = payload.get("task_name") From 7b50040d83760c7fe9ed0c442a58d51509fa34a2 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:15:34 +0530 Subject: [PATCH 12/24] perf: bind cloudpickle once in CloudpickleSerializer --- py_src/taskito/serializers.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/py_src/taskito/serializers.py b/py_src/taskito/serializers.py index 9233b809..e83b0c83 100644 --- a/py_src/taskito/serializers.py +++ b/py_src/taskito/serializers.py @@ -5,6 +5,8 @@ import json from typing import Any, Protocol, runtime_checkable +import cloudpickle + @runtime_checkable class Serializer(Protocol): @@ -23,13 +25,9 @@ class CloudpickleSerializer: """Default serializer using cloudpickle (handles lambdas, closures, etc.).""" def dumps(self, obj: Any) -> bytes: - import cloudpickle - return cloudpickle.dumps(obj) # type: ignore[no-any-return] def loads(self, data: bytes) -> Any: - import cloudpickle - return cloudpickle.loads(data) From c4e94daa393331fd69384d448de2267923923998 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:15:38 +0530 Subject: [PATCH 13/24] perf(interception): cache type resolution by object type --- py_src/taskito/interception/registry.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/py_src/taskito/interception/registry.py b/py_src/taskito/interception/registry.py index cffde9d0..21d53400 100644 --- a/py_src/taskito/interception/registry.py +++ b/py_src/taskito/interception/registry.py @@ -8,6 +8,9 @@ from taskito.interception.strategy import Strategy +# Sentinel distinguishing "not cached" from "cached as no-match (None)". +_UNCACHED = object() + @dataclass class RegistryEntry: @@ -36,6 +39,9 @@ class TypeRegistry: def __init__(self) -> None: self._entries: list[RegistryEntry] = [] self._sorted = False + # Memoize resolution per concrete object type — the result for a given + # type is stable until the registry changes. Cleared on register(). + self._cache: dict[type, RegistryEntry | None] = {} def register( self, @@ -68,20 +74,30 @@ def register( ) self._entries.append(entry) self._sorted = False + self._cache.clear() def resolve(self, obj: Any) -> RegistryEntry | None: """Find the highest-priority entry matching ``obj``.""" if not self._sorted: self._entries.sort(key=lambda e: e.priority, reverse=True) self._sorted = True + + obj_type = type(obj) + cached = self._cache.get(obj_type, _UNCACHED) + if cached is not _UNCACHED: + return cached # type: ignore[return-value] + + result: RegistryEntry | None = None for entry in self._entries: try: if isinstance(obj, entry.types): - return entry + result = entry + break except TypeError: # Entry contains a non-type — skip it continue - return None + self._cache[obj_type] = result + return result def __len__(self) -> int: return len(self._entries) From 4c8a149637c6e41e92b19b9a0bcaf1bcf748179a Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:38:23 +0530 Subject: [PATCH 14/24] perf(storage): skip dep lookup via has_deps; filter reap/purge in SQL --- crates/taskito-core/src/job.rs | 7 + .../src/storage/diesel_common/archival.rs | 2 + .../src/storage/diesel_common/dead_letter.rs | 1 + .../src/storage/diesel_common/jobs.rs | 174 ++++++++---------- .../src/storage/diesel_common/migrations.rs | 23 ++- crates/taskito-core/src/storage/models.rs | 2 + .../taskito-core/src/storage/postgres/mod.rs | 3 + crates/taskito-core/src/storage/schema.rs | 1 + crates/taskito-core/src/storage/sqlite/mod.rs | 3 + .../taskito-core/src/storage/sqlite/tests.rs | 99 ++++++++++ 10 files changed, 221 insertions(+), 94 deletions(-) diff --git a/crates/taskito-core/src/job.rs b/crates/taskito-core/src/job.rs index b02df6a8..ca15c411 100644 --- a/crates/taskito-core/src/job.rs +++ b/crates/taskito-core/src/job.rs @@ -89,6 +89,10 @@ pub struct Job { pub expires_at: Option, pub result_ttl_ms: Option, pub namespace: Option, + /// True when the job was enqueued with at least one dependency. Lets the + /// scheduler skip the dependency lookup entirely for the common case. + #[serde(default)] + pub has_deps: bool, } impl From for Job { @@ -117,6 +121,7 @@ impl From for Job { expires_at: row.expires_at, result_ttl_ms: row.result_ttl_ms, namespace: row.namespace, + has_deps: row.has_deps, } } } @@ -143,6 +148,7 @@ pub struct NewJob { impl NewJob { pub fn into_job(self) -> Job { let now = now_millis(); + let has_deps = !self.depends_on.is_empty(); Job { id: Uuid::now_v7().to_string(), queue: self.queue, @@ -167,6 +173,7 @@ impl NewJob { expires_at: self.expires_at, result_ttl_ms: self.result_ttl_ms, namespace: self.namespace, + has_deps, } } } diff --git a/crates/taskito-core/src/storage/diesel_common/archival.rs b/crates/taskito-core/src/storage/diesel_common/archival.rs index 72d50d26..53105a59 100644 --- a/crates/taskito-core/src/storage/diesel_common/archival.rs +++ b/crates/taskito-core/src/storage/diesel_common/archival.rs @@ -42,6 +42,8 @@ macro_rules! impl_diesel_archival_ops { expires_at: row.expires_at, result_ttl_ms: row.result_ttl_ms, namespace: row.namespace, + // Archived jobs are terminal and never re-dequeued. + has_deps: false, }) .collect()) } diff --git a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs index 495da038..c33f082c 100644 --- a/crates/taskito-core/src/storage/diesel_common/dead_letter.rs +++ b/crates/taskito-core/src/storage/diesel_common/dead_letter.rs @@ -137,6 +137,7 @@ macro_rules! impl_diesel_dead_letter_ops { expires_at: job.expires_at, result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), + has_deps: job.has_deps, }; diesel::insert_into(jobs::table) diff --git a/crates/taskito-core/src/storage/diesel_common/jobs.rs b/crates/taskito-core/src/storage/diesel_common/jobs.rs index f0b6e904..cf6bce14 100644 --- a/crates/taskito-core/src/storage/diesel_common/jobs.rs +++ b/crates/taskito-core/src/storage/diesel_common/jobs.rs @@ -13,27 +13,22 @@ macro_rules! impl_diesel_job_ops { return Ok(()); } + diesel::delete(job_errors::table.filter(job_errors::job_id.eq_any(job_ids))) + .execute(conn)?; + diesel::delete(task_logs::table.filter(task_logs::job_id.eq_any(job_ids))) + .execute(conn)?; + diesel::delete(task_metrics::table.filter(task_metrics::job_id.eq_any(job_ids))) + .execute(conn)?; diesel::delete( - job_errors::table.filter(job_errors::job_id.eq_any(job_ids)), - ) - .execute(conn)?; - diesel::delete( - task_logs::table.filter(task_logs::job_id.eq_any(job_ids)), - ) - .execute(conn)?; - diesel::delete( - task_metrics::table.filter(task_metrics::job_id.eq_any(job_ids)), + job_dependencies::table.filter( + job_dependencies::job_id + .eq_any(job_ids) + .or(job_dependencies::depends_on_job_id.eq_any(job_ids)), + ), ) .execute(conn)?; - diesel::delete(job_dependencies::table.filter( - job_dependencies::job_id - .eq_any(job_ids) - .or(job_dependencies::depends_on_job_id.eq_any(job_ids)), - )) - .execute(conn)?; diesel::delete( - replay_history::table - .filter(replay_history::original_job_id.eq_any(job_ids)), + replay_history::table.filter(replay_history::original_job_id.eq_any(job_ids)), ) .execute(conn)?; @@ -108,6 +103,7 @@ macro_rules! impl_diesel_job_ops { expires_at: job.expires_at, result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), + has_deps: job.has_deps, }; diesel::insert_into(jobs::table) @@ -129,11 +125,9 @@ macro_rules! impl_diesel_job_ops { Ok(()) }) .map_err(|e| match e { - diesel::result::Error::RollbackTransaction => { - QueueError::DependencyNotFound( - "dependency not found or already dead/cancelled".to_string(), - ) - } + diesel::result::Error::RollbackTransaction => QueueError::DependencyNotFound( + "dependency not found or already dead/cancelled".to_string(), + ), other => QueueError::Storage(other), })?; @@ -142,12 +136,19 @@ macro_rules! impl_diesel_job_ops { /// Enqueue multiple jobs in a single transaction. pub fn enqueue_batch(&self, new_jobs: Vec) -> Result> { + // Bound rows-per-INSERT so the bound-parameter count stays + // under SQLite's 999 limit (NewJobRow has ~19 columns; + // 50 * 19 < 999). Postgres tolerates far more, but one shared + // chunk size keeps the macro-generated code identical. + const BATCH_INSERT_CHUNK: usize = 50; + let mut conn = self.conn()?; let jobs: Vec = new_jobs.into_iter().map(|nj| nj.into_job()).collect(); conn.transaction(|conn| { - for job in &jobs { - let row = NewJobRow { + let rows: Vec = jobs + .iter() + .map(|job| NewJobRow { id: &job.id, queue: &job.queue, task_name: &job.task_name, @@ -166,10 +167,15 @@ macro_rules! impl_diesel_job_ops { expires_at: job.expires_at, result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), - }; + has_deps: job.has_deps, + }) + .collect(); + // One multi-row INSERT per chunk instead of N single-row + // INSERTs — far fewer round trips / statement executions. + for chunk in rows.chunks(BATCH_INSERT_CHUNK) { diesel::insert_into(jobs::table) - .values(&row) + .values(chunk) .execute(conn)?; } Ok(jobs) @@ -188,10 +194,8 @@ macro_rules! impl_diesel_job_ops { let existing: Option = jobs::table .filter(jobs::unique_key.eq(uk)) .filter( - jobs::status.eq_any([ - JobStatus::Pending as i32, - JobStatus::Running as i32, - ]), + jobs::status + .eq_any([JobStatus::Pending as i32, JobStatus::Running as i32]), ) .select(JobRow::as_select()) .first(conn) @@ -221,6 +225,7 @@ macro_rules! impl_diesel_job_ops { expires_at: job.expires_at, result_ttl_ms: job.result_ttl_ms, namespace: job.namespace.as_deref(), + has_deps: job.has_deps, }; diesel::insert_into(jobs::table) @@ -251,17 +256,16 @@ macro_rules! impl_diesel_job_ops { )) => { if let Some(ref uk) = job.unique_key { let mut conn = self.conn()?; - let existing: Option = jobs::table - .filter(jobs::unique_key.eq(uk)) - .filter( - jobs::status.eq_any([ + let existing: Option = + jobs::table + .filter(jobs::unique_key.eq(uk)) + .filter(jobs::status.eq_any([ JobStatus::Pending as i32, JobStatus::Running as i32, - ]), - ) - .select(JobRow::as_select()) - .first(&mut conn) - .optional()?; + ])) + .select(JobRow::as_select()) + .first(&mut conn) + .optional()?; if let Some(row) = existing { return Ok(Job::from(row)); } @@ -275,7 +279,12 @@ macro_rules! impl_diesel_job_ops { /// Atomically dequeue the highest-priority ready job from the given queue. /// Skips expired jobs. When `namespace` is `Some`, only jobs in that /// namespace are considered; when `None`, only jobs with no namespace. - pub fn dequeue(&self, queue_name: &str, now: i64, namespace: Option<&str>) -> Result> { + pub fn dequeue( + &self, + queue_name: &str, + now: i64, + namespace: Option<&str>, + ) -> Result> { let mut conn = self.conn()?; conn.transaction(|conn| { @@ -293,9 +302,7 @@ macro_rules! impl_diesel_job_ops { query = query.filter(jobs::namespace.is_null()); } - let candidates: Vec = query - .select(JobRow::as_select()) - .load(conn)?; + let candidates: Vec = query.select(JobRow::as_select()).load(conn)?; for row in candidates { // Skip expired jobs @@ -314,7 +321,9 @@ macro_rules! impl_diesel_job_ops { } } - if !Self::deps_satisfied(conn, &row.id)? { + // Common case: jobs with no dependencies skip the + // job_dependencies lookup entirely. + if row.has_deps && !Self::deps_satisfied(conn, &row.id)? { continue; } @@ -340,7 +349,12 @@ macro_rules! impl_diesel_job_ops { } /// Dequeue from multiple queues, checking each in order. - pub fn dequeue_from(&self, queues: &[String], now: i64, namespace: Option<&str>) -> Result> { + pub fn dequeue_from( + &self, + queues: &[String], + now: i64, + namespace: Option<&str>, + ) -> Result> { for queue_name in queues { if let Some(job) = self.dequeue(queue_name, now, namespace)? { return Ok(Some(job)); @@ -764,10 +778,8 @@ macro_rules! impl_diesel_job_ops { Self::delete_job_children(conn, &job_ids)?; - let affected = diesel::delete( - jobs::table.filter(jobs::id.eq_any(&job_ids)), - ) - .execute(conn)?; + let affected = diesel::delete(jobs::table.filter(jobs::id.eq_any(&job_ids))) + .execute(conn)?; Ok(affected as u64) }) @@ -786,35 +798,27 @@ macro_rules! impl_diesel_job_ops { .select(jobs::id) .load(conn)?; - let rows_with_ttl: Vec = jobs::table + // Push the per-job `completed_at + result_ttl_ms < now` + // check into SQL and select only the id — avoids loading + // full rows (payload + result blobs) just to filter them. + let per_job_ids: Vec = jobs::table .filter(jobs::status.eq(JobStatus::Complete as i32)) .filter(jobs::result_ttl_ms.is_not_null()) - .select(JobRow::as_select()) + .filter(jobs::completed_at.is_not_null()) + .filter( + (jobs::completed_at.assume_not_null() + + jobs::result_ttl_ms.assume_not_null()) + .lt(now), + ) + .select(jobs::id) .load(conn)?; - let per_job_ids: Vec = rows_with_ttl - .into_iter() - .filter(|row| { - matches!( - (row.completed_at, row.result_ttl_ms), - (Some(completed), Some(ttl)) - if completed - .checked_add(ttl) - .is_some_and(|expiry| expiry < now) - ) - }) - .map(|row| row.id) - .collect(); - - let all_ids: Vec = - global_ids.into_iter().chain(per_job_ids).collect(); + let all_ids: Vec = global_ids.into_iter().chain(per_job_ids).collect(); Self::delete_job_children(conn, &all_ids)?; - let affected = diesel::delete( - jobs::table.filter(jobs::id.eq_any(&all_ids)), - ) - .execute(conn)?; + let affected = diesel::delete(jobs::table.filter(jobs::id.eq_any(&all_ids))) + .execute(conn)?; Ok(affected as u64) }) @@ -824,37 +828,21 @@ macro_rules! impl_diesel_job_ops { pub fn reap_stale_jobs(&self, now: i64) -> Result> { let mut conn = self.conn()?; + // Push the `started_at + timeout_ms < now` deadline into SQL so + // only genuinely-stale rows (and their payload blobs) are read, + // instead of every running job. let rows: Vec = jobs::table .filter(jobs::status.eq(JobStatus::Running as i32)) .filter(jobs::started_at.is_not_null()) + .filter((jobs::started_at.assume_not_null() + jobs::timeout_ms).lt(now)) .select(JobRow::as_select()) .load(&mut conn)?; - let stale: Vec = rows - .into_iter() - .filter(|r| { - if let Some(started) = r.started_at { - match started.checked_add(r.timeout_ms) { - Some(deadline) => deadline < now, - None => true, - } - } else { - false - } - }) - .map(Job::from) - .collect(); - - Ok(stale) + Ok(rows.into_iter().map(Job::from).collect()) } /// Record an error for a job attempt. - pub fn record_error( - &self, - job_id: &str, - attempt: i32, - error: &str, - ) -> Result<()> { + pub fn record_error(&self, job_id: &str, attempt: i32, error: &str) -> Result<()> { let mut conn = self.conn()?; let id = uuid::Uuid::now_v7().to_string(); let now = now_millis(); diff --git a/crates/taskito-core/src/storage/diesel_common/migrations.rs b/crates/taskito-core/src/storage/diesel_common/migrations.rs index 5f228d5e..ddc17e2b 100644 --- a/crates/taskito-core/src/storage/diesel_common/migrations.rs +++ b/crates/taskito-core/src/storage/diesel_common/migrations.rs @@ -91,7 +91,8 @@ pub fn create_tables(d: &Dialect) -> Vec { notes TEXT, cancel_requested INTEGER NOT NULL DEFAULT 0, expires_at {bi}, - result_ttl_ms {bi} + result_ttl_ms {bi}, + has_deps {bool_false} )" ), format!( @@ -296,6 +297,7 @@ pub fn alter_statements(d: &Dialect) -> Vec { let bi = d.big_int; let real = d.real; let ife = d.alter_if_not_exists; + let bool_false = d.boolean_default_false; vec![ // jobs ── cancel_requested / expires_at / result_ttl_ms @@ -345,5 +347,24 @@ pub fn alter_statements(d: &Dialect) -> Vec { format!("ALTER TABLE jobs ADD COLUMN {ife}notes TEXT"), format!("ALTER TABLE dead_letter ADD COLUMN {ife}notes TEXT"), format!("ALTER TABLE archived_jobs ADD COLUMN {ife}notes TEXT"), + // has_deps fast-path flag: lets dequeue skip the dependency lookup for + // the common no-dependency job. Defaults false; backfilled below. + format!("ALTER TABLE jobs ADD COLUMN {ife}has_deps {bool_false}"), + ] +} + +/// One-time data backfills run after [`alter_statements`]. Unlike the schema +/// alters these are idempotent UPDATEs guarded so they only write rows whose +/// value actually differs — after the first run they match nothing. +/// +/// Backfills `has_deps` for pre-existing pending jobs: a database upgraded +/// while pending jobs already had dependencies would otherwise default them to +/// `has_deps = false` and let dequeue skip the dependency check. The +/// `id IN (subquery)` expression yields a boolean in Postgres and 0/1 in +/// SQLite, so the same statement works on both backends. +pub fn backfill_statements() -> &'static [&'static str] { + &[ + "UPDATE jobs SET has_deps = (id IN (SELECT job_id FROM job_dependencies)) \ + WHERE status = 0 AND has_deps <> (id IN (SELECT job_id FROM job_dependencies))", ] } diff --git a/crates/taskito-core/src/storage/models.rs b/crates/taskito-core/src/storage/models.rs index e0ecd804..e45c13d4 100644 --- a/crates/taskito-core/src/storage/models.rs +++ b/crates/taskito-core/src/storage/models.rs @@ -34,6 +34,7 @@ pub struct JobRow { pub expires_at: Option, pub result_ttl_ms: Option, pub namespace: Option, + pub has_deps: bool, } /// Insertable struct for creating new jobs. @@ -58,6 +59,7 @@ pub struct NewJobRow<'a> { pub expires_at: Option, pub result_ttl_ms: Option, pub namespace: Option<&'a str>, + pub has_deps: bool, } /// A row in the `dead_letter` table. diff --git a/crates/taskito-core/src/storage/postgres/mod.rs b/crates/taskito-core/src/storage/postgres/mod.rs index 9440445c..a3a1c317 100644 --- a/crates/taskito-core/src/storage/postgres/mod.rs +++ b/crates/taskito-core/src/storage/postgres/mod.rs @@ -124,6 +124,9 @@ impl PostgresStorage { for sql in common_migrations::alter_statements(&common_migrations::POSTGRES) { migration_alter(&mut conn, &sql); } + for sql in common_migrations::backfill_statements() { + migration_alter(&mut conn, sql); + } Ok(()) } diff --git a/crates/taskito-core/src/storage/schema.rs b/crates/taskito-core/src/storage/schema.rs index 12be5dad..27c85f72 100644 --- a/crates/taskito-core/src/storage/schema.rs +++ b/crates/taskito-core/src/storage/schema.rs @@ -23,6 +23,7 @@ diesel::table! { expires_at -> Nullable, result_ttl_ms -> Nullable, namespace -> Nullable, + has_deps -> Bool, } } diff --git a/crates/taskito-core/src/storage/sqlite/mod.rs b/crates/taskito-core/src/storage/sqlite/mod.rs index 8579e6a0..98585c35 100644 --- a/crates/taskito-core/src/storage/sqlite/mod.rs +++ b/crates/taskito-core/src/storage/sqlite/mod.rs @@ -115,6 +115,9 @@ impl SqliteStorage { for sql in common_migrations::alter_statements(&common_migrations::SQLITE) { migration_alter(&mut conn, &sql); } + for sql in common_migrations::backfill_statements() { + migration_alter(&mut conn, sql); + } Ok(()) } diff --git a/crates/taskito-core/src/storage/sqlite/tests.rs b/crates/taskito-core/src/storage/sqlite/tests.rs index 0dba4b13..f004f0ef 100644 --- a/crates/taskito-core/src/storage/sqlite/tests.rs +++ b/crates/taskito-core/src/storage/sqlite/tests.rs @@ -559,3 +559,102 @@ fn test_setting_preserves_unicode_and_json() { Some(payload.to_string()) ); } + +#[test] +fn test_reap_stale_jobs_only_returns_expired() { + let storage = test_storage(); + let t0 = now_millis(); + + let mut short = make_job("short_timeout"); + short.timeout_ms = 1; + storage.enqueue(short).unwrap(); + let mut long = make_job("long_timeout"); + long.timeout_ms = 300_000; + storage.enqueue(long).unwrap(); + + // Run both: started_at = t0 for each. + storage.dequeue("default", t0, None).unwrap(); + storage.dequeue("default", t0, None).unwrap(); + + // Well past the short job's deadline (t0 + 1) but before the long one's. + let stale = storage.reap_stale_jobs(t0 + 1_000).unwrap(); + assert_eq!(stale.len(), 1); + assert_eq!(stale[0].task_name, "short_timeout"); +} + +#[test] +fn test_has_deps_flag_gates_dequeue() { + let storage = test_storage(); + + // No-dependency job: has_deps is false. + let plain = storage.enqueue(make_job("plain")).unwrap(); + assert!(!plain.has_deps); + + // Dependency target and dependent child, each on its own queue so the + // dequeue calls are unambiguous. + let mut target = make_job("target"); + target.queue = "qt".to_string(); + let target = storage.enqueue(target).unwrap(); + let mut child = make_job("child"); + child.queue = "q2".to_string(); + child.depends_on = vec![target.id.clone()]; + let child = storage.enqueue(child).unwrap(); + assert!(child.has_deps); + + let t0 = now_millis(); + // Blocked while the dependency is incomplete. + assert!(storage.dequeue("q2", t0, None).unwrap().is_none()); + + // Complete the dependency, then the child becomes dequeueable. + storage.dequeue("qt", t0, None).unwrap(); + storage.complete(&target.id, None).unwrap(); + let got = storage.dequeue("q2", t0, None).unwrap(); + assert_eq!(got.map(|j| j.id), Some(child.id)); +} + +#[test] +fn test_enqueue_batch_crosses_chunk_boundary() { + let storage = test_storage(); + // More than the 50-row insert chunk so multiple multi-row INSERTs run. + let count = 120; + let jobs: Vec = (0..count) + .map(|i| make_job(&format!("batch_{i}"))) + .collect(); + + let result = storage.enqueue_batch(jobs).unwrap(); + assert_eq!(result.len(), count); + assert_eq!(storage.stats().unwrap().pending, count as i64); +} + +#[test] +fn test_purge_completed_respects_per_job_ttl() { + let storage = test_storage(); + + // Completed with a 1ms TTL — should be purged once that elapses. Each on + // its own queue so the dequeue/complete pair is unambiguous. + let mut expired = make_job("ttl_expired"); + expired.queue = "qa".to_string(); + expired.result_ttl_ms = Some(1); + let expired = storage.enqueue(expired).unwrap(); + // Completed with a far-future TTL — should survive. + let mut kept = make_job("ttl_kept"); + kept.queue = "qb".to_string(); + kept.result_ttl_ms = Some(3_600_000); + let kept = storage.enqueue(kept).unwrap(); + + // Capture `now` after enqueue so both jobs' scheduled_at are eligible. + let now = now_millis(); + storage.dequeue("qa", now, None).unwrap(); + storage.dequeue("qb", now, None).unwrap(); + storage.complete(&expired.id, None).unwrap(); + storage.complete(&kept.id, None).unwrap(); + + // Ensure the 1ms TTL has elapsed relative to purge's `now`. + std::thread::sleep(std::time::Duration::from_millis(5)); + + // global_cutoff = 0 so only the per-job TTL path can match. + storage.purge_completed_with_ttl(0).unwrap(); + + assert!(storage.get_job(&expired.id).unwrap().is_none()); + assert!(storage.get_job(&kept.id).unwrap().is_some()); +} From e6a6715b926b86476f992e51517e19e1738f7f79 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:38:35 +0530 Subject: [PATCH 15/24] perf(redis): batch dequeue loads with MGET, reuse connection --- .../src/storage/redis_backend/jobs/dequeue.rs | 45 ++++++++++++------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs b/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs index adaea51e..86ec6acc 100644 --- a/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs +++ b/crates/taskito-core/src/storage/redis_backend/jobs/dequeue.rs @@ -21,9 +21,16 @@ impl RedisStorage { .zrangebyscore_limit(&queue_key, "-inf", "+inf", 0, 100) .map_err(map_err)?; - for job_id in candidates { - let job_key = self.key(&["job", &job_id]); - let data: Option = conn.get(&job_key).map_err(map_err)?; + if candidates.is_empty() { + return Ok(None); + } + + // Batch-load every candidate's JSON in one MGET instead of one GET per + // candidate. + let job_keys: Vec = candidates.iter().map(|id| self.key(&["job", id])).collect(); + let blobs: Vec> = conn.mget(&job_keys).map_err(map_err)?; + + for (job_id, data) in candidates.into_iter().zip(blobs) { let data = match data { Some(d) => d, None => { @@ -64,24 +71,28 @@ impl RedisStorage { } } - // Check dependencies - let deps_key = self.key(&["job", &job_id, "depends_on"]); - let dep_ids: Vec = conn.smembers(&deps_key).map_err(map_err)?; - if !dep_ids.is_empty() { - let mut all_complete = true; - for dep_id in &dep_ids { - if let Some(dep_job) = self.get_job(dep_id)? { - if dep_job.status != JobStatus::Complete { + // Check dependencies — only for jobs that actually have them, and + // resolve them on the existing connection rather than opening a new + // one per dependency. + if job.has_deps { + let deps_key = self.key(&["job", &job_id, "depends_on"]); + let dep_ids: Vec = conn.smembers(&deps_key).map_err(map_err)?; + if !dep_ids.is_empty() { + let mut all_complete = true; + for dep_id in &dep_ids { + if let Some(dep_job) = self.load_job(&mut conn, dep_id)? { + if dep_job.status != JobStatus::Complete { + all_complete = false; + break; + } + } else { all_complete = false; break; } - } else { - all_complete = false; - break; } - } - if !all_complete { - continue; + if !all_complete { + continue; + } } } From 0b682c0ba962735e80b4e691fcb0dd39fbe1d8c9 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:38:39 +0530 Subject: [PATCH 16/24] perf(worker): collapse double GIL acquire on failure path --- crates/taskito-python/src/async_worker.rs | 17 ++++++++++------- crates/taskito-python/src/py_worker.rs | 18 ++++++++++-------- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/crates/taskito-python/src/async_worker.rs b/crates/taskito-python/src/async_worker.rs index a7bfeff5..c005c816 100644 --- a/crates/taskito-python/src/async_worker.rs +++ b/crates/taskito-python/src/async_worker.rs @@ -91,11 +91,18 @@ impl WorkerDispatcher for AsyncWorkerPool { } } Err(e) => { - let (error_msg, is_cancelled, exc_class_name) = Python::with_gil(|py| { + // Single GIL acquisition: extract the error info and the + // retry decision together instead of taking the GIL twice. + let (error_msg, is_cancelled, should_retry) = Python::with_gil(|py| { let msg = format_python_error(py, &e); let cancelled = is_cancelled_error(py, &e); - let class_name = get_exception_class_name(py, &e); - (msg, cancelled, class_name) + let retry = if cancelled { + false + } else { + let class_name = get_exception_class_name(py, &e); + check_should_retry(py, &filters, &task_name, &class_name, &e) + }; + (msg, cancelled, retry) }); if is_cancelled { @@ -105,10 +112,6 @@ impl WorkerDispatcher for AsyncWorkerPool { wall_time_ns, } } else { - let should_retry = Python::with_gil(|py| { - check_should_retry(py, &filters, &task_name, &exc_class_name, &e) - }); - log::error!("[taskito] Task {task_name}[{job_id}] failed: {error_msg}"); JobResult::Failure { job_id, diff --git a/crates/taskito-python/src/py_worker.rs b/crates/taskito-python/src/py_worker.rs index 864c406f..c8f1a3d8 100644 --- a/crates/taskito-python/src/py_worker.rs +++ b/crates/taskito-python/src/py_worker.rs @@ -85,11 +85,18 @@ fn worker_loop( } } Err(e) => { - let (error_msg, is_cancelled, exc_class_name) = Python::with_gil(|py| { + // Single GIL acquisition: extract the error info and the retry + // decision together instead of taking the GIL twice. + let (error_msg, is_cancelled, should_retry) = Python::with_gil(|py| { let msg = format_python_error(py, &e); let cancelled = is_cancelled_error(py, &e); - let class_name = get_exception_class_name(py, &e); - (msg, cancelled, class_name) + let retry = if cancelled { + false + } else { + let class_name = get_exception_class_name(py, &e); + check_should_retry(py, &retry_filters, &task_name, &class_name, &e) + }; + (msg, cancelled, retry) }); if is_cancelled { @@ -99,11 +106,6 @@ fn worker_loop( wall_time_ns, } } else { - // Determine should_retry based on retry_on/dont_retry_on filters - let should_retry = Python::with_gil(|py| { - check_should_retry(py, &retry_filters, &task_name, &exc_class_name, &e) - }); - log::error!("[taskito] Task {task_name}[{job_id}] failed: {error_msg}"); JobResult::Failure { job_id, From b9998302d12cb311bfb37b9e42a7a6a494ebad39 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 14:38:43 +0530 Subject: [PATCH 17/24] perf(scheduler): borrow queue list when nothing is paused --- crates/taskito-core/src/scheduler/poller.rs | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/crates/taskito-core/src/scheduler/poller.rs b/crates/taskito-core/src/scheduler/poller.rs index 666a6c1a..0023665d 100644 --- a/crates/taskito-core/src/scheduler/poller.rs +++ b/crates/taskito-core/src/scheduler/poller.rs @@ -94,8 +94,10 @@ impl Scheduler { } /// Snapshot the queue list with paused queues filtered out. The paused - /// list is cached for 1s to avoid hammering storage on every tick. - fn active_queues(&self) -> Vec { + /// list is cached for 1s to avoid hammering storage on every tick. Borrows + /// the queue list directly in the common case (nothing paused), allocating + /// only when a filtered copy is actually needed. + fn active_queues(&self) -> std::borrow::Cow<'_, [String]> { let mut cache = self.paused_cache.lock().unwrap_or_else(|poisoned| { warn!("paused_cache mutex was poisoned, recovering"); poisoned.into_inner() @@ -110,13 +112,15 @@ impl Scheduler { cache.1 = std::time::Instant::now(); } if cache.0.is_empty() { - self.queues.clone() + std::borrow::Cow::Borrowed(&self.queues) } else { - self.queues - .iter() - .filter(|q| !cache.0.contains(*q)) - .cloned() - .collect() + std::borrow::Cow::Owned( + self.queues + .iter() + .filter(|q| !cache.0.contains(*q)) + .cloned() + .collect(), + ) } } From 8ec41d28ea4a5c682c7e1f2613f879e8e51f448e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:21:12 +0530 Subject: [PATCH 18/24] fix(storage): fail fast on has_deps backfill errors --- crates/taskito-core/src/storage/postgres/mod.rs | 4 +++- crates/taskito-core/src/storage/sqlite/mod.rs | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/taskito-core/src/storage/postgres/mod.rs b/crates/taskito-core/src/storage/postgres/mod.rs index a3a1c317..2ad6f68a 100644 --- a/crates/taskito-core/src/storage/postgres/mod.rs +++ b/crates/taskito-core/src/storage/postgres/mod.rs @@ -124,8 +124,10 @@ impl PostgresStorage { for sql in common_migrations::alter_statements(&common_migrations::POSTGRES) { migration_alter(&mut conn, &sql); } + // Data backfills must fail loudly — a swallowed failure would leave + // has_deps wrong and let dequeue bypass dependency enforcement. for sql in common_migrations::backfill_statements() { - migration_alter(&mut conn, sql); + diesel::sql_query(*sql).execute(&mut conn)?; } Ok(()) diff --git a/crates/taskito-core/src/storage/sqlite/mod.rs b/crates/taskito-core/src/storage/sqlite/mod.rs index 98585c35..1a37a6c8 100644 --- a/crates/taskito-core/src/storage/sqlite/mod.rs +++ b/crates/taskito-core/src/storage/sqlite/mod.rs @@ -115,8 +115,10 @@ impl SqliteStorage { for sql in common_migrations::alter_statements(&common_migrations::SQLITE) { migration_alter(&mut conn, &sql); } + // Data backfills must fail loudly — a swallowed failure would leave + // has_deps wrong and let dequeue bypass dependency enforcement. for sql in common_migrations::backfill_statements() { - migration_alter(&mut conn, sql); + diesel::sql_query(*sql).execute(&mut conn)?; } Ok(()) From 576f05d47365d52ce2778c65a75b1ed8b80c3ba7 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:21:18 +0530 Subject: [PATCH 19/24] refactor: extract shared middleware_key/class_path helper --- py_src/taskito/dashboard/handlers/middleware.py | 7 ++++--- py_src/taskito/middleware.py | 12 ++++++++++++ py_src/taskito/mixins/decorators.py | 5 ++++- py_src/taskito/mixins/middleware_admin.py | 9 +++++---- 4 files changed, 25 insertions(+), 8 deletions(-) diff --git a/py_src/taskito/dashboard/handlers/middleware.py b/py_src/taskito/dashboard/handlers/middleware.py index e0fd85b3..10ff60ce 100644 --- a/py_src/taskito/dashboard/handlers/middleware.py +++ b/py_src/taskito/dashboard/handlers/middleware.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any from taskito.dashboard.errors import _BadRequest, _NotFound +from taskito.middleware import middleware_class_path, middleware_key if TYPE_CHECKING: from taskito.app import Queue @@ -24,13 +25,13 @@ def handle_get_task_middleware(queue: Queue, _qs: dict, task_name: str) -> dict[ # can render every toggle. base_chain = queue._global_middleware + queue._task_middleware.get(task_name, []) entries: list[dict[str, Any]] = [] - chain_names = {getattr(mw, "name", "") for mw in chain} + chain_names = {middleware_key(mw) for mw in chain} for mw in base_chain: - name = getattr(mw, "name", "") or f"{type(mw).__module__}.{type(mw).__qualname__}" + name = middleware_key(mw) entries.append( { "name": name, - "class_path": f"{type(mw).__module__}.{type(mw).__qualname__}", + "class_path": middleware_class_path(mw), "disabled": name in disabled, "effective": name in chain_names, } diff --git a/py_src/taskito/middleware.py b/py_src/taskito/middleware.py index 077ff33f..6b617e4c 100644 --- a/py_src/taskito/middleware.py +++ b/py_src/taskito/middleware.py @@ -121,3 +121,15 @@ def on_timeout(self, ctx: JobContext) -> None: def on_cancel(self, ctx: JobContext) -> None: """Called when a job is cancelled during execution.""" + + +def middleware_class_path(mw: TaskMiddleware) -> str: + """Fully-qualified class path of a middleware instance.""" + return f"{type(mw).__module__}.{type(mw).__qualname__}" + + +def middleware_key(mw: TaskMiddleware) -> str: + """Stable identifier a middleware is keyed by in the dashboard disable + list: its ``name`` if set, otherwise its class path. This is the single + source of truth shared by chain filtering and admin discovery.""" + return getattr(mw, "name", "") or middleware_class_path(mw) diff --git a/py_src/taskito/mixins/decorators.py b/py_src/taskito/mixins/decorators.py index 8db59335..11f88311 100644 --- a/py_src/taskito/mixins/decorators.py +++ b/py_src/taskito/mixins/decorators.py @@ -27,6 +27,7 @@ from taskito.exceptions import TaskCancelledError from taskito.inject import Inject, _InjectAlias from taskito.interception.reconstruct import reconstruct_args +from taskito.middleware import middleware_key from taskito.predicates.core import coerce_predicate from taskito.proxies import cleanup_proxies, reconstruct_proxies from taskito.task import TaskWrapper @@ -155,7 +156,9 @@ def _get_middleware_chain(self, task_name: str) -> list[TaskMiddleware]: disabled = [] if disabled: disabled_set = set(disabled) - chain = [mw for mw in chain if getattr(mw, "name", "") not in disabled_set] + # ``middleware_key`` matches admin discovery's keying (name, else + # class path) so a dashboard disable on an unnamed middleware works. + chain = [mw for mw in chain if middleware_key(mw) not in disabled_set] self._mw_chain_cache[task_name] = (chain, version, time.monotonic()) return chain diff --git a/py_src/taskito/mixins/middleware_admin.py b/py_src/taskito/mixins/middleware_admin.py index aa5b6c0c..eed1658b 100644 --- a/py_src/taskito/mixins/middleware_admin.py +++ b/py_src/taskito/mixins/middleware_admin.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any from taskito.dashboard.middleware_store import MiddlewareDisableStore +from taskito.middleware import middleware_class_path, middleware_key if TYPE_CHECKING: from taskito.middleware import TaskMiddleware @@ -30,23 +31,23 @@ def list_middleware(self) -> list[dict[str, Any]]: ``name`` is the value the disable list keys on.""" seen: dict[str, dict[str, Any]] = {} for mw in self._global_middleware: - name = getattr(mw, "name", "") or f"{type(mw).__module__}.{type(mw).__qualname__}" + name = middleware_key(mw) seen.setdefault( name, { "name": name, - "class_path": f"{type(mw).__module__}.{type(mw).__qualname__}", + "class_path": middleware_class_path(mw), "scopes": [], }, )["scopes"].append({"kind": "global"}) for task_name, mws in self._task_middleware.items(): for mw in mws: - name = getattr(mw, "name", "") or f"{type(mw).__module__}.{type(mw).__qualname__}" + name = middleware_key(mw) entry = seen.setdefault( name, { "name": name, - "class_path": f"{type(mw).__module__}.{type(mw).__qualname__}", + "class_path": middleware_class_path(mw), "scopes": [], }, ) From f4bbc42295bef1d318295830794ce7647447cd9a Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:22:18 +0530 Subject: [PATCH 20/24] fix: validate per-job list lengths in enqueue_many --- py_src/taskito/app.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/py_src/taskito/app.py b/py_src/taskito/app.py index b2189ea0..60bc0fde 100644 --- a/py_src/taskito/app.py +++ b/py_src/taskito/app.py @@ -631,6 +631,21 @@ def enqueue_many( f"kwargs_list length ({len(kwargs_list)}) must match " f"args_list length ({len(args_list)})" ) + # Validate every per-job list up front so a length mismatch fails here + # with a clear message rather than misaligning the batch arrays later. + for field_name, values in ( + ("delay_list", delay_list), + ("unique_keys", unique_keys), + ("metadata_list", metadata_list), + ("notes_list", notes_list), + ("expires_list", expires_list), + ("result_ttl_list", result_ttl_list), + ("idempotency_keys", idempotency_keys), + ): + if values is not None and len(values) != count: + raise ValueError( + f"{field_name} length ({len(values)}) must match args_list length ({count})" + ) kw_list = kwargs_list or [{}] * count chain = self._get_middleware_chain(task_name) From 8f00b678fecba4bc4c3290b61382133fb59c27d3 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:22:23 +0530 Subject: [PATCH 21/24] fix(proxies): cancel still-queued reconstruction on timeout --- py_src/taskito/proxies/reconstruct.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/py_src/taskito/proxies/reconstruct.py b/py_src/taskito/proxies/reconstruct.py index 8590b835..d3355609 100644 --- a/py_src/taskito/proxies/reconstruct.py +++ b/py_src/taskito/proxies/reconstruct.py @@ -181,6 +181,10 @@ def _reconstruct_one( try: obj = future.result(timeout=max_timeout) except FuturesTimeout: + # Frees the slot if the task hasn't started; a started + # reconstruction can't be interrupted (thread pools don't + # support termination), so handlers should bound their own I/O. + future.cancel() raise ProxyReconstructionError( f"Reconstruction of '{handler_name}' timed out after {max_timeout}s" ) from None From 946b3d106c4f71b9c050c3d2a5163ddaf528bbef Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:22:27 +0530 Subject: [PATCH 22/24] perf(dashboard): align LastRefreshed timeout to label boundary --- dashboard/src/components/layout/last-refreshed.tsx | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/dashboard/src/components/layout/last-refreshed.tsx b/dashboard/src/components/layout/last-refreshed.tsx index 70c9d366..5db7fcb6 100644 --- a/dashboard/src/components/layout/last-refreshed.tsx +++ b/dashboard/src/components/layout/last-refreshed.tsx @@ -25,7 +25,14 @@ export function LastRefreshed({ className }: { className?: string }) { let id: ReturnType; const schedule = () => { const age = Date.now() - lastRefreshedAt; - const delay = age < 60_000 ? 1_000 : age < 3_600_000 ? 60_000 : 3_600_000; + // Fire at the next boundary the label actually changes on, so it never + // lags by almost a full bucket. + const delay = + age < 60_000 + ? 1_000 - (age % 1_000) + : age < 3_600_000 + ? 60_000 - (age % 60_000) + : 3_600_000 - (age % 3_600_000); id = setTimeout(() => { setTick((n) => n + 1); schedule(); From 7b864da2220152c0f13b80dc2fff01264784e4c2 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:22:33 +0530 Subject: [PATCH 23/24] test: pin middleware-chain TTL for deterministic cache assertion --- tests/core/test_middleware_chain_cache.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/core/test_middleware_chain_cache.py b/tests/core/test_middleware_chain_cache.py index 70622885..c6a9d02e 100644 --- a/tests/core/test_middleware_chain_cache.py +++ b/tests/core/test_middleware_chain_cache.py @@ -10,12 +10,19 @@ from pathlib import Path +import pytest + from taskito import Queue -def test_middleware_chain_caches_disable_reads(tmp_path: Path) -> None: +def test_middleware_chain_caches_disable_reads( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: """Repeated dispatches read the disable list from storage at most once within the TTL window.""" + # Pin the TTL large so the assertion can't flake if the loop happens to + # straddle the real 1s window under load. + monkeypatch.setattr("taskito.mixins.decorators._MW_CHAIN_TTL", 1_000_000.0) queue = Queue(db_path=str(tmp_path / "cache.db")) reads = {"count": 0} @@ -32,8 +39,12 @@ def counting_get_setting(key: str) -> str | None: assert reads["count"] == 1, "disable list should be read once, then cached" -def test_disable_version_bump_invalidates_cache(tmp_path: Path) -> None: +def test_disable_version_bump_invalidates_cache( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: """A same-process disable change re-reads the disable list on the next call.""" + # Pin the TTL so only the version bump (not expiry) can trigger a re-read. + monkeypatch.setattr("taskito.mixins.decorators._MW_CHAIN_TTL", 1_000_000.0) queue = Queue(db_path=str(tmp_path / "cache.db")) reads = {"count": 0} From 437337fb9791d02fbe5751912706ca87dffaa67c Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Fri, 29 May 2026 15:38:30 +0530 Subject: [PATCH 24/24] refactor(proxies): bound reconstruction timeout with a daemon thread A shared ThreadPoolExecutor force-joins its workers at interpreter exit (CPython issue 36780), so a hung reconstruction would block shutdown, and a bounded pool can be exhausted by hung tasks. A process pool can't return the non-picklable live objects reconstruction produces. Run each timed reconstruction on a daemon thread joined with the timeout instead: an overrun is abandoned without blocking the worker or shutdown. Adds timeout tests. --- py_src/taskito/proxies/reconstruct.py | 75 +++++++++++++++++++-------- tests/resources/test_proxies.py | 57 ++++++++++++++++++++ 2 files changed, 109 insertions(+), 23 deletions(-) diff --git a/py_src/taskito/proxies/reconstruct.py b/py_src/taskito/proxies/reconstruct.py index d3355609..62f5cd2a 100644 --- a/py_src/taskito/proxies/reconstruct.py +++ b/py_src/taskito/proxies/reconstruct.py @@ -2,15 +2,10 @@ from __future__ import annotations -import atexit import logging +import threading import time -from concurrent.futures import ( - ThreadPoolExecutor, -) -from concurrent.futures import ( - TimeoutError as FuturesTimeout, -) +from collections.abc import Callable from typing import TYPE_CHECKING, Any from taskito.exceptions import ProxyReconstructionError @@ -27,11 +22,49 @@ _PROXY_KEY = "__taskito_proxy__" _REF_KEY = "__taskito_ref__" -# Shared, process-wide pool for timeout-bounded proxy reconstruction. A fresh -# ThreadPoolExecutor per reconstructed argument would spawn and tear down an OS -# thread on the worker hot path; one small pool serves all reconstructions. -_RECONSTRUCT_POOL = ThreadPoolExecutor(max_workers=4, thread_name_prefix="taskito-proxy") -atexit.register(_RECONSTRUCT_POOL.shutdown, wait=False) + +class _ReconstructionTimeout(Exception): + """Internal signal that a reconstruction exceeded its wall-clock budget.""" + + +def _run_with_timeout(fn: Callable[[], Any], timeout: float) -> Any: + """Run ``fn`` with a wall-clock ``timeout`` and return its result. + + ``fn`` runs on a dedicated **daemon** thread so an overrunning reconstruction + is abandoned without blocking either the worker or interpreter shutdown. + This is the only safe in-process option: + + * A process pool could be terminated, but reconstructed resources (sessions, + clients, file handles) are not picklable and so can't be returned across a + process boundary. + * A shared ``ThreadPoolExecutor`` force-joins its workers at interpreter exit + (CPython issue 36780), so a hung task would hang shutdown. + * Python cannot safely kill a running thread. + + The watchdog only bounds how long the *caller* waits; a hung reconstruction + keeps running in the background until it returns. Hard runaway protection + must come from the handler bounding its own I/O (connect/read timeouts). + + Raises :class:`_ReconstructionTimeout` if the budget is exceeded, and + re-raises any exception ``fn`` itself raised. + """ + result: list[Any] = [] + error: list[Exception] = [] + + def runner() -> None: + try: + result.append(fn()) + except Exception as exc: # surfaced to the caller below + error.append(exc) + + worker = threading.Thread(target=runner, name="taskito-proxy-reconstruct", daemon=True) + worker.start() + worker.join(timeout) + if worker.is_alive(): + raise _ReconstructionTimeout + if error: + raise error[0] + return result[0] def reconstruct_proxies( @@ -177,19 +210,15 @@ def _reconstruct_one( start = time.monotonic() try: if max_timeout > 0: - future = _RECONSTRUCT_POOL.submit(handler.reconstruct, recipe, version) - try: - obj = future.result(timeout=max_timeout) - except FuturesTimeout: - # Frees the slot if the task hasn't started; a started - # reconstruction can't be interrupted (thread pools don't - # support termination), so handlers should bound their own I/O. - future.cancel() - raise ProxyReconstructionError( - f"Reconstruction of '{handler_name}' timed out after {max_timeout}s" - ) from None + obj = _run_with_timeout(lambda: handler.reconstruct(recipe, version), max_timeout) else: obj = handler.reconstruct(recipe, version) + except _ReconstructionTimeout: + if metrics is not None: + metrics.record_error(handler_name) + raise ProxyReconstructionError( + f"Reconstruction of '{handler_name}' timed out after {max_timeout}s" + ) from None except ProxyReconstructionError: if metrics is not None: metrics.record_error(handler_name) diff --git a/tests/resources/test_proxies.py b/tests/resources/test_proxies.py index b6c960fd..312ed4f9 100644 --- a/tests/resources/test_proxies.py +++ b/tests/resources/test_proxies.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import time from typing import Any import pytest @@ -12,6 +13,7 @@ from taskito.proxies.built_in import register_builtin_handlers from taskito.proxies.handlers.file import FileHandler from taskito.proxies.handlers.logger import LoggerHandler +from taskito.proxies.reconstruct import _ReconstructionTimeout, _run_with_timeout # --------------------------------------------------------------------------- # FileHandler @@ -427,3 +429,58 @@ def test_logger_proxy_marker_production(tmp_path: Any) -> None: assert args[0].get("__taskito_proxy__") is True assert args[0]["handler"] == "logger" assert args[0]["recipe"]["name"] == "test.proxy.marker" + + +class TestReconstructionTimeout: + """The reconstruction watchdog bounds caller wait time without blocking.""" + + def test_run_with_timeout_returns_and_propagates(self) -> None: + assert _run_with_timeout(lambda: 42, 1.0) == 42 + + def boom() -> None: + raise ValueError("handler failed") + + with pytest.raises(ValueError, match="handler failed"): + _run_with_timeout(boom, 1.0) + + def test_run_with_timeout_unblocks_caller_on_hang(self) -> None: + """A hung callable raises promptly — the caller never waits it out.""" + started = time.monotonic() + with pytest.raises(_ReconstructionTimeout): + _run_with_timeout(lambda: time.sleep(5), 0.05) + # Returned on the ~0.05s budget, not after the 5s sleep. + assert time.monotonic() - started < 1.0 + + def test_slow_reconstruct_times_out_as_proxy_error(self) -> None: + class SlowHandler: + name = "slow" + version = 1 + handled_types: tuple[type, ...] = () + + def detect(self, obj: Any) -> bool: + return False + + def deconstruct(self, obj: Any) -> dict[str, Any]: + return {} + + def reconstruct(self, recipe: dict[str, Any], version: int) -> Any: + time.sleep(5) + return "never" + + def cleanup(self, obj: Any) -> None: + pass + + reg = ProxyRegistry() + reg.register(SlowHandler()) + marker = { + "__taskito_proxy__": True, + "handler": "slow", + "version": 1, + "recipe": {}, + } + + started = time.monotonic() + with pytest.raises(ProxyReconstructionError, match="timed out"): + reconstruct_proxies((marker,), {}, reg, max_timeout=1) + # Surfaced on the 1s budget rather than blocking for the 5s sleep. + assert time.monotonic() - started < 3.0