将读写锁替换为WaitGroup

Author: langhuihui

ring.go

@@ -9,7 +9,7 @@ import (
 
 type RingItem struct {
 	avformat.AVPacket
-	sync.RWMutex
+	sync.WaitGroup
 	*bytes.Buffer
 }
 
@@ -27,7 +27,7 @@ func NewRing(exp int) (r *Ring) {
 	r.Size = 1 << exp
 	r.buffer = make([]RingItem, r.Size)
 	r.RingItem = &r.buffer[0]
-	r.Lock()
+	r.Add(1)
 	return
 }
 func (r *Ring) offset(v int) int {
@@ -71,14 +71,8 @@ func (r *Ring) GoBack() {
 func (r *Ring) NextW() {
 	item := r.RingItem
 	r.GoNext()
-	r.RingItem.Lock()
-	item.Unlock()
-}
-
-// NextR 读下一个
-func (r *Ring) NextR() {
-	r.RingItem.RUnlock()
-	r.GoNext()
+	r.RingItem.Add(1)
+	item.Done()
 }
 
 func (r *Ring) GetBuffer() *bytes.Buffer {

stream.go

@@ -29,10 +29,9 @@ func FindStream(streamPath string) *Stream {
 //GetStream 根据流路径获取流，如果不存在则创建一个新的
 func GetStream(streamPath string) (result *Stream) {
 	item, loaded := streamCollection.LoadOrStore(streamPath, &Stream{
-		Subscribers:  make(map[string]*Subscriber),
-		Control:      make(chan interface{}),
-		AVRing:       NewRing(config.RingSize),
-		WaitingMutex: new(sync.RWMutex),
+		Subscribers: make(map[string]*Subscriber),
+		Control:     make(chan interface{}),
+		AVRing:      NewRing(config.RingSize),
 		StreamInfo: StreamInfo{
 			StreamPath:     streamPath,
 			SubscriberInfo: make([]*SubscriberInfo, 0),
@@ -42,7 +41,7 @@ func GetStream(streamPath string) (result *Stream) {
 	if !loaded {
 		Summary.Streams = append(Summary.Streams, &result.StreamInfo)
 		result.Context, result.Cancel = context.WithCancel(context.Background())
-		result.WaitingMutex.Lock() //等待发布者
+		result.WaitPub.Add(1) //等待发布者
 		go result.Run()
 	}
 	return
@@ -60,7 +59,7 @@ type Stream struct {
 	AudioTag     *AVPacket              // 每个音频包都是这样的结构,区别在于Payload的大小.FMS在发送AAC sequence header,需要加上 AudioTags,这个tag 1个字节(8bits)的数据
 	FirstScreen  *Ring                  //最近的关键帧位置，首屏渲染
 	AVRing       *Ring                  //数据环
-	WaitingMutex *sync.RWMutex          //用于订阅和等待发布者
+	WaitPub      sync.WaitGroup         //用于订阅和等待发布者
 	UseTimestamp bool                   //是否采用数据包中的时间戳
 }
 
@@ -326,7 +325,7 @@ func (r *Stream) PushVideo(timestamp uint32, payload []byte) {
 		}
 		if video.IsKeyFrame {
 			if r.FirstScreen == nil {
-				defer r.WaitingMutex.Unlock()
+				defer r.WaitPub.Done()
 				r.FirstScreen = video.Clone()
 			} else {
 				oldNumber := r.FirstScreen.Number

subscriber.go

@@ -29,6 +29,8 @@ type Subscriber struct {
 	Cancel     context.CancelFunc
 	Sign       string
 	OffsetTime uint32
+	startTime  uint32
+	avformat.SendPacket
 }
 
 // IsClosed 检查订阅者是否已经关闭
@@ -52,55 +54,45 @@ func (s *Subscriber) Subscribe(streamPath string) (err error) {
 	}
 	GetStream(streamPath).Subscribe(s)
 	defer s.UnSubscribe(s)
-	//加锁解锁的目的是等待发布者首屏数据，如果发布者尚为发布，则会等待，否则就会往下执行
-	s.WaitingMutex.RLock()
-	s.WaitingMutex.RUnlock()
-	sendPacket := avformat.NewSendPacket(s.VideoTag, 0)
-	defer sendPacket.Recycle()
-	s.OnData(sendPacket)
+	//等待发布者首屏数据，如果发布者尚为发布，则会等待，否则就会往下执行
+	s.WaitPub.Wait()
+	s.sendAv(s.VideoTag, 0)
 	packet := s.FirstScreen.Clone()
-	startTime := packet.Timestamp
-	packet.RLock()
-	sendPacket.AVPacket = &packet.AVPacket
-	s.OnData(sendPacket)
-	packet.NextR()
-	atsent := false
-	dropping := false
-	droped := 0
-	for {
-		select {
-		case <-s.Done():
-			return s.Err()
-		default:
-			s.TotalPacket++
-			packet.RLock()
-			if !dropping {
-				if packet.Type == avformat.FLV_TAG_TYPE_AUDIO && !atsent {
-					sendPacket.AVPacket = s.AudioTag
-					sendPacket.Timestamp = 0
-					s.OnData(sendPacket)
-					atsent = true
-				}
-				sendPacket.AVPacket = &packet.AVPacket
-				sendPacket.Timestamp = packet.Timestamp - startTime
-				s.OnData(sendPacket)
-				if s.checkDrop(packet) {
-					dropping = true
-					droped = 0
-				}
-				packet.NextR()
-			} else if packet.IsKeyFrame {
-				//遇到关键帧则退出丢帧
-				dropping = false
-				//fmt.Println("drop package ", droped)
-				s.TotalDrop += droped
-				packet.RUnlock()
-			} else {
-				droped++
-				packet.NextR()
+	s.startTime = packet.Timestamp
+	s.send(packet)
+	packet.GoNext()
+	for atsent, dropping, droped := false, false, 0; s.Err() == nil; packet.GoNext() {
+		s.TotalPacket++
+		if !dropping {
+			if packet.Type == avformat.FLV_TAG_TYPE_AUDIO && !atsent {
+				s.sendAv(s.AudioTag, 0)
+				atsent = true
 			}
+			s.send(packet)
+			if s.checkDrop(packet) {
+				dropping = true
+				droped = 0
+			}
+		} else if packet.IsKeyFrame {
+			//遇到关键帧则退出丢帧
+			dropping = false
+			//fmt.Println("drop package ", droped)
+			s.TotalDrop += droped
+			s.send(packet)
+		} else {
+			droped++
 		}
 	}
+	return s.Err()
+}
+func (s *Subscriber) sendAv(packet *avformat.AVPacket, t uint32) {
+	s.AVPacket = packet
+	s.Timestamp = t
+	s.OnData(&s.SendPacket)
+}
+func (s *Subscriber) send(packet *Ring) {
+	packet.Wait()
+	s.sendAv(&packet.AVPacket, packet.Timestamp-s.startTime)
 }
 func (s *Subscriber) checkDrop(packet *Ring) bool {
 	pIndex := s.AVRing.Index