diff --git a/cmd/edge-agent/main.go b/cmd/edge-agent/main.go index c54e29a..4a95902 100644 --- a/cmd/edge-agent/main.go +++ b/cmd/edge-agent/main.go @@ -21,8 +21,8 @@ func main() { flag.Parse() cfg := config.Load(*cfgPath) - log.Printf("edge-agent start id=%s uuid=%s cloud=%s mqtt=%s", - cfg.EdgeID, cfg.GetDeviceIdentity(), cfg.CloudURL, cfg.MqttBroker) + log.Printf("edge-agent start id=%s uuid=%s cloud=%s mqtt=%s auto_pull=%v", + cfg.EdgeID, cfg.GetDeviceIdentity(), cfg.CloudURL, cfg.MqttBroker, cfg.AutoPull) ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer cancel() @@ -43,8 +43,17 @@ func main() { mqttMgr := control.NewMqttManager(cfg) run(mqttMgr.Run) + // 1.5 Stream Puller (Optional: Dynamic URL fetching) + var puller *stream.Puller + if cfg.AutoPull { + puller = stream.NewPuller(cfg) + puller.StartHeartbeatLoop() + defer puller.Stop() + log.Println("puller: auto-pull mode enabled") + } + // 2. Stream Ingestion (Dynamic) - run(stream.NewStreamManager(cfg, frames, mqttMgr.GetUpdates()).Run) + run(stream.NewStreamManager(cfg, puller, frames, mqttMgr.GetUpdates()).Run) // 3. Inference (Dynamic Workers) run(infer.NewClient(cfg, frames, events, mqttMgr.GetUpdates()).Run) diff --git a/config/edge.yaml b/config/edge.yaml index 7a65c61..8576ad2 100644 --- a/config/edge.yaml +++ b/config/edge.yaml @@ -5,8 +5,11 @@ mqtt_broker: tcp://101.36.73.102:1883 mqtt_user: "" mqtt_pass: "" edge_token: "" -rtsp_urls: - - http://101.36.73.102:8080/rtp/34020000002000000003_34020000001310000001.live.flv +rtsp_urls: [] +# 自动按需拉流配置 +stream_pull_url: http://101.36.73.102:9000 +stream_protocol: flv +auto_pull: false infer_socket: /tmp/edge-infer.sock infer_fps: 2 infer_workers: 3 diff --git a/internal/config/config.go b/internal/config/config.go index 1f09103..067c4c1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -19,6 +19,10 @@ type Config struct { EdgeToken string `yaml:"edge_token"` StreamEnabled bool `yaml:"stream_enabled"` RTSPURLs []string `yaml:"rtsp_urls"` + // 新增：自动按需拉流配置 + StreamPullURL string `yaml:"stream_pull_url"` // 云端拉流网关地址 + StreamProtocol string `yaml:"stream_protocol"` // 拉流协议: flv, rtsp, ws_flv + AutoPull bool `yaml:"auto_pull"` // 是否启用自动按需拉流 InferSocket string `yaml:"infer_socket"` InferFPS int `yaml:"infer_fps"` InferWorkers int `yaml:"infer_workers"` diff --git a/internal/stream/ingestor.go b/internal/stream/ingestor.go index 9b693fa..1a2ac15 100644 --- a/internal/stream/ingestor.go +++ b/internal/stream/ingestor.go @@ -29,15 +29,17 @@ type streamProcess struct { type StreamManager struct { cfg *config.Config + puller *Puller // 可选的动态拉流客户端 frames chan<- Frame processes map[string]*streamProcess // url -> {cmd, fps} mu sync.Mutex updates <-chan *config.Config } -func NewStreamManager(cfg *config.Config, frames chan<- Frame, updates <-chan *config.Config) *StreamManager { +func NewStreamManager(cfg *config.Config, puller *Puller, frames chan<- Frame, updates <-chan *config.Config) *StreamManager { return &StreamManager{ cfg: cfg, + puller: puller, frames: frames, processes: make(map[string]*streamProcess), updates: updates, @@ -45,7 +47,14 @@ func NewStreamManager(cfg *config.Config, frames chan<- Frame, updates <-chan *c } func (sm *StreamManager) Run(ctx context.Context) { - // Initial start + // 1. 动态拉流模式 + if sm.puller != nil { + log.Println("stream: running in auto-pull mode") + sm.runAutoPullMode(ctx) + return + } + + // 2. 静态配置模式 sm.applyStreams(ctx, sm.cfg.RTSPURLs, sm.cfg.InferFPS) for { @@ -60,6 +69,39 @@ func (sm *StreamManager) Run(ctx context.Context) { } } +// runAutoPullMode 动态拉流模式主循环 +func (sm *StreamManager) runAutoPullMode(ctx context.Context) { + for { + select { + case <-ctx.Done(): + sm.stopAll() + return + default: + } + + // 从云端获取拉流地址 + url, err := sm.puller.Pull() + if err != nil { + log.Printf("stream: pull failed: %v, retry in 10s", err) + time.Sleep(10 * time.Second) + continue + } + + log.Printf("stream: got dynamic url=%s, starting ingestion", url) + + // 启动该 URL 的拉流 + sm.applyStreams(ctx, []string{url}, sm.cfg.InferFPS) + + // 等待上下文结束或 URL 被移除（通常意味着需要重新 Pull） + // 在 auto-pull 模式下，如果 ffmpeg 断开，streamLoop 会自动重试， + // 但我们需要在这里监听 ctx，或者等待 puller 心跳失败后重新拉取。 + // 简单起见，这里阻塞直到 ctx 取消，streamLoop 内部会处理 ffmpeg 重连。 + // 如果 streamLoop 发现 URL 失效，会退出并删除 process。 + <-ctx.Done() + sm.stopAll() + } +} + func (sm *StreamManager) applyStreams(ctx context.Context, urls []string, fps int) { sm.mu.Lock() defer sm.mu.Unlock() diff --git a/internal/stream/puller.go b/internal/stream/puller.go new file mode 100644 index 0000000..da4dd56 --- /dev/null +++ b/internal/stream/puller.go @@ -0,0 +1,205 @@ +package stream + +import ( + "bytes" + "encoding/json" + "fmt" + "log" + "net/http" + "time" + + "tianyan-edge/internal/config" +) + +// Puller 负责从云端网关动态获取拉流地址并维护心跳 +type Puller struct { + cfg *config.Config + client *http.Client + sessionID string + stopCh chan struct{} +} + +// PullResponse 云端网关返回结构 +type PullResponse struct { + Code int `json:"code"` + Msg string `json:"msg"` + Data struct { + URL string `json:"url"` + Protocol string `json:"protocol"` + TTL int `json:"ttl"` + DeviceID string `json:"device_id"` + ChannelID string `json:"channel_id"` + } `json:"data"` +} + +// HeartbeatPayload 心跳请求结构 +type HeartbeatPayload struct { + EdgeToken string `json:"edge_token"` +} + +// PullPayload 拉流请求结构 +type PullPayload struct { + EdgeToken string `json:"edge_token"` + Protocol string `json:"protocol"` +} + +func NewPuller(cfg *config.Config) *Puller { + return &Puller{ + cfg: cfg, + client: &http.Client{ + Timeout: 60 * time.Second, // WVP SIP 信令较慢，设置较长超时 + }, + stopCh: make(chan struct{}), + } +} + +// Pull 请求云端获取拉流 URL +func (p *Puller) Pull() (string, error) { + if p.cfg.StreamPullURL == "" { + return "", fmt.Errorf("stream_pull_url not configured") + } + + payload := PullPayload{ + EdgeToken: p.cfg.EdgeToken, + Protocol: p.cfg.StreamProtocol, + } + body, err := json.Marshal(payload) + if err != nil { + return "", err + } + + req, err := http.NewRequest("POST", p.cfg.StreamPullURL+"/api/v1/edge/pull", bytes.NewReader(body)) + if err != nil { + return "", err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Edge-Token", p.cfg.EdgeToken) + + resp, err := p.client.Do(req) + if err != nil { + return "", fmt.Errorf("pull request failed: %w", err) + } + defer resp.Body.Close() + + var pullResp PullResponse + if err := json.NewDecoder(resp.Body).Decode(&pullResp); err != nil { + return "", fmt.Errorf("decode pull response failed: %w", err) + } + + if pullResp.Code != 0 { + return "", fmt.Errorf("pull failed: %s", pullResp.Msg) + } + + p.sessionID = p.cfg.EdgeToken + log.Printf("puller: got stream url protocol=%s ttl=%ds", pullResp.Data.Protocol, pullResp.Data.TTL) + return pullResp.Data.URL, nil +} + +// StartHeartbeatLoop 启动后台心跳保活协程 +func (p *Puller) StartHeartbeatLoop() { + go func() { + ticker := time.NewTicker(20 * time.Second) + defer ticker.Stop() + + for { + select { + case <-p.stopCh: + log.Println("puller: heartbeat loop stopped") + return + case <-ticker.C: + if p.sessionID == "" { + continue + } + if err := p.sendHeartbeat(); err != nil { + log.Printf("puller: heartbeat failed: %v (will retry next tick)", err) + } + } + } + }() +} + +// sendHeartbeat 发送单次心跳 +func (p *Puller) sendHeartbeat() error { + payload := HeartbeatPayload{EdgeToken: p.cfg.EdgeToken} + body, err := json.Marshal(payload) + if err != nil { + return err + } + + req, err := http.NewRequest("POST", p.cfg.StreamPullURL+"/api/v1/edge/heartbeat", bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Edge-Token", p.cfg.EdgeToken) + + resp, err := p.client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + var respData struct { + Code int `json:"code"` + Msg string `json:"msg"` + } + if err := json.NewDecoder(resp.Body).Decode(&respData); err != nil { + return err + } + + if respData.Code != 0 { + return fmt.Errorf("heartbeat failed: %s", respData.Msg) + } + return nil +} + +// Release 主动释放云端拉流会话 +func (p *Puller) Release() { + if p.sessionID == "" || p.cfg.StreamPullURL == "" { + return + } + + payload := HeartbeatPayload{EdgeToken: p.cfg.EdgeToken} + body, err := json.Marshal(payload) + if err != nil { + log.Printf("puller: release marshal error: %v", err) + return + } + + req, err := http.NewRequest("POST", p.cfg.StreamPullURL+"/api/v1/edge/release", bytes.NewReader(body)) + if err != nil { + log.Printf("puller: release request error: %v", err) + return + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Edge-Token", p.cfg.EdgeToken) + + resp, err := p.client.Do(req) + if err != nil { + log.Printf("puller: release do error: %v", err) + return + } + defer resp.Body.Close() + + var respData struct { + Code int `json:"code"` + Msg string `json:"msg"` + } + if err := json.NewDecoder(resp.Body).Decode(&respData); err != nil { + log.Printf("puller: release decode error: %v", err) + return + } + + if respData.Code == 0 { + log.Println("puller: session released successfully") + } else { + log.Printf("puller: release failed: %s", respData.Msg) + } + p.sessionID = "" +} + +// Stop 停止拉流管理器 +func (p *Puller) Stop() { + close(p.stopCh) + p.Release() +}