Adds storage policy option to recon

With this patch, recon gets policy-aware in regard to the hosts to retrieve.
If no policy is passed and the server_type is object,
all hosts of all policies are retrieved.
Previously, recon did only retrieve the hosts of the storage-policy 0.

Change-Id: If5735cd6721eac504aed8aaf3884cb91b6a0fcac
Closes-Bug: 1541491

swift/cli/recon.py

@@ -22,6 +22,7 @@ from eventlet.green import urllib2, socket
 from six.moves.urllib.parse import urlparse
 from swift.common.utils import SWIFT_CONF_FILE
 from swift.common.ring import Ring
+from swift.common.storage_policy import POLICIES
 from hashlib import md5
 import eventlet
 import json
@@ -203,18 +204,19 @@ class SwiftRecon(object):
                 block = f.read(4096)
         return md5sum.hexdigest()
 
-    def get_devices(self, region_filter, zone_filter, swift_dir, ring_name):
+    def get_hosts(self, region_filter, zone_filter, swift_dir, ring_names):
         """
-        Get a list of hosts in the ring
+        Get a list of hosts in the rings.
 
         :param region_filter: Only list regions matching given filter
         :param zone_filter: Only list zones matching given filter
         :param swift_dir: Directory of swift config, usually /etc/swift
-        :param ring_name: Name of the ring, such as 'object'
+        :param ring_names: Collection of ring names, such as
+         ['object', 'object-2']
         :returns: a set of tuples containing the ip and port of hosts
         """
-        ring_data = Ring(swift_dir, ring_name=ring_name)
+        rings = [Ring(swift_dir, ring_name=n) for n in ring_names]
-        devs = [d for d in ring_data.devs if d]
+        devs = [d for r in rings for d in r.devs if d]
         if region_filter is not None:
             devs = [d for d in devs if d['region'] == region_filter]
         if zone_filter is not None:
@@ -914,6 +916,26 @@ class SwiftRecon(object):
             matches, len(hosts), errors))
         print("=" * 79)
 
