From a37b736e6526d82e43306c1bccef2f7f19f91fd9 Mon Sep 17 00:00:00 2001
From: vapao <live1989@foxmail.com>
Date: Thu, 25 Nov 2021 17:25:49 +0800
Subject: [PATCH] =?UTF-8?q?U=20=E4=BC=98=E5=8C=96=E8=87=AA=E5=AE=9A?=
 =?UTF-8?q?=E4=B9=89=E5=8F=91=E5=B8=83=E6=95=B0=E6=8D=AE=E6=B8=85=E7=90=86?=
 =?UTF-8?q?=E6=9C=BA=E5=88=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 spug_api/apps/deploy/helper.py | 7 +++++++
 spug_api/apps/deploy/utils.py  | 6 ++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/spug_api/apps/deploy/helper.py b/spug_api/apps/deploy/helper.py
index 75977d4..183244d 100644
--- a/spug_api/apps/deploy/helper.py
+++ b/spug_api/apps/deploy/helper.py
@@ -20,6 +20,7 @@ class Helper:
         self.rds = rds
         self.key = key
         self.rds.delete(self.key)
+        self.callback = []
 
     @classmethod
     def _make_dd_notify(cls, url, action, req, version, host_str):
@@ -201,6 +202,9 @@ class Helper:
             else:
                 raise NotImplementedError
 
+    def add_callback(self, func):
+        self.callback.append(func)
+
     def parse_filter_rule(self, data: str, sep='\n'):
         data, files = data.strip(), []
         if data:
@@ -230,6 +234,9 @@ class Helper:
         # save logs for two weeks
         self.rds.expire(self.key, 14 * 24 * 60 * 60)
         self.rds.close()
+        # callback
+        for func in self.callback:
+            func()
 
     def local(self, command, env=None):
         if env:
diff --git a/spug_api/apps/deploy/utils.py b/spug_api/apps/deploy/utils.py
index 2a7cefa..e3afca6 100644
--- a/spug_api/apps/deploy/utils.py
+++ b/spug_api/apps/deploy/utils.py
@@ -11,6 +11,7 @@ from apps.repository.models import Repository
 from apps.repository.utils import dispatch as build_repository
 from apps.deploy.helper import Helper, SpugError
 from concurrent import futures
+from functools import partial
 import json
 import uuid
 import os
@@ -128,7 +129,6 @@ def _ext2_deploy(req, helper, env):
         step += 1
     helper.send_step('local', 100, '')
 
-    tmp_transfer_file = None
     for action in host_actions:
         if action.get('type') == 'transfer':
             if action.get('src_mode') == '1':
@@ -157,7 +157,7 @@ def _ext2_deploy(req, helper, env):
             tar_gz_file = f'{req.spug_version}.tar.gz'
             helper.local(f'cd {sp_dir} && tar -zcf {tar_gz_file} {exclude} {contain}')
             helper.send_info('local', f'{human_time()} \033[32m完成√\033[0m\r\n')
-            tmp_transfer_file = os.path.join(sp_dir, tar_gz_file)
+            helper.add_callback(partial(os.remove, os.path.join(sp_dir, tar_gz_file)))
             break
     if host_actions:
         if req.deploy.is_parallel:
@@ -175,8 +175,6 @@ def _ext2_deploy(req, helper, env):
                         latest_exception = exception
                         if not isinstance(exception, SpugError):
                             helper.send_error(t.h_id, f'Exception: {exception}', False)
-                if tmp_transfer_file:
-                    os.remove(tmp_transfer_file)
             if latest_exception:
                 raise latest_exception
         else: