Change the eager reference counting of compression dictionaries to lazy

patch by Stefan Miklosovic; reviewed by Yifan Cai, Jyothsna Konisa for CASSANDRA-21074

Author: smiklosovic

CHANGES.txt

@@ -1,4 +1,5 @@
 5.1
+ * Change the eager reference counting of compression dictionaries to lazy (CASSANDRA-21074)
  * Add cursor based optimized compaction path (CASSANDRA-20918)
  * Ensure peers with LEFT status are expired from gossip state (CASSANDRA-21035)
  * Optimize UTF8Validator.validate for ASCII prefixed Strings (CASSANDRA-21075)

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

@@ -36,7 +36,16 @@
 import org.apache.cassandra.utils.concurrent.Ref;
 
 /**
- * Interface for compression dictionaries with reference-counted lifecycle management.
+ * Interface for compression dictionaries with opt-in reference-counted lifecycle management.
+ * <p>
+ * Dictionaries can be used in two modes:
+ * <ul>
+ *   <li><b>Lightweight mode</b>: No native resources allocated. Suitable for export, serialization,
+ *       or scenarios where only the raw dictionary bytes are needed.</li>
+ *   <li><b>Managed mode</b>: Native compression/decompression resources allocated on-demand and
+ *       managed via reference counting. Required for caching and active use.</li>
+ * </ul>
+ * Call {@link #initRefLazily()} or {@link #tryRef()} to transition from lightweight to managed mode.
  *
  * <h2>Reference Counting Model</h2>
  * Compression dictionaries hold native resources that must be explicitly managed. This interface
@@ -109,6 +118,13 @@ default Kind kind()
         return dictId().kind;
     }
 
+    /**
+     * Returns a reference from lazily initialized reference counter.
+     *
+     * @return reference to this dictionary; once initialized, the reference is the same as self-reference
+     */
+    Ref<? extends CompressionDictionary> initRefLazily();
+
     /**
      * Try to acquire a new reference to this dictionary.
      * Returns null if the dictionary is already released.
@@ -125,9 +141,11 @@ default Kind kind()
     /**
      * Get the self-reference of this dictionary.
      * This is used to release the primary reference held by the cache.
+     * Self-reference is initialized after initRefLazily or tryRef
      *
-     * @return the self-reference
+     * @return the self-reference or null if not yet initialized.
      */
