Cosmos: Modernize JSON serialization - Update pipeline (#38024)

Refactors DocumentSource to serialize to stream instead of DOM.
Leverages JsonValueReaderWriter for serialization.
For backwards compatibility:

* Adds custom JsonValueReaderWriter's where found to be needed.
* Use UnsafeRelaxedJsonEscaping (encoding special (html) characters)

Partly removes __jObject, but not fully so the binding and query pipelines still work. Will be fully removed in next PR.
Adds some todo's aswell, these will be picked up in following pr's if they have #34567 linked.

Obsoletes ContentResponseOnWriteEnabled
Simplifies CosmosTransactionalBatchTest and updates values for new (unindented) serialization

Part of: #34567

Author: JoasE

src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs

@@ -205,13 +205,14 @@ public virtual CosmosDbContextOptionsBuilder MaxRequestsPerTcpConnection(int req
     ///     This reduces networking and CPU load by not sending the resource back over the network and serializing it on the client.
     /// </summary>
     /// <remarks>
+    ///     The EntityFrameworkCore default is <see langword="false" /> since 11.0.
     ///     See <see href="https://aka.ms/efcore-docs-dbcontext-options">Using DbContextOptions</see>, and
     ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
     /// </remarks>
     /// <param name="enabled"><see langword="false" /> to have null resource</param>
+    [Obsolete("Enabling ContentResponseOnWrite currently has no benefit for EF Core.")]
     public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool enabled = true)
-        => WithOption(e => e.ContentResponseOnWriteEnabled(Check.NotNull(enabled)));
-
+        => WithOption(e => e.ContentResponseOnWriteEnabled(enabled));
 
     /// <summary>
     ///     Sets the <see cref="Cosmos.Infrastructure.SessionTokenManagementMode"/> to use.

src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs

@@ -74,6 +74,7 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom)
         _gatewayModeMaxConnectionLimit = copyFrom._gatewayModeMaxConnectionLimit;
         _maxTcpConnectionsPerEndpoint = copyFrom._maxTcpConnectionsPerEndpoint;
         _maxRequestsPerTcpConnection = copyFrom._maxRequestsPerTcpConnection;
+        _enableContentResponseOnWrite = copyFrom._enableContentResponseOnWrite;
         _httpClientFactory = copyFrom._httpClientFactory;
         _sessionTokenManagementMode = copyFrom._sessionTokenManagementMode;
         _enableBulkExecution = copyFrom._enableBulkExecution;

src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs

@@ -141,7 +141,7 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public virtual bool? EnableContentResponseOnWrite { get; }
+    public virtual bool? EnableContentResponseOnWrite { get; private set; }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -185,6 +185,7 @@ public virtual void Initialize(IDbContextOptions options)
             GatewayModeMaxConnectionLimit = cosmosOptions.GatewayModeMaxConnectionLimit;
             MaxTcpConnectionsPerEndpoint = cosmosOptions.MaxTcpConnectionsPerEndpoint;
             MaxRequestsPerTcpConnection = cosmosOptions.MaxRequestsPerTcpConnection;
+            EnableContentResponseOnWrite = cosmosOptions.EnableContentResponseOnWrite;
             HttpClientFactory = cosmosOptions.HttpClientFactory;
             EnableBulkExecution = cosmosOptions.EnableBulkExecution;
         }
@@ -216,6 +217,7 @@ public virtual void Validate(IDbContextOptions options)
                 || GatewayModeMaxConnectionLimit != cosmosOptions.GatewayModeMaxConnectionLimit
                 || MaxTcpConnectionsPerEndpoint != cosmosOptions.MaxTcpConnectionsPerEndpoint
                 || MaxRequestsPerTcpConnection != cosmosOptions.MaxRequestsPerTcpConnection
+                || EnableContentResponseOnWrite != cosmosOptions.EnableContentResponseOnWrite
                 || HttpClientFactory != cosmosOptions.HttpClientFactory
                 || EnableBulkExecution != cosmosOptions.EnableBulkExecution
             ))

src/EFCore.Cosmos/Query/Internal/CosmosSerializationUtilities.cs

@@ -17,7 +17,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
 /// <remarks>
 ///     Inspired by RelationalJsonUtilities.
 /// </remarks>
