Expose immediately-executed tasks in the queries virtual table

SEPExecutor.maybeExecuteImmediately() runs a task synchronously on the
calling worker thread, nested within the task the worker is already
running. Such immediate tasks were invisible in system_views.queries, which only exposed each worker's primary running task.
This is common on the coordinator path, where a local read or mutation is executed immediately within the enclosing QUERY task.

Each SEPWorker now also tracks an immediate current task, set around
maybeExecuteImmediately(), and exposes it as an additional
DebuggableTaskRunner, so the queries table reports both the enclosing
task and the immediate one as separate rows.

patch by Dmitry Konstantinov; reviewed by Caleb Rackliffe for CASSANDRA-21471

CHANGES.txt

@@ -1,4 +1,5 @@
 6.0-alpha2
+ * Expose immediately-executed tasks in the queries virtual table (CASSANDRA-21471)
  * Avoid potential deadlock between GlobalLogFollower and GossipStage (CASSANDRA-21384)
  * Add CMS membership as a field in ClusterMetadata (CASSANDRA-20736)
  * Fix maven remote publishing of Accord artifacts (CASSANDRA-21261)

src/java/org/apache/cassandra/concurrent/CassandraThread.java

@@ -27,19 +27,34 @@ public class CassandraThread extends FastThreadLocalThread
     private ThreadLocalMetrics threadLocalMetrics;
     private ExecutorLocals executorLocals;
 
-    public CassandraThread(ThreadGroup group, Runnable target, String name)
+    private final ImmediateTaskHolder immediateTaskHolder;
+
+    public CassandraThread(ThreadGroup group, Runnable target, String name, ImmediateTaskHolder immediateTaskHolder)
     {
         super(group, target, name);
+        assert immediateTaskHolder != null;
+        this.immediateTaskHolder = immediateTaskHolder;
+    }
+    public CassandraThread(ThreadGroup group, Runnable target, String name)
+    {
+        this(group, target, name, ImmediateTaskHolder.NO_OP);
     }
 
     public CassandraThread()
     {
         super();
+        this.immediateTaskHolder = ImmediateTaskHolder.NO_OP;
     }
 
     public CassandraThread(Runnable target)
     {
         super(target);
+        this.immediateTaskHolder = ImmediateTaskHolder.NO_OP;
+    }
+
+    public ImmediateTaskHolder getImmediateTaskHolder()
+    {
+        return immediateTaskHolder;
     }
 
     public ThreadLocalMetrics getThreadLocalMetrics()

src/java/org/apache/cassandra/concurrent/ImmediateTaskHolder.java

@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.cassandra.concurrent;
+
+public interface ImmediateTaskHolder
+{
+    /**
+     * @param currentTask to set as a current immediate task
+     * @return the previously set immediate task, can be null
+     */
+    Runnable setImmediateTask(Runnable currentTask);
+
+    class NoOp implements ImmediateTaskHolder
+    {
+
+        @Override
+        public Runnable setImmediateTask(Runnable currentTask)
+        {
+            // nothing to do
+            return null;
+        }
+    }
+
+    NoOp NO_OP = new NoOp();
+}

src/java/org/apache/cassandra/concurrent/SEPExecutor.java

@@ -211,12 +211,15 @@ public class SEPExecutor implements LocalAwareExecutorPlus, SEPExecutorMBean
         }
         else
         {
+            ImmediateTaskHolder taskHolder = getNestedCurrentTaskHolder();
+            Runnable previousTask = taskHolder.setImmediateTask(task);
             try
             {
                 task.run();
             }
             finally
             {
+                taskHolder.setImmediateTask(previousTask);
                 returnWorkPermit();
                 // we have to maintain our invariant of always scheduling after any work is performed
                 // in this case in particular we are not processing the rest of the queue anyway, and so
@@ -226,6 +229,16 @@ public class SEPExecutor implements LocalAwareExecutorPlus, SEPExecutorMBean
         }
     }
 
