diff --git a/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessage.java b/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessage.java index 7401daea4..f105ab1ad 100644 --- a/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessage.java +++ b/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessage.java @@ -18,6 +18,7 @@ package org.apache.hugegraph.computer.core.sender; import java.nio.ByteBuffer; +import java.util.concurrent.CompletableFuture; import org.apache.hugegraph.computer.core.network.message.MessageType; @@ -26,11 +27,18 @@ public class QueuedMessage { private final int partitionId; private final MessageType type; private final ByteBuffer buffer; + private final CompletableFuture controlFuture; public QueuedMessage(int partitionId, MessageType type, ByteBuffer buffer) { + this(partitionId, type, buffer, null); + } + + public QueuedMessage(int partitionId, MessageType type, ByteBuffer buffer, + CompletableFuture controlFuture) { this.partitionId = partitionId; this.type = type; this.buffer = buffer; + this.controlFuture = controlFuture; } public int partitionId() { @@ -44,4 +52,8 @@ public MessageType type() { public ByteBuffer buffer() { return this.buffer; } + + public CompletableFuture controlFuture() { + return this.controlFuture; + } } diff --git a/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSender.java b/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSender.java index b2006b886..6f452074f 100644 --- a/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSender.java +++ b/computer/computer-core/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSender.java @@ -85,15 +85,20 @@ public void addWorkerClient(int workerId, TransportClient client) { public CompletableFuture send(int workerId, MessageType type) throws InterruptedException { WorkerChannel channel = this.channels[channelId(workerId)]; - CompletableFuture future = channel.newFuture(); - future.whenComplete((r, e) -> { - channel.resetFuture(future); - }); + CompletableFuture future = new CompletableFuture<>(); + if (!channel.setControlFuture(future)) { + return future; + } /* * Control message just need message type is enough, * partitionId = -1 and buffer = null represents a meaningless value */ - channel.queue.put(new QueuedMessage(-1, type, null)); + try { + channel.queue.put(new QueuedMessage(-1, type, null, future)); + } catch (InterruptedException e) { + channel.completeControlFuture(future, e); + throw e; + } return future; } @@ -108,7 +113,7 @@ public void send(int workerId, QueuedMessage message) public void transportExceptionCaught(TransportException cause, ConnectionId connectionId) { for (WorkerChannel channel : this.channels) { if (channel.client.connectionId().equals(connectionId)) { - channel.futureRef.get().completeExceptionally(cause); + channel.transportExceptionCaught(cause); } } } @@ -137,11 +142,19 @@ public void run() { ++emptyQueueCount; continue; } - if (channel.doSend(message)) { - // Only consume the message after it is sent + try { + if (channel.doSend(message)) { + // Only consume the message after it is sent + channel.queue.take(); + } else { + ++busyClientCount; + } + } catch (TransportException | RuntimeException e) { + channel.failDataSend(e); + // Discard the failed data message to keep sending channel.queue.take(); - } else { - ++busyClientCount; + LOG.warn("Failed to send {} message to {}, " + + "discard it", message.type(), channel, e); } } int channelCount = channels.length; @@ -172,9 +185,6 @@ public void run() { "Interrupted when waiting for message " + "queue not empty"); } - } catch (TransportException e) { - // TODO: should handle this in main workflow thread - throw new ComputerException("Failed to send message", e); } } LOG.info("The send-executor is terminated"); @@ -227,75 +237,127 @@ private static class WorkerChannel { private final MessageQueue queue; // Each target worker has a TransportClient private final TransportClient client; - private final AtomicReference> futureRef; + private final AtomicReference> controlFutureRef; + private final AtomicReference dataFailureRef; public WorkerChannel(int workerId, MessageQueue queue, TransportClient client) { this.workerId = workerId; this.queue = queue; this.client = client; - this.futureRef = new AtomicReference<>(); - } - - public CompletableFuture newFuture() { - CompletableFuture future = new CompletableFuture<>(); - if (!this.futureRef.compareAndSet(null, future)) { - throw new ComputerException("The origin future must be null"); - } - return future; - } - - public void resetFuture(CompletableFuture future) { - if (!this.futureRef.compareAndSet(future, null)) { - throw new ComputerException("Failed to reset futureRef, " + - "expect future object is %s, " + - "but some thread modified it", - future); - } + this.controlFutureRef = new AtomicReference<>(); + this.dataFailureRef = new AtomicReference<>(); } public boolean doSend(QueuedMessage message) throws TransportException, InterruptedException { switch (message.type()) { case START: - this.sendStartMessage(); + this.sendStartMessage(message.controlFuture()); return true; case FINISH: - this.sendFinishMessage(); + this.sendFinishMessage(message.controlFuture()); return true; default: return this.sendDataMessage(message); } } - public void sendStartMessage() throws TransportException { - this.client.startSessionAsync().whenComplete((r, e) -> { - CompletableFuture future = this.futureRef.get(); - assert future != null; - - if (e != null) { - LOG.info("Failed to start session connected to {}", this); - future.completeExceptionally(e); - } else { - LOG.info("Start session connected to {}", this); - future.complete(null); - } - }); + public void sendStartMessage(CompletableFuture future) { + if (!this.controlFutureInFlight(future)) { + return; + } + try { + this.client.startSessionAsync().whenComplete((r, e) -> { + if (e != null) { + LOG.info("Failed to start session connected to {}", this); + } else { + LOG.info("Start session connected to {}", this); + } + this.completeControlFuture(future, e); + }); + } catch (TransportException e) { + this.completeControlFuture(future, e); + } catch (RuntimeException e) { + this.completeControlFuture(future, e); + } + } + + public void sendFinishMessage(CompletableFuture future) { + if (!this.controlFutureInFlight(future)) { + return; + } + try { + this.client.finishSessionAsync().whenComplete((r, e) -> { + if (e != null) { + LOG.info("Failed to finish session connected to {}", this); + } else { + LOG.info("Finish session connected to {}", this); + } + this.completeControlFuture(future, e); + }); + } catch (TransportException e) { + this.completeControlFuture(future, e); + } catch (RuntimeException e) { + this.completeControlFuture(future, e); + } + } + + public void transportExceptionCaught(TransportException cause) { + CompletableFuture future = this.controlFutureRef.get(); + if (future == null) { + this.failDataSend(cause); + } else { + this.completeControlFuture(future, cause); + } + } + + public void failControlFuture(Throwable cause) { + CompletableFuture future = this.controlFutureRef.getAndSet(null); + if (future != null) { + future.completeExceptionally(cause); + } + } + + public void failDataSend(Throwable cause) { + this.dataFailureRef.compareAndSet(null, cause); + this.failControlFuture(this.dataFailureRef.get()); } - public void sendFinishMessage() throws TransportException { - this.client.finishSessionAsync().whenComplete((r, e) -> { - CompletableFuture future = this.futureRef.get(); - assert future != null; - - if (e != null) { - LOG.info("Failed to finish session connected to {}", this); - future.completeExceptionally(e); - } else { - LOG.info("Finish session connected to {}", this); - future.complete(null); + private boolean setControlFuture(CompletableFuture future) { + Throwable failure = this.dataFailureRef.get(); + if (failure != null) { + future.completeExceptionally(failure); + return false; + } + if (this.controlFutureRef.compareAndSet(null, future)) { + failure = this.dataFailureRef.get(); + if (failure == null) { + return true; } - }); + this.completeControlFuture(future, failure); + return false; + } + ComputerException e = new ComputerException( + "The origin future must be null"); + future.completeExceptionally(e); + return false; + } + + private boolean controlFutureInFlight(CompletableFuture future) { + return this.controlFutureRef.get() == future; + } + + private void completeControlFuture(CompletableFuture future, + Throwable cause) { + if (!this.controlFutureRef.compareAndSet(future, null)) { + return; + } + if (cause == null) { + future.complete(null); + } else { + future.completeExceptionally(cause); + } } public boolean sendDataMessage(QueuedMessage message) diff --git a/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/network/netty/NettyTransportClientTest.java b/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/network/netty/NettyTransportClientTest.java index fac4046b6..6f254bdad 100644 --- a/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/network/netty/NettyTransportClientTest.java +++ b/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/network/netty/NettyTransportClientTest.java @@ -127,7 +127,6 @@ public void testSend() throws IOException { @Test public void testDataUniformity() throws IOException { - NettyTransportClient client = (NettyTransportClient) this.oneClient(); byte[] sourceBytes1 = StringEncodeUtil.encode("test data message"); byte[] sourceBytes2 = StringEncodeUtil.encode("test data edge"); byte[] sourceBytes3 = StringEncodeUtil.encode("test data vertex"); @@ -165,6 +164,7 @@ public void testDataUniformity() throws IOException { return null; }).when(serverHandler).handle(Mockito.any(), Mockito.eq(1), Mockito.any()); + NettyTransportClient client = (NettyTransportClient) this.oneClient(); client.startSession(); client.send(MessageType.MSG, 1, ByteBuffer.wrap(sourceBytes1)); client.send(MessageType.EDGE, 1, ByteBuffer.wrap(sourceBytes2)); @@ -274,12 +274,12 @@ public void testFlowControl() throws IOException { @Test public void testHandlerException() throws IOException { - NettyTransportClient client = (NettyTransportClient) this.oneClient(); - client.startSession(); - Mockito.doThrow(new RuntimeException("test exception")).when(serverHandler) .handle(Mockito.any(), Mockito.anyInt(), Mockito.any()); + NettyTransportClient client = (NettyTransportClient) this.oneClient(); + client.startSession(); + ByteBuffer buffer = ByteBuffer.wrap(StringEncodeUtil.encode("test data")); boolean send = client.send(MessageType.MSG, 1, buffer); Assert.assertTrue(send); diff --git a/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSenderTest.java b/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSenderTest.java index 07fd15929..d2bda26fa 100644 --- a/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSenderTest.java +++ b/computer/computer-test/src/main/java/org/apache/hugegraph/computer/core/sender/QueuedMessageSenderTest.java @@ -17,8 +17,19 @@ package org.apache.hugegraph.computer.core.sender; +import java.net.InetSocketAddress; +import java.nio.ByteBuffer; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +import org.apache.hugegraph.computer.core.common.exception.TransportException; import org.apache.hugegraph.computer.core.config.ComputerOptions; import org.apache.hugegraph.computer.core.config.Config; +import org.apache.hugegraph.computer.core.network.TransportClient; +import org.apache.hugegraph.computer.core.network.message.MessageType; import org.apache.hugegraph.computer.core.worker.MockComputation2; import org.apache.hugegraph.computer.suite.unit.UnitTestBase; import org.apache.hugegraph.testutil.Assert; @@ -47,12 +58,18 @@ public void setup() { ); } - @Test - public void testInitAndClose() { + private QueuedMessageSender newSender(TransportClient first, TransportClient second) { QueuedMessageSender sender = new QueuedMessageSender(this.config); - sender.addWorkerClient(1, new MockTransportClient()); - sender.addWorkerClient(2, new MockTransportClient()); + sender.addWorkerClient(1, first); + sender.addWorkerClient(2, second); sender.init(); + return sender; + } + + @Test + public void testInitAndClose() { + QueuedMessageSender sender = this.newSender(new MockTransportClient(), + new MockTransportClient()); Thread sendExecutor = Whitebox.getInternalState(sender, "sendExecutor"); Assert.assertTrue(ImmutableSet.of(Thread.State.NEW, @@ -64,4 +81,390 @@ public void testInitAndClose() { Assert.assertTrue(ImmutableSet.of(Thread.State.TERMINATED) .contains(sendExecutor.getState())); } + + @Test + public void testControlBeforeCompletionFinishes() throws Exception { + ControlFutureClient client = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(client, new MockTransportClient()); + + CountDownLatch completionStarted = new CountDownLatch(1); + CountDownLatch allowCompletion = new CountDownLatch(1); + Thread completionThread = null; + try { + CompletableFuture startFuture = sender.send(1, MessageType.START); + Assert.assertTrue(await(client.startCalled)); + startFuture.whenComplete((r, e) -> { + completionStarted.countDown(); + try { + allowCompletion.await(); + } catch (InterruptedException exception) { + Thread.currentThread().interrupt(); + throw new AssertionError(exception); + } + }); + + completionThread = new Thread( + () -> client.startFuture.complete(null)); + completionThread.start(); + Assert.assertTrue(completionStarted.await(1, TimeUnit.SECONDS)); + + CompletableFuture finishFuture = sender.send(1, MessageType.FINISH); + Assert.assertTrue(await(client.finishCalled)); + allowCompletion.countDown(); + completionThread.join(TimeUnit.SECONDS.toMillis(1)); + Assert.assertFalse(completionThread.isAlive()); + client.finishFuture.complete(null); + finishFuture.get(1, TimeUnit.SECONDS); + } finally { + allowCompletion.countDown(); + if (completionThread != null) { + completionThread.join(TimeUnit.SECONDS.toMillis(1)); + } + sender.close(); + } + } + + @Test + public void testTransportExceptionControlFuture() throws Exception { + ControlFutureClient client = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(client, new MockTransportClient()); + + try { + CompletableFuture startFuture = sender.send(1, MessageType.START); + Assert.assertTrue(await(client.startCalled)); + + TransportException cause = new TransportException("connection failed"); + sender.transportExceptionCaught(cause, client.connectionId()); + assertFutureFailedWith(startFuture, cause); + + CompletableFuture finishFuture = sender.send(1, MessageType.FINISH); + Assert.assertTrue(await(client.finishCalled)); + client.startFuture.complete(null); + assertFutureFailedWith(startFuture, cause); + + client.finishFuture.complete(null); + finishFuture.get(1, TimeUnit.SECONDS); + } finally { + sender.close(); + } + } + + @Test + public void testTransportExceptionDispatch() throws Exception { + ControlFutureClient client = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(client, new MockTransportClient()); + + client.blockDataSend = true; + try { + sender.send(1, new QueuedMessage(0, MessageType.MSG, ByteBuffer.allocate(1))); + Assert.assertTrue(await(client.dataSendCalled)); + + CompletableFuture startFuture = sender.send(1, MessageType.START); + TransportException cause = new TransportException("connection failed before start"); + sender.transportExceptionCaught(cause, client.connectionId()); + assertFutureFailedWith(startFuture, cause); + + client.allowDataSend.countDown(); + Assert.assertFalse(await(client.startCalled)); + } finally { + client.allowDataSend.countDown(); + sender.close(); + } + } + + @Test + public void testExecutorAlive() throws Exception { + ControlFutureClient client = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(client, new MockTransportClient()); + + try { + RuntimeException startCause = new IllegalArgumentException("start session failed"); + client.startFailure = startCause; + CompletableFuture startFuture = sender.send(1, MessageType.START); + assertFutureFailedWith(startFuture, startCause); + + RuntimeException finishCause = new IllegalArgumentException("finish session failed"); + client.finishFailure = finishCause; + CompletableFuture finishFuture = sender.send(1, MessageType.FINISH); + assertFutureFailedWith(finishFuture, finishCause); + + Thread sendExecutor = Whitebox.getInternalState(sender, "sendExecutor"); + sendExecutor.join(TimeUnit.SECONDS.toMillis(1)); + Assert.assertTrue(sendExecutor.isAlive()); + + client.startFailure = null; + CompletableFuture nextStartFuture = sender.send(1, MessageType.START); + Assert.assertTrue(await(client.startCalled)); + client.startFuture.complete(null); + nextStartFuture.get(1, TimeUnit.SECONDS); + } finally { + sender.close(); + } + } + + @Test + public void testOtherClients() throws Exception { + ControlFutureClient failedClient = new ControlFutureClient(); + ControlFutureClient activeClient = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(failedClient, activeClient); + + try { + TransportException startCause = + new TransportException("start session failed"); + failedClient.startFailure = startCause; + CompletableFuture failedStart = sender.send(1, MessageType.START); + assertFutureFailedWith(failedStart, startCause); + + CompletableFuture activeStart = sender.send(2, MessageType.START); + Assert.assertTrue(await(activeClient.startCalled)); + activeClient.startFuture.complete(null); + activeStart.get(1, TimeUnit.SECONDS); + + TransportException finishCause = new TransportException("finish session failed"); + failedClient.finishFailure = finishCause; + CompletableFuture failedFinish = sender.send(1, MessageType.FINISH); + assertFutureFailedWith(failedFinish, finishCause); + + CompletableFuture activeFinish = sender.send(2, MessageType.FINISH); + Assert.assertTrue(await(activeClient.finishCalled)); + activeClient.finishFuture.complete(null); + activeFinish.get(1, TimeUnit.SECONDS); + + Thread sendExecutor = Whitebox.getInternalState(sender, "sendExecutor"); + Assert.assertTrue(sendExecutor.isAlive()); + } finally { + sender.close(); + } + } + + @Test + public void testQueuedFinish() throws Exception { + this.assertSynchronousDataFailureCompletesQueuedFinish( + new TransportException("data send failed")); + } + + @Test + public void testCompletesQueuedFinish() throws Exception { + this.assertSynchronousDataFailureCompletesQueuedFinish( + new IllegalStateException("data send failed")); + } + + @Test + public void testFinishFailsFinish() throws Exception { + this.assertSynchronousDataFailureBeforeFinishFailsFinish( + new TransportException("data send failed before finish")); + } + + @Test + public void testDataRuntimeFinish() throws Exception { + this.assertSynchronousDataFailureBeforeFinishFailsFinish( + new IllegalStateException("data send failed before finish")); + } + + @Test + public void testConflictKeepsSendExecutorAlive() throws Exception { + ControlFutureClient client = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(client, new MockTransportClient()); + + try { + CompletableFuture startFuture = sender.send(1, MessageType.START); + Assert.assertTrue(await(client.startCalled)); + + CompletableFuture conflictingFinishFuture = sender.send(1, MessageType.FINISH); + assertFutureFailedWithMessage(conflictingFinishFuture, "The origin future must be null"); + + Thread sendExecutor = Whitebox.getInternalState(sender, "sendExecutor"); + sendExecutor.join(TimeUnit.SECONDS.toMillis(1)); + Assert.assertTrue(sendExecutor.isAlive()); + + client.startFuture.complete(null); + startFuture.get(1, TimeUnit.SECONDS); + + CompletableFuture finishFuture = sender.send(1, MessageType.FINISH); + Assert.assertTrue(await(client.finishCalled)); + client.finishFuture.complete(null); + finishFuture.get(1, TimeUnit.SECONDS); + } finally { + sender.close(); + } + } + + private static void assertFutureFailedWith(CompletableFuture future, Throwable cause) + throws InterruptedException, TimeoutException { + try { + future.get(1, TimeUnit.SECONDS); + Assert.fail("Expected control future to fail"); + } catch (ExecutionException exception) { + Assert.assertSame(cause, exception.getCause()); + } + } + + private static boolean await(CountDownLatch latch) throws InterruptedException { + return latch.await(1, TimeUnit.SECONDS); + } + + private void assertSynchronousDataFailureCompletesQueuedFinish( + Throwable cause) throws Exception { + ControlFutureClient failedClient = new ControlFutureClient(); + ControlFutureClient activeClient = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender(failedClient, activeClient); + + failedClient.blockDataSend = true; + try { + sender.send(1, new QueuedMessage(0, MessageType.MSG, + ByteBuffer.allocate(1))); + Assert.assertTrue(await(failedClient.dataSendCalled)); + + CompletableFuture finishFuture = sender.send(1, MessageType.FINISH); + failedClient.dataFailure = cause; + failedClient.allowDataSend.countDown(); + assertFutureFailedWith(finishFuture, cause); + + CompletableFuture activeStart = sender.send(2, MessageType.START); + Assert.assertTrue(await(activeClient.startCalled)); + activeClient.startFuture.complete(null); + activeStart.get(1, TimeUnit.SECONDS); + } finally { + failedClient.allowDataSend.countDown(); + sender.close(); + } + } + + @Test + public void testTransportExceptionDuringDataSendFailsLaterFinish() + throws Exception { + ControlFutureClient client = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender( + client, new MockTransportClient()); + + client.blockDataSend = true; + try { + sender.send(1, new QueuedMessage(0, MessageType.MSG, + ByteBuffer.allocate(1))); + Assert.assertTrue(await(client.dataSendCalled)); + + TransportException cause = + new TransportException("connection failed during data"); + sender.transportExceptionCaught(cause, client.connectionId()); + client.allowDataSend.countDown(); + waitForQueueEmpty(sender, 1); + + CompletableFuture finishFuture = sender.send( + 1, MessageType.FINISH); + assertFutureFailedWith(finishFuture, cause); + Assert.assertFalse(await(client.finishCalled)); + } finally { + client.allowDataSend.countDown(); + sender.close(); + } + } + + private void assertSynchronousDataFailureBeforeFinishFailsFinish( + Throwable cause) throws Exception { + ControlFutureClient failedClient = new ControlFutureClient(); + QueuedMessageSender sender = this.newSender( + failedClient, new MockTransportClient()); + + failedClient.dataFailure = cause; + try { + sender.send(1, new QueuedMessage(0, MessageType.MSG, + ByteBuffer.allocate(1))); + waitForQueueEmpty(sender, 1); + + CompletableFuture finishFuture = sender.send(1, MessageType.FINISH); + assertFutureFailedWith(finishFuture, cause); + Assert.assertFalse(await(failedClient.finishCalled)); + } finally { + sender.close(); + } + } + + private static void waitForQueueEmpty(QueuedMessageSender sender, + int workerId) + throws InterruptedException { + Object[] channels = Whitebox.getInternalState(sender, "channels"); + MessageQueue queue = Whitebox.getInternalState(channels[workerId - 1], + "queue"); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1L); + while (queue.peek() != null && System.nanoTime() < deadline) { + Thread.sleep(10L); + } + Assert.assertTrue("Timed out to wait for sender queue to be empty", + queue.peek() == null); + } + + private static void assertFutureFailedWithMessage(CompletableFuture future, + String message) + throws InterruptedException, TimeoutException { + try { + future.get(1, TimeUnit.SECONDS); + Assert.fail("Expected control future to fail"); + } catch (ExecutionException exception) { + Assert.assertContains(message, exception.getCause().getMessage()); + } + } + + private static class ControlFutureClient extends MockTransportClient { + + private final CountDownLatch startCalled = new CountDownLatch(1); + private final CountDownLatch finishCalled = new CountDownLatch(1); + private final CountDownLatch dataSendCalled = new CountDownLatch(1); + private final CountDownLatch allowDataSend = new CountDownLatch(1); + private final CompletableFuture startFuture = new CompletableFuture<>(); + private final CompletableFuture finishFuture = new CompletableFuture<>(); + private Throwable startFailure; + private Throwable finishFailure; + private Throwable dataFailure; + private boolean blockDataSend; + + @Override + public CompletableFuture startSessionAsync() throws TransportException { + throwFailure(this.startFailure); + this.startCalled.countDown(); + return this.startFuture; + } + + @Override + public CompletableFuture finishSessionAsync() throws TransportException { + throwFailure(this.finishFailure); + this.finishCalled.countDown(); + return this.finishFuture; + } + + @Override + public boolean send(MessageType messageType, int partition, + ByteBuffer buffer) throws TransportException { + if (this.blockDataSend) { + this.dataSendCalled.countDown(); + try { + this.allowDataSend.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new TransportException("Interrupted data send", e); + } + } + throwFailure(this.dataFailure); + return true; + } + + private static void throwFailure(Throwable failure) throws TransportException { + if (failure instanceof TransportException) { + throw (TransportException) failure; + } + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + } + + @Override + public boolean sessionActive() { + return false; + } + + @Override + public InetSocketAddress remoteAddress() { + return new InetSocketAddress("127.0.0.1", 8080); + } + + } } diff --git a/computer/computer-test/src/main/java/org/apache/hugegraph/computer/suite/integrate/SenderIntegrateTest.java b/computer/computer-test/src/main/java/org/apache/hugegraph/computer/suite/integrate/SenderIntegrateTest.java index 6ae0008be..85d7b1a93 100644 --- a/computer/computer-test/src/main/java/org/apache/hugegraph/computer/suite/integrate/SenderIntegrateTest.java +++ b/computer/computer-test/src/main/java/org/apache/hugegraph/computer/suite/integrate/SenderIntegrateTest.java @@ -17,18 +17,22 @@ package org.apache.hugegraph.computer.suite.integrate; -import java.io.IOException; import java.util.ArrayList; -import java.util.HashMap; +import java.util.Arrays; import java.util.List; -import java.util.Map; -import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; -import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; import java.util.function.Function; +import java.util.function.Supplier; import org.apache.hugegraph.computer.algorithm.centrality.pagerank.PageRankParams; +import org.apache.hugegraph.computer.core.common.exception.ComputerException; import org.apache.hugegraph.computer.core.common.exception.TransportException; import org.apache.hugegraph.computer.core.config.ComputerOptions; import org.apache.hugegraph.computer.core.config.Config; @@ -46,17 +50,20 @@ import org.apache.hugegraph.testutil.Whitebox; import org.apache.hugegraph.util.Log; import org.junit.AfterClass; +import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; import org.slf4j.Logger; -import com.google.common.collect.Sets; - public class SenderIntegrateTest { public static final Logger LOG = Log.logger(SenderIntegrateTest.class); private static final Class COMPUTATION = MockComputation.class; + private static final long BSP_WAIT_TIMEOUT = TimeUnit.MINUTES.toMillis(5L); + private static final long SERVICE_WAIT_TIMEOUT = + BSP_WAIT_TIMEOUT + TimeUnit.SECONDS.toMillis(10L); + private static final long TEST_THREAD_JOIN_TIMEOUT = TimeUnit.SECONDS.toMillis(5L); @BeforeClass public static void init() { @@ -69,227 +76,305 @@ public static void clear() { } @Test - public void testOneWorker() { - AtomicReference masterServiceRef = new AtomicReference<>(); - AtomicReference workerServiceRef = new AtomicReference<>(); - CompletableFuture masterFuture = new CompletableFuture<>(); - Thread masterThread = new Thread(() -> { - String[] args = OptionsBuilder.newInstance() - .withJobId("local_002") - .withAlgorithm(PageRankParams.class) - .withResultName("rank") - .withResultClass(DoubleValue.class) - .withMessageClass(DoubleValue.class) - .withMaxSuperStep(3) - .withComputationClass(COMPUTATION) - .withWorkerCount(1) - .withBufferThreshold(50) - .withBufferCapacity(60) - .withRpcServerHost("127.0.0.1") - .withRpcServerPort(8611) - .withRpcServerPort(0) - .build(); - try (MasterService service = initMaster(args)) { - masterServiceRef.set(service); - service.execute(); - masterFuture.complete(null); - } catch (Exception e) { - LOG.error("Failed to execute master service", e); - masterFuture.completeExceptionally(e); - } - }); - masterThread.setDaemon(true); + public void testWaitForServicesFailsFast() { + CompletableFuture failedWorker = new CompletableFuture<>(); + CompletableFuture waitingMaster = new CompletableFuture<>(); + IllegalStateException cause = new IllegalStateException("worker failed"); + failedWorker.completeExceptionally(cause); - CompletableFuture workerFuture = new CompletableFuture<>(); - Thread workerThread = new Thread(() -> { - String[] args = OptionsBuilder.newInstance() - .withJobId("local_002") - .withAlgorithm(PageRankParams.class) - .withResultName("rank") - .withResultClass(DoubleValue.class) - .withMessageClass(DoubleValue.class) - .withMaxSuperStep(3) - .withComputationClass(COMPUTATION) - .withWorkerCount(1) - .withBufferThreshold(50) - .withBufferCapacity(60) - .withTransoprtServerPort(0) - .build(); - try (WorkerService service = initWorker(args)) { - workerServiceRef.set(service); - service.execute(); - workerFuture.complete(null); - } catch (Throwable e) { - LOG.error("Failed to execute worker service", e); - workerFuture.completeExceptionally(e); - } + try { + waitForServices(Arrays.asList(failedWorker, waitingMaster)); + Assert.fail("Expected worker failure to stop service wait"); + } catch (ComputerException e) { + Assert.assertSame(cause, e.getCause()); + } + } + + @Test + public void testCleanupFailureDoesNotHideWaitFailure() { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + CompletableFuture failedWorker = new CompletableFuture<>(); + IllegalStateException waitFailure = + new IllegalStateException("worker failed"); + IllegalStateException cleanupFailure = + new IllegalStateException("worker close failed"); + failedWorker.completeExceptionally(waitFailure); + lifecycle.registerWorker(() -> { + throw cleanupFailure; }); - workerThread.setDaemon(true); - masterThread.start(); - workerThread.start(); try { - CompletableFuture.allOf(workerFuture, masterFuture).join(); + waitForServicesAndClose(lifecycle, Arrays.asList(failedWorker), + new ArrayList<>(), null); + Assert.fail("Expected worker failure to be preserved"); + } catch (ComputerException e) { + Assert.assertSame(waitFailure, e.getCause()); + Assert.assertEquals(1, e.getSuppressed().length); + Assert.assertSame(cleanupFailure, e.getSuppressed()[0].getCause()); + } + } + + @Test + public void testMasterErrorCompletesServiceFuture() throws Exception { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + Error cause = new AssertionError("master failed"); + ServiceTask master = newServiceTask( + Object::new, lifecycle::registerMaster, + service -> { + throw cause; + }, service -> { + }); + master.thread.start(); + + try { + master.future.get(1, TimeUnit.SECONDS); + Assert.fail("Expected master error to fail the service future"); + } catch (ExecutionException e) { + Assert.assertSame(cause, e.getCause()); + } catch (TimeoutException e) { + Assert.fail("Timed out to wait for master error"); } finally { - workerServiceRef.get().close(); - masterServiceRef.get().close(); + lifecycle.closeAll(); + interruptAndJoinThreads(Arrays.asList(master.thread), + TEST_THREAD_JOIN_TIMEOUT); } } @Test - public void testMultiWorkers() throws IOException { - int workerCount = 3; - int partitionCount = 3; - AtomicReference masterServiceRef = new AtomicReference<>(); - Set workerServices = Sets.newConcurrentHashSet(); - CompletableFuture masterFuture = new CompletableFuture<>(); - Thread masterThread = new Thread(() -> { - String[] args = OptionsBuilder.newInstance() - .withJobId("local_003") - .withAlgorithm(PageRankParams.class) - .withResultName("rank") - .withResultClass(DoubleValue.class) - .withMessageClass(DoubleValue.class) - .withMaxSuperStep(3) - .withComputationClass(COMPUTATION) - .withWorkerCount(workerCount) - .withPartitionCount(partitionCount) - .withRpcServerHost("127.0.0.1") - .withRpcServerPort(0) - .build(); + public void testServiceLifecycleClosesLateRegisteredService() { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + AtomicBoolean closed = new AtomicBoolean(); + + lifecycle.closeAll(); + + Assert.assertFalse(lifecycle.registerWorker(() -> closed.set(true))); + Assert.assertTrue(closed.get()); + } + + @Test + public void testServiceLifecycleClosesWorkersBeforeMasterAfterFailure() { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + AtomicBoolean activeWorkerClosed = new AtomicBoolean(); + AtomicBoolean masterClosed = new AtomicBoolean(); + RuntimeException workerFailure = + new IllegalStateException("worker close failed"); + + lifecycle.registerMaster(() -> { + Assert.assertTrue(activeWorkerClosed.get()); + masterClosed.set(true); + }); + lifecycle.registerWorker(() -> { + throw workerFailure; + }); + lifecycle.registerWorker(() -> activeWorkerClosed.set(true)); + + Throwable failure = lifecycle.closeAll(); + + Assert.assertSame(workerFailure, failure); + Assert.assertTrue(activeWorkerClosed.get()); + Assert.assertTrue(masterClosed.get()); + } + + @Test + public void testInitializeServiceClosesPartiallyInitializedService() { + AtomicBoolean closed = new AtomicBoolean(); + RuntimeException cause = new IllegalStateException("init failed"); + + try { + initializeService(new Object(), service -> { + throw cause; + }, service -> closed.set(true)); + Assert.fail("Expected initialization to fail"); + } catch (RuntimeException e) { + Assert.assertSame(cause, e); + } + + Assert.assertTrue(closed.get()); + } + + @Test + public void testCloseServicesAndJoinStopsSpawnedThreads() throws Exception { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + AtomicBoolean closed = new AtomicBoolean(); + CountDownLatch started = new CountDownLatch(1); + Thread thread = new Thread(() -> { + started.countDown(); try { - MasterService service = initMaster(args); - masterServiceRef.set(service); - service.execute(); - masterFuture.complete(null); - } catch (Throwable e) { - LOG.error("Failed to execute master service", e); - masterFuture.completeExceptionally(e); + Thread.sleep(Long.MAX_VALUE); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); } }); - masterThread.setDaemon(true); + lifecycle.registerWorker(() -> closed.set(true)); + try { + thread.start(); + Assert.assertTrue(started.await(1, TimeUnit.SECONDS)); - Map> workers = new HashMap<>(workerCount); - for (int i = 1; i <= workerCount; i++) { - String dir = "[jobs-" + i + "]"; - - CompletableFuture workerFuture = new CompletableFuture<>(); - Thread thread = new Thread(() -> { - String[] args; - args = OptionsBuilder.newInstance() - .withJobId("local_003") - .withAlgorithm(PageRankParams.class) - .withResultName("rank") - .withResultClass(DoubleValue.class) - .withMessageClass(DoubleValue.class) - .withMaxSuperStep(3) - .withComputationClass(COMPUTATION) - .withWorkerCount(workerCount) - .withPartitionCount(partitionCount) - .withTransoprtServerPort(0) - .withDataDirs(dir) - .build(); - try { - WorkerService service = initWorker(args); - workerServices.add(service); - service.execute(); - workerFuture.complete(null); - } catch (Throwable e) { - LOG.error("Failed to execute worker service", e); - workerFuture.completeExceptionally(e); - } - }); - thread.setDaemon(true); - workers.put(thread, workerFuture); - } + closeServicesAndJoin(lifecycle, Arrays.asList(thread), null); - masterThread.start(); - for (Thread worker : workers.keySet()) { - worker.start(); + Assert.assertTrue(closed.get()); + Assert.assertFalse(thread.isAlive()); + } finally { + lifecycle.closeAll(); + interruptAndJoinThreads(Arrays.asList(thread), + TEST_THREAD_JOIN_TIMEOUT); } + } - List> futures = new ArrayList<>(workers.values()); - futures.add(masterFuture); + @Test + public void testCloseServicesAndJoinUsesOneTimeoutBudget() throws Exception { + long timeout = TimeUnit.SECONDS.toMillis(1L); + long maxElapsed = timeout + timeout / 2L; + AtomicBoolean stopping = new AtomicBoolean(); + CountDownLatch started = new CountDownLatch(3); + Thread firstWorker = newInterruptIgnoringThread(stopping, started); + Thread secondWorker = newInterruptIgnoringThread(stopping, started); + Thread master = newInterruptIgnoringThread(stopping, started); + firstWorker.start(); + secondWorker.start(); + master.start(); + Assert.assertTrue(started.await(1, TimeUnit.SECONDS)); + + long start = System.nanoTime(); + long elapsed = -1L; try { - CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); + closeServicesAndJoin(new ServiceLifecycle(), + Arrays.asList(firstWorker, secondWorker), + master, timeout); + Assert.fail("Expected service threads to time out"); + } catch (ComputerException ignored) { + elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start); } finally { - for (WorkerService workerService : workerServices) { - workerService.close(); - } - masterServiceRef.get().close(); + stopping.set(true); + firstWorker.interrupt(); + secondWorker.interrupt(); + master.interrupt(); + firstWorker.join(TimeUnit.SECONDS.toMillis(1L)); + secondWorker.join(TimeUnit.SECONDS.toMillis(1L)); + master.join(TimeUnit.SECONDS.toMillis(1L)); } + Assert.assertTrue("Cleanup exceeded its shared timeout budget: " + elapsed, + elapsed < maxElapsed); } @Test - public void testOneWorkerWithBusyClient() { - AtomicReference masterServiceRef = new AtomicReference<>(); - AtomicReference workerServiceRef = new AtomicReference<>(); - CompletableFuture masterFuture = new CompletableFuture<>(); - Thread masterThread = new Thread(() -> { - String[] args = OptionsBuilder.newInstance() - .withJobId("local_002") - .withAlgorithm(PageRankParams.class) - .withResultName("rank") - .withResultClass(DoubleValue.class) - .withMessageClass(DoubleValue.class) - .withMaxSuperStep(3) - .withComputationClass(COMPUTATION) - .withWorkerCount(1) - .withWriteBufferHighMark(10) - .withWriteBufferLowMark(5) - .withRpcServerHost("127.0.0.1") - .withRpcServerPort(0) - .build(); - try (MasterService service = initMaster(args)) { - masterServiceRef.set(service); - service.execute(); - masterFuture.complete(null); - } catch (Throwable e) { - LOG.error("Failed to execute master service", e); - masterFuture.completeExceptionally(e); + public void testCloseBeforeMaster() throws Exception { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + CountDownLatch workerStarted = new CountDownLatch(1); + CountDownLatch workerStopped = new CountDownLatch(1); + CountDownLatch masterStarted = new CountDownLatch(1); + AtomicBoolean masterInterruptedBeforeWorkerStopped = new AtomicBoolean(); + Thread workerThread = new Thread(() -> { + workerStarted.countDown(); + try { + Thread.sleep(Long.MAX_VALUE); + } catch (InterruptedException e) { + try { + Thread.sleep(200L); + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } finally { + workerStopped.countDown(); } }); - masterThread.setDaemon(true); - - CompletableFuture workerFuture = new CompletableFuture<>(); - int transoprtServerPort = 8998; - Thread workerThread = new Thread(() -> { - String[] args = OptionsBuilder.newInstance() - .withJobId("local_002") - .withAlgorithm(PageRankParams.class) - .withResultName("rank") - .withResultClass(DoubleValue.class) - .withMessageClass(DoubleValue.class) - .withMaxSuperStep(3) - .withComputationClass(COMPUTATION) - .withWorkerCount(1) - .withWriteBufferHighMark(20) - .withWriteBufferLowMark(10) - .withTransoprtServerPort(transoprtServerPort) - .build(); - try (WorkerService service = initWorker(args)) { - workerServiceRef.set(service); - // Let send rate slowly - this.slowSendFunc(service, transoprtServerPort); - service.execute(); - workerFuture.complete(null); - } catch (Throwable e) { - LOG.error("Failed to execute worker service", e); - workerFuture.completeExceptionally(e); + Thread masterThread = new Thread(() -> { + masterStarted.countDown(); + try { + workerStopped.await(); + } catch (InterruptedException e) { + masterInterruptedBeforeWorkerStopped.set( + workerStopped.getCount() != 0L); + Thread.currentThread().interrupt(); } }); - workerThread.setDaemon(true); - masterThread.start(); - workerThread.start(); - try { - CompletableFuture.allOf(workerFuture, masterFuture).join(); + workerThread.start(); + masterThread.start(); + Assert.assertTrue(workerStarted.await(1, TimeUnit.SECONDS)); + Assert.assertTrue(masterStarted.await(1, TimeUnit.SECONDS)); + + closeServicesAndJoin(lifecycle, Arrays.asList(workerThread), + masterThread); + + Assert.assertFalse(masterInterruptedBeforeWorkerStopped.get()); + Assert.assertFalse(workerThread.isAlive()); + Assert.assertFalse(masterThread.isAlive()); } finally { - workerServiceRef.get().close(); - masterServiceRef.get().close(); + lifecycle.closeAll(); + interruptAndJoinThreads(Arrays.asList(workerThread, masterThread), + TEST_THREAD_JOIN_TIMEOUT); } } + @Test + public void testOneWorker() { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + String[] masterArgs = commonOptions("local_002", 1) + .withBufferThreshold(50) + .withBufferCapacity(60) + .withRpcServerHost("127.0.0.1") + .withRpcServerPort(0) + .build(); + String[] workerArgs = commonOptions("local_002", 1) + .withBufferThreshold(50) + .withBufferCapacity(60) + .withTransoprtServerPort(0) + .build(); + ServiceTask master = masterTask(lifecycle, masterArgs); + ServiceTask worker = workerTask(lifecycle, workerArgs, + WorkerService::execute); + runServices(lifecycle, Arrays.asList(worker), master); + } + + @Test + public void testMultiWorkers() { + int workerCount = 3; + int partitionCount = 3; + ServiceLifecycle lifecycle = new ServiceLifecycle(); + String[] masterArgs = commonOptions("local_003", workerCount) + .withPartitionCount(partitionCount) + .withRpcServerHost("127.0.0.1") + .withRpcServerPort(0) + .build(); + ServiceTask master = masterTask(lifecycle, masterArgs); + List workers = new ArrayList<>(workerCount); + for (int i = 1; i <= workerCount; i++) { + String[] workerArgs = commonOptions("local_003", workerCount) + .withPartitionCount(partitionCount) + .withTransoprtServerPort(0) + .withDataDirs("[jobs-" + i + "]") + .build(); + workers.add(workerTask(lifecycle, workerArgs, + WorkerService::execute)); + } + runServices(lifecycle, workers, master); + } + + @Test + public void testOneWorkerWithBusyClient() { + ServiceLifecycle lifecycle = new ServiceLifecycle(); + int transoprtServerPort = 8998; + String[] masterArgs = commonOptions("local_002", 1) + .withWriteBufferHighMark(10) + .withWriteBufferLowMark(5) + .withRpcServerHost("127.0.0.1") + .withRpcServerPort(0) + .build(); + String[] workerArgs = commonOptions("local_002", 1) + .withWriteBufferHighMark(20) + .withWriteBufferLowMark(10) + .withTransoprtServerPort(transoprtServerPort) + .build(); + ServiceTask master = masterTask(lifecycle, masterArgs); + ServiceTask worker = workerTask(lifecycle, workerArgs, service -> { + // Let send rate slowly + this.slowSendFunc(service, transoprtServerPort); + service.execute(); + }); + runServices(lifecycle, Arrays.asList(worker), master); + } + private void slowSendFunc(WorkerService service, int port) throws TransportException { Managers managers = Whitebox.getInternalState(service, "managers"); DataClientManager clientManager = managers.get( @@ -315,20 +400,319 @@ private void slowSendFunc(WorkerService service, int port) throws TransportExcep Whitebox.setInternalState(clientSession, "sendFunction", sendFunc); } + private static OptionsBuilder commonOptions(String jobId, + int workerCount) { + return OptionsBuilder.newInstance() + .withJobId(jobId) + .withAlgorithm(PageRankParams.class) + .withResultName("rank") + .withResultClass(DoubleValue.class) + .withMessageClass(DoubleValue.class) + .withMaxSuperStep(3) + .withComputationClass(COMPUTATION) + .withWorkerCount(workerCount) + .withTestBspTimeouts(); + } + + private ServiceTask masterTask(ServiceLifecycle lifecycle, String[] args) { + return newServiceTask(() -> this.initMaster(args), + lifecycle::registerMaster, + MasterService::execute, + MasterService::close); + } + + private ServiceTask workerTask(ServiceLifecycle lifecycle, String[] args, + ServiceExecutor executor) { + return newServiceTask(() -> this.initWorker(args), + lifecycle::registerWorker, executor, + WorkerService::close); + } + + private static ServiceTask newServiceTask( + Supplier initializer, Function registrar, + ServiceExecutor executor, Consumer closer) { + CompletableFuture future = new CompletableFuture<>(); + Thread thread = new Thread(() -> { + try { + T service = initializer.get(); + try { + if (!registrar.apply(() -> closer.accept(service))) { + future.cancel(false); + return; + } + executor.execute(service); + future.complete(null); + } finally { + closer.accept(service); + } + } catch (Throwable e) { + LOG.error("Failed to execute service", e); + future.completeExceptionally(e); + } + }); + thread.setDaemon(true); + return new ServiceTask(thread, future); + } + + private static void runServices(ServiceLifecycle lifecycle, + List workers, + ServiceTask master) { + master.thread.start(); + List> futures = + new ArrayList<>(workers.size() + 1); + List workerThreads = new ArrayList<>(workers.size()); + for (ServiceTask worker : workers) { + worker.thread.start(); + futures.add(worker.future); + workerThreads.add(worker.thread); + } + futures.add(master.future); + waitForServicesAndClose(lifecycle, futures, workerThreads, + master.thread); + } + private MasterService initMaster(String[] args) { Config config = ComputerContextUtil.initContext( ComputerContextUtil.convertToMap(args)); MasterService service = new MasterService(); - service.init(config); - return service; + return initializeService(service, s -> s.init(config), + MasterService::close); } private WorkerService initWorker(String[] args) { Config config = ComputerContextUtil.initContext( ComputerContextUtil.convertToMap(args)); WorkerService service = new WorkerService(); - service.init(config); - return service; + return initializeService(service, s -> s.init(config), + WorkerService::close); + } + + private static T initializeService(T service, Consumer initializer, + Consumer closer) { + try { + initializer.accept(service); + return service; + } catch (RuntimeException | Error e) { + try { + closer.accept(service); + } catch (RuntimeException | Error closeFailure) { + e.addSuppressed(closeFailure); + } + throw e; + } + } + + private static void waitForServices(List> futures) { + CompletableFuture result = new CompletableFuture<>(); + for (CompletableFuture future : futures) { + future.whenComplete((r, e) -> { + if (e != null) { + result.completeExceptionally(e); + } + }); + } + CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])) + .thenRun(() -> result.complete(null)); + try { + result.get(SERVICE_WAIT_TIMEOUT, TimeUnit.MILLISECONDS); + } catch (TimeoutException e) { + throw new ComputerException("Timed out to wait for master and " + + "worker services", e); + } catch (ExecutionException e) { + throw new ComputerException("Failed to wait for master and " + + "worker services", e.getCause()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new ComputerException("Interrupted when waiting for master " + + "and worker services", e); + } + } + + private static void waitForServicesAndClose( + ServiceLifecycle lifecycle, List> futures, + List workerThreads, Thread masterThread) { + try { + waitForServices(futures); + } catch (RuntimeException | Error e) { + try { + closeServicesAndJoin(lifecycle, workerThreads, masterThread); + } catch (RuntimeException | Error closeFailure) { + e.addSuppressed(closeFailure); + } + throw e; + } + closeServicesAndJoin(lifecycle, workerThreads, masterThread); + } + + private static void closeServicesAndJoin(ServiceLifecycle lifecycle, + List workerThreads, + Thread masterThread) { + closeServicesAndJoin(lifecycle, workerThreads, masterThread, + SERVICE_WAIT_TIMEOUT); + } + + private static void closeServicesAndJoin(ServiceLifecycle lifecycle, + List workerThreads, + Thread masterThread, + long timeout) { + long deadline = System.nanoTime() + + TimeUnit.MILLISECONDS.toNanos(timeout); + Throwable closeFailure = lifecycle.closeAll(); + Throwable workerFailure = interruptAndJoinThreads( + workerThreads, + remainingTimeout(deadline)); + Throwable masterFailure = null; + if (masterThread != null) { + masterFailure = interruptAndJoinThreads( + Arrays.asList(masterThread), + remainingTimeout(deadline)); + } + if (closeFailure != null) { + addFailure(closeFailure, workerFailure); + addFailure(closeFailure, masterFailure); + throw new ComputerException("Failed to close service", closeFailure); + } + if (workerFailure != null) { + addFailure(workerFailure, masterFailure); + throw new ComputerException("Failed to close worker service thread", + workerFailure); + } + if (masterFailure != null) { + throw new ComputerException("Failed to close master service thread", + masterFailure); + } + } + + private static Throwable interruptAndJoinThreads(List threads, + long timeout) { + long deadline = System.nanoTime() + + TimeUnit.MILLISECONDS.toNanos(timeout); + for (Thread thread : threads) { + thread.interrupt(); + } + Throwable failure = null; + for (Thread thread : threads) { + long remaining = remainingTimeout(deadline); + if (remaining > 0L) { + try { + thread.join(remaining); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + failure = addFailure(failure, new ComputerException( + "Interrupted when waiting for service " + + "thread to stop", e)); + } + } + if (thread.isAlive()) { + failure = addFailure(failure, new ComputerException( + "Timed out to wait for service thread " + + "to stop")); + } + } + return failure; + } + + private static Thread newInterruptIgnoringThread(AtomicBoolean stopping, + CountDownLatch started) { + Thread thread = new Thread(() -> { + started.countDown(); + while (!stopping.get()) { + try { + Thread.sleep(Long.MAX_VALUE); + } catch (InterruptedException ignored) { + // Keep waiting until the test explicitly stops this thread + } + } + }); + thread.setDaemon(true); + return thread; + } + + private static long remainingTimeout(long deadline) { + long remaining = deadline - System.nanoTime(); + return remaining <= 0L ? 0L : + TimeUnit.NANOSECONDS.toMillis(remaining) + 1L; + } + + private static Throwable addFailure(Throwable failure, Throwable cause) { + if (cause == null) { + return failure; + } + if (failure == null) { + return cause; + } + failure.addSuppressed(cause); + return failure; + } + + @FunctionalInterface + private interface ServiceExecutor { + + void execute(T service) throws Throwable; + } + + private static class ServiceTask { + + private final Thread thread; + private final CompletableFuture future; + + public ServiceTask(Thread thread, CompletableFuture future) { + this.thread = thread; + this.future = future; + } + } + + private static class ServiceLifecycle { + + private final List workerClosers = new ArrayList<>(); + private final List masterClosers = new ArrayList<>(); + private boolean closing; + + public boolean registerWorker(Runnable closer) { + return this.register(this.workerClosers, closer); + } + + public boolean registerMaster(Runnable closer) { + return this.register(this.masterClosers, closer); + } + + private boolean register(List closers, Runnable closer) { + synchronized (this) { + if (!this.closing) { + closers.add(closer); + return true; + } + } + closer.run(); + return false; + } + + public Throwable closeAll() { + List workerClosers; + List masterClosers; + synchronized (this) { + this.closing = true; + workerClosers = new ArrayList<>(this.workerClosers); + masterClosers = new ArrayList<>(this.masterClosers); + this.workerClosers.clear(); + this.masterClosers.clear(); + } + workerClosers.addAll(masterClosers); + Throwable failure = closeAll(workerClosers, null); + return failure; + } + + private static Throwable closeAll(List closers, + Throwable failure) { + for (Runnable closer : closers) { + try { + closer.run(); + } catch (Throwable e) { + failure = addFailure(failure, e); + } + } + return failure; + } } private static class OptionsBuilder { @@ -395,6 +779,14 @@ public OptionsBuilder withWorkerCount(int count) { return this; } + public OptionsBuilder withTestBspTimeouts() { + this.options.add(ComputerOptions.BSP_WAIT_WORKERS_TIMEOUT.name()); + this.options.add(String.valueOf(BSP_WAIT_TIMEOUT)); + this.options.add(ComputerOptions.BSP_WAIT_MASTER_TIMEOUT.name()); + this.options.add(String.valueOf(BSP_WAIT_TIMEOUT)); + return this; + } + public OptionsBuilder withPartitionCount(int count) { this.options.add(ComputerOptions.JOB_PARTITIONS_COUNT.name()); this.options.add(String.valueOf(count));