remove UserTimestamp store in favor of simple map to track ts

Author: chenosaurus

libwebrtc/src/native/user_timestamp.rs

@@ -18,8 +18,13 @@
 //! in encoded video frames as trailers. The timestamps are preserved
 //! through the WebRTC pipeline and can be extracted on the receiver side.
 //!
-//! This works independently of e2ee encryption - timestamps can be
-//! embedded even when encryption is disabled.
+//! On the send side, user timestamps are stored in the handler's internal
+//! map keyed by capture timestamp. When the encoder produces a frame,
+//! the transformer looks up the user timestamp via the frame's CaptureTime().
+//!
+//! On the receive side, extracted user timestamps are stored in an
+//! internal map keyed by RTP timestamp. Decoded frames look up their
+//! user timestamp via lookup_user_timestamp(rtp_timestamp).
 
 use cxx::SharedPtr;
 use webrtc_sys::user_timestamp::ffi as sys_ut;
@@ -30,94 +35,15 @@ use crate::{
     rtp_sender::RtpSender,
 };
 
-/// Thread-safe store for mapping capture timestamps to user timestamps.
-///
-/// Used on the sender side to correlate video frame capture time with
-/// the user timestamp that should be embedded in the encoded frame.
-#[derive(Clone)]
-pub struct UserTimestampStore {
-    sys_handle: SharedPtr<sys_ut::UserTimestampStore>,
-}
-
-impl UserTimestampStore {
-    /// Create a new user timestamp store.
-    pub fn new() -> Self {
-        Self {
-            sys_handle: sys_ut::new_user_timestamp_store(),
-        }
-    }
-
-    /// Store a user timestamp associated with a capture timestamp.
-    ///
-    /// Call this when capturing a video frame with a user timestamp.
-    /// The `capture_timestamp_us` should match the `timestamp_us` field
-    /// of the VideoFrame.
-    pub fn store(&self, capture_timestamp_us: i64, user_timestamp_us: i64) {
-        log::info!(
-            target: "user_timestamp",
-            "store: capture_ts_us={}, user_ts_us={}",
-            capture_timestamp_us,
-            user_timestamp_us
-        );
-        self.sys_handle.store(capture_timestamp_us, user_timestamp_us);
-    }
-
-    /// Lookup a user timestamp by capture timestamp (for debugging).
-    /// Returns None if not found.
-    pub fn lookup(&self, capture_timestamp_us: i64) -> Option<i64> {
-        let result = self.sys_handle.lookup(capture_timestamp_us);
-        if result < 0 {
-            None
-        } else {
-            Some(result)
-        }
-    }
-
-    /// Pop the oldest user timestamp from the queue.
-    /// Returns None if the queue is empty.
-    pub fn pop(&self) -> Option<i64> {
-        let result = self.sys_handle.pop();
-        if result < 0 {
-            None
-        } else {
-            Some(result)
-        }
-    }
-
-    /// Peek at the oldest user timestamp without removing it.
-    /// Returns None if the queue is empty.
-    pub fn peek(&self) -> Option<i64> {
-        let result = self.sys_handle.peek();
-        if result < 0 {
-            None
-        } else {
-            Some(result)
-        }
-    }
-
-    /// Clear old entries (older than the given threshold in microseconds).
-    pub fn prune(&self, max_age_us: i64) {
-        self.sys_handle.prune(max_age_us);
-    }
-
-    pub(crate) fn sys_handle(&self) -> SharedPtr<sys_ut::UserTimestampStore> {
-        self.sys_handle.clone()
-    }
-}
-
-impl Default for UserTimestampStore {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 /// Handler for user timestamp embedding/extraction on RTP streams.
 ///
-/// For sender side: Embeds user timestamps as 12-byte trailers on
-/// encoded frames before they are sent.
+/// For sender side: Stores user timestamps keyed by capture timestamp
+/// and embeds them as 12-byte trailers on encoded frames before they
+/// are sent. Use `store_user_timestamp()` to associate a user timestamp
+/// with a captured frame.
 ///
 /// For receiver side: Extracts user timestamps from received frames
-/// and makes them available for retrieval.
+/// and makes them available for retrieval via `lookup_user_timestamp()`.
 #[derive(Clone)]
 pub struct UserTimestampHandler {
     sys_handle: SharedPtr<sys_ut::UserTimestampHandler>,
@@ -164,24 +90,44 @@ impl UserTimestampHandler {
         }
     }
 
+    /// Store a user timestamp for a given capture timestamp (sender side).
+    ///
+    /// The `capture_timestamp_us` must be the TimestampAligner-adjusted
+    /// timestamp (as produced by `VideoTrackSource::on_captured_frame`),
+    /// NOT the original `timestamp_us` from the VideoFrame. The transformer
+    /// looks up the user timestamp by the frame's `CaptureTime()` which is
+    /// derived from the aligned value.
+    ///
+    /// In normal usage this is called automatically by the C++ layer —
+    /// callers should set `user_timestamp_us` on the `VideoFrame` and let
+    /// `capture_frame` / `on_captured_frame` handle the rest.
+    pub fn store_user_timestamp(&self, capture_timestamp_us: i64, user_timestamp_us: i64) {
+        log::info!(
+            target: "user_timestamp",
+            "store: capture_ts_us={}, user_ts_us={}",
+            capture_timestamp_us,
+            user_timestamp_us
+        );
+        self.sys_handle.store_user_timestamp(capture_timestamp_us, user_timestamp_us);
+    }
+
     pub(crate) fn sys_handle(&self) -> SharedPtr<sys_ut::UserTimestampHandler> {
         self.sys_handle.clone()
     }
 }
 
 /// Create a sender-side user timestamp handler.
 ///
