From cc040a9c29361caa7cc02253e529c75b13fb2759 Mon Sep 17 00:00:00 2001 From: Ilya Kharin Date: Thu, 16 May 2013 19:38:42 +0400 Subject: [PATCH] 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 | 25 +++++------ swift/common/ring/builder.py | 7 +++ test/unit/common/ring/test_builder.py | 63 ++++++++++++++++++++++++--- 3 files changed, 75 insertions(+), 20 deletions(-) diff --git a/bin/swift-ring-builder b/bin/swift-ring-builder index 9fd552aed5..c983357f65 100755 --- a/bin/swift-ring-builder +++ b/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 create except OSError, err: if err.errno != EEXIST: raise - pickle.dump(builder.to_dict(), open(pathjoin(backup_dir, - '%d.' % time() + basename(argv[1])), 'wb'), protocol=2) - pickle.dump(builder.to_dict(), open(argv[1], 'wb'), protocol=2) + builder.save(pathjoin(backup_dir, '%d.' % time() + basename(argv[1]))) + builder.save(argv[1]) exit(EXIT_SUCCESS) def default(): @@ -342,7 +340,7 @@ swift-ring-builder 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 set_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 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 remove [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 rebalance ts = time() builder.get_ring().save( pathjoin(backup_dir, '%d.' % ts + basename(ring_file))) - pickle.dump(builder.to_dict(), open(pathjoin(backup_dir, - '%d.' % ts + basename(argv[1])), 'wb'), protocol=2) + builder.save(pathjoin(backup_dir, '%d.' % ts + basename(argv[1]))) 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 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 set_min_part_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 set_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__': diff --git a/swift/common/ring/builder.py b/swift/common/ring/builder.py index 62c1fc696d..6219147565 100644 --- a/swift/common/ring/builder.py +++ b/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 can be of the form:: diff --git a/test/unit/common/ring/test_builder.py b/test/unit/common/ring/test_builder.py index 4b34181ec3..1fd28cd1cf 100644 --- a/test/unit/common/ring/test_builder.py +++ b/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_open = Mock(return_value=None) + fake_pickle = mock.Mock(return_value=rb) + 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,