Discard messages instead of replying later

2 files changed, 7 insertions, 7 deletions

diff --git a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java
index 81eb237713f..0b2664a75ee 100644
--- a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/Destination.java
@@ -22,15 +22,13 @@ import com.yahoo.messagebus.Error;
 import com.yahoo.messagebus.ErrorCode;
 import com.yahoo.messagebus.Message;
 import com.yahoo.messagebus.MessageHandler;
-import com.yahoo.messagebus.Protocol;
 import com.yahoo.messagebus.RPCMessageBus;
 import com.yahoo.messagebus.Reply;
 import com.yahoo.messagebus.network.Identity;
 import com.yahoo.messagebus.network.rpc.RPCNetworkParams;
 
-import java.util.Arrays;
 import java.util.List;
-import java.util.concurrent.Phaser;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Mock-up destination used for testing.
@@ -39,7 +37,7 @@ import java.util.concurrent.Phaser;
  */
 public class Destination implements MessageHandler {
 
-    final Phaser phaser = new Phaser(1);
+    final AtomicBoolean discard = new AtomicBoolean();
 
     private final DestinationSession session;
     private final DocumentAccess access;
@@ -61,8 +59,11 @@ public class Destination implements MessageHandler {
     }
 
     public void handleMessage(Message msg) {
+        if (discard.get()) {
+            msg.discard();
+            return;
+        }
 
-        phaser.arriveAndAwaitAdvance();
         Reply reply = ((DocumentMessage)msg).createReply();
         try {
             switch (msg.getType()) {
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java
index 7ffae650666..d1fbbd74795 100644
--- a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/MessageBusDocumentApiTestCase.java
@@ -95,7 +95,7 @@ public class MessageBusDocumentApiTestCase extends AbstractDocumentApiTestCase {
         DocumentType type = access().getDocumentTypeManager().getDocumentType("music");
         Document doc1 = new Document(type, new DocumentId("id:ns:music::1"));
 
-        destination.phaser.register();
+        destination.discard.set(true);
         assertTrue(session.put(new DocumentPut(doc1),
                                DocumentOperationParameters.parameters()
                                                           .withResponseHandler(result -> {
@@ -106,7 +106,6 @@ public class MessageBusDocumentApiTestCase extends AbstractDocumentApiTestCase {
                           .isSuccess());
         assertTrue(latch.await(60, TimeUnit.SECONDS));
         assertEquals(Response.Outcome.TIMEOUT, response.get().outcome());
-        destination.phaser.arriveAndDeregister();
         session.destroy();
     }