[Core] Wiring up rehydration token logic in data-plane (Azure#49224)

Author: kinelski

sdk/communication/Azure.Communication.Email/src/EmailSendOperation.cs

@@ -184,6 +184,10 @@ async ValueTask<OperationState<EmailSendResult>> IOperation<EmailSendResult>.Upd
             return OperationState<EmailSendResult>.Pending(rawResponse);
         }
 
+        // This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
+        RehydrationToken IOperation<EmailSendResult>.GetRehydrationToken() =>
+            throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
+
         private static IDictionary<string, string> CreateAdditionalInformation(ErrorDetail error)
         {
             if (string.IsNullOrEmpty(error.ToString()))

sdk/confidentialledger/Azure.Security.CodeTransparency/src/CreateEntryOperation.cs

@@ -125,5 +125,9 @@ async ValueTask<OperationState> IOperation.UpdateStateAsync(bool async, Cancella
                 }
              }
         }
+
+        // This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
+        RehydrationToken IOperation.GetRehydrationToken() =>
+            throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
     }
 }

sdk/confidentialledger/Azure.Security.ConfidentialLedger/src/PostLedgerEntryOperation.cs

@@ -77,6 +77,10 @@ async ValueTask<OperationState> IOperation.UpdateStateAsync(bool async, Cancella
             return OperationState.Pending(statusResponse);
         }
 
+        // This method is never invoked since we don't override Operation<T>.GetRehydrationToken.
+        RehydrationToken IOperation.GetRehydrationToken() =>
+            throw new NotSupportedException($"{nameof(GetRehydrationToken)} is not supported.");
+
         /// <summary>
         /// The transactionId of the posted ledger entry.
         /// </summary>

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

@@ -58,17 +58,21 @@ public static IOperation Create(
                 apiVersionStr = !skipApiVersionOverride && TryGetApiVersion(startRequestUri, out ReadOnlySpan<char> apiVersion) ? apiVersion.ToString() : null;
             }
             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));
-            }
 
             string? lastKnownLocation;
             if (!response.Headers.TryGetValue("Location", out lastKnownLocation))
             {
                 lastKnownLocation = null;
             }
-            return new NextLinkOperationImplementation(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, lastKnownLocation, finalStateVia, apiVersionStr, isNextRequestPolling : isNextRequestPolling);
+
+            NextLinkOperationImplementation operation = new(pipeline, requestMethod, startRequestUri, nextRequestUri, headerSource, lastKnownLocation, finalStateVia, apiVersionStr, isNextRequestPolling: isNextRequestPolling);
+
+            if (headerSource == HeaderSource.None && IsFinalState(response, headerSource, out var failureState, out _))
+            {
+                return new CompletedOperation(failureState ?? GetOperationStateFromFinalResponse(requestMethod, response), operation);
+            }
+
+            return operation;
         }
 
         public static IOperation<T> Create<T>(
@@ -647,12 +651,17 @@ private class CompletedOperation : IOperation
         {
             private readonly OperationState _operationState;
 
-            public CompletedOperation(OperationState operationState)
+            private readonly NextLinkOperationImplementation _operation;
+
+            public CompletedOperation(OperationState operationState, NextLinkOperationImplementation operation)
             {
                 _operationState = operationState;
+                _operation = operation;
             }
 
             public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => new(_operationState);
+
+            public RehydrationToken GetRehydrationToken() => _operation.GetRehydrationToken();
         }
 
         private sealed class OperationToOperationOfT<T> : IOperation<T>
@@ -685,6 +694,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,11 @@ internal ProtocolOperation(ClientDiagnostics clientDiagnostics, HttpPipeline pip
         public override string Id => throw new NotSupportedException();
 #pragma warning restore CA1822
 
+        /// <inheritdoc />
+        public override RehydrationToken? GetRehydrationToken() => ((IOperation<T>)this).GetRehydrationToken();
+
+        RehydrationToken IOperation<T>.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,9 @@ public TestOperation(Func<bool, CancellationToken, ValueTask<OperationState<int>
             }
 
             public ValueTask<OperationState<int>> UpdateStateAsync(bool async, CancellationToken cancellationToken) => _updateStateAsyncHandler(async, cancellationToken);
+
+            public RehydrationToken GetRehydrationToken() =>
+                throw new NotImplementedException();
         }
 
         private class CallCountStrategy : DelayStrategy

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

@@ -438,6 +438,9 @@ public TestOperation(Func<bool, CancellationToken, ValueTask<OperationState>> up
             }
 
             public ValueTask<OperationState> UpdateStateAsync(bool async, CancellationToken cancellationToken) => _updateStateAsyncHandler(async, cancellationToken);
+
+            public RehydrationToken GetRehydrationToken() =>
+                throw new NotImplementedException();
         }
 
         private class CallCountStrategy : DelayStrategy