Skip to content

Commit ffaf45e

Browse files
sharwellsharwell
committed
Use a COM message filter for retry semantics
See dotnet/roslyn#25455
1 parent 4965677 commit ffaf45e

File tree

3 files changed

+153
-40
lines changed

3 files changed

+153
-40
lines changed

‎Tvl.VisualStudio.MouseFastScroll.IntegrationTests/AbstractIntegrationTest.cs

Lines changed: 146 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,32 +4,60 @@
44
namespace Tvl.VisualStudio.MouseFastScroll.IntegrationTests
55
{
66
using System;
7+
using System.Runtime.InteropServices;
78
using System.Threading;
89
using System.Threading.Tasks;
910
using System.Windows.Automation;
11+
using Microsoft.VisualStudio;
12+
using Microsoft.Win32.SafeHandles;
1013
using Xunit;
14+
using IMessageFilter = Microsoft.VisualStudio.OLE.Interop.IMessageFilter;
15+
using INTERFACEINFO = Microsoft.VisualStudio.OLE.Interop.INTERFACEINFO;
16+
using PENDINGMSG = Microsoft.VisualStudio.OLE.Interop.PENDINGMSG;
17+
using SERVERCALL = Microsoft.VisualStudio.OLE.Interop.SERVERCALL;
1118

1219
[CaptureTestName]
1320
[Collection(nameof(SharedIntegrationHostFixture))]
1421
public abstract class AbstractIntegrationTest : IAsyncLifetime, IDisposable
1522
{
23+
private readonly MessageFilter _messageFilter;
1624
private readonly VisualStudioInstanceFactory _instanceFactory;
1725
private readonly Version _version;
1826
private VisualStudioInstanceContext _visualStudioContext;
1927

2028
protected AbstractIntegrationTest(VisualStudioInstanceFactory instanceFactory, Version version)
2129
{
2230
Assert.Equal(ApartmentState.STA, Thread.CurrentThread.GetApartmentState());
31+
32+
// Install a COM message filter to handle retry operations when the first attempt fails
33+
_messageFilter = RegisterMessageFilter();
2334
_instanceFactory = instanceFactory;
2435
_version = version;
25-
Automation.TransactionTimeout = 20000;
36+
37+
try
38+
{
39+
Automation.TransactionTimeout = 20000;
40+
}
41+
catch
42+
{
43+
_messageFilter.Dispose();
44+
throw;
45+
}
2646
}
2747

2848
public VisualStudioInstance VisualStudio => _visualStudioContext?.Instance;
2949

3050
public virtual async Task InitializeAsync()
3151
{
32-
_visualStudioContext = await _instanceFactory.GetNewOrUsedInstanceAsync(_version, SharedIntegrationHostFixture.RequiredPackageIds).ConfigureAwait(false);
52+
try
53+
{
54+
_visualStudioContext = await _instanceFactory.GetNewOrUsedInstanceAsync(_version, SharedIntegrationHostFixture.RequiredPackageIds).ConfigureAwait(false);
55+
}
56+
catch
57+
{
58+
_messageFilter.Dispose();
59+
throw;
60+
}
3361
}
3462

3563
public Task DisposeAsync()
@@ -43,11 +71,126 @@ public void Dispose()
4371
GC.SuppressFinalize(this);
4472
}
4573

74+
protected virtual MessageFilter RegisterMessageFilter()
75+
=> new MessageFilter();
76+
4677
protected virtual void Dispose(bool disposing)
4778
{
4879
if (disposing)
4980
{
50-
_visualStudioContext.Dispose();
81+
try
82+
{
83+
_visualStudioContext.Dispose();
84+
}
85+
finally
86+
{
87+
_messageFilter.Dispose();
88+
}
89+
}
90+
}
91+
92+
protected class MessageFilter : IMessageFilter, IDisposable
93+
{
94+
protected const uint CancelCall = ~0U;
95+
96+
private readonly MessageFilterSafeHandle _messageFilterRegistration;
97+
private readonly TimeSpan _timeout;
98+
private readonly TimeSpan _retryDelay;
99+
100+
public MessageFilter()
101+
: this(timeout: TimeSpan.FromSeconds(60), retryDelay: TimeSpan.FromMilliseconds(150))
102+
{
103+
}
104+
105+
public MessageFilter(TimeSpan timeout, TimeSpan retryDelay)
106+
{
107+
_timeout = timeout;
108+
_retryDelay = retryDelay;
109+
_messageFilterRegistration = MessageFilterSafeHandle.Register(this);
110+
}
111+
112+
public virtual uint HandleInComingCall(uint dwCallType, IntPtr htaskCaller, uint dwTickCount, INTERFACEINFO[] lpInterfaceInfo)
113+
{
114+
return (uint)SERVERCALL.SERVERCALL_ISHANDLED;
115+
}
116+
117+
public virtual uint RetryRejectedCall(IntPtr htaskCallee, uint dwTickCount, uint dwRejectType)
118+
{
119+
if ((SERVERCALL)dwRejectType != SERVERCALL.SERVERCALL_RETRYLATER
120+
&& (SERVERCALL)dwRejectType != SERVERCALL.SERVERCALL_REJECTED)
121+
{
122+
return CancelCall;
123+
}
124+
125+
if (dwTickCount >= _timeout.TotalMilliseconds)
126+
{
127+
return CancelCall;
128+
}
129+
130+
return (uint)_retryDelay.TotalMilliseconds;
131+
}
132+
133+
public virtual uint MessagePending(IntPtr htaskCallee, uint dwTickCount, uint dwPendingType)
134+
{
135+
return (uint)PENDINGMSG.PENDINGMSG_WAITDEFPROCESS;
136+
}
137+
138+
protected virtual void Dispose(bool disposing)
139+
{
140+
if (disposing)
141+
{
142+
_messageFilterRegistration.Dispose();
143+
}
144+
}
145+
146+
public void Dispose()
147+
{
148+
Dispose(true);
149+
GC.SuppressFinalize(this);
150+
}
151+
}
152+
153+
private sealed class MessageFilterSafeHandle : SafeHandleMinusOneIsInvalid
154+
{
155+
private readonly IntPtr _oldFilter;
156+
157+
private MessageFilterSafeHandle(IntPtr handle)
158+
: base(true)
159+
{
160+
SetHandle(handle);
161+
162+
try
163+
{
164+
if (CoRegisterMessageFilter(handle, out _oldFilter) != VSConstants.S_OK)
165+
{
166+
throw new InvalidOperationException("Failed to register a new message filter");
167+
}
168+
}
169+
catch
170+
{
171+
SetHandleAsInvalid();
172+
throw;
173+
}
174+
}
175+
176+
[DllImport("ole32", SetLastError = true)]
177+
private static extern int CoRegisterMessageFilter(IntPtr messageFilter, out IntPtr oldMessageFilter);
178+
179+
public static MessageFilterSafeHandle Register<T>(T messageFilter)
180+
where T : IMessageFilter
181+
{
182+
var handle = Marshal.GetComInterfaceForObject<T, IMessageFilter>(messageFilter);
183+
return new MessageFilterSafeHandle(handle);
184+
}
185+
186+
protected override bool ReleaseHandle()
187+
{
188+
if (CoRegisterMessageFilter(_oldFilter, out _) == VSConstants.S_OK)
189+
{
190+
Marshal.Release(handle);
191+
}
192+
193+
return true;
51194
}
52195
}
53196
}

