Skip to content

Commit

Permalink
Update BoltConnectionProvider (#1610)
Browse files Browse the repository at this point in the history 
The BoltConnectionProvider will now accept `address` and some other additional parameters in its methods.
  • Loading branch information
@injectives
injectives authored Jan 29, 2025
1 parent 95d8eb0 commit 9ce4ede
Show file tree
Hide file tree
Showing 13 changed files with 700 additions and 182 deletions.
75 changes: 47 additions & 28 deletions ...y/src/main/java/org/neo4j/driver/internal/bolt/basicimpl/NettyBoltConnectionProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,56 +58,39 @@
public final class NettyBoltConnectionProvider implements BoltConnectionProvider {
private final LoggingProvider logging;
private final System.Logger log;

private final ConnectionProvider connectionProvider;

private BoltServerAddress address;

private RoutingContext routingContext;
private BoltAgent boltAgent;
private String userAgent;
private int connectTimeoutMillis;
private CompletableFuture<Void> closeFuture;
private MetricsListener metricsListener;
private final MetricsListener metricsListener;
private final Clock clock;
private final ValueFactory valueFactory;

private CompletableFuture<Void> closeFuture;

public NettyBoltConnectionProvider(
EventLoopGroup eventLoopGroup,
Clock clock,
DomainNameResolver domainNameResolver,
LocalAddress localAddress,
LoggingProvider logging,
ValueFactory valueFactory) {
ValueFactory valueFactory,
MetricsListener metricsListener) {
Objects.requireNonNull(eventLoopGroup);
this.clock = Objects.requireNonNull(clock);
this.logging = Objects.requireNonNull(logging);
this.log = logging.getLog(getClass());
this.connectionProvider = ConnectionProviders.netty(
eventLoopGroup, clock, domainNameResolver, localAddress, logging, valueFactory);
this.valueFactory = Objects.requireNonNull(valueFactory);
this.metricsListener = NoopMetricsListener.getInstance();
InternalLoggerFactory.setDefaultFactory(new NettyLogging(logging));
}

@Override
public CompletionStage<Void> init(
public CompletionStage<BoltConnection> connect(
BoltServerAddress address,
RoutingContext routingContext,
BoltAgent boltAgent,
String userAgent,
int connectTimeoutMillis,
MetricsListener metricsListener) {
this.address = address;
this.routingContext = routingContext;
this.boltAgent = boltAgent;
this.userAgent = userAgent;
this.connectTimeoutMillis = connectTimeoutMillis;
this.metricsListener = NoopMetricsListener.getInstance();
InternalLoggerFactory.setDefaultFactory(new NettyLogging(logging));
return CompletableFuture.completedStage(null);
}

@Override
public CompletionStage<BoltConnection> connect(
SecurityPlan securityPlan,
DatabaseName databaseName,
Supplier<CompletionStage<AuthToken>> authTokenStageSupplier,
Expand Down Expand Up @@ -180,8 +163,20 @@ public CompletionStage<BoltConnection> connect(
}

@Override
public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, AuthToken authToken) {
public CompletionStage<Void> verifyConnectivity(
BoltServerAddress address,
RoutingContext routingContext,
BoltAgent boltAgent,
String userAgent,
int connectTimeoutMillis,
SecurityPlan securityPlan,
AuthToken authToken) {
return connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
null,
() -> CompletableFuture.completedStage(authToken),
Expand All @@ -196,8 +191,20 @@ public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, AuthT
}

@Override
public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, AuthToken authToken) {
public CompletionStage<Boolean> supportsMultiDb(
BoltServerAddress address,
RoutingContext routingContext,
BoltAgent boltAgent,
String userAgent,
int connectTimeoutMillis,
SecurityPlan securityPlan,
AuthToken authToken) {
return connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
null,
() -> CompletableFuture.completedStage(authToken),
Expand All @@ -215,8 +222,20 @@ public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, AuthT
}

@Override
public CompletionStage<Boolean> supportsSessionAuth(SecurityPlan securityPlan, AuthToken authToken) {
public CompletionStage<Boolean> supportsSessionAuth(
BoltServerAddress address,
RoutingContext routingContext,
BoltAgent boltAgent,
String userAgent,
int connectTimeoutMillis,
SecurityPlan securityPlan,
AuthToken authToken) {
return connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
null,
() -> CompletableFuture.completedStage(authToken),
Expand Down
90 changes: 69 additions & 21 deletions ...src/main/java/org/neo4j/driver/internal/bolt/pooledimpl/PooledBoltConnectionProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,15 @@ public class PooledBoltConnectionProvider implements BoltConnectionProvider {
private final long maxLifetime;
private final long idleBeforeTest;
private final Clock clock;
private MetricsListener metricsListener;
private CompletionStage<Void> closeStage;
private BoltServerAddress address;
private String poolId;
private final MetricsListener metricsListener;
private final BoltServerAddress address;
private final RoutingContext routingContext;
private final BoltAgent boltAgent;
private final String userAgent;
private final int connectTimeoutMillis;
private final String poolId;

private CompletionStage<Void> closeStage;
private long minAuthTimestamp;

public PooledBoltConnectionProvider(
Expand All @@ -80,7 +84,13 @@ public PooledBoltConnectionProvider(
long maxLifetime,
long idleBeforeTest,
Clock clock,
LoggingProvider logging) {
LoggingProvider logging,
MetricsListener metricsListener,
BoltServerAddress address,
RoutingContext routingContext,
BoltAgent boltAgent,
String userAgent,
int connectTimeoutMillis) {
this.boltConnectionProvider = boltConnectionProvider;
this.pooledConnectionEntries = new ArrayList<>();
this.pendingAcquisitions = new ArrayDeque<>(100);
Expand All @@ -90,19 +100,13 @@ public PooledBoltConnectionProvider(
this.idleBeforeTest = idleBeforeTest;
this.clock = Objects.requireNonNull(clock);
this.log = logging.getLog(getClass());
}

@Override
public CompletionStage<Void> init(
BoltServerAddress address,
RoutingContext routingContext,
BoltAgent boltAgent,
String userAgent,
int connectTimeoutMillis,
MetricsListener metricsListener) {
this.metricsListener = Objects.requireNonNull(metricsListener);
this.address = Objects.requireNonNull(address);
this.routingContext = Objects.requireNonNull(routingContext);
this.boltAgent = Objects.requireNonNull(boltAgent);
this.userAgent = Objects.requireNonNull(userAgent);
this.connectTimeoutMillis = connectTimeoutMillis;
this.poolId = poolId(address);
this.metricsListener = Objects.requireNonNull(metricsListener);
metricsListener.registerPoolMetrics(
poolId,
address,
Expand All @@ -120,13 +124,16 @@ public CompletionStage<Void> init(
.count();
}
});
return boltConnectionProvider.init(
address, routingContext, boltAgent, userAgent, connectTimeoutMillis, metricsListener);
}

@SuppressWarnings({"ReassignedVariable"})
@Override
public CompletionStage<BoltConnection> connect(
BoltServerAddress ignoredAddress,
RoutingContext ignoredRoutingContext,
BoltAgent ignoredBoltAgent,
String ignoredUserAgent,
int ignoredConnectTimeoutMillis,
SecurityPlan securityPlan,
DatabaseName databaseName,
Supplier<CompletionStage<AuthToken>> authTokenStageSupplier,
Expand Down Expand Up @@ -334,6 +341,11 @@ private void connect(
var entry = connectionEntryWithMetadata.connectionEntry;
boltConnectionProvider
.connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
databaseName,
empty.get()
Expand Down Expand Up @@ -502,8 +514,20 @@ private CompletionStage<Void> livenessCheckStage(ConnectionEntry entry) {
}

@Override
public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, AuthToken authToken) {
public CompletionStage<Void> verifyConnectivity(
BoltServerAddress ignoredAddress,
RoutingContext ignoredRoutingContext,
BoltAgent ignoredBoltAgent,
String ignoredUserAgent,
int ignoredConnectTimeoutMillis,
SecurityPlan securityPlan,
AuthToken authToken) {
return connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
null,
() -> CompletableFuture.completedStage(authToken),
Expand All @@ -518,8 +542,20 @@ public CompletionStage<Void> verifyConnectivity(SecurityPlan securityPlan, AuthT
}

@Override
public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, AuthToken authToken) {
public CompletionStage<Boolean> supportsMultiDb(
BoltServerAddress ignoredAddress,
RoutingContext ignoredRoutingContext,
BoltAgent ignoredBoltAgent,
String ignoredUserAgent,
int ignoredConnectTimeoutMillis,
SecurityPlan securityPlan,
AuthToken authToken) {
return connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
null,
() -> CompletableFuture.completedStage(authToken),
Expand All @@ -537,8 +573,20 @@ public CompletionStage<Boolean> supportsMultiDb(SecurityPlan securityPlan, AuthT
}

@Override
public CompletionStage<Boolean> supportsSessionAuth(SecurityPlan securityPlan, AuthToken authToken) {
public CompletionStage<Boolean> supportsSessionAuth(
BoltServerAddress ignoredAddress,
RoutingContext ignoredRoutingContext,
BoltAgent ignoredBoltAgent,
String ignoredUserAgent,
int ignoredConnectTimeoutMillis,
SecurityPlan securityPlan,
AuthToken authToken) {
return connect(
address,
routingContext,
boltAgent,
userAgent,
connectTimeoutMillis,
securityPlan,
null,
() -> CompletableFuture.completedStage(authToken),
Expand Down
Loading

0 comments on commit 9ce4ede

Please sign in to comment.