Use just a lookup map of RTP modules in PacketRouter

Since SSRCs of RTP modules are now set at construction time, we can
use just a simple unordered map from SSRC to module in packet router.

Bug: webrtc:11036
Change-Id: I0b3527f17c9ee2df9253c778e5b9e3651a70b355
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/155965
Commit-Queue: Erik Språng <sprang@webrtc.org>
Reviewed-by: Sebastian Jansson <srte@webrtc.org>
Reviewed-by: Danil Chapovalov <danilchap@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#29510}
diff --git a/modules/pacing/packet_router.cc b/modules/pacing/packet_router.cc
index 56922b7..8edfd1f 100644
--- a/modules/pacing/packet_router.cc
+++ b/modules/pacing/packet_router.cc
@@ -44,7 +44,8 @@
       transport_seq_(start_transport_seq) {}
 
 PacketRouter::~PacketRouter() {
-  RTC_DCHECK(rtp_send_modules_.empty());
+  RTC_DCHECK(send_modules_map_.empty());
+  RTC_DCHECK(send_modules_list_.empty());
   RTC_DCHECK(rtcp_feedback_senders_.empty());
   RTC_DCHECK(sender_remb_candidates_.empty());
   RTC_DCHECK(receiver_remb_candidates_.empty());
@@ -53,14 +54,17 @@
 
 void PacketRouter::AddSendRtpModule(RtpRtcp* rtp_module, bool remb_candidate) {
   rtc::CritScope cs(&modules_crit_);
-  RTC_DCHECK(std::find(rtp_send_modules_.begin(), rtp_send_modules_.end(),
-                       rtp_module) == rtp_send_modules_.end());
-  // Put modules which can use regular payload packets (over rtx) instead of
-  // padding first as it's less of a waste
+
+  AddSendRtpModuleToMap(rtp_module, rtp_module->SSRC());
+  if (absl::optional<uint32_t> rtx_ssrc = rtp_module->RtxSsrc()) {
+    AddSendRtpModuleToMap(rtp_module, *rtx_ssrc);
+  }
+  if (absl::optional<uint32_t> flexfec_ssrc = rtp_module->FlexfecSsrc()) {
+    AddSendRtpModuleToMap(rtp_module, *flexfec_ssrc);
+  }
+
   if (rtp_module->SupportsRtxPayloadPadding()) {
-    rtp_send_modules_.push_front(rtp_module);
-  } else {
-    rtp_send_modules_.push_back(rtp_module);
+    last_send_module_ = rtp_module;
   }
 
   if (remb_candidate) {
@@ -68,14 +72,32 @@
   }
 }
 
+void PacketRouter::AddSendRtpModuleToMap(RtpRtcp* rtp_module, uint32_t ssrc) {
+  RTC_DCHECK(send_modules_map_.find(ssrc) == send_modules_map_.end());
+  send_modules_list_.push_front(rtp_module);
+  send_modules_map_[ssrc] = std::pair<RtpRtcp*, std::list<RtpRtcp*>::iterator>(
+      rtp_module, send_modules_list_.begin());
+}
+
+void PacketRouter::RemoveSendRtpModuleFromMap(uint32_t ssrc) {
+  auto kv = send_modules_map_.find(ssrc);
+  RTC_DCHECK(kv != send_modules_map_.end());
+  send_modules_list_.erase(kv->second.second);
+  send_modules_map_.erase(kv);
+}
+
 void PacketRouter::RemoveSendRtpModule(RtpRtcp* rtp_module) {
   rtc::CritScope cs(&modules_crit_);
-  rtp_module_cache_map_.clear();
   MaybeRemoveRembModuleCandidate(rtp_module, /* media_sender = */ true);
-  auto it =
-      std::find(rtp_send_modules_.begin(), rtp_send_modules_.end(), rtp_module);
-  RTC_DCHECK(it != rtp_send_modules_.end());
-  rtp_send_modules_.erase(it);
+
+  RemoveSendRtpModuleFromMap(rtp_module->SSRC());
+  if (absl::optional<uint32_t> rtx_ssrc = rtp_module->RtxSsrc()) {
+    RemoveSendRtpModuleFromMap(*rtx_ssrc);
+  }
+  if (absl::optional<uint32_t> flexfec_ssrc = rtp_module->FlexfecSsrc()) {
+    RemoveSendRtpModuleFromMap(*flexfec_ssrc);
+  }
+
   if (last_send_module_ == rtp_module) {
     last_send_module_ = nullptr;
   }
@@ -105,25 +127,6 @@
   rtcp_feedback_senders_.erase(it);
 }
 
-RtpRtcp* PacketRouter::FindRtpModule(uint32_t ssrc) {
-  auto it = rtp_module_cache_map_.find(ssrc);
-  if (it != rtp_module_cache_map_.end()) {
-    if (ssrc == it->second->SSRC() || ssrc == it->second->FlexfecSsrc()) {
-      return it->second;
-    }
-    // This entry is stale due to a changed ssrc - remove it.
-    rtp_module_cache_map_.erase(it);
-  }
-  // Slow path - find and cache matching module
-  for (RtpRtcp* rtp_module : rtp_send_modules_) {
-    if (ssrc == rtp_module->SSRC() || ssrc == rtp_module->FlexfecSsrc()) {
-      rtp_module_cache_map_[ssrc] = rtp_module;
-      return rtp_module;
-    }
-  }
-  return nullptr;
-}
-
 void PacketRouter::SendPacket(std::unique_ptr<RtpPacketToSend> packet,
                               const PacedPacketInfo& cluster_info) {
   rtc::CritScope cs(&modules_crit_);
@@ -133,26 +136,27 @@
     packet->SetExtension<TransportSequenceNumber>(AllocateSequenceNumber());
   }
 
-  auto it = rtp_module_cache_map_.find(packet->Ssrc());
-  if (it != rtp_module_cache_map_.end()) {
-    if (TrySendPacket(packet.get(), cluster_info, it->second)) {
-      return;
-    }
-    // Entry is stale, remove it.
-    rtp_module_cache_map_.erase(it);
+  uint32_t ssrc = packet->Ssrc();
+  auto kv = send_modules_map_.find(ssrc);
+  if (kv == send_modules_map_.end()) {
+    RTC_LOG(LS_WARNING)
+        << "Failed to send packet, matching RTP module not found "
+           "or transport error. SSRC = "
+        << packet->Ssrc() << ", sequence number " << packet->SequenceNumber();
+    return;
   }
 
-  // Slow path, find the correct send module.
-  for (auto* rtp_module : rtp_send_modules_) {
-    if (TrySendPacket(packet.get(), cluster_info, rtp_module)) {
-      return;
-    }
+  RtpRtcp* rtp_module = kv->second.first;
+  if (!rtp_module->TrySendPacket(packet.get(), cluster_info)) {
+    RTC_LOG(LS_WARNING) << "Failed to send packet, rejected by RTP module.";
+    return;
   }
 
-  RTC_LOG(LS_WARNING) << "Failed to send packet, matching RTP module not found "
-                         "or transport error. SSRC = "
-                      << packet->Ssrc() << ", sequence number "
-                      << packet->SequenceNumber();
+  if (rtp_module->SupportsRtxPayloadPadding()) {
+    // This is now the last module to send media, and has the desired
+    // properties needed for payload based padding. Cache it for later use.
+    last_send_module_ = rtp_module;
+  }
 }
 
 std::vector<std::unique_ptr<RtpPacketToSend>> PacketRouter::GeneratePadding(
@@ -164,25 +168,26 @@
   // will be more skewed towards the highest bitrate stream. At the very least
   // this prevents sending payload padding on a disabled stream where it's
   // guaranteed not to be useful.
+  std::vector<std::unique_ptr<RtpPacketToSend>> padding_packets;
   if (last_send_module_ != nullptr &&
       last_send_module_->SupportsRtxPayloadPadding()) {
-    RTC_DCHECK(std::find(rtp_send_modules_.begin(), rtp_send_modules_.end(),
-                         last_send_module_) != rtp_send_modules_.end());
-    return last_send_module_->GeneratePadding(target_size_bytes);
-  }
-
-  // Rtp modules are ordered by which stream can most benefit from padding.
-  for (RtpRtcp* rtp_module : rtp_send_modules_) {
-    if (rtp_module->SupportsPadding()) {
-      auto padding_packets = rtp_module->GeneratePadding(target_size_bytes);
-      if (!padding_packets.empty()) {
-        last_send_module_ = rtp_module;
-      }
+    padding_packets = last_send_module_->GeneratePadding(target_size_bytes);
+    if (!padding_packets.empty()) {
       return padding_packets;
     }
   }
 
-  return {};
+  for (RtpRtcp* rtp_module : send_modules_list_) {
+    if (rtp_module->SupportsPadding()) {
+      padding_packets = rtp_module->GeneratePadding(target_size_bytes);
+      if (!padding_packets.empty()) {
+        last_send_module_ = rtp_module;
+        break;
+      }
+    }
+  }
+
+  return padding_packets;
 }
 
 void PacketRouter::SetTransportWideSequenceNumber(uint16_t sequence_number) {
@@ -276,7 +281,7 @@
   rtc::CritScope cs(&modules_crit_);
 
   // Prefer send modules.
-  for (auto* rtp_module : rtp_send_modules_) {
+  for (RtpRtcp* rtp_module : send_modules_list_) {
     if (rtp_module->RTCP() == RtcpMode::kOff) {
       continue;
     }
@@ -352,23 +357,4 @@
   active_remb_module_ = new_active_remb_module;
 }
 
-bool PacketRouter::TrySendPacket(RtpPacketToSend* packet,
-                                 const PacedPacketInfo& cluster_info,
-                                 RtpRtcp* rtp_module) {
-  uint32_t ssrc = packet->Ssrc();
-  if (rtp_module->TrySendPacket(packet, cluster_info)) {
-    // Sending succeeded, make sure this SSRC mapping for future use.
-    rtp_module_cache_map_[ssrc] = rtp_module;
-
-    if (rtp_module->SupportsRtxPayloadPadding()) {
-      // This is now the last module to send media, and has the desired
-      // properties needed for payload based padding. Cache it for later use.
-      last_send_module_ = rtp_module;
-    }
-
-    return true;
-  }
-  return false;
-}
-
 }  // namespace webrtc
diff --git a/modules/pacing/packet_router.h b/modules/pacing/packet_router.h
index 3680bce..1359e04 100644
--- a/modules/pacing/packet_router.h
+++ b/modules/pacing/packet_router.h
@@ -17,6 +17,7 @@
 #include <list>
 #include <memory>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "api/transport/network_types.h"
@@ -84,9 +85,6 @@
       std::vector<std::unique_ptr<rtcp::RtcpPacket>> packets) override;
 
  private:
-  RtpRtcp* FindRtpModule(uint32_t ssrc)
-      RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
-
   void AddRembModuleCandidate(RtcpFeedbackSenderInterface* candidate_module,
                               bool media_sender)
       RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
@@ -95,17 +93,17 @@
       bool media_sender) RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
   void UnsetActiveRembModule() RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
   void DetermineActiveRembModule() RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
-  bool TrySendPacket(RtpPacketToSend* packet,
-                     const PacedPacketInfo& cluster_info,
-                     RtpRtcp* rtp_module)
+  void AddSendRtpModuleToMap(RtpRtcp* rtp_module, uint32_t ssrc)
+      RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
+  void RemoveSendRtpModuleFromMap(uint32_t ssrc)
       RTC_EXCLUSIVE_LOCKS_REQUIRED(modules_crit_);
 
   rtc::CriticalSection modules_crit_;
-  // Rtp and Rtcp modules of the rtp senders.
-  std::list<RtpRtcp*> rtp_send_modules_ RTC_GUARDED_BY(modules_crit_);
-  // Ssrc to RtpRtcp module cache.
-  std::unordered_map<uint32_t, RtpRtcp*> rtp_module_cache_map_
-      RTC_GUARDED_BY(modules_crit_);
+  // Ssrc to RtpRtcp module and iterator into |send_modules_list_|;
+  std::unordered_map<uint32_t,
+                     std::pair<RtpRtcp*, std::list<RtpRtcp*>::iterator>>
+      send_modules_map_ RTC_GUARDED_BY(modules_crit_);
+  std::list<RtpRtcp*> send_modules_list_ RTC_GUARDED_BY(modules_crit_);
   // The last module used to send media.
   RtpRtcp* last_send_module_ RTC_GUARDED_BY(modules_crit_);
   // Rtcp modules of the rtp receivers.
diff --git a/modules/pacing/packet_router_unittest.cc b/modules/pacing/packet_router_unittest.cc
index 1239201..0c95e7f 100644
--- a/modules/pacing/packet_router_unittest.cc
+++ b/modules/pacing/packet_router_unittest.cc
@@ -179,8 +179,10 @@
   // and supports rtx.
   EXPECT_CALL(rtp_2, GeneratePadding(kPaddingBytes))
       .Times(1)
-      .WillOnce([](size_t target_size_bytes) {
-        return std::vector<std::unique_ptr<RtpPacketToSend>>();
+      .WillOnce([&](size_t target_size_bytes) {
+        std::vector<std::unique_ptr<RtpPacketToSend>> packets;
+        packets.push_back(BuildRtpPacket(kSsrc2));
+        return packets;
       });
   packet_router_.GeneratePadding(kPaddingBytes);
 
@@ -189,41 +191,45 @@
 
   EXPECT_CALL(rtp_1, GeneratePadding(kPaddingBytes))
       .Times(1)
-      .WillOnce([](size_t target_size_bytes) {
-        return std::vector<std::unique_ptr<RtpPacketToSend>>();
+      .WillOnce([&](size_t target_size_bytes) {
+        std::vector<std::unique_ptr<RtpPacketToSend>> packets;
+        packets.push_back(BuildRtpPacket(kSsrc1));
+        return packets;
       });
   packet_router_.GeneratePadding(kPaddingBytes);
 
   // Send media on second module. Padding should be sent there.
   packet_router_.SendPacket(BuildRtpPacket(kSsrc2), PacedPacketInfo());
 
-  EXPECT_CALL(rtp_2, GeneratePadding(kPaddingBytes))
-      .Times(1)
-      .WillOnce([](size_t target_size_bytes) {
-        return std::vector<std::unique_ptr<RtpPacketToSend>>();
-      });
-  packet_router_.GeneratePadding(kPaddingBytes);
-
-  // Remove second module, padding should now fall back to first module.
+  // If the last active module is removed, and no module sends media before
+  // the next padding request, and arbitrary module will be selected.
   packet_router_.RemoveSendRtpModule(&rtp_2);
+
+  // Send on and then remove all remaining modules.
+  RtpRtcp* last_send_module;
   EXPECT_CALL(rtp_1, GeneratePadding(kPaddingBytes))
       .Times(1)
-      .WillOnce([](size_t target_size_bytes) {
-        return std::vector<std::unique_ptr<RtpPacketToSend>>();
+      .WillOnce([&](size_t target_size_bytes) {
+        last_send_module = &rtp_1;
+        std::vector<std::unique_ptr<RtpPacketToSend>> packets;
+        packets.push_back(BuildRtpPacket(kSsrc1));
+        return packets;
       });
-  packet_router_.GeneratePadding(kPaddingBytes);
-
-  // Remove first module too, leaving only the one without rtx.
-  packet_router_.RemoveSendRtpModule(&rtp_1);
-
   EXPECT_CALL(rtp_3, GeneratePadding(kPaddingBytes))
       .Times(1)
-      .WillOnce([](size_t target_size_bytes) {
-        return std::vector<std::unique_ptr<RtpPacketToSend>>();
+      .WillOnce([&](size_t target_size_bytes) {
+        last_send_module = &rtp_3;
+        std::vector<std::unique_ptr<RtpPacketToSend>> packets;
+        packets.push_back(BuildRtpPacket(kSsrc3));
+        return packets;
       });
-  packet_router_.GeneratePadding(kPaddingBytes);
 
-  packet_router_.RemoveSendRtpModule(&rtp_3);
+  for (int i = 0; i < 2; ++i) {
+    last_send_module = nullptr;
+    packet_router_.GeneratePadding(kPaddingBytes);
+    EXPECT_NE(last_send_module, nullptr);
+    packet_router_.RemoveSendRtpModule(last_send_module);
+  }
 }
 
 TEST_F(PacketRouterTest, AllocatesTransportSequenceNumbers) {
@@ -272,12 +278,11 @@
 }
 
 TEST_F(PacketRouterTest, SendPacketWithoutTransportSequenceNumbers) {
-  NiceMock<MockRtpRtcp> rtp_1;
-  packet_router_.AddSendRtpModule(&rtp_1, false);
-
   const uint16_t kSsrc1 = 1234;
+  NiceMock<MockRtpRtcp> rtp_1;
   ON_CALL(rtp_1, SendingMedia).WillByDefault(Return(true));
   ON_CALL(rtp_1, SSRC).WillByDefault(Return(kSsrc1));
+  packet_router_.AddSendRtpModule(&rtp_1, false);
 
   // Send a packet without TransportSequenceNumber extension registered,
   // packets sent should not have the extension set.
@@ -300,15 +305,15 @@
   NiceMock<MockRtpRtcp> rtp_1;
   NiceMock<MockRtpRtcp> rtp_2;
 
-  packet_router_.AddSendRtpModule(&rtp_1, false);
-  packet_router_.AddSendRtpModule(&rtp_2, false);
-
   const uint16_t kSsrc1 = 1234;
   const uint16_t kSsrc2 = 2345;
 
   ON_CALL(rtp_1, SSRC).WillByDefault(Return(kSsrc1));
   ON_CALL(rtp_2, SSRC).WillByDefault(Return(kSsrc2));
 
+  packet_router_.AddSendRtpModule(&rtp_1, false);
+  packet_router_.AddSendRtpModule(&rtp_2, false);
+
   // Transport sequence numbers start at 1, for historical reasons.
   uint16_t transport_sequence_number = 1;
 
@@ -327,9 +332,6 @@
   packet = BuildRtpPacket(kSsrc2);
   EXPECT_TRUE(packet->ReserveExtension<TransportSequenceNumber>());
 
-  // There will be a failed attempt to send on kSsrc1 before trying
-  // the correct RTP module.
-  EXPECT_CALL(rtp_1, TrySendPacket).WillOnce(Return(false));
   EXPECT_CALL(
       rtp_2,
       TrySendPacket(
diff --git a/modules/rtp_rtcp/include/rtp_rtcp.h b/modules/rtp_rtcp/include/rtp_rtcp.h
index 7682b4a..b877045 100644
--- a/modules/rtp_rtcp/include/rtp_rtcp.h
+++ b/modules/rtp_rtcp/include/rtp_rtcp.h
@@ -226,6 +226,9 @@
   // a combination of values of the enumerator RtxMode.
   virtual int RtxSendStatus() const = 0;
 
+  // Returns the SSRC used for RTX if set, otherwise a nullopt.
+  virtual absl::optional<uint32_t> RtxSsrc() const = 0;
+
   // Sets the payload type to use when sending RTX packets. Note that this
   // doesn't enable RTX, only the payload type is set.
   virtual void SetRtxSendPayloadType(int payload_type,
diff --git a/modules/rtp_rtcp/mocks/mock_rtp_rtcp.h b/modules/rtp_rtcp/mocks/mock_rtp_rtcp.h
index 5b81fe1..332f243 100644
--- a/modules/rtp_rtcp/mocks/mock_rtp_rtcp.h
+++ b/modules/rtp_rtcp/mocks/mock_rtp_rtcp.h
@@ -68,6 +68,7 @@
   MOCK_METHOD1(SetCSRCStatus, int32_t(bool include));
   MOCK_METHOD1(SetRtxSendStatus, void(int modes));
   MOCK_CONST_METHOD0(RtxSendStatus, int());
+  MOCK_CONST_METHOD0(RtxSsrc, absl::optional<uint32_t>());
   MOCK_METHOD1(SetRtxSsrc, void(uint32_t));
   MOCK_METHOD2(SetRtxSendPayloadType, void(int, int));
   MOCK_CONST_METHOD0(FlexfecSsrc, absl::optional<uint32_t>());
diff --git a/modules/rtp_rtcp/source/rtp_rtcp_impl.cc b/modules/rtp_rtcp/source/rtp_rtcp_impl.cc
index 4ff584e..c7cbf50 100644
--- a/modules/rtp_rtcp/source/rtp_rtcp_impl.cc
+++ b/modules/rtp_rtcp/source/rtp_rtcp_impl.cc
@@ -180,6 +180,10 @@
   rtp_sender_->SetRtxPayloadType(payload_type, associated_payload_type);
 }
 
+absl::optional<uint32_t> ModuleRtpRtcpImpl::RtxSsrc() const {
+  return rtp_sender_ ? rtp_sender_->RtxSsrc() : absl::nullopt;
+}
+
 absl::optional<uint32_t> ModuleRtpRtcpImpl::FlexfecSsrc() const {
   if (rtp_sender_)
     return rtp_sender_->FlexfecSsrc();
diff --git a/modules/rtp_rtcp/source/rtp_rtcp_impl.h b/modules/rtp_rtcp/source/rtp_rtcp_impl.h
index 2d6cfff..03dd81c 100644
--- a/modules/rtp_rtcp/source/rtp_rtcp_impl.h
+++ b/modules/rtp_rtcp/source/rtp_rtcp_impl.h
@@ -106,6 +106,7 @@
 
   void SetRtxSendStatus(int mode) override;
   int RtxSendStatus() const override;
+  absl::optional<uint32_t> RtxSsrc() const override;
 
   void SetRtxSendPayloadType(int payload_type,
                              int associated_payload_type) override;
diff --git a/modules/rtp_rtcp/source/rtp_sender.h b/modules/rtp_rtcp/source/rtp_sender.h
index 28512b8..50ece54 100644
--- a/modules/rtp_rtcp/source/rtp_sender.h
+++ b/modules/rtp_rtcp/source/rtp_sender.h
@@ -114,10 +114,7 @@
   // RTX.
   void SetRtxStatus(int mode);
   int RtxStatus() const;
-  uint32_t RtxSsrc() const {
-    RTC_DCHECK(rtx_ssrc_);
-    return *rtx_ssrc_;
-  }
+  absl::optional<uint32_t> RtxSsrc() const { return rtx_ssrc_; }
 
   void SetRtxPayloadType(int payload_type, int associated_payload_type);
 
diff --git a/video/video_send_stream_tests.cc b/video/video_send_stream_tests.cc
index d769bfe..9e96821 100644
--- a/video/video_send_stream_tests.cc
+++ b/video/video_send_stream_tests.cc
@@ -1572,7 +1572,7 @@
         VideoEncoderConfig* encoder_config) override {
       // Turn on RTX.
       send_config->rtp.rtx.payload_type = kFakeVideoSendPayloadType;
-      send_config->rtp.rtx.ssrcs.push_back(kVideoSendSsrcs[0]);
+      send_config->rtp.rtx.ssrcs.push_back(kSendRtxSsrcs[0]);
     }
 
     void PerformTest() override {