Improve get_waiting_tasks
* Switch to a generator in `get_waiting_tasks` * Faster checking for empty task queues
This commit is contained in:
tcely
GitHub
@@ -1,10 +1,14 @@
|
||||
import uuid
|
||||
from datetime import timedelta
|
||||
from itertools import islice
|
||||
|
||||
from django.db import models
|
||||
from django.db import connection, models, transaction
|
||||
from django.utils import timezone
|
||||
|
||||
from ..json import JSONEncoder
|
||||
# from common.json import JSONEncoder
|
||||
from ..utils import is_empty_iterator
|
||||
#from common.utils import is_empty_iterator
|
||||
|
||||
# cls = TaskHistory
|
||||
# TaskHistory is defined below this function in this file
|
||||
@@ -43,9 +47,79 @@ def th_schedule(cls, task_wrapper, /, *args, remove_duplicates=False, vn_args=()
|
||||
task_wrapper.huey.enqueue(task_obj)
|
||||
return True
|
||||
|
||||
# self = TaskHistoryQuerySet
|
||||
# TaskHistoryQuerySet is defined below this function in this file
|
||||
def thqs_from_huey_ids(self, /, huey_task_ids):
|
||||
"""
|
||||
Robustly matches Huey task.id values to TaskHistory records.
|
||||
|
||||
Optimized for large datasets across SQLite, PostgreSQL, and MariaDB.
|
||||
Bypasses SQL variable limits by using request-cycle temporary tables.
|
||||
"""
|
||||
# 1. Guard clause for empty input - returns a valid, empty QuerySet
|
||||
empty, huey_task_ids = is_empty_iterator(huey_task_ids)
|
||||
if empty:
|
||||
return self.none()
|
||||
|
||||
# 2. Dynamic metadata and unique naming
|
||||
task_history_table = self.model._meta.db_table
|
||||
unique_suffix = uuid.uuid4().hex[:8]
|
||||
input_tmp = f"tmp_huey_ids_{unique_suffix}"
|
||||
results_tmp = f"tmp_history_pks_{unique_suffix}"
|
||||
|
||||
# 3. Database operations
|
||||
def validated_id_generator():
|
||||
for tid in huey_task_ids:
|
||||
try:
|
||||
# Validates format and normalizes to lowercase dashed string
|
||||
yield (str(uuid.UUID(str(tid))).lower(),)
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
# Skip malformed IDs and log for troubleshooting
|
||||
from common.logger import log
|
||||
log.warning(f"Skipping malformed Huey task ID: {tid}")
|
||||
continue
|
||||
|
||||
with transaction.atomic():
|
||||
with connection.cursor() as cursor:
|
||||
# Stage 1: Store and normalize input IDs
|
||||
cursor.execute(f"CREATE TEMPORARY TABLE {input_tmp} (tid VARCHAR(40) PRIMARY KEY)")
|
||||
|
||||
# Single iteration: ensures dashed strings and lowercase for case-sensitive DBs
|
||||
# Batch stream
|
||||
batch_size = 40_000
|
||||
stream = validated_id_generator()
|
||||
while batch := list(islice(stream, batch_size)):
|
||||
cursor.executemany(f"INSERT INTO {input_tmp} VALUES (%s)", batch)
|
||||
|
||||
# Stage 2: Filter to a PK-only results table
|
||||
cursor.execute(f"CREATE TEMPORARY TABLE {results_tmp} (id BIGINT)")
|
||||
cursor.execute(f"""
|
||||
INSERT INTO {results_tmp} (id)
|
||||
SELECT id FROM {task_history_table}
|
||||
WHERE task_id IN (SELECT tid FROM {input_tmp})
|
||||
""")
|
||||
|
||||
# Index the result PKs to ensure dashboard pagination and ordering are fast
|
||||
cursor.execute(f"CREATE INDEX idx_{results_tmp} ON {results_tmp}(id)")
|
||||
|
||||
# Immediate cleanup of the input string table to save memory
|
||||
cursor.execute(f"DROP TABLE IF EXISTS {input_tmp}")
|
||||
|
||||
# 4. Return a Lazy QuerySet via RawSQL to bypass RawQuerySet.clone() limitations.
|
||||
# The 'results_tmp' table persists until the connection closes after the request.
|
||||
return self.filter(
|
||||
id__in=models.expressions.RawSQL(
|
||||
f"SELECT id FROM {results_tmp}",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TaskHistoryQuerySet(models.QuerySet):
|
||||
|
||||
def from_huey_ids(self, /, huey_task_ids):
|
||||
return thqs_from_huey_ids(self, huey_task_ids)
|
||||
|
||||
def running(self, now=None, within=None):
|
||||
if now is None:
|
||||
now = timezone.now()
|
||||
@@ -135,7 +209,7 @@ class TaskHistory(models.Model):
|
||||
self.queue = self.queue or None
|
||||
self.verbose_name = self.verbose_name or None
|
||||
return super().save(*args, **kwargs)
|
||||
|
||||
|
||||
def has_error(self):
|
||||
"""
|
||||
Check if the last_error field is empty.
|
||||
|
||||
@@ -9,6 +9,7 @@ import string
|
||||
import time
|
||||
from django.core.paginator import Paginator
|
||||
from functools import partial
|
||||
from itertools import chain
|
||||
from operator import attrgetter, itemgetter
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlunsplit, urlencode, urlparse
|
||||
@@ -71,6 +72,23 @@ def glob_quote(filestr, /):
|
||||
return filestr.translate(str.maketrans(_glob_specials))
|
||||
|
||||
|
||||
def is_empty_iterator(iterator):
|
||||
"""
|
||||
Checks if an iterator is empty without fully consuming it.
|
||||
Returns: (is_empty_boolean, iterator)
|
||||
"""
|
||||
returned_iterator = iterator
|
||||
try:
|
||||
first_item = next(iterator)
|
||||
except StopIteration:
|
||||
return True, iter([])
|
||||
else:
|
||||
# Put the first item back at the start of a new iterator
|
||||
# Chaining the single item with the original remaining stream
|
||||
returned_iterator = chain([first_item], iterator)
|
||||
return False, returned_iterator
|
||||
|
||||
|
||||
def list_of_dictionaries(arg_list, /, arg_function=lambda x: x):
|
||||
assert callable(arg_function)
|
||||
_map_list = arg_list
|
||||
|
||||
@@ -3,7 +3,6 @@ import os
|
||||
from base64 import b64decode
|
||||
import pathlib
|
||||
import sys
|
||||
import uuid
|
||||
from django.conf import settings
|
||||
from django.http import FileResponse, Http404, HttpResponseNotFound, HttpResponseRedirect
|
||||
from django.views.generic import TemplateView, ListView, DetailView
|
||||
@@ -13,9 +12,8 @@ from django.views.generic.detail import SingleObjectMixin
|
||||
from django.core.exceptions import SuspiciousFileOperation
|
||||
from django.http import HttpResponse
|
||||
from django.urls import reverse_lazy
|
||||
from django.db import connection, transaction, IntegrityError
|
||||
from django.db import connection, IntegrityError
|
||||
from django.db.models import F, Q, Count, Sum, When, Case
|
||||
from django.db.models.expressions import RawSQL
|
||||
from django.forms import Form, ValidationError
|
||||
from django.utils.text import slugify
|
||||
from django.utils._os import safe_join
|
||||
@@ -45,73 +43,34 @@ from . import signals # noqa
|
||||
from . import youtube
|
||||
|
||||
|
||||
def get_histories_from_huey_ids(huey_task_ids):
|
||||
"""
|
||||
Robustly matches Huey task.id values to TaskHistory records.
|
||||
|
||||
Optimized for large datasets across SQLite, PostgreSQL, and MariaDB.
|
||||
Bypasses SQL variable limits by using request-cycle temporary tables.
|
||||
"""
|
||||
# 1. Guard clause for empty input - returns a valid, empty QuerySet
|
||||
if not huey_task_ids:
|
||||
return TaskHistory.objects.none()
|
||||
|
||||
# 2. Dynamic metadata and unique naming
|
||||
task_history_table = TaskHistory._meta.db_table
|
||||
unique_suffix = uuid.uuid4().hex[:8]
|
||||
input_tmp = f"tmp_huey_ids_{unique_suffix}"
|
||||
results_tmp = f"tmp_history_pks_{unique_suffix}"
|
||||
|
||||
# 3. Database operations
|
||||
def validated_id_generator():
|
||||
for tid in huey_task_ids:
|
||||
try:
|
||||
# Validates format and normalizes to lowercase dashed string
|
||||
yield (str(uuid.UUID(str(tid))).lower(),)
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
# Skip malformed IDs and log for troubleshooting
|
||||
log.warning(f"Skipping malformed Huey task ID: {tid}")
|
||||
continue
|
||||
|
||||
with transaction.atomic():
|
||||
with connection.cursor() as cursor:
|
||||
# Stage 1: Store and normalize input IDs
|
||||
cursor.execute(f"CREATE TEMPORARY TABLE {input_tmp} (tid VARCHAR(40) PRIMARY KEY)")
|
||||
|
||||
# Single iteration: ensures dashed strings and lowercase for case-sensitive DBs
|
||||
cursor.executemany(f"INSERT INTO {input_tmp} VALUES (%s)", validated_id_generator())
|
||||
|
||||
# Stage 2: Filter to a PK-only results table
|
||||
cursor.execute(f"CREATE TEMPORARY TABLE {results_tmp} (id BIGINT)")
|
||||
cursor.execute(f"""
|
||||
INSERT INTO {results_tmp} (id)
|
||||
SELECT id FROM {task_history_table}
|
||||
WHERE task_id IN (SELECT tid FROM {input_tmp})
|
||||
""")
|
||||
|
||||
# Index the result PKs to ensure dashboard pagination and ordering are fast
|
||||
cursor.execute(f"CREATE INDEX idx_{results_tmp} ON {results_tmp}(id)")
|
||||
|
||||
# Immediate cleanup of the input string table to save memory
|
||||
cursor.execute(f"DROP TABLE IF EXISTS {input_tmp}")
|
||||
|
||||
# 4. Return a Lazy QuerySet via RawSQL to bypass RawQuerySet.clone() limitations.
|
||||
# The 'results_tmp' table persists until the connection closes after the request.
|
||||
return TaskHistory.objects.filter(
|
||||
id__in=RawSQL(f"SELECT id FROM {results_tmp}", [])
|
||||
)
|
||||
|
||||
def get_waiting_tasks():
|
||||
huey_queue_names = (DJANGO_HUEY or {}).get('queues', {})
|
||||
huey_queues = list(map(get_queue, huey_queue_names))
|
||||
huey_task_ids = {
|
||||
str(t.id) for q in huey_queues for t in set(
|
||||
q.pending()
|
||||
).union(
|
||||
q.scheduled()
|
||||
)
|
||||
}
|
||||
return get_histories_from_huey_ids(huey_task_ids)
|
||||
|
||||
# Fast Guard: Check counts across all Huey queues
|
||||
if not any(0 < (q.pending_count() + q.scheduled_count()) for q in huey_queues):
|
||||
return TaskHistory.objects.none()
|
||||
|
||||
def id_generator(queue):
|
||||
# Stream pending tasks
|
||||
for task in queue.pending():
|
||||
yield str(task.id)
|
||||
|
||||
# Stream scheduled tasks
|
||||
for task in queue.scheduled():
|
||||
yield str(task.id)
|
||||
|
||||
def deduplicating_id_generator():
|
||||
seen = set()
|
||||
for q in huey_queues:
|
||||
for tid in id_generator(q):
|
||||
if tid not in seen:
|
||||
seen.add(tid)
|
||||
yield tid
|
||||
seen.clear()
|
||||
|
||||
huey_task_ids = deduplicating_id_generator()
|
||||
return TaskHistory.objects.from_huey_ids(huey_task_ids)
|
||||
|
||||
|
||||
class DashboardView(TemplateView):
|
||||
|
||||
Reference in New Issue
Block a user