From c94f2cb816720fbe57aed5e99c1f851e20a11606 Mon Sep 17 00:00:00 2001
From: Ilan Steemers <koed00@gmail.com>
Date: Sun, 5 Jul 2020 18:18:43 +0200
Subject: [PATCH] Adds psutil support for max_rss on other platforms

---
 django_q/__init__.py | 13 +------------
 django_q/admin.py    |  3 +--
 django_q/cluster.py  | 25 ++++++++++++++++++-------
 django_q/conf.py     | 10 ++++++++++
 django_q/models.py   |  2 +-
 5 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/django_q/__init__.py b/django_q/__init__.py
index f96c270..72938b9 100644
--- a/django_q/__init__.py
+++ b/django_q/__init__.py
@@ -3,15 +3,4 @@ VERSION = (1, 3, 1)
 default_app_config = "django_q.apps.DjangoQConfig"
 
 
-__all__ = ["conf", "cluster", "models", "tasks", "croniter"]
-
-# Optional Imports
-try:
-    from croniter import croniter
-except ImportError:
-    croniter = None
-
-try:
-    import resource
-except ModuleNotFoundError:
-    resource = None
+__all__ = ["conf", "cluster", "models", "tasks"]
diff --git a/django_q/admin.py b/django_q/admin.py
index 2625728..a7484ce 100644
--- a/django_q/admin.py
+++ b/django_q/admin.py
@@ -2,10 +2,9 @@
 from django.contrib import admin
 from django.utils.translation import gettext_lazy as _
 
-from django_q.conf import Conf
+from django_q.conf import Conf, croniter
 from django_q.models import Success, Failure, Schedule, OrmQ
 from django_q.tasks import async_task
-from django_q import croniter
 
 
 class TaskAdmin(admin.ModelAdmin):
diff --git a/django_q/cluster.py b/django_q/cluster.py
index 96da4c5..027b6a5 100644
--- a/django_q/cluster.py
+++ b/django_q/cluster.py
@@ -20,14 +20,21 @@ from django.utils.translation import gettext_lazy as _
 # Local
 import django_q.tasks
 from django_q.brokers import get_broker, Broker
-from django_q.conf import Conf, logger, psutil, get_ppid, error_reporter
+from django_q.conf import (
+    Conf,
+    logger,
+    psutil,
+    get_ppid,
+    error_reporter,
+    croniter,
+    resource,
+)
 from django_q.humanhash import humanize
 from django_q.models import Task, Success, Schedule
 from django_q.queues import Queue
 from django_q.signals import pre_execute
 from django_q.signing import SignedPackage, BadSignature
 from django_q.status import Stat, Status
-from django_q import croniter, resource
 
 
 class Cluster:
@@ -434,11 +441,7 @@ def worker(
             result_queue.put(task)
             timer.value = -1  # Idle
             # Recycle
-            if task_count == Conf.RECYCLE or (
-                resource
-                and Conf.MAX_RSS
-                and resource.getrusage(resource.RUSAGE_SELF).ru_maxrss >= Conf.MAX_RSS
-            ):
+            if task_count == Conf.RECYCLE or rss_check():
                 timer.value = -2  # Recycled
                 break
     logger.info(_(f"{name} stopped doing work"))
@@ -697,3 +700,11 @@ def set_cpu_affinity(n: int, process_ids: list, actual: bool = not Conf.TESTING)
             if actual:
                 p.cpu_affinity(affinity)
             logger.info(_(f"{pid} will use cpu {affinity}"))
+
+
+def rss_check():
+    if Conf.MAX_RSS and resource:
+        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss >= Conf.MAX_RSS
+    elif Conf.MAX_RSS and psutil:
+        return psutil.Process().memory_info().rss >= Conf.MAX_RSS * 1024
+    return False
diff --git a/django_q/conf.py b/django_q/conf.py
index 5b9b5b3..2e9ddc2 100644
--- a/django_q/conf.py
+++ b/django_q/conf.py
@@ -16,6 +16,16 @@ try:
 except ImportError:
     psutil = None
 
+try:
+    from croniter import croniter
+except ImportError:
+    croniter = None
+
+try:
+    import resource
+except ModuleNotFoundError:
+    resource = None
+
 
 class Conf:
     """
diff --git a/django_q/models.py b/django_q/models.py
index f2ba852..7037f12 100644
--- a/django_q/models.py
+++ b/django_q/models.py
@@ -13,8 +13,8 @@ from picklefield import PickledObjectField
 from picklefield.fields import dbsafe_decode
 
 # Local
+from django_q.conf import croniter
 from django_q.signing import SignedPackage
-from django_q import croniter
 
 
 class Task(models.Model):