Fix memory leak in ModelContainer - #1101

src/Microsoft.AspNet.OData.Shared/EnableQueryAttribute.cs

@@ -7,6 +7,7 @@
 using System.Diagnostics.Contracts;
 using System.Linq;
 using System.Net;
+using System.Net.Http;
 using Microsoft.AspNet.OData.Common;
 using Microsoft.AspNet.OData.Formatter;
 using Microsoft.AspNet.OData.Interfaces;
@@ -374,7 +375,7 @@ private object OnActionExecuted(
         {
             if (!_querySettings.PageSize.HasValue && responseValue != null)
             {
-                GetModelBoundPageSize(responseValue, singleResultCollection, actionDescriptor, modelFunction, request.Context.Path, createErrorAction);
+                GetModelBoundPageSize(responseValue, singleResultCollection, actionDescriptor, modelFunction, request, createErrorAction);
             }
 
             // Apply the query if there are any query options, if there is a page size set, in the case of
@@ -486,30 +487,34 @@ public virtual object ApplyQuery(object entity, ODataQueryOptions queryOptions)
         }
 
         /// <summary>
-        /// Get the ODaya query context.
+        /// Get the OData query context.
         /// </summary>
         /// <param name="responseValue">The response value.</param>
         /// <param name="singleResultCollection">The content as SingleResult.Queryable.</param>
         /// <param name="actionDescriptor">The action context, i.e. action and controller name.</param>
         /// <param name="modelFunction">A function to get the model.</param>
-        /// <param name="path">The OData path.</param>
+        /// <param name="request">The request message</param>
         /// <returns></returns>
         private static ODataQueryContext GetODataQueryContext(
             object responseValue,
             IQueryable singleResultCollection,
             IWebApiActionDescriptor actionDescriptor,
             Func<Type, IEdmModel> modelFunction,
-            ODataPath path)
+            IWebApiRequestMessage request)
         {
             Type elementClrType = GetElementType(responseValue, singleResultCollection, actionDescriptor);
+            ODataPath path = request.Context.Path;
 
             IEdmModel model = modelFunction(elementClrType);
             if (model == null)
             {
                 throw Error.InvalidOperation(SRResources.QueryGetModelMustNotReturnNull);
             }
 
-            return new ODataQueryContext(model, elementClrType, path);
+            ODataQueryContext queryContext = new ODataQueryContext(model, elementClrType, path);
+            RegisterContext(request, queryContext);
+
+            return queryContext;
         }
 
         /// <summary>
