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 package com.android.tradefed.service;
17 
18 import com.android.annotations.VisibleForTesting;
19 import com.android.tradefed.command.ICommandScheduler.IScheduledInvocationListener;
20 import com.android.tradefed.config.IConfiguration;
21 import com.android.tradefed.config.IConfigurationReceiver;
22 import com.android.tradefed.invoker.TestInformation;
23 import com.android.tradefed.invoker.logger.CurrentInvocation;
24 import com.android.tradefed.invoker.logger.InvocationMetricLogger;
25 import com.android.tradefed.invoker.tracing.CloseableTraceScope;
26 import com.android.tradefed.invoker.tracing.TracingLogger;
27 import com.android.tradefed.log.LogUtil.CLog;
28 import com.android.tradefed.service.internal.IRemoteScheduledListenersFeature;
29 import com.android.tradefed.testtype.ITestInformationReceiver;
30 import com.android.tradefed.util.StreamUtil;
31 import com.android.tradefed.util.SystemUtil;
32 
33 import com.proto.tradefed.feature.ErrorInfo;
34 import com.proto.tradefed.feature.FeatureRequest;
35 import com.proto.tradefed.feature.FeatureResponse;
36 import com.proto.tradefed.feature.TradefedInformationGrpc.TradefedInformationImplBase;
37 
38 import java.io.IOException;
39 import java.util.List;
40 import java.util.Map;
41 import java.util.ServiceLoader;
42 import java.util.UUID;
43 import java.util.concurrent.ConcurrentHashMap;
44 
45 import io.grpc.Server;
46 import io.grpc.ServerBuilder;
47 import io.grpc.stub.StreamObserver;
48 
49 /** A server that responds to requests for triggering features. */
50 public class TradefedFeatureServer extends TradefedInformationImplBase {
51 
52     public static final String SERVER_REFERENCE = "SERVER_REFERENCE";
53     public static final String TEST_INFORMATION_OBJECT = "TEST_INFORMATION";
54     public static final String TF_SERVICE_PORT = "TF_SERVICE_PORT";
55 
56     private static final int DEFAULT_PORT = 0;
57     private static Integer sInternalPort = null;
58 
59     private Server mServer;
60 
61     private Map<String, IConfiguration> mRegisteredInvocation = new ConcurrentHashMap<>();
62     private Map<String, ThreadGroup> mRegisteredGroup =
63             new ConcurrentHashMap<String, ThreadGroup>();
64     private Map<String, List<IScheduledInvocationListener>>
65             mRegisteredScheduledInvocationListeners = new ConcurrentHashMap<>();
66 
67     /** Returns the port used by the server. */
getPort()68     public static int getPort() {
69         if (sInternalPort != null) {
70             return sInternalPort;
71         }
72         return System.getenv(TF_SERVICE_PORT) != null
73                 ? Integer.parseInt(System.getenv(TF_SERVICE_PORT))
74                 : DEFAULT_PORT;
75     }
76 
TradefedFeatureServer()77     public TradefedFeatureServer() {
78         this(ServerBuilder.forPort(getPort()));
79     }
80 
81     @VisibleForTesting
TradefedFeatureServer(ServerBuilder<?> serverBuilder)82     TradefedFeatureServer(ServerBuilder<?> serverBuilder) {
83         mServer = serverBuilder.addService(this).build();
84     }
85 
86     /** Start the grpc server to listen to requests. */
start()87     public void start() {
88         try {
89             CLog.d("Starting feature server.");
90             mServer.start();
91             sInternalPort = mServer.getPort();
92         } catch (IOException e) {
93             if (SystemUtil.isLocalMode()) {
94                 CLog.w("TradefedFeatureServer already started: %s", e.getMessage());
95             } else {
96                 throw new RuntimeException(e);
97             }
98         }
99     }
100 
101     /** Stop the grpc server. */
shutdown()102     public void shutdown() throws InterruptedException {
103         if (mServer != null) {
104             CLog.d("Stopping feature server.");
105             mServer.shutdown();
106             mServer.awaitTermination();
107         }
108     }
109 
110     @Override
triggerFeature( FeatureRequest request, StreamObserver<FeatureResponse> responseObserver)111     public void triggerFeature(
112             FeatureRequest request, StreamObserver<FeatureResponse> responseObserver) {
113         FeatureResponse response;
114         try {
115             response = createResponse(request);
116         } catch (RuntimeException exception) {
117             response = FeatureResponse.newBuilder()
118                 .setErrorInfo(
119                     ErrorInfo.newBuilder()
120                             .setErrorTrace(StreamUtil.getStackTrace(exception)))
121                     .build();
122         }
123         responseObserver.onNext(response);
124 
125         responseObserver.onCompleted();
126     }
127 
128     /** Register an invocation with a unique reference that can be queried */
registerInvocation( IConfiguration config, ThreadGroup tg, List<IScheduledInvocationListener> listeners)129     public String registerInvocation(
130             IConfiguration config, ThreadGroup tg, List<IScheduledInvocationListener> listeners) {
131         String referenceId = UUID.randomUUID().toString();
132         mRegisteredInvocation.put(referenceId, config);
133         mRegisteredGroup.put(referenceId, tg);
134         mRegisteredScheduledInvocationListeners.put(referenceId, listeners);
135         config.getConfigurationDescription().addMetadata(SERVER_REFERENCE, referenceId);
136         return referenceId;
137     }
138 
139     /** Unregister an invocation by its configuration. */
unregisterInvocation(IConfiguration reference)140     public void unregisterInvocation(IConfiguration reference) {
141         String referenceId =
142                 reference
143                         .getConfigurationDescription()
144                         .getAllMetaData()
145                         .getUniqueMap()
146                         .get(SERVER_REFERENCE);
147         if (referenceId != null) {
148             mRegisteredInvocation.remove(referenceId);
149             mRegisteredGroup.remove(referenceId);
150             mRegisteredScheduledInvocationListeners.remove(referenceId);
151         }
152     }
153 
createResponse(FeatureRequest request)154     private FeatureResponse createResponse(FeatureRequest request) {
155         ServiceLoader<IRemoteFeature> serviceLoader = ServiceLoader.load(IRemoteFeature.class);
156         for (IRemoteFeature feature : serviceLoader) {
157             if (feature.getName().equals(request.getName())) {
158                 CurrentInvocation.setLocalGroup(mRegisteredGroup.get(request.getReferenceId()));
159                 InvocationMetricLogger.setLocalGroup(
160                         mRegisteredGroup.get(request.getReferenceId()));
161                 if (feature instanceof IConfigurationReceiver) {
162                     ((IConfigurationReceiver) feature)
163                             .setConfiguration(mRegisteredInvocation.get(request.getReferenceId()));
164                 }
165                 if (feature instanceof ITestInformationReceiver) {
166                     if (mRegisteredInvocation.get(request.getReferenceId()) != null) {
167                         ((ITestInformationReceiver) feature)
168                                 .setTestInformation(
169                                         (TestInformation) mRegisteredInvocation
170                                             .get(request.getReferenceId())
171                                             .getConfigurationObject(TEST_INFORMATION_OBJECT));
172                     }
173                 }
174                 if (feature instanceof IRemoteScheduledListenersFeature) {
175                     List<IScheduledInvocationListener> listeners =
176                             mRegisteredScheduledInvocationListeners.get(request.getReferenceId());
177                     if (listeners != null) {
178                         ((IRemoteScheduledListenersFeature) feature).setListeners(listeners);
179                     }
180                 }
181                 TracingLogger.setLocalGroup(mRegisteredGroup.get(request.getReferenceId()));
182                 try (CloseableTraceScope ignored =
183                         new CloseableTraceScope(
184                                 TracingLogger.getActiveTraceForGroup(
185                                         mRegisteredGroup.get(request.getReferenceId())),
186                                 feature.getName())) {
187                     FeatureResponse rep = feature.execute(request);
188                     if (rep == null) {
189                         return FeatureResponse.newBuilder()
190                                 .setErrorInfo(
191                                         ErrorInfo.newBuilder()
192                                                 .setErrorTrace(
193                                                         String.format(
194                                                                 "Feature '%s' returned null"
195                                                                         + " response.",
196                                                                 request.getName())))
197                                 .build();
198                     }
199                     return rep;
200                 } finally {
201                     if (feature instanceof IConfigurationReceiver) {
202                         ((IConfigurationReceiver) feature).setConfiguration(null);
203                     }
204                     TracingLogger.resetLocalGroup();
205                     InvocationMetricLogger.resetLocalGroup();
206                     CurrentInvocation.resetLocalGroup();
207                 }
208             }
209         }
210         return FeatureResponse.newBuilder()
211                 .setErrorInfo(
212                         ErrorInfo.newBuilder()
213                                 .setErrorTrace(
214                                         String.format(
215                                                 "No feature matching the requested one '%s'",
216                                                 request.getName())))
217                 .build();
218     }
219 }
220