diff --git a/asterixdb/asterix-replication/src/main/java/org/apache/asterix/replication/messaging/MarkComponentValidTask.java b/asterixdb/asterix-replication/src/main/java/org/apache/asterix/replication/messaging/MarkComponentValidTask.java
index b581321..57474ef 100644
--- a/asterixdb/asterix-replication/src/main/java/org/apache/asterix/replication/messaging/MarkComponentValidTask.java
+++ b/asterixdb/asterix-replication/src/main/java/org/apache/asterix/replication/messaging/MarkComponentValidTask.java
@@ -29,11 +29,11 @@
 
 import org.apache.asterix.common.api.INcApplicationContext;
 import org.apache.asterix.common.exceptions.ReplicationException;
-import org.apache.asterix.replication.api.IReplicationWorker;
 import org.apache.asterix.common.storage.IIndexCheckpointManager;
 import org.apache.asterix.common.storage.IIndexCheckpointManagerProvider;
 import org.apache.asterix.common.storage.ResourceReference;
 import org.apache.asterix.replication.api.IReplicaTask;
+import org.apache.asterix.replication.api.IReplicationWorker;
 import org.apache.hyracks.api.exceptions.HyracksDataException;
 import org.apache.hyracks.storage.am.lsm.common.impls.AbstractLSMIndexFileManager;
 
@@ -71,13 +71,13 @@
         final IIndexCheckpointManagerProvider checkpointManagerProvider = appCtx.getIndexCheckpointManagerProvider();
         final IIndexCheckpointManager indexCheckpointManager = checkpointManagerProvider.get(indexRef);
         long replicationTimeOut = TimeUnit.SECONDS.toMillis(appCtx.getReplicationProperties().getReplicationTimeOut());
-        final long startTime = System.nanoTime();
         synchronized (indexCheckpointManager) {
             // wait until the lsn mapping is flushed to disk
             while (!indexCheckpointManager.isFlushed(masterLsn)) {
                 if (replicationTimeOut <= 0) {
                     throw new ReplicationException(new TimeoutException("Couldn't receive flush lsn from master"));
                 }
+                final long startTime = System.nanoTime();
                 indexCheckpointManager.wait(replicationTimeOut);
                 replicationTimeOut -= TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
             }
diff --git a/hyracks-fullstack/hyracks/hyracks-api/src/main/java/org/apache/hyracks/api/util/InvokeUtil.java b/hyracks-fullstack/hyracks/hyracks-api/src/main/java/org/apache/hyracks/api/util/InvokeUtil.java
index fb2bdeb..b2dd680 100644
--- a/hyracks-fullstack/hyracks/hyracks-api/src/main/java/org/apache/hyracks/api/util/InvokeUtil.java
+++ b/hyracks-fullstack/hyracks/hyracks-api/src/main/java/org/apache/hyracks/api/util/InvokeUtil.java
@@ -242,14 +242,13 @@
      */
     public static void runWithTimeout(ThrowingAction action, BooleanSupplier stopCondition, long timeout, TimeUnit unit)
             throws Exception {
-        long remainingTime = unit.toNanos(timeout);
+        final long waitTime = unit.toNanos(timeout);
         final long startTime = System.nanoTime();
         while (!stopCondition.getAsBoolean()) {
-            if (remainingTime <= 0) {
+            action.run();
+            if (System.nanoTime() - startTime >= waitTime) {
                 throw new TimeoutException("Stop condition was not met after " + unit.toSeconds(timeout) + " seconds.");
             }
-            action.run();
-            remainingTime -= System.nanoTime() - startTime;
         }
     }
 }