Run loops check the interrupted flag (#1215)

* Run loops check the interrupted flag * Run loops check the interrupted flag
nats-io · Aug 21, 2024 · 0e4f39c · 0e4f39c
1 parent 4cf29d4
commit 0e4f39c
Show file tree

Hide file tree

Showing 7 changed files with 16 additions and 28 deletions.
diff --git a/src/main/java/io/nats/client/impl/MessageManager.java b/src/main/java/io/nats/client/impl/MessageManager.java
@@ -149,7 +149,7 @@ public void shutdown() {
 
         @Override
         public void run() {
-            if (alive.get()) {
+            if (alive.get() && !Thread.interrupted()) {
                 long sinceLast = System.currentTimeMillis() - lastMsgReceived.get();
                 if (alive.get() && sinceLast > alarmPeriodSetting) {
                     handleHeartbeatError();

diff --git a/src/main/java/io/nats/client/impl/NatsConnection.java b/src/main/java/io/nats/client/impl/NatsConnection.java
@@ -2178,20 +2178,17 @@ public CompletableFuture<Boolean> drain(Duration timeout) throws TimeoutExceptio
         // Wait for the timeout or the pending count to go to 0
         executor.submit(() -> {
             try {
-                Instant now = Instant.now();
-
-                while (timeout == null || timeout.equals(Duration.ZERO)
-                        || Duration.between(start, now).compareTo(timeout) < 0) {
+                long stop = (timeout == null || timeout.equals(Duration.ZERO))
+                    ? Long.MAX_VALUE
+                    : System.nanoTime() + timeout.toNanos();
+                while (System.nanoTime() < stop && !Thread.interrupted())
+                {
                     consumers.removeIf(NatsConsumer::isDrained);
-
                     if (consumers.isEmpty()) {
                         break;
                     }
-
                     //noinspection BusyWait
                     Thread.sleep(1); // Sleep 1 milli
-
-                    now = Instant.now();
                 }
 
                 // Stop publishing
@@ -2201,16 +2198,13 @@ public CompletableFuture<Boolean> drain(Duration timeout) throws TimeoutExceptio
                 if (timeout == null || timeout.equals(Duration.ZERO)) {
                     this.flush(Duration.ZERO);
                 } else {
-                    now = Instant.now();
-
+                    Instant now = Instant.now();
                     Duration passed = Duration.between(start, now);
                     Duration newTimeout = timeout.minus(passed);
-
                     if (newTimeout.toNanos() > 0) {
                         this.flush(newTimeout);
                     }
                 }
-
                 this.close(false, false); // close the connection after the last flush
                 tracker.complete(consumers.isEmpty());
             } catch (TimeoutException e) {

diff --git a/src/main/java/io/nats/client/impl/NatsConnectionReader.java b/src/main/java/io/nats/client/impl/NatsConnectionReader.java
@@ -128,7 +128,7 @@ public void run() {
             this.gotCR = false;
             this.opPos = 0;
 
-            while (this.running.get()) {
+            while (running.get() && !Thread.interrupted()) {
                 this.bufferPosition = 0;
                 int bytesRead = dataPort.read(this.buffer, 0, this.buffer.length);
 

diff --git a/src/main/java/io/nats/client/impl/NatsConnectionWriter.java b/src/main/java/io/nats/client/impl/NatsConnectionWriter.java
@@ -193,7 +193,7 @@ public void run() {
             dataPort = this.dataPortFuture.get(); // Will wait for the future to complete
             StatisticsCollector stats = this.connection.getNatsStatistics();
 
-            while (this.running.get()) {
+            while (running.get() && !Thread.interrupted()) {
                 NatsMessage msg;
                 if (this.reconnectMode.get()) {
                     msg = this.reconnectOutgoing.accumulate(sendBufferLength.get(), Options.MAX_MESSAGES_IN_NETWORK_BUFFER, reconnectTimeout);

diff --git a/src/main/java/io/nats/client/impl/NatsConsumer.java b/src/main/java/io/nats/client/impl/NatsConsumer.java
@@ -205,17 +205,15 @@ public CompletableFuture<Boolean> drain(Duration timeout) throws InterruptedExce
         // draining
         connection.getExecutor().submit(() -> {
             try {
-                Instant now = Instant.now();
-
-                while (timeout == null || timeout.equals(Duration.ZERO)
-                        || Duration.between(start, now).compareTo(timeout) < 0) {
+                long stop = (timeout == null || timeout.equals(Duration.ZERO))
+                    ? Long.MAX_VALUE
+                    : System.nanoTime() + timeout.toNanos();
+                while (System.nanoTime() < stop && !Thread.interrupted()) {
                     if (this.isDrained()) {
                         break;
                     }
-
+                    //noinspection BusyWait
                     Thread.sleep(1); // Sleep 1 milli
-
-                    now = Instant.now();
                 }
 
                 this.cleanUpAfterDrain();

diff --git a/src/main/java/io/nats/client/impl/NatsDispatcher.java b/src/main/java/io/nats/client/impl/NatsDispatcher.java
@@ -84,10 +84,8 @@ boolean breakRunLoop() {
 
     public void run() {
         try {
-            while (this.running.get()) { // start
-
+            while (running.get() && !Thread.interrupted()) {
                 NatsMessage msg = this.incoming.pop(this.waitForMessage);
-
                 if (msg != null) {
                     NatsSubscription sub = msg.getNatsSubscription();
                     if (sub != null && sub.isActive()) {

diff --git a/src/main/java/io/nats/client/impl/NatsDispatcherWithExecutor.java b/src/main/java/io/nats/client/impl/NatsDispatcherWithExecutor.java
@@ -24,10 +24,8 @@ class NatsDispatcherWithExecutor extends NatsDispatcher {
     @Override
     public void run() {
         try {
-            while (this.running.get()) { // start
-
+            while (running.get() && !Thread.interrupted()) {
                 NatsMessage msg = this.incoming.pop(this.waitForMessage);
-
                 if (msg != null) {
                     NatsSubscription sub = msg.getNatsSubscription();
                     if (sub != null && sub.isActive()) {