Do not change partition table when adding node and reimplement pt.tweak()

neo/admin/handler.py

@@ -78,6 +78,7 @@ class AdminEventHandler(EventHandler):
         conn.answer(Packets.AnswerPrimary(master_node.getUUID()))
 
     addPendingNodes = forward_ask(Packets.AddPendingNodes)
+    tweakPartitionTable = forward_ask(Packets.TweakPartitionTable)
     setClusterState = forward_ask(Packets.SetClusterState)
     checkReplicas = forward_ask(Packets.CheckReplicas)
 

neo/lib/protocol.py

@@ -26,7 +26,7 @@ except ImportError:
     pass
 
 # The protocol version (major, minor).
-PROTOCOL_VERSION = (11, 1)
+PROTOCOL_VERSION = (12, 1)
 
 # Size restrictions.
 MIN_PACKET_SIZE = 10
@@ -1150,6 +1150,16 @@ class AddPendingNodes(Packet):
 
     _answer = Error
 
+class TweakPartitionTable(Packet):
+    """
+    Ask the primary to optimize the partition table. A -> PM.
+    """
+    _fmt = PStruct('tweak_partition_table',
+        PFUUIDList,
+    )
+
+    _answer = Error
+
 class NotifyNodeInformation(Packet):
     """
     Notify information about one or more nodes. PM -> Any.
@@ -1666,6 +1676,8 @@ class Packets(dict):
                     SetNodeState, ignore_when_closed=False)
     AddPendingNodes = register(
                     AddPendingNodes, ignore_when_closed=False)
+    TweakPartitionTable = register(
+                    TweakPartitionTable, ignore_when_closed=False)
     AskNodeInformation, AnswerNodeInformation = register(
                     NodeInformation)
     SetClusterState = register(

neo/master/handlers/administration.py

@@ -130,37 +130,40 @@ class AdministrationHandler(MasterHandler):
         uuids = ', '.join(map(uuid_str, uuid_list))
         logging.debug('Add nodes %s', uuids)
         app = self.app
-        nm = app.nm
-        em = app.em
-        pt = app.pt
-        cell_list = []
-        uuid_set = set()
-        if app.getClusterState() == ClusterStates.RUNNING:
+        state = app.getClusterState()
+        # XXX: Would it be safe to allow more states ?
+        if state not in (ClusterStates.RUNNING,
+                         ClusterStates.STARTING_BACKUP,
+                         ClusterStates.BACKINGUP):
+            raise ProtocolError('Can not add nodes in %s state' % state)
         # take all pending nodes
-            for node in nm.getStorageList():
-                if node.isPending():
-                    uuid_set.add(node.getUUID())
-            # keep only selected nodes
-            if uuid_list:
-                uuid_set = uuid_set.intersection(set(uuid_list))
-        # nothing to do
-        if not uuid_set:
-            logging.warning('No nodes added')
-            conn.answer(Errors.Ack('No nodes added'))
-            return
-        uuids = ', '.join(map(uuid_str, uuid_set))
-        logging.info('Adding nodes %s', uuids)
-        # switch nodes to running state
-        node_list = map(nm.getByUUID, uuid_set)
+        node_list = list(app.pt.addNodeList(node
+            for node in app.nm.getStorageList()
+            if node.isPending() and node.getUUID() in uuid_list))
+        if node_list:
+            p = Packets.StartOperation()
             for node in node_list:
-            new_cells = pt.addNode(node)
-            cell_list.extend(new_cells)
                 node.setRunning()
-            node.getConnection().notify(Packets.StartOperation())
+                node.notify(p)
             app.broadcastNodesInformation(node_list)
-        # broadcast the new partition table
-        app.broadcastPartitionChanges(cell_list)
-        conn.answer(Errors.Ack('Nodes added: %s' % (uuids, )))
+            conn.answer(Errors.Ack('Nodes added: %s' %
+                ', '.join(uuid_str(x.getUUID()) for x in node_list)))
+        else:
+            logging.warning('No node added')
+            conn.answer(Errors.Ack('No node added'))
+
+    def tweakPartitionTable(self, conn, uuid_list):
+        app = self.app
+        state = app.getClusterState()
+        # XXX: Would it be safe to allow more states ?
+        if state not in (ClusterStates.RUNNING,
+                         ClusterStates.STARTING_BACKUP,
+                         ClusterStates.BACKINGUP):
+            raise ProtocolError('Can not tweak partition table in %s state'
+                                % state)
+        app.broadcastPartitionChanges(app.pt.tweak(
+            map(app.nm.getByUUID, uuid_list)))
+        conn.answer(Errors.Ack(''))
 
     def checkReplicas(self, conn, partition_dict, min_tid, max_tid):
         app = self.app

neo/master/pt.py

@@ -35,6 +35,16 @@ class Cell(neo.lib.pt.Cell):
 neo.lib.pt.Cell = Cell
 
 
+class MappedNode(object):
+
+    def __init__(self, node):
+        self.node = node
+        self.assigned = set()
+
+    def __getattr__(self, attr):
+        return getattr(self.node, attr)
+
+
 class PartitionTable(neo.lib.pt.PartitionTable):
     """This class manages a partition table for the primary master node"""
 
@@ -164,130 +174,101 @@ class PartitionTable(neo.lib.pt.PartitionTable):
 
         return cell_list
 
-    def addNode(self, node):
-        """Add a node. Take it into account that it might not be really a new
-        node. The strategy is, if a row does not contain a good number of
-        cells, add this node to the row, unless the node is already present
-        in the same row. Otherwise, check if this node should replace another
-        cell."""
-        cell_list = []
-        node_count = self.count_dict.get(node, 0)
-        for offset, row in enumerate(self.partition_list):
-            max_count = 0
-            max_cell = None
-            num_cells = 0
-            for cell in row:
-                if cell.getNode() is node:
-                    break
-                if not cell.isFeeding():
-                    num_cells += 1
-                    count = self.count_dict[cell.getNode()]
-                    if count > max_count:
-                        max_count = count
-                        max_cell = cell
-
-            else:
-                if self.nr < num_cells:
-                    if node_count + 1 >= max_count:
-                        continue
-                    if max_cell.isReadable():
-                        max_cell.setState(CellStates.FEEDING)
-                        cell_list.append((offset, max_cell.getUUID(),
-                                          CellStates.FEEDING))
-                    else:
-                        row.remove(max_cell)
-                        cell_list.append((offset, max_cell.getUUID(),
-                                          CellStates.DISCARDED))
-                    self.count_dict[max_cell.getNode()] -= 1
-                row.append(Cell(node, CellStates.OUT_OF_DATE))
-                cell_list.append((offset, node.getUUID(),
-                                  CellStates.OUT_OF_DATE))
-                node_count += 1
-
-        self.count_dict[node] = node_count
-        self.log()
-        return cell_list
+    def addNodeList(self, node_list):
+        """Add nodes"""
+        added_list = []
+        for node in node_list:
+            if node not in self.count_dict:
+                self.count_dict[node] = 0
+                added_list.append(node)
+        return added_list
 
-    def tweak(self):
-        """Test if nodes are distributed uniformly. Otherwise, correct the
-        partition table."""
-        changed_cell_list = []
+    def tweak(self, drop_list=()):
+        """Optimize partition table
 
