Send client warnings when writing to a large partition

Patch by Minal Kyada; reviewed by David Capwell and marcuse for CASSANDRA-17258

Author: minal-kyada

CHANGES.txt

@@ -1,4 +1,5 @@
 5.1
+ * Send client warnings when writing to a large partition (CASSANDRA-17258)
  * Harden the possible range of values for max dictionary size and max total sample size for dictionary training (CASSANDRA-21194)
  * Implement a guardrail ensuring that minimum training frequency parameter is provided in ZstdDictionaryCompressor (CASSANDRA-21192)
  * Replace manual referencing with ColumnFamilyStore.selectAndReference when training a dictionary (CASSANDRA-21188)

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

@@ -579,6 +579,11 @@ public static class SSTableConfig
     public volatile int tombstone_warn_threshold = 1000;
     public volatile int tombstone_failure_threshold = 100000;
 
+    public volatile boolean write_thresholds_enabled = false;
+    public volatile DataStorageSpec.LongBytesBound write_size_warn_threshold = null;
+    public volatile DurationSpec.LongMillisecondsBound coordinator_write_warn_interval = new DurationSpec.LongMillisecondsBound("1000ms");
+    public volatile int write_tombstone_warn_threshold = -1;
+
     public TombstonesMetricGranularity tombstone_read_purgeable_metric_granularity = TombstonesMetricGranularity.disabled;
 
     public final ReplicaFilteringProtectionOptions replica_filtering_protection = new ReplicaFilteringProtectionOptions();

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

@@ -931,7 +931,7 @@ else if (conf.repair_session_space.toMebibytes() > (int) (Runtime.getRuntime().m
 
         applyConcurrentValidations(conf);
         applyRepairCommandPoolSize(conf);
-        applyReadThresholdsValidations(conf);
+        applyThresholdsValidations(conf);
 
         if (conf.concurrent_materialized_view_builders <= 0)
             throw new ConfigurationException("concurrent_materialized_view_builders should be strictly greater than 0, but was " + conf.concurrent_materialized_view_builders, false);
@@ -1311,11 +1311,19 @@ static void applyRepairCommandPoolSize(Config config)
     }
 
     @VisibleForTesting
-    static void applyReadThresholdsValidations(Config config)
+    static void applyThresholdsValidations(Config config)
     {
+        // Validate read thresholds
         validateReadThresholds("coordinator_read_size", config.coordinator_read_size_warn_threshold, config.coordinator_read_size_fail_threshold);
         validateReadThresholds("local_read_size", config.local_read_size_warn_threshold, config.local_read_size_fail_threshold);
         validateReadThresholds("row_index_read_size", config.row_index_read_size_warn_threshold, config.row_index_read_size_fail_threshold);
+
+        // Write threshold warning depends on top_partitions tracking
+        if (config.write_thresholds_enabled && !config.top_partitions_enabled)
+            logger.warn("Write thresholds require top partitions tracking to be enabled");
+
+        validateWriteSizeThreshold(config.write_size_warn_threshold, config.min_tracked_partition_size);
+        validateWriteTombstoneThresholdRange(config.write_tombstone_warn_threshold, config.min_tracked_partition_tombstone_count);
     }
 
     private static void validateReadThresholds(String name, DataStorageSpec.LongBytesBound warn, DataStorageSpec.LongBytesBound fail)
@@ -1326,6 +1334,25 @@ private static void validateReadThresholds(String name, DataStorageSpec.LongByte
                                                            name + "_warn_threshold", warn));
     }
 
