From dd30ab2e8da3c3bee7cc5f81d9e8c6a90c6bb54a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Test=C3=A9?= <david.teste@zama.ai>
Date: Wed, 11 Dec 2024 16:27:32 +0100
Subject: [PATCH] feat(action): handle user cancellation

This is useful in start mode where one can cancel before the
instance is actually ready. This way a best-effort is produced to
clean up any spawned instance instead of it turning into a zombie.
---
 dist/index.js | 48 +++++++++++++++++++++++++++++++++++++++++++++++-
 src/index.js  | 17 +++++++++++++++++
 src/slab.js   | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 94 insertions(+), 2 deletions(-)

diff --git a/dist/index.js b/dist/index.js
index 4eb12a6..7331446 100644
--- a/dist/index.js
+++ b/dist/index.js
@@ -49733,7 +49733,8 @@ async function startInstanceRequest() {
   const details = {
     backend: {
       provider,
-      profile: config.input.profile
+      profile: config.input.profile,
+      create_watchdog_task: true
     }
   }
 
@@ -49835,6 +49836,9 @@ async function waitForInstance(taskId, taskName) {
       const task_status = body[taskName].status.toLowerCase()
 
       if (task_status === 'done') {
+        if (taskName === 'start') {
+          await acknowledgeTaskDone(taskId)
+        }
         await removeTask(taskId)
         return body
       } else if (task_status === 'failed') {
@@ -49900,6 +49904,31 @@ async function removeTask(taskId) {
   }
 }
 
+async function acknowledgeTaskDone(taskId) {
+  const url = config.input.slabUrl
+  const route = `task_ack_done/${config.githubContext.repo}/${taskId}`
+  let response
+
+  try {
+    response = await fetch(concat_path(url, route), {
+      method: 'POST'
+    })
+  } catch (error) {
+    core.error(`Failed to acknowledge task done with ID: ${taskId}`)
+    throw error
+  }
+
+  if (response.ok) {
+    core.debug('Instance task successfully acknowledged')
+    return response
+  } else {
+    core.error(
+      `Instance task acknowledgment request has failed (ID: ${taskId}, HTTP status code: ${response.status})`
+    )
+    throw new Error('task acknowledging failed')
+  }
+}
+
 module.exports = {
   startInstanceRequest,
   stopInstanceRequest,
@@ -51855,11 +51884,27 @@ const slab = __nccwpck_require__(4156)
 const config = __nccwpck_require__(4570)
 const core = __nccwpck_require__(2186)
 const { waitForRunnerRegistered } = __nccwpck_require__(6989)
+const utils = __nccwpck_require__(1608)
 
 function setOutput(label) {
   core.setOutput('label', label)
 }
 
+// This variable should only be defined for cleanup purpose.
+let runner_name
+
+async function cleanup() {
+  if (runner_name) {
+    core.info('Stop instance after cancellation')
+    await slab.stopInstanceRequest(runner_name)
+  }
+}
+
+process.on('SIGINT', async function () {
+  await cleanup()
+  process.exit()
+})
+
 async function start() {
   const provider = config.input.backend
 
@@ -51868,6 +51913,7 @@ async function start() {
   for (let i = 1; i <= 3; i++) {
     try {
       start_instance_response = await slab.startInstanceRequest()
+      runner_name = start_instance_response.runner_name
       break
     } catch (error) {
       core.info('Retrying request now...')
diff --git a/src/index.js b/src/index.js
index d5f5327..b3daf6a 100644
--- a/src/index.js
+++ b/src/index.js
@@ -2,11 +2,27 @@ const slab = require('./slab')
 const config = require('./config')
 const core = require('@actions/core')
 const { waitForRunnerRegistered } = require('./gh')
+const utils = require('./utils')
 
 function setOutput(label) {
   core.setOutput('label', label)
 }
 
+// This variable should only be defined for cleanup purpose.
+let runner_name
+
+async function cleanup() {
+  if (runner_name) {
+    core.info('Stop instance after cancellation')
+    await slab.stopInstanceRequest(runner_name)
+  }
+}
+
+process.on('SIGINT', async function () {
+  await cleanup()
+  process.exit()
+})
+
 async function start() {
   const provider = config.input.backend
 
@@ -15,6 +31,7 @@ async function start() {
   for (let i = 1; i <= 3; i++) {
     try {
       start_instance_response = await slab.startInstanceRequest()
+      runner_name = start_instance_response.runner_name
       break
     } catch (error) {
       core.info('Retrying request now...')
diff --git a/src/slab.js b/src/slab.js
index e07586b..b8e84bd 100644
--- a/src/slab.js
+++ b/src/slab.js
@@ -25,7 +25,8 @@ async function startInstanceRequest() {
   const details = {
     backend: {
       provider,
-      profile: config.input.profile
+      profile: config.input.profile,
+      create_watchdog_task: true
     }
   }
 
@@ -127,6 +128,9 @@ async function waitForInstance(taskId, taskName) {
       const task_status = body[taskName].status.toLowerCase()
 
       if (task_status === 'done') {
+        if (taskName === 'start') {
+          await acknowledgeTaskDone(taskId)
+        }
         await removeTask(taskId)
         return body
       } else if (task_status === 'failed') {
@@ -192,6 +196,31 @@ async function removeTask(taskId) {
   }
 }
 
+async function acknowledgeTaskDone(taskId) {
+  const url = config.input.slabUrl
+  const route = `task_ack_done/${config.githubContext.repo}/${taskId}`
+  let response
+
+  try {
+    response = await fetch(concat_path(url, route), {
+      method: 'POST'
+    })
+  } catch (error) {
+    core.error(`Failed to acknowledge task done with ID: ${taskId}`)
+    throw error
+  }
+
+  if (response.ok) {
+    core.debug('Instance task successfully acknowledged')
+    return response
+  } else {
+    core.error(
+      `Instance task acknowledgment request has failed (ID: ${taskId}, HTTP status code: ${response.status})`
+    )
+    throw new Error('task acknowledging failed')
+  }
+}
+
 module.exports = {
   startInstanceRequest,
   stopInstanceRequest,