Update Azure.Core shared for LRO rehydration

sdk/core/Azure.Core/src/Shared/NextLinkOperationImplementation.cs

@@ -4,8 +4,10 @@
 #nullable enable
 
 using System;
+using System.ClientModel.Primitives;
+using System.Collections.Generic;
 using System.Linq;
-using System.IO;
+using System.Reflection;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -38,7 +40,7 @@ public static IOperation Create(
             bool skipApiVersionOverride = false,
             string? apiVersionOverrideValue = null)
         {
-            string? apiVersionStr = null;
+            string? apiVersionStr;
             if (apiVersionOverrideValue is not null)
             {
                 apiVersionStr = apiVersionOverrideValue;
@@ -70,6 +72,41 @@ public static IOperation<T> Create<T>(
             return new OperationToOperationOfT<T>(operationSource, operation);
         }
 
+        public static IOperation Create(
+            HttpPipeline pipeline,
+            RehydrationToken? rehydrationToken,
+            string? apiVersionOverride = null)
+        {
+            Argument.AssertNotNull(rehydrationToken, nameof(rehydrationToken));
+
+            var lroDetails = ModelReaderWriter.Write(rehydrationToken!, new ModelReaderWriterOptions("J")).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"]);
+            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, lastKnownLocation, finalStateVia, apiVersionStr);
+        }
+
+        public static IOperation<T> Create<T>(
+            IOperationSource<T> operationSource,
+            HttpPipeline pipeline,
+            RehydrationToken? rehydrationToken,
+            string? apiVersionOverride = null)
+        {
+            Argument.AssertNotNull(rehydrationToken, nameof(rehydrationToken));
+
+            var operation = Create(pipeline, rehydrationToken, apiVersionOverride);
+            return new OperationToOperationOfT<T>(operationSource, operation);
+        }
+
         private NextLinkOperationImplementation(
             HttpPipeline pipeline,
             RequestMethod requestMethod,
@@ -90,6 +127,41 @@ private NextLinkOperationImplementation(
             _apiVersion = apiVersion;
         }
 
+        public static RehydrationToken GetRehydrationToken(
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            Response response,
+            OperationFinalStateVia finalStateVia,
+            bool skipApiVersionOverride = false,
+            string? apiVersionOverrideValue = null)
+        {
+            string? apiVersionStr;
+            if (apiVersionOverrideValue is not null)
+            {
+                apiVersionStr = apiVersionOverrideValue;
+            }
+            else
+            {
+                apiVersionStr = !skipApiVersionOverride && TryGetApiVersion(startRequestUri, out ReadOnlySpan<char> apiVersion) ? apiVersion.ToString() : null;
+            }
+            var headerSource = GetHeaderSource(requestMethod, startRequestUri, response, apiVersionStr, out var nextRequestUri);
+            response.Headers.TryGetValue("Location", out var lastKnownLocation);
+            return GetRehydrationToken(requestMethod, startRequestUri, nextRequestUri, headerSource.ToString(), lastKnownLocation, finalStateVia.ToString());
+        }
+
+        public static RehydrationToken GetRehydrationToken(
+            RequestMethod requestMethod,
+            Uri startRequestUri,
+            string nextRequestUri,
+            string headerSource,
+            string? lastKnownLocation,
+            string finalStateVia)
+        {
+            var parameters = new object?[] { null, null, headerSource, nextRequestUri, startRequestUri.AbsoluteUri, requestMethod, lastKnownLocation, finalStateVia };
+            var rehydrationToken = Activator.CreateInstance(typeof(RehydrationToken), BindingFlags.NonPublic | BindingFlags.Instance, null, parameters, null);
+            return (RehydrationToken)rehydrationToken!;
+        }
+
         public async ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken)
         {
             Response response = await GetResponseAsync(async, _nextRequestUri, cancellationToken).ConfigureAwait(false);
@@ -307,7 +379,7 @@ private static bool IsFinalState(Response response, HeaderSource headerSource, o
 
             if (response.Status is >= 200 and <= 204)
             {
-                if (response.ContentStream is {Length: > 0})
+                if (response.ContentStream is { Length: > 0 })
                 {
                     try
                     {
@@ -399,7 +471,7 @@ private static HeaderSource GetHeaderSource(RequestMethod requestMethod, Uri req
             return HeaderSource.None;
         }
 
-        private enum HeaderSource
+        internal enum HeaderSource
         {
             None,
             OperationLocation,

sdk/core/Azure.Core/src/Shared/OperationInternal.cs

@@ -51,14 +51,16 @@ internal class OperationInternal : OperationInternalBase
         /// Initializes a new instance of the <see cref="OperationInternal"/> class in a final successful state.
         /// </summary>
         /// <param name="rawResponse">The final value of <see cref="OperationInternalBase.RawResponse"/>.</param>
-        public static OperationInternal Succeeded(Response rawResponse) => new(OperationState.Success(rawResponse));
+        /// <param name="rehydrationToken">The rehydration token.</param>
+        public static OperationInternal Succeeded(Response rawResponse, RehydrationToken? rehydrationToken = null) => new(OperationState.Success(rawResponse), rehydrationToken);
 
         /// <summary>
         /// Initializes a new instance of the <see cref="OperationInternal"/> class in a final failed state.
         /// </summary>
         /// <param name="rawResponse">The final value of <see cref="OperationInternalBase.RawResponse"/>.</param>
         /// <param name="operationFailedException">The exception that will be thrown by <c>UpdateStatusAsync</c>.</param>
-        public static OperationInternal Failed(Response rawResponse, RequestFailedException operationFailedException) => new(OperationState.Failure(rawResponse, operationFailedException));
+        /// <param name="rehydrationToken">rehydration token</param>
+        public static OperationInternal Failed(Response rawResponse, RequestFailedException operationFailedException, RehydrationToken? rehydrationToken = null) => new(OperationState.Failure(rawResponse, operationFailedException), rehydrationToken);
 
         /// <summary>
         /// Initializes a new instance of the <see cref="OperationInternal"/> class.
@@ -83,23 +85,25 @@ internal class OperationInternal : OperationInternalBase
         /// </param>
         /// <param name="scopeAttributes">The attributes to use during diagnostic scope creation.</param>
         /// <param name="fallbackStrategy"> The delay strategy to use. Default is <see cref="FixedDelayWithNoJitterStrategy"/>.</param>
+        /// <param name="rehydrationToken">The rehydration token.</param>
         public OperationInternal(IOperation operation,
             ClientDiagnostics clientDiagnostics,
-            Response rawResponse,
+            Response? rawResponse,
             string? operationTypeName = null,
             IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null,
-            DelayStrategy? fallbackStrategy = null)
-            :base(clientDiagnostics, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy)
+            DelayStrategy? fallbackStrategy = null,
+            RehydrationToken? rehydrationToken = null)
+            : base(clientDiagnostics, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy, rehydrationToken)
         {
-            _internalOperation = new OperationInternal<VoidValue>(new OperationToOperationOfTProxy(operation), clientDiagnostics, rawResponse, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy);
+            _internalOperation = new OperationInternal<VoidValue>(new OperationToOperationOfTProxy(operation), clientDiagnostics, rawResponse, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy, rehydrationToken);
         }
 
-        private OperationInternal(OperationState finalState)
-            :base(finalState.RawResponse)
+        private OperationInternal(OperationState finalState, RehydrationToken? rehydrationToken)
+            : base(finalState.RawResponse, rehydrationToken)
         {
             _internalOperation = finalState.HasSucceeded
-                ? OperationInternal<VoidValue>.Succeeded(finalState.RawResponse, default)
-                : OperationInternal<VoidValue>.Failed(finalState.RawResponse, finalState.OperationFailedException!);
+                ? OperationInternal<VoidValue>.Succeeded(finalState.RawResponse, default, rehydrationToken)
+                : OperationInternal<VoidValue>.Failed(finalState.RawResponse, finalState.OperationFailedException!, rehydrationToken);
         }
 
         public override Response RawResponse => _internalOperation.RawResponse;

sdk/core/Azure.Core/src/Shared/OperationInternalBase.cs

@@ -22,8 +22,9 @@ internal abstract class OperationInternalBase
         private readonly string _waitForCompletionResponseScopeName;
         protected readonly string _updateStatusScopeName;
         protected readonly string _waitForCompletionScopeName;
+        protected readonly RehydrationToken? _rehydrationToken;
 
-        protected OperationInternalBase(Response rawResponse)
+        protected OperationInternalBase(Response rawResponse, RehydrationToken? rehydrationToken)
         {
             _diagnostics = new ClientDiagnostics(ClientOptions.Default);
             _updateStatusScopeName = string.Empty;
@@ -32,9 +33,10 @@ protected OperationInternalBase(Response rawResponse)
             _scopeAttributes = default;
             _fallbackStrategy = default;
             _responseLock = new AsyncLockWithValue<Response>(rawResponse);
+            _rehydrationToken = rehydrationToken;
         }
 
-        protected OperationInternalBase(ClientDiagnostics clientDiagnostics, string operationTypeName, IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null, DelayStrategy? fallbackStrategy = null)
+        protected OperationInternalBase(ClientDiagnostics clientDiagnostics, string operationTypeName, IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null, DelayStrategy? fallbackStrategy = null, RehydrationToken? rehydrationToken = null)
         {
             _diagnostics = clientDiagnostics;
             _updateStatusScopeName = $"{operationTypeName}.{nameof(UpdateStatus)}";
@@ -43,6 +45,7 @@ protected OperationInternalBase(ClientDiagnostics clientDiagnostics, string oper
             _scopeAttributes = scopeAttributes?.ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
             _fallbackStrategy = fallbackStrategy;
             _responseLock = new AsyncLockWithValue<Response>();
+            _rehydrationToken = rehydrationToken;
         }
 
         /// <summary>

sdk/core/Azure.Core/src/Shared/OperationInternalOfT.cs

@@ -4,7 +4,10 @@
 #nullable enable
 
 using System;
+using System.ClientModel.Primitives;
 using System.Collections.Generic;
+using System.IO;
+using System.Net;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core.Pipeline;
@@ -53,21 +56,23 @@ 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.
         /// </summary>
         /// <param name="rawResponse">The final value of <see cref="OperationInternalBase.RawResponse"/>.</param>
         /// <param name="value">The final result of the long-running operation.</param>
-        public static OperationInternal<T> Succeeded(Response rawResponse, T value) => new(OperationState<T>.Success(rawResponse, value));
+        /// <param name="rehydrationToken">rehydration token</param>
+        public static OperationInternal<T> Succeeded(Response rawResponse, T value, RehydrationToken? rehydrationToken = null) => new(OperationState<T>.Success(rawResponse, value), rehydrationToken);
 
         /// <summary>
         /// Initializes a new instance of the <see cref="OperationInternal"/> class in a final failed state.
         /// </summary>
         /// <param name="rawResponse">The final value of <see cref="OperationInternalBase.RawResponse"/>.</param>
         /// <param name="operationFailedException">The exception that will be thrown by <c>UpdateStatusAsync</c>.</param>
-        public static OperationInternal<T> Failed(Response rawResponse, RequestFailedException operationFailedException) => new(OperationState<T>.Failure(rawResponse, operationFailedException));
+        /// <param name="rehydrationToken">rehydration token</param>
+        public static OperationInternal<T> Failed(Response rawResponse, RequestFailedException operationFailedException, RehydrationToken? rehydrationToken = null) => new(OperationState<T>.Failure(rawResponse, operationFailedException), rehydrationToken);
 
         /// <summary>
         /// Initializes a new instance of the <see cref="OperationInternal{T}"/> class.
@@ -93,21 +98,23 @@ internal class OperationInternal<T> : OperationInternalBase
         /// <param name="scopeAttributes">The attributes to use during diagnostic scope creation.</param>
         /// <param name="fallbackStrategy">The delay strategy when Retry-After header is not present.  When it is present, the longer of the two delays will be used.
         ///     Default is <see cref="FixedDelayWithNoJitterStrategy"/>.</param>
+        /// <param name="rehydrationToken">The rehydration token.</param>
         public OperationInternal(IOperation<T> operation,
             ClientDiagnostics clientDiagnostics,
-            Response rawResponse,
+            Response? rawResponse,
             string? operationTypeName = null,
             IEnumerable<KeyValuePair<string, string>>? scopeAttributes = null,
-            DelayStrategy? fallbackStrategy = null)
-            : base(clientDiagnostics, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy)
+            DelayStrategy? fallbackStrategy = null,
+            RehydrationToken? rehydrationToken = null)
+            : base(clientDiagnostics, operationTypeName ?? operation.GetType().Name, scopeAttributes, fallbackStrategy, rehydrationToken)
         {
             _operation = operation;
             _rawResponse = rawResponse;
             _stateLock = new AsyncLockWithValue<OperationState<T>>();
         }
 
-        private OperationInternal(OperationState<T> finalState)
-            : base(finalState.RawResponse)
+        private OperationInternal(OperationState<T> finalState, RehydrationToken? rehydrationToken)
+            : base(finalState.RawResponse, rehydrationToken)
         {
             // FinalOperation represents operation that is in final state and can't be updated.
             // It implements IOperation<T> and throws exception when UpdateStateAsync is called.
@@ -116,7 +123,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;
 
@@ -265,7 +272,7 @@ protected override async ValueTask<Response> UpdateStatusAsync(bool async, Cance
                 }
 
                 asyncLock.SetValue(state);
-                return GetResponseFromState(state);
+                return GetResponseFromState(state, GetRequestMethod(_rehydrationToken));
             }
             catch (Exception e)
             {
@@ -274,16 +281,89 @@ protected override async ValueTask<Response> UpdateStatusAsync(bool async, Cance
             }
         }
 
-        private static Response GetResponseFromState(OperationState<T> state)
+        private RequestMethod? GetRequestMethod(RehydrationToken? rehydrationToken)
+        {
+            if (rehydrationToken is null)
+            {
+                return null;
+            }
+            var lroDetails = ModelReaderWriter.Write(rehydrationToken, new ModelReaderWriterOptions("J")).ToObjectFromJson<Dictionary<string, string>>();
+            return new RequestMethod(lroDetails["requestMethod"]);
+        }
+
+        private static Response GetResponseFromState(OperationState<T> state, RequestMethod? requestmethod = null)
         {
             if (state.HasSucceeded)
             {
                 return state.RawResponse;
             }
 
+            // if this is a fake delete lro with 404, just return empty response with 200
+            if (RequestMethod.Delete == requestmethod && state.RawResponse.Status == 404)
+            {
+                return new EmptyResponse(HttpStatusCode.OK);
+            }
+
             throw state.OperationFailedException!;
         }
 
+        private class EmptyResponse : Response
+        {
+            public EmptyResponse(HttpStatusCode status)
+            {
+                Status = (int)status;
+                ReasonPhrase = status.ToString();
+            }
+
+            public override int Status { get; }
+
+            public override string ReasonPhrase { get; }
+
+            public override Stream? ContentStream { get => null; set => throw new System.NotImplementedException(); }
+            public override string ClientRequestId { get => string.Empty; set => throw new System.NotImplementedException(); }
+
+            public override void Dispose()
+            {
+                throw new System.NotImplementedException();
+            }
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override bool ContainsHeader(string name)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override IEnumerable<HttpHeader> EnumerateHeaders()
+            {
+                throw new System.NotImplementedException();
+            }
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override bool TryGetHeader(string name, out string value)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            /// <inheritdoc />
+#if HAS_INTERNALS_VISIBLE_CORE
+            internal
+#endif
+            protected override bool TryGetHeaderValues(string name, out IEnumerable<string> values)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
         private class FinalOperation : IOperation<T>
         {
             public ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken)