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