Switched to a nicer contextmanager implemenation for role hierarchy rebuild batching

#1206
This commit is contained in:
Akita Noek
2016-03-15 15:30:43 -04:00
parent e45982b011
commit ce669b03ad
3 changed files with 258 additions and 256 deletions

View File

@@ -23,7 +23,7 @@ from django.db.transaction import TransactionManagementError
# AWX # AWX
from awx.main.models.rbac import RolePermission, Role from awx.main.models.rbac import RolePermission, Role, batch_role_ancestor_rebuilding
__all__ = ['AutoOneToOneField', 'ImplicitRoleField'] __all__ = ['AutoOneToOneField', 'ImplicitRoleField']
@@ -280,12 +280,11 @@ class ImplicitRoleField(models.ForeignKey):
for parent in parents: for parent in parents:
new_parents.add(parent) new_parents.add(parent)
Role.pause_role_ancestor_rebuilding() with batch_role_ancestor_rebuilding():
for role in original_parents - new_parents: for role in original_parents - new_parents:
this_role.parents.remove(role) this_role.parents.remove(role)
for role in new_parents - original_parents: for role in new_parents - original_parents:
this_role.parents.add(role) this_role.parents.add(role)
Role.unpause_role_ancestor_rebuilding()
setattr(self, '__original_parent_roles', new_parents) setattr(self, '__original_parent_roles', new_parents)

View File

