[BUG] azure-core-amqp: token-manager close or terminal error can interrupt refresh and close the shared dispatcher pipe #51112

import com.azure.core.amqp.ClaimsBasedSecurityNode;
import com.azure.core.amqp.implementation.ActiveClientTokenManager;
import com.azure.core.amqp.implementation.ReactorDispatcher;
import java.nio.channels.Pipe;
import java.time.OffsetDateTime;
import java.time.Duration;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
import org.apache.qpid.proton.Proton;
import org.apache.qpid.proton.reactor.Reactor;
import reactor.core.publisher.Mono;
import reactor.core.publisher.MonoSink;

public class RefreshInterruptProbe {
  public static void main(String[] args) throws Exception {
    boolean expectInterrupted = args.length == 0 || Boolean.parseBoolean(args[0]);
    String mode = args.length < 2 ? "success" : args[1];
    Pipe pipe = Pipe.open();
    Reactor reactor = Proton.reactor();
    ReactorDispatcher dispatcher = new ReactorDispatcher("independent-refresh-probe", reactor, pipe);
    AtomicReference<Throwable> shutdown = new AtomicReference<>();
    dispatcher.getShutdownSignal().subscribe(x -> {}, shutdown::set);
    AtomicInteger calls = new AtomicInteger();
    AtomicReference<MonoSink<OffsetDateTime>> refreshSink = new AtomicReference<>();
    AtomicBoolean release = new AtomicBoolean();
    AtomicBoolean interrupted = new AtomicBoolean();
    CountDownLatch entered = new CountDownLatch(1), finished = new CountDownLatch(1);
    ClaimsBasedSecurityNode node = new ClaimsBasedSecurityNode() {
      public Mono<OffsetDateTime> authorize(String a, String s) {
        if (calls.incrementAndGet() == 1) return Mono.just(OffsetDateTime.now().plusSeconds(3));
        return Mono.create(sink -> {
          refreshSink.set(sink);
          entered.countDown();
          if (mode.equals("same")) { sink.success(OffsetDateTime.now().plusHours(1)); release.set(true); }
          // Model the initiating thread still inside CBS dispatch while a reply completes elsewhere.
          long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
          while (!release.get() && System.nanoTime() < deadline) Thread.onSpinWait();
          interrupted.set(Thread.currentThread().isInterrupted());
          try { dispatcher.invoke(() -> {}); } catch (Exception e) { throw new RuntimeException(e); }
          finally { finished.countDown(); }
        });
      }
      public void close() {}
    };
    ActiveClientTokenManager manager = new ActiveClientTokenManager(Mono.just(node), "audience", "scope");
    manager.authorize().block(Duration.ofSeconds(2));
    if (!entered.await(3, TimeUnit.SECONDS)) throw new AssertionError("No refresh");
    if (mode.equals("same")) { /* source completes on its own thread */ }
    else if (mode.equals("close")) manager.close();
    else if (mode.equals("terminal")) refreshSink.get().error(new IllegalStateException("injected terminal CBS failure"));
    else if (mode.equals("transient")) refreshSink.get().error(new com.azure.core.amqp.exception.AmqpException(true, "injected transient CBS failure", new com.azure.core.amqp.exception.AmqpErrorContext("test")));
    else refreshSink.get().success(OffsetDateTime.now().plusHours(1));
    if (!mode.equals("same")) release.set(true);
    if (!finished.await(3, TimeUnit.SECONDS)) throw new AssertionError("Refresh did not return");
    Throwable cause = shutdown.get() == null ? null : shutdown.get().getCause();
    System.out.printf("interrupted=%s pipeOpen=%s shutdownCause=%s calls=%d%n",
      interrupted.get(), pipe.sink().isOpen(), cause, calls.get());
    manager.close(); pipe.sink().close(); pipe.source().close(); reactor.free();
    if (interrupted.get() != expectInterrupted) throw new AssertionError("Unexpected interruption");
    if (expectInterrupted && !(cause instanceof java.nio.channels.ClosedByInterruptException))
      throw new AssertionError("Did not reproduce actual dispatcher pipe failure");
    if (!expectInterrupted && cause != null) throw new AssertionError("Dispatcher failed");
  }
}