fix: Populate final stauts to initial request during connection shutdown (#2228)

* chore(main): release 2.41.1 (#2222)

Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>

* chore(main): release 2.41.1 (#2222)

Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: release-please[bot] <55107282+release-please[bot]@users.noreply.github.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: 3 people

README.md

@@ -50,20 +50,20 @@ If you are using Maven without the BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.21.0')
+implementation platform('com.google.cloud:libraries-bom:26.22.0')
 
 implementation 'com.google.cloud:google-cloud-bigquerystorage'
 ```
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-bigquerystorage:2.41.0'
+implementation 'com.google.cloud:google-cloud-bigquerystorage:2.41.1'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-bigquerystorage" % "2.41.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-bigquerystorage" % "2.41.1"
 ```
 <!-- {x-version-update-end} -->
 
@@ -220,7 +220,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-bigquerystorage/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-bigquerystorage.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-bigquerystorage/2.41.0
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-bigquerystorage/2.41.1
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles

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

@@ -839,8 +839,24 @@ private void cleanupInflightRequests() {
             + streamName
             + " id: "
             + writerId);
-    while (!localQueue.isEmpty()) {
-      localQueue.pollFirst().appendResult.setException(finalStatus);
+    int sizeOfQueue = localQueue.size();
+    for (int i = 0; i < sizeOfQueue; i++) {
+      if (i == 0) {
+        localQueue.pollFirst().appendResult.setException(finalStatus);
+      } else {
+        localQueue
+            .pollFirst()
+            .appendResult
+            .setException(
+                new Exceptions.StreamWriterClosedException(
+                    Status.fromCode(Code.ABORTED)
+                        .withDescription(
+                            "Connection is aborted due to an unrecoverable failure of "
+                                + "another request sharing the connection. Please retry this "
+                                + "request."),
+                    streamName,
+                    writerId));
+      }
     }
   }
 

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

@@ -418,7 +418,11 @@ public void testThrowExceptionWhileWithinAppendLoop() throws Exception {
           assertThrows(
               ExecutionException.class,
               () -> futures.get(finalI).get().getAppendResult().getOffset().getValue());
-      assertThat(ex.getCause()).hasMessageThat().contains("Any exception can happen.");
+      if (i == 0) {
+        assertThat(ex.getCause()).hasMessageThat().contains("Any exception can happen.");
+      } else {
+        assertThat(ex.getCause()).hasMessageThat().contains("Connection is aborted due to ");
+      }
     }
 
     // The future append will directly fail.
@@ -654,7 +658,11 @@ public void testThrowExceptionWhileWithinAppendLoop_MaxWaitTimeExceed() throws E
           assertThrows(
               ExecutionException.class,
               () -> futures.get(finalI).get().getAppendResult().getOffset().getValue());
-      assertThat(ex.getCause()).hasMessageThat().contains("Request has waited in inflight queue");
+      if (i == 0) {
+        assertThat(ex.getCause()).hasMessageThat().contains("Request has waited in inflight queue");
+      } else {
+        assertThat(ex.getCause()).hasMessageThat().contains("Connection is aborted due to ");
+      }
     }
 
     // The future append will directly fail.

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

@@ -39,6 +39,7 @@
 import com.google.api.gax.rpc.UnknownException;
 import com.google.cloud.bigquery.storage.test.Test.FooType;
 import com.google.cloud.bigquery.storage.v1.ConnectionWorkerPool.Settings;
+import com.google.cloud.bigquery.storage.v1.Exceptions.StreamWriterClosedException;
 import com.google.cloud.bigquery.storage.v1.StorageError.StorageErrorCode;
 import com.google.cloud.bigquery.storage.v1.StreamWriter.SingleConnectionOrConnectionPool.Kind;
 import com.google.common.base.Strings;
@@ -666,8 +667,12 @@ public void serverCloseWhileRequestsInflight() throws Exception {
 
     // Server close should properly handle all inflight requests.
     for (int i = 0; i < appendCount; i++) {
-      ApiException actualError = assertFutureException(ApiException.class, futures.get(i));
-      assertEquals(Code.INVALID_ARGUMENT, actualError.getStatusCode().getCode());
+      if (i == 0) {
+        ApiException actualError = assertFutureException(ApiException.class, futures.get(i));
+        assertEquals(Code.INVALID_ARGUMENT, actualError.getStatusCode().getCode());
+      } else {
+        assertFutureException(StreamWriterClosedException.class, futures.get(i));
+      }
     }
 
     writer.close();
@@ -988,7 +993,13 @@ public void testThrowExceptionWhileWithinAppendLoop_MaxWaitTimeExceed() throws E
           assertThrows(
               ExecutionException.class,
               () -> futures.get(finalI).get().getAppendResult().getOffset().getValue());
-      assertThat(ex.getCause()).hasMessageThat().contains("Request has waited in inflight queue");
+      if (i == 0) {
+        assertThat(ex.getCause()).hasMessageThat().contains("Request has waited in inflight queue");
+      } else {
+        assertThat(ex.getCause())
+            .hasMessageThat()
+            .contains("Connection is aborted due to an unrecoverable");
+      }
     }
   }
 
@@ -1027,7 +1038,11 @@ public void testAppendWithResetNeverSuccess() throws Exception {
       assertEquals(futures.get(0).get().getAppendResult().getOffset().getValue(), 0);
       // after 5 seconds, the requests will bail out.
       for (int i = 1; i < appendCount; i++) {
-        assertFutureException(AbortedException.class, futures.get(i));
+        if (i == 1) {
+          assertFutureException(AbortedException.class, futures.get(i));
+        } else {
+          assertFutureException(StreamWriterClosedException.class, futures.get(i));
+        }
       }
     }
   }
@@ -1048,7 +1063,11 @@ public void testAppendWithResetNeverSuccessWithMultiplexing() throws Exception {
       assertEquals(futures.get(0).get().getAppendResult().getOffset().getValue(), 0);
       // after 5 seconds, the requests will bail out.
       for (int i = 1; i < appendCount; i++) {
-        assertFutureException(AbortedException.class, futures.get(i));
+        if (i == 1) {
+          assertFutureException(AbortedException.class, futures.get(i));
+        } else {
+          assertFutureException(StreamWriterClosedException.class, futures.get(i));
+        }
       }
     }
   }