From 84c3d4dc3dd9b63b3341b31d8e2d751f0e7f3192 Mon Sep 17 00:00:00 2001
From: uubulb <uub@kuzu.uk>
Date: Fri, 20 Dec 2024 02:11:23 +0800
Subject: [PATCH] cleanup

---
 model/common.go           | 16 ++++++++++++
 service/singleton/user.go | 55 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/model/common.go b/model/common.go
index 05344f2..d70317b 100644
--- a/model/common.go
+++ b/model/common.go
@@ -27,6 +27,10 @@ func (c *Common) GetID() uint64 {
 	return c.ID
 }
 
+func (c *Common) GetUserID() uint64 {
+	return c.UserID
+}
+
 func (c *Common) HasPermission(ctx *gin.Context) bool {
 	auth, ok := ctx.Get(CtxKeyAuthorizedUser)
 	if !ok {
@@ -43,9 +47,21 @@ func (c *Common) HasPermission(ctx *gin.Context) bool {
 
 type CommonInterface interface {
 	GetID() uint64
+	GetUserID() uint64
 	HasPermission(*gin.Context) bool
 }
 
+func FindUserID[S ~[]E, E CommonInterface](s S, uid uint64) []uint64 {
+	var list []uint64
+	for _, v := range s {
+		if v.GetUserID() == uid {
+			list = append(list, v.GetID())
+		}
+	}
+
+	return list
+}
+
 type Response struct {
 	Code    int         `json:"code,omitempty"`
 	Message string      `json:"message,omitempty"`
diff --git a/service/singleton/user.go b/service/singleton/user.go
index 763344e..88c333d 100644
--- a/service/singleton/user.go
+++ b/service/singleton/user.go
@@ -4,6 +4,7 @@ import (
 	"sync"
 
 	"github.com/nezhahq/nezha/model"
+	"gorm.io/gorm"
 )
 
 var (
@@ -46,9 +47,63 @@ func OnUserDelete(id []uint64) {
 		return
 	}
 
+	var (
+		cron   bool
+		server bool
+	)
+
 	for _, uid := range id {
 		secret := UserIdToAgentSecret[uid]
 		delete(AgentSecretToUserId, secret)
 		delete(UserIdToAgentSecret, uid)
+
+		CronLock.RLock()
+		crons := model.FindUserID(CronList, uid)
+		CronLock.RUnlock()
+
+		cron = len(crons) > 0
+		if cron {
+			DB.Unscoped().Delete(&model.Cron{}, "id in (?)", crons)
+			OnDeleteCron(crons)
+		}
+
+		SortedServerLock.RLock()
+		servers := model.FindUserID(SortedServerList, uid)
+		SortedServerLock.RUnlock()
+
+		server = len(servers) > 0
+		if server {
+			DB.Transaction(func(tx *gorm.DB) error {
+				if err := tx.Unscoped().Delete(&model.Server{}, "id in (?)", servers).Error; err != nil {
+					return err
+				}
+				if err := tx.Unscoped().Delete(&model.ServerGroupServer{}, "server_id in (?)", servers).Error; err != nil {
+					return err
+				}
+				return nil
+			})
+
+			AlertsLock.Lock()
+			for _, sid := range servers {
+				for _, alert := range Alerts {
+					if AlertsCycleTransferStatsStore[alert.ID] != nil {
+						delete(AlertsCycleTransferStatsStore[alert.ID].ServerName, sid)
+						delete(AlertsCycleTransferStatsStore[alert.ID].Transfer, sid)
+						delete(AlertsCycleTransferStatsStore[alert.ID].NextUpdate, sid)
+					}
+				}
+			}
+			DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers)
+			AlertsLock.Unlock()
+			OnServerDelete(servers)
+		}
+	}
+
+	if cron {
+		UpdateCronList()
+	}
+
+	if server {
+		ReSortServer()
 	}
 }