revert unnecessary change

Author: nytian

src/DurableTask.AzureStorage/Partitioning/TablePartitionManager.cs

@@ -412,24 +412,6 @@ public async Task<ReadTableReponse> ReadAndWriteTableAsync(bool isShuttingDown,
                         try
                         {
                             await this.partitionTable.ReplaceEntityAsync(partition, etag, forcefulShutdownToken);
-
-                            // Ensure worker is listening to the control queue iff either:
-                            // 1) worker just claimed the lease,
-                            // 2) worker was already the owner in the partitions table and is not actively draining the queue.
-                            //    Note that during draining, we renew the lease but do not want to listen to new messages.
-                            //    Otherwise, we'll never finish draining our in-memory messages.
-                            // When drain completes, and the worker may decide to release the lease. In that moment,
-                            // IsDrainingPartition can still be true but renewedLease is false — without checking
-                            // !releasedLease, the worker could incorrectly resume listening just before releasing the lease.
-                            bool isRenewingToDrainQueue = renewedLease && response.IsDrainingPartition && !releasedLease;
-                            if (claimedLease || !isRenewingToDrainQueue)
-                            {
-                                // Notify the orchestration session manager that we acquired a lease for one of the partitions.
-                                // This will cause it to start reading control queue messages for that partition.
-                                await this.service.OnTableLeaseAcquiredAsync(partition);
-                            }
-
-                            this.LogHelper(partition, claimedLease, stoleLease, renewedLease, drainedLease, releasedLease, previousOwner);
                         }
                         catch (DurableTaskStorageException ex) when (ex.HttpStatusCode == (int)HttpStatusCode.PreconditionFailed)
                         {
@@ -441,6 +423,24 @@ public async Task<ReadTableReponse> ReadAndWriteTableAsync(bool isShuttingDown,
                                 $"Failed to update table entry due to an Etag mismatch. Failed ETag value: '{etag}'.");
                             throw;
                         }
+
+                        // Ensure worker is listening to the control queue iff either:
+                        // 1) worker just claimed the lease,
+                        // 2) worker was already the owner in the partitions table and is not actively draining the queue.
+                        //    Note that during draining, we renew the lease but do not want to listen to new messages.
+                        //    Otherwise, we'll never finish draining our in-memory messages.
+                        // When drain completes, and the worker may decide to release the lease. In that moment,
+                        // IsDrainingPartition can still be true but renewedLease is false — without checking
+                        // !releasedLease, the worker could incorrectly resume listening just before releasing the lease.
+                        bool isRenewingToDrainQueue = renewedLease && response.IsDrainingPartition && !releasedLease;
+                        if (claimedLease || !isRenewingToDrainQueue)
+                        {
+                            // Notify the orchestration session manager that we acquired a lease for one of the partitions.
+                            // This will cause it to start reading control queue messages for that partition.
+                            await this.service.OnTableLeaseAcquiredAsync(partition);
+                        }
+
+                        this.LogHelper(partition, claimedLease, stoleLease, renewedLease, drainedLease, releasedLease, previousOwner);
                     }
                 }