+        This is done by computing a minimal diff between current partition table
+        and what make() would do.
+        """
+        assigned_dict = dict((x, {}) for x in self.count_dict)
+        readable_dict = dict((i, set()) for i in xrange(self.np))
         for offset, row in enumerate(self.partition_list):
-            removed_cell_list = []
-            feeding_cell = None
-            out_of_date_cell_list = []
-            up_to_date_cell_list = []
             for cell in row:
-                if cell.getNode().isBroken():
-                    # Remove a broken cell.
-                    removed_cell_list.append(cell)
-                elif cell.isFeeding():
-                    if feeding_cell is None:
-                        feeding_cell = cell
-                    else:
-                        # Remove an excessive feeding cell.
-                        removed_cell_list.append(cell)
-                elif cell.isUpToDate():
-                    up_to_date_cell_list.append(cell)
-                else:
-                    out_of_date_cell_list.append(cell)
-
-            # If all cells are up-to-date, a feeding cell is not required.
-            if len(out_of_date_cell_list) == 0 and feeding_cell is not None:
-                removed_cell_list.append(feeding_cell)
-
-            ideal_num = self.nr + 1
-            while len(out_of_date_cell_list) + len(up_to_date_cell_list) > \
-                    ideal_num:
-                # This row contains too many cells.
-                if len(up_to_date_cell_list) > 1:
-                    # There are multiple up-to-date cells, so choose whatever
-                    # used too much.
-                    cell_list = out_of_date_cell_list + up_to_date_cell_list
-                else:
-                    # Drop an out-of-date cell.
-                    cell_list = out_of_date_cell_list
-
-                max_count = 0
-                chosen_cell = None
-                for cell in cell_list:
-                    count = self.count_dict[cell.getNode()]
-                    if max_count < count:
-                        max_count = count
-                        chosen_cell = cell
-                removed_cell_list.append(chosen_cell)
-                try:
-                    out_of_date_cell_list.remove(chosen_cell)
-                except ValueError:
-                    up_to_date_cell_list.remove(chosen_cell)
-
-            # Now remove cells really.
-            for cell in removed_cell_list:
-                row.remove(cell)
-                if not cell.isFeeding():
-                    self.count_dict[cell.getNode()] -= 1
-                changed_cell_list.append((offset, cell.getUUID(),
-                    CellStates.DISCARDED))
-
-        # Add cells, if a row contains less than the number of replicas.
-        for offset, row in enumerate(self.partition_list):
-            num_cells = 0
+                if cell.isReadable():
+                    readable_dict[offset].add(cell)
+                assigned_dict[cell.getNode()][offset] = cell
+        pt = PartitionTable(self.np, self.nr)
+        drop_list = set(x for x in drop_list if x in assigned_dict)
+        node_set = set(MappedNode(x) for x in assigned_dict
+                                     if x not in drop_list)
+        pt.make(node_set)
+        for offset, row in enumerate(pt.partition_list):
             for cell in row:
+                if cell.isReadable():
+                    cell.getNode().assigned.add(offset)
+        def map_nodes():
+            node_list = []
+            for node, assigned in assigned_dict.iteritems():
+                if node in drop_list:
+                    yield node, frozenset()
+                    continue
+                readable = set(offset for offset, cell in assigned.iteritems()
+                                      if cell.isReadable())
+                # the criterion on UUID is purely cosmetic
+                node_list.append((len(readable), len(assigned),
+                                  -node.getUUID(), readable, node))
+            node_list.sort(reverse=1)
+            for _, _, _, readable, node in node_list:
+                assigned = assigned_dict[node]
+                mapped = min(node_set, key=lambda m: (
+                    len(m.assigned.symmetric_difference(assigned)),
+                    len(m.assigned ^ readable)))
+                node_set.remove(mapped)
+                yield node, mapped.assigned
+            assert not node_set
+        changed_list = []
+        uptodate_set = set()
+        remove_dict = dict((i, []) for i in xrange(self.np))
+        for node, mapped in map_nodes():
+            uuid = node.getUUID()
+            assigned = assigned_dict[node]
+            for offset, cell in assigned.iteritems():
+                if offset in mapped:
+                    if cell.isReadable():
+                        uptodate_set.add(offset)
+                        readable_dict[offset].remove(cell)
+                        if cell.isFeeding():
+                            self.count_dict[node] += 1
+                            state = CellStates.UP_TO_DATE
+                            cell.setState(state)
+                            changed_list.append((offset, uuid, state))
+                else:
                     if not cell.isFeeding():
-                    num_cells += 1
-            while num_cells <= self.nr:
-                node = self.findLeastUsedNode([cell.getNode() for cell in row])
-                if node is None:
-                    break
-                row.append(Cell(node, CellStates.OUT_OF_DATE))
-                changed_cell_list.append((offset, node.getUUID(),
-                    CellStates.OUT_OF_DATE))
+                        self.count_dict[node] -= 1
+                    remove_dict[offset].append(cell)
+            for offset in mapped.difference(assigned):
                 self.count_dict[node] += 1
-                num_cells += 1
-
-        self.log()
-        return changed_cell_list
+                state = CellStates.OUT_OF_DATE
+                self.partition_list[offset].append(Cell(node, state))
+                changed_list.append((offset, uuid, state))
+        count_dict = self.count_dict.copy()
+        for offset, cell_list in remove_dict.iteritems():
+            if not cell_list:
+                continue
+            row = self.partition_list[offset]
+            feeding = None if offset in uptodate_set else min(
+                readable_dict[offset], key=lambda x: count_dict[x.getNode()])
+            for cell in cell_list:
+                if cell is feeding:
+                    count_dict[cell.getNode()] += 1
+                    if cell.isFeeding():
+                        continue
+                    state = CellStates.FEEDING
+                    cell.setState(state)
+                else:
+                    state = CellStates.DISCARDED
+                    row.remove(cell)
+                changed_list.append((offset, cell.getUUID(), state))
+        return changed_list
 
     def outdate(self, lost_node=None):
         """Outdate all non-working nodes

