Skip to content

Commit ee26791

Browse files
Merge pull request #52260 from CyrusNajmabadi/cloudcacheCrash
Address crash in the cloud cache system when doing two single byte reads.
2 parents 35997a1 + f8e3f7c commit ee26791

File tree

6 files changed

+304
-5
lines changed

6 files changed

+304
-5
lines changed

src/Tools/IdeCoreBenchmarks/IdeCoreBenchmarks.csproj

+2
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
<ItemGroup>
1717
<Compile Include="..\..\VisualStudio\Core\Def\Storage\AbstractCloudCachePersistentStorageService.cs" Link="CloudCache\AbstractCloudCachePersistentStorageService.cs" />
1818
<Compile Include="..\..\VisualStudio\Core\Def\Storage\CloudCachePersistentStorage.cs" Link="CloudCache\CloudCachePersistentStorage.cs" />
19+
<Compile Include="..\..\VisualStudio\Core\Def\Storage\Nerdbank\ReadOnlySequenceStream.cs" Link="CloudCache\ReadOnlySequenceStream.cs" />
1920
<Compile Include="..\..\VisualStudio\Core\Def\Storage\ProjectContainerKeyCache.cs" Link="CloudCache\ProjectContainerKeyCache.cs" />
2021
<Compile Include="..\..\VisualStudio\CSharp\Test\PersistentStorage\Mocks\AuthorizationServiceMock.cs" Link="CloudCache\AuthorizationServiceMock.cs" />
2122
<Compile Include="..\..\VisualStudio\CSharp\Test\PersistentStorage\Mocks\FileSystemServiceMock.cs" Link="CloudCache\FileSystemServiceMock.cs" />
@@ -38,6 +39,7 @@
3839
<PackageReference Include="Microsoft.Win32.Registry" Version="$(MicrosoftWin32RegistryVersion)" />
3940
<PackageReference Include="Newtonsoft.Json" Version="$(NewtonsoftJsonVersion)" />
4041
<PackageReference Include="Microsoft.VisualStudio.Cache" Version="$(MicrosoftVisualStudioCacheVersion)" />
42+
<PackageReference Include="Nerdbank.Streams" Version="$(NerdbankStreamsVersion)" />
4143
</ItemGroup>
4244

4345
<ItemGroup>

src/VisualStudio/CSharp/Test/PersistentStorage/AbstractPersistentStorageTests.cs

+20
Original file line numberDiff line numberDiff line change
@@ -786,6 +786,26 @@ public void CacheDirectoryShouldNotBeAtRoot()
786786
Assert.False(location?.StartsWith("/") ?? false);
787787
}
788788

789+
[Theory]
790+
[CombinatorialData]
791+
public async Task PersistentService_ReadByteTwice(Size size, bool withChecksum)
792+
{
793+
var solution = CreateOrOpenSolution();
794+
var streamName1 = "PersistentService_ReadByteTwice";
795+
796+
await using (var storage = await GetStorageAsync(solution))
797+
{
798+
Assert.True(await storage.WriteStreamAsync(streamName1, EncodeString(GetData1(size)), GetChecksum1(withChecksum)));
799+
}
800+
801+
await using (var storage = await GetStorageAsync(solution))
802+
{
803+
using var stream = await storage.ReadStreamAsync(streamName1, GetChecksum1(withChecksum));
804+
stream.ReadByte();
805+
stream.ReadByte();
806+
}
807+
}
808+
789809
[PartNotDiscoverable]
790810
[ExportWorkspaceService(typeof(IPersistentStorageLocationService), layer: ServiceLayer.Test), Shared]
791811
private class TestPersistentStorageLocationService : DefaultPersistentStorageLocationService

src/VisualStudio/Core/Def/Microsoft.VisualStudio.LanguageServices.csproj

+1
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,7 @@
169169
<PackageReference Include="NuGet.VisualStudio" Version="$(NuGetVisualStudioVersion)" />
170170
<PackageReference Include="Newtonsoft.Json" Version="$(NewtonsoftJsonVersion)" />
171171
<PackageReference Include="System.Threading.Tasks.Dataflow" Version="$(SystemThreadingTasksDataflowVersion)" />
172+
<PackageReference Include="Nerdbank.Streams" Version="$(NerdbankStreamsVersion)" />
172173

173174
<!-- By default build assets that define embedded interop types do not flow. Set PrivateAssets to none to make them flow. -->
174175
<PackageReference Include="Microsoft.VisualStudio.SDK.EmbedInteropTypes" Version="$(MicrosoftVisualStudioSDKEmbedInteropTypesVersion)" PrivateAssets="none" />

src/VisualStudio/Core/Def/Storage/CloudCachePersistentStorage.cs

