Update NextLinkOperationImplementation to support LRO rehydration

eng/DownloadSharedSource.ps1

@@ -19,7 +19,8 @@ $files = @('AsyncLockWithValue.cs', 'ClientDiagnostics.cs', 'DiagnosticScope.cs'
     'Multipart/MultipartContent.cs', 'AzureKeyCredentialPolicy.cs', 'AppContextSwitchHelper.cs',
     'ConstantDelayStrategy.cs', 'DelayStrategy.cs', 'ExponentialDelayStrategy.cs', 'OperationPoller.cs', 'RetryAfterDelayStrategy.cs',
     'ForwardsClientCallsAttribute.cs', 'AsyncLockWithValue.cs', 'VoidValue.cs')
-$baseUrl = 'https://raw.githubusercontent.com/Azure/azure-sdk-for-net/main/sdk/core/Azure.Core/src/Shared/'
+# TODO: temporary change to target at support_lro_rehydration branch
+$baseUrl = 'https://raw.githubusercontent.com/Azure/azure-sdk-for-net/support_lro_rehydration/sdk/core/Azure.Core/src/Shared/'
 DownloadAll $files $baseUrl $downloadPath
 
 #Download management Shared

src/assets/Azure.Core.Shared/OperationInternal.cs

@@ -86,7 +86,7 @@ internal class OperationInternal : OperationInternalBase
         public OperationInternal(
             ClientDiagnostics clientDiagnostics,
             IOperation operation,
-            Response rawResponse,
+            Response? rawResponse,
             string? operationTypeName = null,
             IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null,
             DelayStrategy? fallbackStrategy = null)
@@ -110,6 +110,11 @@ private OperationInternal(OperationState finalState)
         protected override async ValueTask<Response> UpdateStatusAsync(bool async, CancellationToken cancellationToken) =>
             async ? await _internalOperation.UpdateStatusAsync(cancellationToken).ConfigureAwait(false) : _internalOperation.UpdateStatus(cancellationToken);
 
+        public virtual string GetOperationId()
+        {
+            return _internalOperation.GetOperationId();
+        }
+
         // Wrapper type that converts OperationState to OperationState<T> and can be passed to `OperationInternal<T>` constructor.
         private class OperationToOperationOfTProxy : IOperation<VoidValue>
         {
@@ -135,6 +140,11 @@ public async ValueTask<OperationState<VoidValue>> UpdateStateAsync(bool async, C
 
                 return OperationState<VoidValue>.Failure(state.RawResponse, state.OperationFailedException);
             }
+
+            public string GetOperationId()
+            {
+                return _operation.GetOperationId();
+            }
         }
     }
 
@@ -173,6 +183,11 @@ internal interface IOperation
         /// </list>
         /// </returns>
         ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// To get the Id of the operation for rehydration purpose.
+        /// </summary>
+        string GetOperationId();
     }
 
     /// <summary>

src/assets/Azure.Core.Shared/OperationInternalOfT.cs

@@ -5,6 +5,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core.Pipeline;
@@ -53,7 +54,7 @@ internal class OperationInternal<T> : OperationInternalBase
     {
         private readonly IOperation<T> _operation;
         private readonly AsyncLockWithValue<OperationState<T>> _stateLock;
-        private Response _rawResponse;
+        private Response? _rawResponse;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="OperationInternal"/> class in a final successful state.
@@ -95,7 +96,7 @@ internal class OperationInternal<T> : OperationInternalBase
         public OperationInternal(
             ClientDiagnostics clientDiagnostics,
             IOperation<T> operation,
-            Response rawResponse,
+            Response? rawResponse,
             string? operationTypeName = null,
             IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null,
             DelayStrategy? fallbackStrategy = null)
@@ -116,7 +117,7 @@ private OperationInternal(OperationState<T> finalState)
             _stateLock = new AsyncLockWithValue<OperationState<T>>(finalState);
         }
 
-        public override Response RawResponse => _stateLock.TryGetValue(out var state) ? state.RawResponse : _rawResponse;
+        public override Response RawResponse => (_stateLock.TryGetValue(out var state) ? state.RawResponse : _rawResponse) ?? throw new InvalidOperationException("The operation does not have a response yet. Please call UpdateStatus or WaitForCompletion first.");
 
         public override bool HasCompleted => _stateLock.HasValue;
 
@@ -279,6 +280,11 @@ protected override async ValueTask<Response> UpdateStatusAsync(bool async, Cance
             }
         }
 
+        public virtual string GetOperationId()
+        {
+            return _operation.GetOperationId();
+        }
+
         private static Response GetResponseFromState(OperationState<T> state)
         {
             if (state.HasSucceeded)
@@ -293,6 +299,8 @@ private class FinalOperation : IOperation<T>
         {
             public ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
                 => throw new NotSupportedException("The operation has already completed");
+
+            public string GetOperationId() => string.Empty;
         }
     }
 
@@ -333,6 +341,11 @@ internal interface IOperation<T>
         /// </list>
         /// </returns>
         ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// To get the Id of the operation for rehydration purpose.
+        /// </summary>
+        string GetOperationId();
     }
 
     /// <summary>

