Fix/background identity interrupt

core/src/main/java/tech/ydb/core/auth/BackgroundIdentity.java

@@ -2,6 +2,7 @@
 
 import java.time.Clock;
 import java.time.Instant;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
@@ -72,6 +73,7 @@ public void close() {
     }
 
     private State updateState(State current, State next) {
+        Objects.requireNonNull(next, "next state cannot be null");
         if (state.compareAndSet(current, next)) {
             next.init();
         }
@@ -92,7 +94,8 @@ private <T> T unwrap(CompletableFuture<T> future) {
         } catch (InterruptedException ex) {
             logger.error("updating of authentication token was interrupted", ex);
             Thread.currentThread().interrupt();
-            return null;
+            // returning null here would poison the state reference and break every following getToken()
+            throw new RuntimeException("authentication update was interrupted", ex);
         }
     }
 

core/src/main/java/tech/ydb/core/auth/StaticCredentials.java

@@ -117,26 +117,28 @@ private void tryLogin(CompletableFuture<Token> future) {
             }
 
             rpc.getExecutor().submit(() -> {
-                try (GrpcTransport transport = rpc.createTransport()) {
-                    GrpcRequestSettings grpcSettings = GrpcRequestSettings.newBuilder()
-                            .withDeadline(Duration.ofSeconds(LOGIN_TIMEOUT_SECONDS))
-                            .build();
-
-                    transport.unaryCall(AuthServiceGrpc.getLoginMethod(), grpcSettings, request)
-                            .thenApply(OperationBinder.bindSync(
-                                    YdbAuth.LoginResponse::getOperation,
-                                    YdbAuth.LoginResult.class
-                            ))
-                            .whenComplete((resp, th) -> {
+                GrpcTransport transport = rpc.createTransport();
+                GrpcRequestSettings grpcSettings = GrpcRequestSettings.newBuilder()
+                        .withDeadline(Duration.ofSeconds(LOGIN_TIMEOUT_SECONDS))
+                        .build();
+
+                transport.unaryCall(AuthServiceGrpc.getLoginMethod(), grpcSettings, request)
+                        .thenApply(OperationBinder.bindSync(
+                                YdbAuth.LoginResponse::getOperation,
+                                YdbAuth.LoginResult.class
+                        ))
+                        .whenComplete((resp, th) -> {
+                            try {
                                 if (resp != null) {
                                     handleResult(future, resp);
                                 }
                                 if (th != null) {
                                     handleException(future, th);
                                 }
-                            })
-                            .join();
-                }
+                            } finally {
+                                transport.close();
+                            }
+                        });
             });
         }
 

core/src/test/java/tech/ydb/core/auth/BackgroundIdentityTest.java

@@ -0,0 +1,68 @@
+package tech.ydb.core.auth;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class BackgroundIdentityTest {
+    private final Instant now = Instant.EPOCH;
+    private final Clock clock = Clock.fixed(now, ZoneId.of("UTC"));
+
+    private static class MockedRpc implements BackgroundIdentity.Rpc {
+        private final CompletableFuture<Token> tokenFuture = new CompletableFuture<>();
+
+        @Override
+        public CompletableFuture<Token> getTokenAsync() {
+            return tokenFuture;
+        }
+
+        @Override
+        public int getTimeoutSeconds() {
+            return 60;
+        }
+    }
+
+    @Test(timeout = 30_000)
+    public void interruptedGetTokenDoesNotBreakIdentity() throws InterruptedException {
+        MockedRpc rpc = new MockedRpc();
+        BackgroundIdentity identity = new BackgroundIdentity(clock, rpc);
+
+        AtomicReference<Throwable> caught = new AtomicReference<>();
+        Thread reader = new Thread(() -> {
+            try {
+                identity.getToken();
+            } catch (Throwable th) {
+                caught.set(th);
+            }
+        });
+
+        // the login never answers, so the reader blocks in the sync state and gets interrupted there
+        reader.start();
+        // the reader may not have reached the await yet, interrupting early is handled the same way
+        reader.interrupt();
+        reader.join();
+
+        Assert.assertNotNull("interrupted getToken must report a failure", caught.get());
+        Assert.assertFalse(
+                "interrupt must not surface as a NullPointerException",
+                caught.get() instanceof NullPointerException
+        );
+
+        // the identity must recover once the login completes
+        rpc.getTokenAsync().complete(
+                new BackgroundIdentity.Rpc.Token(
+                        "token-value",
+                        now.plus(Duration.ofHours(2)),
+                        now.plus(Duration.ofHours(1))
+                )
+        );
+
+        Assert.assertEquals("token-value", identity.getToken());
+    }
+}

core/src/test/java/tech/ydb/core/auth/CredentialsAuthProviderTest.java

@@ -4,6 +4,7 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.protobuf.Any;
@@ -202,6 +203,39 @@ public void syncRefreshTokenTest() {
         identity.close();
     }
 
+    @Test
+    public void interruptedGetTokenDoesNotBreakIdentity() throws InterruptedException {
+        Mockito.when(clock.instant()).thenReturn(now);
+
+        CompletableFuture<Result<YdbAuth.LoginResponse>> result = new CompletableFuture<>();
+        Mockito.when(transport.unaryCall(Mockito.eq(AuthServiceGrpc.getLoginMethod()), Mockito.any(), Mockito.any()))
+                .thenReturn(result);
+
+        tech.ydb.auth.AuthIdentity identity = createAuth("user", "password");
+
+        AtomicReference<Throwable> caught = new AtomicReference<>();
+        Thread reader = new Thread(() -> {
+            try {
+                identity.getToken();
+            } catch (Throwable th) {
+                caught.set(th);
+            }
+        });
+
+        reader.start();
+        reader.interrupt();
+        reader.join();
+
+        Assert.assertNotNull("interrupted getToken must report a failure", caught.get());
+        Assert.assertEquals("authentication update was interrupted", caught.get().getMessage());
+
+        String token = JwtBuilder.create(now.plus(Duration.ofHours(2)), now);
+        result.complete(Result.success(responseOk(token)));
+
+        // the identity must recover once the login completes
+        Assert.assertEquals(token, identity.getToken());
+    }
+
     private tech.ydb.auth.AuthIdentity createAuth(String login, String password) {
         return new StaticCredentials(clock, login, password)
                 .createAuthIdentity(rpc);

core/src/test/resources/log4j2.xml

@@ -18,6 +18,9 @@
         </Logger>
         <Logger name="tech.ydb.core.impl.pool.GrpcChannel" level="error"  additivity="false">
         </Logger>
+        <Logger name="tech.ydb.core.auth.BackgroundIdentity" level="off" additivity="false">
+            <AppenderRef ref="Console"/>
+        </Logger>
 
         <Root level="debug" >
             <AppenderRef ref="Console"/>