neo/master/verification.py

@@ -93,7 +93,6 @@ class VerificationManager(BaseServiceHandler):
         return state, self
 
     def run(self):
-
         self.app.changeClusterState(ClusterStates.VERIFYING)
         while True:
             try:
@@ -102,14 +101,7 @@ class VerificationManager(BaseServiceHandler):
                 continue
             break
         # At this stage, all non-working nodes are out-of-date.
-        cell_list = self.app.pt.outdate()
-
-        # Tweak the partition table, if the distribution of storage nodes
-        # is not uniform.
-        cell_list.extend(self.app.pt.tweak())
-
-        # If anything changed, send the changes.
-        self.app.broadcastPartitionChanges(cell_list)
+        self.app.broadcastPartitionChanges(self.app.pt.outdate())
 
     def verifyData(self):
         """Verify the data in storage nodes and clean them up, if necessary."""

neo/neoctl/app.py

@@ -34,6 +34,7 @@ action_dict = {
     'check': 'checkReplicas',
     'start': 'startCluster',
     'add': 'enableStorageList',
+    'tweak': 'tweakPartitionTable',
     'drop': 'dropNode',
 }
 
@@ -169,6 +170,14 @@ class TerminalNeoCTL(object):
             uuid_list = map(self.asNode, params)
         return self.neoctl.enableStorageList(uuid_list)
 
