增加推流码流统计

api.go

@@ -405,7 +405,7 @@ func (api *ApiServer) OnSourceList(w http.ResponseWriter, r *http.Request) {
 			Protocol:  source.GetType().String(),
 			Time:      source.CreateTime(),
 			SinkCount: source.SinkCount(),
-			Bitrate:   "", // 后续开发
+			Bitrate:   strconv.Itoa(source.GetBitrateStatistics().PreviousSecond()/1024) + "KBS", // 后续开发
 			Tracks:    tracks,
 		})
 	}

main.go

@@ -135,6 +135,7 @@ func init() {
 		"record":  &config.Record,
 	}
 
+	// 读取运行参数
 	disableOptions, enableOptions := readRunArgs()
 	mergeArgs(options, disableOptions, enableOptions)
 

stream/bitrate_statistics.go

@@ -0,0 +1,56 @@
+package stream
+
+import "time"
+
+// BitrateStatistics 码流统计, 单位Byte
+type BitrateStatistics struct {
+	totalBytes     int64 // 总共传输的字节数
+	elapsedSeconds int   // 经过的秒数
+	currentSecond  int   // 当前秒数
+
+	previousSecondBytes int // 前一秒传输的字节数
+	latestSecondBytes   int // 当前秒正在传输的字节数
+}
+
+func (b *BitrateStatistics) Input(size int) {
+	b.totalBytes += int64(size)
+
+	second := time.Now().Second()
+	if b.currentSecond == -1 {
+		b.currentSecond = second
+	}
+
+	if second != b.currentSecond {
+		b.elapsedSeconds++
+		b.currentSecond = second
+		b.previousSecondBytes = b.latestSecondBytes
+		b.latestSecondBytes = 0
+	}
+
+	b.latestSecondBytes += size
+}
+
+// Average 返回每秒平均码流大小
+func (b *BitrateStatistics) Average() int {
+	if b.elapsedSeconds < 1 {
+		return b.latestSecondBytes
+	}
+
+	return int((b.totalBytes - int64(b.latestSecondBytes)) / int64(b.elapsedSeconds))
+}
+
+// Total 返回总码流大小
+func (b *BitrateStatistics) Total() int64 {
+	return b.totalBytes
+}
+
+// PreviousSecond 返回前一秒的码流大小
+func (b *BitrateStatistics) PreviousSecond() int {
+	return b.previousSecondBytes
+}
+
+func NewBitrateStatistics() *BitrateStatistics {
+	return &BitrateStatistics{
+		currentSecond: -1,
+	}
+}

stream/source.go

@@ -117,6 +117,8 @@ type Source interface {
 	SetCreateTime(time time.Time)
 
 	Sinks() []Sink
+
+	GetBitrateStatistics() *BitrateStatistics
 }
 
 type PublishSource struct {
@@ -154,7 +156,8 @@ type PublishSource struct {
 	sinkCount         int        // 拉流端计数
 	urlValues         url.Values // 推流url携带的参数
 	timeoutTracks     []int
-	createTime        time.Time // source创建时间
+	createTime        time.Time          // source创建时间
+	statistics        *BitrateStatistics // 码流统计
 }
 
 func (s *PublishSource) SetLastPacketTime(time2 time.Time) {
@@ -204,6 +207,7 @@ func (s *PublishSource) Init(receiveQueueSize int) {
 	s.TransStreams = make(map[TransStreamID]TransStream, 10)
 	s.sinks = make(map[SinkID]Sink, 128)
 	s.TransStreamSinks = make(map[TransStreamID]map[SinkID]Sink, len(transStreamFactories)+1)
+	s.statistics = NewBitrateStatistics()
 }
 
 func (s *PublishSource) CreateDefaultOutStreams() {
@@ -263,6 +267,7 @@ func (s *PublishSource) FindOrCreatePacketBuffer(index int, mediaType utils.AVMe
 
 func (s *PublishSource) Input(data []byte) error {
 	s.streamPipe <- data
+	s.statistics.Input(len(data))
 	return nil
 }
 
@@ -348,7 +353,7 @@ func (s *PublishSource) write(sink Sink, index int, data [][]byte, timestamp int
 		//return
 	}
 
-	// 推流失败, 可能是服务器或拉流端带宽不够、拉流端不读取数据等情况造成内核发送缓冲区满, 进而阻塞.
+	// 推流超时, 可能是服务器或拉流端带宽不够、拉流端不读取数据等情况造成内核发送缓冲区满, 进而阻塞.
 	// 直接关闭连接. 当然也可以将sink先挂起, 后续再继续推流.
 	_, ok := err.(*transport.ZeroWindowSizeError)
 	if ok {
@@ -845,3 +850,7 @@ func (s *PublishSource) Sinks() []Sink {
 	group.Wait()
 	return sinks
 }
+
+func (s *PublishSource) GetBitrateStatistics() *BitrateStatistics {
+	return s.statistics
+}