check_repeaters() task iterates repeaters

corehq/apps/app_manager/tests/test_repeater.py

@@ -3,8 +3,6 @@
 from django.test import TestCase
 from django.test.client import Client
 
-from unittest.mock import patch
-
 from corehq.apps.accounting.models import SoftwarePlanEdition
 from corehq.apps.accounting.tests.utils import DomainSubscriptionMixin
 from corehq.apps.accounting.utils import clear_plan_version_cache
@@ -64,18 +62,3 @@ def test_repeat_record_created(self):
         later = datetime.utcnow() + timedelta(hours=48 + 1)
         repeat_records = RepeatRecord.objects.filter(domain=self.domain, next_check__lt=later)
         self.assertEqual(len(repeat_records), 1)
-
-    def test_repeat_record_forwarded(self):
-        """
-        When an application with a repeater is saved, then HQ should try to forward the repeat record
-        """
-        self.app_structure_repeater = AppStructureRepeater(domain=self.domain, connection_settings=self.conn)
-        self.app_structure_repeater.save()
-        self.addCleanup(self.app_structure_repeater.delete)
-
-        with patch('corehq.motech.repeaters.models.simple_request') as mock_fire:
-            self.application = Application(domain=self.domain)
-            self.application.save()
-            self.addCleanup(self.application.delete)
-
-            self.assertEqual(mock_fire.call_count, 1)

corehq/apps/domain/tests/test_deletion_models.py

@@ -25,6 +25,7 @@
     'icds_reports',
     'notifications',
     'oauth2_provider',
+    'pg_lock',
     'phonelog',  # these are deleted after 60 days regardless
     'pillow_retry',
     'pillowtop',

corehq/apps/dump_reload/tests/test_dump_models.py

@@ -79,6 +79,7 @@
     "otp_static.StaticDevice",
     "otp_static.StaticToken",
     "otp_totp.TOTPDevice",
+    "pg_lock.PGLock",
     "phone.SyncLogSQL",  # not required and can be a lot of data
     "pillow_retry.PillowError",
     "pillowtop.DjangoPillowCheckpoint",

corehq/apps/pg_lock/apps.py

@@ -0,0 +1,5 @@
+from django.apps import AppConfig
+
+
+class PGLockConfig(AppConfig):
+    name = 'corehq.apps.pg_lock'

corehq/apps/pg_lock/migrations/0001_initial.py

@@ -0,0 +1,23 @@
+# Generated by Django 4.2.14 on 2024-08-09 18:49
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = []
+
+    operations = [
+        migrations.CreateModel(
+            name="PGLock",
+            fields=[
+                (
+                    "key",
+                    models.CharField(max_length=255, primary_key=True, serialize=False),
+                ),
+                ("expires_at", models.DateTimeField(blank=True, null=True)),
+            ],
+        ),
+    ]

corehq/apps/pg_lock/models.py

@@ -0,0 +1,82 @@
+"""
+A partial implementation of a lock using PostgreSQL. It does not
+implement blocking.
+
+The intention of using Postgres for locking is to have a lock that can
+be used for coordinating across multiple worker processes.
+"""
+from contextlib import contextmanager
+from datetime import datetime, timedelta
+
+from django.db import models, IntegrityError
+
+
+class PGLock(models.Model):
+    key = models.CharField(max_length=255, primary_key=True)
+    expires_at = models.DateTimeField(null=True, blank=True)
+
+
+class Lock:
+    def __init__(self, key):
+        self.key = key
+
+    def __str__(self):
+        return self.key
+
+    @property
+    def name(self):  # Used by MeteredLock
+        return self.key
+
+    def acquire(self, blocking=True, timeout=-1):
+        if blocking:
+            raise NotImplementedError("Blocking is not supported")
+
+        if timeout >= 0:
+            expires_at = datetime.utcnow() + timedelta(seconds=timeout)
+        else:
+            expires_at = None
+
+        try:
+            pg_lock, created = PGLock.objects.get_or_create(
+                key=self.key,
+                defaults={'expires_at': expires_at},
+            )
+            if created:
+                return True
+            if (
+                pg_lock.expires_at is not None
+                and pg_lock.expires_at <= datetime.utcnow()
+            ):
+                # Lock has expired
+                pg_lock.expires_at = expires_at
+                pg_lock.save()
+                return True
+            return False
+        except IntegrityError:
+            return False
+
+    def release(self):
+        PGLock.objects.filter(key=self.key).delete()
+
+    def locked(self):
+        return (
+            PGLock.objects
+            .filter(key=self.key)
+            .filter(
+                models.Q(expires_at__isnull=True)
+                | models.Q(expires_at__gt=datetime.utcnow())
+            )
+            .exists()
+        )
+
+
+@contextmanager
+def lock(key, timeout=-1):
+    lock = Lock(key)
+    acquired = False
+    try:
+        acquired = lock.acquire(blocking=False, timeout=timeout)
+        yield acquired
+    finally:
+        if acquired:
+            lock.release()

