From 264545eb9dd5089e711cbc228a60573e6adabfd5 Mon Sep 17 00:00:00 2001 From: Mike Lang Date: Sun, 26 Nov 2023 16:02:52 +1100 Subject: [PATCH] CachedIterator: Fix bug where state can change while taking the lock Resulting in a case where we grab the wrong result, or even try to get the next item after the iterator has already been discarded. --- common/common/cached_iterator.py | 43 +++++++++++++++++--------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/common/common/cached_iterator.py b/common/common/cached_iterator.py index bc3011c..1ea218f 100644 --- a/common/common/cached_iterator.py +++ b/common/common/cached_iterator.py @@ -25,28 +25,31 @@ class CachedIterator(): # If we're more than 1 beyond the end, something has gone horribly wrong. # We should've already lengthened it last iteration assert len(self.items) == i, "CachedIterator logic error: {} != {}".format(len(self.items), i) - # Check if the iterator is still active. If not, we've reached the end or an error. - if self.iterator is None: - if self.error is not None: - raise self.error - return # Note we don't need the lock up until now because we're only trying to be gevent-safe, # not thread-safe. Simple operations like checking lengths can't be interrupted. # However calling next on the iterator may cause a switch. with self.lock: - try: - item = next(self.iterator) - except StopIteration: - # We've reached the end. Discard the iterator (in theory an iterator that - # has raised StopIteration once will keep raising it every time thereafter, - # but best not to rely on that). - self.iterator = None - # And we're done. - return - except Exception as ex: - # Discard the iterator and save the error to be re-served - self.iterator = None - self.error = ex - raise - self.items.append(item) + # Taking the lock may have also caused a switch, so we need to re-check + # our conditions. Someone else may have already added the item we need. + if len(self.items) <= i: + # Check if the iterator is still active. If not, we've reached the end or an error. + if self.iterator is None: + if self.error is not None: + raise self.error + return + try: + item = next(self.iterator) + except StopIteration: + # We've reached the end. Discard the iterator (in theory an iterator that + # has raised StopIteration once will keep raising it every time thereafter, + # but best not to rely on that). + self.iterator = None + # And we're done. + return + except Exception as ex: + # Discard the iterator and save the error to be re-served + self.iterator = None + self.error = ex + raise + self.items.append(item) yield self.items[i]