1 /* 2 * Copyright (C) 2019 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 package android.device.collectors; 17 18 import android.app.StatsManager; 19 import android.app.StatsManager.StatsUnavailableException; 20 import android.content.Context; 21 import android.content.res.AssetManager; 22 import android.os.Bundle; 23 import android.os.Environment; 24 import android.util.Log; 25 import androidx.annotation.VisibleForTesting; 26 import androidx.test.InstrumentationRegistry; 27 28 import com.android.internal.os.StatsdConfigProto.StatsdConfig; 29 import com.android.os.StatsLog.ConfigMetricsReportList; 30 import com.google.protobuf.InvalidProtocolBufferException; 31 32 import org.junit.runner.Description; 33 import org.junit.runner.Result; 34 35 import java.io.File; 36 import java.io.IOException; 37 import java.io.InputStream; 38 import java.nio.file.Files; 39 import java.nio.file.Path; 40 import java.nio.file.Paths; 41 import java.util.Arrays; 42 import java.util.HashMap; 43 import java.util.List; 44 import java.util.Map; 45 import java.util.UUID; 46 import java.util.function.Function; 47 import java.util.stream.Collectors; 48 49 /** A device-side metric listener that collects statsd-based metrics using bundled config files. */ 50 public class StatsdListener extends BaseMetricListener { 51 private static final String LOG_TAG = StatsdListener.class.getSimpleName(); 52 53 // TODO(harrytczhang): Add option and support for per-test collection. 54 static final String OPTION_CONFIGS_TEST_RUN = "statsd-configs-per-run"; 55 56 // Sub-directory within the test APK's assets/ directory to look for configs. 57 static final String CONFIG_SUB_DIRECTORY = "statsd-configs"; 58 // File extension for all statsd configs. 59 static final String PROTO_EXTENSION = ".pb"; 60 61 // Parent directory for all statsd reports. 62 static final String REPORT_PATH_ROOT = "statsd-reports"; 63 // Sub-directory for test run reports. 64 static final String REPORT_PATH_TEST_RUN = "test-run"; 65 66 // Configs used for tests and test runs, respectively. 67 private Map<String, StatsdConfig> mTestRunConfigs = new HashMap<String, StatsdConfig>(); 68 69 // Map to associate config names with their config Ids. 70 private Map<String, Long> mTestRunConfigIds = new HashMap<String, Long>(); 71 72 // Cached stats manager instance. 73 private StatsManager mStatsManager; 74 75 /** Registers the test run configs with {@link StatsManager} before the test run starts. */ 76 @Override onTestRunStart(DataRecord runData, Description description)77 public void onTestRunStart(DataRecord runData, Description description) { 78 // The argument parsing has to be performed here as the instrumentation has not yet been 79 // registered when the constructor of this class is called. 80 mTestRunConfigs.putAll(getConfigsFromOption(OPTION_CONFIGS_TEST_RUN)); 81 82 mTestRunConfigIds = registerConfigsWithStatsManager(mTestRunConfigs); 83 } 84 85 /** 86 * Dumps the test run stats reports to the test run subdirectory after the test run ends. 87 * 88 * <p>Dumps the stats regardless of whether all the tests pass. 89 */ 90 @Override onTestRunEnd(DataRecord runData, Result result)91 public void onTestRunEnd(DataRecord runData, Result result) { 92 Map<String, File> configReports = 93 pullReportsAndRemoveConfigs( 94 mTestRunConfigIds, Paths.get(REPORT_PATH_ROOT, REPORT_PATH_TEST_RUN)); 95 for (String configName : configReports.keySet()) { 96 runData.addFileMetric(configName, configReports.get(configName)); 97 } 98 } 99 100 /** 101 * Register a set of statsd configs and return their config IDs in a {@link Map}. 102 * 103 * @param configs Map of (config name, config proto message) 104 * @return Map of (config name, config id) 105 */ registerConfigsWithStatsManager( final Map<String, StatsdConfig> configs)106 private Map<String, Long> registerConfigsWithStatsManager( 107 final Map<String, StatsdConfig> configs) { 108 Map<String, Long> configIds = new HashMap<String, Long>(); 109 adoptShellPermissionIdentity(); 110 for (String configName : configs.keySet()) { 111 long configId = getUniqueIdForConfig(configs.get(configName)); 112 StatsdConfig newConfig = configs.get(configName).toBuilder().setId(configId).build(); 113 try { 114 addStatsConfig(configId, newConfig.toByteArray()); 115 configIds.put(configName, configId); 116 } catch (StatsUnavailableException e) { 117 Log.e( 118 LOG_TAG, 119 String.format( 120 "Failed to add statsd config %s due to %s.", 121 configName, e.toString())); 122 } 123 } 124 dropShellPermissionIdentity(); 125 return configIds; 126 } 127 128 /** 129 * For a set of statsd config ids, retrieve the config reports from {@link StatsManager}, remove 130 * the config and dump the reports into the designated directory on the device's external 131 * storage. 132 * 133 * @param configIds Map of (config name, config Id) 134 * @param directory relative directory on external storage to dump the report in. Each report 135 * will be named after its config. 136 * @return Map of (config name, config report file) 137 */ pullReportsAndRemoveConfigs( final Map<String, Long> configIds, Path directory)138 private Map<String, File> pullReportsAndRemoveConfigs( 139 final Map<String, Long> configIds, Path directory) { 140 File externalStorage = Environment.getExternalStorageDirectory(); 141 File saveDirectory = new File(externalStorage, directory.toString()); 142 if (!saveDirectory.isDirectory()) { 143 saveDirectory.mkdirs(); 144 } 145 Map<String, File> savedConfigFiles = new HashMap<String, File>(); 146 adoptShellPermissionIdentity(); 147 for (String configName : configIds.keySet()) { 148 // Dump the metric report to external storage. 149 ConfigMetricsReportList reportList; 150 try { 151 reportList = 152 ConfigMetricsReportList.parseFrom( 153 getStatsReports(configIds.get(configName))); 154 File reportFile = new File(saveDirectory, configName + PROTO_EXTENSION); 155 writeToFile(reportFile, reportList.toByteArray()); 156 savedConfigFiles.put(configName, reportFile); 157 } catch (StatsUnavailableException e) { 158 Log.e( 159 LOG_TAG, 160 String.format( 161 "Failed to retrieve metrics for config %s due to %s.", 162 configName, e.toString())); 163 } catch (InvalidProtocolBufferException e) { 164 Log.e( 165 LOG_TAG, 166 String.format( 167 "Unable to parse report for config %s. Details: %s.", 168 configName, e.toString())); 169 } catch (IOException e) { 170 Log.e( 171 LOG_TAG, 172 String.format( 173 "Failed to write metric report for config %s to device. " 174 + "Details: %s.", 175 configName, e.toString())); 176 } 177 178 // Remove the statsd config. 179 try { 180 removeStatsConfig(configIds.get(configName)); 181 } catch (StatsUnavailableException e) { 182 Log.e( 183 LOG_TAG, 184 String.format( 185 "Unable to remove config %s due to %s.", configName, e.toString())); 186 } 187 } 188 dropShellPermissionIdentity(); 189 return savedConfigFiles; 190 } 191 192 /** 193 * Adopt shell permission identity to communicate with {@link StatsManager}. 194 * 195 * @hide 196 */ 197 @VisibleForTesting adoptShellPermissionIdentity()198 protected void adoptShellPermissionIdentity() { 199 InstrumentationRegistry.getInstrumentation() 200 .getUiAutomation() 201 .adoptShellPermissionIdentity(); 202 } 203 204 /** 205 * Drop shell permission identity once communication with {@link StatsManager} is done. 206 * 207 * @hide 208 */ 209 @VisibleForTesting dropShellPermissionIdentity()210 protected void dropShellPermissionIdentity() { 211 InstrumentationRegistry.getInstrumentation() 212 .getUiAutomation() 213 .dropShellPermissionIdentity(); 214 } 215 216 /** Returns the cached {@link StatsManager} instance; if none exists, request and cache it. */ getStatsManager()217 private StatsManager getStatsManager() { 218 if (mStatsManager == null) { 219 mStatsManager = 220 (StatsManager) 221 InstrumentationRegistry.getTargetContext() 222 .getSystemService(Context.STATS_MANAGER); 223 } 224 return mStatsManager; 225 } 226 227 /** 228 * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked. 229 * 230 * @hide 231 */ 232 @VisibleForTesting addStatsConfig(long configKey, byte[] config)233 protected void addStatsConfig(long configKey, byte[] config) throws StatsUnavailableException { 234 getStatsManager().addConfig(configKey, config); 235 } 236 237 /** 238 * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked. 239 * 240 * @hide 241 */ 242 @VisibleForTesting removeStatsConfig(long configKey)243 protected void removeStatsConfig(long configKey) throws StatsUnavailableException { 244 mStatsManager.removeConfig(configKey); 245 } 246 247 /** 248 * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked. 249 * 250 * @hide 251 */ 252 @VisibleForTesting getStatsReports(long configKey)253 protected byte[] getStatsReports(long configKey) throws StatsUnavailableException { 254 return mStatsManager.getReports(configKey); 255 } 256 257 /** 258 * Allow tests to stub out getting instrumentation arguments. 259 * 260 * @hide 261 */ 262 @VisibleForTesting getArguments()263 protected Bundle getArguments() { 264 return InstrumentationRegistry.getArguments(); 265 } 266 267 /** 268 * Allow tests to stub out file I/O. 269 * 270 * @hide 271 */ 272 @VisibleForTesting writeToFile(File f, byte[] content)273 protected File writeToFile(File f, byte[] content) throws IOException { 274 Files.write(f.toPath(), content); 275 return f; 276 } 277 278 /** 279 * Allow tests to override the random ID generation. The config is passed in to allow a specific 280 * ID to be associated with a config in the test. 281 * 282 * @hide 283 */ 284 @VisibleForTesting getUniqueIdForConfig(StatsdConfig config)285 protected long getUniqueIdForConfig(StatsdConfig config) { 286 return (long) UUID.randomUUID().hashCode(); 287 } 288 289 /** 290 * Allow tests to stub out {@link AssetManager} interactions as that class is final and cannot . 291 * be mocked. 292 * 293 * @hide 294 */ 295 @VisibleForTesting openConfigWithAssetManager(AssetManager manager, String configName)296 protected InputStream openConfigWithAssetManager(AssetManager manager, String configName) 297 throws IOException { 298 String configFilePath = 299 Paths.get(CONFIG_SUB_DIRECTORY, configName + PROTO_EXTENSION).toString(); 300 return manager.open(configFilePath); 301 } 302 303 /** 304 * Parse a config from its name using {@link AssetManager}. 305 * 306 * <p>The option name is passed in for better error messaging. 307 */ parseConfigFromName( final AssetManager manager, String optionName, String configName)308 private StatsdConfig parseConfigFromName( 309 final AssetManager manager, String optionName, String configName) { 310 try (InputStream configStream = openConfigWithAssetManager(manager, configName)) { 311 try { 312 return StatsdConfig.parseFrom(configStream); 313 } catch (IOException e) { 314 throw new RuntimeException( 315 String.format( 316 "Cannot parse profile %s in option %s.", configName, optionName), 317 e); 318 } 319 } catch (IOException e) { 320 throw new IllegalArgumentException( 321 String.format( 322 "Config name %s in option %s does not exist", configName, optionName)); 323 } 324 } 325 326 /** 327 * Parse the suppplied option to get a set of statsd configs keyed by their names. 328 * 329 * @hide 330 */ 331 @VisibleForTesting getConfigsFromOption(String optionName)332 protected Map<String, StatsdConfig> getConfigsFromOption(String optionName) { 333 List<String> configNames = 334 Arrays.asList(getArguments().getString(optionName, "").split(",")) 335 .stream() 336 .map(s -> s.trim()) 337 .filter(s -> !s.isEmpty()) 338 .distinct() 339 .collect(Collectors.toList()); 340 // Look inside the APK assets for the configuration file. 341 final AssetManager manager = InstrumentationRegistry.getContext().getAssets(); 342 return configNames 343 .stream() 344 .collect( 345 Collectors.toMap( 346 Function.identity(), 347 configName -> 348 parseConfigFromName(manager, optionName, configName))); 349 } 350 } 351