+    def _get_ring_names(self, policy=None):
+        '''
+        Retrieve name of ring files.
+
+        If no policy is passed and the server type is object,
+        the ring names of all storage-policies are retrieved.
+
+        :param policy: name or index of storage policy, only applicable
+         with server_type==object.
+         :returns: list of ring names.
+        '''
+        if self.server_type == 'object':
+            ring_names = [p.ring_name for p in POLICIES if (
+                p.name == policy or not policy or (
+                    policy.isdigit() and int(policy) == int(p)))]
+        else:
+            ring_names = [self.server_type]
+
+        return ring_names
+
     def main(self):
         """
         Retrieve and report cluster info from hosts running recon middleware.
@@ -983,6 +1005,9 @@ class SwiftRecon(object):
                         default=5)
         args.add_option('--swiftdir', default="/etc/swift",
                         help="Default = /etc/swift")
+        args.add_option('--policy', '-p',
+                        help='Only query object servers in specified '
+                        'storage policy (specified as name or index).')
         options, arguments = args.parse_args()
 
         if len(sys.argv) <= 1 or len(arguments) > 1:
@@ -1004,8 +1029,14 @@ class SwiftRecon(object):
         self.suppress_errors = options.suppress
         self.timeout = options.timeout
 
-        hosts = self.get_devices(options.region, options.zone,
+        ring_names = self._get_ring_names(options.policy)
-                                 swift_dir, self.server_type)
+        if not ring_names:
+            print('Invalid Storage Policy')
+            args.print_help()
+            sys.exit(0)
+
+        hosts = self.get_hosts(options.region, options.zone,
+                               swift_dir, ring_names)
 
         print("--> Starting reconnaissance on %s hosts" % len(hosts))
         print("=" * 79)

test/unit/cli/test_recon.py

@@ -16,12 +16,12 @@
 import json
 import mock
 import os
-import random
 import re
-import string
 import tempfile
 import time
 import unittest
+import shutil
+import sys
 
 from eventlet.green import urllib2, socket
 from six import StringIO
@@ -30,6 +30,9 @@ from six.moves import urllib
 from swift.cli import recon
 from swift.common import utils
 from swift.common.ring import builder
+from swift.common.ring import utils as ring_utils
+from swift.common.storage_policy import StoragePolicy, POLICIES
+from test.unit import patch_policies
 
 
 class TestHelpers(unittest.TestCase):
@@ -135,22 +138,50 @@ class TestScout(unittest.TestCase):
         self.assertEqual(status, -1)
 
 
+@patch_policies
 class TestRecon(unittest.TestCase):
     def setUp(self, *_args, **_kwargs):
         self.recon_instance = recon.SwiftRecon()
-        self.swift_dir = tempfile.gettempdir()
+        self.swift_dir = tempfile.mkdtemp()
-        self.ring_name = "test_object_%s" % (
+        self.ring_name = POLICIES.legacy.ring_name
-            ''.join(random.choice(string.digits) for x in range(6)))
+        self.tmpfile_name = os.path.join(
-        self.tmpfile_name = "%s/%s.ring.gz" % (self.swift_dir, self.ring_name)
+            self.swift_dir, self.ring_name + '.ring.gz')
+        self.ring_name2 = POLICIES[1].ring_name
+        self.tmpfile_name2 = os.path.join(
+            self.swift_dir, self.ring_name2 + '.ring.gz')
 
         utils.HASH_PATH_SUFFIX = 'endcap'
         utils.HASH_PATH_PREFIX = 'startcap'
 
     def tearDown(self, *_args, **_kwargs):
-        try:
+        shutil.rmtree(self.swift_dir, ignore_errors=True)
-            os.remove(self.tmpfile_name)
+
-        except OSError:
+    def _make_object_rings(self):
-            pass
+        ringbuilder = builder.RingBuilder(2, 3, 1)
+        devs = [
+            'r0z0-127.0.0.1:10000/sda1',
+            'r0z1-127.0.0.1:10001/sda1',
+            'r1z0-127.0.0.1:10002/sda1',
+            'r1z1-127.0.0.1:10003/sda1',
+        ]
+        for raw_dev_str in devs:
+            dev = ring_utils.parse_add_value(raw_dev_str)
+            dev['weight'] = 1.0
+            ringbuilder.add_dev(dev)
+        ringbuilder.rebalance()
+        ringbuilder.get_ring().save(self.tmpfile_name)
+
+        ringbuilder = builder.RingBuilder(2, 2, 1)
+        devs = [
+            'r0z0-127.0.0.1:10000/sda1',
+            'r0z1-127.0.0.2:10004/sda1',
+        ]
+        for raw_dev_str in devs:
+            dev = ring_utils.parse_add_value(raw_dev_str)
+            dev['weight'] = 1.0
+            ringbuilder.add_dev(dev)
+        ringbuilder.rebalance()
+        ringbuilder.get_ring().save(self.tmpfile_name2)
 
     def test_gen_stats(self):
         stats = self.recon_instance._gen_stats((1, 4, 10, None), 'Sample')
@@ -176,47 +207,56 @@ class TestRecon(unittest.TestCase):
             self.assertEqual(timestamp2, "2013-12-17 10:00:00")
             mock_gmtime.assert_called_with()
 
-    def test_get_devices(self):
+    def test_get_hosts(self):
-        ringbuilder = builder.RingBuilder(2, 3, 1)
+        self._make_object_rings()
-        ringbuilder.add_dev({'id': 0, 'zone': 0, 'weight': 1,
-                             'ip': '127.0.0.1', 'port': 10000,
-                             'device': 'sda1', 'region': 0})
-        ringbuilder.add_dev({'id': 1, 'zone': 1, 'weight': 1,
-                             'ip': '127.0.0.1', 'port': 10001,
-                             'device': 'sda1', 'region': 0})
-        ringbuilder.add_dev({'id': 2, 'zone': 0, 'weight': 1,
-                             'ip': '127.0.0.1', 'port': 10002,
-                             'device': 'sda1', 'region': 1})
-        ringbuilder.add_dev({'id': 3, 'zone': 1, 'weight': 1,
-                             'ip': '127.0.0.1', 'port': 10003,
-                             'device': 'sda1', 'region': 1})
-        ringbuilder.rebalance()
-        ringbuilder.get_ring().save(self.tmpfile_name)
 
-        ips = self.recon_instance.get_devices(
+        ips = self.recon_instance.get_hosts(
-            None, None, self.swift_dir, self.ring_name)
+            None, None, self.swift_dir, [self.ring_name])
         self.assertEqual(
             set([('127.0.0.1', 10000), ('127.0.0.1', 10001),
                  ('127.0.0.1', 10002), ('127.0.0.1', 10003)]), ips)
 
-        ips = self.recon_instance.get_devices(
+        ips = self.recon_instance.get_hosts(
-            0, None, self.swift_dir, self.ring_name)
+            0, None, self.swift_dir, [self.ring_name])
         self.assertEqual(
             set([('127.0.0.1', 10000), ('127.0.0.1', 10001)]), ips)
 
-        ips = self.recon_instance.get_devices(
+        ips = self.recon_instance.get_hosts(
-            1, None, self.swift_dir, self.ring_name)
+            1, None, self.swift_dir, [self.ring_name])
         self.assertEqual(
             set([('127.0.0.1', 10002), ('127.0.0.1', 10003)]), ips)
 
-        ips = self.recon_instance.get_devices(
+        ips = self.recon_instance.get_hosts(
-            0, 0, self.swift_dir, self.ring_name)
+            0, 0, self.swift_dir, [self.ring_name])
         self.assertEqual(set([('127.0.0.1', 10000)]), ips)
 
-        ips = self.recon_instance.get_devices(
+        ips = self.recon_instance.get_hosts(
-            1, 1, self.swift_dir, self.ring_name)
+            1, 1, self.swift_dir, [self.ring_name])
         self.assertEqual(set([('127.0.0.1', 10003)]), ips)
 
+        ips = self.recon_instance.get_hosts(
+            None, None, self.swift_dir, [self.ring_name, self.ring_name2])
+        self.assertEqual(
+            set([('127.0.0.1', 10000), ('127.0.0.1', 10001),
+                 ('127.0.0.1', 10002), ('127.0.0.1', 10003),
+                 ('127.0.0.2', 10004)]), ips)
+
+        ips = self.recon_instance.get_hosts(
+            0, None, self.swift_dir, [self.ring_name, self.ring_name2])
+        self.assertEqual(
+            set([('127.0.0.1', 10000), ('127.0.0.1', 10001),
+                 ('127.0.0.2', 10004)]), ips)
+
+        ips = self.recon_instance.get_hosts(
+            1, None, self.swift_dir, [self.ring_name, self.ring_name2])
+        self.assertEqual(
+            set([('127.0.0.1', 10002), ('127.0.0.1', 10003)]), ips)
+
+        ips = self.recon_instance.get_hosts(
+            0, 1, self.swift_dir, [self.ring_name, self.ring_name2])
+        self.assertEqual(set([('127.0.0.1', 10001),
+                              ('127.0.0.2', 10004)]), ips)
+
     def test_get_ringmd5(self):
         for server_type in ('account', 'container', 'object', 'object-1'):
             ring_name = '%s.ring.gz' % server_type
@@ -343,6 +383,89 @@ class TestRecon(unittest.TestCase):
                                  " Failed: %s%%, no_result: %s, reported: %s"
                                  % expected)
 
+    def test_get_ring_names(self):
+        self.recon_instance.server_type = 'not-object'
+        self.assertEqual(self.recon_instance._get_ring_names(), ['not-object'])
+
+        self.recon_instance.server_type = 'object'
+
+        with patch_policies([StoragePolicy(0, 'zero', is_default=True)]):
+            self.assertEqual(self.recon_instance._get_ring_names(),
+                             ['object'])
+
+        with patch_policies([StoragePolicy(0, 'zero', is_default=True),
+                             StoragePolicy(1, 'one')]):
+            self.assertEqual(self.recon_instance._get_ring_names(),
+                             ['object', 'object-1'])
+            self.assertEqual(self.recon_instance._get_ring_names('0'),
+                             ['object'])
+            self.assertEqual(self.recon_instance._get_ring_names('zero'),
+                             ['object'])
+            self.assertEqual(self.recon_instance._get_ring_names('1'),
+                             ['object-1'])
+            self.assertEqual(self.recon_instance._get_ring_names('one'),
+                             ['object-1'])
+
+            self.assertEqual(self.recon_instance._get_ring_names('3'), [])
+            self.assertEqual(self.recon_instance._get_ring_names('wrong'),
+                             [])
+
+    def test_main_object_hosts_default_all_policies(self):
+        self._make_object_rings()
+        discovered_hosts = set()
+
+        def server_type_check(hosts):
+            for h in hosts:
+                discovered_hosts.add(h)
+
+        self.recon_instance.server_type_check = server_type_check
+        with mock.patch.object(sys, 'argv', [
+                "prog", "object", "--swiftdir=%s" % self.swift_dir,
+                "--validate-servers"]):
+            self.recon_instance.main()
+
+        expected = set([
+            ('127.0.0.1', 10000),
+            ('127.0.0.1', 10001),
+            ('127.0.0.1', 10002),
+            ('127.0.0.1', 10003),
+            ('127.0.0.2', 10004),
+        ])
+
+        self.assertEqual(expected, discovered_hosts)
+
+    def test_main_object_hosts_default_unu(self):
+        self._make_object_rings()
+        discovered_hosts = set()
+
+        def server_type_check(hosts):
+            for h in hosts:
+                discovered_hosts.add(h)
+
+        self.recon_instance.server_type_check = server_type_check
+
+        with mock.patch.object(sys, 'argv', [
+                "prog", "object", "--swiftdir=%s" % self.swift_dir,
+                "--validate-servers", '--policy=unu']):
+
+            self.recon_instance.main()
+
+        expected = set([
+            ('127.0.0.1', 10000),
+            ('127.0.0.2', 10004),
+        ])
+        self.assertEqual(expected, discovered_hosts)
+
+    def test_main_object_hosts_default_invalid(self):
+        self._make_object_rings()
+        stdout = StringIO()
+        with mock.patch.object(sys, 'argv', [
+                "prog", "object", "--swiftdir=%s" % self.swift_dir,
+                "--validate-servers", '--policy=invalid']),\
+                mock.patch('sys.stdout', stdout):
+            self.assertRaises(SystemExit, recon.main)
+            self.assertIn('Invalid Storage Policy', stdout.getvalue())
+
 
 class TestReconCommands(unittest.TestCase):
     def setUp(self):