Better handling of zeromq connection close

almond-sh · Jun 19, 2024 · 0a4622c · 0a4622c
1 parent 575ed93
commit 0a4622c
Showing 1 changed file with 22 additions and 11 deletions.
diff --git a/modules/shared/channels/src/main/scala/almond/channels/zeromq/ZeromqConnection.scala b/modules/shared/channels/src/main/scala/almond/channels/zeromq/ZeromqConnection.scala
@@ -1,6 +1,6 @@
 package almond.channels.zeromq
 
-import java.nio.channels.{ClosedByInterruptException, Selector}
+import java.nio.channels.{ClosedByInterruptException, ClosedSelectorException, Selector}
 import java.nio.charset.StandardCharsets.UTF_8
 
 import almond.channels._
@@ -190,18 +190,29 @@ final class ZeromqConnection(
           (channel, new PollItem(socket.channel, Poller.POLLIN))
         }
 
-      withSelector { selector =>
-        ZMQ.poll(selector, pollItems.map(_._2).toArray, pollingDelay.toMillis)
+      val doRead = withSelector { selector =>
+        try {
+          ZMQ.poll(selector, pollItems.map(_._2).toArray, pollingDelay.toMillis)
+          true
+        }
+        catch {
+          case _: ClosedSelectorException if selectorOpt.isEmpty =>
+            // channel was closed
+            false
+        }
       }
 
-      pollItems
-        .collectFirst {
-          case (channel, pi) if pi.isReadable =>
-            channelSocket0(channel)
-              .read
-              .map(_.map((channel, _)))
-        }
-        .getOrElse(IO.pure(None))
+      if (doRead)
+        pollItems
+          .collectFirst {
+            case (channel, pi) if pi.isReadable =>
+              channelSocket0(channel)
+                .read
+                .map(_.map((channel, _)))
+          }
+          .getOrElse(IO.pure(None))
+      else
+        IO.pure(None)
     }.evalOn(threads.pollingEc).flatMap(identity)
 
   def close(partial: Boolean, lingerDuration: Duration): IO[Unit] = {