blob: 289c2d5c1837faef1650b144708a42d2cc84556c [file] [log] [blame] [edit]
# -*- coding: utf-8 -*-
# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""A class that sets up the environment for telemetry testing."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from autotest_lib.client.common_lib.cros import dev_server
import contextlib
import errno
import fcntl
import logging
import os
import shutil
import subprocess
import tempfile
import requests
_READ_BUFFER_SIZE_BYTES = 1024 * 1024 # 1 MB
@contextlib.contextmanager
def lock_dir(dir_name):
"""Lock a directory exclusively by placing a file lock in it.
Args:
dir_name: the directory name to be locked.
"""
lock_file = os.path.join(dir_name, '.lock')
with open(lock_file, 'w+') as f:
fcntl.flock(f, fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(f, fcntl.LOCK_UN)
class TelemetrySetupError(Exception):
"""Exception class used by this module."""
pass
class TelemetrySetup(object):
"""Class that sets up the environment for telemetry testing."""
# Relevant directory paths.
_BASE_DIR_PATH = '/tmp/telemetry-workdir'
_PARTIAL_DEPENDENCY_DIR_PATH = 'autotest/packages'
# Relevant directory names.
_TELEMETRY_SRC_DIR_NAME = 'telemetry_src'
_TEST_SRC_DIR_NAME = 'test_src'
_SRC_DIR_NAME = 'src'
# Names of the telemetry dependency tarballs.
_DEPENDENCIES = [
'dep-telemetry_dep.tar.bz2',
'dep-page_cycler_dep.tar.bz2',
'dep-chrome_test.tar.bz2',
'dep-perf_data_dep.tar.bz2',
]
# Partial devserver URLs.
_STATIC_URL_TEMPLATE = '%s/static/%s/autotest/packages/%s'
def __init__(self, hostname, build):
"""Initializes the TelemetrySetup class.
Args:
hostname: The host for which telemetry environment should be setup. This
is important for devserver resolution.
build: The build for which telemetry environment should be setup. It is
typically in the format <board>/<version>.
"""
self._build = build
self._ds = dev_server.ImageServer.resolve(self._build,
hostname=hostname)
self._setup_dir_path = tempfile.mkdtemp(prefix='telemetry-setupdir_')
self._tmp_build_dir = os.path.join(self._BASE_DIR_PATH, self._build)
self._tlm_src_dir_path = os.path.join(self._tmp_build_dir,
self._TELEMETRY_SRC_DIR_NAME)
def Setup(self):
"""Sets up the environment for telemetry testing.
This method downloads the telemetry dependency tarballs and extracts
them into a 'src' directory.
Returns:
Path to the src directory where the telemetry dependencies have been
downloaded and extracted.
"""
src_folder = os.path.join(self._tlm_src_dir_path, self._SRC_DIR_NAME)
test_src = os.path.join(self._tlm_src_dir_path,
self._TEST_SRC_DIR_NAME)
self._MkDirP(self._tlm_src_dir_path)
with lock_dir(self._tlm_src_dir_path):
if not os.path.exists(src_folder):
# Download the required dependency tarballs.
for dep in self._DEPENDENCIES:
dep_path = self._DownloadFilesFromDevserver(
dep, self._setup_dir_path)
if os.path.exists(dep_path):
self._ExtractTarball(dep_path, self._tlm_src_dir_path)
# By default all the tarballs extract to test_src but some parts
# of the telemetry code specifically hardcoded to exist inside
# of 'src'.
try:
shutil.move(test_src, src_folder)
except shutil.Error:
raise TelemetrySetupError(
'Failure in telemetry setup for build %s. Appears '
'that the test_src to src move failed.' %
self._build)
return src_folder
def _DownloadFilesFromDevserver(self, filename, dest_path):
"""Downloads the given tar.bz2 file from the devserver.
Args:
filename: Name of the tar.bz2 file to be downloaded.
dest_path: Full path to the directory where it should be downloaded.
Returns:
Full path to the downloaded file.
Raises:
TelemetrySetupError when the download cannot be completed for any
reason.
"""
dep_path = os.path.join(dest_path, filename)
url = (self._STATIC_URL_TEMPLATE %
(self._ds.url(), self._build, filename))
try:
resp = requests.get(url)
resp.raise_for_status()
with open(dep_path, 'w') as f:
for content in resp.iter_content(_READ_BUFFER_SIZE_BYTES):
f.write(content)
except Exception as e:
if (isinstance(e, requests.exceptions.HTTPError)
and resp.status_code == 404):
logging.error(
'The request %s returned a 404 Not Found status.'
'This dependency could be new and therefore does not '
'exist. Hence, squashing the exception and proceeding.',
url)
elif isinstance(e, requests.exceptions.ConnectionError):
logging.warning(
'The request failed because a connection to the devserver '
'%s could not be established. Attempting to execute the '
'request %s once by SSH-ing into the devserver.',
self._ds.url(), url)
return self._DownloadFilesFromDevserverViaSSH(url, dep_path)
else:
raise TelemetrySetupError(
'An error occurred while trying to complete %s: %s' %
(url, e))
return dep_path
def _DownloadFilesFromDevserverViaSSH(self, url, dep_path):
"""Downloads the file at the URL from the devserver by SSH-ing into it.
Args:
url: URL of the location of the tar.bz2 file on the devserver.
dep_path: Full path to the file where it will be downloaded.
Returns:
Full path to the downloaded file.
Raises:
TelemetrySetupError when the download cannot be completed for any
reason.
"""
cmd = ['ssh', self._ds.hostname, 'curl', url]
with open(dep_path, 'w') as f:
proc = subprocess.Popen(cmd, stdout=f, stderr=subprocess.PIPE)
_, err = proc.communicate()
if proc.returncode != 0:
raise TelemetrySetupError(
'The command: %s finished with returncode %s and '
'errors as following: %s. The telemetry dependency '
'could not be downloaded.' %
(' '.join(cmd), proc.returncode, err))
return dep_path
def _ExtractTarball(self, tarball_path, dest_path):
"""Extracts the given tarball into the destination directory.
Args:
tarball_path: Full path to the tarball to be extracted.
dest_path: Full path to the directory where the tarball should be
extracted.
Raises:
TelemetrySetupError if the method is unable to extract the tarball for
any reason.
"""
cmd = ['tar', 'xf', tarball_path, '--directory', dest_path]
try:
proc = subprocess.Popen(cmd,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
proc.communicate()
except Exception as e:
shutil.rmtree(dest_path)
raise TelemetrySetupError(
'An exception occurred while trying to untar %s into %s: %s'
% (tarball_path, dest_path, str(e)))
def _MkDirP(self, path):
"""Recursively creates the given directory.
Args:
path: Full path to the directory that needs to the created.
Raises:
TelemetrySetupError is the method is unable to create directories for
any reason except OSError EEXIST which indicates that the
directory already exists.
"""
try:
os.makedirs(path)
except Exception as e:
if not isinstance(e, OSError) or e.errno != errno.EEXIST:
raise TelemetrySetupError(
'Could not create directory %s due to %s.' %
(path, str(e)))
def Cleanup(self):
"""Cleans up telemetry setup and work environment."""
try:
shutil.rmtree(self._setup_dir_path)
except Exception as e:
logging.error('Something went wrong. Could not delete %s: %s',
self._setup_dir_path, e)
try:
shutil.rmtree(self._tlm_src_dir_path)
except Exception as e:
logging.error('Something went wrong. Could not delete %s: %s',
self._tlm_src_dir_path, e)