-/// This handler will embed user timestamps from the provided store
-/// into encoded frames before they are packetized and sent.
+/// This handler will embed user timestamps into encoded frames before
+/// they are packetized and sent. Use `store_user_timestamp()` to
+/// associate a user timestamp with a captured frame's capture timestamp.
 pub fn create_sender_handler(
     peer_factory: &PeerConnectionFactory,
-    store: &UserTimestampStore,
     sender: &RtpSender,
 ) -> UserTimestampHandler {
     UserTimestampHandler {
         sys_handle: sys_ut::new_user_timestamp_sender(
             peer_factory.handle.sys_handle.clone(),
-            store.sys_handle(),
             sender.handle.sys_handle.clone(),
         ),
     }
@@ -195,13 +141,11 @@ pub fn create_sender_handler(
 /// timestamp for a specific decoded frame.
 pub fn create_receiver_handler(
     peer_factory: &PeerConnectionFactory,
-    store: &UserTimestampStore,
     receiver: &RtpReceiver,
 ) -> UserTimestampHandler {
     UserTimestampHandler {
         sys_handle: sys_ut::new_user_timestamp_receiver(
             peer_factory.handle.sys_handle.clone(),
-            store.sys_handle(),
             receiver.handle.sys_handle.clone(),
         ),
     }

libwebrtc/src/native/video_source.rs

@@ -23,7 +23,7 @@ use parking_lot::Mutex;
 use webrtc_sys::{video_frame as vf_sys, video_frame::ffi::VideoRotation, video_track as vt_sys};
 
 use crate::{
-    native::user_timestamp::UserTimestampStore,
+    native::user_timestamp::UserTimestampHandler,
     video_frame::{I420Buffer, VideoBuffer, VideoFrame},
     video_source::VideoResolution,
 };
@@ -48,7 +48,6 @@ pub struct NativeVideoSource {
 
 struct VideoSourceInner {
     captured_frames: usize,
-    user_timestamp_store: Option<UserTimestampStore>,
 }
 
 impl NativeVideoSource {
@@ -59,7 +58,6 @@ impl NativeVideoSource {
             )),
             inner: Arc::new(Mutex::new(VideoSourceInner {
                 captured_frames: 0,
-                user_timestamp_store: None,
             })),
         };
 
@@ -84,7 +82,11 @@ impl NativeVideoSource {
                     let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
                     builder.pin_mut().set_timestamp_us(now.as_micros() as i64);
 
-                    source.sys_handle.on_captured_frame(&builder.pin_mut().build());
+                    source.sys_handle.on_captured_frame(
+                        &builder.pin_mut().build(),
+                        false,
+                        0,
+                    );
                 }
             }
         });
@@ -97,41 +99,43 @@ impl NativeVideoSource {
     }
 
     pub fn capture_frame<T: AsRef<dyn VideoBuffer>>(&self, frame: &VideoFrame<T>) {
-        let mut inner = self.inner.lock();
-        inner.captured_frames += 1;
-
         let mut builder = vf_sys::ffi::new_video_frame_builder();
         builder.pin_mut().set_rotation(frame.rotation.into());
         builder.pin_mut().set_video_frame_buffer(frame.buffer.as_ref().sys_handle());
 
         let capture_ts = if frame.timestamp_us == 0 {
-            // If the timestamp is set to 0, default to now
             let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
             now.as_micros() as i64
         } else {
             frame.timestamp_us
         };
         builder.pin_mut().set_timestamp_us(capture_ts);
 
-        // If a user timestamp is provided and a store is available, record
-        // the mapping so the UserTimestampTransformer can embed it into the
-        // encoded RTP frame.
-        if let Some(user_ts) = frame.user_timestamp_us {
-            if let Some(store) = &inner.user_timestamp_store {
-                store.store(capture_ts, user_ts);
-            }
-        }
+        // Pass the user timestamp to the C++ on_captured_frame so it can
+        // store the mapping keyed by the TimestampAligner-adjusted capture
+        // timestamp.  This is the only correct key because the aligner runs
+        // inside on_captured_frame and replaces timestamp_us with a value
+        // derived from rtc::TimeMicros() (monotonic), which is what
+        // CaptureTime() returns in TransformSend.
+        let (has_user_ts, user_ts) = match frame.user_timestamp_us {
+            Some(ts) => (true, ts),
+            None => (false, 0),
+        };
+
+        self.inner.lock().captured_frames += 1;
 
-        self.sys_handle.on_captured_frame(&builder.pin_mut().build());
+        self.sys_handle.on_captured_frame(&builder.pin_mut().build(), has_user_ts, user_ts);
     }
 
