Cache based on checking all weak references

src/jvmMain/java/ru/langrafhomyak/db/jdbc_resources_manager/PreparedStatementsCompilationLazyWeakCache.java

@@ -0,0 +1,115 @@
+package ru.langrafhomyak.db.jdbc_resources_manager;
+
+import org.intellij.lang.annotations.Language;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.lang.ref.WeakReference;
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.util.function.Predicate;
+
+import static java.lang.System.identityHashCode;
+
+public class PreparedStatementsCompilationLazyWeakCache implements PreparedStatementsCompilationCache {
+	@NotNull
+	private final String _statementSource;
+	@NotNull
+	private final _SynchronizedEmbedKeyHashMap<ConnectionsMapNode> _map;
+	@NotNull
+	private final CleanupFunction _cleanupFn;
+
+	public PreparedStatementsCompilationLazyWeakCache(@NotNull @Language("SQL") String statementSource) {
+		this._statementSource = statementSource;
+		this._map = new _SynchronizedEmbedKeyHashMap<>();
+		this._cleanupFn = new CleanupFunction();
+	}
+
+	@Override
+	@NotNull
+	public PreparedStatement preparedStatementForConnection(@NotNull Connection connection) throws SQLException {
+		final int key = identityHashCode(connection);
+		ConnectionsMapNode node = this._map.get(key, connection);
+		if (node == null) {
+			/*
+				Assuming that user will clever enough to not use the same connection in different threads
+				and there wouldn't be a race condition for adding instance of this connection to map.
+				Anyway, technically there is no error if the same connection is added twice.
+			 */
+			this._map.cleanup(this._cleanupFn);
+			node = new ConnectionsMapNode(this, connection);
+			this._map.add(node);
+		}
+		return node.createWrapper();
+	}
+
+	private static class CleanupFunction implements Predicate<ConnectionsMapNode> {
+		@Override
+		public boolean test(ConnectionsMapNode node) {
+			return node._connectionRef.get() == null;
+		}
+	}
+
+	private static class ConnectionsMapNode extends _SynchronizedEmbedKeyHashMap.Node<ConnectionsMapNode> {
+		@NotNull
+		private final WeakReference<Object> _connectionRef;
+		@NotNull
+		private final WeakReference<PreparedStatement> _statement;
+		@Nullable
+		private PreparedStatement _wrapper;
+		@NotNull
+		private final Object _wrapperSync;
+
+		public ConnectionsMapNode(@NotNull PreparedStatementsCompilationLazyWeakCache owner, @NotNull Connection connection) throws SQLException {
+			super(identityHashCode(connection));
+			this._connectionRef = new WeakReference<>(connection);
+			this._statement = new WeakReference<>(connection.prepareStatement(owner._statementSource));
+			this._wrapper = null;
+			this._wrapperSync = new Object();
+		}
+
+		@Override
+		protected boolean _checkKeyReference(@NotNull Object ref) {
+			return this._connectionRef.refersTo(ref);
+		}
+
+		@NotNull
+		public PreparedStatement createWrapper() {
+			synchronized (this._wrapperSync) {
+				if (this._wrapper != null)
+					throw new IllegalStateException("Requested prepared statement already used by this connection (or not closed after using)");
+				PreparedStatement orig = this._statement.get();
+				if (orig == null)
+					throw new RuntimeException("");
+				PreparedStatement wrapper = new BoundPreparedStatementWrapper(orig, this);
+				this._wrapper = wrapper;
+				return wrapper;
+			}
+		}
+
+		public void unbindWrapper() {
+			synchronized (this._wrapperSync) {
+				this._wrapper = null;
+			}
+		}
+	}
+
+	private static class BoundPreparedStatementWrapper extends _PreparedStatementWrapper_Delegates {
+		@NotNull
+		private final ConnectionsMapNode _node;
+
+		BoundPreparedStatementWrapper(@NotNull PreparedStatement orig, @NotNull ConnectionsMapNode node) {
+			super(orig);
+			this._node = node;
+		}
+
+		@Override
+		protected void _onClose() throws SQLException {
+			this._orig.clearParameters();
+			this._orig.clearBatch();
+			this._orig.clearWarnings();
+			this._node.unbindWrapper();
+		}
+	}
+}

src/jvmMain/java/ru/langrafhomyak/db/jdbc_resources_manager/_SynchronizedEmbedKeyHashMap.java

@@ -1,9 +1,13 @@
 package ru.langrafhomyak.db.jdbc_resources_manager;
 
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.Iterator;
+import java.util.function.Predicate;
+
 /* package */ class _SynchronizedEmbedKeyHashMap<T extends _SynchronizedEmbedKeyHashMap.Node<T>> {
 	public static abstract class Node<T extends _SynchronizedEmbedKeyHashMap.Node<T>> {
 		public final int keyHashCode;
@@ -70,4 +74,29 @@ import org.jetbrains.annotations.Nullable;
 			this._impl.trim();
 		}
 	}
+
+	/**
+	 * @param nodePredicate {@code true} if should be removed
+	 */
+	public void cleanup(Predicate<T> nodePredicate) {
+		synchronized (this._sync) {
+			final Iterator<Int2ObjectMap.Entry<T>> it = this._impl.int2ObjectEntrySet().fastIterator();
+			while (it.hasNext()) {
+				final Int2ObjectMap.Entry<T> entry = it.next();
+				final int key = entry.getIntKey();
+				T node = entry.getValue();
+
+				while (node != null && nodePredicate.test(node)) {
+					this._impl.put(key, node._nextNode);
+					node = node._nextNode;
+				}
+				while (node != null) {
+					while (node._nextNode != null && nodePredicate.test(node._nextNode)) {
+						node._nextNode = node._nextNode._nextNode;
+					}
+					node = node._nextNode;
+				}
+			}
+		}
+	}
 }