From 64deb61f15d0832192af7f0f446997dfc3eda1b6 Mon Sep 17 00:00:00 2001
From: EGOIST <0x142857@gmail.com>
Date: Sat, 30 Jul 2022 05:32:07 +0000
Subject: [PATCH] fix: return a function in `onSuccess` to do cleanup

---
 docs/README.md | 25 +++++++++++++++++++++++--
 src/index.ts   | 35 +++++++++++++++--------------------
 src/options.ts |  4 +++-
 3 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/docs/README.md b/docs/README.md
index d738cb387..0e4a7679c 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -331,13 +331,34 @@ tsup src/index.ts --watch --onSuccess "node dist/index.js"
 
 > Warning: You should not use shell scripts, if you need to specify shell scripts you can add it in your "scripts" field and set for example `tsup src/index.ts --watch --onSuccess \"npm run dev\"`
 
-`onSuccess` can also be a `function` that returns `Promise`. For this to work, you need to use `tsup.config.ts` instead of the cli flag: 
+`onSuccess` can also be a `function` that returns `Promise`. For this to work, you need to use `tsup.config.ts` instead of the cli flag:
 
 ```ts
 import { defineConfig } from 'tsup'
 
 export default defineConfig({
-  onSuccess: async () => { ... }
+  async onSuccess() {
+    // Start some long running task
+    // Like a server
+  },
+})
+```
+
+You can return a cleanup function in `onSuccess`:
+
+```ts
+import { defineConfig } from 'tsup'
+
+export default defineConfig({
+  async onSuccess() {
+    const server = http.createServer((req, res) => {
+      res.end('Hello World!')
+    })
+    server.listen(3000)
+    return () => {
+      server.close()
+    }
+  },
 })
 ```
 
diff --git a/src/index.ts b/src/index.ts
index d6824c76f..fe4695773 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -179,29 +179,24 @@ export async function build(_options: Options) {
           }
         }
 
-        const otherTasks = async () => {
+        const mainTasks = async () => {
           if (!options.dts?.only) {
-            let existingOnSuccess: ChildProcess | undefined
-            let existingOnSuccessFnPromise: Promise<any> | undefined
+            let onSuccessProcess: ChildProcess | undefined
+            let onSuccessCleanup: (() => any) | undefined | void
             /** Files imported by the entry */
             const buildDependencies: Set<string> = new Set()
 
-            const killPreviousProcessOrPromise = async () => {
-              if (existingOnSuccess) {
+            const doOnSuccessCleanup = async () => {
+              if (onSuccessProcess) {
                 await killProcess({
-                  pid: existingOnSuccess.pid,
+                  pid: onSuccessProcess.pid,
                 })
-              } else if (existingOnSuccessFnPromise) {
-                await Promise.race([
-                  existingOnSuccessFnPromise,
-                  // cancel existingOnSuccessFnPromise if it is still running,
-                  // using a promise that's been already resolved
-                  Promise.resolve(),
-                ])
+              } else if (onSuccessCleanup) {
+                await onSuccessCleanup()
               }
               // reset them in all occassions anyway
-              existingOnSuccess = undefined
-              existingOnSuccessFnPromise = undefined
+              onSuccessProcess = undefined
+              onSuccessCleanup = undefined
             }
 
             const debouncedBuildAll = debouncePromise(
@@ -213,7 +208,7 @@ export async function build(_options: Options) {
             )
 
             const buildAll = async () => {
-              const killPromise = killPreviousProcessOrPromise()
+              await doOnSuccessCleanup()
               // Store previous build dependencies in case the build failed
               // So we can restore it
               const previousBuildDependencies = new Set(buildDependencies)
@@ -258,12 +253,12 @@ export async function build(_options: Options) {
                   })
                 }),
               ])
-              await killPromise
+
               if (options.onSuccess) {
                 if (typeof options.onSuccess === 'function') {
-                  existingOnSuccessFnPromise = options.onSuccess()
+                  onSuccessCleanup = await options.onSuccess()
                 } else {
-                  existingOnSuccess = execa(options.onSuccess, {
+                  onSuccessProcess = execa(options.onSuccess, {
                     shell: true,
                     stdio: 'inherit',
                   })
@@ -338,7 +333,7 @@ export async function build(_options: Options) {
           }
         }
 
-        await Promise.all([dtsTask(), otherTasks()])
+        await Promise.all([dtsTask(), mainTasks()])
       }
     )
   )
diff --git a/src/options.ts b/src/options.ts
index bdec9dcad..4fd71a028 100644
--- a/src/options.ts
+++ b/src/options.ts
@@ -72,7 +72,9 @@ export type Options = {
   keepNames?: boolean
   watch?: boolean | string | (string | boolean)[]
   ignoreWatch?: string[] | string
-  onSuccess?: string | ((...params: any[]) => Promise<any>),
+  onSuccess?:
+    | string
+    | (() => Promise<void | undefined | (() => void | Promise<void>)>)
   jsxFactory?: string
   jsxFragment?: string
   outDir?: string