Shared server sockets

core/shared/src/main/scala/fs2/Shared.scala

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2013 Functional Streams for Scala
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of
+ * this software and associated documentation files (the "Software"), to deal in
+ * the Software without restriction, including without limitation the rights to
+ * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+ * the Software, and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+ * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+ * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+ * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+ * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+package fs2
+
+import cats.effect._
+import cats.syntax.all._
+
+sealed trait Shared[F[_], A] {
+  def resource: Resource[F, A]
+}
+
+object Shared {
+  def allocate[F[_], A](
+      resource: Resource[F, A]
+  )(implicit F: Concurrent[F]): Resource[F, (Shared[F, A], A)] = {
+    final case class State(value: A, finalizer: F[Unit], permits: Int) {
+      def addPermit: State = copy(permits = permits + 1)
+      def releasePermit: State = copy(permits = permits - 1)
+    }
+
+    MonadCancel[Resource[F, *]].uncancelable { _ =>
+      for {
+        underlying <- Resource.eval(resource.allocated)
+        state <- Resource.eval(F.ref[Option[State]](Some(State(underlying._1, underlying._2, 0))))
+        shared = new Shared[F, A] {
+          def acquire: F[A] =
+            state.modify {
+              case Some(st) => (Some(st.addPermit), F.pure(st.value))
+              case None =>
+                (None, F.raiseError[A](new Throwable("finalization has already occurred")))
+            }.flatten
+
+          def release: F[Unit] =
+            state.modify {
+              case Some(st) if st.permits > 1 => (Some(st.releasePermit), F.unit)
+              case Some(st)                   => (None, st.finalizer)
+              case None                       => (None, F.raiseError[Unit](new Throwable("can't finalize")))
+            }.flatten
+
+          override def resource: Resource[F, A] =
+            Resource.make(acquire)(_ => release)
+        }
+        _ <- shared.resource
+      } yield shared -> underlying._1
+    }
+  }
+}

io/src/main/scala/fs2/io/net/Network.scala

@@ -132,13 +132,13 @@ object Network {
           address: Option[Host],
           port: Option[Port],
           options: List[SocketOption]
-      ): Stream[F, Socket[F]] = globalSocketGroup.server(address, port, options)
+      ): Stream[F, Shared[F, Socket[F]]] = globalSocketGroup.server(address, port, options)
 
       def serverResource(
           address: Option[Host],
           port: Option[Port],
           options: List[SocketOption]
-      ): Resource[F, (SocketAddress[IpAddress], Stream[F, Socket[F]])] =
+      ): Resource[F, (SocketAddress[IpAddress], Stream[F, Shared[F, Socket[F]]])] =
         globalSocketGroup.serverResource(address, port, options)
 
       def openDatagramSocket(

io/src/main/scala/fs2/io/net/Socket.scala

@@ -24,7 +24,7 @@ package io
 package net
 
 import com.comcast.ip4s.{IpAddress, SocketAddress}
-import cats.effect.{Async, Resource}
+import cats.effect.Async
 import cats.effect.std.Semaphore
 import cats.syntax.all._
 
@@ -79,12 +79,10 @@ trait Socket[F[_]] {
 object Socket {
   private[net] def forAsync[F[_]: Async](
       ch: AsynchronousSocketChannel
-  ): Resource[F, Socket[F]] =
-    Resource.make {
-      (Semaphore[F](1), Semaphore[F](1)).mapN { (readSemaphore, writeSemaphore) =>
-        new AsyncSocket[F](ch, readSemaphore, writeSemaphore)
-      }
-    }(_ => Async[F].delay(if (ch.isOpen) ch.close else ()))
+  ): F[Socket[F]] =
+    (Semaphore[F](1), Semaphore[F](1)).mapN { (readSemaphore, writeSemaphore) =>
+      new AsyncSocket[F](ch, readSemaphore, writeSemaphore)
+    }
 
   private final class AsyncSocket[F[_]](
       ch: AsynchronousSocketChannel,

io/src/main/scala/fs2/io/net/SocketGroup.scala

@@ -68,15 +68,15 @@ trait SocketGroup[F[_]] {
       address: Option[Host] = None,
       port: Option[Port] = None,
       options: List[SocketOption] = List.empty
-  ): Stream[F, Socket[F]]
+  ): Stream[F, Shared[F, Socket[F]]]
 
   /** Like [[server]] but provides the `SocketAddress` of the bound server socket before providing accepted sockets.
     */
   def serverResource(
       address: Option[Host] = None,
       port: Option[Port] = None,
       options: List[SocketOption] = List.empty
-  ): Resource[F, (SocketAddress[IpAddress], Stream[F, Socket[F]])]
+  ): Resource[F, (SocketAddress[IpAddress], Stream[F, Shared[F, Socket[F]]])]
 }
 
 private[net] object SocketGroup {
@@ -114,14 +114,16 @@ private[net] object SocketGroup {
           }
         }
 
