Simplify implementation

Author: yash2710

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedProcessor.cs

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Collections.Generic;
+    using System.Text.Json;
     using System.Threading;
     using System.Threading.Tasks;
 
@@ -27,53 +28,45 @@ public abstract class ChangeFeedProcessor
         public abstract Task StopAsync();
 
         /// <summary>
-        /// Exports all leases from the lease container to a list of <see cref="LeaseExportData"/> objects.
+        /// Exports all leases from the lease container.
         /// </summary>
         /// <param name="cancellationToken">A cancellation token to observe.</param>
-        /// <returns>A list of exported lease data.</returns>
+        /// <returns>A list of lease objects as JSON elements.</returns>
         /// <remarks>
         /// <para>
-        /// The processor must be stopped before calling this method. If the processor is running,
-        /// an <see cref="InvalidOperationException"/> will be thrown.
+        /// Each exported lease is a JSON object representing the serialized lease state.
+        /// The payload should not be modified and is intended to be passed directly to <see cref="ImportLeasesAsync"/>.
         /// </para>
         /// <para>
-        /// The export includes all lease metadata including continuation tokens, ownership history,
-        /// and custom properties. The exported data can be used to restore leases using <see cref="ImportLeasesAsync"/>.
+        /// This operation can be performed while the processor is running.
         /// </para>
         /// </remarks>
-        /// <exception cref="InvalidOperationException">Thrown when the processor is currently running.</exception>
         /// <exception cref="NotSupportedException">Thrown when the implementation does not support export.</exception>
-        public virtual Task<IReadOnlyList<LeaseExportData>> ExportLeasesAsync(CancellationToken cancellationToken = default)
+        public virtual Task<IReadOnlyList<JsonElement>> ExportLeasesAsync(CancellationToken cancellationToken = default)
         {
             throw new NotSupportedException("This ChangeFeedProcessor implementation does not support lease export.");
         }
 
         /// <summary>
-        /// Imports leases from a list of <see cref="LeaseExportData"/> objects into the lease container.
+        /// Imports leases into the lease container.
         /// </summary>
-        /// <param name="leases">The list of lease data to import.</param>
-        /// <param name="overwriteExisting">Whether to overwrite existing leases with the same ID. Default is false.</param>
+        /// <param name="leases">The list of lease objects as JSON elements to import.</param>
+        /// <param name="overwriteExisting">Whether to overwrite existing leases with the same token. Default is false.</param>
         /// <param name="cancellationToken">A cancellation token to observe.</param>
         /// <returns>A task representing the asynchronous operation.</returns>
         /// <remarks>
         /// <para>
-        /// The processor must be stopped before calling this method. If the processor is running,
-        /// an <see cref="InvalidOperationException"/> will be thrown.
-        /// </para>
-        /// <para>
-        /// When importing, the ownership is transferred to the current processor instance,
-        /// and the ownership history is updated to reflect the import action.
+        /// The lease objects should be the opaque JSON elements obtained from <see cref="ExportLeasesAsync"/>.
         /// </para>
         /// <para>
         /// If <paramref name="overwriteExisting"/> is false (default), existing leases will not be modified.
         /// If true, existing leases will be replaced with the imported data.
         /// </para>
         /// </remarks>
         /// <exception cref="ArgumentNullException">Thrown when <paramref name="leases"/> is null.</exception>
-        /// <exception cref="InvalidOperationException">Thrown when the processor is currently running.</exception>
         /// <exception cref="NotSupportedException">Thrown when the implementation does not support import.</exception>
         public virtual Task ImportLeasesAsync(
-            IReadOnlyList<LeaseExportData> leases,
+            IReadOnlyList<JsonElement> leases,
             bool overwriteExisting = false,
             CancellationToken cancellationToken = default)
         {

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedProcessorCore.cs

@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos.ChangeFeed
 {
     using System;
     using System.Collections.Generic;
+    using System.Text.Json;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.ChangeFeed.Bootstrapping;
@@ -101,44 +102,37 @@ public override async Task StopAsync()
             }
         }
 