-public static class CosmosSerializationUtilities
+public static class CosmosSerializationUtilities // @TODO: Can this be removed? Use document source instead? #34567
 {
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to

src/EFCore.Cosmos/Storage/Internal/ByteArrayConverter.cs

@@ -56,7 +56,7 @@ public override object ReadJson(
     {
         if (reader.TokenType != JsonToken.StartArray)
         {
-            throw new Exception(reader.TokenType.ToString());
+            throw new InvalidOperationException(CoreStrings.JsonReaderInvalidTokenType(reader.TokenType));
         }
 
         var byteList = new List<byte>();

src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs

@@ -5,7 +5,7 @@
 using System.Collections.ObjectModel;
 using System.Net;
 using System.Runtime.CompilerServices;
-using System.Text;
+using System.Runtime.InteropServices;
 using Microsoft.Azure.Cosmos.Scripts;
 using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
@@ -47,8 +47,6 @@ public class CosmosClientWrapper : ICosmosClientWrapper
     private readonly string _databaseId;
     private readonly IExecutionStrategy _executionStrategy;
     private readonly IDiagnosticsLogger<DbLoggerCategory.Database.Command> _commandLogger;
-    private readonly IDiagnosticsLogger<DbLoggerCategory.Database> _databaseLogger;
-    private readonly bool? _enableContentResponseOnWrite;
 
     static CosmosClientWrapper()
     {
@@ -68,34 +66,14 @@ public CosmosClientWrapper(
         ISingletonCosmosClientWrapper singletonWrapper,
         IDbContextOptions dbContextOptions,
         IExecutionStrategy executionStrategy,
-        IDiagnosticsLogger<DbLoggerCategory.Database.Command> commandLogger,
-        IDiagnosticsLogger<DbLoggerCategory.Database> databaseLogger)
+        IDiagnosticsLogger<DbLoggerCategory.Database.Command> commandLogger)
     {
         var options = dbContextOptions.FindExtension<CosmosOptionsExtension>();
 
         _singletonWrapper = singletonWrapper;
         _databaseId = options!.DatabaseName;
         _executionStrategy = executionStrategy;
         _commandLogger = commandLogger;
-        _databaseLogger = databaseLogger;
-        _enableContentResponseOnWrite = options.EnableContentResponseOnWrite;
-    }
-
-    /// <summary>
-    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
-    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
-    ///     any release. You should only use it directly in your code with extreme caution and knowing that
-    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
-    /// </summary>
-    public static Stream Serialize(JToken document)
-    {
-        var stream = new MemoryStream();
-        using var writer = new StreamWriter(stream, new UTF8Encoding(), bufferSize: 1024, leaveOpen: true);
-
-        using var jsonWriter = new JsonTextWriter(writer);
-        CosmosClientWrapper.Serializer.Serialize(jsonWriter, document);
-        jsonWriter.Flush();
-        return stream;
     }
 
     /// <summary>
@@ -335,25 +313,25 @@ private static string GetPathFromRoot(IReadOnlyEntityType entityType)
     /// </summary>
     public virtual Task<bool> CreateItemAsync(
         string containerId,
-        JToken document,
+        string documentId,
+        ReadOnlyMemory<byte> document,
         IUpdateEntry updateEntry,
         ISessionTokenStorage sessionTokenStorage,
         CancellationToken cancellationToken = default)
-        => _executionStrategy.ExecuteAsync((containerId, document, updateEntry, sessionTokenStorage, this), CreateItemOnceAsync, null, cancellationToken);
+        => _executionStrategy.ExecuteAsync((containerId, documentId, document, updateEntry, sessionTokenStorage, this), CreateItemOnceAsync, null, cancellationToken);
 
     private static async Task<bool> CreateItemOnceAsync(
         DbContext _,
-        (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters,
+        (string ContainerId, string DocumentId, ReadOnlyMemory<byte> Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters,
         CancellationToken cancellationToken = default)
     {
-        using var stream = Serialize(parameters.Document);
-
         var containerId = parameters.ContainerId;
+        var documentId = parameters.DocumentId;
         var entry = parameters.Entry;
         var wrapper = parameters.Wrapper;
         var sessionTokenStorage = parameters.SessionTokenStorage;
         var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId);
-        var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite, sessionTokenStorage.GetSessionToken(containerId));
+        var itemRequestOptions = CreateItemRequestOptions(entry, sessionTokenStorage.GetSessionToken(containerId));
         var partitionKeyValue = ExtractPartitionKeyValue(entry);
         var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Create);
         var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Create);
@@ -370,6 +348,12 @@ private static async Task<bool> CreateItemOnceAsync(
             }
         }
 
+        if (!MemoryMarshal.TryGetArray(parameters.Document, out var segment) || segment.Array == null)
+        {
+            throw new UnreachableException("ReadOnlyMemory should have an underlying array.");
+        }
+
+        using var stream = new MemoryStream(segment.Array, segment.Offset, segment.Count);
         using var response = await container.CreateItemStreamAsync(
                 stream,
                 partitionKeyValue,
@@ -381,7 +365,7 @@ private static async Task<bool> CreateItemOnceAsync(
             response.Diagnostics.GetClientElapsedTime(),
             response.Headers.RequestCharge,
             response.Headers.ActivityId,
-            parameters.Document["id"]!.ToString(),
+            documentId,
             containerId,
             partitionKeyValue);
 
@@ -399,7 +383,7 @@ private static async Task<bool> CreateItemOnceAsync(
     public virtual Task<bool> ReplaceItemAsync(
         string collectionId,
         string documentId,
-        JObject document,
+        ReadOnlyMemory<byte> document,
         IUpdateEntry updateEntry,
         ISessionTokenStorage sessionTokenStorage,
         CancellationToken cancellationToken = default)
@@ -408,17 +392,15 @@ public virtual Task<bool> ReplaceItemAsync(
 
     private static async Task<bool> ReplaceItemOnceAsync(
         DbContext _,
-        (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters,
+        (string ContainerId, string ResourceId, ReadOnlyMemory<byte> Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters,
         CancellationToken cancellationToken = default)
     {
-        using var stream = Serialize(parameters.Document);
-
         var containerId = parameters.ContainerId;
         var entry = parameters.Entry;
         var wrapper = parameters.Wrapper;
         var sessionTokenStorage = parameters.SessionTokenStorage;
         var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId);
-        var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite, sessionTokenStorage.GetSessionToken(containerId));
+        var itemRequestOptions = CreateItemRequestOptions(entry, sessionTokenStorage.GetSessionToken(containerId));
         var partitionKeyValue = ExtractPartitionKeyValue(entry);
         var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Replace);
         var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Replace);
@@ -435,6 +417,12 @@ private static async Task<bool> ReplaceItemOnceAsync(
             }
         }
 
+        if (!MemoryMarshal.TryGetArray(parameters.Document, out var segment) || segment.Array == null)
+        {
+            throw new UnreachableException("ReadOnlyMemory should have an underlying array.");
+        }
+
+        using var stream = new MemoryStream(segment.Array, segment.Offset, segment.Count);
         using var response = await container.ReplaceItemStreamAsync(
                 stream,
                 parameters.ResourceId,
@@ -481,7 +469,7 @@ private static async Task<bool> DeleteItemOnceAsync(
         var sessionTokenStorage = parameters.SessionTokenStorage;
         var items = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId);
 
-        var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite, sessionTokenStorage.GetSessionToken(containerId));
+        var itemRequestOptions = CreateItemRequestOptions(entry, sessionTokenStorage.GetSessionToken(containerId));
         var partitionKeyValue = ExtractPartitionKeyValue(entry);
         var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Delete);
         var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Delete);
