Examples with CheckedFunction org.opensearch.common.CheckedFunction used on opensource projects

Example 21 with CheckedFunction

use of org.opensearch.common.CheckedFunction in project OpenSearch by opensearch-project.

the class IndexReaderWrapperTests method testReaderCloseListenerIsCalled.

public void testReaderCloseListenerIsCalled() throws IOException {
    Directory dir = newDirectory();
    IndexWriterConfig iwc = newIndexWriterConfig();
    IndexWriter writer = new IndexWriter(dir, iwc);
    Document doc = new Document();
    doc.add(new StringField("id", "1", random().nextBoolean() ? Field.Store.YES : Field.Store.NO));
    doc.add(new TextField("field", "doc", random().nextBoolean() ? Field.Store.YES : Field.Store.NO));
    writer.addDocument(doc);
    DirectoryReader open = OpenSearchDirectoryReader.wrap(DirectoryReader.open(writer), new ShardId("foo", "_na_", 1));
    IndexSearcher searcher = new IndexSearcher(open);
    assertEquals(1, searcher.search(new TermQuery(new Term("field", "doc")), 1).totalHits.value);
    final AtomicInteger closeCalls = new AtomicInteger(0);
    CheckedFunction<DirectoryReader, DirectoryReader, IOException> wrapper = reader -> new FieldMaskingReader("field", reader, closeCalls);
    final int sourceRefCount = open.getRefCount();
    final AtomicInteger count = new AtomicInteger();
    final AtomicInteger outerCount = new AtomicInteger();
    final AtomicBoolean closeCalled = new AtomicBoolean(false);
    final Engine.Searcher wrap = IndexShard.wrapSearcher(new Engine.Searcher("foo", open, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), () -> closeCalled.set(true)), wrapper);
    assertEquals(1, wrap.getIndexReader().getRefCount());
    OpenSearchDirectoryReader.addReaderCloseListener(wrap.getDirectoryReader(), key -> {
        if (key == open.getReaderCacheHelper().getKey()) {
            count.incrementAndGet();
        }
        outerCount.incrementAndGet();
    });
    assertEquals(0, wrap.search(new TermQuery(new Term("field", "doc")), 1).totalHits.value);
    wrap.close();
    assertFalse("wrapped reader is closed", wrap.getIndexReader().tryIncRef());
    assertEquals(sourceRefCount, open.getRefCount());
    assertTrue(closeCalled.get());
    assertEquals(1, closeCalls.get());
    IOUtils.close(open, writer, dir);
    assertEquals(1, outerCount.get());
    assertEquals(1, count.get());
    assertEquals(0, open.getRefCount());
    assertEquals(1, closeCalls.get());
}

Example 22 with CheckedFunction

use of org.opensearch.common.CheckedFunction in project OpenSearch by opensearch-project.

the class IndexReaderWrapperTests method testIsCacheable.

public void testIsCacheable() throws IOException {
    Directory dir = newDirectory();
    IndexWriterConfig iwc = newIndexWriterConfig();
    IndexWriter writer = new IndexWriter(dir, iwc);
    Document doc = new Document();
    doc.add(new StringField("id", "1", random().nextBoolean() ? Field.Store.YES : Field.Store.NO));
    doc.add(new TextField("field", "doc", random().nextBoolean() ? Field.Store.YES : Field.Store.NO));
    writer.addDocument(doc);
    DirectoryReader open = OpenSearchDirectoryReader.wrap(DirectoryReader.open(writer), new ShardId("foo", "_na_", 1));
    IndexSearcher searcher = new IndexSearcher(open);
    assertEquals(1, searcher.search(new TermQuery(new Term("field", "doc")), 1).totalHits.value);
    searcher.setSimilarity(iwc.getSimilarity());
    final AtomicInteger closeCalls = new AtomicInteger(0);
    CheckedFunction<DirectoryReader, DirectoryReader, IOException> wrapper = reader -> new FieldMaskingReader("field", reader, closeCalls);
    final ConcurrentHashMap<Object, TopDocs> cache = new ConcurrentHashMap<>();
    AtomicBoolean closeCalled = new AtomicBoolean(false);
    try (Engine.Searcher wrap = IndexShard.wrapSearcher(new Engine.Searcher("foo", open, IndexSearcher.getDefaultSimilarity(), IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), () -> closeCalled.set(true)), wrapper)) {
        OpenSearchDirectoryReader.addReaderCloseListener(wrap.getDirectoryReader(), key -> {
            cache.remove(key);
        });
        TopDocs search = wrap.search(new TermQuery(new Term("field", "doc")), 1);
        cache.put(wrap.getIndexReader().getReaderCacheHelper().getKey(), search);
    }
    assertTrue(closeCalled.get());
    assertEquals(1, closeCalls.get());
    assertEquals(1, cache.size());
    IOUtils.close(open, writer, dir);
    assertEquals(0, cache.size());
    assertEquals(1, closeCalls.get());
}