blob: 8b1bcb65158bf27ad117c8d020460d5f516a2ff2 [file] [log] [blame]
# Copyright 2021 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
from collections import OrderedDict
from recipe_engine import recipe_api
from recipe_engine.recipe_api import Property
# Windows command helpers
from . import offline_winpe_customization as offwinpecust
from . import online_windows_customization as onwincust
from . import windows_iso_customization as winiso
from . import sources
from . import helper
from PB.recipes.infra.windows_image_builder import windows_image_builder as wib
from PB.recipes.infra.windows_image_builder import sources as src_pb
from PB.recipes.infra.windows_image_builder import dest as dest_pb
# Customization to Builder mapping
CUST_BUILDER = {
'offline_winpe_customization': 'Wim Customization Builder',
'online_windows_customization': 'Windows Customization Builder',
'windows_iso_customization': 'Windows Customization Builder',
}
BUILDERS = set(CUST_BUILDER.values())
class WindowsPSExecutorAPI(recipe_api.RecipeApi):
"""API for using Windows PowerShell scripts."""
def __init__(self, envvars, **kwargs):
super(WindowsPSExecutorAPI, self).__init__(**kwargs)
self._scripts = self.resource
self._workdir = ''
self._sources = None
self._configs_dir = None
self._customizations = []
self._executable_cust = []
self._config = wib.Image()
self._try_job = False
if envvars.MAX_CUST_BATCH_SIZE:
self._max_cust_batch_size = int(envvars.MAX_CUST_BATCH_SIZE)
else:
# Batch 1 cust per job by default
self._max_cust_batch_size = 1
def init(self, try_job=False):
""" init initializes all the dirs and sub modules required."""
self._try_job = try_job
with self.m.step.nest('Initialize the config engine'):
self._sources = sources.Source(self.m.path.cache_dir / 'Pkgs', self.m)
self._configs_dir = self.m.path.cleanup_dir / 'configs'
helper.ensure_dirs(self.m.file, [self._configs_dir])
def init_customizations(self, config):
""" init_customizations initializes the given config and returns
list of customizations
Args:
* config: wib.Image proto config
"""
# ensure that arch is specified in the image
if config.arch == wib.Arch.ARCH_UNSPECIFIED:
raise self.m.step.StepFailure('Missing arch in config')
arch = wib.Arch.Name(config.arch).replace('ARCH_', '').lower()
custs = [] # list of customizations from customizations.py
# initialize all customizations
for cust in config.customizations:
if cust.WhichOneof('customization') == 'offline_winpe_customization':
custs.append(
offwinpecust.OfflineWinPECustomization(
image=config,
cust=cust,
arch=arch,
scripts=self._scripts,
configs=self._configs_dir,
module=self.m,
source=self._sources,
try_job=self._try_job,
))
if cust.WhichOneof('customization') == 'online_windows_customization':
custs.append(
onwincust.OnlineWindowsCustomization(
image=config,
cust=cust,
arch=arch,
scripts=self._scripts,
configs=self._configs_dir,
module=self.m,
source=self._sources,
try_job=self._try_job,
))
if cust.WhichOneof('customization') == 'windows_iso_customization':
custs.append(
winiso.WinISOCustomization(
image=config,
cust=cust,
arch=arch,
scripts=self._scripts,
configs=self._configs_dir,
module=self.m,
source=self._sources,
try_job=self._try_job,
))
return custs
def process_customizations(self, custs, ctx, inputs=()):
""" process_customizations pins all the volatile srcs and generates
canonnical configs.
Args:
* custs: List of customizations from customization.py
* ctx: dict containing the context for the customization
* inputs: List of inputs that are required
Returns list of customizations in order that they were processed
"""
with self.m.step.nest('Process the customizations'):
resolved_cust = []
while len(custs) != len(resolved_cust):
pinnable_cust = []
for cust in custs:
if cust not in resolved_cust:
if cust.pinnable(ctx):
pinnable_cust.append(cust)
if pinnable_cust:
self.pin_customizations(pinnable_cust, ctx)
self.gen_canonical_configs(pinnable_cust)
# If the cust is online, inject cache upload
# request for drives
for cust in pinnable_cust:
cust_type = cust.customization().WhichOneof('customization')
if cust_type == 'online_windows_customization':
cust.inject_cache_upload(inputs)
ctx = self.update_context(pinnable_cust, ctx)
resolved_cust.extend(pinnable_cust)
else:
unpinned_custs = ''
for cust in custs:
if cust not in resolved_cust:
unpinned_custs += cust.id + '<br>'
unpins = '<ul>'
for pin in cust.unpinnable(ctx):
unpins += '<li>{}</li>'.format(pin.local_src)
unpins += '</ul> <br>'
reason = 'Cannot pin custs [<br> {} <br>].'.format(unpinned_custs)
raise self.m.step.StepFailure(
'Cyclical dependency?: {}'.format(reason))
return resolved_cust
def pin_customizations(self, customizations, ctx):
""" pin_customizations pins all the sources in the customizations
Args:
* customizations: List of Customizations object from customizations.py
* ctx: dict containing the context for the customization
"""
for cust in customizations:
with self.m.step.nest('Pin resources from {}'.format(cust.name())):
cust.pin_sources(ctx)
def trim_uploads(self, customizations):
""" trim_uploads removes the user specified uploads from a config.
"""
for cust in customizations:
cust.remove_upload_dests()
def update_context(self, custs, ctx):
""" update_context returns an updated dict with all the contexts
updated
Args:
* custs: List of customizations from customization.py
* ctx: Current context
Returns updated context dict
"""
for cust in custs:
ctx.update(cust.context)
return ctx
def gen_canonical_configs(self, customizations):
""" gen_canonical_configs strips all the names in the config and returns
individual configs containing one customization per image.
Example:
Given an Image
Image{
arch: x86,
name: "windows10_x86_GCE",
customizations: [
Customization{
OfflineWinPECustomization{
name: "winpe_networking"
image_dest: GCSSrc{
bucket: "chrome-win-wim"
source: "rel/win10_networking.wim"
}
...
}
},
Customization{
OfflineWinPECustomization{
name: "winpe_diskpart"
image_src: Src{
gcs_src: GCSSrc{
bucket: "chrome-win-wim"
source: "rel/win10_networking.wim"
}
}
...
}
}
]
}
Writes two configs: windows10_x86_GCE-winpe_networking.cfg with
Image{
arch: x86,
name: "",
customizations: [
Customization{
OfflineWinPECustomization{
name: ""
image_dest: GCSSrc{
bucket: "chrome-win-wim"
source: "rel/win10_networking.wim"
}
...
}
}
]
}
and windows10_x86_GCE-winpe_diskpart.cfg with
Image{
arch: x86,
name: "",
customizations: [
Customization{
OfflineWinPECustomization{
name: ""
image_src: Src{
gcs_src: GCSSrc{
bucket: "chrome-win-wim"
source: "rel/win10_networking.wim"
}
}
...
}
}
]
}
to disk, calculates the hash for each config and sets the key for each
of them. The strings representing name of the image, customization,...
etc,. are set to empty before calculating the hash to maintain the
uniqueness of the hash.
Args:
* customizations: List of Customizations object from customizations.py
"""
for cust in customizations:
# create a new image object, with same arch and containing only one
# customization
canon_image = wib.Image(
arch=cust.image().arch, customizations=[cust.get_canonical_cfg()])
name = cust.name()
# write the config to disk
cfg_file = self._configs_dir.joinpath('{}-{}.cfg'.format(
cust.image().name, name))
self.m.file.write_proto(
'Write config {}'.format(cfg_file),
cfg_file,
canon_image,
codec='BINARY')
# estimate the unique hash for the config (identifier for the image built
# by this config)
key = self.m.file.file_hash(cfg_file)
cust.set_key(key)
# save the config to disk as <key>.cfg
key_file = self._configs_dir.joinpath('{}.cfg'.format(key))
self.m.file.copy('Copy {} to {}'.format(cfg_file, key_file), cfg_file,
key_file)
return customizations
def filter_executable_customizations(self, customizations):
""" filter_executable_customizations generates a list of customizations
that need to be executed.
Args:
* customizations: List of Customizations object from customizations.py
"""
exec_customizations = []
for cust in customizations:
if cust.needs_build:
exec_customizations.append(cust)
return exec_customizations
def download_all_packages(self, custs):
""" download_all_packages downloads all the packages referenced by given
custs.
Args:
* custs: List of Customizations object from customizations.py
"""
for cust in custs:
with self.m.step.nest('Download resources for {}'.format(cust.name())):
cust.download_sources()
def execute_customizations(self, custs):
""" Executes the windows image builder user config.
Args:
* custs: List of Customizations object from customizations.py
"""
with self.m.step.nest('execute config {}'.format(custs[0].image().name)):
for cust in custs:
cust.execute_customization()
def get_executable_configs(self, custs):
""" get_executable_configs returns a list of images that can be executed at
this time.
image is generated after determining what customizations can be executed as
part of the image. A list of keys representing the customizations to be
executed are also returned with the image.
Args:
* custs: List of customizations to be processed.
Returns a dict mapping builder name to the image-key_list tuple
"""
# executions is a dict that will be returned
executions = OrderedDict()
for builder in BUILDERS:
# filter all the customizations for the builder
builder_custs = [
cust for cust in custs if CUST_BUILDER[cust.type] == builder
]
# Get the image, key_list tuple for the given builder
images = self.gen_executable_configs(builder_custs)
if images:
executions[builder] = images
return executions
def gen_executable_configs(self, custs):
""" gen_executable_configs generates wib.Image configs that can be executed.
Given a list of custs that can be run on a builder. Generates wib.Image
proto configs making sure that they can be executed independently. Some
customizations are dependent on others for inputs and can only be executed
if the dependent input is generated in the same wib.Image proto or is
already available.
Args:
* custs: list of customization objects from customizations.py that can be
executed on the same builder
Returns a list of tuples containing config and set of customization hash
that can be executed at the time
"""
# Dict mapping the customization object key to itself
key_cust_map = OrderedDict()
for cust in custs:
# Update the key map
key_cust_map[cust.get_key()] = cust
# List of customization keys that are already processed
processed_cust = []
# List of images. This will be returned by the method
images = []
for key, cust in key_cust_map.items():
# First find a customization that can be executed without any dependency
if key not in processed_cust and cust.executable():
# image is list of customizations that can be executed on the builder
image = [key]
# ensure that image contains customizations for the same arch
arch = cust.arch
# Add the image to the list of images to return
images.append(image)
# record the customization key that is already being processed
processed_cust.append(key)
# get the list of outputs that are generated by this customization. This
# will be used to determine if we can execute the following
# customizations
outputs = [self._sources.get_url(op) for op in cust.outputs]
for nkey, node in key_cust_map.items():
# If we can execute this customization then add it to the image
if node.arch == arch and nkey not in processed_cust and \
node.executable(
inputs=outputs):
# Add the node to list of processed_cust
processed_cust.append(nkey)
# Add the node to the image
image.append(nkey)
# Extend the outputs to include the ones generated by this node
outputs.extend([self._sources.get_url(op) for op in node.outputs])
if len(image) >= self._max_cust_batch_size:
break
# Convert all the images to wib.Image configs and return back
configs = []
for image in images:
config = wib.Image()
config.name = '_'.join(
sorted(
set([key_cust_map[cust_key].image().name for cust_key in image])))
keys = set()
for cust_key in image:
customization = key_cust_map[cust_key].customization()
# Include all the image names this config represents
config.customizations.append(customization)
keys.add(cust_key)
# Add an image name and arch
config.arch = key_cust_map[image[0]].image().arch
configs.append((config, keys))
return configs