From 1cfd679e16c4b391c95001c39276a9ccc3b0b2b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Test=C3=A9?= <david.teste@zama.ai>
Date: Tue, 10 Dec 2024 14:38:04 +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 | 45 +++++++++++++++++++++++++++++++++++++++++++++
 src/index.js  | 17 +++++++++++++++++
 src/slab.js   | 28 ++++++++++++++++++++++++++++
 3 files changed, 90 insertions(+)

diff --git a/dist/index.js b/dist/index.js
index 4eb12a6..8bb3f90 100644
--- a/dist/index.js
+++ b/dist/index.js
@@ -49835,6 +49835,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 +49903,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 +51883,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 +51912,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..71e4261 100644
--- a/src/slab.js
+++ b/src/slab.js
@@ -127,6 +127,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 +195,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,