+    @Nullable
     Ref<? extends CompressionDictionary> selfRef();
 
     /**
@@ -137,9 +155,12 @@ default Kind kind()
      * 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
+     * There is no need to call this method in the context of releasing references when an instance
+     * is never put to a cache. That might be the case when we are constructing a dictionary object
+     * just for the purpose of passing it to a user (e.g. on exporting and similar) where reference counting
+     * is not necessary nor needed.
+     * <p>
+     * For dictionaries managed by the cache, the cache's removal listener handles cleanup via
      * {@code selfRef().release()}.
      *
      * @see #selfRef()
@@ -148,7 +169,9 @@ default Kind kind()
     @VisibleForTesting
     default void close()
     {
-        selfRef().close();
+        Ref<?> selfRef = selfRef();
+        if (selfRef != null)
+            selfRef.close();
     }
 
     /**

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

@@ -32,6 +32,7 @@
 import com.github.benmanes.caffeine.cache.RemovalCause;
 import org.apache.cassandra.config.DatabaseDescriptor;
 import org.apache.cassandra.db.compression.CompressionDictionary.DictId;
+import org.apache.cassandra.utils.concurrent.Ref;
 
 /**
  * Manages caching and current dictionary state for compression dictionaries.
@@ -68,7 +69,9 @@ public CompressionDictionaryCache()
                                  {
                                      try
                                      {
-                                         dictionary.selfRef().release();
+                                         // The dictionary's selfRef should never be null when evicting from cache
+                                         // but using close() to have better resiliency
+                                         dictionary.close();
                                      }
                                      catch (Exception e)
                                      {
@@ -102,7 +105,14 @@ public void add(@Nullable CompressionDictionary compressionDictionary)
 
         // Only update cache if not already in the cache
         DictId newDictId = compressionDictionary.dictId();
-        cache.get(newDictId, id -> compressionDictionary);
+        cache.get(newDictId, id -> {
+            Ref<?> ref = compressionDictionary.initRefLazily();
+            if (ref == null)
+            {
+                throw new IllegalStateException("Failed to acquire reference to compression dictionary");
+            }
+            return compressionDictionary;
+        });
 
         // Update current dictionary if we don't have one or the new one has a higher ID (newer)
         DictId currentId = currentDictId.get();

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

@@ -43,7 +43,7 @@ public class ZstdCompressionDictionary implements CompressionDictionary, SelfRef
     // 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 final AtomicReference<ZstdDictDecompress> dictDecompress = new AtomicReference<>();
-    private final Ref<ZstdCompressionDictionary> selfRef;
+    private volatile Ref<ZstdCompressionDictionary> selfRef;
 
     @VisibleForTesting
     public ZstdCompressionDictionary(DictId dictId, byte[] rawDictionary)
@@ -58,7 +58,7 @@ public ZstdCompressionDictionary(DictId dictId, byte[] rawDictionary, int checks
         this.dictId = dictId;
         this.rawDictionary = rawDictionary;
         this.checksum = checksum;
-        this.selfRef = new Ref<>(this, new Tidy(zstdDictCompressPerLevel, dictDecompress));
+        this.selfRef = null;
     }
 
     @Override
@@ -180,7 +180,7 @@ public ZstdDictDecompress dictionaryForDecompression()
     @Override
     public Ref<ZstdCompressionDictionary> tryRef()
     {
-        return selfRef.tryRef();
+        return initRefLazily().tryRef();
     }
 
     @Override
@@ -192,11 +192,31 @@ public Ref<ZstdCompressionDictionary> selfRef()
     @Override
     public Ref<ZstdCompressionDictionary> ref()
     {
-        return selfRef.ref();
+        return initRefLazily().ref();
+    }
+
+    @Override
+    public Ref<ZstdCompressionDictionary> initRefLazily()
+    {
+        if (selfRef == null)
+        {
+            synchronized (this)
+            {
+                if (selfRef == null)
+                {
+                    selfRef = new Ref<>(this, new Tidy(zstdDictCompressPerLevel, dictDecompress));
+                }
+            }
+        }
+        return selfRef;
     }
 
     private void ensureNotReleased()
     {
+        if (selfRef == null)
+            throw new IllegalStateException("Dictionary ref is not initialized. " +
+                                            "Call initRefLazily() or tryRef() first: " + dictId);
+
         if (selfRef.globalCount() <= 0)
             throw new IllegalStateException("Dictionary has been released: " + dictId);
     }

test/unit/org/apache/cassandra/db/compression/CompressionDictionaryCacheTest.java

@@ -32,6 +32,7 @@
 
 import com.github.luben.zstd.ZstdDictTrainer;
 import org.apache.cassandra.config.DatabaseDescriptor;
+import org.apache.cassandra.utils.concurrent.Ref;
 
 import static org.apache.cassandra.db.compression.CompressionDictionary.DictId;
 import static org.apache.cassandra.db.compression.CompressionDictionary.Kind;
@@ -427,7 +428,11 @@ private static void closeQuietly(ZstdCompressionDictionary resource)
         {
             try
             {
-                resource.selfRef().release();
+                Ref<ZstdCompressionDictionary> selfRef = resource.selfRef();
+                if (selfRef != null && selfRef.globalCount() > 0)
+                {
+                    selfRef.release();
+                }
             }
             catch (Exception e)
             {

test/unit/org/apache/cassandra/db/compression/ZstdCompressionDictionaryTest.java

@@ -61,6 +61,7 @@ public static void setUpClass()
     public void setUp()
     {
         dictionary = new ZstdCompressionDictionary(SAMPLE_DICT_ID, SAMPLE_DICT_DATA);
+        dictionary.initRefLazily();
     }
 
     @Test
@@ -81,9 +82,6 @@ public void testEqualsAndHashCode()
         assertThat(dictionary)
         .as("Dictionaries with different IDs should not be equal")
         .isNotEqualTo(differentIdDict);
-
-        dictionary2.selfRef().release();
-        differentIdDict.selfRef().release();
     }
 
     @Test
@@ -187,6 +185,8 @@ public void testDictionaryAccessWithoutReference()
             SAMPLE_DICT_DATA
         );
 
+        testDict.initRefLazily();
+
         // Access some dictionaries first to initialize them
         testDict.dictionaryForCompression(3);
         testDict.dictionaryForDecompression();
@@ -260,7 +260,6 @@ public void testMultipleReferences()
     @Test
     public void testReferenceAfterClose()
     {
-        // Release the self-reference
         dictionary.selfRef().release();
 
         assertThatThrownBy(() -> dictionary.ref())
@@ -386,9 +385,6 @@ public void testSerializeDeserializeWithManager() throws Exception
         .as("Both deserializations should return identical dictionary")
         .isNotNull()
         .isEqualTo(dict2);
-
-        dict1.selfRef().release();
-        dict2.selfRef().release();
     }
 
     @Test