+17-5
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
using Microsoft.CodeAnalysis.PersistentStorage;
1414
using Microsoft.CodeAnalysis.PooledObjects;
1515
using Microsoft.VisualStudio.RpcContracts.Caching;
16+
using Nerdbank.Streams;
1617
using Roslyn.Utilities;
1718

1819
namespace Microsoft.VisualStudio.LanguageServices.Storage
@@ -163,16 +164,27 @@ private async Task<bool> ChecksumMatchesAsync(string name, Checksum checksum, Ca
163164
// and then pass that out. This should not be a problem in practice as PipeReader internally intelligently
164165
// uses and pools reasonable sized buffers, preventing us from exacerbating the GC or causing LOH
165166
// allocations.
167+
return await AsPrebufferedStreamAsync(pipe.Reader, cancellationToken).ConfigureAwait(false);
168+
}
169+
170+
private static async Task<Stream> AsPrebufferedStreamAsync(PipeReader pipeReader, CancellationToken cancellationToken = default)
171+
{
166172
while (true)
167173
{
168-
var readResult = await pipe.Reader.ReadAsync(cancellationToken).ConfigureAwait(false);
169-
pipe.Reader.AdvanceTo(readResult.Buffer.Start, readResult.Buffer.End);
174+
// Read and immediately report all bytes as "examined" so that the next ReadAsync call will block till more bytes come in.
175+
// The goal here is to force the PipeReader to buffer everything internally (even if it were to exceed its natural writer threshold limit).
176+
ReadResult readResult = await pipeReader.ReadAsync(cancellationToken).ConfigureAwait(false);
177+
pipeReader.AdvanceTo(readResult.Buffer.Start, readResult.Buffer.End);
170178

171179
if (readResult.IsCompleted)
172-
break;
180+
{
181+
// After having buffered and "examined" all the bytes, the stream returned from PipeReader.AsStream() would fail
182+
// because it may not "examine" all bytes at once.
183+
// Instead, we'll create our own Stream over just the buffer itself, and recycle the buffers when the stream is disposed
184+
// the way the stream returned from PipeReader.AsStream() would have.
185+
return new ReadOnlySequenceStream(readResult.Buffer, reader => ((PipeReader)reader!).Complete(), pipeReader);
186+
}
173187
}
174-
175-
return pipe.Reader.AsStream();
176188
}
177189

