weijuesen
493 lines
14 KiB
Go
493 lines
14 KiB
Go
package main
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"os"
|
||
"strconv"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/aws/aws-sdk-go-v2/aws"
|
||
"github.com/aws/aws-sdk-go-v2/config"
|
||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||
"github.com/aws/aws-sdk-go-v2/service/s3/types"
|
||
)
|
||
|
||
// ===== 配置 =====
|
||
|
||
var (
|
||
zlmBase = getEnv("ZLM_BASE", "http://localhost:8081")
|
||
s3Endpoint = getEnv("S3_ENDPOINT", "http://localhost:7480")
|
||
s3AccessKey = getEnv("S3_ACCESS_KEY", "silk-app")
|
||
s3SecretKey = getEnv("S3_SECRET_KEY", "Silk-App-Secret-2026!")
|
||
s3Bucket = getEnv("S3_BUCKET", "silk-video-events")
|
||
s3Region = getEnv("S3_REGION", "us-east-1")
|
||
nestjsBase = getEnv("NESTJS_BASE", "http://localhost:3000/api/v1")
|
||
nestjsUser = getEnv("NESTJS_USER", "admin")
|
||
nestjsPass = getEnv("NESTJS_PASSWORD", "silk@123")
|
||
internalKey = getEnv("INTERNAL_API_KEY", "silk-internal-2026")
|
||
segmentDur = getEnvInt("SEGMENT_DURATION", 300)
|
||
chunkSize = getEnvInt("CHUNK_SIZE", 5*1024*1024)
|
||
httpPort = getEnv("HTTP_PORT", "9090")
|
||
)
|
||
|
||
func getEnv(key, def string) string {
|
||
if v := os.Getenv(key); v != "" {
|
||
return v
|
||
}
|
||
return def
|
||
}
|
||
|
||
func getEnvInt(key string, def int) int {
|
||
if v := os.Getenv(key); v != "" {
|
||
if n, err := strconv.Atoi(v); err == nil {
|
||
return n
|
||
}
|
||
}
|
||
return def
|
||
}
|
||
|
||
// ===== S3 客户端 =====
|
||
|
||
var s3Client *s3.Client
|
||
|
||
func initS3() {
|
||
cfg, err := config.LoadDefaultConfig(context.Background(),
|
||
config.WithRegion(s3Region),
|
||
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(s3AccessKey, s3SecretKey, "")),
|
||
)
|
||
if err != nil {
|
||
log.Fatalf("S3 config error: %v", err)
|
||
}
|
||
s3Client = s3.NewFromConfig(cfg, func(o *s3.Options) {
|
||
o.BaseEndpoint = aws.String(s3Endpoint)
|
||
o.UsePathStyle = true
|
||
})
|
||
}
|
||
|
||
// ===== StreamRecorder =====
|
||
//
|
||
// 连续录制模式：循环录制 5 分钟分段，直到 stop() 或流 EOF。
|
||
// 每个分段独立写入 S3 并创建 DB 记录，实现"一直录制、按 5 分钟落地"。
|
||
|
||
type StreamRecorder struct {
|
||
Stream string
|
||
CameraID string
|
||
DeviceID string
|
||
ChannelID string
|
||
StartedAt time.Time // 当前分段开始时间（每段重置）
|
||
UploadID string
|
||
S3Key string
|
||
Parts []types.CompletedPart
|
||
PartNum int32
|
||
TotalBytes int64
|
||
Running bool
|
||
SegmentCount int
|
||
streamCtx context.Context
|
||
streamCancel context.CancelFunc
|
||
}
|
||
|
||
var (
|
||
recordings = make(map[string]*StreamRecorder)
|
||
recMu sync.RWMutex
|
||
)
|
||
|
||
// start 启动录制线程
|
||
func (r *StreamRecorder) start() {
|
||
r.Running = true
|
||
go r.record()
|
||
}
|
||
|
||
// stop 设置 Running=false，当前分段会在下一次循环检查时收尾，且不再开始新段
|
||
func (r *StreamRecorder) stop() {
|
||
r.Running = false
|
||
// 取消拉流 context 以便及时退出阻塞的 Read（S3 操作使用独立 context，不受影响）
|
||
if r.streamCancel != nil {
|
||
r.streamCancel()
|
||
}
|
||
}
|
||
|
||
// record 核心录制逻辑：循环录制分段
|
||
func (r *StreamRecorder) record() {
|
||
defer func() {
|
||
r.Running = false
|
||
recMu.Lock()
|
||
delete(recordings, r.Stream)
|
||
recMu.Unlock()
|
||
log.Printf("[%s] 录制线程退出（共完成 %d 段）", r.Stream, r.SegmentCount)
|
||
|
||
// 通知 Go 后端录制已结束，清理 activeRecordings
|
||
notifyRecordingEnd(r.CameraID, r.Stream)
|
||
}()
|
||
|
||
for r.Running {
|
||
reason := r.recordOneSegment()
|
||
if reason == "eof" || reason == "error" || reason == "stopped" {
|
||
break
|
||
}
|
||
// reason == "segment_done"：继续下一段
|
||
if r.Running {
|
||
log.Printf("[%s] 立即开始下一段录制（已累计 %d 段）", r.Stream, r.SegmentCount)
|
||
}
|
||
}
|
||
}
|
||
|
||
// segmentReason 分段结束原因
|
||
type segmentReason string
|
||
|
||
const (
|
||
segmentDone segmentReason = "segment_done" // 达到分段时长，正常结束
|
||
reasonEOF segmentReason = "eof" // 流结束
|
||
reasonError segmentReason = "error" // 出错
|
||
reasonStopped segmentReason = "stopped" // 收到 stop
|
||
)
|
||
|
||
// recordOneSegment 录制单个 5 分钟分段
|
||
func (r *StreamRecorder) recordOneSegment() segmentReason {
|
||
segStart := time.Now()
|
||
r.StartedAt = segStart
|
||
r.Parts = nil
|
||
r.PartNum = 0
|
||
r.TotalBytes = 0
|
||
r.S3Key = ""
|
||
r.UploadID = ""
|
||
|
||
// 拉流 context（stop 时取消以中断阻塞的 Read）
|
||
streamCtx, streamCancel := context.WithCancel(context.Background())
|
||
r.streamCtx = streamCtx
|
||
r.streamCancel = streamCancel
|
||
defer streamCancel()
|
||
|
||
// 1. CreateMultipartUpload（用独立 context，stop 不影响）
|
||
r.S3Key = fmt.Sprintf("%s/%d.mp4", segStart.Format("2006/01/02"), segStart.UnixMilli())
|
||
createOut, err := s3Client.CreateMultipartUpload(context.Background(), &s3.CreateMultipartUploadInput{
|
||
Bucket: aws.String(s3Bucket),
|
||
Key: aws.String(r.S3Key),
|
||
ContentType: aws.String("video/mp4"),
|
||
})
|
||
if err != nil {
|
||
log.Printf("[%s] CreateMultipartUpload 失败: %v", r.Stream, err)
|
||
return reasonError
|
||
}
|
||
r.UploadID = *createOut.UploadId
|
||
log.Printf("[%s] 开始分段 #%d CreateMultipartUpload: s3://%s/%s",
|
||
r.Stream, r.SegmentCount+1, s3Bucket, r.S3Key)
|
||
|
||
// 2. HTTP 拉流（用 streamCtx，stop 时可中断）
|
||
url := fmt.Sprintf("%s/rtp/%s.live.mp4", zlmBase, r.Stream)
|
||
req, err := http.NewRequestWithContext(streamCtx, "GET", url, nil)
|
||
if err != nil {
|
||
log.Printf("[%s] 创建请求失败: %v", r.Stream, err)
|
||
r.abortMultipart()
|
||
return reasonError
|
||
}
|
||
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
if streamCtx.Err() != nil {
|
||
// stop 导致的
|
||
r.abortMultipart()
|
||
return reasonStopped
|
||
}
|
||
log.Printf("[%s] 拉流失败: %v", r.Stream, err)
|
||
r.abortMultipart()
|
||
return reasonError
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
// 3. 循环读取并上传
|
||
buf := make([]byte, 0, chunkSize)
|
||
tmp := make([]byte, 64*1024)
|
||
stopped := false
|
||
|
||
for {
|
||
if !r.Running {
|
||
stopped = true
|
||
break
|
||
}
|
||
if time.Since(segStart) >= time.Duration(segmentDur)*time.Second {
|
||
log.Printf("[%s] 达到分段时长 %ds，结束本段", r.Stream, segmentDur)
|
||
break
|
||
}
|
||
|
||
n, err := resp.Body.Read(tmp)
|
||
if n > 0 {
|
||
buf = append(buf, tmp[:n]...)
|
||
}
|
||
if err == io.EOF {
|
||
log.Printf("[%s] 流结束（EOF）", r.Stream)
|
||
r.finalizeSegment(buf, segStart)
|
||
return reasonEOF
|
||
}
|
||
if err != nil {
|
||
if streamCtx.Err() != nil {
|
||
// stop 导致的中断
|
||
stopped = true
|
||
break
|
||
}
|
||
log.Printf("[%s] 读取流错误: %v", r.Stream, err)
|
||
r.finalizeSegment(buf, segStart)
|
||
return reasonError
|
||
}
|
||
|
||
if len(buf) >= chunkSize {
|
||
if err := r.uploadPart(buf[:chunkSize]); err != nil {
|
||
log.Printf("[%s] UploadPart 失败: %v", r.Stream, err)
|
||
r.abortMultipart()
|
||
return reasonError
|
||
}
|
||
buf = make([]byte, 0, chunkSize)
|
||
}
|
||
}
|
||
|
||
// 收尾：上传最后一块 + CompleteMultipartUpload + DB 记录
|
||
r.finalizeSegment(buf, segStart)
|
||
|
||
if stopped {
|
||
return reasonStopped
|
||
}
|
||
return segmentDone
|
||
}
|
||
|
||
// finalizeSegment 上传最后一块并完成 MultipartUpload + 创建 DB 记录
|
||
func (r *StreamRecorder) finalizeSegment(buf []byte, segStart time.Time) {
|
||
if len(buf) > 0 {
|
||
if err := r.uploadPart(buf); err != nil {
|
||
log.Printf("[%s] 最后一块上传失败: %v", r.Stream, err)
|
||
}
|
||
}
|
||
|
||
if len(r.Parts) == 0 {
|
||
log.Printf("[%s] 本段无数据，AbortMultipartUpload", r.Stream)
|
||
r.abortMultipart()
|
||
return
|
||
}
|
||
|
||
// CompleteMultipartUpload（用独立 context）
|
||
_, err := s3Client.CompleteMultipartUpload(context.Background(), &s3.CompleteMultipartUploadInput{
|
||
Bucket: aws.String(s3Bucket),
|
||
Key: aws.String(r.S3Key),
|
||
UploadId: aws.String(r.UploadID),
|
||
MultipartUpload: &types.CompletedMultipartUpload{
|
||
Parts: r.Parts,
|
||
},
|
||
})
|
||
if err != nil {
|
||
log.Printf("[%s] CompleteMultipartUpload 失败: %v", r.Stream, err)
|
||
r.abortMultipart()
|
||
return
|
||
}
|
||
log.Printf("[%s] CompleteMultipartUpload parts=%d bytes=%d", r.Stream, len(r.Parts), r.TotalBytes)
|
||
|
||
// 创建数据库记录
|
||
dur := int(time.Since(segStart).Seconds())
|
||
r.createClipRecord(dur)
|
||
r.SegmentCount++
|
||
log.Printf("[%s] 分段 #%d 完成 ✓ bytes=%d dur=%ds", r.Stream, r.SegmentCount, r.TotalBytes, dur)
|
||
}
|
||
|
||
func (r *StreamRecorder) uploadPart(data []byte) error {
|
||
r.PartNum++
|
||
out, err := s3Client.UploadPart(context.Background(), &s3.UploadPartInput{
|
||
Bucket: aws.String(s3Bucket),
|
||
Key: aws.String(r.S3Key),
|
||
UploadId: aws.String(r.UploadID),
|
||
PartNumber: aws.Int32(r.PartNum),
|
||
Body: bytes.NewReader(data),
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
r.Parts = append(r.Parts, types.CompletedPart{
|
||
PartNumber: aws.Int32(r.PartNum),
|
||
ETag: out.ETag,
|
||
})
|
||
r.TotalBytes += int64(len(data))
|
||
log.Printf("[%s] UploadPart #%d size=%d", r.Stream, r.PartNum, len(data))
|
||
return nil
|
||
}
|
||
|
||
func (r *StreamRecorder) abortMultipart() {
|
||
if r.UploadID == "" {
|
||
return
|
||
}
|
||
_, err := s3Client.AbortMultipartUpload(context.Background(), &s3.AbortMultipartUploadInput{
|
||
Bucket: aws.String(s3Bucket),
|
||
Key: aws.String(r.S3Key),
|
||
UploadId: aws.String(r.UploadID),
|
||
})
|
||
if err != nil {
|
||
log.Printf("[%s] AbortMultipartUpload 失败: %v", r.Stream, err)
|
||
}
|
||
}
|
||
|
||
func (r *StreamRecorder) createClipRecord(dur int) {
|
||
body := map[string]interface{}{
|
||
"cameraId": r.CameraID,
|
||
"trigger": "schedule",
|
||
"format": "mp4",
|
||
"startAt": r.StartedAt.Format(time.RFC3339),
|
||
"durationSec": dur,
|
||
"sizeBytes": strconv.FormatInt(r.TotalBytes, 10),
|
||
"s3Bucket": s3Bucket,
|
||
"s3Key": r.S3Key,
|
||
"notes": "自动录制(不落地-Go)",
|
||
}
|
||
jsonBody, _ := json.Marshal(body)
|
||
req, _ := http.NewRequest("POST", nestjsBase+"/video/clips/internal", bytes.NewReader(jsonBody))
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("x-api-key", internalKey)
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
log.Printf("[%s] 创建数据库记录失败: %v", r.Stream, err)
|
||
return
|
||
}
|
||
resp.Body.Close()
|
||
log.Printf("[%s] 创建数据库记录: %s", r.Stream, nestjsBase+"/video/clips/internal")
|
||
}
|
||
|
||
// notifyRecordingEnd 通知 Go 后端录制已结束（流 EOF / 错误 / 停止）
|
||
// Go 后端收到后清理 activeRecordings 并停止 WVP play session
|
||
func notifyRecordingEnd(cameraID, stream string) {
|
||
body := map[string]string{
|
||
"cameraId": cameraID,
|
||
"stream": stream,
|
||
}
|
||
jsonBody, _ := json.Marshal(body)
|
||
req, _ := http.NewRequest("POST", nestjsBase+"/video/recordings/internal/end", bytes.NewReader(jsonBody))
|
||
req.Header.Set("Content-Type", "application/json")
|
||
req.Header.Set("x-api-key", internalKey)
|
||
resp, err := http.DefaultClient.Do(req)
|
||
if err != nil {
|
||
log.Printf("[%s] 通知录制结束失败: %v", stream, err)
|
||
return
|
||
}
|
||
resp.Body.Close()
|
||
log.Printf("[%s] 已通知 Go 后端录制结束", stream)
|
||
}
|
||
|
||
// ===== HTTP 控制服务 =====
|
||
|
||
func startHTTPServer() {
|
||
mux := http.NewServeMux()
|
||
|
||
mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
|
||
json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
|
||
})
|
||
|
||
mux.HandleFunc("/record/start", func(w http.ResponseWriter, r *http.Request) {
|
||
if r.Method != "POST" {
|
||
http.Error(w, "Method not allowed", 405)
|
||
return
|
||
}
|
||
var body struct {
|
||
Stream string `json:"stream"`
|
||
CameraID string `json:"cameraId"`
|
||
DeviceID string `json:"deviceId"`
|
||
ChannelID string `json:"channelId"`
|
||
}
|
||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||
http.Error(w, "Bad request", 400)
|
||
return
|
||
}
|
||
|
||
recMu.Lock()
|
||
if _, exists := recordings[body.Stream]; exists {
|
||
recMu.Unlock()
|
||
http.Error(w, "already recording", 409)
|
||
return
|
||
}
|
||
rec := &StreamRecorder{
|
||
Stream: body.Stream,
|
||
CameraID: body.CameraID,
|
||
DeviceID: body.DeviceID,
|
||
ChannelID: body.ChannelID,
|
||
}
|
||
recordings[body.Stream] = rec
|
||
recMu.Unlock()
|
||
|
||
rec.start()
|
||
log.Printf("HTTP /record/start: %s (camera=%s) — 连续录制模式", body.Stream, body.CameraID)
|
||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||
"ok": true,
|
||
"stream": body.Stream,
|
||
"startedAt": time.Now().Format(time.RFC3339),
|
||
})
|
||
})
|
||
|
||
mux.HandleFunc("/record/stop", func(w http.ResponseWriter, r *http.Request) {
|
||
if r.Method != "POST" {
|
||
http.Error(w, "Method not allowed", 405)
|
||
return
|
||
}
|
||
var body struct {
|
||
Stream string `json:"stream"`
|
||
}
|
||
json.NewDecoder(r.Body).Decode(&body)
|
||
|
||
recMu.RLock()
|
||
rec, exists := recordings[body.Stream]
|
||
recMu.RUnlock()
|
||
|
||
if !exists {
|
||
http.Error(w, "not recording", 404)
|
||
return
|
||
}
|
||
log.Printf("HTTP /record/stop: %s (已完成 %d 段，等待当前分段收尾)", body.Stream, rec.SegmentCount)
|
||
rec.stop()
|
||
json.NewEncoder(w).Encode(map[string]interface{}{
|
||
"ok": true,
|
||
"stream": body.Stream,
|
||
"segmentCount": rec.SegmentCount,
|
||
})
|
||
})
|
||
|
||
mux.HandleFunc("/record/status", func(w http.ResponseWriter, r *http.Request) {
|
||
type StatusItem struct {
|
||
Stream string `json:"stream"`
|
||
CameraID string `json:"cameraId"`
|
||
StartedAt string `json:"startedAt"`
|
||
TotalBytes int64 `json:"totalBytes"`
|
||
SegmentCount int `json:"segmentCount"`
|
||
}
|
||
items := make([]StatusItem, 0)
|
||
recMu.RLock()
|
||
for _, rec := range recordings {
|
||
items = append(items, StatusItem{
|
||
Stream: rec.Stream,
|
||
CameraID: rec.CameraID,
|
||
StartedAt: rec.StartedAt.Format(time.RFC3339),
|
||
TotalBytes: rec.TotalBytes,
|
||
SegmentCount: rec.SegmentCount,
|
||
})
|
||
}
|
||
recMu.RUnlock()
|
||
json.NewEncoder(w).Encode(map[string]interface{}{"recordings": items})
|
||
})
|
||
|
||
log.Printf("HTTP 控制服务监听 0.0.0.0:%s", httpPort)
|
||
log.Fatal(http.ListenAndServe(":"+httpPort, mux))
|
||
}
|
||
|
||
// ===== 主函数 =====
|
||
|
||
func main() {
|
||
log.SetFlags(log.LstdFlags)
|
||
log.Println("=== Go 录制服务启动 ===")
|
||
log.Printf("ZLM: %s, S3: %s, Backend: %s", zlmBase, s3Endpoint, nestjsBase)
|
||
log.Printf("分段时长: %ds, 分块大小: %dMB", segmentDur, chunkSize/1024/1024)
|
||
log.Printf("录制模式: 连续录制，每 %ds 落地一个 S3 分段，由 /record/stop 停止", segmentDur)
|
||
|
||
initS3()
|
||
|
||
// 注：autoCheck 自动检查已移除——录制完全由前端"开始/停止录制"按钮通过 Go 后端驱动
|
||
startHTTPServer()
|
||
}
|