Merge branch 'main' into migrate-java-shared-config-v2

Author: jinseopkim0

google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/UserCredentials.java

@@ -328,7 +328,7 @@ private InputStream getUserCredentialsStream() throws IOException {
       json.put("client_secret", clientSecret);
     }
     if (quotaProjectId != null) {
-      json.put("quota_project", clientSecret);
+      json.put("quota_project_id", quotaProjectId);
     }
     json.setFactory(JSON_FACTORY);
     String text = json.toPrettyString();

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/DefaultCredentialsProviderTest.java

@@ -387,6 +387,7 @@ void getDefaultCredentials_GdchServiceAccount() throws IOException {
     assertNotNull(((GdchCredentials) defaultCredentials).getApiAudience());
   }
 
+  @Test
   void getDefaultCredentials_quota_project() throws IOException {
     InputStream userStream =
         UserCredentialsTest.writeUserStream(

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountCredentialsTest.java

@@ -1762,6 +1762,17 @@ void createScopes_existingAccessTokenInvalidated() throws IOException {
     assertNull(newAccessToken);
   }
 
+  @Test
+  void getRequestMetadata_withMultipleScopes_selfSignedJWT() throws IOException {
+    List<String> scopes = Arrays.asList("scope1", "scope2");
+    ServiceAccountCredentials credentials =
+        createDefaultBuilderWithKey(OAuth2Utils.privateKeyFromPkcs8(PRIVATE_KEY_PKCS8))
+            .setScopes(scopes)
+            .setUseJwtAccessWithScope(true)
+            .build();
+    verifyJwtAccess(credentials.getRequestMetadata(CALL_URI), "scope1 scope2");
+  }
+
   private void verifyJwtAccess(Map<String, List<String>> metadata, String expectedScopeClaim)
       throws IOException {
     assertNotNull(metadata);
@@ -1777,6 +1788,8 @@ private void verifyJwtAccess(Map<String, List<String>> metadata, String expected
     assertNotNull(assertion, "Bearer assertion not found");
     JsonWebSignature signature =
         JsonWebSignature.parse(GsonFactory.getDefaultInstance(), assertion);
+    assertEquals("RS256", signature.getHeader().getAlgorithm());
+    assertEquals("JWT", signature.getHeader().getType());
     assertEquals(CLIENT_EMAIL, signature.getPayload().getIssuer());
     assertEquals(CLIENT_EMAIL, signature.getPayload().getSubject());
     if (expectedScopeClaim != null) {
@@ -1787,6 +1800,14 @@ private void verifyJwtAccess(Map<String, List<String>> metadata, String expected
       assertFalse(signature.getPayload().containsKey("scope"));
     }
     assertEquals(PRIVATE_KEY_ID, signature.getHeader().getKeyId());
+
+    Long iat = signature.getPayload().getIssuedAtTimeSeconds();
+    Long exp = signature.getPayload().getExpirationTimeSeconds();
+    assertNotNull(iat);
+    assertNotNull(exp);
+    assertEquals(3600L, exp - iat);
+    long currentTimeSecs = System.currentTimeMillis() / 1000;
+    assertTrue(Math.abs(currentTimeSecs - iat) < 60);
   }
 
   static GenericJson writeServiceAccountJson(

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountJwtAccessCredentialsTest.java

@@ -914,6 +914,8 @@ private void verifyJwtAccess(Map<String, List<String>> metadata, URI expectedAud
     }
     assertNotNull(assertion, "Bearer assertion not found");
     JsonWebSignature signature = JsonWebSignature.parse(JSON_FACTORY, assertion);
+    assertEquals("RS256", signature.getHeader().getAlgorithm());
+    assertEquals("JWT", signature.getHeader().getType());
     assertEquals(
         ServiceAccountJwtAccessCredentialsTest.SA_CLIENT_EMAIL, signature.getPayload().getIssuer());
     assertEquals(
@@ -922,6 +924,14 @@ private void verifyJwtAccess(Map<String, List<String>> metadata, URI expectedAud
     assertEquals(expectedAudience.toString(), signature.getPayload().getAudience());
     assertEquals(
         ServiceAccountJwtAccessCredentialsTest.SA_PRIVATE_KEY_ID, signature.getHeader().getKeyId());
+
+    Long iat = signature.getPayload().getIssuedAtTimeSeconds();
+    Long exp = signature.getPayload().getExpirationTimeSeconds();
+    assertNotNull(iat);
+    assertNotNull(exp);
+    assertEquals(3600L, exp - iat);
+    long currentTimeSecs = System.currentTimeMillis() / 1000;
+    assertTrue(Math.abs(currentTimeSecs - iat) < 60);
   }
 
   private static void testFromStreamException(InputStream stream, String expectedMessageContent) {

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/UserCredentialsTest.java

@@ -635,6 +635,7 @@ void saveAndRestoreUserCredential_saveAndRestored_doesNotThrow() throws IOExcept
             .setClientId(CLIENT_ID)
             .setClientSecret(CLIENT_SECRET)
             .setRefreshToken(REFRESH_TOKEN)
+            .setQuotaProjectId(QUOTA_PROJECT)
             .build();
 
     File file = File.createTempFile("GOOGLE_APPLICATION_CREDENTIALS", null, null);
@@ -649,6 +650,7 @@ void saveAndRestoreUserCredential_saveAndRestored_doesNotThrow() throws IOExcept
       assertEquals(userCredentials.getClientId(), restoredCredentials.getClientId());
       assertEquals(userCredentials.getClientSecret(), restoredCredentials.getClientSecret());
       assertEquals(userCredentials.getRefreshToken(), restoredCredentials.getRefreshToken());
+      assertEquals(userCredentials.getQuotaProjectId(), restoredCredentials.getQuotaProjectId());
     }
   }
 

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

@@ -52,6 +52,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
@@ -140,13 +141,16 @@ class ConnectionWorker implements AutoCloseable {
    * Tracks current inflight requests in the stream.
    */
   @GuardedBy("lock")
-  private long inflightRequests = 0;
+  private final AtomicLong inflightRequests = new AtomicLong(0);
 
   /*
    * Tracks current inflight bytes in the stream.
    */
   @GuardedBy("lock")
-  private long inflightBytes = 0;
+  private final AtomicLong inflightBytes = new AtomicLong(0);
+
+  private final TrackRequestQueueEarliestSendTime trackRequestQueueEarliestSendTime =
+      new TrackRequestQueueEarliestSendTime();
 
   /*
    * Tracks how often the stream was closed due to a retriable error. Streaming will stop when the
@@ -395,7 +399,7 @@ private void gatherHealthCheckMetrics(HealthCheckFields healthCheckFields) {
       healthCheckFields.queuedRequestCountMax = windowedQueuedRequestsMax;
       healthCheckFields.queuedRetryCountMax = windowedQueuedRetriesMax;
       healthCheckFields.msecLongestResponseWaitTime = windowedMilliResponseWaitTimeMax;
-      healthCheckFields.inflightBytes = inflightBytes;
+      healthCheckFields.inflightBytes = inflightBytes.get();
       healthCheckFields.requestsSentCount = windowedRequestsSent;
       healthCheckFields.responseCount = windowedResponsesAcked;
       if (HEALTH_CHECK_INTERVAL.toMillis() > 0) {
@@ -779,8 +783,8 @@ private void addMessageToFrontOfWaitingQueue(AppendRequestAndResponse requestWra
   @GuardedBy("lock")
   private void addMessageToWaitingQueue(
       AppendRequestAndResponse requestWrapper, boolean addToFront) {
-    ++this.inflightRequests;
-    this.inflightBytes += requestWrapper.messageSize;
+    this.inflightRequests.incrementAndGet();
+    this.inflightBytes.addAndGet(requestWrapper.messageSize);
     hasMessageInWaitingQueue.signal();
     requestProfilerHook.startOperation(
         RequestProfiler.OperationName.WAIT_QUEUE, requestWrapper.requestUniqueId);
@@ -896,11 +900,11 @@ private ApiFuture<AppendRowsResponse> appendInternal(
       }
       // Check if queue is going to be full before adding the request.
       if (this.limitExceededBehavior == FlowController.LimitExceededBehavior.ThrowException) {
-        if (this.inflightRequests + 1 >= this.maxInflightRequests) {
+        if (this.inflightRequests.get() + 1 >= this.maxInflightRequests) {
           throw new Exceptions.InflightRequestsLimitExceededException(
               writerId, this.maxInflightRequests);
         }
-        if (this.inflightBytes + requestWrapper.messageSize >= this.maxInflightBytes) {
+        if (this.inflightBytes.get() + requestWrapper.messageSize >= this.maxInflightBytes) {
           throw new Exceptions.InflightBytesLimitExceededException(writerId, this.maxInflightBytes);
         }
       }
@@ -926,8 +930,8 @@ private ApiFuture<AppendRowsResponse> appendInternal(
         return requestWrapper.appendResult;
       }
       requestProfilerHook.startOperation(RequestProfiler.OperationName.WAIT_QUEUE, requestUniqueId);
-      ++this.inflightRequests;
-      this.inflightBytes += requestWrapper.messageSize;
+      this.inflightRequests.incrementAndGet();
+      this.inflightBytes.addAndGet(requestWrapper.messageSize);
       requestWrapper.placedInWaitingQueueTime = Instant.now();
       waitingRequestQueue.addLast(requestWrapper);
       healthCheckMetrics.updateWindowedQueuedRequestsMax(
@@ -938,9 +942,9 @@ private ApiFuture<AppendRowsResponse> appendInternal(
       try {
         maybeWaitForInflightQuota();
       } catch (StatusRuntimeException ex) {
-        --this.inflightRequests;
+        this.inflightRequests.decrementAndGet();
         waitingRequestQueue.pollLast();
-        this.inflightBytes -= requestWrapper.messageSize;
+        this.inflightBytes.addAndGet(-requestWrapper.messageSize);
         throw ex;
       }
       requestProfilerHook.endOperation(
@@ -954,8 +958,8 @@ private ApiFuture<AppendRowsResponse> appendInternal(
   @GuardedBy("lock")
   private void maybeWaitForInflightQuota() {
     long start_time = System.currentTimeMillis();
-    while (this.inflightRequests >= this.maxInflightRequests
-        || this.inflightBytes >= this.maxInflightBytes) {
+    while (this.inflightRequests.get() >= this.maxInflightRequests
+        || this.inflightBytes.get() >= this.maxInflightBytes) {
       try {
         inflightReduced.await(100, TimeUnit.MILLISECONDS);
       } catch (InterruptedException e) {
@@ -998,6 +1002,11 @@ void setTestOnlyRunTimeExceptionInAppendLoop(
     this.testOnlyRunTimeExceptionInAppendLoop = testOnlyRunTimeExceptionInAppendLoop;
   }
 
+  @VisibleForTesting
+  Instant getEarliestSendTime() {
+    return trackRequestQueueEarliestSendTime.getEarliestSendTime();
+  }
+
   @VisibleForTesting()
   HealthCheckMetrics.HealthCheckFields gatherTestOnlyHealthCheckMetrics() {
     this.lock.lock();
@@ -1229,7 +1238,9 @@ private void appendLoop() {
         firstRequestForTableOrSchemaSwitch = true;
       }
       while (!localQueue.isEmpty()) {
-        localQueue.peekFirst().setRequestSendQueueTime();
+        AppendRequestAndResponse head = localQueue.peekFirst();
+        head.setRequestSendQueueTime();
+        trackRequestQueueEarliestSendTime.captureEarliest(head.requestSendTimeStamp);
         AppendRequestAndResponse wrapper = localQueue.pollFirst();
         AppendRowsRequest originalRequest = wrapper.message;
         String requestUniqueId = wrapper.requestUniqueId;
@@ -1642,6 +1653,9 @@ private void requestCallback(AppendRowsResponse response) {
     if (response.hasError()) {
       if (retryOnRetryableError(Code.values()[response.getError().getCode()], requestWrapper)) {
         log.info("Attempting to retry on error: " + response.getError().toString());
+        // Note that if we are retrying a request it is still in the system so we don't refresh the
+        // earliest send time. That way we can keep track of the earliest send time based on the
+        // first time the request was sent, which gives us a better idea of load on this worker.
         return;
       }
     }
@@ -1653,6 +1667,10 @@ private void requestCallback(AppendRowsResponse response) {
         this.lock.unlock();
       }
     }
+    // Since we have processed a response and have now removed that request from the system, go
+    // ahead and refresh the earliest send time, based on the remaining requests that are
+    // outstanding.
+    trackRequestQueueEarliestSendTime.discardAndRefresh();
 
     // We need a separate thread pool to unblock the next request callback.
     // Otherwise user may call append inside request callback, which may be blocked on waiting
@@ -1788,8 +1806,8 @@ private AppendRequestAndResponse pollInflightRequestQueue(boolean pollLast) {
     AppendRequestAndResponse requestWrapper =
         pollLast ? inflightRequestQueue.pollLast() : inflightRequestQueue.poll();
     requestWrapper.requestSendTimeStamp = null;
-    --this.inflightRequests;
-    this.inflightBytes -= requestWrapper.messageSize;
+    this.inflightRequests.decrementAndGet();
+    this.inflightBytes.addAndGet(-requestWrapper.messageSize);
     this.inflightReduced.signal();
     return requestWrapper;
   }
@@ -1881,9 +1899,15 @@ void setRequestSendQueueTime() {
 
   /** Returns the current workload of this worker. */
   public Load getLoad() {
+    Duration timeSinceLastCallback = Duration.ZERO;
+    Instant earliestSendTime = trackRequestQueueEarliestSendTime.getEarliestSendTime();
+    if (earliestSendTime != null) {
+      timeSinceLastCallback = Duration.between(earliestSendTime, Instant.now());
+    }
     return Load.create(
-        inflightBytes,
-        inflightRequests,
+        timeSinceLastCallback,
+        inflightBytes.get(),
+        inflightRequests.get(),
         destinationSet.size(),
         maxInflightBytes,
         maxInflightRequests);
@@ -1896,11 +1920,15 @@ public Load getLoad() {
   @AutoValue
   public abstract static class Load {
 
-    // Consider the load on this worker to be overwhelmed when above some percentage of
-    // in-flight bytes or in-flight requests count.
+    // Consider the load on this worker to be overwhelmed when above some inflight latency or
+    // percentage of in-flight bytes or in-flight requests count.
+    private static Duration overwhelmedTimeSinceLastCallback = Duration.ofSeconds(3);
     private static double overwhelmedInflightCount = 0.2;
     private static double overwhelmedInflightBytes = 0.2;
 
+    // Time we have spent waiting for a response in the worker.
+    abstract Duration timeSinceLastCallback();
+
     // Number of in-flight requests bytes in the worker.
     abstract long inFlightRequestsBytes();
 
@@ -1917,12 +1945,14 @@ public abstract static class Load {
     abstract long maxInflightCount();
 
     static Load create(
+        Duration timeSinceLastCallback,
         long inFlightRequestsBytes,
         long inFlightRequestsCount,
         long destinationCount,
         long maxInflightBytes,
         long maxInflightCount) {
       return new AutoValue_ConnectionWorker_Load(
+          timeSinceLastCallback,
           inFlightRequestsBytes,
           inFlightRequestsCount,
           destinationCount,
@@ -1934,20 +1964,29 @@ boolean isOverwhelmed() {
       // Consider only in flight bytes and count for now, as by experiment those two are the most
       // efficient and has great simplity.
       return inFlightRequestsCount() > overwhelmedInflightCount * maxInflightCount()
-          || inFlightRequestsBytes() > overwhelmedInflightBytes * maxInflightBytes();
+          || inFlightRequestsBytes() > overwhelmedInflightBytes * maxInflightBytes()
+          || timeSinceLastCallback().compareTo(overwhelmedTimeSinceLastCallback) > 0;
     }
 
-    // Compares two different load. First compare in flight request bytes split by size 1024 bucket.
+    // Compares two different load. First compare the timeSinceLastCallback bucketed into 1 second
+    // intervals.
+    // Then compare in flight request bytes split by size 1024 bucket.
     // Then compare the inflight requests count.
     // Then compare destination count of the two connections.
     public static final Comparator<Load> LOAD_COMPARATOR =
-        Comparator.comparing((Load key) -> (int) (key.inFlightRequestsBytes() / 1024))
+        Comparator.comparing((Load key) -> (int) key.timeSinceLastCallback().getSeconds())
+            .thenComparing((Load key) -> (int) (key.inFlightRequestsBytes() / 1024))
             .thenComparing((Load key) -> (int) (key.inFlightRequestsCount() / 100))
             .thenComparing(Load::destinationCount);
 
     // Compares two different load without bucket, used in smaller scale unit testing.
+    // First compare the timeSinceLastCallback.
+    // Then compare in flight request bytes.
+    // Then compare the inflight requests count.
+    // Then compare destination count of the two connections.
     public static final Comparator<Load> TEST_LOAD_COMPARATOR =
-        Comparator.comparing((Load key) -> (int) key.inFlightRequestsBytes())
+        Comparator.comparing(Load::timeSinceLastCallback)
+            .thenComparing((Load key) -> (int) key.inFlightRequestsBytes())
             .thenComparing((Load key) -> (int) key.inFlightRequestsCount())
             .thenComparing(Load::destinationCount);
 
@@ -1960,6 +1999,11 @@ public static void setOverwhelmedBytesThreshold(double newThreshold) {
     public static void setOverwhelmedCountsThreshold(double newThreshold) {
       overwhelmedInflightCount = newThreshold;
     }
+
+    @VisibleForTesting
+    public static void setOverwhelmedTimeSinceLastCallbackThreshold(Duration newThreshold) {
+      overwhelmedTimeSinceLastCallback = newThreshold;
+    }
   }
 
   @VisibleForTesting
@@ -1985,4 +2029,36 @@ static TableSchemaAndTimestamp create(long updateTimeStamp, TableSchema updatedS
       return new AutoValue_ConnectionWorker_TableSchemaAndTimestamp(updateTimeStamp, updatedSchema);
     }
   }
+
+  class TrackRequestQueueEarliestSendTime {
+    private final AtomicReference<Instant> earliestSendTime = new AtomicReference<>(null);
+
+    public void captureEarliest(Instant sendTime) {
+      // This method records the given sendTime only if earliestSendTime is currently NULL.
+      if (sendTime == null) {
+        return;
+      }
+      earliestSendTime.compareAndSet(null, sendTime);
+    }
+
+    public void discardAndRefresh() {
+      Instant newEarliestSendTime = null;
+      lock.lock();
+      try {
+        if (!inflightRequestQueue.isEmpty()) {
+          AppendRequestAndResponse head = inflightRequestQueue.peekFirst();
+          if (head != null) {
+            newEarliestSendTime = head.requestSendTimeStamp;
+          }
+        }
+      } finally {
+        lock.unlock();
+      }
+      earliestSendTime.set(newEarliestSendTime);
+    }
+
+    public Instant getEarliestSendTime() {
+      return earliestSendTime.get();
+    }
+  }
 }