Initial BookmarkManager implementation

Author: injectives

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver;
+
+import java.io.Serializable;
+import java.util.Set;
+
+/**
+ * Keeps track of database 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.
+ *
+ * @see org.neo4j.driver.Config.ConfigBuilder#withBookmarkManager(BookmarkManager)
+ */
+public interface BookmarkManager extends Serializable {
+    /**
+     * Fallback database name used by the driver when session has no database name configured and database discovery is unavailable.
+     */
+    String FALLBACK_DATABASE_NAME = "";
+
+    /**
+     * Updates database bookmarks by deleting the given previous bookmarks and adding the new bookmarks.
+     *
+     * @param database          the database name
+     * @param previousBookmarks the previous bookmarks
+     * @param newBookmarks      the new bookmarks
+     */
+    void updateBookmarks(String database, Set<Bookmark> previousBookmarks, Set<Bookmark> newBookmarks);
+
+    /**
+     * Gets an immutable set of bookmarks for a given database.
+     *
+     * @param database the database name
+     * @return the set of bookmarks or an empty set if the database name is unknown to the bookmark manager
+     */
+    Set<Bookmark> getBookmarks(String database);
+
+    /**
+     * Gets an immutable set of bookmarks for all databases.
+     * <p>
+     * The additional set of database names provided to this method allows bookmark manager implementations to supply
+     * bookmarks for databases that are not directly managed by them. For instance, a bookmark manager implementation
+     * may use an external supplier of bookmarks for databases.
+     *
+     * @param additionalDatabases the set of additional database names
+     * @return the set of bookmarks
+     */
+    Set<Bookmark> getAllBookmarks(Set<String> additionalDatabases);
+
+    /**
+     * Deletes bookmarks for the given databases.
+     *
+     * @param databases the set of database names
+     */
+    void forget(Set<String> databases);
+}

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

@@ -0,0 +1,81 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.neo4j.driver;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+
+public final class BookmarkManagerConfig {
+    private final Map<String, Set<Bookmark>> initialBookmarks;
+    private final BiConsumer<String, Set<Bookmark>> updateListener;
+    private final Function<String, Set<Bookmark>> extraBookmarksProvider;
+
+    private BookmarkManagerConfig(BookmarkManagerConfigBuilder builder) {
+        this.initialBookmarks = builder.initialBookmarks;
+        this.updateListener = builder.updateListener;
+        this.extraBookmarksProvider = builder.extraBookmarksProvider;
+    }
+
+    public Map<String, Set<Bookmark>> initialBookmarks() {
+        return initialBookmarks;
+    }
+
+    public BiConsumer<String, Set<Bookmark>> updateListener() {
+        return updateListener;
+    }
+
+    public Function<String, Set<Bookmark>> extraBookmarksProvider() {
+        return extraBookmarksProvider;
+    }
+
+    public static BookmarkManagerConfigBuilder builder() {
+        return new BookmarkManagerConfigBuilder();
+    }
+
+    public static class BookmarkManagerConfigBuilder {
+        private Map<String, Set<Bookmark>> initialBookmarks = Collections.emptyMap();
+        private BiConsumer<String, Set<Bookmark>> updateListener;
+        private Function<String, Set<Bookmark>> extraBookmarksProvider;
+
+        private BookmarkManagerConfigBuilder() {}
+
+        public BookmarkManagerConfigBuilder withInitialBookmarks(Map<String, Set<Bookmark>> initialBookmarks) {
+            this.initialBookmarks = initialBookmarks;
+            return this;
+        }
+
+        public BookmarkManagerConfigBuilder withUpdateListener(BiConsumer<String, Set<Bookmark>> updateListener) {
+            this.updateListener = updateListener;
+            return this;
+        }
+
+        public BookmarkManagerConfigBuilder withExtraBookmarksProvider(
+                Function<String, Set<Bookmark>> extraBookmarksProvider) {
+            this.extraBookmarksProvider = extraBookmarksProvider;
+            return this;
+        }
+
+        public BookmarkManagerConfig build() {
+            return new BookmarkManagerConfig(this);
+        }
+    }
+}

driver/src/main/java/org/neo4j/driver/internal/ReadOnlyBookmarksHolder.java renamed to driver/src/main/java/org/neo4j/driver/BookmarkManagers.java

