From a3d40299374ce25ef56749249e8ea49c89de6f19 Mon Sep 17 00:00:00 2001
From: codyjlin <31944154+codyjlin@users.noreply.github.com>
Date: Tue, 6 Apr 2021 15:55:50 -0400
Subject: [PATCH] Add graceful shutdown via sigterm (#42)

* Add graceful shutdown via sigterm

* Add documentation

* minor fixes; search port number

* Add documentation and fix CI

* fix test

Co-authored-by: Cody Lin <clin@ibm.com>
---
 Dockerfile       |  2 ++
 README.md        |  6 ++++++
 run-presto.sh    | 23 +++++++++++++++++++++++
 test_config.yaml |  2 +-
 4 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100755 run-presto.sh

diff --git a/Dockerfile b/Dockerfile
index 385910f..47143e7 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -32,3 +32,5 @@ USER root
 USER presto:presto
 # Add Db2 connector
 COPY --from=builder --chown=presto:presto presto-db2-* /usr/lib/presto/plugin/db2
+
+CMD ["./run-presto.sh"]
diff --git a/README.md b/README.md
index 71cb027..d110693 100644
--- a/README.md
+++ b/README.md
@@ -40,3 +40,9 @@ Then:
 ```SHELL
 docker run -d -p 8080:8080 -v /foo/bar/db2.properties:/usr/lib/presto/default/etc/catalog/db2.properties:ro shawnzhu/prestodb:latest
 ```
+
+## Features
+
+### Graceful Shutdown
+
+It adds the [graceful shutdown](https://trino.io/docs/current/admin/graceful-shutdown.html) feature from Trino such that on a SIGTERM signal to the container, the worker will have a grace period before interrupting active queries.
diff --git a/run-presto.sh b/run-presto.sh
new file mode 100755
index 0000000..bcffa27
--- /dev/null
+++ b/run-presto.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+set -xeuo pipefail
+
+graceful_shutdown() {
+    echo "calling graceful shutdown"
+    PORT=$(grep "^http-server.http.port=" /etc/presto/config.properties | cut -d '=' -f2)
+    curl -v -X PUT -d '"SHUTTING_DOWN"' -H "Content-type: application/json" http://localhost:$PORT\/v1/info/state -H X-Presto-User:name
+}
+trap graceful_shutdown TERM
+
+set +e
+grep -s -q 'node.id' /etc/presto/node.properties
+NODE_ID_EXISTS=$?
+set -e
+
+NODE_ID=""
+if [[ ${NODE_ID_EXISTS} != 0 ]] ; then
+    NODE_ID="-Dnode.id=${HOSTNAME}"
+fi
+
+exec /usr/lib/presto/bin/launcher run --etc-dir /etc/presto ${NODE_ID} "$@"
+
diff --git a/test_config.yaml b/test_config.yaml
index c953337..241c805 100644
--- a/test_config.yaml
+++ b/test_config.yaml
@@ -19,4 +19,4 @@ metadataTest:
     value: /usr/lib/jvm/zulu11
   exposedPorts: ["8080"]
   cmd:
-  - /usr/lib/presto/bin/run-presto
+  - ./run-presto.sh