Add attribute for debug purposes.

jpountz · Mar 14, 2024 · 27f44b5 · 27f44b5
1 parent cce0fe3
commit 27f44b5
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/codec/Elasticsearch814Codec.java b/server/src/main/java/org/elasticsearch/index/codec/Elasticsearch814Codec.java
@@ -23,7 +23,7 @@
 import org.elasticsearch.index.codec.zstd.Zstd814StoredFieldsFormat;
 
 /**
- * Elasticsearch codec as of 8.13. This extends the Lucene 9.9 codec to compressed stored fields with ZSTD instead of LZ4/DEFLATE. See
+ * Elasticsearch codec as of 8.14. This extends the Lucene 9.9 codec to compressed stored fields with ZSTD instead of LZ4/DEFLATE. See
  * {@link Zstd814StoredFieldsFormat}.
  */
 public class Elasticsearch814Codec extends FilterCodec {

diff --git a/server/src/main/java/org/elasticsearch/index/codec/zstd/Zstd814StoredFieldsFormat.java b/server/src/main/java/org/elasticsearch/index/codec/zstd/Zstd814StoredFieldsFormat.java
@@ -8,14 +8,18 @@
 
 package org.elasticsearch.index.codec.zstd;
 
+import org.apache.lucene.codecs.StoredFieldsWriter;
 import org.apache.lucene.codecs.compressing.CompressionMode;
 import org.apache.lucene.codecs.compressing.Compressor;
 import org.apache.lucene.codecs.compressing.Decompressor;
 import org.apache.lucene.codecs.lucene90.compressing.Lucene90CompressingStoredFieldsFormat;
 import org.apache.lucene.index.CorruptIndexException;
+import org.apache.lucene.index.SegmentInfo;
 import org.apache.lucene.store.ByteBuffersDataInput;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.store.DataOutput;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.IOContext;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.nativeaccess.CloseableByteBuffer;
@@ -39,6 +43,9 @@ public final class Zstd814StoredFieldsFormat extends Lucene90CompressingStoredFi
     private static final int BEST_SPEED_BLOCK_SIZE = (16 - 2) * 1_024;
     private static final int BEST_COMPRESSION_BLOCK_SIZE = (256 - 16) * 1_024;
 
+    /** Attribute key for compression mode. */
+    public static final String MODE_KEY = Zstd814StoredFieldsFormat.class.getSimpleName() + ".mode";
+
     public enum Mode {
         BEST_SPEED(0, BEST_SPEED_BLOCK_SIZE, 128),
         BEST_COMPRESSION(3, BEST_COMPRESSION_BLOCK_SIZE, 2048);
@@ -52,12 +59,23 @@ private Mode(int level, int blockSizeInBytes, int blockDocCount) {
         }
     }
 
+    private final Mode mode;
+
     public Zstd814StoredFieldsFormat(Mode mode) {
-        this(mode.level, mode.blockSizeInBytes, mode.blockDocCount);
+        super("ZstdStoredFields814", new ZstdCompressionMode(mode.level), mode.blockSizeInBytes, mode.blockDocCount, 10);
+        this.mode = mode;
     }
 
-    Zstd814StoredFieldsFormat(int level, int blockSizeInBytes, int blockDocCount) {
-        super("ZstdStoredFields814", new ZstdCompressionMode(level), blockSizeInBytes, blockDocCount, 10);
+    @Override
+    public StoredFieldsWriter fieldsWriter(Directory directory, SegmentInfo si, IOContext context) throws IOException {
+        // Both modes are compatible, we only put an attribute for debug purposes.
+        String previous = si.putAttribute(MODE_KEY, mode.name());
+        if (previous != null && previous.equals(mode.name()) == false) {
+            throw new IllegalStateException(
+                "found existing value for " + MODE_KEY + " for segment: " + si.name + "old=" + previous + ", new=" + mode.name()
+            );
+        }
+        return super.fieldsWriter(directory, si, context);
     }
 
     private static class ZstdCompressionMode extends CompressionMode {