feat: judge stream end

moqt-core/src/modules/messages/data_streams/object_stream_subgroup.rs

@@ -47,6 +47,10 @@ impl ObjectStreamSubgroup {
     pub fn object_id(&self) -> u64 {
         self.object_id
     }
+
+    pub fn object_status(&self) -> Option<ObjectStatus> {
+        self.object_status
+    }
 }
 
 impl DataStreams for ObjectStreamSubgroup {

moqt-core/src/modules/messages/data_streams/object_stream_track.rs

@@ -54,6 +54,10 @@ impl ObjectStreamTrack {
     pub fn object_id(&self) -> u64 {
         self.object_id
     }
+
+    pub fn object_status(&self) -> Option<ObjectStatus> {
+        self.object_status
+    }
 }
 
 impl DataStreams for ObjectStreamTrack {

moqt-core/src/modules/pubsub_relation_manager_repository.rs

@@ -53,6 +53,11 @@ pub trait PubSubRelationManagerRepository: Send + Sync {
         track_namespace: Vec<String>,
         track_name: String,
     ) -> Result<Option<Subscription>>;
+    async fn get_upstream_subscription_by_ids(
+        &self,
+        upstream_session_id: usize,
+        upstream_subscribe_id: u64,
+    ) -> Result<Option<Subscription>>;
     async fn get_downstream_subscription_by_ids(
         &self,
         downstream_session_id: usize,

moqt-server/src/modules/message_handlers/object_stream.rs

@@ -17,7 +17,7 @@ use std::io::Cursor;
 
 #[derive(Debug, PartialEq)]
 pub enum ObjectStreamProcessResult {
-    Success,
+    Success(CacheObject),
     IncompleteMessage,
     Failure(TerminationErrorCode, String),
 }
@@ -61,16 +61,18 @@ pub async fn object_stream_handler(
                 Ok(object) => {
                     read_buf.advance(read_cur.position() as usize);
 
-                    let cache_object = CacheObject::Track(object);
+                    let received_object = CacheObject::Track(object);
                     object_cache_storage
-                        .set_object(client.id(), subscribe_id, cache_object, duration)
+                        .set_object(client.id(), subscribe_id, received_object.clone(), duration)
                         .await
                         .unwrap();
+
+                    ObjectStreamProcessResult::Success(received_object)
                 }
                 Err(err) => {
                     tracing::warn!("{:#?}", err);
                     read_cur.set_position(0);
-                    return ObjectStreamProcessResult::IncompleteMessage;
+                    ObjectStreamProcessResult::IncompleteMessage
                 }
             }
         }
@@ -80,26 +82,24 @@ pub async fn object_stream_handler(
                 Ok(object) => {
                     read_buf.advance(read_cur.position() as usize);
 
-                    let cache_object = CacheObject::Subgroup(object);
+                    let received_object = CacheObject::Subgroup(object);
                     object_cache_storage
-                        .set_object(client.id(), subscribe_id, cache_object, duration)
+                        .set_object(client.id(), subscribe_id, received_object.clone(), duration)
                         .await
                         .unwrap();
+
+                    ObjectStreamProcessResult::Success(received_object)
                 }
                 Err(err) => {
                     tracing::warn!("{:#?}", err);
                     read_cur.set_position(0);
-                    return ObjectStreamProcessResult::IncompleteMessage;
+                    ObjectStreamProcessResult::IncompleteMessage
                 }
             }
         }
-        unknown => {
-            return ObjectStreamProcessResult::Failure(
-                TerminationErrorCode::ProtocolViolation,
-                format!("Unknown message type: {:?}", unknown),
-            );
-        }
-    };
-
-    ObjectStreamProcessResult::Success
+        unknown => ObjectStreamProcessResult::Failure(
+            TerminationErrorCode::ProtocolViolation,
+            format!("Unknown message type: {:?}", unknown),
+        ),
+    }
 }

moqt-server/src/modules/message_handlers/stream_header.rs

@@ -9,7 +9,7 @@ use crate::{
             stream_track_subgroup::process_stream_header_subgroup,
         },
         moqt_client::{MOQTClient, MOQTClientStatus},
-        object_cache_storage::ObjectCacheStorageWrapper,
+        object_cache_storage::{CacheHeader, ObjectCacheStorageWrapper},
     },
 };
 use anyhow::{bail, Result};
