Refactor import_authorized_keys() function for performance

The main change is to fetch all of the relation_data() at once, and then
iterate through the python dictionary.  This speeds up processing of
potentially hundreds of hosts and authorized_keys.

Change-Id: I095104f535c1eae1554f842502ae93ebb92e44fe
Related-Bug: #1833420
This commit is contained in:
Alex Kavanagh 2019-06-19 17:08:45 +01:00
parent af051c4fba
commit 2e2d19f7c7
2 changed files with 59 additions and 64 deletions

View File

@ -598,39 +598,33 @@ def import_authorized_keys(user='root', prefix=None):
"""Import SSH authorized_keys + known_hosts from a cloud-compute relation.
Store known_hosts in user's $HOME/.ssh and authorized_keys in a path
specified using authorized-keys-path config option.
"""
known_hosts = []
authorized_keys = []
if prefix:
known_hosts_index = relation_get(
'{}_known_hosts_max_index'.format(prefix))
if known_hosts_index:
for index in range(0, int(known_hosts_index)):
known_hosts.append(relation_get(
'{}_known_hosts_{}'.format(prefix, index)))
authorized_keys_index = relation_get(
'{}_authorized_keys_max_index'.format(prefix))
if authorized_keys_index:
for index in range(0, int(authorized_keys_index)):
authorized_keys.append(relation_get(
'{}_authorized_keys_{}'.format(prefix, index)))
else:
# XXX: Should this be managed via templates + contexts?
known_hosts_index = relation_get('known_hosts_max_index')
if known_hosts_index:
for index in range(0, int(known_hosts_index)):
known_hosts.append(relation_get(
'known_hosts_{}'.format(index)))
authorized_keys_index = relation_get('authorized_keys_max_index')
if authorized_keys_index:
for index in range(0, int(authorized_keys_index)):
authorized_keys.append(relation_get(
'authorized_keys_{}'.format(index)))
# XXX: Should partial return of known_hosts or authorized_keys
# be allowed ?
if not len(known_hosts) or not len(authorized_keys):
The relation_get data is a series of key values of the form:
[prefix_]known_hosts_max_index: <int>
[prefix_]authorized_keys_max_index: <int>
[prefix_]known_hosts_[n]: <str>
[prefix_]authorized_keys_[n]: <str>
:param user: the user to write the known hosts and keys for (default 'root)
:type user: str
:param prefix: A prefix to add to the relation data keys (default None)
:type prefix: Option[str, None]
"""
_prefix = "{}_".format(prefix) if prefix else ""
# get all the data at once with one relation_get call
rdata = relation_get() or {}
known_hosts_index = int(
rdata.get('{}known_hosts_max_index'.format(_prefix), '0'))
authorized_keys_index = int(
rdata.get('{}authorized_keys_max_index'.format(_prefix), '0'))
if known_hosts_index == 0 or authorized_keys_index == 0:
return
homedir = pwd.getpwnam(user).pw_dir
dest_auth_keys = config('authorized-keys-path').format(
homedir=homedir, username=user)
@ -638,12 +632,17 @@ def import_authorized_keys(user='root', prefix=None):
log('Saving new known_hosts file to %s and authorized_keys file to: %s.' %
(dest_known_hosts, dest_auth_keys))
with open(dest_known_hosts, 'wt') as _hosts:
for index in range(0, int(known_hosts_index)):
_hosts.write('{}\n'.format(known_hosts[index]))
with open(dest_auth_keys, 'wt') as _keys:
for index in range(0, int(authorized_keys_index)):
_keys.write('{}\n'.format(authorized_keys[index]))
# write known hosts using data from relation_get
with open(dest_known_hosts, 'wt') as f:
for index in range(known_hosts_index):
f.write("{}\n".format(
rdata.get("{}known_hosts_{}".format(_prefix, index))))
# write authorized keys using data from relation_get
with open(dest_auth_keys, 'wt') as f:
for index in range(authorized_keys_index):
f.write("{}\n".format(
rdata.get('{}authorized_keys_{}'.format(_prefix, index))))
def do_openstack_upgrade(configs):

View File

@ -553,16 +553,27 @@ class NovaComputeUtilsTests(CharmTestCase):
auth_key_path='/home/foo/.ssh/'
'authorized_keys'):
getpwnam.return_value = self.fake_user('foo')
self.relation_get.side_effect = [
3, # relation_get('known_hosts_max_index')
'k_h_0', # relation_get_('known_hosts_0')
'k_h_1', # relation_get_('known_hosts_1')
'k_h_2', # relation_get_('known_hosts_2')
3, # relation_get('authorized_keys_max_index')
'auth_0', # relation_get('authorized_keys_0')
'auth_1', # relation_get('authorized_keys_1')
'auth_2', # relation_get('authorized_keys_2')
]
d = {
'known_hosts_max_index': 3,
'known_hosts_0': 'k_h_0',
'known_hosts_1': 'k_h_1',
'known_hosts_2': 'k_h_2',
'authorized_keys_max_index': 3,
'authorized_keys_0': 'auth_0',
'authorized_keys_1': 'auth_1',
'authorized_keys_2': 'auth_2',
}
if prefix:
for k, v in d.copy().items():
d["{}_{}".format(prefix, k)] = v
def _relation_get(scope=None, *args, **kwargs):
if scope is not None:
return d.get(scope, None)
return d
self.relation_get.side_effect = _relation_get
ex_open = [
call('/home/foo/.ssh/known_hosts', 'wt'),
@ -577,27 +588,12 @@ class NovaComputeUtilsTests(CharmTestCase):
call('auth_2\n')
]
# we only have to verify that the files are writen as expected as this
# implicitly checks that the relation_get calls have occurred.
with patch_open() as (_open, _file):
utils.import_authorized_keys(user='foo', prefix=prefix)
self.assertEqual(ex_open, _open.call_args_list)
self.assertEqual(ex_write, _file.write.call_args_list)
authkey_root = 'authorized_keys_'
known_hosts_root = 'known_hosts_'
if prefix:
authkey_root = prefix + '_authorized_keys_'
known_hosts_root = prefix + '_known_hosts_'
expected_relations = [
call(known_hosts_root + 'max_index'),
call(known_hosts_root + '0'),
call(known_hosts_root + '1'),
call(known_hosts_root + '2'),
call(authkey_root + 'max_index'),
call(authkey_root + '0'),
call(authkey_root + '1'),
call(authkey_root + '2')
]
self.assertEqual(sorted(self.relation_get.call_args_list),
sorted(expected_relations))
def test_import_authorized_keys_noprefix(self):
self._test_import_authorized_keys_base()