diff --git a/Runtime/Scripts/DataStreams/ByteDataStream.cs b/Runtime/Scripts/DataStreams/ByteDataStream.cs index 3823a976..73d6c971 100644 --- a/Runtime/Scripts/DataStreams/ByteDataStream.cs +++ b/Runtime/Scripts/DataStreams/ByteDataStream.cs @@ -145,8 +145,9 @@ internal ReadAllInstruction(ulong asyncId) /// YieldInstruction for . /// /// - /// Usage: while is false (i.e. the stream has not ended), - /// call , yield the instruction, and then access . + /// Usage: loop yielding the instruction; when it resumes, break if + /// is true (the stream has ended and every chunk has been drained), otherwise read + /// and call to advance to the next chunk. /// public sealed class ReadIncrementalInstruction : ReadIncrementalInstructionBase { diff --git a/Runtime/Scripts/DataStreams/DataStream.cs b/Runtime/Scripts/DataStreams/DataStream.cs index 68ee95a9..fd5f7eab 100644 --- a/Runtime/Scripts/DataStreams/DataStream.cs +++ b/Runtime/Scripts/DataStreams/DataStream.cs @@ -72,19 +72,43 @@ internal StreamError(Proto.StreamError proto) : base(proto.Description) { } /// /// Shared state and helpers for incremental stream reader yield instructions. - /// Holds the latest chunk, end-of-stream flag, and error; subclasses own the - /// typed event subscription and convert raw event payloads via - /// and . + /// Models the incoming stream as a single ordered queue whose items are either a + /// chunk or a the end-of-stream marker (carrying an optional ). + /// The consumer advances through the queue one item at a time via ; + /// end-of-stream is only observed once every buffered chunk has been drained, mirroring + /// the JS and Python SDKs. Subclasses own the typed event subscription and convert raw + /// event payloads via and . /// public abstract class ReadIncrementalInstructionBase : StreamYieldInstruction { private readonly ulong _handleValue; - private readonly Queue _pendingChunks = new(); - private TContent _latestChunk; - // Chunk events arrive on the FFI thread; Reset() and the LatestChunk getter - // run on the main-thread coroutine. _gate serializes mutations of the queue, - // _latestChunk, IsCurrentReadDone, IsEos, and Error across both sides. + // A chunk or the end-of-stream marker. Chunks and the marker share one + // FIFO so the consumer can never observe end-of-stream while chunks remain buffered. + // The head of the queue is the item the consumer is currently positioned on; Reset() + // dequeues it to advance to the next. + private readonly struct Item + { + public readonly TContent Chunk; + public readonly bool IsEos; + public readonly StreamError Error; + + private Item(TContent chunk, bool isEos, StreamError error) + { + Chunk = chunk; + IsEos = isEos; + Error = error; + } + + public static Item ForChunk(TContent chunk) => new Item(chunk, false, null); + public static Item ForEos(StreamError error) => new Item(default, true, error); + } + + private readonly Queue _queue = new(); + + // Chunk/EOS events arrive on the FFI thread; Reset() and the LatestChunk getter run on + // the main-thread coroutine. _gate serializes mutations of the queue, IsCurrentReadDone, + // IsEos, and Error across both sides. private readonly object _gate = new(); /// @@ -99,10 +123,11 @@ public abstract class ReadIncrementalInstructionBase : StreamYieldInst /// /// The chunk from the most recent completed read. Throws the captured - /// if the last read errored. Internal so the optional - /// UniTask async-enumerable adapter (which has InternalsVisibleTo access) can read - /// it generically; the typed Bytes/Text accessors on the concrete - /// readers delegate here. + /// if the stream ended with an error. Returns the default + /// value once positioned on a normal end-of-stream marker (there is no chunk to read). + /// Internal so the optional UniTask async-enumerable adapter (which has + /// InternalsVisibleTo access) can read it generically; the typed Bytes/Text + /// accessors on the concrete readers delegate here. /// internal TContent LatestChunk { @@ -111,7 +136,7 @@ internal TContent LatestChunk lock (_gate) { if (Error != null) throw Error; - return _latestChunk; + return _queue.Count > 0 ? _queue.Peek().Chunk : default; } } } @@ -123,55 +148,57 @@ protected ReadIncrementalInstructionBase(FfiHandle readerHandle) protected bool MatchesHandle(ulong eventHandle) => eventHandle == _handleValue; - protected void OnChunk(TContent content) + protected void OnChunk(TContent content) => Enqueue(Item.ForChunk(content)); + + protected void OnEos(Proto.StreamError protoError) => + Enqueue(Item.ForEos(protoError != null ? new StreamError(protoError) : null)); + + private void Enqueue(Item item) { lock (_gate) { - if (IsCurrentReadDone) - { - // Consumer hasn't yielded since the last chunk; buffer until Reset(). - _pendingChunks.Enqueue(content); - } - else - { - _latestChunk = content; - IsCurrentReadDone = true; - } + // When the queue was empty this item becomes the head, so publish it now; + // otherwise it stays buffered in order behind the current head. + bool wasEmpty = _queue.Count == 0; + _queue.Enqueue(item); + if (wasEmpty) SettleHead(); } } public override void Reset() { - // base.Reset() must run under the same lock as OnChunk, otherwise the - // window between IsCurrentReadDone=false (from base) and the dequeue - // below lets a producer race in, write _latestChunk, and have its - // chunk immediately overwritten by the dequeue. That race lost ~4% of - // chunks under stress before this fix. + // base.Reset() must run under the same lock as OnChunk/OnEos so a producer can't + // race between the flag clear and the dequeue below. It also throws when the + // consumer tries to advance past the Eos marker (IsEos), the intended guard. lock (_gate) { base.Reset(); - if (_pendingChunks.Count > 0) - { - _latestChunk = _pendingChunks.Dequeue(); - IsCurrentReadDone = true; - } + if (_queue.Count > 0) _queue.Dequeue(); // drop the item just consumed + SettleHead(); } } - protected void OnEos(Proto.StreamError protoError) + // Reflects the queue head in the base completion flags: a chunk becomes readable + // (IsCurrentReadDone), the Eos marker (IsEos) ends the stream. An empty queue + // parks the consumer (keepWaiting stays true) until the next item arrives. + // Caller must hold _gate. + private void SettleHead() { - lock (_gate) + if (_queue.Count == 0) return; + var head = _queue.Peek(); + if (head.IsEos) { // Assign Error before flipping IsEos. The IsEos setter fires the awaiter - // continuation, which inspects IsError/Error on resume; when completion runs - // inline on the main thread, setting IsEos first would let the continuation - // observe IsError == false and silently swallow the stream error. - if (protoError != null) - { - Error = new StreamError(protoError); - } + // continuation, which inspects IsError/Error on resume; setting IsEos first + // would let the continuation observe IsError == false and silently swallow + // the stream error. + if (head.Error != null) Error = head.Error; IsEos = true; } + else + { + IsCurrentReadDone = true; + } } } diff --git a/Runtime/Scripts/DataStreams/TextDataStream.cs b/Runtime/Scripts/DataStreams/TextDataStream.cs index f6181a2c..bddcbb77 100644 --- a/Runtime/Scripts/DataStreams/TextDataStream.cs +++ b/Runtime/Scripts/DataStreams/TextDataStream.cs @@ -136,8 +136,9 @@ public ReadIncrementalInstruction ReadIncremental() /// YieldInstruction for . /// /// - /// Usage: while is false (i.e. the stream has not ended), - /// call , yield the instruction, and then access . + /// Usage: loop yielding the instruction; when it resumes, break if + /// is true (the stream has ended and every chunk has been drained), otherwise read + /// and call to advance to the next chunk. /// public sealed class ReadIncrementalInstruction : ReadIncrementalInstructionBase { diff --git a/Runtime/Scripts/UniTask/StreamReaderUniTaskExtensions.cs b/Runtime/Scripts/UniTask/StreamReaderUniTaskExtensions.cs index 45ea1284..5eb4f144 100644 --- a/Runtime/Scripts/UniTask/StreamReaderUniTaskExtensions.cs +++ b/Runtime/Scripts/UniTask/StreamReaderUniTaskExtensions.cs @@ -18,17 +18,13 @@ public static class StreamReaderUniTaskExtensions /// (string). /// /// - /// Iteration ends when the stream reaches end-of-stream. If the stream ends with an - /// error, the enumerable throws that (idiomatic for - /// await foreach; this is the one place the UniTask surface throws rather than - /// exposing IsError). Cancellation (via the token or the enumerator) surfaces as - /// with abandon-awaiter semantics — the - /// underlying FFI read is not cancelled on the wire. - /// - /// The current chunk is delivered on the iteration where end-of-stream is also observed, - /// then iteration stops. Chunks buffered beyond the current one when - /// end-of-stream arrives are not drainable, because the reader disallows Reset() - /// past end-of-stream. + /// Iteration ends when the stream reaches end-of-stream. Every buffered chunk is + /// drained before iteration stops, since end-of-stream is an ordered item at the tail + /// of the chunk queue. If the stream ends with an error, the enumerable throws that + /// (idiomatic for await foreach; this is the one place + /// the UniTask surface throws rather than exposing IsError). Cancellation (via + /// the token or the enumerator) surfaces as + /// with abandon-awaiter semantics — the underlying FFI read is not cancelled on the wire. /// public static IUniTaskAsyncEnumerable AsAsyncEnumerable( this ReadIncrementalInstructionBase instruction, @@ -44,30 +40,20 @@ public static IUniTaskAsyncEnumerable AsAsyncEnumerable( while (true) { - // Completes when a chunk is ready (IsCurrentReadDone) or the stream ends (IsEos). + // Completes when a chunk is ready (IsCurrentReadDone) or the consumer has + // advanced onto the terminal marker (IsEos) — the two are mutually exclusive. await instruction.AsUniTask(ct); if (instruction.IsCurrentReadDone) { var chunk = instruction.LatestChunk; await writer.YieldAsync(chunk); - - // Re-check IsEos AFTER yielding: end-of-stream may have arrived while - // the consumer was suspended. Reset() throws once IsEos is set, so this - // re-check (not a value captured before the yield) is what keeps the - // loop safe — mirroring the coroutine consumer's "if (IsEos) break; - // else Reset()" ordering. - if (instruction.IsEos) - { - if (instruction.IsError) throw instruction.Error; - return; - } - - instruction.Reset(); + instruction.Reset(); // advance to the next chunk or the terminal marker continue; } - // Not IsCurrentReadDone => end-of-stream with nothing left to read. + // Positioned on the terminal marker: the queue is fully drained. Surface an + // error close by throwing, otherwise end iteration cleanly. if (instruction.IsError) throw instruction.Error; return; } diff --git a/Samples~/Agents/Assets/Runtime/Agent/TranscriptionReader.cs b/Samples~/Agents/Assets/Runtime/Agent/TranscriptionReader.cs index aaa77186..e8e104c0 100644 --- a/Samples~/Agents/Assets/Runtime/Agent/TranscriptionReader.cs +++ b/Samples~/Agents/Assets/Runtime/Agent/TranscriptionReader.cs @@ -69,24 +69,18 @@ IEnumerator PumpNpc(TextStreamReader reader, string identity) var append = _beginReply(Speaker.Npc); var filter = new AnnotationFilter(); var r = reader.ReadIncremental(); - string lastSeen = null; while (true) { yield return r; + if (r.IsEos) break; - // EOS ticks don't refresh Text — only forward when the reference actually changed. var chunk = r.Text; - if (!ReferenceEquals(chunk, lastSeen)) + if (!string.IsNullOrEmpty(chunk)) { - if (!string.IsNullOrEmpty(chunk)) - { - var filtered = filter.Filter(chunk); - if (!string.IsNullOrEmpty(filtered)) append(filtered); - } - lastSeen = chunk; + var filtered = filter.Filter(chunk); + if (!string.IsNullOrEmpty(filtered)) append(filtered); } - if (r.IsEos) break; r.Reset(); } } diff --git a/Tests/EditMode/DataStreamIncrementalReadTests.cs b/Tests/EditMode/DataStreamIncrementalReadTests.cs index 0129f6d0..f3aba861 100644 --- a/Tests/EditMode/DataStreamIncrementalReadTests.cs +++ b/Tests/EditMode/DataStreamIncrementalReadTests.cs @@ -86,12 +86,11 @@ public void OnChunk_ConcurrentProducerAndConsumer_AllChunksObservedInOrder() Assert.AreEqual(i.ToString(), observed[i], $"Chunk reordering at index {i}."); } - // OnEos races with OnChunk on the FFI thread. After both complete, the lock - // must leave the instruction in a consistent state: IsEos visible, no torn - // reads of LatestChunk, and the chunks pushed before EOS still drainable - // (up to the point where Reset() is correctly disallowed past EOS). + // OnEos races with OnChunk on the FFI thread. EOS is an ordered item behind the + // chunks, so after both complete the consumer must be able to drain every chunk in + // FIFO order before observing end-of-stream — never seeing EOS while data remains. [Test] - public void OnEos_RacingWithChunks_FinalStateConsistent() + public void OnEos_RacingWithChunks_AllChunksDrainThenEos() { for (int trial = 0; trial < 50; trial++) { @@ -109,20 +108,128 @@ public void OnEos_RacingWithChunks_FinalStateConsistent() start.Set(); producer.Wait(TimeSpan.FromSeconds(5)); - Assert.IsTrue(reader.IsEos, $"Trial {trial}: EOS not observed after producer finished."); - // Reading LatestChunk after EOS must never throw an unexpected exception - // — the only allowed throw is the StreamError on protocol error, which - // we don't simulate here. Specifically: no NullReferenceException, no - // collection-modified exception, no torn-read. + var observed = new List(100); + var deadline = DateTime.UtcNow.AddSeconds(5); + while (true) + { + while (reader.keepWaiting) + { + if (DateTime.UtcNow > deadline) + Assert.Fail($"Trial {trial}: drain stalled at {observed.Count}/100 — chunk lost or deadlock."); + Thread.Yield(); + } + if (reader.IsEos) break; + observed.Add(reader.Value); + reader.Reset(); + } + + Assert.AreEqual(100, observed.Count, + $"Trial {trial}: EOS observed before all chunks were drained."); + for (int i = 0; i < 100; i++) + Assert.AreEqual(i.ToString(), observed[i], $"Trial {trial}: chunk reordering at index {i}."); + + // A clean end-of-stream leaves no chunk to read: LatestChunk returns the + // default value and never throws (a torn read would surface here). Assert.DoesNotThrow(() => { var _ = reader.Value; }, $"Trial {trial}: reading LatestChunk after EOS race threw unexpectedly."); + Assert.IsNull(reader.Value, $"Trial {trial}: expected null chunk at end-of-stream."); + } + } + + // The whole stream — every chunk AND end-of-stream — arrives before the consumer + // reads anything (the common small-stream / already-buffered case). Every chunk must + // still be delivered before EOS is observed. This is the core regression the + // single-queue / EOS-as-terminal-item fix addresses: previously IsEos short-circuited + // the buffered chunks and the whole stream was dropped. + [Test] + public void Burst_AllChunksThenEosBeforeConsumer_AllDrainedThenEos() + { + using var handle = new FfiHandle(IntPtr.Zero); + var reader = new TestIncrementalReader(handle); + + reader.PushChunk("A"); + reader.PushChunk("B"); + reader.PushChunk("C"); + reader.PushEos(); + + CollectionAssert.AreEqual(new[] { "A", "B", "C" }, Drain(reader)); + Assert.IsTrue(reader.IsEos); + } + + // A single chunk followed immediately by end-of-stream — the previous implementation + // dropped the chunk because IsEos short-circuited the read on the first iteration. + [Test] + public void Burst_SingleChunkThenEos_ChunkDelivered() + { + using var handle = new FfiHandle(IntPtr.Zero); + var reader = new TestIncrementalReader(handle); + + reader.PushChunk("only"); + reader.PushEos(); + + CollectionAssert.AreEqual(new[] { "only" }, Drain(reader)); + Assert.IsTrue(reader.IsEos); + } + + // An empty stream (end-of-stream with no chunks) drains to zero chunks and reports EOS. + [Test] + public void Burst_EosWithNoChunks_ReportsEosWithNoChunks() + { + using var handle = new FfiHandle(IntPtr.Zero); + var reader = new TestIncrementalReader(handle); + + reader.PushEos(); + + CollectionAssert.IsEmpty(Drain(reader)); + Assert.IsTrue(reader.IsEos); + } + + // Chunks that precede an error close are still delivered in order; the error then + // surfaces via IsError/Error once the consumer reaches the terminal marker. + [Test] + public void Burst_ChunksThenErrorEos_ChunksDrainedThenErrorSurfaces() + { + using var handle = new FfiHandle(IntPtr.Zero); + var reader = new TestIncrementalReader(handle); + + reader.PushChunk("A"); + reader.PushChunk("B"); + reader.PushEos(new LiveKit.Proto.StreamError { Description = "boom" }); + + // Read-first drain; break on EOS before reading (an error terminal has no chunk, + // and reading it would rethrow the StreamError). + var observed = new List(); + while (!reader.keepWaiting) + { + if (reader.IsEos) break; + observed.Add(reader.Value); + reader.Reset(); + } + + CollectionAssert.AreEqual(new[] { "A", "B" }, observed); + Assert.IsTrue(reader.IsEos); + Assert.IsTrue(reader.IsError); + Assert.AreEqual("boom", reader.Error.Message); + } + + // Canonical read-first drain used by the tests above: spin until the reader is + // positioned on a chunk or the terminal marker, break on EOS, else read and advance. + private static List Drain(TestIncrementalReader reader) + { + var observed = new List(); + while (!reader.keepWaiting) + { + if (reader.IsEos) break; + observed.Add(reader.Value); + reader.Reset(); } + return observed; } // OnChunk-after-Reset is the most common interleaving in production: the // consumer's coroutine just yielded and called Reset, and an FFI thread // pushes a chunk before the consumer wakes. The lock must serialize them - // so the new chunk is correctly placed (either as _latestChunk if the + // so the new chunk is correctly placed (either as the current item if the // queue was empty, or appended to the queue). [Test] public void OnChunk_RacingWithReset_NoChunkLost() diff --git a/Tests/PlayMode/UniTask/StreamUniTaskTests.cs b/Tests/PlayMode/UniTask/StreamUniTaskTests.cs index eaa067eb..84f05f80 100644 --- a/Tests/PlayMode/UniTask/StreamUniTaskTests.cs +++ b/Tests/PlayMode/UniTask/StreamUniTaskTests.cs @@ -55,9 +55,8 @@ public System.Collections.IEnumerator AsAsyncEnumerable_DeliversChunksInOrder_Th } }); - // The current chunk is delivered even when EoS is already set at the time it is read, - // then the sequence ends. (Chunks buffered beyond the current one when EoS arrives are - // not drainable, because the reader disallows Reset() past EoS.) + // The final chunk is delivered even when EoS is already set at the time it is read, + // then the sequence ends. [UnityTest] public System.Collections.IEnumerator AsAsyncEnumerable_DeliversFinalChunkThenEos() => UniTask.ToCoroutine(async () => { @@ -74,6 +73,27 @@ public System.Collections.IEnumerator AsAsyncEnumerable_DeliversFinalChunkThenEo CollectionAssert.AreEqual(new[] { "only" }, observed); }); + // Every chunk buffered before EoS is drained, even when all of them plus EoS arrive + // before the consumer reads anything (the already-buffered burst case). EoS is an + // ordered item at the tail of the chunk queue, so it can never short-circuit the data. + [UnityTest] + public System.Collections.IEnumerator AsAsyncEnumerable_DrainsAllBufferedChunks_WhenEosArrivesFirst() => UniTask.ToCoroutine(async () => + { + using var handle = new FfiHandle(IntPtr.Zero); + var reader = new TestIncrementalReader(handle); + + reader.PushChunk("A"); + reader.PushChunk("B"); + reader.PushChunk("C"); + reader.PushEos(); + + var observed = new List(); + await foreach (var chunk in reader.AsAsyncEnumerable()) + observed.Add(chunk); + + CollectionAssert.AreEqual(new[] { "A", "B", "C" }, observed); + }); + // A chunk delivered before the stream errors is observed; the subsequent error EoS // then surfaces as a thrown StreamError. Manual enumeration models the real timeline // (chunk arrives, is consumed, THEN the error ends the stream) — note that once the