fix: scheduling race with the QueryScheduler (#1347)

bdraco · web-flow · commit cf40470b89f9 · 2023-12-15T17:39:51.000-10:00
diff --git a/src/zeroconf/_services/browser.pxd b/src/zeroconf/_services/browser.pxd
@@ -73,8 +73,6 @@ cdef class QueryScheduler:
     cdef double _clock_resolution_millis
     cdef object _question_type
 
-    cpdef void schedule_ptr_first_refresh(self, DNSPointer pointer)
-
     cdef void _schedule_ptr_refresh(self, DNSPointer pointer, double expire_time_millis, double refresh_time_millis)
 
     cdef void _schedule_ptr_query(self, _ScheduledPTRQuery scheduled_query)
diff --git a/src/zeroconf/_services/browser.py b/src/zeroconf/_services/browser.py
@@ -374,12 +374,6 @@ def stop(self) -> None:
         self._next_scheduled_for_alias.clear()
         self._query_heap.clear()
 
-    def schedule_ptr_first_refresh(self, pointer: DNSPointer) -> None:
-        """Schedule a query for a pointer."""
-        expire_time_millis = pointer.get_expiration_time(100)
-        refresh_time_millis = pointer.get_expiration_time(_EXPIRE_REFRESH_TIME_PERCENT)
-        self._schedule_ptr_refresh(pointer, expire_time_millis, refresh_time_millis)
-
     def _schedule_ptr_refresh(
         self, pointer: DNSPointer, expire_time_millis: float_, refresh_time_millis: float_
     ) -> None:
@@ -415,6 +409,7 @@ def reschedule_ptr_first_refresh(self, pointer: DNSPointer) -> None:
             ):
                 return
             current.cancelled = True
+            del self._next_scheduled_for_alias[pointer.alias]
         expire_time_millis = pointer.get_expiration_time(100)
         self._schedule_ptr_refresh(pointer, expire_time_millis, refresh_time_millis)
 
@@ -490,10 +485,8 @@ def _process_ready_types(self) -> None:
             if query.when_millis > end_time_millis:
                 next_scheduled = query
                 break
-
+            query = heappop(self._query_heap)
             ready_types.add(query.name)
-
-            heappop(self._query_heap)
             del self._next_scheduled_for_alias[query.alias]
             # If there is still more than 10% of the TTL remaining
             # schedule a query again to try to rescue the record
@@ -670,7 +663,7 @@ def async_update_records(self, zc: 'Zeroconf', now: float_, records: List[Record
                 for type_ in self.types.intersection(cached_possible_types(pointer.name)):
                     if old_record is None:
                         self._enqueue_callback(SERVICE_STATE_CHANGE_ADDED, type_, pointer.alias)
-                        self.query_scheduler.schedule_ptr_first_refresh(pointer)
+                        self.query_scheduler.reschedule_ptr_first_refresh(pointer)
                     elif pointer.is_expired(now):
                         self._enqueue_callback(SERVICE_STATE_CHANGE_REMOVED, type_, pointer.alias)
                         self.query_scheduler.cancel_ptr_refresh(pointer)
diff --git a/tests/services/test_browser.py b/tests/services/test_browser.py
@@ -1185,15 +1185,15 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
             "disappear._hap._tcp.local.",
         )
 
-        query_scheduler.schedule_ptr_first_refresh(ptr_record)
+        query_scheduler.reschedule_ptr_first_refresh(ptr_record)
         expected_when_time = ptr_record.get_expiration_time(const._EXPIRE_REFRESH_TIME_PERCENT)
         expected_expire_time = ptr_record.get_expiration_time(100)
         ptr_query = _ScheduledPTRQuery(
             ptr_record.alias, ptr_record.name, int(ptr_record.ttl), expected_expire_time, expected_when_time
         )
         assert query_scheduler._query_heap == [ptr_query]
 
-        query_scheduler.schedule_ptr_first_refresh(ptr2_record)
+        query_scheduler.reschedule_ptr_first_refresh(ptr2_record)
         expected_when_time = ptr2_record.get_expiration_time(const._EXPIRE_REFRESH_TIME_PERCENT)
         expected_expire_time = ptr2_record.get_expiration_time(100)
         ptr2_query = _ScheduledPTRQuery(
@@ -1268,7 +1268,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
             "zoomer._hap._tcp.local.",
         )
 
-        query_scheduler.schedule_ptr_first_refresh(ptr_record)
+        query_scheduler.reschedule_ptr_first_refresh(ptr_record)
         expected_when_time = ptr_record.get_expiration_time(const._EXPIRE_REFRESH_TIME_PERCENT)
         expected_expire_time = ptr_record.get_expiration_time(100)
         ptr_query = _ScheduledPTRQuery(