@@ -519,21 +524,21 @@ private static ODataQueryContext GetODataQueryContext(
         /// <param name="singleResultCollection">The content as SingleResult.Queryable.</param>
         /// <param name="actionDescriptor">The action context, i.e. action and controller name.</param>
         /// <param name="modelFunction">A function to get the model.</param>
-        /// <param name="path">The OData path.</param>
+        /// <param name="request">The request message.</param>
         /// <param name="createErrorAction">A function used to generate error response.</param>
         private void GetModelBoundPageSize(
             object responseValue,
             IQueryable singleResultCollection,
             IWebApiActionDescriptor actionDescriptor,
             Func<Type, IEdmModel> modelFunction,
-            ODataPath path,
+            IWebApiRequestMessage request,
             Action<HttpStatusCode, string, Exception> createErrorAction)
         {
             ODataQueryContext queryContext = null;
 
             try
             {
-                queryContext = GetODataQueryContext(responseValue, singleResultCollection, actionDescriptor, modelFunction, path);
+                queryContext = GetODataQueryContext(responseValue, singleResultCollection, actionDescriptor, modelFunction, request);
             }
             catch (InvalidOperationException e)
             {
@@ -571,7 +576,8 @@ private object ExecuteQuery(
             IWebApiRequestMessage request,
             Func<ODataQueryContext, ODataQueryOptions> createQueryOptionFunction)
         {
-            ODataQueryContext queryContext = GetODataQueryContext(responseValue, singleResultCollection, actionDescriptor, modelFunction, request.Context.Path);
+            ODataQueryContext queryContext = GetODataQueryContext(responseValue, singleResultCollection, actionDescriptor, modelFunction, request);
+            RegisterContext(request, queryContext);
 
             // Create and validate the query options.
             ODataQueryOptions queryOptions = createQueryOptionFunction(queryContext);

src/Microsoft.AspNet.OData.Shared/ODataQueryContext.cs

@@ -8,6 +8,7 @@
 using Microsoft.AspNet.OData.Common;
 using Microsoft.AspNet.OData.Formatter;
 using Microsoft.AspNet.OData.Query;
+using Microsoft.AspNet.OData.Query.Expressions;
 using Microsoft.AspNet.OData.Routing;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.OData.Edm;
@@ -20,6 +21,7 @@ namespace Microsoft.AspNet.OData
     public class ODataQueryContext
     {
         private DefaultQuerySettings _defaultQuerySettings;
+        private string modelId = null;
 
         /// <summary>
         /// Constructs an instance of <see cref="ODataQueryContext"/> with <see cref="IEdmModel" />, element CLR type,
@@ -72,18 +74,30 @@ public ODataQueryContext(IEdmModel model, IEdmType elementType, ODataPath path)
             {
                 throw Error.ArgumentNull("model");
             }
+
             if (elementType == null)
             {
                 throw Error.ArgumentNull("elementType");
             }
 
-            Model = model;
             ElementType = elementType;
+            Model = model;
             Path = path;
             NavigationSource = GetNavigationSource(Model, ElementType, path);
             GetPathContext();
         }
 
+        /// <summary>
+        /// Destructor called to clean up reference in ModelContainer
+        /// </summary>
+        ~ODataQueryContext()
+        {
+            if (this.modelId != null)
+            {
+                ModelContainer.RemoveModelId(ModelId);
+            }
+        }
+
         internal ODataQueryContext(IEdmModel model, Type elementClrType)
             : this(model, elementClrType, path: null)
         {
@@ -152,6 +166,19 @@ public DefaultQuerySettings DefaultQuerySettings
 
         internal string TargetName { get; private set; }
 
+        internal string ModelId
+        {
+            get
+            {
+                if (this.modelId == null)
+                {
+                    this.modelId = ModelContainer.GetModelID(this.Model);
+                }
+
+                return this.modelId;
+            }
+        }
+
         private static IEdmNavigationSource GetNavigationSource(IEdmModel model, IEdmType elementType, ODataPath odataPath)
         {
             Contract.Assert(model != null);

src/Microsoft.AspNet.OData.Shared/Query/Expressions/ModelContainer.cs

@@ -14,19 +14,24 @@ namespace Microsoft.AspNet.OData.Query.Expressions
     /// </summary>
     internal static class ModelContainer
     {
-        private static ConcurrentDictionary<IEdmModel, string> _map = new ConcurrentDictionary<IEdmModel, string>();
-        private static ConcurrentDictionary<string, IEdmModel> _reverseMap = new ConcurrentDictionary<string, IEdmModel>();
+        private static ConcurrentDictionary<string, IEdmModel> _map = new ConcurrentDictionary<string, IEdmModel>();
 
         public static string GetModelID(IEdmModel model)
         {
-            string index = _map.GetOrAdd(model, m => Guid.NewGuid().ToString());
-            _reverseMap.TryAdd(index, model);
+            string index = Guid.NewGuid().ToString();
+            _map.TryAdd(index, model);
             return index;
         }
 
         public static IEdmModel GetModel(string id)
         {
-            return _reverseMap[id];
+            return _map[id];
+        }
+
+        public static void RemoveModelId(string id)
+        {
+            IEdmModel model;
+            _map.TryRemove(id, out model);
         }
     }
 }

src/Microsoft.AspNet.OData.Shared/Query/Expressions/SelectExpandBinder.cs

@@ -42,7 +42,7 @@ public SelectExpandBinder(ODataQuerySettings settings, SelectExpandQueryOption s
             _selectExpandQuery = selectExpandQuery;
             _context = selectExpandQuery.Context;
             _model = _context.Model;
-            _modelID = ModelContainer.GetModelID(_model);
+            _modelID = _context.ModelId;
             _settings = settings;
         }
 

src/Microsoft.AspNet.OData/EnableQueryAttribute.cs

@@ -16,6 +16,7 @@
 using Microsoft.AspNet.OData.Common;
 using Microsoft.AspNet.OData.Extensions;
 using Microsoft.AspNet.OData.Formatter;
+using Microsoft.AspNet.OData.Interfaces;
 using Microsoft.AspNet.OData.Query;
 using Microsoft.OData.Edm;
 
@@ -193,5 +194,21 @@ public virtual IEdmModel GetModel(Type elementClrType, HttpRequestMessage reques
             Contract.Assert(model != null);
             return model;
         }
+
+        /// <summary>
+        /// Register the QueryContext with the Request so that it doesn't get garbage collected early
+        /// </summary>
+        /// <param name="request">The request message to register with</param>
+        /// <param name="queryContext">The ODataQueryContext to be registered</param>
+        private static void RegisterContext(IWebApiRequestMessage request, ODataQueryContext queryContext)
+        {
+            HttpRequestScope httpRequestScope = request.RequestContainer.GetService(typeof(HttpRequestScope)) as HttpRequestScope;
+            HttpRequestMessage httpRequest = httpRequestScope == null ? null : httpRequestScope.HttpRequest;
+
+            if (httpRequest != null)
+            {
+                httpRequest.ODataProperties().QueryContexts.Add(queryContext);
+            }
+        }
     }
 }

src/Microsoft.AspNet.OData/Extensions/HttpRequestMessageProperties.cs

@@ -38,6 +38,7 @@ public class HttpRequestMessageProperties
         private const string TotalCountFuncKey = "Microsoft.AspNet.OData.TotalCountFunc";
 
         private HttpRequestMessage _request;
+        private IList<ODataQueryContext> _queryContexts = new List<ODataQueryContext>();
 
         internal HttpRequestMessageProperties(HttpRequestMessage request)
         {
@@ -229,6 +230,14 @@ private set
             }
         }
 
+        internal IList<ODataQueryContext> QueryContexts
+        {
+            get
+            {
+                return _queryContexts;
+            }
+        }
+
         internal ODataVersion? ODataServiceVersion
         {
             get

src/Microsoft.AspNet.OData/GlobalSuppressions.cs

@@ -28,3 +28,4 @@
 [assembly: SuppressMessage("Microsoft.Maintainability", "CA1506:AvoidExcessiveClassCoupling", Scope = "member", Target = "Microsoft.AspNet.OData.Builder.ODataConventionModelBuilder.#.cctor()")]
 [assembly: SuppressMessage("Microsoft.Maintainability", "CA1506:AvoidExcessiveClassCoupling", Scope = "type", Target = "Microsoft.AspNet.OData.Query.Expressions.PropertyContainer", Justification = "Using generated classes to simulate b-tree.")]
 [assembly: SuppressMessage("Microsoft.Maintainability", "CA1506:AvoidExcessiveClassCoupling", Scope = "member", Target = "Microsoft.AspNet.OData.Query.Expressions.PropertyContainer.#.cctor()", Justification = "Using generated classes to simulate b-tree.")]
+[assembly: SuppressMessage("Microsoft.Maintainability", "CA1506:AvoidExcessiveClassCoupling", Scope = "type", Target = "Microsoft.AspNet.OData.EnableQueryAttribute")]

src/Microsoft.AspNetCore.OData/EnableQueryAttribute.cs

@@ -11,6 +11,7 @@
 using Microsoft.AspNet.OData.Common;
 using Microsoft.AspNet.OData.Extensions;
 using Microsoft.AspNet.OData.Formatter;
+using Microsoft.AspNet.OData.Interfaces;
 using Microsoft.AspNet.OData.Query;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
@@ -247,5 +248,41 @@ public virtual IEdmModel GetModel(
             Contract.Assert(model != null);
             return model;
         }
+
+        /// <summary>
+        /// Register the QueryContext with the Request so that it doesn't get garbage collected early
+        /// </summary>
+        /// <param name="request">The HTTPRequestMessage to register with</param>
+        /// <param name="queryContext">The ODataQueryContext to be registered</param>
+        private static void RegisterContext(IWebApiRequestMessage request, ODataQueryContext queryContext)
+        {
+            HttpRequestScope httpRequestScope = request.RequestContainer.GetService(typeof(HttpRequestScope)) as HttpRequestScope;
+            HttpRequest httpRequest = httpRequestScope == null ? null : httpRequestScope.HttpRequest;
+
+
+            if (httpRequest != null)
+            {
+                HttpContext context = httpRequest.HttpContext;
+                ODataQueryContextCacheFeature queryCache = context.Features.Get<ODataQueryContextCacheFeature>();
+                if (queryCache == null)
+                {
+                    lock (context.Features)
+                    {
+                        queryCache = context.Features.Get<ODataQueryContextCacheFeature>();
+                        if (queryCache == null)
+                        {
+                            queryCache = new ODataQueryContextCacheFeature();
+                            context.Features.Set<ODataQueryContextCacheFeature>(queryCache);
+                        }
+                    }
+                }
+
+                queryCache.Add(queryContext);
+            }
+        }
+
+        private class ODataQueryContextCacheFeature : List<ODataQueryContext>
+        {
+        }
     }
 }

test/UnitTest/Microsoft.AspNet.OData.Test.Shared/Query/EnableQueryAttributeTest.cs

@@ -393,6 +393,7 @@ public async Task NonGenericEnumerableReturnType_ReturnsBadRequest()
             EnableQueryAttribute attribute = new EnableQueryAttribute();
             HttpRequestMessage request = new HttpRequestMessage(HttpMethod.Get, "http://localhost/Customer/?$skip=1");
             var config = RoutingConfigurationFactory.Create();
+            config.EnableDependencyInjection();
             config.IncludeErrorDetailPolicy = IncludeErrorDetailPolicy.Always;
             request.SetConfiguration(config);
             HttpControllerContext controllerContext = new HttpControllerContext(config, new HttpRouteData(new HttpRoute()), request);

test/UnitTest/Microsoft.AspNet.OData.Test.Shared/Query/Expressions/ModelContainerTest.cs

@@ -18,14 +18,6 @@ public void GetModelID_Returns_DifferentIDForDifferentModels()
             Assert.NotEqual(ModelContainer.GetModelID(model1), ModelContainer.GetModelID(model2));
         }
 