-    /// Set the user timestamp store used by this source.
+    /// Set the user timestamp handler used by this source.
     ///
     /// When set, any frame captured with a `user_timestamp_us` value will
-    /// automatically have its timestamp pushed into the store so the
+    /// automatically have its timestamp stored in the handler so the
     /// `UserTimestampTransformer` can embed it into the encoded frame.
-    pub fn set_user_timestamp_store(&self, store: UserTimestampStore) {
-        self.inner.lock().user_timestamp_store = Some(store);
+    /// The handler is set on the C++ VideoTrackSource so it has access to
+    /// the TimestampAligner-adjusted capture timestamp for correct keying.
+    pub fn set_user_timestamp_handler(&self, handler: UserTimestampHandler) {
+        self.sys_handle.set_user_timestamp_handler(handler.sys_handle());
     }
 
     pub fn video_resolution(&self) -> VideoResolution {

libwebrtc/src/video_source.rs

@@ -50,7 +50,7 @@ pub mod native {
     use std::fmt::{Debug, Formatter};
 
     use super::*;
-    use crate::native::user_timestamp::UserTimestampStore;
+    use crate::native::user_timestamp::UserTimestampHandler;
     use crate::video_frame::{VideoBuffer, VideoFrame};
 
     #[derive(Clone)]
@@ -79,13 +79,14 @@ pub mod native {
             self.handle.capture_frame(frame)
         }
 
-        /// Set the user timestamp store used by this source.
+        /// Set the user timestamp handler used by this source.
         ///
         /// When set, any frame captured with a `user_timestamp_us` value will
-        /// automatically have its timestamp pushed into the store so the
+        /// automatically have its timestamp stored in the handler (keyed by
+        /// the TimestampAligner-adjusted capture timestamp) so the
         /// `UserTimestampTransformer` can embed it into the encoded frame.
-        pub fn set_user_timestamp_store(&self, store: UserTimestampStore) {
-            self.handle.set_user_timestamp_store(store)
+        pub fn set_user_timestamp_handler(&self, handler: UserTimestampHandler) {
+            self.handle.set_user_timestamp_handler(handler)
         }
 
         pub fn video_resolution(&self) -> VideoResolution {

livekit/src/room/e2ee/manager.rs

@@ -19,7 +19,7 @@ use libwebrtc::{
         frame_cryptor::{
             DataPacketCryptor, EncryptedPacket, EncryptionAlgorithm, EncryptionState, FrameCryptor,
         },
-        user_timestamp::{self, UserTimestampStore},
+        user_timestamp,
     },
     rtp_receiver::RtpReceiver,
     rtp_sender::RtpSender,
@@ -107,10 +107,8 @@ impl E2eeManager {
 
         // Always set up user timestamp extraction for remote video tracks.
         if let RemoteTrack::Video(video_track) = &track {
-            let store = UserTimestampStore::new();
             let handler = user_timestamp::create_receiver_handler(
                 LkRuntime::instance().pc_factory(),
-                &store,
                 &receiver,
             );
             video_track.set_user_timestamp_handler(handler.clone());
@@ -143,21 +141,19 @@ impl E2eeManager {
 
         // Always set up user timestamp embedding for local video tracks.
         if let LocalTrack::Video(video_track) = &track {
-            let store = UserTimestampStore::new();
-            video_track.set_user_timestamp_store(store.clone());
+            let handler = user_timestamp::create_sender_handler(
+                LkRuntime::instance().pc_factory(),
+                &sender,
+            );
+            video_track.set_user_timestamp_handler(handler.clone());
 
-            // Also set the store on the video source so that capture_frame()
-            // can automatically push user timestamps into it.
+            // Also set the handler on the video source so that capture_frame()
+            // can automatically store user timestamps into it.
             #[cfg(not(target_arch = "wasm32"))]
             if let RtcVideoSource::Native(ref native_source) = video_track.rtc_source() {
-                native_source.set_user_timestamp_store(store.clone());
+                native_source.set_user_timestamp_handler(handler.clone());
             }
 
-            let handler = user_timestamp::create_sender_handler(
-                LkRuntime::instance().pc_factory(),
-                &store,
-                &sender,
-            );
             user_timestamp_handler = Some(handler);
         }