From ae9ac8270d4ce12388148136fae7ab6bb17c220a Mon Sep 17 00:00:00 2001
From: Denis Baryshev <dennybaa@gmail.com>
Date: Fri, 17 Dec 2021 16:58:01 +0300
Subject: [PATCH] add graceful-run.sh entrypoint

Signed-off-by: Denis Baryshev <dennybaa@gmail.com>
---
 docker/prover/Dockerfile      |  4 ++--
 docker/prover/graceful-run.sh | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100755 docker/prover/graceful-run.sh

diff --git a/docker/prover/Dockerfile b/docker/prover/Dockerfile
index 9409450105..7e308c7acb 100644
--- a/docker/prover/Dockerfile
+++ b/docker/prover/Dockerfile
@@ -29,10 +29,10 @@ RUN apt-get update && apt-get install -y axel ca-certificates tini && rm -rf /va
 #ENV MISC_DOCKER_DUMMY_PROVER $MISC_DOCKER_DUMMY_PROVER
 COPY --from=builder /usr/src/zksync/target/release/plonk_step_by_step_prover /bin/
 COPY --from=builder /usr/src/zksync/target/release/dummy_prover /bin/
-COPY docker/prover/prover-entry.sh /bin/
+COPY docker/prover/prover-entry.sh docker/prover/graceful-run.sh /bin/
 COPY keys/packed /keys/packed
 COPY contracts/artifacts/ /contracts/artifacts/
 
 ## Mitigate memory leaks when app exists it's restarted
 ENTRYPOINT [ "/usr/bin/tini", "-g", "--" ]
-CMD while :; do /bin/prover-entry.sh || exit $?; done
+CMD ["/bin/graceful-run.sh"]
diff --git a/docker/prover/graceful-run.sh b/docker/prover/graceful-run.sh
new file mode 100755
index 0000000000..17d4e4d292
--- /dev/null
+++ b/docker/prover/graceful-run.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+terminating=""
+
+# exit when the second signal received
+graceful_exit() {
+  if [ -n "$terminating" ]; then
+    exit $code
+  fi
+  terminating="yes"
+}
+
+trap graceful_exit SIGINT SIGTERM SIGHUP
+
+while : ; do
+  /bin/prover-entry.sh; code=$?
+  [ "$terminating" = "yes" ] && exit $code
+  # restart prover on failure
+  [ "$code" -eq 0 ] || exit $code
+done