<commit_before>package org.labrad

import java.io.{PrintWriter, StringWriter}
import java.util.regex.Pattern
import org.labrad.annotations.IsServer
import org.labrad.data._
import org.labrad.errors.LabradException
import org.labrad.util.{AsyncSemaphore, Logging, Util}
import scala.collection.mutable
import scala.concurrent.{Await, ExecutionContext, Future, Promise}
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration._
import scala.reflect.ClassTag
import scala.reflect.runtime.universe.TypeTag
import scala.util.{Success, Failure}

trait ServerContext {
  def init(): Unit
  def expire(): Unit
}

abstract class Server[S <: Server[S, _] : TypeTag, T <: ServerContext : TypeTag] extends Logging {

  val name: String
  val doc: String

  private var _cxn: ServerConnection = _

  /**
   * Allow subclasses to access (but not modify) the server connection
   */
  protected def cxn: ServerConnection = _cxn

  protected val contexts = mutable.Map.empty[Context, ContextState]

  private val (srvSettings, srvHandlerFactory) = Reflect.makeHandler[S]
  private val (ctxSettings, ctxHandlerFactory) = Reflect.makeHandler[T]
  private val settings = srvSettings ++ ctxSettings
  private val srvHandler = srvHandlerFactory(this.asInstanceOf[S])

  private val srvSettingIds = srvSettings.map(_.id).toSet
  private val ctxSettingIds = ctxSettings.map(_.id).toSet

  /**
   * Start this server by connecting to the manager at the given location.
   */
  def start(host: String, port: Int, password: Array[Char], nameOpt: Option[String] = None): Unit = {
    val name = nameOpt.getOrElse(this.name)
    _cxn = new ServerConnection(name, doc, host, port, password, packet => handleRequest(packet))
    cxn.connect()

    // if the vm goes down or we lose the connection, shutdown
    sys.ShutdownHookThread { stop() }
    cxn.addConnectionListener { case false => stop() }

    init()

    val msgId = cxn.getMessageId
    cxn.addMessageListener {
      case Message(`msgId`, context, _, _) =>
        val stateOpt = contexts.synchronized { contexts.remove(context) }
        stateOpt.foreach { _.expire() }
    }

    val registrations = settings.sortBy(_.id).map(s =>
      Cluster(
        UInt(s.id),
        Str(s.name),
        Str(s.doc),
        Arr(s.accepts.strs.map(Str(_))), // TODO: when manager supports patterns, just use .pat here
        Arr(s.returns.strs.map(Str(_))),
        Str(""))
    )

    val registerF = for {
      _ <- cxn.send("Manager", registrations.map("S: Register Setting" -> _): _*)
      _ <- cxn.send("Manager", "S: Notify on Context Expiration" -> Cluster(UInt(msgId), Bool(false)))
      _ <- cxn.send("Manager", "S: Start Serving" -> Data.NONE)
    } yield ()
    Await.result(registerF, 30.seconds)

    log.info("Now serving...")
  }

  /**
   * A promise that will be completed when this server shuts down.
   */
  private val shutdownPromise = Promise[Unit]

  /**
   * Stop this server
   */
  def stop(): Unit = {
    try {
      val ctxs = contexts.synchronized {
        val ctxs = contexts.values.toVector
        contexts.clear()
        ctxs
      }
      val expirations = ctxs.map { _.expire() }
      Await.result(Future.sequence(expirations), 60.seconds)
      shutdown()
      cxn.close()
      shutdownPromise.trySuccess(())
    } catch {
      case e: Exception =>
        shutdownPromise.tryFailure(e)
    }
  }

  /**
   * Wait for this server to shutdown
   */
  def awaitShutdown(timeout: Duration = Duration.Inf): Unit = {
    Await.result(shutdownPromise.future, timeout)
  }

  // user-overidden methods
  def init(): Unit
  def shutdown(): Unit
  def newContext(context: Context): T


  /**
   * Helper class containing per-context state and a semaphore to enable
   * serializing requests made within a given context.
   */
  class ContextState {
    val sem: AsyncSemaphore = new AsyncSemaphore(1)
    var handler: Option[(T, RequestContext => Data)] = None

