remove TagsToFile

Remove one layer of abstraction by moving the code into the DataStore.

data-store/src/main/java/org/lucares/pdb/datastore/internal/DataStore.java

@@ -12,8 +12,8 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.SortedSet;
 import java.util.TreeSet;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 
 import org.lucares.collections.LongList;
 import org.lucares.pdb.api.RuntimeIOException;
@@ -74,6 +74,8 @@ public class DataStore implements AutoCloseable {
 	// easily.
 	private final HotEntryCache<Long, Doc> docIdToDocCache = new HotEntryCache<>(Duration.ofMillis(30), 100_000);
 
+	private final HotEntryCache<Tags, PdbWriter> writerCache;
+
 	private final DiskStorage diskStorage;
 	private final Path diskStorageFilePath;
 	private final Path storageBasePath;
@@ -102,6 +104,9 @@ public class DataStore implements AutoCloseable {
 		docIdToDoc = new PersistentMap<>(docIdToDocIndexPath, PersistentMap.LONG_CODER, new DocEncoderDecoder());
 
 		queryCompletionIndex = new QueryCompletionIndex(storageBasePath);
+
+		writerCache = new HotEntryCache<>(Duration.ofSeconds(10), 1000);
+		writerCache.addListener((k, v) -> v.close());
 	}
 
 	private Path keyCompressionFile(final Path dataDirectory) throws IOException {
@@ -160,6 +165,30 @@ public class DataStore implements AutoCloseable {
 		return NEXT_DOC_ID.getAndIncrement();
 	}
 
+	public List<PdbFile> getFilesForQuery(final String query) {
+
+		final List<Doc> searchResult = search(query);
+		if (searchResult.size() > 500_000) {
+			throw new IllegalStateException("Too many results.");
+		}
+
+		final List<PdbFile> result = toPdbFiles(searchResult);
+		return result;
+	}
+
+	private List<PdbFile> toPdbFiles(final List<Doc> searchResult) {
+		final List<PdbFile> result = new ArrayList<>(searchResult.size());
+		for (final Doc document : searchResult) {
+
+			final long rootBlockNumber = document.getRootBlockNumber();
+			final Tags tags = document.getTags();
+			final PdbFile pdbFile = new PdbFile(rootBlockNumber, tags);
+
+			result.add(pdbFile);
+		}
+		return result;
+	}
+
 	public List<Doc> search(final String query) {
 		try {
 			final LongList docIdsList = executeQuery(query);
@@ -262,26 +291,13 @@ public class DataStore implements AutoCloseable {
 	}
 
 	private Doc getDocByDocId(final Long docId) {
-		try {
-			return docIdToDocCache.putIfAbsent(docId, () -> {
-				try {
-					return docIdToDoc.getValue(docId);
-				} catch (final IOException e) {
-					throw new RuntimeIOException(e);
-				}
-			});
-		} catch (final ExecutionException e) {
-			throw new RuntimeException(e);
-		}
-	}
-
-	@Override
-	public void close() throws IOException {
-		try {
-			diskStorage.close();
-		} finally {
-			tagToDocsId.close();
-		}
+		return docIdToDocCache.putIfAbsent(docId, () -> {
+			try {
+				return docIdToDoc.getValue(docId);
+			} catch (final IOException e) {
+				throw new RuntimeIOException(e);
+			}
+		});
 	}
 
 	public List<Proposal> propose(final String query, final int caretIndex) {
@@ -296,7 +312,12 @@ public class DataStore implements AutoCloseable {
 		return diskStorage;
 	}
 
-	public PdbWriter getWriter(final Tags tags) {
+	public PdbWriter getWriter(final long dateAsEpochMilli, final Tags tags) throws ReadException, WriteException {
+
+		return writerCache.putIfAbsent(tags, () -> getWriter(tags));
+	}
+
+	private PdbWriter getWriter(final Tags tags) {
 		final Optional<Doc> docsForTags = getByTags(tags);
 		PdbWriter writer;
 		if (docsForTags.isPresent()) {
@@ -336,4 +357,50 @@ public class DataStore implements AutoCloseable {
 		return result;
 	}
 
+	@Override
+	public void close() throws RuntimeIOException {
+		try {
+			// we cannot simply clear the cache, because the cache implementation (Guava at
+			// the time of writing) handles eviction events asynchronously.
+			forEachWriter(cachedWriter -> {
+				try {
+					cachedWriter.close();
+				} catch (final Exception e) {
+					throw new WriteException(e);
+				}
+			});
+		} finally {
+			try {
+				diskStorage.close();
+			} catch (final IOException e) {
+				throw new RuntimeIOException(e);
+			} finally {
+				try {
+					tagToDocsId.close();
+				} catch (final IOException e) {
+					throw new RuntimeIOException(e);
+				}
+			}
+		}
+	}
+
+	private void forEachWriter(final Consumer<PdbWriter> consumer) {
+		writerCache.forEach(writer -> {
+			try {
+				consumer.accept(writer);
+			} catch (final RuntimeException e) {
+				LOGGER.warn("Exception while applying consumer to PdbWriter for " + writer.getPdbFile(), e);
+			}
+		});
+	}
+
+	public void flush() {
+		forEachWriter(t -> {
+			try {
+				t.flush();
+			} catch (final Exception e) {
+				throw new WriteException(e);
+			}
+		});
+	}
 }