@@ -23,7 +23,7 @@ use std::io::Cursor;
 
 #[derive(Debug, PartialEq)]
 pub enum StreamHeaderProcessResult {
-    Success((u64, DataStreamType)),
+    Success(CacheHeader),
     IncompleteMessage,
     Failure(TerminationErrorCode, String),
 }
@@ -92,7 +92,7 @@ pub async fn stream_header_handler(
         );
     }
 
-    let subscribe_id = match header_type {
+    match header_type {
         DataStreamType::StreamHeaderTrack => {
             match process_stream_header_track(
                 &mut read_cur,
@@ -102,16 +102,17 @@ pub async fn stream_header_handler(
             )
             .await
             {
-                Ok(subscribe_id) => {
+                Ok(received_header) => {
                     read_buf.advance(read_cur.position() as usize);
-                    subscribe_id
+
+                    StreamHeaderProcessResult::Success(received_header)
                 }
                 Err(err) => {
                     read_buf.advance(read_cur.position() as usize);
-                    return StreamHeaderProcessResult::Failure(
+                    StreamHeaderProcessResult::Failure(
                         TerminationErrorCode::InternalError,
                         err.to_string(),
-                    );
+                    )
                 }
             }
         }
@@ -124,26 +125,23 @@ pub async fn stream_header_handler(
             )
             .await
             {
-                Ok(subscribe_id) => {
+                Ok(received_header) => {
                     read_buf.advance(read_cur.position() as usize);
-                    subscribe_id
+
+                    StreamHeaderProcessResult::Success(received_header)
                 }
                 Err(err) => {
                     read_buf.advance(read_cur.position() as usize);
-                    return StreamHeaderProcessResult::Failure(
+                    StreamHeaderProcessResult::Failure(
                         TerminationErrorCode::InternalError,
                         err.to_string(),
-                    );
+                    )
                 }
             }
         }
-        unknown => {
-            return StreamHeaderProcessResult::Failure(
-                TerminationErrorCode::ProtocolViolation,
-                format!("Unknown message type: {:?}", unknown),
-            );
-        }
-    };
-
-    StreamHeaderProcessResult::Success((subscribe_id, header_type))
+        unknown => StreamHeaderProcessResult::Failure(
+            TerminationErrorCode::ProtocolViolation,
+            format!("Unknown message type: {:?}", unknown),
+        ),
+    }
 }

moqt-server/src/modules/message_handlers/stream_header/handlers/stream_subgroup_header_handler.rs

@@ -16,7 +16,7 @@ pub(crate) async fn stream_header_subgroup_handler(
     pubsub_relation_manager_repository: &mut dyn PubSubRelationManagerRepository,
     object_cache_storage: &mut ObjectCacheStorageWrapper,
     client: &MOQTClient,
-) -> Result<u64> {
+) -> Result<CacheHeader> {
     tracing::trace!("stream_header_subgroup_handler start.");
 
     tracing::debug!(
@@ -37,8 +37,12 @@ pub(crate) async fn stream_header_subgroup_handler(
 
     let cache_header = CacheHeader::Subgroup(stream_header_subgroup_message);
     object_cache_storage
-        .set_subscription(upstream_session_id, upstream_subscribe_id, cache_header)
+        .set_subscription(
+            upstream_session_id,
+            upstream_subscribe_id,
+            cache_header.clone(),
+        )
         .await?;
 
-    Ok(upstream_subscribe_id)
+    Ok(cache_header)
 }

moqt-server/src/modules/message_handlers/stream_header/handlers/stream_track_header_handler.rs

@@ -16,7 +16,7 @@ pub(crate) async fn stream_header_track_handler(
     pubsub_relation_manager_repository: &mut dyn PubSubRelationManagerRepository,
     object_cache_storage: &mut ObjectCacheStorageWrapper,
     client: &MOQTClient,
-) -> Result<u64> {
+) -> Result<CacheHeader> {
     tracing::trace!("stream_header_track_handler start.");
 
     tracing::debug!(
@@ -37,8 +37,12 @@ pub(crate) async fn stream_header_track_handler(
 
     let cache_header = CacheHeader::Track(stream_header_track_message);
     object_cache_storage
-        .set_subscription(upstream_session_id, upstream_subscribe_id, cache_header)
+        .set_subscription(
+            upstream_session_id,
+            upstream_subscribe_id,
+            cache_header.clone(),
+        )
         .await?;
 
-    Ok(upstream_subscribe_id)
+    Ok(cache_header)
 }

moqt-server/src/modules/message_handlers/stream_header/server_processes/stream_track_header.rs

@@ -7,15 +7,16 @@ use moqt_core::{
 
 use crate::modules::{
     message_handlers::stream_header::handlers::stream_track_header_handler::stream_header_track_handler,
-    moqt_client::MOQTClient, object_cache_storage::ObjectCacheStorageWrapper,
+    moqt_client::MOQTClient,
+    object_cache_storage::{CacheHeader, ObjectCacheStorageWrapper},
 };
 
 pub(crate) async fn process_stream_header_track(
     read_cur: &mut std::io::Cursor<&[u8]>,
     pubsub_relation_manager_repository: &mut dyn PubSubRelationManagerRepository,
     object_cache_storage: &mut ObjectCacheStorageWrapper,
     client: &MOQTClient,
-) -> Result<u64> {
+) -> Result<CacheHeader> {
     let stream_header_track = match StreamHeaderTrack::depacketize(read_cur) {
         Ok(stream_header_track) => stream_header_track,
         Err(err) => {

moqt-server/src/modules/message_handlers/stream_header/server_processes/stream_track_subgroup.rs

@@ -7,15 +7,16 @@ use moqt_core::{
 
 use crate::modules::{
     message_handlers::stream_header::handlers::stream_subgroup_header_handler::stream_header_subgroup_handler,
-    moqt_client::MOQTClient, object_cache_storage::ObjectCacheStorageWrapper,
+    moqt_client::MOQTClient,
+    object_cache_storage::{CacheHeader, ObjectCacheStorageWrapper},
 };
 
 pub(crate) async fn process_stream_header_subgroup(
     read_cur: &mut std::io::Cursor<&[u8]>,
     pubsub_relation_manager_repository: &mut dyn PubSubRelationManagerRepository,
     object_cache_storage: &mut ObjectCacheStorageWrapper,
     client: &MOQTClient,
-) -> Result<u64> {
+) -> Result<CacheHeader> {
     let stream_header_subgroup = match StreamHeaderSubgroup::depacketize(read_cur) {
         Ok(stream_header_subgroup) => stream_header_subgroup,
         Err(err) => {

moqt-server/src/modules/object_cache_storage.rs

@@ -22,7 +22,7 @@ pub(crate) enum CacheHeader {
 }
 
 #[allow(dead_code)]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub(crate) enum CacheObject {
     Datagram(ObjectDatagram),
     Track(ObjectStreamTrack),

moqt-server/src/modules/pubsub_relation_manager/commands.rs

@@ -59,6 +59,11 @@ pub(crate) enum PubSubRelationCommand {
         track_name: String,
         resp: oneshot::Sender<Result<Option<Subscription>>>,
     },
+    GetUpstreamSubscriptionBySessionIdAndSubscribeId {
+        upstream_session_id: usize,
+        upstream_subscribe_id: u64,
+        resp: oneshot::Sender<Result<Option<Subscription>>>,
+    },
     GetDownstreamSubscriptionBySessionIdAndSubscribeId {
         downstream_session_id: usize,
         downstream_subscribe_id: u64,

moqt-server/src/modules/pubsub_relation_manager/manager.rs

@@ -288,6 +288,16 @@ pub(crate) async fn pubsub_relation_manager(rx: &mut mpsc::Receiver<PubSubRelati
 
                 resp.send(result).unwrap();
             }
+            GetUpstreamSubscriptionBySessionIdAndSubscribeId {
+                upstream_session_id,
+                upstream_subscribe_id,
+                resp,
+            } => {
+                let consumer = consumers.get(&upstream_session_id).unwrap();
+                let result = consumer.get_subscription(upstream_subscribe_id);
+
+                resp.send(result).unwrap();
+            }
             GetDownstreamSubscriptionBySessionIdAndSubscribeId {
                 downstream_session_id,
                 downstream_subscribe_id,