支持录制流

config.json

@@ -20,7 +20,7 @@
   },
 
   "hls": {
-    "enable": false,
+    "enable": true,
     "segment_duration": 2,
     "playlist_length": 10,
     "dir": "../tmp"
@@ -49,11 +49,12 @@
   },
 
   "record": {
-    "format": "mp4",
+    "enable": true,
-    "path": ""
+    "format": "flv",
+    "dir": "../record"
   },
 
-  "hook": {
+  "hooks": {
     "enable": true,
     "timeout": 10,
 
@@ -64,7 +65,7 @@
     "on_play" : "http://localhost:9000/api/v1/hook/on_play",
     "on_play_done" : "http://localhost:9000/api/v1/hook/on_play_done",
 
-    "on_record": "http://localhost:9000/api/v1/hook/on_reocrd",
+    "on_record": "http://localhost:9000/api/v1/hook/on_record",
     "on_idle_timeout": "http://localhost:9000/api/v1/hook/on_idle_timeout",
     "on_receive_timeout": "http://localhost:9000/api/v1/hook/on_receive_timeout"
   },

gb28181/source.go

@@ -196,7 +196,7 @@ func (source *BaseGBSource) PreparePublish(conn net.Conn, ssrc uint32, source_ G
 	source.SetSSRC(ssrc)
 	source.SetState(stream.SessionStateTransferring)
 
-	if stream.AppConfig.Hook.IsEnablePublishEvent() {
+	if stream.AppConfig.Hooks.IsEnablePublishEvent() {
 		go func() {
 			_, state := stream.HookPublishEvent(source_)
 			if utils.HookStateOK != state {

main.go

@@ -8,6 +8,7 @@ import (
 	"github.com/lkmio/lkm/hls"
 	"github.com/lkmio/lkm/jt1078"
 	"github.com/lkmio/lkm/log"
+	"github.com/lkmio/lkm/record"
 	"github.com/lkmio/lkm/rtc"
 	"github.com/lkmio/lkm/rtsp"
 	"go.uber.org/zap/zapcore"
@@ -26,6 +27,7 @@ func init() {
 	stream.RegisterTransStreamFactory(stream.ProtocolFlv, flv.TransStreamFactory)
 	stream.RegisterTransStreamFactory(stream.ProtocolRtsp, rtsp.TransStreamFactory)
 	stream.RegisterTransStreamFactory(stream.ProtocolRtc, rtc.TransStreamFactory)
+	stream.SetRecordStreamFactory(record.NewFLVFileSink)
 
 	config, err := stream.LoadConfigFile("./config.json")
 	if err != nil {
@@ -126,7 +128,7 @@ func main() {
 		log.Sugar.Info("启动jt1078服务成功 addr:", jtAddr.String())
 	}
 
-	if stream.AppConfig.Hook.IsEnableOnStarted() {
+	if stream.AppConfig.Hooks.IsEnableOnStarted() {
 		go func() {
 			_, _ = stream.Hook(stream.HookEventStarted, "", nil)
 		}()

record/record_flv.go

@@ -0,0 +1,69 @@
+package record
+
+import (
+	"github.com/lkmio/lkm/stream"
+	"os"
+	"path/filepath"
+	"time"
+)
+
+type FLVFileSink struct {
+	stream.BaseSink
+	file *os.File
+	fail bool
+}
+
+// Input 输入http-flv数据
+func (f *FLVFileSink) Input(data []byte) error {
+	if f.fail {
+		return nil
+	}
+
+	//去掉不需要的换行符
+	var offset int
+	for i := 2; i < len(data); i++ {
+		if data[i-2] == 0x0D && data[i-1] == 0x0A {
+			offset = i
+			break
+		}
+	}
+
+	_, err := f.file.Write(data[offset : len(data)-2])
+	if err != nil {
+		//只要写入失败一次，后续不再允许写入, 不影响拉流
+		f.fail = true
+	}
+
+	return err
+}
+
+func (f *FLVFileSink) Close() {
+	if f.file != nil {
+		f.file.Close()
+		f.file = nil
+	}
+}
+
+// NewFLVFileSink 创建FLV文件录制流Sink
+// 保存path: dir/sourceId/yyyy-MM-dd/HH-mm-ss.flv
+func NewFLVFileSink(sourceId string) (stream.Sink, string, error) {
+	now := time.Now().Format("2006-01-02/15-04-05")
+	path := filepath.Join(stream.AppConfig.Record.Dir, sourceId, now+".flv")
+
+	//创建目录
+	dir := filepath.Dir(path)
+	if err := os.MkdirAll(dir, 0666); err != nil {
+		return nil, "", err
+	}
+
+	//创建flv文件
+	file, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0666)
+	if err != nil {
+		return nil, "", err
+	}
+
+	return &FLVFileSink{
+		BaseSink: stream.BaseSink{Id_: "record-sink-flv", SourceId_: sourceId, Protocol_: stream.ProtocolFlv},
+		file:     file,
+	}, path, nil
+}

stream/config.go

@@ -49,6 +49,7 @@ type RtspConfig struct {
 type RecordConfig struct {
 	Enable bool   `json:"enable"`
 	Format string `json:"format"`
+	Dir    string `json:"dir"`
 }
 
 type LogConfig struct {
@@ -120,7 +121,7 @@ func (c HlsConfig) TSFormat(sourceId string) string {
 	return split[len(split)-1] + "_%d.ts"
 }
 
-type HookConfig struct {
+type HooksConfig struct {
 	Enable              bool   `json:"enable"`
 	Timeout             int64  `json:"timeout"`
 	OnStartedUrl        string `json:"on_started"`         //应用启动后回调
@@ -133,35 +134,35 @@ type HookConfig struct {
 	OnReceiveTimeoutUrl string `json:"on_receive_timeout"` //没有推流回调
 }
 
-func (hook *HookConfig) IsEnablePublishEvent() bool {
+func (hook *HooksConfig) IsEnablePublishEvent() bool {
 	return hook.Enable && hook.OnPublishUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnPublishDone() bool {
+func (hook *HooksConfig) IsEnableOnPublishDone() bool {
 	return hook.Enable && hook.OnPublishDoneUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnPlay() bool {
+func (hook *HooksConfig) IsEnableOnPlay() bool {
 	return hook.Enable && hook.OnPlayUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnPlayDone() bool {
+func (hook *HooksConfig) IsEnableOnPlayDone() bool {
 	return hook.Enable && hook.OnPlayDoneUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnRecord() bool {
+func (hook *HooksConfig) IsEnableOnRecord() bool {
 	return hook.Enable && hook.OnRecordUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnIdleTimeout() bool {
+func (hook *HooksConfig) IsEnableOnIdleTimeout() bool {
 	return hook.Enable && hook.OnIdleTimeoutUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnReceiveTimeout() bool {
+func (hook *HooksConfig) IsEnableOnReceiveTimeout() bool {
 	return hook.Enable && hook.OnReceiveTimeoutUrl != ""
 }
 
-func (hook *HookConfig) IsEnableOnStarted() bool {
+func (hook *HooksConfig) IsEnableOnStarted() bool {
 	return hook.Enable && hook.OnStartedUrl != ""
 }
 
@@ -251,7 +252,7 @@ type AppConfig_ struct {
 	GB28181           GB28181Config
 	WebRtc            WebRtcConfig
 
-	Hook   HookConfig
+	Hooks  HooksConfig
 	Record RecordConfig
 	Log    LogConfig
 	Http   HttpConfig
@@ -292,7 +293,7 @@ func SetDefaultConfig(config_ *AppConfig_) {
 
 	config_.IdleTimeout *= int64(time.Second)
 	config_.ReceiveTimeout *= int64(time.Second)
-	config_.Hook.Timeout *= int64(time.Second)
+	config_.Hooks.Timeout *= int64(time.Second)
 }
 
 func limitMin(min, value int) int {

stream/hook.go

@@ -29,7 +29,7 @@ func responseBodyToString(resp *http.Response) string {
 
 func sendHookEvent(url string, body []byte) (*http.Response, error) {
 	client := &http.Client{
-		Timeout: time.Duration(AppConfig.Hook.Timeout),
+		Timeout: time.Duration(AppConfig.Hooks.Timeout),
 	}
 	request, err := http.NewRequest("post", url, bytes.NewBuffer(body))
 	if err != nil {
@@ -77,3 +77,15 @@ func NewHookPlayEventInfo(sink Sink) eventInfo {
 func NewHookPublishEventInfo(source Source) eventInfo {
 	return eventInfo{Stream: source.Id(), Protocol: source.Type().ToString(), RemoteAddr: source.RemoteAddr()}
 }
+
+func NewRecordEventInfo(source Source, path string) interface{} {
+	data := struct {
+		eventInfo
+		Path string `json:"path"`
+	}{
+		eventInfo: NewHookPublishEventInfo(source),
+		Path:      path,
+	}
+
+	return data
+}

stream/hook_event.go

@@ -21,14 +21,14 @@ var (
 
 func InitHookUrl() {
 	hookUrls = map[HookEvent]string{
-		HookEventPublish:        AppConfig.Hook.OnPublishUrl,
+		HookEventPublish:        AppConfig.Hooks.OnPublishUrl,
-		HookEventPublishDone:    AppConfig.Hook.OnPublishDoneUrl,
+		HookEventPublishDone:    AppConfig.Hooks.OnPublishDoneUrl,
-		HookEventPlay:           AppConfig.Hook.OnPlayUrl,
+		HookEventPlay:           AppConfig.Hooks.OnPlayUrl,
-		HookEventPlayDone:       AppConfig.Hook.OnPlayDoneUrl,
+		HookEventPlayDone:       AppConfig.Hooks.OnPlayDoneUrl,
-		HookEventRecord:         AppConfig.Hook.OnRecordUrl,
+		HookEventRecord:         AppConfig.Hooks.OnRecordUrl,
-		HookEventIdleTimeout:    AppConfig.Hook.OnIdleTimeoutUrl,
+		HookEventIdleTimeout:    AppConfig.Hooks.OnIdleTimeoutUrl,
-		HookEventReceiveTimeout: AppConfig.Hook.OnReceiveTimeoutUrl,
+		HookEventReceiveTimeout: AppConfig.Hooks.OnReceiveTimeoutUrl,
-		HookEventStarted:        AppConfig.Hook.OnStartedUrl,
+		HookEventStarted:        AppConfig.Hooks.OnStartedUrl,
 	}
 }
 

stream/hook_sink.go

@@ -9,7 +9,7 @@ import (
 func PreparePlaySink(sink Sink) (*http.Response, utils.HookState) {
 	var response *http.Response
 
-	if AppConfig.Hook.IsEnableOnPlay() {
+	if AppConfig.Hooks.IsEnableOnPlay() {
 		hook, err := Hook(HookEventPlay, sink.UrlValues().Encode(), NewHookPlayEventInfo(sink))
 		if err != nil {
 			log.Sugar.Errorf("通知播放事件失败 err:%s sink:%s-%v source:%s", err.Error(), sink.Protocol().ToString(), sink.Id(), sink.SourceId())
@@ -46,7 +46,7 @@ func PreparePlaySink(sink Sink) (*http.Response, utils.HookState) {
 func HookPlayDoneEvent(sink Sink) (*http.Response, bool) {
 	var response *http.Response
 
-	if AppConfig.Hook.IsEnableOnPlayDone() {
+	if AppConfig.Hooks.IsEnableOnPlayDone() {
 		hook, err := Hook(HookEventPlayDone, sink.UrlValues().Encode(), NewHookPlayEventInfo(sink))
 		if err != nil {
 			log.Sugar.Errorf("通知播放结束事件失败 err:%s sink:%s-%v source:%s", err.Error(), sink.Protocol().ToString(), sink.Id(), sink.SourceId())

stream/hook_source.go

@@ -10,7 +10,7 @@ import (
 func PreparePublishSource(source Source, hook bool) (*http.Response, utils.HookState) {
 	var response *http.Response
 
-	if hook && AppConfig.Hook.IsEnablePublishEvent() {
+	if hook && AppConfig.Hooks.IsEnablePublishEvent() {
 		rep, state := HookPublishEvent(source)
 		if utils.HookStateOK != state {
 			return rep, state
@@ -41,7 +41,7 @@ func PreparePublishSource(source Source, hook bool) (*http.Response, utils.HookS
 func HookPublishEvent(source Source) (*http.Response, utils.HookState) {
 	var response *http.Response
 
-	if AppConfig.Hook.IsEnablePublishEvent() {
+	if AppConfig.Hooks.IsEnablePublishEvent() {
 		hook, err := Hook(HookEventPublish, source.UrlValues().Encode(), NewHookPublishEventInfo(source))
 		if err != nil {
 			return hook, utils.HookStateFailure
@@ -54,7 +54,7 @@ func HookPublishEvent(source Source) (*http.Response, utils.HookState) {
 }
 
 func HookPublishDoneEvent(source Source) {
-	if AppConfig.Hook.IsEnablePublishEvent() {
+	if AppConfig.Hooks.IsEnablePublishEvent() {
 		_, _ = Hook(HookEventPublishDone, source.UrlValues().Encode(), NewHookPublishEventInfo(source))
 	}
 }
@@ -62,7 +62,7 @@ func HookPublishDoneEvent(source Source) {
 func HookReceiveTimeoutEvent(source Source) (*http.Response, utils.HookState) {
 	var response *http.Response
 
-	if AppConfig.Hook.IsEnableOnReceiveTimeout() {
+	if AppConfig.Hooks.IsEnableOnReceiveTimeout() {
 		resp, err := Hook(HookEventReceiveTimeout, source.UrlValues().Encode(), NewHookPublishEventInfo(source))
 		if err != nil {
 			return resp, utils.HookStateFailure
@@ -77,7 +77,7 @@ func HookReceiveTimeoutEvent(source Source) (*http.Response, utils.HookState) {
 func HookIdleTimeoutEvent(source Source) (*http.Response, utils.HookState) {
 	var response *http.Response
 
-	if AppConfig.Hook.IsEnableOnIdleTimeout() {
+	if AppConfig.Hooks.IsEnableOnIdleTimeout() {
 		resp, err := Hook(HookEventIdleTimeout, source.UrlValues().Encode(), NewHookPublishEventInfo(source))
 		if err != nil {
 			return resp, utils.HookStateFailure
@@ -88,3 +88,9 @@ func HookIdleTimeoutEvent(source Source) (*http.Response, utils.HookState) {
 
 	return response, utils.HookStateOK
 }
+
+func HookRecordEvent(source Source, path string) {
+	if AppConfig.Hooks.IsEnableOnRecord() {
+		_, _ = Hook(HookEventRecord, "", NewRecordEventInfo(source, path))
+	}
+}

stream/source.go

@@ -143,6 +143,7 @@ type PublishSource struct {
 
 	TransDeMuxer     stream.DeMuxer            //负责从推流协议中解析出AVStream和AVPacket
 	recordSink       Sink                      //每个Source的录制流
+	recordFilePath   string                    //录制流文件路径
 	hlsStream        TransStream               //HLS传输流, 如果开启, 在@seee writeHeader 直接创建, 如果等拉流时再创建, 会进一步加大HLS延迟.
 	audioTranscoders []transcode.Transcoder    //音频解码器
 	videoTranscoders []transcode.Transcoder    //视频解码器
@@ -210,7 +211,13 @@ func (s *PublishSource) CreateDefaultOutStreams() {
 
 	//创建录制流
 	if AppConfig.Record.Enable {
-
+		sink, path, err := CreateRecordStream(s.Id_)
+		if err != nil {
+			log.Sugar.Errorf("创建录制sink失败 source:%s err:%s", s.Id_, err.Error())
+		} else {
+			s.recordSink = sink
+			s.recordFilePath = path
+		}
 	}
 
 	//创建HLS输出流
@@ -415,8 +422,10 @@ func (s *PublishSource) AddSink(sink Sink) bool {
 		sink.SetState(SessionStateTransferring)
 	}
 
+	if s.recordSink != sink {
 		s.sinkCount++
 		log.Sugar.Infof("sink count:%d source:%s", s.sinkCount, s.Id_)
+	}
 
 	//新的传输流，发送缓存的音视频帧
 	if !ok && AppConfig.GOPCache && s.existVideo {
@@ -497,6 +506,10 @@ func (s *PublishSource) doClose() {
 		s.idleTimer.Stop()
 	}
 
+	if s.recordSink != nil {
+		s.recordSink.Close()
+	}
+
 	//释放解复用器
 	//释放转码器
 	//释放每路转协议流， 将所有sink添加到等待队列
@@ -510,6 +523,10 @@ func (s *PublishSource) doClose() {
 
 		transStream.PopAllSink(func(sink Sink) {
 			sink.SetTransStreamId(0)
+			if s.recordSink == sink {
+				return
+			}
+
 			{
 				sink.Lock()
 				defer sink.UnLock()
@@ -537,6 +554,10 @@ func (s *PublishSource) doClose() {
 		}
 
 		HookPublishDoneEvent(s)
+
+		if s.recordSink != nil {
+			HookRecordEvent(s, s.recordFilePath)
+		}
 	}()
 }
 
@@ -599,6 +620,10 @@ func (s *PublishSource) writeHeader() {
 	s.CreateDefaultOutStreams()
 
 	sinks := PopWaitingSinks(s.Id_)
+	if s.recordSink != nil {
+		sinks = append(sinks, s.recordSink)
+	}
+
 	for _, sink := range sinks {
 		if !s.AddSink(sink) {
 			sink.Close()

stream/stream_factory.go

@@ -7,8 +7,11 @@ import (
 
 type TransStreamFactory func(source Source, protocol Protocol, streams []utils.AVStream) (TransStream, error)
 
+type RecordStreamFactory func(source string) (Sink, string, error)
+
 var (
 	transStreamFactories map[Protocol]TransStreamFactory
+	recordStreamFactory  RecordStreamFactory
 )
 
 func init() {
@@ -41,3 +44,11 @@ func CreateTransStream(source Source, protocol Protocol, streams []utils.AVStrea
 
 	return factory(source, protocol, streams)
 }
+
+func SetRecordStreamFactory(factory RecordStreamFactory) {
+	recordStreamFactory = factory
+}
+
+func CreateRecordStream(sourceId string) (Sink, string, error) {
+	return recordStreamFactory(sourceId)
+}

stream/trans_stream.go

@@ -109,7 +109,9 @@ func (t *TCPTransStream) AddSink(sink Sink) error {
 		return err
 	}
 
+	if sink.GetConn() != nil {
 		sink.GetConn().(*transport.Conn).EnableAsyncWriteMode(AppConfig.WriteBufferNumber - 1)
+	}
 	return nil
 }