Handle UnpicklingError in persistent scheduler initialization (celery#9952)

* Handle UnpicklingError in persistent scheduler initialization

* Update t/unit/app/test_beat.py


* Fix incorrect comment referring to dbm.error instead of UnpicklingError

* Fix broken tests

* Fix lint offense

* Dummy commit to trigger pipeline

---------

Co-authored-by: Asif Saif Uddin {"Auvi":"অভি"} <[email protected]>

Author: pakallis

celery/beat.py

@@ -12,6 +12,7 @@
 from calendar import timegm
 from collections import namedtuple
 from functools import total_ordering
+from pickle import UnpicklingError
 from threading import Event, Thread
 
 from billiard import ensure_multiprocessing
@@ -569,11 +570,11 @@ def _create_schedule(self):
         for _ in (1, 2):
             try:
                 self._store['entries']
-            except (KeyError, UnicodeDecodeError, TypeError):
+            except (KeyError, UnicodeDecodeError, TypeError, UnpicklingError):
                 # new schedule db
                 try:
                     self._store['entries'] = {}
-                except (KeyError, UnicodeDecodeError, TypeError) + dbm.error as exc:
+                except (KeyError, UnicodeDecodeError, TypeError, UnpicklingError) + dbm.error as exc:
                     self._store = self._destroy_open_corrupted_schedule(exc)
                     continue
             else:

t/unit/app/test_beat.py

@@ -1,5 +1,6 @@
 import dbm
 import errno
+import pickle
 import sys
 from datetime import datetime, timedelta, timezone
 from pickle import dumps, loads
@@ -708,6 +709,26 @@ def test_create_schedule_corrupted_dbm_error(self):
         s._create_schedule()
         s._destroy_open_corrupted_schedule.assert_called_with(expected_error)
 
+    def test_create_schedule_corrupted_pickle_error(self):
+        """
+        Test that any UnpicklingError that might happen when opening beat-schedule.db is caught
+        """
+        s = create_persistent_scheduler()[0](app=self.app,
+                                             schedule_filename='schedule')
+        s._store = MagicMock()
+        s._destroy_open_corrupted_schedule = Mock()
+        s._destroy_open_corrupted_schedule.return_value = MagicMock()
+
+        # self._store['entries'] = {} will throw a pickle.UnpicklingError
+        s._store.__getitem__.side_effect = pickle.UnpicklingError("test")
+        # then, when _create_schedule tries to reset _store['entries'],
+        # throw another error, specifically pickle.UnpicklingError
+        expected_error = pickle.UnpicklingError("test")
+        s._store.__setitem__.side_effect = expected_error
+
+        s._create_schedule()
+        s._destroy_open_corrupted_schedule.assert_called_with(expected_error)
+
     def test_create_schedule_missing_entries(self):
         """
         Test that if _create_schedule can't find the key "entries" in _store it will recreate it