AwaitingLeadershipProxy extends AbstractInvocationHandler { public static <U> U newProxyInstance(Class<U> interfaceClass, Supplier<U> delegateSupplier, LeaderElectionService leaderElectionService) { AwaitingLeadershipProxy<U> proxy = new AwaitingLeadershipProxy<>( delegateSupplier, leaderElectionService, interfaceClass); proxy.tryToGainLeadership(); return (U) Proxy.newProxyInstance( interfaceClass.getClassLoader(), new Class<?>[] { interfaceClass, Closeable.class }, proxy); } private  AwaitingLeadershipProxy(
            Supplier<T> delegateSupplier,
            LeaderElectionService leaderElectionService,
            Class<T> interfaceClass); static U newProxyInstance(Class<U> interfaceClass,
                                         Supplier<U> delegateSupplier,
                                         LeaderElectionService leaderElectionService);  }
@Test @SuppressWarnings("SelfEquals") public void shouldAllowObjectMethodsWhenLeading() { Runnable proxy = AwaitingLeadershipProxy.newProxyInstance( Runnable.class, delegateSupplier, leaderElectionService); assertThat(proxy.hashCode()).isNotNull(); assertThat(proxy.equals(proxy)).isTrue(); assertThat(proxy.equals(null)).isFalse(); assertThat(proxy.toString()).startsWith("com.palantir.leader.proxy.AwaitingLeadershipProxy@"); }
@Test public void listenableFutureMethodsDoNotBlockWhenNotLeading() throws ExecutionException, InterruptedException { ReturnsListenableFutureImpl listenableFuture = new ReturnsListenableFutureImpl(); ReturnsListenableFuture proxy = AwaitingLeadershipProxy.newProxyInstance( ReturnsListenableFuture.class, () -> listenableFuture, leaderElectionService); waitForLeadershipToBeGained(); SettableFuture<StillLeadingStatus> inProgressCheck = SettableFuture.create(); when(leaderElectionService.isStillLeading(any(LeadershipToken.class))).thenReturn(inProgressCheck); ListenableFuture<?> future = proxy.future(); assertThat(future).isNotDone(); inProgressCheck.set(StillLeadingStatus.NOT_LEADING); assertThat(future.isDone()); expect.expectCause(isA(NotCurrentLeaderException.class)); future.get(); }
@Test public void listenableFutureMethodsDoNotBlockWhenLeading() throws InterruptedException, ExecutionException { ReturnsListenableFutureImpl listenableFuture = new ReturnsListenableFutureImpl(); ReturnsListenableFuture proxy = AwaitingLeadershipProxy.newProxyInstance( ReturnsListenableFuture.class, () -> listenableFuture, leaderElectionService); waitForLeadershipToBeGained(); SettableFuture<StillLeadingStatus> inProgressCheck = SettableFuture.create(); when(leaderElectionService.isStillLeading(any(LeadershipToken.class))).thenReturn(inProgressCheck); ListenableFuture<?> future = proxy.future(); assertThat(future).isNotDone(); inProgressCheck.set(StillLeadingStatus.LEADING); assertThat(future).isNotDone(); listenableFuture.future.set(null); future.get(); }
@Test public void listenableFutureMethodsRetryProxyFailures() throws InterruptedException, ExecutionException { ReturnsListenableFutureImpl listenableFuture = new ReturnsListenableFutureImpl(); ReturnsListenableFuture proxy = AwaitingLeadershipProxy.newProxyInstance( ReturnsListenableFuture.class, () -> listenableFuture, leaderElectionService); waitForLeadershipToBeGained(); SettableFuture<StillLeadingStatus> inProgressCheck = SettableFuture.create(); when(leaderElectionService.isStillLeading(any(LeadershipToken.class))) .thenAnswer($ -> { Uninterruptibles.sleepUninterruptibly(37, TimeUnit.MILLISECONDS); return Futures.immediateFuture(StillLeadingStatus.NO_QUORUM); }