    def expire(): Future[Unit] = sem.map {<commit_msg>Ensure that server shutdown code only runs once.<commit_after>package org.labrad

import java.io.{PrintWriter, StringWriter}
import java.util.concurrent.atomic.AtomicBoolean
import java.util.regex.Pattern
import org.labrad.annotations.IsServer
import org.labrad.data._
import org.labrad.errors.LabradException
import org.labrad.util.{AsyncSemaphore, Logging, Util}
import scala.collection.mutable
import scala.concurrent.{Await, ExecutionContext, Future, Promise}
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.duration._
import scala.reflect.ClassTag
import scala.reflect.runtime.universe.TypeTag
import scala.util.{Success, Failure}

trait ServerContext {
  def init(): Unit
  def expire(): Unit
}

abstract class Server[S <: Server[S, _] : TypeTag, T <: ServerContext : TypeTag] extends Logging {

  val name: String
  val doc: String

  private var _cxn: ServerConnection = _
  private val shutdownCalled = new AtomicBoolean(false)

  /**
   * Allow subclasses to access (but not modify) the server connection
   */
  protected def cxn: ServerConnection = _cxn

  protected val contexts = mutable.Map.empty[Context, ContextState]

  private val (srvSettings, srvHandlerFactory) = Reflect.makeHandler[S]
  private val (ctxSettings, ctxHandlerFactory) = Reflect.makeHandler[T]
  private val settings = srvSettings ++ ctxSettings
  private val srvHandler = srvHandlerFactory(this.asInstanceOf[S])

  private val srvSettingIds = srvSettings.map(_.id).toSet
  private val ctxSettingIds = ctxSettings.map(_.id).toSet

  /**
   * Start this server by connecting to the manager at the given location.
   */
  def start(host: String, port: Int, password: Array[Char], nameOpt: Option[String] = None): Unit = {
    val name = nameOpt.getOrElse(this.name)
    _cxn = new ServerConnection(name, doc, host, port, password, packet => handleRequest(packet))
    cxn.connect()

    // if the vm goes down or we lose the connection, shutdown
    sys.ShutdownHookThread { stop() }
    cxn.addConnectionListener { case false => stop() }

    init()

    val msgId = cxn.getMessageId
    cxn.addMessageListener {
      case Message(`msgId`, context, _, _) =>
        val stateOpt = contexts.synchronized { contexts.remove(context) }
        stateOpt.foreach { _.expire() }
    }

    val registrations = settings.sortBy(_.id).map(s =>
      Cluster(
        UInt(s.id),
        Str(s.name),
        Str(s.doc),
        Arr(s.accepts.strs.map(Str(_))), // TODO: when manager supports patterns, just use .pat here
        Arr(s.returns.strs.map(Str(_))),
        Str(""))
    )

    val registerF = for {
      _ <- cxn.send("Manager", registrations.map("S: Register Setting" -> _): _*)
      _ <- cxn.send("Manager", "S: Notify on Context Expiration" -> Cluster(UInt(msgId), Bool(false)))
      _ <- cxn.send("Manager", "S: Start Serving" -> Data.NONE)
    } yield ()
    Await.result(registerF, 30.seconds)

    log.info("Now serving...")
  }

  /**
   * A promise that will be completed when this server shuts down.
   */
  private val shutdownPromise = Promise[Unit]

  /**
   * Stop this server
   */
  def stop(): Unit = {
    if (shutdownCalled.compareAndSet(false, true)) {
      try {
        val ctxs = contexts.synchronized {
          val ctxs = contexts.values.toVector
          contexts.clear()
          ctxs
        }
        val expirations = ctxs.map { _.expire() }
        Await.result(Future.sequence(expirations), 10.seconds)
        shutdown()
        cxn.close()
        shutdownPromise.trySuccess(())
      } catch {
        case e: Exception =>
          shutdownPromise.tryFailure(e)
      }
    }
  }

  /**
   * Wait for this server to shutdown
   */
  def awaitShutdown(timeout: Duration = Duration.Inf): Unit = {
    Await.result(shutdownPromise.future, timeout)
  }

  // user-overidden methods
  def init(): Unit
  def shutdown(): Unit
  def newContext(context: Context): T


  /**
   * Helper class containing per-context state and a semaphore to enable
   * serializing requests made within a given context.
   */
  class ContextState {
    val sem: AsyncSemaphore = new AsyncSemaphore(1)
    var handler: Option[(T, RequestContext => Data)] = None

    def expire(): Future[Unit] = sem.map {
{'fraction_non_alphanumeric': 0.11672115268192275, 'fraction_numerical': 0.001210800339024095, 'mean_word_length': 3.203562340966921, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1544815', 'n_tokens_mistral': 2578, 'n_tokens_neox': 2491, 'n_words': 665}