Merge pull request #9691 from kitbellew/server-keeps-accepting

[2.x] fix: Keep accepting after a client fails

main-command/src/main/scala/sbt/internal/AtomicCloseable.scala

@@ -18,7 +18,7 @@ private[sbt] class AtomicCloseable[A >: Null <: AutoCloseable](val ref: AtomicRe
     extends AnyVal:
   def get: A = ref.get
   def set(c: A): Unit = AtomicCloseable.close(ref.getAndSet(c))
-  def close(): Unit = AtomicCloseable.close(ref.getAndSet(null))
+  def close(): Unit = AtomicCloseable.close(AtomicCloseable.release(ref))
 
   /** Keeps the value another caller put here, and closes the one this caller built. */
   def setIfEmpty(ctor: => A): A =
@@ -35,6 +35,13 @@ private[sbt] object AtomicCloseable:
   def apply[A >: Null <: AutoCloseable](): AtomicCloseable[A] =
     new AtomicCloseable(new AtomicReference[A])
 
+  def apply[A >: Null <: AutoCloseable](c: A): AtomicCloseable[A] =
+    new AtomicCloseable(new AtomicReference[A](c))
+
   def close(obj: AutoCloseable): Unit =
     if obj ne null then Util.ignoreTry(obj.close())
+
+  def release[A >: Null <: AutoCloseable](ref: AtomicReference[A]): A =
+    ref.getAndSet(null)
+
 end AtomicCloseable

main-command/src/main/scala/sbt/internal/server/Server.scala

@@ -12,7 +12,7 @@ package server
 
 import java.io.{ File, IOException }
 import java.net.{ InetAddress, ServerSocket, Socket, SocketException, SocketTimeoutException }
-import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.{ AtomicBoolean, AtomicReference }
 import java.security.SecureRandom
 import java.math.BigInteger
 
@@ -45,7 +45,7 @@ private[sbt] object Server {
 
   def start(
       connection: ServerConnection,
-      onIncomingSocket: (Socket, ServerInstance) => Unit,
+      onIncomingSocket: (AtomicReference[Socket], ServerInstance) => Unit,
       log: Logger
   ): ServerInstance =
     new ServerInstance { self =>
@@ -105,14 +105,19 @@ private[sbt] object Server {
               running.set(true)
               p.success(())
               while (running.get()) {
+                val clientSocket = AtomicCloseable[Socket]()
                 try {
-                  val socket = serverSocket.accept()
-                  onIncomingSocket(socket, self)
+                  clientSocket.set(serverSocket.accept())
+                  onIncomingSocket(clientSocket.ref, self)
                 } catch {
+                  case scala.util.control.NonFatal(e) if clientSocket.get ne null =>
+                    log.error(s"sbt server failed to serve a client: $e")
+                    log.trace(e)
                   case e: IOException if Option(e.getMessage).exists(_.contains("connect")) =>
                   case _: SocketTimeoutException          => // its ok
                   case _: SocketException if !running.get => // the server is shutting down
                 }
+                clientSocket.close()
               }
               serverSocketHolder.close()
           }

main-command/src/test/scala/sbt/internal/AtomicCloseableSpec.scala

@@ -17,6 +17,14 @@ object AtomicCloseableSpec extends BasicTestSuite:
     override def close(): Unit = closed = true
   end Probe
 
+  test("a holder built around a value"):
+    val probe = new Probe
+    val holder = AtomicCloseable(probe)
+    assert(holder.get == probe)
+    assert(!probe.closed)
+    holder.close()
+    assert(probe.closed)
+
   test("a value that replaces another closes it"):
     val holder = AtomicCloseable[Probe]()
     val first, second = new Probe

main-command/src/test/scala/sbt/internal/server/ServerAcceptSpec.scala

@@ -0,0 +1,96 @@
+/*
+ * sbt
+ * Copyright 2023, Scala center
+ * Copyright 2011 - 2022, Lightbend, Inc.
+ * Copyright 2008 - 2010, Mark Harrah
+ * Licensed under Apache License 2.0 (see LICENSE)
+ */
+
+package sbt
+package internal
+package server
+
+import java.io.File
+import java.net.Socket
+import java.util.concurrent.atomic.AtomicReference
+import java.nio.file.{ Files, Paths }
+import java.util.concurrent.atomic.AtomicInteger
+
+import scala.concurrent.Await
+import scala.concurrent.duration.*
+
+import org.scalasbt.ipcsocket.UnixDomainSocket
+import sbt.internal.util.Util.isWindows
+import verify.BasicTestSuite
+
+object ServerAcceptSpec extends BasicTestSuite:
+  private def withServer(
+      onIncomingSocket: (AtomicReference[Socket], ServerInstance) => Unit
+  )(f: (ServerInstance, File) => Unit): Unit =
+    // the socket path has a length limit, so keep the directory short
+    val dir = Files.createTempDirectory(Paths.get("/tmp"), "sbtsrv").toFile
+    val connection = ServerConnection(
+      connectionType = ConnectionType.Local,
+      host = "127.0.0.1",
+      port = 0,
+      auth = Set.empty,
+      portfile = new File(dir, "active.json"),
+      tokenfile = new File(dir, "token.json"),
+      socketfile = new File(dir, "sock"),
+      pipeName = "sbt-test-" + dir.getName,
+      appConfiguration = null, // only a bsp connection file reads it, and bsp is off here
+      windowsServerSecurityLevel = 0,
+      useJni = false,
+      bspEnabled = false,
+    )
+    val instance = Server.start(connection, onIncomingSocket, sbt.util.Logger.Null)
+    Await.ready(instance.ready, 10.seconds)
+    try f(instance, connection.socketfile)
+    finally
+      instance.shutdown()
+      sbt.io.IO.delete(dir)
+
+  private def waitUntil(p: => Boolean): Boolean =
+    val deadline = 10.seconds.fromNow
+    while !p && deadline.hasTimeLeft() do Thread.sleep(20)
+    p
+
+  test("a client that the server fails to serve"):
+    if !isWindows then
+      val served = new AtomicInteger
+      val handler: (AtomicReference[Socket], ServerInstance) => Unit = (_, _) =>
+        served.incrementAndGet()
+        throw new RuntimeException("this client cannot be served")
+      withServer(handler): (_, socketfile) =>
+        new UnixDomainSocket(socketfile.getAbsolutePath, false)
+        assert(waitUntil(served.get == 1))
+        new UnixDomainSocket(socketfile.getAbsolutePath, false)
+        // the loop accepted a second client, so the first one did not end it
+        assert(waitUntil(served.get == 2))
+
+  test("a socket the callback takes over"):
+    if !isWindows then
+      val served = new AtomicInteger
+      val first = new AtomicReference[Socket]
+      val handler: (AtomicReference[Socket], ServerInstance) => Unit = (socket, _) =>
+        if served.getAndIncrement == 0 then
+          first.set(socket.get)
+          AtomicCloseable.release(socket)
+      withServer(handler): (_, socketfile) =>
+        new UnixDomainSocket(socketfile.getAbsolutePath, false)
+        new UnixDomainSocket(socketfile.getAbsolutePath, false)
+        // the second client proves the loop went round, so it has passed its close
+        assert(waitUntil(served.get == 2))
+        assert(!first.get.isClosed)
+
+  test("a socket the callback leaves"):
+    if !isWindows then
+      val left = new AtomicReference[Socket]
+      val handler: (AtomicReference[Socket], ServerInstance) => Unit =
+        (socket, _) => left.set(socket.get)
+      withServer(handler): (_, socketfile) =>
+        new UnixDomainSocket(socketfile.getAbsolutePath, false)
+        assert(waitUntil(left.get ne null))
+        assert(waitUntil(left.get.isClosed))
+
+end ServerAcceptSpec

main/src/main/scala/sbt/internal/CommandExchange.scala

@@ -207,19 +207,20 @@ private[sbt] final class CommandExchange {
     lazy val enableBsp = s.get(bspEnabled).getOrElse(true)
     lazy val portfile = s.baseDir / "project" / "target" / "active.json"
 
-    def onIncomingSocket(socket: Socket, instance: ServerInstance): Unit = {
+    def onIncomingSocket(socket: AtomicReference[Socket], instance: ServerInstance): Unit = {
       val name = newNetworkName
       Terminal.consoleLog(s"new client connected: $name")
       val channel =
         new NetworkChannel(
           name,
-          socket,
+          socket.get,
           auth,
           instance,
           handlers,
           mkAskUser(name),
         )
       subscribe(channel)
+      AtomicCloseable.release(socket) // i took over
     }
     if (server.isEmpty && firstInstance.get) {
       val h = Hash.halfHashString(IO.toURI(portfile).toString)