From d9e24486d21decd9c47d79afbee8c8ff5ed5b37c Mon Sep 17 00:00:00 2001
From: Chuan Xu <xuchuan@xuchuan.net>
Date: Mon, 29 Aug 2022 19:15:40 +0800
Subject: [PATCH] enhance(client): use destructor instead of atexit

---
 client/starwhale/api/_impl/data_store.py | 10 ++++++----
 client/tests/sdk/test_base.py            |  6 ------
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/client/starwhale/api/_impl/data_store.py b/client/starwhale/api/_impl/data_store.py
index 5e7a403b08..0a35c9f05e 100644
--- a/client/starwhale/api/_impl/data_store.py
+++ b/client/starwhale/api/_impl/data_store.py
@@ -2,7 +2,6 @@
 import re
 import sys
 import json
-import atexit
 import base64
 import struct
 import urllib
@@ -713,7 +712,6 @@ def get_instance() -> "LocalDataStore":
                 ensure_dir(ds_path)
 
                 LocalDataStore._instance = LocalDataStore(str(ds_path))
-                atexit.register(LocalDataStore._instance.dump)
             return LocalDataStore._instance
 
     def __init__(self, root_path: str) -> None:
@@ -721,6 +719,9 @@ def __init__(self, root_path: str) -> None:
         self.name_pattern = re.compile(r"^[A-Za-z0-9-_/: ]+$")
         self.tables: Dict[str, MemoryTable] = {}
 
+    def __del__(self) -> None:
+        self.dump()
+
     def update_table(
         self,
         table_name: str,
@@ -1028,7 +1029,6 @@ def __init__(self, table_name: str, key_column: str = "id") -> None:
         self.data_store = get_data_store()
         self.cond = threading.Condition()
         self.setDaemon(True)
-        atexit.register(self.close)
         self.start()
 
     def __enter__(self) -> Any:
@@ -1037,10 +1037,12 @@ def __enter__(self) -> Any:
     def __exit__(self, type: Any, value: Any, tb: Any) -> None:
         self.close()
 
+    def __del__(self) -> None:
+        self.close()
+
     def close(self) -> None:
         with self.cond:
             if not self.stopped:
-                atexit.unregister(self.close)
                 self.stopped = True
                 self.cond.notify()
         self.join()
diff --git a/client/tests/sdk/test_base.py b/client/tests/sdk/test_base.py
index 96978e4232..c1272376e7 100644
--- a/client/tests/sdk/test_base.py
+++ b/client/tests/sdk/test_base.py
@@ -1,7 +1,6 @@
 import os
 import tempfile
 import unittest
-from unittest.mock import patch, MagicMock
 
 from starwhale.utils import config as sw_config
 from starwhale.consts import ENV_SW_CLI_CONFIG, ENV_SW_LOCAL_STORAGE
@@ -18,12 +17,7 @@ def setUp(self) -> None:
         self.datastore_root = str(sw_config.SWCliConfigMixed().datastore_dir)
         ensure_dir(self.datastore_root)
 
-        self.mock_atexit = patch("starwhale.api._impl.data_store.atexit", MagicMock())
-        self.mock_atexit.start()
-
     def tearDown(self) -> None:
         empty_dir(self.local_storage)
         os.environ.pop(ENV_SW_CLI_CONFIG, "")
         os.environ.pop(ENV_SW_LOCAL_STORAGE, "")
-
-        self.mock_atexit.stop()