use a lock to keep incrementals from running during full reindex

Author: ewdurbin

tests/unit/search/test_tasks.py

@@ -12,10 +12,12 @@
 
 import os
 
+import celery
 import elasticsearch
 import packaging.version
 import pretend
 import pytest
+import redis
 
 from first import first
 
@@ -25,6 +27,7 @@
     reindex_project,
     unindex_project,
     _project_docs,
+    SearchLock,
 )
 
 
@@ -125,6 +128,34 @@ def __init__(self):
         self.indices = FakeESIndices()
 
 
+class NotLock:
+    def __init__(*a, **kw):
+        pass
+
+    def acquire(self):
+        return True
+
+    def release(self):
+        return True
+
+
+class TestSearchLock:
+    def test_success(self):
+        lock_stub = pretend.stub(acquire=pretend.call_recorder(lambda: True))
+        r = pretend.stub(lock=lambda *a, **kw: lock_stub)
+        test_lock = SearchLock(r)
+        test_lock.__enter__()
+        assert lock_stub.acquire.calls == [pretend.call()]
+
+    def test_failure(self):
+        lock_stub = pretend.stub(acquire=pretend.call_recorder(lambda: False))
+        r = pretend.stub(lock=lambda *a, **kw: lock_stub)
+        test_lock = SearchLock(r)
+        with pytest.raises(redis.exceptions.LockError):
+            test_lock.__enter__()
+        assert lock_stub.acquire.calls == [pretend.call()]
+
+
 class TestReindex:
     def test_fails_when_raising(self, db_request, monkeypatch):
         docs = pretend.stub()
@@ -134,10 +165,14 @@ def project_docs(db):
 
         monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)
 
+        task = pretend.stub()
         es_client = FakeESClient()
 
         db_request.registry.update({"elasticsearch.index": "warehouse"})
-        db_request.registry.settings = {"elasticsearch.url": "http://some.url"}
+        db_request.registry.settings = {
+            "elasticsearch.url": "http://some.url",
+            "celery.scheduler_url": "redis://redis:6379/0",
+        }
         monkeypatch.setattr(
             warehouse.search.tasks.elasticsearch,
             "Elasticsearch",
@@ -153,18 +188,41 @@ def parallel_bulk(client, iterable, index=None):
             assert index == "warehouse-cbcbcbcbcb"
             raise TestException
 
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
+
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
 
         monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)
 
         with pytest.raises(TestException):
-            reindex(db_request)
+            reindex(task, db_request)
 
         assert es_client.indices.delete.calls == [
             pretend.call(index="warehouse-cbcbcbcbcb")
         ]
         assert es_client.indices.put_settings.calls == []
 
+    def test_retry_on_lock(self, db_request, monkeypatch):
+        task = pretend.stub(
+            retry=pretend.call_recorder(pretend.raiser(celery.exceptions.Retry))
+        )
+
+        db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}
+
+        le = redis.exceptions.LockError()
+        monkeypatch.setattr(
+            redis.StrictRedis,
+            "from_url",
+            lambda *a, **kw: pretend.stub(lock=pretend.raiser(le)),
+        )
+
+        with pytest.raises(celery.exceptions.Retry):
+            reindex(task, db_request)
+
+        assert task.retry.calls == [pretend.call(countdown=60, exc=le)]
+
     def test_successfully_indexes_and_adds_new(self, db_request, monkeypatch):
 
         docs = pretend.stub()
@@ -174,24 +232,31 @@ def project_docs(db):
 
         monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)
 
+        task = pretend.stub()
         es_client = FakeESClient()
 
         db_request.registry.update(
             {"elasticsearch.index": "warehouse", "elasticsearch.shards": 42}
         )
-        db_request.registry.settings = {"elasticsearch.url": "http://some.url"}
+        db_request.registry.settings = {
+            "elasticsearch.url": "http://some.url",
+            "celery.scheduler_url": "redis://redis:6379/0",
+        }
         monkeypatch.setattr(
             warehouse.search.tasks.elasticsearch,
             "Elasticsearch",
             lambda *a, **kw: es_client,
         )
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
 
         parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
 
         monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)
 
-        reindex(db_request)
+        reindex(task, db_request)
 
         assert parallel_bulk.calls == [
             pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")
@@ -220,6 +285,7 @@ def project_docs(db):
 
     def test_successfully_indexes_and_replaces(self, db_request, monkeypatch):
         docs = pretend.stub()
+        task = pretend.stub()
 
         def project_docs(db):
             return docs
@@ -238,19 +304,25 @@ def project_docs(db):
                 "sqlalchemy.engine": db_engine,
             }
         )
