Wait until other thread is waiting for new context to stabilize unit test

2 files changed, 30 insertions, 14 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManager.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManager.java
index 808a63eb130..016d60b1354 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManager.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManager.java
@@ -22,6 +22,7 @@ public class NodeAgentContextManager implements NodeAgentContextSupplier, NodeAg
     private boolean wantFrozen = false;
     private boolean isFrozen = true;
     private boolean pendingInterrupt = false;
+    private boolean isWaitingForNextContext = false;
 
     public NodeAgentContextManager(Clock clock, NodeAgentContext context) {
         this.clock = clock;
@@ -62,6 +63,7 @@ public class NodeAgentContextManager implements NodeAgentContextSupplier, NodeAg
     public NodeAgentContext nextContext() throws InterruptedException {
         synchronized (monitor) {
             nextContext = null; // Reset any previous context and wait for the next one
+            isWaitingForNextContext = true;
             Duration untilNextContext = Duration.ZERO;
             while (setAndGetIsFrozen(wantFrozen) ||
                     nextContext == null ||
@@ -76,6 +78,7 @@ public class NodeAgentContextManager implements NodeAgentContextSupplier, NodeAg
                 } catch (InterruptedException ignored) { }
             }
 
+            isWaitingForNextContext = false;
             currentContext = nextContext;
             return currentContext;
         }
@@ -105,4 +108,11 @@ public class NodeAgentContextManager implements NodeAgentContextSupplier, NodeAg
             return this.isFrozen;
         }
     }
+
+    /** FOR TESTING ONLY */
+    boolean isWaitingForNextContext() {
+        synchronized (monitor) {
+            return isWaitingForNextContext;
+        }
+    }
 }
\ No newline at end of file
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManagerTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManagerTest.java
index 195b07bd73f..240a29080cc 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManagerTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextManagerTest.java
@@ -27,13 +27,13 @@ public class NodeAgentContextManagerTest {
     private final NodeAgentContextManager manager = new NodeAgentContextManager(clock, initialContext);
 
     @Test(timeout = TIMEOUT)
-    public void context_is_ignored_unless_scheduled_while_waiting() throws InterruptedException {
+    public void context_is_ignored_unless_scheduled_while_waiting() {
         NodeAgentContext context1 = generateContext();
         manager.scheduleTickWith(context1, clock.instant());
         assertSame(initialContext, manager.currentContext());
 
         AsyncExecutor<NodeAgentContext> async = new AsyncExecutor<>(manager::nextContext);
-        Thread.sleep(20);
+        waitUntilWaitingForNextContext();
         assertFalse(async.isCompleted());
 
         NodeAgentContext context2 = generateContext();
@@ -44,9 +44,9 @@ public class NodeAgentContextManagerTest {
     }
 
     @Test(timeout = TIMEOUT)
-    public void returns_no_earlier_than_at_given_time() throws InterruptedException {
+    public void returns_no_earlier_than_at_given_time() {
         AsyncExecutor<NodeAgentContext> async = new AsyncExecutor<>(manager::nextContext);
-        Thread.sleep(20);
+        waitUntilWaitingForNextContext();
 
         NodeAgentContext context1 = generateContext();
         Instant returnAt = clock.instant().plusMillis(500);
@@ -59,10 +59,9 @@ public class NodeAgentContextManagerTest {
     }
 
     @Test(timeout = TIMEOUT)
-    public void blocks_in_nextContext_until_one_is_scheduled() throws InterruptedException {
+    public void blocks_in_nextContext_until_one_is_scheduled() {
         AsyncExecutor<NodeAgentContext> async = new AsyncExecutor<>(manager::nextContext);
-        assertFalse(async.isCompleted());
-        Thread.sleep(10);
+        waitUntilWaitingForNextContext();
         assertFalse(async.isCompleted());
 
         NodeAgentContext context1 = generateContext();
@@ -74,10 +73,9 @@ public class NodeAgentContextManagerTest {
     }
 
     @Test(timeout = TIMEOUT)
-    public void blocks_in_nextContext_until_interrupt() throws InterruptedException {
+    public void blocks_in_nextContext_until_interrupt() {
         AsyncExecutor<NodeAgentContext> async = new AsyncExecutor<>(manager::nextContext);
-        assertFalse(async.isCompleted());
-        Thread.sleep(10);
+        waitUntilWaitingForNextContext();
         assertFalse(async.isCompleted());
 
         manager.interrupt();
@@ -88,13 +86,13 @@ public class NodeAgentContextManagerTest {
     }
 
     @Test(timeout = TIMEOUT)
-    public void setFrozen_does_not_block_with_no_timeout() throws InterruptedException {
+    public void setFrozen_does_not_block_with_no_timeout() {
         assertFalse(manager.setFrozen(false, Duration.ZERO));
 
         // Generate new context and get it from the supplier, this completes the unfreeze
         NodeAgentContext context1 = generateContext();
         AsyncExecutor<NodeAgentContext> async = new AsyncExecutor<>(manager::nextContext);
-        Thread.sleep(20);
+        waitUntilWaitingForNextContext();
         manager.scheduleTickWith(context1, clock.instant());
         assertSame(context1, async.awaitResult().response.get());
 
@@ -118,7 +116,7 @@ public class NodeAgentContextManagerTest {
             Thread.sleep(200); // Simulate running NodeAgent::converge
             return context;
         });
-        Thread.sleep(20);
+        waitUntilWaitingForNextContext();
 
         NodeAgentContext context1 = generateContext();
         manager.scheduleTickWith(context1, clock.instant());
@@ -132,7 +130,7 @@ public class NodeAgentContextManagerTest {
         assertFalse(asyncScheduler.isCompleted()); // Still waiting for consumer to converge to frozen
 
         AsyncExecutor<NodeAgentContext> asyncConsumer2 = new AsyncExecutor<>(manager::nextContext);
-        Thread.sleep(20);
+        waitUntilWaitingForNextContext();
         assertFalse(asyncConsumer2.isCompleted()); // Waiting for next context
         assertTrue(asyncScheduler.isCompleted()); // While consumer is waiting, it has converged to frozen
 
@@ -143,6 +141,14 @@ public class NodeAgentContextManagerTest {
         assertEquals(Optional.of(true), asyncScheduler.response);
     }
 
+    private void waitUntilWaitingForNextContext() {
+        while (!manager.isWaitingForNextContext()) {
+            try {
+                Thread.sleep(10);
+            } catch (InterruptedException ignored) { }
+        }
+    }
+
     private static NodeAgentContext generateContext() {
         return new NodeAgentContextImpl.Builder("container-123.domain.tld").build();
     }