CavalliumDBEngine/src/main/java/it/cavallium/dbengine/database/disk/LLLocalLuceneIndex.java

715 lines
25 KiB
Java
Raw Normal View History

2020-12-07 22:15:18 +01:00
package it.cavallium.dbengine.database.disk;
2021-02-28 10:57:16 +01:00
import com.google.common.base.Suppliers;
2021-03-02 01:53:36 +01:00
import it.cavallium.dbengine.client.query.QueryParser;
import it.cavallium.dbengine.client.query.current.data.QueryParams;
2021-02-28 14:52:11 +01:00
import it.cavallium.dbengine.database.EnglishItalianStopFilter;
import it.cavallium.dbengine.database.LLDocument;
import it.cavallium.dbengine.database.LLKeyScore;
import it.cavallium.dbengine.database.LLLuceneIndex;
import it.cavallium.dbengine.database.LLSearchCollectionStatisticsGetter;
import it.cavallium.dbengine.database.LLSearchResult;
2021-03-03 15:03:25 +01:00
import it.cavallium.dbengine.database.LLSignal;
import it.cavallium.dbengine.database.LLSnapshot;
import it.cavallium.dbengine.database.LLTerm;
2021-03-03 15:03:25 +01:00
import it.cavallium.dbengine.database.LLTotalHitsCount;
import it.cavallium.dbengine.database.LLUtils;
import it.cavallium.dbengine.lucene.LuceneUtils;
2021-02-03 13:48:30 +01:00
import it.cavallium.dbengine.lucene.ScheduledTaskLifecycle;
import it.cavallium.dbengine.lucene.analyzer.TextFieldsAnalyzer;
import it.cavallium.dbengine.lucene.analyzer.TextFieldsSimilarity;
2021-01-30 22:14:48 +01:00
import it.cavallium.dbengine.lucene.searcher.AdaptiveStreamSearcher;
import it.cavallium.dbengine.lucene.searcher.AllowOnlyQueryParsingCollectorStreamSearcher;
2021-01-30 22:14:48 +01:00
import it.cavallium.dbengine.lucene.searcher.LuceneStreamSearcher;
2021-03-03 00:13:57 +01:00
import it.cavallium.dbengine.lucene.searcher.LuceneStreamSearcher.HandleResult;
2020-12-07 22:15:18 +01:00
import java.io.IOException;
import java.nio.file.Path;
import java.time.Duration;
import java.util.HashMap;
2020-12-07 22:15:18 +01:00
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
2020-12-07 22:15:18 +01:00
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
2021-03-03 15:03:25 +01:00
import java.util.concurrent.Semaphore;
2020-12-07 22:15:18 +01:00
import java.util.concurrent.TimeUnit;
2021-03-03 15:03:25 +01:00
import java.util.concurrent.atomic.AtomicBoolean;
2020-12-07 22:15:18 +01:00
import java.util.concurrent.atomic.AtomicLong;
2021-03-03 10:57:45 +01:00
import java.util.function.Function;
2021-02-28 10:57:16 +01:00
import java.util.function.Supplier;
2020-12-07 22:15:18 +01:00
import org.apache.lucene.index.IndexCommit;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.KeepOnlyLastCommitDeletionPolicy;
import org.apache.lucene.index.SnapshotDeletionPolicy;
import org.apache.lucene.queries.mlt.MoreLikeThis;
import org.apache.lucene.search.BooleanClause.Occur;
import org.apache.lucene.search.BooleanQuery;
2021-02-27 19:05:13 +01:00
import org.apache.lucene.search.ConstantScoreQuery;
2020-12-07 22:15:18 +01:00
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreMode;
2020-12-07 22:15:18 +01:00
import org.apache.lucene.search.SearcherManager;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.similarities.Similarity;
2021-02-28 14:52:11 +01:00
import org.apache.lucene.search.similarities.TFIDFSimilarity;
2020-12-07 22:15:18 +01:00
import org.apache.lucene.store.Directory;
2020-12-31 20:10:47 +01:00
import org.apache.lucene.store.FSDirectory;
2020-12-07 22:15:18 +01:00
import org.jetbrains.annotations.Nullable;
2021-02-25 00:00:16 +01:00
import org.warp.commonutils.log.Logger;
import org.warp.commonutils.log.LoggerFactory;
2020-12-07 22:15:18 +01:00
import reactor.core.publisher.Flux;
2021-03-03 21:32:45 +01:00
import reactor.core.publisher.FluxSink.OverflowStrategy;
import reactor.core.publisher.GroupedFlux;
2020-12-07 22:15:18 +01:00
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Scheduler;
2020-12-07 22:15:18 +01:00
import reactor.core.scheduler.Schedulers;
import reactor.util.function.Tuple2;
2020-12-07 22:15:18 +01:00
import reactor.util.function.Tuples;
public class LLLocalLuceneIndex implements LLLuceneIndex {
2021-02-25 00:00:16 +01:00
protected static final Logger logger = LoggerFactory.getLogger(LLLocalLuceneIndex.class);
2020-12-07 22:15:18 +01:00
private static final LuceneStreamSearcher streamSearcher = new AdaptiveStreamSearcher();
private static final AllowOnlyQueryParsingCollectorStreamSearcher allowOnlyQueryParsingCollectorStreamSearcher = new AllowOnlyQueryParsingCollectorStreamSearcher();
/**
* Global lucene index scheduler.
* There is only a single thread globally to not overwhelm the disk with
2021-02-03 14:37:02 +01:00
* concurrent commits or concurrent refreshes.
*/
2021-03-03 10:57:45 +01:00
private static final Scheduler luceneHeavyTasksScheduler = Schedulers.newBoundedElastic(1,
2021-02-03 13:48:30 +01:00
Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE,
2021-02-28 10:57:16 +01:00
"lucene",
2021-03-03 10:57:45 +01:00
Integer.MAX_VALUE,
2021-02-03 13:48:30 +01:00
true
);
2021-03-03 10:57:45 +01:00
private final Scheduler luceneBlockingScheduler;
private static final Function<String, Scheduler> boundedSchedulerSupplier = name -> Schedulers.newBoundedElastic(Schedulers.DEFAULT_BOUNDED_ELASTIC_SIZE,
Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE,
"lucene-" + name,
60
);
2021-03-06 17:28:33 +01:00
private final Supplier<Scheduler> lowMemorySchedulerSupplier = Suppliers.memoize(() ->
Schedulers.newBoundedElastic(1, Schedulers.DEFAULT_BOUNDED_ELASTIC_QUEUESIZE, "lucene-low-memory", Integer.MAX_VALUE))::get;
2021-03-06 17:28:33 +01:00
private final Supplier<Scheduler> querySchedulerSupplier = Suppliers.memoize(() -> boundedSchedulerSupplier.apply("query"))::get;
private final Supplier<Scheduler> blockingSchedulerSupplier = Suppliers.memoize(() -> boundedSchedulerSupplier.apply("blocking"))::get;
private final Supplier<Scheduler> blockingLuceneSearchSchedulerSupplier = Suppliers.memoize(() -> boundedSchedulerSupplier.apply("search-blocking"))::get;
/**
* Lucene query scheduler.
*/
2021-02-28 10:57:16 +01:00
private final Scheduler luceneQueryScheduler;
2021-03-04 22:01:50 +01:00
private final Scheduler blockingLuceneSearchScheduler;
2020-12-07 22:15:18 +01:00
private final String luceneIndexName;
private final SnapshotDeletionPolicy snapshotter;
private final IndexWriter indexWriter;
private final SearcherManager searcherManager;
private final Directory directory;
/**
* Last snapshot sequence number. 0 is not used
*/
private final AtomicLong lastSnapshotSeqNo = new AtomicLong(0);
/**
2021-03-02 01:53:36 +01:00
* LLSnapshot seq no to index commit point
2020-12-07 22:15:18 +01:00
*/
private final ConcurrentHashMap<Long, LuceneIndexSnapshot> snapshots = new ConcurrentHashMap<>();
private final boolean lowMemory;
private final TextFieldsSimilarity similarity;
2020-12-07 22:15:18 +01:00
private final ScheduledTaskLifecycle scheduledTasksLifecycle;
private final @Nullable LLSearchCollectionStatisticsGetter distributedCollectionStatisticsGetter;
2020-12-07 22:15:18 +01:00
public LLLocalLuceneIndex(Path luceneBasePath,
String name,
TextFieldsAnalyzer analyzer,
TextFieldsSimilarity similarity,
2020-12-07 22:15:18 +01:00
Duration queryRefreshDebounceTime,
Duration commitDebounceTime,
boolean lowMemory,
@Nullable LLSearchCollectionStatisticsGetter distributedCollectionStatisticsGetter) throws IOException {
2020-12-07 22:15:18 +01:00
if (name.length() == 0) {
throw new IOException("Empty lucene database name");
}
Path directoryPath = luceneBasePath.resolve(name + ".lucene.db");
2020-12-31 20:10:47 +01:00
this.directory = FSDirectory.open(directoryPath);
2020-12-07 22:15:18 +01:00
this.luceneIndexName = name;
this.snapshotter = new SnapshotDeletionPolicy(new KeepOnlyLastCommitDeletionPolicy());
this.lowMemory = lowMemory;
this.similarity = similarity;
this.distributedCollectionStatisticsGetter = distributedCollectionStatisticsGetter;
2020-12-07 22:15:18 +01:00
IndexWriterConfig indexWriterConfig = new IndexWriterConfig(LuceneUtils.getAnalyzer(analyzer));
indexWriterConfig.setOpenMode(IndexWriterConfig.OpenMode.CREATE_OR_APPEND);
indexWriterConfig.setIndexDeletionPolicy(snapshotter);
indexWriterConfig.setCommitOnClose(true);
if (lowMemory) {
indexWriterConfig.setRAMBufferSizeMB(32);
indexWriterConfig.setRAMPerThreadHardLimitMB(32);
} else {
indexWriterConfig.setRAMBufferSizeMB(128);
2021-03-20 12:41:11 +01:00
//indexWriterConfig.setRAMPerThreadHardLimitMB(512);
2020-12-07 22:15:18 +01:00
}
indexWriterConfig.setSimilarity(getSimilarity());
2020-12-07 22:15:18 +01:00
this.indexWriter = new IndexWriter(directory, indexWriterConfig);
this.searcherManager = new SearcherManager(indexWriter, false, false, null);
2021-02-28 10:57:16 +01:00
if (lowMemory) {
2021-03-03 10:57:45 +01:00
this.luceneQueryScheduler = this.luceneBlockingScheduler = lowMemorySchedulerSupplier.get();
2021-02-28 10:57:16 +01:00
} else {
2021-03-03 10:57:45 +01:00
this.luceneBlockingScheduler = blockingSchedulerSupplier.get();
this.luceneQueryScheduler = querySchedulerSupplier.get();
2021-02-28 10:57:16 +01:00
}
2021-03-04 22:01:50 +01:00
this.blockingLuceneSearchScheduler = blockingLuceneSearchSchedulerSupplier.get();
// Create scheduled tasks lifecycle manager
this.scheduledTasksLifecycle = new ScheduledTaskLifecycle();
// Start scheduled tasks
registerScheduledFixedTask(this::scheduledCommit, commitDebounceTime);
registerScheduledFixedTask(this::scheduledQueryRefresh, queryRefreshDebounceTime);
}
private Similarity getSimilarity() {
return LuceneUtils.getSimilarity(similarity);
}
private void registerScheduledFixedTask(Runnable task, Duration duration) {
2021-03-03 10:57:45 +01:00
scheduledTasksLifecycle.registerScheduledTask(luceneHeavyTasksScheduler.schedulePeriodically(() -> {
2020-12-31 12:05:04 +01:00
scheduledTasksLifecycle.startScheduledTask();
try {
task.run();
} finally {
scheduledTasksLifecycle.endScheduledTask();
}
}, duration.toMillis(), duration.toMillis(), TimeUnit.MILLISECONDS));
2020-12-07 22:15:18 +01:00
}
@Override
public String getLuceneIndexName() {
return luceneIndexName;
}
@Override
public Mono<LLSnapshot> takeSnapshot() {
2021-03-03 10:57:45 +01:00
return takeLuceneSnapshot()
.flatMap(snapshot -> Mono
.fromCallable(() -> {
2021-03-03 10:57:45 +01:00
var snapshotSeqNo = lastSnapshotSeqNo.incrementAndGet();
this.snapshots.put(snapshotSeqNo, new LuceneIndexSnapshot(snapshot));
return new LLSnapshot(snapshotSeqNo);
})
2021-02-03 14:37:02 +01:00
.subscribeOn(luceneBlockingScheduler)
2021-03-03 10:57:45 +01:00
);
2020-12-07 22:15:18 +01:00
}
/**
* Use internally. This method commits before taking the snapshot if there are no commits in a new database,
* avoiding the exception.
*/
private Mono<IndexCommit> takeLuceneSnapshot() {
2021-02-03 14:37:02 +01:00
return Mono
2021-03-03 10:57:45 +01:00
.fromCallable(snapshotter::snapshot)
.subscribeOn(luceneBlockingScheduler)
.onErrorResume(ex -> Mono
.defer(() -> {
if (ex instanceof IllegalStateException && "No index commit to snapshot".equals(ex.getMessage())) {
return Mono.fromCallable(() -> {
//noinspection BlockingMethodInNonBlockingContext
indexWriter.commit();
//noinspection BlockingMethodInNonBlockingContext
return snapshotter.snapshot();
}).subscribeOn(luceneHeavyTasksScheduler);
} else {
return Mono.error(ex);
}
})
);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<Void> releaseSnapshot(LLSnapshot snapshot) {
return Mono.<Void>fromCallable(() -> {
var indexSnapshot = this.snapshots.remove(snapshot.getSequenceNumber());
if (indexSnapshot == null) {
2021-03-02 01:53:36 +01:00
throw new IOException("LLSnapshot " + snapshot.getSequenceNumber() + " not found!");
}
2020-12-07 22:15:18 +01:00
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexSnapshot.close();
2020-12-07 22:15:18 +01:00
var luceneIndexSnapshot = indexSnapshot.getSnapshot();
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
snapshotter.release(luceneIndexSnapshot);
// Delete unused files after releasing the snapshot
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.deleteUnusedFiles();
return null;
2021-02-03 14:37:02 +01:00
}).subscribeOn(luceneBlockingScheduler);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<Void> addDocument(LLTerm key, LLDocument doc) {
return Mono.<Void>fromCallable(() -> {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.addDocument(LLUtils.toDocument(doc));
return null;
2021-02-03 14:37:02 +01:00
}).subscribeOn(luceneBlockingScheduler);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<Void> addDocuments(Flux<GroupedFlux<LLTerm, LLDocument>> documents) {
return documents
.flatMap(group -> group
.collectList()
.flatMap(docs -> Mono
.<Void>fromCallable(() -> {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.addDocuments(LLUtils.toDocuments(docs));
return null;
})
2021-02-03 14:37:02 +01:00
.subscribeOn(luceneBlockingScheduler))
)
.then();
2020-12-07 22:15:18 +01:00
}
2020-12-07 22:15:18 +01:00
@Override
public Mono<Void> deleteDocument(LLTerm id) {
return Mono.<Void>fromCallable(() -> {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.deleteDocuments(LLUtils.toTerm(id));
return null;
2021-02-03 14:37:02 +01:00
}).subscribeOn(luceneBlockingScheduler);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<Void> updateDocument(LLTerm id, LLDocument document) {
return Mono.<Void>fromCallable(() -> {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.updateDocument(LLUtils.toTerm(id), LLUtils.toDocument(document));
return null;
2021-02-03 14:37:02 +01:00
}).subscribeOn(luceneBlockingScheduler);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<Void> updateDocuments(Flux<GroupedFlux<LLTerm, LLDocument>> documents) {
return documents.flatMap(this::updateDocuments).then();
}
private Mono<Void> updateDocuments(GroupedFlux<LLTerm, LLDocument> documents) {
return documents
.map(LLUtils::toDocument)
.collectList()
.flatMap(luceneDocuments -> Mono
.<Void>fromCallable(() -> {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.updateDocuments(LLUtils.toTerm(documents.key()), luceneDocuments);
return null;
})
2021-02-03 14:37:02 +01:00
.subscribeOn(luceneBlockingScheduler)
);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<Void> deleteAll() {
return Mono.<Void>fromCallable(() -> {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.deleteAll();
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.forceMergeDeletes(true);
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.commit();
return null;
2021-03-03 10:57:45 +01:00
}).subscribeOn(luceneHeavyTasksScheduler);
2020-12-07 22:15:18 +01:00
}
private Mono<IndexSearcher> acquireSearcherWrapper(LLSnapshot snapshot, boolean distributedPre, long actionId) {
return Mono.fromCallable(() -> {
IndexSearcher indexSearcher;
if (snapshot == null) {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexSearcher = searcherManager.acquire();
indexSearcher.setSimilarity(getSimilarity());
} else {
indexSearcher = resolveSnapshot(snapshot).getIndexSearcher();
}
if (distributedCollectionStatisticsGetter != null && actionId != -1) {
return new LLIndexSearcherWithCustomCollectionStatistics(indexSearcher,
distributedCollectionStatisticsGetter,
distributedPre,
actionId
);
} else {
return indexSearcher;
}
2021-02-03 14:37:02 +01:00
}).subscribeOn(luceneBlockingScheduler);
}
2020-12-07 22:15:18 +01:00
private Mono<Void> releaseSearcherWrapper(LLSnapshot snapshot, IndexSearcher indexSearcher) {
return Mono.<Void>fromRunnable(() -> {
if (snapshot == null) {
try {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
searcherManager.release(indexSearcher);
} catch (IOException e) {
e.printStackTrace();
}
}
2021-02-03 14:37:02 +01:00
}).subscribeOn(luceneBlockingScheduler);
2020-12-07 22:15:18 +01:00
}
@Override
public Mono<LLSearchResult> moreLikeThis(@Nullable LLSnapshot snapshot,
2021-03-02 01:53:36 +01:00
QueryParams queryParams,
String keyFieldName,
Flux<Tuple2<String, Set<String>>> mltDocumentFieldsFlux) {
2021-02-28 16:50:59 +01:00
return moreLikeThis(snapshot,
2021-03-02 01:53:36 +01:00
queryParams,
2021-02-28 16:50:59 +01:00
keyFieldName,
2021-03-02 01:53:36 +01:00
mltDocumentFieldsFlux,
2021-02-28 16:50:59 +01:00
false,
0,
1
);
}
public Mono<LLSearchResult> distributedMoreLikeThis(@Nullable LLSnapshot snapshot,
2021-03-02 01:53:36 +01:00
QueryParams queryParams,
String keyFieldName,
2021-03-02 01:53:36 +01:00
Flux<Tuple2<String, Set<String>>> mltDocumentFieldsFlux,
long actionId,
int scoreDivisor) {
2021-02-28 16:50:59 +01:00
return moreLikeThis(snapshot,
2021-03-02 01:53:36 +01:00
queryParams,
2021-02-28 16:50:59 +01:00
keyFieldName,
2021-03-02 01:53:36 +01:00
mltDocumentFieldsFlux,
2021-02-28 16:50:59 +01:00
false,
actionId,
scoreDivisor
);
}
public Mono<Void> distributedPreMoreLikeThis(@Nullable LLSnapshot snapshot,
2021-03-02 01:53:36 +01:00
QueryParams queryParams,
String keyFieldName,
Flux<Tuple2<String, Set<String>>> mltDocumentFieldsFlux,
2021-03-02 01:53:36 +01:00
long actionId) {
2021-02-28 16:50:59 +01:00
return moreLikeThis(snapshot,
2021-03-02 01:53:36 +01:00
queryParams,
2021-02-28 16:50:59 +01:00
keyFieldName,
2021-03-02 01:53:36 +01:00
mltDocumentFieldsFlux,
2021-02-28 16:50:59 +01:00
true,
actionId,
1
)
.flatMap(LLSearchResult::completion);
}
@SuppressWarnings({"unchecked", "rawtypes"})
private Mono<LLSearchResult> moreLikeThis(@Nullable LLSnapshot snapshot,
2021-03-02 01:53:36 +01:00
QueryParams queryParams,
String keyFieldName,
2021-03-02 01:53:36 +01:00
Flux<Tuple2<String, Set<String>>> mltDocumentFieldsFlux,
boolean doDistributedPre,
long actionId,
int scoreDivisor) {
Query luceneAdditionalQuery;
try {
2021-03-02 01:53:36 +01:00
luceneAdditionalQuery = QueryParser.toQuery(queryParams.getQuery());
} catch (Exception e) {
return Mono.error(e);
}
return mltDocumentFieldsFlux
.collectMap(Tuple2::getT1, Tuple2::getT2, HashMap::new)
.flatMap(mltDocumentFields -> {
2021-02-27 19:05:13 +01:00
mltDocumentFields.entrySet().removeIf(entry -> entry.getValue().isEmpty());
if (mltDocumentFields.isEmpty()) {
return Mono.just(LLSearchResult.empty());
}
2021-02-28 14:52:11 +01:00
return acquireSearcherWrapper(snapshot, doDistributedPre, actionId).flatMap(indexSearcher -> Mono
.fromCallable(() -> {
var mlt = new MoreLikeThis(indexSearcher.getIndexReader());
mlt.setAnalyzer(indexWriter.getAnalyzer());
mlt.setFieldNames(mltDocumentFields.keySet().toArray(String[]::new));
mlt.setMinTermFreq(1);
mlt.setMinDocFreq(3);
mlt.setMaxDocFreqPct(20);
2021-03-02 01:53:36 +01:00
mlt.setBoost(QueryParser.isScoringEnabled(queryParams));
2021-02-28 14:52:11 +01:00
mlt.setStopWords(EnglishItalianStopFilter.getStopWordsString());
var similarity = getSimilarity();
if (similarity instanceof TFIDFSimilarity) {
mlt.setSimilarity((TFIDFSimilarity) similarity);
} else {
logger.trace("Using an unsupported similarity algorithm for MoreLikeThis: {}. You must use a similarity instance based on TFIDFSimilarity!", similarity);
}
// Get the reference doc and apply it to MoreLikeThis, to generate the query
//noinspection BlockingMethodInNonBlockingContext
var mltQuery = mlt.like((Map) mltDocumentFields);
Query luceneQuery;
if (luceneAdditionalQuery != null) {
luceneQuery = new BooleanQuery.Builder()
.add(mltQuery, Occur.MUST)
.add(new ConstantScoreQuery(luceneAdditionalQuery), Occur.MUST)
.build();
} else {
luceneQuery = mltQuery;
}
2021-03-03 20:00:58 +01:00
return luceneQuery;
2021-02-28 14:52:11 +01:00
})
.subscribeOn(luceneQueryScheduler)
2021-03-03 20:00:58 +01:00
.map(luceneQuery -> luceneSearch(doDistributedPre,
indexSearcher,
queryParams.getLimit(),
queryParams.getMinCompetitiveScore().getNullable(),
keyFieldName,
scoreDivisor,
luceneQuery,
QueryParser.toSort(queryParams.getSort()),
QueryParser.toScoreMode(queryParams.getScoreMode())
))
2021-02-28 14:52:11 +01:00
.materialize()
.flatMap(signal -> {
if (signal.isOnComplete() || signal.isOnError()) {
return releaseSearcherWrapper(snapshot, indexSearcher).thenReturn(signal);
} else {
return Mono.just(signal);
}
2021-03-06 17:28:33 +01:00
}).dematerialize());
});
2020-12-07 22:15:18 +01:00
}
private LLKeyScore fixKeyScore(LLKeyScore keyScore, int scoreDivisor) {
return scoreDivisor == 1 ? keyScore : new LLKeyScore(keyScore.getKey(), keyScore.getScore() / (float) scoreDivisor);
}
2020-12-07 22:15:18 +01:00
@Override
2021-03-02 01:53:36 +01:00
public Mono<LLSearchResult> search(@Nullable LLSnapshot snapshot, QueryParams queryParams, String keyFieldName) {
return search(snapshot, queryParams, keyFieldName, false, 0, 1);
}
2021-03-02 01:53:36 +01:00
public Mono<LLSearchResult> distributedSearch(@Nullable LLSnapshot snapshot, QueryParams queryParams, String keyFieldName, long actionId, int scoreDivisor) {
return search(snapshot, queryParams, keyFieldName, false, actionId, scoreDivisor);
}
2021-03-02 01:53:36 +01:00
public Mono<Void> distributedPreSearch(@Nullable LLSnapshot snapshot, QueryParams queryParams, String keyFieldName, long actionId) {
2021-02-14 13:46:11 +01:00
return this
2021-03-02 01:53:36 +01:00
.search(snapshot, queryParams, keyFieldName, true, actionId, 1)
.flatMap(LLSearchResult::completion);
}
private Mono<LLSearchResult> search(@Nullable LLSnapshot snapshot,
2021-03-02 01:53:36 +01:00
QueryParams queryParams, String keyFieldName,
boolean doDistributedPre, long actionId, int scoreDivisor) {
return acquireSearcherWrapper(snapshot, doDistributedPre, actionId)
.flatMap(indexSearcher -> Mono
.fromCallable(() -> {
2021-03-02 01:53:36 +01:00
Objects.requireNonNull(queryParams.getScoreMode(), "ScoreMode must not be null");
Query luceneQuery = QueryParser.toQuery(queryParams.getQuery());
Sort luceneSort = QueryParser.toSort(queryParams.getSort());
org.apache.lucene.search.ScoreMode luceneScoreMode = QueryParser.toScoreMode(queryParams.getScoreMode());
2021-02-03 13:48:30 +01:00
return Tuples.of(luceneQuery, Optional.ofNullable(luceneSort), luceneScoreMode);
})
.subscribeOn(luceneQueryScheduler)
.flatMap(tuple -> Mono
2021-03-03 20:00:58 +01:00
.fromSupplier(() -> {
2021-02-03 13:48:30 +01:00
Query luceneQuery = tuple.getT1();
Sort luceneSort = tuple.getT2().orElse(null);
ScoreMode luceneScoreMode = tuple.getT3();
2021-02-27 19:05:13 +01:00
return luceneSearch(doDistributedPre,
indexSearcher,
2021-03-02 01:53:36 +01:00
queryParams.getLimit(),
queryParams.getMinCompetitiveScore().getNullable(),
2021-02-27 19:05:13 +01:00
keyFieldName,
scoreDivisor,
luceneQuery,
luceneSort,
luceneScoreMode
2021-02-25 00:00:16 +01:00
);
2021-03-03 20:00:58 +01:00
})
)
.materialize()
2021-02-17 13:59:35 +01:00
.flatMap(signal -> {
if (signal.isOnComplete() || signal.isOnError()) {
return releaseSearcherWrapper(snapshot, indexSearcher).thenReturn(signal);
} else {
return Mono.just(signal);
}
})
2021-03-06 17:28:33 +01:00
.dematerialize()
);
2020-12-07 22:15:18 +01:00
}
2021-02-27 19:05:13 +01:00
private LLSearchResult luceneSearch(boolean doDistributedPre,
IndexSearcher indexSearcher,
long limit,
@Nullable Float minCompetitiveScore,
String keyFieldName,
int scoreDivisor,
Query luceneQuery,
Sort luceneSort,
ScoreMode luceneScoreMode) {
2021-03-03 20:00:58 +01:00
return new LLSearchResult(Flux.just(Flux.defer(() -> Flux.<LLSignal>create(sink -> {
2021-03-03 15:03:25 +01:00
AtomicBoolean cancelled = new AtomicBoolean();
2021-03-24 00:02:47 +01:00
Semaphore requests = new Semaphore(0);
2021-03-03 15:03:25 +01:00
sink.onDispose(() -> {
cancelled.set(true);
});
2021-03-03 20:00:58 +01:00
sink.onCancel(() -> {
cancelled.set(true);
});
2021-03-03 15:03:25 +01:00
sink.onRequest(delta -> {
2021-03-24 00:02:47 +01:00
requests.release((int) Math.min(delta, Integer.MAX_VALUE));
2021-03-03 15:03:25 +01:00
});
2021-02-27 19:05:13 +01:00
2021-03-03 21:32:45 +01:00
try {
2021-03-06 17:28:33 +01:00
if (!cancelled.get()) {
if (doDistributedPre) {
//noinspection BlockingMethodInNonBlockingContext
allowOnlyQueryParsingCollectorStreamSearcher.search(indexSearcher, luceneQuery);
sink.next(new LLTotalHitsCount(0L));
} else {
int boundedLimit = Math.max(0, limit > Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) limit);
//noinspection BlockingMethodInNonBlockingContext
streamSearcher.search(indexSearcher,
luceneQuery,
boundedLimit,
luceneSort,
luceneScoreMode,
minCompetitiveScore,
keyFieldName,
keyScore -> {
try {
if (cancelled.get()) {
return HandleResult.HALT;
}
2021-03-24 00:02:47 +01:00
while (!requests.tryAcquire(500, TimeUnit.MILLISECONDS)) {
2021-03-06 17:28:33 +01:00
if (cancelled.get()) {
2021-03-03 20:00:58 +01:00
return HandleResult.HALT;
}
2021-03-06 17:28:33 +01:00
}
sink.next(fixKeyScore(keyScore, scoreDivisor));
2021-03-24 00:02:47 +01:00
if (cancelled.get()) {
return HandleResult.HALT;
} else {
return HandleResult.CONTINUE;
}
2021-03-06 17:28:33 +01:00
} catch (Exception ex) {
sink.error(ex);
cancelled.set(true);
2021-03-24 00:02:47 +01:00
requests.release(Integer.MAX_VALUE);
2021-03-06 17:28:33 +01:00
return HandleResult.HALT;
}
},
totalHitsCount -> {
try {
if (cancelled.get()) {
return;
}
2021-03-24 00:02:47 +01:00
while (!requests.tryAcquire(500, TimeUnit.MILLISECONDS)) {
2021-03-06 17:28:33 +01:00
if (cancelled.get()) {
return;
2021-03-03 20:00:58 +01:00
}
2021-03-03 15:03:25 +01:00
}
2021-03-06 17:28:33 +01:00
sink.next(new LLTotalHitsCount(totalHitsCount));
} catch (Exception ex) {
sink.error(ex);
cancelled.set(true);
2021-03-24 00:02:47 +01:00
requests.release(Integer.MAX_VALUE);
2021-03-06 17:28:33 +01:00
}
}
);
2021-03-03 15:03:25 +01:00
}
2021-03-06 17:28:33 +01:00
sink.complete();
}
2021-03-03 21:32:45 +01:00
} catch (Exception ex) {
sink.error(ex);
}
2021-03-06 17:28:33 +01:00
}, OverflowStrategy.ERROR).subscribeOn(blockingLuceneSearchScheduler).publishOn(luceneQueryScheduler))));
2021-02-27 19:05:13 +01:00
}
2020-12-07 22:15:18 +01:00
@Override
public Mono<Void> close() {
return Mono
.<Void>fromCallable(() -> {
2021-03-19 20:55:38 +01:00
logger.debug("Closing IndexWriter...");
2021-03-04 22:01:50 +01:00
this.blockingLuceneSearchScheduler.dispose();
scheduledTasksLifecycle.cancelAndWait();
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
indexWriter.close();
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
directory.close();
2021-03-19 20:55:38 +01:00
logger.debug("IndexWriter closed");
return null;
})
2021-03-03 10:57:45 +01:00
.subscribeOn(luceneHeavyTasksScheduler);
2020-12-07 22:15:18 +01:00
}
2021-02-03 13:48:30 +01:00
@Override
public Mono<Void> flush() {
return Mono
.<Void>fromCallable(() -> {
scheduledTasksLifecycle.startScheduledTask();
try {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
2021-02-03 13:48:30 +01:00
indexWriter.commit();
} finally {
scheduledTasksLifecycle.endScheduledTask();
}
return null;
})
2021-03-03 10:57:45 +01:00
.subscribeOn(luceneHeavyTasksScheduler);
2021-02-03 13:48:30 +01:00
}
@Override
public Mono<Void> refresh() {
return Mono
.<Void>fromCallable(() -> {
scheduledTasksLifecycle.startScheduledTask();
try {
2021-02-03 14:37:02 +01:00
//noinspection BlockingMethodInNonBlockingContext
2021-02-03 13:48:30 +01:00
searcherManager.maybeRefreshBlocking();
} finally {
scheduledTasksLifecycle.endScheduledTask();
}
return null;
})
2021-03-03 10:57:45 +01:00
.subscribeOn(luceneHeavyTasksScheduler);
2021-02-03 13:48:30 +01:00
}
2020-12-07 22:15:18 +01:00
private void scheduledCommit() {
try {
if (indexWriter.hasUncommittedChanges()) {
indexWriter.commit();
}
} catch (IOException ex) {
ex.printStackTrace();
}
}
@SuppressWarnings("unused")
2020-12-07 22:15:18 +01:00
private void scheduledQueryRefresh() {
try {
boolean refreshStarted = searcherManager.maybeRefresh();
// if refreshStarted == false, another thread is currently already refreshing
2020-12-07 22:15:18 +01:00
} catch (IOException ex) {
ex.printStackTrace();
}
}
private LuceneIndexSnapshot resolveSnapshot(@Nullable LLSnapshot snapshot) {
if (snapshot == null) {
return null;
}
return Objects.requireNonNull(snapshots.get(snapshot.getSequenceNumber()),
() -> "Can't resolve snapshot " + snapshot.getSequenceNumber()
);
}
@Override
public boolean isLowMemoryMode() {
return lowMemory;
}
}