diff --git a/src/typstsharp.tests/Tests.cs b/src/typstsharp.tests/Tests.cs index 4a5b967..cb046a0 100644 --- a/src/typstsharp.tests/Tests.cs +++ b/src/typstsharp.tests/Tests.cs @@ -1,3 +1,4 @@ +using System.Buffers; using System.Text; using UglyToad.PdfPig; using UglyToad.PdfPig.Content; @@ -721,6 +722,97 @@ public async Task OutputStreamRejectsReadsAfterTheDocumentIsDisposed() await Assert.That(() => stream.Read(new byte[16].AsSpan())).Throws(); } + /// + /// A span read must not route through , which serves the read + /// from an array rented from the process-wide and returns it without + /// clearing it. That would leave the rendered document readable by the next unrelated component to + /// rent from the same pool, which is exactly what PooledBuffer goes out of its way to prevent. + /// + [Test] + public async Task SpanReadLeavesNoDocumentBytesInTheSharedArrayPool() + { + using var compiler = TypstCompiler.FromSource("= Not for the next renter"); + using var document = compiler.CompileToDocument(); + + int length = (int)document.GetOutputLength(); + + // Prime the pool so that a rent of this size is served from a known array rather than a + // fresh allocation. Nothing may await between here and the final rent, or the thread-local + // pool slot the priming lands in may not be the one the read and the check see. + byte[] primed = ArrayPool.Shared.Rent(length); + primed.AsSpan().Clear(); + ArrayPool.Shared.Return(primed); + + using (var stream = document.OpenOutputStream()) + { + stream.ReadExactly(new byte[length]); + } + + byte[] afterwards = ArrayPool.Shared.Rent(length); + bool carriesTheDocument = afterwards.AsSpan(0, 5).SequenceEqual("%PDF-"u8); + ArrayPool.Shared.Return(afterwards); + + await Assert.That(carriesTheDocument).IsFalse(); + } + + [Test] + public async Task SpanReadReturnsTheWholeBufferAndThenReportsEndOfStream() + { + using var compiler = TypstCompiler.FromSource("= Read by span"); + using var document = compiler.CompileToDocument(); + var expected = document.GetOutputBytes(); + + using var stream = document.OpenOutputStream(); + + // A span larger than the document is the boundary worth pinning: the read must stop at the + // end of the buffer rather than at the end of the span. + var oversized = new byte[expected.Length + 64]; + int read = stream.Read(oversized.AsSpan()); + + await Assert.That(read).IsEqualTo(expected.Length); + await Assert.That(oversized.AsSpan(0, expected.Length).SequenceEqual(expected)).IsTrue(); + await Assert.That(stream.Read(oversized.AsSpan())).IsEqualTo(0); + await Assert.That(stream.Read(Span.Empty)).IsEqualTo(0); + } + + [Test] + public async Task ChunkedSpanReadsReassembleTheDocument() + { + using var compiler = TypstCompiler.FromSource(TwoPageSource); + using var document = compiler.CompileToDocument(); + var expected = document.GetOutputBytes(); + + using var stream = document.OpenOutputStream(); + var reassembled = new byte[expected.Length]; + + // Chunks that do not divide the length evenly, so the final short read is covered too. + const int chunk = 1000; + int offset = 0; + int read; + while ((read = stream.Read(reassembled.AsSpan(offset, Math.Min(chunk, reassembled.Length - offset)))) > 0) + { + offset += read; + } + + await Assert.That(offset).IsEqualTo(expected.Length); + await Assert.That(reassembled.SequenceEqual(expected)).IsTrue(); + } + + [Test] + public async Task AsyncReadMatchesOutputBytes() + { + using var compiler = TypstCompiler.FromSource("= Read asynchronously"); + using var document = compiler.CompileToDocument(); + var expected = document.GetOutputBytes(); + + using var stream = document.OpenOutputStream(); + var destination = new byte[expected.Length]; + int read = await stream.ReadAsync(destination.AsMemory()); + + await Assert.That(read).IsEqualTo(expected.Length); + await Assert.That(destination.SequenceEqual(expected)).IsTrue(); + } + [Test] public async Task OutputLengthMatchesOutputBytes() { diff --git a/src/typstsharp/TypstDocument.cs b/src/typstsharp/TypstDocument.cs index 5813764..84046b4 100644 --- a/src/typstsharp/TypstDocument.cs +++ b/src/typstsharp/TypstDocument.cs @@ -326,24 +326,53 @@ private unsafe void Free() private sealed unsafe class OutputStream : UnmanagedMemoryStream { private readonly TypstDocument _owner; + private readonly byte* _pointer; internal OutputStream(TypstDocument owner, byte* pointer, long length) : base(pointer, length, length, FileAccess.Read) { _owner = owner; + _pointer = pointer; } - // Only the byte-array reads are overridden. Every other read path on Stream and - // UnmanagedMemoryStream, span and async alike, ends up calling one of these two, so this - // covers them all. Overriding Read(Span) as well would recurse: the UnmanagedMemoryStream - // override delegates to Stream.Read(Span) for any derived type, and that implementation - // calls back into Read(byte[]). - public override int Read(byte[] buffer, int offset, int count) + /// + /// Copies straight out of the native buffer, which is what keeps every span-shaped and + /// asynchronous read off the managed heap. + /// + /// + /// only takes its own direct path when the runtime type is + /// exactly ; for a derived type it defers to + /// , which rents an array the size of the caller's span + /// from the shared , reads into it, copies it out, and returns it to + /// the pool without clearing it. That costs a second copy of every byte and, worse, leaves the + /// rendered document in a process-wide pool for whatever rents from it next, which is the very + /// thing PooledBuffer clears its buffer to avoid. Reading the pointer here sidesteps + /// both. It does not recurse: recursion would need a call back into base.Read(Span). + /// + public override int Read(Span buffer) { ObjectDisposedException.ThrowIf(_owner.IsDisposed, _owner); - int read = base.Read(buffer, offset, count); + + long position = Position; + int count = (int)Math.Min((long)buffer.Length, Length - position); + if (count <= 0) + { + GC.KeepAlive(_owner); + return 0; + } + + new ReadOnlySpan(_pointer + position, count).CopyTo(buffer); + Position = position + count; GC.KeepAlive(_owner); - return read; + return count; + } + + // Every other read path on Stream and UnmanagedMemoryStream, byte array and async alike, + // funnels into Read(Span) above. + public override int Read(byte[] buffer, int offset, int count) + { + ValidateBufferArguments(buffer, offset, count); + return Read(new Span(buffer, offset, count)); } public override int ReadByte()