+    private static void validateWriteSizeThreshold(DataStorageSpec.LongBytesBound writeSizeWarn, DataStorageSpec.LongBytesBound minTrackedSize)
+    {
+        if (writeSizeWarn != null && minTrackedSize != null)
+        {
+            if (writeSizeWarn.toBytes() < minTrackedSize.toBytes())
+                throw new ConfigurationException(String.format("write_size_warn_threshold (%s) cannot be less than min_tracked_partition_size (%s)", writeSizeWarn, minTrackedSize));
+        }
+    }
+
+    private static void validateWriteTombstoneThresholdRange(int writeTombstoneWarn, long minTrackedTombstoneCount)
+    {
+        if (writeTombstoneWarn < -1)
+            throw new ConfigurationException(String.format("write_tombstone_warn_threshold (%d) must be -1 (disabled) or >= 0", writeTombstoneWarn));
+
+        if (writeTombstoneWarn != -1 && writeTombstoneWarn < minTrackedTombstoneCount)
+            throw new ConfigurationException(String.format("write_tombstone_warn_threshold (%d) cannot be less than min_tracked_partition_tombstone_count (%d)",
+                                                           writeTombstoneWarn, minTrackedTombstoneCount));
+    }
+
     public static GuardrailsOptions getGuardrailsConfig()
     {
         return guardrails;
@@ -5482,6 +5509,55 @@ public static void setRowIndexReadSizeFailThreshold(@Nullable DataStorageSpec.Lo
         conf.row_index_read_size_fail_threshold = value;
     }
 
+    public static boolean getWriteThresholdsEnabled()
+    {
+        return conf.write_thresholds_enabled;
+    }
+
+    public static void setWriteThresholdsEnabled(boolean enabled)
+    {
+        if (enabled && !conf.top_partitions_enabled)
+            logger.warn("Write thresholds require top partitions tracking to be enabled");
+        logger.info("updating write_thresholds_enabled to {}", enabled);
+        conf.write_thresholds_enabled = enabled;
+    }
+
+    @Nullable
+    public static DataStorageSpec.LongBytesBound getWriteSizeWarnThreshold()
+    {
+        return conf.write_size_warn_threshold;
+    }
+
+    public static void setWriteSizeWarnThreshold(@Nullable DataStorageSpec.LongBytesBound value)
+    {
+        validateWriteSizeThreshold(value, conf.min_tracked_partition_size);
+        logger.info("updating write_size_warn_threshold to {}", value);
+        conf.write_size_warn_threshold = value;
+    }
+
+    public static DurationSpec.LongMillisecondsBound getCoordinatorWriteWarnInterval()
+    {
+        return conf.coordinator_write_warn_interval;
+    }
+
+    public static void setCoordinatorWriteWarnInterval(DurationSpec.LongMillisecondsBound ms)
+    {
+        logger.info("updating coordinator_write_warn_interval to {}", ms);
+        conf.coordinator_write_warn_interval = ms;
+    }
+
+    public static int getWriteTombstoneWarnThreshold()
+    {
+        return conf.write_tombstone_warn_threshold;
+    }
+
+    public static void setWriteTombstoneWarnThreshold(int threshold)
+    {
+        validateWriteTombstoneThresholdRange(threshold, conf.min_tracked_partition_tombstone_count);
+        logger.info("updating write_tombstone_warn_threshold to {}", threshold);
+        conf.write_tombstone_warn_threshold = threshold;
+    }
+
     public static int getDefaultKeyspaceRF()
     {
         return conf.default_keyspace_rf;

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

@@ -18,6 +18,7 @@
 
 package org.apache.cassandra.db;
 
+import java.util.Collections;
 import java.util.EnumMap;
 import java.util.Map;
 
@@ -66,6 +67,22 @@ public static void reset()
         get().clear();
     }
 
+    /**
+     * Capture the current MessageParams for use across threads.
+     * This returns a copy of the current ThreadLocal params, which can be safely
+     * passed to async callbacks that may run on different threads.
+     *
+     * @return immutable copy of current params, or empty map if none
+     */
+    public static Map<ParamType, Object> capture()
+    {
+        Map<ParamType, Object> current = local.get();
+        if (current == null || current.isEmpty())
+            return Collections.emptyMap();
+
+        return new EnumMap<>(current);
+    }
+
     public static <T> Message<T> addToMessage(Message<T> message)
     {
         return message.withParams(get());

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

@@ -17,6 +17,8 @@
  */
 package org.apache.cassandra.db;
 
+import java.util.Map;
+
 import org.apache.cassandra.exceptions.WriteTimeoutException;
 import org.apache.cassandra.locator.InetAddressAndPort;
 import org.apache.cassandra.net.ForwardingInfo;
@@ -33,10 +35,13 @@ public class MutationVerbHandler extends AbstractMutationVerbHandler<Mutation>
 {
     public static final MutationVerbHandler instance = new MutationVerbHandler();
 
-    private void respond(Message<?> respondTo, InetAddressAndPort respondToAddress)
+    private void respond(Message<?> respondTo, InetAddressAndPort respondToAddress, Map<ParamType, Object> params)
     {
         Tracing.trace("Enqueuing response to {}", respondToAddress);
-        MessagingService.instance().send(respondTo.emptyResponse(), respondToAddress);
+        Message<?> response = respondTo.emptyResponse();
+        if (!params.isEmpty())
+            response = response.withParams(params);
+        MessagingService.instance().send(response, respondToAddress);
     }
 
     private void failed()
@@ -53,9 +58,10 @@ public void doVerb(Message<Mutation> message)
             return;
         }
 
+        MessageParams.reset();
         message.payload.validateSize(MessagingService.current_version, ENTRY_OVERHEAD_SIZE);
+        WriteThresholds.checkWriteThresholds(message.payload);
 
-        // Check if there were any forwarding headers in this message
         ForwardingInfo forwardTo = message.forwardTo();
         if (forwardTo != null)
             forwardToLocalNodes(message, forwardTo);
@@ -73,7 +79,8 @@ public void doVerb(Message<Mutation> message)
 
     protected void applyMutation(Message<Mutation> message, InetAddressAndPort respondToAddress)
     {
-        message.payload.applyFuture().addCallback(o -> respond(message, respondToAddress), wto -> failed());
+        Map<ParamType, Object> params = MessageParams.capture();
+        message.payload.applyFuture().addCallback(o -> respond(message, respondToAddress, params), wto -> failed());
     }
 
     private static void forwardToLocalNodes(Message<Mutation> originalMessage, ForwardingInfo forwardTo)

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

@@ -0,0 +1,112 @@
+/*
+ * 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.db;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.config.DataStorageSpec;
+import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.metrics.TopPartitionTracker;
+import org.apache.cassandra.net.ParamType;
+import org.apache.cassandra.schema.Schema;
+import org.apache.cassandra.schema.TableId;
+import org.apache.cassandra.schema.TableMetadata;
+import org.apache.cassandra.utils.NoSpamLogger;
+
+/**
+ * Utility class for checking write threshold warnings on replicas.
+ * CASSANDRA-17258: paxos and accord do complex thread hand off and custom write logic which makes this patch complex, so was deferred
+ */
+public class WriteThresholds
+{
+    private static final Logger logger = LoggerFactory.getLogger(WriteThresholds.class);
+    private static final NoSpamLogger noSpamLogger = NoSpamLogger.getLogger(logger, 1, TimeUnit.MINUTES);
+
+    /**
+     * Check write thresholds for a mutation by comparing the estimated partition size and tombstone count
+     * from {@link TopPartitionTracker} against configured warn thresholds. If a threshold is breached,
+     * a warning is logged and the corresponding {@link ParamType} is added to {@link MessageParams}
+     * for propagation back to the coordinator.
+     *
+     * @param mutation the mutation containing one or more partition updates
+     */
+    public static void checkWriteThresholds(Mutation mutation)
+    {
+        if (!DatabaseDescriptor.isDaemonInitialized() || !DatabaseDescriptor.getWriteThresholdsEnabled())
+            return;
+
+        DataStorageSpec.LongBytesBound sizeWarnThreshold = DatabaseDescriptor.getWriteSizeWarnThreshold();
+        int tombstoneWarnThreshold = DatabaseDescriptor.getWriteTombstoneWarnThreshold();
+
+        if (sizeWarnThreshold == null && tombstoneWarnThreshold == -1)
+            return;
+
+        long sizeWarnBytes = sizeWarnThreshold != null ? sizeWarnThreshold.toBytes() : -1;
+        DecoratedKey key = mutation.key();
+
+        Map<TableId, Long> sizeWarnings = null;
+        Map<TableId, Long> tombstoneWarnings = null;
+
+        for (TableId tableId : mutation.getTableIds())
+        {
+            ColumnFamilyStore cfs = Schema.instance.getColumnFamilyStoreInstance(tableId);
+            if (cfs == null || cfs.topPartitions == null)
+                continue;
+
+            TableMetadata metadata = cfs.metadata();
+            if (sizeWarnBytes != -1)
+            {
+                long estimatedSize = cfs.topPartitions.topSizes().getEstimate(key);
+                if (estimatedSize > sizeWarnBytes)
+                {
+                    if (sizeWarnings == null)
+                        sizeWarnings = new HashMap<>();
+                    sizeWarnings.put(tableId, estimatedSize);
+                    noSpamLogger.warn("Write to {} partition {} triggered size warning; " +
+                                      "estimated size is {} bytes, threshold is {} bytes (see write_size_warn_threshold)",
+                                      metadata, metadata.partitionKeyType.toCQLString(key.getKey()), estimatedSize, sizeWarnBytes);
+                }
+            }
+
+            if (tombstoneWarnThreshold != -1)
+            {
+                long estimatedTombstones = cfs.topPartitions.topTombstones().getEstimate(key);
+                if (estimatedTombstones > tombstoneWarnThreshold)
+                {
+                    if (tombstoneWarnings == null)
+                        tombstoneWarnings = new HashMap<>();
+                    tombstoneWarnings.put(tableId, estimatedTombstones);
+                    noSpamLogger.warn("Write to {} partition {} triggered tombstone warning; " +
+                                      "estimated tombstone count is {}, threshold is {} (see write_tombstone_warn_threshold)",
+                                      metadata, metadata.partitionKeyType.toCQLString(key.getKey()), estimatedTombstones, tombstoneWarnThreshold);
+                }
+            }
+        }
+
+        if (sizeWarnings != null)
+            MessageParams.add(ParamType.WRITE_SIZE_WARN, sizeWarnings);
+        if (tombstoneWarnings != null)
+            MessageParams.add(ParamType.WRITE_TOMBSTONE_WARN, tombstoneWarnings);
+    }
+}

src/java/org/apache/cassandra/metrics/KeyspaceMetrics.java

@@ -191,6 +191,10 @@ public class KeyspaceMetrics
 
     public final Meter tooManySSTableIndexesReadWarnings;
     public final Meter tooManySSTableIndexesReadAborts;
+
+    public final Meter writeSizeWarnings;
+    public final Meter writeTombstoneWarnings;
+
     public final Meter bytesAnticompacted;
     public final Meter bytesMutatedAnticompaction;
     public final Meter bytesPreviewed;
@@ -309,6 +313,9 @@ public KeyspaceMetrics(final Keyspace ks)
         tooManySSTableIndexesReadWarnings = createKeyspaceMeter("TooManySSTableIndexesReadWarnings");
         tooManySSTableIndexesReadAborts = createKeyspaceMeter("TooManySSTableIndexesReadAborts");
 
+        writeSizeWarnings = createKeyspaceMeter("WriteSizeWarnings");
+        writeTombstoneWarnings = createKeyspaceMeter("WriteTombstoneWarnings");
+
         formatSpecificGauges = createFormatSpecificGauges(keyspace);
 
         outOfRangeTokenReads = createKeyspaceCounter("ReadOutOfRangeToken");

src/java/org/apache/cassandra/metrics/TableMetrics.java

@@ -307,6 +307,9 @@ public class TableMetrics
     public final TableMeter tooManySSTableIndexesReadWarnings;
     public final TableMeter tooManySSTableIndexesReadAborts;
 
+    public final TableMeter writeSizeWarnings;
+    public final TableMeter writeTombstoneWarnings;
+
     public final ImmutableMap<SSTableFormat<?, ?>, ImmutableMap<String, Gauge<? extends Number>>> formatSpecificGauges;
 
     // Time spent building SSTableIntervalTree when constructing a new View under the Tracker lock
@@ -907,6 +910,9 @@ public Long getValue()
         tooManySSTableIndexesReadWarnings = createTableMeter("TooManySSTableIndexesReadWarnings", cfs.keyspace.metric.tooManySSTableIndexesReadWarnings);
         tooManySSTableIndexesReadAborts = createTableMeter("TooManySSTableIndexesReadAborts", cfs.keyspace.metric.tooManySSTableIndexesReadAborts);
 
+        writeSizeWarnings = createTableMeter("WriteSizeWarnings", cfs.keyspace.metric.writeSizeWarnings);
+        writeTombstoneWarnings = createTableMeter("WriteTombstoneWarnings", cfs.keyspace.metric.writeTombstoneWarnings);
+
         viewSSTableIntervalTree = createLatencyMetrics("ViewSSTableIntervalTree", cfs.keyspace.metric.viewSSTableIntervalTree);
 
         formatSpecificGauges = createFormatSpecificGauges(cfs);