src/assets/Generator.Shared/NextLinkOperationImplementation.cs

@@ -4,9 +4,11 @@
 #nullable enable
 
 using System;
-using System.Linq;
+using System.Collections.Generic;
 using System.IO;
+using System.Linq;
 using System.Text.Json;
+using System.Text.Json.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core.Pipeline;
@@ -62,6 +64,38 @@ public static IOperation<T> Create<T>(
             return new OperationToOperationOfT<T>(operationSource, operation);
         }
 
+        public static IOperation? Create(
+            HttpPipeline pipeline,
+            string id,
+            string? apiVersionOverride = null)
+        {
+            var lroDetails = BinaryData.FromBytes(Convert.FromBase64String(id)).ToObjectFromJson<Dictionary<string, string>>();
+            if (!Uri.TryCreate(lroDetails["InitialUri"], UriKind.Absolute, out var startRequestUri))
+                throw new InvalidOperationException("Invalid initial URI");
+            if (!lroDetails.TryGetValue("NextRequestUri", out var nextRequestUri))
+                throw new InvalidOperationException("Invalid next request URI");
+            RequestMethod requestMethod = new RequestMethod(lroDetails["RequestMethod"]);
+            bool originalResponseHasLocation = bool.Parse(lroDetails["OriginalResponseHasLocation"]);
+            string lastKnownLocation = lroDetails["LastKnownLocation"];
+            if (!Enum.TryParse(lroDetails["FinalStateVia"], out OperationFinalStateVia finalStateVia))
+                finalStateVia = OperationFinalStateVia.Location;
+            string? apiVersionStr = apiVersionOverride ?? (TryGetApiVersion(startRequestUri, out ReadOnlySpan<char> apiVersion) ? apiVersion.ToString() : null);
+            if (!Enum.TryParse(lroDetails["HeaderSource"], out HeaderSource headerSource))
+                headerSource = HeaderSource.None;
+
+            return new NextLinkOperationImplementation(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, originalResponseHasLocation, lastKnownLocation, finalStateVia, apiVersionStr);
+        }
+
+        public static IOperation<T>? Create<T>(
+            IOperationSource<T> operationSource,
+            HttpPipeline pipeline,
+            string id,
+            string? apiVersionOverride = null)
+        {
+            var operation = Create(pipeline, id, apiVersionOverride);
+            return new OperationToOperationOfT<T>(operationSource, operation!);
+        }
+
         private NextLinkOperationImplementation(
             HttpPipeline pipeline,
             RequestMethod requestMethod,
@@ -84,6 +118,22 @@ private NextLinkOperationImplementation(
             _apiVersion = apiVersion;
         }
 
+        public string GetOperationId()
+        {
+            var lroDetails = new Dictionary<string, string?>()
+            {
+                ["HeaderSource"] = _headerSource.ToString(),
+                ["NextRequestUri"] = _nextRequestUri,
+                ["InitialUri"] = _startRequestUri.AbsoluteUri,
+                ["RequestMethod"] = _requestMethod.ToString(),
+                ["OriginalResponseHasLocation"] = _originalResponseHasLocation.ToString(),
+                ["LastKnownLocation"] = _lastKnownLocation,
+                ["FinalStateVia"] = _finalStateVia.ToString()
+            };
+            var lroData = BinaryData.FromObjectAsJson(lroDetails);
+            return Convert.ToBase64String(lroData.ToArray());
+        }
+
         public async ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken)
         {
             Response response = await GetResponseAsync(async, _nextRequestUri, cancellationToken).ConfigureAwait(false);
@@ -402,6 +452,11 @@ public CompletedOperation(OperationState operationState)
             }
 
             public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(_operationState);
+
+            public string GetOperationId()
+            {
+                return string.Empty;
+            }
         }
 
         private sealed class OperationToOperationOfT<T> : IOperation<T>
@@ -434,6 +489,11 @@ public async ValueTask<OperationState<T>> UpdateStateAsync(bool async, Cancellat
 
                 return OperationState<T>.Pending(state.RawResponse);
             }
+
+            public string GetOperationId()
+            {
+                return _operation.GetOperationId();
+            }
         }
     }
 }

src/assets/Generator.Shared/ProtocolOperation.cs

@@ -23,12 +23,8 @@ internal ProtocolOperation(ClientDiagnostics clientDiagnostics, HttpPipeline pip
             _operation = new OperationInternal<T>(clientDiagnostics, this, response, scopeName);
         }
 
-#pragma warning disable CA1822
-        // This scenario is currently unsupported.
-        // See: https://github.com/Azure/autorest.csharp/issues/2158.
         /// <inheritdoc />
-        public override string Id => throw new NotSupportedException();
-#pragma warning restore CA1822
+        public override string Id => _nextLinkOperation.GetOperationId();
 
         /// <inheritdoc />
         public override T Value => _operation.Value;
@@ -69,5 +65,10 @@ async ValueTask<OperationState<T>> IOperation<T>.UpdateStateAsync(bool async, Ca
 
             return OperationState<T>.Pending(state.RawResponse);
         }
+
+        public string GetOperationId()
+        {
+            return _nextLinkOperation.GetOperationId();
+        }
     }
 }