+    def tweakPartitionTable(self, params):
+        """
+          Optimize partition table.
+          No partitition will be assigned to specified storage nodes.
+          Parameters: [node [...]]
+        """
+        return self.neoctl.tweakPartitionTable(map(self.asNode, params))
+
     def dropNode(self, params):
         """
           Set node into DOWN state.

neo/neoctl/neoctl.py

@@ -84,6 +84,12 @@ class NeoCTL(object):
             raise RuntimeError(response)
         return response[2]
 
+    def tweakPartitionTable(self, uuid_list=()):
+        response = self.__ask(Packets.TweakPartitionTable(uuid_list))
+        if response[0] != Packets.Error or response[1] != ErrorCodes.ACK:
+            raise RuntimeError(response)
+        return response[2]
+
     def setClusterState(self, state):
         """
           Set cluster state.

neo/tests/functional/testStorage.py

@@ -160,6 +160,7 @@ class StorageTests(NEOFunctionalTest):
         # add it to the partition table
         self.neo.neoctl.enableStorageList([stopped[0].getUUID()])
         self.neo.expectRunning(stopped[0])
+        self.neo.neoctl.tweakPartitionTable()
         self.neo.expectAssignedCells(stopped[0], number=10)
         self.neo.expectClusterRunning()
 
@@ -298,6 +299,7 @@ class StorageTests(NEOFunctionalTest):
         stopped[0].start()
         self.neo.expectPending(stopped[0])
         self.neo.neoctl.enableStorageList([stopped[0].getUUID()])
+        self.neo.neoctl.tweakPartitionTable()
         self.neo.expectRunning(stopped[0])
         self.neo.expectClusterRunning()
         self.neo.expectOudatedCells(number=0)
@@ -356,6 +358,7 @@ class StorageTests(NEOFunctionalTest):
         stopped[0].start()
         self.neo.expectPending(stopped[0])
         self.neo.neoctl.enableStorageList([stopped[0].getUUID()])
+        self.neo.neoctl.tweakPartitionTable()
         self.neo.expectRunning(stopped[0])
         self.neo.expectClusterRunning()
         self.neo.expectAssignedCells(started[0], 10)
@@ -469,6 +472,7 @@ class StorageTests(NEOFunctionalTest):
         stopped[0].start()
         self.neo.expectPending(stopped[0])
         self.neo.neoctl.enableStorageList([stopped[0].getUUID()])
+        self.neo.neoctl.tweakPartitionTable()
         self.neo.expectRunning(stopped[0])
         self.neo.expectClusterRunning()
         self.neo.expectAssignedCells(started[0], 10)

neo/tests/master/testMasterPT.py

@@ -15,6 +15,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import unittest
+from collections import defaultdict
 from mock import Mock
 from .. import NeoUnitTestBase
 from neo.lib.protocol import NodeStates, CellStates
