xg_dba/internal/middleware/ProgressManager.go

package middleware

import (
	"sync"

	"github.com/google/uuid"
)

// ProgressManager 用于管理所有请求的进度
type ProgressManager struct {
	progressMap map[string]chan string
	mu          sync.RWMutex
}

func NewProgressManager() *ProgressManager {
	return &ProgressManager{
		progressMap: make(map[string]chan string),
	}
}

func (pm *ProgressManager) CreateProgress() string {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	// 生成一个唯一的请求 ID
	id := uuid.New().String()
	// 创建一个新的 progress channel
	pm.progressMap[id] = make(chan string)
	return id
}

func (pm *ProgressManager) GetProgressChan(id string) (chan string, bool) {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	progressChan, ok := pm.progressMap[id]
	return progressChan, ok
}

func (pm *ProgressManager) CloseProgress(id string) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	if progressChan, ok := pm.progressMap[id]; ok {
		close(progressChan)
		delete(pm.progressMap, id)
	}
}