Skip to content

Commit 10a44ad

Browse files
akarnokdOren Novotny
akarnokd
authored and
Oren Novotny
committed
4.x: Make TailRecursiveSink lock-free and have less allocations (#499)
1 parent eb64149 commit 10a44ad

File tree

4 files changed

+134
-127
lines changed

4 files changed

+134
-127
lines changed

‎Rx.NET/Source/src/System.Reactive/Internal/ConcatSink.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,6 @@ public ConcatSink(IObserver<TSource> observer, IDisposable cancel)
1515

1616
protected override IEnumerable<IObservable<TSource>> Extract(IObservable<TSource> source) => (source as IConcatenatable<TSource>)?.GetSources();
1717

18-
public override void OnCompleted() => _recurse();
18+
public override void OnCompleted() => Recurse();
1919
}
2020
}

‎Rx.NET/Source/src/System.Reactive/Internal/TailRecursiveSink.cs‎

Lines changed: 130 additions & 123 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using System.Collections.Generic;
66
using System.Reactive.Concurrency;
77
using System.Reactive.Disposables;
8+
using System.Threading;
89

910
namespace System.Reactive
1011
{
@@ -15,164 +16,176 @@ public TailRecursiveSink(IObserver<TSource> observer, IDisposable cancel)
1516
{
1617
}
1718

18-
private bool _isDisposed;
19-
private SerialDisposable _subscription;
20-
private AsyncLock _gate;
21-
private Stack<IEnumerator<IObservable<TSource>>> _stack;
22-
private Stack<int?> _length;
23-
protected Action _recurse;
19+
bool _isDisposed;
20+
21+
int trampoline;
22+
23+
IDisposable currentSubscription;
24+
25+
Stack<IEnumerator<IObservable<TSource>>> stack;
2426

2527
public IDisposable Run(IEnumerable<IObservable<TSource>> sources)
2628
{
27-
_isDisposed = false;
28-
_subscription = new SerialDisposable();
29-
_gate = new AsyncLock();
30-
_stack = new Stack<IEnumerator<IObservable<TSource>>>();
31-
_length = new Stack<int?>();
32-
33-
if (!TryGetEnumerator(sources, out var e))
29+
if (!TryGetEnumerator(sources, out var current))
3430
return Disposable.Empty;
3531

36-
_stack.Push(e);
37-
_length.Push(Helpers.GetLength(sources));
32+
stack = new Stack<IEnumerator<IObservable<TSource>>>();
33+
stack.Push(current);
3834

39-
var cancelable = SchedulerDefaults.TailRecursion.Schedule(self =>
40-
{
41-
_recurse = self;
42-
_gate.Wait(MoveNext);
43-
});
35+
Drain();
4436

45-
return StableCompositeDisposable.Create(_subscription, cancelable, Disposable.Create(() => _gate.Wait(Dispose)));
37+
return new RecursiveSinkDisposable(this);
4638
}
4739

48-
protected abstract IEnumerable<IObservable<TSource>> Extract(IObservable<TSource> source);
49-
50-
private void MoveNext()
40+
sealed class RecursiveSinkDisposable : IDisposable
5141
{
52-
var hasNext = false;
53-
var next = default(IObservable<TSource>);
42+
readonly TailRecursiveSink<TSource> parent;
5443

55-
do
44+
public RecursiveSinkDisposable(TailRecursiveSink<TSource> parent)
5645
{
57-
if (_stack.Count == 0)
58-
break;
46+
this.parent = parent;
47+
}
5948

60-
if (_isDisposed)
61-
return;
49+
public void Dispose()
50+
{
51+
parent.DisposeAll();
52+
}
53+
}
6254

63-
var e = _stack.Peek();
64-
var l = _length.Peek();
55+
void Drain()
56+
{
57+
if (Interlocked.Increment(ref trampoline) != 1)
58+
{
59+
return;
60+
}
6561

66-
var current = default(IObservable<TSource>);
67-
try
62+
for (; ; )
63+
{
64+
if (Volatile.Read(ref _isDisposed))
6865
{
69-
hasNext = e.MoveNext();
70-
if (hasNext)
66+
while (stack.Count != 0)
7167
{
72-
current = e.Current;
68+
var enumerator = stack.Pop();
69+
enumerator.Dispose();
70+
}
71+
if (Volatile.Read(ref currentSubscription) != BooleanDisposable.True)
72+
{
73+
Interlocked.Exchange(ref currentSubscription, BooleanDisposable.True)?.Dispose();
7374
}
74-
}
75-
catch (Exception ex)
76-
{
77-
e.Dispose();
78-
79-
//
80-
// Failure to enumerate the sequence cannot be handled, even by
81-
// operators like Catch, because it'd lead to another attempt at
82-
// enumerating to find the next observable sequence. Therefore,
83-
// we feed those errors directly to the observer.
84-
//
85-
_observer.OnError(ex);
86-
base.Dispose();
87-
return;
88-
}
89-
90-
if (!hasNext)
91-
{
92-
e.Dispose();
93-
94-
_stack.Pop();
95-
_length.Pop();
9675
}
9776
else
9877
{
99-
var r = l - 1;
100-
_length.Pop();
101-
_length.Push(r);
102-
103-
try
78+
if (stack.Count != 0)
10479
{
105-
next = Helpers.Unpack(current);
106-
}
107-
catch (Exception exception)
108-
{
109-
//
110-
// Errors from unpacking may produce side-effects that normally
111-
// would occur during a SubscribeSafe operation. Those would feed
112-
// back into the observer and be subject to the operator's error
113-
// handling behavior. For example, Catch would allow to handle
114-
// the error using a handler function.
115-
//
116-
if (!Fail(exception))
80+
var currentEnumerator = stack.Peek();
81+
82+
var currentObservable = default(IObservable<TSource>);
83+
var next = default(IObservable<TSource>);
84+
85+
try
86+
{
87+
if (currentEnumerator.MoveNext())
88+
{
89+
currentObservable = currentEnumerator.Current;
90+
}
91+
}
92+
catch (Exception ex)
11793
{
118-
e.Dispose();
94+
currentEnumerator.Dispose();
95+
_observer.OnError(ex);
96+
base.Dispose();
97+
Volatile.Write(ref _isDisposed, true);
98+
continue;
11999
}
120100

121-
return;
122-
}
101+
try
102+
{
103+
next = Helpers.Unpack(currentObservable);
123104

124-
//
125-
// Tail recursive case; drop the current frame.
126-
//
127-
if (r == 0)
128-
{
129-
e.Dispose();
105+
}
106+
catch (Exception ex)
107+
{
108+
next = null;
109+
if (!Fail(ex))
110+
{
111+
Volatile.Write(ref _isDisposed, true);
112+
}
113+
continue;
114+
}
130115

131-
_stack.Pop();
132-
_length.Pop();
116+
if (next != null)
117+
{
118+
var nextSeq = Extract(next);
119+
if (nextSeq != null)
120+
{
121+
if (TryGetEnumerator(nextSeq, out var nextEnumerator))
122+
{
123+
stack.Push(nextEnumerator);
124+
continue;
125+
}
126+
else
127+
{
128+
Volatile.Write(ref _isDisposed, true);
129+
continue;
130+
}
131+
}
132+
else
133+
{
134+
var sad = new SingleAssignmentDisposable();
135+
if (Interlocked.CompareExchange(ref currentSubscription, sad, null) == null)
136+
{
137+
sad.Disposable = next.SubscribeSafe(this);
138+
}
139+
else
140+
{
141+
continue;
142+
}
143+
}
144+
}
145+
else
146+
{
147+
stack.Pop();
148+
currentEnumerator.Dispose();
149+
continue;
150+
}
133151
}
134-
135-
//
136-
// Flattening of nested sequences. Prevents stack overflow in observers.
137-
//
138-
var nextSeq = Extract(next);
139-
if (nextSeq != null)
152+
else
140153
{
141-
if (!TryGetEnumerator(nextSeq, out var nextEnumerator))
142-
return;
143-
144-
_stack.Push(nextEnumerator);
145-
_length.Push(Helpers.GetLength(nextSeq));
146-
147-
hasNext = false;
154+
Volatile.Write(ref _isDisposed, true);
155+
Done();
148156
}
149157
}
150-
} while (!hasNext);
151158

152-
if (!hasNext)
153-
{
154-
Done();
155-
return;
159+
if (Interlocked.Decrement(ref trampoline) == 0)
160+
{
161+
break;
162+
}
156163
}
164+
}
157165

158-
var d = new SingleAssignmentDisposable();
159-
_subscription.Disposable = d;
160-
d.Disposable = next.SubscribeSafe(this);
166+
void DisposeAll()
167+
{
168+
Volatile.Write(ref _isDisposed, true);
169+
// the disposing of currentSubscription is deferred to drain due to some ObservableExTest.Iterate_Complete()
170+
// Interlocked.Exchange(ref currentSubscription, BooleanDisposable.True)?.Dispose();
171+
Drain();
161172
}
162173

163-
private new void Dispose()
174+
protected void Recurse()
164175
{
165-
while (_stack.Count > 0)
176+
var d = Volatile.Read(ref currentSubscription);
177+
if (d != BooleanDisposable.True)
166178
{
167-
var e = _stack.Pop();
168-
_length.Pop();
169-
170-
e.Dispose();
179+
d?.Dispose();
180+
if (Interlocked.CompareExchange(ref currentSubscription, null, d) == d)
181+
{
182+
Drain();
183+
}
171184
}
172-
173-
_isDisposed = true;
174185
}
175186

187+
protected abstract IEnumerable<IObservable<TSource>> Extract(IObservable<TSource> source);
188+
176189
private bool TryGetEnumerator(IEnumerable<IObservable<TSource>> sources, out IEnumerator<IObservable<TSource>> result)
177190
{
178191
try
@@ -182,12 +195,6 @@ private bool TryGetEnumerator(IEnumerable<IObservable<TSource>> sources, out IEn
182195
}
183196
catch (Exception exception)
184197
{
185-
//
186-
// Failure to enumerate the sequence cannot be handled, even by
187-
// operators like Catch, because it'd lead to another attempt at
188-
// enumerating to find the next observable sequence. Therefore,
189-
// we feed those errors directly to the observer.
190-
//
191198
_observer.OnError(exception);
192199
base.Dispose();
193200

‎Rx.NET/Source/src/System.Reactive/Linq/Observable/Catch.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ public override void OnNext(TSource value)
4545
public override void OnError(Exception error)
4646
{
4747
_lastException = error;
48-
_recurse();
48+
Recurse();
4949
}
5050

5151
public override void OnCompleted()

‎Rx.NET/Source/src/System.Reactive/Linq/Observable/OnErrorResumeNext.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,12 @@ public override void OnNext(TSource value)
4141

4242
public override void OnError(Exception error)
4343
{
44-
_recurse();
44+
Recurse();
4545
}
4646

4747
public override void OnCompleted()
4848
{
49-
_recurse();
49+
Recurse();
5050
}
5151

5252
protected override bool Fail(Exception error)

0 commit comments

Comments
 (0)