-      Resource.eval(setup.flatMap(connect)).flatMap(Socket.forAsync(_))
+      Resource
+        .make(setup.flatMap(connect))(ch => Async[F].delay(if (ch.isOpen()) ch.close() else ()))
+        .evalMap(Socket.forAsync(_))
     }
 
     def server(
         address: Option[Host],
         port: Option[Port],
         options: List[SocketOption]
-    ): Stream[F, Socket[F]] =
+    ): Stream[F, Shared[F, Socket[F]]] =
       Stream
         .resource(
           serverResource(
@@ -136,7 +138,7 @@ private[net] object SocketGroup {
         address: Option[Host],
         port: Option[Port],
         options: List[SocketOption]
-    ): Resource[F, (SocketAddress[IpAddress], Stream[F, Socket[F]])] = {
+    ): Resource[F, (SocketAddress[IpAddress], Stream[F, Shared[F, Socket[F]]])] = {
 
       val setup: F[AsynchronousServerSocketChannel] =
         address.traverse(_.resolve[F]).flatMap { addr =>
@@ -161,8 +163,8 @@ private[net] object SocketGroup {
 
       def acceptIncoming(
           sch: AsynchronousServerSocketChannel
-      ): Stream[F, Socket[F]] = {
-        def go: Stream[F, Socket[F]] = {
+      ): Stream[F, Shared[F, Socket[F]]] = {
+        def go: Stream[F, Shared[F, Socket[F]]] = {
           def acceptChannel: F[AsynchronousSocketChannel] =
             Async[F].async_[AsynchronousSocketChannel] { cb =>
               sch.accept(
@@ -177,8 +179,17 @@ private[net] object SocketGroup {
             }
 
           Stream.eval(acceptChannel.attempt).flatMap {
-            case Left(_)         => Stream.empty[F]
-            case Right(accepted) => Stream.resource(Socket.forAsync(accepted))
+            case Left(_) => Stream.empty[F]
+            case Right(accepted) =>
+              val socketResource = Resource
+                .make(Async[F].pure(accepted))(ch =>
+                  Async[F].delay(
+                    if (ch.isOpen()) ch.close()
+                    else ()
+                  )
+                )
+                .evalMap(Socket.forAsync(_))
+              Stream.resource(Shared.allocate(socketResource)).map(_._1)
           } ++ go
         }
 

io/src/test/scala/fs2/io/net/tcp/SocketSuite.scala

@@ -50,10 +50,12 @@ class SocketSuite extends Fs2Suite {
       Stream
         .resource(setup)
         .flatMap { case (server, clients) =>
-          val echoServer = server.map { socket =>
-            socket.reads
-              .through(socket.writes)
-              .onFinalize(socket.endOfOutput)
+          val echoServer = server.map { shared =>
+            Stream.resource(shared.resource).flatMap { socket =>
+              socket.reads
+                .through(socket.writes)
+                .onFinalize(socket.endOfOutput)
+            }
           }.parJoinUnbounded
 
           val msgClients = clients
@@ -86,11 +88,13 @@ class SocketSuite extends Fs2Suite {
       Stream
         .resource(setup)
         .flatMap { case (server, clients) =>
-          val junkServer = server.map { socket =>
-            Stream
-              .chunk(message)
-              .through(socket.writes)
-              .onFinalize(socket.endOfOutput)
+          val junkServer = server.map { shared =>
+            Stream.resource(shared.resource).flatMap { socket =>
+              Stream
+                .chunk(message)
+                .through(socket.writes)
+                .onFinalize(socket.endOfOutput)
+            }
           }.parJoinUnbounded
 
           val client =
@@ -117,7 +121,9 @@ class SocketSuite extends Fs2Suite {
       Stream
         .resource(setup)
         .flatMap { case (server, clients) =>
-          val readOnlyServer = server.map(_.reads).parJoinUnbounded
+          val readOnlyServer = server.map { shared =>
+            Stream.resource(shared.resource).flatMap(_.reads)
+          }.parJoinUnbounded
           val client =
             clients.take(1).flatMap { socket =>
               // concurrent writes

io/src/test/scala/fs2/io/net/tls/TLSSocketSuite.scala

@@ -117,7 +117,12 @@ class TLSSocketSuite extends TLSSuite {
         addressAndConnections <- Network[IO].serverResource(Some(ip"127.0.0.1"))
         (serverAddress, server) = addressAndConnections
         client <- Network[IO].client(serverAddress).flatMap(tlsContext.client(_))
-      } yield server.flatMap(s => Stream.resource(tlsContext.server(s))) -> client
+      } yield {
+        val tlsServer = server.flatMap { shared =>
+          Stream.resource(shared.resource.flatMap(tlsContext.server(_)))
+        }
+        tlsServer -> client
+      }
 
       Stream
         .resource(setup)