Don't fail in `afterExecute` if context is already closed (#21563)

We run an assert on an potentially closed thread context. this should
not bubble up the `IllegalStateException`.

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

@@ -19,7 +19,6 @@
 
 package org.elasticsearch.common.util.concurrent;
 
-
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -112,8 +111,22 @@ public class EsThreadPoolExecutor extends ThreadPoolExecutor {
     @Override
     protected void afterExecute(Runnable r, Throwable t) {
         super.afterExecute(r, t);
-        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 + "]";
+        assert assertDefaultContext(r);
+    }
+
+    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;
+            }
+        }
+        return true;
     }
 
     /**

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

@@ -253,6 +253,13 @@ public final class ThreadContext implements Closeable, Writeable {
         return threadLocal.get() == DEFAULT_CONTEXT;
     }
 
+    /**
+     * 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