ensure garbage collection

Author: Neenu1995

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java

@@ -840,6 +840,8 @@ public void close() throws SQLException {
       throw new BigQueryJdbcException(ex);
     } catch (InterruptedException e) {
       throw new BigQueryJdbcRuntimeException(e);
+    } finally {
+      BigQueryJdbcMdc.removeInstance(this);
     }
     this.isClosed = true;
   }

java-bigquery/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcMdc.java

@@ -32,6 +32,8 @@ public class BigQueryJdbcMdc {
       new ConcurrentHashMap<>();
 
   /** Allocates an exclusive InheritableThreadLocal and registers the connection mapping. */
+  private static final InheritableThreadLocal<String> currentConnectionId = new InheritableThreadLocal<>();
+
   public static void registerInstance(BigQueryConnection connection, String id) {
     if (connection != null) {
       String cleanId =
@@ -44,6 +46,7 @@ public static void registerInstance(BigQueryConnection connection, String id) {
                 return "JdbcConnection-" + numericPart;
               });
 
+      currentConnectionId.set(cleanId);
       InheritableThreadLocal<String> threadLocal =
           instanceLocals.computeIfAbsent(connection, k -> new InheritableThreadLocal<>());
       threadLocal.set(cleanId);
@@ -69,19 +72,26 @@ public static String getConnectionId(BigQueryConnection connection) {
    * Returns the connection ID carried by any registered active connection on the current thread.
    */
   public static String getConnectionId() {
-    for (InheritableThreadLocal<String> local : instanceLocals.values()) {
-      String val = local.get();
-      if (val != null) {
-        return val;
+    return currentConnectionId.get();
+  }
+
+  /** Clears the connection ID context from all active connection contexts on the current thread. */
+  public static void removeInstance(BigQueryConnection connection) {
+    if (connection != null) {
+      InheritableThreadLocal<String> local = instanceLocals.remove(connection);
+      if (local != null) {
+        local.remove();
       }
+      instanceIds.remove(connection);
     }
-    return null;
   }
 
-  /** Clears the connection ID context from all active connection contexts on the current thread. */
   public static void clear() {
+    currentConnectionId.remove();
     for (InheritableThreadLocal<String> local : instanceLocals.values()) {
       local.remove();
     }
+    instanceLocals.clear();
+    instanceIds.clear();
   }
 }