Give workFolder a local group so it continue to be used

Ensure workfolder is respected in feature server.

Test: presubmit
Bug: 332119563
Change-Id: I7e00311a63258fad6ab902fc91643e050a516bc1

2 files changed, 18 insertions, 0 deletions

diff --git a/invocation_interfaces/com/android/tradefed/invoker/logger/CurrentInvocation.java b/invocation_interfaces/com/android/tradefed/invoker/logger/CurrentInvocation.java
index 22cc9e734..8408e9bdb 100644
--- a/invocation_interfaces/com/android/tradefed/invoker/logger/CurrentInvocation.java
+++ b/invocation_interfaces/com/android/tradefed/invoker/logger/CurrentInvocation.java
@@ -82,6 +82,18 @@ public class CurrentInvocation {
     private static final Map<ThreadGroup, Map<InvocationLocal<?>, Optional<?>>> mInvocationLocals =
             new ConcurrentHashMap<>();
 
+    private static ThreadLocal<ThreadGroup> sLocal = new ThreadLocal<>();
+
+    /** Tracks a localized context when using the properties inside the gRPC server */
+    public static void setLocalGroup(ThreadGroup tg) {
+        sLocal.set(tg);
+    }
+
+    /** Resets the localized context. */
+    public static void resetLocalGroup() {
+        sLocal.remove();
+    }
+
     /**
      * Add one key-value to be tracked at the invocation level.
      *
@@ -102,6 +114,9 @@ public class CurrentInvocation {
     public static File getInfo(InvocationInfo key) {
         ThreadGroup group = Thread.currentThread().getThreadGroup();
         synchronized (mPerGroupInfo) {
+            if (sLocal.get() != null) {
+                group = sLocal.get();
+            }
             if (mPerGroupInfo.get(group) == null) {
                 mPerGroupInfo.put(group, new InternalInvocationTracking());
             }
diff --git a/src/com/android/tradefed/service/TradefedFeatureServer.java b/src/com/android/tradefed/service/TradefedFeatureServer.java
index 67cff2169..e0e9a59dc 100644
--- a/src/com/android/tradefed/service/TradefedFeatureServer.java
+++ b/src/com/android/tradefed/service/TradefedFeatureServer.java
@@ -20,6 +20,7 @@ import com.android.tradefed.command.ICommandScheduler.IScheduledInvocationListen
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.IConfigurationReceiver;
 import com.android.tradefed.invoker.TestInformation;
+import com.android.tradefed.invoker.logger.CurrentInvocation;
 import com.android.tradefed.invoker.logger.InvocationMetricLogger;
 import com.android.tradefed.invoker.tracing.CloseableTraceScope;
 import com.android.tradefed.invoker.tracing.TracingLogger;
@@ -154,6 +155,7 @@ public class TradefedFeatureServer extends TradefedInformationImplBase {
         ServiceLoader<IRemoteFeature> serviceLoader = ServiceLoader.load(IRemoteFeature.class);
         for (IRemoteFeature feature : serviceLoader) {
             if (feature.getName().equals(request.getName())) {
+                CurrentInvocation.setLocalGroup(mRegisteredGroup.get(request.getReferenceId()));
                 InvocationMetricLogger.setLocalGroup(
                         mRegisteredGroup.get(request.getReferenceId()));
                 if (feature instanceof IConfigurationReceiver) {
@@ -199,6 +201,7 @@ public class TradefedFeatureServer extends TradefedInformationImplBase {
                         ((IConfigurationReceiver) feature).setConfiguration(null);
                     }
                     InvocationMetricLogger.resetLocalGroup();
+                    CurrentInvocation.resetLocalGroup();
                 }
             }
         }