+    private static ImmediateTaskHolder getNestedCurrentTaskHolder()
+    {
+        Thread currentThread = Thread.currentThread();
+        if (currentThread instanceof CassandraThread)
+        {
+            return ((CassandraThread) currentThread).getImmediateTaskHolder();
+        }
+        return ImmediateTaskHolder.NO_OP;
+    }
+
     @Override
     public void execute(Runnable run)
     {

src/java/org/apache/cassandra/concurrent/SEPWorker.java

@@ -51,6 +51,35 @@ final class SEPWorker extends AtomicReference<SEPWorker.Work> implements Runnabl
 
     private final AtomicReference<Runnable> currentTask = new AtomicReference<>();
 
+    private class ImmediateDebuggableTaskRunner implements DebuggableTask.DebuggableTaskRunner, ImmediateTaskHolder
+    {
+        private final AtomicReference<Runnable> immediateCurrentTask = new AtomicReference<>();
+
+        @Override
+        public DebuggableTask running()
+        {
+            return getDebuggableTask(immediateCurrentTask.get());
+        }
+
+        @Override
+        public String id()
+        {
+            // derive from the current thread name so the nested row tracks renames and correlates with the worker's main row
+            return thread.getName() + "(immediate)";
+        }
+
+        @Override
+        public Runnable setImmediateTask(Runnable currentTask)
+        {
+            // plain is used to reduce overheads, the method is expected to be invoked only by a single thread
+            Runnable previousTask = immediateCurrentTask.getPlain();
+            immediateCurrentTask.lazySet(currentTask);
+            return previousTask;
+        }
+    }
+
+    private final ImmediateDebuggableTaskRunner immediateDebuggableTaskRunner;
+
     private String lastUsedExecutorName;
 
     SEPWorker(ThreadGroup threadGroup, Long workerId, Work initialState, SharedExecutorPool pool)
@@ -58,18 +87,27 @@ final class SEPWorker extends AtomicReference<SEPWorker.Work> implements Runnabl
         this.pool = pool;
         this.workerId = workerId;
         this.workerIdThreadSuffix = '-' + workerId.toString();
-        thread = new CassandraThread(threadGroup, this, threadGroup.getName() + "-Worker-" + workerId);
+        String threadName = threadGroup.getName() + "-Worker-" + workerId;
+        this.immediateDebuggableTaskRunner = new ImmediateDebuggableTaskRunner();
+        thread = new CassandraThread(threadGroup, this, threadName, immediateDebuggableTaskRunner);
         thread.setDaemon(true);
         set(initialState);
         thread.start();
     }
 
+    public DebuggableTask.DebuggableTaskRunner immediateRunner()
+    {
+        return immediateDebuggableTaskRunner;
+    }
+
     @Override
     public DebuggableTask running()
     {
-        // can change after null check so go off local reference
+        return getDebuggableTask(currentTask.get());
-        Runnable task = currentTask.get();
+    }
 
+    private static DebuggableTask getDebuggableTask(Runnable task)
+    {
         // Local read and mutation Runnables are themselves debuggable
         if (task instanceof DebuggableTask)
             return (DebuggableTask) task;

src/java/org/apache/cassandra/concurrent/SharedExecutorPool.java

@@ -123,7 +123,7 @@ public class SharedExecutorPool
 
     public Stream<? extends DebuggableTaskRunner> workers()
     {
-        return allWorkers.stream();
+        return allWorkers.stream().flatMap(worker -> Stream.of(worker, worker.immediateRunner()));
     }
 
     void maybeStartSpinningWorker()

test/distributed/org/apache/cassandra/distributed/test/QueriesTableTest.java

@@ -121,13 +121,14 @@ public class QueriesTableTest extends TestBaseImpl
             String task = row.get("task").toString();
 
             boolean localReaderThread = threadId.contains("Read") || threadId.contains("SharedPool-Worker");
-            readVisible |= localReaderThread && task.contains("SELECT");
+            // QUERY in a task name is used as a way to differentiate a local replica read vs a coordinator read
+            readVisible |= localReaderThread && task.contains("SELECT") && !task.contains("QUERY");
             boolean coordReaderThread = threadId.contains("Native-Transport-Requests") || threadId.contains("SharedPool-Worker");
-            coordinatorReadVisible |= coordReaderThread && task.contains("SELECT");
+            coordinatorReadVisible |= coordReaderThread && task.contains("SELECT") && task.contains("QUERY");
             boolean localWriterThread = threadId.contains("Mutation") || threadId.contains("SharedPool-Worker");
-            writeVisible |= localWriterThread && task.contains("Mutation");
+            writeVisible |= localWriterThread && task.contains("Mutation") && !task.contains("QUERY");
             boolean coordWriterThread = threadId.contains("Native-Transport-Requests") || threadId.contains("SharedPool-Worker");
-            coordinatorWriteVisible |= coordWriterThread && task.contains("INSERT");
+            coordinatorWriteVisible |= coordWriterThread && task.contains("INSERT") && task.contains("QUERY");
         }
 
         assertTrue(readVisible);
@@ -169,9 +170,10 @@ public class QueriesTableTest extends TestBaseImpl
             String task = row.get("task").toString();
 
             boolean localReaderThread = threadId.contains("Read") || threadId.contains("SharedPool-Worker");
-            readVisible |= localReaderThread && task.contains("SELECT");
+            // QUERY in a task name is used as a way to differentiate a local replica read vs a coordinator read
+            readVisible |= localReaderThread && task.contains("SELECT") && !task.contains("QUERY");
             boolean coordUpdateThread = threadId.contains("Native-Transport-Requests") || threadId.contains("SharedPool-Worker");
-            coordinatorUpdateVisible |= coordUpdateThread && task.contains("UPDATE");
+            coordinatorUpdateVisible |= coordUpdateThread && task.contains("UPDATE") && task.contains("QUERY");
         }
 
         assertTrue(readVisible);