1# -*- coding: utf-8 -*- 2# Copyright 2021 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5"""A class that sets up the environment for telemetry testing.""" 6 7from __future__ import absolute_import 8from __future__ import division 9from __future__ import print_function 10 11from autotest_lib.client.common_lib.cros import dev_server 12 13import contextlib 14import errno 15import fcntl 16import logging 17import os 18import shutil 19import subprocess 20import tempfile 21 22import requests 23 24_READ_BUFFER_SIZE_BYTES = 1024 * 1024 # 1 MB 25 26 27@contextlib.contextmanager 28def lock_dir(dir_name): 29 """Lock a directory exclusively by placing a file lock in it. 30 31 Args: 32 dir_name: the directory name to be locked. 33 """ 34 lock_file = os.path.join(dir_name, '.lock') 35 with open(lock_file, 'w+') as f: 36 fcntl.flock(f, fcntl.LOCK_EX) 37 try: 38 yield 39 finally: 40 fcntl.flock(f, fcntl.LOCK_UN) 41 42 43class TelemetrySetupError(Exception): 44 """Exception class used by this module.""" 45 pass 46 47 48class TelemetrySetup(object): 49 """Class that sets up the environment for telemetry testing.""" 50 51 # Relevant directory paths. 52 _BASE_DIR_PATH = '/tmp/telemetry-workdir' 53 _PARTIAL_DEPENDENCY_DIR_PATH = 'autotest/packages' 54 55 # Relevant directory names. 56 _TELEMETRY_SRC_DIR_NAME = 'telemetry_src' 57 _TEST_SRC_DIR_NAME = 'test_src' 58 _SRC_DIR_NAME = 'src' 59 60 # Names of the telemetry dependency tarballs. 61 _DEPENDENCIES = [ 62 'dep-telemetry_dep.tar.bz2', 63 'dep-page_cycler_dep.tar.bz2', 64 'dep-chrome_test.tar.bz2', 65 'dep-perf_data_dep.tar.bz2', 66 ] 67 68 # Partial devserver URLs. 69 _STATIC_URL_TEMPLATE = '%s/static/%s/autotest/packages/%s' 70 71 def __init__(self, hostname, build): 72 """Initializes the TelemetrySetup class. 73 74 Args: 75 hostname: The host for which telemetry environment should be setup. This 76 is important for devserver resolution. 77 build: The build for which telemetry environment should be setup. It is 78 typically in the format <board>/<version>. 79 """ 80 self._build = build 81 self._ds = dev_server.ImageServer.resolve(self._build, 82 hostname=hostname) 83 self._setup_dir_path = tempfile.mkdtemp(prefix='telemetry-setupdir_') 84 self._tmp_build_dir = os.path.join(self._BASE_DIR_PATH, self._build) 85 self._tlm_src_dir_path = os.path.join(self._tmp_build_dir, 86 self._TELEMETRY_SRC_DIR_NAME) 87 88 def Setup(self): 89 """Sets up the environment for telemetry testing. 90 91 This method downloads the telemetry dependency tarballs and extracts 92 them into a 'src' directory. 93 94 Returns: 95 Path to the src directory where the telemetry dependencies have been 96 downloaded and extracted. 97 """ 98 src_folder = os.path.join(self._tlm_src_dir_path, self._SRC_DIR_NAME) 99 test_src = os.path.join(self._tlm_src_dir_path, 100 self._TEST_SRC_DIR_NAME) 101 self._MkDirP(self._tlm_src_dir_path) 102 with lock_dir(self._tlm_src_dir_path): 103 if not os.path.exists(src_folder): 104 # Download the required dependency tarballs. 105 for dep in self._DEPENDENCIES: 106 dep_path = self._DownloadFilesFromDevserver( 107 dep, self._setup_dir_path) 108 if os.path.exists(dep_path): 109 self._ExtractTarball(dep_path, self._tlm_src_dir_path) 110 111 # By default all the tarballs extract to test_src but some parts 112 # of the telemetry code specifically hardcoded to exist inside 113 # of 'src'. 114 try: 115 shutil.move(test_src, src_folder) 116 except shutil.Error: 117 raise TelemetrySetupError( 118 'Failure in telemetry setup for build %s. Appears ' 119 'that the test_src to src move failed.' % 120 self._build) 121 return src_folder 122 123 def _DownloadFilesFromDevserver(self, filename, dest_path): 124 """Downloads the given tar.bz2 file from the devserver. 125 126 Args: 127 filename: Name of the tar.bz2 file to be downloaded. 128 dest_path: Full path to the directory where it should be downloaded. 129 130 Returns: 131 Full path to the downloaded file. 132 133 Raises: 134 TelemetrySetupError when the download cannot be completed for any 135 reason. 136 """ 137 dep_path = os.path.join(dest_path, filename) 138 url = (self._STATIC_URL_TEMPLATE % 139 (self._ds.url(), self._build, filename)) 140 try: 141 resp = requests.get(url) 142 resp.raise_for_status() 143 with open(dep_path, 'wb') as f: 144 for content in resp.iter_content(_READ_BUFFER_SIZE_BYTES): 145 f.write(content) 146 except Exception as e: 147 if (isinstance(e, requests.exceptions.HTTPError) 148 and resp.status_code == 404): 149 logging.error( 150 'The request %s returned a 404 Not Found status.' 151 'This dependency could be new and therefore does not ' 152 'exist. Hence, squashing the exception and proceeding.', 153 url) 154 elif isinstance(e, requests.exceptions.ConnectionError): 155 logging.warning( 156 'The request failed because a connection to the devserver ' 157 '%s could not be established. Attempting to execute the ' 158 'request %s once by SSH-ing into the devserver.', 159 self._ds.url(), url) 160 return self._DownloadFilesFromDevserverViaSSH(url, dep_path) 161 else: 162 raise TelemetrySetupError( 163 'An error occurred while trying to complete %s: %s' % 164 (url, e)) 165 return dep_path 166 167 def _DownloadFilesFromDevserverViaSSH(self, url, dep_path): 168 """Downloads the file at the URL from the devserver by SSH-ing into it. 169 170 Args: 171 url: URL of the location of the tar.bz2 file on the devserver. 172 dep_path: Full path to the file where it will be downloaded. 173 174 Returns: 175 Full path to the downloaded file. 176 177 Raises: 178 TelemetrySetupError when the download cannot be completed for any 179 reason. 180 """ 181 cmd = ['ssh', self._ds.hostname, 'curl', url] 182 with open(dep_path, 'w') as f: 183 proc = subprocess.Popen(cmd, stdout=f, stderr=subprocess.PIPE) 184 _, err = proc.communicate() 185 if proc.returncode != 0: 186 raise TelemetrySetupError( 187 'The command: %s finished with returncode %s and ' 188 'errors as following: %s. The telemetry dependency ' 189 'could not be downloaded.' % 190 (' '.join(cmd), proc.returncode, err)) 191 return dep_path 192 193 def _ExtractTarball(self, tarball_path, dest_path): 194 """Extracts the given tarball into the destination directory. 195 196 Args: 197 tarball_path: Full path to the tarball to be extracted. 198 dest_path: Full path to the directory where the tarball should be 199 extracted. 200 201 Raises: 202 TelemetrySetupError if the method is unable to extract the tarball for 203 any reason. 204 """ 205 cmd = ['tar', 'xf', tarball_path, '--directory', dest_path] 206 try: 207 proc = subprocess.Popen(cmd, 208 stdout=subprocess.PIPE, 209 stderr=subprocess.PIPE) 210 proc.communicate() 211 except Exception as e: 212 shutil.rmtree(dest_path) 213 raise TelemetrySetupError( 214 'An exception occurred while trying to untar %s into %s: %s' 215 % (tarball_path, dest_path, str(e))) 216 217 def _MkDirP(self, path): 218 """Recursively creates the given directory. 219 220 Args: 221 path: Full path to the directory that needs to the created. 222 223 Raises: 224 TelemetrySetupError is the method is unable to create directories for 225 any reason except OSError EEXIST which indicates that the 226 directory already exists. 227 """ 228 try: 229 os.makedirs(path) 230 except Exception as e: 231 if not isinstance(e, OSError) or e.errno != errno.EEXIST: 232 raise TelemetrySetupError( 233 'Could not create directory %s due to %s.' % 234 (path, str(e))) 235 236 def Cleanup(self): 237 """Cleans up telemetry setup and work environment.""" 238 try: 239 shutil.rmtree(self._setup_dir_path) 240 except Exception as e: 241 logging.error('Something went wrong. Could not delete %s: %s', 242 self._setup_dir_path, e) 243 try: 244 shutil.rmtree(self._tlm_src_dir_path) 245 except Exception as e: 246 logging.error('Something went wrong. Could not delete %s: %s', 247 self._tlm_src_dir_path, e) 248