batch_image_cleanup.go 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "strings"
  6. "sync"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  11. "go.uber.org/zap"
  12. )
  13. const (
  14. defaultBatchImageInputRetentionAfterTerminal = 24 * time.Hour
  15. defaultBatchImageOutputRetentionAfterTerminal = 72 * time.Hour
  16. defaultBatchImageCleanupInterval = 30 * time.Minute
  17. defaultBatchImageCleanupBatchSize = 100
  18. )
  19. type BatchImageCleanupService struct {
  20. Repo BatchImageRepository
  21. ProviderRegistry *BatchImageProviderRegistry
  22. AccountResolver BatchImageAccountResolver
  23. Config *config.Config
  24. cancel context.CancelFunc
  25. done chan struct{}
  26. mu sync.Mutex
  27. }
  28. func NewBatchImageCleanupService(repo BatchImageRepository, accountRepo AccountRepository, cfg *config.Config) *BatchImageCleanupService {
  29. return &BatchImageCleanupService{
  30. Repo: repo,
  31. ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg),
  32. AccountResolver: &BatchImageAccountRepositoryResolver{Repo: accountRepo},
  33. Config: cfg,
  34. }
  35. }
  36. // appendCleanupEvent 追加清理审计事件;事件写入失败不阻断清理流程,但必须留痕。
  37. func (s *BatchImageCleanupService) appendCleanupEvent(ctx context.Context, batchID, eventType string, payload any) {
  38. if err := s.Repo.AppendBatchImageEvent(ctx, batchID, eventType, payload); err != nil {
  39. logger.L().Warn("batch_image.cleanup_event_failed",
  40. zap.String("batch_id", batchID),
  41. zap.String("event_type", eventType),
  42. zap.Error(err),
  43. )
  44. }
  45. }
  46. func (s *BatchImageCleanupService) DeleteOutputsForOwner(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImagePublicBatch, error) {
  47. job, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  48. if err != nil {
  49. return nil, err
  50. }
  51. if job.Status == BatchImageJobStatusOutputDeleted || job.OutputDeletedAt != nil {
  52. return BatchImageJobToPublic(job), nil
  53. }
  54. if job.Status != BatchImageJobStatusCompleted {
  55. return nil, ErrBatchImageOutputDeleteNotReady
  56. }
  57. s.appendCleanupEvent(ctx, job.BatchID, "manual_output_delete_requested", map[string]any{
  58. "batch_id": job.BatchID,
  59. "cleanup_target": "output",
  60. "reason": "manual",
  61. })
  62. if err := s.cleanupJob(ctx, job, CleanupTargetOutput, "manual"); err != nil {
  63. return nil, err
  64. }
  65. updated, err := s.Repo.GetBatchImageJobByBatchIDForOwner(ctx, owner.UserID, owner.APIKeyID, batchID)
  66. if err != nil {
  67. return nil, err
  68. }
  69. return BatchImageJobToPublic(updated), nil
  70. }
  71. func (s *BatchImageCleanupService) CleanupInput(ctx context.Context, batchID string) error {
  72. job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID)
  73. if err != nil {
  74. return err
  75. }
  76. return s.cleanupJob(ctx, job, CleanupTargetInput, "ttl")
  77. }
  78. func (s *BatchImageCleanupService) CleanupOutput(ctx context.Context, batchID string, reason string) error {
  79. job, err := s.Repo.GetBatchImageJobByBatchID(ctx, batchID)
  80. if err != nil {
  81. return err
  82. }
  83. return s.cleanupJob(ctx, job, CleanupTargetOutput, reason)
  84. }
  85. func (s *BatchImageCleanupService) RunOnce(ctx context.Context, now time.Time) (BatchImageCleanupRunResult, error) {
  86. if s == nil || s.Repo == nil {
  87. return BatchImageCleanupRunResult{}, ErrBatchImageCleanupFailed
  88. }
  89. if now.IsZero() {
  90. now = time.Now()
  91. }
  92. limit := s.cleanupBatchSize()
  93. result := BatchImageCleanupRunResult{}
  94. inputCutoff := now.Add(-s.inputRetentionAfterTerminal())
  95. inputJobs, err := s.Repo.ListBatchImageJobsDueForInputCleanup(ctx, inputCutoff, limit)
  96. if err != nil {
  97. return result, err
  98. }
  99. for _, job := range inputJobs {
  100. if job == nil {
  101. continue
  102. }
  103. if err := s.cleanupJob(ctx, job, CleanupTargetInput, "ttl"); err != nil {
  104. result.Failures++
  105. continue
  106. }
  107. result.InputCleaned++
  108. }
  109. outputJobs, err := s.Repo.ListBatchImageJobsDueForOutputCleanup(ctx, now, limit)
  110. if err != nil {
  111. return result, err
  112. }
  113. for _, job := range outputJobs {
  114. if job == nil {
  115. continue
  116. }
  117. if err := s.cleanupJob(ctx, job, CleanupTargetOutput, "expired"); err != nil {
  118. result.Failures++
  119. continue
  120. }
  121. result.OutputCleaned++
  122. }
  123. return result, nil
  124. }
  125. func (s *BatchImageCleanupService) Start() {
  126. if s == nil || s.Repo == nil || s.Config == nil || !s.Config.BatchImage.Enabled || s.cleanupInterval() <= 0 {
  127. return
  128. }
  129. s.mu.Lock()
  130. defer s.mu.Unlock()
  131. if s.cancel != nil {
  132. return
  133. }
  134. ctx, cancel := context.WithCancel(context.Background())
  135. s.cancel = cancel
  136. s.done = make(chan struct{})
  137. go func() {
  138. defer close(s.done)
  139. ticker := time.NewTicker(s.cleanupInterval())
  140. defer ticker.Stop()
  141. for {
  142. _, _ = s.RunOnce(ctx, time.Now())
  143. select {
  144. case <-ctx.Done():
  145. return
  146. case <-ticker.C:
  147. }
  148. }
  149. }()
  150. }
  151. func (s *BatchImageCleanupService) Stop() {
  152. if s == nil {
  153. return
  154. }
  155. s.mu.Lock()
  156. cancel := s.cancel
  157. done := s.done
  158. s.cancel = nil
  159. s.done = nil
  160. s.mu.Unlock()
  161. if cancel != nil {
  162. cancel()
  163. }
  164. if done != nil {
  165. <-done
  166. }
  167. }
  168. func (s *BatchImageCleanupService) cleanupJob(ctx context.Context, job *BatchImageJob, target CleanupTarget, reason string) error {
  169. if job == nil {
  170. return ErrBatchImageJobNotFound
  171. }
  172. switch target {
  173. case CleanupTargetInput:
  174. if job.InputDeletedAt != nil {
  175. return nil
  176. }
  177. if !IsTerminalBatchImageJobStatus(job.Status) {
  178. return ErrBatchImageCleanupFailed
  179. }
  180. s.appendCleanupEvent(ctx, job.BatchID, "input_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
  181. case CleanupTargetOutput:
  182. if job.OutputDeletedAt != nil || job.Status == BatchImageJobStatusOutputDeleted {
  183. return nil
  184. }
  185. if job.Status != BatchImageJobStatusCompleted && job.Status != BatchImageJobStatusFailed && job.Status != BatchImageJobStatusCancelled {
  186. return ErrBatchImageOutputDeleteNotReady
  187. }
  188. s.appendCleanupEvent(ctx, job.BatchID, "output_cleanup_started", cleanupEventPayload(job.BatchID, target, reason, nil))
  189. default:
  190. return ErrUnsupportedCleanupTarget
  191. }
  192. if err := s.callProviderCleanup(ctx, job, target); err != nil {
  193. code := cleanupFailureCode(err)
  194. msg := sanitizeBatchImagePublicMessage(err.Error())
  195. if recordErr := s.Repo.RecordBatchImageCleanupFailure(ctx, job.BatchID, code, msg); recordErr != nil {
  196. logger.L().Warn("batch_image.cleanup_failure_record_failed",
  197. zap.String("batch_id", job.BatchID),
  198. zap.Error(recordErr),
  199. )
  200. }
  201. event := string(target) + "_cleanup_failed"
  202. s.appendCleanupEvent(ctx, job.BatchID, event, map[string]any{"batch_id": job.BatchID, "cleanup_target": string(target), "reason": reason, "error_code": code})
  203. if errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) {
  204. return ErrBatchImageCleanupUnsafePath
  205. }
  206. return ErrBatchImageProviderCleanupFailed
  207. }
  208. deletedAt := time.Now()
  209. if target == CleanupTargetInput {
  210. return s.Repo.MarkBatchImageInputDeleted(ctx, job.BatchID, deletedAt)
  211. }
  212. return s.Repo.MarkBatchImageOutputDeleted(ctx, job.BatchID, deletedAt)
  213. }
  214. func (s *BatchImageCleanupService) callProviderCleanup(ctx context.Context, job *BatchImageJob, target CleanupTarget) error {
  215. if s == nil || s.ProviderRegistry == nil || s.AccountResolver == nil {
  216. return ErrBatchImageCleanupFailed
  217. }
  218. provider, ok := s.ProviderRegistry.Get(job.Provider)
  219. if !ok || provider == nil {
  220. return ErrBatchImageUnsupportedProvider
  221. }
  222. if job.AccountID == nil || *job.AccountID <= 0 {
  223. return ErrBatchImageMissingAccountID
  224. }
  225. account, err := s.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID)
  226. if err != nil {
  227. return err
  228. }
  229. if err := provider.Cleanup(ctx, job, account, target); err != nil {
  230. if cleanupErrorIsNotFound(err) {
  231. return nil
  232. }
  233. return err
  234. }
  235. return nil
  236. }
  237. func (s *BatchImageCleanupService) inputRetentionAfterTerminal() time.Duration {
  238. if s != nil && s.Config != nil && s.Config.BatchImage.InputRetentionAfterTerminalHours > 0 {
  239. return time.Duration(s.Config.BatchImage.InputRetentionAfterTerminalHours) * time.Hour
  240. }
  241. return defaultBatchImageInputRetentionAfterTerminal
  242. }
  243. func (s *BatchImageCleanupService) cleanupInterval() time.Duration {
  244. if s != nil && s.Config != nil && s.Config.BatchImage.CleanupIntervalMinutes > 0 {
  245. return time.Duration(s.Config.BatchImage.CleanupIntervalMinutes) * time.Minute
  246. }
  247. return defaultBatchImageCleanupInterval
  248. }
  249. func (s *BatchImageCleanupService) cleanupBatchSize() int {
  250. if s != nil && s.Config != nil && s.Config.BatchImage.CleanupBatchSize > 0 {
  251. return s.Config.BatchImage.CleanupBatchSize
  252. }
  253. return defaultBatchImageCleanupBatchSize
  254. }
  255. type BatchImageCleanupRunResult struct {
  256. InputCleaned int
  257. OutputCleaned int
  258. Failures int
  259. }
  260. func cleanupEventPayload(batchID string, target CleanupTarget, reason string, deletedAt *time.Time) map[string]any {
  261. payload := map[string]any{
  262. "batch_id": batchID,
  263. "cleanup_target": string(target),
  264. "reason": reason,
  265. }
  266. if deletedAt != nil {
  267. payload["deleted_at"] = deletedAt.UTC().Format(time.RFC3339)
  268. }
  269. return payload
  270. }
  271. func cleanupErrorIsNotFound(err error) bool {
  272. if err == nil {
  273. return false
  274. }
  275. reason := strings.ToUpper(infraerrors.Reason(err))
  276. msg := strings.ToUpper(err.Error())
  277. return strings.Contains(reason, "NOT_FOUND") || strings.Contains(msg, "NOT FOUND") || strings.Contains(msg, "404")
  278. }
  279. func cleanupFailureCode(err error) string {
  280. if errors.Is(err, ErrBatchImageProviderUnsafeCleanupPath) {
  281. return "BATCH_IMAGE_CLEANUP_UNSAFE_PATH"
  282. }
  283. reason := strings.TrimSpace(infraerrors.Reason(err))
  284. if reason != "" {
  285. return reason
  286. }
  287. return "BATCH_IMAGE_PROVIDER_CLEANUP_FAILED"
  288. }