@@ -539,7 +527,7 @@ public virtual ICosmosTransactionalBatchWrapper CreateTransactionalBatch(string
 
         var batch = container.CreateTransactionalBatch(partitionKeyValue);
 
-        return new CosmosTransactionalBatchWrapper(batch, containerId, partitionKeyValue, checkSize, _enableContentResponseOnWrite);
+        return new CosmosTransactionalBatchWrapper(batch, containerId, partitionKeyValue, checkSize);
     }
 
     /// <summary>
@@ -578,9 +566,9 @@ private static async Task<CosmosTransactionalBatchResult> ExecuteTransactionalBa
         return ProcessBatchResponse(batch.CollectionId, response, batch.Entries, sessionTokenStorage);
     }
 
-    private static ItemRequestOptions CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite, string? sessionToken)
+    private static ItemRequestOptions CreateItemRequestOptions(IUpdateEntry entry, string? sessionToken)
     {
-        var helper = RequestOptionsHelper.Create(entry, enableContentResponseOnWrite);
+        var helper = RequestOptionsHelper.Create(entry);
 
         var itemRequestOptions = new ItemRequestOptions
         {
@@ -590,7 +578,6 @@ private static ItemRequestOptions CreateItemRequestOptions(IUpdateEntry entry, b
         if (helper != null)
         {
             itemRequestOptions.IfMatchEtag = helper.IfMatchEtag;
-            itemRequestOptions.EnableContentResponseOnWrite = helper.EnableContentResponseOnWrite;
         }
 
         return itemRequestOptions;
@@ -681,31 +668,23 @@ private static CosmosTransactionalBatchResult ProcessBatchResponse(string contai
             var entry = entries[i];
             var item = response[i];
 
-            ProcessWriteResponse(entry.Entry, (string)item.ETag, (Stream)item.ResourceStream);
+            ProcessWriteResponse(entry.Entry, item.ETag, item.ResourceStream);
         }
 
         return CosmosTransactionalBatchResult.Success;
     }
 
     private static void ProcessWriteResponse(IUpdateEntry entry, string eTag, Stream? content)
     {
-        var etagProperty = entry.EntityType.GetETagProperty();
-        if (etagProperty != null && entry.EntityState != EntityState.Deleted)
+        if (entry.EntityState == EntityState.Deleted)
         {
-            entry.SetStoreGeneratedValue(etagProperty, eTag);
+            return;
         }
 
-        var jObjectProperty = entry.EntityType.FindProperty(CosmosPartitionKeyInPrimaryKeyConvention.JObjectPropertyName);
-        if (jObjectProperty is { ValueGenerated: ValueGenerated.OnAddOrUpdate }
-            && content != null)
+        var etagProperty = entry.EntityType.GetETagProperty();
+        if (etagProperty != null)
         {
-            using var responseStream = content;
-            using var reader = new StreamReader(responseStream);
-            using var jsonReader = new JsonTextReader(reader);
-
-            var createdDocument = Serializer.Deserialize<JObject>(jsonReader);
-
-            entry.SetStoreGeneratedValue(jObjectProperty, createdDocument);
+            entry.SetStoreGeneratedValue(etagProperty, eTag);
         }
     }