Fix In extension in async queries

Author: alex-kulakov

Orm/Xtensive.Orm/Orm/Providers/SqlIncludeProvider.cs

@@ -14,6 +14,8 @@
 using Tuple = Xtensive.Tuples.Tuple;
 using Xtensive.Orm.Rse;
 using Xtensive.Orm.Rse.Providers;
+using System.Threading.Tasks;
+using System.Threading;
 
 namespace Xtensive.Orm.Providers
 {
@@ -51,6 +53,29 @@ protected override void OnBeforeEnumerate(Rse.Providers.EnumerationContext conte
       }
     }
 
+    protected override async Task OnBeforeEnumerateAsync(Rse.Providers.EnumerationContext context, CancellationToken token)
+    {
+      await base.OnBeforeEnumerateAsync(context, token).ConfigureAwait(false);
+
+      switch (Origin.Algorithm) {
+        case IncludeAlgorithm.Auto:
+          var filterData = filterDataSource.Invoke().ToList();
+          if (filterData.Count > WellKnown.MaxNumberOfConditions)
+            await LockAndStoreAsync(context, filterData, token).ConfigureAwait(false);
+          else
+            context.SetValue(filterDataSource, RowFilterDataName, filterData);
+          break;
+        case IncludeAlgorithm.ComplexCondition:
+          // nothing
+          break;
+        case IncludeAlgorithm.TemporaryTable:
+          await LockAndStoreAsync(context, filterDataSource.Invoke(), token);
+          break;
+        default:
+          throw new ArgumentOutOfRangeException("Origin.Algorithm");
+      }
+    }
+
     /// <inheritdoc/>
     protected override void OnAfterEnumerate(Rse.Providers.EnumerationContext context)
     {

Orm/Xtensive.Orm/Orm/Rse/RecordSet.cs

@@ -111,8 +111,15 @@ private async Task<IEnumerator<Tuple>> GetGreedyEnumeratorAsync(CancellationToke
     /// </summary>
     private async Task<IEnumerator<Tuple>> GetBatchedEnumeratorAsync(CancellationToken token)
     {
-      EnumerationScope currentScope = null;
-      var enumerator = await Source.GetEnumeratorAsync(Context, token).ConfigureAwait(false);
+      var currentScope = Context.Activate();
+      IEnumerator<Tuple> enumerator = null;
+      try {
+        enumerator = await Source.GetEnumeratorAsync(Context, token).ConfigureAwait(false);
+      }
+      finally {
+        currentScope.DisposeSafely();
+      }
+        
       var batched = enumerator.ToEnumerable().Batch(2);
 
       var cs = Context.BeginEnumeration();
@@ -124,10 +131,11 @@ private async Task<IEnumerator<Tuple>> GetBatchedEnumeratorAsync(CancellationTok
           completableScope.Complete();
           completableScope.Dispose();
         };
-      return GetTwoLevelEnumerator(batched, afterEnumerationAction, cs);
+        return GetTwoLevelEnumerator(batched, afterEnumerationAction, cs);
     }
 
-    private static IEnumerator<Tuple> GetTwoLevelEnumerator(IEnumerable<IEnumerable<Tuple>> enumerable, Action<object> afterEnumerationAction, object parameterForAction)
+    private static IEnumerator<Tuple> GetTwoLevelEnumerator(IEnumerable<IEnumerable<Tuple>> enumerable,
+      Action<object> afterEnumerationAction, ICompletableScope parameterForAction)
     {
       try {
         foreach (var batch in enumerable)