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 com.android.tradefed.cluster;
17 
18 import com.android.tradefed.util.UniqueMultiMap;
19 
20 import com.google.common.base.Strings;
21 import java.util.ArrayList;
22 import java.util.List;
23 import java.util.UUID;
24 import org.json.JSONArray;
25 import org.json.JSONException;
26 import org.json.JSONObject;
27 
28 /** A class that represents a task fetched from TF Cluster. */
29 public class ClusterCommand {
30 
31     public static enum RequestType {
32         /** An unmanaged request: the command line will run as is by the current TF process. */
33         UNMANAGED,
34         /** A managed request: the command line will run by a new TF process. */
35         MANAGED;
36     }
37 
38     /** Command's status in the TF cluster. */
39     public static enum State {
40         /** Initial state, or failed to determine state. */
41         UNKNOWN,
42         /** Inserted into the cluster's queue. */
43         QUEUED,
44         /** Currently being executed. */
45         RUNNING,
46         /** Canceled by user, or failed to allocate a device. */
47         CANCELED,
48         /** Completed successfully. */
49         COMPLETED,
50         /** Completed exceptionally. */
51         ERROR,
52         /** Non-retryable error, e.g. invalid configuration. */
53         FATAL
54     }
55 
56     private final String mTaskId;
57     private final String mRequestId;
58     private final String mCommandId;
59     private final String mCommandLine;
60     private final RequestType mRequestType;
61     private final Integer mShardCount;
62     private final Integer mShardIndex;
63     private final String mAttemptId;
64     // Devices try to match the current command.
65     private List<String> mTargetDeviceSerials = new ArrayList<>();
66     // Additional options to inject
67     private UniqueMultiMap<String, String> mExtraOptions = new UniqueMultiMap<>();
68 
ClusterCommand(String commandId, String taskId, String cmdLine)69     public ClusterCommand(String commandId, String taskId, String cmdLine) {
70         this(null, commandId, taskId, cmdLine, null, RequestType.UNMANAGED, null, null);
71     }
72 
73     /**
74      * Constructor.
75      *
76      * @param requestId A request ID
77      * @param commandId The ID of the command that issued this task
78      * @param taskId The ID of this task
79      * @param cmdLine The command line to run
80      * @param requestType A request type
81      * @param shardCount A shard count
82      * @param shardIndex A shard index
83      */
ClusterCommand( String requestId, String commandId, String taskId, String cmdLine, String attemptId, RequestType requestType, Integer shardCount, Integer shardIndex)84     public ClusterCommand(
85             String requestId,
86             String commandId,
87             String taskId,
88             String cmdLine,
89             String attemptId,
90             RequestType requestType,
91             Integer shardCount,
92             Integer shardIndex) {
93         mTaskId = taskId;
94         mRequestId = requestId;
95         mCommandId = commandId;
96         mCommandLine = cmdLine;
97         mRequestType = requestType;
98         mShardCount = shardCount;
99         mShardIndex = shardIndex;
100         if (!Strings.isNullOrEmpty(attemptId)) {
101             mAttemptId = attemptId;
102         } else {
103             // TODO(b/123294120): Remove creating attemptId on TF side once
104             // b/123294120 rolls out.
105             mAttemptId = UUID.randomUUID().toString();
106         }
107     }
108 
109     /**
110      * Returns the task ID.
111      *
112      * @return task ID.
113      */
getTaskId()114     public String getTaskId() {
115         return mTaskId;
116     }
117 
118     /**
119      * Returns the request ID.
120      *
121      * @return the request ID
122      */
getRequestId()123     public String getRequestId() {
124         return mRequestId;
125     }
126 
127     /**
128      * Returns the command ID.
129      *
130      * @return the command ID
131      */
getCommandId()132     public String getCommandId() {
133         return mCommandId;
134     }
135 
136     /**
137      * Returns the attempt ID. The attempt is randomly generated GUID used to distinguish multiple
138      * command runs.
139      *
140      * @return the attempt ID
141      */
getAttemptId()142     public String getAttemptId() {
143         return mAttemptId;
144     }
145 
146     /**
147      * Returns the command line string.
148      *
149      * @return the command line string.
150      */
getCommandLine()151     public String getCommandLine() {
152         return mCommandLine;
153     }
154 
155     /**
156      * Returns a request type
157      *
158      * @return a request type
159      */
getRequestType()160     public RequestType getRequestType() {
161         return mRequestType;
162     }
163 
164     /**
165      * Returns the list of target device serials on which this command will attempt to run.
166      *
167      * @return the list of target device serials
168      */
getTargetDeviceSerials()169     public List<String> getTargetDeviceSerials() {
170         return mTargetDeviceSerials;
171     }
172 
173     /**
174      * Sets the list of target device serials on which the command will try to run.
175      *
176      * @param targetDeviceSerials the list of device serials to set
177      */
setTargetDeviceSerials(List<String> targetDeviceSerials)178     public void setTargetDeviceSerials(List<String> targetDeviceSerials) {
179         this.mTargetDeviceSerials = targetDeviceSerials;
180     }
181 
182     /** @return multimap of additional options to inject */
getExtraOptions()183     public UniqueMultiMap<String, String> getExtraOptions() {
184         return mExtraOptions;
185     }
186 
187     /**
188      * Returns a shard count.
189      *
190      * @return a shard count.
191      */
getShardCount()192     public Integer getShardCount() {
193         return mShardCount;
194     }
195 
196     /**
197      * Returns a shard index.
198      *
199      * @return a shard index.
200      */
getShardIndex()201     public Integer getShardIndex() {
202         return mShardIndex;
203     }
204 
optInteger(JSONObject json, String name)205     private static Integer optInteger(JSONObject json, String name) throws JSONException {
206         if (json.isNull(name)) {
207             return null;
208         }
209         return json.getInt(name);
210     }
211 
fromJson(JSONObject json)212     public static ClusterCommand fromJson(JSONObject json) throws JSONException {
213         ClusterCommand command =
214                 new ClusterCommand(
215                         json.getString("request_id"),
216                         json.getString("command_id"),
217                         json.getString("task_id"),
218                         json.getString("command_line"),
219                         json.optString("attempt_id", null),
220                         RequestType.valueOf(
221                                 json.optString("request_type", RequestType.UNMANAGED.name())),
222                         optInteger(json, "shard_count"),
223                         optInteger(json, "shard_index"));
224         JSONArray jsonDeviceSerials = json.optJSONArray("device_serials");
225         if (jsonDeviceSerials != null) {
226             final List<String> deviceSerials = new ArrayList<>();
227             for (int j = 0; j < jsonDeviceSerials.length(); j++) {
228                 deviceSerials.add(jsonDeviceSerials.getString(j));
229             }
230             command.setTargetDeviceSerials(deviceSerials);
231         }
232         JSONArray extraOptions = json.optJSONArray("extra_options");
233         if (extraOptions != null) {
234             for (int i = 0; i < extraOptions.length(); i++) {
235                 JSONObject entry = extraOptions.getJSONObject(i);
236                 String key = entry.getString("key");
237                 JSONArray values = entry.getJSONArray("values");
238                 for (int j = 0; j < values.length(); j++) {
239                     command.mExtraOptions.put(key, values.getString(j));
240                 }
241             }
242         }
243         return command;
244     }
245 }
246