From 7f3877e4378285f25c2fa87885497cc644ac1618 Mon Sep 17 00:00:00 2001
From: xiongziliang <771730766@qq.com>
Date: Thu, 1 Oct 2020 11:02:00 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E7=BD=91=E9=80=9F=E7=BB=9F?=
 =?UTF-8?q?=E8=AE=A1=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/WebApi.cpp           |  1 +
 src/Common/MediaSource.cpp  |  4 +++
 src/Common/MediaSource.h    | 52 +++++++++++++++++++++++++++++++++++++
 src/FMP4/FMP4MediaSource.h  |  1 +
 src/Record/HlsMakerImp.cpp  |  3 +++
 src/Record/HlsMediaSource.h |  4 +++
 src/Rtmp/RtmpMediaSource.h  |  1 +
 src/Rtsp/RtspMediaSource.h  |  1 +
 src/TS/TSMediaSource.h      |  1 +
 9 files changed, 68 insertions(+)

diff --git a/server/WebApi.cpp b/server/WebApi.cpp
index 4fe3eef8..44c4549d 100644
--- a/server/WebApi.cpp
+++ b/server/WebApi.cpp
@@ -395,6 +395,7 @@ void installWebApi() {
         item["vhost"] = media->getVhost();
         item["app"] = media->getApp();
         item["stream"] = media->getId();
+        item["bytes_speed"] = media->getBytesSpeed();
         item["readerCount"] = media->readerCount();
         item["totalReaderCount"] = media->totalReaderCount();
         item["originType"] = (int) media->getOriginType();
diff --git a/src/Common/MediaSource.cpp b/src/Common/MediaSource.cpp
index 6f5f564e..295f554f 100644
--- a/src/Common/MediaSource.cpp
+++ b/src/Common/MediaSource.cpp
@@ -66,6 +66,10 @@ const string& MediaSource::getId() const {
     return _stream_id;
 }
 
+int MediaSource::getBytesSpeed(){
+    return _speed.getSpeed();
+}
+
 vector<Track::Ptr> MediaSource::getTracks(bool ready) const {
     auto listener = _listener.lock();
     if(!listener){
diff --git a/src/Common/MediaSource.h b/src/Common/MediaSource.h
index 039cf127..bd43b04c 100644
--- a/src/Common/MediaSource.h
+++ b/src/Common/MediaSource.h
@@ -137,6 +137,52 @@ public:
     string _param_strs;
 };
 
+class BytesSpeed {
+public:
+    BytesSpeed() = default;
+    ~BytesSpeed() = default;
+
+    /**
+     * 添加统计字节
+     */
+    BytesSpeed& operator += (uint64_t bytes) {
+        _bytes += bytes;
+        if (_bytes > 1024 * 1024) {
+            //数据大于1MB就计算一次网速
+            computeSpeed();
+        }
+        return *this;
+    }
+
+    /**
+     * 获取速度，单位bytes/s
+     */
+    int getSpeed() {
+        if (_ticker.elapsedTime() < 1000) {
+            //获取频率小于1秒，那么返回上次计算结果
+            return _speed;
+        }
+        return computeSpeed();
+    }
+
+private:
+    uint64_t computeSpeed() {
+        auto elapsed = _ticker.elapsedTime();
+        if (!elapsed) {
+            return _speed;
+        }
+        _speed = _bytes * 1000 / elapsed;
+        _ticker.resetTime();
+        _bytes = 0;
+        return _speed;
+    }
+
+private:
+    int _speed = 0;
+    uint64_t _bytes = 0;
+    Ticker _ticker;
+};
+
 /**
  * 媒体源，任何rtsp/rtmp的直播流都源自该对象
  */
@@ -170,6 +216,9 @@ public:
     // 设置时间戳
     virtual void setTimeStamp(uint32_t stamp) {};
 
+    // 获取数据速率，单位bytes/s
+    int getBytesSpeed();
+
     ////////////////MediaSourceEvent相关接口实现////////////////
 
     // 设置监听者
@@ -229,6 +278,9 @@ private:
     //触发媒体事件
     void emitEvent(bool regist);
 
+protected:
+    BytesSpeed _speed;
+
 private:
     string _schema;
     string _vhost;
diff --git a/src/FMP4/FMP4MediaSource.h b/src/FMP4/FMP4MediaSource.h
index 277d4bf8..173f12e4 100644
--- a/src/FMP4/FMP4MediaSource.h
+++ b/src/FMP4/FMP4MediaSource.h
@@ -99,6 +99,7 @@ public:
         if (key) {
             _have_video = true;
         }
+        _speed += packet->size();
         PacketCache<FMP4Packet, FMP4FlushPolicy>::inputPacket(true, packet, key);
     }
 
diff --git a/src/Record/HlsMakerImp.cpp b/src/Record/HlsMakerImp.cpp
index 93959c1b..4cd9cf76 100644
--- a/src/Record/HlsMakerImp.cpp
+++ b/src/Record/HlsMakerImp.cpp
@@ -92,6 +92,9 @@ void HlsMakerImp::onWriteSegment(const char *data, int len) {
     if (_file) {
         fwrite(data, len, 1, _file.get());
     }
+    if (_media_src) {
+        _media_src->onSegmentSize(len);
+    }
 }
 
 void HlsMakerImp::onWriteHls(const char *data, int len) {
diff --git a/src/Record/HlsMediaSource.h b/src/Record/HlsMediaSource.h
index b5b8357f..32e46cc5 100644
--- a/src/Record/HlsMediaSource.h
+++ b/src/Record/HlsMediaSource.h
@@ -79,6 +79,10 @@ public:
         _list_cb.emplace_back(std::move(cb));
     }
 
+    void onSegmentSize(uint64_t bytes) {
+        _speed += bytes;
+    }
+
 private:
     bool _is_regist = false;
     RingType::Ptr _ring;
diff --git a/src/Rtmp/RtmpMediaSource.h b/src/Rtmp/RtmpMediaSource.h
index cd3f6b5c..977f77de 100644
--- a/src/Rtmp/RtmpMediaSource.h
+++ b/src/Rtmp/RtmpMediaSource.h
@@ -119,6 +119,7 @@ public:
      * @param pkt rtmp包
      */
     void onWrite(const RtmpPacket::Ptr &pkt, bool = true) override {
+        _speed += pkt->size();
         //保存当前时间戳
         switch (pkt->type_id) {
             case MSG_VIDEO : _track_stamps[TrackVideo] = pkt->time_stamp, _have_video = true; break;
diff --git a/src/Rtsp/RtspMediaSource.h b/src/Rtsp/RtspMediaSource.h
index a67a4b34..341aca5c 100644
--- a/src/Rtsp/RtspMediaSource.h
+++ b/src/Rtsp/RtspMediaSource.h
@@ -157,6 +157,7 @@ public:
      * @param keyPos 该包是否为关键帧的第一个包
      */
     void onWrite(const RtpPacket::Ptr &rtp, bool keyPos) override {
+        _speed += rtp->size();
         assert(rtp->type >= 0 && rtp->type < TrackMax);
         auto track = _tracks[rtp->type];
         if (track) {
diff --git a/src/TS/TSMediaSource.h b/src/TS/TSMediaSource.h
index c7b05f53..2a5a7080 100644
--- a/src/TS/TSMediaSource.h
+++ b/src/TS/TSMediaSource.h
@@ -76,6 +76,7 @@ public:
      * @param key 是否为关键帧第一个包
      */
     void onWrite(const TSPacket::Ptr &packet, bool key) override {
+        _speed += packet->size();
         if (!_ring) {
             createRing();
         }