Skip to content

Commit

Permalink
feat: add shutdown() to ConnectorRegistry
Browse files Browse the repository at this point in the history
  • Loading branch information
ttosta-google committed Dec 8, 2023
1 parent 5c6c85f commit dfdd571
Show file tree
Hide file tree
Showing 3 changed files with 49 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -45,4 +45,12 @@ public static void close(String name) {
public static void reset() {
InternalConnectorRegistry.INSTANCE.resetInstance();
}

/**
* Shutdown the entire AlloyDB JDBC Connector. This will stop all background threads. All future
* attempts to connect to the database will fail.
*/
public static void shutdown() {
InternalConnectorRegistry.INSTANCE.shutdownInstance();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,9 @@ enum InternalConnectorRegistry implements Closeable {
@SuppressWarnings("ImmutableEnumChecker")
private ConcurrentHashMap<String, Connector> namedConnectors;

@SuppressWarnings("ImmutableEnumChecker")
private boolean shutdown = false;

InternalConnectorRegistry() {
// During refresh, each instance consumes 2 threads from the thread pool. By using 8 threads,
// there should be enough free threads so that there will not be a deadlock. Most users
Expand All @@ -63,6 +66,10 @@ enum InternalConnectorRegistry implements Closeable {
* @throws IOException if error occurs during socket creation.
*/
public Socket connect(ConnectionConfig config) throws IOException {
if (shutdown) {
throw new IllegalStateException("ConnectorRegistry was shut down.");
}

if (config.getNamedConnector() != null) {
Connector connector = getNamedConnector(config.getNamedConnector());
return connector.connect(config.withConnectorConfig(connector.getConfig()));
Expand All @@ -79,6 +86,10 @@ public Socket connect(ConnectionConfig config) throws IOException {

/** Register the configuration for a named connector. */
public void register(String name, ConnectorConfig config) {
if (shutdown) {
throw new IllegalStateException("ConnectorRegistry was shut down.");
}

if (this.namedConnectors.containsKey(name)) {
throw new IllegalArgumentException("Named connection " + name + " exists.");
}
Expand All @@ -87,15 +98,19 @@ public void register(String name, ConnectorConfig config) {

/** Close a named connector, stopping the refresh process and removing it from the registry. */
public void close(String name) {
if (shutdown) {
throw new IllegalStateException("ConnectorRegistry was shut down.");
}

Connector connector = namedConnectors.remove(name);
if (connector == null) {
throw new IllegalArgumentException("Named connection " + name + " does not exist.");
}
connector.close();
}

/** Shutdown all connectors and remove the singleton instance. */
private void shutdown() {
/** Shutdown all connectors. */
private void shutdownConnectors() {
this.unnamedConnectors.forEach((key, c) -> c.close());
this.unnamedConnectors.clear();
this.namedConnectors.forEach((key, c) -> c.close());
Expand All @@ -104,13 +119,19 @@ private void shutdown() {

@Override
public void close() {
shutdown();
this.executor.shutdown();
shutdownInstance();
}

/** Calls shutdown on the singleton. */
public void resetInstance() {
shutdown();
shutdownConnectors();
}

/** Calls shutdown on the singleton. */
public void shutdownInstance() {
shutdown = true;
shutdownConnectors();
this.executor.shutdown();
}

private Connector getConnector(ConnectionConfig config) {
Expand Down
18 changes: 15 additions & 3 deletions docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -182,9 +182,21 @@ ConnectorRegistry.reset();
```

After calling `ConnectorRegistry.reset()`, the next attempt to connect to a
database using a SocketFactory or R2DBC ConnectionFactory, or
to `ConnectorRegistry.register()` will start a new connector registry, restart
the background threads, and create a new connector.
database, or to `ConnectorRegistry.register()` will start a new connector
registry, restart the background threads, and create a new connector.

### Shutdown The Connector Registry

The application may shut down the ConnectorRegistry. This closes all existing
named and unnamed connectors, and stops internal background threads.

```java
ConnectorRegistry.shutdown();
```

After calling `ConnectorRegistry.shutdown()`, subsequent attempts to connect to
a database, or to `ConnectorRegistry.register()` will fail,
throwing `IllegalStateException`.

## Configuring Google Credentials

Expand Down

0 comments on commit dfdd571

Please sign in to comment.