corehq/apps/pg_lock/tests/test_models.py

@@ -0,0 +1,85 @@
+from datetime import datetime
+from unittest import expectedFailure
+from unittest.mock import patch
+
+from django.test import TestCase
+
+from corehq.apps.pg_lock.models import Lock, lock
+from dimagi.utils.couch import get_redis_lock, get_pg_lock
+
+
+class PGLockTests(TestCase):
+
+    def test_acquires_lock_when_not_locked(self):
+        lock_instance = Lock('test_lock')
+        acquired = lock_instance.acquire(blocking=False)
+        self.assertTrue(acquired)
+        self.assertTrue(lock_instance.locked())
+        lock_instance.release()
+
+    def test_does_not_acquire_lock_when_already_locked(self):
+        lock_instance = Lock('test_lock')
+        lock_instance.acquire(blocking=False)
+        another_lock_instance = Lock('test_lock')
+        acquired = another_lock_instance.acquire(blocking=False)
+        self.assertFalse(acquired)
+        lock_instance.release()
+
+    def test_releases_lock(self):
+        lock_instance = Lock('test_lock')
+        lock_instance.acquire(blocking=False)
+        lock_instance.release()
+        self.assertFalse(lock_instance.locked())
+
+    def test_releases_lock_not_acquired(self):
+        lock_instance = Lock('test_lock')
+        lock_instance.acquire(blocking=False)
+        another_lock_instance = Lock('test_lock')
+        another_lock_instance.release()
+        self.assertFalse(lock_instance.locked())
+
+    @patch('corehq.apps.pg_lock.models.datetime')
+    def test_acquires_lock_after_expiration(self, mock_datetime):
+        lock_instance = Lock('test_lock')
+        mock_datetime.utcnow.return_value = datetime(2023, 1, 1, 12, 0, 0)
+        lock_instance.acquire(blocking=False, timeout=1)
+        mock_datetime.utcnow.return_value = datetime(2023, 1, 1, 12, 0, 2)
+        acquired = lock_instance.acquire(blocking=False)
+        self.assertTrue(acquired)
+        lock_instance.release()
+
+    def test_context_manager_acquires_and_releases_lock(self):
+        with lock('test_lock') as acquired:
+            self.assertTrue(acquired)
+            lock_instance = Lock('test_lock')
+            self.assertTrue(lock_instance.locked())
+        lock_instance = Lock('test_lock')
+        self.assertFalse(lock_instance.locked())
+
+
+class TestLockWorkers(TestCase):
+
+    @expectedFailure
+    def test_release_redis_lock_not_acquired(self):
+        # Worker 1:
+        lock1 = get_redis_lock('test-key', timeout=1, name='test-name')
+        self.assertTrue(lock1.acquire(blocking=False))
+        self.assertTrue(lock1.locked())
+
+        # Worker 2:
+        lock2 = get_redis_lock('test-key', timeout=1, name='test-name')
+        lock2.release()  # redis.exceptions.LockError: Cannot release an unlocked lock
+
+        self.assertFalse(lock1.locked())
+
+    def test_release_pg_lock_not_acquired(self):
+        # Worker 1:
+        lock1 = get_pg_lock('test-key', name='test-name')
+        self.assertTrue(lock1.acquire(blocking=False))
+        self.assertTrue(lock1.locked())
+
+        # Worker 2:
+        lock2 = get_pg_lock('test-key', name='test-name')
+        lock2.release()
+
+        self.assertFalse(lock1.locked())

corehq/apps/reports/filters/select.py

@@ -144,6 +144,7 @@ def options(self):
             State.Pending,
             State.Cancelled,
             State.Fail,
+            State.InvalidPayload,
         ]]
 
 