-        [Fact]
-        public void GetModelID_Returns_SameIDForSameModel()
-        {
-            EdmModel model = new EdmModel();
-
-            Assert.Equal(ModelContainer.GetModelID(model), ModelContainer.GetModelID(model));
-        }
-
         [Fact]
         public void GetModelID_AndThen_GetModel_ReturnsOriginalModel()
         {

test/UnitTest/Microsoft.AspNet.OData.Test/PublicApi/Microsoft.AspNet.OData.PublicApi.bsl

@@ -483,6 +483,8 @@ public class Microsoft.AspNet.OData.ODataQueryContext {
 	Microsoft.OData.Edm.IEdmNavigationSource NavigationSource  { public get; }
 	ODataPath Path  { public get; }
 	System.IServiceProvider RequestContainer  { public get; }
+
+	protected virtual void Finalize ()
 }
 
 public class Microsoft.AspNet.OData.ODataSwaggerConverter {

test/UnitTest/Microsoft.AspNetCore.OData.Test/PublicApi/Microsoft.AspNetCore.OData.PublicApi.bsl

@@ -530,6 +530,8 @@ public class Microsoft.AspNet.OData.ODataQueryContext {
 	Microsoft.OData.Edm.IEdmNavigationSource NavigationSource  { public get; }
 	ODataPath Path  { public get; }
 	System.IServiceProvider RequestContainer  { public get; }
+
+	protected virtual void Finalize ()
 }
 
 public class Microsoft.AspNet.OData.ODataSwaggerConverter {