-        db_request.registry.settings = {"elasticsearch.url": "http://some.url"}
+        db_request.registry.settings = {
+            "elasticsearch.url": "http://some.url",
+            "celery.scheduler_url": "redis://redis:6379/0",
+        }
         monkeypatch.setattr(
             warehouse.search.tasks.elasticsearch,
             "Elasticsearch",
             lambda *a, **kw: es_client,
         )
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
 
         parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
 
         monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)
 
-        reindex(db_request)
+        reindex(task, db_request)
 
         assert parallel_bulk.calls == [
             pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")
@@ -281,6 +353,7 @@ def project_docs(db):
 class TestPartialReindex:
     def test_reindex_fails_when_raising(self, db_request, monkeypatch):
         docs = pretend.stub()
+        task = pretend.stub()
 
         def project_docs(db, project_name=None):
             return docs
@@ -302,44 +375,96 @@ def parallel_bulk(client, iterable, index=None):
             raise TestException
 
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
 
         with pytest.raises(TestException):
-            reindex_project(db_request, "foo")
+            reindex_project(task, db_request, "foo")
 
         assert es_client.indices.put_settings.calls == []
 
-    def test_unindex_fails_when_raising(self, db_request):
+    def test_unindex_fails_when_raising(self, db_request, monkeypatch):
+        task = pretend.stub()
+
         class TestException(Exception):
             pass
 
         es_client = FakeESClient()
         es_client.delete = pretend.raiser(TestException)
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
 
         db_request.registry.update(
             {"elasticsearch.client": es_client, "elasticsearch.index": "warehouse"}
         )
 
         with pytest.raises(TestException):
-            unindex_project(db_request, "foo")
+            unindex_project(task, db_request, "foo")
+
+    def test_unindex_accepts_defeat(self, db_request, monkeypatch):
+        task = pretend.stub()
 
-    def test_unindex_accepts_defeat(self, db_request):
         es_client = FakeESClient()
         es_client.delete = pretend.call_recorder(
             pretend.raiser(elasticsearch.exceptions.NotFoundError)
         )
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
 
         db_request.registry.update(
             {"elasticsearch.client": es_client, "elasticsearch.index": "warehouse"}
         )
 
-        unindex_project(db_request, "foo")
+        unindex_project(task, db_request, "foo")
 
         assert es_client.delete.calls == [
             pretend.call(index="warehouse", doc_type="doc", id="foo")
         ]
 
+    def test_unindex_retry_on_lock(self, db_request, monkeypatch):
+        task = pretend.stub(
+            retry=pretend.call_recorder(pretend.raiser(celery.exceptions.Retry))
+        )
+
+        db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}
+
+        le = redis.exceptions.LockError()
+        monkeypatch.setattr(
+            redis.StrictRedis,
+            "from_url",
+            lambda *a, **kw: pretend.stub(lock=pretend.raiser(le)),
+        )
+
+        with pytest.raises(celery.exceptions.Retry):
+            unindex_project(task, db_request, "foo")
+
+        assert task.retry.calls == [pretend.call(countdown=60, exc=le)]
+
+    def test_reindex_retry_on_lock(self, db_request, monkeypatch):
+        task = pretend.stub(
+            retry=pretend.call_recorder(pretend.raiser(celery.exceptions.Retry))
+        )
+
+        db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}
+
+        le = redis.exceptions.LockError()
+        monkeypatch.setattr(
+            redis.StrictRedis,
+            "from_url",
+            lambda *a, **kw: pretend.stub(lock=pretend.raiser(le)),
+        )
+
+        with pytest.raises(celery.exceptions.Retry):
+            reindex_project(task, db_request, "foo")
+
+        assert task.retry.calls == [pretend.call(countdown=60, exc=le)]
+
     def test_successfully_indexes(self, db_request, monkeypatch):
         docs = pretend.stub()
+        task = pretend.stub()
 
         def project_docs(db, project_name=None):
             return docs
@@ -364,8 +489,11 @@ def project_docs(db, project_name=None):
             lambda client, iterable, index=None: [None]
         )
         monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)
+        monkeypatch.setattr(
+            redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)
+        )
 
-        reindex_project(db_request, "foo")
+        reindex_project(task, db_request, "foo")
 
         assert parallel_bulk.calls == [pretend.call(es_client, docs, index="warehouse")]
         assert es_client.indices.create.calls == []