Cached factory thread safety

pom.xml

@@ -11,7 +11,7 @@
 
 	<groupId>org.janelia.saalfeldlab</groupId>
 	<artifactId>n5-universe</artifactId>
-	<version>3.0.3-SNAPSHOT</version>
+	<version>3.1.0-SNAPSHOT</version>
 
 	<name>N5-Universe</name>
 	<description>Utilities spanning all of the N5 repositories</description>

src/main/java/org/janelia/saalfeldlab/n5/universe/N5Factory.java

@@ -34,6 +34,7 @@
 import java.util.List;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
 
 import static org.janelia.saalfeldlab.n5.universe.StorageFormat.*;
@@ -690,6 +691,56 @@ public N5Writer openWriter(@Nullable final StorageFormat storage, final KeyValue
 		return null;
 	}
 
+	/**
+	 * Get a new {@link N5Writer} IFF the container already exists. If one does not already exist,
+	 * a new one will NOT be created.
+	 * This is achieved by first opening a reader at the location, and only then opening
+	 * a writer if the reader successfully opens.
+	 *
+	 * @param storage the storage format, or null
+	 * @param access to the key-value backend
+	 * @param location root URI of the container
+	 * @return a writer for the existing container
+	 *
+	 * @throws N5IOException if the container does not exist, no attempt will be made to create it
+	 */
+	public N5Writer openExistingWriter(@Nullable final StorageFormat storage, final KeyValueAccess access, final URI location) throws N5IOException {
+
+		requireContainerExists(() -> openReader(storage, access, location));
+		return openWriter(storage, access, location);
+	}
+
+	/**
+	 * Get a new {@link N5Writer} IFF the container already exists. If one does not already exist,
+	 * a new one will NOT be created.
+	 * This is achieved by first opening a reader at the location, and only then opening
+	 * a writer if the reader successfully opens.
+	 *
+	 * @param uri the container location, optionally prefixed with a storage format
+	 * @return a writer for the existing container
+	 *
+	 * @throws N5IOException if the container does not exist, no attempt will be made to create it
+	 */
+	public N5Writer openExistingWriter(final String uri) throws N5IOException {
+
+		requireContainerExists(() -> openReader(uri));
+		return openWriter(uri);
+	}
+
+    /**
+	 * Test if a reader can be opened from given supplier
+	 *
+	 * @throws N5IOException if the container does not exist
+     */
+	protected void requireContainerExists(Supplier<N5Reader> getReader) throws N5IOException {
+		//noinspection EmptyTryBlock
+		try (N5Reader ignored = getReader.get()) {
+			/* Just want to verify we can open, then we close to get the writer*/
+		} catch (final Exception e) {
+			throw new N5IOException("Existing container could not be opened, or does not exist.", e);
+		}
+	}
+
     /**
 	 * Try to get a zarr writer at the given location and access.
 	 * If a container exists at the location, load that version as a writer if possible.

src/main/java/org/janelia/saalfeldlab/n5/universe/N5FactoryWithCache.java

@@ -6,62 +6,152 @@
 import org.janelia.saalfeldlab.n5.hdf5.N5HDF5Reader;
 import org.janelia.saalfeldlab.n5.zarr.ZarrKeyValueReader;
 import org.janelia.saalfeldlab.n5.zarr.v3.ZarrV3KeyValueReader;
-import org.jspecify.annotations.Nullable;
 
 import java.net.URI;
 import java.nio.file.Paths;
-import java.util.HashMap;
 import java.util.UUID;
-
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
+
+/**
+ * An {@link N5Factory} that caches opened readers and writers by normalized URI.
+ * <p>
+ * `openReader`/`openWriter` calls on cached containers is non-blocking and thread-safe.
+ * If the container is not cached yet, the open call syncronizes on `this` instance to ensure
+ * thread safe.
+ */
 public class N5FactoryWithCache extends N5Factory {
 
-	private final HashMap<URI, N5Reader> readerCache = new HashMap<>();
-	private final HashMap<URI, N5Writer> writerCache = new HashMap<>();
+	private final ConcurrentHashMap<URI, N5Reader> readerCache = new ConcurrentHashMap<>();
+	private final ConcurrentHashMap<URI, N5Writer> writerCache = new ConcurrentHashMap<>();
 
-	@Override
-	public N5Reader openReader(StorageFormat storage, KeyValueAccess access, URI location) {
+	/**
+	 * Open a reader for the container. If a cached reader is present, return it.
+	 * If no cached reader, but a cached writer, return it as an N5Reader.
+	 * Finally, open a new reader if not present in the cache.
+	 */
+	public N5Reader openReaderAllowCachedWriter(StorageFormat storage, KeyValueAccess access, URI location) {
 
-		final N5Reader reader = getReaderFromCache(storage, location);
-		if (reader != null)
-			return reader;
-		return openAndCacheReader(storage, access, location);
+		N5Reader cachedReader = getCachedReaderAllowWriter(storage, location);
+		if (cachedReader != null)
+			return cachedReader;
+
+		return cachedOpenReader(storage, access, location);
 	}
 
-	private N5Reader openAndCacheReader(StorageFormat storageFormat, KeyValueAccess access, URI uri) {
+	/**
+	 * Open a reader for the container. If a cached reader is present, return it.
+	 * If no cached reader, but a cached writer, return it as an N5Reader.
+	 * Finally, open a new reader if not present in the cache.
+	 */
+	public N5Reader openReaderAllowCachedWriter(String uri) {
 
-		final N5Reader reader = super.openReader(storageFormat, access, uri);
-		if (reader == null)
-			return null;
+		final Pair<StorageFormat, URI> parsed = StorageFormat.parseUri(uri);
+		final StorageFormat storage = parsed.getA();
+		final URI location = parsed.getB();
 
-		URI normalUri = normalizeUri(uri);
-		readerCache.put(normalUri, reader);
-		return reader;
+		N5Reader cachedReader = getCachedReaderAllowWriter(storage, location);
+		if (cachedReader != null)
+			return cachedReader;
+
+		return openReader(uri);
+	}
+
+	private N5Reader getCachedReaderAllowWriter(StorageFormat storage, URI location) {
+		/* try to get a cached reader*/
+		final N5Reader readerFromCache = getReaderFromCache(storage, location);
+		if (readerFromCache != null)
+			return readerFromCache;
+
+		/* else see if we have a cached writer we can return as an N5Reader*/
+		final boolean dontTestWrite = false;
+        return getWriterFromCache(storage, location, dontTestWrite);
+    }
+
+	private N5Reader cachedOpenReader(StorageFormat storage, KeyValueAccess access, URI location) {
+
+		/* if possible, return cached reader without synchronizing */
+		final N5Reader cached = getReaderFromCache(storage, location);
+		if (cached != null)
+			return cached;
+
+		synchronized (this) {
+
+			/* another thread may have opened it while we waited */
+			final N5Reader reopened = getReaderFromCache(storage, location);
+			if (reopened != null)
+				return reopened;
+
+			final N5Reader reader = super.openReader(storage, access, location);
+			if (reader != null)
+				readerCache.put(normalizeUri(location), reader);
+			return reader;
+		}
 	}
 
-	protected synchronized N5Reader getReaderFromCache(StorageFormat format, URI location) {
+	@Override
+	public N5Reader openReader(StorageFormat storage, KeyValueAccess access, URI location) {
+
+		return cachedOpenReader(storage, access, location);
+	}
+
+	protected N5Reader getReaderFromCache(StorageFormat format, URI location) {
 
 		final URI uri = normalizeUri(location);
 		final N5Reader reader = readerCache.get(uri);
 		if (reader == null)
 			return null;
 
-		if (!n5MatchesFormat(reader, format) || !canRead(reader)) {
-			readerCache.remove(uri);
+		boolean readerIsValid = n5MatchesFormat(reader, format) && canRead(reader);
+		if (!readerIsValid) {
+			readerCache.remove(uri, reader);
 			return null;
 		}
-
 		return reader;
 	}
 
 	@Override
 	public N5Writer openWriter(StorageFormat storage, KeyValueAccess access, URI location) {
 
-		final N5Writer writer = getWriterFromCache(storage, location);
-		if (writer != null)
+		final boolean testWrite = true;
+		final N5Writer cached = getWriterFromCache(storage, location, testWrite);
+		if (cached != null)
+			return cached;
+
+        synchronized (this) {
+			/* see if someone opened this while we were waiting. */
+			final N5Writer reopened = getWriterFromCache(storage, location, testWrite);
+			if (reopened != null)
+				return reopened;
+
+			final N5Writer writer = super.openWriter(storage, access, location);
+			if (writer != null)
+				writerCache.put(normalizeUri(location), writer);
 			return writer;
-		return openAndCacheWriter(storage, access, location);
+		}
+	}
+
+	protected N5Writer getWriterFromCache(StorageFormat format, URI location, boolean testWrite) {
+
+		final URI uri = normalizeUri(location);
+		final N5Writer writer = writerCache.get(uri);
+		if (writer == null)
+			return null;
+
+		boolean writerIsValid = n5MatchesFormat(writer, format) && (!testWrite || canWrite(writer));
+
+		if (!writerIsValid) {
+			writerCache.remove(uri, writer);
+			return null;
+		}
+		return writer;
 	}
 
+	/**
+	 * Test if a read attempt on this {@link N5Reader} will succeed.
+	 *
+     * @return if read was successful
+     */
 	private boolean canRead(N5Reader reader) {
 
 		try {
@@ -76,56 +166,56 @@ private boolean canRead(N5Reader reader) {
 		}
 	}
 
-	private N5Writer openAndCacheWriter(StorageFormat storageFormat, KeyValueAccess access, URI uri) {
-
-		final N5Writer writer = super.openWriter(storageFormat, access, uri);
-		if (writer == null)
-			return null;
-
-		URI normalUri = normalizeUri(uri);
-		writerCache.put(normalUri, writer);
-		return writer;
-	}
-
-	protected synchronized N5Writer getWriterFromCache(StorageFormat format, URI uri) {
-
-		URI normalUri = normalizeUri(uri);
-		final N5Writer writer = writerCache.get(normalUri);
-		if (writer == null)
-			return null;
-
-		if (!n5MatchesFormat(writer, format) || !canWrite(writer)) {
-			writerCache.remove(normalUri);
-			return null;
-		}
-
-		return writer;
-	}
+	/**
+	 * Test if a write attempt on this {@link N5Writer} will succeed.
+	 * Writes a dummy placeholder metadata key/value that is immediately removed if successful.
+	 *
+	 * @return if write was successful
+	 */
 
 	private boolean canWrite(N5Writer writer) {
 
+		/* If we already have a writer, and we `canRead`, it means we passed the `canWrite` check the first time.
+		* In this case, that's enough evidence that we `canWrite`. It doesn't guarantee, e.g someone didn't change
+		* permissions underneath us, but we never guarded against that anyway. */
 		if (canRead(writer))
 			return true;
 
-		final String uuid = UUID.randomUUID().toString();
-		try {
-
-
-			/* Not ideal, but we want to guarantee that the writer is write-able.
-			 * Normally this is guaranteed at construction, but it isn't being constructed here. */
-			writer.setAttribute("/", uuid, "CACHE_CHECK");
-			return true;
-		} catch (Exception e) {
-			return false;
-		} finally {
+		synchronized (this) {
+			final String uuid = UUID.randomUUID().toString();
 			try {
-				/* If it fails, defer to the result of the outer try/catch*/
-				writer.removeAttribute("/", uuid);
-			} catch (Exception ignored) {
+				/* Not ideal, but we want to guarantee that the writer is write-able.
+				 * Normally this is guaranteed at construction, but it isn't being constructed here. */
+				writer.setAttribute("/", uuid, "CACHE_CHECK");
+				return true;
+			} catch (Exception e) {
+				return false;
+			} finally {
+				try {
+					/* If it fails, defer to the result of the outer try/catch*/
+					writer.removeAttribute("/", uuid);
+				} catch (Exception ignored) {
+				}
 			}
 		}
 	}
 
+	@Override
+	protected void requireContainerExists(Supplier<N5Reader> getReader) throws N5Exception.N5IOException {
+		N5Reader reader = null;
+		try {
+			reader = getReader.get();
+		} catch (final Exception e) {
+			throw new N5Exception.N5IOException("Existing container could not be opened, or does not exist.", e);
+		} finally {
+			/* For the N5FactoryWithCache, we don't want to close the reader necessarily, since w
+			* we may have it cached for later re-use. However, for HDF5 specifically, we must close it
+			* since it is not valid to have a reader and writer open at the same time. */
+			if (reader instanceof N5HDF5Reader)
+				reader.close();
+		}
+	}
+
 	private boolean n5MatchesFormat(N5Reader reader, StorageFormat format) {
 
 		// succeed if no format to compare against