Add ability to save builder data to a disk file

Instances of the RingBuilder class can store its data to a disk file by
the save method and load it by the load method.

blueprint argparse-in-swift-ring-builder

Change-Id: I69fdf0693ca9f520d235a795ecdd2da310dcd5d3

bin/swift-ring-builder

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import cPickle as pickle
 from array import array
 from errno import EEXIST
 from itertools import islice, izip
@@ -74,9 +73,8 @@ swift-ring-builder <builder_file> create <part_power> <replicas>
         except OSError, err:
             if err.errno != EEXIST:
                 raise
-        pickle.dump(builder.to_dict(), open(pathjoin(backup_dir,
+        builder.save(pathjoin(backup_dir, '%d.' % time() + basename(argv[1])))
-                    '%d.' % time() + basename(argv[1])), 'wb'), protocol=2)
+        builder.save(argv[1])
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
         exit(EXIT_SUCCESS)
 
     def default():
@@ -342,7 +340,7 @@ swift-ring-builder <builder_file> add
                 (region, zone, ip, port, device_name))[0]
             print('Device %s with %s weight got id %s' %
                   (format_device(new_dev), weight, new_dev['id']))
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
     def set_weight():
@@ -382,7 +380,7 @@ swift-ring-builder <builder_file> set_weight <search-value> <weight>
                 builder.set_dev_weight(dev['id'], weight)
                 print '%s weight set to %s' % (format_device(dev),
                                                dev['weight'])
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
     def set_info():
@@ -500,7 +498,7 @@ swift-ring-builder <builder_file> set_info
                     dev[key] = value
                 print 'Device %s is now %s' % (orig_dev_string,
                                                format_device(dev))
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
     def remove():
@@ -552,7 +550,7 @@ swift-ring-builder <builder_file> remove <search-value> [search-value ...]
 
                 print '%s marked for removal and will ' \
                       'be removed next rebalance.' % format_device(dev)
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
     def rebalance():
@@ -619,10 +617,9 @@ swift-ring-builder <builder_file> rebalance <seed>
         ts = time()
         builder.get_ring().save(
             pathjoin(backup_dir, '%d.' % ts + basename(ring_file)))
-        pickle.dump(builder.to_dict(), open(pathjoin(backup_dir,
+        builder.save(pathjoin(backup_dir, '%d.' % ts + basename(argv[1])))
-                    '%d.' % ts + basename(argv[1])), 'wb'), protocol=2)
         builder.get_ring().save(ring_file)
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(status)
 
     def validate():
@@ -656,7 +653,7 @@ swift-ring-builder <builder_file> write_ring
 
     def pretend_min_part_hours_passed():
         builder.pretend_min_part_hours_passed()
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
     def set_min_part_hours():
@@ -672,7 +669,7 @@ swift-ring-builder <builder_file> set_min_part_hours <hours>
         builder.change_min_part_hours(int(argv[3]))
         print 'The minimum number of hours before a partition can be ' \
               'reassigned is now set to %s' % argv[3]
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
     def set_replicas():
@@ -704,7 +701,7 @@ swift-ring-builder <builder_file> set_replicas <replicas>
         builder.set_replicas(new_replicas)
         print 'The replica count is now %.6f.' % builder.replicas
         print 'The change will take effect after the next rebalance.'
-        pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2)
+        builder.save(argv[1])
         exit(EXIT_SUCCESS)
 
 if __name__ == '__main__':

swift/common/ring/builder.py

@@ -1000,6 +1000,13 @@ class RingBuilder(object):
                     dev.setdefault('replication_port', dev['port'])
         return builder
 
