package it.cavallium.dbengine.database.disk; import io.net5.buffer.api.Drop; import io.net5.buffer.api.Owned; import io.net5.buffer.api.Resource; import io.net5.buffer.api.Send; import io.net5.buffer.api.internal.ResourceSupport; import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; import java.io.IOException; import java.io.UncheckedIOException; import java.util.ArrayList; import java.util.Collections; import java.util.Comparator; import java.util.List; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.index.IndexReader; import org.apache.lucene.index.MultiReader; import org.apache.lucene.search.IndexSearcher; public interface LLIndexSearchers extends Resource { static LLIndexSearchers of(List> indexSearchers) { return new ShardedIndexSearchers(indexSearchers, null); } static UnshardedIndexSearchers unsharded(Send indexSearcher) { return new UnshardedIndexSearchers(indexSearcher, null); } List shards(); IndexSearcher shard(int shardIndex); IndexReader allShards(); class UnshardedIndexSearchers extends ResourceSupport implements LLIndexSearchers { private static final Logger logger = LogManager.getLogger(UnshardedIndexSearchers.class); private static final Drop DROP = new Drop<>() { @Override public void drop(UnshardedIndexSearchers obj) { try { if (obj.indexSearcher != null) { obj.indexSearcher.close(); } } catch (Throwable ex) { logger.error("Failed to close indexSearcher", ex); } try { if (obj.onClose != null) { obj.onClose.run(); } } catch (Throwable ex) { logger.error("Failed to close onClose", ex); } } @Override public Drop fork() { return this; } @Override public void attach(UnshardedIndexSearchers obj) { } }; private LLIndexSearcher indexSearcher; private Runnable onClose; public UnshardedIndexSearchers(Send indexSearcher, Runnable onClose) { super(DROP); this.indexSearcher = indexSearcher.receive(); this.onClose = onClose; } @Override public List shards() { return List.of(indexSearcher.getIndexSearcher()); } @Override public IndexSearcher shard(int shardIndex) { if (!isOwned()) { throw attachTrace(new IllegalStateException("UnshardedIndexSearchers must be owned to be used")); } if (shardIndex != -1) { throw new IndexOutOfBoundsException("Shard index " + shardIndex + " is invalid, this is a unsharded index"); } return indexSearcher.getIndexSearcher(); } @Override public IndexReader allShards() { return indexSearcher.getIndexReader(); } public IndexSearcher shard() { return this.shard(-1); } @Override protected RuntimeException createResourceClosedException() { return new IllegalStateException("Closed"); } @Override protected Owned prepareSend() { Send indexSearcher = this.indexSearcher.send(); var onClose = this.onClose; return drop -> { var instance = new UnshardedIndexSearchers(indexSearcher, onClose); drop.attach(instance); return instance; }; } protected void makeInaccessible() { this.indexSearcher = null; this.onClose = null; } } class ShardedIndexSearchers extends ResourceSupport implements LLIndexSearchers { private static final Logger logger = LogManager.getLogger(ShardedIndexSearchers.class); private static final Drop DROP = new Drop<>() { @Override public void drop(ShardedIndexSearchers obj) { try { for (LLIndexSearcher indexSearcher : obj.indexSearchers) { indexSearcher.close(); } } catch (Throwable ex) { logger.error("Failed to close indexSearcher", ex); } try { if (obj.onClose != null) { obj.onClose.run(); } } catch (Throwable ex) { logger.error("Failed to close onClose", ex); } } @Override public Drop fork() { return this; } @Override public void attach(ShardedIndexSearchers obj) { } }; private List indexSearchers; private List indexSearchersVals; private Runnable onClose; public ShardedIndexSearchers(List> indexSearchers, Runnable onClose) { super(DROP); this.indexSearchers = new ArrayList<>(indexSearchers.size()); this.indexSearchersVals = new ArrayList<>(indexSearchers.size()); for (Send llIndexSearcher : indexSearchers) { var indexSearcher = llIndexSearcher.receive(); this.indexSearchers.add(indexSearcher); this.indexSearchersVals.add(indexSearcher.getIndexSearcher()); } this.onClose = onClose; } @Override public List shards() { if (!isOwned()) { throw attachTrace(new IllegalStateException("ShardedIndexSearchers must be owned to be used")); } return Collections.unmodifiableList(indexSearchersVals); } @Override public IndexSearcher shard(int shardIndex) { if (!isOwned()) { throw attachTrace(new IllegalStateException("ShardedIndexSearchers must be owned to be used")); } if (shardIndex < 0) { throw new IndexOutOfBoundsException("Shard index " + shardIndex + " is invalid"); } return indexSearchersVals.get(shardIndex); } @Override public IndexReader allShards() { if (!isOwned()) { throw attachTrace(new IllegalStateException("ShardedIndexSearchers must be owned to be used")); } var irs = new IndexReader[indexSearchersVals.size()]; for (int i = 0, s = indexSearchersVals.size(); i < s; i++) { irs[i] = indexSearchersVals.get(i).getIndexReader(); } Object2IntOpenHashMap indexes = new Object2IntOpenHashMap<>(); for (int i = 0; i < irs.length; i++) { indexes.put(irs[i], i); } try { return new MultiReader(irs, Comparator.comparingInt(indexes::getInt), false); } catch (IOException ex) { // This shouldn't happen throw new UncheckedIOException(ex); } } @Override protected RuntimeException createResourceClosedException() { return new IllegalStateException("Closed"); } @Override protected Owned prepareSend() { List> indexSearchers = new ArrayList<>(this.indexSearchers.size()); for (LLIndexSearcher indexSearcher : this.indexSearchers) { indexSearchers.add(indexSearcher.send()); } var onClose = this.onClose; return drop -> { var instance = new ShardedIndexSearchers(indexSearchers, onClose); drop.attach(instance); return instance; }; } protected void makeInaccessible() { this.indexSearchers = null; this.indexSearchersVals = null; this.onClose = null; } } }