From 941d13ea3e144ea5d4e2acf5dd5c1a65574cc209 Mon Sep 17 00:00:00 2001
From: Jake Moshenko <jake.moshenko@coreos.com>
Date: Tue, 10 Nov 2015 16:14:44 -0500
Subject: [PATCH] Fix an off by one error in the common backfill code

---
 test/test_backfill_allocator.py | 27 ++++++++++++++++++++++++++-
 util/migrate/allocator.py       | 12 ++++++------
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/test/test_backfill_allocator.py b/test/test_backfill_allocator.py
index 40000cdb5..54258cb35 100644
--- a/test/test_backfill_allocator.py
+++ b/test/test_backfill_allocator.py
@@ -3,7 +3,7 @@ import logging
 import random
 
 from datetime import datetime, timedelta
-from util.migrate.allocator import CompletedKeys, NoAvailableKeysError
+from util.migrate.allocator import CompletedKeys, NoAvailableKeysError, yield_random_entries
 
 
 class CompletedTestCase(unittest.TestCase):
@@ -127,6 +127,31 @@ class CompletedTestCase(unittest.TestCase):
     self.assertGreater(iterations, 1024)
 
 
+class FakeQuery(object):
+  def __init__(self, result_list):
+    self._result_list = result_list
+
+  def limit(self, *args, **kwargs):
+    return self
+
+  def where(self, *args, **kwargs):
+    return self
+
+  def __iter__(self):
+    return self._result_list.__iter__()
+
+
+class QueryAllocatorTest(unittest.TestCase):
+  FAKE_PK_FIELD = 10  # Must be able to compare to integers
+
+  def test_no_work(self):
+    def create_empty_query():
+      return FakeQuery([])
+
+    for _ in yield_random_entries(create_empty_query, self.FAKE_PK_FIELD, 1, 10):
+      self.fail('There should never be any actual work!')
+
+
 if __name__ == '__main__':
   logging.basicConfig(level=logging.DEBUG)
   unittest.main()
diff --git a/util/migrate/allocator.py b/util/migrate/allocator.py
index 999416ace..64fea9ea4 100644
--- a/util/migrate/allocator.py
+++ b/util/migrate/allocator.py
@@ -136,21 +136,21 @@ def yield_random_entries(batch_query, primary_key_field, batch_size, max_id):
                             .where(primary_key_field >= start_index))
 
       if len(all_candidates) == 0:
-        logger.debug('No candidates, new highest id: %s', start_index)
-        allocator.mark_completed(start_index, max_id)
+        logger.info('No candidates, new highest id: %s', start_index)
+        allocator.mark_completed(start_index, max_id + 1)
         continue
 
-      logger.debug('Found %s candidates, processing block')
+      logger.info('Found %s candidates, processing block')
       for candidate in all_candidates:
         abort_early = Event()
         yield candidate, abort_early
         if abort_early.is_set():
-          logger.debug('Overlap with another worker, aborting')
+          logger.info('Overlap with another worker, aborting')
           break
 
       completed_through = candidate.id + 1
-      logger.debug('Marking id range as completed: %s-%s', start_index, completed_through)
+      logger.info('Marking id range as completed: %s-%s', start_index, completed_through)
       allocator.mark_completed(start_index, completed_through)
 
   except NoAvailableKeysError:
-    logger.debug('No more work')
+    logger.info('No more work')