Adjust Lifecycle of Executors

master/src/main/scala/Master.scala

@@ -11,6 +11,7 @@ import kr.ac.postech.paranode.utils.Progress._
 import org.apache.logging.log4j.scala.Logging
 
 import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
 import scala.concurrent.Await
 import scala.concurrent.ExecutionContext
 import scala.concurrent.ExecutionContextExecutor
@@ -28,10 +29,18 @@ object Master extends Logging {
       masterArguments.numberOfWorkers
     )
 
+    val executor = Executors.newCachedThreadPool()
+
+    val executionContext: ExecutionContext =
+      ExecutionContext.fromExecutor(executor)
+
     Await.result(
-      master.run()(ExecutionContext.global),
+      master.run()(executionContext),
       scala.concurrent.duration.Duration.Inf
     )
+
+    executor.shutdown()
+    executor.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)
   }
 
 }
@@ -54,10 +63,10 @@ class Master(host: String, port: Int, numberOfWorkers: Int) extends Logging {
           s"numberOfWorkers: ${numberOfWorkers}\n"
       )
 
+      val serviceExecutor = Executors.newCachedThreadPool()
+
       val serviceExecutionContext: ExecutionContextExecutor =
-        ExecutionContext.fromExecutor(
-          Executors.newCachedThreadPool()
-        )
+        ExecutionContext.fromExecutor(serviceExecutor)
 
       val server =
         new GrpcServer(
@@ -87,10 +96,12 @@ class Master(host: String, port: Int, numberOfWorkers: Int) extends Logging {
         WorkerClient(worker.host, worker.port)
       }
 
+      val requestExecutor = Executors
+        .newFixedThreadPool(registeredWorkers.size)
+
       val requestExecutionContext: ExecutionContextExecutor =
-        scala.concurrent.ExecutionContext.fromExecutor(
-          java.util.concurrent.Executors
-            .newFixedThreadPool(registeredWorkers.size)
+        ExecutionContext.fromExecutor(
+          requestExecutor
         )
 
       logger.info(s"[Master] Clients: $clients")
@@ -138,11 +149,29 @@ class Master(host: String, port: Int, numberOfWorkers: Int) extends Logging {
 
       logger.info("[Master] Merge finished")
 
+      logger.info("[Master] Terminate started")
+
+      clients.terminate()(requestExecutionContext)
+
+      logger.info("[Master] Terminate finished")
+
       clients.foreach(_.shutdown())
 
       server.stop()
 
       serverState.update(_ => Progress.Finished)
+
+      serviceExecutor.shutdown()
+      requestExecutor.shutdown()
+
+      serviceExecutor.awaitTermination(
+        5,
+        TimeUnit.SECONDS
+      )
+      requestExecutor.awaitTermination(
+        5,
+        TimeUnit.SECONDS
+      )
     }
 
   def blockUntilRunning(): Unit = {

rpc/src/main/protobuf/worker.proto

@@ -11,6 +11,7 @@ service Worker {
   rpc Exchange (ExchangeRequest) returns (ExchangeReply) {}
   rpc SaveBlock (SaveBlockRequest) returns (SaveBlockReply) {}
   rpc Merge (MergeRequest) returns (MergeReply) {}
+  rpc Terminate (TerminateRequest) returns (TerminateReply) {}
 }
 
 message SampleRequest {
@@ -46,3 +47,8 @@ message SaveBlockReply {}
 message MergeRequest {}
 
 message MergeReply {}
+
+message TerminateRequest {}
+
+message TerminateReply {}
+

rpc/src/main/scala/GrpcServer.scala

@@ -5,6 +5,8 @@ import io.grpc.ServerBuilder
 import io.grpc.ServerServiceDefinition
 import org.apache.logging.log4j.scala.Logging
 
+import java.util.concurrent.TimeUnit
+
 class GrpcServer[T <: ServerBuilder[T]](
     service: ServerServiceDefinition,
     port: Int
@@ -36,7 +38,8 @@ class GrpcServer[T <: ServerBuilder[T]](
 
   def stop(): Unit = {
     if (server != null) {
-      server.shutdown()
+      server.shutdown().awaitTermination(5, TimeUnit.SECONDS)
+
     }
   }
 

rpc/src/main/scala/WorkerClient.scala

@@ -68,6 +68,17 @@ object WorkerClient {
         ),
         scala.concurrent.duration.Duration.Inf
       )
+
+    def terminate()(implicit
+        executionContext: ExecutionContext
+    ): List[TerminateReply] =
+      Await.result(
+        Future.traverse(clients)(_.terminate())(
+          GenericBuildFrom[WorkerClient, TerminateReply],
+          executionContext
+        ),
+        scala.concurrent.duration.Duration.Inf
+      )
   }
 
   def apply(host: String, port: Int): WorkerClient = {
@@ -137,4 +148,9 @@ class WorkerClient private (
     val request = MergeRequest()
     stub.merge(request)
   }
+
+  def terminate(): Future[TerminateReply] = {
+    val request = TerminateRequest()
+    stub.terminate(request)
+  }
 }

worker/src/main/scala/Worker.scala

@@ -7,6 +7,7 @@ import kr.ac.postech.paranode.utils.Hooks
 import org.apache.logging.log4j.scala.Logging
 
 import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
 import scala.concurrent.Await
 import scala.concurrent.ExecutionContext
 import scala.concurrent.Future
@@ -28,14 +29,27 @@ object Worker extends Logging {
       workerArguments.outputDirectory
     )
 
+    val executor = Executors.newCachedThreadPool()
+
+    val executionContext: ExecutionContext =
+      ExecutionContext.fromExecutor(
+        executor
+      )
+
     try {
       Await.result(
-        worker.run()(ExecutionContext.global),
+        worker.run()(executionContext),
         scala.concurrent.duration.Duration.Inf
       )
     } catch {
-      case _: io.grpc.StatusRuntimeException => System.exit(0)
-      case _: Exception                      => System.exit(1)
+      case statusRuntimeException: io.grpc.StatusRuntimeException if {
+            val status = statusRuntimeException.getStatus
+            status.getCode == io.grpc.Status.Code.UNAVAILABLE && status.getDescription
+              .contains("debug data: app_requested")
+          } => // Suppress
+    } finally {
+      worker.shutdown()
+      executor.shutdown()
     }
   }
 
@@ -51,15 +65,21 @@ class Worker(
 ) extends Logging {
   private val workerMetadata = WorkerMetadata(host, port, None)
 
+  private val serviceExecutor = Executors.newCachedThreadPool()
+
   private val serviceExecutionContext: ExecutionContext =
     ExecutionContext.fromExecutor(
-      Executors.newCachedThreadPool()
+      serviceExecutor
     )
 
   private val server = new GrpcServer(
     WorkerService(
       inputDirectories,
-      outputDirectory
+      outputDirectory,
+      onFinished = () => {
+        serviceExecutor.shutdown()
+        serviceExecutor.awaitTermination(3, TimeUnit.SECONDS)
+      }
     )(serviceExecutionContext),
     port
   )