Fix Rust DataConsumer

### Details * [x] Rust `set_subchannels()` method was only implemented in `DirectDataConsumer` which is wrong. * [ ] Rust `DirectDataConsumer` lacks ALL methods of `DataConsumer` such as `pause()`, `closed()`, etc!!! * [x] Bonus track: Use an ordered set in `DataConsumer` in worker so `subchannels` will be always shown in numerical order (no need to sort them later in Node/Rust layer).
versatica · Dec 12, 2023 · aca18a3 · aca18a3
1 parent 042b7f5
commit aca18a3
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 25 deletions.
diff --git a/node/src/tests/test-DataConsumer.ts b/node/src/tests/test-DataConsumer.ts
@@ -71,7 +71,7 @@ test('transport.consumeData() succeeds', async () =>
 	expect(dataConsumer1.label).toBe('foo');
 	expect(dataConsumer1.protocol).toBe('bar');
 	expect(dataConsumer1.paused).toBe(false);
-	expect(dataConsumer1.subchannels.sort((a, b) => a - b)).toEqual([ 0, 1, 2, 100, 65535 ]);
+	expect(dataConsumer1.subchannels).toEqual([ 0, 1, 2, 100, 65535 ]);
 	expect(dataConsumer1.appData).toEqual({ baz: 'LOL' });
 
 	const dump = await router.dump();
@@ -134,7 +134,7 @@ test('dataConsumer.setSubchannels() succeeds', async () =>
 {
 	await dataConsumer1.setSubchannels([ 999, 999, 998, 65536 ]);
 
-	expect(dataConsumer1.subchannels.sort((a, b) => a - b)).toEqual([ 0, 998, 999 ]);
+	expect(dataConsumer1.subchannels).toEqual([ 0, 998, 999 ]);
 }, 2000);
 
 test('transport.consumeData() on a DirectTransport succeeds', async () =>

diff --git a/rust/src/router/data_consumer.rs b/rust/src/router/data_consumer.rs
@@ -411,6 +411,7 @@ impl fmt::Debug for DirectDataConsumer {
             .field("data_producer_id", &self.inner.data_producer_id)
             .field("paused", &self.inner.paused)
             .field("data_producer_paused", &self.inner.data_producer_paused)
+            .field("subchannels", &self.inner.subchannels)
             .field("transport", &self.inner.transport)
             .field("closed", &self.inner.closed)
             .finish()
@@ -788,6 +789,22 @@ impl DataConsumer {
             .await
     }
 
+    /// Sets subchannels to the worker DataConsumer.
+        pub async fn set_subchannels(&self, subchannels: Vec<u16>) -> Result<(), RequestError> {
+            let response = self
+                .inner()
+                .channel
+                .request(
+                    self.id(),
+                    DataConsumerSetSubchannelsRequest { subchannels },
+                )
+                .await?;
+
+            *self.inner().subchannels.lock() = response.subchannels;
+
+            Ok(())
+        }
+
     /// Callback is called when a message has been received from the corresponding data producer.
     ///
     /// # Notes on usage
@@ -918,22 +935,6 @@ impl DirectDataConsumer {
             )
             .await
     }
-
-    /// Sets subchannels to the worker DataConsumer.
-    pub async fn set_subchannels(&self, subchannels: Vec<u16>) -> Result<(), RequestError> {
-        let response = self
-            .inner
-            .channel
-            .request(
-                self.inner.id,
-                DataConsumerSetSubchannelsRequest { subchannels },
-            )
-            .await?;
-
-        *self.inner.subchannels.lock() = response.subchannels;
-
-        Ok(())
-    }
 }
 
 /// [`WeakDataConsumer`] doesn't own data consumer instance on mediasoup-worker and will not prevent

diff --git a/rust/tests/integration/data_consumer.rs b/rust/tests/integration/data_consumer.rs
@@ -153,11 +153,7 @@ fn consume_data_succeeds() {
         }
         assert_eq!(data_consumer.label().as_str(), "foo");
         assert_eq!(data_consumer.protocol().as_str(), "bar");
-
-        let mut sorted_subchannels = data_consumer.subchannels();
-        sorted_subchannels.sort();
-
-        assert_eq!(sorted_subchannels, [0, 1, 2, 100, 65535]);
+        assert_eq!(data_consumer.subchannels(), [0, 1, 2, 100, 65535]);
         assert_eq!(
             data_consumer
                 .app_data()
@@ -321,6 +317,32 @@ fn get_stats_succeeds() {
     });
 }
 
+#[test]
+fn set_subchannels() {
+    future::block_on(async move {
+        let (_worker, router, transport1, data_producer) = init().await;
+
+        let data_consumer = transport1
+            .consume_data({
+                let options = DataConsumerOptions::new_sctp_unordered_with_life_time(
+                    data_producer.id(),
+                    4000,
+                );
+
+                options
+            })
+            .await
+            .expect("Failed to consume data");
+
+        data_consumer
+            .set_subchannels([ 999, 999, 998, 65536 ].to_vec())
+            .await
+            .expect("Failed to set data consumer subchannels");
+
+        assert_eq!(data_consumer.subchannels(), [ 0, 998, 999 ]);
+    });
+}
+
 #[test]
 fn consume_data_on_direct_transport_succeeds() {
     future::block_on(async move {

diff --git a/worker/include/RTC/DataConsumer.hpp b/worker/include/RTC/DataConsumer.hpp
@@ -6,7 +6,7 @@
 #include "Channel/ChannelSocket.hpp"
 #include "RTC/SctpDictionaries.hpp"
 #include "RTC/Shared.hpp"
-#include <absl/container/flat_hash_set.h>
+#include <absl/container/btree_set.h>
 #include <string>
 
 namespace RTC
@@ -126,7 +126,7 @@ namespace RTC
 		RTC::SctpStreamParameters sctpStreamParameters;
 		std::string label;
 		std::string protocol;
-		absl::flat_hash_set<uint16_t> subchannels;
+		absl::btree_set<uint16_t> subchannels;
 		bool transportConnected{ false };
 		bool sctpAssociationConnected{ false };
 		bool paused{ false };