1 /*
2  * Copyright (C) 2021 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 
17 package com.android.bedstead.metricsrecorder;
18 
19 import static com.android.os.nano.AtomsProto.Atom.DEVICE_POLICY_EVENT_FIELD_NUMBER;
20 
21 import com.android.bedstead.nene.exceptions.AdbException;
22 import com.android.bedstead.nene.exceptions.NeneException;
23 import com.android.bedstead.nene.utils.ShellCommand;
24 import com.android.framework.protobuf.nano.CodedOutputByteBufferNano;
25 import com.android.framework.protobuf.nano.InvalidProtocolBufferNanoException;
26 import com.android.framework.protobuf.nano.MessageNano;
27 import com.android.internal.os.nano.StatsdConfigProto.AtomMatcher;
28 import com.android.internal.os.nano.StatsdConfigProto.EventMetric;
29 import com.android.internal.os.nano.StatsdConfigProto.FieldValueMatcher;
30 import com.android.internal.os.nano.StatsdConfigProto.SimpleAtomMatcher;
31 import com.android.internal.os.nano.StatsdConfigProto.StatsdConfig;
32 import com.android.os.nano.AtomsProto;
33 import com.android.os.nano.StatsLog;
34 import com.android.os.nano.StatsLog.ConfigMetricsReportList;
35 
36 import java.util.ArrayList;
37 import java.util.Arrays;
38 import java.util.Comparator;
39 import java.util.List;
40 import java.util.Objects;
41 import java.util.stream.Collectors;
42 
43 /**
44  * Metrics testing utility
45  *
46  * <p>Example usage:
47  * <pre>{@code
48  *     try (EnterpriseMetricsRecorder r = EnterpriseMetricsRecorder.create() {
49  *         // Call code which generates metrics
50  *
51  *         assertThat(r.query().poll()).isNotNull();
52  *     }
53  *
54  * }</pre>
55  */
56 public class EnterpriseMetricsRecorder implements AutoCloseable {
57 
58     /** Create a {@link EnterpriseMetricsRecorder} and begin listening for metrics. */
create()59     public static EnterpriseMetricsRecorder create() {
60         EnterpriseMetricsRecorder r = new EnterpriseMetricsRecorder();
61         r.start(DEVICE_POLICY_EVENT_FIELD_NUMBER);
62 
63         return r;
64     }
65 
66     private static final long CONFIG_ID = "cts_config".hashCode();
67 
68     private final List<EnterpriseMetricInfo> mData = new ArrayList<>();
69 
EnterpriseMetricsRecorder()70     private EnterpriseMetricsRecorder() {
71 
72     }
73 
start(int atomTag)74     private void start(int atomTag) {
75         cleanLogs();
76         createAndUploadConfig(atomTag);
77     }
78 
79     /**
80      * Begin querying the recorded metrics.
81      */
query()82     public MetricQueryBuilder query() {
83         return new MetricQueryBuilder(this);
84     }
85 
fetchLatestData()86     List<EnterpriseMetricInfo> fetchLatestData() {
87         mData.addAll(getEventMetricDataList(getReportList()));
88         return mData;
89     }
90 
91     @Override
close()92     public void close() {
93         cleanLogs();
94     }
95 
createAndUploadConfig(int atomTag)96     private void createAndUploadConfig(int atomTag) {
97         StatsdConfig conf = new StatsdConfig();
98         conf.id = CONFIG_ID;
99         conf.allowedLogSource = new String[]{"AID_SYSTEM"};
100 
101         addAtomEvent(conf, atomTag);
102         uploadConfig(conf);
103     }
104 
addAtomEvent(StatsdConfig conf, int atomTag)105     private void addAtomEvent(StatsdConfig conf, int atomTag) {
106         addAtomEvent(conf, atomTag, new ArrayList<>());
107     }
108 
addAtomEvent(StatsdConfig conf, int atomTag, List<FieldValueMatcher> fvms)109     private void addAtomEvent(StatsdConfig conf, int atomTag,
110             List<FieldValueMatcher> fvms) {
111         String atomName = "Atom" + System.nanoTime();
112         String eventName = "Event" + System.nanoTime();
113 
114         SimpleAtomMatcher sam = new SimpleAtomMatcher();
115         sam.atomId = atomTag;
116         if (fvms != null) {
117             sam.fieldValueMatcher = fvms.toArray(new FieldValueMatcher[]{});
118         }
119 
120         AtomMatcher atomMatcher = new AtomMatcher();
121         atomMatcher.id = atomName.hashCode();
122         atomMatcher.setSimpleAtomMatcher(sam);
123 
124         conf.atomMatcher = new AtomMatcher[]{
125                 atomMatcher
126         };
127 
128         EventMetric eventMetric = new EventMetric();
129         eventMetric.id = eventName.hashCode();
130         eventMetric.what = atomName.hashCode();
131 
132         conf.eventMetric = new EventMetric[]{
133                 eventMetric
134         };
135     }
136 
uploadConfig(StatsdConfig config)137     private void uploadConfig(StatsdConfig config) {
138         byte[] bytes = new byte[config.getSerializedSize()];
139         CodedOutputByteBufferNano b = CodedOutputByteBufferNano.newInstance(bytes);
140         try {
141             ShellCommand.builder("cmd stats config update")
142                     .addOperand(CONFIG_ID)
143                     .writeToStdIn(MessageNano.toByteArray(config))
144                     .validate(String::isEmpty)
145                     .execute();
146         } catch (AdbException e) {
147             throw new NeneException("Error uploading config", e);
148         }
149     }
150 
cleanLogs()151     private void cleanLogs() {
152         removeConfig(CONFIG_ID);
153         getReportList(); // Clears data.
154     }
155 
removeConfig(long configId)156     private void removeConfig(long configId) {
157         try {
158             ShellCommand.builder("cmd stats config remove").addOperand(configId)
159                     .validate(String::isEmpty).execute();
160         } catch (AdbException e) {
161             throw new NeneException("Error removing config " + configId, e);
162         }
163     }
164 
getReportList()165     private ConfigMetricsReportList getReportList() {
166         try {
167             byte[] bytes = ShellCommand.builder("cmd stats dump-report")
168                     .addOperand(CONFIG_ID)
169                     .addOperand("--include_current_bucket")
170                     .addOperand("--proto")
171                     .forBytes()
172                     .execute();
173 
174             return ConfigMetricsReportList.parseFrom(bytes);
175         } catch (AdbException e) {
176             throw new NeneException("Error getting stat report list", e);
177         } catch (InvalidProtocolBufferNanoException e) {
178             throw new NeneException("Invalid proto", e);
179         }
180     }
181 
getEventMetricDataList( ConfigMetricsReportList reportList)182     private List<EnterpriseMetricInfo> getEventMetricDataList(
183             ConfigMetricsReportList reportList) {
184         return Arrays.stream(reportList.reports)
185                 .flatMap(s -> Arrays.stream(s.metrics.clone()))
186                 .filter(s -> s.getEventMetrics() != null && s.getEventMetrics().data != null)
187                 .flatMap(statsLogReport -> Arrays.stream(
188                         statsLogReport.getEventMetrics().data.clone()))
189                 .flatMap(eventMetricData -> Arrays.stream(
190                         backfillAggregatedAtomsinEventMetric(eventMetricData)))
191                 .sorted(Comparator.comparing(e -> e.elapsedTimestampNanos))
192                 .map(e -> e.atom)
193                 .filter((Objects::nonNull))
194                 .map(AtomsProto.Atom::getDevicePolicyEvent)
195                 .filter((Objects::nonNull))
196                 .map(EnterpriseMetricInfo::new)
197                 .collect(Collectors.toList());
198     }
199 
backfillAggregatedAtomsinEventMetric( StatsLog.EventMetricData metricData)200     private StatsLog.EventMetricData[] backfillAggregatedAtomsinEventMetric(
201             StatsLog.EventMetricData metricData) {
202         if (metricData.aggregatedAtomInfo == null) {
203             return new StatsLog.EventMetricData[]{metricData};
204         }
205         List<StatsLog.EventMetricData> data = new ArrayList<>();
206         StatsLog.AggregatedAtomInfo atomInfo = metricData.aggregatedAtomInfo;
207         for (long timestamp : atomInfo.elapsedTimestampNanos) {
208             StatsLog.EventMetricData newMetricData = new StatsLog.EventMetricData();
209             newMetricData.atom = atomInfo.atom;
210             newMetricData.elapsedTimestampNanos = timestamp;
211             data.add(newMetricData);
212         }
213         return data.toArray(new StatsLog.EventMetricData[0]);
214     }
215 }
216