Add support for a per-operation context (#1619)

The goal of this commit is to allow low-level code, typically at the
level of file reads (say, a custom `FileChannel` implementation), to
obtain context on which higher level operation they correspond to. A
typical use case may for tiered storage where `FileChannel`
implementations may forward reads to remote storage, and where having
context on the overall operation could help, say, take prior work done
to set timeout, or aggregate metrics per "user-level" operations (rather
than per `FileChannel` operation).

To do this, we reuse the existing `ExecutorLocals` mechanism, and simply
have the top-level operation setup the proper context as an
`ExecutorLocal`, allowing it to be accessed by any low-level operations
operating on behalf of that operation. This is, in many way, similar to
tracing, but instead of a `TraceState` that collect what happens during
the operation, it is a relatively flexible notion of `OperationContext`.

As of this patch, this feature is fairly barebone and mostly exists for
extensions in the following sense:
1. only user reads (`ReadCommand` execution) currently sets up such a
   context. The code is written in such a way that adding support for
   other operations should be easy, but this is not done.
2. the context set by reads by default has barely any information: it
   merely has a `toString()` method that can roughly tell what the
   operation itself is, and so could have use for debugging, but not
   much more. Further, that context is not read by anything in this
   patch. However, said context are created through a "factory" and
   the factory class is configurable through a system property. So
   extension can override the factory in order to create contexts
   with more information/state, and they fetch/use those context where
   they see fit.

Author: pcmanus

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

@@ -21,6 +21,7 @@
 import io.netty.util.concurrent.FastThreadLocal;
 import org.apache.cassandra.sensors.RequestSensors;
 import org.apache.cassandra.service.ClientWarn;
+import org.apache.cassandra.service.context.OperationContext;
 import org.apache.cassandra.tracing.TraceState;
 import org.apache.cassandra.utils.Closeable;
 import org.apache.cassandra.utils.WithResources;
@@ -33,7 +34,7 @@
  */
 public class ExecutorLocals implements WithResources, Closeable
 {
-    private static final ExecutorLocals none = new ExecutorLocals(null, null, null);
+    private static final ExecutorLocals none = new ExecutorLocals(null, null, null, null);
     private static final FastThreadLocal<ExecutorLocals> locals = new FastThreadLocal<ExecutorLocals>()
     {
         @Override
@@ -45,22 +46,24 @@ protected ExecutorLocals initialValue()
 
     public static class Impl
     {
-        public static void set(TraceState traceState, ClientWarn.State clientWarnState, RequestSensors sensors)
+        public static void set(TraceState traceState, ClientWarn.State clientWarnState, RequestSensors sensors, OperationContext operationContext)
         {
-            if (traceState == null && clientWarnState == null && sensors == null) locals.set(none);
-            else locals.set(new ExecutorLocals(traceState, clientWarnState, sensors));
+            if (traceState == null && clientWarnState == null && sensors == null && operationContext == null) locals.set(none);
+            else locals.set(new ExecutorLocals(traceState, clientWarnState, sensors, operationContext));
         }
     }
 
     public final TraceState traceState;
     public final ClientWarn.State clientWarnState;
     public final RequestSensors sensors;
+    public final OperationContext operationContext;
 
-    protected ExecutorLocals(TraceState traceState, ClientWarn.State clientWarnState, RequestSensors sensors)
+    protected ExecutorLocals(TraceState traceState, ClientWarn.State clientWarnState, RequestSensors sensors, OperationContext operationContext)
     {
         this.traceState = traceState;
         this.clientWarnState = clientWarnState;
         this.sensors = sensors;
+        this.operationContext = operationContext;
     }
 
     /**
@@ -84,13 +87,13 @@ public static WithResources propagate()
     public static ExecutorLocals create(TraceState traceState)
     {
         ExecutorLocals current = locals.get();
-        return current.traceState == traceState ? current : new ExecutorLocals(traceState, current.clientWarnState, current.sensors);
+        return current.traceState == traceState ? current : new ExecutorLocals(traceState, current.clientWarnState, current.sensors, current.operationContext);
     }
 
     public static ExecutorLocals create(RequestSensors sensors)
     {
         ExecutorLocals current = locals.get();
-        return current.sensors == sensors ? current : new ExecutorLocals(current.traceState, current.clientWarnState, sensors);
+        return current.sensors == sensors ? current : new ExecutorLocals(current.traceState, current.clientWarnState, sensors, current.operationContext);
     }
 
     public static void clear()

src/java/org/apache/cassandra/config/CassandraRelevantProperties.java

@@ -35,6 +35,7 @@
 import org.apache.cassandra.net.MessagingService;
 import org.apache.cassandra.sensors.SensorsFactory;
 import org.apache.cassandra.service.FileSystemOwnershipCheck;
+import org.apache.cassandra.service.context.OperationContext;
 import org.apache.cassandra.service.reads.range.EndpointGroupingRangeCommandIterator;
 import org.apache.cassandra.utils.FBUtilities;
 import org.apache.cassandra.utils.StorageCompatibilityMode;
@@ -574,6 +575,11 @@ public enum CassandraRelevantProperties
      * when the JVM terminates. Therefore, we can use such optimization and not wait unnecessarily. */
     NON_GRACEFUL_CLOSE("cassandra.messagingService.nonGracefulClose"),
     NON_GRACEFUL_SHUTDOWN("cassandra.test.messagingService.nonGracefulShutdown"),
+    /**
+     * Allows custom implementation of {@link OperationContext.Factory} to optionally create and configure custom
+     * {@link OperationContext} instances.
+     */
+    OPERATION_CONTEXT_FACTORY("cassandra.operation_context_factory_class"),
     /** for specific tests */
     /** This property indicates whether disable_mbean_registration is true */
     ORG_APACHE_CASSANDRA_DISABLE_MBEAN_REGISTRATION("org.apache.cassandra.disable_mbean_registration"),

src/java/org/apache/cassandra/db/ReadExecutionController.java

@@ -27,6 +27,8 @@
 import org.apache.cassandra.index.Index;
 import org.apache.cassandra.metrics.DecayingEstimatedHistogramReservoir;
 import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.service.context.OperationContext;
+import org.apache.cassandra.service.context.OperationContextTracker;
 import org.apache.cassandra.tracing.Tracing;
 import org.apache.cassandra.utils.MonotonicClock;
 import org.apache.cassandra.utils.concurrent.OpOrder;
@@ -143,6 +145,8 @@ static ReadExecutionController forCommand(ReadCommand command, boolean trackRepa
 
         long createdAtNanos = baseCfs.metric.topLocalReadQueryTime.isEnabled() ? clock.now() : NO_SAMPLING;
 
+        OperationContextTracker.start(OperationContext.FACTORY.forRead(command, baseCfs));
+
         if (indexCfs == null)
             return new ReadExecutionController(command, baseCfs.readOrdering.start(), baseCfs.metadata(), null, null, createdAtNanos, trackRepairedStatus);
 
@@ -176,6 +180,7 @@ static ReadExecutionController forCommand(ReadCommand command, boolean trackRepa
                 if (indexController != null)
                     indexController.close();
             }
+            OperationContextTracker.endCurrent();
             throw e;
         }
     }
@@ -217,6 +222,8 @@ public void close()
             }
         }
 
