Fix CompressionDictionary being closed while still in use

Patch by Yifan Cai; Reviewed by Stefan Miklosovic for CASSANDRA-21047

Author: yifan-c

CHANGES.txt

@@ -1,4 +1,5 @@
 5.1
+ * Fix CompressionDictionary being closed while still in use (CASSANDRA-21047)
  * When updating a multi cell collection element, if the update is rejected then the shared Row.Builder is not freed causing all future mutations to be rejected (CASSANDRA-21055)
  * Schema annotations escape validation on CREATE and ALTER DDL statements (CASSANDRA-21046)
  * Calculate once and cache the result of ModificationStatement#requiresRead as a perf optimization (CASSANDRA-21040)

src/java/org/apache/cassandra/db/compression/CompressionDictionary.java

@@ -25,15 +25,49 @@
 import java.util.Objects;
 import javax.annotation.Nullable;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.hash.Hasher;
 import com.google.common.hash.Hashing;
 
 import org.apache.cassandra.cql3.UntypedResultSet;
 import org.apache.cassandra.io.compress.ICompressor;
 import org.apache.cassandra.io.compress.ZstdDictionaryCompressor;
+import org.apache.cassandra.utils.concurrent.Ref;
 
-public interface CompressionDictionary extends AutoCloseable
+/**
+ * Interface for compression dictionaries with reference-counted lifecycle management.
+ *
+ * <h2>Reference Counting Model</h2>
+ * Compression dictionaries hold native resources that must be explicitly managed. This interface
+ * uses {@link Ref} for safe lifecycle management across multiple concurrent users.
+ *
+ * <h3>Ownership and Usage in Cassandra</h3>
+ * <ul>
+ *   <li><b>CompressionDictionaryManager</b>: Holds the primary reference ({@link #selfRef()}) for cached dictionaries</li>
+ *   <li><b>CompressionMetadata.Writer</b>: Acquires a reference during SSTable write, held for the writer's lifetime</li>
+ *   <li><b>CompressionMetadata</b>: Acquires a reference when created (via {@link #tryRef()}), held for the SSTable reader's lifetime.
+ *       All copies created via sharedCopy() share this single reference through WrappedSharedCloseable</li>
+ * </ul>
+ *
+ * <h3>Correctness Guarantee</h3>
+ * The reference counting prevents premature cleanup of native resources:
+ * <ol>
+ *   <li>CompressionMetadata acquires a reference when an SSTable is opened</li>
+ *   <li>Native resources remain valid as long as any reference exists (refcount &gt; 0)</li>
+ *   <li>Even if the cache evicts the dictionary, the SSTable's reference keeps resources alive</li>
+ *   <li>Cleanup runs exactly once when the last reference is released (refcount goes 0 → -1)</li>
+ *   <li>After cleanup, {@link #tryRef()} returns null, preventing new references to released resources</li>
+ * </ol>
+ *
+ * This ensures dictionaries cannot be freed while SSTables are using them for compression/decompression,
+ * even when the cache evicts the dictionary concurrently.
+ *
+ * @see Ref for reference counting implementation
+ * @see CompressionDictionaryManager for cache management
+ * @see org.apache.cassandra.io.compress.CompressionMetadata for SSTable usage
+ */
+public interface CompressionDictionary
 {
     /**
      * Get the dictionary id
@@ -75,6 +109,48 @@ default Kind kind()
         return dictId().kind;
     }
 
+    /**
+     * Try to acquire a new reference to this dictionary.
+     * Returns null if the dictionary is already released.
+     * <p>
+     * The caller must ensure the returned reference is released when no longer needed,
+     * either by calling {@code ref.release()} or {@code ref.close()} (they are equivalent).
+     * Failing to release the reference will prevent cleanup of native resources and cause
+     * a memory leak.
+     *
+     * @return a new reference to this dictionary, or null if already released
+     */
+    Ref<? extends CompressionDictionary> tryRef();
+
+    /**
+     * Get the self-reference of this dictionary.
+     * This is used to release the primary reference held by the cache.
+     *
+     * @return the self-reference
+     */
+    Ref<? extends CompressionDictionary> selfRef();
+
+    /**
+     * Releases the self-reference of this dictionary.
+     * This is a convenience method equivalent to calling {@code selfRef().close()}.
+     * <p>
+     * This method is idempotent - calling it multiple times is safe and will only
+     * release the self-reference once. Subsequent calls have no effect.
+     * <p>
+     * This method is typically used when creating a dictionary outside the cache
+     * (e.g., in tests or temporary usage) and needing to clean it up. For dictionaries
+     * managed by the cache, the cache's removal listener handles cleanup via
+     * {@code selfRef().release()}.
+     *
+     * @see #selfRef()
+     * @see #tryRef()
+     */
+    @VisibleForTesting
+    default void close()
+    {
+        selfRef().close();
+    }
+
     /**
      * Write compression dictionary to file
      *
@@ -192,15 +268,15 @@ static CompressionDictionary createFromRow(UntypedResultSet.Row row)
             if (dict.length != storedLength)
             {
                 throw new IllegalStateException(String.format("Dictionary length mismatch for %s dict id %d. Expected: %d, actual: %d",
-                                                               kindStr, dictId, storedLength, dict.length));
+                                                              kindStr, dictId, storedLength, dict.length));
             }
 
             // Validate checksum
             int calculatedChecksum = calculateChecksum((byte) kind.ordinal(), dictId, dict);
             if (calculatedChecksum != storedChecksum)
             {
                 throw new IllegalStateException(String.format("Dictionary checksum mismatch for %s dict id %d. Expected: %d, actual: %d",
-                                                               kindStr, dictId, storedChecksum, calculatedChecksum));
+                                                              kindStr, dictId, storedChecksum, calculatedChecksum));
             }
 
             return kind.createDictionary(new DictId(kind, dictId), row.getByteArray("dict"), storedChecksum);

src/java/org/apache/cassandra/db/compression/CompressionDictionaryCache.java

@@ -62,17 +62,17 @@ public CompressionDictionaryCache()
                              .removalListener((DictId dictId,
                                                CompressionDictionary dictionary,
                                                RemovalCause cause) -> {
-                                 // Close dictionary when evicted from cache to free native resources
-                                 // SelfRefCounted ensures dictionary won't be actually closed if still referenced by compressors
+                                 // Release the cache's reference to the dictionary when evicted
+                                 // The dictionary will only be truly cleaned up when all references are released
                                  if (dictionary != null)
                                  {
                                      try
                                      {
-                                         dictionary.close();
+                                         dictionary.selfRef().release();
                                      }
                                      catch (Exception e)
                                      {
-                                         logger.warn("Failed to close compression dictionary {}", dictId, e);
+                                         logger.warn("Failed to release compression dictionary {}", dictId, e);
                                      }
                                  }
                              })

src/java/org/apache/cassandra/db/compression/ZstdCompressionDictionary.java

@@ -20,7 +20,7 @@
 
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.slf4j.Logger;
@@ -42,8 +42,7 @@ public class ZstdCompressionDictionary implements CompressionDictionary, SelfRef
     private final int checksum;
     // One ZstdDictDecompress and multiple ZstdDictCompress (per level) can be derived from the same raw dictionary content
     private final ConcurrentHashMap<Integer, ZstdDictCompress> zstdDictCompressPerLevel = new ConcurrentHashMap<>();
-    private volatile ZstdDictDecompress dictDecompress;
-    private final AtomicBoolean closed = new AtomicBoolean(false);
+    private final AtomicReference<ZstdDictDecompress> dictDecompress = new AtomicReference<>();
     private final Ref<ZstdCompressionDictionary> selfRef;
 
     @VisibleForTesting
@@ -90,7 +89,7 @@ public int checksum()
     public int estimatedOccupiedMemoryBytes()
     {
         int occupied = rawDictionary.length;
-        occupied += dictDecompress != null ? rawDictionary.length : 0;
+        occupied += dictDecompress.get() != null ? rawDictionary.length : 0;
         occupied += zstdDictCompressPerLevel.size() * rawDictionary.length;
 
         return occupied;
@@ -114,50 +113,65 @@ public int hashCode()
      * Get a pre-processed compression tables that is optimized for compression.
      * It is derived/computed from dictionary bytes.
      * The internal data structure is different from the tables for decompression.
-     *
+     * <br>
+     * IMPORTANT: Caller MUST hold a valid reference (via tryRef/ref) to this dictionary.
+     * The reference counting mechanism ensures tidy() cannot run while references exist,
+     * making synchronization unnecessary. This method is safe to call concurrently as long
+     * as each caller holds a reference.
+     * <br>
      * @param compressionLevel compression level to create the compression table
-     * @return ZstdDictCompress
+     * @return ZstdDictCompress for the specified compression level
+     * @throws IllegalStateException if called without holding a valid reference
      */
     public ZstdDictCompress dictionaryForCompression(int compressionLevel)
     {
-        if (closed.get())
-            throw new IllegalStateException("Dictionary has been closed. " + dictId);
-
+        ensureNotReleased();
         ZstdCompressorBase.validateCompressionLevel(compressionLevel);
 
-        return zstdDictCompressPerLevel.computeIfAbsent(compressionLevel, level -> {
-            if (closed.get())
-                throw new IllegalStateException("Dictionary has been closed");
-            return new ZstdDictCompress(rawDictionary, level);
-        });
+        // Fast path: check if already exists to avoid locking the bin
+        ZstdDictCompress existing = zstdDictCompressPerLevel.get(compressionLevel);
+        if (existing != null)
+            return existing;
+
+        // A little slow path: create new dictionary for this compression level
+        // No additional synchronization needed - reference counting prevents tidy() while in use
+        return zstdDictCompressPerLevel.computeIfAbsent(compressionLevel, level ->
+            new ZstdDictCompress(rawDictionary, level));
     }
 
     /**
      * Get a pre-processed decompression tables that is optimized for decompression.
      * It is derived/computed from dictionary bytes.
      * The internal data structure is different from the tables for compression.
+     * <br>
+     * IMPORTANT: Caller MUST hold a valid reference (via tryRef/ref) to this dictionary.
+     * The reference counting mechanism ensures tidy() cannot run while references exist,
+     * making synchronization unnecessary. This method is safe to call concurrently as long
+     * as each caller holds a reference.
+     * <br>
+     * Thread-safe: Multiple threads can safely call this method concurrently.
+     * The decompression dictionary will be created exactly once on first access.
      *
-     * @return ZstdDictDecompress
+     * @return ZstdDictDecompress for decompression operations
+     * @throws IllegalStateException if called without holding a valid reference
      */
     public ZstdDictDecompress dictionaryForDecompression()
     {
-        if (closed.get())
-            throw new IllegalStateException("Dictionary has been closed");
-
-        ZstdDictDecompress result = dictDecompress;
+        ensureNotReleased();
+        // Fast path: if already initialized, return immediately
+        ZstdDictDecompress result = dictDecompress.get();
         if (result != null)
             return result;
 
+        // Slow path: need to initialize with proper double-checked locking
+        // Reference counting guarantees tidy() won't run during this operation
         synchronized (this)
         {
-            if (closed.get())
-                throw new IllegalStateException("Dictionary has been closed");
-
-            result = dictDecompress;
+            result = dictDecompress.get();
             if (result == null)
             {
                 result = new ZstdDictDecompress(rawDictionary);
-                dictDecompress = result;
+                dictDecompress.set(result);
             }
             return result;
         }
@@ -181,30 +195,48 @@ public Ref<ZstdCompressionDictionary> ref()
         return selfRef.ref();
     }
 
-    @Override
-    public void close()
+    private void ensureNotReleased()
     {
-        if (closed.compareAndSet(false, true))
-        {
-            selfRef.release();
-        }
+        if (selfRef.globalCount() <= 0)
+            throw new IllegalStateException("Dictionary has been released: " + dictId);
     }
 
+    /**
+     * Tidy implementation for cleaning up native Zstd resources.
+     *
+     * This class holds direct references to the resources that need cleanup,
+     * avoiding a circular reference pattern where Tidy would hold a reference
+     * to the parent dictionary object.
+     */
     private static class Tidy implements RefCounted.Tidy
     {
         private final ConcurrentHashMap<Integer, ZstdDictCompress> zstdDictCompressPerLevel;
-        private volatile ZstdDictDecompress dictDecompress;
+        private final AtomicReference<ZstdDictDecompress> dictDecompress;
 
-        Tidy(ConcurrentHashMap<Integer, ZstdDictCompress> zstdDictCompressPerLevel, ZstdDictDecompress dictDecompress)
+        Tidy(ConcurrentHashMap<Integer, ZstdDictCompress> zstdDictCompressPerLevel,
+             AtomicReference<ZstdDictDecompress> dictDecompress)
         {
             this.zstdDictCompressPerLevel = zstdDictCompressPerLevel;
             this.dictDecompress = dictDecompress;
         }
 
+        /**
+         * Clean up native resources when reference count reaches zero.
+         *
+         * IMPORTANT: This method is called exactly once when the last reference is released.
+         * Reference counting guarantees that no other thread can be executing
+         * dictionaryForCompression/Decompression when this runs, because:
+         * 1. Those methods require holding a valid reference
+         * 2. This only runs when refcount goes from 0 to -1
+         * 3. Once refcount is negative, tryRef() returns null, preventing new references
+         *
+         * Therefore, no synchronization is needed - we have exclusive access to clean up.
+         */
         @Override
         public void tidy()
         {
             // Close all compression dictionaries
+            // No synchronization needed - reference counting ensures exclusive access
             for (ZstdDictCompress compressDict : zstdDictCompressPerLevel.values())
             {
                 try
@@ -220,7 +252,7 @@ public void tidy()
             zstdDictCompressPerLevel.clear();
 
             // Close decompression dictionary
-            ZstdDictDecompress decompressDict = dictDecompress;
+            ZstdDictDecompress decompressDict = dictDecompress.get();
             if (decompressDict != null)
             {
                 try
@@ -231,7 +263,7 @@ public void tidy()
                 {
                     logger.warn("Failed to close ZstdDictDecompress", e);
                 }
-                dictDecompress = null;
+                dictDecompress.set(null);
             }
         }