From fd65fc2db4958e77f38a66907d7cd3b8de670cca Mon Sep 17 00:00:00 2001 From: Arpit Bandejiya Date: Tue, 24 Mar 2026 17:28:39 +0530 Subject: [PATCH 1/3] Initial commit for CatalogSnapshot Signed-off-by: Arpit Bandejiya --- .../exec/DefaultPlanExecutorTests.java | 10 +- .../index/engine/exec/CatalogSnapshot.java | 42 ++- .../exec/coord/CatalogSnapshotManager.java | 163 +++++++++ .../coord/DataformatAwareCatalogSnapshot.java | 231 +++++++++++++ .../coord/SegmentInfosCatalogSnapshot.java | 119 +++++++ .../index/engine/exec/coord/package-info.java | 14 + .../dataformat/DataFormatPluginTests.java | 134 ++++++++ .../coord/CatalogSnapshotManagerTests.java | 322 ++++++++++++++++++ .../DataformatAwareCatalogSnapshotTests.java | 284 +++++++++++++++ .../SegmentInfosCatalogSnapshotTests.java | 110 ++++++ 10 files changed, 1415 insertions(+), 14 deletions(-) create mode 100644 server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java create mode 100644 server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java create mode 100644 server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java create mode 100644 server/src/main/java/org/opensearch/index/engine/exec/coord/package-info.java create mode 100644 server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java create mode 100644 server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java create mode 100644 server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java index d623993227d95..c641f77f6a4b8 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java @@ -288,16 +288,18 @@ public String serializeToString() { } @Override - public void setCatalogSnapshotMap(Map map) {} - - @Override - public void setUserData(Map userData, boolean b) {} + public void setUserData(Map userData) {} @Override public Object getReader(DataFormat dataFormat) { return null; } + @Override + public MockCatalogSnapshot clone() { + return new MockCatalogSnapshot(generation, segments, format); + } + @Override protected void closeInternal() {} } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java index 80abcb59eccbe..b7b32a00905b4 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java @@ -10,6 +10,9 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.util.concurrent.AbstractRefCounted; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.index.engine.dataformat.DataFormat; import java.io.IOException; @@ -25,7 +28,7 @@ * Subclasses must implement methods for accessing file metadata, segments, and user data. */ @ExperimentalApi -public abstract class CatalogSnapshot extends AbstractRefCounted { +public abstract class CatalogSnapshot extends AbstractRefCounted implements Writeable, Cloneable { /** * Key for storing catalog snapshot in user data. @@ -49,6 +52,24 @@ public CatalogSnapshot(String name, long generation, long version) { this.version = version; } + /** + * Constructs a CatalogSnapshot from a {@link StreamInput}. + * + * @param in the stream input to read from + * @throws IOException if an I/O error occurs + */ + public CatalogSnapshot(StreamInput in) throws IOException { + super("catalog_snapshot"); + this.generation = in.readLong(); + this.version = in.readLong(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeLong(generation); + out.writeLong(version); + } + public long getGeneration() { return generation; } @@ -108,13 +129,6 @@ public long getVersion() { */ public abstract String serializeToString() throws IOException; - /** - * Sets the catalog snapshot map for tracking multiple snapshots. - * - * @param catalogSnapshotMap map of generation to catalog snapshots - */ - public abstract void setCatalogSnapshotMap(Map catalogSnapshotMap); - /** * Creates a clone without acquiring a reference count. * Used for Lucene compatibility where clone is required. @@ -131,9 +145,17 @@ public CatalogSnapshot cloneNoAcquire() { * Sets user-defined metadata for this catalog snapshot. * * @param userData map of user data key-value pairs - * @param b additional boolean parameter for implementation-specific behavior */ - public abstract void setUserData(Map userData, boolean b); + public abstract void setUserData(Map userData); + + /** + * Creates a deep copy of this catalog snapshot. The cloned snapshot starts with a fresh reference count + * of 1 (from {@link AbstractRefCounted}). + * Subclasses must ensure all mutable state is properly copied. + * + * @return a new {@link CatalogSnapshot} with the same logical state + */ + public abstract CatalogSnapshot clone(); public abstract Object getReader(DataFormat dataFormat); } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java new file mode 100644 index 0000000000000..d20f0880f8fbf --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java @@ -0,0 +1,163 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec.coord; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.engine.exec.CatalogSnapshot; + +import java.io.Closeable; +import java.util.Collections; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Supplier; + +/** + * Manages the lifecycle of {@link CatalogSnapshot} instances for the composite multi-format engine. + * + *

Tracks all live snapshots in a map keyed by generation. When a snapshot's reference count reaches + * zero (via {@link #decRefAndRemove}), it is automatically removed from the map. All {@code decRef} + * calls on managed snapshots go through this method to ensure consistent cleanup.

+ * + *

The write path (commit) is single-threaded (refresh is serialized per shard), while the read + * path (acquireSnapshot) is safe for concurrent access via volatile reads and {@code tryIncRef}.

+ */ +@ExperimentalApi +public class CatalogSnapshotManager implements Closeable { + + private volatile CatalogSnapshot latestCatalogSnapshot; + private final AtomicLong generation; + private final AtomicBoolean closed = new AtomicBoolean(false); + private final Map catalogSnapshotMap = new ConcurrentHashMap<>(); + + /** + * Constructs a new CatalogSnapshotManager. The supplier is invoked exactly once to produce the + * initial snapshot, which is then tracked in the live snapshot map. + * + * @param initialSnapshotSupplier supplier for the initial snapshot; must not be null or return null + */ + public CatalogSnapshotManager(Supplier initialSnapshotSupplier) { + Objects.requireNonNull(initialSnapshotSupplier, "initialSnapshotSupplier must not be null"); + CatalogSnapshot initialSnapshot = Objects.requireNonNull( + initialSnapshotSupplier.get(), + "initialSnapshotSupplier must not return null" + ); + this.latestCatalogSnapshot = initialSnapshot; + this.generation = new AtomicLong(initialSnapshot.getGeneration()); + catalogSnapshotMap.put(initialSnapshot.getGeneration(), initialSnapshot); + } + + /** + * Acquires the current snapshot with an incremented reference count, wrapped in a {@link ReleasableRef} + * that calls {@link #decRefAndRemove} on close. + * + * @return a {@link ReleasableRef} wrapping the current {@link CatalogSnapshot} + * @throws IllegalStateException if the manager or snapshot is already closed + */ + public ReleasableRef acquireSnapshot() { + if (closed.get()) { + throw new IllegalStateException("CatalogSnapshotManager is closed"); + } + final CatalogSnapshot snapshot = latestCatalogSnapshot; + if (snapshot.tryIncRef() == false) { + throw new IllegalStateException("CatalogSnapshot [gen=" + snapshot.getGeneration() + "] is already closed"); + } + return new ReleasableRef<>(snapshot) { + @Override + public void close() { + decRefAndRemove(snapshot); + } + }; + } + + /** + * Commits a new snapshot, replacing the current one. The old snapshot is decRef'd and removed + * from the map if its count reaches zero. + * + * @param newSnapshot the new catalog snapshot to commit + */ + public void commitNewSnapshot(CatalogSnapshot newSnapshot) { + assert closed.get() == false : "Cannot commit to a closed CatalogSnapshotManager"; + assert newSnapshot.getGeneration() > latestCatalogSnapshot.getGeneration() : "New snapshot generation must be greater than current"; + + catalogSnapshotMap.put(newSnapshot.getGeneration(), newSnapshot); + generation.set(newSnapshot.getGeneration()); + CatalogSnapshot oldSnapshot = latestCatalogSnapshot; + latestCatalogSnapshot = newSnapshot; + decRefAndRemove(oldSnapshot); + } + + /** + * Decrements the reference count and removes the snapshot from the tracking map if it reaches zero. + * Generation is captured before decRef to avoid accessing the snapshot after closeInternal. + */ + private void decRefAndRemove(CatalogSnapshot snapshot) { + final long gen = snapshot.getGeneration(); + if (snapshot.decRef()) { + catalogSnapshotMap.remove(gen); + } + } + + /** + * Returns an unmodifiable view of all live snapshots keyed by generation. + * + * @return unmodifiable map of generation to catalog snapshot + */ + public Map getCatalogSnapshotMap() { + return Collections.unmodifiableMap(catalogSnapshotMap); + } + + /** + * Returns the current generation counter value. + * + * @return the current generation + */ + public long getCurrentGeneration() { + return generation.get(); + } + + /** + * Returns the current snapshot. Note: this does not increment the reference count. + * Use {@link #acquireSnapshot()} for safe concurrent access. + * + * @return the current {@link CatalogSnapshot} + */ + public CatalogSnapshot getCurrentSnapshot() { + return latestCatalogSnapshot; + } + + /** + * Closes this manager. Idempotent. DecRefs the current snapshot and removes it if count reaches zero. + */ + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + decRefAndRemove(latestCatalogSnapshot); + } + } + + /** + * A generic reference wrapper for safe resource management via try-with-resources. + */ + @ExperimentalApi + public abstract static class ReleasableRef implements AutoCloseable { + + private final T ref; + + public ReleasableRef(T ref) { + this.ref = ref; + } + + public T getRef() { + return ref; + } + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java new file mode 100644 index 0000000000000..9d9770a73ce3e --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java @@ -0,0 +1,231 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec.coord; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.CatalogSnapshot; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Base64; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Concrete implementation of {@link CatalogSnapshot} for the composite multi-format engine. + * Holds segments grouped by data format, supports searchable file lookups across formats, + * and tracks snapshot metadata including user data and writer generation. + */ +@ExperimentalApi +public class DataformatAwareCatalogSnapshot extends CatalogSnapshot { + + private final long id; + private final List segments; + private final long lastWriterGeneration; + private Map userData; + + /** + * Constructs a new DataformatAwareCatalogSnapshot. + * + * @param id the unique snapshot identifier + * @param generation the monotonically increasing generation number + * @param version the schema version for serialization compatibility + * @param segments the list of segments in this snapshot + * @param lastWriterGeneration the generation of the last writer that contributed to this snapshot + * @param userData user-defined metadata key-value pairs + */ + public DataformatAwareCatalogSnapshot( + long id, + long generation, + long version, + List segments, + long lastWriterGeneration, + Map userData + ) { + super("dataformat_aware_catalog_snapshot", generation, version); + this.id = id; + this.segments = Collections.unmodifiableList(new ArrayList<>(segments)); + this.lastWriterGeneration = lastWriterGeneration; + this.userData = Map.copyOf(userData); + } + + /** + * Constructs a DataformatAwareCatalogSnapshot from a {@link StreamInput}. + * + * @param in the stream input to read from + * @throws IOException if an I/O error occurs + */ + public DataformatAwareCatalogSnapshot(StreamInput in) throws IOException { + super(in); + + // Read userData map + int userDataSize = in.readVInt(); + this.userData = new HashMap<>(); + for (int i = 0; i < userDataSize; i++) { + String key = in.readString(); + String value = in.readString(); + userData.put(key, value); + } + + this.id = in.readLong(); + this.lastWriterGeneration = in.readLong(); + + int segmentCount = in.readVInt(); + List segmentList = new ArrayList<>(segmentCount); + for (int i = 0; i < segmentCount; i++) { + long segGeneration = in.readLong(); + int formatCount = in.readVInt(); + Map dfGrouped = new HashMap<>(formatCount); + for (int j = 0; j < formatCount; j++) { + String formatName = in.readString(); + String directory = in.readString(); + long writerGeneration = in.readLong(); + List fileList = in.readStringList(); + long numRows = in.readLong(); + dfGrouped.put(formatName, new WriterFileSet(directory, writerGeneration, new HashSet<>(fileList), numRows)); + } + segmentList.add(new Segment(segGeneration, dfGrouped)); + } + this.segments = Collections.unmodifiableList(segmentList); + } + + @Override + public long getId() { + return id; + } + + @Override + public List getSegments() { + return segments; + } + + @Override + public Collection getSearchableFiles(String dataFormat) { + List result = new ArrayList<>(); + for (Segment segment : segments) { + WriterFileSet writerFileSet = segment.dfGroupedSearchableFiles().get(dataFormat); + if (writerFileSet != null) { + result.add(writerFileSet); + } + } + return result; + } + + @Override + public Set getDataFormats() { + Set formats = new HashSet<>(); + for (Segment segment : segments) { + formats.addAll(segment.dfGroupedSearchableFiles().keySet()); + } + return formats; + } + + @Override + public long getLastWriterGeneration() { + return lastWriterGeneration; + } + + @Override + public Map getUserData() { + return userData; + } + + @Override + public void setUserData(Map userData) { + this.userData = Map.copyOf(userData); + } + + @Override + public String serializeToString() throws IOException { + try (BytesStreamOutput out = new BytesStreamOutput()) { + this.writeTo(out); + return Base64.getEncoder().encodeToString(out.bytes().toBytesRef().bytes); + } + } + + /** + * Deserializes a {@link DataformatAwareCatalogSnapshot} from a Base64-encoded binary string. + * + * @param serializedData the Base64 string produced by {@link #serializeToString()} + * @return a reconstructed {@link DataformatAwareCatalogSnapshot} + * @throws IOException if the data is malformed or missing required fields + */ + public static DataformatAwareCatalogSnapshot deserializeFromString(String serializedData) throws IOException { + try { + byte[] bytes = Base64.getDecoder().decode(serializedData); + try (BytesStreamInput in = new BytesStreamInput(bytes)) { + return new DataformatAwareCatalogSnapshot(in); + } + } catch (IOException e) { + throw e; + } catch (Exception e) { + throw new IOException("Failed to deserialize DataformatAwareCatalogSnapshot: " + e.getMessage(), e); + } + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + + // Write userData map + if (userData == null) { + out.writeVInt(0); + } else { + out.writeVInt(userData.size()); + for (Map.Entry entry : userData.entrySet()) { + out.writeString(entry.getKey()); + out.writeString(entry.getValue()); + } + } + + out.writeLong(id); + out.writeLong(lastWriterGeneration); + + out.writeVInt(segments.size()); + for (Segment seg : segments) { + out.writeLong(seg.generation()); + out.writeVInt(seg.dfGroupedSearchableFiles().size()); + for (Map.Entry dfEntry : seg.dfGroupedSearchableFiles().entrySet()) { + out.writeString(dfEntry.getKey()); + WriterFileSet wfs = dfEntry.getValue(); + out.writeString(wfs.directory()); + out.writeLong(wfs.writerGeneration()); + out.writeStringCollection(wfs.files()); + out.writeLong(wfs.numRows()); + } + } + } + + @Override + public DataformatAwareCatalogSnapshot clone() { + return new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData); + } + + @Override + protected void closeInternal() { + // Subclass-specific resource cleanup. Map removal is handled by CatalogSnapshotManager.decRefAndRemove. + } + + @Override + public Object getReader(DataFormat dataFormat) { + throw new UnsupportedOperationException("Not implemented"); + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java new file mode 100644 index 0000000000000..28dde74c8a817 --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java @@ -0,0 +1,119 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec.coord; + +import org.apache.lucene.index.SegmentInfos; +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.CatalogSnapshot; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; + +import java.io.IOException; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * A thin adapter that wraps Lucene's {@link SegmentInfos} as a {@link CatalogSnapshot}. + * Used by {@code InternalEngine} (the standard single-format Lucene engine) to participate + * in the {@link CatalogSnapshot} abstraction without requiring composite engine infrastructure. + * + *

Multi-format methods ({@link #getSegments()}, {@link #getSearchableFiles(String)}, + * {@link #getDataFormats()}, {@link #serializeToString()}) throw {@link UnsupportedOperationException} + * since Lucene-only engines do not use composite segments.

+ */ +@ExperimentalApi +public class SegmentInfosCatalogSnapshot extends CatalogSnapshot { + + private static final String CATALOG_SNAPSHOT_KEY = "_segment_infos_catalog_snapshot_"; + + private final SegmentInfos segmentInfos; + + /** + * Constructs a new SegmentInfosCatalogSnapshot wrapping the given SegmentInfos. + * + * @param segmentInfos the Lucene SegmentInfos to wrap + */ + public SegmentInfosCatalogSnapshot(SegmentInfos segmentInfos) { + super(CATALOG_SNAPSHOT_KEY + segmentInfos.getGeneration(), segmentInfos.getGeneration(), segmentInfos.getVersion()); + this.segmentInfos = segmentInfos; + } + + /** + * Returns the wrapped Lucene SegmentInfos instance. + * + * @return the SegmentInfos + */ + public SegmentInfos getSegmentInfos() { + return segmentInfos; + } + + @Override + public long getId() { + return generation; + } + + @Override + public Map getUserData() { + return segmentInfos.getUserData(); + } + + @Override + public long getLastWriterGeneration() { + return -1; + } + + @Override + public List getSegments() { + throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support getSegments()"); + } + + @Override + public Collection getSearchableFiles(String dataFormat) { + throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support getSearchableFiles()"); + } + + @Override + public Set getDataFormats() { + throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support getDataFormats()"); + } + + @Override + public String serializeToString() throws IOException { + throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support serializeToString()"); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support writeTo()"); + } + + @Override + public void setUserData(Map userData) { + // No-op for SegmentInfosCatalogSnapshot + } + + @Override + public Object getReader(DataFormat dataFormat) { + throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support getReader()"); + } + + @Override + protected void closeInternal() { + // No resources to release for SegmentInfos wrapper. + } + + @Override + public SegmentInfosCatalogSnapshot clone() { + return new SegmentInfosCatalogSnapshot(segmentInfos); + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/package-info.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/package-info.java new file mode 100644 index 0000000000000..53ae20e9b9aef --- /dev/null +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/package-info.java @@ -0,0 +1,14 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/** + * Coordination layer for the composite multi-format engine. + * Contains the CatalogSnapshotManager and related snapshot implementations + * for managing immutable point-in-time views of index data across multiple data formats. + */ +package org.opensearch.index.engine.exec.coord; diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java b/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java index c6d8a01d448e8..a15fdc7997d3d 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java @@ -407,4 +407,138 @@ public void testFileLifecycleNotifications() throws IOException { assertEquals(1, rm.deletedFiles.size()); assertTrue(rm.deletedFiles.contains("a.parquet")); } + + static class MockReader { + final List fileNames; + final long totalRows; + boolean closed; + + MockReader(List fileNames, long totalRows) { + this.fileNames = fileNames; + this.totalRows = totalRows; + } + + void close() { + closed = true; + } + } + + static class MockReaderManager implements EngineReaderManager { + private final String formatName; + private final Map readers = new HashMap<>(); + final List addedFiles = new ArrayList<>(); + final List deletedFiles = new ArrayList<>(); + + MockReaderManager(String formatName) { + this.formatName = formatName; + } + + @Override + public MockReader getReader(CatalogSnapshot snapshot) { + return readers.get(snapshot); + } + + int readerCount() { + return readers.size(); + } + + @Override + public void beforeRefresh() {} + + @Override + public void afterRefresh(boolean didRefresh, CatalogSnapshot snapshot) { + if (didRefresh == false || readers.containsKey(snapshot)) return; + Collection files = snapshot.getSearchableFiles(formatName); + List allFiles = new ArrayList<>(); + long totalRows = 0; + for (WriterFileSet wfs : files) { + allFiles.addAll(wfs.files()); + totalRows += wfs.numRows(); + } + readers.put(snapshot, new MockReader(allFiles, totalRows)); + } + + @Override + public void onDeleted(CatalogSnapshot snapshot) { + MockReader reader = readers.remove(snapshot); + if (reader != null) reader.close(); + } + + @Override + public void onFilesDeleted(Collection files) { + deletedFiles.addAll(files); + } + + @Override + public void onFilesAdded(Collection files) { + addedFiles.addAll(files); + } + } + + static class MockCatalogSnapshot extends CatalogSnapshot { + private final List segments; + private final MockDataFormat format; + + MockCatalogSnapshot(long generation, List segments, MockDataFormat format) { + super("mock-snapshot", generation, 1L); + this.segments = segments; + this.format = format; + } + + @Override + public Map getUserData() { + return Map.of(); + } + + @Override + public long getId() { + return generation; + } + + @Override + public List getSegments() { + return segments; + } + + @Override + public Collection getSearchableFiles(String dataFormat) { + List result = new ArrayList<>(); + for (Segment seg : segments) { + WriterFileSet wfs = seg.dfGroupedSearchableFiles().get(dataFormat); + if (wfs != null) result.add(wfs); + } + return result; + } + + @Override + public Set getDataFormats() { + return Set.of(format.name()); + } + + @Override + public long getLastWriterGeneration() { + return generation; + } + + @Override + public String serializeToString() { + return "mock-snapshot-" + generation; + } + + @Override + public void setUserData(Map userData) {} + + @Override + public Object getReader(DataFormat dataFormat) { + return null; + } + + @Override + public MockCatalogSnapshot clone() { + return new MockCatalogSnapshot(generation, segments, format); + } + + @Override + protected void closeInternal() {} + } } diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java new file mode 100644 index 0000000000000..0a8ad9e18fb53 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java @@ -0,0 +1,322 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec.coord; + +import org.opensearch.index.engine.exec.CatalogSnapshot; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Property-based tests for {@link CatalogSnapshotManager}. + */ +public class CatalogSnapshotManagerTests extends OpenSearchTestCase { + + private WriterFileSet randomWriterFileSet() { + String directory = "/tmp/" + randomAlphaOfLength(8); + long writerGeneration = randomNonNegativeLong(); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom("cfs", "si", "parquet", "dat")); + } + return new WriterFileSet(directory, writerGeneration, files, randomIntBetween(0, 10000)); + } + + private Segment randomSegment() { + long generation = randomNonNegativeLong(); + int formatCount = randomIntBetween(1, 4); + Map dfGrouped = new HashMap<>(); + for (int i = 0; i < formatCount; i++) { + dfGrouped.put(randomFrom("lucene", "parquet", "arrow", "custom_" + randomAlphaOfLength(3)), randomWriterFileSet()); + } + return new Segment(generation, dfGrouped); + } + + private List randomSegments() { + int segmentCount = randomIntBetween(0, 5); + List segments = new ArrayList<>(); + for (int i = 0; i < segmentCount; i++) { + segments.add(randomSegment()); + } + return segments; + } + + private DataformatAwareCatalogSnapshot buildSnapshot( + long generation, + List segments, + long lastWriterGeneration, + Map userData + ) { + return new DataformatAwareCatalogSnapshot(generation, generation, 0L, segments, lastWriterGeneration, userData); + } + + private CatalogSnapshotManager createRandomManager() { + return new CatalogSnapshotManager( + () -> new DataformatAwareCatalogSnapshot( + randomNonNegativeLong(), + randomIntBetween(0, 100), + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + Map.of() + ) + ); + } + + public void testCommitProducesCorrectNewSnapshot() { + for (int iter = 0; iter < 100; iter++) { + CatalogSnapshotManager manager = createRandomManager(); + try { + long previousGeneration = manager.getCurrentGeneration(); + Set seenIds = new HashSet<>(); + seenIds.add(manager.getCurrentSnapshot().getId()); + + int numCommits = randomIntBetween(1, 10); + for (int c = 0; c < numCommits; c++) { + List newSegments = randomSegments(); + long newWriterGeneration = randomNonNegativeLong(); + long newGeneration = previousGeneration + 1; + + manager.commitNewSnapshot(buildSnapshot(newGeneration, newSegments, newWriterGeneration, Map.of())); + + assertEquals(previousGeneration + 1, manager.getCurrentGeneration()); + assertTrue(seenIds.add(manager.getCurrentSnapshot().getId())); + assertEquals(newSegments, manager.getCurrentSnapshot().getSegments()); + assertTrue(manager.getCatalogSnapshotMap().containsKey(newGeneration)); + + previousGeneration = manager.getCurrentGeneration(); + } + } finally { + manager.close(); + } + } + } + + public void testUserDataPreservationOnCommit() { + for (int iter = 0; iter < 100; iter++) { + int initialEntries = randomIntBetween(1, 5); + Map initialUserData = new HashMap<>(); + for (int i = 0; i < initialEntries; i++) { + initialUserData.put("init_" + randomAlphaOfLength(4), randomAlphaOfLength(8)); + } + long initGen = randomIntBetween(0, 100); + CatalogSnapshotManager manager = new CatalogSnapshotManager( + () -> new DataformatAwareCatalogSnapshot( + randomNonNegativeLong(), + initGen, + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + initialUserData + ) + ); + try { + long gen1 = initGen + 1; + manager.commitNewSnapshot( + buildSnapshot(gen1, randomSegments(), randomNonNegativeLong(), manager.getCurrentSnapshot().getUserData()) + ); + assertEquals(initialUserData, manager.getCurrentSnapshot().getUserData()); + + Map newUserData = new HashMap<>(); + for (int i = 0; i < randomIntBetween(1, 5); i++) { + newUserData.put("new_" + randomAlphaOfLength(4), randomAlphaOfLength(8)); + } + long gen2 = gen1 + 1; + manager.commitNewSnapshot(buildSnapshot(gen2, randomSegments(), randomNonNegativeLong(), newUserData)); + assertEquals(newUserData, manager.getCurrentSnapshot().getUserData()); + } finally { + manager.close(); + } + } + } + + public void testReferenceCountingLifecycle() { + for (int iter = 0; iter < 100; iter++) { + AtomicBoolean initialCloseInternalCalled = new AtomicBoolean(false); + long initGen = randomIntBetween(0, 100); + CatalogSnapshotManager manager = new CatalogSnapshotManager( + () -> new TrackableSnapshot( + randomNonNegativeLong(), + initGen, + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + Collections.emptyMap(), + initialCloseInternalCalled + ) + ); + + CatalogSnapshot initialSnapshot = manager.getCurrentSnapshot(); + assertEquals(1, initialSnapshot.refCount()); + + long gen1 = initGen + 1; + manager.commitNewSnapshot(buildSnapshot(gen1, randomSegments(), randomNonNegativeLong(), Map.of())); + + assertEquals(0, initialSnapshot.refCount()); + assertTrue(initialCloseInternalCalled.get()); + assertFalse(initialSnapshot.tryIncRef()); + + int numCommits = randomIntBetween(1, 8); + for (int c = 0; c < numCommits; c++) { + CatalogSnapshot prev = manager.getCurrentSnapshot(); + assertEquals(1, prev.refCount()); + long nextGen = manager.getCurrentGeneration() + 1; + manager.commitNewSnapshot(buildSnapshot(nextGen, randomSegments(), randomNonNegativeLong(), Map.of())); + assertEquals(0, prev.refCount()); + } + + CatalogSnapshot finalSnapshot = manager.getCurrentSnapshot(); + assertEquals(1, finalSnapshot.refCount()); + manager.close(); + assertEquals(0, finalSnapshot.refCount()); + } + } + + public void testAcquireAndReleaseViaReleasableRef() throws Exception { + for (int iter = 0; iter < 100; iter++) { + CatalogSnapshotManager manager = createRandomManager(); + try { + CatalogSnapshot currentSnap = manager.getCurrentSnapshot(); + assertEquals(1, currentSnap.refCount()); + + int numAcquires = randomIntBetween(1, 5); + List> refs = new ArrayList<>(); + for (int a = 0; a < numAcquires; a++) { + refs.add(manager.acquireSnapshot()); + assertEquals(1 + (a + 1), currentSnap.refCount()); + } + + for (int r = 0; r < numAcquires; r++) { + refs.get(r).close(); + assertEquals(1 + numAcquires - r - 1, currentSnap.refCount()); + } + assertEquals(1, currentSnap.refCount()); + + CatalogSnapshotManager.ReleasableRef heldRef = manager.acquireSnapshot(); + CatalogSnapshot heldSnapshot = heldRef.getRef(); + assertEquals(2, heldSnapshot.refCount()); + + long nextGen = manager.getCurrentGeneration() + 1; + manager.commitNewSnapshot(buildSnapshot(nextGen, randomSegments(), randomNonNegativeLong(), Map.of())); + assertEquals(1, heldSnapshot.refCount()); + + heldRef.close(); + assertEquals(0, heldSnapshot.refCount()); + } finally { + manager.close(); + } + } + } + + public void testClosedManagerRejectsAcquisition() { + for (int iter = 0; iter < 100; iter++) { + CatalogSnapshotManager manager = createRandomManager(); + int numCommits = randomIntBetween(0, 5); + for (int c = 0; c < numCommits; c++) { + long nextGen = manager.getCurrentGeneration() + 1; + manager.commitNewSnapshot(buildSnapshot(nextGen, randomSegments(), randomNonNegativeLong(), Map.of())); + } + manager.close(); + expectThrows(IllegalStateException.class, manager::acquireSnapshot); + } + } + + public void testInitialSnapshotRecovery() throws Exception { + for (int iter = 0; iter < 100; iter++) { + long id = randomNonNegativeLong(); + long generation = randomIntBetween(0, 100); + long version = randomNonNegativeLong(); + long lastWriterGeneration = randomNonNegativeLong(); + List segments = randomIntBetween(1, 5) == 1 ? Collections.emptyList() : randomSegments(); + Map userData = new HashMap<>(); + for (int i = 0; i < randomIntBetween(0, 4); i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); + } + + CatalogSnapshotManager manager = new CatalogSnapshotManager( + () -> new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData) + ); + CatalogSnapshotManager.ReleasableRef ref = null; + try { + ref = manager.acquireSnapshot(); + CatalogSnapshot acquired = ref.getRef(); + assertEquals(id, acquired.getId()); + assertEquals(generation, acquired.getGeneration()); + assertEquals(segments, acquired.getSegments()); + assertEquals(userData, acquired.getUserData()); + assertEquals(lastWriterGeneration, acquired.getLastWriterGeneration()); + assertSame(acquired, manager.getCurrentSnapshot()); + } finally { + if (ref != null) ref.close(); + manager.close(); + } + } + } + + public void testDecRefAndRemoveFromMap() throws Exception { + for (int iter = 0; iter < 100; iter++) { + CatalogSnapshotManager manager = createRandomManager(); + long initGen = manager.getCurrentGeneration(); + assertEquals(1, manager.getCatalogSnapshotMap().size()); + + long gen1 = initGen + 1; + manager.commitNewSnapshot(buildSnapshot(gen1, randomSegments(), randomNonNegativeLong(), Map.of())); + // Old snapshot had no extra refs, so decRef brought it to 0 and it was removed + assertEquals(1, manager.getCatalogSnapshotMap().size()); + assertFalse(manager.getCatalogSnapshotMap().containsKey(initGen)); + assertTrue(manager.getCatalogSnapshotMap().containsKey(gen1)); + + // Acquire a ref, commit, then the old snapshot stays in the map until the ref is released + CatalogSnapshotManager.ReleasableRef ref = manager.acquireSnapshot(); + long gen2 = gen1 + 1; + manager.commitNewSnapshot(buildSnapshot(gen2, randomSegments(), randomNonNegativeLong(), Map.of())); + assertEquals(2, manager.getCatalogSnapshotMap().size()); + assertTrue(manager.getCatalogSnapshotMap().containsKey(gen1)); + + ref.close(); + assertEquals(1, manager.getCatalogSnapshotMap().size()); + assertFalse(manager.getCatalogSnapshotMap().containsKey(gen1)); + + manager.close(); + assertEquals(0, manager.getCatalogSnapshotMap().size()); + } + } + + private static class TrackableSnapshot extends DataformatAwareCatalogSnapshot { + private final AtomicBoolean closeInternalCalled; + + TrackableSnapshot( + long id, + long generation, + long version, + List segments, + long lastWriterGeneration, + Map userData, + AtomicBoolean closeInternalCalled + ) { + super(id, generation, version, segments, lastWriterGeneration, userData); + this.closeInternalCalled = closeInternalCalled; + } + + @Override + protected void closeInternal() { + closeInternalCalled.set(true); + } + } +} diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java new file mode 100644 index 0000000000000..f077b48e421e0 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java @@ -0,0 +1,284 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec.coord; + +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Property-based tests for {@link DataformatAwareCatalogSnapshot}. + * Uses OpenSearch randomization utilities to generate random inputs across many iterations. + */ +public class DataformatAwareCatalogSnapshotTests extends OpenSearchTestCase { + + // Feature: catalog-snapshot-manager, Property 1: Snapshot field access consistency + + /** + * Generates a random {@link WriterFileSet} with random directory, writer generation, files, and row count. + */ + private WriterFileSet randomWriterFileSet() { + String directory = "/tmp/" + randomAlphaOfLength(8); + long writerGeneration = randomNonNegativeLong(); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom("cfs", "si", "parquet", "dat")); + } + long numRows = randomIntBetween(0, 10000); + return new WriterFileSet(directory, writerGeneration, files, numRows); + } + + /** + * Generates a random {@link Segment} with a random generation and random data format keys. + */ + private Segment randomSegment() { + long generation = randomNonNegativeLong(); + int formatCount = randomIntBetween(1, 4); + Map dfGrouped = new HashMap<>(); + for (int i = 0; i < formatCount; i++) { + String formatKey = randomFrom("lucene", "parquet", "arrow", "custom_" + randomAlphaOfLength(3)); + dfGrouped.put(formatKey, randomWriterFileSet()); + } + return new Segment(generation, dfGrouped); + } + + /** + * Generates a random {@link DataformatAwareCatalogSnapshot} with random fields. + */ + private DataformatAwareCatalogSnapshot randomSnapshot() { + long id = randomLong(); + long generation = randomNonNegativeLong(); + long version = randomNonNegativeLong(); + int segmentCount = randomIntBetween(0, 5); + List segments = new ArrayList<>(); + for (int i = 0; i < segmentCount; i++) { + segments.add(randomSegment()); + } + long lastWriterGeneration = randomNonNegativeLong(); + int userDataEntries = randomIntBetween(0, 4); + Map userData = new HashMap<>(); + for (int i = 0; i < userDataEntries; i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); + } + return new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData); + } + + /** + * Property 1: Snapshot field access consistency. + * For any valid combination of inputs, constructing a DataformatAwareCatalogSnapshot and querying + * its accessors should return consistent results. + * + * Validates: Requirements 1.1, 1.2, 1.3, 1.4, 1.6 + */ + public void testSnapshotFieldAccessConsistency() { + for (int iter = 0; iter < 100; iter++) { + long id = randomLong(); + long generation = randomNonNegativeLong(); + long version = randomNonNegativeLong(); + int segmentCount = randomIntBetween(0, 5); + List segments = new ArrayList<>(); + for (int i = 0; i < segmentCount; i++) { + segments.add(randomSegment()); + } + long lastWriterGeneration = randomNonNegativeLong(); + int userDataEntries = randomIntBetween(0, 4); + Map userData = new HashMap<>(); + for (int i = 0; i < userDataEntries; i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); + } + + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot( + id, + generation, + version, + segments, + lastWriterGeneration, + userData + ); + + try { + // Verify getId() + assertEquals("getId() should return the snapshot id", id, snapshot.getId()); + + // Verify getGeneration() and getVersion() from parent + assertEquals("getGeneration() should return the generation", generation, snapshot.getGeneration()); + assertEquals("getVersion() should return the version", version, snapshot.getVersion()); + + // Verify getSegments() returns equal content + assertEquals("getSegments() should return segments equal to input", segments, snapshot.getSegments()); + + // Verify getLastWriterGeneration() + assertEquals( + "getLastWriterGeneration() should return the writer generation", + lastWriterGeneration, + snapshot.getLastWriterGeneration() + ); + + // Verify getUserData() + assertEquals("getUserData() should return the user data", userData, snapshot.getUserData()); + + // Verify getSearchableFiles() returns exactly the WriterFileSets matching the queried format + Set expectedFormats = new HashSet<>(); + for (Segment seg : segments) { + expectedFormats.addAll(seg.dfGroupedSearchableFiles().keySet()); + } + for (String format : expectedFormats) { + Collection searchableFiles = snapshot.getSearchableFiles(format); + List expected = new ArrayList<>(); + for (Segment seg : segments) { + WriterFileSet wfs = seg.dfGroupedSearchableFiles().get(format); + if (wfs != null) { + expected.add(wfs); + } + } + assertEquals( + "getSearchableFiles('" + format + "') should return matching WriterFileSets", + expected, + new ArrayList<>(searchableFiles) + ); + } + + // Verify getSearchableFiles() for a non-existent format returns empty + Collection emptyResult = snapshot.getSearchableFiles("nonexistent_format_" + randomAlphaOfLength(5)); + assertTrue("getSearchableFiles for unknown format should be empty", emptyResult.isEmpty()); + + // Verify getDataFormats() returns the union of all format keys + assertEquals("getDataFormats() should return union of all format keys", expectedFormats, snapshot.getDataFormats()); + + // Verify getSegments() returns an unmodifiable list + expectThrows(UnsupportedOperationException.class, () -> snapshot.getSegments().add(randomSegment())); + } finally { + snapshot.decRef(); + } + } + } + + // Feature: catalog-snapshot-manager, Property 2: Serialization round-trip + + /** + * Property 2: Serialization round-trip. + * For any valid DataformatAwareCatalogSnapshot, serializing via serializeToString(), + * then deserializing via deserializeFromString(), then serializing again should produce + * a binary-encoded string that, when deserialized, yields an equivalent snapshot. + * + * Validates: Requirements 2.1, 2.2, 2.3 + */ + public void testSerializationRoundTrip() throws Exception { + for (int iter = 0; iter < 100; iter++) { + DataformatAwareCatalogSnapshot original = randomSnapshot(); + try { + // First serialization + String json1 = original.serializeToString(); + + // Deserialize from the first JSON + DataformatAwareCatalogSnapshot deserialized1 = DataformatAwareCatalogSnapshot.deserializeFromString(json1); + try { + // Second serialization from the deserialized snapshot + String json2 = deserialized1.serializeToString(); + + // Deserialize from the second JSON to verify round-trip equivalence + DataformatAwareCatalogSnapshot deserialized2 = DataformatAwareCatalogSnapshot.deserializeFromString(json2); + try { + // Verify deserialized1 fields match the original + assertSnapshotFieldsEqual("first deserialization", original, deserialized1); + + // Verify deserialized2 fields match deserialized1 (round-trip stability) + assertSnapshotFieldsEqual("second deserialization", deserialized1, deserialized2); + } finally { + deserialized2.decRef(); + } + } finally { + deserialized1.decRef(); + } + } finally { + original.decRef(); + } + } + } + + /** + * Asserts that two {@link DataformatAwareCatalogSnapshot} instances have equivalent fields. + */ + private void assertSnapshotFieldsEqual(String context, DataformatAwareCatalogSnapshot expected, DataformatAwareCatalogSnapshot actual) { + assertEquals(context + ": id should match", expected.getId(), actual.getId()); + assertEquals(context + ": generation should match", expected.getGeneration(), actual.getGeneration()); + assertEquals(context + ": version should match", expected.getVersion(), actual.getVersion()); + assertEquals(context + ": segments should match", expected.getSegments(), actual.getSegments()); + assertEquals(context + ": lastWriterGeneration should match", expected.getLastWriterGeneration(), actual.getLastWriterGeneration()); + assertEquals(context + ": userData should match", expected.getUserData(), actual.getUserData()); + } + + // Feature: catalog-snapshot-manager, Property 3: Deserialization rejects invalid input + + /** + * Property 3: Deserialization rejects invalid input. + * For any string that is not valid Base64 or is truncated/corrupted binary data, + * calling deserializeFromString() should throw an IOException. + * + * Validates: Requirements 2.4 + */ + public void testDeserializationRejectsInvalidInput() { + for (int iter = 0; iter < 100; iter++) { + String input = generateInvalidInput(iter); + expectThrows(IOException.class, () -> DataformatAwareCatalogSnapshot.deserializeFromString(input)); + } + } + + /** + * Generates an invalid input string for deserialization testing. + * Mixes different categories of invalid input across iterations. + */ + private String generateInvalidInput(int iter) { + int category = iter % 6; + switch (category) { + case 0: + // Completely random alphanumeric strings (not valid Base64 payload) + return randomAlphaOfLengthBetween(1, 200); + case 1: + // Valid Base64 but random bytes (not a valid serialized snapshot) + byte[] randomBytes = new byte[randomIntBetween(1, 100)]; + random().nextBytes(randomBytes); + return java.util.Base64.getEncoder().encodeToString(randomBytes); + case 2: + // Truncated valid Base64: serialize a valid snapshot, then truncate the Base64 string + DataformatAwareCatalogSnapshot snap = randomSnapshot(); + try { + String validBase64 = snap.serializeToString(); + int truncateAt = randomIntBetween(1, Math.max(1, validBase64.length() / 2)); + return validBase64.substring(0, truncateAt); + } catch (IOException e) { + return "AAAA"; + } finally { + snap.decRef(); + } + case 3: + // Empty string + return ""; + case 4: + // Strings with invalid Base64 characters + return randomFrom("not-base64!!!", "===", "@@@@", "hello world", "{\"json\":true}"); + case 5: + // Null-like strings + return randomFrom("null", "undefined", "None", "nil", "NaN"); + default: + return randomAlphaOfLength(10); + } + } + +} diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java new file mode 100644 index 0000000000000..7b469f68c57fd --- /dev/null +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java @@ -0,0 +1,110 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec.coord; + +import org.apache.lucene.index.SegmentInfos; +import org.apache.lucene.util.Version; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.HashMap; +import java.util.Map; + +/** + * Property-based tests for {@link SegmentInfosCatalogSnapshot}. + * Uses OpenSearch randomization utilities to generate random inputs across many iterations. + */ +public class SegmentInfosCatalogSnapshotTests extends OpenSearchTestCase { + + // Feature: catalog-snapshot-manager, Property 10: SegmentInfosCatalogSnapshot delegates to SegmentInfos + + /** + * Creates a random {@link SegmentInfos} instance with random user data. + */ + private SegmentInfos randomSegmentInfos() { + SegmentInfos segmentInfos = new SegmentInfos(Version.LATEST.major); + int userDataEntries = randomIntBetween(0, 5); + Map userData = new HashMap<>(); + for (int i = 0; i < userDataEntries; i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); + } + segmentInfos.setUserData(userData, false); + return segmentInfos; + } + + /** + * Property 10: SegmentInfosCatalogSnapshot delegates to SegmentInfos. + * For any valid Lucene SegmentInfos instance, constructing a SegmentInfosCatalogSnapshot + * and querying its accessors should return consistent results delegated from SegmentInfos. + * + * Validates: Requirements 7.1, 7.2, 7.3, 7.4 + */ + public void testSegmentInfosDelegation() { + for (int iter = 0; iter < 100; iter++) { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); + + try { + // Verify getId() returns the SegmentInfos generation + assertEquals("getId() should return the SegmentInfos generation", segmentInfos.getGeneration(), snapshot.getId()); + + // Verify getGeneration() returns the SegmentInfos generation + assertEquals( + "getGeneration() should return the SegmentInfos generation", + segmentInfos.getGeneration(), + snapshot.getGeneration() + ); + + // Verify getVersion() returns the SegmentInfos version + assertEquals("getVersion() should return the SegmentInfos version", segmentInfos.getVersion(), snapshot.getVersion()); + + // Verify getUserData() returns the SegmentInfos user data + assertEquals("getUserData() should return the SegmentInfos user data", segmentInfos.getUserData(), snapshot.getUserData()); + + // Verify getLastWriterGeneration() returns -1 + assertEquals("getLastWriterGeneration() should return -1", -1L, snapshot.getLastWriterGeneration()); + + // Verify getSegmentInfos() returns the wrapped instance + assertSame("getSegmentInfos() should return the wrapped SegmentInfos", segmentInfos, snapshot.getSegmentInfos()); + + // Verify getSegments() throws UnsupportedOperationException + expectThrows(UnsupportedOperationException.class, () -> snapshot.getSegments()); + + // Verify getSearchableFiles() throws UnsupportedOperationException + String randomFormat = randomAlphaOfLength(5); + expectThrows(UnsupportedOperationException.class, () -> snapshot.getSearchableFiles(randomFormat)); + + // Verify getDataFormats() throws UnsupportedOperationException + expectThrows(UnsupportedOperationException.class, () -> snapshot.getDataFormats()); + + // Verify serializeToString() throws UnsupportedOperationException + expectThrows(UnsupportedOperationException.class, () -> snapshot.serializeToString()); + + // Verify writeTo() throws UnsupportedOperationException + expectThrows(UnsupportedOperationException.class, () -> snapshot.writeTo(new BytesStreamOutput())); + + // Verify clone() returns a new instance wrapping the same SegmentInfos + SegmentInfosCatalogSnapshot cloned = snapshot.clone(); + try { + assertNotSame("clone() should return a new instance", snapshot, cloned); + assertSame("clone() should wrap the same SegmentInfos", segmentInfos, cloned.getSegmentInfos()); + assertEquals("clone() should have the same generation", snapshot.getGeneration(), cloned.getGeneration()); + assertEquals("clone() should have the same version", snapshot.getVersion(), cloned.getVersion()); + } finally { + cloned.decRef(); + } + + // Verify setUserData() is a no-op (does not throw) + snapshot.setUserData(Map.of("key", "value")); + } finally { + snapshot.decRef(); + } + } + } +} From a176b5e9bbefafe032cbd80d7607e683f95ba0b4 Mon Sep 17 00:00:00 2001 From: Arpit Bandejiya Date: Tue, 31 Mar 2026 14:58:44 +0530 Subject: [PATCH 2/3] Address comments Signed-off-by: Arpit Bandejiya --- CHANGELOG.md | 1 + .../index/engine/exec/CatalogSnapshot.java | 5 + .../opensearch/index/engine/exec/Segment.java | 24 +- .../index/engine/exec/WriterFileSet.java | 22 +- .../exec/coord/CatalogSnapshotManager.java | 65 +--- .../coord/DataformatAwareCatalogSnapshot.java | 54 +--- .../coord/SegmentInfosCatalogSnapshot.java | 29 +- .../index/engine/exec/SegmentTests.java | 69 ++++ .../index/engine/exec/WriterFileSetTests.java | 40 +++ .../coord/CatalogSnapshotManagerTests.java | 208 +++++------- .../DataformatAwareCatalogSnapshotTests.java | 296 ++++++------------ .../SegmentInfosCatalogSnapshotTests.java | 123 +++----- 12 files changed, 445 insertions(+), 491 deletions(-) create mode 100644 server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java create mode 100644 server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 356bc0f8cd93d..fe0ecf102f6b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add support for enabling pluggable data formats, starting with phase-1 of decoupling shard from engine, and introducing basic abstractions ([#20675](https://github.com/opensearch-project/OpenSearch/pull/20675)) - Add concurrent queue in libs and composite engine sandbox plugin ([#20909](https://github.com/opensearch-project/OpenSearch/pull/20909)) - Add interface for the Multi format merge flow ([#20908](https://github.com/opensearch-project/OpenSearch/pull/20908)) +- Add CatalogSnapshotManager lifecycle management with reference-counted snapshot tracking and serialization support for Segment and WriterFileSet ([#20982](https://github.com/opensearch-project/OpenSearch/pull/20982)) - Add warmup phase to wait for lag to catch up in pull-based ingestion before serving ([#20526](https://github.com/opensearch-project/OpenSearch/pull/20526)) - Add a new static method to IndicesOptions API to expose `STRICT_EXPAND_OPEN_HIDDEN_FORBID_CLOSED` index option ([#20980](https://github.com/opensearch-project/OpenSearch/pull/20980)) diff --git a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java index b7b32a00905b4..81c3e756ec977 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java @@ -26,6 +26,11 @@ * Maintains versioned information about segments, files, and metadata for index operations. * Extends AbstractRefCounted to support reference counting for safe concurrent access. * Subclasses must implement methods for accessing file metadata, segments, and user data. + * + *

Important: Do not call {@code incRef()}, {@code decRef()}, or {@code tryIncRef()} directly. + * Use {@link org.opensearch.index.engine.exec.coord.CatalogSnapshotManager#acquireSnapshot()} to obtain + * a reference-counted handle, and close the returned {@link org.opensearch.common.concurrent.GatedCloseable} + * when done. The manager handles all reference counting internally.

*/ @ExperimentalApi public abstract class CatalogSnapshot extends AbstractRefCounted implements Writeable, Cloneable { diff --git a/server/src/main/java/org/opensearch/index/engine/exec/Segment.java b/server/src/main/java/org/opensearch/index/engine/exec/Segment.java index 5a5811804aff3..974b607e3ea92 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/Segment.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/Segment.java @@ -9,23 +9,43 @@ package org.opensearch.index.engine.exec; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.index.engine.dataformat.DataFormat; +import java.io.IOException; import java.util.HashMap; import java.util.Map; /** * Represents a segment in the catalog snapshot containing files grouped by data format. * Each segment has a unique generation number and maintains searchable files organized by their data format type. - * This class is serializable and can be transmitted across nodes for replication and recovery operations. */ @ExperimentalApi -public record Segment(long generation, Map dfGroupedSearchableFiles) { +public record Segment(long generation, Map dfGroupedSearchableFiles) implements Writeable { public Segment { dfGroupedSearchableFiles = Map.copyOf(dfGroupedSearchableFiles); } + /** + * Constructs a Segment by deserializing from a {@link StreamInput}. + */ + public Segment(StreamInput in) throws IOException { + this(in.readLong(), in.readMap(StreamInput::readString, WriterFileSet::new)); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeLong(generation); + out.writeVInt(dfGroupedSearchableFiles.size()); + for (Map.Entry entry : dfGroupedSearchableFiles.entrySet()) { + out.writeString(entry.getKey()); + entry.getValue().writeTo(out); + } + } + public static Builder builder(long generation) { return new Builder(generation); } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java b/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java index 65fe34410e9d0..d36ddbfba2262 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java @@ -9,6 +9,9 @@ package org.opensearch.index.engine.exec; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; import java.io.IOException; import java.nio.file.Path; @@ -20,12 +23,19 @@ * Groups files by directory and writer generation, tracking metadata such as row count and total size. */ @ExperimentalApi -public record WriterFileSet(String directory, long writerGeneration, Set files, long numRows) { +public record WriterFileSet(String directory, long writerGeneration, Set files, long numRows) implements Writeable { public WriterFileSet { files = Set.copyOf(files); } + /** + * Constructs a WriterFileSet by deserializing from a {@link StreamInput}. + */ + public WriterFileSet(StreamInput in) throws IOException { + this(in.readString(), in.readLong(), new HashSet<>(in.readStringList()), in.readLong()); + } + public long getTotalSize() { return files.stream().mapToLong(file -> { try { @@ -41,6 +51,16 @@ public String toString() { return "WriterFileSet{" + "directory=" + directory + ", writerGeneration=" + writerGeneration + ", files=" + files + '}'; } + /** + * Serializes this WriterFileSet to the given stream output. + */ + public void writeTo(StreamOutput out) throws IOException { + out.writeString(directory); + out.writeLong(writerGeneration); + out.writeStringCollection(files); + out.writeLong(numRows); + } + /** * Creates a new builder for constructing WriterFileSet instances. * diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java index d20f0880f8fbf..a23a03cf6eebb 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java @@ -9,15 +9,14 @@ package org.opensearch.index.engine.exec.coord; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.index.engine.exec.CatalogSnapshot; import java.io.Closeable; -import java.util.Collections; import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.concurrent.atomic.AtomicLong; import java.util.function.Supplier; /** @@ -34,7 +33,6 @@ public class CatalogSnapshotManager implements Closeable { private volatile CatalogSnapshot latestCatalogSnapshot; - private final AtomicLong generation; private final AtomicBoolean closed = new AtomicBoolean(false); private final Map catalogSnapshotMap = new ConcurrentHashMap<>(); @@ -51,18 +49,21 @@ public CatalogSnapshotManager(Supplier initialSnapshotSupplier) "initialSnapshotSupplier must not return null" ); this.latestCatalogSnapshot = initialSnapshot; - this.generation = new AtomicLong(initialSnapshot.getGeneration()); - catalogSnapshotMap.put(initialSnapshot.getGeneration(), initialSnapshot); + if (catalogSnapshotMap.putIfAbsent(initialSnapshot.getGeneration(), initialSnapshot) != null) { + throw new IllegalStateException( + "Duplicate snapshot generation [" + initialSnapshot.getGeneration() + "] in catalog snapshot map" + ); + } } /** - * Acquires the current snapshot with an incremented reference count, wrapped in a {@link ReleasableRef} + * Acquires the current snapshot with an incremented reference count, wrapped in a {@link GatedCloseable} * that calls {@link #decRefAndRemove} on close. * - * @return a {@link ReleasableRef} wrapping the current {@link CatalogSnapshot} + * @return a {@link GatedCloseable} wrapping the current {@link CatalogSnapshot} * @throws IllegalStateException if the manager or snapshot is already closed */ - public ReleasableRef acquireSnapshot() { + public GatedCloseable acquireSnapshot() { if (closed.get()) { throw new IllegalStateException("CatalogSnapshotManager is closed"); } @@ -70,12 +71,7 @@ public ReleasableRef acquireSnapshot() { if (snapshot.tryIncRef() == false) { throw new IllegalStateException("CatalogSnapshot [gen=" + snapshot.getGeneration() + "] is already closed"); } - return new ReleasableRef<>(snapshot) { - @Override - public void close() { - decRefAndRemove(snapshot); - } - }; + return new GatedCloseable<>(snapshot, () -> decRefAndRemove(snapshot)); } /** @@ -84,12 +80,13 @@ public void close() { * * @param newSnapshot the new catalog snapshot to commit */ - public void commitNewSnapshot(CatalogSnapshot newSnapshot) { + public synchronized void commitNewSnapshot(CatalogSnapshot newSnapshot) { assert closed.get() == false : "Cannot commit to a closed CatalogSnapshotManager"; assert newSnapshot.getGeneration() > latestCatalogSnapshot.getGeneration() : "New snapshot generation must be greater than current"; - catalogSnapshotMap.put(newSnapshot.getGeneration(), newSnapshot); - generation.set(newSnapshot.getGeneration()); + if (catalogSnapshotMap.putIfAbsent(newSnapshot.getGeneration(), newSnapshot) != null) { + throw new IllegalStateException("Duplicate snapshot generation [" + newSnapshot.getGeneration() + "] in catalog snapshot map"); + } CatalogSnapshot oldSnapshot = latestCatalogSnapshot; latestCatalogSnapshot = newSnapshot; decRefAndRemove(oldSnapshot); @@ -106,24 +103,6 @@ private void decRefAndRemove(CatalogSnapshot snapshot) { } } - /** - * Returns an unmodifiable view of all live snapshots keyed by generation. - * - * @return unmodifiable map of generation to catalog snapshot - */ - public Map getCatalogSnapshotMap() { - return Collections.unmodifiableMap(catalogSnapshotMap); - } - - /** - * Returns the current generation counter value. - * - * @return the current generation - */ - public long getCurrentGeneration() { - return generation.get(); - } - /** * Returns the current snapshot. Note: this does not increment the reference count. * Use {@link #acquireSnapshot()} for safe concurrent access. @@ -144,20 +123,4 @@ public void close() { } } - /** - * A generic reference wrapper for safe resource management via try-with-resources. - */ - @ExperimentalApi - public abstract static class ReleasableRef implements AutoCloseable { - - private final T ref; - - public ReleasableRef(T ref) { - this.ref = ref; - } - - public T getRef() { - return ref; - } - } } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java index 9d9770a73ce3e..c52cbb3c590fe 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java @@ -10,6 +10,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.common.io.stream.BytesStreamInput; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; @@ -23,7 +24,6 @@ import java.util.Base64; import java.util.Collection; import java.util.Collections; -import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -76,14 +76,7 @@ public DataformatAwareCatalogSnapshot( public DataformatAwareCatalogSnapshot(StreamInput in) throws IOException { super(in); - // Read userData map - int userDataSize = in.readVInt(); - this.userData = new HashMap<>(); - for (int i = 0; i < userDataSize; i++) { - String key = in.readString(); - String value = in.readString(); - userData.put(key, value); - } + this.userData = in.readMap(StreamInput::readString, StreamInput::readString); this.id = in.readLong(); this.lastWriterGeneration = in.readLong(); @@ -91,18 +84,7 @@ public DataformatAwareCatalogSnapshot(StreamInput in) throws IOException { int segmentCount = in.readVInt(); List segmentList = new ArrayList<>(segmentCount); for (int i = 0; i < segmentCount; i++) { - long segGeneration = in.readLong(); - int formatCount = in.readVInt(); - Map dfGrouped = new HashMap<>(formatCount); - for (int j = 0; j < formatCount; j++) { - String formatName = in.readString(); - String directory = in.readString(); - long writerGeneration = in.readLong(); - List fileList = in.readStringList(); - long numRows = in.readLong(); - dfGrouped.put(formatName, new WriterFileSet(directory, writerGeneration, new HashSet<>(fileList), numRows)); - } - segmentList.add(new Segment(segGeneration, dfGrouped)); + segmentList.add(new Segment(in)); } this.segments = Collections.unmodifiableList(segmentList); } @@ -157,7 +139,7 @@ public void setUserData(Map userData) { public String serializeToString() throws IOException { try (BytesStreamOutput out = new BytesStreamOutput()) { this.writeTo(out); - return Base64.getEncoder().encodeToString(out.bytes().toBytesRef().bytes); + return Base64.getEncoder().encodeToString(BytesReference.toBytes(out.bytes())); } } @@ -169,6 +151,9 @@ public String serializeToString() throws IOException { * @throws IOException if the data is malformed or missing required fields */ public static DataformatAwareCatalogSnapshot deserializeFromString(String serializedData) throws IOException { + if (serializedData == null || serializedData.isEmpty()) { + throw new IOException("Cannot deserialize DataformatAwareCatalogSnapshot: input is null or empty"); + } try { byte[] bytes = Base64.getDecoder().decode(serializedData); try (BytesStreamInput in = new BytesStreamInput(bytes)) { @@ -184,33 +169,12 @@ public static DataformatAwareCatalogSnapshot deserializeFromString(String serial @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); - - // Write userData map - if (userData == null) { - out.writeVInt(0); - } else { - out.writeVInt(userData.size()); - for (Map.Entry entry : userData.entrySet()) { - out.writeString(entry.getKey()); - out.writeString(entry.getValue()); - } - } - + out.writeMap(userData, StreamOutput::writeString, StreamOutput::writeString); out.writeLong(id); out.writeLong(lastWriterGeneration); - out.writeVInt(segments.size()); for (Segment seg : segments) { - out.writeLong(seg.generation()); - out.writeVInt(seg.dfGroupedSearchableFiles().size()); - for (Map.Entry dfEntry : seg.dfGroupedSearchableFiles().entrySet()) { - out.writeString(dfEntry.getKey()); - WriterFileSet wfs = dfEntry.getValue(); - out.writeString(wfs.directory()); - out.writeLong(wfs.writerGeneration()); - out.writeStringCollection(wfs.files()); - out.writeLong(wfs.numRows()); - } + seg.writeTo(out); } } diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java index 28dde74c8a817..83642edb10872 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java @@ -9,7 +9,12 @@ package org.opensearch.index.engine.exec.coord; import org.apache.lucene.index.SegmentInfos; +import org.apache.lucene.store.BufferedChecksumIndexInput; +import org.apache.lucene.store.ByteBuffersDataOutput; +import org.apache.lucene.store.ByteBuffersIndexOutput; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.common.lucene.store.ByteArrayIndexInput; +import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.exec.CatalogSnapshot; @@ -48,6 +53,23 @@ public SegmentInfosCatalogSnapshot(SegmentInfos segmentInfos) { this.segmentInfos = segmentInfos; } + /** + * Constructs a SegmentInfosCatalogSnapshot from a {@link StreamInput} by deserializing the + * SegmentInfos binary representation. + * + * @param in the stream input to read from + * @throws IOException if an I/O error occurs + */ + public SegmentInfosCatalogSnapshot(StreamInput in) throws IOException { + super(in); + byte[] segmentInfosBytes = in.readByteArray(); + this.segmentInfos = SegmentInfos.readCommit( + null, + new BufferedChecksumIndexInput(new ByteArrayIndexInput("SegmentInfos", segmentInfosBytes)), + 0L + ); + } + /** * Returns the wrapped Lucene SegmentInfos instance. * @@ -94,7 +116,12 @@ public String serializeToString() throws IOException { @Override public void writeTo(StreamOutput out) throws IOException { - throw new UnsupportedOperationException("SegmentInfosCatalogSnapshot does not support writeTo()"); + super.writeTo(out); + ByteBuffersDataOutput buffer = new ByteBuffersDataOutput(); + try (ByteBuffersIndexOutput indexOutput = new ByteBuffersIndexOutput(buffer, "", null)) { + segmentInfos.write(indexOutput); + } + out.writeByteArray(buffer.toArrayCopy()); } @Override diff --git a/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java b/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java new file mode 100644 index 0000000000000..5a0a82ac30457 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java @@ -0,0 +1,69 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec; + +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +/** + * Tests for {@link Segment} serialization. + */ +public class SegmentTests extends OpenSearchTestCase { + + public void testCopyWriteable() throws Exception { + Segment original = randomSegment(); + Segment copy = copyWriteable(original, new NamedWriteableRegistry(Collections.emptyList()), Segment::new); + assertEquals(original, copy); + } + + public void testCopyWriteableEmpty() throws Exception { + Segment empty = new Segment(0L, Map.of()); + Segment copy = copyWriteable(empty, new NamedWriteableRegistry(Collections.emptyList()), Segment::new); + assertEquals(empty, copy); + } + + public void testCopyWriteableMultiFormat() throws Exception { + Map dfGrouped = new HashMap<>(); + dfGrouped.put("lucene", randomWriterFileSet("lucene")); + dfGrouped.put("parquet", randomWriterFileSet("parquet")); + Segment original = new Segment(randomNonNegativeLong(), dfGrouped); + + Segment copy = copyWriteable(original, new NamedWriteableRegistry(Collections.emptyList()), Segment::new); + assertEquals(original, copy); + assertEquals(2, copy.dfGroupedSearchableFiles().size()); + } + + // --- helpers --- + + private WriterFileSet randomWriterFileSet(String format) { + String directory = "/tmp/" + randomAlphaOfLength(8); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + String[] extensions = "lucene".equals(format) ? new String[] { "cfs", "si", "dat" } : new String[] { "parquet" }; + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom(extensions)); + } + return new WriterFileSet(directory, randomNonNegativeLong(), files, randomIntBetween(0, 10000)); + } + + private Segment randomSegment() { + Map dfGrouped = new HashMap<>(); + for (int i = 0; i < randomIntBetween(1, 3); i++) { + String format = randomFrom("lucene", "parquet"); + dfGrouped.put(format, randomWriterFileSet(format)); + } + return new Segment(randomNonNegativeLong(), dfGrouped); + } +} diff --git a/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java b/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java new file mode 100644 index 0000000000000..1a34b530a1032 --- /dev/null +++ b/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java @@ -0,0 +1,40 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.index.engine.exec; + +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; + +/** + * Tests for {@link WriterFileSet}. + */ +public class WriterFileSetTests extends OpenSearchTestCase { + + public void testCopyWriteable() throws Exception { + WriterFileSet original = randomWriterFileSet(); + WriterFileSet copy = copyWriteable(original, new NamedWriteableRegistry(Collections.emptyList()), WriterFileSet::new); + assertEquals(original, copy); + } + + // --- helpers --- + + private WriterFileSet randomWriterFileSet() { + String directory = "/tmp/" + randomAlphaOfLength(8); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom("cfs", "si", "dat", "parquet")); + } + return new WriterFileSet(directory, randomNonNegativeLong(), files, randomIntBetween(0, 10000)); + } +} diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java index 0a8ad9e18fb53..6a17af1c5d8a5 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java @@ -8,6 +8,7 @@ package org.opensearch.index.engine.exec.coord; +import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; @@ -23,84 +24,30 @@ import java.util.concurrent.atomic.AtomicBoolean; /** - * Property-based tests for {@link CatalogSnapshotManager}. + * Tests for {@link CatalogSnapshotManager}. */ public class CatalogSnapshotManagerTests extends OpenSearchTestCase { - private WriterFileSet randomWriterFileSet() { - String directory = "/tmp/" + randomAlphaOfLength(8); - long writerGeneration = randomNonNegativeLong(); - int fileCount = randomIntBetween(1, 5); - Set files = new HashSet<>(); - for (int i = 0; i < fileCount; i++) { - files.add(randomAlphaOfLength(6) + "." + randomFrom("cfs", "si", "parquet", "dat")); - } - return new WriterFileSet(directory, writerGeneration, files, randomIntBetween(0, 10000)); - } - - private Segment randomSegment() { - long generation = randomNonNegativeLong(); - int formatCount = randomIntBetween(1, 4); - Map dfGrouped = new HashMap<>(); - for (int i = 0; i < formatCount; i++) { - dfGrouped.put(randomFrom("lucene", "parquet", "arrow", "custom_" + randomAlphaOfLength(3)), randomWriterFileSet()); - } - return new Segment(generation, dfGrouped); - } - - private List randomSegments() { - int segmentCount = randomIntBetween(0, 5); - List segments = new ArrayList<>(); - for (int i = 0; i < segmentCount; i++) { - segments.add(randomSegment()); - } - return segments; - } - - private DataformatAwareCatalogSnapshot buildSnapshot( - long generation, - List segments, - long lastWriterGeneration, - Map userData - ) { - return new DataformatAwareCatalogSnapshot(generation, generation, 0L, segments, lastWriterGeneration, userData); - } - - private CatalogSnapshotManager createRandomManager() { - return new CatalogSnapshotManager( - () -> new DataformatAwareCatalogSnapshot( - randomNonNegativeLong(), - randomIntBetween(0, 100), - randomNonNegativeLong(), - randomSegments(), - randomNonNegativeLong(), - Map.of() - ) - ); - } - public void testCommitProducesCorrectNewSnapshot() { for (int iter = 0; iter < 100; iter++) { CatalogSnapshotManager manager = createRandomManager(); try { - long previousGeneration = manager.getCurrentGeneration(); + long previousGeneration = manager.getCurrentSnapshot().getGeneration(); Set seenIds = new HashSet<>(); seenIds.add(manager.getCurrentSnapshot().getId()); int numCommits = randomIntBetween(1, 10); for (int c = 0; c < numCommits; c++) { List newSegments = randomSegments(); - long newWriterGeneration = randomNonNegativeLong(); long newGeneration = previousGeneration + 1; - manager.commitNewSnapshot(buildSnapshot(newGeneration, newSegments, newWriterGeneration, Map.of())); + manager.commitNewSnapshot(buildSnapshot(newGeneration, newSegments, randomNonNegativeLong(), Map.of())); - assertEquals(previousGeneration + 1, manager.getCurrentGeneration()); + assertEquals(previousGeneration + 1, manager.getCurrentSnapshot().getGeneration()); assertTrue(seenIds.add(manager.getCurrentSnapshot().getId())); assertEquals(newSegments, manager.getCurrentSnapshot().getSegments()); - assertTrue(manager.getCatalogSnapshotMap().containsKey(newGeneration)); - previousGeneration = manager.getCurrentGeneration(); + previousGeneration = manager.getCurrentSnapshot().getGeneration(); } } finally { manager.close(); @@ -110,11 +57,7 @@ public void testCommitProducesCorrectNewSnapshot() { public void testUserDataPreservationOnCommit() { for (int iter = 0; iter < 100; iter++) { - int initialEntries = randomIntBetween(1, 5); - Map initialUserData = new HashMap<>(); - for (int i = 0; i < initialEntries; i++) { - initialUserData.put("init_" + randomAlphaOfLength(4), randomAlphaOfLength(8)); - } + Map initialUserData = randomUserData(randomIntBetween(1, 5)); long initGen = randomIntBetween(0, 100); CatalogSnapshotManager manager = new CatalogSnapshotManager( () -> new DataformatAwareCatalogSnapshot( @@ -133,10 +76,7 @@ public void testUserDataPreservationOnCommit() { ); assertEquals(initialUserData, manager.getCurrentSnapshot().getUserData()); - Map newUserData = new HashMap<>(); - for (int i = 0; i < randomIntBetween(1, 5); i++) { - newUserData.put("new_" + randomAlphaOfLength(4), randomAlphaOfLength(8)); - } + Map newUserData = randomUserData(randomIntBetween(1, 5)); long gen2 = gen1 + 1; manager.commitNewSnapshot(buildSnapshot(gen2, randomSegments(), randomNonNegativeLong(), newUserData)); assertEquals(newUserData, manager.getCurrentSnapshot().getUserData()); @@ -148,7 +88,7 @@ public void testUserDataPreservationOnCommit() { public void testReferenceCountingLifecycle() { for (int iter = 0; iter < 100; iter++) { - AtomicBoolean initialCloseInternalCalled = new AtomicBoolean(false); + AtomicBoolean closeInternalCalled = new AtomicBoolean(false); long initGen = randomIntBetween(0, 100); CatalogSnapshotManager manager = new CatalogSnapshotManager( () -> new TrackableSnapshot( @@ -158,26 +98,24 @@ public void testReferenceCountingLifecycle() { randomSegments(), randomNonNegativeLong(), Collections.emptyMap(), - initialCloseInternalCalled + closeInternalCalled ) ); CatalogSnapshot initialSnapshot = manager.getCurrentSnapshot(); assertEquals(1, initialSnapshot.refCount()); - long gen1 = initGen + 1; - manager.commitNewSnapshot(buildSnapshot(gen1, randomSegments(), randomNonNegativeLong(), Map.of())); - + manager.commitNewSnapshot(buildSnapshot(initGen + 1, randomSegments(), randomNonNegativeLong(), Map.of())); assertEquals(0, initialSnapshot.refCount()); - assertTrue(initialCloseInternalCalled.get()); - assertFalse(initialSnapshot.tryIncRef()); + assertTrue(closeInternalCalled.get()); int numCommits = randomIntBetween(1, 8); for (int c = 0; c < numCommits; c++) { CatalogSnapshot prev = manager.getCurrentSnapshot(); assertEquals(1, prev.refCount()); - long nextGen = manager.getCurrentGeneration() + 1; - manager.commitNewSnapshot(buildSnapshot(nextGen, randomSegments(), randomNonNegativeLong(), Map.of())); + manager.commitNewSnapshot( + buildSnapshot(manager.getCurrentSnapshot().getGeneration() + 1, randomSegments(), randomNonNegativeLong(), Map.of()) + ); assertEquals(0, prev.refCount()); } @@ -188,7 +126,7 @@ public void testReferenceCountingLifecycle() { } } - public void testAcquireAndReleaseViaReleasableRef() throws Exception { + public void testAcquireAndReleaseViaGatedCloseable() throws Exception { for (int iter = 0; iter < 100; iter++) { CatalogSnapshotManager manager = createRandomManager(); try { @@ -196,24 +134,24 @@ public void testAcquireAndReleaseViaReleasableRef() throws Exception { assertEquals(1, currentSnap.refCount()); int numAcquires = randomIntBetween(1, 5); - List> refs = new ArrayList<>(); + List> refs = new ArrayList<>(); for (int a = 0; a < numAcquires; a++) { refs.add(manager.acquireSnapshot()); assertEquals(1 + (a + 1), currentSnap.refCount()); } - for (int r = 0; r < numAcquires; r++) { refs.get(r).close(); assertEquals(1 + numAcquires - r - 1, currentSnap.refCount()); } assertEquals(1, currentSnap.refCount()); - CatalogSnapshotManager.ReleasableRef heldRef = manager.acquireSnapshot(); - CatalogSnapshot heldSnapshot = heldRef.getRef(); + GatedCloseable heldRef = manager.acquireSnapshot(); + CatalogSnapshot heldSnapshot = heldRef.get(); assertEquals(2, heldSnapshot.refCount()); - long nextGen = manager.getCurrentGeneration() + 1; - manager.commitNewSnapshot(buildSnapshot(nextGen, randomSegments(), randomNonNegativeLong(), Map.of())); + manager.commitNewSnapshot( + buildSnapshot(manager.getCurrentSnapshot().getGeneration() + 1, randomSegments(), randomNonNegativeLong(), Map.of()) + ); assertEquals(1, heldSnapshot.refCount()); heldRef.close(); @@ -227,10 +165,10 @@ public void testAcquireAndReleaseViaReleasableRef() throws Exception { public void testClosedManagerRejectsAcquisition() { for (int iter = 0; iter < 100; iter++) { CatalogSnapshotManager manager = createRandomManager(); - int numCommits = randomIntBetween(0, 5); - for (int c = 0; c < numCommits; c++) { - long nextGen = manager.getCurrentGeneration() + 1; - manager.commitNewSnapshot(buildSnapshot(nextGen, randomSegments(), randomNonNegativeLong(), Map.of())); + for (int c = 0; c < randomIntBetween(0, 5); c++) { + manager.commitNewSnapshot( + buildSnapshot(manager.getCurrentSnapshot().getGeneration() + 1, randomSegments(), randomNonNegativeLong(), Map.of()) + ); } manager.close(); expectThrows(IllegalStateException.class, manager::acquireSnapshot); @@ -244,18 +182,13 @@ public void testInitialSnapshotRecovery() throws Exception { long version = randomNonNegativeLong(); long lastWriterGeneration = randomNonNegativeLong(); List segments = randomIntBetween(1, 5) == 1 ? Collections.emptyList() : randomSegments(); - Map userData = new HashMap<>(); - for (int i = 0; i < randomIntBetween(0, 4); i++) { - userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); - } + Map userData = randomUserData(randomIntBetween(0, 4)); CatalogSnapshotManager manager = new CatalogSnapshotManager( () -> new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData) ); - CatalogSnapshotManager.ReleasableRef ref = null; - try { - ref = manager.acquireSnapshot(); - CatalogSnapshot acquired = ref.getRef(); + try (GatedCloseable ref = manager.acquireSnapshot()) { + CatalogSnapshot acquired = ref.get(); assertEquals(id, acquired.getId()); assertEquals(generation, acquired.getGeneration()); assertEquals(segments, acquired.getSegments()); @@ -263,39 +196,64 @@ public void testInitialSnapshotRecovery() throws Exception { assertEquals(lastWriterGeneration, acquired.getLastWriterGeneration()); assertSame(acquired, manager.getCurrentSnapshot()); } finally { - if (ref != null) ref.close(); manager.close(); } } } - public void testDecRefAndRemoveFromMap() throws Exception { - for (int iter = 0; iter < 100; iter++) { - CatalogSnapshotManager manager = createRandomManager(); - long initGen = manager.getCurrentGeneration(); - assertEquals(1, manager.getCatalogSnapshotMap().size()); - - long gen1 = initGen + 1; - manager.commitNewSnapshot(buildSnapshot(gen1, randomSegments(), randomNonNegativeLong(), Map.of())); - // Old snapshot had no extra refs, so decRef brought it to 0 and it was removed - assertEquals(1, manager.getCatalogSnapshotMap().size()); - assertFalse(manager.getCatalogSnapshotMap().containsKey(initGen)); - assertTrue(manager.getCatalogSnapshotMap().containsKey(gen1)); - - // Acquire a ref, commit, then the old snapshot stays in the map until the ref is released - CatalogSnapshotManager.ReleasableRef ref = manager.acquireSnapshot(); - long gen2 = gen1 + 1; - manager.commitNewSnapshot(buildSnapshot(gen2, randomSegments(), randomNonNegativeLong(), Map.of())); - assertEquals(2, manager.getCatalogSnapshotMap().size()); - assertTrue(manager.getCatalogSnapshotMap().containsKey(gen1)); - - ref.close(); - assertEquals(1, manager.getCatalogSnapshotMap().size()); - assertFalse(manager.getCatalogSnapshotMap().containsKey(gen1)); + // --- helpers --- - manager.close(); - assertEquals(0, manager.getCatalogSnapshotMap().size()); + private WriterFileSet randomWriterFileSet(String format) { + String directory = "/tmp/" + randomAlphaOfLength(8); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + String[] extensions = "lucene".equals(format) ? new String[] { "cfs", "si", "dat" } : new String[] { "parquet" }; + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom(extensions)); + } + return new WriterFileSet(directory, randomNonNegativeLong(), files, randomIntBetween(0, 10000)); + } + + private Segment randomSegment() { + Map dfGrouped = new HashMap<>(); + for (int i = 0; i < randomIntBetween(1, 2); i++) { + String format = randomFrom("lucene", "parquet"); + dfGrouped.put(format, randomWriterFileSet(format)); + } + return new Segment(randomNonNegativeLong(), dfGrouped); + } + + private List randomSegments() { + List segments = new ArrayList<>(); + for (int i = 0; i < randomIntBetween(0, 5); i++) { + segments.add(randomSegment()); + } + return segments; + } + + private Map randomUserData(int entries) { + Map userData = new HashMap<>(); + for (int i = 0; i < entries; i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); } + return userData; + } + + private DataformatAwareCatalogSnapshot buildSnapshot(long gen, List segments, long writerGen, Map userData) { + return new DataformatAwareCatalogSnapshot(gen, gen, 0L, segments, writerGen, userData); + } + + private CatalogSnapshotManager createRandomManager() { + return new CatalogSnapshotManager( + () -> new DataformatAwareCatalogSnapshot( + randomNonNegativeLong(), + randomIntBetween(0, 100), + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + Map.of() + ) + ); } private static class TrackableSnapshot extends DataformatAwareCatalogSnapshot { @@ -303,14 +261,14 @@ private static class TrackableSnapshot extends DataformatAwareCatalogSnapshot { TrackableSnapshot( long id, - long generation, + long gen, long version, List segments, - long lastWriterGeneration, + long writerGen, Map userData, AtomicBoolean closeInternalCalled ) { - super(id, generation, version, segments, lastWriterGeneration, userData); + super(id, gen, version, segments, writerGen, userData); this.closeInternalCalled = closeInternalCalled; } diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java index f077b48e421e0..1e866cea75d91 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java @@ -8,13 +8,14 @@ package org.opensearch.index.engine.exec.coord; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; import java.util.ArrayList; -import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -22,86 +23,18 @@ import java.util.Set; /** - * Property-based tests for {@link DataformatAwareCatalogSnapshot}. - * Uses OpenSearch randomization utilities to generate random inputs across many iterations. + * Tests for {@link DataformatAwareCatalogSnapshot}. */ public class DataformatAwareCatalogSnapshotTests extends OpenSearchTestCase { - // Feature: catalog-snapshot-manager, Property 1: Snapshot field access consistency - - /** - * Generates a random {@link WriterFileSet} with random directory, writer generation, files, and row count. - */ - private WriterFileSet randomWriterFileSet() { - String directory = "/tmp/" + randomAlphaOfLength(8); - long writerGeneration = randomNonNegativeLong(); - int fileCount = randomIntBetween(1, 5); - Set files = new HashSet<>(); - for (int i = 0; i < fileCount; i++) { - files.add(randomAlphaOfLength(6) + "." + randomFrom("cfs", "si", "parquet", "dat")); - } - long numRows = randomIntBetween(0, 10000); - return new WriterFileSet(directory, writerGeneration, files, numRows); - } - - /** - * Generates a random {@link Segment} with a random generation and random data format keys. - */ - private Segment randomSegment() { - long generation = randomNonNegativeLong(); - int formatCount = randomIntBetween(1, 4); - Map dfGrouped = new HashMap<>(); - for (int i = 0; i < formatCount; i++) { - String formatKey = randomFrom("lucene", "parquet", "arrow", "custom_" + randomAlphaOfLength(3)); - dfGrouped.put(formatKey, randomWriterFileSet()); - } - return new Segment(generation, dfGrouped); - } - - /** - * Generates a random {@link DataformatAwareCatalogSnapshot} with random fields. - */ - private DataformatAwareCatalogSnapshot randomSnapshot() { - long id = randomLong(); - long generation = randomNonNegativeLong(); - long version = randomNonNegativeLong(); - int segmentCount = randomIntBetween(0, 5); - List segments = new ArrayList<>(); - for (int i = 0; i < segmentCount; i++) { - segments.add(randomSegment()); - } - long lastWriterGeneration = randomNonNegativeLong(); - int userDataEntries = randomIntBetween(0, 4); - Map userData = new HashMap<>(); - for (int i = 0; i < userDataEntries; i++) { - userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); - } - return new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData); - } - - /** - * Property 1: Snapshot field access consistency. - * For any valid combination of inputs, constructing a DataformatAwareCatalogSnapshot and querying - * its accessors should return consistent results. - * - * Validates: Requirements 1.1, 1.2, 1.3, 1.4, 1.6 - */ public void testSnapshotFieldAccessConsistency() { for (int iter = 0; iter < 100; iter++) { long id = randomLong(); long generation = randomNonNegativeLong(); long version = randomNonNegativeLong(); - int segmentCount = randomIntBetween(0, 5); - List segments = new ArrayList<>(); - for (int i = 0; i < segmentCount; i++) { - segments.add(randomSegment()); - } + List segments = randomSegments(); long lastWriterGeneration = randomNonNegativeLong(); - int userDataEntries = randomIntBetween(0, 4); - Map userData = new HashMap<>(); - for (int i = 0; i < userDataEntries; i++) { - userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); - } + Map userData = randomUserData(); DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot( id, @@ -112,127 +45,56 @@ public void testSnapshotFieldAccessConsistency() { userData ); - try { - // Verify getId() - assertEquals("getId() should return the snapshot id", id, snapshot.getId()); - - // Verify getGeneration() and getVersion() from parent - assertEquals("getGeneration() should return the generation", generation, snapshot.getGeneration()); - assertEquals("getVersion() should return the version", version, snapshot.getVersion()); + assertEquals(id, snapshot.getId()); + assertEquals(generation, snapshot.getGeneration()); + assertEquals(version, snapshot.getVersion()); + assertEquals(segments, snapshot.getSegments()); + assertEquals(lastWriterGeneration, snapshot.getLastWriterGeneration()); + assertEquals(userData, snapshot.getUserData()); - // Verify getSegments() returns equal content - assertEquals("getSegments() should return segments equal to input", segments, snapshot.getSegments()); - - // Verify getLastWriterGeneration() - assertEquals( - "getLastWriterGeneration() should return the writer generation", - lastWriterGeneration, - snapshot.getLastWriterGeneration() - ); - - // Verify getUserData() - assertEquals("getUserData() should return the user data", userData, snapshot.getUserData()); - - // Verify getSearchableFiles() returns exactly the WriterFileSets matching the queried format - Set expectedFormats = new HashSet<>(); + Set expectedFormats = new HashSet<>(); + for (Segment seg : segments) { + expectedFormats.addAll(seg.dfGroupedSearchableFiles().keySet()); + } + for (String format : expectedFormats) { + List expected = new ArrayList<>(); for (Segment seg : segments) { - expectedFormats.addAll(seg.dfGroupedSearchableFiles().keySet()); - } - for (String format : expectedFormats) { - Collection searchableFiles = snapshot.getSearchableFiles(format); - List expected = new ArrayList<>(); - for (Segment seg : segments) { - WriterFileSet wfs = seg.dfGroupedSearchableFiles().get(format); - if (wfs != null) { - expected.add(wfs); - } - } - assertEquals( - "getSearchableFiles('" + format + "') should return matching WriterFileSets", - expected, - new ArrayList<>(searchableFiles) - ); + WriterFileSet wfs = seg.dfGroupedSearchableFiles().get(format); + if (wfs != null) expected.add(wfs); } - - // Verify getSearchableFiles() for a non-existent format returns empty - Collection emptyResult = snapshot.getSearchableFiles("nonexistent_format_" + randomAlphaOfLength(5)); - assertTrue("getSearchableFiles for unknown format should be empty", emptyResult.isEmpty()); - - // Verify getDataFormats() returns the union of all format keys - assertEquals("getDataFormats() should return union of all format keys", expectedFormats, snapshot.getDataFormats()); - - // Verify getSegments() returns an unmodifiable list - expectThrows(UnsupportedOperationException.class, () -> snapshot.getSegments().add(randomSegment())); - } finally { - snapshot.decRef(); + assertEquals(expected, new ArrayList<>(snapshot.getSearchableFiles(format))); } + + assertTrue(snapshot.getSearchableFiles("nonexistent_" + randomAlphaOfLength(5)).isEmpty()); + assertEquals(expectedFormats, snapshot.getDataFormats()); + expectThrows(UnsupportedOperationException.class, () -> snapshot.getSegments().add(randomSegment())); } } - // Feature: catalog-snapshot-manager, Property 2: Serialization round-trip - - /** - * Property 2: Serialization round-trip. - * For any valid DataformatAwareCatalogSnapshot, serializing via serializeToString(), - * then deserializing via deserializeFromString(), then serializing again should produce - * a binary-encoded string that, when deserialized, yields an equivalent snapshot. - * - * Validates: Requirements 2.1, 2.2, 2.3 - */ public void testSerializationRoundTrip() throws Exception { for (int iter = 0; iter < 100; iter++) { DataformatAwareCatalogSnapshot original = randomSnapshot(); - try { - // First serialization - String json1 = original.serializeToString(); + String serialized = original.serializeToString(); - // Deserialize from the first JSON - DataformatAwareCatalogSnapshot deserialized1 = DataformatAwareCatalogSnapshot.deserializeFromString(json1); - try { - // Second serialization from the deserialized snapshot - String json2 = deserialized1.serializeToString(); - - // Deserialize from the second JSON to verify round-trip equivalence - DataformatAwareCatalogSnapshot deserialized2 = DataformatAwareCatalogSnapshot.deserializeFromString(json2); - try { - // Verify deserialized1 fields match the original - assertSnapshotFieldsEqual("first deserialization", original, deserialized1); + DataformatAwareCatalogSnapshot deserialized = DataformatAwareCatalogSnapshot.deserializeFromString(serialized); + assertSnapshotFieldsEqual("round-trip", original, deserialized); - // Verify deserialized2 fields match deserialized1 (round-trip stability) - assertSnapshotFieldsEqual("second deserialization", deserialized1, deserialized2); - } finally { - deserialized2.decRef(); - } - } finally { - deserialized1.decRef(); - } - } finally { - original.decRef(); - } + String reserialized = deserialized.serializeToString(); + DataformatAwareCatalogSnapshot deserialized2 = DataformatAwareCatalogSnapshot.deserializeFromString(reserialized); + assertSnapshotFieldsEqual("double round-trip", deserialized, deserialized2); } } - /** - * Asserts that two {@link DataformatAwareCatalogSnapshot} instances have equivalent fields. - */ - private void assertSnapshotFieldsEqual(String context, DataformatAwareCatalogSnapshot expected, DataformatAwareCatalogSnapshot actual) { - assertEquals(context + ": id should match", expected.getId(), actual.getId()); - assertEquals(context + ": generation should match", expected.getGeneration(), actual.getGeneration()); - assertEquals(context + ": version should match", expected.getVersion(), actual.getVersion()); - assertEquals(context + ": segments should match", expected.getSegments(), actual.getSegments()); - assertEquals(context + ": lastWriterGeneration should match", expected.getLastWriterGeneration(), actual.getLastWriterGeneration()); - assertEquals(context + ": userData should match", expected.getUserData(), actual.getUserData()); + public void testCopyWriteable() throws Exception { + DataformatAwareCatalogSnapshot original = randomSnapshot(); + DataformatAwareCatalogSnapshot copy = copyWriteable( + original, + new NamedWriteableRegistry(Collections.emptyList()), + DataformatAwareCatalogSnapshot::new + ); + assertSnapshotFieldsEqual("copyWriteable", original, copy); } - // Feature: catalog-snapshot-manager, Property 3: Deserialization rejects invalid input - - /** - * Property 3: Deserialization rejects invalid input. - * For any string that is not valid Base64 or is truncated/corrupted binary data, - * calling deserializeFromString() should throw an IOException. - * - * Validates: Requirements 2.4 - */ public void testDeserializationRejectsInvalidInput() { for (int iter = 0; iter < 100; iter++) { String input = generateInvalidInput(iter); @@ -240,45 +102,93 @@ public void testDeserializationRejectsInvalidInput() { } } - /** - * Generates an invalid input string for deserialization testing. - * Mixes different categories of invalid input across iterations. - */ + // --- helpers --- + + private WriterFileSet randomWriterFileSet(String format) { + String directory = "/tmp/" + randomAlphaOfLength(8); + long writerGeneration = randomNonNegativeLong(); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + String[] extensions = "lucene".equals(format) ? new String[] { "cfs", "si", "dat" } : new String[] { "parquet" }; + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom(extensions)); + } + return new WriterFileSet(directory, writerGeneration, files, randomIntBetween(0, 10000)); + } + + private Segment randomSegment() { + long generation = randomNonNegativeLong(); + int formatCount = randomIntBetween(1, 2); + Map dfGrouped = new HashMap<>(); + for (int i = 0; i < formatCount; i++) { + String format = randomFrom("lucene", "parquet"); + dfGrouped.put(format, randomWriterFileSet(format)); + } + return new Segment(generation, dfGrouped); + } + + private List randomSegments() { + int count = randomIntBetween(0, 5); + List segments = new ArrayList<>(); + for (int i = 0; i < count; i++) { + segments.add(randomSegment()); + } + return segments; + } + + private Map randomUserData() { + int entries = randomIntBetween(0, 4); + Map userData = new HashMap<>(); + for (int i = 0; i < entries; i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); + } + return userData; + } + + private DataformatAwareCatalogSnapshot randomSnapshot() { + return new DataformatAwareCatalogSnapshot( + randomLong(), + randomNonNegativeLong(), + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + randomUserData() + ); + } + + private void assertSnapshotFieldsEqual(String context, DataformatAwareCatalogSnapshot expected, DataformatAwareCatalogSnapshot actual) { + assertEquals(context + ": id", expected.getId(), actual.getId()); + assertEquals(context + ": generation", expected.getGeneration(), actual.getGeneration()); + assertEquals(context + ": version", expected.getVersion(), actual.getVersion()); + assertEquals(context + ": segments", expected.getSegments(), actual.getSegments()); + assertEquals(context + ": lastWriterGeneration", expected.getLastWriterGeneration(), actual.getLastWriterGeneration()); + assertEquals(context + ": userData", expected.getUserData(), actual.getUserData()); + } + private String generateInvalidInput(int iter) { - int category = iter % 6; - switch (category) { + switch (iter % 6) { case 0: - // Completely random alphanumeric strings (not valid Base64 payload) return randomAlphaOfLengthBetween(1, 200); case 1: - // Valid Base64 but random bytes (not a valid serialized snapshot) byte[] randomBytes = new byte[randomIntBetween(1, 100)]; random().nextBytes(randomBytes); return java.util.Base64.getEncoder().encodeToString(randomBytes); case 2: - // Truncated valid Base64: serialize a valid snapshot, then truncate the Base64 string DataformatAwareCatalogSnapshot snap = randomSnapshot(); try { String validBase64 = snap.serializeToString(); - int truncateAt = randomIntBetween(1, Math.max(1, validBase64.length() / 2)); - return validBase64.substring(0, truncateAt); + return validBase64.substring(0, randomIntBetween(1, Math.max(1, validBase64.length() / 2))); } catch (IOException e) { return "AAAA"; - } finally { - snap.decRef(); } case 3: - // Empty string return ""; case 4: - // Strings with invalid Base64 characters return randomFrom("not-base64!!!", "===", "@@@@", "hello world", "{\"json\":true}"); case 5: - // Null-like strings return randomFrom("null", "undefined", "None", "nil", "NaN"); default: return randomAlphaOfLength(10); } } - } diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java index 7b469f68c57fd..0c8986d3a1043 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshotTests.java @@ -10,101 +10,78 @@ import org.apache.lucene.index.SegmentInfos; import org.apache.lucene.util.Version; -import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.test.OpenSearchTestCase; +import java.util.Collections; import java.util.HashMap; import java.util.Map; /** - * Property-based tests for {@link SegmentInfosCatalogSnapshot}. - * Uses OpenSearch randomization utilities to generate random inputs across many iterations. + * Tests for {@link SegmentInfosCatalogSnapshot}. */ public class SegmentInfosCatalogSnapshotTests extends OpenSearchTestCase { - // Feature: catalog-snapshot-manager, Property 10: SegmentInfosCatalogSnapshot delegates to SegmentInfos - - /** - * Creates a random {@link SegmentInfos} instance with random user data. - */ - private SegmentInfos randomSegmentInfos() { - SegmentInfos segmentInfos = new SegmentInfos(Version.LATEST.major); - int userDataEntries = randomIntBetween(0, 5); - Map userData = new HashMap<>(); - for (int i = 0; i < userDataEntries; i++) { - userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); - } - segmentInfos.setUserData(userData, false); - return segmentInfos; - } - - /** - * Property 10: SegmentInfosCatalogSnapshot delegates to SegmentInfos. - * For any valid Lucene SegmentInfos instance, constructing a SegmentInfosCatalogSnapshot - * and querying its accessors should return consistent results delegated from SegmentInfos. - * - * Validates: Requirements 7.1, 7.2, 7.3, 7.4 - */ - public void testSegmentInfosDelegation() { + public void testDelegation() { for (int iter = 0; iter < 100; iter++) { SegmentInfos segmentInfos = randomSegmentInfos(); SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); - try { - // Verify getId() returns the SegmentInfos generation - assertEquals("getId() should return the SegmentInfos generation", segmentInfos.getGeneration(), snapshot.getId()); + assertEquals(segmentInfos.getGeneration(), snapshot.getId()); + assertEquals(segmentInfos.getGeneration(), snapshot.getGeneration()); + assertEquals(segmentInfos.getVersion(), snapshot.getVersion()); + assertEquals(segmentInfos.getUserData(), snapshot.getUserData()); + assertEquals(-1L, snapshot.getLastWriterGeneration()); + assertSame(segmentInfos, snapshot.getSegmentInfos()); - // Verify getGeneration() returns the SegmentInfos generation - assertEquals( - "getGeneration() should return the SegmentInfos generation", - segmentInfos.getGeneration(), - snapshot.getGeneration() - ); + expectThrows(UnsupportedOperationException.class, snapshot::getSegments); + expectThrows(UnsupportedOperationException.class, () -> snapshot.getSearchableFiles(randomAlphaOfLength(5))); + expectThrows(UnsupportedOperationException.class, snapshot::getDataFormats); + expectThrows(UnsupportedOperationException.class, snapshot::serializeToString); - // Verify getVersion() returns the SegmentInfos version - assertEquals("getVersion() should return the SegmentInfos version", segmentInfos.getVersion(), snapshot.getVersion()); - - // Verify getUserData() returns the SegmentInfos user data - assertEquals("getUserData() should return the SegmentInfos user data", segmentInfos.getUserData(), snapshot.getUserData()); - - // Verify getLastWriterGeneration() returns -1 - assertEquals("getLastWriterGeneration() should return -1", -1L, snapshot.getLastWriterGeneration()); - - // Verify getSegmentInfos() returns the wrapped instance - assertSame("getSegmentInfos() should return the wrapped SegmentInfos", segmentInfos, snapshot.getSegmentInfos()); + // setUserData is a no-op, should not throw + snapshot.setUserData(Map.of("key", "value")); + } + } - // Verify getSegments() throws UnsupportedOperationException - expectThrows(UnsupportedOperationException.class, () -> snapshot.getSegments()); + public void testClone() { + for (int iter = 0; iter < 100; iter++) { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfos); - // Verify getSearchableFiles() throws UnsupportedOperationException - String randomFormat = randomAlphaOfLength(5); - expectThrows(UnsupportedOperationException.class, () -> snapshot.getSearchableFiles(randomFormat)); + SegmentInfosCatalogSnapshot cloned = snapshot.clone(); + assertNotSame(snapshot, cloned); + assertSame(segmentInfos, cloned.getSegmentInfos()); + assertEquals(snapshot.getGeneration(), cloned.getGeneration()); + assertEquals(snapshot.getVersion(), cloned.getVersion()); + } + } - // Verify getDataFormats() throws UnsupportedOperationException - expectThrows(UnsupportedOperationException.class, () -> snapshot.getDataFormats()); + public void testCopyWriteable() throws Exception { + SegmentInfos segmentInfos = randomSegmentInfos(); + SegmentInfosCatalogSnapshot original = new SegmentInfosCatalogSnapshot(segmentInfos); - // Verify serializeToString() throws UnsupportedOperationException - expectThrows(UnsupportedOperationException.class, () -> snapshot.serializeToString()); + SegmentInfosCatalogSnapshot copy = copyWriteable( + original, + new NamedWriteableRegistry(Collections.emptyList()), + SegmentInfosCatalogSnapshot::new + ); - // Verify writeTo() throws UnsupportedOperationException - expectThrows(UnsupportedOperationException.class, () -> snapshot.writeTo(new BytesStreamOutput())); + assertEquals(original.getGeneration(), copy.getGeneration()); + assertEquals(original.getVersion(), copy.getVersion()); + assertEquals(original.getUserData(), copy.getUserData()); + } - // Verify clone() returns a new instance wrapping the same SegmentInfos - SegmentInfosCatalogSnapshot cloned = snapshot.clone(); - try { - assertNotSame("clone() should return a new instance", snapshot, cloned); - assertSame("clone() should wrap the same SegmentInfos", segmentInfos, cloned.getSegmentInfos()); - assertEquals("clone() should have the same generation", snapshot.getGeneration(), cloned.getGeneration()); - assertEquals("clone() should have the same version", snapshot.getVersion(), cloned.getVersion()); - } finally { - cloned.decRef(); - } + // --- helpers --- - // Verify setUserData() is a no-op (does not throw) - snapshot.setUserData(Map.of("key", "value")); - } finally { - snapshot.decRef(); - } + private SegmentInfos randomSegmentInfos() { + SegmentInfos segmentInfos = new SegmentInfos(Version.LATEST.major); + int userDataEntries = randomIntBetween(0, 5); + Map userData = new HashMap<>(); + for (int i = 0; i < userDataEntries; i++) { + userData.put(randomAlphaOfLength(5), randomAlphaOfLength(10)); } + segmentInfos.setUserData(userData, false); + return segmentInfos; } } From f8c7c70f9427f2dd5b9915574f2f619710e44d3a Mon Sep 17 00:00:00 2001 From: Arpit Bandejiya Date: Wed, 1 Apr 2026 15:48:12 +0530 Subject: [PATCH 3/3] Address comment part2 Signed-off-by: Arpit Bandejiya --- .gitignore | 1 + .../datafusion/DatafusionReaderManager.java | 2 +- .../be/lucene/LuceneReaderManager.java | 2 +- .../exec/DefaultPlanExecutorTests.java | 14 +- .../index/engine/DataFormatAwareEngine.java | 69 ++-- .../index/engine/EngineBackedIndexer.java | 2 +- .../engine/dataformat/merge/MergeHandler.java | 2 +- .../CatalogSnapshotLifecycleListener.java | 1 + ...taFormatEngineCatalogSnapshotListener.java | 1 + .../engine/exec/EngineReaderManager.java | 1 + .../index/engine/exec/IndexFileDeleter.java | 1 + .../engine/exec/IndexReaderProvider.java | 1 + .../opensearch/index/engine/exec/Indexer.java | 1 + .../exec/IndexerLifecycleOperations.java | 1 + .../opensearch/index/engine/exec/Segment.java | 19 +- .../index/engine/exec/WriterFileSet.java | 5 +- .../exec/{ => coord}/CatalogSnapshot.java | 65 +++- .../exec/coord/CatalogSnapshotManager.java | 73 ++-- .../coord/DataformatAwareCatalogSnapshot.java | 27 +- .../coord/SegmentInfosCatalogSnapshot.java | 1 - .../dataformat/DataFormatPluginTests.java | 226 +++--------- .../engine/dataformat/merge/MergeTests.java | 2 +- .../dataformat/stub/MockCatalogSnapshot.java | 19 +- .../dataformat/stub/MockReaderManager.java | 2 +- .../index/engine/exec/SegmentTests.java | 23 +- .../index/engine/exec/WriterFileSetTests.java | 25 +- .../coord/CatalogSnapshotManagerTests.java | 213 +++++------ .../DataformatAwareCatalogSnapshotTests.java | 330 +++++++++++++++++- 28 files changed, 736 insertions(+), 393 deletions(-) rename server/src/main/java/org/opensearch/index/engine/exec/{ => coord}/CatalogSnapshot.java (71%) diff --git a/.gitignore b/.gitignore index 4662ff4c5a4d1..f5637f04366fc 100644 --- a/.gitignore +++ b/.gitignore @@ -70,3 +70,4 @@ testfixtures_shared/ # build files generated doc-tools/missing-doclet/bin/ /sandbox/plugins/engine-datafusion/target/ +**/Cargo.lock \ No newline at end of file diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReaderManager.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReaderManager.java index f97f11f78b743..591ca9f26135e 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReaderManager.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReaderManager.java @@ -10,8 +10,8 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.EngineReaderManager; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.shard.ShardPath; import java.io.IOException; diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneReaderManager.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneReaderManager.java index c46d480bccfb3..d3fe2c6338089 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneReaderManager.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneReaderManager.java @@ -12,8 +12,8 @@ import org.apache.lucene.search.ReferenceManager; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.EngineReaderManager; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.io.IOException; import java.util.Collection; diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java index c641f77f6a4b8..b7c832438092f 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java @@ -30,15 +30,17 @@ import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.core.index.Index; import org.opensearch.index.IndexService; import org.opensearch.index.engine.DataFormatAwareEngine; import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.dataformat.FieldTypeCapabilities; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.EngineReaderManager; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; import org.opensearch.index.shard.IndexShard; import org.opensearch.indices.IndicesService; import org.opensearch.test.OpenSearchTestCase; @@ -104,13 +106,15 @@ public void testEndToEndExecuteWithMockBackend() throws IOException { Segment seg1 = Segment.builder(0L).addSearchableFiles(format, fs1).build(); Segment seg2 = Segment.builder(1L).addSearchableFiles(format, fs2).build(); - MockCatalogSnapshot snapshot = new MockCatalogSnapshot(1L, List.of(seg1, seg2), format); + + CatalogSnapshotManager snapshotManager = new CatalogSnapshotManager(1L, 1L, 0L, List.of(seg1, seg2), 2L, Map.of()); MockReaderManager readerManager = new MockReaderManager(format.name()); - readerManager.afterRefresh(true, snapshot); + try (GatedCloseable ref = snapshotManager.acquireSnapshot()) { + readerManager.afterRefresh(true, ref.get()); + } - DataFormatAwareEngine engine = new DataFormatAwareEngine(Map.of(format, readerManager)); - engine.setLatestSnapshot(snapshot); + DataFormatAwareEngine engine = new DataFormatAwareEngine(Map.of(format, readerManager), snapshotManager); // Mock shard + cluster wiring IndexShard shard = mock(IndexShard.class); diff --git a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java index 1d55e84d68b37..03d4354f9d03d 100644 --- a/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java +++ b/server/src/main/java/org/opensearch/index/engine/DataFormatAwareEngine.java @@ -11,10 +11,11 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.DataFormatAwareEngineFactory; import org.opensearch.index.engine.exec.EngineReaderManager; import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; import java.io.Closeable; import java.io.IOException; @@ -35,51 +36,45 @@ public class DataFormatAwareEngine implements IndexReaderProvider, Closeable { private final Map> readerManagers; - private volatile CatalogSnapshot latestSnapshot; + private volatile CatalogSnapshotManager catalogSnapshotManager; /** - * Constructs a new DataFormatAwareEngine with pre-built maps. + * Constructs a new DataFormatAwareEngine. * Prefer using {@link DataFormatAwareEngineFactory#create()}. */ + public DataFormatAwareEngine(Map> readerManagers, CatalogSnapshotManager catalogSnapshotManager) { + this.readerManagers = readerManagers; + this.catalogSnapshotManager = catalogSnapshotManager; + } + + /** + * Constructs a new DataFormatAwareEngine without a snapshot manager. + * The manager must be set via {@link #setCatalogSnapshotManager} before acquiring readers. + */ public DataFormatAwareEngine(Map> readerManagers) { this.readerManagers = readerManagers; } - public EngineReaderManager getReaderManager(DataFormat format) { - return readerManagers.get(format); + public void setCatalogSnapshotManager(CatalogSnapshotManager catalogSnapshotManager) { + this.catalogSnapshotManager = catalogSnapshotManager; } - /** - * Called by the catalog snapshot lifecycle listener after a refresh - * to update the latest searchable snapshot. - */ - public void setLatestSnapshot(CatalogSnapshot snapshot) { - CatalogSnapshot prev = this.latestSnapshot; - this.latestSnapshot = snapshot; - if (prev != null) { - prev.decRef(); - } + public EngineReaderManager getReaderManager(DataFormat format) { + return readerManagers.get(format); } /** * Acquires a DataFormatAwareReader on the latest catalog snapshot. - * The snapshot is incRef'd; the caller MUST close the returned - * {@link DataFormatAwareReader} when done, which decRef's the snapshot. + * The caller MUST close the returned {@link DataFormatAwareReader} when done, + * which releases the snapshot reference. */ public GatedCloseable acquireReader() throws IOException { - CatalogSnapshot snapshot = latestSnapshot; - if (snapshot == null) { - throw new IllegalStateException("No catalog snapshot available"); + if (catalogSnapshotManager == null) { + throw new IllegalStateException("CatalogSnapshotManager not set"); } - return acquireReader(snapshot); - } - - /** - * Acquires a dataFormatAwareReader on a specific catalog snapshot. - */ - private GatedCloseable acquireReader(CatalogSnapshot catalogSnapshot) throws IOException { - catalogSnapshot.incRef(); + GatedCloseable snapshotRef = catalogSnapshotManager.acquireSnapshot(); try { + CatalogSnapshot catalogSnapshot = snapshotRef.get(); Map readers = new HashMap<>(); for (Map.Entry> entry : readerManagers.entrySet()) { Object reader = entry.getValue().getReader(catalogSnapshot); @@ -87,10 +82,10 @@ private GatedCloseable acquireReader(CatalogSnapshot catalogSnapshot) th readers.put(entry.getKey(), reader); } } - DataFormatAwareReader reader = new DataFormatAwareReader(catalogSnapshot, readers); + DataFormatAwareReader reader = new DataFormatAwareReader(catalogSnapshot, snapshotRef, readers); return new GatedCloseable<>(reader, reader::close); } catch (Exception e) { - catalogSnapshot.decRef(); + snapshotRef.close(); throw e; } } @@ -102,10 +97,16 @@ private GatedCloseable acquireReader(CatalogSnapshot catalogSnapshot) th @ExperimentalApi public static class DataFormatAwareReader implements IndexReaderProvider.Reader { private final CatalogSnapshot catalogSnapshot; + private final GatedCloseable snapshotRef; private final Map readers; - DataFormatAwareReader(CatalogSnapshot catalogSnapshot, Map readers) { + DataFormatAwareReader( + CatalogSnapshot catalogSnapshot, + GatedCloseable snapshotRef, + Map readers + ) { this.catalogSnapshot = catalogSnapshot; + this.snapshotRef = snapshotRef; this.readers = readers; } @@ -121,7 +122,11 @@ public CatalogSnapshot catalogSnapshot() { @Override public void close() { - catalogSnapshot.decRef(); + try { + snapshotRef.close(); + } catch (IOException e) { + throw new RuntimeException("Failed to release catalog snapshot reference", e); + } } } diff --git a/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java b/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java index 4fd056d97c762..eb03ff54c0e11 100644 --- a/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java +++ b/server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java @@ -14,8 +14,8 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.index.VersionType; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Indexer; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.mapper.DocumentMapperForType; import org.opensearch.index.mapper.SourceToParse; import org.opensearch.index.merge.MergeStats; diff --git a/server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java b/server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java index 48dfe22952202..7c6b2e3cb657d 100644 --- a/server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java +++ b/server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java @@ -15,9 +15,9 @@ import org.opensearch.common.logging.Loggers; import org.opensearch.core.index.shard.ShardId; import org.opensearch.index.engine.dataformat.MergeResult; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Indexer; import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.util.ArrayDeque; import java.util.Collection; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshotLifecycleListener.java b/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshotLifecycleListener.java index e0a40709acf33..ee0431c3d2a63 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshotLifecycleListener.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshotLifecycleListener.java @@ -9,6 +9,7 @@ package org.opensearch.index.engine.exec; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/DataFormatEngineCatalogSnapshotListener.java b/server/src/main/java/org/opensearch/index/engine/exec/DataFormatEngineCatalogSnapshotListener.java index 85e247bd29fd1..c8c6a2dc2f002 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/DataFormatEngineCatalogSnapshotListener.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/DataFormatEngineCatalogSnapshotListener.java @@ -10,6 +10,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.io.IOException; import java.util.Collection; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/EngineReaderManager.java b/server/src/main/java/org/opensearch/index/engine/exec/EngineReaderManager.java index b420dd6299471..b7d412850bd3d 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/EngineReaderManager.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/EngineReaderManager.java @@ -9,6 +9,7 @@ package org.opensearch.index.engine.exec; import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/IndexFileDeleter.java b/server/src/main/java/org/opensearch/index/engine/exec/IndexFileDeleter.java index 61507b7ffe9d7..214402e4064b5 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/IndexFileDeleter.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/IndexFileDeleter.java @@ -11,6 +11,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.index.engine.DataFormatAwareEngine; import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.shard.ShardPath; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/IndexReaderProvider.java b/server/src/main/java/org/opensearch/index/engine/exec/IndexReaderProvider.java index 5ecd9317f40a9..0662fc40c4833 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/IndexReaderProvider.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/IndexReaderProvider.java @@ -11,6 +11,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.io.Closeable; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java b/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java index 39b6ea4e80a88..d2025a2f0c22e 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/Indexer.java @@ -16,6 +16,7 @@ import org.opensearch.index.engine.EngineException; import org.opensearch.index.engine.LifecycleAware; import org.opensearch.index.engine.SafeCommitInfo; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.translog.Translog; import org.opensearch.index.translog.TranslogManager; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/IndexerLifecycleOperations.java b/server/src/main/java/org/opensearch/index/engine/exec/IndexerLifecycleOperations.java index e26189444d33a..6d51587e49def 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/IndexerLifecycleOperations.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/IndexerLifecycleOperations.java @@ -12,6 +12,7 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.core.common.unit.ByteSizeValue; import org.opensearch.index.engine.EngineException; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.shard.ShardPath; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/index/engine/exec/Segment.java b/server/src/main/java/org/opensearch/index/engine/exec/Segment.java index 974b607e3ea92..576d871832dde 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/Segment.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/Segment.java @@ -17,6 +17,7 @@ import java.io.IOException; import java.util.HashMap; import java.util.Map; +import java.util.function.Function; /** * Represents a segment in the catalog snapshot containing files grouped by data format. @@ -31,9 +32,23 @@ public record Segment(long generation, Map dfGroupedSearc /** * Constructs a Segment by deserializing from a {@link StreamInput}. + * + * @param in the stream input to read from + * @param directoryResolver function that maps a data format name to its directory path */ - public Segment(StreamInput in) throws IOException { - this(in.readLong(), in.readMap(StreamInput::readString, WriterFileSet::new)); + public Segment(StreamInput in, Function directoryResolver) throws IOException { + this(in.readLong(), readWriterFileSets(in, directoryResolver)); + } + + private static Map readWriterFileSets(StreamInput in, Function directoryResolver) + throws IOException { + int size = in.readVInt(); + Map map = new HashMap<>(size); + for (int i = 0; i < size; i++) { + String key = in.readString(); + map.put(key, new WriterFileSet(in, directoryResolver.apply(key))); + } + return map; } @Override diff --git a/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java b/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java index d36ddbfba2262..5c497d21cfc74 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java @@ -32,8 +32,8 @@ public record WriterFileSet(String directory, long writerGeneration, Set /** * Constructs a WriterFileSet by deserializing from a {@link StreamInput}. */ - public WriterFileSet(StreamInput in) throws IOException { - this(in.readString(), in.readLong(), new HashSet<>(in.readStringList()), in.readLong()); + public WriterFileSet(StreamInput in, String directory) throws IOException { + this(directory, in.readLong(), new HashSet<>(in.readStringList()), in.readLong()); } public long getTotalSize() { @@ -55,7 +55,6 @@ public String toString() { * Serializes this WriterFileSet to the given stream output. */ public void writeTo(StreamOutput out) throws IOException { - out.writeString(directory); out.writeLong(writerGeneration); out.writeStringCollection(files); out.writeLong(numRows); diff --git a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java similarity index 71% rename from server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java rename to server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java index 81c3e756ec977..f10cd55a075e3 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.index.engine.exec; +package org.opensearch.index.engine.exec.coord; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.util.concurrent.AbstractRefCounted; @@ -14,6 +14,8 @@ import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.Segment; +import org.opensearch.index.engine.exec.WriterFileSet; import java.io.IOException; import java.util.Collection; @@ -24,8 +26,9 @@ /** * Abstract base class representing a snapshot of the catalog state at a specific point in time. * Maintains versioned information about segments, files, and metadata for index operations. - * Extends AbstractRefCounted to support reference counting for safe concurrent access. - * Subclasses must implement methods for accessing file metadata, segments, and user data. + * Uses an internal reference counter for safe concurrent access. + * Subclasses must implement {@link #closeInternal()} for resource cleanup and methods for + * accessing file metadata, segments, and user data. * *

Important: Do not call {@code incRef()}, {@code decRef()}, or {@code tryIncRef()} directly. * Use {@link org.opensearch.index.engine.exec.coord.CatalogSnapshotManager#acquireSnapshot()} to obtain @@ -33,7 +36,7 @@ * when done. The manager handles all reference counting internally.

*/ @ExperimentalApi -public abstract class CatalogSnapshot extends AbstractRefCounted implements Writeable, Cloneable { +public abstract class CatalogSnapshot implements Writeable, Cloneable { /** * Key for storing catalog snapshot in user data. @@ -51,10 +54,17 @@ public abstract class CatalogSnapshot extends AbstractRefCounted implements Writ protected final long generation; protected long version; - public CatalogSnapshot(String name, long generation, long version) { - super(name); + private final AbstractRefCounted refCounter; + + protected CatalogSnapshot(String name, long generation, long version) { this.generation = generation; this.version = version; + this.refCounter = new AbstractRefCounted(name) { + @Override + protected void closeInternal() { + CatalogSnapshot.this.closeInternal(); + } + }; } /** @@ -63,10 +73,15 @@ public CatalogSnapshot(String name, long generation, long version) { * @param in the stream input to read from * @throws IOException if an I/O error occurs */ - public CatalogSnapshot(StreamInput in) throws IOException { - super("catalog_snapshot"); + protected CatalogSnapshot(StreamInput in) throws IOException { this.generation = in.readLong(); this.version = in.readLong(); + this.refCounter = new AbstractRefCounted("catalog_snapshot") { + @Override + protected void closeInternal() { + CatalogSnapshot.this.closeInternal(); + } + }; } @Override @@ -83,6 +98,35 @@ public long getVersion() { return version; } + // Package-private ref counting — only accessible within exec.coord (i.e., CatalogSnapshotManager) + + /** + * Decrements the reference count. Returns {@code true} if the count reached zero + * and {@link #closeInternal()} was invoked. + */ + boolean decRef() { + return refCounter.decRef(); + } + + /** + * Tries to increment the reference count. Returns {@code false} if the snapshot is already closed. + */ + boolean tryIncRef() { + return refCounter.tryIncRef(); + } + + /** + * Returns the current reference count. + */ + int refCount() { + return refCounter.refCount(); + } + + /** + * Called when the reference count reaches zero. Subclasses should release any resources here. + */ + protected abstract void closeInternal(); + /** * Gets user-defined metadata associated with this catalog snapshot. * @@ -141,8 +185,6 @@ public long getVersion() { * @return this catalog snapshot instance */ public CatalogSnapshot cloneNoAcquire() { - // Still using the clone call since Lucene call requires clone. This will allow a SegmentsInfos backed CatalogSnapshot to use the - // same method in calls. return this; } @@ -154,8 +196,7 @@ public CatalogSnapshot cloneNoAcquire() { public abstract void setUserData(Map userData); /** - * Creates a deep copy of this catalog snapshot. The cloned snapshot starts with a fresh reference count - * of 1 (from {@link AbstractRefCounted}). + * Creates a deep copy of this catalog snapshot. The cloned snapshot starts with a fresh reference count of 1. * Subclasses must ensure all mutable state is properly copied. * * @return a new {@link CatalogSnapshot} with the same logical state diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java index a23a03cf6eebb..1005c5e2b9eb6 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java @@ -10,14 +10,13 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.common.concurrent.GatedCloseable; -import org.opensearch.index.engine.exec.CatalogSnapshot; +import org.opensearch.index.engine.exec.Segment; import java.io.Closeable; +import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; -import java.util.function.Supplier; /** * Manages the lifecycle of {@link CatalogSnapshot} instances for the composite multi-format engine. @@ -37,23 +36,33 @@ public class CatalogSnapshotManager implements Closeable { private final Map catalogSnapshotMap = new ConcurrentHashMap<>(); /** - * Constructs a new CatalogSnapshotManager. The supplier is invoked exactly once to produce the - * initial snapshot, which is then tracked in the live snapshot map. + * Constructs a new CatalogSnapshotManager with an initial snapshot built from the given parameters. * - * @param initialSnapshotSupplier supplier for the initial snapshot; must not be null or return null + * @param id the unique snapshot identifier + * @param generation the initial generation number + * @param version the schema version + * @param segments the initial segments + * @param lastWriterGeneration the last writer generation + * @param userData user-defined metadata */ - public CatalogSnapshotManager(Supplier initialSnapshotSupplier) { - Objects.requireNonNull(initialSnapshotSupplier, "initialSnapshotSupplier must not be null"); - CatalogSnapshot initialSnapshot = Objects.requireNonNull( - initialSnapshotSupplier.get(), - "initialSnapshotSupplier must not return null" + public CatalogSnapshotManager( + long id, + long generation, + long version, + List segments, + long lastWriterGeneration, + Map userData + ) { + DataformatAwareCatalogSnapshot initialSnapshot = new DataformatAwareCatalogSnapshot( + id, + generation, + version, + segments, + lastWriterGeneration, + userData ); this.latestCatalogSnapshot = initialSnapshot; - if (catalogSnapshotMap.putIfAbsent(initialSnapshot.getGeneration(), initialSnapshot) != null) { - throw new IllegalStateException( - "Duplicate snapshot generation [" + initialSnapshot.getGeneration() + "] in catalog snapshot map" - ); - } + catalogSnapshotMap.put(initialSnapshot.getGeneration(), initialSnapshot); } /** @@ -75,18 +84,24 @@ public GatedCloseable acquireSnapshot() { } /** - * Commits a new snapshot, replacing the current one. The old snapshot is decRef'd and removed - * from the map if its count reaches zero. + * Commits a new snapshot built from the given refreshed segments, replacing the current one. + * The new snapshot inherits user data from the current snapshot and increments the generation. + * The old snapshot is decRef'd and removed from the map if its count reaches zero. * - * @param newSnapshot the new catalog snapshot to commit + * @param refreshedSegments the segments produced by the latest refresh */ - public synchronized void commitNewSnapshot(CatalogSnapshot newSnapshot) { + public synchronized void commitNewSnapshot(List refreshedSegments) { assert closed.get() == false : "Cannot commit to a closed CatalogSnapshotManager"; - assert newSnapshot.getGeneration() > latestCatalogSnapshot.getGeneration() : "New snapshot generation must be greater than current"; - if (catalogSnapshotMap.putIfAbsent(newSnapshot.getGeneration(), newSnapshot) != null) { - throw new IllegalStateException("Duplicate snapshot generation [" + newSnapshot.getGeneration() + "] in catalog snapshot map"); - } + DataformatAwareCatalogSnapshot newSnapshot = new DataformatAwareCatalogSnapshot( + latestCatalogSnapshot.getId() + 1, + latestCatalogSnapshot.getGeneration() + 1, + latestCatalogSnapshot.getVersion(), + refreshedSegments, + latestCatalogSnapshot.getLastWriterGeneration() + 1, + latestCatalogSnapshot.getUserData() + ); + CatalogSnapshot oldSnapshot = latestCatalogSnapshot; latestCatalogSnapshot = newSnapshot; decRefAndRemove(oldSnapshot); @@ -103,16 +118,6 @@ private void decRefAndRemove(CatalogSnapshot snapshot) { } } - /** - * Returns the current snapshot. Note: this does not increment the reference count. - * Use {@link #acquireSnapshot()} for safe concurrent access. - * - * @return the current {@link CatalogSnapshot} - */ - public CatalogSnapshot getCurrentSnapshot() { - return latestCatalogSnapshot; - } - /** * Closes this manager. Idempotent. DecRefs the current snapshot and removes it if count reaches zero. */ diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java index c52cbb3c590fe..9426bbeaad47b 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshot.java @@ -15,7 +15,6 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; @@ -28,6 +27,8 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Function; /** * Concrete implementation of {@link CatalogSnapshot} for the composite multi-format engine. @@ -41,6 +42,7 @@ public class DataformatAwareCatalogSnapshot extends CatalogSnapshot { private final List segments; private final long lastWriterGeneration; private Map userData; + private final AtomicBoolean closed = new AtomicBoolean(false); /** * Constructs a new DataformatAwareCatalogSnapshot. @@ -52,7 +54,7 @@ public class DataformatAwareCatalogSnapshot extends CatalogSnapshot { * @param lastWriterGeneration the generation of the last writer that contributed to this snapshot * @param userData user-defined metadata key-value pairs */ - public DataformatAwareCatalogSnapshot( + DataformatAwareCatalogSnapshot( long id, long generation, long version, @@ -71,9 +73,10 @@ public DataformatAwareCatalogSnapshot( * Constructs a DataformatAwareCatalogSnapshot from a {@link StreamInput}. * * @param in the stream input to read from + * @param directoryResolver function that maps a data format name to its directory path * @throws IOException if an I/O error occurs */ - public DataformatAwareCatalogSnapshot(StreamInput in) throws IOException { + DataformatAwareCatalogSnapshot(StreamInput in, Function directoryResolver) throws IOException { super(in); this.userData = in.readMap(StreamInput::readString, StreamInput::readString); @@ -84,7 +87,7 @@ public DataformatAwareCatalogSnapshot(StreamInput in) throws IOException { int segmentCount = in.readVInt(); List segmentList = new ArrayList<>(segmentCount); for (int i = 0; i < segmentCount; i++) { - segmentList.add(new Segment(in)); + segmentList.add(new Segment(in, directoryResolver)); } this.segments = Collections.unmodifiableList(segmentList); } @@ -147,17 +150,19 @@ public String serializeToString() throws IOException { * Deserializes a {@link DataformatAwareCatalogSnapshot} from a Base64-encoded binary string. * * @param serializedData the Base64 string produced by {@link #serializeToString()} + * @param directoryResolver function that maps a data format name to its directory path * @return a reconstructed {@link DataformatAwareCatalogSnapshot} * @throws IOException if the data is malformed or missing required fields */ - public static DataformatAwareCatalogSnapshot deserializeFromString(String serializedData) throws IOException { + public static DataformatAwareCatalogSnapshot deserializeFromString(String serializedData, Function directoryResolver) + throws IOException { if (serializedData == null || serializedData.isEmpty()) { throw new IOException("Cannot deserialize DataformatAwareCatalogSnapshot: input is null or empty"); } try { byte[] bytes = Base64.getDecoder().decode(serializedData); try (BytesStreamInput in = new BytesStreamInput(bytes)) { - return new DataformatAwareCatalogSnapshot(in); + return new DataformatAwareCatalogSnapshot(in, directoryResolver); } } catch (IOException e) { throw e; @@ -185,7 +190,15 @@ public DataformatAwareCatalogSnapshot clone() { @Override protected void closeInternal() { - // Subclass-specific resource cleanup. Map removal is handled by CatalogSnapshotManager.decRefAndRemove. + closed.set(true); + } + + /** + * Returns {@code true} if {@link #closeInternal()} has been invoked (ref count reached zero). + * This method is intended for testing only. + */ + public boolean isClosed() { + return closed.get(); } @Override diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java index 83642edb10872..82cf48f07d804 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/SegmentInfosCatalogSnapshot.java @@ -17,7 +17,6 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java b/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java index a15fdc7997d3d..4c5191ac1af2c 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/DataFormatPluginTests.java @@ -10,6 +10,7 @@ import org.opensearch.Version; import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.common.settings.Settings; import org.opensearch.core.index.shard.ShardId; import org.opensearch.index.IndexSettings; @@ -23,6 +24,9 @@ import org.opensearch.index.engine.dataformat.stub.MockReaderManager; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager; +import org.opensearch.index.engine.exec.coord.DataformatAwareCatalogSnapshot; import org.opensearch.index.mapper.MappedFieldType; import org.opensearch.index.mapper.MapperService; import org.opensearch.index.shard.ShardPath; @@ -30,7 +34,6 @@ import java.io.IOException; import java.nio.file.Path; -import java.util.Collection; import java.util.List; import java.util.Map; import java.util.Optional; @@ -230,14 +233,8 @@ public void testRefreshInput() { /** * Search holds snapshot alive while refresh replaces it. - *

- * Timeline: - * 1. new s1 → refcount = 1 (construction) - * 2. setLatestSnapshot(s1) → refcount = 1 (engine takes over construction ref) - * 3. acquireReader() → refcount = 2 (search adds ref) - * 4. setLatestSnapshot(s2) → s1 refcount = 1 (engine releases s1) - * 5. readerManager.onDeleted(s1) → reader closed, but s1 alive (search ref) - * 6. compositeReader.close() → s1 refcount = 0 → dead + * CatalogSnapshotManager handles ref counting: acquireReader increments, + * commitNewSnapshot replaces the latest, and closing the reader releases the old snapshot. */ public void testSearchHoldsSnapshotAliveWhileRefreshDeletesFiles() throws IOException { MockDataFormat format = new MockDataFormat(); @@ -253,20 +250,24 @@ public void testSearchHoldsSnapshotAliveWhileRefreshDeletesFiles() throws IOExce w1.close(); RefreshResult rr1 = indexEngine.refresh(RefreshInput.builder().addWriterFileSet(fs1).build()); - MockCatalogSnapshot snapshot1 = new MockCatalogSnapshot(1L, rr1.refreshedSegments(), format); + + CatalogSnapshotManager manager = new CatalogSnapshotManager(1L, 1L, 0L, rr1.refreshedSegments(), 1L, Map.of()); MockReaderManager readerManager = new MockReaderManager(format.name()); - readerManager.afterRefresh(true, snapshot1); + try (GatedCloseable ref = manager.acquireSnapshot()) { + readerManager.afterRefresh(true, ref.get()); + } - DataFormatAwareEngine dataFormatAwareEngine = new DataFormatAwareEngine(Map.of(format, readerManager)); - dataFormatAwareEngine.setLatestSnapshot(snapshot1); // takes over construction ref, refcount: 1 + DataFormatAwareEngine dataFormatAwareEngine = new DataFormatAwareEngine(Map.of(format, readerManager), manager); - // Search acquires reader — refcount: 2 + // Search acquires reader on snapshot1 — holds a ref var dataFormatAwareReader = dataFormatAwareEngine.acquireReader(); + CatalogSnapshot snapshot1 = dataFormatAwareReader.get().catalogSnapshot(); MockReader searchReader = (MockReader) dataFormatAwareReader.get().reader(format); assertEquals(1, searchReader.totalRows); + assertEquals(1L, snapshot1.getGeneration()); - // New refresh arrives — setLatestSnapshot(s2) decRefs s1 → refcount: 1 + // New refresh arrives — commit replaces snapshot Writer w2 = indexEngine.createWriter(2L); MockDocumentInput d2 = indexEngine.newDocumentInput(); d2.addField(mock(MappedFieldType.class), "Bob"); @@ -276,27 +277,41 @@ public void testSearchHoldsSnapshotAliveWhileRefreshDeletesFiles() throws IOExce w2.close(); RefreshResult rr2 = indexEngine.refresh(RefreshInput.builder().addWriterFileSet(fs1).addWriterFileSet(fs2).build()); - MockCatalogSnapshot snapshot2 = new MockCatalogSnapshot(2L, rr2.refreshedSegments(), format); - readerManager.afterRefresh(true, snapshot2); - dataFormatAwareEngine.setLatestSnapshot(snapshot2); // s1 refcount: 1 (only search ref) + manager.commitNewSnapshot(rr2.refreshedSegments()); + + try (GatedCloseable ref = manager.acquireSnapshot()) { + readerManager.afterRefresh(true, ref.get()); + } - // Old snapshot deleted from reader manager — reader closes - readerManager.onDeleted(snapshot1); - assertTrue("Reader for snapshot1 closed in reader manager", searchReader.closed); + // Snapshot1 still alive — search reader still works because the ref is held + assertFalse("Snapshot1 should still be alive while search holds ref", ((DataformatAwareCatalogSnapshot) snapshot1).isClosed()); + assertEquals(1, searchReader.totalRows); + assertSame(snapshot1, dataFormatAwareReader.get().catalogSnapshot()); - // But snapshot1 still alive — search holds the last ref - assertTrue("Snapshot1 alive while search holds ref", snapshot1.tryIncRef()); - snapshot1.decRef(); // undo probe + // New acquireSnapshot returns snapshot2, not snapshot1 + try (GatedCloseable ref = manager.acquireSnapshot()) { + assertEquals(2L, ref.get().getGeneration()); + assertNotSame(snapshot1, ref.get()); + } - // Search completes — s1 refcount: 0 → dead + // Search completes — releases the old snapshot ref dataFormatAwareReader.close(); - assertFalse("Snapshot1 dead after search releases", snapshot1.tryIncRef()); - // Snapshot 2 still works + // Snapshot1 is now dead + assertTrue( + "Snapshot1 should be closed after search releases the last ref", + ((DataformatAwareCatalogSnapshot) snapshot1).isClosed() + ); + + // Snapshot1 is now dead — tryIncRef would fail (verified via new acquire returning snapshot2) + // Snapshot 2 works try (var cr2 = dataFormatAwareEngine.acquireReader()) { MockReader r2 = (MockReader) cr2.get().reader(format); assertEquals(2, r2.totalRows); + assertEquals(2L, cr2.get().catalogSnapshot().getGeneration()); } + + manager.close(); } /** @@ -328,24 +343,15 @@ public Set supportedFields() { WriterFileSet wfs1 = WriterFileSet.builder().directory(dir).writerGeneration(1L).addFile("data.parquet").addNumRows(10).build(); WriterFileSet wfs2 = WriterFileSet.builder().directory(dir).writerGeneration(1L).addFile("data.lucene").addNumRows(10).build(); Segment seg = Segment.builder(0L).addSearchableFiles(format1, wfs1).addSearchableFiles(format2, wfs2).build(); - MockCatalogSnapshot snapshot = new MockCatalogSnapshot(1L, List.of(seg), format1) { - @Override - public Collection getSearchableFiles(String dataFormat) { - if ("mock-lucene".equals(dataFormat)) return List.of(wfs2); - return super.getSearchableFiles(dataFormat); - } - @Override - public Set getDataFormats() { - return Set.of(format1.name(), format2.name()); - } - }; + CatalogSnapshotManager manager = new CatalogSnapshotManager(1L, 1L, 0L, List.of(seg), 1L, Map.of()); - rm1.afterRefresh(true, snapshot); - rm2.afterRefresh(true, snapshot); + try (GatedCloseable ref = manager.acquireSnapshot()) { + rm1.afterRefresh(true, ref.get()); + rm2.afterRefresh(true, ref.get()); + } - DataFormatAwareEngine dataFormatAwareEngine = new DataFormatAwareEngine(Map.of(format1, rm1, format2, rm2)); - dataFormatAwareEngine.setLatestSnapshot(snapshot); + DataFormatAwareEngine dataFormatAwareEngine = new DataFormatAwareEngine(Map.of(format1, rm1, format2, rm2), manager); try (var cr = dataFormatAwareEngine.acquireReader()) { MockReader r1 = (MockReader) cr.get().reader(format1); @@ -357,6 +363,8 @@ public Set getDataFormats() { assertTrue(r1.fileNames.contains("data.parquet")); assertTrue(r2.fileNames.contains("data.lucene")); } + + manager.close(); } /** @@ -407,138 +415,4 @@ public void testFileLifecycleNotifications() throws IOException { assertEquals(1, rm.deletedFiles.size()); assertTrue(rm.deletedFiles.contains("a.parquet")); } - - static class MockReader { - final List fileNames; - final long totalRows; - boolean closed; - - MockReader(List fileNames, long totalRows) { - this.fileNames = fileNames; - this.totalRows = totalRows; - } - - void close() { - closed = true; - } - } - - static class MockReaderManager implements EngineReaderManager { - private final String formatName; - private final Map readers = new HashMap<>(); - final List addedFiles = new ArrayList<>(); - final List deletedFiles = new ArrayList<>(); - - MockReaderManager(String formatName) { - this.formatName = formatName; - } - - @Override - public MockReader getReader(CatalogSnapshot snapshot) { - return readers.get(snapshot); - } - - int readerCount() { - return readers.size(); - } - - @Override - public void beforeRefresh() {} - - @Override - public void afterRefresh(boolean didRefresh, CatalogSnapshot snapshot) { - if (didRefresh == false || readers.containsKey(snapshot)) return; - Collection files = snapshot.getSearchableFiles(formatName); - List allFiles = new ArrayList<>(); - long totalRows = 0; - for (WriterFileSet wfs : files) { - allFiles.addAll(wfs.files()); - totalRows += wfs.numRows(); - } - readers.put(snapshot, new MockReader(allFiles, totalRows)); - } - - @Override - public void onDeleted(CatalogSnapshot snapshot) { - MockReader reader = readers.remove(snapshot); - if (reader != null) reader.close(); - } - - @Override - public void onFilesDeleted(Collection files) { - deletedFiles.addAll(files); - } - - @Override - public void onFilesAdded(Collection files) { - addedFiles.addAll(files); - } - } - - static class MockCatalogSnapshot extends CatalogSnapshot { - private final List segments; - private final MockDataFormat format; - - MockCatalogSnapshot(long generation, List segments, MockDataFormat format) { - super("mock-snapshot", generation, 1L); - this.segments = segments; - this.format = format; - } - - @Override - public Map getUserData() { - return Map.of(); - } - - @Override - public long getId() { - return generation; - } - - @Override - public List getSegments() { - return segments; - } - - @Override - public Collection getSearchableFiles(String dataFormat) { - List result = new ArrayList<>(); - for (Segment seg : segments) { - WriterFileSet wfs = seg.dfGroupedSearchableFiles().get(dataFormat); - if (wfs != null) result.add(wfs); - } - return result; - } - - @Override - public Set getDataFormats() { - return Set.of(format.name()); - } - - @Override - public long getLastWriterGeneration() { - return generation; - } - - @Override - public String serializeToString() { - return "mock-snapshot-" + generation; - } - - @Override - public void setUserData(Map userData) {} - - @Override - public Object getReader(DataFormat dataFormat) { - return null; - } - - @Override - public MockCatalogSnapshot clone() { - return new MockCatalogSnapshot(generation, segments, format); - } - - @Override - protected void closeInternal() {} - } } diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/merge/MergeTests.java b/server/src/test/java/org/opensearch/index/engine/dataformat/merge/MergeTests.java index fb7cf71caa84f..9444d0d6d11f8 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/merge/MergeTests.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/merge/MergeTests.java @@ -16,10 +16,10 @@ import org.opensearch.index.engine.dataformat.DataFormat; import org.opensearch.index.engine.dataformat.MergeResult; import org.opensearch.index.engine.dataformat.stub.MockDataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Indexer; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.test.OpenSearchTestCase; import java.nio.file.Path; diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java index d5af44b775abf..9d619d95ccbcb 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockCatalogSnapshot.java @@ -8,11 +8,13 @@ package org.opensearch.index.engine.dataformat.stub; +import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.index.engine.dataformat.DataFormat; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; +import java.io.IOException; import java.util.ArrayList; import java.util.Collection; import java.util.List; @@ -73,16 +75,23 @@ public String serializeToString() { } @Override - public void setCatalogSnapshotMap(Map map) {} - - @Override - public void setUserData(Map userData, boolean b) {} + public void setUserData(Map userData) {} @Override public Object getReader(DataFormat dataFormat) { return null; } + @Override + public CatalogSnapshot clone() { + return new MockCatalogSnapshot(generation, segments, format); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + } + @Override protected void closeInternal() {} } diff --git a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockReaderManager.java b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockReaderManager.java index bfb16bbf2d329..7f628e73d26fa 100644 --- a/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockReaderManager.java +++ b/server/src/test/java/org/opensearch/index/engine/dataformat/stub/MockReaderManager.java @@ -8,9 +8,9 @@ package org.opensearch.index.engine.dataformat.stub; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.EngineReaderManager; import org.opensearch.index.engine.exec.WriterFileSet; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import java.util.ArrayList; import java.util.Collection; diff --git a/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java b/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java index 5a0a82ac30457..d5afc4257c4a1 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/SegmentTests.java @@ -22,15 +22,25 @@ */ public class SegmentTests extends OpenSearchTestCase { + private static final String TEST_DIRECTORY = "/tmp/test-segment"; + public void testCopyWriteable() throws Exception { Segment original = randomSegment(); - Segment copy = copyWriteable(original, new NamedWriteableRegistry(Collections.emptyList()), Segment::new); + Segment copy = copyWriteable( + original, + new NamedWriteableRegistry(Collections.emptyList()), + in -> new Segment(in, key -> TEST_DIRECTORY) + ); assertEquals(original, copy); } public void testCopyWriteableEmpty() throws Exception { Segment empty = new Segment(0L, Map.of()); - Segment copy = copyWriteable(empty, new NamedWriteableRegistry(Collections.emptyList()), Segment::new); + Segment copy = copyWriteable( + empty, + new NamedWriteableRegistry(Collections.emptyList()), + in -> new Segment(in, key -> TEST_DIRECTORY) + ); assertEquals(empty, copy); } @@ -40,7 +50,11 @@ public void testCopyWriteableMultiFormat() throws Exception { dfGrouped.put("parquet", randomWriterFileSet("parquet")); Segment original = new Segment(randomNonNegativeLong(), dfGrouped); - Segment copy = copyWriteable(original, new NamedWriteableRegistry(Collections.emptyList()), Segment::new); + Segment copy = copyWriteable( + original, + new NamedWriteableRegistry(Collections.emptyList()), + in -> new Segment(in, key -> TEST_DIRECTORY) + ); assertEquals(original, copy); assertEquals(2, copy.dfGroupedSearchableFiles().size()); } @@ -48,14 +62,13 @@ public void testCopyWriteableMultiFormat() throws Exception { // --- helpers --- private WriterFileSet randomWriterFileSet(String format) { - String directory = "/tmp/" + randomAlphaOfLength(8); int fileCount = randomIntBetween(1, 5); Set files = new HashSet<>(); String[] extensions = "lucene".equals(format) ? new String[] { "cfs", "si", "dat" } : new String[] { "parquet" }; for (int i = 0; i < fileCount; i++) { files.add(randomAlphaOfLength(6) + "." + randomFrom(extensions)); } - return new WriterFileSet(directory, randomNonNegativeLong(), files, randomIntBetween(0, 10000)); + return new WriterFileSet(TEST_DIRECTORY, randomNonNegativeLong(), files, randomIntBetween(0, 10000)); } private Segment randomSegment() { diff --git a/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java b/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java index 1a34b530a1032..2eb0d82b92728 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/WriterFileSetTests.java @@ -22,10 +22,33 @@ public class WriterFileSetTests extends OpenSearchTestCase { public void testCopyWriteable() throws Exception { WriterFileSet original = randomWriterFileSet(); - WriterFileSet copy = copyWriteable(original, new NamedWriteableRegistry(Collections.emptyList()), WriterFileSet::new); + String directory = original.directory(); + WriterFileSet copy = copyWriteable( + original, + new NamedWriteableRegistry(Collections.emptyList()), + in -> new WriterFileSet(in, directory) + ); assertEquals(original, copy); } + public void testDirectoryNotSerialized() throws Exception { + String originalDirectory = "/tmp/original"; + String differentDirectory = "/tmp/different"; + WriterFileSet original = new WriterFileSet(originalDirectory, 1L, Set.of("a.dat"), 10); + + WriterFileSet deserialized = copyWriteable( + original, + new NamedWriteableRegistry(Collections.emptyList()), + in -> new WriterFileSet(in, differentDirectory) + ); + + assertEquals(differentDirectory, deserialized.directory()); + assertNotEquals(originalDirectory, deserialized.directory()); + assertEquals(original.writerGeneration(), deserialized.writerGeneration()); + assertEquals(original.files(), deserialized.files()); + assertEquals(original.numRows(), deserialized.numRows()); + } + // --- helpers --- private WriterFileSet randomWriterFileSet() { diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java index 6a17af1c5d8a5..5814def67f666 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManagerTests.java @@ -9,7 +9,6 @@ package org.opensearch.index.engine.exec.coord; import org.opensearch.common.concurrent.GatedCloseable; -import org.opensearch.index.engine.exec.CatalogSnapshot; import org.opensearch.index.engine.exec.Segment; import org.opensearch.index.engine.exec.WriterFileSet; import org.opensearch.test.OpenSearchTestCase; @@ -21,33 +20,34 @@ import java.util.List; import java.util.Map; import java.util.Set; -import java.util.concurrent.atomic.AtomicBoolean; /** * Tests for {@link CatalogSnapshotManager}. */ public class CatalogSnapshotManagerTests extends OpenSearchTestCase { - public void testCommitProducesCorrectNewSnapshot() { + public void testCommitProducesCorrectNewSnapshot() throws Exception { for (int iter = 0; iter < 100; iter++) { CatalogSnapshotManager manager = createRandomManager(); try { - long previousGeneration = manager.getCurrentSnapshot().getGeneration(); + long previousGeneration; Set seenIds = new HashSet<>(); - seenIds.add(manager.getCurrentSnapshot().getId()); + try (GatedCloseable ref = manager.acquireSnapshot()) { + previousGeneration = ref.get().getGeneration(); + seenIds.add(ref.get().getId()); + } int numCommits = randomIntBetween(1, 10); for (int c = 0; c < numCommits; c++) { List newSegments = randomSegments(); - long newGeneration = previousGeneration + 1; - - manager.commitNewSnapshot(buildSnapshot(newGeneration, newSegments, randomNonNegativeLong(), Map.of())); - - assertEquals(previousGeneration + 1, manager.getCurrentSnapshot().getGeneration()); - assertTrue(seenIds.add(manager.getCurrentSnapshot().getId())); - assertEquals(newSegments, manager.getCurrentSnapshot().getSegments()); - - previousGeneration = manager.getCurrentSnapshot().getGeneration(); + manager.commitNewSnapshot(newSegments); + + try (GatedCloseable ref = manager.acquireSnapshot()) { + assertEquals(previousGeneration + 1, ref.get().getGeneration()); + assertTrue(seenIds.add(ref.get().getId())); + assertEquals(newSegments, ref.get().getSegments()); + previousGeneration = ref.get().getGeneration(); + } } } finally { manager.close(); @@ -55,71 +55,73 @@ public void testCommitProducesCorrectNewSnapshot() { } } - public void testUserDataPreservationOnCommit() { + public void testUserDataPreservationOnCommit() throws Exception { for (int iter = 0; iter < 100; iter++) { Map initialUserData = randomUserData(randomIntBetween(1, 5)); long initGen = randomIntBetween(0, 100); CatalogSnapshotManager manager = new CatalogSnapshotManager( - () -> new DataformatAwareCatalogSnapshot( - randomNonNegativeLong(), - initGen, - randomNonNegativeLong(), - randomSegments(), - randomNonNegativeLong(), - initialUserData - ) + randomNonNegativeLong(), + initGen, + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + initialUserData ); try { - long gen1 = initGen + 1; - manager.commitNewSnapshot( - buildSnapshot(gen1, randomSegments(), randomNonNegativeLong(), manager.getCurrentSnapshot().getUserData()) - ); - assertEquals(initialUserData, manager.getCurrentSnapshot().getUserData()); - - Map newUserData = randomUserData(randomIntBetween(1, 5)); - long gen2 = gen1 + 1; - manager.commitNewSnapshot(buildSnapshot(gen2, randomSegments(), randomNonNegativeLong(), newUserData)); - assertEquals(newUserData, manager.getCurrentSnapshot().getUserData()); + manager.commitNewSnapshot(randomSegments()); + try (GatedCloseable ref = manager.acquireSnapshot()) { + assertEquals(initialUserData, ref.get().getUserData()); + } + + manager.commitNewSnapshot(randomSegments()); + try (GatedCloseable ref = manager.acquireSnapshot()) { + assertEquals(initialUserData, ref.get().getUserData()); + } } finally { manager.close(); } } } - public void testReferenceCountingLifecycle() { + public void testReferenceCountingLifecycle() throws Exception { for (int iter = 0; iter < 100; iter++) { - AtomicBoolean closeInternalCalled = new AtomicBoolean(false); long initGen = randomIntBetween(0, 100); CatalogSnapshotManager manager = new CatalogSnapshotManager( - () -> new TrackableSnapshot( - randomNonNegativeLong(), - initGen, - randomNonNegativeLong(), - randomSegments(), - randomNonNegativeLong(), - Collections.emptyMap(), - closeInternalCalled - ) + randomNonNegativeLong(), + initGen, + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + Collections.emptyMap() ); - CatalogSnapshot initialSnapshot = manager.getCurrentSnapshot(); + CatalogSnapshot initialSnapshot; + try (GatedCloseable ref = manager.acquireSnapshot()) { + initialSnapshot = ref.get(); + assertEquals(2, initialSnapshot.refCount()); + } assertEquals(1, initialSnapshot.refCount()); - manager.commitNewSnapshot(buildSnapshot(initGen + 1, randomSegments(), randomNonNegativeLong(), Map.of())); + manager.commitNewSnapshot(randomSegments()); assertEquals(0, initialSnapshot.refCount()); - assertTrue(closeInternalCalled.get()); int numCommits = randomIntBetween(1, 8); for (int c = 0; c < numCommits; c++) { - CatalogSnapshot prev = manager.getCurrentSnapshot(); + CatalogSnapshot prev; + try (GatedCloseable ref = manager.acquireSnapshot()) { + prev = ref.get(); + assertEquals(2, prev.refCount()); + } assertEquals(1, prev.refCount()); - manager.commitNewSnapshot( - buildSnapshot(manager.getCurrentSnapshot().getGeneration() + 1, randomSegments(), randomNonNegativeLong(), Map.of()) - ); + manager.commitNewSnapshot(randomSegments()); assertEquals(0, prev.refCount()); } - CatalogSnapshot finalSnapshot = manager.getCurrentSnapshot(); + CatalogSnapshot finalSnapshot; + try (GatedCloseable ref = manager.acquireSnapshot()) { + finalSnapshot = ref.get(); + assertEquals(2, finalSnapshot.refCount()); + } assertEquals(1, finalSnapshot.refCount()); manager.close(); assertEquals(0, finalSnapshot.refCount()); @@ -130,7 +132,11 @@ public void testAcquireAndReleaseViaGatedCloseable() throws Exception { for (int iter = 0; iter < 100; iter++) { CatalogSnapshotManager manager = createRandomManager(); try { - CatalogSnapshot currentSnap = manager.getCurrentSnapshot(); + CatalogSnapshot currentSnap; + try (GatedCloseable initialRef = manager.acquireSnapshot()) { + currentSnap = initialRef.get(); + assertEquals(2, currentSnap.refCount()); + } assertEquals(1, currentSnap.refCount()); int numAcquires = randomIntBetween(1, 5); @@ -149,9 +155,7 @@ public void testAcquireAndReleaseViaGatedCloseable() throws Exception { CatalogSnapshot heldSnapshot = heldRef.get(); assertEquals(2, heldSnapshot.refCount()); - manager.commitNewSnapshot( - buildSnapshot(manager.getCurrentSnapshot().getGeneration() + 1, randomSegments(), randomNonNegativeLong(), Map.of()) - ); + manager.commitNewSnapshot(randomSegments()); assertEquals(1, heldSnapshot.refCount()); heldRef.close(); @@ -162,13 +166,11 @@ public void testAcquireAndReleaseViaGatedCloseable() throws Exception { } } - public void testClosedManagerRejectsAcquisition() { + public void testClosedManagerRejectsAcquisition() throws Exception { for (int iter = 0; iter < 100; iter++) { CatalogSnapshotManager manager = createRandomManager(); for (int c = 0; c < randomIntBetween(0, 5); c++) { - manager.commitNewSnapshot( - buildSnapshot(manager.getCurrentSnapshot().getGeneration() + 1, randomSegments(), randomNonNegativeLong(), Map.of()) - ); + manager.commitNewSnapshot(randomSegments()); } manager.close(); expectThrows(IllegalStateException.class, manager::acquireSnapshot); @@ -184,9 +186,7 @@ public void testInitialSnapshotRecovery() throws Exception { List segments = randomIntBetween(1, 5) == 1 ? Collections.emptyList() : randomSegments(); Map userData = randomUserData(randomIntBetween(0, 4)); - CatalogSnapshotManager manager = new CatalogSnapshotManager( - () -> new DataformatAwareCatalogSnapshot(id, generation, version, segments, lastWriterGeneration, userData) - ); + CatalogSnapshotManager manager = new CatalogSnapshotManager(id, generation, version, segments, lastWriterGeneration, userData); try (GatedCloseable ref = manager.acquireSnapshot()) { CatalogSnapshot acquired = ref.get(); assertEquals(id, acquired.getId()); @@ -194,13 +194,58 @@ public void testInitialSnapshotRecovery() throws Exception { assertEquals(segments, acquired.getSegments()); assertEquals(userData, acquired.getUserData()); assertEquals(lastWriterGeneration, acquired.getLastWriterGeneration()); - assertSame(acquired, manager.getCurrentSnapshot()); } finally { manager.close(); } } } + public void testCloseInternalInvokedOnCommit() throws Exception { + CatalogSnapshotManager manager = createRandomManager(); + + CatalogSnapshot initialSnapshot; + try (GatedCloseable ref = manager.acquireSnapshot()) { + initialSnapshot = ref.get(); + } + assertFalse(((DataformatAwareCatalogSnapshot) initialSnapshot).isClosed()); + + manager.commitNewSnapshot(randomSegments()); + assertTrue( + "snapshot should be closed when commit replaces the last ref", + ((DataformatAwareCatalogSnapshot) initialSnapshot).isClosed() + ); + manager.close(); + } + + public void testCloseInternalInvokedOnManagerClose() throws Exception { + CatalogSnapshotManager manager = createRandomManager(); + + CatalogSnapshot snapshot; + try (GatedCloseable ref = manager.acquireSnapshot()) { + snapshot = ref.get(); + } + assertFalse(((DataformatAwareCatalogSnapshot) snapshot).isClosed()); + + manager.close(); + assertTrue("snapshot should be closed when manager releases the last ref", ((DataformatAwareCatalogSnapshot) snapshot).isClosed()); + } + + public void testCloseInternalNotInvokedWhileRefsHeld() throws Exception { + CatalogSnapshotManager manager = createRandomManager(); + + GatedCloseable heldRef = manager.acquireSnapshot(); + CatalogSnapshot heldSnapshot = heldRef.get(); + assertFalse(((DataformatAwareCatalogSnapshot) heldSnapshot).isClosed()); + + manager.commitNewSnapshot(randomSegments()); + assertFalse("snapshot should not be closed while a ref is still held", ((DataformatAwareCatalogSnapshot) heldSnapshot).isClosed()); + + heldRef.close(); + assertTrue("snapshot should be closed after the last ref is released", ((DataformatAwareCatalogSnapshot) heldSnapshot).isClosed()); + + manager.close(); + } + // --- helpers --- private WriterFileSet randomWriterFileSet(String format) { @@ -239,42 +284,14 @@ private Map randomUserData(int entries) { return userData; } - private DataformatAwareCatalogSnapshot buildSnapshot(long gen, List segments, long writerGen, Map userData) { - return new DataformatAwareCatalogSnapshot(gen, gen, 0L, segments, writerGen, userData); - } - private CatalogSnapshotManager createRandomManager() { return new CatalogSnapshotManager( - () -> new DataformatAwareCatalogSnapshot( - randomNonNegativeLong(), - randomIntBetween(0, 100), - randomNonNegativeLong(), - randomSegments(), - randomNonNegativeLong(), - Map.of() - ) + randomNonNegativeLong(), + randomIntBetween(0, 100), + randomNonNegativeLong(), + randomSegments(), + randomNonNegativeLong(), + Map.of() ); } - - private static class TrackableSnapshot extends DataformatAwareCatalogSnapshot { - private final AtomicBoolean closeInternalCalled; - - TrackableSnapshot( - long id, - long gen, - long version, - List segments, - long writerGen, - Map userData, - AtomicBoolean closeInternalCalled - ) { - super(id, gen, version, segments, writerGen, userData); - this.closeInternalCalled = closeInternalCalled; - } - - @Override - protected void closeInternal() { - closeInternalCalled.set(true); - } - } } diff --git a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java index 1e866cea75d91..3b6089ab4729e 100644 --- a/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java +++ b/server/src/test/java/org/opensearch/index/engine/exec/coord/DataformatAwareCatalogSnapshotTests.java @@ -76,21 +76,30 @@ public void testSerializationRoundTrip() throws Exception { DataformatAwareCatalogSnapshot original = randomSnapshot(); String serialized = original.serializeToString(); - DataformatAwareCatalogSnapshot deserialized = DataformatAwareCatalogSnapshot.deserializeFromString(serialized); - assertSnapshotFieldsEqual("round-trip", original, deserialized); + // Directory is not serialized; pass a placeholder for deserialization + String directory = "/tmp/deserialized"; + DataformatAwareCatalogSnapshot deserialized = DataformatAwareCatalogSnapshot.deserializeFromString( + serialized, + key -> directory + ); + assertSnapshotMetadataEqual("round-trip", original, deserialized); String reserialized = deserialized.serializeToString(); - DataformatAwareCatalogSnapshot deserialized2 = DataformatAwareCatalogSnapshot.deserializeFromString(reserialized); + DataformatAwareCatalogSnapshot deserialized2 = DataformatAwareCatalogSnapshot.deserializeFromString( + reserialized, + key -> directory + ); assertSnapshotFieldsEqual("double round-trip", deserialized, deserialized2); } } public void testCopyWriteable() throws Exception { - DataformatAwareCatalogSnapshot original = randomSnapshot(); + String directory = "/tmp/" + randomAlphaOfLength(8); + DataformatAwareCatalogSnapshot original = randomSnapshotWithDirectory(directory); DataformatAwareCatalogSnapshot copy = copyWriteable( original, new NamedWriteableRegistry(Collections.emptyList()), - DataformatAwareCatalogSnapshot::new + in -> new DataformatAwareCatalogSnapshot(in, key -> directory) ); assertSnapshotFieldsEqual("copyWriteable", original, copy); } @@ -98,8 +107,242 @@ public void testCopyWriteable() throws Exception { public void testDeserializationRejectsInvalidInput() { for (int iter = 0; iter < 100; iter++) { String input = generateInvalidInput(iter); - expectThrows(IOException.class, () -> DataformatAwareCatalogSnapshot.deserializeFromString(input)); + expectThrows(IOException.class, () -> DataformatAwareCatalogSnapshot.deserializeFromString(input, key -> "/tmp/test")); + } + } + + public void testInitialRefCountIsOne() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertEquals(1, snapshot.refCount()); + } + + public void testAcquireRefIncrementsCount() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertEquals(1, snapshot.refCount()); + + snapshot.tryIncRef(); + assertEquals(2, snapshot.refCount()); + + snapshot.tryIncRef(); + assertEquals(3, snapshot.refCount()); + + snapshot.decRef(); + snapshot.decRef(); + snapshot.decRef(); + } + + public void testReleaseRefDecrementsAndTriggersCloseAtZero() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertEquals(1, snapshot.refCount()); + + snapshot.tryIncRef(); + assertEquals(2, snapshot.refCount()); + + assertFalse(snapshot.decRef()); + assertEquals(1, snapshot.refCount()); + + assertTrue(snapshot.decRef()); + assertEquals(0, snapshot.refCount()); + } + + public void testTryAcquireRefSucceedsWhenOpen() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertTrue(snapshot.tryIncRef()); + assertEquals(2, snapshot.refCount()); + + snapshot.decRef(); + snapshot.decRef(); + } + + public void testTryAcquireRefFailsWhenClosed() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertTrue(snapshot.decRef()); + assertEquals(0, snapshot.refCount()); + + assertFalse(snapshot.tryIncRef()); + } + + public void testCloseInternalCalledOnceAtZeroRefCount() { + final java.util.concurrent.atomic.AtomicInteger closeCount = new java.util.concurrent.atomic.AtomicInteger(0); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of()) { + @Override + protected void closeInternal() { + closeCount.incrementAndGet(); + } + }; + + snapshot.tryIncRef(); + snapshot.tryIncRef(); + assertEquals(0, closeCount.get()); + + snapshot.decRef(); + assertEquals(0, closeCount.get()); + + snapshot.decRef(); + assertEquals(0, closeCount.get()); + + snapshot.decRef(); + assertEquals(1, closeCount.get()); + } + + public void testClonedSnapshotHasFreshRefCount() { + DataformatAwareCatalogSnapshot original = randomSnapshot(); + original.tryIncRef(); + assertEquals(2, original.refCount()); + + DataformatAwareCatalogSnapshot cloned = original.clone(); + assertEquals(1, cloned.refCount()); + assertEquals(2, original.refCount()); + + original.decRef(); + original.decRef(); + cloned.decRef(); + } + + public void testRefCounterDelegatesCloseInternalToSubclass() { + // Verifies the anonymous AbstractRefCounted bridge calls the subclass closeInternal, not a default + final List events = new ArrayList<>(); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of()) { + @Override + protected void closeInternal() { + events.add("subclass-closed"); + } + }; + + assertTrue(events.isEmpty()); + snapshot.decRef(); + assertEquals(List.of("subclass-closed"), events); + } + + public void testEachSnapshotHasIndependentRefCounter() { + DataformatAwareCatalogSnapshot snap1 = randomSnapshot(); + DataformatAwareCatalogSnapshot snap2 = randomSnapshot(); + + snap1.tryIncRef(); + assertEquals(2, snap1.refCount()); + assertEquals(1, snap2.refCount()); + + snap2.decRef(); + assertEquals(0, snap2.refCount()); + assertEquals(2, snap1.refCount()); + + snap1.decRef(); + snap1.decRef(); + } + + public void testRefCounterSurvivesMultipleIncDecCycles() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + + for (int cycle = 0; cycle < 10; cycle++) { + int refs = randomIntBetween(1, 20); + for (int i = 0; i < refs; i++) { + snapshot.tryIncRef(); + } + assertEquals(1 + refs, snapshot.refCount()); + for (int i = 0; i < refs; i++) { + assertFalse(snapshot.decRef()); + } + assertEquals(1, snapshot.refCount()); + } + + assertTrue(snapshot.decRef()); + assertEquals(0, snapshot.refCount()); + assertFalse(snapshot.tryIncRef()); + } + + public void testCloseInternalNotCalledOnIntermediateDecRef() { + final java.util.concurrent.atomic.AtomicBoolean closed = new java.util.concurrent.atomic.AtomicBoolean(false); + DataformatAwareCatalogSnapshot snapshot = new DataformatAwareCatalogSnapshot(1L, 1L, 1L, List.of(), 0L, Map.of()) { + @Override + protected void closeInternal() { + closed.set(true); + } + }; + + // Acquire several refs + int extraRefs = randomIntBetween(2, 10); + for (int i = 0; i < extraRefs; i++) { + snapshot.tryIncRef(); + } + + // Release all but one — closeInternal must NOT fire + for (int i = 0; i < extraRefs; i++) { + snapshot.decRef(); + assertFalse("closeInternal should not fire while refs remain", closed.get()); + } + + // Release the last ref — now it fires + snapshot.decRef(); + assertTrue("closeInternal should fire when last ref is released", closed.get()); + } + + public void testDeserializedSnapshotHasIndependentRefCounter() throws Exception { + String directory = "/tmp/" + randomAlphaOfLength(8); + DataformatAwareCatalogSnapshot original = randomSnapshotWithDirectory(directory); + String serialized = original.serializeToString(); + + DataformatAwareCatalogSnapshot deserialized = DataformatAwareCatalogSnapshot.deserializeFromString(serialized, key -> directory); + + // Each has its own ref counter starting at 1 + assertEquals(1, original.refCount()); + assertEquals(1, deserialized.refCount()); + + original.tryIncRef(); + assertEquals(2, original.refCount()); + assertEquals(1, deserialized.refCount()); + + original.decRef(); + original.decRef(); + deserialized.decRef(); + } + + public void testIsClosedInitiallyFalse() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertFalse(snapshot.isClosed()); + snapshot.decRef(); + } + + public void testIsClosedTrueAfterLastDecRef() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + assertFalse(snapshot.isClosed()); + + snapshot.decRef(); + assertTrue(snapshot.isClosed()); + } + + public void testIsClosedFalseWhileRefsRemain() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + int extraRefs = randomIntBetween(1, 10); + for (int i = 0; i < extraRefs; i++) { + snapshot.tryIncRef(); + } + + for (int i = 0; i < extraRefs; i++) { + snapshot.decRef(); + assertFalse("isClosed should be false while refs remain", snapshot.isClosed()); } + + snapshot.decRef(); + assertTrue(snapshot.isClosed()); + } + + public void testIsClosedAfterCloneIndependent() { + DataformatAwareCatalogSnapshot original = randomSnapshot(); + DataformatAwareCatalogSnapshot cloned = original.clone(); + + original.decRef(); + assertTrue(original.isClosed()); + assertFalse(cloned.isClosed()); + + cloned.decRef(); + assertTrue(cloned.isClosed()); + } + + public void testTryIncRefFailsAfterClosed() { + DataformatAwareCatalogSnapshot snapshot = randomSnapshot(); + snapshot.decRef(); + assertTrue(snapshot.isClosed()); + assertFalse(snapshot.tryIncRef()); } // --- helpers --- @@ -156,6 +399,48 @@ private DataformatAwareCatalogSnapshot randomSnapshot() { ); } + private DataformatAwareCatalogSnapshot randomSnapshotWithDirectory(String directory) { + return new DataformatAwareCatalogSnapshot( + randomLong(), + randomNonNegativeLong(), + randomNonNegativeLong(), + randomSegmentsWithDirectory(directory), + randomNonNegativeLong(), + randomUserData() + ); + } + + private List randomSegmentsWithDirectory(String directory) { + int count = randomIntBetween(0, 5); + List segments = new ArrayList<>(); + for (int i = 0; i < count; i++) { + segments.add(randomSegmentWithDirectory(directory)); + } + return segments; + } + + private Segment randomSegmentWithDirectory(String directory) { + long generation = randomNonNegativeLong(); + int formatCount = randomIntBetween(1, 2); + Map dfGrouped = new HashMap<>(); + for (int i = 0; i < formatCount; i++) { + String format = randomFrom("lucene", "parquet"); + dfGrouped.put(format, randomWriterFileSetWithDirectory(format, directory)); + } + return new Segment(generation, dfGrouped); + } + + private WriterFileSet randomWriterFileSetWithDirectory(String format, String directory) { + long writerGeneration = randomNonNegativeLong(); + int fileCount = randomIntBetween(1, 5); + Set files = new HashSet<>(); + String[] extensions = "lucene".equals(format) ? new String[] { "cfs", "si", "dat" } : new String[] { "parquet" }; + for (int i = 0; i < fileCount; i++) { + files.add(randomAlphaOfLength(6) + "." + randomFrom(extensions)); + } + return new WriterFileSet(directory, writerGeneration, files, randomIntBetween(0, 10000)); + } + private void assertSnapshotFieldsEqual(String context, DataformatAwareCatalogSnapshot expected, DataformatAwareCatalogSnapshot actual) { assertEquals(context + ": id", expected.getId(), actual.getId()); assertEquals(context + ": generation", expected.getGeneration(), actual.getGeneration()); @@ -165,6 +450,39 @@ private void assertSnapshotFieldsEqual(String context, DataformatAwareCatalogSna assertEquals(context + ": userData", expected.getUserData(), actual.getUserData()); } + /** + * Asserts metadata equality between two snapshots, ignoring directory (which is not serialized). + */ + private void assertSnapshotMetadataEqual( + String context, + DataformatAwareCatalogSnapshot expected, + DataformatAwareCatalogSnapshot actual + ) { + assertEquals(context + ": id", expected.getId(), actual.getId()); + assertEquals(context + ": generation", expected.getGeneration(), actual.getGeneration()); + assertEquals(context + ": version", expected.getVersion(), actual.getVersion()); + assertEquals(context + ": segment count", expected.getSegments().size(), actual.getSegments().size()); + for (int i = 0; i < expected.getSegments().size(); i++) { + Segment expectedSeg = expected.getSegments().get(i); + Segment actualSeg = actual.getSegments().get(i); + assertEquals(context + ": segment[" + i + "].generation", expectedSeg.generation(), actualSeg.generation()); + assertEquals( + context + ": segment[" + i + "].formats", + expectedSeg.dfGroupedSearchableFiles().keySet(), + actualSeg.dfGroupedSearchableFiles().keySet() + ); + for (String format : expectedSeg.dfGroupedSearchableFiles().keySet()) { + WriterFileSet expectedWfs = expectedSeg.dfGroupedSearchableFiles().get(format); + WriterFileSet actualWfs = actualSeg.dfGroupedSearchableFiles().get(format); + assertEquals(context + ": writerGeneration", expectedWfs.writerGeneration(), actualWfs.writerGeneration()); + assertEquals(context + ": files", expectedWfs.files(), actualWfs.files()); + assertEquals(context + ": numRows", expectedWfs.numRows(), actualWfs.numRows()); + } + } + assertEquals(context + ": lastWriterGeneration", expected.getLastWriterGeneration(), actual.getLastWriterGeneration()); + assertEquals(context + ": userData", expected.getUserData(), actual.getUserData()); + } + private String generateInvalidInput(int iter) { switch (iter % 6) { case 0: