ThreadPool and ThreadContext are not closeable

This commit changes the ThreadContext to just use a regular ThreadLocal
over the lucene CloseableThreadLocal. The CloseableThreadLocal solves
issues with ThreadLocals that are no longer needed during runtime but
in the case of the ThreadContext, we need it for the runtime of the
node and it is typically not closed until the node closes, so we miss
out on the benefits that this class provides.

Additionally by removing the close logic, we simplify code in other
places that deal with exceptions and tracking to see if it happens when
the node is closing.

Closes elastic#42577

Author: jaymode

server/src/main/java/org/elasticsearch/common/util/concurrent/EsThreadPoolExecutor.java

@@ -106,17 +106,8 @@ protected void afterExecute(Runnable r, Throwable t) {
     }
 
     private boolean assertDefaultContext(Runnable r) {
-        try {
-            assert contextHolder.isDefaultContext() : "the thread context is not the default context and the thread [" +
-                Thread.currentThread().getName() + "] is being returned to the pool after executing [" + r + "]";
-        } catch (IllegalStateException ex) {
-            // sometimes we execute on a closed context and isDefaultContext doen't bypass the ensureOpen checks
-            // this must not trigger an exception here since we only assert if the default is restored and
-            // we don't really care if we are closed
-            if (contextHolder.isClosed() == false) {
-                throw ex;
-            }
-        }
+        assert contextHolder.isDefaultContext() : "the thread context is not the default context and the thread [" +
+            Thread.currentThread().getName() + "] is being returned to the pool after executing [" + r + "]";
         return true;
     }
 

server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java

@@ -20,7 +20,6 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.apache.lucene.util.CloseableThreadLocal;
 import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.client.OriginSettingClient;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -31,7 +30,6 @@
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.http.HttpTransportSettings;
 
-import java.io.Closeable;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
@@ -41,7 +39,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
 import java.util.function.BinaryOperator;
 import java.util.function.Function;
@@ -81,7 +78,7 @@
  * </pre>
  *
  */
-public final class ThreadContext implements Closeable, Writeable {
+public final class ThreadContext implements Writeable {
 
     public static final String PREFIX = "request.headers";
     public static final Setting<Settings> DEFAULT_HEADERS_SETTING = Setting.groupSetting(PREFIX + ".", Property.NodeScope);
@@ -94,7 +91,7 @@ public final class ThreadContext implements Closeable, Writeable {
     private static final Logger logger = LogManager.getLogger(ThreadContext.class);
     private static final ThreadContextStruct DEFAULT_CONTEXT = new ThreadContextStruct();
     private final Map<String, String> defaultHeader;
-    private final ContextThreadLocal threadLocal;
+    private final ThreadLocal<ThreadContextStruct> threadLocal;
     private final int maxWarningHeaderCount;
     private final long maxWarningHeaderSize;
 
@@ -113,34 +110,23 @@ public ThreadContext(Settings settings) {
             }
             this.defaultHeader = Collections.unmodifiableMap(defaultHeader);
         }
-        threadLocal = new ContextThreadLocal();
+        threadLocal = ThreadLocal.withInitial(() -> DEFAULT_CONTEXT);
         this.maxWarningHeaderCount = SETTING_HTTP_MAX_WARNING_HEADER_COUNT.get(settings);
         this.maxWarningHeaderSize = SETTING_HTTP_MAX_WARNING_HEADER_SIZE.get(settings).getBytes();
     }
 
-    @Override
-    public void close() {
-        threadLocal.close();
-    }
-
     /**
      * Removes the current context and resets a default context. The removed context can be
      * restored by closing the returned {@link StoredContext}.
      */
     public StoredContext stashContext() {
         final ThreadContextStruct context = threadLocal.get();
-        threadLocal.set(null);
+        threadLocal.set(DEFAULT_CONTEXT);
         return () -> {
             // If the node and thus the threadLocal get closed while this task
             // is still executing, we don't want this runnable to fail with an
             // uncaught exception
-            try {
-                threadLocal.set(context);
-            } catch (IllegalStateException e) {
-                if (isClosed() == false) {
-                    throw e;
-                }
-            }
+            threadLocal.set(context);
         };
     }
 
@@ -399,13 +385,6 @@ public boolean isSystemContext() {
         return threadLocal.get().isSystemContext;
     }
 
-    /**
-     * Returns <code>true</code> if the context is closed, otherwise <code>true</code>
-     */
-    boolean isClosed() {
-        return threadLocal.closed.get();
-    }
-
     @FunctionalInterface
     public interface StoredContext extends AutoCloseable {
         @Override
@@ -617,55 +596,6 @@ private void writeTo(StreamOutput out, Map<String, String> defaultHeaders) throw
         }
     }
 
