CASSANDRA-17258: Add write threshold warnings for large partitions and tombstones

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

@@ -583,6 +583,10 @@ 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 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

@@ -5443,6 +5443,40 @@ 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)
+    {
+        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)
+    {
+        logger.info("updating write_size_warn_threshold to {}", value);
+        conf.write_size_warn_threshold = value;
+    }
+
+    public static int getWriteTombstoneWarnThreshold()
+    {
+        return conf.write_tombstone_warn_threshold;
+    }
+
+    public static void setWriteTombstoneWarnThreshold(int threshold)
+    {
+        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,120 @@
+/*
+ * 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.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.db.partitions.PartitionUpdate;
+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 all partition updates in a mutation.
+     * This method iterates through all partition updates in the mutation.
+     *
+     * @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;
+
+        for (PartitionUpdate update : mutation.getPartitionUpdates())
+        {
+            checkWriteThresholdsInternal(update, update.partitionKey(), sizeWarnBytes, tombstoneWarnThreshold);
+        }
+    }
+
+    /**
+     * Internal method to check write thresholds for a single partition update.
+     * This method looks up the partition in TopPartitionTracker and adds
+     * warning params to MessageParams if thresholds are exceeded.
+     *
+     * @param update                 the partition update being written
+     * @param key                    the partition key being written
+     * @param sizeWarnBytes          size threshold in bytes, or -1 if disabled
+     * @param tombstoneWarnThreshold tombstone count threshold, or -1 if disabled
+     */
+    private static void checkWriteThresholdsInternal(PartitionUpdate update, DecoratedKey key,
+                                                     long sizeWarnBytes, int tombstoneWarnThreshold)
+    {
+        TableId tableId = update.metadata().id;
+        ColumnFamilyStore cfs = Schema.instance.getColumnFamilyStoreInstance(tableId);
+
+        if (cfs == null || cfs.topPartitions == null)
+            return;
+
+        long estimatedSize = cfs.topPartitions.topSizes().getEstimate(key);
+        long estimatedTombstones = cfs.topPartitions.topTombstones().getEstimate(key);
+        TableMetadata meta = update.metadata();
+
+        if (sizeWarnBytes != -1 && estimatedSize > sizeWarnBytes)
+        {
+            Number currentValue = MessageParams.get(ParamType.WRITE_SIZE_WARN);
+            long currentLong = currentValue != null ? currentValue.longValue() : -1;
+
+            if (currentLong < estimatedSize)
+            {
+                MessageParams.add(ParamType.WRITE_SIZE_WARN, estimatedSize);
+                noSpamLogger.warn("Write to {} partition {} triggered size warning; " +
+                                  "estimated size is {} bytes, threshold is {} bytes (see write_size_warn_threshold)",
+                                  meta, meta.partitionKeyType.toCQLString(key.getKey()), estimatedSize, sizeWarnBytes);
+            }
+        }
+
+        if (tombstoneWarnThreshold != -1 && estimatedTombstones > tombstoneWarnThreshold)
+        {
+            Number currentValue = MessageParams.get(ParamType.WRITE_TOMBSTONE_WARN);
+            long currentLong = currentValue != null ? currentValue.longValue() : -1;
+
+            if (currentLong < estimatedTombstones)
+            {
+                MessageParams.add(ParamType.WRITE_TOMBSTONE_WARN, (int) estimatedTombstones);
+                noSpamLogger.warn("Write to {} partition {} triggered tombstone warning; " +
+                                  "estimated tombstone count is {}, threshold is {} (see write_tombstone_warn_threshold)",
+                                  meta, meta.partitionKeyType.toCQLString(key.getKey()), estimatedTombstones, tombstoneWarnThreshold);
+            }
+        }
+    }
+}

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

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

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

@@ -307,6 +307,10 @@ public class TableMetrics
     public final TableMeter tooManySSTableIndexesReadWarnings;
     public final TableMeter tooManySSTableIndexesReadAborts;
 
+    public final TableMeter writeSizeWarnings;
+    public final TableHistogram writeSize;
+    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 +911,10 @@ public Long getValue()
         tooManySSTableIndexesReadWarnings = createTableMeter("TooManySSTableIndexesReadWarnings", cfs.keyspace.metric.tooManySSTableIndexesReadWarnings);
         tooManySSTableIndexesReadAborts = createTableMeter("TooManySSTableIndexesReadAborts", cfs.keyspace.metric.tooManySSTableIndexesReadAborts);
 
