385 lines
9.1 KiB
Go
385 lines
9.1 KiB
Go
/**
|
||
* package mediaserver
|
||
*
|
||
* FFmpeg RTMP 接收器
|
||
* 使用 FFmpeg -listen 模式接收小程序 RTMP 推流，转换为 FLV 输出
|
||
* 解决自定义 RTMP 协议解析器无法正确处理微信小程序 live-pusher 的问题
|
||
*/
|
||
package mediaserver
|
||
|
||
import (
|
||
"bytes"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"os/exec"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// FFmpegRTMPReceiver FFmpeg RTMP 接收器
|
||
type FFmpegRTMPReceiver struct {
|
||
streamID string
|
||
port int
|
||
ffmpegPath string
|
||
rtmpServer *RTMPServer
|
||
|
||
cmd *exec.Cmd
|
||
stdout io.ReadCloser
|
||
stderr io.ReadCloser
|
||
|
||
outputChan chan []byte
|
||
stopChan chan struct{}
|
||
|
||
mu sync.Mutex
|
||
running bool
|
||
|
||
// 统计信息
|
||
startTime time.Time
|
||
outputBytes int64
|
||
}
|
||
|
||
// FFmpegRTMPReceiverPool 接收器池，管理多个 FFmpeg RTMP 接收器
|
||
type FFmpegRTMPReceiverPool struct {
|
||
receivers map[string]*FFmpegRTMPReceiver
|
||
mu sync.RWMutex
|
||
rtmpServer *RTMPServer
|
||
ffmpegPath string
|
||
basePort int
|
||
nextPort int
|
||
maxPort int
|
||
usedPorts map[int]bool
|
||
}
|
||
|
||
// NewFFmpegRTMPReceiverPool 创建接收器池
|
||
func NewFFmpegRTMPReceiverPool(rtmpServer *RTMPServer, basePort int) (*FFmpegRTMPReceiverPool, error) {
|
||
ffmpegPath, err := EnsureFFmpeg()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("FFmpeg 不可用: %w", err)
|
||
}
|
||
|
||
return &FFmpegRTMPReceiverPool{
|
||
receivers: make(map[string]*FFmpegRTMPReceiver),
|
||
rtmpServer: rtmpServer,
|
||
ffmpegPath: ffmpegPath,
|
||
basePort: basePort,
|
||
nextPort: basePort,
|
||
maxPort: basePort + 100, // 最多支持 100 个并发流
|
||
usedPorts: make(map[int]bool),
|
||
}, nil
|
||
}
|
||
|
||
// allocatePortLocked 分配一个可用端口（调用者必须持有锁）
|
||
func (p *FFmpegRTMPReceiverPool) allocatePortLocked() (int, error) {
|
||
// 注意：此方法假设调用者已持有 p.mu 锁
|
||
// 查找可用端口
|
||
for port := p.basePort; port < p.maxPort; port++ {
|
||
if !p.usedPorts[port] {
|
||
p.usedPorts[port] = true
|
||
return port, nil
|
||
}
|
||
}
|
||
|
||
return 0, fmt.Errorf("没有可用端口")
|
||
}
|
||
|
||
// releasePortLocked 释放端口（调用者必须持有锁）
|
||
func (p *FFmpegRTMPReceiverPool) releasePortLocked(port int) {
|
||
// 注意：此方法假设调用者已持有 p.mu 锁
|
||
delete(p.usedPorts, port)
|
||
}
|
||
|
||
// GetOrCreate 获取或创建接收器
|
||
func (p *FFmpegRTMPReceiverPool) GetOrCreate(streamID string) (*FFmpegRTMPReceiver, error) {
|
||
p.mu.Lock()
|
||
defer p.mu.Unlock()
|
||
|
||
// 检查是否已存在
|
||
if r, exists := p.receivers[streamID]; exists && r.IsRunning() {
|
||
return r, nil
|
||
}
|
||
|
||
// 分配端口（使用不加锁版本，因为我们已经持有锁）
|
||
port, err := p.allocatePortLocked()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 创建新的接收器
|
||
r := &FFmpegRTMPReceiver{
|
||
streamID: streamID,
|
||
port: port,
|
||
ffmpegPath: p.ffmpegPath,
|
||
rtmpServer: p.rtmpServer,
|
||
outputChan: make(chan []byte, 100),
|
||
stopChan: make(chan struct{}),
|
||
}
|
||
|
||
p.receivers[streamID] = r
|
||
return r, nil
|
||
}
|
||
|
||
// Release 释放接收器
|
||
func (p *FFmpegRTMPReceiverPool) Release(streamID string) {
|
||
p.mu.Lock()
|
||
defer p.mu.Unlock()
|
||
|
||
if r, exists := p.receivers[streamID]; exists {
|
||
p.releasePortLocked(r.port)
|
||
r.Stop()
|
||
delete(p.receivers, streamID)
|
||
}
|
||
}
|
||
|
||
// ReleaseAll 释放所有接收器
|
||
func (p *FFmpegRTMPReceiverPool) ReleaseAll() {
|
||
p.mu.Lock()
|
||
defer p.mu.Unlock()
|
||
|
||
for id, r := range p.receivers {
|
||
p.releasePortLocked(r.port)
|
||
r.Stop()
|
||
delete(p.receivers, id)
|
||
}
|
||
}
|
||
|
||
// GetRTMPPort 获取接收器的 RTMP 端口
|
||
func (r *FFmpegRTMPReceiver) GetRTMPPort() int {
|
||
return r.port
|
||
}
|
||
|
||
// GetRTMPURL 获取完整的 RTMP 推流地址
|
||
func (r *FFmpegRTMPReceiver) GetRTMPURL() string {
|
||
return fmt.Sprintf("rtmp://127.0.0.1:%d/live/%s", r.port, r.streamID)
|
||
}
|
||
|
||
// Start 启动接收器
|
||
func (r *FFmpegRTMPReceiver) Start() error {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
|
||
if r.running {
|
||
return fmt.Errorf("接收器已在运行")
|
||
}
|
||
|
||
// 构建 FFmpeg 命令
|
||
// ffmpeg -listen 1 -i rtmp://0.0.0.0:{port}/live/{stream_id} -c copy -f flv pipe:1
|
||
args := []string{
|
||
"-listen", "1", // 作为服务器监听
|
||
"-timeout", "30000000", // 超时时间 30 秒（微秒）
|
||
"-i", fmt.Sprintf("rtmp://0.0.0.0:%d/live/%s", r.port, r.streamID), // 监听地址
|
||
"-c", "copy", // 直接复制，不重新编码
|
||
"-fflags", "nobuffer", // 禁用输入缓冲
|
||
"-flags", "low_delay", // 低延迟模式
|
||
"-f", "flv", // 输出格式
|
||
"pipe:1", // 输出到 stdout
|
||
}
|
||
|
||
r.cmd = exec.Command(r.ffmpegPath, args...)
|
||
|
||
var err error
|
||
|
||
// 获取 stdout
|
||
r.stdout, err = r.cmd.StdoutPipe()
|
||
if err != nil {
|
||
return fmt.Errorf("获取 stdout 失败: %w", err)
|
||
}
|
||
|
||
// 获取 stderr
|
||
r.stderr, err = r.cmd.StderrPipe()
|
||
if err != nil {
|
||
return fmt.Errorf("获取 stderr 失败: %w", err)
|
||
}
|
||
|
||
// 启动 FFmpeg 进程
|
||
if err := r.cmd.Start(); err != nil {
|
||
return fmt.Errorf("启动 FFmpeg 失败: %w", err)
|
||
}
|
||
|
||
r.running = true
|
||
r.startTime = time.Now()
|
||
|
||
log.Printf("🎬 [FFmpegRTMPReceiver] 启动成功 | Stream:%s Port:%d PID:%d",
|
||
r.streamID, r.port, r.cmd.Process.Pid)
|
||
|
||
// 启动输出读取协程
|
||
go r.readOutput()
|
||
|
||
// 启动 stderr 读取协程
|
||
go r.readStderr()
|
||
|
||
// 启动进程监控协程
|
||
go r.monitor()
|
||
|
||
return nil
|
||
}
|
||
|
||
// Stop 停止接收器
|
||
// 与 FFmpegTranscoder.Stop 相同的并发安全约束：
|
||
// 1. 不调用 cmd.Wait()——monitor 协程是唯一的 Wait 回收方，重复 Wait 行为未定义；
|
||
// 2. 不 close(outputChan)——readOutput（唯一发送方）可能仍在发送，
|
||
// 向已关闭通道发送会 panic 崩溃进程；通道由 readOutput 退出时统一关闭。
|
||
func (r *FFmpegRTMPReceiver) Stop() error {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
|
||
if !r.running {
|
||
return nil
|
||
}
|
||
|
||
r.running = false
|
||
close(r.stopChan)
|
||
|
||
// 杀掉进程即可：stdout 管道随之 EOF，readOutput 退出并关闭 outputChan；
|
||
// 僵尸进程由 monitor 协程的 cmd.Wait() 统一回收
|
||
if r.cmd != nil && r.cmd.Process != nil {
|
||
r.cmd.Process.Kill()
|
||
}
|
||
|
||
duration := time.Since(r.startTime)
|
||
log.Printf("🛑 [FFmpegRTMPReceiver] 已停止 | Stream:%s Port:%d Duration:%v Output:%d bytes",
|
||
r.streamID, r.port, duration, r.outputBytes)
|
||
|
||
return nil
|
||
}
|
||
|
||
// IsRunning 检查是否运行中
|
||
func (r *FFmpegRTMPReceiver) IsRunning() bool {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
return r.running
|
||
}
|
||
|
||
// Output 获取输出通道
|
||
func (r *FFmpegRTMPReceiver) Output() <-chan []byte {
|
||
return r.outputChan
|
||
}
|
||
|
||
// readOutput 读取 FFmpeg FLV 输出
|
||
// 本协程是 outputChan 的唯一发送方，由它在退出时关闭通道：
|
||
// 既保证 range Output() 的消费者能正常结束，又杜绝"向已关闭通道发送"的 panic
|
||
func (r *FFmpegRTMPReceiver) readOutput() {
|
||
defer close(r.outputChan)
|
||
log.Printf("▶️ [FFmpegRTMPReceiver] 开始读取 FLV 输出 | Stream:%s", r.streamID)
|
||
|
||
buf := make([]byte, 64*1024) // 64KB 缓冲区
|
||
flvHeaderSent := false
|
||
|
||
for {
|
||
select {
|
||
case <-r.stopChan:
|
||
return
|
||
default:
|
||
}
|
||
|
||
n, err := r.stdout.Read(buf)
|
||
if err != nil {
|
||
if err != io.EOF {
|
||
log.Printf("⚠️ [FFmpegRTMPReceiver] 读取输出错误: %v | Stream:%s", err, r.streamID)
|
||
}
|
||
return
|
||
}
|
||
|
||
if n > 0 {
|
||
r.outputBytes += int64(n)
|
||
|
||
// 复制数据
|
||
data := make([]byte, n)
|
||
copy(data, buf[:n])
|
||
|
||
// 跳过 FLV 头（前 13 字节）- 只跳过一次
|
||
if !flvHeaderSent && len(data) >= 13 {
|
||
// 验证 FLV 头
|
||
if data[0] == 'F' && data[1] == 'L' && data[2] == 'V' {
|
||
log.Printf("✅ [FFmpegRTMPReceiver] 收到 FLV 头 | Stream:%s", r.streamID)
|
||
flvHeaderSent = true
|
||
// 广播 FLV 数据（包括头）
|
||
r.broadcastFLVData(data)
|
||
} else {
|
||
// 不是 FLV 头，直接广播
|
||
r.broadcastFLVData(data)
|
||
}
|
||
} else {
|
||
// 广播 FLV 数据
|
||
r.broadcastFLVData(data)
|
||
}
|
||
|
||
// 发送到输出通道
|
||
select {
|
||
case r.outputChan <- data:
|
||
default:
|
||
// 通道满，丢弃
|
||
log.Printf("⚠️ [FFmpegRTMPReceiver] 输出通道满，丢弃 %d bytes | Stream:%s", n, r.streamID)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// broadcastFLVData 广播 FLV 数据给所有订阅者
|
||
func (r *FFmpegRTMPReceiver) broadcastFLVData(data []byte) {
|
||
if r.rtmpServer == nil {
|
||
return
|
||
}
|
||
|
||
// 解析 FLV tag 并广播
|
||
r.parseFLVAndBroadcast(data)
|
||
}
|
||
|
||
// parseFLVAndBroadcast 解析 FLV 数据并广播
|
||
func (r *FFmpegRTMPReceiver) parseFLVAndBroadcast(data []byte) {
|
||
// 获取流
|
||
stream := r.rtmpServer.GetStream(r.streamID)
|
||
if stream == nil {
|
||
return
|
||
}
|
||
|
||
// 简单广播原始数据给订阅者
|
||
r.rtmpServer.subMu.RLock()
|
||
subs := r.rtmpServer.subscribers[r.streamID]
|
||
r.rtmpServer.subMu.RUnlock()
|
||
|
||
if subs == nil {
|
||
return
|
||
}
|
||
|
||
for _, sub := range subs {
|
||
select {
|
||
case sub.DataChan <- data:
|
||
default:
|
||
// 订阅者通道满，跳过
|
||
}
|
||
}
|
||
}
|
||
|
||
// readStderr 读取 FFmpeg 错误输出
|
||
func (r *FFmpegRTMPReceiver) readStderr() {
|
||
buf := new(bytes.Buffer)
|
||
io.Copy(buf, r.stderr)
|
||
|
||
if buf.Len() > 0 {
|
||
output := buf.String()
|
||
if len(output) > 500 {
|
||
output = output[:500] + "..."
|
||
}
|
||
log.Printf("📋 [FFmpegRTMPReceiver] FFmpeg 输出: %s | Stream:%s", output, r.streamID)
|
||
}
|
||
}
|
||
|
||
// monitor 监控 FFmpeg 进程
|
||
func (r *FFmpegRTMPReceiver) monitor() {
|
||
err := r.cmd.Wait()
|
||
|
||
r.mu.Lock()
|
||
wasRunning := r.running
|
||
r.running = false
|
||
r.mu.Unlock()
|
||
|
||
if wasRunning {
|
||
if err != nil {
|
||
log.Printf("⚠️ [FFmpegRTMPReceiver] FFmpeg 进程异常退出: %v | Stream:%s", err, r.streamID)
|
||
} else {
|
||
log.Printf("ℹ️ [FFmpegRTMPReceiver] FFmpeg 进程正常退出 | Stream:%s", r.streamID)
|
||
}
|
||
}
|
||
}
|