-    private static class ContextThreadLocal extends CloseableThreadLocal<ThreadContextStruct> {
-        private final AtomicBoolean closed = new AtomicBoolean(false);
-
-        @Override
-        public void set(ThreadContextStruct object) {
-            try {
-                if (object == DEFAULT_CONTEXT) {
-                    super.set(null);
-                } else {
-                    super.set(object);
-                }
-            } catch (NullPointerException ex) {
-                /* This is odd but CloseableThreadLocal throws a NPE if it was closed but still accessed.
-                   to get a real exception we call ensureOpen() to tell the user we are already closed.*/
-                ensureOpen();
-                throw ex;
-            }
-        }
-
-        @Override
-        public ThreadContextStruct get() {
-            try {
-                ThreadContextStruct threadContextStruct = super.get();
-                if (threadContextStruct != null) {
-                    return threadContextStruct;
-                }
-                return DEFAULT_CONTEXT;
-            } catch (NullPointerException ex) {
-                /* This is odd but CloseableThreadLocal throws a NPE if it was closed but still accessed.
-                   to get a real exception we call ensureOpen() to tell the user we are already closed.*/
-                ensureOpen();
-                throw ex;
-            }
-        }
-
-        private void ensureOpen() {
-            if (closed.get()) {
-                throw new IllegalStateException("threadcontext is already closed");
-            }
-        }
-
-        @Override
-        public void close() {
-            if (closed.compareAndSet(false, true)) {
-                super.close();
-            }
-        }
-    }
-
     /**
      * Wraps a Runnable to preserve the thread context.
      */
@@ -680,19 +610,9 @@ private ContextPreservingRunnable(Runnable in) {
 
         @Override
         public void run() {
-            boolean whileRunning = false;
             try (ThreadContext.StoredContext ignore = stashContext()){
                 ctx.restore();
-                whileRunning = true;
                 in.run();
-                whileRunning = false;
-            } catch (IllegalStateException ex) {
-                if (whileRunning || threadLocal.closed.get() == false) {
-                    throw ex;
-                }
-                // if we hit an ISE here we have been shutting down
-                // this comes from the threadcontext and barfs if
-                // our threadpool has been shutting down
             }
         }
 
@@ -749,21 +669,9 @@ public void onRejection(Exception e) {
 
         @Override
         protected void doRun() throws Exception {
-            boolean whileRunning = false;
             threadsOriginalContext = stashContext();
-            try {
-                creatorsContext.restore();
-                whileRunning = true;
-                in.doRun();
-                whileRunning = false;
-            } catch (IllegalStateException ex) {
-                if (whileRunning || threadLocal.closed.get() == false) {
-                    throw ex;
-                }
-                // if we hit an ISE here we have been shutting down
-                // this comes from the threadcontext and barfs if
-                // our threadpool has been shutting down
-            }
+            creatorsContext.restore();
+            in.doRun();
         }
 
         @Override

server/src/main/java/org/elasticsearch/threadpool/ThreadPool.java

@@ -39,7 +39,6 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.node.Node;
 
-import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -62,7 +61,7 @@
 import static java.util.Collections.unmodifiableMap;
 import static java.util.Map.entry;
 
-public class ThreadPool implements Scheduler, Closeable {
+public class ThreadPool implements Scheduler {
 
     private static final Logger logger = LogManager.getLogger(ThreadPool.class);
 
@@ -704,15 +703,13 @@ private static boolean awaitTermination(
     public static boolean terminate(ThreadPool pool, long timeout, TimeUnit timeUnit) {
         if (pool != null) {
             // Leverage try-with-resources to close the threadpool
-            try (ThreadPool c = pool) {
-                pool.shutdown();
-                if (awaitTermination(pool, timeout, timeUnit)) {
-                    return true;
-                }
-                // last resort
-                pool.shutdownNow();
-                return awaitTermination(pool, timeout, timeUnit);
+            pool.shutdown();
+            if (awaitTermination(pool, timeout, timeUnit)) {
+                return true;
             }
+            // last resort
+            pool.shutdownNow();
+            return awaitTermination(pool, timeout, timeUnit);
         }
         return false;
     }
@@ -731,11 +728,6 @@ private static boolean awaitTermination(
         return false;
     }
 
-    @Override
-    public void close() {
-        threadContext.close();
-    }
-
     public ThreadContext getThreadContext() {
         return threadContext;
     }

server/src/main/java/org/elasticsearch/transport/TransportLogger.java

@@ -97,9 +97,8 @@ private String format(TcpChannel channel, BytesReference message, String event)
                         streamInput = compressor.streamInput(streamInput);
                     }
 
-                    try (ThreadContext context = new ThreadContext(Settings.EMPTY)) {
-                        context.readHeaders(streamInput);
-                    }
+                    ThreadContext context = new ThreadContext(Settings.EMPTY);
+                    context.readHeaders(streamInput);
                     // now we decode the features
                     streamInput.readStringArray();
                     sb.append(", action: ").append(streamInput.readString());