@@ -142,130 +143,6 @@ class MasterPartitionTableTests(NeoUnitTestBase):
         cell = cells[0]
         self.assertEqual(cell.getState(), CellStates.UP_TO_DATE)
 
-    def test_14_addNode(self):
-        num_partitions = 5
-        num_replicas = 2
-        pt = PartitionTable(num_partitions, num_replicas)
-        # add nodes
-        uuid1 = self.getStorageUUID()
-        server1 = ("127.0.0.1", 19001)
-        sn1 = StorageNode(Mock(), server1, uuid1)
-        # add it to an empty pt
-        cell_list = pt.addNode(sn1)
-        self.assertEqual(len(cell_list), 5)
-        # it must be added to all partitions
-        for x in xrange(num_partitions):
-            self.assertEqual(len(pt.getCellList(x)), 1)
-            self.assertEqual(pt.getCellList(x)[0].getState(), CellStates.OUT_OF_DATE)
-            self.assertEqual(pt.getCellList(x)[0].getNode(), sn1)
-        self.assertEqual(pt.count_dict[sn1], 5)
-        # add same node again, must remain the same
-        cell_list = pt.addNode(sn1)
-        self.assertEqual(len(cell_list), 0)
-        for x in xrange(num_partitions):
-            self.assertEqual(len(pt.getCellList(x)), 1)
-            self.assertEqual(pt.getCellList(x)[0].getState(), CellStates.OUT_OF_DATE)
-            self.assertEqual(pt.getCellList(x)[0].getNode(), sn1)
-        self.assertEqual(pt.count_dict[sn1], 5)
-        # add a second node to fill the partition table
-        uuid2 = self.getStorageUUID()
-        server2 = ("127.0.0.2", 19002)
-        sn2 = StorageNode(Mock(), server2, uuid2)
-        # add it
-        cell_list = pt.addNode(sn2)
-        self.assertEqual(len(cell_list), 5)
-        for x in xrange(num_partitions):
-            self.assertEqual(len(pt.getCellList(x)), 2)
-            self.assertEqual(pt.getCellList(x)[0].getState(), CellStates.OUT_OF_DATE)
-            self.assertTrue(pt.getCellList(x)[0].getNode() in (sn1, sn2))
-        # test the most used node is remove from some partition
-        uuid3 = self.getStorageUUID()
-        server3 = ("127.0.0.3", 19001)
-        sn3 = StorageNode(Mock(), server3, uuid3)
-        uuid4 = self.getStorageUUID()
-        server4 = ("127.0.0.4", 19001)
-        sn4 = StorageNode(Mock(), server4, uuid4)
-        uuid5 = self.getStorageUUID()
-        server5 = ("127.0.0.5", 1900)
-        sn5 = StorageNode(Mock(), server5, uuid5)
-        # partition looks like:
-        # 0 : sn1, sn2
-        # 1 : sn1, sn3
-        # 2 : sn1, sn4
-        # 3 : sn1, sn5
-        num_partitions = 4
-        num_replicas = 1
-        pt = PartitionTable(num_partitions, num_replicas)
-        # node most used is out of date, just dropped
-        pt.setCell(0, sn1, CellStates.OUT_OF_DATE)
-        pt.setCell(0, sn2, CellStates.UP_TO_DATE)
-        pt.setCell(1, sn1, CellStates.OUT_OF_DATE)
-        pt.setCell(1, sn3, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn1, CellStates.OUT_OF_DATE)
-        pt.setCell(2, sn4, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn1, CellStates.OUT_OF_DATE)
-        pt.setCell(3, sn5, CellStates.UP_TO_DATE)
-        uuid6 = self.getStorageUUID()
-        server6 = ("127.0.0.6", 19006)
-        sn6 = StorageNode(Mock(), server6, uuid6)
-        # sn1 is removed twice and sn6 is added twice
-        self.assertEqual(sorted(pt.addNode(sn6)), [
-            (0, uuid1, CellStates.DISCARDED),
-            (0, uuid6, CellStates.OUT_OF_DATE),
-            (1, uuid1, CellStates.DISCARDED),
-            (1, uuid6, CellStates.OUT_OF_DATE)])
-        for x in xrange(num_partitions):
-            self.assertEqual(len(pt.getCellList(x)), 2)
-        # there is a feeding cell, just dropped
-        pt.clear()
-        pt.setCell(0, sn1, CellStates.OUT_OF_DATE)
-        pt.setCell(0, sn2, CellStates.UP_TO_DATE)
-        pt.setCell(0, sn3, CellStates.FEEDING)
-        pt.setCell(1, sn1, CellStates.OUT_OF_DATE)
-        pt.setCell(1, sn2, CellStates.FEEDING)
-        pt.setCell(1, sn3, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn4, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn5, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn4, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn5, CellStates.UP_TO_DATE)
-        # sn1 is removed twice and sn6 is added twice
-        self.assertEqual(sorted(pt.addNode(sn6)), [
-            (0, uuid1, CellStates.DISCARDED),
-            (0, uuid6, CellStates.OUT_OF_DATE),
-            (1, uuid1, CellStates.DISCARDED),
-            (1, uuid6, CellStates.OUT_OF_DATE)])
-        for x in xrange(num_partitions):
-            self.assertEqual(len(pt.getCellList(x)), 3)
-        # there is no feeding cell, marked as feeding
-        pt.clear()
-        pt.setCell(0, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(0, sn2, CellStates.UP_TO_DATE)
-        pt.setCell(1, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(1, sn3, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn4, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn5, CellStates.UP_TO_DATE)
-        # sn1 is removed twice and sn6 is added twice
-        self.assertEqual(sorted(pt.addNode(sn6)), [
-            (0, uuid1, CellStates.FEEDING),
-            (0, uuid6, CellStates.OUT_OF_DATE),
-            (1, uuid1, CellStates.FEEDING),
-            (1, uuid6, CellStates.OUT_OF_DATE)])
-        pt.setUpToDate(sn6, 0)
-        pt.setUpToDate(sn6, 1)
-        for x in xrange(num_partitions):
-            self.assertEqual(len(pt.getCellList(x)), 2)
-        pt = PartitionTable(12, 0)
-        node_count = 0
-        for node in sn1, sn2, sn3, sn4:
-            pt.addNode(node)
-            node_count += 1
-            self.assertEqual(pt.count_dict.values(),
-                             [12/node_count] * node_count)
-
     def test_15_dropNode(self):
         num_partitions = 4
         num_replicas = 2
@@ -376,108 +253,100 @@ class MasterPartitionTableTests(NeoUnitTestBase):
         self.assertTrue(pt.filled())
         self.assertTrue(pt.operational())
 
-    def test_17_tweak(self):
-        # remove broken node
-        # remove if too many feeding nodes
-        # remove feeding if all cells are up to date
-        # if too many cells, remove most used cell
-        # if not enought cell, add least used node
+    def _pt_states(self, pt):
+        node_dict = defaultdict(list)
+        for offset, row in enumerate(pt.partition_list):
+            for cell in row:
+                state_list = node_dict[cell.getNode()]
+                if state_list:
+                    self.assertTrue(state_list[-1][0] < offset)
+                state_list.append((offset, str(cell.getState())[0]))
+        return map(dict, sorted(node_dict.itervalues()))
 
-        # create nodes
-        uuid1 = self.getStorageUUID()
-        server1 = ("127.0.0.1", 19001)
-        sn1 = StorageNode(Mock(), server1, uuid1, NodeStates.RUNNING)
-        uuid2 = self.getStorageUUID()
-        server2 = ("127.0.0.2", 19002)
-        sn2 = StorageNode(Mock(), server2, uuid2, NodeStates.RUNNING)
-        uuid3 = self.getStorageUUID()
-        server3 = ("127.0.0.3", 19003)
-        sn3 = StorageNode(Mock(), server3, uuid3, NodeStates.RUNNING)
-        uuid4 = self.getStorageUUID()
-        server4 = ("127.0.0.4", 19004)
-        sn4 = StorageNode(Mock(), server4, uuid4, NodeStates.RUNNING)
-        uuid5 = self.getStorageUUID()
-        server5 = ("127.0.0.5", 19005)
-        sn5 = StorageNode(Mock(), server5, uuid5, NodeStates.RUNNING)
-        # create partition table
-        # 0 : sn1(discarded), sn2(up), -> sn2 must remain
-        # 1 : sn1(feeding), sn2(feeding), sn3(up) -> one feeding and sn3 must remain
-        # 2 : sn1(feeding), sn2(up), sn3(up) -> sn2 and sn3 must remain, feeding must go away
-        # 3 : sn1(up), sn2(up), sn3(up), sn4(up) -> only 3 cell must remain
-        # 4 : sn1(up), sn5(up) -> one more cell must be added
-        num_partitions = 5
-        num_replicas = 2
-        pt = PartitionTable(num_partitions, num_replicas)
+    def checkPT(self, pt, exclude_empty=False):
+        new_pt = PartitionTable(pt.np, pt.nr)
+        new_pt.make(pt.getNodeList() if exclude_empty else pt.count_dict)
+        self.assertEqual(self._pt_states(pt), self._pt_states(new_pt))
+
+    def update(self, pt, change_list=None):
+        if change_list is None:
+            for offset, row in enumerate(pt.partition_list):
+                for cell in list(row):
+                    if cell.isOutOfDate():
+                        pt.setUpToDate(cell.getNode(), offset)
+        else:
+            node_dict = dict((x.getUUID(), x) for x in pt.count_dict)
+            for offset, uuid, state in change_list:
+                if state is CellStates.OUT_OF_DATE:
+                    pt.setUpToDate(node_dict[uuid], offset)
+
+    def tweak(self, pt, drop_list=()):
+        change_list = pt.tweak(drop_list)
+        self.assertFalse(pt.tweak(drop_list))
+        return change_list
+
+    def test_17_tweak(self):
+        sn = [StorageNode(Mock(), None, i + 1, NodeStates.RUNNING)
+              for i in xrange(5)]
+        pt = PartitionTable(5, 2)
         # part 0
-        pt.setCell(0, sn1, CellStates.DISCARDED)
-        pt.setCell(0, sn2, CellStates.UP_TO_DATE)
+        pt.setCell(0, sn[0], CellStates.DISCARDED)
+        pt.setCell(0, sn[1], CellStates.UP_TO_DATE)
         # part 1
-        pt.setCell(1, sn1, CellStates.FEEDING)
-        pt.setCell(1, sn2, CellStates.FEEDING)
-        pt.setCell(1, sn3, CellStates.OUT_OF_DATE)
+        pt.setCell(1, sn[0], CellStates.FEEDING)
+        pt.setCell(1, sn[1], CellStates.FEEDING)
+        pt.setCell(1, sn[2], CellStates.OUT_OF_DATE)
         # part 2
-        pt.setCell(2, sn1, CellStates.FEEDING)
-        pt.setCell(2, sn2, CellStates.UP_TO_DATE)
-        pt.setCell(2, sn3, CellStates.UP_TO_DATE)
+        pt.setCell(2, sn[0], CellStates.FEEDING)
+        pt.setCell(2, sn[1], CellStates.UP_TO_DATE)
+        pt.setCell(2, sn[2], CellStates.UP_TO_DATE)
         # part 3
-        pt.setCell(3, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn2, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn3, CellStates.UP_TO_DATE)
-        pt.setCell(3, sn4, CellStates.UP_TO_DATE)
+        pt.setCell(3, sn[0], CellStates.UP_TO_DATE)
+        pt.setCell(3, sn[1], CellStates.UP_TO_DATE)
+        pt.setCell(3, sn[2], CellStates.UP_TO_DATE)
+        pt.setCell(3, sn[3], CellStates.UP_TO_DATE)
         # part 4
-        pt.setCell(4, sn1, CellStates.UP_TO_DATE)
-        pt.setCell(4, sn5, CellStates.UP_TO_DATE)
-        # now tweak the table
-        pt.tweak()
-        # check part 1
-        cells =  pt.getCellList(0)
-        self.assertEqual(len(cells), 3)
-        for cell in cells:
-            self.assertNotEqual(cell.getState(), CellStates.DISCARDED)
-            if cell.getNode() == sn2:
-                self.assertEqual(cell.getState(), CellStates.UP_TO_DATE)
-            else:
-                self.assertEqual(cell.getState(), CellStates.OUT_OF_DATE)
-        self.assertTrue(sn2 in [x.getNode() for x in cells])
+        pt.setCell(4, sn[0], CellStates.UP_TO_DATE)
+        pt.setCell(4, sn[4], CellStates.UP_TO_DATE)
 
-        # check part 2
-        cells =  pt.getCellList(1)
-        self.assertEqual(len(cells), 4)
-        for cell in cells:
-            if cell.getNode() == sn1:
-                self.assertEqual(cell.getState(), CellStates.FEEDING)
-            else:
-                self.assertEqual(cell.getState(), CellStates.OUT_OF_DATE)
-        self.assertTrue(sn3 in [x.getNode() for x in cells])
-        self.assertTrue(sn1 in [x.getNode() for x in cells])
+        count_dict = defaultdict(int)
+        change_list = self.tweak(pt)
+        for offset, uuid, state in change_list:
+            count_dict[state] += 1
+        self.assertEqual(count_dict, {CellStates.DISCARDED: 3,
+                                      CellStates.OUT_OF_DATE: 5,
+                                      CellStates.UP_TO_DATE: 3})
+        self.update(pt, change_list)
+        self.checkPT(pt)
 
-        # check part 3
-        cells =  pt.getCellList(2)
-        self.assertEqual(len(cells), 3)
-        for cell in cells:
-            if cell.getNode() in (sn2, sn3):
-                self.assertEqual(cell.getState(), CellStates.UP_TO_DATE)
-            else:
-                self.assertEqual(cell.getState(), CellStates.OUT_OF_DATE)
-        self.assertTrue(sn3 in [x.getNode() for x in cells])
-        self.assertTrue(sn2 in [x.getNode() for x in cells])
+        for offset in pt.getAssignedPartitionList(sn[1].getUUID()):
+            pt.removeCell(offset, sn[1])
+        change_list = self.tweak(pt)
+        self.assertEqual(3, len(change_list))
+        self.update(pt, change_list)
+        self.checkPT(pt)
 
-        # check part 4
-        cells =  pt.getCellList(3)
-        self.assertEqual(len(cells), 3)
-        for cell in cells:
-            self.assertEqual(cell.getState(), CellStates.UP_TO_DATE)
+        for np, i in (12, 0), (12, 1), (13, 2):
+            pt = PartitionTable(np, i)
+            i += 1
+            pt.make(sn[:i])
+            for n in sn[i:i+3]:
+                self.assertEqual([n], pt.addNodeList([n]))
+                self.update(pt, self.tweak(pt))
+                self.checkPT(pt)
+            pt.clear()
+            pt.make(sn[:i])
+            for n in sn[i:i+3]:
+                self.assertEqual([n], pt.addNodeList([n]))
+                self.tweak(pt)
+            self.update(pt)
+            self.checkPT(pt)
 
-        # check part 5
-        cells =  pt.getCellList(4)
-        self.assertEqual(len(cells), 3)
-        for cell in cells:
-            if cell.getNode() in (sn1, sn5):
-                self.assertEqual(cell.getState(), CellStates.UP_TO_DATE)
-            else:
-                self.assertEqual(cell.getState(), CellStates.OUT_OF_DATE)
-        self.assertTrue(sn1 in [x.getNode() for x in cells])
-        self.assertTrue(sn5 in [x.getNode() for x in cells])
+        pt = PartitionTable(7, 0)
+        pt.make(sn[:1])
+        pt.addNodeList(sn[1:3])
+        self.update(pt, self.tweak(pt, sn[:1]))
+        self.checkPT(pt, True)
 
 
 if __name__ == '__main__':

neo/tests/threaded/testReplication.py

@@ -184,13 +184,12 @@ class ReplicationTests(NEOThreadedTest):
         another source.
         Here are the different states of partitions over time:
           pt: 0: U|U|U
-          pt: 0: UO|UO|UO
-          pt: 0: FOO|UO.|U.O # node 1 replicates from node 0
-          pt: 0: .OU|UO.|U.O # here node 0 lost partition 0
+          pt: 0: UO.|U.O|FOO
+          pt: 0: UU.|U.O|FOO
+          pt: 0: UU.|U.U|FOO # nodes 1 & 2 replicate from node 0
+          pt: 0: UU.|U.U|.OU # here node 0 lost partition 2
                              # and node 1 must switch to node 2
-          pt: 0: .OU|UO.|U.U
-          pt: 0: .OU|UU.|U.U
-          pt: 0: .UU|UU.|U.U
+          pt: 0: UU.|U.U|.UU
         """
         def connected(orig, *args, **kw):
             patch[0] = s1.filterConnection(s0)
@@ -218,6 +217,7 @@ class ReplicationTests(NEOThreadedTest):
                 s2.start()
                 cluster.tic()
                 cluster.neoctl.enableStorageList([s1.uuid, s2.uuid])
+                cluster.neoctl.tweakPartitionTable()
                 offset, = [offset for offset, row in enumerate(
                                       cluster.master.pt.partition_list)
                                   for cell in row if cell.isFeeding()]