[Core] Wiring up rehydration token logic in data-plane

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

@@ -60,7 +60,7 @@ public static IOperation Create(
             var headerSource = GetHeaderSource(requestMethod, startRequestUri, response, apiVersionStr, out string nextRequestUri, out bool isNextRequestPolling);
             if (headerSource == HeaderSource.None && IsFinalState(response, headerSource, out var failureState, out _))
             {
-                return new CompletedOperation(failureState ?? GetOperationStateFromFinalResponse(requestMethod, response));
+                return new CompletedOperation(failureState ?? GetOperationStateFromFinalResponse(requestMethod, response), startRequestUri);
             }
 
             string? lastKnownLocation;
@@ -647,12 +647,18 @@ private class CompletedOperation : IOperation
         {
             private readonly OperationState _operationState;
 
-            public CompletedOperation(OperationState operationState)
+            private readonly Uri _startRequestUri;
+
+            public CompletedOperation(OperationState operationState, Uri startRequestUri)
             {
                 _operationState = operationState;
+                _startRequestUri = startRequestUri;
             }
 
             public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(_operationState);
+
+            public RehydrationToken GetRehydrationToken() =>
+                NextLinkOperationImplementation.GetRehydrationToken(RequestMethod.Put, _startRequestUri, _startRequestUri.AbsoluteUri, "None", null, OperationFinalStateVia.OriginalUri.ToString());
         }
 
         private sealed class OperationToOperationOfT<T> : IOperation<T>
@@ -685,6 +691,8 @@ public async ValueTask<OperationState<T>> UpdateStateAsync(bool async, Cancellat
 
                 return OperationState<T>.Pending(state.RawResponse);
             }
+
+            public RehydrationToken? GetRehydrationToken() => _operation.GetRehydrationToken();
         }
     }
 }

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

@@ -119,6 +119,8 @@ public OperationToOperationOfTProxy(IOperation operation)
                 _operation = operation;
             }
 
+            public RehydrationToken? GetRehydrationToken() => _operation.GetRehydrationToken();
+
             public async ValueTask<OperationState<VoidValue>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
             {
                 var state = await _operation.UpdateStateAsync(async, cancellationToken).ConfigureAwait(false);
@@ -172,6 +174,11 @@ internal interface IOperation
         /// </list>
         /// </returns>
         ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Get a token that can be used to rehydrate the operation.
+        /// </summary>
+        RehydrationToken GetRehydrationToken();
     }
 
     /// <summary>

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

@@ -288,6 +288,10 @@ private class FinalOperation : IOperation<T>
         {
             public ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken)
                 => throw new NotSupportedException("The operation has already completed");
+
+            // Unreachable path. _operation.GetRehydrationToken() is never invoked.
+            public RehydrationToken? GetRehydrationToken()
+                => throw new NotSupportedException($"Getting the rehydration token of a {nameof(FinalOperation)} is not supported");
         }
     }
 
@@ -328,6 +332,11 @@ internal interface IOperation<T>
         /// </list>
         /// </returns>
         ValueTask<OperationState<T>> UpdateStateAsync(bool async, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Get a token that can be used to rehydrate the operation.
+        /// </summary>
+        RehydrationToken? GetRehydrationToken();
     }
 
     /// <summary>

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

@@ -30,6 +30,9 @@ internal ProtocolOperation(ClientDiagnostics clientDiagnostics, HttpPipeline pip
         public override string Id => throw new NotSupportedException();
 #pragma warning restore CA1822
 
+        /// <inheritdoc />
+        public override RehydrationToken? GetRehydrationToken() => _nextLinkOperation.GetRehydrationToken();
+
         /// <inheritdoc />
         public override T Value => _operation.Value;
 

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

@@ -76,6 +76,7 @@ public ConvertOperation(Operation<TFrom> operation, ClientDiagnostics diagnostic
             public override TTo Value => GetOrCreateValue();
             public override bool HasValue => _operation.HasValue;
             public override bool HasCompleted => _operation.HasCompleted;
+            public override RehydrationToken? GetRehydrationToken() => _operation.GetRehydrationToken();
             public override Response GetRawResponse() => _operation.GetRawResponse();
 
             public override Response UpdateStatus(CancellationToken cancellationToken = default)

sdk/core/Azure.Core/tests/OperationInternalOfTTests.cs

@@ -518,6 +518,8 @@ public TestOperation(Func<bool, CancellationToken, ValueTask<OperationState<int>
             }
 
             public ValueTask<OperationState<int>> UpdateStateAsync(bool async, CancellationToken cancellationToken) => _updateStateAsyncHandler(async, cancellationToken);
+
+            public RehydrationToken? GetRehydrationToken() => null;
         }
 
         private class CallCountStrategy : DelayStrategy

sdk/core/Azure.Core/tests/OperationInternalTests.cs

@@ -390,7 +390,7 @@
             }
         }
 
         private readonly struct TestOperation : IOperation
         {
             private readonly Func<bool, CancellationToken, ValueTask<OperationState>> _updateStateAsyncHandler;
 
@@ -438,6 +438,8 @@
             }
 
             public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => _updateStateAsyncHandler(async, cancellationToken);
