From 89145dedee6e51c533238fff5f189237998e5a87 Mon Sep 17 00:00:00 2001
From: "Lauesen.Uffe UFE" <UFE@kmd.dk>
Date: Fri, 13 Jan 2023 11:43:13 +0100
Subject: [PATCH] feat: allow @odata.count to be written after enumeration
 #2594

---
 .../JsonLight/ODataJsonLightWriter.cs         | 60 ++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

diff --git a/src/Microsoft.OData.Core/JsonLight/ODataJsonLightWriter.cs b/src/Microsoft.OData.Core/JsonLight/ODataJsonLightWriter.cs
index eefdd245f3..b5788b439c 100644
--- a/src/Microsoft.OData.Core/JsonLight/ODataJsonLightWriter.cs
+++ b/src/Microsoft.OData.Core/JsonLight/ODataJsonLightWriter.cs
@@ -1668,6 +1668,10 @@ await this.instanceAnnotationWriter.WriteInstanceAnnotationsAsync(
 
                 if (this.writingResponse)
                 {
+                    // Write the count if it's available.
+                    await this.WriteResourceSetCountAsync(resourceSet.Count, null)
+                        .ConfigureAwait(false);
+
                     // Write the next link if it's available.
                     await this.WriteResourceSetNextLinkAsync(resourceSet.NextPageLink, /*propertyName*/null)
                         .ConfigureAwait(false);
@@ -1704,6 +1708,9 @@ await this.asynchronousJsonWriter.EndArrayScopeAsync()
                     // Write the next link if it's available.
                     await this.WriteResourceSetNextLinkAsync(resourceSet.NextPageLink, propertyName)
                         .ConfigureAwait(false);
+
+                    await this.WriteResourceSetCountAsync(resourceSet.Count, propertyName)
+                        .ConfigureAwait(false);
                 }
             }
 
@@ -2736,7 +2743,11 @@ await this.asynchronousJsonWriter.EndObjectScopeAsync()
         /// <returns>A task that represents the asynchronous write operation.</returns>
         private Task WriteResourceSetCountAsync(long? count, string propertyName)
         {
-            if (count.HasValue)
+            bool countWritten = this.State == WriterState.ResourceSet ?
+                this.CurrentResourceSetScope.CountWritten :
+                this.CurrentDeltaResourceSetScope.CountWritten;
+
+            if (count.HasValue && !countWritten)
             {
                 return WriteResourceSetCountInnerAsync(count.Value, propertyName);
 
@@ -2756,6 +2767,15 @@ await this.asynchronousODataAnnotationWriter.WritePropertyAnnotationNameAsync(
 
                     await this.asynchronousJsonWriter.WriteValueAsync(innerCount)
                         .ConfigureAwait(false);
+
+                    if (this.State == WriterState.ResourceSet)
+                    {
+                        this.CurrentResourceSetScope.CountWritten = true;
+                    }
+                    else
+                    {
+                        this.CurrentDeltaResourceSetScope.CountWritten = true;
+                    }
                 }
             }
 
@@ -3060,6 +3080,9 @@ await this.asynchronousJsonWriter.WriteValueAsync(UriUtils.UriToString(link.Targ
         /// </summary>
         private sealed class JsonLightResourceSetScope : ResourceSetScope
         {
+            /// <summary>true if the odata.count was already written, false otherwise.</summary>
+            private bool countWritten;
+
             /// <summary>true if the odata.nextLink was already written, false otherwise.</summary>
             private bool nextLinkWritten;
 
@@ -3085,6 +3108,22 @@ internal JsonLightResourceSetScope(ODataResourceSet resourceSet, IEdmNavigationS
                 this.isUndeclared = isUndeclared;
             }
 
+            /// <summary>
+            /// true if the odata.count annotation was already written, false otherwise.
+            /// </summary>
+            internal bool CountWritten
+            {
+                get
+                {
+                    return this.countWritten;
+                }
+
+                set
+                {
+                    this.countWritten = value;
+                }
+            }
+
             /// <summary>
             /// true if the odata.nextLink annotation was already written, false otherwise.
             /// </summary>
@@ -3373,6 +3412,9 @@ public JsonLightDeltaLinkScope(WriterState state, ODataItem link, ODataResourceS
         /// </summary>
         private sealed class JsonLightDeltaResourceSetScope : DeltaResourceSetScope
         {
+            /// <summary>true if the odata.count was already written, false otherwise.</summary>
+            private bool countWritten;
+
             /// <summary>true if the odata.nextLink was already written, false otherwise.</summary>
             private bool nextLinkWritten;
 
@@ -3392,6 +3434,22 @@ public JsonLightDeltaResourceSetScope(ODataDeltaResourceSet resourceSet, IEdmNav
             {
             }
 
+            /// <summary>
+            /// true if the odata.count annotation was already written, false otherwise.
+            /// </summary>
+            internal bool CountWritten
+            {
+                get
+                {
+                    return this.countWritten;
+                }
+
+                set
+                {
+                    this.countWritten = value;
+                }
+            }
+
             /// <summary>
             /// true if the odata.nextLink annotation was already written, false otherwise.
             /// </summary>