corehq/ex-submodules/dimagi/utils/couch/__init__.py

@@ -97,6 +97,12 @@ def get_redis_lock(key, timeout=None, name=None, track_unreleased=True, **kw):
     return MeteredLock(lock, name, track_unreleased)
 
 
+def get_pg_lock(key, name, track_unreleased=True):
+    from corehq.apps.pg_lock.models import Lock
+    lock = Lock(key)
+    return MeteredLock(lock, name, track_unreleased)
+
+
 def acquire_lock(lock, degrade_gracefully, **kwargs):
     acquired = False
     try:

corehq/motech/repeaters/const.py

@@ -8,12 +8,18 @@
 MIN_RETRY_WAIT = timedelta(minutes=60)
 RATE_LIMITER_DELAY_RANGE = (timedelta(minutes=0), timedelta(minutes=15))
 CHECK_REPEATERS_INTERVAL = timedelta(minutes=5)
-CHECK_REPEATERS_PARTITION_COUNT = settings.CHECK_REPEATERS_PARTITION_COUNT
+CHECK_REPEATERS_PARTITION_COUNT = settings.CHECK_REPEATERS_PARTITION_COUNT  # TODO: Drop
 CHECK_REPEATERS_KEY = 'check-repeaters-key'
 # Number of attempts to an online endpoint before cancelling payload
 MAX_ATTEMPTS = 3
 # Number of exponential backoff attempts to an offline endpoint
 MAX_BACKOFF_ATTEMPTS = 6
+# The maximum number of workers that one repeater can use to send repeat
+# records at the same time. (In other words, HQ's capacity to DDOS
+# attack a remote API endpoint.) This is a guardrail to prevent one
+# domain from hogging repeat record queue workers and to ensure that
+# repeaters are iterated fairly.
+MAX_REPEATER_WORKERS = 7
 
 
 class State(IntegerChoices):
@@ -23,13 +29,15 @@ class State(IntegerChoices):
     Success = 4, _('Succeeded')
     Cancelled = 8, _('Cancelled')
     Empty = 16, _('Empty')  # There was nothing to send. Implies Success.
+    InvalidPayload = 32, _('Invalid Payload')  # Implies Cancelled.
 
 
 RECORD_PENDING_STATE = State.Pending
 RECORD_SUCCESS_STATE = State.Success
 RECORD_FAILURE_STATE = State.Fail
 RECORD_CANCELLED_STATE = State.Cancelled
 RECORD_EMPTY_STATE = State.Empty
+RECORD_INVALIDPAYLOAD_STATE = State.InvalidPayload
 
 
 class UCRRestrictionFFStatus(IntegerChoices):

corehq/motech/repeaters/management/commands/find_missing_repeat_records.py

@@ -527,7 +527,8 @@ def create_case_repeater_register(repeater, domain, payload):
         'domain': domain,
         'doc_type': repeater.repeater_type
     })
-    repeat_record.attempt_forward_now()
+    # TODO: No, the repeater sends the repeat record.
+    # repeat_record.attempt_forward_now()
     return repeat_record
 
 

corehq/motech/repeaters/migrations/0013_repeater_max_workers_alter_repeater_is_paused_and_more.py

@@ -0,0 +1,46 @@
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("repeaters", "0012_formexpressionrepeater_arcgisformexpressionrepeater"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="repeater",
+            name="max_workers",
+            field=models.IntegerField(default=7),
+        ),
+        migrations.AlterField(
+            model_name="repeater",
+            name="is_paused",
+            field=models.BooleanField(db_index=True, default=False),
+        ),
+        migrations.AlterField(
+            model_name="repeater",
+            name="next_attempt_at",
+            field=models.DateTimeField(blank=True, db_index=True, null=True),
+        ),
+        migrations.AlterField(
+            model_name="repeatrecord",
+            name="domain",
+            field=models.CharField(db_index=True, max_length=126),
+        ),
+        migrations.AlterField(
+            model_name="repeatrecord",
+            name="state",
+            field=models.PositiveSmallIntegerField(
+                choices=[
+                    (1, "Pending"),
+                    (2, "Failed"),
+                    (4, "Succeeded"),
+                    (8, "Cancelled"),
+                    (16, "Empty"),
+                ],
+                db_index=True,
+                default=1,
+            ),
+        ),
+    ]