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.os.SystemClock;
25 import android.util.Log;
26 import android.util.StatsLog;
27 import androidx.annotation.VisibleForTesting;
28 import androidx.test.InstrumentationRegistry;
29 
30 import com.android.internal.os.StatsdConfigProto.StatsdConfig;
31 import com.android.os.StatsLog.ConfigMetricsReportList;
32 import com.google.protobuf.InvalidProtocolBufferException;
33 
34 import org.junit.runner.Description;
35 import org.junit.runner.Result;
36 
37 import java.io.File;
38 import java.io.IOException;
39 import java.io.InputStream;
40 import java.nio.file.Files;
41 import java.nio.file.Path;
42 import java.nio.file.Paths;
43 import java.util.Arrays;
44 import java.util.HashMap;
45 import java.util.List;
46 import java.util.Map;
47 import java.util.UUID;
48 import java.util.concurrent.TimeUnit;
49 import java.util.function.Function;
50 import java.util.stream.Collectors;
51 
52 /** A device-side metric listener that collects statsd-based metrics using bundled config files. */
53 public class StatsdListener extends BaseMetricListener {
54     private static final String LOG_TAG = StatsdListener.class.getSimpleName();
55 
56     static final String OPTION_CONFIGS_RUN_LEVEL = "statsd-configs-run-level";
57     static final String OPTION_CONFIGS_TEST_LEVEL = "statsd-configs-test-level";
58 
59     // Sub-directory within the test APK's assets/ directory to look for configs.
60     static final String CONFIG_SUB_DIRECTORY = "statsd-configs";
61     // File extension for all statsd configs.
62     static final String PROTO_EXTENSION = ".pb";
63 
64     // Parent directory for all statsd reports.
65     static final String REPORT_PATH_ROOT = "statsd-reports";
66     // Sub-directory for test run reports.
67     static final String REPORT_PATH_RUN_LEVEL = "run-level";
68     // Sub-directory for test-level reports.
69     static final String REPORT_PATH_TEST_LEVEL = "test-level";
70     // Suffix template for test-level metric report files.
71     static final String TEST_SUFFIX_TEMPLATE = "_%s-%d";
72 
73     // Common prefix for the metric key pointing to the report path.
74     static final String REPORT_KEY_PREFIX = "statsd-";
75     // Common prefix for the metric file.
76     static final String REPORT_FILENAME_PREFIX = "statsd-";
77 
78     // Labels used to signify test events to statsd with the AppBreadcrumbReported atom.
79     static final int RUN_EVENT_LABEL = 7;
80     static final int TEST_EVENT_LABEL = 11;
81     // A short delay after pushing the AppBreadcrumbReported event so that metrics can be dumped.
82     static final long METRIC_PULL_DELAY = TimeUnit.SECONDS.toMillis(1);
83 
84     // Configs used for the test run and each test, respectively.
85     private Map<String, StatsdConfig> mRunLevelConfigs = new HashMap<String, StatsdConfig>();
86     private Map<String, StatsdConfig> mTestLevelConfigs = new HashMap<String, StatsdConfig>();
87 
88     // Map to associate config names with their config Ids.
89     private Map<String, Long> mRunLevelConfigIds = new HashMap<String, Long>();
90     private Map<String, Long> mTestLevelConfigIds = new HashMap<String, Long>();
91 
92     // "Counter" for test iterations, keyed by the display name of each test's description.
93     private Map<String, Integer> mTestIterations = new HashMap<String, Integer>();
94 
95     // Cached stats manager instance.
96     private StatsManager mStatsManager;
97 
98     /** Register the test run configs with {@link StatsManager} before the test run starts. */
99     @Override
onTestRunStart(DataRecord runData, Description description)100     public void onTestRunStart(DataRecord runData, Description description) {
101         // The argument parsing has to be performed here as the instrumentation has not yet been
102         // registered when the constructor of this class is called.
103         mRunLevelConfigs.putAll(getConfigsFromOption(OPTION_CONFIGS_RUN_LEVEL));
104         mTestLevelConfigs.putAll(getConfigsFromOption(OPTION_CONFIGS_TEST_LEVEL));
105 
106         mRunLevelConfigIds = registerConfigsWithStatsManager(mRunLevelConfigs);
107 
108         if (!logStart(RUN_EVENT_LABEL)) {
109             Log.w(LOG_TAG, "Failed to log a test run start event. Metrics might be incomplete.");
110         }
111     }
112 
113     /**
114      * Dump the test run stats reports to the test run subdirectory after the test run ends.
115      *
116      * <p>Dumps the stats regardless of whether all the tests pass.
117      */
118     @Override
onTestRunEnd(DataRecord runData, Result result)119     public void onTestRunEnd(DataRecord runData, Result result) {
120         if (!logStop(RUN_EVENT_LABEL)) {
121             Log.w(LOG_TAG, "Failed to log a test run end event. Metrics might be incomplete.");
122         }
123         SystemClock.sleep(METRIC_PULL_DELAY);
124 
125         Map<String, File> configReports =
126                 pullReportsAndRemoveConfigs(
127                         mRunLevelConfigIds, Paths.get(REPORT_PATH_ROOT, REPORT_PATH_RUN_LEVEL), "");
128         for (String configName : configReports.keySet()) {
129             runData.addFileMetric(REPORT_KEY_PREFIX + configName, configReports.get(configName));
130         }
131     }
132 
133     /** Register the test-level configs with {@link StatsManager} before each test starts. */
134     @Override
onTestStart(DataRecord testData, Description description)135     public void onTestStart(DataRecord testData, Description description) {
136         mTestIterations.computeIfPresent(description.getDisplayName(), (name, count) -> count + 1);
137         mTestIterations.computeIfAbsent(description.getDisplayName(), name -> 1);
138         mTestLevelConfigIds = registerConfigsWithStatsManager(mTestLevelConfigs);
139 
140         if (!logStart(TEST_EVENT_LABEL)) {
141             Log.w(LOG_TAG, "Failed to log a test start event. Metrics might be incomplete.");
142         }
143     }
144 
145     /**
146      * Dump the test-level stats reports to the test-specific subdirectory after the test ends.
147      *
148      * <p>Dumps the stats regardless of whether the test passes.
149      */
150     @Override
onTestEnd(DataRecord testData, Description description)151     public void onTestEnd(DataRecord testData, Description description) {
152         if (!logStop(TEST_EVENT_LABEL)) {
153             Log.w(LOG_TAG, "Failed to log a test end event. Metrics might be incomplete.");
154         }
155         SystemClock.sleep(METRIC_PULL_DELAY);
156 
157         Map<String, File> configReports =
158                 pullReportsAndRemoveConfigs(
159                         mTestLevelConfigIds,
160                         Paths.get(REPORT_PATH_ROOT, REPORT_PATH_TEST_LEVEL),
161                         getTestSuffix(description));
162         for (String configName : configReports.keySet()) {
163             testData.addFileMetric(REPORT_KEY_PREFIX + configName, configReports.get(configName));
164         }
165     }
166 
167     /**
168      * Register a set of statsd configs and return their config IDs in a {@link Map}.
169      *
170      * @param configs Map of (config name, config proto message)
171      * @return Map of (config name, config id)
172      */
registerConfigsWithStatsManager( final Map<String, StatsdConfig> configs)173     private Map<String, Long> registerConfigsWithStatsManager(
174             final Map<String, StatsdConfig> configs) {
175         Map<String, Long> configIds = new HashMap<String, Long>();
176         adoptShellPermissionIdentity();
177         for (String configName : configs.keySet()) {
178             long configId = getUniqueIdForConfig(configs.get(configName));
179             StatsdConfig newConfig = configs.get(configName).toBuilder().setId(configId).build();
180             try {
181                 addStatsConfig(configId, newConfig.toByteArray());
182                 configIds.put(configName, configId);
183             } catch (StatsUnavailableException e) {
184                 Log.e(
185                         LOG_TAG,
186                         String.format(
187                                 "Failed to add statsd config %s due to %s.",
188                                 configName, e.toString()));
189             }
190         }
191         dropShellPermissionIdentity();
192         return configIds;
193     }
194 
195     /**
196      * For a set of statsd config ids, retrieve the config reports from {@link StatsManager}, remove
197      * the config and dump the reports into the designated directory on the device's external
198      * storage.
199      *
200      * @param configIds Map of (config name, config Id)
201      * @param directory relative directory on external storage to dump the report in. Each report
202      *     will be named after its config.
203      * @param suffix a suffix to append to the metric report file name, used to differentiate
204      *     between tests and left empty for the test run.
205      * @return Map of (config name, config report file)
206      */
pullReportsAndRemoveConfigs( final Map<String, Long> configIds, Path directory, String suffix)207     private Map<String, File> pullReportsAndRemoveConfigs(
208             final Map<String, Long> configIds, Path directory, String suffix) {
209         File externalStorage = Environment.getExternalStorageDirectory();
210         File saveDirectory = new File(externalStorage, directory.toString());
211         if (!saveDirectory.isDirectory()) {
212             saveDirectory.mkdirs();
213         }
214         Map<String, File> savedConfigFiles = new HashMap<String, File>();
215         adoptShellPermissionIdentity();
216         for (String configName : configIds.keySet()) {
217             // Dump the metric report to external storage.
218             ConfigMetricsReportList reportList;
219             try {
220                 reportList =
221                         ConfigMetricsReportList.parseFrom(
222                                 getStatsReports(configIds.get(configName)));
223                 File reportFile =
224                         new File(
225                                 saveDirectory,
226                                 REPORT_FILENAME_PREFIX + configName + suffix + PROTO_EXTENSION);
227                 writeToFile(reportFile, reportList.toByteArray());
228                 savedConfigFiles.put(configName, reportFile);
229             } catch (StatsUnavailableException e) {
230                 Log.e(
231                         LOG_TAG,
232                         String.format(
233                                 "Failed to retrieve metrics for config %s due to %s.",
234                                 configName, e.toString()));
235             } catch (InvalidProtocolBufferException e) {
236                 Log.e(
237                         LOG_TAG,
238                         String.format(
239                                 "Unable to parse report for config %s. Details: %s.",
240                                 configName, e.toString()));
241             } catch (IOException e) {
242                 Log.e(
243                         LOG_TAG,
244                         String.format(
245                                 "Failed to write metric report for config %s to device. "
246                                         + "Details: %s.",
247                                 configName, e.toString()));
248             }
249 
250             // Remove the statsd config.
251             try {
252                 removeStatsConfig(configIds.get(configName));
253             } catch (StatsUnavailableException e) {
254                 Log.e(
255                         LOG_TAG,
256                         String.format(
257                                 "Unable to remove config %s due to %s.", configName, e.toString()));
258             }
259         }
260         dropShellPermissionIdentity();
261         return savedConfigFiles;
262     }
263 
264     /**
265      * Adopt shell permission identity to communicate with {@link StatsManager}.
266      *
267      * @hide
268      */
269     @VisibleForTesting
adoptShellPermissionIdentity()270     protected void adoptShellPermissionIdentity() {
271         InstrumentationRegistry.getInstrumentation()
272                 .getUiAutomation()
273                 .adoptShellPermissionIdentity();
274     }
275 
276     /**
277      * Drop shell permission identity once communication with {@link StatsManager} is done.
278      *
279      * @hide
280      */
281     @VisibleForTesting
dropShellPermissionIdentity()282     protected void dropShellPermissionIdentity() {
283         InstrumentationRegistry.getInstrumentation()
284                 .getUiAutomation()
285                 .dropShellPermissionIdentity();
286     }
287 
288     /** Returns the cached {@link StatsManager} instance; if none exists, request and cache it. */
getStatsManager()289     private StatsManager getStatsManager() {
290         if (mStatsManager == null) {
291             mStatsManager =
292                     (StatsManager)
293                             InstrumentationRegistry.getTargetContext()
294                                     .getSystemService(Context.STATS_MANAGER);
295         }
296         return mStatsManager;
297     }
298 
299     /** Get the suffix for a test + iteration combination to differentiate it from other files. */
300     @VisibleForTesting
getTestSuffix(Description description)301     String getTestSuffix(Description description) {
302         return String.format(
303                 TEST_SUFFIX_TEMPLATE,
304                 formatDescription(description),
305                 mTestIterations.get(description.getDisplayName()));
306     }
307 
308     /** Format a JUnit {@link Description} to a desired string format. */
309     @VisibleForTesting
formatDescription(Description description)310     String formatDescription(Description description) {
311         // Use String.valueOf() to guard agaist a null class name. This normally should not happen
312         // but the Description class does not explicitly guarantee it.
313         String className = String.valueOf(description.getClassName());
314         String methodName = description.getMethodName();
315         return methodName == null ? className : String.join("#", className, methodName);
316     }
317 
318     /**
319      * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked.
320      *
321      * @hide
322      */
323     @VisibleForTesting
addStatsConfig(long configKey, byte[] config)324     protected void addStatsConfig(long configKey, byte[] config) throws StatsUnavailableException {
325         getStatsManager().addConfig(configKey, config);
326     }
327 
328     /**
329      * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked.
330      *
331      * @hide
332      */
333     @VisibleForTesting
removeStatsConfig(long configKey)334     protected void removeStatsConfig(long configKey) throws StatsUnavailableException {
335         mStatsManager.removeConfig(configKey);
336     }
337 
338     /**
339      * Forwarding logic for {@link StatsManager} as it is final and cannot be mocked.
340      *
341      * @hide
342      */
343     @VisibleForTesting
getStatsReports(long configKey)344     protected byte[] getStatsReports(long configKey) throws StatsUnavailableException {
345         return mStatsManager.getReports(configKey);
346     }
347 
348     /**
349      * Allow tests to stub out getting instrumentation arguments.
350      *
351      * @hide
352      */
353     @VisibleForTesting
getArguments()354     protected Bundle getArguments() {
355         return InstrumentationRegistry.getArguments();
356     }
357 
358     /**
359      * Allow tests to stub out file I/O.
360      *
361      * @hide
362      */
363     @VisibleForTesting
writeToFile(File f, byte[] content)364     protected File writeToFile(File f, byte[] content) throws IOException {
365         Files.write(f.toPath(), content);
366         return f;
367     }
368 
369     /**
370      * Allow tests to override the random ID generation. The config is passed in to allow a specific
371      * ID to be associated with a config in the test.
372      *
373      * @hide
374      */
375     @VisibleForTesting
getUniqueIdForConfig(StatsdConfig config)376     protected long getUniqueIdForConfig(StatsdConfig config) {
377         return (long) UUID.randomUUID().hashCode();
378     }
379 
380     /**
381      * Allow tests to stub out {@link AssetManager} interactions as that class is final and cannot .
382      * be mocked.
383      *
384      * @hide
385      */
386     @VisibleForTesting
openConfigWithAssetManager(AssetManager manager, String configName)387     protected InputStream openConfigWithAssetManager(AssetManager manager, String configName)
388             throws IOException {
389         String configFilePath =
390                 Paths.get(CONFIG_SUB_DIRECTORY, configName + PROTO_EXTENSION).toString();
391         return manager.open(configFilePath);
392     }
393 
394     /**
395      * Parse a config from its name using {@link AssetManager}.
396      *
397      * <p>The option name is passed in for better error messaging.
398      */
parseConfigFromName( final AssetManager manager, String optionName, String configName)399     private StatsdConfig parseConfigFromName(
400             final AssetManager manager, String optionName, String configName) {
401         try (InputStream configStream = openConfigWithAssetManager(manager, configName)) {
402             try {
403                 return StatsdConfig.parseFrom(configStream);
404             } catch (IOException e) {
405                 throw new RuntimeException(
406                         String.format(
407                                 "Cannot parse profile %s in option %s.", configName, optionName),
408                         e);
409             }
410         } catch (IOException e) {
411             throw new IllegalArgumentException(
412                     String.format(
413                             "Config name %s in option %s does not exist", configName, optionName));
414         }
415     }
416 
417     /**
418      * Parse the suppplied option to get a set of statsd configs keyed by their names.
419      *
420      * @hide
421      */
422     @VisibleForTesting
getConfigsFromOption(String optionName)423     protected Map<String, StatsdConfig> getConfigsFromOption(String optionName) {
424         List<String> configNames =
425                 Arrays.asList(getArguments().getString(optionName, "").split(","))
426                         .stream()
427                         .map(s -> s.trim())
428                         .filter(s -> !s.isEmpty())
429                         .distinct()
430                         .collect(Collectors.toList());
431         // Look inside the APK assets for the configuration file.
432         final AssetManager manager = InstrumentationRegistry.getContext().getAssets();
433         return configNames
434                 .stream()
435                 .collect(
436                         Collectors.toMap(
437                                 Function.identity(),
438                                 configName ->
439                                         parseConfigFromName(manager, optionName, configName)));
440     }
441 
442     /**
443      * Log a "start" AppBreadcrumbReported event to statsd. Wraps a static method for testing.
444      *
445      * @hide
446      */
447     @VisibleForTesting
logStart(int label)448     protected boolean logStart(int label) {
449         return StatsLog.logStart(label);
450     }
451 
452     /**
453      * Log a "stop" AppBreadcrumbReported event to statsd. Wraps a static method for testing.
454      *
455      * @hide
456      */
457     @VisibleForTesting
logStop(int label)458     protected boolean logStop(int label) {
459         return StatsLog.logStop(label);
460     }
461 }
462