add autoclosable for asyncresultset

googleapis · Jan 6, 2025 · 13ff2db · 13ff2db
1 parent f72a36a
commit 13ff2db
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 20 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AsyncResultSetImpl.java
@@ -489,8 +489,8 @@ public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
       // Start to fetch data and buffer these.
       this.result = SettableApiFuture.create();
       this.state = State.STREAMING_INITIALIZED;
-      initiateProduceRows();
-      //      this.service.execute(new InitiateStreamingRunnable());
+      this.service.execute(new InitiateStreamingRunnable());
+//      initiateProduceRows();
       this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec));
       this.callback = Preconditions.checkNotNull(cb);
       pausedLatch.countDown();

diff --git a/...e-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java b/...e-cloud-spanner/src/test/java/com/google/cloud/spanner/RetryOnInvalidatedSessionTest.java
@@ -1593,25 +1593,26 @@ private void asyncTransactionManager_readAsync(
           AsyncTransactionStep<Void, Long> count =
               context.then(
                   (transaction, ignored) -> {
-                    AsyncResultSet rs = fn.apply(transaction);
-                    ApiFuture<Void> fut =
-                        rs.setCallback(
-                            queryExecutor,
-                            resultSet -> {
-                              while (true) {
-                                switch (resultSet.tryNext()) {
-                                  case OK:
-                                    counter.incrementAndGet();
-                                    break;
-                                  case DONE:
-                                    return CallbackResponse.DONE;
-                                  case NOT_READY:
-                                    return CallbackResponse.CONTINUE;
+                    try(AsyncResultSet rs = fn.apply(transaction)) {
+                      ApiFuture<Void> fut =
+                          rs.setCallback(
+                              queryExecutor,
+                              resultSet -> {
+                                while (true) {
+                                  switch (resultSet.tryNext()) {
+                                    case OK:
+                                      counter.incrementAndGet();
+                                      break;
+                                    case DONE:
+                                      return CallbackResponse.DONE;
+                                    case NOT_READY:
+                                      return CallbackResponse.CONTINUE;
+                                  }
                                 }
-                              }
-                            });
-                    return ApiFutures.transform(
-                        fut, input -> counter.get(), MoreExecutors.directExecutor());
+                              });
+                      return ApiFutures.transform(
+                          fut, input -> counter.get(), MoreExecutors.directExecutor());
+                    }
                   },
                   executor);
           CommitTimestampFuture ts = count.commitAsync();