+    def save(self, builder_file):
+        """Serialize this RingBuilder instance to disk.
+
+        :param builder_file: path to builder file to save
+        """
+        pickle.dump(self.to_dict(), open(builder_file, 'wb'), protocol=2)
+
     def search_devs(self, search_value):
         """
     The <search-value> can be of the form::

test/unit/common/ring/test_builder.py

@@ -13,13 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import mock
 import operator
 import os
 import unittest
 import cPickle as pickle
 from collections import defaultdict
 from shutil import rmtree
-from mock import Mock, call as mock_call
 
 from swift.common import exceptions
 from swift.common import ring
@@ -666,12 +666,12 @@ class TestRingBuilder(unittest.TestCase):
         real_pickle = pickle.load
         try:
             #test a legit builder
-            fake_pickle = Mock(return_value=rb)
+            fake_pickle = mock.Mock(return_value=rb)
-            fake_open = Mock(return_value=None)
+            fake_open = mock.Mock(return_value=None)
             pickle.load = fake_pickle
             builder = ring.RingBuilder.load('fake.builder', open=fake_open)
             self.assertEquals(fake_pickle.call_count, 1)
-            fake_open.assert_has_calls([mock_call('fake.builder', 'rb')])
+            fake_open.assert_has_calls([mock.call('fake.builder', 'rb')])
             self.assertEquals(builder, rb)
             fake_pickle.reset_mock()
             fake_open.reset_mock()
@@ -680,7 +680,7 @@ class TestRingBuilder(unittest.TestCase):
             fake_pickle.return_value = rb.to_dict()
             pickle.load = fake_pickle
             builder = ring.RingBuilder.load('fake.builder', open=fake_open)
-            fake_open.assert_has_calls([mock_call('fake.builder', 'rb')])
+            fake_open.assert_has_calls([mock.call('fake.builder', 'rb')])
             self.assertEquals(builder.devs, rb.devs)
             fake_pickle.reset_mock()
             fake_open.reset_mock()
@@ -692,12 +692,63 @@ class TestRingBuilder(unittest.TestCase):
             fake_pickle.return_value = no_meta_builder
             pickle.load = fake_pickle
             builder = ring.RingBuilder.load('fake.builder', open=fake_open)
-            fake_open.assert_has_calls([mock_call('fake.builder', 'rb')])
+            fake_open.assert_has_calls([mock.call('fake.builder', 'rb')])
             self.assertEquals(builder.devs, rb.devs)
             fake_pickle.reset_mock()
         finally:
             pickle.load = real_pickle
 
+    def test_save_load(self):
+        rb = ring.RingBuilder(8, 3, 1)
+        devs = [{'id': 0, 'region': 0, 'zone': 0, 'weight': 1,
+                 'ip': '127.0.0.0', 'port': 10000,
+                 'replication_ip': '127.0.0.0', 'replication_port': 10000,
+                 'device': 'sda1', 'meta': 'meta0'},
+                {'id': 1, 'region': 0, 'zone': 1, 'weight': 1,
+                 'ip': '127.0.0.1', 'port': 10001,
+                 'replication_ip': '127.0.0.1', 'replication_port': 10001,
+                 'device': 'sdb1', 'meta': 'meta1'},
+                {'id': 2, 'region': 0, 'zone': 2, 'weight': 2,
+                 'ip': '127.0.0.2', 'port': 10002,
+                 'replication_ip': '127.0.0.2', 'replication_port': 10002,
+                 'device': 'sdc1', 'meta': 'meta2'},
+                {'id': 3, 'region': 0, 'zone': 3, 'weight': 2,
+                 'ip': '127.0.0.3', 'port': 10003,
+                 'replication_ip': '127.0.0.3', 'replication_port': 10003,
+                 'device': 'sdd1', 'meta': ''}]
+        for d in devs:
+            rb.add_dev(d)
+        rb.rebalance()
+        builder_file = os.path.join(self.testdir, 'test_save.builder')
+        rb.save(builder_file)
+        loaded_rb = ring.RingBuilder.load(builder_file)
+        self.maxDiff = None
+        self.assertEquals(loaded_rb.to_dict(), rb.to_dict())
+
+    @mock.patch('__builtin__.open', autospec=True)
+    @mock.patch('swift.common.ring.builder.pickle.dump', autospec=True)
+    def test_save(self, mock_pickle_dump, mock_open):
+        mock_open.return_value = mock_fh = mock.Mock()
+        rb = ring.RingBuilder(8, 3, 1)
+        devs = [{'id': 0, 'region': 0, 'zone': 0, 'weight': 1,
+                 'ip': '127.0.0.0', 'port': 10000, 'device': 'sda1',
+                 'meta': 'meta0'},
+                {'id': 1, 'region': 0, 'zone': 1, 'weight': 1,
+                 'ip': '127.0.0.1', 'port': 10001, 'device': 'sdb1',
+                 'meta': 'meta1'},
+                {'id': 2, 'region': 0, 'zone': 2, 'weight': 2,
+                 'ip': '127.0.0.2', 'port': 10002, 'device': 'sdc1',
+                 'meta': 'meta2'},
+                {'id': 3, 'region': 0, 'zone': 3, 'weight': 2,
+                 'ip': '127.0.0.3', 'port': 10003, 'device': 'sdd1'}]
+        for d in devs:
+            rb.add_dev(d)
+        rb.rebalance()
+        rb.save('some.builder')
+        mock_open.assert_called_once_with('some.builder', 'wb')
+        mock_pickle_dump.assert_called_once_with(rb.to_dict(), mock_fh,
+                                                 protocol=2)
+
     def test_search_devs(self):
         rb = ring.RingBuilder(8, 3, 1)
         devs = [{'id': 0, 'region': 0, 'zone': 0, 'weight': 1,