-        public override async Task<IReadOnlyList<LeaseExportData>> ExportLeasesAsync(CancellationToken cancellationToken = default)
+        public override async Task<IReadOnlyList<JsonElement>> ExportLeasesAsync(CancellationToken cancellationToken = default)
         {
-            // Wait for any ongoing start/stop operations to complete
-            // If processor is running, this will wait until StopAsync is called by the user
-            await this.runningLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-            try 
+            // Initialize if needed to access the lease container
+            if (!this.initialized)
             {
-                if (this.isRunning)
+                await this.runningLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+                try
                 {
-                    // Release lock and throw - user must stop the processor first
-                    throw new InvalidOperationException(
-                        "Cannot export leases while the ChangeFeedProcessor is running. " +
-                        "Please call StopAsync() before exporting leases.");
+                    if (!this.initialized)
+                    {
+                        await this.InitializeAsync().ConfigureAwait(false);
+                    }
                 }
-
-                // Initialize if needed to access the lease container
-                if (!this.initialized)
+                finally
                 {
-                    await this.InitializeAsync().ConfigureAwait(false);
+                    this.runningLock.Release();
                 }
+            }
 
-                DefaultTrace.TraceInformation("Exporting leases...");
-                IReadOnlyList<LeaseExportData> exportedLeases = await this.documentServiceLeaseStoreManager
-                    .LeaseContainer
-                    .ExportLeasesAsync(this.instanceName, cancellationToken)
-                    .ConfigureAwait(false);
+            DefaultTrace.TraceInformation("Exporting leases...");
+            IReadOnlyList<JsonElement> exportedLeases = await this.documentServiceLeaseStoreManager
+                .LeaseContainer
+                .ExportLeasesAsync(cancellationToken)
+                .ConfigureAwait(false);
 
-                DefaultTrace.TraceInformation("Exported {0} leases.", exportedLeases.Count);
-                return exportedLeases;
-            }
-            finally
-            {
-                this.runningLock.Release();
-            }
+            DefaultTrace.TraceInformation("Exported {0} leases.", exportedLeases.Count);
+            return exportedLeases;
         }
 
         public override async Task ImportLeasesAsync(
-            IReadOnlyList<LeaseExportData> leases,
+            IReadOnlyList<JsonElement> leases,
             bool overwriteExisting = false,
             CancellationToken cancellationToken = default)
         {
@@ -147,36 +141,30 @@ public override async Task ImportLeasesAsync(
                 throw new ArgumentNullException(nameof(leases));
             }
 
-            // Wait for any ongoing start/stop operations to complete
-            await this.runningLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-            try
+            // Initialize if needed to access the lease container
+            if (!this.initialized)
             {
-                if (this.isRunning)
+                await this.runningLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+                try
                 {
-                    // Release lock and throw - user must stop the processor first
-                    throw new InvalidOperationException(
-                        "Cannot import leases while the ChangeFeedProcessor is running. " +
-                        "Please call StopAsync() before importing leases.");
+                    if (!this.initialized)
+                    {
+                        await this.InitializeAsync().ConfigureAwait(false);
+                    }
                 }
-
-                // Initialize if needed to access the lease container
-                if (!this.initialized)
+                finally
                 {
-                    await this.InitializeAsync().ConfigureAwait(false);
+                    this.runningLock.Release();
                 }
+            }
 
-                DefaultTrace.TraceInformation("Importing {0} leases (overwriteExisting={1})...", leases.Count, overwriteExisting);
-                await this.documentServiceLeaseStoreManager
-                    .LeaseContainer
-                    .ImportLeasesAsync(leases, this.instanceName, overwriteExisting, cancellationToken)
-                    .ConfigureAwait(false);
+            DefaultTrace.TraceInformation("Importing {0} leases (overwriteExisting={1})...", leases.Count, overwriteExisting);
+            await this.documentServiceLeaseStoreManager
+                .LeaseContainer
+                .ImportLeasesAsync(leases, overwriteExisting, cancellationToken)
+                .ConfigureAwait(false);
 
-                DefaultTrace.TraceInformation("Imported {0} leases.", leases.Count);
-            }
-            finally
-            {
-                this.runningLock.Release();
-            }
+            DefaultTrace.TraceInformation("Imported {0} leases.", leases.Count);
         }
 
         private async Task InitializeAsync()

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/LeaseExportData.cs

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.ChangeFeed.LeaseManagement
 {
     using System.Collections.Generic;
+    using System.Text.Json;
     using System.Threading;
     using System.Threading.Tasks;
 
@@ -26,26 +27,22 @@ internal abstract class DocumentServiceLeaseContainer
         public abstract Task<IEnumerable<DocumentServiceLease>> GetOwnedLeasesAsync();
 
         /// <summary>
-        /// Exports all leases to a list of <see cref="LeaseExportData"/> objects.
+        /// Exports all leases as JSON elements.
         /// </summary>
-        /// <param name="exportedBy">The name of the instance performing the export.</param>
         /// <param name="cancellationToken">A cancellation token to observe.</param>
-        /// <returns>A list of exported lease data.</returns>
-        public abstract Task<IReadOnlyList<LeaseExportData>> ExportLeasesAsync(
-            string exportedBy,
+        /// <returns>A list of lease objects as JSON elements.</returns>
+        public abstract Task<IReadOnlyList<JsonElement>> ExportLeasesAsync(
             CancellationToken cancellationToken = default);
 
         /// <summary>
-        /// Imports leases from a list of <see cref="LeaseExportData"/> objects.
+        /// Imports leases from a list of JSON elements.
         /// </summary>
-        /// <param name="leases">The list of lease data to import.</param>
-        /// <param name="importedBy">The name of the instance performing the import.</param>
+        /// <param name="leases">The list of lease objects as JSON elements to import.</param>
         /// <param name="overwriteExisting">Whether to overwrite existing leases with the same ID.</param>
         /// <param name="cancellationToken">A cancellation token to observe.</param>
         /// <returns>A task representing the asynchronous operation.</returns>
         public abstract Task ImportLeasesAsync(
-            IReadOnlyList<LeaseExportData> leases,
-            string importedBy,
+            IReadOnlyList<JsonElement> leases,
             bool overwriteExisting = false,
             CancellationToken cancellationToken = default);
     }