github.com/cnotch/ipchub@v1.1.0/media/cache/hevccache.go (about)

     1  // Copyright (c) 2019,CAOHONGJU All rights reserved.
     2  // Use of this source code is governed by a MIT-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package cache
     6  
     7  import (
     8  	"sync"
     9  
    10  	"github.com/cnotch/ipchub/av/codec/hevc"
    11  	"github.com/cnotch/ipchub/av/format/rtp"
    12  	"github.com/cnotch/queue"
    13  )
    14  
    15  // HevcCache 画面组缓存(Group of Pictures).
    16  type HevcCache struct {
    17  	cacheGop bool
    18  	l        sync.RWMutex
    19  	gop      queue.Queue
    20  	vps      *rtp.Packet // 视频参数集包
    21  	sps      *rtp.Packet // 序列参数集包
    22  	pps      *rtp.Packet // 图像参数集包
    23  }
    24  
    25  // NewHevcCache 创建 HEVC 缓存
    26  func NewHevcCache(cacheGop bool) *HevcCache {
    27  	return &HevcCache{
    28  		cacheGop: cacheGop,
    29  	}
    30  }
    31  
    32  // CachePack 向HevcCache中缓存包
    33  func (cache *HevcCache) CachePack(pack Pack) bool {
    34  	rtppack := pack.(*rtp.Packet)
    35  
    36  	if rtppack.Channel != rtp.ChannelVideo {
    37  		return false
    38  	}
    39  
    40  	// 判断是否是参数和关键帧包
    41  	vps, sps, pps, islice := cache.getPalyloadType(rtppack.Payload())
    42  
    43  	cache.l.Lock()
    44  	defer cache.l.Unlock()
    45  
    46  	if vps { // 视频参数
    47  		cache.vps = rtppack
    48  		return false
    49  	}
    50  
    51  	if sps { // 序列头参数
    52  		cache.sps = rtppack
    53  		return false
    54  	}
    55  
    56  	if pps { // 图像参数
    57  		cache.pps = rtppack
    58  		return false
    59  	}
    60  
    61  	if cache.cacheGop { // 需要缓存 GOP
    62  		if islice { // 关键帧
    63  			cache.gop.Reset()
    64  			cache.gop.Push(rtppack)
    65  		} else if cache.gop.Len() > 0 {
    66  			cache.gop.Push(rtppack)
    67  		}
    68  	}
    69  	return islice
    70  }
    71  
    72  // Reset 重置HevcCache缓存
    73  func (cache *HevcCache) Reset() {
    74  	cache.l.Lock()
    75  	defer cache.l.Unlock()
    76  
    77  	cache.vps = nil
    78  	cache.sps = nil
    79  	cache.pps = nil
    80  	cache.gop.Reset()
    81  }
    82  
    83  // PushTo 入列到指定的队列
    84  func (cache *HevcCache) PushTo(q *queue.SyncQueue) int {
    85  	bytes := 0
    86  	cache.l.RLock()
    87  	defer cache.l.RUnlock()
    88  
    89  	// 写参数包
    90  	if cache.vps != nil {
    91  		q.Queue().Push(cache.vps)
    92  		bytes += cache.vps.Size()
    93  	}
    94  
    95  	if cache.sps != nil {
    96  		q.Queue().Push(cache.sps)
    97  		bytes += cache.sps.Size()
    98  	}
    99  
   100  	if cache.pps != nil {
   101  		q.Queue().Push(cache.pps)
   102  		bytes += cache.pps.Size()
   103  	}
   104  
   105  	// 如果必要，写 GopCache
   106  	if cache.cacheGop {
   107  		packs := cache.gop.Elems()
   108  		q.Queue().PushN(packs) // 启动阶段调用，无需加锁
   109  		for _, p := range packs {
   110  			bytes += p.(Pack).Size()
   111  		}
   112  	}
   113  
   114  	return bytes
   115  }
   116  
   117  func (cache *HevcCache) getPalyloadType(payload []byte) (vps, sps, pps, islice bool) {
   118  	if len(payload) < 3 {
   119  		return
   120  	}
   121  	naluType := (payload[0] >> 1) & 0x3f
   122  
   123  	switch naluType {
   124  	case hevc.NalStapInRtp: // 在RTP中的聚合（AP）
   125  		off := 2
   126  		// 循环读取被封装的NAL
   127  		for {
   128  			// nal长度
   129  			nalSize := ((uint16(payload[off])) << 8) | uint16(payload[off+1])
   130  			if nalSize < 1 {
   131  				return
   132  			}
   133  
   134  			off += 2
   135  			naluType = (payload[off] >> 1) & 0x3f
   136  			cache.nalType(naluType, &vps, &sps, &pps, &islice)
   137  			off += int(nalSize)
   138  
   139  			if off >= len(payload) { // 扫描完成
   140  				break
   141  			}
   142  		}
   143  		return
   144  	case hevc.NalFuInRtp: // 在RTP中的扩展,分片(FU)
   145  		naluType = payload[2] & 0x3f
   146  		if (payload[2]>>7)&1 == 1 { // 第一个分片
   147  			cache.nalType(naluType, &vps, &sps, &pps, &islice)
   148  		}
   149  		return
   150  	default:
   151  		cache.nalType(naluType, &vps, &sps, &pps, &islice)
   152  		return
   153  	}
   154  }
   155  
   156  func (cache *HevcCache) nalType(nalType byte, vps, sps, pps, islice *bool) {
   157  	if nalType >= hevc.NalBlaWLp && nalType <= hevc.NalCraNut {
   158  		*islice = true
   159  		return
   160  	}
   161  
   162  	switch nalType {
   163  	case hevc.NalVps:
   164  		*vps = true
   165  	case hevc.NalSps:
   166  		*sps = true
   167  	case hevc.NalPps:
   168  		*pps = true
   169  	}
   170  	return
   171  }