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.
pull/389/head
Mike Lang 1 year ago committed by Mike Lang
parent 0895ce44ff
commit 264545eb9d

@ -25,15 +25,18 @@ class CachedIterator():
# If we're more than 1 beyond the end, something has gone horribly wrong. # If we're more than 1 beyond the end, something has gone horribly wrong.
# We should've already lengthened it last iteration # We should've already lengthened it last iteration
assert len(self.items) == i, "CachedIterator logic error: {} != {}".format(len(self.items), i) assert len(self.items) == i, "CachedIterator logic error: {} != {}".format(len(self.items), i)
# 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:
# 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. # Check if the iterator is still active. If not, we've reached the end or an error.
if self.iterator is None: if self.iterator is None:
if self.error is not None: if self.error is not None:
raise self.error raise self.error
return 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: try:
item = next(self.iterator) item = next(self.iterator)
except StopIteration: except StopIteration:

Loading…
Cancel
Save