feat: add exit logic

betarixm · Dec 10, 2023 · 7c946e3 · 7c946e3
1 parent d0b77c7
commit 7c946e3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/master/src/main/scala/Master.scala b/master/src/main/scala/Master.scala
@@ -32,6 +32,7 @@ object Master extends Logging {
       master.run()(ExecutionContext.global),
       scala.concurrent.duration.Duration.Inf
     )
+    System.exit(0)
   }
 
 }

diff --git a/worker/src/main/scala/Worker.scala b/worker/src/main/scala/Worker.scala
@@ -28,10 +28,15 @@ object Worker extends Logging {
       workerArguments.outputDirectory
     )
 
-    Await.result(
-      worker.run()(ExecutionContext.global),
-      scala.concurrent.duration.Duration.Inf
-    )
+    try {
+      Await.result(
+        worker.run()(ExecutionContext.global),
+        scala.concurrent.duration.Duration.Inf
+      )
+    } catch {
+      case _: io.grpc.StatusRuntimeException => System.exit(0)
+      case _: Exception => System.exit(1)
+    }
   }
 
 }