1 /*
2  * Copyright (C) 2023 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 com.android.tradefed.device.internal;
17 
18 import com.android.annotations.VisibleForTesting;
19 import com.android.tradefed.device.DeviceNotAvailableException;
20 import com.android.tradefed.device.ITestDevice;
21 import com.android.tradefed.device.NativeDevice;
22 import com.android.tradefed.device.StubDevice;
23 import com.android.tradefed.error.HarnessRuntimeException;
24 import com.android.tradefed.error.IHarnessException;
25 import com.android.tradefed.invoker.IInvocationContext;
26 import com.android.tradefed.invoker.logger.CurrentInvocation;
27 import com.android.tradefed.invoker.logger.CurrentInvocation.IsolationGrade;
28 import com.android.tradefed.invoker.logger.InvocationMetricLogger;
29 import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
30 import com.android.tradefed.log.LogUtil.CLog;
31 import com.android.tradefed.result.error.InfraErrorIdentifier;
32 import com.android.tradefed.service.TradefedFeatureClient;
33 import com.android.tradefed.util.SerializationUtil;
34 
35 import com.proto.tradefed.feature.FeatureResponse;
36 
37 import java.io.IOException;
38 import java.util.HashMap;
39 import java.util.Map;
40 import java.util.regex.Matcher;
41 import java.util.regex.Pattern;
42 
43 /**
44  * Utility handling Cuttlefish snapshot. This is meant to only be used internally to the test
45  * harness. This shouldn't be called during a test.
46  */
47 public class DeviceSnapshotHandler {
48 
49     private final TradefedFeatureClient mClient;
50     private final IInvocationContext mContext;
51 
DeviceSnapshotHandler()52     public DeviceSnapshotHandler() {
53         this(new TradefedFeatureClient(), CurrentInvocation.getInvocationContext());
54     }
55 
56     @VisibleForTesting
DeviceSnapshotHandler(TradefedFeatureClient client, IInvocationContext context)57     DeviceSnapshotHandler(TradefedFeatureClient client, IInvocationContext context) {
58         mClient = client;
59         mContext = context;
60     }
61 
62     /**
63      * Calls snapshot of the given device.
64      *
65      * @param device The device to snapshot.
66      * @param snapshotId Snapshot ID for the device to be saved to.
67      * @return True if snapshot was successful, false otherwise.
68      * @throws DeviceNotAvailableException
69      */
snapshotDevice(ITestDevice device, String snapshotId)70     public void snapshotDevice(ITestDevice device, String snapshotId)
71             throws DeviceNotAvailableException {
72         if (device.getIDevice() instanceof StubDevice) {
73             CLog.d("Device '%s' is a stub device. skipping snapshot.", device.getSerialNumber());
74             return;
75         }
76         FeatureResponse response;
77         try {
78             Map<String, String> args = new HashMap<>();
79             args.put(DeviceSnapshotFeature.DEVICE_NAME, mContext.getDeviceName(device));
80             args.put(DeviceSnapshotFeature.SNAPSHOT_ID, snapshotId);
81             response =
82                     mClient.triggerFeature(
83                             DeviceSnapshotFeature.DEVICE_SNAPSHOT_FEATURE_NAME, args);
84             CLog.d("Response from snapshot request: %s", response.getResponse());
85         } finally {
86             mClient.close();
87         }
88         if (response.hasErrorInfo()) {
89             String trace = response.getErrorInfo().getErrorTrace();
90             // Handle if it's an exception error.
91             Object o = null;
92             try {
93                 o = SerializationUtil.deserialize(trace);
94             } catch (IOException | RuntimeException e) {
95                 CLog.e("Failed to deserialize snapshot error response: %s", e.getMessage());
96             }
97             if (o instanceof DeviceNotAvailableException) {
98                 throw (DeviceNotAvailableException) o;
99             } else if (o instanceof IHarnessException) {
100                 IHarnessException exception = (IHarnessException) o;
101                 throw new HarnessRuntimeException(
102                         "Exception while snapshotting the device.", exception);
103             } else if (o instanceof Exception) {
104                 throw new HarnessRuntimeException(
105                         "Exception while snapshotting the device.",
106                         (Exception) o,
107                         InfraErrorIdentifier.UNDETERMINED);
108             }
109             throw new HarnessRuntimeException(
110                     "Exception while snapshotting the device. Unserialized error response: "
111                             + trace,
112                     InfraErrorIdentifier.UNDETERMINED);
113         }
114 
115         // TODO: parse snapshot ID from response, and save it to mContext.
116 
117         // Save snapshot performance data
118         Pattern durationPattern = Pattern.compile("Snapshot\\sfinished\\sin (\\d+)\\sms");
119         Matcher matcher;
120         matcher = durationPattern.matcher(response.getResponse());
121         if (matcher.find()) {
122             InvocationMetricLogger.addInvocationMetrics(
123                     InvocationMetricKey.DEVICE_SNAPSHOT_SUCCESS_COUNT, 1);
124             InvocationMetricLogger.addInvocationMetrics(
125                     InvocationMetricKey.DEVICE_SNAPSHOT_DURATIONS, matcher.group(1));
126         }
127     }
128 
129     /**
130      * Calls restore snapshot of the given device.
131      *
132      * @param device The device to restore.
133      * @param snapshotId Snapshot ID for the device to be restored to.
134      * @return True if restore was successful, false otherwise.
135      * @throws DeviceNotAvailableException
136      */
restoreSnapshotDevice(ITestDevice device, String snapshotId)137     public void restoreSnapshotDevice(ITestDevice device, String snapshotId)
138             throws DeviceNotAvailableException {
139         if (device.getIDevice() instanceof StubDevice) {
140             CLog.d(
141                     "Device '%s' is a stub device. skipping restoring snapshot.",
142                     device.getSerialNumber());
143             return;
144         }
145         FeatureResponse response;
146         try {
147             Map<String, String> args = new HashMap<>();
148             args.put(DeviceSnapshotFeature.SNAPSHOT_ID, snapshotId);
149             args.put(DeviceSnapshotFeature.RESTORE_FLAG, "true");
150             args.put(DeviceSnapshotFeature.DEVICE_NAME, mContext.getDeviceName(device));
151             response =
152                     mClient.triggerFeature(
153                             DeviceSnapshotFeature.DEVICE_SNAPSHOT_FEATURE_NAME, args);
154             CLog.d(
155                     "Response from restoring snapshot(%s) request: %s",
156                     snapshotId, response.getResponse());
157         } finally {
158             mClient.close();
159         }
160         if (response.hasErrorInfo()) {
161             String trace = response.getErrorInfo().getErrorTrace();
162             // Handle if it's an exception error.
163             Object o = null;
164             try {
165                 o = SerializationUtil.deserialize(trace);
166             } catch (IOException | RuntimeException e) {
167                 CLog.e("Failed to deserialize snapshot error response: %s", e.getMessage());
168             }
169             if (o instanceof DeviceNotAvailableException) {
170                 throw (DeviceNotAvailableException) o;
171             } else if (o instanceof IHarnessException) {
172                 IHarnessException exception = (IHarnessException) o;
173                 throw new HarnessRuntimeException(
174                         "Exception while restoring snapshot of the device.", exception);
175             } else if (o instanceof Exception) {
176                 throw new HarnessRuntimeException(
177                         "Exception while restoring snapshot of the device.",
178                         (Exception) o,
179                         InfraErrorIdentifier.UNDETERMINED);
180             }
181             throw new HarnessRuntimeException(
182                     "Exception while restoring snapshot of the device. Unserialized error response:"
183                             + " "
184                             + trace,
185                     InfraErrorIdentifier.UNDETERMINED);
186         }
187         if (device instanceof NativeDevice) {
188             ((NativeDevice) device).resetContentProviderSetup();
189         }
190         CurrentInvocation.setModuleIsolation(IsolationGrade.FULLY_ISOLATED);
191         CurrentInvocation.setRunIsolation(IsolationGrade.FULLY_ISOLATED);
192 
193         // Save snapshot performance data
194         Pattern durationPattern = Pattern.compile("Restoring snapshot\\sfinished\\sin (\\d+)\\sms");
195         Matcher matcher;
196         matcher = durationPattern.matcher(response.getResponse());
197         if (matcher.find()) {
198             InvocationMetricLogger.addInvocationMetrics(
199                     InvocationMetricKey.DEVICE_SNAPSHOT_RESTORE_SUCCESS_COUNT, 1);
200             InvocationMetricLogger.addInvocationMetrics(
201                     InvocationMetricKey.DEVICE_SNAPSHOT_RESTORE_DURATIONS, matcher.group(1));
202         }
203     }
204 }
205