Properly close Storage API batch connections (#31710)

* properly close connections; add active connection counter

* only invalidate stream at teardown for PENDING type

* cleanup

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/AppendClientInfo.java

@@ -30,6 +30,8 @@
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
 
 /**
  * Container class used by {@link StorageApiWritesShardedRecords} and {@link
@@ -38,6 +40,9 @@
  */
 @AutoValue
 abstract class AppendClientInfo {
+  private final Counter activeConnections =
+      Metrics.counter(AppendClientInfo.class, "activeConnections");
+
   abstract @Nullable BigQueryServices.StreamAppendClient getStreamAppendClient();
 
   abstract TableSchema getTableSchema();
@@ -114,19 +119,21 @@ public AppendClientInfo withAppendClient(
       return this;
     } else {
       String streamName = getStreamName.get();
-      return toBuilder()
-          .setStreamName(streamName)
-          .setStreamAppendClient(
-              writeStreamService.getStreamAppendClient(
-                  streamName, getDescriptor(), useConnectionPool, missingValueInterpretation))
-          .build();
+      BigQueryServices.StreamAppendClient client =
+          writeStreamService.getStreamAppendClient(
+              streamName, getDescriptor(), useConnectionPool, missingValueInterpretation);
+
+      activeConnections.inc();
+
+      return toBuilder().setStreamName(streamName).setStreamAppendClient(client).build();
     }
   }
 
   public void close() {
     BigQueryServices.StreamAppendClient client = getStreamAppendClient();
     if (client != null) {
       getCloseAppendClient().accept(client);
+      activeConnections.dec();
     }
   }
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -334,6 +334,11 @@ void teardown() {
           if (client != null) {
             runAsyncIgnoreFailure(closeWriterExecutor, client::unpin);
           }
+          // if this is a PENDING stream, we won't be using it again after cleaning up this
+          // destination state, so clear it from the cache
+          if (!useDefaultStream) {
+            APPEND_CLIENTS.invalidate(streamName);
+          }
           appendClientInfo = null;
         }
       }