RDCT-1013 Propagate client id to ProcessIOExecutorService tasks

fix wrapping order

GitOrigin-RevId: 42f20125fbf85a5fd4529eca68879bab7b2b1d98

platform/util/src/com/intellij/codeWithMe/ClientIdPropagation.kt

@@ -7,6 +7,8 @@ package com.intellij.codeWithMe
 import com.intellij.util.Processor
 import org.jetbrains.annotations.ApiStatus
 import java.util.concurrent.Callable
+import java.util.concurrent.FutureTask
+import java.util.concurrent.TimeUnit
 import java.util.function.BiConsumer
 import java.util.function.Function
 
@@ -91,3 +93,35 @@ fun <T, U> decorateBiConsumer(biConsumer: BiConsumer<T, U>): BiConsumer<T, U> {
     }
   }
 }
+
+fun <T> decorateFutureTask(futureTask: FutureTask<T>): FutureTask<T> {
+  if (!propagateClientIdAcrossThreads) return futureTask
+  val currentId = currentClientIdString
+  return object : FutureTask<T>({ null }) {
+    override fun run() {
+      withClientId(currentId) {
+        futureTask.run()
+      }
+    }
+
+    override fun cancel(mayInterruptIfRunning: Boolean): Boolean {
+      return futureTask.cancel(mayInterruptIfRunning)
+    }
+
+    override fun isCancelled(): Boolean {
+      return futureTask.isCancelled
+    }
+
+    override fun isDone(): Boolean {
+      return futureTask.isDone
+    }
+
+    override fun get(): T {
+      return futureTask.get()
+    }
+
+    override fun get(timeout: Long, unit: TimeUnit): T {
+      return futureTask.get(timeout, unit)
+    }
+  }
+}

platform/util/src/com/intellij/util/concurrency/AppScheduledExecutorService.java

@@ -13,7 +13,7 @@ import java.util.List;
 import java.util.concurrent.*;
 import java.util.function.BiConsumer;
 
-import static com.intellij.codeWithMe.ClientIdPropagation.decorateCallable;
+import static com.intellij.codeWithMe.ClientIdPropagation.decorateFutureTask;
 import static com.intellij.codeWithMe.ClientIdPropagation.decorateRunnable;
 import static com.intellij.util.concurrency.AppExecutorUtil.propagateContextOrCancellation;
 
@@ -234,13 +234,13 @@ public final class AppScheduledExecutorService extends SchedulingWrapper {
     if (!propagateContextOrCancellation()) {
       return command;
     }
-    return Propagation.capturePropagationAndCancellationContext(decorateRunnable(command));
+    return decorateRunnable(Propagation.capturePropagationAndCancellationContext(command));
   }
 
   public static <T> @NotNull FutureTask<T> capturePropagationAndCancellationContext(@NotNull Callable<T> callable) {
     if (!propagateContextOrCancellation()) {
       return new FutureTask<>(callable);
     }
-    return Propagation.capturePropagationAndCancellationContext(decorateCallable(callable));
+    return decorateFutureTask(Propagation.capturePropagationAndCancellationContext(callable));
   }
 }