‎Tvl.VisualStudio.MouseFastScroll.IntegrationTests/TrivialIntegrationTest.cs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ protected TrivialIntegrationTest(VisualStudioInstanceFactory instanceFactory, Ve
2424
[WpfFact]
2525
public void TestOpenAndCloseIDE()
2626
{
27-
var currentVersion = VisualStudioInstance.RetryRpcCall(() => VisualStudio.Dte.Version);
27+
var currentVersion = VisualStudio.Dte.Version;
2828
var expectedVersion = VisualStudio.Version;
2929
if (expectedVersion.Major >= 15)
3030
{
@@ -38,7 +38,7 @@ public void TestOpenAndCloseIDE()
3838
[WpfFact]
3939
public void BasicScrollingBehavior()
4040
{
41-
var window = VisualStudioInstance.RetryRpcCall(() => VisualStudio.Dte.ItemOperations.NewFile(Name: Guid.NewGuid() + ".txt"));
41+
var window = VisualStudio.Dte.ItemOperations.NewFile(Name: Guid.NewGuid() + ".txt");
4242

4343
string initialText = string.Join(string.Empty, Enumerable.Range(0, 400).Select(i => Guid.NewGuid() + Environment.NewLine));
4444
VisualStudio.Editor.SetText(initialText);
@@ -124,7 +124,7 @@ public void BasicScrollingBehavior()
124124
Assert.Equal(0, VisualStudio.Editor.GetCaretPosition());
125125
Assert.Equal(0, VisualStudio.Editor.GetFirstVisibleLine());
126126

127-
VisualStudioInstance.RetryRpcCall(() => window.Close(vsSaveChanges.vsSaveChangesNo));
127+
window.Close(vsSaveChanges.vsSaveChangesNo);
128128
}
129129

130130
/// <summary>
@@ -133,7 +133,7 @@ public void BasicScrollingBehavior()
133133
[WpfFact]
134134
public void ZoomDisabled()
135135
{
136-
var window = VisualStudioInstance.RetryRpcCall(() => VisualStudio.Dte.ItemOperations.NewFile(Name: Guid.NewGuid() + ".txt"));
136+
var window = VisualStudio.Dte.ItemOperations.NewFile(Name: Guid.NewGuid() + ".txt");
137137

138138
string initialText = string.Join(string.Empty, Enumerable.Range(0, 400).Select(i => Guid.NewGuid() + Environment.NewLine));
139139
VisualStudio.Editor.SetText(initialText);
@@ -195,7 +195,7 @@ public void ZoomDisabled()
195195
Assert.Equal(0, VisualStudio.Editor.GetFirstVisibleLine());
196196
Assert.Equal(zoomLevel, VisualStudio.Editor.GetZoomLevel());
197197

198-
VisualStudioInstance.RetryRpcCall(() => window.Close(vsSaveChanges.vsSaveChangesNo));
198+
window.Close(vsSaveChanges.vsSaveChangesNo);
199199
}
200200

201201
[VersionTrait(typeof(VS2012))]

‎Tvl.VisualStudio.MouseFastScroll.IntegrationTests/VisualStudioInstance.cs

Lines changed: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -177,9 +177,9 @@ private void CloseRemotingService()
177177
private void StartRemoteIntegrationService(DTE dte)
178178
{
179179
// We use DTE over RPC to start the integration service. All other DTE calls should happen in the host process.
180-
if (RetryRpcCall(() => dte.Commands.Item(WellKnownCommandNames.IntegrationTestServiceStart).IsAvailable))
180+
if (dte.Commands.Item(WellKnownCommandNames.IntegrationTestServiceStart).IsAvailable)
181181
{
182-
RetryRpcCall(() => dte.ExecuteCommand(WellKnownCommandNames.IntegrationTestServiceStart));
182+
dte.ExecuteCommand(WellKnownCommandNames.IntegrationTestServiceStart);
183183
}
184184
}
185185

@@ -190,35 +190,5 @@ private void StopRemoteIntegrationService()
190190
_inProc.ExecuteCommand(WellKnownCommandNames.IntegrationTestServiceStop);
191191
}
192192
}
193-
194-
public static void RetryRpcCall(Action action)
195-
{
196-
do
197-
{
198-
try
199-
{
200-
action();
201-
return;
202-
}
203-
catch (COMException exception) when ((exception.HResult == NativeMethods.RPC_E_CALL_REJECTED) ||
204-
(exception.HResult == NativeMethods.RPC_E_SERVERCALL_RETRYLATER))
205-
{
206-
// We'll just try again in this case
207-
}
208-
}
209-
while (true);
210-
}
211-
212-
public static T RetryRpcCall<T>(Func<T> action)
213-
{
214-
var result = default(T);
215-
216-
RetryRpcCall(() =>
217-
{
218-
result = action();
219-
});
220-
221-
return result;
222-
}
223193
}
224194
}

0 commit comments

Comments
 (0)