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.tradefed.config.IConfiguration; 19 import com.android.tradefed.config.IConfigurationReceiver; 20 import com.android.tradefed.config.IDeviceConfiguration; 21 import com.android.tradefed.device.DeviceNotAvailableException; 22 import com.android.tradefed.device.ITestDevice; 23 import com.android.tradefed.device.cloud.GceAvdInfo; 24 import com.android.tradefed.device.cloud.RemoteAndroidVirtualDevice; 25 import com.android.tradefed.device.connection.AbstractConnection; 26 import com.android.tradefed.device.connection.AdbSshConnection; 27 import com.android.tradefed.invoker.TestInformation; 28 import com.android.tradefed.result.error.DeviceErrorIdentifier; 29 import com.android.tradefed.service.IRemoteFeature; 30 import com.android.tradefed.targetprep.TargetSetupError; 31 import com.android.tradefed.testtype.ITestInformationReceiver; 32 import com.android.tradefed.util.CommandResult; 33 import com.android.tradefed.util.CommandStatus; 34 import com.android.tradefed.util.SerializationUtil; 35 36 import com.proto.tradefed.feature.ErrorInfo; 37 import com.proto.tradefed.feature.FeatureRequest; 38 import com.proto.tradefed.feature.FeatureResponse; 39 40 import java.io.IOException; 41 42 /** Server side implementation of device snapshot. */ 43 public class DeviceSnapshotFeature 44 implements IRemoteFeature, IConfigurationReceiver, ITestInformationReceiver { 45 46 public static final String DEVICE_SNAPSHOT_FEATURE_NAME = "snapshotDevice"; 47 public static final String DEVICE_NAME = "device_name"; 48 public static final String SNAPSHOT_ID = "snapshot_id"; 49 public static final String RESTORE_FLAG = "restore_flag"; 50 public static final String DELETE_FLAG = "delete_flag"; 51 52 private IConfiguration mConfig; 53 private TestInformation mTestInformation; 54 55 @Override getName()56 public String getName() { 57 return DEVICE_SNAPSHOT_FEATURE_NAME; 58 } 59 60 @Override setConfiguration(IConfiguration configuration)61 public void setConfiguration(IConfiguration configuration) { 62 mConfig = configuration; 63 } 64 65 @Override setTestInformation(TestInformation testInformation)66 public void setTestInformation(TestInformation testInformation) { 67 mTestInformation = testInformation; 68 } 69 70 @Override getTestInformation()71 public TestInformation getTestInformation() { 72 return mTestInformation; 73 } 74 75 @Override execute(FeatureRequest request)76 public FeatureResponse execute(FeatureRequest request) { 77 FeatureResponse.Builder responseBuilder = FeatureResponse.newBuilder(); 78 String deviceName = request.getArgsMap().get(DEVICE_NAME); 79 if (deviceName == null) { 80 responseBuilder.setErrorInfo( 81 ErrorInfo.newBuilder().setErrorTrace("No device_name args specified.")); 82 return responseBuilder.build(); 83 } 84 85 IDeviceConfiguration configHolder = mConfig.getDeviceConfigByName(deviceName); 86 int index = 0; 87 for (IDeviceConfiguration deviceConfig : mConfig.getDeviceConfig()) { 88 if (deviceConfig == configHolder) { 89 break; 90 } 91 index++; 92 } 93 94 try { 95 mTestInformation.setActiveDeviceIndex(index); 96 AbstractConnection connection = mTestInformation.getDevice().getConnection(); 97 // TODO: Support NestedRemoteDevice 98 if ((mTestInformation.getDevice() instanceof RemoteAndroidVirtualDevice) 99 || (connection instanceof AdbSshConnection)) { 100 GceAvdInfo info = getAvdInfo(mTestInformation.getDevice(), connection); 101 if (info == null) { 102 throw new RuntimeException("GceAvdInfo was null. skipping"); 103 } 104 Integer offset = info.getDeviceOffset(); 105 String user = info.getInstanceUser(); 106 107 String snapshotId = request.getArgsMap().get(SNAPSHOT_ID); 108 boolean deleteFlag = Boolean.parseBoolean(request.getArgsMap().get(DELETE_FLAG)); 109 boolean restoreFlag = Boolean.parseBoolean(request.getArgsMap().get(RESTORE_FLAG)); 110 if (deleteFlag) { 111 deleteSnapshot(responseBuilder, connection, user, snapshotId); 112 } else if (restoreFlag) { 113 restoreSnapshot(responseBuilder, connection, user, offset, snapshotId); 114 } else { 115 snapshot(responseBuilder, connection, user, offset, snapshotId); 116 } 117 } else { 118 String error = 119 String.format( 120 "Device type %s with connection type %s doesn't support" 121 + " snapshotting", 122 mTestInformation.getDevice().getClass().getSimpleName(), 123 connection != null 124 ? connection.getClass().getSimpleName() 125 : "[null]"); 126 responseBuilder.setErrorInfo(ErrorInfo.newBuilder().setErrorTrace(error)); 127 return responseBuilder.build(); 128 } 129 } catch (DeviceNotAvailableException | TargetSetupError e) { 130 String error = "Failed to snapshot device."; 131 try { 132 error = SerializationUtil.serializeToString(e); 133 } catch (RuntimeException | IOException serializationError) { 134 // Ignore 135 } 136 responseBuilder.setErrorInfo(ErrorInfo.newBuilder().setErrorTrace(error)); 137 } finally { 138 mTestInformation.setActiveDeviceIndex(0); 139 } 140 return responseBuilder.build(); 141 } 142 snapshot( FeatureResponse.Builder responseBuilder, AbstractConnection connection, String user, Integer offset, String snapshotId)143 private void snapshot( 144 FeatureResponse.Builder responseBuilder, 145 AbstractConnection connection, 146 String user, 147 Integer offset, 148 String snapshotId) 149 throws DeviceNotAvailableException, TargetSetupError { 150 String response = 151 String.format( 152 "Attempting snapshot device on %s (%s).", 153 mTestInformation.getDevice().getSerialNumber(), 154 mTestInformation.getDevice().getClass().getSimpleName()); 155 try { 156 long startTime = System.currentTimeMillis(); 157 CommandResult result = snapshotGce(connection, user, offset, snapshotId); 158 if (!CommandStatus.SUCCESS.equals(result.getStatus())) { 159 throw new DeviceNotAvailableException( 160 String.format( 161 "Failed to snapshot device: %s. status:%s\n" 162 + "stdout: %s\n" 163 + "stderr:%s", 164 mTestInformation.getDevice().getSerialNumber(), 165 result.getStatus(), 166 result.getStdout(), 167 result.getStderr()), 168 mTestInformation.getDevice().getSerialNumber(), 169 DeviceErrorIdentifier.DEVICE_FAILED_TO_SNAPSHOT); 170 } 171 response += 172 String.format( 173 " Snapshot finished in %d ms.", System.currentTimeMillis() - startTime); 174 } finally { 175 responseBuilder.setResponse(response); 176 } 177 } 178 restoreSnapshot( FeatureResponse.Builder responseBuilder, AbstractConnection connection, String user, Integer offset, String snapshotId)179 private void restoreSnapshot( 180 FeatureResponse.Builder responseBuilder, 181 AbstractConnection connection, 182 String user, 183 Integer offset, 184 String snapshotId) 185 throws DeviceNotAvailableException, TargetSetupError { 186 String response = 187 String.format( 188 "Attempting restore device snapshot on %s (%s) to %s.", 189 mTestInformation.getDevice().getSerialNumber(), 190 mTestInformation.getDevice().getClass().getSimpleName(), 191 snapshotId); 192 try { 193 long startTime = System.currentTimeMillis(); 194 CommandResult result = restoreSnapshotGce(connection, user, offset, snapshotId); 195 if (!CommandStatus.SUCCESS.equals(result.getStatus())) { 196 throw new DeviceNotAvailableException( 197 String.format( 198 "Failed to restore snapshot on device: %s. status:%s\n" 199 + "stdout: %s\n" 200 + "stderr:%s", 201 mTestInformation.getDevice().getSerialNumber(), 202 result.getStatus(), 203 result.getStdout(), 204 result.getStderr()), 205 mTestInformation.getDevice().getSerialNumber(), 206 DeviceErrorIdentifier.DEVICE_FAILED_TO_RESTORE_SNAPSHOT); 207 } 208 response += 209 String.format( 210 " Restoring snapshot finished in %d ms.", 211 System.currentTimeMillis() - startTime); 212 } finally { 213 responseBuilder.setResponse(response); 214 } 215 } 216 deleteSnapshot( FeatureResponse.Builder responseBuilder, AbstractConnection connection, String user, String snapshotId)217 private void deleteSnapshot( 218 FeatureResponse.Builder responseBuilder, 219 AbstractConnection connection, 220 String user, 221 String snapshotId) 222 throws DeviceNotAvailableException, TargetSetupError { 223 String response = 224 String.format( 225 "Attempting delete device snapshot on %s (%s) to %s.", 226 mTestInformation.getDevice().getSerialNumber(), 227 mTestInformation.getDevice().getClass().getSimpleName(), 228 snapshotId); 229 try { 230 long startTime = System.currentTimeMillis(); 231 CommandResult result = deleteSnapshotGce(connection, user, snapshotId); 232 if (!CommandStatus.SUCCESS.equals(result.getStatus())) { 233 throw new DeviceNotAvailableException( 234 String.format( 235 "Failed to delete snapshot on device: %s. status:%s\n" 236 + "stdout: %s\n" 237 + "stderr:%s", 238 mTestInformation.getDevice().getSerialNumber(), 239 result.getStatus(), 240 result.getStdout(), 241 result.getStderr()), 242 mTestInformation.getDevice().getSerialNumber(), 243 DeviceErrorIdentifier.DEVICE_FAILED_TO_DELETE_SNAPSHOT); 244 } 245 response += 246 String.format( 247 " Deleting snapshot finished in %d ms.", 248 System.currentTimeMillis() - startTime); 249 } finally { 250 responseBuilder.setResponse(response); 251 } 252 } 253 getAvdInfo(ITestDevice device, AbstractConnection connection)254 private GceAvdInfo getAvdInfo(ITestDevice device, AbstractConnection connection) { 255 if (connection instanceof AdbSshConnection) { 256 return ((AdbSshConnection) connection).getAvdInfo(); 257 } 258 if (device instanceof RemoteAndroidVirtualDevice) { 259 return ((RemoteAndroidVirtualDevice) device).getAvdInfo(); 260 } 261 return null; 262 } 263 snapshotGce( AbstractConnection connection, String user, Integer offset, String snapshotId)264 private CommandResult snapshotGce( 265 AbstractConnection connection, String user, Integer offset, String snapshotId) 266 throws TargetSetupError { 267 if (connection instanceof AdbSshConnection) { 268 return ((AdbSshConnection) connection).snapshotGce(user, offset, snapshotId); 269 } 270 CommandResult res = new CommandResult(CommandStatus.EXCEPTION); 271 res.setStderr("Incorrect connection type while attempting device snapshot"); 272 return res; 273 } 274 restoreSnapshotGce( AbstractConnection connection, String user, Integer offset, String snapshotId)275 private CommandResult restoreSnapshotGce( 276 AbstractConnection connection, String user, Integer offset, String snapshotId) 277 throws TargetSetupError { 278 if (connection instanceof AdbSshConnection) { 279 return ((AdbSshConnection) connection).restoreSnapshotGce(user, offset, snapshotId); 280 } 281 CommandResult res = new CommandResult(CommandStatus.EXCEPTION); 282 res.setStderr("Incorrect connection type while attempting device restore"); 283 return res; 284 } 285 deleteSnapshotGce( AbstractConnection connection, String user, String snapshotId)286 private CommandResult deleteSnapshotGce( 287 AbstractConnection connection, String user, String snapshotId) throws TargetSetupError { 288 if (connection instanceof AdbSshConnection) { 289 return ((AdbSshConnection) connection).deleteSnapshotGce(user, snapshotId); 290 } 291 CommandResult res = new CommandResult(CommandStatus.EXCEPTION); 292 res.setStderr("Incorrect connection type while attempting device delete"); 293 return res; 294 } 295 } 296