Update Bookmark Manager for no longer tracking per database (#1335)

* Update Bookmark Manager for no longer tracking per database

BookmarkManager interface changes:
* The methods `getBookmarks` and `getAllBookmarks` were merge into `getBookmarks`.
* `forget` was removed.
* The database param was removed from the `updateBookmarks` method.

BookmarkManagerConfig changes:
* The BookmarksSupplier interface was replaced by Supplier<Set<Bookmark>>.
* The bookmarksConsumer changed its type to Consumer<Set<Bookmark>>.
* `initialBookmarks` changed its type to Set<Bookmark>

* Address suggestions from the review

driver/src/main/java/org/neo4j/driver/BookmarkManager.java

@@ -23,7 +23,7 @@
 import org.neo4j.driver.util.Experimental;
 
 /**
- * Keeps track of database bookmarks and is used by the driver to ensure causal consistency between sessions and query executions.
+ * Keeps track of bookmarks and is used by the driver to ensure causal consistency between sessions and query executions.
  * <p>
  * Please note that implementations of this interface MUST NOT block for extended periods of time.
  * <p>
@@ -34,35 +34,17 @@
 @Experimental
 public interface BookmarkManager extends Serializable {
     /**
-     * Updates database bookmarks by deleting the given previous bookmarks and adding the new bookmarks.
+     * Updates bookmarks by deleting the given previous bookmarks and adding the new bookmarks.
      *
-     * @param database          the database name, this might be an empty string when session has no database name configured and database discovery is unavailable
      * @param previousBookmarks the previous bookmarks
      * @param newBookmarks      the new bookmarks
      */
-    void updateBookmarks(String database, Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks);
+    void updateBookmarks(Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks);
 
     /**
-     * Gets an immutable set of bookmarks for a given database.
+     * Gets an immutable set of bookmarks.
      *
-     * @param database the database name
-     * @return the set of bookmarks or an empty set if the database name is unknown to the bookmark manager
+     * @return the set of bookmarks.
      */
-    Set<Bookmark> getBookmarks(String database);
-
-    /**
-     * Gets an immutable set of bookmarks for all databases.
-     *
-     * @return the set of bookmarks or an empty set
-     */
-    Set<Bookmark> getAllBookmarks();
-
-    /**
-     * Deletes bookmarks for the given databases.
-     * <p>
-     * This method should be called by driver users if data deletion is desired when bookmarks for the given databases are no longer needed.
-     *
-     * @param databases the set of database names
-     */
-    void forget(Set<String> databases);
+    Set<Bookmark> getBookmarks();
 }

driver/src/main/java/org/neo4j/driver/BookmarkManagerConfig.java