@@ -104,249 +104,245 @@ class Command(BaseCommand):
try: try:
with transaction.atomic(): with transaction.atomic():
Role.pause_role_ancestor_rebuilding() with batch_role_ancestor_rebuilding():
print('# Creating %d organizations' % n_organizations) print('# Creating %d organizations' % n_organizations)
for i in xrange(n_organizations): for i in xrange(n_organizations):
sys.stdout.write('\r%d ' % (i + 1)) sys.stdout.write('\r%d ' % (i + 1))
sys.stdout.flush()
organizations.append(Organization.objects.create(name='%s Organization %d' % (prefix, i)))
print('')
print('# Creating %d users' % n_users)
org_idx = 0
for n in spread(n_users, n_organizations):
for i in range(n):
ids['user'] += 1
user_id = ids['user']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, organizations[org_idx].name, i+ 1))
sys.stdout.flush() sys.stdout.flush()
user = User.objects.create(username='%suser-%d' % (prefix, user_id)) organizations.append(Organization.objects.create(name='%s Organization %d' % (prefix, i)))
organizations[org_idx].member_role.members.add(user)
users.append(user)
org_idx += 1
print('') print('')
print('# Creating %d teams' % n_teams) print('# Creating %d users' % n_users)
org_idx = 0 org_idx = 0
for n in spread(n_teams, n_organizations): for n in spread(n_users, n_organizations):
org = organizations[org_idx]
for i in range(n):
ids['team'] += 1
team_id = ids['team']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, org.name, i+ 1))
sys.stdout.flush()
team = Team.objects.create(name='%s Team %d Org %d' % (prefix, team_id, org_idx), organization=org)
teams.append(team)
org_idx += 1
print('')
print('# Adding users to teams')
for org in organizations:
org_teams = [t for t in org.teams.all()]
org_users = [u for u in org.member_role.members.all()]
print(' Spreading %d users accross %d teams for %s' % (len(org_users), len(org_teams), org.name))
# Our normal spread for most users
cur_user_idx = 0
cur_team_idx = 0
for n in spread(len(org_users), len(org_teams)):
team = org_teams[cur_team_idx]
for i in range(n): for i in range(n):
if cur_user_idx < len(org_users): ids['user'] += 1
user = org_users[cur_user_idx] user_id = ids['user']
team.member_role.members.add(user) sys.stdout.write('\r Assigning %d to %s: %d ' % (n, organizations[org_idx].name, i+ 1))
cur_user_idx += 1 sys.stdout.flush()
cur_team_idx += 1 user = User.objects.create(username='%suser-%d' % (prefix, user_id))
organizations[org_idx].member_role.members.add(user)
users.append(user)
org_idx += 1
print('')
# First user gets added to all teams print('# Creating %d teams' % n_teams)
for team in org_teams: org_idx = 0
team.member_role.members.add(org_users[0]) for n in spread(n_teams, n_organizations):
org = organizations[org_idx]
for i in range(n):
ids['team'] += 1
team_id = ids['team']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, org.name, i+ 1))
sys.stdout.flush()
team = Team.objects.create(name='%s Team %d Org %d' % (prefix, team_id, org_idx), organization=org)
teams.append(team)
org_idx += 1
print('')
print('# Adding users to teams')
for org in organizations:
org_teams = [t for t in org.teams.all()]
org_users = [u for u in org.member_role.members.all()]
print(' Spreading %d users accross %d teams for %s' % (len(org_users), len(org_teams), org.name))
# Our normal spread for most users
cur_user_idx = 0
cur_team_idx = 0
for n in spread(len(org_users), len(org_teams)):
team = org_teams[cur_team_idx]
for i in range(n):
if cur_user_idx < len(org_users):
user = org_users[cur_user_idx]
team.member_role.members.add(user)
cur_user_idx += 1
cur_team_idx += 1
# First user gets added to all teams
for team in org_teams:
team.member_role.members.add(org_users[0])
print('# Creating %d credentials for users' % (n_credentials - n_credentials // 2)) print('# Creating %d credentials for users' % (n_credentials - n_credentials // 2))
user_idx = 0 user_idx = 0
for n in spread(n_credentials - n_credentials // 2, n_users): for n in spread(n_credentials - n_credentials // 2, n_users):
user = users[user_idx] user = users[user_idx]
for i in range(n): for i in range(n):
ids['credential'] += 1 ids['credential'] += 1
sys.stdout.write('\r %d ' % (ids['credential'])) sys.stdout.write('\r %d ' % (ids['credential']))
sys.stdout.flush() sys.stdout.flush()
credential_id = ids['credential'] credential_id = ids['credential']
credential = Credential.objects.create(name='%s Credential %d User %d' % (prefix, credential_id, user_idx), user=user) credential = Credential.objects.create(name='%s Credential %d User %d' % (prefix, credential_id, user_idx), user=user)
credentials.append(credential) credentials.append(credential)
user_idx += 1 user_idx += 1
print('')
print('# Creating %d credentials for teams' % (n_credentials // 2))
team_idx = 0
starting_credential_id = ids['credential']
for n in spread(n_credentials - n_credentials // 2, n_teams):
team = teams[team_idx]
for i in range(n):
ids['credential'] += 1
sys.stdout.write('\r %d ' % (ids['credential'] - starting_credential_id))
sys.stdout.flush()
credential_id = ids['credential']
credential = Credential.objects.create(name='%s Credential %d team %d' % (prefix, credential_id, team_idx), team=team)
credentials.append(credential)
team_idx += 1
print('')
print('# Creating %d projects' % n_projects)
org_idx = 0
for n in spread(n_projects, n_organizations):
org = organizations[org_idx]
for i in range(n):
ids['project'] += 1
project_id = ids['project']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, org.name, i+ 1))
sys.stdout.flush()
project = Project.objects.create(name='%s Project %d Org %d' % (prefix, project_id, org_idx), organization=org)
projects.append(project)
org_idx += 1
print('') print('')
print('# Creating %d credentials for teams' % (n_credentials // 2))
print('# Creating %d inventories' % n_inventories) team_idx = 0
org_idx = 0 starting_credential_id = ids['credential']
for n in spread(n_inventories, min(n_inventories // 4 + 1, n_organizations)): for n in spread(n_credentials - n_credentials // 2, n_teams):
org = organizations[org_idx] team = teams[team_idx]
for i in range(n): for i in range(n):
ids['inventory'] += 1 ids['credential'] += 1
inventory_id = ids['inventory'] sys.stdout.write('\r %d ' % (ids['credential'] - starting_credential_id))
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, org.name, i+ 1)) sys.stdout.flush()
sys.stdout.flush() credential_id = ids['credential']
inventory = Inventory.objects.create(name='%s Inventory %d Org %d' % (prefix, inventory_id, org_idx), organization=org) credential = Credential.objects.create(name='%s Credential %d team %d' % (prefix, credential_id, team_idx), team=team)
inventories.append(inventory) credentials.append(credential)
team_idx += 1
org_idx += 1
print('') print('')
print('# Creating %d projects' % n_projects)
org_idx = 0
for n in spread(n_projects, n_organizations):
org = organizations[org_idx]
for i in range(n):
ids['project'] += 1
project_id = ids['project']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, org.name, i+ 1))
sys.stdout.flush()
project = Project.objects.create(name='%s Project %d Org %d' % (prefix, project_id, org_idx), organization=org)
projects.append(project)
print('# Creating %d inventory_groups' % n_inventory_groups) org_idx += 1
inv_idx = 0 print('')
for n in spread(n_inventory_groups, n_inventories):
inventory = inventories[inv_idx]
parent_list = [None] * 3
for i in range(n):
ids['group'] += 1
group_id = ids['group']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, inventory.name, i+ 1))
sys.stdout.flush()
group = Group.objects.create(
name='%s Group %d Inventory %d' % (prefix, group_id, inv_idx),
inventory=inventory,
)
# Have each group have up to 3 parent groups
for parent_n in range(3):
if i // 4 + parent_n < len(parent_list) and parent_list[i // 4 + parent_n]:
group.parents.add(parent_list[i // 4 + parent_n])
if parent_list[i // 4] is None:
parent_list[i // 4] = group
else:
parent_list.append(group)
inventory_groups.append(group)
inv_idx += 1
print('')
print('# Creating %d inventory_hosts' % n_inventory_hosts) print('# Creating %d inventories' % n_inventories)
group_idx = 0 org_idx = 0
for n in spread(n_inventory_hosts, n_inventory_groups): for n in spread(n_inventories, min(n_inventories // 4 + 1, n_organizations)):
group = inventory_groups[group_idx] org = organizations[org_idx]
for i in range(n): for i in range(n):
ids['host'] += 1 ids['inventory'] += 1
host_id = ids['host'] inventory_id = ids['inventory']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, group.name, i+ 1)) sys.stdout.write('\r Assigning %d to %s: %d ' % (n, org.name, i+ 1))
sys.stdout.flush() sys.stdout.flush()
host = Host.objects.create(name='%s Host %d Group %d' % (prefix, host_id, group_idx), inventory=group.inventory) inventory = Inventory.objects.create(name='%s Inventory %d Org %d' % (prefix, inventory_id, org_idx), organization=org)
# Add the host to up to 3 groups inventories.append(inventory)
host.groups.add(group)
for m in range(2):
if group_idx + m < len(inventory_groups) and group.inventory.id == inventory_groups[group_idx + m].inventory.id:
host.groups.add(inventory_groups[group_idx + m])
inventory_hosts.append(host) org_idx += 1
print('')
group_idx += 1
print('')
print('# Creating %d job_templates' % n_job_templates) print('# Creating %d inventory_groups' % n_inventory_groups)
project_idx = 0 inv_idx = 0
inv_idx = 0 for n in spread(n_inventory_groups, n_inventories):
for n in spread(n_job_templates, n_projects): inventory = inventories[inv_idx]
project = projects[project_idx] parent_list = [None] * 3
for i in range(n): for i in range(n):
ids['job_template'] += 1 ids['group'] += 1
job_template_id = ids['job_template'] group_id = ids['group']
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, project.name, i+ 1)) sys.stdout.write('\r Assigning %d to %s: %d ' % (n, inventory.name, i+ 1))
sys.stdout.flush() sys.stdout.flush()
group = Group.objects.create(
name='%s Group %d Inventory %d' % (prefix, group_id, inv_idx),
inventory=inventory,
)
# Have each group have up to 3 parent groups
for parent_n in range(3):
if i // 4 + parent_n < len(parent_list) and parent_list[i // 4 + parent_n]:
group.parents.add(parent_list[i // 4 + parent_n])
if parent_list[i // 4] is None:
parent_list[i // 4] = group
else:
parent_list.append(group)
inventory_groups.append(group)
inventory = None
org_inv_count = project.organization.inventories.count()
if org_inv_count > 0:
inventory = project.organization.inventories.all()[inv_idx % org_inv_count]
job_template = JobTemplate.objects.create(
name='%s Job Template %d Project %d' % (prefix, job_template_id, project_idx),
inventory=inventory,
project=project,
)
job_templates.append(job_template)
inv_idx += 1 inv_idx += 1
project_idx += 1 print('')
print('')
print('# Creating %d jobs' % n_jobs)
group_idx = 0
job_template_idx = 0
for n in spread(n_jobs, n_job_templates):
job_template = job_templates[job_template_idx]
for i in range(n):
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, job_template.name, i+ 1))
sys.stdout.flush()
job = Job.objects.create(job_template=job_template)
jobs.append(job)
if job_template.inventory: print('# Creating %d inventory_hosts' % n_inventory_hosts)
inv_groups = [g for g in job_template.inventory.groups.all()] group_idx = 0
if len(inv_groups): for n in spread(n_inventory_hosts, n_inventory_groups):
JobHostSummary.objects.bulk_create([ group = inventory_groups[group_idx]
JobHostSummary( for i in range(n):
job=job, host=h, host_name=h.name, processed=1, ids['host'] += 1
created=now(), modified=now() host_id = ids['host']
) sys.stdout.write('\r Assigning %d to %s: %d ' % (n, group.name, i+ 1))
for h in inv_groups[group_idx % len(inv_groups)].hosts.all()[:100] sys.stdout.flush()
]) host = Host.objects.create(name='%s Host %d Group %d' % (prefix, host_id, group_idx), inventory=group.inventory)
# Add the host to up to 3 groups
host.groups.add(group)
for m in range(2):
if group_idx + m < len(inventory_groups) and group.inventory.id == inventory_groups[group_idx + m].inventory.id:
host.groups.add(inventory_groups[group_idx + m])
inventory_hosts.append(host)
group_idx += 1 group_idx += 1
job_template_idx += 1
if n:
print('') print('')
print('# Creating %d job events' % n_job_events) print('# Creating %d job_templates' % n_job_templates)
job_idx = 0 project_idx = 0
for n in spread(n_job_events, n_jobs): inv_idx = 0
job = jobs[job_idx] for n in spread(n_job_templates, n_projects):
sys.stdout.write('\r Creating %d job events for job %d' % (n, job.id)) project = projects[project_idx]
sys.stdout.flush() for i in range(n):
JobEvent.objects.bulk_create([ ids['job_template'] += 1
JobEvent( job_template_id = ids['job_template']
created=now(), sys.stdout.write('\r Assigning %d to %s: %d ' % (n, project.name, i+ 1))
modified=now(), sys.stdout.flush()
job=job,
event='runner_on_ok' inventory = None
) org_inv_count = project.organization.inventories.count()
for i in range(n) if org_inv_count > 0:
]) inventory = project.organization.inventories.all()[inv_idx % org_inv_count]
job_idx += 1
if n: job_template = JobTemplate.objects.create(
name='%s Job Template %d Project %d' % (prefix, job_template_id, project_idx),
inventory=inventory,
project=project,
)
job_templates.append(job_template)
inv_idx += 1
project_idx += 1
print('') print('')
print('# Creating %d jobs' % n_jobs)
group_idx = 0
job_template_idx = 0
for n in spread(n_jobs, n_job_templates):
job_template = job_templates[job_template_idx]
for i in range(n):
sys.stdout.write('\r Assigning %d to %s: %d ' % (n, job_template.name, i+ 1))
sys.stdout.flush()
job = Job.objects.create(job_template=job_template)
jobs.append(job)
if job_template.inventory:
inv_groups = [g for g in job_template.inventory.groups.all()]
if len(inv_groups):
JobHostSummary.objects.bulk_create([
JobHostSummary(
job=job, host=h, host_name=h.name, processed=1,
created=now(), modified=now()
)
for h in inv_groups[group_idx % len(inv_groups)].hosts.all()[:100]
])
group_idx += 1
job_template_idx += 1
if n:
print('')
Role.unpause_role_ancestor_rebuilding() print('# Creating %d job events' % n_job_events)
job_idx = 0
for n in spread(n_job_events, n_jobs):
job = jobs[job_idx]
sys.stdout.write('\r Creating %d job events for job %d' % (n, job.id))
sys.stdout.flush()
JobEvent.objects.bulk_create([
JobEvent(
created=now(),
modified=now(),
job=job,
event='runner_on_ok'
)
for i in range(n)
])
job_idx += 1
if n:
print('')
if options['pretend']: if options['pretend']:
raise Rollback() raise Rollback()

View File

@@ -3,9 +3,11 @@
# Python # Python
import logging import logging
import threading
import contextlib
# Django # Django
from django.db import models from django.db import models, transaction
from django.db.models import Q from django.db.models import Q
from django.db.models.aggregates import Max from django.db.models.aggregates import Max
from django.core.urlresolvers import reverse from django.core.urlresolvers import reverse
@@ -19,6 +21,7 @@ from awx.main.models.base import * # noqa
__all__ = [ __all__ = [
'Role', 'Role',
'RolePermission', 'RolePermission',
'batch_role_ancestor_rebuilding',
'get_user_permissions_on_resource', 'get_user_permissions_on_resource',
'get_role_permissions_on_resource', 'get_role_permissions_on_resource',
'ROLE_SINGLETON_SYSTEM_ADMINISTRATOR', 'ROLE_SINGLETON_SYSTEM_ADMINISTRATOR',
@@ -30,12 +33,43 @@ logger = logging.getLogger('awx.main.models.rbac')
ROLE_SINGLETON_SYSTEM_ADMINISTRATOR='System Administrator' ROLE_SINGLETON_SYSTEM_ADMINISTRATOR='System Administrator'
ROLE_SINGLETON_SYSTEM_AUDITOR='System Auditor' ROLE_SINGLETON_SYSTEM_AUDITOR='System Auditor'
role_rebuilding_paused = False
roles_needing_rebuilding = set()
ALL_PERMISSIONS = {'create': True, 'read': True, 'update': True, 'delete': True, ALL_PERMISSIONS = {'create': True, 'read': True, 'update': True, 'delete': True,
'write': True, 'scm_update': True, 'use': True, 'execute': True} 'write': True, 'scm_update': True, 'use': True, 'execute': True}
tls = threading.local() # thread local storage
@contextlib.contextmanager
def batch_role_ancestor_rebuilding(allow_nesting=False):
'''
Batches the role ancestor rebuild work necessary whenever role-role
relations change. This can result in a big speedup when performing
any bulk manipulation.
WARNING: Calls to anything related to checking access/permissions
while within the context of the batch_role_ancestor_rebuilding will
likely not work.
'''
batch_role_rebuilding = getattr(tls, 'batch_role_rebuilding', False)
try:
setattr(tls, 'batch_role_rebuilding', True)
if not batch_role_rebuilding:
setattr(tls, 'roles_needing_rebuilding', set())
yield
finally:
setattr(tls, 'batch_role_rebuilding', batch_role_rebuilding)
if not batch_role_rebuilding:
rebuild_set = getattr(tls, 'roles_needing_rebuilding')
with transaction.atomic():
for role in Role.objects.filter(id__in=list(rebuild_set)).all():
# TODO: We can reduce this to one rebuild call with our new upcoming rebuild method.. do this
role.rebuild_role_ancestor_list()
delattr(tls, 'roles_needing_rebuilding')
class Role(CommonModelNameNotUnique): class Role(CommonModelNameNotUnique):
''' '''
Role model Role model
@@ -61,35 +95,6 @@ class Role(CommonModelNameNotUnique):
def get_absolute_url(self): def get_absolute_url(self):
return reverse('api:role_detail', args=(self.pk,)) return reverse('api:role_detail', args=(self.pk,))
@staticmethod
def pause_role_ancestor_rebuilding():
'''
Pauses role ancestor list updating. This is useful when you're making
many changes to the same roles, for example doing bulk inserts or
making many changes to the same object in succession.
Note that the unpause_role_ancestor_rebuilding MUST be called within
the same execution context (preferably within the same transaction),
otherwise the RBAC role ancestor hierarchy will not be properly
updated.
'''
global role_rebuilding_paused
role_rebuilding_paused = True
@staticmethod
def unpause_role_ancestor_rebuilding():
'''
Unpauses the role ancestor list updating. This will will rebuild all
roles that need updating since the last call to
pause_role_ancestor_rebuilding and bring everything back into sync.
'''
global role_rebuilding_paused
global roles_needing_rebuilding
role_rebuilding_paused = False
for role in Role.objects.filter(id__in=list(roles_needing_rebuilding)).all():
role.rebuild_role_ancestor_list()
roles_needing_rebuilding = set()
def rebuild_role_ancestor_list(self): def rebuild_role_ancestor_list(self):
''' '''
@@ -100,9 +105,11 @@ class Role(CommonModelNameNotUnique):
Note that this method relies on any parents' ancestor list being correct. Note that this method relies on any parents' ancestor list being correct.
''' '''
global role_rebuilding_paused, roles_needing_rebuilding global tls
batch_role_rebuilding = getattr(tls, 'batch_role_rebuilding', False)
if role_rebuilding_paused: if batch_role_rebuilding:
roles_needing_rebuilding = getattr(tls, 'roles_needing_rebuilding')
roles_needing_rebuilding.add(self.id) roles_needing_rebuilding.add(self.id)
return return