@@ -16,28 +16,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.neo4j.driver.internal;
+package org.neo4j.driver;
 
-import java.util.Set;
-import org.neo4j.driver.Bookmark;
+import org.neo4j.driver.internal.Neo4jBookmarkManager;
 
-/**
- * @since 2.0
- */
-public class ReadOnlyBookmarksHolder implements BookmarksHolder {
-    private final Set<Bookmark> bookmarks;
-
-    public ReadOnlyBookmarksHolder(Set<Bookmark> bookmarks) {
-        this.bookmarks = bookmarks;
-    }
-
-    @Override
-    public Set<Bookmark> getBookmarks() {
-        return bookmarks;
-    }
+public interface BookmarkManagers {
 
-    @Override
-    public void setBookmark(Bookmark bookmark) {
-        // NO_OP
+    static BookmarkManager bookmarkManager(BookmarkManagerConfig config) {
+        return new Neo4jBookmarkManager(
+                config.initialBookmarks(), config.updateListener(), config.extraBookmarksProvider());
     }
 }

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

@@ -98,6 +98,7 @@ public class Config implements Serializable {
     private final int eventLoopThreads;
     private final String userAgent;
     private final MetricsAdapter metricsAdapter;
+    private final BookmarkManager bookmarkManager;
 
     private Config(ConfigBuilder builder) {
         this.logging = builder.logging;
@@ -119,6 +120,8 @@ private Config(ConfigBuilder builder) {
 
         this.eventLoopThreads = builder.eventLoopThreads;
         this.metricsAdapter = builder.metricsAdapter;
+
+        this.bookmarkManager = builder.bookmarkManager;
     }
 
     /**
@@ -252,6 +255,15 @@ public String userAgent() {
         return userAgent;
     }
 
+    /**
+     * A {@link BookmarkManager} implementation for the driver to use.
+     *
+     * @return bookmark implementation or {@code null}.
+     */
+    public BookmarkManager bookmarkManager() {
+        return bookmarkManager;
+    }
+
     /**
      * Used to build new config instances
      */
@@ -272,6 +284,7 @@ public static class ConfigBuilder {
         private MetricsAdapter metricsAdapter = MetricsAdapter.DEV_NULL;
         private long fetchSize = FetchSizeUtil.DEFAULT_FETCH_SIZE;
         private int eventLoopThreads = 0;
+        private BookmarkManager bookmarkManager;
 
         private ConfigBuilder() {}
 
@@ -645,6 +658,19 @@ public ConfigBuilder withUserAgent(String userAgent) {
             return this;
         }
 
+        /**
+         * Sets a {@link BookmarkManager} implementation for the driver to use.
+         * <p>
+         * By default, bookmark manager is effectively disabled.
+         *
+         * @param bookmarkManager bookmark manager implementation. Providing {@code null} effectively disables bookmark manager.
+         * @return this builder.
+         */
+        public ConfigBuilder withBookmarkManager(BookmarkManager bookmarkManager) {
+            this.bookmarkManager = bookmarkManager;
+            return this;
+        }
+
         /**
          * Extracts the driver version from the driver jar MANIFEST.MF file.
          */

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

@@ -45,13 +45,15 @@ public class SessionConfig implements Serializable {
     private final String database;
     private final Long fetchSize;
     private final String impersonatedUser;
+    private final boolean bypassBookmarkManager;
 
     private SessionConfig(Builder builder) {
         this.bookmarks = builder.bookmarks;
         this.defaultAccessMode = builder.defaultAccessMode;
         this.database = builder.database;
         this.fetchSize = builder.fetchSize;
         this.impersonatedUser = builder.impersonatedUser;
+        this.bypassBookmarkManager = builder.bypassBookmarkManager;
     }
 
     /**
@@ -130,6 +132,15 @@ public Optional<String> impersonatedUser() {
         return Optional.ofNullable(impersonatedUser);
     }
 
+    /**
+     * Determines if {@link BookmarkManager} configured at driver level should be bypassed.
+     *
+     * @return {@code true} if bookmark manager should be bypassed and not otherwise.
+     */
+    public boolean bypassBookmarkManager() {
+        return bypassBookmarkManager;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {
@@ -143,12 +154,13 @@ public boolean equals(Object o) {
                 && defaultAccessMode == that.defaultAccessMode
                 && Objects.equals(database, that.database)
                 && Objects.equals(fetchSize, that.fetchSize)
-                && Objects.equals(impersonatedUser, that.impersonatedUser);
+                && Objects.equals(impersonatedUser, that.impersonatedUser)
+                && Objects.equals(bypassBookmarkManager, that.bypassBookmarkManager);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(bookmarks, defaultAccessMode, database, impersonatedUser);
+        return Objects.hash(bookmarks, defaultAccessMode, database, impersonatedUser, bypassBookmarkManager);
     }
 
     @Override
@@ -167,6 +179,7 @@ public static class Builder {
         private AccessMode defaultAccessMode = AccessMode.WRITE;
         private String database = null;
         private String impersonatedUser = null;
+        private boolean bypassBookmarkManager = false;
 
         private Builder() {}
 
@@ -292,6 +305,17 @@ public Builder withImpersonatedUser(String impersonatedUser) {
             return this;
         }
 
+        /**
+         * Bypass {@link BookmarkManager} configured at driver level using {@link org.neo4j.driver.Config.ConfigBuilder#withBookmarkManager(BookmarkManager)}.
+         *
+         * @param bypass bypass if {@code true}, use otherwise.
+         * @return this builder.
+         */
+        public Builder withBypassedBookmarkManager(boolean bypass) {
+            this.bypassBookmarkManager = bypass;
+            return this;
+        }
+
         public SessionConfig build() {
             return new SessionConfig(this);
         }

driver/src/main/java/org/neo4j/driver/internal/BookmarksHolder.java renamed to driver/src/main/java/org/neo4j/driver/internal/DatabaseBookmark.java

@@ -18,22 +18,6 @@
  */
 package org.neo4j.driver.internal;
 
-import java.util.Collections;
-import java.util.Set;
 import org.neo4j.driver.Bookmark;
 
-public interface BookmarksHolder {
-    Set<Bookmark> getBookmarks();
-
-    void setBookmark(Bookmark bookmark);
-
-    BookmarksHolder NO_OP = new BookmarksHolder() {
-        @Override
-        public Set<Bookmark> getBookmarks() {
-            return Collections.emptySet();
-        }
-
-        @Override
-        public void setBookmark(Bookmark bookmark) {}
-    };
-}
+public record DatabaseBookmark(String databaseName, Bookmark bookmark) {}

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

@@ -260,7 +260,8 @@ protected InternalDriver createRoutingDriver(
      */
     protected InternalDriver createDriver(
             SecurityPlan securityPlan, SessionFactory sessionFactory, MetricsProvider metricsProvider, Config config) {
-        return new InternalDriver(securityPlan, sessionFactory, metricsProvider, config.logging());
+        return new InternalDriver(
+                securityPlan, sessionFactory, metricsProvider, config.logging(), config.bookmarkManager());
     }
 
     /**

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

@@ -22,6 +22,7 @@
 
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicBoolean;
+import org.neo4j.driver.BookmarkManager;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.Logger;
 import org.neo4j.driver.Logging;
@@ -50,15 +51,19 @@ public class InternalDriver implements Driver {
     private AtomicBoolean closed = new AtomicBoolean(false);
     private final MetricsProvider metricsProvider;
 
+    private final BookmarkManager bookmarkManager;
+
     InternalDriver(
             SecurityPlan securityPlan,
             SessionFactory sessionFactory,
             MetricsProvider metricsProvider,
-            Logging logging) {
+            Logging logging,
+            BookmarkManager bookmarkManager) {
         this.securityPlan = securityPlan;
         this.sessionFactory = sessionFactory;
         this.metricsProvider = metricsProvider;
         this.log = logging.getLog(getClass());
+        this.bookmarkManager = bookmarkManager != null ? bookmarkManager : new NoOpBookmarkManager();
     }
 
     @Override
@@ -164,7 +169,8 @@ private static RuntimeException driverCloseException() {
 
     public NetworkSession newSession(SessionConfig config) {
         assertOpen();
-        NetworkSession session = sessionFactory.newInstance(config);
+        var bookmarkManager = config.bypassBookmarkManager() ? new NoOpBookmarkManager() : this.bookmarkManager;
+        NetworkSession session = sessionFactory.newInstance(config, bookmarkManager);
         if (closed.get()) {
             // session does not immediately acquire connection, it is fine to just throw
             throw driverCloseException();