@@ -19,19 +19,19 @@
 package org.neo4j.driver;
 
 import java.util.Collections;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * Bookmark configuration used to configure bookmark manager supplied by {@link BookmarkManagers#defaultManager(BookmarkManagerConfig)}.
  */
 public final class BookmarkManagerConfig {
-    private final Map<String, Set<Bookmark>> initialBookmarks;
-    private final BiConsumer<String, Set<Bookmark>> bookmarksConsumer;
-    private final BookmarksSupplier bookmarksSupplier;
+    private final Set<Bookmark> initialBookmarks;
+    private final Consumer<Set<Bookmark>> bookmarksConsumer;
+    private final Supplier<Set<Bookmark>> bookmarksSupplier;
 
     private BookmarkManagerConfig(BookmarkManagerConfigBuilder builder) {
         this.initialBookmarks = builder.initialBookmarks;
@@ -53,16 +53,16 @@ public static BookmarkManagerConfigBuilder builder() {
      *
      * @return the map of bookmarks
      */
-    public Map<String, Set<Bookmark>> initialBookmarks() {
+    public Set<Bookmark> initialBookmarks() {
         return initialBookmarks;
     }
 
     /**
-     * Returns bookmarks consumer that will be notified when database bookmarks are updated.
+     * Returns bookmarks consumer that will be notified when bookmarks are updated.
      *
      * @return the bookmarks consumer
      */
-    public Optional<BiConsumer<String, Set<Bookmark>>> bookmarksConsumer() {
+    public Optional<Consumer<Set<Bookmark>>> bookmarksConsumer() {
         return Optional.ofNullable(bookmarksConsumer);
     }
 
@@ -71,29 +71,29 @@ public Optional<BiConsumer<String, Set<Bookmark>>> bookmarksConsumer() {
      *
      * @return the bookmark supplier
      */
-    public Optional<BookmarksSupplier> bookmarksSupplier() {
+    public Optional<Supplier<Set<Bookmark>>> bookmarksSupplier() {
         return Optional.ofNullable(bookmarksSupplier);
     }
 
     /**
      * Builder used to configure {@link BookmarkManagerConfig} which will be used to create a bookmark manager.
      */
     public static final class BookmarkManagerConfigBuilder {
-        private Map<String, Set<Bookmark>> initialBookmarks = Collections.emptyMap();
-        private BiConsumer<String, Set<Bookmark>> bookmarksConsumer;
-        private BookmarksSupplier bookmarksSupplier;
+        private Set<Bookmark> initialBookmarks = Collections.emptySet();
+        private Consumer<Set<Bookmark>> bookmarksConsumer;
+        private Supplier<Set<Bookmark>> bookmarksSupplier;
 
         private BookmarkManagerConfigBuilder() {}
 
         /**
          * Provide a map of initial bookmarks to initialise the bookmark manager.
          *
-         * @param databaseToBookmarks database to bookmarks map
+         * @param initialBookmarks initial set of bookmarks
          * @return this builder
          */
-        public BookmarkManagerConfigBuilder withInitialBookmarks(Map<String, Set<Bookmark>> databaseToBookmarks) {
-            Objects.requireNonNull(databaseToBookmarks, "databaseToBookmarks must not be null");
-            this.initialBookmarks = databaseToBookmarks;
+        public BookmarkManagerConfigBuilder withInitialBookmarks(Set<Bookmark> initialBookmarks) {
+            Objects.requireNonNull(initialBookmarks, "initialBookmarks must not be null");
+            this.initialBookmarks = initialBookmarks;
             return this;
         }
 
@@ -105,22 +105,23 @@ public BookmarkManagerConfigBuilder withInitialBookmarks(Map<String, Set<Bookmar
          * @param bookmarksConsumer bookmarks consumer
          * @return this builder
          */
-        public BookmarkManagerConfigBuilder withBookmarksConsumer(BiConsumer<String, Set<Bookmark>> bookmarksConsumer) {
+        public BookmarkManagerConfigBuilder withBookmarksConsumer(Consumer<Set<Bookmark>> bookmarksConsumer) {
             this.bookmarksConsumer = bookmarksConsumer;
             return this;
         }
 
         /**
          * Provide bookmarks supplier.
          * <p>
-         * The supplied bookmarks will be served alongside the bookmarks served by the bookmark manager. The supplied bookmarks will not be stored nor managed by the bookmark manager.
+         * The supplied bookmarks will be served alongside the bookmarks served by the bookmark manager. The supplied bookmarks will not be stored nor managed
+         * by the bookmark manager.
          * <p>
          * The supplier will be called outside bookmark manager's synchronisation lock.
          *
          * @param bookmarksSupplier the bookmarks supplier
          * @return this builder
          */
-        public BookmarkManagerConfigBuilder withBookmarksSupplier(BookmarksSupplier bookmarksSupplier) {
+        public BookmarkManagerConfigBuilder withBookmarksSupplier(Supplier<Set<Bookmark>> bookmarksSupplier) {
             this.bookmarksSupplier = bookmarksSupplier;
             return this;
         }

driver/src/main/java/org/neo4j/driver/internal/Neo4jBookmarkManager.java

@@ -21,20 +21,16 @@
 import static org.neo4j.driver.internal.util.LockUtil.executeWithLock;
 
 import java.io.Serial;
-import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
-import java.util.function.BiConsumer;
-import java.util.stream.Collectors;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 import org.neo4j.driver.Bookmark;
 import org.neo4j.driver.BookmarkManager;
-import org.neo4j.driver.BookmarksSupplier;
 
 /**
  * A basic {@link BookmarkManager} implementation.
@@ -45,66 +41,40 @@ public final class Neo4jBookmarkManager implements BookmarkManager {
 
     private final ReadWriteLock rwLock = new ReentrantReadWriteLock();
 
-    private final Map<String, Set<Bookmark>> databaseToBookmarks = new HashMap<>();
-    private final BiConsumer<String, Set<Bookmark>> updateListener;
-    private final BookmarksSupplier bookmarksSupplier;
+    private final Set<Bookmark> bookmarks;
+    private final Consumer<Set<Bookmark>> updateListener;
+    private final Supplier<Set<Bookmark>> bookmarksSupplier;
 
     public Neo4jBookmarkManager(
-            Map<String, Set<Bookmark>> initialBookmarks,
-            BiConsumer<String, Set<Bookmark>> updateListener,
-            BookmarksSupplier bookmarksSupplier) {
+            Set<Bookmark> initialBookmarks,
+            Consumer<Set<Bookmark>> updateListener,
+            Supplier<Set<Bookmark>> bookmarksSupplier) {
         Objects.requireNonNull(initialBookmarks, "initialBookmarks must not be null");
-        this.databaseToBookmarks.putAll(initialBookmarks);
+        this.bookmarks = new HashSet<>(initialBookmarks);
         this.updateListener = updateListener;
         this.bookmarksSupplier = bookmarksSupplier;
     }
 
     @Override
-    public void updateBookmarks(String database, Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks) {
-        var immutableBookmarks = executeWithLock(
-                rwLock.writeLock(),
-                () -> databaseToBookmarks.compute(database, (ignored, bookmarks) -> {
-                    var updatedBookmarks = new HashSet<Bookmark>();
-                    if (bookmarks != null) {
-                        bookmarks.stream()
-                                .filter(bookmark -> !previousBookmarks.contains(bookmark))
-                                .forEach(updatedBookmarks::add);
-                    }
-                    updatedBookmarks.addAll(newBookmarks);
-                    return Collections.unmodifiableSet(updatedBookmarks);
-                }));
+    public void updateBookmarks(Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks) {
+        var immutableBookmarks = executeWithLock(rwLock.writeLock(), () -> {
+            this.bookmarks.removeAll(previousBookmarks);
+            this.bookmarks.addAll(newBookmarks);
+            return Collections.unmodifiableSet(this.bookmarks);
+        });
         if (updateListener != null) {
-            updateListener.accept(database, immutableBookmarks);
+            updateListener.accept(immutableBookmarks);
         }
     }
 
     @Override
-    public Set<Bookmark> getBookmarks(String database) {
-        var immutableBookmarks = executeWithLock(
-                rwLock.readLock(), () -> databaseToBookmarks.getOrDefault(database, Collections.emptySet()));
+    public Set<Bookmark> getBookmarks() {
+        var immutableBookmarks = executeWithLock(rwLock.readLock(), () -> Collections.unmodifiableSet(this.bookmarks));
         if (bookmarksSupplier != null) {
             var bookmarks = new HashSet<>(immutableBookmarks);
-            bookmarks.addAll(bookmarksSupplier.getBookmarks(database));
+            bookmarks.addAll(bookmarksSupplier.get());
             immutableBookmarks = Collections.unmodifiableSet(bookmarks);
         }
         return immutableBookmarks;
     }
-
-    @Override
-    public Set<Bookmark> getAllBookmarks() {
-        var immutableBookmarks = executeWithLock(rwLock.readLock(), () -> databaseToBookmarks.values().stream()
-                        .flatMap(Collection::stream))
-                .collect(Collectors.toUnmodifiableSet());
-        if (bookmarksSupplier != null) {
-            var bookmarks = new HashSet<>(immutableBookmarks);
-            bookmarks.addAll(bookmarksSupplier.getAllBookmarks());
-            immutableBookmarks = Collections.unmodifiableSet(bookmarks);
-        }
-        return immutableBookmarks;
-    }
-
-    @Override
-    public void forget(Set<String> databases) {
-        executeWithLock(rwLock.writeLock(), () -> databases.forEach(databaseToBookmarks::remove));
-    }
 }

driver/src/main/java/org/neo4j/driver/internal/NoOpBookmarkManager.java

@@ -34,22 +34,16 @@ public class NoOpBookmarkManager implements BookmarkManager {
     private static final Set<Bookmark> EMPTY = Collections.emptySet();
 
     @Override
-    public void updateBookmarks(String database, Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks) {
+    public void updateBookmarks(Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks) {
         // ignored
     }
 
     @Override
-    public Set<Bookmark> getBookmarks(String database) {
+    public Set<Bookmark> getBookmarks() {
         return EMPTY;
     }
 
-    @Override
-    public Set<Bookmark> getAllBookmarks() {
+    private Set<Bookmark> getAllBookmarks() {
         return EMPTY;
     }
-
-    @Override
-    public void forget(Set<String> databases) {
-        // ignored
-    }
 }

driver/src/main/java/org/neo4j/driver/internal/async/NetworkSession.java

@@ -43,7 +43,6 @@
 import org.neo4j.driver.exceptions.TransactionNestingException;
 import org.neo4j.driver.internal.DatabaseBookmark;
 import org.neo4j.driver.internal.DatabaseName;
-import org.neo4j.driver.internal.DatabaseNameUtil;
 import org.neo4j.driver.internal.FailableCursor;
 import org.neo4j.driver.internal.ImpersonationUtil;
 import org.neo4j.driver.internal.cursor.AsyncResultCursor;
@@ -100,7 +99,7 @@ public NetworkSession(
         this.bookmarkManager = bookmarkManager;
         this.lastReceivedBookmarks = bookmarks;
         this.connectionContext =
-                new NetworkSessionConnectionContext(databaseNameFuture, determineBookmarks(true), impersonatedUser);
+                new NetworkSessionConnectionContext(databaseNameFuture, determineBookmarks(false), impersonatedUser);
         this.fetchSize = fetchSize;
     }
 
@@ -145,7 +144,7 @@ public CompletionStage<UnmanagedTransaction> beginTransactionAsync(
                         ImpersonationUtil.ensureImpersonationSupport(connection, connection.impersonatedUser()))
                 .thenCompose(connection -> {
                     UnmanagedTransaction tx = new UnmanagedTransaction(connection, this::handleNewBookmark, fetchSize);
-                    return tx.beginAsync(determineBookmarks(false), config, txType);
+                    return tx.beginAsync(determineBookmarks(true), config, txType);
                 });
 
         // update the reference to the only known transaction
@@ -240,7 +239,7 @@ private CompletionStage<ResultCursorFactory> buildResultCursorFactory(Query quer
                                 .runInAutoCommitTransaction(
                                         connection,
                                         query,
-                                        determineBookmarks(false),
+                                        determineBookmarks(true),
                                         this::handleNewBookmark,
                                         config,
                                         fetchSize);
@@ -342,21 +341,14 @@ private void handleNewBookmark(DatabaseBookmark databaseBookmark) {
         var bookmark = databaseBookmark.bookmark();
         if (bookmark != null) {
             var bookmarks = Set.of(bookmark);
-            String databaseName = databaseBookmark.databaseName();
-            if (databaseName == null || databaseName.isEmpty()) {
-                databaseName = getDatabaseNameNow().orElse(FALLBACK_DATABASE_NAME);
-            }
             lastReceivedBookmarks = bookmarks;
-            bookmarkManager.updateBookmarks(databaseName, lastUsedBookmarks, bookmarks);
+            bookmarkManager.updateBookmarks(lastUsedBookmarks, bookmarks);
         }
     }
 
-    private Set<Bookmark> determineBookmarks(boolean useSystemOnly) {
-        var bookmarks = new HashSet<Bookmark>();
-        if (useSystemOnly) {
-            bookmarks.addAll(bookmarkManager.getBookmarks(DatabaseNameUtil.SYSTEM_DATABASE_NAME));
-        } else {
-            bookmarks.addAll(bookmarkManager.getAllBookmarks());
+    private Set<Bookmark> determineBookmarks(boolean updateLastUsed) {
+        var bookmarks = new HashSet<>(bookmarkManager.getBookmarks());
+        if (updateLastUsed) {
             lastUsedBookmarks = Collections.unmodifiableSet(bookmarks);
         }
         bookmarks.addAll(lastReceivedBookmarks);