Flush cache after the current transaction is committed (#296)

waffle/models.py

@@ -6,7 +6,7 @@
 
 from django.conf import settings
 from django.contrib.auth.models import Group
-from django.db import models, router
+from django.db import models, router, transaction
 from django.utils import timezone
 from django.utils.encoding import python_2_unicode_compatible
 from django.utils.translation import ugettext_lazy as _
@@ -98,12 +98,18 @@ def flush(self):
     def save(self, *args, **kwargs):
         self.modified = timezone.now()
         ret = super(BaseModel, self).save(*args, **kwargs)
-        self.flush()
+        if hasattr(transaction, 'on_commit'):
+            transaction.on_commit(self.flush)
+        else:
+            self.flush()
         return ret
 
     def delete(self, *args, **kwargs):
         ret = super(BaseModel, self).delete(*args, **kwargs)
-        self.flush()
+        if hasattr(transaction, 'on_commit'):
+            transaction.on_commit(self.flush)
+        else:
+            self.flush()
         return ret
 
 

waffle/tests/test_testutils.py

@@ -3,14 +3,14 @@
 from decimal import Decimal
 
 from django.contrib.auth.models import AnonymousUser
-from django.test import TestCase, RequestFactory
+from django.test import TransactionTestCase, RequestFactory
 
 import waffle
 from waffle.models import Switch, Flag, Sample
 from waffle.testutils import override_switch, override_flag, override_sample
 
 
-class OverrideSwitchTests(TestCase):
+class OverrideSwitchTests(TransactionTestCase):
     def test_switch_existed_and_was_active(self):
         Switch.objects.create(name='foo', active=True)
 
@@ -94,7 +94,7 @@ def req():
     return r
 
 
-class OverrideFlagTests(TestCase):
+class OverrideFlagTests(TransactionTestCase):
     def test_flag_existed_and_was_active(self):
         Flag.objects.create(name='foo', everyone=True)
 
@@ -140,7 +140,7 @@ def test_flag_did_not_exist(self):
         assert not Flag.objects.filter(name='foo').exists()
 
 
-class OverrideSampleTests(TestCase):
+class OverrideSampleTests(TransactionTestCase):
     def test_sample_existed_and_was_100(self):
         Sample.objects.create(name='foo', percent='100.0')
 
@@ -190,7 +190,7 @@ def test_sample_did_not_exist(self):
 
 
 @override_switch('foo', active=False)
-class OverrideSwitchOnClassTests(TestCase):
+class OverrideSwitchOnClassTests(TransactionTestCase):
     def setUp(self):
         assert not Switch.objects.filter(name='foo').exists()
         Switch.objects.create(name='foo', active=True)
@@ -200,7 +200,7 @@ def test_undecorated_method_is_set_properly_for_switch(self):
 
 
 @override_flag('foo', active=False)
-class OverrideFlagOnClassTests(TestCase):
+class OverrideFlagOnClassTests(TransactionTestCase):
     def setUp(self):
         assert not Flag.objects.filter(name='foo').exists()
         Flag.objects.create(name='foo', everyone=True)
@@ -210,7 +210,7 @@ def test_undecorated_method_is_set_properly_for_flag(self):
 
 
 @override_sample('foo', active=False)
-class OverrideSampleOnClassTests(TestCase):
+class OverrideSampleOnClassTests(TransactionTestCase):
     def setUp(self):
         assert not Sample.objects.filter(name='foo').exists()
         Sample.objects.create(name='foo', percent='100.0')

waffle/tests/test_waffle.py

@@ -1,10 +1,13 @@
 from __future__ import unicode_literals
 
 import random
+import threading
+import unittest
 
+from django.conf import settings
 from django.contrib.auth.models import AnonymousUser, Group, User
-from django.db import connection
-from django.test import RequestFactory
+from django.db import connection, transaction
+from django.test import RequestFactory, TransactionTestCase
 from django.test.utils import override_settings
 
 import mock
@@ -394,3 +397,47 @@ def test_read_from_write_db(self):
             # The next read should now be directed to the write DB, ensuring
             # the cache and DB are in sync.
             assert waffle.sample_is_active(sample.name)
+
+
+class TransactionTests(TransactionTestCase):
+    @unittest.skipIf('sqlite3' in settings.DATABASES['default']['ENGINE'],
+                     'This test uses threads, which the sqlite3 DB engine '
+                     'does not support.')
+    def test_update_switch_in_transaction(self):
+        """Wait to invalidate the cache until after the current transaction."""
+
+        switch_name = 'transaction-switch-name'
+        switch = Switch.objects.create(name=switch_name, active=False)
+        self.addCleanup(switch.delete)
+
+        switch_written_in_background_thread = threading.Event()
+        switch_read_in_main_thread = threading.Event()
+
+        @transaction.atomic
+        def update_switch():
+            switch.active = True
+            switch.save()
+
+            # Signal to the main thread that the switch has been updated, but
+            # the transaction is not yet committed.
+            switch_written_in_background_thread.set()
+
+            # Pause here to allow the main thread to make an assertion.
+            switch_read_in_main_thread.wait(timeout=1)
+
+        # Start a background thread to update the switch in a transaction.
+        t = threading.Thread(target=update_switch)
+        t.daemon = True
+        t.start()
+
+        # After the switch is updated but before the transaction is committed,
+        # the cache will still have the previous value.
+        switch_written_in_background_thread.wait(timeout=1)
+        assert not waffle.switch_is_active(switch_name)
+
+        # After the transaction is committed, the cache should have been
+        # invalidated, hence the next call to switch_is_active should have the
+        # correct value.
+        switch_read_in_main_thread.set()
+        t.join(timeout=1)
+        assert waffle.switch_is_active(switch_name)