From 8fd0e07a1c4990e98e9e3fb7fa98028f0dad4b5d Mon Sep 17 00:00:00 2001
From: Thomas Waldmann <tw@waldmann-edv.de>
Date: Thu, 29 Sep 2016 04:41:59 +0200
Subject: [PATCH] hashindex: fix iterator implementation

NSKeyIterator and ChunkKeyIterator raised StopIteration once only when they reached their end.

But they did not raise StopIteration if one called next() again after they were exhausted,
so they did not comply to the standard iterator protocol.

AFAIK, this did not cause actual problems due to the way these iterators are used,
but when I tried to use itertools.islice() to get n-long sequences from these iterators,
it failed / went into an endless loop.
---
 borg/hashindex.pyx          | 10 ++++++++++
 borg/testsuite/hashindex.py |  5 ++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/borg/hashindex.pyx b/borg/hashindex.pyx
index c32c4dd1a..09ec89614 100644
--- a/borg/hashindex.pyx
+++ b/borg/hashindex.pyx
@@ -168,17 +168,22 @@ cdef class NSKeyIterator:
     cdef HashIndex *index
     cdef const void *key
     cdef int key_size
+    cdef int exhausted
 
     def __cinit__(self, key_size):
         self.key = NULL
         self.key_size = key_size
+        self.exhausted = 0
 
     def __iter__(self):
         return self
 
     def __next__(self):
+        if self.exhausted:
+            raise StopIteration
         self.key = hashindex_next_key(self.index, <char *>self.key)
         if not self.key:
+            self.exhausted = 1
             raise StopIteration
         cdef uint32_t *value = <uint32_t *>(self.key + self.key_size)
         cdef uint32_t segment = _le32toh(value[0])
@@ -330,17 +335,22 @@ cdef class ChunkKeyIterator:
     cdef HashIndex *index
     cdef const void *key
     cdef int key_size
+    cdef int exhausted
 
     def __cinit__(self, key_size):
         self.key = NULL
         self.key_size = key_size
+        self.exhausted = 0
 
     def __iter__(self):
         return self
 
     def __next__(self):
+        if self.exhausted:
+            raise StopIteration
         self.key = hashindex_next_key(self.index, <char *>self.key)
         if not self.key:
+            self.exhausted = 1
             raise StopIteration
         cdef uint32_t *value = <uint32_t *>(self.key + self.key_size)
         cdef uint32_t refcount = _le32toh(value[0])
diff --git a/borg/testsuite/hashindex.py b/borg/testsuite/hashindex.py
index 629ae4e57..5fd9e304e 100644
--- a/borg/testsuite/hashindex.py
+++ b/borg/testsuite/hashindex.py
@@ -83,8 +83,11 @@ class HashIndexTestCase(BaseTestCase):
         idx = NSIndex()
         for x in range(100):
             idx[H(x)] = x, x
-        all = list(idx.iteritems())
+        iterator = idx.iteritems()
+        all = list(iterator)
         self.assert_equal(len(all), 100)
+        # iterator is already exhausted by list():
+        self.assert_raises(StopIteration, next, iterator)
         second_half = list(idx.iteritems(marker=all[49][0]))
         self.assert_equal(len(second_half), 50)
         self.assert_equal(second_half, all[50:])