+        writeSizeWarnings = createTableMeter("WriteSizeWarnings", cfs.keyspace.metric.writeSizeWarnings);
+        writeSize = createTableHistogram("WriteSize", cfs.keyspace.metric.writeSize, false);
+        writeTombstoneWarnings = createTableMeter("WriteTombstoneWarnings", cfs.keyspace.metric.writeTombstoneWarnings);
+
         viewSSTableIntervalTree = createLatencyMetrics("ViewSSTableIntervalTree", cfs.keyspace.metric.viewSSTableIntervalTree);
 
         formatSpecificGauges = createFormatSpecificGauges(cfs);

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

@@ -32,6 +32,8 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
+import com.carrotsearch.hppc.ObjectLongHashMap;
+import com.carrotsearch.hppc.ObjectLongMap;
 import com.google.common.annotations.VisibleForTesting;
 
 import org.apache.cassandra.concurrent.ScheduledExecutors;
@@ -197,6 +199,7 @@ public String toString()
     public static class TopHolder
     {
         public final NavigableSet<TopPartition> top;
+        public final ObjectLongMap<DecoratedKey> lookup;
         private final int maxTopPartitionCount;
         private final long minTrackedValue;
         private final Collection<Range<Token>> ranges;
@@ -205,14 +208,15 @@ public static class TopHolder
 
         private TopHolder(int maxTopPartitionCount, long minTrackedValue, Collection<Range<Token>> ranges)
         {
-            this(maxTopPartitionCount, minTrackedValue, new TreeSet<>(), ranges, 0);
+            this(maxTopPartitionCount, minTrackedValue, new TreeSet<>(), new ObjectLongHashMap<>(), ranges, 0);
         }
 
-        private TopHolder(int maxTopPartitionCount, long minTrackedValue, NavigableSet<TopPartition> top, Collection<Range<Token>> ranges, long lastUpdate)
+        private TopHolder(int maxTopPartitionCount, long minTrackedValue, NavigableSet<TopPartition> top, ObjectLongMap<DecoratedKey> lookup, Collection<Range<Token>> ranges, long lastUpdate)
         {
             this.maxTopPartitionCount = maxTopPartitionCount;
             this.minTrackedValue = minTrackedValue;
             this.top = top;
+            this.lookup = lookup;
             this.ranges = ranges;
             this.lastUpdate = lastUpdate;
         }
@@ -224,6 +228,7 @@ private TopHolder(StoredTopPartitions storedTopPartitions,
             this.maxTopPartitionCount = maxTopPartitionCount;
             this.minTrackedValue = minTrackedValue;
             top = new TreeSet<>();
+            this.lookup = new ObjectLongHashMap<>();
             this.ranges = null;
             this.lastUpdate = storedTopPartitions.lastUpdated;
 
@@ -243,9 +248,11 @@ public void track(DecoratedKey key, long value)
         private void track(TopPartition tp)
         {
             top.add(tp);
+            lookup.put(tp.key, tp.value);
             while (top.size() > maxTopPartitionCount)
             {
-                top.pollLast();
+                TopPartition p = top.pollLast();
+                lookup.remove(p.key);
                 currentMinValue = top.last().value;
             }
             currentMinValue = Math.min(tp.value, currentMinValue);
@@ -274,7 +281,12 @@ public TopHolder merge(TopHolder holder, Collection<Range<Token>> ownedRanges)
 
         private TopHolder cloneForMerging(long lastUpdate)
         {
-            return new TopHolder(maxTopPartitionCount, minTrackedValue, new TreeSet<>(top), ranges, lastUpdate);
+            return new TopHolder(maxTopPartitionCount, minTrackedValue, new TreeSet<>(top), new ObjectLongHashMap<>(lookup), ranges, lastUpdate);
+        }
+
+        public long getEstimate(DecoratedKey dk)
+        {
+            return lookup.getOrDefault(dk, 0L);
         }
 
         public String toString()

src/java/org/apache/cassandra/net/ParamType.java

@@ -54,7 +54,9 @@ public enum ParamType
     ROW_INDEX_READ_SIZE_WARN         (13, Int64Serializer.serializer),
     CUSTOM_MAP                       (14, CustomParamsSerializer.serializer),
     TOO_MANY_REFERENCED_INDEXES_WARN (16, Int32Serializer.serializer),
-    TOO_MANY_REFERENCED_INDEXES_FAIL (17, Int32Serializer.serializer);
+    TOO_MANY_REFERENCED_INDEXES_FAIL (17, Int32Serializer.serializer),
+    WRITE_SIZE_WARN                  (18, Int64Serializer.serializer),
+    WRITE_TOMBSTONE_WARN             (19, Int32Serializer.serializer);
 
     final int id;
     final IVersionedSerializer serializer;