fix: make location cache expire after 10 minutes

Author: agrawal-siddharth

google-cloud-bigquerystorage/pom.xml

@@ -28,6 +28,11 @@
         <artifactId>arrow-memory-core</artifactId>
         <version>17.0.0</version>
       </dependency>
+      <dependency>
+        <groupId>com.github.ben-manes.caffeine</groupId>
+        <artifactId>caffeine</artifactId>
+        <version>3.2.2</version>
+      </dependency>
     </dependencies>
   </dependencyManagement>
   <build>
@@ -175,6 +180,10 @@
       <groupId>org.apache.arrow</groupId>
       <artifactId>arrow-vector</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.github.ben-manes.caffeine</groupId>
+      <artifactId>caffeine</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java

@@ -15,6 +15,8 @@
  */
 package com.google.cloud.bigquery.storage.v1;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.api.core.ApiFuture;
 import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.CredentialsProvider;
@@ -75,8 +77,15 @@ public class StreamWriter implements AutoCloseable {
   private static Pattern streamPatternDefaultStream = Pattern.compile(defaultStreamMatching);
 
   // Cache of location info for a given dataset.
-  private static Map<String, String> projectAndDatasetToLocation = new ConcurrentHashMap<>();
+  private static long LOCATION_CACHE_EXPIRE_MILLIS = 10 * 60 * 1000; // 10 minutes
 
+  private static Cache<String, String> allocateProjectLocationCache() {
+    return Caffeine.newBuilder()
+        .expireAfterWrite(LOCATION_CACHE_EXPIRE_MILLIS, TimeUnit.MILLISECONDS)
+        .build();
+  }
+
+  private static Cache<String, String> projectAndDatasetToLocation = allocateProjectLocationCache();
   /*
    * The identifier of stream to write to.
    */
@@ -288,7 +297,7 @@ private StreamWriter(Builder builder) throws IOException {
         // Location is not passed in, try to fetch from RPC
         String datasetAndProjectName = extractDatasetAndProjectName(builder.streamName);
         location =
-            projectAndDatasetToLocation.computeIfAbsent(
+            projectAndDatasetToLocation.get(
                 datasetAndProjectName,
                 (key) -> {
                   GetWriteStreamRequest writeStreamRequest =
@@ -371,6 +380,12 @@ static boolean isDefaultStream(String streamName) {
     return streamMatcher.find();
   }
 
+  @VisibleForTesting
+  static void recreateProjectLocationCache(long durationExpireMillis) {
+    LOCATION_CACHE_EXPIRE_MILLIS = durationExpireMillis;
+    projectAndDatasetToLocation = allocateProjectLocationCache();
+  }
+
   String getFullTraceId() {
     return fullTraceId;
   }

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java

@@ -56,6 +56,7 @@
 import com.google.protobuf.Descriptors;
 import com.google.protobuf.Descriptors.DescriptorValidationException;
 import com.google.protobuf.Int64Value;
+import com.google.protobuf.Timestamp;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.opentelemetry.api.common.Attributes;
@@ -2551,4 +2552,66 @@ public void testGetDefaultStreamName() {
     assertEquals(
         "projects/projectId/datasets/datasetId/tables/tableId/_default", actualDefaultName);
   }
+
+  @Test
+  public void testLocationCacheIsHit() throws Exception {
+    WriteStream expectedResponse =
+        WriteStream.newBuilder()
+            .setName(WriteStreamName.of("[PROJECT]", "[DATASET]", "[TABLE]", "[STREAM]").toString())
+            .setCreateTime(Timestamp.newBuilder().build())
+            .setCommitTime(Timestamp.newBuilder().build())
+            .setTableSchema(TableSchema.newBuilder().build())
+            .setLocation("oklahoma")
+            .build();
+    testBigQueryWrite.addResponse(expectedResponse);
+
+    // first stream will result in call to getWriteStream for location lookup
+    StreamWriter writer1 =
+        StreamWriter.newBuilder(TEST_STREAM_1, client)
+            .setWriterSchema(createProtoSchema())
+            .setEnableConnectionPool(true)
+            .build();
+
+    // second stream will hit the location cache
+    StreamWriter writer2 =
+        StreamWriter.newBuilder(TEST_STREAM_1, client)
+            .setWriterSchema(createProtoSchema())
+            .setEnableConnectionPool(true)
+            .build();
+    assertEquals(1, testBigQueryWrite.getWriteStreamRequests().size());
+  }
+
+  @Test
+  public void testLocationCacheExpires() throws Exception {
+    // force cache to expire in 1000 millis
+    StreamWriter.recreateProjectLocationCache(1000);
+    WriteStream expectedResponse =
+        WriteStream.newBuilder()
+            .setName(WriteStreamName.of("[PROJECT]", "[DATASET]", "[TABLE]", "[STREAM]").toString())
+            .setCreateTime(Timestamp.newBuilder().build())
+            .setCommitTime(Timestamp.newBuilder().build())
+            .setTableSchema(TableSchema.newBuilder().build())
+            .setLocation("oklahoma")
+            .build();
+    testBigQueryWrite.addResponse(expectedResponse);
+    testBigQueryWrite.addResponse(expectedResponse);
+
+    // first stream will result in call to getWriteStream for location lookup
+    StreamWriter writer1 =
+        StreamWriter.newBuilder(TEST_STREAM_1, client)
+            .setWriterSchema(createProtoSchema())
+            .setEnableConnectionPool(true)
+            .build();
+
+    // force cache to expire
+    TimeUnit.SECONDS.sleep(2);
+
+    // second stream will result in call to getWriteStream for location lookup
+    StreamWriter writer2 =
+        StreamWriter.newBuilder(TEST_STREAM_1, client)
+            .setWriterSchema(createProtoSchema())
+            .setEnableConnectionPool(true)
+            .build();
+    assertEquals(2, testBigQueryWrite.getWriteStreamRequests().size());
+  }
 }