AsyncEnumerable.Exceptions.cs 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381
  1. // Copyright (c) Microsoft Open Technologies, Inc. All rights reserved. See License.txt in the project root for license information.
  2. using System;
  3. using System.Collections.Generic;
  4. using System.Linq;
  5. using System.Threading;
  6. using System.Threading.Tasks;
  7. namespace System.Linq
  8. {
  9. public static partial class AsyncEnumerable
  10. {
  11. public static IAsyncEnumerable<TSource> Catch<TSource, TException>(this IAsyncEnumerable<TSource> source, Func<TException, IAsyncEnumerable<TSource>> handler)
  12. where TException : Exception
  13. {
  14. if (source == null)
  15. throw new ArgumentNullException("source");
  16. if (handler == null)
  17. throw new ArgumentNullException("handler");
  18. return Create(() =>
  19. {
  20. var e = source.GetEnumerator();
  21. var cts = new CancellationTokenDisposable();
  22. var a = new AssignableDisposable { Disposable = e };
  23. var d = new CompositeDisposable(cts, a);
  24. var done = false;
  25. var f = default(Action<TaskCompletionSource<bool>, CancellationToken>);
  26. f = (tcs, ct) =>
  27. {
  28. if (!done)
  29. {
  30. e.MoveNext(ct).ContinueWith(t =>
  31. {
  32. t.Handle(tcs,
  33. res =>
  34. {
  35. tcs.TrySetResult(res);
  36. },
  37. ex =>
  38. {
  39. var err = default(IAsyncEnumerator<TSource>);
  40. try
  41. {
  42. ex.Flatten().Handle(ex_ =>
  43. {
  44. var exx = ex_ as TException;
  45. if (exx != null)
  46. {
  47. err = handler(exx).GetEnumerator();
  48. return true;
  49. }
  50. return false;
  51. });
  52. }
  53. catch (Exception ex2)
  54. {
  55. tcs.TrySetException(ex2);
  56. return;
  57. }
  58. if (err != null)
  59. {
  60. e = err;
  61. a.Disposable = e;
  62. done = true;
  63. f(tcs, ct);
  64. }
  65. }
  66. );
  67. });
  68. }
  69. else
  70. {
  71. e.MoveNext(ct).ContinueWith(t =>
  72. {
  73. t.Handle(tcs, res =>
  74. {
  75. tcs.TrySetResult(res);
  76. });
  77. });
  78. }
  79. };
  80. return Create(
  81. (ct, tcs) =>
  82. {
  83. f(tcs, cts.Token);
  84. return tcs.Task.UsingEnumerator(a);
  85. },
  86. () => e.Current,
  87. d.Dispose
  88. );
  89. });
  90. }
  91. public static IAsyncEnumerable<TSource> Catch<TSource>(this IEnumerable<IAsyncEnumerable<TSource>> sources)
  92. {
  93. if (sources == null)
  94. throw new ArgumentNullException("sources");
  95. return sources.Catch_();
  96. }
  97. public static IAsyncEnumerable<TSource> Catch<TSource>(params IAsyncEnumerable<TSource>[] sources)
  98. {
  99. if (sources == null)
  100. throw new ArgumentNullException("sources");
  101. return sources.Catch_();
  102. }
  103. public static IAsyncEnumerable<TSource> Catch<TSource>(this IAsyncEnumerable<TSource> first, IAsyncEnumerable<TSource> second)
  104. {
  105. if (first == null)
  106. throw new ArgumentNullException("first");
  107. if (second == null)
  108. throw new ArgumentNullException("second");
  109. return new[] { first, second }.Catch_();
  110. }
  111. private static IAsyncEnumerable<TSource> Catch_<TSource>(this IEnumerable<IAsyncEnumerable<TSource>> sources)
  112. {
  113. return Create(() =>
  114. {
  115. var se = sources.GetEnumerator();
  116. var e = default(IAsyncEnumerator<TSource>);
  117. var cts = new CancellationTokenDisposable();
  118. var a = new AssignableDisposable();
  119. var d = new CompositeDisposable(cts, se, a);
  120. var error = default(Exception);
  121. var f = default(Action<TaskCompletionSource<bool>, CancellationToken>);
  122. f = (tcs, ct) =>
  123. {
  124. if (e == null)
  125. {
  126. var b = false;
  127. try
  128. {
  129. b = se.MoveNext();
  130. if (b)
  131. e = se.Current.GetEnumerator();
  132. }
  133. catch (Exception ex)
  134. {
  135. tcs.TrySetException(ex);
  136. return;
  137. }
  138. if (!b)
  139. {
  140. if (error != null)
  141. {
  142. tcs.TrySetException(error);
  143. return;
  144. }
  145. tcs.TrySetResult(false);
  146. return;
  147. }
  148. error = null;
  149. a.Disposable = e;
  150. }
  151. e.MoveNext(ct).ContinueWith(t =>
  152. {
  153. t.Handle(tcs,
  154. res =>
  155. {
  156. tcs.TrySetResult(res);
  157. },
  158. ex =>
  159. {
  160. e.Dispose();
  161. e = null;
  162. error = ex;
  163. f(tcs, ct);
  164. }
  165. );
  166. });
  167. };
  168. return Create(
  169. (ct, tcs) =>
  170. {
  171. f(tcs, cts.Token);
  172. return tcs.Task.UsingEnumerator(a);
  173. },
  174. () => e.Current,
  175. d.Dispose
  176. );
  177. });
  178. }
  179. public static IAsyncEnumerable<TSource> Finally<TSource>(this IAsyncEnumerable<TSource> source, Action finallyAction)
  180. {
  181. if (source == null)
  182. throw new ArgumentNullException("source");
  183. if (finallyAction == null)
  184. throw new ArgumentNullException("finallyAction");
  185. return Create(() =>
  186. {
  187. var e = source.GetEnumerator();
  188. var cts = new CancellationTokenDisposable();
  189. var r = new Disposable(finallyAction);
  190. var d = new CompositeDisposable(cts, e, r);
  191. var f = default(Action<TaskCompletionSource<bool>, CancellationToken>);
  192. f = (tcs, ct) =>
  193. {
  194. e.MoveNext(ct).ContinueWith(t =>
  195. {
  196. t.Handle(tcs, res =>
  197. {
  198. tcs.TrySetResult(res);
  199. });
  200. });
  201. };
  202. return Create(
  203. (ct, tcs) =>
  204. {
  205. f(tcs, cts.Token);
  206. return tcs.Task.UsingEnumeratorSync(r);
  207. },
  208. () => e.Current,
  209. d.Dispose
  210. );
  211. });
  212. }
  213. public static IAsyncEnumerable<TSource> OnErrorResumeNext<TSource>(this IAsyncEnumerable<TSource> first, IAsyncEnumerable<TSource> second)
  214. {
  215. if (first == null)
  216. throw new ArgumentNullException("first");
  217. if (second == null)
  218. throw new ArgumentNullException("second");
  219. return OnErrorResumeNext_(new[] { first, second });
  220. }
  221. public static IAsyncEnumerable<TSource> OnErrorResumeNext<TSource>(params IAsyncEnumerable<TSource>[] sources)
  222. {
  223. if (sources == null)
  224. throw new ArgumentNullException("sources");
  225. return OnErrorResumeNext_(sources);
  226. }
  227. public static IAsyncEnumerable<TSource> OnErrorResumeNext<TSource>(this IEnumerable<IAsyncEnumerable<TSource>> sources)
  228. {
  229. if (sources == null)
  230. throw new ArgumentNullException("sources");
  231. return OnErrorResumeNext_(sources);
  232. }
  233. private static IAsyncEnumerable<TSource> OnErrorResumeNext_<TSource>(IEnumerable<IAsyncEnumerable<TSource>> sources)
  234. {
  235. return Create(() =>
  236. {
  237. var se = sources.GetEnumerator();
  238. var e = default(IAsyncEnumerator<TSource>);
  239. var cts = new CancellationTokenDisposable();
  240. var a = new AssignableDisposable();
  241. var d = new CompositeDisposable(cts, se, a);
  242. var f = default(Action<TaskCompletionSource<bool>, CancellationToken>);
  243. f = (tcs, ct) =>
  244. {
  245. if (e == null)
  246. {
  247. var b = false;
  248. try
  249. {
  250. b = se.MoveNext();
  251. if (b)
  252. e = se.Current.GetEnumerator();
  253. }
  254. catch (Exception ex)
  255. {
  256. tcs.TrySetException(ex);
  257. return;
  258. }
  259. if (!b)
  260. {
  261. tcs.TrySetResult(false);
  262. return;
  263. }
  264. a.Disposable = e;
  265. }
  266. e.MoveNext(ct).ContinueWith(t =>
  267. {
  268. t.Handle(tcs,
  269. res =>
  270. {
  271. if (res)
  272. {
  273. tcs.TrySetResult(true);
  274. }
  275. else
  276. {
  277. e.Dispose();
  278. e = null;
  279. f(tcs, ct);
  280. }
  281. },
  282. ex =>
  283. {
  284. e.Dispose();
  285. e = null;
  286. f(tcs, ct);
  287. }
  288. );
  289. });
  290. };
  291. return Create(
  292. (ct, tcs) =>
  293. {
  294. f(tcs, cts.Token);
  295. return tcs.Task.UsingEnumerator(a);
  296. },
  297. () => e.Current,
  298. d.Dispose
  299. );
  300. });
  301. }
  302. public static IAsyncEnumerable<TSource> Retry<TSource>(this IAsyncEnumerable<TSource> source)
  303. {
  304. if (source == null)
  305. throw new ArgumentNullException("source");
  306. return new[] { source }.Repeat().Catch();
  307. }
  308. public static IAsyncEnumerable<TSource> Retry<TSource>(this IAsyncEnumerable<TSource> source, int retryCount)
  309. {
  310. if (source == null)
  311. throw new ArgumentNullException("source");
  312. if (retryCount < 0)
  313. throw new ArgumentOutOfRangeException("retryCount");
  314. return new[] { source }.Repeat(retryCount).Catch();
  315. }
  316. private static IEnumerable<TSource> Repeat<TSource>(this IEnumerable<TSource> source)
  317. {
  318. while (true)
  319. foreach (var item in source)
  320. yield return item;
  321. }
  322. private static IEnumerable<TSource> Repeat<TSource>(this IEnumerable<TSource> source, int count)
  323. {
  324. for (var i = 0; i < count; i++)
  325. foreach (var item in source)
  326. yield return item;
  327. }
  328. }
  329. }