Add Channel#closeWithElement

core/shared/src/main/scala/fs2/concurrent/Channel.scala

@@ -110,6 +110,12 @@ sealed trait Channel[F[_], A] {
     */
   def close: F[Either[Channel.Closed, Unit]]
 
+  /** Sends an element through this channel, and closes it right after.
+    *
+    * No-op if the channel is closed, see [[close]] for further info.
+    */
+  def sendAndClose(a: A): F[Either[Channel.Closed, Unit]]
+
   /** Returns true if this channel is closed */
   def isClosed: F[Boolean]
 
@@ -151,7 +157,7 @@ object Channel {
             .drain
         }
 
-        def send(a: A) =
+        def sendImpl(a: A, close: Boolean) =
           F.deferred[Unit].flatMap { producer =>
             state.flatModifyFull { case (poll, state) =>
               state match {
@@ -162,17 +168,23 @@ object Channel {
                   if (size < capacity)
                     (
                       State(a :: values, size + 1, None, producers, false),
-                      notifyStream(waiting).as(rightUnit)
+                      signalClosure.whenA(close) *> notifyStream(waiting).as(rightUnit)
                     )
                   else
                     (
                       State(values, size, None, (a, producer) :: producers, false),
-                      notifyStream(waiting).as(rightUnit) <* waitOnBound(producer, poll)
+                      signalClosure.whenA(close) *>
+                        notifyStream(waiting).as(rightUnit) <*
+                        waitOnBound(producer, poll)
                     )
               }
             }
           }
 
+        def send(a: A) = sendImpl(a, false)
+
+        def sendAndClose(a: A) = sendImpl(a, true)
+
         def trySend(a: A) =
           state.flatModify {
             case s @ State(_, _, _, _, closed @ true) =>

core/shared/src/test/scala/fs2/concurrent/ChannelSuite.scala

@@ -251,4 +251,28 @@ class ChannelSuite extends Fs2Suite {
 
     TestControl.executeEmbed(test)
   }
+
+  test("sendAndClose closes right after sending the last element") {
+    val result = Channel.bounded[IO, Int](1).flatMap { ch =>
+      ch.sendAndClose(0) *> ch.isClosed
+    }
+
+    result.assertEquals(true)
+  }
+
+  test("racing send and sendAndClose should work in bounded case") {
+    val test = Channel.bounded[IO, Int](2).flatMap { ch =>
+      ch.send(0).both(ch.sendAndClose(1))
+    }
+    test.assertEquals((Right(()), Right(())))
+  }
+
+  test("racing send and sendAndClose should work in unbounded case") {
+    val test = Channel.unbounded[IO, Int].flatMap { ch =>
+      ch.send(0).both(ch.sendAndClose(1))
+    }
+
+    test.assertEquals((Right(()), Right(())))
+  }
+
 }