private async Task UpdateReplicaStoreAfterUpgradeStartAsync(
            FabricUpgradeProgress upgradeProgress,
            CancellationToken token)
        {
            try
            {
                await this.WaitForPartitionReadyAsync(token);

                using (var tx = await this.store.CreateTransactionWithRetryAsync(this.exceptionPolicy, token))
                {
                    KeyValueStoreItem         storeItem         = null;
                    CurrentUpgradeInformation updateInformation = null;
                    if (this.store.Contains(tx, CurrentUpgradeKey))
                    {
                        try
                        {
                            storeItem         = this.store.Get(tx, CurrentUpgradeKey);
                            updateInformation = CurrentUpgradeInformation.Deserialize(storeItem.Value);
                            if (!updateInformation.UpgradeStarted)
                            {
                                updateInformation.UpgradeStarted = true;
                            }
                        }
                        catch (Exception ex)
                        {
                            Trace.WriteError(TraceType, "Update state deserialize failed with exception {0}", ex);
                            updateInformation = null;
                        }
                    }

                    if (updateInformation != null &&
                        !updateInformation.UpgradeStarted &&
                        updateInformation.TargetCodeVersion == upgradeProgress.TargetCodeVersion &&
                        (updateInformation.TargetConfigVersion == null ||
                         updateInformation.TargetConfigVersion == upgradeProgress.TargetConfigVersion))
                    {
                        updateInformation.UpgradeStarted = true;
                        this.store.Update(tx, CurrentUpgradeKey, updateInformation.Serialize(),
                                          storeItem.Metadata.SequenceNumber);
                        await tx.CommitAsync();
                    }
                }
            }
            catch (Exception ex)
            {
                Trace.WriteError(TraceType, "UpdateReplicaStoreAfterUpgradeStartAsync: Error : {0}", ex);
            }
        }
        private async Task <TimeSpan> GetWaitDurationAndUpdateHealth(FabricUpgradeProgress upgradeProgress, CancellationToken token)
        {
            TimeSpan waitDuration = TimeSpan.MinValue;

            using (var tx = await this.store.CreateTransactionWithRetryAsync(this.exceptionPolicy, token))
            {
                KeyValueStoreItem         storeItem         = null;
                CurrentUpgradeInformation updateInformation = null;
                if (this.store.Contains(tx, this.CurrentUpgradeKey))
                {
                    try
                    {
                        storeItem         = this.store.Get(tx, this.CurrentUpgradeKey);
                        updateInformation = CurrentUpgradeInformation.Deserialize(storeItem.Value);
                    }
                    catch (Exception ex)
                    {
                        Trace.WriteError(TraceType, "Update state deserialize failed with exception {0}", ex);
                        updateInformation = null;
                    }
                }

                if (updateInformation != null &&
                    (upgradeProgress.UpgradeState == FabricUpgradeState.Failed ||
                     upgradeProgress.UpgradeState == FabricUpgradeState.RollingBackCompleted))
                {
                    if (updateInformation.CurrentCodeVersion == upgradeProgress.TargetCodeVersion &&
                        updateInformation.CurrentConfigVersion == upgradeProgress.TargetConfigVersion)
                    {
                        // Error so report Warning health for the service
                        var description = string.Format(
                            CultureInfo.InvariantCulture,
                            "Upgrade failed for Update Id:{0} Target code version: {1}",
                            updateInformation.UpdateId,
                            updateInformation.TargetCodeVersion);
                        this.ReportHealthEvent(description, HealthState.Warning);
                        waitDuration = this.waitTimeBeforePolling;
                    }
                }
                else
                {
                    this.ReportHealthEvent("Health", HealthState.Ok);
                }
            }

            return(waitDuration);
        }
        private async Task <ClusterUpgradeCommandParameter> UpdateReplicaStoreBeforeUpgradeStartAsync(
            ClusterUpgradeCommandParameter commandParameter,
            FabricUpgradeProgress upgradeProgress,
            string updateId,
            CancellationToken token)
        {
            try
            {
                await this.WaitForPartitionReadyAsync(token);

                Trace.WriteInfo(TraceType, "UpdateReplicaStoreBeforeUpgradeStartAsync: Update replica store.");
                KeyValueStoreItem storeItem = null;
                using (var tx = await this.store.CreateTransactionWithRetryAsync(this.exceptionPolicy, token))
                {
                    CurrentUpgradeInformation updateInformation = null;
                    if (this.store.Contains(tx, CurrentUpgradeKey))
                    {
                        storeItem = this.store.Get(tx, CurrentUpgradeKey);
                        try
                        {
                            updateInformation = CurrentUpgradeInformation.Deserialize(storeItem.Value);
                            if (updateInformation.UpgradeStarted && updateId == updateInformation.UpdateId)
                            {
                                updateInformation.Retry++;
                                updateInformation.UpgradeStarted = false;
                            }
                        }
                        catch (Exception ex)
                        {
                            Trace.WriteError(TraceType, "Deserialize failed with exception {0}", ex);
                            updateInformation = null;
                        }
                    }

                    if (updateInformation == null || updateId != updateInformation.UpdateId)
                    {
                        updateInformation = new CurrentUpgradeInformation
                        {
                            Retry                = 0,
                            TargetCodeVersion    = commandParameter.CodeVersion,
                            TargetConfigVersion  = commandParameter.ConfigVersion,
                            CurrentCodeVersion   = upgradeProgress.TargetCodeVersion,
                            CurrentConfigVersion = upgradeProgress.TargetConfigVersion,
                            UpdateId             = updateId,
                            UpgradeStarted       = false
                        };
                    }

                    Trace.WriteInfo(TraceType, "Updating store for update Id: {0}, Retry: {1} ", updateId, updateInformation.Retry);
                    if (storeItem != null)
                    {
                        this.store.Update(tx, CurrentUpgradeKey, updateInformation.Serialize(), storeItem.Metadata.SequenceNumber);
                    }
                    else
                    {
                        this.store.Add(tx, CurrentUpgradeKey, updateInformation.Serialize());
                    }

                    await tx.CommitAsync();

                    Trace.WriteInfo(TraceType, "UpdateReplicaStoreBeforeUpgradeStartAsync: replica store updated");
                }
            }
            catch (Exception ex)
            {
                Trace.WriteError(TraceType, "UpdateReplicaStoreBeforeUpgradeStartAsync: Error : {0}", ex);
                commandParameter = null;
            }

            if (commandParameter != null)
            {
                if (commandParameter.UpgradeDescription == null)
                {
                    commandParameter.UpgradeDescription = new CommandProcessorClusterUpgradeDescription();
                }

                commandParameter.UpgradeDescription.ForceRestart = true;
            }

            return(commandParameter);
        }