178190
public sealed override Task<bool> WriteStreamAsync(string name, Stream stream, Checksum? checksum, CancellationToken cancellationToken)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
// Copied from https://raw.githubusercontent.com/AArnott/Nerdbank.Streams/2b142fa6a38b15e4b06ecc53bf073aa49fd1de34/src/Nerdbank.Streams/ReadOnlySequenceStream.cs
6+
// Remove once we move to Nerdbank.Streams 2.7.62-alpha
7+
8+
namespace Nerdbank.Streams
9+
{
10+
using System;
11+
using System.Buffers;
12+
using System.IO;
13+
using System.Runtime.InteropServices;
14+
using System.Threading;
15+
using System.Threading.Tasks;
16+
using Microsoft;
17+
18+
internal class ReadOnlySequenceStream : Stream, IDisposableObservable
19+
{
20+
private static readonly Task<int> TaskOfZero = Task.FromResult(0);
21+
22+
private readonly Action<object?>? disposeAction;
23+
private readonly object? disposeActionArg;
24+
25+
/// <summary>
26+
/// A reusable task if two consecutive reads return the same number of bytes.
27+
/// </summary>
28+
private Task<int>? lastReadTask;
29+
30+
private readonly ReadOnlySequence<byte> readOnlySequence;
31+
32+
private SequencePosition position;
33+
34+
internal ReadOnlySequenceStream(ReadOnlySequence<byte> readOnlySequence, Action<object?>? disposeAction, object? disposeActionArg)
35+
{
36+
this.readOnlySequence = readOnlySequence;
37+
this.disposeAction = disposeAction;
38+
this.disposeActionArg = disposeActionArg;
39+
this.position = readOnlySequence.Start;
40+
}
41+
42+
/// <inheritdoc/>
43+
public override bool CanRead => !this.IsDisposed;
44+
45+
/// <inheritdoc/>
46+
public override bool CanSeek => !this.IsDisposed;
47+
48+
/// <inheritdoc/>
49+
public override bool CanWrite => false;
50+
51+
/// <inheritdoc/>
52+
public override long Length => this.ReturnOrThrowDisposed(this.readOnlySequence.Length);
53+
54+
/// <inheritdoc/>
55+
public override long Position
56+
{
57+
get => this.readOnlySequence.Slice(0, this.position).Length;
58+
set
59+
{
60+
Requires.Range(value >= 0, nameof(value));
61+
this.position = this.readOnlySequence.GetPosition(value, this.readOnlySequence.Start);
62+
}
63+
}
64+
65+
/// <inheritdoc/>
66+
public bool IsDisposed { get; private set; }
67+
68+
/// <inheritdoc/>
69+
public override void Flush() => this.ThrowDisposedOr(new NotSupportedException());
70+
71+
/// <inheritdoc/>
72+
public override Task FlushAsync(CancellationToken cancellationToken) => throw this.ThrowDisposedOr(new NotSupportedException());
73+
74+
/// <inheritdoc/>
75+
public override int Read(byte[] buffer, int offset, int count)
76+
{
77+
ReadOnlySequence<byte> remaining = this.readOnlySequence.Slice(this.position);
78+
ReadOnlySequence<byte> toCopy = remaining.Slice(0, Math.Min(count, remaining.Length));
79+
this.position = toCopy.End;
80+
toCopy.CopyTo(buffer.AsSpan(offset, count));
81+
return (int)toCopy.Length;
82+
}
83+
84+
/// <inheritdoc/>
85+
public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
86+
{
87+
cancellationToken.ThrowIfCancellationRequested();
88+
int bytesRead = this.Read(buffer, offset, count);
89+
if (bytesRead == 0)
90+
{
91+
return TaskOfZero;
92+
}
93+
94+
if (this.lastReadTask?.Result == bytesRead)
95+
{
96+
return this.lastReadTask;
97+
}
98+
else
99+
{
100+
return this.lastReadTask = Task.FromResult(bytesRead);
101+
}
102+
}
103+
104+
/// <inheritdoc/>
105+
public override int ReadByte()
106+
{
107+
ReadOnlySequence<byte> remaining = this.readOnlySequence.Slice(this.position);
108+
if (remaining.Length > 0)
109+
{
110+
byte result = remaining.First.Span[0];
111+
this.position = this.readOnlySequence.GetPosition(1, this.position);
112+
return result;
113+
}
114+
else
115+
{
116+
return -1;
117+
}
118+
}
119+
120+
/// <inheritdoc/>
121+
public override long Seek(long offset, SeekOrigin origin)
122+
{
123+
Verify.NotDisposed(this);
124+
125+
SequencePosition relativeTo;
126+
switch (origin)
127+
{
128+
case SeekOrigin.Begin:
129+
relativeTo = this.readOnlySequence.Start;
130+
break;
131+
case SeekOrigin.Current:
132+
if (offset >= 0)
133+
{
134+
relativeTo = this.position;
135+
}
136+
else
137+
{
138+
relativeTo = this.readOnlySequence.Start;
139+
offset += this.Position;
140+
}
141+
142+
break;
143+
case SeekOrigin.End:
144+
if (offset >= 0)
145+
{
146+
relativeTo = this.readOnlySequence.End;
147+
}
148+
else
149+
{
150+
relativeTo = this.readOnlySequence.Start;
151+
offset += this.Position;
152+
}
153+
154+
break;
155+
default:
156+
throw new ArgumentOutOfRangeException(nameof(origin));
157+
}
158+
159+
this.position = this.readOnlySequence.GetPosition(offset, relativeTo);
160+
return this.Position;
161+
}
162+
163+
/// <inheritdoc/>
164+
public override void SetLength(long value) => this.ThrowDisposedOr(new NotSupportedException());
165+
166+
/// <inheritdoc/>
167+
public override void Write(byte[] buffer, int offset, int count) => this.ThrowDisposedOr(new NotSupportedException());
168+
169+
/// <inheritdoc/>
170+
public override void WriteByte(byte value) => this.ThrowDisposedOr(new NotSupportedException());
171+
172+
/// <inheritdoc/>
173+
public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => throw this.ThrowDisposedOr(new NotSupportedException());
174+
175+
/// <inheritdoc/>
176+
public override async Task CopyToAsync(Stream destination, int bufferSize, CancellationToken cancellationToken)
177+
{
178+
foreach (var segment in this.readOnlySequence)
179+
{
180+
await WriteAsync(destination, segment, cancellationToken).ConfigureAwait(false);
181+
}
182+
}
183+
184+
private static ValueTask WriteAsync(Stream stream, ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
185+
{
186+
Requires.NotNull(stream, nameof(stream));
187+
188+
if (MemoryMarshal.TryGetArray(buffer, out ArraySegment<byte> array))
189+
{
190+
return new ValueTask(stream.WriteAsync(array.Array!, array.Offset, array.Count, cancellationToken));
191+
}
192+
else
193+
{
194+
byte[] sharedBuffer = ArrayPool<byte>.Shared.Rent(buffer.Length);
195+
buffer.Span.CopyTo(sharedBuffer);
196+
return new ValueTask(FinishWriteAsync(stream.WriteAsync(sharedBuffer, 0, buffer.Length, cancellationToken), sharedBuffer));
197+
}
198+
199+
async Task FinishWriteAsync(Task writeTask, byte[] localBuffer)
200+
{
201+
try
202+
{
203+
await writeTask.ConfigureAwait(false);
204+
}
205+
finally
206+
{
207+
ArrayPool<byte>.Shared.Return(localBuffer);
208+
}
209+
}
210+
}
211+
212+
#if SPAN_BUILTIN
213+
214+
/// <inheritdoc/>
215+
public override int Read(Span<byte> buffer)
216+
{
217+
ReadOnlySequence<byte> remaining = this.readOnlySequence.Slice(this.position);
218+
ReadOnlySequence<byte> toCopy = remaining.Slice(0, Math.Min(buffer.Length, remaining.Length));
219+
this.position = toCopy.End;
220+
toCopy.CopyTo(buffer);
221+
return (int)toCopy.Length;
222+
}
223+
224+
/// <inheritdoc/>
225+
public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
226+
{
227+
cancellationToken.ThrowIfCancellationRequested();
228+
return new ValueTask<int>(this.Read(buffer.Span));
229+
}
230+
231+
/// <inheritdoc/>
232+
public override void Write(ReadOnlySpan<byte> buffer) => throw this.ThrowDisposedOr(new NotSupportedException());
233+
234+
/// <inheritdoc/>
235+
public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default) => throw this.ThrowDisposedOr(new NotSupportedException());
236+
237+
#endif
238+
239+
/// <inheritdoc/>
240+
protected override void Dispose(bool disposing)
241+
{
242+
if (!this.IsDisposed)
243+
{
244+
this.IsDisposed = true;
245+
this.disposeAction?.Invoke(this.disposeActionArg);
246+
base.Dispose(disposing);
247+
}
248+
}
249+
250+
private T ReturnOrThrowDisposed<T>(T value)
251+
{
252+
Verify.NotDisposed(this);
253+
return value;
254+
}
255+
256+
private Exception ThrowDisposedOr(Exception ex)
257+
{
258+
Verify.NotDisposed(this);
259+
throw ex;
260+
}
261+
}
262+
}

