shudao-main/shudao-chat-go/models/chroma.go

package models

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/beego/beego/v2/server/web"
)

// HeartbeatTask 心跳任务结构体
type HeartbeatTask struct {
	URL        string
	Interval   time.Duration
	HTTPClient *http.Client
	ctx        context.Context
	cancel     context.CancelFunc
}

// NewHeartbeatTask 创建新的心跳任务
func NewHeartbeatTask(url string, interval time.Duration) *HeartbeatTask {
	ctx, cancel := context.WithCancel(context.Background())
	return &HeartbeatTask{
		URL:        url,
		Interval:   interval,
		HTTPClient: &http.Client{Timeout: 30 * time.Second},
		ctx:        ctx,
		cancel:     cancel,
	}
}

// Start 启动心跳任务
func (h *HeartbeatTask) Start() {
	go h.run()
}

// Stop 停止心跳任务
func (h *HeartbeatTask) Stop() {
	h.cancel()
}

// run 运行心跳任务
func (h *HeartbeatTask) run() {
	ticker := time.NewTicker(h.Interval)
	defer ticker.Stop()

	// 立即执行一次
	h.sendHeartbeat()

	for {
		select {
		case <-h.ctx.Done():
			fmt.Println("心跳任务已停止")
			return
		case <-ticker.C:
			h.sendHeartbeat()
		}
	}
}

// sendHeartbeat 发送心跳请求
func (h *HeartbeatTask) sendHeartbeat() {
	req, err := http.NewRequestWithContext(h.ctx, "GET", h.URL, nil)
	if err != nil {
		fmt.Printf("创建心跳请求失败: %v\n", err)
		return
	}

	resp, err := h.HTTPClient.Do(req)
	if err != nil {
		fmt.Printf("心跳请求失败: %v\n", err)
		return
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		fmt.Printf("读取心跳响应失败: %v\n", err)
		return
	}

	if resp.StatusCode == http.StatusOK {
		fmt.Printf("心跳成功 [%s] - 状态码: %d, 响应: %s\n",
			time.Now().Format("2006-01-02 15:04:05"), resp.StatusCode, string(body))
	} else {
		fmt.Printf("心跳失败 [%s] - 状态码: %d, 响应: %s\n",
			time.Now().Format("2006-01-02 15:04:05"), resp.StatusCode, string(body))
	}
}

// StartHeartbeatTask 启动心跳任务（从配置中读取URL）
func StartHeartbeatTask() {
	// 从配置文件中读取心跳API URL
	heartbeatURL, err := web.AppConfig.String("heartbeat_api_url")
	if err != nil || heartbeatURL == "" {
		fmt.Println("未配置心跳API URL，跳过心跳任务")
		return
	}

	// 创建心跳任务，每10分钟执行一次
	heartbeatTask := NewHeartbeatTask(heartbeatURL, 10*time.Minute)

	fmt.Printf("启动心跳任务，目标URL: %s，间隔: 10分钟\n", heartbeatURL)
	heartbeatTask.Start()
}