From e3a01a580af9713822cbda93415f9af8ce91b76e Mon Sep 17 00:00:00 2001
From: xia-chu <771730766@qq.com>
Date: Sat, 9 Dec 2023 16:48:21 +0800
Subject: [PATCH] =?UTF-8?q?ps/ts=E8=A7=A3=E6=9E=90=E6=94=AF=E6=8C=81?=
 =?UTF-8?q?=E5=A4=9Atrack?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/Extension/Frame.h |  7 +++++++
 src/Rtp/Decoder.cpp   | 39 ++++++++++++++++++++++-----------------
 src/Rtp/Decoder.h     | 13 +++++++++----
 3 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/src/Extension/Frame.h b/src/Extension/Frame.h
index 1bdf526c..e691b1e3 100644
--- a/src/Extension/Frame.h
+++ b/src/Extension/Frame.h
@@ -124,6 +124,13 @@ public:
      */
     TrackType getTrackType() const;
     std::string getTrackTypeStr() const;
+
+    void setIndex(int index) { _index = index; }
+
+    int getIndex() const { return _index < 0 ? (int)getTrackType() : _index; }
+
+private:
+    int _index = -1;
 };
 
 /**
diff --git a/src/Rtp/Decoder.cpp b/src/Rtp/Decoder.cpp
index dc0ef749..840bac24 100644
--- a/src/Rtp/Decoder.cpp
+++ b/src/Rtp/Decoder.cpp
@@ -62,7 +62,9 @@ DecoderImp::Ptr DecoderImp::createDecoder(Type type, MediaSinkInterface *sink){
 }
 
 void DecoderImp::flush() {
-    _merger.flush();
+    for (auto &pr : _tracks) {
+        pr.second.second.flush();
+    }
 }
 
 ssize_t DecoderImp::input(const uint8_t *data, size_t bytes){
@@ -88,9 +90,9 @@ void DecoderImp::onStream(int stream, int codecid, const void *extra, size_t byt
     if (!track) {
         return;
     }
-    onTrack(std::move(track));
+    onTrack(stream, std::move(track));
     // 防止未获取视频track提前complete导致忽略后续视频的问题，用于兼容一些不太规范的ps流
-    if (finish && _tracks[TrackVideo]) {
+    if (finish && _tracks[stream].first) {
         _sink->addTrackCompleted();
         InfoL << "Add track finished";
     }
@@ -104,37 +106,40 @@ void DecoderImp::onDecode(int stream, int codecid, int flags, int64_t pts, int64
     if (codec == CodecInvalid) {
         return;
     }
-    if (!_tracks[getTrackType(codec)]) {
-        onTrack(Factory::getTrackByCodecId(codec, 8000, 1, 16));
+    auto &ref = _tracks[stream];
+    if (!ref.first) {
+        onTrack(stream, Factory::getTrackByCodecId(codec, 8000, 1, 16));
     }
-    // TODO 支持多track
-    auto frame = Factory::getFrameFromPtr(codec, (char *) data, bytes, dts, pts);
-    if (getTrackType(codec) == TrackVideo) {
-        _merger.inputFrame(frame, [&](uint64_t dts, uint64_t pts, const Buffer::Ptr &buffer, bool) {
-            onFrame(Factory::getFrameFromBuffer(codec, buffer, dts, pts));
-        });
-    } else {
-        onFrame(frame);
+    auto frame = Factory::getFrameFromPtr(codec, (char *)data, bytes, dts, pts);
+    if (getTrackType(codec) != TrackVideo) {
+        onFrame(stream, frame);
+        return;
     }
+    ref.second.inputFrame(frame, [&](uint64_t dts, uint64_t pts, const Buffer::Ptr &buffer, bool) {
+        onFrame(stream, Factory::getFrameFromBuffer(codec, buffer, dts, pts));
+    });
 }
 #else
 void DecoderImp::onDecode(int stream,int codecid,int flags,int64_t pts,int64_t dts,const void *data,size_t bytes) {}
 void DecoderImp::onStream(int stream,int codecid,const void *extra,size_t bytes,int finish) {}
 #endif
 
-void DecoderImp::onTrack(const Track::Ptr &track) {
+void DecoderImp::onTrack(int index, const Track::Ptr &track) {
     if (!track) {
         return;
     }
-    if (!_tracks[track->getTrackType()]) {
-        _tracks[track->getTrackType()] = track;
+    track->setIndex(index);
+    auto &ref = _tracks[index];
+    if (!ref.first) {
+        ref.first = track;
         _sink->addTrack(track);
         InfoL << "got track: " << track->getCodecName();
     }
 }
 
-void DecoderImp::onFrame(const Frame::Ptr &frame) {
+void DecoderImp::onFrame(int index, const Frame::Ptr &frame) {
     if (frame) {
+        frame->setIndex(index);
         _sink->inputFrame(frame);
     }
 }
diff --git a/src/Rtp/Decoder.h b/src/Rtp/Decoder.h
index 55012870..48a72a78 100644
--- a/src/Rtp/Decoder.h
+++ b/src/Rtp/Decoder.h
@@ -11,6 +11,7 @@
 #ifndef ZLMEDIAKIT_DECODER_H
 #define ZLMEDIAKIT_DECODER_H
 
+#include <map>
 #include <stdint.h>
 #include <memory>
 #include <functional>
@@ -48,8 +49,8 @@ public:
     void flush();
 
 protected:
-    void onTrack(const Track::Ptr &track);
-    void onFrame(const Frame::Ptr &frame);
+    void onTrack(int index, const Track::Ptr &track);
+    void onFrame(int index, const Frame::Ptr &frame);
 
 private:
     DecoderImp(const Decoder::Ptr &decoder, MediaSinkInterface *sink);
@@ -59,8 +60,12 @@ private:
 private:
     Decoder::Ptr _decoder;
     MediaSinkInterface *_sink;
-    FrameMerger _merger{FrameMerger::none};
-    Track::Ptr _tracks[TrackMax];
+
+    class FrameMergerImp : public FrameMerger {
+    public:
+        FrameMergerImp() : FrameMerger(FrameMerger::none) {}
+    };
+    std::map<int, std::pair<Track::Ptr, FrameMergerImp> > _tracks;
 };
 
 }//namespace mediakit