+
+            public RehydrationToken? GetRehydrationToken() => null;
         }
 
         private class CallCountStrategy : DelayStrategy

sdk/core/Azure.Core/tests/OperationPollerTests.cs

@@ -54,9 +54,11 @@
             Assert.Catch<OperationCanceledException>(() => poller.WaitForCompletionResponse(operation, null, cts.Token));
         }
 
         private class EndlessOperation : IOperation
         {
             public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(OperationState.Pending(new MockResponse(200)));
+
+            public RehydrationToken? GetRehydrationToken() => default;
         }
 
         private class TestDelayStrategy : DelayStrategy

sdk/core/Azure.Core/tests/ProtocolOperationTests.cs

@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.IO;
+using System.Text;
+using System.Threading.Tasks;
+using Azure.Core.Pipeline;
+using Azure.Core.TestFramework;
+using NUnit.Framework;
+
+namespace Azure.Core.Tests
+{
+    public class ProtocolOperationTests
+    {
+        [Test]
+        public void GetRehydrationToken()
+        {
+            const string OperationId = "oper-ati-oni-d00";
+            const string RequestUri = "https://www.example.com/request";
+            const string LocationUri = $"https://www.example.com/{OperationId}";
+
+            using var request = new MockRequest() { Method = RequestMethod.Post };
+            using var response = new MockResponse(200);
+            var transport = new MockTransport();
+            var pipeline = new HttpPipeline(transport);
+
+            request.Uri.Reset(new Uri(RequestUri));
+            response.AddHeader("Location", LocationUri);
+
+            var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
+            var token = operation.GetRehydrationToken();
+
+            Assert.True(token.HasValue);
+            Assert.AreEqual(OperationId, token.Value.Id);
+            Assert.AreEqual(NextLinkOperationImplementation.RehydrationTokenVersion, token.Value.Version);
+            Assert.AreEqual("Location", token.Value.HeaderSource);
+            Assert.AreEqual(LocationUri, token.Value.NextRequestUri);
+            Assert.AreEqual(RequestUri, token.Value.InitialUri);
+            Assert.AreEqual(RequestMethod.Post, token.Value.RequestMethod);
+            Assert.AreEqual(LocationUri, token.Value.LastKnownLocation);
+            Assert.AreEqual("Location", token.Value.FinalStateVia);
+        }
+
+        [Test]
+        [TestCase(true)]
+        [TestCase(false)]
+        public async Task RehydrationOperationCanPoll(bool isAsync)
+        {
+            const string RequestUri = "https://www.example.com/request";
+            const string LocationUri = "https://www.example.com/oper-ati-oni-d00";
+
+            using var request = new MockRequest() { Method = RequestMethod.Post };
+            using var response = new MockResponse(200);
+            var transport = new MockTransport(response);
+            var pipeline = new HttpPipeline(transport);
+
+            request.Uri.Reset(new Uri(RequestUri));
+            response.AddHeader("Location", LocationUri);
+
+            var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
+            var token = operation.GetRehydrationToken();
+            _ = isAsync
+                ? await Operation.RehydrateAsync(pipeline, token.Value)
+                : Operation.Rehydrate(pipeline, token.Value);
+            var rehydrationUpdateRequest = transport.SingleRequest;
+
+            Assert.AreEqual(LocationUri, rehydrationUpdateRequest.Uri.ToString());
+            Assert.AreEqual(RequestMethod.Get, rehydrationUpdateRequest.Method);
+        }
+
+        [Test]
+        [TestCase(true)]
+        [TestCase(false)]
+        public async Task RehydrationOperationOfTCanPoll(bool isAsync)
+        {
+            const string RequestUri = "https://www.example.com/request";
+            const string LocationUri = "https://www.example.com/oper-ati-oni-d00";
+
+            using var request = new MockRequest() { Method = RequestMethod.Post };
+            using var response = new MockResponse(200);
+            var transport = new MockTransport(response);
+            var pipeline = new HttpPipeline(transport);
+            using var responseContentStream = new MemoryStream(Encoding.UTF8.GetBytes("""
+                {
+                    "IntValue": 1,
+                    "StringValue": "one"
+                }
+                """));
+
+            request.Uri.Reset(new Uri(RequestUri));
+            response.AddHeader("Location", LocationUri);
+            response.ContentStream = responseContentStream;
+
+            var operation = new ProtocolOperation<MockJsonModel>(null, pipeline, request, response, OperationFinalStateVia.Location, null, null);
+            var token = operation.GetRehydrationToken();
+            var rehydrationOperation = isAsync
+                ? await Operation.RehydrateAsync<MockJsonModel>(pipeline, token.Value)
+                : Operation.Rehydrate<MockJsonModel>(pipeline, token.Value);
+            var rehydrationUpdateRequest = transport.SingleRequest;
+
+            Assert.AreEqual(LocationUri, rehydrationUpdateRequest.Uri.ToString());
+            Assert.AreEqual(RequestMethod.Get, rehydrationUpdateRequest.Method);
+
+            Assert.True(rehydrationOperation.HasValue);
+            Assert.AreEqual(rehydrationOperation.Value.IntValue, 1);
+            Assert.AreEqual(rehydrationOperation.Value.StringValue, "one");
+        }
+    }
+}