src/Workspaces/Remote/ServiceHub/Microsoft.CodeAnalysis.Remote.ServiceHub.csproj

+2
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
<PackageReference Include="Microsoft.VisualStudio.Threading" Version="$(MicrosoftVisualStudioThreadingVersion)" />
3030
<PackageReference Include="Microsoft.ServiceHub.Framework" Version="$(MicrosoftServiceHubFrameworkVersion)" />
3131
<PackageReference Include="Microsoft.VisualStudio.RpcContracts" Version="$(MicrosoftVisualStudioRpcContractsVersion)" />
32+
<PackageReference Include="Nerdbank.Streams" Version="$(NerdbankStreamsVersion)" />
3233
</ItemGroup>
3334
<ItemGroup>
3435
<PublicAPI Include="PublicAPI.Shipped.txt" />
@@ -39,6 +40,7 @@
3940
<Compile Include="..\..\..\VisualStudio\Core\Def\Implementation\Watson\WatsonReporter.cs" Link="Services\WorkspaceTelemetry\WatsonReporter.cs" />
4041
<Compile Include="..\..\..\VisualStudio\Core\Def\Storage\AbstractCloudCachePersistentStorageService.cs" Link="Host\Storage\AbstractCloudCachePersistentStorageService.cs" />
4142
<Compile Include="..\..\..\VisualStudio\Core\Def\Storage\CloudCachePersistentStorage.cs" Link="Host\Storage\CloudCachePersistentStorage.cs" />
43+
<Compile Include="..\..\..\VisualStudio\Core\Def\Storage\Nerdbank\ReadOnlySequenceStream.cs" Link="Host\Storage\ReadOnlySequenceStream.cs" />
4244
<Compile Include="..\..\..\VisualStudio\Core\Def\Storage\ProjectContainerKeyCache.cs" Link="Host\Storage\ProjectContainerKeyCache.cs" />
4345
<Compile Include="..\..\..\VisualStudio\Core\Def\Telemetry\VSTelemetryCache.cs" Link="Services\WorkspaceTelemetry\VSTelemetryCache.cs" />
4446
<Compile Include="..\..\..\VisualStudio\Core\Def\Telemetry\VSTelemetryLogger.cs" Link="Services\WorkspaceTelemetry\VSTelemetryLogger.cs" />

0 commit comments

Comments
 (0)