+        OperationContextTracker.endCurrent();
+
         if (createdAtNanos != NO_SAMPLING)
             addSample();
 

src/java/org/apache/cassandra/sensors/RequestTracker.java

@@ -38,6 +38,6 @@ public RequestSensors get()
     public void set(RequestSensors sensors)
     {
         ExecutorLocals current = ExecutorLocals.current();
-        ExecutorLocals.Impl.set(current.traceState, current.clientWarnState, sensors);
+        ExecutorLocals.Impl.set(current.traceState, current.clientWarnState, sensors, current.operationContext);
     }
 }

src/java/org/apache/cassandra/service/ClientWarn.java

@@ -42,7 +42,7 @@ public State get()
     public void set(State value)
     {
         ExecutorLocals current = ExecutorLocals.current();
-        ExecutorLocals.Impl.set(current.traceState, value, current.sensors);
+        ExecutorLocals.Impl.set(current.traceState, value, current.sensors, current.operationContext);
     }
 
     public void warn(String text)

src/java/org/apache/cassandra/service/context/DefaultOperationContext.java

@@ -0,0 +1,64 @@
+/*
+ * 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.service.context;
+
+import java.util.function.Supplier;
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+
+/**
+ * Default implementation of {@link OperationContext}.
+ * <p>
+ * This default implementation is mostly only useful for debugging as the only concrete method is provices is a
+ * {@link #toString()} method giving details on the operation the context corresponds to (though the context object
+ * also identify the operation, so it could also theoretically be used from 2 separate place in the code to decide
+ * if they execute as part of the same operation).
+ */
+public class DefaultOperationContext implements OperationContext
+{
+    private final Supplier<String> toDebugString;
+
+    private DefaultOperationContext(Supplier<String> toDebugString)
+    {
+        this.toDebugString = toDebugString;
+    }
+
+    @Override
+    public void close()
+    {
+    }
+
+    @Override
+    public String toString()
+    {
+        return String.format("[%d] %s", System.identityHashCode(this), toDebugString.get());
+    }
+
+    /**
+     * Simple default implementation of {@link OperationContext.Factory} that creates {@link DefaultOperationContext}.
+     */
+    static class Factory implements OperationContext.Factory
+    {
+        @Override
+        public OperationContext forRead(ReadCommand command, ColumnFamilyStore cfs)
+        {
+            return new DefaultOperationContext(command::toCQLString);
+        }
+    }
+}

src/java/org/apache/cassandra/service/context/OperationContext.java

