Restore ThreadLocal context in coroutine

Author: kacgal

core/pom.xml

@@ -64,6 +64,12 @@
             <version>1.4.10</version>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.jetbrains.kotlinx</groupId>
+            <artifactId>kotlinx-coroutines-core</artifactId>
+            <version>1.3.9</version>
+            <scope>provided</scope>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

core/src/main/java/co/aikar/commands/RegisteredCommand.java

@@ -176,7 +176,8 @@ void invoke(CommandIssuer sender, List<String> args, CommandOperationContext con
                 Object[] methodArgsCopy = new Object[methodArgs.length + 1];
                 System.arraycopy(methodArgs, 0, methodArgsCopy, 0, methodArgs.length);
                 methodArgs = methodArgsCopy;
-                methodArgs[methodArgs.length - 1] = new JavaContinuation<Object>() {
+                ThreadLocalRestorer threadLocalRestorer = new ThreadLocalRestorer(CommandManager.commandOperationContext.get().peek());
+                methodArgs[methodArgs.length - 1] = new JavaContinuation<Object>(threadLocalRestorer) {
                     @Override
                     public void resumeWithException(@NotNull Throwable exception) {
                         handleException(sender, args, exception);

core/src/main/java/co/aikar/commands/ThreadLocalRestorer.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright (c) 2016-2020 Daniel Ennis (Aikar) - MIT License
+ *
+ *  Permission is hereby granted, free of charge, to any person obtaining
+ *  a copy of this software and associated documentation files (the
+ *  "Software"), to deal in the Software without restriction, including
+ *  without limitation the rights to use, copy, modify, merge, publish,
+ *  distribute, sublicense, and/or sell copies of the Software, and to
+ *  permit persons to whom the Software is furnished to do so, subject to
+ *  the following conditions:
+ *
+ *  The above copyright notice and this permission notice shall be
+ *  included in all copies or substantial portions of the Software.
+ *
+ *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ *  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ *  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ *  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+ *  LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+ *  OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+ *  WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+package co.aikar.commands;
+
+import kotlin.coroutines.CoroutineContext;
+import kotlin.jvm.functions.Function2;
+import kotlinx.coroutines.ThreadContextElement;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public final class ThreadLocalRestorer implements ThreadContextElement<Void> {
+
+    private final CommandOperationContext context;
+    private final Key key = new Key();
+
+    public ThreadLocalRestorer(CommandOperationContext context) {
+        this.context = context;
+    }
+
+    @Override
+    public void restoreThreadContext(@NotNull CoroutineContext coroutineContext, Void unused) {
+        CommandManager.commandOperationContext.get().pop();
+    }
+
+    @Override
+    public Void updateThreadContext(@NotNull CoroutineContext coroutineContext) {
+        CommandManager.commandOperationContext.get().push(this.context);
+        return null;
+    }
+
+    @NotNull
+    @Override
+    public CoroutineContext plus(@NotNull CoroutineContext coroutineContext) {
+        return DefaultImpls.plus(this, coroutineContext);
+    }
+
+    @NotNull
+    @Override
+    public Key getKey() {
+        return this.key;
+    }
+
+    @Override
+    public <R> R fold(R r, @NotNull Function2<? super R, ? super Element, ? extends R> function2) {
+        return DefaultImpls.fold(this, r, function2);
+    }
+
+    @Nullable
+    @Override
+    public <E extends Element> E get(@NotNull CoroutineContext.Key<E> key) {
+        return DefaultImpls.get(this, key);
+    }
+
+    @NotNull
+    @Override
+    public CoroutineContext minusKey(@NotNull CoroutineContext.Key key) {
+        return DefaultImpls.minusKey(this, key);
+    }
+
+    public static final class Key implements kotlin.coroutines.CoroutineContext.Key<ThreadLocalRestorer> {
+    }
+}

core/src/main/java/co/aikar/commands/kotlin/JavaContinuation.kt

@@ -27,14 +27,11 @@ import kotlin.coroutines.Continuation
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 
-abstract class JavaContinuation<T> : Continuation<T> {
-    override val context: CoroutineContext
-        get() = EmptyCoroutineContext
-
+abstract class JavaContinuation<T>(override val context: CoroutineContext = EmptyCoroutineContext) : Continuation<T> {
     override fun resumeWith(result: Result<T>) {
         result.fold(::resume, ::resumeWithException)
     }
 
     abstract fun resume(result: T)
     abstract fun resumeWithException(exception: Throwable)
-}
+}