Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@
*
* @see <a href="https://cloud.google.com/bigquery/what-is-bigquery">Google Cloud BigQuery</a>
*/
public interface BigQuery extends Service<BigQueryOptions> {
public interface BigQuery extends Service<BigQueryOptions>, AutoCloseable {

/**
* Fields of a BigQuery Dataset resource.
Expand Down Expand Up @@ -1821,4 +1821,14 @@ Object queryWithTimeout(
* represents the subset of granted permissions.
*/
List<String> testIamPermissions(TableId table, List<String> permissions, IAMOption... options);

/**
* Closes any background resources and transport channels held by this service.
*
* <p>The default implementation does nothing. Implementations that manage background resources
* (such as gRPC channels or storage clients) should override this method to release them
* deterministically.
*/
@Override
default void close() {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -504,6 +504,7 @@ private void closeClient() {

private transient ConcurrentHashMap<String, BigQueryReadClient> bqReadClients;
private transient boolean isGlobalClientUserProvided;
private transient volatile boolean closed = false;

/**
* Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming
Expand All @@ -527,6 +528,9 @@ BigQueryReadClient getBigQueryReadClient() {
* @throws BigQueryException if initializing the storage read client fails
*/
BigQueryReadClient getBigQueryReadClient(String location) {
if (closed) {
throw new IllegalStateException("BigQuery service has been closed");
}
String cacheKey = location != null ? location.toLowerCase() : "global";
if (bqReadClients == null) {
synchronized (this) {
Expand All @@ -543,6 +547,9 @@ BigQueryReadClient getBigQueryReadClient(String location) {
return client;
}
synchronized (this) {
if (closed) {
throw new IllegalStateException("BigQuery service has been closed");
}
client = bqReadClients.get(cacheKey);
if (client == null && isGlobalClientUserProvided) {
client = bqReadClients.get("global");
Expand All @@ -552,6 +559,10 @@ BigQueryReadClient getBigQueryReadClient(String location) {
configureReadSettings(settingsBuilder, getOptions());
try {
client = BigQueryReadClient.create(settingsBuilder.build());
if (closed) {
client.close();
throw new IllegalStateException("BigQuery service has been closed");
}
if (bqReadClients.size() < MAX_CACHED_READ_CLIENTS) {
bqReadClients.put(cacheKey, client);
}
Expand Down Expand Up @@ -583,6 +594,30 @@ void setBigQueryReadClient(String location, BigQueryReadClient client) {
}
}

/**
* Closes any background resources and transport channels held by this {@link BigQueryImpl},
* including the underlying {@link BigQueryReadClient} instances used for Arrow query streaming.
*/
@Override
public void close() {
synchronized (this) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you double check this: I think we only need to have a mutex on the closed var (e.g. mutex to set close = true).

Afterwards, we can begin the orderly shutdown of clients as new regions won't be able to create new bqReadClients.

I think the thing that comes to mind, is that most managed resources in the clients will invoke close() which invokes shutdown(). shutdown does not immediately force all resources to shutdown (any in flight processes will continue, but new requests will not be closed)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the feedback, I agree. I've updated the code to narrow the mutex strictly to setting closed = true and initiate orderly shutdown of bqReadClients directly.

if (closed) {
return;
}
closed = true;
}
if (bqReadClients != null) {
for (BigQueryReadClient client : bqReadClients.values()) {
try {
client.close();
} catch (Exception e) {
// Ignore exceptions during teardown
}
}
bqReadClients.clear();
}
}
Comment thread
jinseopkim0 marked this conversation as resolved.
Comment thread
jinseopkim0 marked this conversation as resolved.

/**
* Configures a {@link BigQueryReadSettings.Builder} with credentials, header provider, and
* universe domain mapped from the given {@link BigQueryOptions}.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4260,4 +4260,70 @@ void testTestIamPermissionsWhenNoPermissionsGranted() throws IOException {
.testIamPermissionsSkipExceptionTranslation(
resourceId, checkedPermissions, EMPTY_RPC_OPTIONS);
}

@Test
void testCloseClosesBigQueryReadClient() {
BigQueryReadClient mockReadClient =
mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
bigquery = options.getService();
((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient);

bigquery.close();

verify(mockReadClient, times(1)).close();
}

@Test
void testCloseIsIdempotent() {
BigQueryReadClient mockReadClient =
mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
bigquery = options.getService();
((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient);

bigquery.close();
bigquery.close();

verify(mockReadClient, times(1)).close();
}

@Test
void testCloseWithoutReadClientDoesNotThrow() {
bigquery = options.getService();
bigquery.close();
}

@Test
void testTryWithResources() {
BigQueryReadClient mockReadClient =
mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
try (BigQuery bq = options.getService()) {
((BigQueryImpl) bq).setBigQueryReadClient(mockReadClient);
assertNotNull(bq);
}
verify(mockReadClient, times(1)).close();
}

@Test
void testGetBigQueryReadClientAfterCloseThrows() {
bigquery = options.getService();
bigquery.close();
assertThrows(
IllegalStateException.class, () -> ((BigQueryImpl) bigquery).getBigQueryReadClient());
}

@Test
void testCloseClosesAllRegionalBigQueryReadClients() {
BigQueryReadClient mockReadClientUs =
mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
BigQueryReadClient mockReadClientEu =
mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
bigquery = options.getService();
((BigQueryImpl) bigquery).setBigQueryReadClient("us", mockReadClientUs);
((BigQueryImpl) bigquery).setBigQueryReadClient("eu", mockReadClientEu);

bigquery.close();

verify(mockReadClientUs, times(1)).close();
verify(mockReadClientEu, times(1)).close();
}
}
Loading