@@ -0,0 +1,73 @@
+/*
+ * 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.service.context;
+
+import org.apache.cassandra.db.ColumnFamilyStore;
+import org.apache.cassandra.db.ReadCommand;
+import org.apache.cassandra.utils.FBUtilities;
+
+import static org.apache.cassandra.config.CassandraRelevantProperties.OPERATION_CONTEXT_FACTORY;
+
+/**
+ * Represents some context about a "top-level" operation.
+ * <p>
+ * This interface is fairly open on purpose, as implementations for different operations could look fairly different.
+ * But it is also open-ended as it is an extension point: the {@link #FACTORY} used to create the context instances
+ * is configurable, and meant to allow extensions to add whatever information they need to the context.
+ * <p>
+ * Also note that what consistute a "top-level" operation is not strictly defined. At the time of this writing, those
+ * context are not serialized across nodes, so "top-level" is understood as "for a node", and so correspond to
+ * operations like "a `ReadCommand` execution on a replica".
+ * <p>
+ * The context of executing operation is tracked by {@link OperationContextTracker} which use the {@link ExecutorLocal}
+ * concept to make that context available to any methods that execute as part of the operation. Basically, this is a way
+ * to make the context available everwhere along the path of execution of the operation, without needing to pass that
+ * context as argument of every single method that could be involved by the operation execution (which in most cases
+ * would be <b>a lot of methods</b>).
+*/
+public interface OperationContext extends AutoCloseable
+{
+    Factory FACTORY = OPERATION_CONTEXT_FACTORY.getString() == null
+                      ? new DefaultOperationContext.Factory()
+                      : FBUtilities.construct(OPERATION_CONTEXT_FACTORY.getString(), "operation context factory");
+
+
+    /**
+     * Called when the operation this is a context of terminates, and thus when the context will not be used/retrieved
+     * anymore.
+     */
+    @Override
+    void close();
+
+    /**
+     * Factory used to create {@link OperationContext} instances.
+     * <p>
+     * The intent is that every operation that wants to set a context should have its own method in this interface, but
+     * operations are added as needed (instead of trying to cover every possible operation upfront).
+     * <p>
+     * Do note however that there can only be one operation context "active" at any given time (meaning, any thread
+     * execute can only see at most one context), so the context should be set at the higher level that make sense
+     * (and if necessary, sub-operations can enrich the context of their parent, assuming the parent context make room
+     * for this).
+     */
+    interface Factory
+    {
+        OperationContext forRead(ReadCommand command, ColumnFamilyStore cfs);
+    }
+}

src/java/org/apache/cassandra/service/context/OperationContextTracker.java

@@ -0,0 +1,55 @@
+/*
+ * 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.service.context;
+
+import org.apache.cassandra.concurrent.ExecutorLocals;
+
+public class OperationContextTracker extends ExecutorLocals.Impl
+{
+    public static final OperationContextTracker instance = new OperationContextTracker();
+
+    private OperationContextTracker()
+    {}
+
+    public OperationContext get()
+    {
+        return ExecutorLocals.current().operationContext;
+    }
+
+    public void set(OperationContext operationContext)
+    {
+        ExecutorLocals current = ExecutorLocals.current();
+        ExecutorLocals.Impl.set(current.traceState, current.clientWarnState, current.sensors, operationContext);
+    }
+
+    public static void start(OperationContext context)
+    {
+        instance.set(context);
+    }
+
+    public static void endCurrent()
+    {
+        OperationContext ctx = instance.get();
+        if (ctx != null)
+        {
+            ctx.close();
+            instance.set(null);
+        }
+    }
+}

src/java/org/apache/cassandra/tracing/Tracing.java

@@ -277,7 +277,7 @@ public TraceState get(TimeUUID sessionId)
     public void set(TraceState tls)
     {
         ExecutorLocals current = ExecutorLocals.current();
-        ExecutorLocals.Impl.set(tls, current.clientWarnState, current.sensors);
+        ExecutorLocals.Impl.set(tls, current.clientWarnState, current.sensors, current.operationContext);
     }
 
     public TraceState begin(final String request, final Map<String, String> parameters)

src/java/org/apache/cassandra/utils/concurrent/Timer.java

@@ -30,6 +30,7 @@
 import org.apache.cassandra.concurrent.ExecutorLocals;
 import org.apache.cassandra.sensors.RequestSensors;
 import org.apache.cassandra.service.ClientWarn;
+import org.apache.cassandra.service.context.OperationContext;
 import org.apache.cassandra.tracing.TraceState;
 
 /**
@@ -72,11 +73,12 @@ public Future<Void> onTimeout(Runnable task, long timeout, TimeUnit unit, Execut
         ClientWarn.State clientWarnState = executorLocals == null ? null : executorLocals.clientWarnState;
         TraceState traceState = executorLocals == null ? null : executorLocals.traceState;
         RequestSensors sensors = executorLocals == null ? null : executorLocals.sensors;
+        OperationContext operationContext = executorLocals == null ? null : executorLocals.operationContext;
         AsyncPromise<Void> result = new AsyncPromise<>();
         Timeout handle = timer.newTimeout(ignored ->
                                           {
 
-                                              ExecutorLocals.Impl.set(traceState, clientWarnState, sensors);
+                                              ExecutorLocals.Impl.set(traceState, clientWarnState, sensors, operationContext);
                                               try
                                               {
                                                   task.run();